ripple-api/app/websockets/main_handler.go

package websockets

import (
	"encoding/json"
	"sync"
	"sync/atomic"
	"time"

	"github.com/leavengood/websocket"
)

var stepNumber uint64

func handler(rawConn *websocket.Conn) {
	defer catchPanic()
	defer rawConn.Close()

	step := atomic.AddUint64(&stepNumber, 1)

	// 5 is a security margin in case
	if step == (1<<10 - 5) {
		atomic.StoreUint64(&stepNumber, 0)
	}

	c := &conn{
		rawConn,
		sync.Mutex{},
		step | uint64(time.Now().UnixNano()<<10),
	}

	c.WriteJSON(TypeConnected, nil)

	defer cleanup(c.ID)

	for {
		var i incomingMessage
		err := c.Conn.ReadJSON(&i)
		if _, ok := err.(*websocket.CloseError); ok {
			return
		}
		if err != nil {
			c.WriteJSON(TypeInvalidMessage, err.Error())
			continue
		}
		f, ok := messageHandler[i.Type]
		if !ok {
			c.WriteJSON(TypeInvalidMessage, "invalid message type")
			continue
		}
		if f != nil {
			f(c, i)
		}
	}
}

type conn struct {
	Conn *websocket.Conn
	Mtx  sync.Mutex
	ID   uint64
}

func (c *conn) WriteJSON(t string, data interface{}) error {
	c.Mtx.Lock()
	err := c.Conn.WriteJSON(newMessage(t, data))
	c.Mtx.Unlock()
	return err
}

var messageHandler = map[string]func(c *conn, message incomingMessage){
	TypeSubscribeScores: SubscribeScores,
	TypePing:            pingHandler,
}

// Server Message Types
const (
	TypeConnected          = "connected"
	TypeInvalidMessage     = "invalid_message_type"
	TypeSubscribedToScores = "subscribed_to_scores"
	TypeNewScore           = "new_score"
	TypePong               = "pong"
)

// Client Message Types
const (
	TypeSubscribeScores = "subscribe_scores"
	TypePing            = "ping"
)

func pingHandler(c *conn, message incomingMessage) {
	c.WriteJSON(TypePong, nil)
}

// Message is the wrapped information for a message sent to the client.
type Message struct {
	Type string      `json:"type"`
	Data interface{} `json:"data,omitempty"`
}

func newMessage(t string, data interface{}) Message {
	return Message{
		Type: t,
		Data: data,
	}
}

type incomingMessage struct {
	Type string          `json:"type"`
	Data json.RawMessage `json:"data"`
}