esign/pkg/api/zap_server.go

// Copyright 2026 Lux Industries Inc. and the luxfi/esign authors.
// Licensed under the Apache License, Version 2.0 (see LICENSE).

package api

import (
	"context"
	"encoding/binary"
	"encoding/json"
	"errors"
	"io"
	"net"
	"sync"
	"time"
)

// ZAPServer serves the eSign RPC surface over the ZAP wire:
//
//	[4 byte BE length] [2 byte LE opcode] [JSON body]
//
// Each request/response pair is one length-prefixed frame in each
// direction. Connections are persistent — clients pipeline multiple
// requests on the same socket.
//
// ZAPServer is the same shape as luxfi/mpc and luxfi/kms expose, so the
// downstream ZAP node-discovery + ML-KEM handshake can wrap this server
// transparently; this package owns only the opcode dispatch + framing.
type ZAPServer struct {
	addr     string
	handler  *Handler
	listener net.Listener

	wg     sync.WaitGroup
	mu     sync.Mutex
	closed bool
}

// NewZAPServer binds to addr but does not yet Accept. Call Serve to
// start the accept loop; call Close to tear down cleanly.
func NewZAPServer(addr string, handler *Handler) (*ZAPServer, error) {
	if handler == nil {
		return nil, errors.New("api: zap server: handler required")
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	return &ZAPServer{
		addr:     ln.Addr().String(),
		handler:  handler,
		listener: ln,
	}, nil
}

// Addr returns the bound listen address. Useful when the caller passed
// ":0" to get an ephemeral port.
func (s *ZAPServer) Addr() string { return s.addr }

// Serve runs the accept loop. Returns when Close is called.
func (s *ZAPServer) Serve() error {
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			s.mu.Lock()
			done := s.closed
			s.mu.Unlock()
			if done {
				return nil
			}
			return err
		}
		s.wg.Add(1)
		go func() {
			defer s.wg.Done()
			defer conn.Close()
			s.handleConn(conn)
		}()
	}
}

// Close stops accepting and waits for inflight conns to drain.
func (s *ZAPServer) Close() error {
	s.mu.Lock()
	s.closed = true
	s.mu.Unlock()
	_ = s.listener.Close()
	s.wg.Wait()
	return nil
}

// handleConn pipelines requests on a single connection until the client
// closes the socket or sends a malformed frame.
func (s *ZAPServer) handleConn(conn net.Conn) {
	// Per-conn read deadline kicks in only if the client goes idle mid-frame.
	const idleTimeout = 60 * time.Second
	for {
		_ = conn.SetReadDeadline(time.Now().Add(idleTimeout))

		// Length prefix.
		var lenBuf [4]byte
		if _, err := io.ReadFull(conn, lenBuf[:]); err != nil {
			return
		}
		frameLen := binary.BigEndian.Uint32(lenBuf[:])
		if frameLen < 2 || frameLen > 64*1024*1024 {
			return // malformed
		}

		// Body.
		body := make([]byte, frameLen)
		if _, err := io.ReadFull(conn, body); err != nil {
			return
		}
		op := binary.LittleEndian.Uint16(body[:2])
		payload := body[2:]

		respBody, respErr := s.dispatch(context.Background(), op, payload)

		out := encodeResponse(op, respBody, respErr)
		_ = conn.SetWriteDeadline(time.Now().Add(idleTimeout))
		if _, err := conn.Write(out); err != nil {
			return
		}
	}
}

// dispatch routes one opcode to its handler. Returns either a JSON-encoded
// body or an error message; the framer wraps both into the same wire shape.
func (s *ZAPServer) dispatch(ctx context.Context, op uint16, payload []byte) ([]byte, error) {
	switch op {
	case OpPing:
		return []byte(`{"ok":true}`), nil
	case OpEnvelopeCreate:
		var req CreateRequest
		if err := json.Unmarshal(payload, &req); err != nil {
			return nil, err
		}
		resp, err := s.handler.Create(ctx, req)
		if err != nil {
			return nil, err
		}
		return json.Marshal(resp)
	case OpEnvelopeGet:
		var req struct{ ID string `json:"id"` }
		if err := json.Unmarshal(payload, &req); err != nil {
			return nil, err
		}
		e, err := s.handler.Get(ctx, req.ID)
		if err != nil {
			return nil, err
		}
		return json.Marshal(e)
	case OpEnvelopeSign:
		var req SignRequest
		if err := json.Unmarshal(payload, &req); err != nil {
			return nil, err
		}
		resp, err := s.handler.Sign(ctx, req)
		if err != nil {
			return nil, err
		}
		return json.Marshal(resp)
	case OpEnvelopeTransition:
		var req TransitionRequest
		if err := json.Unmarshal(payload, &req); err != nil {
			return nil, err
		}
		resp, err := s.handler.Transition(ctx, req)
		if err != nil {
			return nil, err
		}
		return json.Marshal(resp)
	case OpEnvelopeVerify:
		var req struct{ ID string `json:"id"` }
		if err := json.Unmarshal(payload, &req); err != nil {
			return nil, err
		}
		resp, err := s.handler.Verify(ctx, req.ID)
		if err != nil {
			return nil, err
		}
		return json.Marshal(resp)
	default:
		return nil, errors.New("api: zap server: unknown opcode")
	}
}

// encodeResponse builds the reply frame: same length-prefix shape as the
// request. On error, body is {"error":"..."} so clients can branch.
func encodeResponse(op uint16, body []byte, err error) []byte {
	if err != nil {
		body, _ = json.Marshal(struct {
			Error string `json:"error"`
		}{err.Error()})
	}
	// 2 byte opcode + body.
	inner := make([]byte, 2+len(body))
	binary.LittleEndian.PutUint16(inner[:2], op)
	copy(inner[2:], body)
	out := make([]byte, 4+len(inner))
	binary.BigEndian.PutUint32(out[:4], uint32(len(inner)))
	copy(out[4:], inner)
	return out
}