focalboard/server/ws/websockets.go
2021-02-09 12:27:34 -08:00

307 lines
7.4 KiB
Go

package ws
import (
"encoding/json"
"log"
"net/http"
"sync"
"time"
"github.com/gorilla/mux"
"github.com/gorilla/websocket"
"github.com/mattermost/focalboard/server/auth"
"github.com/mattermost/focalboard/server/model"
)
// IsValidSessionToken authenticates session tokens
type IsValidSessionToken func(token string) bool
// Server is a WebSocket server.
type Server struct {
upgrader websocket.Upgrader
listeners map[string][]*websocket.Conn
mu sync.RWMutex
auth *auth.Auth
singleUserToken string
}
// UpdateMsg is sent on block updates
type UpdateMsg struct {
Action string `json:"action"`
Block model.Block `json:"block"`
}
// ErrorMsg is sent on errors
type ErrorMsg struct {
Error string `json:"error"`
}
// WebsocketCommand is an incoming command from the client.
type WebsocketCommand struct {
Action string `json:"action"`
Token string `json:"token"`
ReadToken string `json:"readToken"`
BlockIDs []string `json:"blockIds"`
}
type websocketSession struct {
client *websocket.Conn
isAuthenticated bool
}
// NewServer creates a new Server.
func NewServer(auth *auth.Auth, singleUserToken string) *Server {
return &Server{
listeners: make(map[string][]*websocket.Conn),
upgrader: websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
},
auth: auth,
singleUserToken: singleUserToken,
}
}
// RegisterRoutes registers routes.
func (ws *Server) RegisterRoutes(r *mux.Router) {
r.HandleFunc("/ws/onchange", ws.handleWebSocketOnChange)
}
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()
}()
wsSession := websocketSession{
client: client,
isAuthenticated: false,
}
// 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 "AUTH":
log.Printf(`Command: AUTH, client: %s`, client.RemoteAddr())
ws.authenticateListener(&wsSession, command.Token, command.ReadToken)
case "ADD":
log.Printf(`Command: Add blockID: %v, client: %s`, command.BlockIDs, client.RemoteAddr())
ws.addListener(&wsSession, &command)
case "REMOVE":
log.Printf(`Command: Remove blockID: %v, client: %s`, command.BlockIDs, client.RemoteAddr())
ws.removeListenerFromBlocks(&wsSession, &command)
default:
log.Printf(`ERROR webSocket command, invalid action: %v`, command.Action)
}
}
}
func (ws *Server) isValidSessionToken(token string) bool {
if len(ws.singleUserToken) > 0 {
return token == ws.singleUserToken
}
session, err := ws.auth.GetSession(token)
if session != nil && err == nil {
return true
}
return false
}
func (ws *Server) authenticateListener(wsSession *websocketSession, token string, readToken string) {
// Authenticate session
isValidSession := ws.isValidSessionToken(token)
if !isValidSession {
wsSession.client.Close()
return
}
// Authenticated
wsSession.isAuthenticated = true
log.Printf("authenticateListener: Authenticated")
}
func (ws *Server) checkAuthentication(wsSession *websocketSession, command *WebsocketCommand) bool {
if wsSession.isAuthenticated {
return true
}
if len(command.ReadToken) > 0 {
// Read token must be valid for all block IDs
for _, blockID := range command.BlockIDs {
isValid, _ := ws.auth.IsValidReadToken(blockID, command.ReadToken)
if !isValid {
return false
}
}
return true
}
return false
}
// addListener adds a listener for a block's change.
func (ws *Server) addListener(wsSession *websocketSession, command *WebsocketCommand) {
if !ws.checkAuthentication(wsSession, command) {
log.Printf("addListener: NOT AUTHENTICATED")
sendError(wsSession.client, "not authenticated")
return
}
ws.mu.Lock()
for _, blockID := range command.BlockIDs {
if ws.listeners[blockID] == nil {
ws.listeners[blockID] = []*websocket.Conn{}
}
ws.listeners[blockID] = append(ws.listeners[blockID], wsSession.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(wsSession *websocketSession, command *WebsocketCommand) {
if !ws.checkAuthentication(wsSession, command) {
log.Printf("removeListenerFromBlocks: NOT AUTHENTICATED")
sendError(wsSession.client, "not authenticated")
return
}
ws.mu.Lock()
for _, blockID := range command.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 wsSession.client == listener {
newListeners := append(listeners[:index], listeners[index+1:]...)
ws.listeners[blockID] = newListeners
break
}
}
}
ws.mu.Unlock()
}
func sendError(conn *websocket.Conn, message string) {
errorMsg := ErrorMsg{
Error: message,
}
err := conn.WriteJSON(errorMsg)
if err != nil {
log.Printf("sendError error: %v", err)
conn.Close()
}
}
// 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
}
// 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()
}
}
}
}
}