focalboard/server/main/websockets.go

package main

import (
	"log"
	"net/http"
	"sync"

	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
)

func (ws *WSServer) RegisterRoutes(r *mux.Router) {
	r.HandleFunc("/ws/onchange", wsServer.handleWebSocketOnChange)
}

// AddListener adds a listener for a blockID's change
func (ws *WSServer) AddListener(client *websocket.Conn, blockID string) {
	ws.mu.Lock()
	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
func (ws *WSServer) RemoveListener(client *websocket.Conn) {
	ws.mu.Lock()
	for key, clients := range ws.listeners {
		var listeners = []*websocket.Conn{}
		for _, existingClient := range clients {
			if client != existingClient {
				listeners = append(listeners, existingClient)
			}
		}
		ws.listeners[key] = listeners
	}
	ws.mu.Unlock()
}

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

	return listeners
}

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

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

// WebsocketMsg is sent on block changes
type WebsocketMsg struct {
	Action  string `json:"action"`
	BlockID string `json:"blockId"`
}

func (ws *WSServer) 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

	query := r.URL.Query()
	blockID := query.Get("id")
	log.Printf("CONNECT WebSocket onChange, blockID: %s, client: %s", blockID, client.RemoteAddr())

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

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

		client.Close()
	}()

	// Register our new client
	ws.AddListener(client, blockID)

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

func (ws *WSServer) broadcastBlockChangeToWebsocketClients(blockIDs []string) {
	for _, blockID := range blockIDs {
		listeners := ws.GetListeners(blockID)
		log.Printf("%d listener(s) for blockID: %s", len(listeners), blockID)

		if listeners != nil {
			var message = WebsocketMsg{
				Action:  "UPDATE_BLOCK",
				BlockID: blockID,
			}
			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()
				}
			}
		}
	}
}