kms/pkg/zapserver/handshake_e2e_test.go

// End-to-end tests for the LP-022 hybrid PQ handshake wired through
// pkg/zapserver and pkg/zapclient. We boot a real zap.Node listening on
// a loopback port, register the secret server (which advertises
// CapMLKEM768), then point a zapclient at it and assert:
//
//  1. After Dial, the client has an active hybrid session.
//  2. A Get round-trips through sealed payloads — the server
//     authenticates, opens the request body, fetches the secret, and
//     returns a sealed reply that the client opens to recover the
//     plaintext value.
//  3. With LocalCaps=0 the client forces classical-only: the server's
//     handshake handler emits a WARN and the client reports
//     hybrid=false. The session still works (32-byte X25519 key).
package zapserver

import (
	"context"
	"crypto/rand"
	"net"
	"strconv"
	"testing"
	"time"

	"github.com/luxfi/keys"
	"github.com/luxfi/kms/pkg/envelope"
	"github.com/luxfi/kms/pkg/store"
	kmszap "github.com/luxfi/kms/pkg/zap"
	"github.com/luxfi/kms/pkg/zapclient"
	"github.com/luxfi/log"
	"github.com/luxfi/zap"
	badger "github.com/luxfi/zapdb"
)

// allowAuthorizer is the test-only ConsensusAuthorizer that lets every
// signed envelope through. The handshake e2e suite is about the PQ
// transport, not the authority predicate.
type allowAuthorizer struct{}

// Authorize always allows. Returns a single-reason allow so audit logs
// can still attribute the decision.
func (allowAuthorizer) Authorize(_ context.Context, _ Identity, _ string, _ Op) (Decision, error) {
	return Allow("test-allow-all"), nil
}

// newE2EIdentity builds a *keys.ServiceIdentity + envelope.IdentityHeader
// pair for the handshake tests. Uses the canonical abandon... about
// test mnemonic so derivation is deterministic.
func newE2EIdentity(t *testing.T, servicePath string) (*keys.ServiceIdentity, envelope.IdentityHeader) {
	t.Helper()
	ident, err := keys.NewServiceIdentity(testMnemonic, servicePath)
	if err != nil {
		t.Fatalf("NewServiceIdentity(%q): %v", servicePath, err)
	}
	hdr := envelope.IdentityHeader{
		NodeID:      ident.NodeID,
		FullDigest:  ident.FullDigest,
		ServicePath: ident.ServicePath,
		PublicKey:   ident.PublicKey,
	}
	return ident, hdr
}

// freePort grabs an OS-assigned loopback port and returns it. We need a
// known port in the client config, and zap.Node always advertises the
// configured port on its listener (it does not expose ln.Addr() back).
func freePort(t *testing.T) int {
	t.Helper()
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	addr := l.Addr().(*net.TCPAddr)
	l.Close()
	return addr.Port
}

// bootServer brings up a real zap.Node + Server with a seeded SecretStore
// and an open ACL (so ACL is not the variable under test). Returns the
// listen addr and a cleanup. Caller MUST defer cleanup() and t.Cleanup.
func bootServer(t *testing.T) (addr string, peerNodeID string) {
	t.Helper()

	port := freePort(t)
	nodeID := "kms-test-" + strconv.Itoa(port)
	n := zap.NewNode(zap.NodeConfig{
		NodeID:      nodeID,
		ServiceType: "_kms._tcp",
		Port:        port,
		NoDiscovery: true, // disable mDNS for hermetic tests
	})
	if err := n.Start(); err != nil {
		t.Fatalf("zap.Node start: %v", err)
	}
	t.Cleanup(func() { n.Stop() })

	opts := badger.DefaultOptions("").WithInMemory(true)
	db, err := badger.Open(opts)
	if err != nil {
		t.Fatalf("zapdb open: %v", err)
	}
	t.Cleanup(func() { db.Close() })

	mk := make([]byte, 32)
	if _, err := rand.Read(mk); err != nil {
		t.Fatalf("rand: %v", err)
	}
	secStore := store.NewSecretStore(db)

	// Seed a value so Get returns OK rather than NotFound.
	sec, err := store.Seal(mk, "ats", "settlement-key", "dev", []byte("hunter2"))
	if err != nil {
		t.Fatalf("store.Seal: %v", err)
	}
	if err := secStore.Put(sec); err != nil {
		t.Fatalf("store.Put: %v", err)
	}

	// Consensus-authorizer that allows any envelope through; the
	// handshake tests are about the PQ session, not the authority
	// predicate. allowAuthorizer is implemented at the bottom of this
	// file.
	srv := New(Config{
		Store:      secStore,
		MasterKey:  mk,
		Authorizer: allowAuthorizer{},
		Logger:     log.NewNoOpLogger(),
	})
	srv.Register(n)

	// luxfi/zap.Node listens on 0.0.0.0:port.
	return "127.0.0.1:" + strconv.Itoa(port), nodeID
}

// TestHandshakeE2E_HybridSecretRoundTrip is the full integration: client
// hello with CapMLKEM768, server replies with same, both derive a
// session, then the client runs a real Get that arrives sealed and the
// server returns a sealed reply.
func TestHandshakeE2E_HybridSecretRoundTrip(t *testing.T) {
	addr, _ := bootServer(t)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	ident, hdr := newE2EIdentity(t, "ats/test-service")
	defer ident.Wipe()
	c, err := zapclient.DialWithConfig(ctx, zapclient.Config{
		NodeID:         "test-client",
		PeerAddr:       addr,
		DefaultPath:    "ats",
		IdentityHeader: hdr,
		Signer:         ident,
	})
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer c.Close()

	// Drive a real Get. Under the hood the client seals
	// {"path":"ats","name":"settlement-key","env":"dev"} and the server
	// opens it before the ACL/store path. The server then seals
	// {"value":"<base64>"} and the client opens it.
	value, err := c.Get(ctx, "settlement-key", "dev")
	if err != nil {
		t.Fatalf("Get: %v", err)
	}
	if value != "hunter2" {
		t.Fatalf("Get value mismatch: got %q want %q", value, "hunter2")
	}
}

// TestHandshakeE2E_FallbackClassicalOnly drives the client with caps=0
// (LocalCaps explicitly set) so the negotiated session is classical only.
// Get still works because the X25519-only key is still 32 bytes and
// drives the same AEAD layer.
func TestHandshakeE2E_FallbackClassicalOnly(t *testing.T) {
	addr, _ := bootServer(t)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	ident, hdr := newE2EIdentity(t, "ats/test-service-classical")
	defer ident.Wipe()
	c, err := zapclient.DialWithConfig(ctx, zapclient.Config{
		NodeID:         "test-client-classical",
		PeerAddr:       addr,
		DefaultPath:    "ats",
		LocalCaps:      0,
		CapsExplicit:   true,
		IdentityHeader: hdr,
		Signer:         ident,
	})
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer c.Close()

	value, err := c.Get(ctx, "settlement-key", "dev")
	if err != nil {
		t.Fatalf("Get (classical): %v", err)
	}
	if value != "hunter2" {
		t.Fatalf("Get value mismatch: got %q want %q", value, "hunter2")
	}
}

// A peer that will not agree keys cannot be talked to at all. There is
// no plaintext path to fall back to, so the dial fails and no client is
// returned — nothing exists that could have carried a secret in the
// clear. Here the peer is a bare ZAP node with no KMS server on it, so
// nothing answers OpClientHello.
func TestHandshakeE2E_PeerThatAgreesNoKeysIsRefused(t *testing.T) {
	port := freePort(t)
	bare := zap.NewNode(zap.NodeConfig{
		NodeID:      "bare-" + strconv.Itoa(port),
		ServiceType: "_kms._tcp",
		Port:        port,
		NoDiscovery: true,
	})
	if err := bare.Start(); err != nil {
		t.Fatalf("zap.Node start: %v", err)
	}
	t.Cleanup(func() { bare.Stop() })

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	ident, hdr := newE2EIdentity(t, "ats/test-service-bare")
	defer ident.Wipe()
	c, err := zapclient.DialWithConfig(ctx, zapclient.Config{
		NodeID:         "test-client-bare",
		PeerAddr:       "127.0.0.1:" + strconv.Itoa(port),
		DefaultPath:    "ats",
		IdentityHeader: hdr,
		Signer:         ident,
	})
	if err == nil {
		c.Close()
		t.Fatal("dialled a peer that agreed no keys")
	}
}

// TestHandshakeE2E_SessionInstalled — once the handshake completes, the
// server has a Session entry under the client's NodeID. We probe via the
// exported helper API: send a fresh ClientHello to a new client and
// verify the server reuses the same NodeID slot (last-write-wins).
func TestHandshakeE2E_SessionInstalled(t *testing.T) {
	addr, _ := bootServer(t)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	ident, hdr := newE2EIdentity(t, "ats/test-service-session")
	defer ident.Wipe()
	c, err := zapclient.DialWithConfig(ctx, zapclient.Config{
		NodeID:         "test-client-session",
		PeerAddr:       addr,
		DefaultPath:    "ats",
		IdentityHeader: hdr,
		Signer:         ident,
	})
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer c.Close()

	// One Get is enough to demonstrate the session is hot — the call
	// would fail with "session decrypt failed" if the server did not
	// hold the matching key.
	if _, err := c.Get(ctx, "settlement-key", "dev"); err != nil {
		t.Fatalf("Get: %v", err)
	}

	// Belt-and-braces: confirm the session opcode constants the client
	// uses match what the server module exports — drift here would
	// silently break the handshake on prod.
	if kmszap.OpClientHello != 0x00F0 || kmszap.OpServerHello != 0x00F1 {
		t.Fatalf("opcode drift: client=0x%04x/0x%04x", kmszap.OpClientHello, kmszap.OpServerHello)
	}
}