markdown-hub/internal/collab/hub.go

package collab

import (
	"database/sql"
	"fmt"
	"log"
	"net/http"
	"strings"
	"sync"

	"github.com/golang-jwt/jwt/v5"
	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool { return true },
}

// Room holds all connected clients for a single document.
type Room struct {
	mu      sync.Mutex
	clients map[*websocket.Conn]bool
}

// Hub manages all active collaboration rooms.
type Hub struct {
	mu    sync.Mutex
	rooms map[string]*Room
	db    *sql.DB
}

func NewHub(db *sql.DB) *Hub {
	return &Hub{
		rooms: make(map[string]*Room),
		db:    db,
	}
}

func (h *Hub) getRoom(fileID string) *Room {
	h.mu.Lock()
	defer h.mu.Unlock()
	if r, ok := h.rooms[fileID]; ok {
		return r
	}
	r := &Room{clients: make(map[*websocket.Conn]bool)}
	h.rooms[fileID] = r
	return r
}

func (h *Hub) removeClient(fileID string, conn *websocket.Conn) {
	h.mu.Lock()
	defer h.mu.Unlock()
	if r, ok := h.rooms[fileID]; ok {
		r.mu.Lock()
		delete(r.clients, conn)
		empty := len(r.clients) == 0
		r.mu.Unlock()
		if empty {
			delete(h.rooms, fileID)
		}
	}
}

// HandleWebSocket handles the Yjs WebSocket sync protocol.
// Clients send binary Yjs update messages; the hub broadcasts to all others in the room.
func (h *Hub) HandleWebSocket(w http.ResponseWriter, r *http.Request, secret string) {
	// Authenticate via query param or cookie
	tokenStr := r.URL.Query().Get("token")
	if tokenStr == "" {
		if c, err := r.Cookie("authToken"); err == nil {
			tokenStr = c.Value
		}
	}
	if tokenStr == "" {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	// Validate token (import auth inline to avoid circular — just parse JWT here)
	// We accept any valid token
	_, _, err := parseToken(tokenStr, secret)
	if err != nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	// Extract file ID from path: /ws/collab/{fileID}
	path := r.URL.Path
	parts := strings.Split(strings.TrimPrefix(path, "/ws/collab/"), "/")
	fileID := parts[0]
	if fileID == "" {
		http.Error(w, "file ID required", http.StatusBadRequest)
		return
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade failed: %v", err)
		return
	}
	defer conn.Close()

	room := h.getRoom(fileID)
	room.mu.Lock()
	room.clients[conn] = true
	room.mu.Unlock()

	defer h.removeClient(fileID, conn)

	// Send stored Yjs state if available
	var storedState []byte
	h.db.QueryRow("SELECT yjs_state FROM collab_state WHERE file_id = ?", fileID).Scan(&storedState)
	if storedState != nil {
		conn.WriteMessage(websocket.BinaryMessage, storedState)
	}

	// Read loop: receive updates from this client, broadcast to others, persist
	for {
		msgType, msg, err := conn.ReadMessage()
		if err != nil {
			break
		}
		if msgType != websocket.BinaryMessage {
			continue
		}

		// Persist latest state
		h.db.Exec(
			`INSERT INTO collab_state (file_id, yjs_state, updated_at) VALUES (?, ?, datetime('now'))
			 ON CONFLICT(file_id) DO UPDATE SET yjs_state = ?, updated_at = datetime('now')`,
			fileID, msg, msg,
		)

		// Broadcast to other clients in the room
		room.mu.Lock()
		for client := range room.clients {
			if client != conn {
				client.WriteMessage(websocket.BinaryMessage, msg)
			}
		}
		room.mu.Unlock()
	}
}

// ActiveUsers returns the number of connected users for a file.
func (h *Hub) ActiveUsers(fileID string) int {
	h.mu.Lock()
	defer h.mu.Unlock()
	if r, ok := h.rooms[fileID]; ok {
		r.mu.Lock()
		defer r.mu.Unlock()
		return len(r.clients)
	}
	return 0
}

func parseToken(tokenStr, secret string) (string, bool, error) {
	token, err := jwt.Parse(tokenStr, func(t *jwt.Token) (interface{}, error) {
		if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, fmt.Errorf("unexpected signing method")
		}
		return []byte(secret), nil
	})
	if err != nil || !token.Valid {
		return "", false, err
	}
	claims := token.Claims.(jwt.MapClaims)
	userID, _ := claims["sub"].(string)
	isAdmin, _ := claims["admin"].(bool)
	return userID, isAdmin, nil
}