niko_trust/internal/server/ws.go

package server

// WebSocket streaming. A session holder may subscribe to the same channels
// the REST list endpoints expose — claims by subject, requests by recipient,
// responses by request hash, revocations by claim id — and receive every
// newly stored envelope on those channels as raw {tce, signature}, exactly
// what a GET would have returned. The relay still decides nothing: pushed
// bytes are verified locally like any other envelope (INV-5).
//
// Hygiene: bounded send buffer per subscriber (a slow client is disconnected
// rather than allowed to stall broadcasts), a cap on subscriptions per
// connection, and sessions expire mid-stream just as they would for REST.

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

	"github.com/coder/websocket"

	"git.n1ko.dev/Niko/niko_trust/pkg/protocol"
	"git.n1ko.dev/Niko/niko_trust/pkg/transport"
)

const (
	wsSendBuffer     = 64
	wsMaxSubsPerConn = 16
)

type wsSubscription struct {
	channel string // claims | requests | responses | revocations
	key     string // subject / recipient / request hash / claim id
}

type wsSubscriber struct {
	subs []wsSubscription
	send chan wsEvent
	ses  session

	// done is closed exactly once, under the hub lock, to signal eviction
	// (slow client or disconnect). Closing a signalling channel that nobody
	// sends on cannot race the broadcaster, which sends only while holding
	// the same lock and skips evicted subscribers.
	done    chan struct{}
	evicted bool

	hub *wsHubState
}

// evict marks the subscriber dead and wakes its writer. Safe to call twice.
func evict(sub *wsSubscriber) {
	h := sub.hub
	if h == nil {
		return
	}
	h.mu.Lock()
	if !sub.evicted {
		sub.evicted = true
		close(sub.done)
	}
	h.mu.Unlock()
}

// wsEvent is one pushed object.
type wsEvent struct {
	Event    string       `json:"event"` // "object" | "error" | "bye"
	Channel  string       `json:"channel,omitempty"`
	Key      string       `json:"key,omitempty"`
	ObjectID string       `json:"object_id,omitempty"`
	Envelope *wireRequest `json:"envelope,omitempty"`
	Message  string       `json:"message,omitempty"`
}

type wsHubState struct {
	mu   sync.Mutex
	subs map[*wsSubscriber]struct{}
}

// register adds a subscriber with its done channel wired to this hub.
func (h *wsHubState) register(sub *wsSubscriber) {
	sub.hub = h
	h.mu.Lock()
	h.subs[sub] = struct{}{}
	h.mu.Unlock()
}

func newWSHub() *wsHubState { return &wsHubState{subs: make(map[*wsSubscriber]struct{})} }

func (h *wsHubState) remove(sub *wsSubscriber) {
	h.mu.Lock()
	delete(h.subs, sub)
	evicted := !sub.evicted
	if evicted {
		sub.evicted = true
		close(sub.done)
	}
	h.mu.Unlock()
}

func (h *wsHubState) broadcast(channel, key, objectID string, tceBytes, sig []byte) {
	ev := wsEvent{
		Event:    "object",
		Channel:  channel,
		Key:      key,
		ObjectID: objectID,
		Envelope: &wireRequest{TCE: tceBytes, Signature: sig},
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	for sub := range h.subs {
		if sub.evicted || !sub.wants(channel, key) {
			continue
		}
		select {
		case sub.send <- ev:
		default:
			// Slow client: drop it rather than stall everyone. The done
			// channel wakes the writer; no send ever races a close because
			// eviction happens under the same lock as this broadcast.
			evictLocked(sub)
		}
	}
}

// evictLocked is evict for callers already holding the hub lock.
func evictLocked(sub *wsSubscriber) {
	if sub.evicted {
		return
	}
	sub.evicted = true
	close(sub.done)
}

func (sub *wsSubscriber) wants(channel, key string) bool {
	for _, s := range sub.subs {
		if s.channel == channel && s.key == key {
			return true
		}
	}
	return false
}

// handleWS upgrades an authenticated session into a stream.
func (s *Server) handleWS(w http.ResponseWriter, r *http.Request) {
	ses, ok := s.sessionByToken(r)
	if !ok {
		writeErr(w, http.StatusUnauthorized, "authentication required")
		return
	}
	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
		// Browsers cannot set Authorization headers on WebSocket opens; the
		// bearer token arrives in the URL. Origin enforcement adds nothing
		// on top of token possession, so any origin may connect.
		OriginPatterns: []string{"*"},
	})
	if err != nil {
		return
	}
	defer conn.Close(websocket.StatusInternalError, "server shutting down")

	s.metrics.wsConnections.Add(1)
	defer s.metrics.wsConnections.Add(-1)

	sub := &wsSubscriber{
		ses:  ses,
		send: make(chan wsEvent, wsSendBuffer),
		done: make(chan struct{}),
	}
	s.ws.register(sub)
	defer s.ws.remove(sub)

	writeFail := make(chan error, 1)
	done := make(chan struct{})

	// Writer loop: forwards events until the connection dies or the session
	// expires. Session lifetime applies here exactly as to REST reads.
	go func() {
		ctx := r.Context()
		ticker := time.NewTicker(20 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-done:
				return
			case <-sub.done:
				conn.Close(websocket.StatusPolicyViolation, "evicted")
				return
			case <-ticker.C:
				if now().After(sub.ses.expiry) {
					conn.Close(websocket.StatusNormalClosure, "session expired")
					return
				}
			case ev := <-sub.send:
				raw, err := json.Marshal(ev)
				if err == nil {
					err = conn.Write(ctx, websocket.MessageText, raw)
					if err == nil {
						s.metrics.inc(&s.metrics.wsMessagesSent)
					}
				}
				if err != nil {
					select {
					case writeFail <- err:
					default:
					}
					return
				}
			}
		}
	}()

	// Reader loop: subscription control frames only.
	fail := func(status websocket.StatusCode, msg string) {
		ev := wsEvent{Event: "error", Message: msg}
		select {
		case sub.send <- ev:
		default:
		}
		_ = status
		_ = msg
	}
	for {
		typ, data, err := conn.Read(r.Context())
		if err != nil {
			break
		}
		if typ != websocket.MessageText {
			continue
		}
		var msg struct {
			Op      string `json:"op"`
			Channel string `json:"channel"`
			Key     string `json:"key"`
		}
		if json.Unmarshal(data, &msg) != nil || msg.Op != "subscribe" && msg.Op != "unsubscribe" {
			fail(0, "bad control frame")
			continue
		}
		if !validWSChannel(msg.Channel) {
			fail(0, "unknown channel")
			continue
		}
		requiredScope := wsChannelScope(msg.Channel)
		if !scopeOK(ses, requiredScope) {
			fail(0, "insufficient scope for "+msg.Channel)
			continue
		}
		if len(sub.subs) >= wsMaxSubsPerConn && msg.Op == "subscribe" {
			fail(0, "too many subscriptions")
			continue
		}

		switch msg.Op {
		case "subscribe":
			if msg.Key == "" {
				fail(0, "missing key")
				continue
			}
			sub.subs = append(sub.subs, wsSubscription{channel: msg.Channel, key: msg.Key})
			ev := wsEvent{Event: "subscribed", Channel: msg.Channel, Key: msg.Key}
			select {
			case sub.send <- ev:
			default:
			}
		case "unsubscribe":
			kept := sub.subs[:0]
			for _, x := range sub.subs {
				if !(x.channel == msg.Channel && (msg.Key == "" || x.key == msg.Key)) {
					kept = append(kept, x)
				}
			}
			sub.subs = kept
		}
	}

	close(done)
	conn.Close(websocket.StatusNormalClosure, "")
	_ = writeFail
}

var _ = context.Background

// validWSChannel names the four channels mirroring the REST list endpoints.
func validWSChannel(ch string) bool {
	switch ch {
	case "claims", "requests", "responses", "revocations":
		return true
	}
	return false
}

// wsChannelScope maps a channel onto its REST scope so that a session's
// grants mean identical things over both transports.
func wsChannelScope(ch string) string {
	return "read:" + ch
}

// hubBroadcastObj fans out a freshly stored object to matching subscribers.
// The object is already decoded by the caller.
func (s *Server) hubBroadcastObj(typ string, obj any, tceBytes, sig []byte, objectID string) {
	if s.ws == nil || len(s.ws.subs) == 0 {
		return
	}
	switch o := obj.(type) {
	case *protocol.Claim:
		s.ws.broadcast("claims", transport.AddrOf(o.Subject), objectID, tceBytes, sig)
	case *protocol.ApprovalRequest:
		s.ws.broadcast("requests", transport.AddrOf(o.Recipient), objectID, tceBytes, sig)
	case *protocol.Revocation:
		s.ws.broadcast("revocations", o.ClaimID.String(), objectID, tceBytes, sig)
	case *protocol.ApprovalResponse:
		s.ws.broadcast("responses", o.RequestHash.String(), objectID, tceBytes, sig)
	}
}