MosswartOverlord/go-services/tracker-go/wsposition.go

package main

import (
	"context"
	"crypto/hmac"
	"encoding/json"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"github.com/coder/websocket"
)

// pluginRegistry maps character_name -> plugin connection for backend->plugin
// command routing (main.py plugin_conns).
type pluginRegistry struct {
	mu    sync.RWMutex
	conns map[string]*websocket.Conn
	log   *slog.Logger
}

func newPluginRegistry(log *slog.Logger) *pluginRegistry {
	return &pluginRegistry{conns: map[string]*websocket.Conn{}, log: log}
}

func (p *pluginRegistry) register(name string, c *websocket.Conn) {
	p.mu.Lock()
	p.conns[name] = c
	p.mu.Unlock()
}

// removeConn drops every name bound to this connection (on disconnect).
func (p *pluginRegistry) removeConn(c *websocket.Conn) {
	p.mu.Lock()
	for n, cc := range p.conns {
		if cc == c {
			delete(p.conns, n)
		}
	}
	p.mu.Unlock()
}

func (p *pluginRegistry) isConnected(name string) bool {
	p.mu.RLock()
	defer p.mu.RUnlock()
	_, ok := p.conns[name]
	return ok
}

// send routes an opaque {player_name, command} envelope to a plugin; evicts the
// connection on write failure (main.py command-forward semantics).
func (p *pluginRegistry) send(name string, payload map[string]any) {
	p.mu.RLock()
	c := p.conns[name]
	p.mu.RUnlock()
	if c == nil {
		return
	}
	b, _ := json.Marshal(payload)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := c.Write(ctx, websocket.MessageText, b); err != nil {
		p.mu.Lock()
		if p.conns[name] == c {
			delete(p.conns, name)
		}
		p.mu.Unlock()
	}
}

// pluginAuthOK constant-time-compares the supplied secret to SHARED_SECRET (and
// the optional rotation fallback). Fails closed when unset or left at the
// placeholder, matching main.py.
func (s *Server) pluginAuthOK(key string) bool {
	ok := s.sharedSecret != "" && s.sharedSecret != "your_shared_secret" &&
		hmac.Equal([]byte(key), []byte(s.sharedSecret))
	if !ok && s.sharedSecretLegacy != "" {
		ok = hmac.Equal([]byte(key), []byte(s.sharedSecretLegacy))
	}
	return ok
}

func (s *Server) handleWSPosition(w http.ResponseWriter, r *http.Request) {
	if s.ingestor == nil {
		http.Error(w, "ingest disabled on this instance", http.StatusServiceUnavailable)
		return
	}
	key := r.URL.Query().Get("secret")
	if key == "" {
		key = r.Header.Get("X-Plugin-Secret")
	}
	if !s.pluginAuthOK(key) {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{InsecureSkipVerify: true})
	if err != nil {
		return
	}
	defer conn.CloseNow()
	defer s.plugins.removeConn(conn)
	conn.SetReadLimit(32 << 20)

	ctx := r.Context()
	for {
		_, raw, err := conn.Read(ctx)
		if err != nil {
			return
		}
		var m map[string]any
		if json.Unmarshal(raw, &m) != nil {
			continue
		}
		if toStr(m["type"]) == "register" {
			name := toStr(m["character_name"])
			if name == "" {
				name = toStr(m["player_name"])
			}
			if name != "" {
				s.plugins.register(name, conn)
				s.ingestor.clearEquipmentCantrip(name)
				s.log.Info("plugin registered", "character", name)
			}
			continue
		}
		s.ingestor.dispatch(ctx, m)
	}
}