focalboard/server/ws/websockets.go

package ws

import (
	"encoding/json"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"github.com/mattermost/mattermost-octo-tasks/server/model"
)

// RegisterRoutes registers routes.
func (ws *Server) RegisterRoutes(r *mux.Router) {
	r.HandleFunc("/ws/onchange", ws.handleWebSocketOnChange)
}

// AddListener adds a listener for a block's change.
func (ws *Server) AddListener(client *websocket.Conn, blockIDs []string) {
	ws.mu.Lock()
	for _, blockID := range blockIDs {
		if ws.listeners[blockID] == nil {
			ws.listeners[blockID] = []*websocket.Conn{}
		}

		ws.listeners[blockID] = append(ws.listeners[blockID], client)
	}
	ws.mu.Unlock()
}

// RemoveListener removes a webSocket listener from all blocks.
func (ws *Server) RemoveListener(client *websocket.Conn) {
	ws.mu.Lock()
	for key, clients := range ws.listeners {
		listeners := []*websocket.Conn{}

		for _, existingClient := range clients {
			if client != existingClient {
				listeners = append(listeners, existingClient)
			}
		}

		ws.listeners[key] = listeners
	}
	ws.mu.Unlock()
}

// RemoveListenerFromBlocks removes a webSocket listener from a set of block.
func (ws *Server) RemoveListenerFromBlocks(client *websocket.Conn, blockIDs []string) {
	ws.mu.Lock()

	for _, blockID := range blockIDs {
		listeners := ws.listeners[blockID]
		if listeners == nil {
			return
		}

		// Remove the first instance of this client that's listening to this block
		// Note: A client can listen multiple times to the same block
		for index, listener := range listeners {
			if client == listener {
				newListeners := append(listeners[:index], listeners[index+1:]...)
				ws.listeners[blockID] = newListeners

				break
			}
		}
	}

	ws.mu.Unlock()
}

// GetListeners returns the listeners to a blockID's changes.
func (ws *Server) GetListeners(blockID string) []*websocket.Conn {
	ws.mu.Lock()
	listeners := ws.listeners[blockID]
	ws.mu.Unlock()

	return listeners
}

// Server is a WebSocket server.
type Server struct {
	upgrader  websocket.Upgrader
	listeners map[string][]*websocket.Conn
	mu        sync.RWMutex
}

// NewServer creates a new Server.
func NewServer() *Server {
	return &Server{
		listeners: make(map[string][]*websocket.Conn),
		upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool {
				return true
			},
		},
	}
}

// UpdateMsg is sent on block updates
type UpdateMsg struct {
	Action string      `json:"action"`
	Block  model.Block `json:"block"`
}

// WebsocketCommand is an incoming command from the client.
type WebsocketCommand struct {
	Action   string   `json:"action"`
	BlockIDs []string `json:"blockIds"`
}

func (ws *Server) handleWebSocketOnChange(w http.ResponseWriter, r *http.Request) {
	// Upgrade initial GET request to a websocket
	client, err := ws.upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Fatal(err)
	}

	// TODO: Auth

	log.Printf("CONNECT WebSocket onChange, client: %s", client.RemoteAddr())

	// Make sure we close the connection when the function returns
	defer func() {
		log.Printf("DISCONNECT WebSocket onChange, client: %s", client.RemoteAddr())

		// Remove client from listeners
		ws.RemoveListener(client)

		client.Close()
	}()

	// Simple message handling loop
	for {
		_, p, err := client.ReadMessage()
		if err != nil {
			log.Printf("ERROR WebSocket onChange, client: %s, err: %v", client.RemoteAddr(), err)
			ws.RemoveListener(client)

			break
		}

		var command WebsocketCommand

		err = json.Unmarshal(p, &command)
		if err != nil {
			// handle this error
			log.Printf(`ERROR webSocket parsing command JSON: %v`, string(p))

			continue
		}

		switch command.Action {
		case "ADD":
			log.Printf(`Command: Add blockID: %v, client: %s`, command.BlockIDs, client.RemoteAddr())
			ws.AddListener(client, command.BlockIDs)

		case "REMOVE":
			log.Printf(`Command: Remove blockID: %v, client: %s`, command.BlockIDs, client.RemoteAddr())
			ws.RemoveListenerFromBlocks(client, command.BlockIDs)

		default:
			log.Printf(`ERROR webSocket command, invalid action: %v`, command.Action)
		}
	}
}

// BroadcastBlockDelete broadcasts delete messages to clients
func (ws *Server) BroadcastBlockDelete(blockID string, parentID string) {
	now := time.Now().Unix()
	block := model.Block{}
	block.ID = blockID
	block.ParentID = parentID
	block.UpdateAt = now
	block.DeleteAt = now

	ws.BroadcastBlockChange(block)
}

// BroadcastBlockChange broadcasts update messages to clients
func (ws *Server) BroadcastBlockChange(block model.Block) {
	blockIDsToNotify := []string{block.ID, block.ParentID}

	for _, blockID := range blockIDsToNotify {
		listeners := ws.GetListeners(blockID)
		log.Printf("%d listener(s) for blockID: %s", len(listeners), blockID)

		if listeners != nil {
			message := UpdateMsg{
				Action: "UPDATE_BLOCK",
				Block:  block,
			}

			for _, listener := range listeners {
				log.Printf("Broadcast change, blockID: %s, remoteAddr: %s", blockID, listener.RemoteAddr())

				err := listener.WriteJSON(message)
				if err != nil {
					log.Printf("broadcast error: %v", err)
					listener.Close()
				}
			}
		}
	}
}