niko_trust/internal/lightnode/handler.go

package lightnode

import (
	"context"

	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"github.com/coder/websocket"

	"git.n1ko.dev/Niko/niko_trust/internal/bft"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

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

// errNoDecision is returned until the first quorum-approved head exists.
var errNoDecision = errors.New("lightnode: no agreed checkpoint yet")

func (n *Node) pickPeer() string {
	n.mu.Lock()
	defer n.mu.Unlock()
	for _, p := range n.cfg.Peers {
		return p
	}
	return ""
}

// objectFromCache returns a cached envelope body, if present.
func (n *Node) objectFromCache(id string) ([]byte, bool) {
	if n.cfg.CacheDir == "" {
		return nil, false
	}
	raw, err := os.ReadFile(filepath.Join(n.cfg.CacheDir, "objects", id+".json"))
	return raw, err == nil
}

func (n *Node) cacheObject(id string, body []byte) {
	if n.cfg.CacheDir == "" {
		return
	}
	dir := filepath.Join(n.cfg.CacheDir, "objects")
	if os.MkdirAll(dir, 0o700) == nil {
		_ = os.WriteFile(filepath.Join(dir, id+".json"), body, 0o600)
	}
	n.evictLocked()
}

// evictLocked drops oldest cached objects while over budget.
func (n *Node) evictLocked() {
	dir := filepath.Join(n.cfg.CacheDir, "objects")
	entries, err := os.ReadDir(dir)
	if err != nil {
		return
	}
	type item struct {
		name string
		size int64
		mod  int64
	}
	var items []item
	var total int64
	for _, e := range entries {
		fi, err := e.Info()
		if err != nil {
			continue
		}
		items = append(items, item{e.Name(), fi.Size(), fi.ModTime().UnixNano()})
		total += fi.Size()
	}
	if total <= n.cfg.CacheMaxBytes {
		return
	}
	sort.Slice(items, func(i, j int) bool { return items[i].mod < items[j].mod })
	for _, it := range items {
		if total <= n.cfg.CacheMaxBytes {
			break
		}
		if os.Remove(filepath.Join(dir, it.name)) == nil {
			total -= it.size
		}
	}
}

// decisionRoot returns the root every served proof must match.
func (n *Node) decisionRoot() ([32]byte, error) {
	n.mu.Lock()
	defer n.mu.Unlock()
	if n.decision == nil {
		return [32]byte{}, errNoDecision
	}
	return n.decision.Root, nil
}

// decisionRootOK is decisionRoot with a boolean for callers that treat "no
// decision yet" as a plain no-op.
func (n *Node) decisionRootOK() ([32]byte, bool) {
	r, err := n.decisionRoot()
	return r, err == nil
}

// inclusionProofFromPeer fetches an inclusion proof from any peer. The
// caller verifies it against the root it trusts.
func (n *Node) inclusionProofFromPeer(idHex string) ([]byte, error) {
	pr, err := n.verifiedProof(context.Background(), idHex, true)
	if err != nil {
		return nil, err
	}
	return base64.StdEncoding.DecodeString(pr.Proof)
}

type proofResponse struct {
	Proof string `json:"proof"`
	Root  string `json:"root"`
}

// verifiedProof fetches an inclusion/absence proof from any peer and accepts
// it only if it verifies against the quorum-agreed root.
func (n *Node) verifiedProof(ctx context.Context, idHex string, wantPresent bool) (*proofResponse, error) {
	root, err := n.decisionRoot()
	if err != nil {
		return nil, err
	}
	kind := "absent"
	if wantPresent {
		kind = "object"
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet,
		n.pickPeer()+"/v1/proof/"+kind+"/"+idHex, nil)
	if err != nil {
		return nil, err
	}
	resp, err := n.http.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, errors.New("lightnode: proof status " + resp.Status)
	}
	var pr proofResponse
	if json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&pr) != nil {
		return nil, errors.New("lightnode: bad proof json")
	}
	proof, err := base64.StdEncoding.DecodeString(pr.Proof)
	if err != nil {
		return nil, err
	}
	var id [32]byte
	idb, err := hex.DecodeString(idHex)
	if err != nil || len(idb) != 32 {
		return nil, errors.New("lightnode: bad id")
	}
	copy(id[:], idb)
	if wantPresent && !smt.VerifyInclusion(root, id, proof) {
		return nil, errors.New("lightnode: inclusion proof rejected against quorum root")
	}
	if !wantPresent && !smt.VerifyAbsence(root, id, proof) {
		return nil, errors.New("lightnode: absence proof rejected against quorum root")
	}
	return &pr, nil
}

// Handler returns the relay-compatible read API plus node status.
func (n *Node) Handler() http.Handler {
	mux := http.NewServeMux()

	mux.HandleFunc("GET /v1/objects/{id}", func(w http.ResponseWriter, r *http.Request) {
		id := r.PathValue("id")
		if len(id) != 64 {
			writeErr(w, http.StatusBadRequest, "bad object id")
			return
		}
		if body, ok := n.objectFromCache(id); ok {
			w.Header().Set("Content-Type", "application/json")
			w.Header().Set("X-Cache", "hit")
			w.Write(body)
			return
		}
		pr, err := n.verifiedProof(r.Context(), id, true)
		if err != nil {
			writeErr(w, http.StatusConflict, err.Error())
			return
		}
		req, err := http.NewRequestWithContext(r.Context(), http.MethodGet,
			n.pickPeer()+"/v1/objects/"+id, nil)
		if err != nil {
			writeErr(w, http.StatusInternalServerError, err.Error())
			return
		}
		resp, err := n.http.Do(req)
		if err != nil {
			writeErr(w, http.StatusBadGateway, err.Error())
			return
		}
		defer resp.Body.Close()
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
		if resp.StatusCode != http.StatusOK {
			writeErr(w, resp.StatusCode, string(body))
			return
		}
		n.cacheObject(id, body)
		w.Header().Set("Content-Type", "application/json")
		w.Header().Set("X-Proof-Root", pr.Root)
		w.Write(body)
	})

	mux.HandleFunc("GET /v1/checkpoint/latest", func(w http.ResponseWriter, r *http.Request) {
		n.mu.Lock()
		d := n.decision
		n.mu.Unlock()
		if d == nil {
			writeErr(w, http.StatusNotFound, "no agreed checkpoint yet")
			return
		}
		writeJSON(w, http.StatusOK, map[string]any{
			"bytes":      base64.StdEncoding.EncodeToString(d.Bytes),
			"signature":  base64.StdEncoding.EncodeToString(d.Signature),
			"id":         d.ID,
			"public_key": d.PublicKey,
		})
	})

	mux.HandleFunc("GET /v1/proof/object/{id}", func(w http.ResponseWriter, r *http.Request) {
		serveVerifiedProof(n, w, r, true)
	})
	mux.HandleFunc("GET /v1/proof/absent/{id}", func(w http.ResponseWriter, r *http.Request) {
		serveVerifiedProof(n, w, r, false)
	})

	mux.HandleFunc("GET /v1/config", func(w http.ResponseWriter, r *http.Request) {
		keys := make([]string, 0, len(n.pins))
		for k := range n.pins {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		writeJSON(w, http.StatusOK, map[string]any{
			"mode":        "lightnode",
			"peers":       n.cfg.Peers,
			"pinned_keys": keys,
			"quorum":      n.cfg.Quorum,
		})
	})

	mux.HandleFunc("GET /v1/bft/certificate/{height}", func(w http.ResponseWriter, r *http.Request) {
		if len(n.cfg.BFTValidators) == 0 {
			writeErr(w, http.StatusNotImplemented, "no validator set pinned")
			return
		}
		h := r.PathValue("height")
		for _, peer := range n.cfg.Peers {
			req, err := http.NewRequestWithContext(r.Context(), http.MethodGet,
				strings.TrimSuffix(peer, "/")+"/v1/bft/certificate/"+h, nil)
			if err != nil {
				continue
			}
			resp, err := n.http.Do(req)
			if err != nil {
				continue
			}
			raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
			resp.Body.Close()
			if resp.StatusCode != http.StatusOK {
				continue
			}
			var cert bft.Certificate
			if json.Unmarshal(raw, &cert) != nil {
				continue
			}
			set, err := bft.NewSet(n.cfg.BFTValidators, make([]string, len(n.cfg.BFTValidators)))
			if err != nil || set.VerifyCertificate(&cert) != nil {
				writeErr(w, http.StatusConflict, "certificate fails against pinned validators")
				return
			}
			w.Header().Set("Content-Type", "application/json")
			w.Write(raw)
			return
		}
		writeErr(w, http.StatusBadGateway, "no peer served a certificate")
	})

	mux.HandleFunc("GET /v1/ws", func(w http.ResponseWriter, r *http.Request) {
		serveDownstreamWS(n, w, r)
	})

	mux.HandleFunc("GET /v1/healthz", func(w http.ResponseWriter, r *http.Request) {
		writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
	})
	mux.HandleFunc("GET /v1/readyz", func(w http.ResponseWriter, r *http.Request) {
		n.mu.Lock()
		d := n.decision
		n.mu.Unlock()
		if d == nil {
			writeErr(w, http.StatusServiceUnavailable, "no agreed checkpoint yet")
			return
		}
		writeJSON(w, http.StatusOK, map[string]string{"status": "ready"})
	})

	return mux
}

func serveVerifiedProof(n *Node, w http.ResponseWriter, r *http.Request, wantPresent bool) {
	id := r.PathValue("id")
	if len(id) != 64 {
		writeErr(w, http.StatusBadRequest, "bad object id")
		return
	}
	pr, err := n.verifiedProof(r.Context(), id, wantPresent)
	if err != nil {
		status := http.StatusBadGateway
		if errors.Is(err, errNoDecision) {
			status = http.StatusServiceUnavailable
		} else if strings.Contains(err.Error(), "rejected") {
			status = http.StatusConflict
		}
		writeErr(w, status, err.Error())
		return
	}
	writeJSON(w, http.StatusOK, pr)
}

func writeErr(w http.ResponseWriter, code int, msg string) {
	writeJSON(w, code, map[string]string{"error": msg})
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(code)
	_ = json.NewEncoder(w).Encode(v)
}

// Run refreshes on a ticker until ctx ends.
func (n *Node) Run(ctx context.Context, interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			_ = n.Refresh(ctx)
		}
	}
}

// serveDownstreamWS upgrades a local client and mirrors the relay's frame
// protocol. No authentication: the light node is an operator-run local
// mirror, and its own credentials never delegate.
func serveDownstreamWS(n *Node, w http.ResponseWriter, r *http.Request) {
	wsn, err := n.WS()
	if err != nil {
		writeErr(w, http.StatusInternalServerError, "ws unavailable")
		return
	}
	conn, err := websocketAccept(w, r)
	if err != nil {
		return
	}
	defer conn.Close(websocket.StatusInternalError, "")

	client := &lnClient{send: make(chan []byte, lnSendBuffer), done: make(chan struct{})}
	wsn.hub.mu.Lock()
	if len(wsn.hub.clients) >= lnMaxConns {
		wsn.hub.mu.Unlock()
		conn.Close(websocket.StatusPolicyViolation, "too many connections")
		return
	}
	wsn.hub.clients[client] = struct{}{}
	wsn.hub.mu.Unlock()
	defer wsn.hub.evict(client)

	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			select {
			case <-r.Context().Done():
				return
			case <-client.done:
				conn.Close(websocket.StatusPolicyViolation, "evicted")
				return
			case payload := <-client.send:
				if err := conn.Write(r.Context(), websocket.MessageText, payload); err != nil {
					return
				}
			}
		}
	}()

	for {
		typ, raw, 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(raw, &msg) != nil || (msg.Op != "subscribe" && msg.Op != "unsubscribe") {
			safeSend(conn, r, map[string]string{"event": "error", "message": "bad control frame"})
			continue
		}
		if !lnChannelKnown(msg.Channel) {
			safeSend(conn, r, map[string]string{"event": "error", "message": "unknown channel"})
			continue
		}

		wsn.hub.mu.Lock()
		switch msg.Op {
		case "subscribe":
			if msg.Key == "" || len(client.subs) >= lnMaxSubsPerConn {
				wsn.hub.mu.Unlock()
				safeSend(conn, r, map[string]string{"event": "error", "message": "bad key or too many subscriptions"})
				continue
			}
			client.subs = append(client.subs, lnSub{msg.Channel, msg.Key})
			wsn.hub.demand[lnSub{msg.Channel, msg.Key}]++
		case "unsubscribe":
			kept := client.subs[:0]
			for _, x := range client.subs {
				if x.channel == msg.Channel && (msg.Key == "" || x.key == msg.Key) {
					k := lnSub{x.channel, x.key}
					wsn.hub.demand[k]--
					if wsn.hub.demand[k] <= 0 {
						delete(wsn.hub.demand, k)
					}
					continue
				}
				kept = append(kept, x)
			}
			client.subs = kept
		}
		wsn.hub.mu.Unlock()
		wsn.hub.signal()
		safeSend(conn, r, map[string]string{"event": "subscribed", "channel": msg.Channel, "key": msg.Key})
	}
	<-done // let the writer drain before the deferred close kills it
}

func safeSend(conn *websocket.Conn, r *http.Request, v any) {
	raw, _ := json.Marshal(v)
	_ = conn.Write(r.Context(), websocket.MessageText, raw)
}

func lnChannelKnown(ch string) bool {
	switch ch {
	case "claims", "requests", "responses", "revocations":
		return true
	}
	return false
}