secure-messaging/go/internal/leaf/leaf.go

// Package leaf starts an in-process nats-server and optional hub leaf.
package leaf

import (
	"encoding/json"
	"errors"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	natsserver "github.com/nats-io/nats-server/v2/server"
	nats "github.com/nats-io/nats.go"
)

// Node is an in-process broker plus client.
type Node struct {
	ns *natsserver.Server
	nc *nats.Conn
}

// ParseHub rewrites nats-leaf:// to nats:// then parses the URL.
func ParseHub(hub string) (*url.URL, error) {
	raw := strings.Replace(hub, "nats-leaf://", "nats://", 1)
	return url.Parse(raw)
}

// Start binds 127.0.0.1:0 and optionally leaf-connects to hub.
func Start(hub string) (*Node, error) {
	opts := &natsserver.Options{Host: "127.0.0.1", Port: -1}
	if hub != "" {
		u, err := ParseHub(hub)
		if err != nil {
			return nil, err
		}
		opts.LeafNode.Remotes = []*natsserver.RemoteLeafOpts{{URLs: []*url.URL{u}}}
	}
	ns, err := natsserver.NewServer(opts)
	if err != nil {
		return nil, err
	}
	go ns.Start()
	if !ns.ReadyForConnections(5 * time.Second) {
		return nil, errors.New("nats not ready")
	}
	nc, err := nats.Connect(ns.ClientURL())
	if err != nil {
		ns.Shutdown()
		return nil, err
	}
	n := &Node{ns: ns, nc: nc}
	if err := n.subscribe(); err != nil {
		n.Shutdown()
		return nil, err
	}
	return n, nil
}

func (n *Node) subscribe() error {
	ack := func(msg *nats.Msg, body []byte) {
		if msg.Reply != "" {
			_ = n.nc.Publish(msg.Reply, body)
		}
	}
	if _, err := n.nc.Subscribe("verae.sm.send", func(msg *nats.Msg) {
		n.handleSend(msg)
	}); err != nil {
		return err
	}
	if _, err := n.nc.Subscribe("verae.sm.dead", func(msg *nats.Msg) {
		ack(msg, []byte(`{"queued":true}`))
	}); err != nil {
		return err
	}
	if _, err := n.nc.Subscribe("verae.sm.error", func(msg *nats.Msg) {
		ack(msg, []byte(`{"emitted":true}`))
	}); err != nil {
		return err
	}
	if _, err := n.nc.Subscribe("verae.sm.log.summary", func(msg *nats.Msg) {
		var hdr map[string]any
		_ = json.Unmarshal(msg.Data, &hdr)
		// Summary only: never log ciphertext or recipient payload.
		if _, ok := hdr["ct"]; ok {
			log.Printf("sm-summary rejected body-like field")
			ack(msg, []byte(`{"logged":false,"error":"body-like fields present"}`))
			return
		}
		log.Printf("sm-summary code=%v lookup=%v class=%v", hdr["error_code"], hdr["lookup_id"], hdr["dest_class"])
		ack(msg, []byte(`{"logged":true}`))
	}); err != nil {
		return err
	}
	return nil
}

func (n *Node) handleSend(msg *nats.Msg) {
	var env map[string]any
	if err := json.Unmarshal(msg.Data, &env); err != nil {
		n.fail(msg, "SM-BAD-JSON", "parse", "")
		return
	}
	to, _ := env["to"].(string)
	alg, _ := env["alg"].(string)
	ct, _ := env["ct"].(string)
	lid, _ := env["from_lookup_id"].(string)
	if _, hasBody := env["body"]; hasBody {
		n.fail(msg, "SM-PLAINTEXT-BODY", "mailbox", lid)
		return
	}
	if to == "" {
		n.fail(msg, "SM-MISSING-TO", "mailbox", lid)
		return
	}
	if alg != "" && alg != "npe" && alg != "lab-xor" && alg != "plain-lab" {
		n.fail(msg, "SM-BAD-ALG", "mailbox", lid)
		return
	}
	if alg != "plain-lab" && ct == "" {
		n.fail(msg, "SM-EMPTY-CT", "mailbox", lid)
		return
	}
	ack, _ := json.Marshal(map[string]any{"accepted": true, "lookup_id": lid})
	if msg.Reply != "" {
		_ = n.nc.Publish(msg.Reply, ack)
	}
}

func (n *Node) fail(msg *nats.Msg, code, destClass, lid string) {
	summary, _ := json.Marshal(map[string]any{
		"error_code": code,
		"lookup_id":  lid,
		"dest_class": destClass,
	})
	_ = n.nc.Publish("verae.sm.log.summary", summary)
	_ = n.nc.Publish("verae.sm.error", summary)
	_ = n.nc.Publish("verae.sm.dead", summary)
	ack, _ := json.Marshal(map[string]any{
		"accepted":   false,
		"error_code": code,
		"lookup_id":  lid,
	})
	if msg.Reply != "" {
		_ = n.nc.Publish(msg.Reply, ack)
	}
}

// HealthHandler is the loopback JSON health mux (does not expose NATS).
func HealthHandler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"ok":true,"service":"sm-leaf"}`))
	})
	return mux
}

// ServeHealth binds a loopback health listener (does not expose NATS).
func ServeHealth(addr string) error {
	return http.ListenAndServe(addr, HealthHandler())
}

// Addr is the in-process client URL.
func (n *Node) Addr() string { return n.ns.ClientURL() }

// Shutdown stops client and server.
func (n *Node) Shutdown() error {
	if n.nc != nil {
		n.nc.Close()
	}
	if n.ns != nil {
		n.ns.Shutdown()
	}
	return nil
}