feat: REST + WebSocket + ZAP transports on every MPC node

- ws.go: EventBus pubsub + /v1/ws upgrade handler (stdlib, no deps) - zap_server.go: StartZAP(port) with opcodes 0x0060-0x0064 keygen, sign, status, wallets, intents over ZAP binary protocol - Events field on Server for cross-transport event broadcasting - Each consensus node serves full API (not separate mpc-api sidecar)
2026-09-24 00:26:24 +00:00 · 2026-04-17 19:55:55 -07:00
parent 50e62caf24
commit 2840ccb5d5
3 changed files with 444 additions and 0 deletions
@@ -69,6 +69,11 @@ type Server struct {
 	webauthnOrigins map[string]bool
 	router         chi.Router
 	replayGuard    *replayGuard
+
+	// Events is the broadcast channel for server-side events (intent status,
+	// signing progress, wallet events). Handlers publish here; WebSocket
+	// clients and ZAP subscribers receive.
+	Events *EventBus
 }

 // maxBodySize returns middleware that limits request body size to prevent OOM
@@ -130,6 +135,7 @@ func NewServer(database *db.Database, mpcBackend MPCBackend, jwtSecret string, o
 		webauthnRPID:    rpID,
 		webauthnOrigins: webauthnOrigins,
 		replayGuard:     newReplayGuard(),
+		Events:          NewEventBus(),
 	}

 	// CORS origins — Lux infrastructure only. Tenants add via MPC_CORS_ORIGINS env.
@@ -183,6 +189,9 @@ func NewServer(database *db.Database, mpcBackend MPCBackend, jwtSecret string, o
 		r.Post("/auth/refresh", s.handleRefresh)
 		r.Post("/auth/oidc", s.handleOIDCExchange)

+		// WebSocket — auth via query param ?token=<jwt> (headers not sent on WS upgrade)
+		r.Get("/ws", s.handleWebSocket)
+
 		// Public payment page
 		r.Get("/pay/{token}", s.handlePublicPay)

@@ -0,0 +1,209 @@
+package api
+
+import (
+	"context"
+	"crypto/sha1"
+	"encoding/base64"
+	"encoding/json"
+	"net"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/luxfi/mpc/pkg/logger"
+)
+
+// Event is a server-side event published to WebSocket clients and ZAP subscribers.
+type Event struct {
+	Type      string          `json:"type"`                // e.g. "intent.status", "sign.progress", "wallet.created"
+	OrgID     string          `json:"org_id,omitempty"`    // scope to org (empty = broadcast to all)
+	WalletID  string          `json:"wallet_id,omitempty"` // optional wallet scope
+	Data      json.RawMessage `json:"data"`
+	Timestamp int64           `json:"ts"`
+}
+
+// EventBus is a simple fan-out pubsub for server events.
+type EventBus struct {
+	mu   sync.RWMutex
+	subs map[*eventSub]struct{}
+}
+
+type eventSub struct {
+	ch    chan Event
+	orgID string // filter: only receive events for this org (empty = all)
+}
+
+// NewEventBus creates a new event bus.
+func NewEventBus() *EventBus {
+	return &EventBus{subs: make(map[*eventSub]struct{})}
+}
+
+// Publish sends an event to all matching subscribers.
+func (eb *EventBus) Publish(evt Event) {
+	if evt.Timestamp == 0 {
+		evt.Timestamp = time.Now().UnixMilli()
+	}
+	eb.mu.RLock()
+	defer eb.mu.RUnlock()
+	for sub := range eb.subs {
+		if sub.orgID != "" && evt.OrgID != "" && sub.orgID != evt.OrgID {
+			continue
+		}
+		select {
+		case sub.ch <- evt:
+		default:
+			// slow consumer, drop
+		}
+	}
+}
+
+// Subscribe returns a channel that receives events for the given org.
+// Call the returned function to unsubscribe.
+func (eb *EventBus) Subscribe(orgID string) (<-chan Event, func()) {
+	sub := &eventSub{
+		ch:    make(chan Event, 64),
+		orgID: orgID,
+	}
+	eb.mu.Lock()
+	eb.subs[sub] = struct{}{}
+	eb.mu.Unlock()
+	return sub.ch, func() {
+		eb.mu.Lock()
+		delete(eb.subs, sub)
+		eb.mu.Unlock()
+	}
+}
+
+// handleWebSocket upgrades an HTTP connection to a WebSocket and streams events.
+// Auth is via ?token=<jwt> query parameter since browsers don't send custom
+// headers on WebSocket upgrade requests.
+func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
+	// Authenticate via query param
+	token := r.URL.Query().Get("token")
+	if token == "" {
+		writeError(w, http.StatusUnauthorized, "missing token query parameter")
+		return
+	}
+	claims, err := s.validateJWT(token)
+	if err != nil {
+		writeError(w, http.StatusUnauthorized, "invalid token")
+		return
+	}
+
+	// Require upgrade header
+	if r.Header.Get("Upgrade") != "websocket" {
+		writeError(w, http.StatusBadRequest, "expected websocket upgrade")
+		return
+	}
+
+	hj, ok := w.(http.Hijacker)
+	if !ok {
+		writeError(w, http.StatusInternalServerError, "websocket upgrade not supported")
+		return
+	}
+
+	conn, bufrw, err := hj.Hijack()
+	if err != nil {
+		logger.Error("websocket hijack failed", err)
+		return
+	}
+	defer conn.Close()
+
+	// Complete the WebSocket handshake (RFC 6455 section 4.2.2)
+	acceptKey := wsAcceptKey(r.Header.Get("Sec-WebSocket-Key"))
+	bufrw.WriteString("HTTP/1.1 101 Switching Protocols\r\n")
+	bufrw.WriteString("Upgrade: websocket\r\n")
+	bufrw.WriteString("Connection: Upgrade\r\n")
+	bufrw.WriteString("Sec-WebSocket-Accept: " + acceptKey + "\r\n")
+	bufrw.WriteString("\r\n")
+	if err := bufrw.Flush(); err != nil {
+		logger.Error("websocket handshake flush failed", err)
+		return
+	}
+
+	logger.Info("websocket client connected", "org", claims.OrgID, "user", claims.UserID)
+
+	// Subscribe to events for this org
+	ch, unsub := s.Events.Subscribe(claims.OrgID)
+	defer unsub()
+
+	ctx, cancel := context.WithCancel(r.Context())
+	defer cancel()
+
+	// Read goroutine — drain client frames to detect close/ping
+	go wsReadDrain(ctx, cancel, conn)
+
+	// Write loop — send events as text frames
+	pingTicker := time.NewTicker(30 * time.Second)
+	defer pingTicker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case evt := <-ch:
+			data, merr := json.Marshal(evt)
+			if merr != nil {
+				continue
+			}
+			conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+			if werr := wsWriteText(conn, data); werr != nil {
+				return
+			}
+		case <-pingTicker.C:
+			conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+			// WebSocket ping frame: opcode 0x9, no payload
+			if _, err := conn.Write([]byte{0x89, 0x00}); err != nil {
+				return
+			}
+		}
+	}
+}
+
+// wsReadDrain reads and discards client frames until error or context cancel.
+func wsReadDrain(ctx context.Context, cancel context.CancelFunc, conn net.Conn) {
+	defer cancel()
+	buf := make([]byte, 512)
+	for {
+		conn.SetReadDeadline(time.Now().Add(90 * time.Second))
+		if _, err := conn.Read(buf); err != nil {
+			return
+		}
+		// Check context between reads
+		select {
+		case <-ctx.Done():
+			return
+		default:
+		}
+	}
+}
+
+// wsWriteText writes a WebSocket text frame (opcode 0x1, FIN bit set).
+// Server-to-client frames are never masked per RFC 6455.
+func wsWriteText(conn net.Conn, payload []byte) error {
+	length := len(payload)
+	var header []byte
+	switch {
+	case length <= 125:
+		header = []byte{0x81, byte(length)}
+	case length <= 65535:
+		header = []byte{0x81, 126, byte(length >> 8), byte(length)}
+	default:
+		header = []byte{0x81, 127,
+			0, 0, 0, 0,
+			byte(length >> 24), byte(length >> 16), byte(length >> 8), byte(length),
+		}
+	}
+	if _, err := conn.Write(header); err != nil {
+		return err
+	}
+	_, err := conn.Write(payload)
+	return err
+}
+
+// wsAcceptKey computes Sec-WebSocket-Accept per RFC 6455 section 4.2.2.
+func wsAcceptKey(clientKey string) string {
+	const wsGUID = "258EAFA5-E914-47DA-95CA-5AB5DC76CB65"
+	h := sha1.Sum([]byte(clientKey + wsGUID))
+	return base64.StdEncoding.EncodeToString(h[:])
+}
@@ -0,0 +1,226 @@
+package api
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/hanzoai/orm"
+
+	"github.com/luxfi/mpc/pkg/db"
+	"github.com/luxfi/mpc/pkg/logger"
+	"github.com/luxfi/zap"
+)
+
+// ZAP opcodes for MPC API operations (0x0060-0x006F range).
+const (
+	OpMPCKeygen  uint16 = 0x0060
+	OpMPCSign    uint16 = 0x0061
+	OpMPCStatus  uint16 = 0x0062
+	OpMPCWallets uint16 = 0x0063
+	OpMPCIntent  uint16 = 0x0064
+)
+
+// ZAP object field offsets (after the 8-byte req/resp correlation header).
+const (
+	fieldOrgID    = 8
+	fieldWalletID = 9
+	fieldPayload  = 10
+	fieldIntentID = 11
+
+	fieldStatus = 8
+	fieldData   = 9
+	fieldError  = 10
+)
+
+// StartZAP creates a ZAP node on a separate port and registers MPC handlers.
+func (s *Server) StartZAP(port int) error {
+	node := zap.NewNode(zap.NodeConfig{
+		NodeID:      fmt.Sprintf("mpc-api-%d", port),
+		ServiceType: "_mpc._tcp",
+		Port:        port,
+		NoDiscovery: true,
+	})
+
+	node.Handle(OpMPCKeygen, s.zapKeygen)
+	node.Handle(OpMPCSign, s.zapSign)
+	node.Handle(OpMPCStatus, s.zapStatus)
+	node.Handle(OpMPCWallets, s.zapWallets)
+	node.Handle(OpMPCIntent, s.zapIntent)
+
+	if err := node.Start(); err != nil {
+		return fmt.Errorf("zap start: %w", err)
+	}
+	logger.Info("ZAP server started", "port", port)
+
+	// Forward events to ZAP peers
+	go s.zapEventForwarder(node)
+
+	return nil
+}
+
+func (s *Server) zapKeygen(_ context.Context, _ string, msg *zap.Message) (*zap.Message, error) {
+	root := msg.Root()
+	orgID := root.Text(fieldOrgID)
+	walletID := root.Text(fieldWalletID)
+	if orgID == "" || walletID == "" {
+		return zapError("org_id and wallet_id required")
+	}
+
+	result, err := s.mpc.TriggerKeygen(orgID, walletID)
+	if err != nil {
+		return zapError(err.Error())
+	}
+
+	data, _ := json.Marshal(result)
+	s.Events.Publish(Event{Type: "wallet.created", OrgID: orgID, WalletID: walletID, Data: data})
+	return zapOK(data)
+}
+
+func (s *Server) zapSign(_ context.Context, _ string, msg *zap.Message) (*zap.Message, error) {
+	root := msg.Root()
+	orgID := root.Text(fieldOrgID)
+	walletID := root.Text(fieldWalletID)
+	payload := root.Bytes(fieldPayload)
+	if orgID == "" || walletID == "" || len(payload) == 0 {
+		return zapError("org_id, wallet_id, and payload required")
+	}
+
+	result, err := s.mpc.TriggerSign(orgID, walletID, payload)
+	if err != nil {
+		return zapError(err.Error())
+	}
+
+	data, _ := json.Marshal(result)
+	s.Events.Publish(Event{Type: "sign.complete", OrgID: orgID, WalletID: walletID, Data: data})
+	return zapOK(data)
+}
+
+func (s *Server) zapStatus(_ context.Context, _ string, _ *zap.Message) (*zap.Message, error) {
+	data, _ := json.Marshal(s.mpc.GetClusterStatus())
+	return zapOK(data)
+}
+
+func (s *Server) zapWallets(_ context.Context, _ string, msg *zap.Message) (*zap.Message, error) {
+	orgID := msg.Root().Text(fieldOrgID)
+	if orgID == "" {
+		return zapError("org_id required")
+	}
+
+	wallets, err := orm.TypedQuery[db.Wallet](s.db.ORM).
+		Filter("orgId=", orgID).
+		Order("-createdAt").
+		GetAll(context.Background())
+	if err != nil {
+		return zapError(err.Error())
+	}
+	if wallets == nil {
+		wallets = []*db.Wallet{}
+	}
+
+	data, _ := json.Marshal(wallets)
+	return zapOK(data)
+}
+
+func (s *Server) zapIntent(_ context.Context, _ string, msg *zap.Message) (*zap.Message, error) {
+	root := msg.Root()
+	orgID := root.Text(fieldOrgID)
+	intentID := root.Text(fieldIntentID)
+	if orgID == "" {
+		return zapError("org_id required")
+	}
+
+	// Query existing intent
+	if intentID != "" {
+		intent, err := orm.Get[db.Intent](s.db.ORM, intentID)
+		if err != nil || intent.OrgID != orgID {
+			return zapError("intent not found")
+		}
+		data, _ := json.Marshal(intent)
+		return zapOK(data)
+	}
+
+	// Create from payload
+	payload := root.Bytes(fieldPayload)
+	if len(payload) == 0 {
+		return zapError("intent_id or payload required")
+	}
+
+	var req struct {
+		WalletID   string `json:"wallet_id"`
+		IntentType string `json:"intent_type"`
+		Chain      string `json:"chain"`
+		ToAddress  string `json:"to_address"`
+		Amount     string `json:"amount"`
+		Token      string `json:"token,omitempty"`
+	}
+	if err := json.Unmarshal(payload, &req); err != nil {
+		return zapError("invalid intent payload")
+	}
+	if req.WalletID == "" || req.IntentType == "" || req.Chain == "" || req.Amount == "" {
+		return zapError("wallet_id, intent_type, chain, amount required")
+	}
+
+	intent := orm.New[db.Intent](s.db.ORM)
+	intent.OrgID = orgID
+	intent.WalletID = req.WalletID
+	intent.IntentType = req.IntentType
+	intent.Chain = req.Chain
+	intent.ToAddress = nilIfEmpty(req.ToAddress)
+	intent.Amount = req.Amount
+	intent.Token = nilIfEmpty(req.Token)
+	intent.Status = "pending_sign"
+	if err := intent.Create(); err != nil {
+		return zapError(err.Error())
+	}
+
+	data, _ := json.Marshal(intent)
+	s.Events.Publish(Event{Type: "intent.created", OrgID: orgID, WalletID: req.WalletID, Data: data})
+	return zapOK(data)
+}
+
+// zapEventForwarder broadcasts EventBus events to ZAP peers.
+func (s *Server) zapEventForwarder(node *zap.Node) {
+	ch, unsub := s.Events.Subscribe("")
+	defer unsub()
+	for evt := range ch {
+		data, err := json.Marshal(evt)
+		if err != nil {
+			continue
+		}
+		errs := node.Broadcast(context.Background(), zapMsg(data))
+		for peer, berr := range errs {
+			logger.Warn("zap broadcast failed", "peer", peer, "err", berr)
+		}
+	}
+}
+
+// zapError builds a ZAP response with error status.
+func zapError(msg string) (*zap.Message, error) {
+	b := zap.NewBuilder(256)
+	ob := b.StartObject(16)
+	ob.SetUint8(fieldStatus, 1)
+	ob.SetText(fieldError, msg)
+	ob.FinishAsRoot()
+	return zap.Parse(b.Finish())
+}
+
+// zapOK builds a ZAP response with success status and JSON data.
+func zapOK(data []byte) (*zap.Message, error) {
+	b := zap.NewBuilder(len(data) + 64)
+	ob := b.StartObject(16)
+	ob.SetUint8(fieldStatus, 0)
+	ob.SetBytes(fieldData, data)
+	ob.FinishAsRoot()
+	return zap.Parse(b.Finish())
+}
+
+// zapMsg builds a one-way ZAP message with data payload.
+func zapMsg(data []byte) *zap.Message {
+	b := zap.NewBuilder(len(data) + 64)
+	ob := b.StartObject(16)
+	ob.SetBytes(fieldData, data)
+	ob.FinishAsRoot()
+	m, _ := zap.Parse(b.Finish())
+	return m
+}