Add initial websocket implementation

app/websockets/entry.go

@@ -0,0 +1,18 @@
+package websockets
+
+import (
+	"github.com/leavengood/websocket"
+	"github.com/valyala/fasthttp"
+)
+
+var upgrader = websocket.FastHTTPUpgrader{
+	Handler: handler,
+	CheckOrigin: func(ctx *fasthttp.RequestCtx) bool {
+		return true
+	},
+}
+
+// WebsocketV1Entry upgrades a connection to a websocket.
+func WebsocketV1Entry(ctx *fasthttp.RequestCtx) {
+	upgrader.UpgradeHandler(ctx)
+}

app/websockets/main_handler.go

@@ -0,0 +1,99 @@
+package websockets
+
+import (
+	"encoding/json"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/leavengood/websocket"
+)
+
+var stepNumber uint64
+
+func handler(rawConn *websocket.Conn) {
+	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
+		}
+		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,
+}
+
+// Server Message Types
+const (
+	TypeConnected      = "connected"
+	TypeInvalidMessage = "invalid_message_type"
+	TypeSubscribed     = "subscribed"
+	TypeNewScore       = "new_score"
+)
+
+// Client Message Types
+const (
+	TypeSubscribeScores = "subscribe_scores"
+)
+
+// 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"`
+}

app/websockets/scores.go

@@ -0,0 +1,121 @@
+package websockets
+
+import (
+	"encoding/json"
+	"fmt"
+	"sync"
+
+	"zxq.co/ripple/rippleapi/app/v1"
+)
+
+type subscribeScoresUser struct {
+	User  int   `json:"user"`
+	Modes []int `json:"modes"`
+}
+
+// SubscribeScores subscribes a connection to score updates.
+func SubscribeScores(c *conn, message incomingMessage) {
+	var ssu []subscribeScoresUser
+	err := json.Unmarshal(message.Data, &ssu)
+	if err != nil {
+		c.WriteJSON(TypeInvalidMessage, err.Error())
+		return
+	}
+
+	scoreSubscriptionsMtx.Lock()
+
+	var found bool
+	for idx, el := range scoreSubscriptions {
+		// already exists, change the users
+		if el.Conn.ID == c.ID {
+			found = true
+			scoreSubscriptions[idx].Users = ssu
+		}
+	}
+
+	// if it was not found, we need to add it
+	if !found {
+		scoreSubscriptions = append(scoreSubscriptions, scoreSubscription{c, ssu})
+	}
+
+	scoreSubscriptionsMtx.Unlock()
+
+	c.WriteJSON(TypeSubscribed, message)
+}
+
+type scoreSubscription struct {
+	Conn  *conn
+	Users []subscribeScoresUser
+}
+
+var scoreSubscriptions []scoreSubscription
+var scoreSubscriptionsMtx = new(sync.RWMutex)
+
+func scoreRetriever() {
+	ps, err := red.Subscribe("api:score_submission")
+	if err != nil {
+		fmt.Println(err)
+	}
+	for {
+		msg, err := ps.ReceiveMessage()
+		if err != nil {
+			fmt.Println(err.Error())
+			return
+		}
+		go handleNewScore(msg.Payload)
+	}
+}
+
+type score struct {
+	v1.Score
+	UserID int `json:"user_id"`
+}
+
+func handleNewScore(id string) {
+	var s score
+	err := db.Get(&s, `
+SELECT
+	id, beatmap_md5, score, max_combo, full_combo, mods,
+	300_count, 100_count, 50_count, gekis_count, katus_count, misses_count,
+	time, play_mode, accuracy, pp, completed, userid AS user_id
+FROM scores WHERE id = ?`, id)
+	if err != nil {
+		fmt.Println(err)
+		return
+	}
+	scoreSubscriptionsMtx.RLock()
+	cp := make([]scoreSubscription, len(scoreSubscriptions))
+	copy(cp, scoreSubscriptions)
+	scoreSubscriptionsMtx.RUnlock()
+
+	for _, el := range cp {
+		if len(el.Users) > 0 && !scoreUserValid(el.Users, s) {
+			continue
+		}
+
+		el.Conn.WriteJSON(TypeNewScore, s)
+	}
+}
+
+func scoreUserValid(users []subscribeScoresUser, s score) bool {
+	for _, u := range users {
+		if u.User == s.UserID {
+			if len(u.Modes) > 0 {
+				if !inModes(u.Modes, s.PlayMode) {
+					return false
+				}
+			}
+			return true
+		}
+	}
+	return false
+}
+
+func inModes(modes []int, i int) bool {
+	for _, m := range modes {
+		if m == i {
+			return true
+		}
+	}
+	return false
+}

app/websockets/websockets.go

@@ -0,0 +1,32 @@
+// Package websockets implements functionality related to the API websockets.
+package websockets
+
+import (
+	"github.com/jmoiron/sqlx"
+	"gopkg.in/redis.v5"
+)
+
+var (
+	red *redis.Client
+	db  *sqlx.DB
+)
+
+// Start begins websocket functionality
+func Start(r *redis.Client, _db *sqlx.DB) error {
+	red = r
+	db = _db
+	go scoreRetriever()
+	return nil
+}
+
+func cleanup(connID uint64) {
+	scoreSubscriptionsMtx.Lock()
+	for idx, el := range scoreSubscriptions {
+		if el.Conn.ID == connID {
+			scoreSubscriptions[idx] = scoreSubscriptions[len(scoreSubscriptions)-1]
+			scoreSubscriptions = scoreSubscriptions[:len(scoreSubscriptions)-1]
+			break
+		}
+	}
+	scoreSubscriptionsMtx.Unlock()
+}