niko_trust/internal/lightnode/ws.go

package lightnode

// WebSocket mirroring. The node exposes the relay's streaming protocol to
// its own clients and feeds it from upstream relays:
//
//	client ──ws──▶ light node ──ws──▶ pinned relays
//
// Every envelope received upstream is signature-checked before forwarding
// and deduplicated across peers by content ID. A freshly streamed object is
// usually younger than the newest quorum-agreed checkpoint, so it cannot yet
// carry an inclusion proof; such objects are forwarded marked
// "verified": false and queued. The queue is re-proven against each new
// agreed root purely as bookkeeping (and dropped once covered or exhausted):
// clients that need proof-backed bytes should re-fetch /v1/objects/{id},
// which only ever serves proof-verified content. Streaming is news,
// fetching is evidence.
//
// Subscriptions demanded downstream are reconciled onto upstream
// connections when they are (re)established; a subscription added while a
// connection is already up takes effect at its next reconnect. Documented
// v1 simplification.

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

	"github.com/coder/websocket"

	"git.n1ko.dev/Niko/niko_trust/internal/smt"
	"git.n1ko.dev/Niko/niko_trust/internal/transport"
)

const (
	lnSendBuffer     = 64
	lnMaxConns       = 64
	lnMaxSubsPerConn = 16
	lnPendingCap     = 512
	lnSeenCap        = 10000
)

type lnSub struct {
	channel, key string
}

type lnClient struct {
	subs    []lnSub
	send    chan []byte
	done    chan struct{}
	evicted bool
}

type lnHub struct {
	mu      sync.Mutex
	clients map[*lnClient]struct{}
	demand  map[lnSub]int // refcounted subscriptions demanded downstream
	kick    chan struct{} // wakes the manager so it can dial missing peers
}

func newLnHub() *lnHub {
	return &lnHub{
		clients: make(map[*lnClient]struct{}),
		demand:  make(map[lnSub]int),
		kick:    make(chan struct{}, 1),
	}
}

// evict removes a client and drops its demand contributions.
func (h *lnHub) evict(c *lnClient) {
	h.mu.Lock()
	if !c.evicted {
		c.evicted = true
		close(c.done)
	}
	delete(h.clients, c)
	for _, s := range c.subs {
		k := lnSub{s.channel, s.key}
		h.demand[k]--
		if h.demand[k] <= 0 {
			delete(h.demand, k)
		}
	}
	h.mu.Unlock()
	h.signal()
}

func (h *lnHub) signal() {
	select {
	case h.kick <- struct{}{}:
	default:
	}
}

type pendingEvent struct {
	channel, key, objID string
	attempts            int
}

type wsNode struct {
	node *Node
	auth *upstreamAuth

	hub *lnHub

	mu               sync.Mutex
	seen             map[string]struct{}
	pending          []pendingEvent
	upstreamLive     map[string]bool
	upstreamStarting map[string]bool
}

func newWSNode(n *Node) (*wsNode, error) {
	auth, err := newUpstreamAuth(n.cfg.CacheDir, n.http)
	if err != nil {
		return nil, err
	}
	return &wsNode{
		node:             n,
		auth:             auth,
		hub:              newLnHub(),
		seen:             make(map[string]struct{}),
		upstreamLive:     make(map[string]bool),
		upstreamStarting: make(map[string]bool),
	}, nil
}

// RunWS maintains one upstream stream per peer while subscriptions are
// demanded downstream. It blocks until ctx ends.
func (w *wsNode) RunWS(ctx context.Context) {
	ticker := time.NewTicker(3 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-w.hub.kick:
		case <-ticker.C:
		}
		w.flushPending()
		w.connectAll(ctx)
	}
}

// UpstreamReady reports whether at least one upstream stream is live. It is
// a convenience for operators and tests: events published before this point
// were never broadcast anywhere, so waiting on it avoids missing early ones.
func (w *wsNode) UpstreamReady() bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	return len(w.upstreamLive) > 0
}

func (w *wsNode) hasDemand() bool {
	w.hub.mu.Lock()
	defer w.hub.mu.Unlock()
	return len(w.hub.demand) > 0
}

func (w *wsNode) connectAll(ctx context.Context) {
	if !w.hasDemand() {
		return
	}
	for _, peer := range w.node.cfg.Peers {
		w.mu.Lock()
		live := w.upstreamLive[peer] || w.upstreamStarting[peer]
		if !live {
			w.upstreamStarting[peer] = true
		}
		w.mu.Unlock()
		if live {
			continue
		}

		go func(peer string) {
			defer func() {
				w.mu.Lock()
				delete(w.upstreamStarting, peer)
				delete(w.upstreamLive, peer)
				w.mu.Unlock()
			}()
			w.serveUpstream(ctx, peer)
		}(peer)
	}
}

// serveUpstream runs one connection lifecycle: authenticate over HTTP,
// upgrade with the token, subscribe to everything currently demanded, pump
// events until the connection or ctx dies. Reconnection happens on the next
// manager tick because demand persists.
func (w *wsNode) serveUpstream(ctx context.Context, peer string) {
	token, err := w.auth.session(ctx, peer)
	if err != nil {
		return
	}
	wsURL := strings.Replace(peer, "http", "ws", 1) + "/v1/ws?token=" + token
	conn, _, err := websocket.Dial(ctx, wsURL, nil)
	if err != nil {
		w.auth.invalidate(peer)
		return
	}
	defer conn.Close(websocket.StatusNormalClosure, "")

	w.hub.mu.Lock()
	for k := range w.hub.demand {
		raw, _ := json.Marshal(map[string]string{"op": "subscribe", "channel": k.channel, "key": k.key})
		if err := conn.Write(ctx, websocket.MessageText, raw); err != nil {
			w.hub.mu.Unlock()
			return
		}
	}
	w.hub.mu.Unlock()

	w.mu.Lock()
	w.upstreamLive[peer] = true
	w.mu.Unlock()

	for {
		typ, raw, err := conn.Read(ctx)
		if err != nil {
			return
		}
		if typ != websocket.MessageText {
			continue
		}
		w.handleUpstreamEvent(raw)
	}
}

func handleUpstreamEvent(w *wsNode, raw []byte) { w.handleUpstreamEvent(raw) }

func (w *wsNode) handleUpstreamEvent(raw []byte) {
	var ev struct {
		Event    string `json:"event"`
		Channel  string `json:"channel"`
		Key      string `json:"key"`
		ObjectID string `json:"object_id"`
		Envelope *struct {
			TCE       []byte `json:"tce"`
			Signature []byte `json:"signature"`
		} `json:"envelope"`
	}
	if json.Unmarshal(raw, &ev) != nil || ev.Event != "object" || ev.Envelope == nil {
		return
	}
	objID := ev.ObjectID
	if len(objID) != 64 || ev.Channel == "" || ev.Key == "" {
		return
	}

	w.mu.Lock()
	if _, dup := w.seen[objID]; dup {
		w.mu.Unlock()
		return // another peer already delivered this exact object
	}
	w.seen[objID] = struct{}{}
	if len(w.seen) > lnSeenCap {
		w.seen = make(map[string]struct{}) // coarse reset; dedupe is best-effort
	}
	w.mu.Unlock()

	env := &transport.Envelope{TCE: ev.Envelope.TCE, Signature: ev.Envelope.Signature}
	typName, err := env.Verify()
	if err != nil || !validLNChannel(typName) {
		return // never forward an envelope that does not verify locally
	}
	w.dispatch(ev.Channel, ev.Key, objID, ev.Envelope.TCE, ev.Envelope.Signature)
}

func validLNChannel(typeName string) bool {
	switch typeName {
	case "claim", "request", "response", "revocation":
		return true
	}
	return false
}

// dispatch pushes to matching downstream clients and queues the object id
// for proof bookkeeping against the next quorum-agreed root.
func (w *wsNode) dispatch(channel, key, objID string, tceBytes, sig []byte) {
	payload, _ := json.Marshal(map[string]any{
		"event":     "object",
		"channel":   channel,
		"key":       key,
		"object_id": objID,
		"envelope":  map[string]any{"tce": tceBytes, "signature": sig},
		"verified":  false,
	})

	w.hub.mu.Lock()
	for c := range w.hub.clients {
		match := false
		for _, s := range c.subs {
			if s.channel == channel && s.key == key {
				match = true
				break
			}
		}
		if !match || c.evicted {
			continue
		}
		select {
		case c.send <- payload:
		default:
			// Slow client under lock: close-once signalling, no send/close race.
			c.evicted = true
			close(c.done)
			delete(w.hub.clients, c)
		}
	}
	w.hub.mu.Unlock()

	w.mu.Lock()
	if len(w.pending) < lnPendingCap {
		w.pending = append(w.pending, pendingEvent{channel: channel, key: key, objID: objID})
	}
	w.mu.Unlock()
}

// flushPending re-proves queued object ids against the newest agreed root,
// keeping only ids that are still unproven and under the attempt budget.
func (w *wsNode) flushPending() {
	root, ok := w.node.decisionRootOK()
	if !ok {
		return
	}
	w.mu.Lock()
	defer w.mu.Unlock()
	kept := w.pending[:0]
	for _, p := range w.pending {
		idb, err := hex.DecodeString(p.objID)
		var id [32]byte
		if err == nil && len(idb) == 32 {
			copy(id[:], idb)
			if proof, err := w.node.inclusionProofFromPeer(p.objID); err == nil &&
				smt.VerifyInclusion(root, id, proof) {
				continue // now covered by the quorum-agreed checkpoint
			}
		}
		p.attempts++
		if p.attempts < 3 {
			kept = append(kept, p)
		}
	}
	w.pending = kept
}

// websocketAccept upgrades with permissive origins: the light node is a
// local mirror whose authority comes from its pinned peers, not from the
// browser that talks to it.
func websocketAccept(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) {
	return websocket.Accept(w, r, &websocket.AcceptOptions{OriginPatterns: []string{"*"}})
}

var _ = http.StatusText