From 57daf19ae87b2abe594d95b75ecdb05eb876875c Mon Sep 17 00:00:00 2001 From: TickTockBent Date: Tue, 12 May 2026 08:30:25 -0400 Subject: [PATCH 1/5] phase 1 (#135): WS transport package Ports repram-mcp/src/node/ws-transport.ts to Go in internal/transport/ws/. Substrate/transient WebSocket attachments share the AttachmentMessage envelope; gossip-typed payloads serialize to the same JSON shape as the HTTP gossip endpoint, so handlers process WS and HTTP frames identically. - Connection: heartbeat, HMAC sign/verify (shares gossip.SignBody/VerifyBody), handler dispatch for messages / attachments / close / error, write serialization through writeMu. - Client: ConnectToSubstrate dials /v1/ws with a configurable timeout. - Server: Handler upgrades incoming requests and hands the wrapped Connection to an onAccept callback. - 25 unit tests at parity with ws-transport.test.ts (gossip round-trip, hello/welcome/goodbye, HMAC accept/reject paths, heartbeat send + pong-reset + missed-pong termination, dial-timeout, post-close behavior). Race-clean. Tree manager, relay, and --mcp wiring are deferred to subsequent PRs as specified in the phase plan. // ticktockbent --- go.mod | 3 +- go.sum | 2 + internal/transport/ws/client.go | 41 ++ internal/transport/ws/connection.go | 363 +++++++++++ internal/transport/ws/connection_test.go | 781 +++++++++++++++++++++++ internal/transport/ws/messages.go | 174 +++++ internal/transport/ws/server.go | 38 ++ 7 files changed, 1401 insertions(+), 1 deletion(-) create mode 100644 internal/transport/ws/client.go create mode 100644 internal/transport/ws/connection.go create mode 100644 internal/transport/ws/connection_test.go create mode 100644 internal/transport/ws/messages.go create mode 100644 internal/transport/ws/server.go diff --git a/go.mod b/go.mod index fa98881..0a0a2ce 100644 --- a/go.mod +++ b/go.mod @@ -6,7 +6,9 @@ toolchain go1.22.2 require ( github.com/gorilla/mux v1.8.0 + github.com/gorilla/websocket v1.5.3 github.com/prometheus/client_golang v1.17.0 + github.com/prometheus/client_model v0.4.1-0.20230718164431-9a2bf3000d16 ) require ( @@ -14,7 +16,6 @@ require ( github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/golang/protobuf v1.5.3 // indirect github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect - github.com/prometheus/client_model v0.4.1-0.20230718164431-9a2bf3000d16 // indirect github.com/prometheus/common v0.44.0 // indirect github.com/prometheus/procfs v0.11.1 // indirect golang.org/x/sys v0.15.0 // indirect diff --git a/go.sum b/go.sum index 485de69..f1f5eb6 100644 --- a/go.sum +++ b/go.sum @@ -13,6 +13,8 @@ github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI= github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So= +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo= github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4= github.com/prometheus/client_golang v1.17.0 h1:rl2sfwZMtSthVU752MqfjQozy7blglC+1SOtjMAMh+Q= diff --git a/internal/transport/ws/client.go b/internal/transport/ws/client.go new file mode 100644 index 0000000..815bd91 --- /dev/null +++ b/internal/transport/ws/client.go @@ -0,0 +1,41 @@ +package ws + +import ( + "context" + "errors" + "fmt" + "net/http" + "time" + + "github.com/gorilla/websocket" +) + +// ConnectToSubstrate opens an outbound WebSocket to address:port on the +// /v1/ws endpoint and returns a wrapped Connection. The connection will +// HMAC-sign outgoing frames when clusterSecret is non-empty. A zero timeout +// applies a 10s default. +func ConnectToSubstrate(ctx context.Context, address string, port int, clusterSecret string, timeout time.Duration) (*Connection, error) { + if timeout <= 0 { + timeout = 10 * time.Second + } + url := fmt.Sprintf("ws://%s:%d/v1/ws", address, port) + + dialCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + dialer := &websocket.Dialer{HandshakeTimeout: timeout} + ws, resp, err := dialer.DialContext(dialCtx, url, http.Header{}) + if err != nil { + if resp != nil { + _ = resp.Body.Close() + } + if errors.Is(dialCtx.Err(), context.DeadlineExceeded) { + return nil, fmt.Errorf("WebSocket connection to %s timed out: %w", url, err) + } + return nil, fmt.Errorf("WebSocket connection to %s failed: %w", url, err) + } + if resp != nil { + _ = resp.Body.Close() + } + return NewConnection(ws, clusterSecret), nil +} diff --git a/internal/transport/ws/connection.go b/internal/transport/ws/connection.go new file mode 100644 index 0000000..e5feef9 --- /dev/null +++ b/internal/transport/ws/connection.go @@ -0,0 +1,363 @@ +package ws + +import ( + "encoding/json" + "errors" + "fmt" + "net" + "sync" + "sync/atomic" + "time" + + "github.com/gorilla/websocket" + + "repram/internal/gossip" + "repram/internal/logging" +) + +const ( + // HeartbeatInterval is the default interval between WS pings while the + // heartbeat loop is running. The TS reference uses 30s; matching it keeps + // cross-validation simple. + HeartbeatInterval = 30 * time.Second + + // MaxMissedPongs is how many consecutive intervals without a pong trigger + // termination. With 30s intervals this gives the same ~90s dead-peer + // timeout as the rest of the cluster's ping-based eviction. + MaxMissedPongs = 3 + + // writeTimeout caps individual frame writes. Concurrent writes are + // serialized through writeMu. + writeTimeout = 10 * time.Second +) + +// Connection wraps a *websocket.Conn with the AttachmentMessage envelope, +// optional HMAC signing/verification, and a ping/pong heartbeat. A single +// internal goroutine reads frames and dispatches handlers; writes are +// serialized through writeMu. All public methods are safe for concurrent use. +type Connection struct { + ws *websocket.Conn + clusterSecret string + + writeMu sync.Mutex + closeOnce sync.Once + closed atomic.Bool + + missedPongs atomic.Int32 + heartbeatStop chan struct{} + heartbeatStarted atomic.Bool + heartbeatPeriod time.Duration + + handlersMu sync.RWMutex + onMessage func(*gossip.Message) + onAttachment func(*AttachmentMessage) + onClose func(code int, reason string) + onError func(error) + + remoteMu sync.RWMutex + remoteNodeID string + remoteEnclave string +} + +// NewConnection wraps an open *websocket.Conn. The Connection takes ownership +// of ws — call Close or Terminate to release it. The clusterSecret enables +// HMAC signing of every outgoing payload and verification of every incoming +// payload; passing "" disables signing (open-cluster mode). +func NewConnection(ws *websocket.Conn, clusterSecret string) *Connection { + c := &Connection{ + ws: ws, + clusterSecret: clusterSecret, + heartbeatStop: make(chan struct{}), + heartbeatPeriod: HeartbeatInterval, + } + ws.SetPongHandler(func(string) error { + c.missedPongs.Store(0) + return nil + }) + go c.readLoop() + return c +} + +// OnMessage registers a handler for gossip-typed AttachmentMessages, decoded +// back into the internal Message form. +func (c *Connection) OnMessage(fn func(*gossip.Message)) { + c.handlersMu.Lock() + c.onMessage = fn + c.handlersMu.Unlock() +} + +// OnAttachment registers a handler that fires for every parsed AttachmentMessage +// (gossip + hello/welcome/goodbye). Gossip frames fire both OnAttachment and +// OnMessage; lifecycle frames fire only OnAttachment. +func (c *Connection) OnAttachment(fn func(*AttachmentMessage)) { + c.handlersMu.Lock() + c.onAttachment = fn + c.handlersMu.Unlock() +} + +// OnClose fires once when the underlying connection closes for any reason +// (graceful close, peer hangup, RST, or local Terminate). +func (c *Connection) OnClose(fn func(code int, reason string)) { + c.handlersMu.Lock() + c.onClose = fn + c.handlersMu.Unlock() +} + +// OnError fires for non-fatal read-loop errors that did not terminate the +// connection. Fatal errors surface through OnClose instead. +func (c *Connection) OnError(fn func(error)) { + c.handlersMu.Lock() + c.onError = fn + c.handlersMu.Unlock() +} + +// RemoteNodeID returns the peer's node ID once the hello/welcome handshake +// has populated it. Empty before the handshake completes. +func (c *Connection) RemoteNodeID() string { + c.remoteMu.RLock() + defer c.remoteMu.RUnlock() + return c.remoteNodeID +} + +// RemoteEnclave returns the peer's enclave, populated alongside RemoteNodeID. +func (c *Connection) RemoteEnclave() string { + c.remoteMu.RLock() + defer c.remoteMu.RUnlock() + return c.remoteEnclave +} + +// SetRemote records the peer's identity after a successful hello/welcome. +// Intended to be called by the tree manager once it has processed the +// handshake payload. +func (c *Connection) SetRemote(nodeID, enclave string) { + c.remoteMu.Lock() + c.remoteNodeID = nodeID + c.remoteEnclave = enclave + c.remoteMu.Unlock() +} + +// IsClosed reports whether the connection has been torn down. +func (c *Connection) IsClosed() bool { return c.closed.Load() } + +// SendGossip serializes msg into a SimpleMessage payload, wraps it in an +// AttachmentMessage of the matching gossip type, and writes it. +func (c *Connection) SendGossip(msg *gossip.Message) error { + if c.closed.Load() { + return nil + } + return c.SendAttachment(gossipTypeFor(msg.Type), messageToWire(msg)) +} + +// SendAttachment marshals payload, signs it with the cluster secret when set, +// and writes the framed envelope. +func (c *Connection) SendAttachment(t AttachmentType, payload any) error { + if c.closed.Load() { + return nil + } + payloadBytes, err := json.Marshal(payload) + if err != nil { + return fmt.Errorf("marshal payload: %w", err) + } + env := AttachmentMessage{Type: t, Payload: payloadBytes} + if c.clusterSecret != "" { + env.Signature = gossip.SignBody(c.clusterSecret, payloadBytes) + } + frame, err := json.Marshal(env) + if err != nil { + return fmt.Errorf("marshal envelope: %w", err) + } + return c.writeFrame(frame) +} + +func (c *Connection) writeFrame(frame []byte) error { + c.writeMu.Lock() + defer c.writeMu.Unlock() + if c.closed.Load() { + return nil + } + _ = c.ws.SetWriteDeadline(time.Now().Add(writeTimeout)) + return c.ws.WriteMessage(websocket.TextMessage, frame) +} + +// StartHeartbeat begins sending WS pings at the configured period. After +// MaxMissedPongs consecutive intervals without a pong the connection is +// terminated. Idempotent — repeated calls are no-ops. +func (c *Connection) StartHeartbeat() { + if !c.heartbeatStarted.CompareAndSwap(false, true) { + return + } + go c.heartbeatLoop() +} + +// StopHeartbeat halts the heartbeat goroutine. Idempotent. +func (c *Connection) StopHeartbeat() { + if c.heartbeatStarted.Load() { + c.closeHeartbeat() + } +} + +func (c *Connection) closeHeartbeat() { + defer func() { _ = recover() }() + select { + case <-c.heartbeatStop: + default: + close(c.heartbeatStop) + } +} + +func (c *Connection) heartbeatLoop() { + c.missedPongs.Store(0) + ticker := time.NewTicker(c.heartbeatPeriod) + defer ticker.Stop() + for { + select { + case <-c.heartbeatStop: + return + case <-ticker.C: + if c.closed.Load() { + return + } + // Increment first; an arriving pong resets to zero. Matches TS. + missed := c.missedPongs.Add(1) + if missed > MaxMissedPongs { + c.remoteMu.RLock() + id := c.remoteNodeID + c.remoteMu.RUnlock() + logging.Warn("WebSocket to %q: %d missed pongs, terminating", id, missed) + c.Terminate() + return + } + // WriteControl is documented safe alongside WriteMessage, so we + // do not hold writeMu here — that would block heartbeat on slow + // large-frame writes. + if err := c.ws.WriteControl(websocket.PingMessage, nil, time.Now().Add(writeTimeout)); err != nil { + logging.Debug("WebSocket ping failed: %v", err) + } + } + } +} + +// Close performs a graceful close, sending a close frame with the given code +// and reason. Idempotent. +func (c *Connection) Close(code int, reason string) { + c.closeWith(code, reason, false) +} + +// Terminate tears the connection down immediately without a close frame. +func (c *Connection) Terminate() { + c.closeWith(websocket.CloseAbnormalClosure, "", true) +} + +func (c *Connection) closeWith(code int, reason string, force bool) { + c.closeOnce.Do(func() { + c.closed.Store(true) + c.closeHeartbeat() + if !force { + msg := websocket.FormatCloseMessage(code, reason) + _ = c.ws.WriteControl(websocket.CloseMessage, msg, time.Now().Add(writeTimeout)) + } + _ = c.ws.Close() + }) +} + +func (c *Connection) readLoop() { + closeCode := websocket.CloseAbnormalClosure + closeReason := "" + + c.ws.SetCloseHandler(func(code int, text string) error { + closeCode = code + closeReason = text + msg := websocket.FormatCloseMessage(code, "") + _ = c.ws.WriteControl(websocket.CloseMessage, msg, time.Now().Add(writeTimeout)) + return nil + }) + + for { + _, data, err := c.ws.ReadMessage() + if err != nil { + var closeErr *websocket.CloseError + if errors.As(err, &closeErr) { + closeCode = closeErr.Code + closeReason = closeErr.Text + } else if errors.Is(err, net.ErrClosed) { + // local Close() already recorded the intent + } + c.closed.Store(true) + c.closeHeartbeat() + _ = c.ws.Close() + c.fireClose(closeCode, closeReason) + return + } + c.handleFrame(data) + } +} + +func (c *Connection) handleFrame(data []byte) { + var env AttachmentMessage + if err := json.Unmarshal(data, &env); err != nil { + logging.Warn("WebSocket received invalid JSON, ignoring") + return + } + if env.Type == "" || len(env.Payload) == 0 { + logging.Warn("WebSocket received malformed AttachmentMessage, ignoring") + return + } + + if c.clusterSecret != "" { + if env.Signature == "" { + logging.Warn("WebSocket message missing signature, rejecting") + return + } + if !gossip.VerifyBody(c.clusterSecret, env.Payload, env.Signature) { + logging.Warn("WebSocket message signature invalid, rejecting") + return + } + } + + c.fireAttachment(&env) + + if isGossipType(env.Type) { + var wire gossip.SimpleMessage + if err := json.Unmarshal(env.Payload, &wire); err != nil { + c.fireError(fmt.Errorf("decode gossip payload: %w", err)) + return + } + c.fireMessage(wireToMessage(&wire)) + } +} + +func (c *Connection) fireMessage(msg *gossip.Message) { + c.handlersMu.RLock() + fn := c.onMessage + c.handlersMu.RUnlock() + if fn != nil { + fn(msg) + } +} + +func (c *Connection) fireAttachment(msg *AttachmentMessage) { + c.handlersMu.RLock() + fn := c.onAttachment + c.handlersMu.RUnlock() + if fn != nil { + fn(msg) + } +} + +func (c *Connection) fireClose(code int, reason string) { + c.handlersMu.RLock() + fn := c.onClose + c.handlersMu.RUnlock() + if fn != nil { + fn(code, reason) + } +} + +func (c *Connection) fireError(err error) { + c.handlersMu.RLock() + fn := c.onError + c.handlersMu.RUnlock() + if fn != nil { + fn(err) + } +} diff --git a/internal/transport/ws/connection_test.go b/internal/transport/ws/connection_test.go new file mode 100644 index 0000000..d4a4dd3 --- /dev/null +++ b/internal/transport/ws/connection_test.go @@ -0,0 +1,781 @@ +package ws + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + + "repram/internal/gossip" +) + +// ---- helpers ---------------------------------------------------------- + +func newPair(t *testing.T, clusterSecret string) (server *Connection, client *Connection, cleanup func()) { + t.Helper() + var serverConn *Connection + ready := make(chan struct{}) + + srv := httptest.NewServer(Handler(clusterSecret, nil, func(c *Connection) { + serverConn = c + close(ready) + })) + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/v1/ws" + u, err := url.Parse(wsURL) + if err != nil { + srv.Close() + t.Fatalf("parse ws url: %v", err) + } + dialer := &websocket.Dialer{HandshakeTimeout: 2 * time.Second} + clientWS, resp, err := dialer.Dial(u.String(), nil) + if err != nil { + if resp != nil { + _ = resp.Body.Close() + } + srv.Close() + t.Fatalf("dial: %v", err) + } + clientConn := NewConnection(clientWS, clusterSecret) + + select { + case <-ready: + case <-time.After(2 * time.Second): + clientConn.Close(websocket.CloseNormalClosure, "") + srv.Close() + t.Fatalf("server side never accepted") + } + + return serverConn, clientConn, func() { + clientConn.Close(websocket.CloseNormalClosure, "") + serverConn.Close(websocket.CloseNormalClosure, "") + srv.Close() + } +} + +func sampleMessage(overrides func(*gossip.Message)) *gossip.Message { + m := &gossip.Message{ + Type: gossip.MessageTypePut, + From: "node-a", + Key: "test-key", + Data: []byte("hello world"), + TTL: 300, + Timestamp: time.Unix(1735689600, 0), + MessageID: "test-msg-1", + } + if overrides != nil { + overrides(m) + } + return m +} + +// waitOnMessage blocks until ch receives or the timeout fires. +func waitOnMessage[T any](t *testing.T, ch <-chan T, timeout time.Duration) T { + t.Helper() + select { + case v := <-ch: + return v + case <-time.After(timeout): + t.Fatalf("timed out waiting for message") + var zero T + return zero + } +} + +// writeRaw bypasses the AttachmentMessage envelope to inject arbitrary bytes. +// White-box helper for testing rejection of malformed frames. +func (c *Connection) writeRaw(t *testing.T, data []byte) { + t.Helper() + c.writeMu.Lock() + defer c.writeMu.Unlock() + _ = c.ws.SetWriteDeadline(time.Now().Add(time.Second)) + if err := c.ws.WriteMessage(websocket.TextMessage, data); err != nil { + t.Fatalf("writeRaw: %v", err) + } +} + +// ---- gossip round-trip tests ----------------------------------------- + +func TestSendAndReceiveGossipMessage(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + got := make(chan *gossip.Message, 1) + server.OnMessage(func(m *gossip.Message) { got <- m }) + + if err := client.SendGossip(sampleMessage(nil)); err != nil { + t.Fatalf("send: %v", err) + } + m := waitOnMessage(t, got, 2*time.Second) + + if m.Type != gossip.MessageTypePut { + t.Errorf("type: got %q want PUT", m.Type) + } + if m.From != "node-a" { + t.Errorf("from: got %q", m.From) + } + if m.Key != "test-key" { + t.Errorf("key: got %q", m.Key) + } + if string(m.Data) != "hello world" { + t.Errorf("data: got %q", m.Data) + } + if m.TTL != 300 { + t.Errorf("ttl: got %d", m.TTL) + } + if m.MessageID != "test-msg-1" { + t.Errorf("message_id: got %q", m.MessageID) + } +} + +func TestAllMessageTypesRoundTrip(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + for _, mt := range []gossip.MessageType{ + gossip.MessageTypePut, + gossip.MessageTypeAck, + gossip.MessageTypePing, + gossip.MessageTypePong, + gossip.MessageTypeSync, + } { + t.Run(string(mt), func(t *testing.T) { + got := make(chan *gossip.Message, 1) + server.OnMessage(func(m *gossip.Message) { got <- m }) + + id := "msg-" + string(mt) + msg := sampleMessage(func(m *gossip.Message) { + m.Type = mt + m.MessageID = id + }) + if err := client.SendGossip(msg); err != nil { + t.Fatalf("send: %v", err) + } + r := waitOnMessage(t, got, 2*time.Second) + if r.Type != mt { + t.Errorf("type: got %q want %q", r.Type, mt) + } + if r.MessageID != id { + t.Errorf("id: got %q want %q", r.MessageID, id) + } + }) + } +} + +func TestPreservesBinaryData(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + binary := []byte{0x00, 0xff, 0x42, 0xde, 0xad, 0xbe, 0xef} + got := make(chan *gossip.Message, 1) + server.OnMessage(func(m *gossip.Message) { got <- m }) + + msg := sampleMessage(func(m *gossip.Message) { m.Data = binary }) + if err := client.SendGossip(msg); err != nil { + t.Fatalf("send: %v", err) + } + r := waitOnMessage(t, got, 2*time.Second) + if !bytes.Equal(r.Data, binary) { + t.Errorf("binary mismatch: got %x want %x", r.Data, binary) + } +} + +func TestPreservesNodeInfoInSync(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + got := make(chan *gossip.Message, 1) + server.OnMessage(func(m *gossip.Message) { got <- m }) + + msg := sampleMessage(func(m *gossip.Message) { + m.Type = gossip.MessageTypeSync + m.NodeInfo = &gossip.Node{ + ID: "node-a", + Address: "192.168.1.1", + Port: 9090, + HTTPPort: 8080, + Enclave: "acme-corp", + } + }) + if err := client.SendGossip(msg); err != nil { + t.Fatalf("send: %v", err) + } + r := waitOnMessage(t, got, 2*time.Second) + if r.NodeInfo == nil { + t.Fatal("NodeInfo missing") + } + if r.NodeInfo.ID != "node-a" { + t.Errorf("id: got %q", r.NodeInfo.ID) + } + if r.NodeInfo.Enclave != "acme-corp" { + t.Errorf("enclave: got %q", r.NodeInfo.Enclave) + } + if r.NodeInfo.HTTPPort != 8080 { + t.Errorf("http_port: got %d", r.NodeInfo.HTTPPort) + } +} + +// ---- wire format parity with HTTP gossip ------------------------------ + +func TestWirePayloadMatchesHTTPGossipFormat(t *testing.T) { + // The AttachmentMessage payload for a gossip frame must marshal to the + // same bytes that the HTTP gossip endpoint expects to receive. This is + // the invariant that lets handlers process WS and HTTP frames identically. + msg := sampleMessage(nil) + wire := messageToWire(msg) + wsPayload, err := json.Marshal(wire) + if err != nil { + t.Fatalf("marshal wire: %v", err) + } + + // Reconstruct what the HTTP transport would send for the same Message. + httpWire := &gossip.SimpleMessage{ + Type: string(msg.Type), + From: string(msg.From), + Key: msg.Key, + Data: msg.Data, + TTL: int32(msg.TTL), + Timestamp: msg.Timestamp.Unix(), + MessageID: msg.MessageID, + } + httpBytes, err := json.Marshal(httpWire) + if err != nil { + t.Fatalf("marshal http: %v", err) + } + if !bytes.Equal(wsPayload, httpBytes) { + t.Errorf("wire mismatch:\nws: %s\nhttp: %s", wsPayload, httpBytes) + } +} + +// ---- bidirectional ---------------------------------------------------- + +func TestBidirectionalGossip(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + srvGot := make(chan *gossip.Message, 1) + cliGot := make(chan *gossip.Message, 1) + server.OnMessage(func(m *gossip.Message) { srvGot <- m }) + client.OnMessage(func(m *gossip.Message) { cliGot <- m }) + + if err := client.SendGossip(sampleMessage(func(m *gossip.Message) { + m.From = "client" + m.MessageID = "c1" + })); err != nil { + t.Fatalf("client send: %v", err) + } + if err := server.SendGossip(sampleMessage(func(m *gossip.Message) { + m.From = "server" + m.MessageID = "s1" + })); err != nil { + t.Fatalf("server send: %v", err) + } + + if m := waitOnMessage(t, srvGot, 2*time.Second); m.From != "client" { + t.Errorf("server received from %q", m.From) + } + if m := waitOnMessage(t, cliGot, 2*time.Second); m.From != "server" { + t.Errorf("client received from %q", m.From) + } +} + +// ---- attachment messages (hello / welcome / goodbye) ----------------- + +func TestHelloRoundTrip(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + got := make(chan *AttachmentMessage, 1) + server.OnAttachment(func(m *AttachmentMessage) { got <- m }) + + hello := HelloPayload{ + NodeID: "mcp-node-1", + Enclave: "acme-corp", + Address: "192.168.1.50", + HTTPPort: 8080, + Capabilities: Capabilities{Inbound: "false"}, + } + if err := client.SendAttachment(AttachmentTypeHello, hello); err != nil { + t.Fatalf("send: %v", err) + } + r := waitOnMessage(t, got, 2*time.Second) + if r.Type != AttachmentTypeHello { + t.Errorf("type: got %q", r.Type) + } + var p HelloPayload + if err := json.Unmarshal(r.Payload, &p); err != nil { + t.Fatalf("payload decode: %v", err) + } + if p.NodeID != "mcp-node-1" { + t.Errorf("node_id: got %q", p.NodeID) + } + if p.Enclave != "acme-corp" { + t.Errorf("enclave: got %q", p.Enclave) + } + if p.Capabilities.Inbound != "false" { + t.Errorf("inbound: got %q", p.Capabilities.Inbound) + } +} + +func TestGoodbyeRoundTrip(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + got := make(chan *AttachmentMessage, 1) + client.OnAttachment(func(m *AttachmentMessage) { got <- m }) + + bye := GoodbyePayload{ + Reason: "shutdown", + AlternativeParents: []AlternativeParent{ + {ID: "cloud-1", Address: "10.0.0.1", HTTPPort: 8080, Enclave: "default"}, + {ID: "cloud-2", Address: "10.0.0.2", HTTPPort: 8080}, + }, + } + if err := server.SendAttachment(AttachmentTypeGoodbye, bye); err != nil { + t.Fatalf("send: %v", err) + } + r := waitOnMessage(t, got, 2*time.Second) + if r.Type != AttachmentTypeGoodbye { + t.Errorf("type: got %q", r.Type) + } + var p GoodbyePayload + if err := json.Unmarshal(r.Payload, &p); err != nil { + t.Fatalf("payload decode: %v", err) + } + if p.Reason != "shutdown" { + t.Errorf("reason: got %q", p.Reason) + } + if len(p.AlternativeParents) != 2 { + t.Errorf("alts: got %d", len(p.AlternativeParents)) + } + if p.AlternativeParents[0].ID != "cloud-1" { + t.Errorf("alt[0].id: got %q", p.AlternativeParents[0].ID) + } +} + +func TestWelcomeRoundTrip(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + got := make(chan *AttachmentMessage, 1) + client.OnAttachment(func(m *AttachmentMessage) { got <- m }) + + welcome := WelcomePayload{ + Topology: []gossip.SimpleMessage{ + {Type: "SYNC", From: "sub-1", MessageID: "topo-1", Timestamp: 1, NodeInfo: &gossip.SimpleNodeInfo{ + ID: "peer-x", Address: "10.0.0.5", Port: 9090, HTTPPort: 8080, Enclave: "default", + }}, + }, + YourPosition: WirePosition{Depth: 1, ParentID: "sub-1"}, + } + if err := server.SendAttachment(AttachmentTypeWelcome, welcome); err != nil { + t.Fatalf("send: %v", err) + } + r := waitOnMessage(t, got, 2*time.Second) + if r.Type != AttachmentTypeWelcome { + t.Errorf("type: got %q", r.Type) + } + var p WelcomePayload + if err := json.Unmarshal(r.Payload, &p); err != nil { + t.Fatalf("payload decode: %v", err) + } + if p.YourPosition.ParentID != "sub-1" || p.YourPosition.Depth != 1 { + t.Errorf("position: got %+v", p.YourPosition) + } + if len(p.Topology) != 1 || p.Topology[0].NodeInfo == nil || p.Topology[0].NodeInfo.ID != "peer-x" { + t.Errorf("topology decode mismatch: %+v", p.Topology) + } +} + +func TestHelloDoesNotFireOnMessage(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + var msgCalls atomic.Int32 + server.OnMessage(func(*gossip.Message) { msgCalls.Add(1) }) + + attachCh := make(chan struct{}, 1) + server.OnAttachment(func(*AttachmentMessage) { attachCh <- struct{}{} }) + + hello := HelloPayload{ + NodeID: "n1", Enclave: "default", Address: "127.0.0.1", + HTTPPort: 8080, Capabilities: Capabilities{Inbound: "false"}, + } + if err := client.SendAttachment(AttachmentTypeHello, hello); err != nil { + t.Fatalf("send: %v", err) + } + <-attachCh + time.Sleep(50 * time.Millisecond) + if n := msgCalls.Load(); n != 0 { + t.Errorf("OnMessage fired for lifecycle frame: %d times", n) + } +} + +// ---- close and post-close behavior ----------------------------------- + +func TestReportsClosedStateAfterClose(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + if client.IsClosed() { + t.Fatal("client closed before Close()") + } + + srvClosed := make(chan struct{}, 1) + server.OnClose(func(int, string) { srvClosed <- struct{}{} }) + + client.Close(websocket.CloseNormalClosure, "") + + select { + case <-srvClosed: + case <-time.After(2 * time.Second): + t.Fatal("server never observed close") + } + if !client.IsClosed() { + t.Error("client.IsClosed() still false") + } +} + +func TestSendAfterCloseIsSilent(t *testing.T) { + _, client, cleanup := newPair(t, "") + defer cleanup() + + client.Close(websocket.CloseNormalClosure, "") + + // Neither call should panic or return an error. + if err := client.SendGossip(sampleMessage(nil)); err != nil { + t.Errorf("SendGossip after close: %v", err) + } + if err := client.SendAttachment(AttachmentTypeHello, HelloPayload{ + NodeID: "n", Enclave: "d", Address: "x", HTTPPort: 0, + Capabilities: Capabilities{Inbound: "false"}, + }); err != nil { + t.Errorf("SendAttachment after close: %v", err) + } +} + +// ---- invalid frames --------------------------------------------------- + +func TestIgnoresInvalidJSON(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + var calls atomic.Int32 + server.OnAttachment(func(*AttachmentMessage) { calls.Add(1) }) + + client.writeRaw(t, []byte("not json {{{")) + time.Sleep(100 * time.Millisecond) + if n := calls.Load(); n != 0 { + t.Errorf("attachment fired for invalid JSON: %d times", n) + } +} + +func TestIgnoresMessagesMissingTypeOrPayload(t *testing.T) { + server, client, cleanup := newPair(t, "") + defer cleanup() + + var calls atomic.Int32 + server.OnAttachment(func(*AttachmentMessage) { calls.Add(1) }) + + client.writeRaw(t, []byte(`{"type":"put"}`)) // missing payload + client.writeRaw(t, []byte(`{"payload":{}}`)) // missing type + time.Sleep(100 * time.Millisecond) + if n := calls.Load(); n != 0 { + t.Errorf("attachment fired for malformed envelope: %d times", n) + } +} + +// ---- HMAC ------------------------------------------------------------- + +func TestHMACAcceptsValidSignature(t *testing.T) { + server, client, cleanup := newPair(t, "test-secret-42") + defer cleanup() + + got := make(chan *gossip.Message, 1) + server.OnMessage(func(m *gossip.Message) { got <- m }) + + if err := client.SendGossip(sampleMessage(nil)); err != nil { + t.Fatalf("send: %v", err) + } + m := waitOnMessage(t, got, 2*time.Second) + if m.Type != gossip.MessageTypePut || m.Key != "test-key" { + t.Errorf("unexpected message: %+v", m) + } +} + +func TestHMACRejectsTamperedSignature(t *testing.T) { + server, client, cleanup := newPair(t, "test-secret-42") + defer cleanup() + + var calls atomic.Int32 + server.OnMessage(func(*gossip.Message) { calls.Add(1) }) + + wire := messageToWire(sampleMessage(nil)) + wireBytes, _ := json.Marshal(wire) + env := AttachmentMessage{ + Type: AttachmentTypePut, + Signature: strings.Repeat("deadbeef", 8), + Payload: wireBytes, + } + frame, _ := json.Marshal(env) + client.writeRaw(t, frame) + + time.Sleep(100 * time.Millisecond) + if n := calls.Load(); n != 0 { + t.Errorf("OnMessage fired for tampered frame: %d times", n) + } +} + +func TestHMACRejectsMissingSignature(t *testing.T) { + server, client, cleanup := newPair(t, "test-secret-42") + defer cleanup() + + var calls atomic.Int32 + server.OnMessage(func(*gossip.Message) { calls.Add(1) }) + + wire := messageToWire(sampleMessage(nil)) + wireBytes, _ := json.Marshal(wire) + env := AttachmentMessage{Type: AttachmentTypePut, Payload: wireBytes} + frame, _ := json.Marshal(env) + client.writeRaw(t, frame) + + time.Sleep(100 * time.Millisecond) + if n := calls.Load(); n != 0 { + t.Errorf("OnMessage fired for unsigned frame: %d times", n) + } +} + +func TestNoSigningWithoutSecret(t *testing.T) { + // Capture the raw bytes that hit the server-side WS by registering an + // onMessage handler — the AttachmentMessage's Signature field tells us + // whether the sender added one. + server, client, cleanup := newPair(t, "") + defer cleanup() + + got := make(chan *AttachmentMessage, 1) + server.OnAttachment(func(m *AttachmentMessage) { got <- m }) + + if err := client.SendGossip(sampleMessage(nil)); err != nil { + t.Fatalf("send: %v", err) + } + r := waitOnMessage(t, got, 2*time.Second) + if r.Signature != "" { + t.Errorf("unexpected signature on no-secret frame: %q", r.Signature) + } +} + +// ---- heartbeat -------------------------------------------------------- + +func newPairWithHeartbeat(t *testing.T, period time.Duration) (*Connection, *Connection, func()) { + t.Helper() + server, client, cleanup := newPair(t, "") + server.heartbeatPeriod = period + client.heartbeatPeriod = period + return server, client, cleanup +} + +// pingCountingConn wraps a Connection and counts ping frames received by the +// underlying gorilla.*Conn — used to verify the heartbeat actually sends pings. +type pingCountingHook struct { + count atomic.Int32 +} + +func installPingHook(t *testing.T, c *Connection) *pingCountingHook { + t.Helper() + h := &pingCountingHook{} + c.ws.SetPingHandler(func(appData string) error { + h.count.Add(1) + // Respond with pong as gorilla's default ping handler would. + _ = c.ws.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(time.Second)) + return nil + }) + return h +} + +func TestHeartbeatSendsPings(t *testing.T) { + server, client, cleanup := newPairWithHeartbeat(t, 50*time.Millisecond) + defer cleanup() + + hook := installPingHook(t, server) + client.StartHeartbeat() + defer client.StopHeartbeat() + + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if hook.count.Load() >= 2 { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Errorf("expected >= 2 pings, got %d", hook.count.Load()) +} + +func TestHeartbeatTerminatesAfterMaxMissed(t *testing.T) { + server, client, cleanup := newPairWithHeartbeat(t, 30*time.Millisecond) + defer cleanup() + + // Silence the server's pong replies by replacing the ping handler with + // a no-op. The client will never receive pongs and should self-terminate + // after MaxMissedPongs intervals. + server.ws.SetPingHandler(func(string) error { return nil }) + + closed := make(chan struct{}, 1) + client.OnClose(func(int, string) { closed <- struct{}{} }) + + client.StartHeartbeat() + defer client.StopHeartbeat() + + // Allow MaxMissedPongs+1 ticks plus slack. + timeout := time.Duration(MaxMissedPongs+2) * 30 * time.Millisecond * 3 + select { + case <-closed: + case <-time.After(timeout): + t.Fatalf("heartbeat did not terminate after missed pongs (%v)", timeout) + } + if !client.IsClosed() { + t.Error("client.IsClosed() false after heartbeat termination") + } +} + +func TestHeartbeatResetsCounterOnPong(t *testing.T) { + server, client, cleanup := newPairWithHeartbeat(t, 30*time.Millisecond) + defer cleanup() + + // Server responds to pings as normal (default handler). The client + // should continue to receive pongs and never terminate. + installPingHook(t, server) + + closed := make(chan struct{}, 1) + client.OnClose(func(int, string) { closed <- struct{}{} }) + + client.StartHeartbeat() + defer client.StopHeartbeat() + + select { + case <-closed: + t.Fatal("client terminated even though pongs were arriving") + case <-time.After(time.Duration(MaxMissedPongs+2) * 30 * time.Millisecond): + // expected — connection remained healthy + } +} + +func TestStopHeartbeatOnClose(t *testing.T) { + _, client, cleanup := newPairWithHeartbeat(t, 20*time.Millisecond) + defer cleanup() + + client.StartHeartbeat() + client.Close(websocket.CloseNormalClosure, "") + // If the heartbeat goroutine kept running, it would panic on write to a + // closed conn or busy-loop. Sleep gives it time to react. + time.Sleep(80 * time.Millisecond) +} + +// ---- ConnectToSubstrate factory -------------------------------------- + +func TestConnectToSubstrateConnectsToV1WS(t *testing.T) { + accepted := make(chan struct{}, 1) + mux := http.NewServeMux() + mux.Handle("/v1/ws", Handler("", nil, func(*Connection) { accepted <- struct{}{} })) + srv := httptest.NewServer(mux) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + host, port := splitHostPort(t, u.Host) + + conn, err := ConnectToSubstrate(context.Background(), host, port, "", time.Second) + if err != nil { + t.Fatalf("connect: %v", err) + } + defer conn.Close(websocket.CloseNormalClosure, "") + + select { + case <-accepted: + case <-time.After(2 * time.Second): + t.Fatal("server never accepted /v1/ws") + } + if conn.IsClosed() { + t.Error("conn already closed") + } +} + +func TestConnectToSubstrateTimesOut(t *testing.T) { + // TCP listener that accepts but never speaks HTTP, so the WS handshake + // stalls until the dial timeout fires. + srv := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + // Block — but keep the request alive by sleeping past the dial timeout. + time.Sleep(time.Second) + })) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + host, port := splitHostPort(t, u.Host) + + _, err := ConnectToSubstrate(context.Background(), host, port, "", 100*time.Millisecond) + if err == nil { + t.Fatal("expected timeout error") + } +} + +func TestConnectToSubstrateExchangesGossip(t *testing.T) { + got := make(chan *gossip.Message, 1) + var serverConn *Connection + var mu sync.Mutex + + mux := http.NewServeMux() + mux.Handle("/v1/ws", Handler("", nil, func(c *Connection) { + mu.Lock() + serverConn = c + mu.Unlock() + c.OnMessage(func(m *gossip.Message) { got <- m }) + })) + srv := httptest.NewServer(mux) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + host, port := splitHostPort(t, u.Host) + + client, err := ConnectToSubstrate(context.Background(), host, port, "", time.Second) + if err != nil { + t.Fatalf("connect: %v", err) + } + defer client.Close(websocket.CloseNormalClosure, "") + + if err := client.SendGossip(sampleMessage(nil)); err != nil { + t.Fatalf("send: %v", err) + } + r := waitOnMessage(t, got, 2*time.Second) + if r.Type != gossip.MessageTypePut || r.Key != "test-key" { + t.Errorf("unexpected message: %+v", r) + } + mu.Lock() + if serverConn != nil { + serverConn.Close(websocket.CloseNormalClosure, "") + } + mu.Unlock() +} + +// ---- helpers ---------------------------------------------------------- + +func splitHostPort(t *testing.T, hostPort string) (string, int) { + t.Helper() + idx := strings.LastIndex(hostPort, ":") + if idx == -1 { + t.Fatalf("no port in %q", hostPort) + } + host := hostPort[:idx] + port, err := strconv.Atoi(hostPort[idx+1:]) + if err != nil { + t.Fatalf("parse port: %v", err) + } + return host, port +} diff --git a/internal/transport/ws/messages.go b/internal/transport/ws/messages.go new file mode 100644 index 0000000..a1f77fa --- /dev/null +++ b/internal/transport/ws/messages.go @@ -0,0 +1,174 @@ +// Package ws implements the WebSocket transport for substrate-transient +// attachments described in docs/internal/REPRAM-Discovery-Protocol-v2.md. +// +// Substrate nodes (REPRAM_INBOUND=true) accept inbound WS connections at +// /v1/ws. Transient nodes (REPRAM_INBOUND=false, the default) dial a +// substrate's /v1/ws after HTTP bootstrap and stay attached. Gossip +// frames sent over WS carry the same JSON wire format as HTTP gossip; +// a node processes a PUT identically regardless of the arrival transport. +// +// The AttachmentMessage envelope adds three lifecycle types (hello, +// welcome, goodbye) that HTTP gossip does not need. +package ws + +import ( + "encoding/json" + "time" + + "repram/internal/gossip" +) + +// AttachmentType identifies the kind of message carried on a substrate- +// transient WebSocket attachment. The gossip-shaped types (put/ack/ping/ +// pong/topology_sync) wrap the same payload as the HTTP gossip endpoint. +type AttachmentType string + +const ( + AttachmentTypePut AttachmentType = "put" + AttachmentTypeAck AttachmentType = "ack" + AttachmentTypePing AttachmentType = "ping" + AttachmentTypePong AttachmentType = "pong" + AttachmentTypeTopologySync AttachmentType = "topology_sync" + AttachmentTypeHello AttachmentType = "hello" + AttachmentTypeWelcome AttachmentType = "welcome" + AttachmentTypeGoodbye AttachmentType = "goodbye" +) + +// AttachmentMessage is the on-wire envelope. Payload is preserved as raw +// JSON bytes so the receiver verifies the HMAC against the exact bytes the +// sender signed — re-marshaling on either side risks key-ordering drift. +type AttachmentMessage struct { + Type AttachmentType `json:"type"` + Signature string `json:"signature,omitempty"` + Payload json.RawMessage `json:"payload"` +} + +// Capabilities declares whether the announcing node accepts inbound WS +// attachments. Maps directly to the REPRAM_INBOUND env var. +type Capabilities struct { + Inbound string `json:"inbound"` // "true" | "false" +} + +// HelloPayload announces a transient attaching to a substrate. +type HelloPayload struct { + NodeID string `json:"node_id"` + Enclave string `json:"enclave"` + Address string `json:"address"` + HTTPPort int `json:"http_port"` + Capabilities Capabilities `json:"capabilities"` +} + +// WirePosition describes the transient's location in the substrate tree. +type WirePosition struct { + Depth int `json:"depth"` + ParentID string `json:"parent_id"` +} + +// WelcomePayload is the substrate's reply to hello. Topology entries are +// SimpleMessage-shaped SYNC announcements so the transient can seed its +// peer map without a separate format. +type WelcomePayload struct { + Topology []gossip.SimpleMessage `json:"topology"` + YourPosition WirePosition `json:"your_position"` +} + +// AlternativeParent is a fallback substrate the transient can attach to +// when its current parent goes away. +type AlternativeParent struct { + ID string `json:"id"` + Address string `json:"address"` + HTTPPort int `json:"http_port"` + Enclave string `json:"enclave,omitempty"` +} + +// GoodbyePayload is sent by a substrate before shutting down to keep +// transient reattachment latency in the 3-5s range instead of waiting +// the full heartbeat timeout. +type GoodbyePayload struct { + Reason string `json:"reason"` + AlternativeParents []AlternativeParent `json:"alternative_parents"` +} + +// gossipTypeFor maps an internal MessageType to its on-wire AttachmentType. +func gossipTypeFor(t gossip.MessageType) AttachmentType { + switch t { + case gossip.MessageTypePut: + return AttachmentTypePut + case gossip.MessageTypeAck: + return AttachmentTypeAck + case gossip.MessageTypePing: + return AttachmentTypePing + case gossip.MessageTypePong: + return AttachmentTypePong + case gossip.MessageTypeSync: + return AttachmentTypeTopologySync + default: + return AttachmentTypePut + } +} + +// isGossipType reports whether the AttachmentType carries a gossip Message +// payload (versus a lifecycle hello/welcome/goodbye). +func isGossipType(t AttachmentType) bool { + switch t { + case AttachmentTypePut, AttachmentTypeAck, + AttachmentTypePing, AttachmentTypePong, + AttachmentTypeTopologySync: + return true + } + return false +} + +// messageToWire converts a gossip Message to its SimpleMessage wire form. +// encoding/json base64-encodes []byte automatically — matches the TS +// reference's explicit base64 encoding of the data field. +func messageToWire(msg *gossip.Message) *gossip.SimpleMessage { + wire := &gossip.SimpleMessage{ + Type: string(msg.Type), + From: string(msg.From), + To: string(msg.To), + Key: msg.Key, + Data: msg.Data, + TTL: int32(msg.TTL), + Timestamp: msg.Timestamp.Unix(), + MessageID: msg.MessageID, + } + if msg.NodeInfo != nil { + wire.NodeInfo = &gossip.SimpleNodeInfo{ + ID: string(msg.NodeInfo.ID), + Address: msg.NodeInfo.Address, + Port: msg.NodeInfo.Port, + HTTPPort: msg.NodeInfo.HTTPPort, + Enclave: msg.NodeInfo.Enclave, + } + } + return wire +} + +// wireToMessage is the inverse of messageToWire. +func wireToMessage(wire *gossip.SimpleMessage) *gossip.Message { + msg := &gossip.Message{ + Type: gossip.MessageType(wire.Type), + From: gossip.NodeID(wire.From), + To: gossip.NodeID(wire.To), + Key: wire.Key, + Data: wire.Data, + TTL: int(wire.TTL), + Timestamp: time.Unix(wire.Timestamp, 0), + MessageID: wire.MessageID, + } + if wire.NodeInfo != nil { + enclave := wire.NodeInfo.Enclave + if enclave == "" { + enclave = "default" + } + msg.NodeInfo = &gossip.Node{ + ID: gossip.NodeID(wire.NodeInfo.ID), + Address: wire.NodeInfo.Address, + Port: wire.NodeInfo.Port, + HTTPPort: wire.NodeInfo.HTTPPort, + Enclave: enclave, + } + } + return msg +} diff --git a/internal/transport/ws/server.go b/internal/transport/ws/server.go new file mode 100644 index 0000000..a9b52e5 --- /dev/null +++ b/internal/transport/ws/server.go @@ -0,0 +1,38 @@ +package ws + +import ( + "net/http" + + "github.com/gorilla/websocket" + + "repram/internal/logging" +) + +// Handler returns the HTTP handler for the /v1/ws endpoint. On successful +// upgrade, onAccept is invoked with the wrapped Connection; the handler then +// returns, and the connection's read loop drives lifecycle from there. +// +// allowOrigin matches gorilla's Upgrader.CheckOrigin contract — pass nil to +// accept any origin (REPRAM's CORS policy is intentionally permissive; see #38). +func Handler(clusterSecret string, allowOrigin func(*http.Request) bool, onAccept func(*Connection)) http.HandlerFunc { + upgrader := &websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { + if allowOrigin != nil { + return allowOrigin(r) + } + return true + }, + } + return func(w http.ResponseWriter, r *http.Request) { + ws, err := upgrader.Upgrade(w, r, nil) + if err != nil { + // Upgrader.Upgrade has already written the HTTP error response. + logging.Warn("WebSocket upgrade failed: %v", err) + return + } + conn := NewConnection(ws, clusterSecret) + if onAccept != nil { + onAccept(conn) + } + } +} From ca7e1cebe7b08a04466c0b1d287f2316f9a6aa01 Mon Sep 17 00:00:00 2001 From: TickTockBent Date: Tue, 12 May 2026 08:36:44 -0400 Subject: [PATCH 2/5] phase 2 (#135): tree manager attach/detach lifecycle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports the lifecycle half of repram-mcp/src/node/tree.ts to Go in internal/tree/. Adds the substrate/transient role model on top of the WS transport from phase 1. Substrate side: HandleHello validates capacity and registers the transient as a child, sends welcome with the current peer topology snapshot, and installs a close handler that auto-removes the child on disconnect. Capacity rejections and MaxChildren=0 both yield a goodbye-with-alternatives followed by a short delayed close. Transient side: Attach drives the hello/welcome handshake with the substrate, populates lastKnownAlts from welcome.topology (excluding self), and installs long-lived goodbye + close handlers with identity guards that prevent a stale handler from clobbering an active reattach. Reattach: three-layer loop runs as a single-flight goroutine — goodbye-supplied alts → cached welcome topology (per-attempt 5s, total 30s deadline) → seed list (per-attempt 10s) — with exponential backoff between full cycles capped at 60s. Self-skip is enforced inside tryAlternatives by address+http_port literal match (regression for #120). Stop() wakes the backoff sleep via stopCh so shutdown is prompt. Multi-subscriber refactor: ws.Connection now offers AddAttachmentHandler / AddCloseHandler returning a remove function. The TS reference uses EventEmitter add/remove semantics for the attach handshake; the prior single-slot setter API couldn't express the temporary-welcome-listener pattern. OnMessage / OnError remain single-slot since there's one application-level consumer for those. 17 tree tests at parity with the attach/detach portion of tree.test.ts (role detection, accept + welcome, child registration + auto-removal on close, capacity rejection, MaxChildren=0, attach handshake + timeout, alternatives ordering, goodbye to children, parent clear after substrate goodbye, topology cache with self exclusion, parseSeedAddress edge cases, stale-connection close guard, #120 self-skip timing assertion, ungraceful close → seed fallback, Stop() unblocks sleep). Race-clean. Relay forwarding (substrate → enclave peers) and ACK reverse-routing land in phases 3-4. // ticktockbent --- internal/transport/ws/connection.go | 90 ++- internal/transport/ws/connection_test.go | 20 +- internal/tree/manager.go | 698 +++++++++++++++++++++ internal/tree/manager_test.go | 765 +++++++++++++++++++++++ 4 files changed, 1537 insertions(+), 36 deletions(-) create mode 100644 internal/tree/manager.go create mode 100644 internal/tree/manager_test.go diff --git a/internal/transport/ws/connection.go b/internal/transport/ws/connection.go index e5feef9..daa3f2d 100644 --- a/internal/transport/ws/connection.go +++ b/internal/transport/ws/connection.go @@ -48,11 +48,12 @@ type Connection struct { heartbeatStarted atomic.Bool heartbeatPeriod time.Duration - handlersMu sync.RWMutex - onMessage func(*gossip.Message) - onAttachment func(*AttachmentMessage) - onClose func(code int, reason string) - onError func(error) + handlersMu sync.RWMutex + onMessage func(*gossip.Message) + onError func(error) + attachmentHandlers []handlerEntry[func(*AttachmentMessage)] + closeHandlers []handlerEntry[func(int, string)] + handlerSeq atomic.Uint64 remoteMu sync.RWMutex remoteNodeID string @@ -78,37 +79,74 @@ func NewConnection(ws *websocket.Conn, clusterSecret string) *Connection { return c } -// OnMessage registers a handler for gossip-typed AttachmentMessages, decoded -// back into the internal Message form. +// handlerEntry wraps a multi-subscriber lifecycle callback with a stable ID +// so a registration can be removed later (matches the EventEmitter +// addListener/removeListener pattern in the TS reference). +type handlerEntry[T any] struct { + id uint64 + fn T +} + +// OnMessage installs the single application-level gossip-message handler. +// Subsequent calls replace the previous handler. Lifecycle events use the +// multi-subscriber AddOn* methods instead. func (c *Connection) OnMessage(fn func(*gossip.Message)) { c.handlersMu.Lock() c.onMessage = fn c.handlersMu.Unlock() } -// OnAttachment registers a handler that fires for every parsed AttachmentMessage -// (gossip + hello/welcome/goodbye). Gossip frames fire both OnAttachment and -// OnMessage; lifecycle frames fire only OnAttachment. -func (c *Connection) OnAttachment(fn func(*AttachmentMessage)) { +// OnError registers the read-loop error handler. Subsequent calls replace. +// Fatal errors surface through close handlers instead. +func (c *Connection) OnError(fn func(error)) { c.handlersMu.Lock() - c.onAttachment = fn + c.onError = fn c.handlersMu.Unlock() } -// OnClose fires once when the underlying connection closes for any reason -// (graceful close, peer hangup, RST, or local Terminate). -func (c *Connection) OnClose(fn func(code int, reason string)) { +// AddAttachmentHandler subscribes fn to every parsed AttachmentMessage +// (gossip + hello/welcome/goodbye). The returned function removes this +// subscription; multiple subscribers fire in registration order. +// +// The tree-manager attach handshake uses this to register a temporary +// welcome/goodbye-waiting handler that it removes once welcome arrives, then +// registers a long-lived goodbye-handling subscription. +func (c *Connection) AddAttachmentHandler(fn func(*AttachmentMessage)) (remove func()) { + id := c.handlerSeq.Add(1) c.handlersMu.Lock() - c.onClose = fn + c.attachmentHandlers = append(c.attachmentHandlers, handlerEntry[func(*AttachmentMessage)]{id: id, fn: fn}) c.handlersMu.Unlock() + return func() { + c.handlersMu.Lock() + defer c.handlersMu.Unlock() + for i, h := range c.attachmentHandlers { + if h.id == id { + c.attachmentHandlers = append(c.attachmentHandlers[:i], c.attachmentHandlers[i+1:]...) + return + } + } + } } -// OnError fires for non-fatal read-loop errors that did not terminate the -// connection. Fatal errors surface through OnClose instead. -func (c *Connection) OnError(fn func(error)) { +// AddCloseHandler subscribes fn to the one-shot close event. Multiple +// subscribers fire in registration order. The returned function removes +// the subscription (e.g., when a temporary handler should not run if +// close happens later). +func (c *Connection) AddCloseHandler(fn func(code int, reason string)) (remove func()) { + id := c.handlerSeq.Add(1) c.handlersMu.Lock() - c.onError = fn + c.closeHandlers = append(c.closeHandlers, handlerEntry[func(int, string)]{id: id, fn: fn}) c.handlersMu.Unlock() + return func() { + c.handlersMu.Lock() + defer c.handlersMu.Unlock() + for i, h := range c.closeHandlers { + if h.id == id { + c.closeHandlers = append(c.closeHandlers[:i], c.closeHandlers[i+1:]...) + return + } + } + } } // RemoteNodeID returns the peer's node ID once the hello/welcome handshake @@ -337,19 +375,19 @@ func (c *Connection) fireMessage(msg *gossip.Message) { func (c *Connection) fireAttachment(msg *AttachmentMessage) { c.handlersMu.RLock() - fn := c.onAttachment + handlers := append([]handlerEntry[func(*AttachmentMessage)](nil), c.attachmentHandlers...) c.handlersMu.RUnlock() - if fn != nil { - fn(msg) + for _, h := range handlers { + h.fn(msg) } } func (c *Connection) fireClose(code int, reason string) { c.handlersMu.RLock() - fn := c.onClose + handlers := append([]handlerEntry[func(int, string)](nil), c.closeHandlers...) c.handlersMu.RUnlock() - if fn != nil { - fn(code, reason) + for _, h := range handlers { + h.fn(code, reason) } } diff --git a/internal/transport/ws/connection_test.go b/internal/transport/ws/connection_test.go index d4a4dd3..7cfd1e9 100644 --- a/internal/transport/ws/connection_test.go +++ b/internal/transport/ws/connection_test.go @@ -296,7 +296,7 @@ func TestHelloRoundTrip(t *testing.T) { defer cleanup() got := make(chan *AttachmentMessage, 1) - server.OnAttachment(func(m *AttachmentMessage) { got <- m }) + server.AddAttachmentHandler(func(m *AttachmentMessage) { got <- m }) hello := HelloPayload{ NodeID: "mcp-node-1", @@ -332,7 +332,7 @@ func TestGoodbyeRoundTrip(t *testing.T) { defer cleanup() got := make(chan *AttachmentMessage, 1) - client.OnAttachment(func(m *AttachmentMessage) { got <- m }) + client.AddAttachmentHandler(func(m *AttachmentMessage) { got <- m }) bye := GoodbyePayload{ Reason: "shutdown", @@ -368,7 +368,7 @@ func TestWelcomeRoundTrip(t *testing.T) { defer cleanup() got := make(chan *AttachmentMessage, 1) - client.OnAttachment(func(m *AttachmentMessage) { got <- m }) + client.AddAttachmentHandler(func(m *AttachmentMessage) { got <- m }) welcome := WelcomePayload{ Topology: []gossip.SimpleMessage{ @@ -405,7 +405,7 @@ func TestHelloDoesNotFireOnMessage(t *testing.T) { server.OnMessage(func(*gossip.Message) { msgCalls.Add(1) }) attachCh := make(chan struct{}, 1) - server.OnAttachment(func(*AttachmentMessage) { attachCh <- struct{}{} }) + server.AddAttachmentHandler(func(*AttachmentMessage) { attachCh <- struct{}{} }) hello := HelloPayload{ NodeID: "n1", Enclave: "default", Address: "127.0.0.1", @@ -432,7 +432,7 @@ func TestReportsClosedStateAfterClose(t *testing.T) { } srvClosed := make(chan struct{}, 1) - server.OnClose(func(int, string) { srvClosed <- struct{}{} }) + server.AddCloseHandler(func(int, string) { srvClosed <- struct{}{} }) client.Close(websocket.CloseNormalClosure, "") @@ -471,7 +471,7 @@ func TestIgnoresInvalidJSON(t *testing.T) { defer cleanup() var calls atomic.Int32 - server.OnAttachment(func(*AttachmentMessage) { calls.Add(1) }) + server.AddAttachmentHandler(func(*AttachmentMessage) { calls.Add(1) }) client.writeRaw(t, []byte("not json {{{")) time.Sleep(100 * time.Millisecond) @@ -485,7 +485,7 @@ func TestIgnoresMessagesMissingTypeOrPayload(t *testing.T) { defer cleanup() var calls atomic.Int32 - server.OnAttachment(func(*AttachmentMessage) { calls.Add(1) }) + server.AddAttachmentHandler(func(*AttachmentMessage) { calls.Add(1) }) client.writeRaw(t, []byte(`{"type":"put"}`)) // missing payload client.writeRaw(t, []byte(`{"payload":{}}`)) // missing type @@ -563,7 +563,7 @@ func TestNoSigningWithoutSecret(t *testing.T) { defer cleanup() got := make(chan *AttachmentMessage, 1) - server.OnAttachment(func(m *AttachmentMessage) { got <- m }) + server.AddAttachmentHandler(func(m *AttachmentMessage) { got <- m }) if err := client.SendGossip(sampleMessage(nil)); err != nil { t.Fatalf("send: %v", err) @@ -630,7 +630,7 @@ func TestHeartbeatTerminatesAfterMaxMissed(t *testing.T) { server.ws.SetPingHandler(func(string) error { return nil }) closed := make(chan struct{}, 1) - client.OnClose(func(int, string) { closed <- struct{}{} }) + client.AddCloseHandler(func(int, string) { closed <- struct{}{} }) client.StartHeartbeat() defer client.StopHeartbeat() @@ -656,7 +656,7 @@ func TestHeartbeatResetsCounterOnPong(t *testing.T) { installPingHook(t, server) closed := make(chan struct{}, 1) - client.OnClose(func(int, string) { closed <- struct{}{} }) + client.AddCloseHandler(func(int, string) { closed <- struct{}{} }) client.StartHeartbeat() defer client.StopHeartbeat() diff --git a/internal/tree/manager.go b/internal/tree/manager.go new file mode 100644 index 0000000..b75d3c3 --- /dev/null +++ b/internal/tree/manager.go @@ -0,0 +1,698 @@ +// Package tree manages substrate-transient attachment lifecycle for +// Discovery Protocol v2. +// +// Substrate nodes (REPRAM_INBOUND=true) accept inbound WS attachments, register +// transients as children, and send goodbye-with-alternatives during shutdown. +// +// Transient nodes (REPRAM_INBOUND=false) dial a substrate's /v1/ws after HTTP +// bootstrap, send hello, parse welcome, and cache the substrate's topology as +// fallback candidates. When the parent connection drops, a three-layer reattach +// loop tries goodbye-supplied alternatives → cached topology → seed list, with +// exponential backoff between full cycles. +// +// Relay forwarding (substrate fans out child PUTs to enclave peers) and ACK +// reverse-routing land in subsequent phases. +package tree + +import ( + "context" + "encoding/json" + "errors" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + "repram/internal/gossip" + "repram/internal/logging" + "repram/internal/transport/ws" +) + +const ( + // DefaultMaxChildren is the default cap on transient attachments per + // substrate. 0 disables inbound attachments entirely. + DefaultMaxChildren = 100 + + // ReattachBackoffInitial is the wait before the first reattach retry. + ReattachBackoffInitial = 5 * time.Second + + // ReattachBackoffMax caps the exponential backoff between full cycles. + ReattachBackoffMax = 60 * time.Second + + // CachedAltConnectTimeout is the per-attempt timeout for cached + // welcome-topology alternatives — short because entries may be stale. + CachedAltConnectTimeout = 5 * time.Second + + // SeedConnectTimeout is the per-attempt timeout for fresh seed-list + // alternatives (and goodbye-supplied alternatives). + SeedConnectTimeout = 10 * time.Second + + // CachedLayerDeadline caps the total time spent in the cached-topology + // layer before falling through to the seed list (#108). + CachedLayerDeadline = 30 * time.Second + + // AttachTimeout is the default timeout for the hello/welcome handshake. + AttachTimeout = 10 * time.Second + + // redirectCloseDelay gives a rejected hello time to receive the goodbye + // frame before the substrate hangs up the socket. + redirectCloseDelay = 500 * time.Millisecond +) + +// Role identifies the resolved node role. +type Role string + +const ( + RoleSubstrate Role = "substrate" + RoleTransient Role = "transient" +) + +// InboundCapability mirrors the REPRAM_INBOUND env var: "true" makes the node +// a substrate, "false" (default) makes it a transient. +type InboundCapability string + +const ( + InboundTrue InboundCapability = "true" + InboundFalse InboundCapability = "false" +) + +// Peerer is the subset of gossip.Protocol that Manager depends on. +// *gossip.Protocol satisfies it; tests use a fake implementation. +type Peerer interface { + GetPeers() []*gossip.Node +} + +// Dialer opens a new outbound WS attachment. Defaults to ws.ConnectToSubstrate. +// Injected to keep reattach tests fast (no need to spin up live servers for +// each cached/seed alternative). +type Dialer func(ctx context.Context, address string, port int, clusterSecret string, timeout time.Duration) (*ws.Connection, error) + +// Options configures Manager construction. +type Options struct { + Inbound InboundCapability + MaxChildren int // 0 disables inbound attachments + ClusterSecret string // empty → open-cluster mode (no HMAC) + Dialer Dialer // nil → ws.ConnectToSubstrate +} + +// Manager tracks the substrate/transient tree state for one node. +// +// All public methods are safe for concurrent use. Children are tracked by +// node ID; close-of-child cleans up automatically through an AddCloseHandler +// installed during HandleHello. +type Manager struct { + local *gossip.Node + gossip Peerer + opts Options + dialer Dialer + + role Role + inboundCapable bool + + mu sync.Mutex + parent *ws.Connection + children map[string]*ws.Connection + lastKnownAlts []ws.AlternativeParent + reattachInFlight bool + onReattach func(*ws.Connection) + seedProvider func() []string + + stopOnce sync.Once + stopping atomic.Bool + stopCh chan struct{} +} + +// NewManager constructs a Manager. The Inbound option resolves the role +// immediately; subsequent calls to HandleHello / Attach honor that role. +func NewManager(local *gossip.Node, peers Peerer, opts Options) *Manager { + if opts.Dialer == nil { + opts.Dialer = ws.ConnectToSubstrate + } + m := &Manager{ + local: local, + gossip: peers, + opts: opts, + dialer: opts.Dialer, + children: make(map[string]*ws.Connection), + stopCh: make(chan struct{}), + } + if opts.Inbound == InboundTrue { + m.role = RoleSubstrate + m.inboundCapable = true + } else { + m.role = RoleTransient + m.inboundCapable = false + } + return m +} + +// Role reports the node's resolved tree role. +func (m *Manager) Role() Role { return m.role } + +// IsInboundCapable reports whether this node accepts inbound attachments. +func (m *Manager) IsInboundCapable() bool { return m.inboundCapable } + +// Parent returns the active outbound substrate attachment, or nil. +func (m *Manager) Parent() *ws.Connection { + m.mu.Lock() + defer m.mu.Unlock() + return m.parent +} + +// ChildCount returns the number of attached transients. +func (m *Manager) ChildCount() int { + m.mu.Lock() + defer m.mu.Unlock() + return len(m.children) +} + +// HasChild reports whether nodeID is currently attached. +func (m *Manager) HasChild(nodeID string) bool { + m.mu.Lock() + defer m.mu.Unlock() + _, ok := m.children[nodeID] + return ok +} + +// Children returns a snapshot of the attached-children map. The returned map +// is a copy — callers may inspect freely without holding any lock. +func (m *Manager) Children() map[string]*ws.Connection { + m.mu.Lock() + defer m.mu.Unlock() + out := make(map[string]*ws.Connection, len(m.children)) + for k, v := range m.children { + out[k] = v + } + return out +} + +// SetSeedProvider wires the freshest reattach fallback. For private clusters +// snapshot REPRAM_PEERS; for public, close over the omega refresher's current +// signed root list. Nil disables the seed-list layer. +func (m *Manager) SetSeedProvider(fn func() []string) { + m.mu.Lock() + m.seedProvider = fn + m.mu.Unlock() +} + +// SetReattachCallback registers a hook fired after a successful reattach. +// The application uses this to rewire its parent-message router to the new +// connection. +func (m *Manager) SetReattachCallback(fn func(*ws.Connection)) { + m.mu.Lock() + m.onReattach = fn + m.mu.Unlock() +} + +// HandleHello processes an incoming hello frame on a freshly-accepted child +// connection. On accept it sends a welcome with the current peer topology and +// registers a close handler that removes the child from the map. On reject +// (capacity, attachments disabled) it sends a goodbye-with-alternatives and +// closes the connection shortly afterward. +func (m *Manager) HandleHello(conn *ws.Connection, hello *ws.HelloPayload) bool { + m.mu.Lock() + if m.opts.MaxChildren == 0 { + m.mu.Unlock() + logging.Info("Rejecting attachment from %s: attachments disabled (MaxChildren=0)", hello.NodeID) + m.sendRedirect(conn, hello.Enclave, "at capacity") + return false + } + if m.opts.MaxChildren > 0 && len(m.children) >= m.opts.MaxChildren { + count := len(m.children) + m.mu.Unlock() + logging.Info("Rejecting attachment from %s: at capacity (%d/%d)", hello.NodeID, count, m.opts.MaxChildren) + m.sendRedirect(conn, hello.Enclave, "at capacity") + return false + } + m.children[hello.NodeID] = conn + count := len(m.children) + m.mu.Unlock() + + conn.SetRemote(hello.NodeID, hello.Enclave) + + peers := m.gossip.GetPeers() + topology := buildWelcomeTopology(peers, m.local) + welcome := ws.WelcomePayload{ + Topology: topology, + YourPosition: ws.WirePosition{Depth: 1, ParentID: string(m.local.ID)}, + } + if err := conn.SendAttachment(ws.AttachmentTypeWelcome, welcome); err != nil { + logging.Warn("Send welcome to %s failed: %v", hello.NodeID, err) + } + + nodeID := hello.NodeID + conn.AddCloseHandler(func(int, string) { + m.mu.Lock() + if cur, ok := m.children[nodeID]; ok && cur == conn { + delete(m.children, nodeID) + } + remaining := len(m.children) + m.mu.Unlock() + logging.Info("Transient %s detached (%d remaining)", nodeID, remaining) + }) + + logging.Info("Transient %s attached (enclave: %s, children: %d)", hello.NodeID, hello.Enclave, count) + return true +} + +func (m *Manager) sendRedirect(conn *ws.Connection, requestedEnclave, reason string) { + alts := m.getAlternativeSubstrates(requestedEnclave) + goodbye := ws.GoodbyePayload{Reason: reason, AlternativeParents: alts} + if err := conn.SendAttachment(ws.AttachmentTypeGoodbye, goodbye); err != nil { + logging.Debug("Send redirect goodbye failed: %v", err) + } + go func() { + time.Sleep(redirectCloseDelay) + if !conn.IsClosed() { + conn.Close(1000, "redirected") + } + }() +} + +// Attach completes the hello/welcome handshake from the transient side over +// conn. On success the connection becomes the active parent; on timeout or +// goodbye-during-handshake it returns an error and conn is left untouched +// (the caller decides whether to close it). +// +// After Attach returns successfully it installs long-lived goodbye and +// close handlers that fire reattach when the parent goes away. +func (m *Manager) Attach(ctx context.Context, conn *ws.Connection) (*ws.WelcomePayload, error) { + return m.attach(ctx, conn, AttachTimeout) +} + +func (m *Manager) attach(ctx context.Context, conn *ws.Connection, timeout time.Duration) (*ws.WelcomePayload, error) { + hello := ws.HelloPayload{ + NodeID: string(m.local.ID), + Enclave: m.local.Enclave, + Address: m.local.Address, + HTTPPort: m.local.HTTPPort, + Capabilities: ws.Capabilities{Inbound: string(m.opts.Inbound)}, + } + if err := conn.SendAttachment(ws.AttachmentTypeHello, hello); err != nil { + return nil, err + } + + // Race a temporary attachment handler against the close event and the + // timeout. The handler is removed in every exit path so it can't fire + // after the long-lived handlers below take over. + welcomeCh := make(chan *ws.WelcomePayload, 1) + rejectedCh := make(chan struct{}, 1) + removeAttach := conn.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + switch msg.Type { + case ws.AttachmentTypeWelcome: + var w ws.WelcomePayload + if err := json.Unmarshal(msg.Payload, &w); err != nil { + logging.Warn("Welcome payload decode failed: %v", err) + rejectedCh <- struct{}{} + return + } + welcomeCh <- &w + case ws.AttachmentTypeGoodbye: + rejectedCh <- struct{}{} + } + }) + removeClose := conn.AddCloseHandler(func(int, string) { + rejectedCh <- struct{}{} + }) + + var welcome *ws.WelcomePayload + select { + case welcome = <-welcomeCh: + case <-rejectedCh: + removeAttach() + removeClose() + return nil, errors.New("substrate attachment failed (goodbye or close)") + case <-time.After(timeout): + removeAttach() + removeClose() + return nil, errors.New("substrate attachment timed out") + case <-ctx.Done(): + removeAttach() + removeClose() + return nil, ctx.Err() + } + removeAttach() + removeClose() + + // Promote to active parent. The role is "transient" by definition when + // we reach here — inbound-capable nodes never call Attach. + m.mu.Lock() + m.parent = conn + m.role = RoleTransient + m.lastKnownAlts = buildAltsFromTopology(welcome.Topology, string(m.local.ID)) + m.mu.Unlock() + conn.SetRemote(welcome.YourPosition.ParentID, m.local.Enclave) + + // Install long-lived handlers. The identity guards prevent a stale + // handler (from a connection that was replaced by a successful + // reattach) from clobbering the active parent. + conn.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + if msg.Type != ws.AttachmentTypeGoodbye { + return + } + m.mu.Lock() + stale := m.parent != conn + m.mu.Unlock() + if stale { + return + } + var p ws.GoodbyePayload + if err := json.Unmarshal(msg.Payload, &p); err == nil { + logging.Info("Substrate sent goodbye: %s (%d alternatives)", p.Reason, len(p.AlternativeParents)) + m.mu.Lock() + m.parent = nil + m.mu.Unlock() + go m.triggerReattach(p.AlternativeParents) + } + }) + conn.AddCloseHandler(func(int, string) { + m.mu.Lock() + stale := m.parent != conn + m.mu.Unlock() + if stale { + return + } + logging.Warn("Substrate attachment to %s lost", conn.RemoteNodeID()) + m.mu.Lock() + m.parent = nil + m.mu.Unlock() + go m.triggerReattach(nil) + }) + + logging.Info("Attached to substrate %s (depth %d, topology %d nodes)", + welcome.YourPosition.ParentID, welcome.YourPosition.Depth, len(welcome.Topology)) + return welcome, nil +} + +// triggerReattach is the single-flight entry point for the reattach loop. +// Concurrent invocations (goodbye and close firing back-to-back) collapse +// to one running loop. +func (m *Manager) triggerReattach(supplied []ws.AlternativeParent) { + if m.stopping.Load() { + return + } + m.mu.Lock() + if m.reattachInFlight { + m.mu.Unlock() + return + } + m.reattachInFlight = true + m.mu.Unlock() + + defer func() { + m.mu.Lock() + m.reattachInFlight = false + m.mu.Unlock() + }() + + m.reattachLoop(supplied) +} + +func (m *Manager) reattachLoop(supplied []ws.AlternativeParent) { + backoff := ReattachBackoffInitial + for !m.stopping.Load() { + // Layer 1: goodbye-supplied alternatives — single-shot. + if len(supplied) > 0 { + if m.tryAlternatives(supplied, SeedConnectTimeout, time.Time{}) { + return + } + supplied = nil + } + + // Layer 2: cached welcome topology. + m.mu.Lock() + cached := append([]ws.AlternativeParent(nil), m.lastKnownAlts...) + m.mu.Unlock() + if len(cached) > 0 { + deadline := time.Now().Add(CachedLayerDeadline) + if m.tryAlternatives(cached, CachedAltConnectTimeout, deadline) { + return + } + } + + // Layer 3: seed list. + m.mu.Lock() + provider := m.seedProvider + m.mu.Unlock() + var seeds []string + if provider != nil { + seeds = provider() + } + seedAlts := make([]ws.AlternativeParent, 0, len(seeds)) + for _, s := range seeds { + if alt, ok := parseSeedAddress(s); ok { + seedAlts = append(seedAlts, alt) + } + } + if len(seedAlts) > 0 { + if m.tryAlternatives(seedAlts, SeedConnectTimeout, time.Time{}) { + return + } + } + + logging.Warn("All reattach paths failed — sleeping %v before retry (local store still serves reads)", backoff) + if !m.sleep(backoff) { + return + } + backoff *= 2 + if backoff > ReattachBackoffMax { + backoff = ReattachBackoffMax + } + } +} + +// tryAlternatives walks alts in order, attempting attach to each. Returns +// true on the first success. Honors stopping and an optional layer deadline. +// Self-matching entries (same address+http_port) are skipped — this is the +// regression guard for #120. +func (m *Manager) tryAlternatives(alts []ws.AlternativeParent, perAttemptTimeout time.Duration, layerDeadline time.Time) bool { + for _, alt := range alts { + if m.stopping.Load() { + return false + } + if !layerDeadline.IsZero() && !time.Now().Before(layerDeadline) { + logging.Warn("Cached-alternatives layer hit deadline; falling through to seed list") + return false + } + if alt.Address == m.local.Address && alt.HTTPPort == m.local.HTTPPort { + continue + } + logging.Info("Attempting reattach to %s (%s:%d)", alt.ID, alt.Address, alt.HTTPPort) + ctx, cancel := context.WithTimeout(context.Background(), perAttemptTimeout) + conn, err := m.dialer(ctx, alt.Address, alt.HTTPPort, m.opts.ClusterSecret, perAttemptTimeout) + cancel() + if err != nil { + logging.Warn("Reattach to %s failed: %v", alt.ID, err) + continue + } + welcome, err := m.attach(context.Background(), conn, perAttemptTimeout) + if err != nil { + logging.Warn("Reattach to %s handshake failed: %v", alt.ID, err) + if !conn.IsClosed() { + conn.Close(1000, "") + } + continue + } + // Race guard: the conn might have dropped between welcome and now. + m.mu.Lock() + stillActive := m.parent == conn && !conn.IsClosed() + cb := m.onReattach + m.mu.Unlock() + if !stillActive { + logging.Warn("Reattach to %s dropped before activation; trying next", alt.ID) + continue + } + _ = welcome // attach() already populated lastKnownAlts and installed handlers + logging.Info("Reattached to %s — gossip resumed", alt.ID) + if cb != nil { + cb(conn) + } + conn.StartHeartbeat() + return true + } + return false +} + +// sleep blocks for d unless the manager is stopped first. Returns false if +// the wait was cut short by stop(). +func (m *Manager) sleep(d time.Duration) bool { + if m.stopping.Load() { + return false + } + t := time.NewTimer(d) + defer t.Stop() + select { + case <-t.C: + return !m.stopping.Load() + case <-m.stopCh: + return false + } +} + +// SendGoodbyeToChildren broadcasts a goodbye-with-alternatives to every +// attached transient. Used during graceful shutdown so transients can +// reattach within seconds instead of waiting the heartbeat timeout. +func (m *Manager) SendGoodbyeToChildren(reason string) { + if reason == "" { + reason = "shutdown" + } + m.mu.Lock() + if len(m.children) == 0 { + m.mu.Unlock() + return + } + conns := make([]*ws.Connection, 0, len(m.children)) + ids := make([]string, 0, len(m.children)) + for id, c := range m.children { + conns = append(conns, c) + ids = append(ids, id) + } + m.mu.Unlock() + + alts := m.getAlternativeSubstrates("") + payload := ws.GoodbyePayload{Reason: reason, AlternativeParents: alts} + + logging.Info("Sending goodbye to %d attached transients (%d alternatives)", len(conns), len(alts)) + for i, c := range conns { + if err := c.SendAttachment(ws.AttachmentTypeGoodbye, payload); err != nil { + logging.Debug("Failed to send goodbye to %s: %v", ids[i], err) + } + } +} + +// GetAlternativeSubstrates returns up to 5 candidate substrates for redirects +// and goodbyes, preferring same-enclave peers. Pass empty enclave to use the +// local node's enclave as the preference. +func (m *Manager) GetAlternativeSubstrates(enclave string) []ws.AlternativeParent { + return m.getAlternativeSubstrates(enclave) +} + +func (m *Manager) getAlternativeSubstrates(enclave string) []ws.AlternativeParent { + if enclave == "" { + enclave = m.local.Enclave + } + peers := m.gossip.GetPeers() + var sameEnclave, other []*gossip.Node + for _, p := range peers { + if p.Enclave == enclave { + sameEnclave = append(sameEnclave, p) + } else { + other = append(other, p) + } + } + candidates := append(sameEnclave, other...) + if len(candidates) > 5 { + candidates = candidates[:5] + } + out := make([]ws.AlternativeParent, 0, len(candidates)) + for _, p := range candidates { + out = append(out, ws.AlternativeParent{ + ID: string(p.ID), + Address: p.Address, + HTTPPort: p.HTTPPort, + Enclave: p.Enclave, + }) + } + return out +} + +// Stop tears the manager down: signals the reattach loop to exit, sends +// goodbye to children, and closes the active parent connection. Safe to +// call more than once. Children are *not* closed here — the HTTP server +// layer owns those sockets. +func (m *Manager) Stop() { + m.stopOnce.Do(func() { + m.stopping.Store(true) + close(m.stopCh) + m.SendGoodbyeToChildren("shutdown") + m.mu.Lock() + parent := m.parent + m.parent = nil + m.children = make(map[string]*ws.Connection) + m.mu.Unlock() + if parent != nil && !parent.IsClosed() { + parent.Close(1000, "shutting down") + } + }) +} + +// Stopping reports whether Stop has been invoked. Exposed for tests. +func (m *Manager) Stopping() bool { return m.stopping.Load() } + +// LastKnownAlternatives returns a snapshot of the cached topology used as the +// reattach fallback. Exposed for testing. +func (m *Manager) LastKnownAlternatives() []ws.AlternativeParent { + m.mu.Lock() + defer m.mu.Unlock() + return append([]ws.AlternativeParent(nil), m.lastKnownAlts...) +} + +// buildWelcomeTopology builds the SimpleMessage topology entries advertised +// to a newly-attached transient. The list is the local node plus all known +// peers, each wrapped as a SYNC announcement (matching what the transient +// would learn through HTTP gossip's SYNC propagation path). +func buildWelcomeTopology(peers []*gossip.Node, local *gossip.Node) []gossip.SimpleMessage { + now := time.Now().Unix() + out := make([]gossip.SimpleMessage, 0, len(peers)+1) + nodes := append(peers, local) + for _, n := range nodes { + out = append(out, gossip.SimpleMessage{ + Type: string(gossip.MessageTypeSync), + From: string(local.ID), + Timestamp: now, + MessageID: "", + NodeInfo: &gossip.SimpleNodeInfo{ + ID: string(n.ID), + Address: n.Address, + Port: n.Port, + HTTPPort: n.HTTPPort, + Enclave: n.Enclave, + }, + }) + } + return out +} + +// buildAltsFromTopology distills welcome.topology into AlternativeParent +// entries for the reattach cache. Self is excluded — the address+port skip +// inside tryAlternatives is the literal-match safety net, but stripping self +// up front keeps the cache honest. +func buildAltsFromTopology(topology []gossip.SimpleMessage, selfID string) []ws.AlternativeParent { + out := make([]ws.AlternativeParent, 0, len(topology)) + for _, sync := range topology { + if sync.NodeInfo == nil { + continue + } + if sync.NodeInfo.ID == selfID { + continue + } + out = append(out, ws.AlternativeParent{ + ID: sync.NodeInfo.ID, + Address: sync.NodeInfo.Address, + HTTPPort: sync.NodeInfo.HTTPPort, + Enclave: sync.NodeInfo.Enclave, + }) + } + return out +} + +// parseSeedAddress turns "host:port" into an AlternativeParent. Empty host or +// port, missing colon, non-numeric port, or out-of-range port → false. IPv6 +// in bracket-less form isn't supported (TS reference behaves the same way). +func parseSeedAddress(seed string) (ws.AlternativeParent, bool) { + idx := strings.LastIndex(seed, ":") + if idx <= 0 || idx == len(seed)-1 { + return ws.AlternativeParent{}, false + } + address := seed[:idx] + port, err := strconv.Atoi(seed[idx+1:]) + if err != nil || port <= 0 || port > 65535 { + return ws.AlternativeParent{}, false + } + return ws.AlternativeParent{ + ID: "seed-" + seed, + Address: address, + HTTPPort: port, + }, true +} diff --git a/internal/tree/manager_test.go b/internal/tree/manager_test.go new file mode 100644 index 0000000..3b98ea8 --- /dev/null +++ b/internal/tree/manager_test.go @@ -0,0 +1,765 @@ +package tree + +import ( + "context" + "encoding/json" + "errors" + "net/http/httptest" + "net/url" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/gorilla/websocket" + + "repram/internal/gossip" + "repram/internal/transport/ws" +) + +// ---- helpers ---------------------------------------------------------- + +// fakePeerer satisfies the Peerer interface with a static peer list. +type fakePeerer struct { + peers []*gossip.Node + mu sync.Mutex +} + +func (f *fakePeerer) GetPeers() []*gossip.Node { + f.mu.Lock() + defer f.mu.Unlock() + out := make([]*gossip.Node, len(f.peers)) + copy(out, f.peers) + return out +} + +func (f *fakePeerer) Add(n *gossip.Node) { + f.mu.Lock() + defer f.mu.Unlock() + f.peers = append(f.peers, n) +} + +func makeNode(id string, opts ...func(*gossip.Node)) *gossip.Node { + n := &gossip.Node{ + ID: gossip.NodeID(id), + Address: "127.0.0.1", + Port: 9090, + HTTPPort: 8080, + Enclave: "default", + } + for _, opt := range opts { + opt(n) + } + return n +} + +// pairServer hosts a /v1/ws endpoint and exposes both sides of a connection. +type pairServer struct { + srv *httptest.Server + serverCh chan *ws.Connection + clientWS *ws.Connection + cleanupFn func() +} + +func newPair(t *testing.T) *pairServer { + t.Helper() + ch := make(chan *ws.Connection, 1) + srv := httptest.NewServer(ws.Handler("", nil, func(c *ws.Connection) { + ch <- c + })) + + wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/v1/ws" + u, _ := url.Parse(wsURL) + d := &websocket.Dialer{HandshakeTimeout: 2 * time.Second} + raw, resp, err := d.Dial(u.String(), nil) + if err != nil { + if resp != nil { + _ = resp.Body.Close() + } + srv.Close() + t.Fatalf("dial: %v", err) + } + clientConn := ws.NewConnection(raw, "") + + p := &pairServer{srv: srv, serverCh: ch, clientWS: clientConn} + p.cleanupFn = func() { + clientConn.Close(websocket.CloseNormalClosure, "") + srv.Close() + } + return p +} + +// serverConn returns the accepted server-side Connection, waiting up to 2s. +func (p *pairServer) serverConn(t *testing.T) *ws.Connection { + t.Helper() + select { + case c := <-p.serverCh: + return c + case <-time.After(2 * time.Second): + t.Fatalf("server never accepted") + return nil + } +} + +func (p *pairServer) Address() (string, int) { + u, _ := url.Parse(p.srv.URL) + host := u.Hostname() + port, _ := strconv.Atoi(u.Port()) + return host, port +} + +// wireHelloHandler installs a server-side handler that runs HandleHello on +// any incoming hello frame, returning the accepted/rejected verdict. +func wireHelloHandler(server *ws.Connection, mgr *Manager) { + server.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + if msg.Type != ws.AttachmentTypeHello { + return + } + var h ws.HelloPayload + if err := json.Unmarshal(msg.Payload, &h); err != nil { + return + } + mgr.HandleHello(server, &h) + }) +} + +// ---- role detection --------------------------------------------------- + +func TestRoleSubstrate(t *testing.T) { + local := makeNode("substrate-1") + peers := &fakePeerer{} + m := NewManager(local, peers, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + if m.Role() != RoleSubstrate { + t.Errorf("role: got %q want substrate", m.Role()) + } + if !m.IsInboundCapable() { + t.Error("substrate should be inbound capable") + } + if m.Parent() != nil { + t.Error("substrate must not have a parent") + } +} + +func TestRoleTransient(t *testing.T) { + local := makeNode("transient-1") + peers := &fakePeerer{} + m := NewManager(local, peers, Options{Inbound: InboundFalse, MaxChildren: DefaultMaxChildren}) + if m.Role() != RoleTransient { + t.Errorf("role: got %q want transient", m.Role()) + } + if m.IsInboundCapable() { + t.Error("transient must not be inbound capable") + } +} + +// ---- server-side handshake ------------------------------------------- + +func TestHandleHelloAcceptsAndSendsWelcome(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + local := makeNode("substrate-1") + peers := &fakePeerer{peers: []*gossip.Node{ + makeNode("peer-1", func(n *gossip.Node) { n.Address = "10.0.0.2"; n.HTTPPort = 8081 }), + }} + m := NewManager(local, peers, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + gotWelcome := make(chan ws.WelcomePayload, 1) + p.clientWS.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + if msg.Type != ws.AttachmentTypeWelcome { + return + } + var w ws.WelcomePayload + if err := json.Unmarshal(msg.Payload, &w); err != nil { + t.Errorf("decode welcome: %v", err) + return + } + gotWelcome <- w + }) + + hello := &ws.HelloPayload{ + NodeID: "transient-1", Enclave: "default", + Address: "192.168.1.100", HTTPPort: 8080, + Capabilities: ws.Capabilities{Inbound: "false"}, + } + if !m.HandleHello(srv, hello) { + t.Fatal("HandleHello rejected unexpectedly") + } + + select { + case w := <-gotWelcome: + if w.YourPosition.ParentID != "substrate-1" { + t.Errorf("parent_id: %q", w.YourPosition.ParentID) + } + if w.YourPosition.Depth != 1 { + t.Errorf("depth: %d", w.YourPosition.Depth) + } + if len(w.Topology) != 2 { + t.Errorf("topology len: %d want 2", len(w.Topology)) + } + case <-time.After(2 * time.Second): + t.Fatal("no welcome received") + } +} + +func TestHandleHelloRegistersChild(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + m := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + hello := &ws.HelloPayload{ + NodeID: "transient-1", Enclave: "default", + Address: "192.168.1.100", HTTPPort: 8080, + Capabilities: ws.Capabilities{Inbound: "false"}, + } + if !m.HandleHello(srv, hello) { + t.Fatal("HandleHello rejected") + } + if m.ChildCount() != 1 { + t.Errorf("child count: %d want 1", m.ChildCount()) + } + if !m.HasChild("transient-1") { + t.Error("transient-1 not registered as child") + } + if srv.RemoteNodeID() != "transient-1" { + t.Errorf("remoteNodeID: %q", srv.RemoteNodeID()) + } + if srv.RemoteEnclave() != "default" { + t.Errorf("remoteEnclave: %q", srv.RemoteEnclave()) + } +} + +func TestChildRemovedOnConnectionClose(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + m := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + hello := &ws.HelloPayload{NodeID: "transient-1", Enclave: "default", Capabilities: ws.Capabilities{Inbound: "false"}} + m.HandleHello(srv, hello) + if m.ChildCount() != 1 { + t.Fatalf("setup: child count %d", m.ChildCount()) + } + + p.clientWS.Close(websocket.CloseNormalClosure, "") + + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if m.ChildCount() == 0 { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Errorf("child still present after close: count %d", m.ChildCount()) +} + +// ---- max children ----------------------------------------------------- + +func TestRejectsAtCapacity(t *testing.T) { + p1 := newPair(t) + defer p1.cleanupFn() + srv1 := p1.serverConn(t) + + m := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: 1}) + defer m.Stop() + + if !m.HandleHello(srv1, &ws.HelloPayload{NodeID: "t1", Enclave: "default"}) { + t.Fatal("first hello rejected") + } + if m.ChildCount() != 1 { + t.Fatalf("after first: count %d", m.ChildCount()) + } + + p2 := newPair(t) + defer p2.cleanupFn() + srv2 := p2.serverConn(t) + + gotBye := make(chan ws.GoodbyePayload, 1) + p2.clientWS.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + if msg.Type != ws.AttachmentTypeGoodbye { + return + } + var g ws.GoodbyePayload + if err := json.Unmarshal(msg.Payload, &g); err == nil { + gotBye <- g + } + }) + + if m.HandleHello(srv2, &ws.HelloPayload{NodeID: "t2", Enclave: "default"}) { + t.Fatal("second hello unexpectedly accepted") + } + if m.ChildCount() != 1 { + t.Errorf("after reject: count %d want 1", m.ChildCount()) + } + select { + case g := <-gotBye: + if g.Reason != "at capacity" { + t.Errorf("reason: %q", g.Reason) + } + case <-time.After(2 * time.Second): + t.Fatal("no goodbye received") + } +} + +func TestMaxChildrenZeroRejectsAll(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + m := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: 0}) + defer m.Stop() + + gotBye := make(chan ws.GoodbyePayload, 1) + p.clientWS.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + if msg.Type != ws.AttachmentTypeGoodbye { + return + } + var g ws.GoodbyePayload + _ = json.Unmarshal(msg.Payload, &g) + gotBye <- g + }) + + if m.HandleHello(srv, &ws.HelloPayload{NodeID: "t1", Enclave: "default"}) { + t.Fatal("hello unexpectedly accepted with MaxChildren=0") + } + select { + case g := <-gotBye: + if g.Reason != "at capacity" { + t.Errorf("reason: %q", g.Reason) + } + case <-time.After(2 * time.Second): + t.Fatal("no goodbye received") + } +} + +// ---- transient-side attach ------------------------------------------- + +func TestAttachCompletesHandshake(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + subMgr := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer subMgr.Stop() + wireHelloHandler(srv, subMgr) + + tranMgr := NewManager( + makeNode("transient-1", func(n *gossip.Node) { n.Address = "192.168.1.100" }), + &fakePeerer{}, + Options{Inbound: InboundFalse, MaxChildren: 0}, + ) + defer tranMgr.Stop() + + welcome, err := tranMgr.Attach(context.Background(), p.clientWS) + if err != nil { + t.Fatalf("attach: %v", err) + } + if welcome.YourPosition.ParentID != "substrate-1" { + t.Errorf("parent_id: %q", welcome.YourPosition.ParentID) + } + if welcome.YourPosition.Depth != 1 { + t.Errorf("depth: %d", welcome.YourPosition.Depth) + } + if tranMgr.Role() != RoleTransient { + t.Errorf("role: %q", tranMgr.Role()) + } + if tranMgr.Parent() != p.clientWS { + t.Error("parent not registered after attach") + } +} + +func TestAttachTimesOut(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + _ = p.serverConn(t) // no welcome-emitting handler installed + + tranMgr := NewManager(makeNode("transient-1"), &fakePeerer{}, Options{Inbound: InboundFalse}) + defer tranMgr.Stop() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + _, err := tranMgr.attach(ctx, p.clientWS, 500*time.Millisecond) + if err == nil { + t.Fatal("expected timeout error") + } + if tranMgr.Parent() != nil { + t.Error("parent set despite timeout") + } +} + +// ---- alternative substrates ------------------------------------------ + +func TestGetAlternativeSubstratesPrefersSameEnclave(t *testing.T) { + local := makeNode("substrate-1") + peers := &fakePeerer{} + peers.Add(makeNode("same-1", func(n *gossip.Node) { n.Enclave = "default" })) + peers.Add(makeNode("other-1", func(n *gossip.Node) { n.Enclave = "other" })) + peers.Add(makeNode("same-2", func(n *gossip.Node) { n.Enclave = "default" })) + + m := NewManager(local, peers, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + alts := m.GetAlternativeSubstrates("default") + if len(alts) != 3 { + t.Fatalf("alt count: %d", len(alts)) + } + if alts[0].ID != "same-1" || alts[1].ID != "same-2" { + t.Errorf("same-enclave alts not first: %+v", alts) + } + if alts[2].ID != "other-1" { + t.Errorf("other-enclave alt not last: %+v", alts) + } +} + +func TestGetAlternativeSubstratesLimitsToFive(t *testing.T) { + local := makeNode("substrate-1") + peers := &fakePeerer{} + for i := 0; i < 10; i++ { + peers.Add(makeNode("peer-" + strconv.Itoa(i))) + } + m := NewManager(local, peers, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + if alts := m.GetAlternativeSubstrates(""); len(alts) != 5 { + t.Errorf("alt count: %d want 5", len(alts)) + } +} + +// ---- goodbye on shutdown --------------------------------------------- + +func TestSendGoodbyeToChildren(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + m := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + gotBye := make(chan ws.GoodbyePayload, 1) + p.clientWS.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + if msg.Type == ws.AttachmentTypeGoodbye { + var g ws.GoodbyePayload + _ = json.Unmarshal(msg.Payload, &g) + gotBye <- g + } + }) + + hello := &ws.HelloPayload{NodeID: "t1", Enclave: "default"} + if !m.HandleHello(srv, hello) { + t.Fatal("hello rejected") + } + + m.SendGoodbyeToChildren("") + select { + case g := <-gotBye: + if g.Reason != "shutdown" { + t.Errorf("reason: %q want shutdown", g.Reason) + } + case <-time.After(2 * time.Second): + t.Fatal("no goodbye received") + } +} + +func TestParentClearedAfterGoodbye(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + subMgr := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer subMgr.Stop() + wireHelloHandler(srv, subMgr) + + tranMgr := NewManager( + makeNode("transient-1", func(n *gossip.Node) { n.Address = "192.168.1.100" }), + &fakePeerer{}, + Options{Inbound: InboundFalse}, + ) + defer tranMgr.Stop() + // Empty seed provider keeps reattach loop spinning harmlessly with no + // candidates — what we care about is the parent clear. + tranMgr.SetSeedProvider(func() []string { return nil }) + + if _, err := tranMgr.Attach(context.Background(), p.clientWS); err != nil { + t.Fatalf("attach: %v", err) + } + if tranMgr.Parent() == nil { + t.Fatal("parent not set after attach") + } + + subMgr.SendGoodbyeToChildren("maintenance") + + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + if tranMgr.Parent() == nil { + return + } + time.Sleep(20 * time.Millisecond) + } + t.Error("parent not cleared after goodbye") +} + +// ---- topology caching -------------------------------------------------- + +func TestAttachCachesTopologyExcludingSelf(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + subPeers := &fakePeerer{} + subPeers.Add(makeNode("peer-a", func(n *gossip.Node) { n.HTTPPort = 8001 })) + subPeers.Add(makeNode("peer-b", func(n *gossip.Node) { n.HTTPPort = 8002 })) + subMgr := NewManager(makeNode("substrate-1"), subPeers, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer subMgr.Stop() + wireHelloHandler(srv, subMgr) + + tranMgr := NewManager( + makeNode("transient-1", func(n *gossip.Node) { n.Address = "10.0.0.1" }), + &fakePeerer{}, + Options{Inbound: InboundFalse}, + ) + defer tranMgr.Stop() + + if _, err := tranMgr.Attach(context.Background(), p.clientWS); err != nil { + t.Fatalf("attach: %v", err) + } + + cached := tranMgr.LastKnownAlternatives() + gotIDs := make([]string, 0, len(cached)) + for _, a := range cached { + gotIDs = append(gotIDs, a.ID) + } + want := []string{"peer-a", "peer-b", "substrate-1"} + // order is insertion order; sort for stable comparison + sortStrings(gotIDs) + if !equalStrings(gotIDs, want) { + t.Errorf("cached alts: got %v want %v", gotIDs, want) + } +} + +// ---- seed parsing ---------------------------------------------------- + +func TestParseSeedAddress(t *testing.T) { + cases := []struct { + in string + ok bool + addr string + port int + }{ + {"10.0.0.5:8080", true, "10.0.0.5", 8080}, + {"host.example:443", true, "host.example", 443}, + {"no-colon", false, "", 0}, + {":8080", false, "", 0}, + {"host:", false, "", 0}, + {"host:notaport", false, "", 0}, + {"host:0", false, "", 0}, + {"host:99999", false, "", 0}, + } + for _, c := range cases { + t.Run(c.in, func(t *testing.T) { + alt, ok := parseSeedAddress(c.in) + if ok != c.ok { + t.Fatalf("ok: got %v want %v", ok, c.ok) + } + if !ok { + return + } + if alt.Address != c.addr || alt.HTTPPort != c.port { + t.Errorf("got %s:%d want %s:%d", alt.Address, alt.HTTPPort, c.addr, c.port) + } + if alt.ID != "seed-"+c.in { + t.Errorf("id: %q", alt.ID) + } + }) + } +} + +// ---- stale-connection close guard (matches tree.test.ts "stale conn" case) ----- + +func TestStaleConnectionCloseDoesNotClobberParent(t *testing.T) { + pA := newPair(t) + defer pA.cleanupFn() + srvA := pA.serverConn(t) + subA := NewManager(makeNode("substrate-a"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer subA.Stop() + wireHelloHandler(srvA, subA) + + pB := newPair(t) + defer pB.cleanupFn() + srvB := pB.serverConn(t) + subB := NewManager(makeNode("substrate-b"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer subB.Stop() + wireHelloHandler(srvB, subB) + + tranMgr := NewManager( + makeNode("transient-1", func(n *gossip.Node) { n.Address = "10.0.0.1" }), + &fakePeerer{}, + Options{Inbound: InboundFalse}, + ) + defer tranMgr.Stop() + tranMgr.SetSeedProvider(func() []string { return nil }) + + if _, err := tranMgr.Attach(context.Background(), pA.clientWS); err != nil { + t.Fatalf("attach A: %v", err) + } + if tranMgr.Parent() != pA.clientWS { + t.Fatal("parent != A after first attach") + } + if _, err := tranMgr.Attach(context.Background(), pB.clientWS); err != nil { + t.Fatalf("attach B: %v", err) + } + if tranMgr.Parent() != pB.clientWS { + t.Fatal("parent != B after second attach") + } + + // Closing A should not clobber the active parent (B). + pA.clientWS.Close(websocket.CloseNormalClosure, "") + time.Sleep(80 * time.Millisecond) + if tranMgr.Parent() != pB.clientWS { + t.Errorf("stale close clobbered active parent: now %v want B", tranMgr.Parent()) + } +} + +// ---- self-skip (#120 regression) ------------------------------------- + +// TestTryAlternativesSkipsSelf is the timing-assertion regression for #120. +// If self-skip works the call returns false almost instantly. If it's +// broken, the dialer would try a 5-second connect to the unreachable +// self-address and the elapsed time would explode. +func TestTryAlternativesSkipsSelf(t *testing.T) { + local := makeNode("self", func(n *gossip.Node) { + n.Address = "10.0.10.104" + n.HTTPPort = 18080 + }) + var dialerCalls atomic.Int32 + dialer := Dialer(func(ctx context.Context, address string, port int, secret string, timeout time.Duration) (*ws.Connection, error) { + dialerCalls.Add(1) + // If self-skip is broken this would be called and stall for + // `timeout`. Return error fast so test failure is fast too. + return nil, errors.New("dialer should not have been called") + }) + m := NewManager(local, &fakePeerer{}, Options{Inbound: InboundFalse, Dialer: dialer}) + defer m.Stop() + + alts := []ws.AlternativeParent{ + {ID: "self-seed", Address: "10.0.10.104", HTTPPort: 18080}, + } + start := time.Now() + ok := m.tryAlternatives(alts, 5*time.Second, time.Time{}) + elapsed := time.Since(start) + + if ok { + t.Error("tryAlternatives unexpectedly succeeded against self-only alts") + } + if dialerCalls.Load() != 0 { + t.Errorf("dialer called %d times; self-skip is broken", dialerCalls.Load()) + } + if elapsed > 500*time.Millisecond { + t.Errorf("self-skip too slow: %v (want < 500ms)", elapsed) + } +} + +// ---- ungraceful disconnect → seed fallback --------------------------- + +func TestUngracefulCloseTriggersReattachLoop(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + subMgr := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer subMgr.Stop() + wireHelloHandler(srv, subMgr) + + // Use a dialer that always fails fast so the reattach loop reaches the + // seed provider quickly. + dialer := Dialer(func(ctx context.Context, address string, port int, secret string, timeout time.Duration) (*ws.Connection, error) { + return nil, errors.New("unreachable") + }) + + var seedCalls atomic.Int32 + seedCh := make(chan struct{}, 1) + + tranMgr := NewManager( + makeNode("transient-1", func(n *gossip.Node) { n.Address = "10.0.0.1" }), + &fakePeerer{}, + Options{Inbound: InboundFalse, Dialer: dialer}, + ) + defer tranMgr.Stop() + tranMgr.SetSeedProvider(func() []string { + seedCalls.Add(1) + select { + case seedCh <- struct{}{}: + default: + } + return nil + }) + + if _, err := tranMgr.Attach(context.Background(), p.clientWS); err != nil { + t.Fatalf("attach: %v", err) + } + + // Ungracefully close the parent. + p.clientWS.Close(websocket.CloseNormalClosure, "") + + select { + case <-seedCh: + // seed provider was hit — reattach loop is running + case <-time.After(5 * time.Second): + t.Fatalf("seed provider never called (calls=%d)", seedCalls.Load()) + } +} + +// ---- stop() unblocks pending backoff sleeps -------------------------- + +func TestStopUnblocksSleep(t *testing.T) { + m := NewManager(makeNode("t-1"), &fakePeerer{}, Options{Inbound: InboundFalse}) + + done := make(chan bool, 1) + go func() { done <- m.sleep(60 * time.Second) }() + + time.Sleep(10 * time.Millisecond) + m.Stop() + + select { + case ok := <-done: + if ok { + t.Error("sleep returned true (not interrupted)") + } + case <-time.After(time.Second): + t.Fatal("Stop() did not wake sleep within 1s") + } + if !m.Stopping() { + t.Error("Stopping() returned false after Stop()") + } +} + +// ---- helpers ---------------------------------------------------------- + +func sortStrings(s []string) { + for i := 1; i < len(s); i++ { + for j := i; j > 0 && s[j-1] > s[j]; j-- { + s[j-1], s[j] = s[j], s[j-1] + } + } +} + +func equalStrings(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} From fb40ef7473be8262657302093654234bc8ef4f76 Mon Sep 17 00:00:00 2001 From: TickTockBent Date: Tue, 12 May 2026 08:53:43 -0400 Subject: [PATCH 3/5] phase 3-5 prep (#135): tree manager ACK routing + child broadcast MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the data structures and helpers the cluster integration will use to wire WS frames into the existing gossip dispatch path: - RecordAckRoute / LookupAckRoute / ClearAckRoute — substrate-side table that maps a relayed PUT's messageId to the originating child connection so an enclave peer's HTTP ACK can be routed back through WS. Entries auto-evict after the configured TTL (matches REPRAM_WRITE_TIMEOUT). Stop() now cancels in-flight timers and clears the table. - BroadcastToChildren — fans a gossip message out to every attached transient whose enclave matches. Substrate calls this from the cluster's PUT handler so transients see other agents' writes in their local store. Cross-enclave traffic is dropped. - Lifecycle hardening: a stop-aware context is threaded into the reattach loop's dialer, the reattach goroutine is tracked in a WaitGroup, and Stop() waits for it. Eliminates a goroutine leak that surfaced as test pollution under -race + -count=N. - Race fix in attach(): the temporary welcome/goodbye/close handlers must be installed BEFORE hello is sent. A fast substrate could answer welcome inside SendAttachment's return path, the handlers weren't registered yet, the event was dropped, and the select below waited the full AttachTimeout. The TS reference's EventEmitter pattern hid this by allocating the listener synchronously around the same event loop tick; the Go port needs to be explicit. Found by 10x race iteration of the suite. 5 new tests: - TestRecordAndLookupAckRoute - TestAckRouteAutoEvicts (TTL eviction) - TestClearAckRoute - TestBroadcastToChildren (delivers to all attached children) - TestBroadcastToChildrenSkipsOtherEnclave (cross-enclave gate) Cluster wiring (phase 3 fan-out + phase 4 ACK reverse + phase 5 child broadcast hookup) lands in the next commit. Tests pass under go test -race -count=10. // ticktockbent --- internal/tree/manager.go | 172 ++++++++++++++++++++++++++++------ internal/tree/manager_test.go | 129 +++++++++++++++++++++++++ 2 files changed, 270 insertions(+), 31 deletions(-) diff --git a/internal/tree/manager.go b/internal/tree/manager.go index b75d3c3..9314cb7 100644 --- a/internal/tree/manager.go +++ b/internal/tree/manager.go @@ -118,9 +118,20 @@ type Manager struct { onReattach func(*ws.Connection) seedProvider func() []string - stopOnce sync.Once - stopping atomic.Bool - stopCh chan struct{} + // ackRoutes tracks (messageId → child connection) for PUTs relayed + // through this substrate. When an enclave peer ACKs a relayed PUT, the + // substrate forwards the ACK back through ackRoutes[messageId] so the + // originating transient observes quorum confirmation. Entries are + // evicted by ackRouteTimers or explicit ClearAckRoute. + ackRoutes map[string]*ws.Connection + ackRouteTimers map[string]*time.Timer + + stopOnce sync.Once + stopping atomic.Bool + stopCh chan struct{} + stopCtx context.Context + stopCancel context.CancelFunc + wg sync.WaitGroup } // NewManager constructs a Manager. The Inbound option resolves the role @@ -129,13 +140,18 @@ func NewManager(local *gossip.Node, peers Peerer, opts Options) *Manager { if opts.Dialer == nil { opts.Dialer = ws.ConnectToSubstrate } + stopCtx, stopCancel := context.WithCancel(context.Background()) m := &Manager{ - local: local, - gossip: peers, - opts: opts, - dialer: opts.Dialer, - children: make(map[string]*ws.Connection), - stopCh: make(chan struct{}), + local: local, + gossip: peers, + opts: opts, + dialer: opts.Dialer, + children: make(map[string]*ws.Connection), + ackRoutes: make(map[string]*ws.Connection), + ackRouteTimers: make(map[string]*time.Timer), + stopCh: make(chan struct{}), + stopCtx: stopCtx, + stopCancel: stopCancel, } if opts.Inbound == InboundTrue { m.role = RoleSubstrate @@ -289,13 +305,9 @@ func (m *Manager) attach(ctx context.Context, conn *ws.Connection, timeout time. HTTPPort: m.local.HTTPPort, Capabilities: ws.Capabilities{Inbound: string(m.opts.Inbound)}, } - if err := conn.SendAttachment(ws.AttachmentTypeHello, hello); err != nil { - return nil, err - } - - // Race a temporary attachment handler against the close event and the - // timeout. The handler is removed in every exit path so it can't fire - // after the long-lived handlers below take over. + // Install the temporary handlers BEFORE sending hello — a fast substrate + // can answer with welcome before SendAttachment returns, and missing the + // event causes the select below to wait the full timeout. welcomeCh := make(chan *ws.WelcomePayload, 1) rejectedCh := make(chan struct{}, 1) removeAttach := conn.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { @@ -304,18 +316,36 @@ func (m *Manager) attach(ctx context.Context, conn *ws.Connection, timeout time. var w ws.WelcomePayload if err := json.Unmarshal(msg.Payload, &w); err != nil { logging.Warn("Welcome payload decode failed: %v", err) - rejectedCh <- struct{}{} + select { + case rejectedCh <- struct{}{}: + default: + } return } - welcomeCh <- &w + select { + case welcomeCh <- &w: + default: + } case ws.AttachmentTypeGoodbye: - rejectedCh <- struct{}{} + select { + case rejectedCh <- struct{}{}: + default: + } } }) removeClose := conn.AddCloseHandler(func(int, string) { - rejectedCh <- struct{}{} + select { + case rejectedCh <- struct{}{}: + default: + } }) + if err := conn.SendAttachment(ws.AttachmentTypeHello, hello); err != nil { + removeAttach() + removeClose() + return nil, err + } + var welcome *ws.WelcomePayload select { case welcome = <-welcomeCh: @@ -363,7 +393,7 @@ func (m *Manager) attach(ctx context.Context, conn *ws.Connection, timeout time. m.mu.Lock() m.parent = nil m.mu.Unlock() - go m.triggerReattach(p.AlternativeParents) + m.triggerReattach(p.AlternativeParents) } }) conn.AddCloseHandler(func(int, string) { @@ -377,7 +407,7 @@ func (m *Manager) attach(ctx context.Context, conn *ws.Connection, timeout time. m.mu.Lock() m.parent = nil m.mu.Unlock() - go m.triggerReattach(nil) + m.triggerReattach(nil) }) logging.Info("Attached to substrate %s (depth %d, topology %d nodes)", @@ -387,7 +417,8 @@ func (m *Manager) attach(ctx context.Context, conn *ws.Connection, timeout time. // triggerReattach is the single-flight entry point for the reattach loop. // Concurrent invocations (goodbye and close firing back-to-back) collapse -// to one running loop. +// to one running loop. The loop runs in its own goroutine tracked by m.wg +// so Stop() can wait for clean exit before returning. func (m *Manager) triggerReattach(supplied []ws.AlternativeParent) { if m.stopping.Load() { return @@ -398,15 +429,18 @@ func (m *Manager) triggerReattach(supplied []ws.AlternativeParent) { return } m.reattachInFlight = true + m.wg.Add(1) m.mu.Unlock() - defer func() { - m.mu.Lock() - m.reattachInFlight = false - m.mu.Unlock() + go func() { + defer m.wg.Done() + defer func() { + m.mu.Lock() + m.reattachInFlight = false + m.mu.Unlock() + }() + m.reattachLoop(supplied) }() - - m.reattachLoop(supplied) } func (m *Manager) reattachLoop(supplied []ws.AlternativeParent) { @@ -479,14 +513,14 @@ func (m *Manager) tryAlternatives(alts []ws.AlternativeParent, perAttemptTimeout continue } logging.Info("Attempting reattach to %s (%s:%d)", alt.ID, alt.Address, alt.HTTPPort) - ctx, cancel := context.WithTimeout(context.Background(), perAttemptTimeout) + ctx, cancel := context.WithTimeout(m.stopCtx, perAttemptTimeout) conn, err := m.dialer(ctx, alt.Address, alt.HTTPPort, m.opts.ClusterSecret, perAttemptTimeout) cancel() if err != nil { logging.Warn("Reattach to %s failed: %v", alt.ID, err) continue } - welcome, err := m.attach(context.Background(), conn, perAttemptTimeout) + welcome, err := m.attach(m.stopCtx, conn, perAttemptTimeout) if err != nil { logging.Warn("Reattach to %s handshake failed: %v", alt.ID, err) if !conn.IsClosed() { @@ -530,6 +564,75 @@ func (m *Manager) sleep(d time.Duration) bool { } } +// RecordAckRoute remembers that messageID's ACK should be forwarded back to +// conn when it arrives. The mapping is auto-evicted after ttl; call +// ClearAckRoute on quorum success to free the entry sooner. +func (m *Manager) RecordAckRoute(messageID string, conn *ws.Connection, ttl time.Duration) { + m.mu.Lock() + if existing, ok := m.ackRouteTimers[messageID]; ok { + existing.Stop() + } + m.ackRoutes[messageID] = conn + m.ackRouteTimers[messageID] = time.AfterFunc(ttl, func() { + m.mu.Lock() + delete(m.ackRoutes, messageID) + delete(m.ackRouteTimers, messageID) + m.mu.Unlock() + }) + m.mu.Unlock() +} + +// LookupAckRoute returns the child connection that should receive an ACK for +// messageID, or nil if no route exists (PUT was not relayed through here). +func (m *Manager) LookupAckRoute(messageID string) *ws.Connection { + m.mu.Lock() + defer m.mu.Unlock() + return m.ackRoutes[messageID] +} + +// ClearAckRoute removes the ACK route for messageID and stops its eviction +// timer. Called by the cluster layer when the originator has accumulated +// enough ACKs for quorum, or when the message times out from the write side. +func (m *Manager) ClearAckRoute(messageID string) { + m.mu.Lock() + if t, ok := m.ackRouteTimers[messageID]; ok { + t.Stop() + delete(m.ackRouteTimers, messageID) + } + delete(m.ackRoutes, messageID) + m.mu.Unlock() +} + +// BroadcastToChildren fans msg out to every attached child whose enclave +// matches. Substrate uses this to deliver PUT replicas received from the +// HTTP gossip mesh down to its attached transients, so transients see +// other agents' writes in their local store. Errors from a single child do +// not abort the broadcast — best-effort like HTTP gossip. +func (m *Manager) BroadcastToChildren(msg *gossip.Message) { + m.mu.Lock() + if len(m.children) == 0 { + m.mu.Unlock() + return + } + conns := make([]*ws.Connection, 0, len(m.children)) + for _, c := range m.children { + conns = append(conns, c) + } + m.mu.Unlock() + + for _, c := range conns { + // Enclave gating: substrate must not leak cross-enclave writes to + // transients. The child's enclave is recorded on the Connection by + // SetRemote during HandleHello. + if c.RemoteEnclave() != "" && c.RemoteEnclave() != m.local.Enclave { + continue + } + if err := c.SendGossip(msg); err != nil { + logging.Debug("BroadcastToChildren: send to %s failed: %v", c.RemoteNodeID(), err) + } + } +} + // SendGoodbyeToChildren broadcasts a goodbye-with-alternatives to every // attached transient. Used during graceful shutdown so transients can // reattach within seconds instead of waiting the heartbeat timeout. @@ -605,15 +708,22 @@ func (m *Manager) Stop() { m.stopOnce.Do(func() { m.stopping.Store(true) close(m.stopCh) + m.stopCancel() m.SendGoodbyeToChildren("shutdown") m.mu.Lock() parent := m.parent m.parent = nil m.children = make(map[string]*ws.Connection) + for _, t := range m.ackRouteTimers { + t.Stop() + } + m.ackRoutes = make(map[string]*ws.Connection) + m.ackRouteTimers = make(map[string]*time.Timer) m.mu.Unlock() if parent != nil && !parent.IsClosed() { parent.Close(1000, "shutting down") } + m.wg.Wait() }) } diff --git a/internal/tree/manager_test.go b/internal/tree/manager_test.go index 3b98ea8..5a1370e 100644 --- a/internal/tree/manager_test.go +++ b/internal/tree/manager_test.go @@ -742,6 +742,135 @@ func TestStopUnblocksSleep(t *testing.T) { } } +// ---- ACK routing (phase-4 prep) -------------------------------------- + +func TestRecordAndLookupAckRoute(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + m := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + m.RecordAckRoute("msg-1", srv, 5*time.Second) + if got := m.LookupAckRoute("msg-1"); got != srv { + t.Errorf("LookupAckRoute: got %v want %v", got, srv) + } + if got := m.LookupAckRoute("msg-unknown"); got != nil { + t.Errorf("LookupAckRoute(unknown): got %v want nil", got) + } +} + +func TestAckRouteAutoEvicts(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + m := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + m.RecordAckRoute("msg-1", srv, 80*time.Millisecond) + if got := m.LookupAckRoute("msg-1"); got != srv { + t.Fatalf("LookupAckRoute before TTL: got %v want %v", got, srv) + } + time.Sleep(200 * time.Millisecond) + if got := m.LookupAckRoute("msg-1"); got != nil { + t.Errorf("LookupAckRoute after TTL: got %v want nil", got) + } +} + +func TestClearAckRoute(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + m := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + m.RecordAckRoute("msg-1", srv, 5*time.Second) + if m.LookupAckRoute("msg-1") == nil { + t.Fatal("setup: route missing") + } + m.ClearAckRoute("msg-1") + if got := m.LookupAckRoute("msg-1"); got != nil { + t.Errorf("LookupAckRoute after Clear: got %v want nil", got) + } +} + +// ---- child broadcast (phase-5 prep) ---------------------------------- + +func TestBroadcastToChildren(t *testing.T) { + p1 := newPair(t) + defer p1.cleanupFn() + srv1 := p1.serverConn(t) + p2 := newPair(t) + defer p2.cleanupFn() + srv2 := p2.serverConn(t) + + m := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + m.HandleHello(srv1, &ws.HelloPayload{NodeID: "t1", Enclave: "default"}) + m.HandleHello(srv2, &ws.HelloPayload{NodeID: "t2", Enclave: "default"}) + + got1 := make(chan *gossip.Message, 1) + got2 := make(chan *gossip.Message, 1) + p1.clientWS.OnMessage(func(msg *gossip.Message) { got1 <- msg }) + p2.clientWS.OnMessage(func(msg *gossip.Message) { got2 <- msg }) + + msg := &gossip.Message{ + Type: gossip.MessageTypePut, + From: "peer-x", + Key: "k", + Data: []byte("v"), + TTL: 60, + Timestamp: time.Now(), + MessageID: "broadcast-test", + } + m.BroadcastToChildren(msg) + + for i, ch := range []<-chan *gossip.Message{got1, got2} { + select { + case r := <-ch: + if r.MessageID != "broadcast-test" { + t.Errorf("child %d: id=%q want broadcast-test", i+1, r.MessageID) + } + case <-time.After(2 * time.Second): + t.Fatalf("child %d never received broadcast", i+1) + } + } +} + +func TestBroadcastToChildrenSkipsOtherEnclave(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + m := NewManager(makeNode("substrate-1"), &fakePeerer{}, Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}) + defer m.Stop() + + m.HandleHello(srv, &ws.HelloPayload{NodeID: "t1", Enclave: "other-enclave"}) + + var calls atomic.Int32 + p.clientWS.OnMessage(func(*gossip.Message) { calls.Add(1) }) + + msg := &gossip.Message{ + Type: gossip.MessageTypePut, + From: "peer-x", + Key: "k", + Data: []byte("v"), + TTL: 60, + Timestamp: time.Now(), + MessageID: "cross-enclave", + } + m.BroadcastToChildren(msg) + + time.Sleep(100 * time.Millisecond) + if n := calls.Load(); n != 0 { + t.Errorf("child in other enclave received broadcast: %d times", n) + } +} + // ---- helpers ---------------------------------------------------------- func sortStrings(s []string) { From db14061ba0e8df25c2bf052aa17a4588e78eb563 Mon Sep 17 00:00:00 2001 From: TickTockBent Date: Tue, 12 May 2026 09:02:21 -0400 Subject: [PATCH 4/5] phase 3-6 (#135): cluster integration + --mcp WS attach MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the WS transport and tree manager into the running binary so substrate nodes accept inbound attachments and transient nodes attach outbound after HTTP bootstrap. cluster.ClusterNode - AckRouter / ChildBroadcaster interfaces — substrate's tree manager satisfies both. SetAckRouter / SetChildBroadcaster installs them; transient nodes leave them nil. - handlePutMessage now consults ackRouter when msg.From is not in the HTTP peer list. Routes the substrate's own immediate ACK back over the originating WS pipe (phase 4 — substrate local store is the first quorum vote). - handlePutMessage also calls childBroadcaster.BroadcastToChildren after storing, fanning replicas out to attached enclave-matched transients (phase 5). - handleAckMessage now falls through to ackRouter when the message isn't ours — forwards enclave-peer ACKs back through the WS pipe so the transient's quorum tally advances beyond the substrate's own vote. - WriteTimeout accessor for the WS dispatch's ACK-route eviction TTL. tree.Manager - RouteAck implements cluster.AckRouter — looks up the route by messageId, writes the ACK to the child connection. cmd/repram/main.go - Always constructs a tree.Manager. Substrate (REPRAM_INBOUND=true) accepts /v1/ws; transient (default) attempts outbound attach after HTTP bootstrap (skipped when REPRAM_PEERS is empty — preserves "single-agent local scratchpad" mode). - /v1/ws is bound on an outer http.ServeMux that bypasses the gorilla router's TimeoutHandler + MaxRequestSize wrappers (both fight WebSocket Hijack / long-lived connections). - wsHandler 404s on transient nodes — avoids leaking role to scanners. - bindWSConnection installs a one-shot hello gate; on accept it wires conn.OnMessage to dispatch incoming gossip into the cluster handler and records an ACK route per relayed PUT. 30s no-hello timeout closes silent connections. - SetReattachCallback re-binds the parent dispatch on the new connection after a successful reattach. - SetSeedProvider feeds the bootstrap list into the tree-side reattach loop. - /v1/topology now exposes role, attached children, and parent_id (acceptance: "substrate's HTTP topology endpoint shows transient as attached child"). - clusterPeerer adapter so *cluster.ClusterNode satisfies tree.Peerer via its Topology method. cmd/repram/ws_integration_test.go (new) - TestWSAttachHandshake — hello/welcome round-trip through the real HTTP server stack - TestWSPutStoresLocallyAndAcks — relay round-trip: PUT over WS, substrate stores, ACK back over WS - TestWSReceivePathFanout — HTTP-gossip arrival fans out to attached transients (phase 5) - TestWSRejectIfTransient — transients return 404 on /v1/ws Full repo passes go test ./... and go test -race ./... — no regressions in existing 118+ tests. Phase 7 (24h burn-in 2.2 on real cluster infra) is deferred to a dedicated follow-up: needs a multi-substrate + multi-transient docker-compose setup and 24h+ of k6 workload, which doesn't belong in a unit-test PR. // ticktockbent --- cmd/repram/main.go | 213 ++++++++++++++++++- cmd/repram/ws_integration_test.go | 341 ++++++++++++++++++++++++++++++ internal/cluster/node.go | 94 +++++++- internal/tree/manager.go | 20 ++ 4 files changed, 654 insertions(+), 14 deletions(-) create mode 100644 cmd/repram/ws_integration_test.go diff --git a/cmd/repram/main.go b/cmd/repram/main.go index 76625d6..92d580e 100644 --- a/cmd/repram/main.go +++ b/cmd/repram/main.go @@ -13,6 +13,7 @@ import ( "sort" "strconv" "strings" + "sync" "syscall" "time" @@ -33,6 +34,8 @@ import ( mcprpc "repram/internal/mcp" "repram/internal/node" "repram/internal/storage" + "repram/internal/transport/ws" + "repram/internal/tree" "repram/internal/trust" ) @@ -230,8 +233,90 @@ func main() { clusterNode.SetSeedProvider(func() []string { return seeds }) } + // Tree manager owns substrate-transient attachment state. Substrate + // nodes (REPRAM_INBOUND=true) accept inbound WS attachments and act as + // AckRouter + ChildBroadcaster for the cluster node. Transients + // (default, REPRAM_INBOUND=false) attach outbound after bootstrap. + inbound := tree.InboundFalse + if strings.EqualFold(os.Getenv("REPRAM_INBOUND"), "true") { + inbound = tree.InboundTrue + } + maxChildren := envInt("REPRAM_MAX_CHILDREN", tree.DefaultMaxChildren) + treeMgr := tree.NewManager( + &gossip.Node{ + ID: gossip.NodeID(nodeID), Address: address, + Port: gossipPort, HTTPPort: httpPort, Enclave: enclave, + }, + clusterPeerer{cn: clusterNode}, + tree.Options{ + Inbound: inbound, + MaxChildren: maxChildren, + ClusterSecret: clusterSecret, + }, + ) + clusterNode.SetAckRouter(treeMgr) + clusterNode.SetChildBroadcaster(treeMgr) + + // Wire reattach: when the parent connection rolls over to a new + // substrate, re-bind the gossip-message dispatch so incoming PUTs + // and ACKs from the new parent flow into the cluster handler. + bindParentDispatch := func(conn *ws.Connection) { + conn.OnMessage(func(m *gossip.Message) { + if err := clusterNode.HandleGossipMessage(m); err != nil { + logging.Debug("Parent WS dispatch: %v", err) + } + }) + } + treeMgr.SetReattachCallback(bindParentDispatch) + + // Transient bootstrap: if this node accepts no inbound, kick off a + // best-effort outbound WS attach to one of the seed substrates after + // HTTP bootstrap is done. Failure falls back to HTTP-only operation + // (writes still propagate via HTTP gossip; reads of other agents' + // writes don't reach this node until reattach succeeds). + if inbound == tree.InboundFalse && len(bootstrapNodes) > 0 { + go func(seeds []string) { + // Give the gossip bootstrap a moment to settle so the peer + // list reflects the actual cluster before we pick an attach + // target. 500ms is enough for the bootstrap response round-trip. + time.Sleep(500 * time.Millisecond) + for _, seed := range seeds { + idx := strings.LastIndex(seed, ":") + if idx <= 0 { + continue + } + host := seed[:idx] + port, err := strconv.Atoi(seed[idx+1:]) + if err != nil || port <= 0 { + continue + } + if host == address && port == httpPort { + continue + } + conn, err := ws.ConnectToSubstrate(ctx, host, port, clusterSecret, 10*time.Second) + if err != nil { + logging.Warn("WS attach to %s failed: %v — trying next seed", seed, err) + continue + } + if _, err := treeMgr.Attach(ctx, conn); err != nil { + logging.Warn("WS attach handshake to %s failed: %v", seed, err) + conn.Close(1000, "") + continue + } + bindParentDispatch(conn) + conn.StartHeartbeat() + logging.Info("Transient mode: attached to substrate at %s", seed) + return + } + logging.Warn("Transient mode: no seed accepted WS attach (degraded — HTTP gossip only)") + }(bootstrapNodes) + } + // Seed provider for tree-side reattach mirrors the cluster's recovery seeds. + treeMgr.SetSeedProvider(func() []string { return bootstrapNodes }) + server := &HTTPServer{ clusterNode: clusterNode, + treeManager: treeMgr, nodeID: nodeID, network: network, minTTL: minTTL, @@ -274,7 +359,13 @@ func main() { httpPort = tcpAddr.Port logging.Info(" HTTP listener bound to :%d", httpPort) } - httpServer := &http.Server{Handler: server.Router()} + // Outer mux routes /v1/ws directly (bypassing data-plane middleware) + // and delegates everything else to the gorilla router. http.NewServeMux + // longest-prefix match means /v1/ws is consumed here, "/" catches the rest. + outerMux := http.NewServeMux() + outerMux.HandleFunc("/v1/ws", server.wsHandler) + outerMux.Handle("/", server.Router()) + httpServer := &http.Server{Handler: outerMux} // Optional pprof server on a separate listener (diagnostic plane). // Uses http.DefaultServeMux which has pprof handlers auto-registered @@ -317,6 +408,7 @@ func main() { } securityMW.Close() + treeMgr.Stop() clusterNode.Stop() cancel() } @@ -442,6 +534,10 @@ type HTTPServer struct { maxTTL int startTime time.Time securityMW *node.SecurityMiddleware + // treeManager owns substrate-transient attachment state. Always non-nil + // — the constructor wires one up regardless of inbound capability so + // transients can also call Attach when they have a substrate peer. + treeManager *tree.Manager } func (s *HTTPServer) Router() *mux.Router { @@ -475,6 +571,10 @@ func (s *HTTPServer) Router() *mux.Router { // Internal gossip endpoints r.HandleFunc("/v1/gossip/message", s.gossipHandler).Methods("POST", "OPTIONS") r.HandleFunc("/v1/bootstrap", s.bootstrapHandler).Methods("POST", "OPTIONS") + // /v1/ws is intentionally NOT registered here. The WebSocket upgrade + // needs Hijack() and a long-lived connection, both of which fight the + // http.TimeoutHandler + MaxRequestSize wrappers above. It's wired + // directly on the outer mux below. r.NotFoundHandler = http.HandlerFunc(s.notFoundHandler) @@ -547,12 +647,34 @@ func (s *HTTPServer) topologyHandler(w http.ResponseWriter, r *http.Request) { }) } - w.Header().Set("Content-Type", "application/json") - json.NewEncoder(w).Encode(map[string]interface{}{ + // Attached children (transients) — visible only on substrate nodes + // that have accepted WS attachments. + type childInfo struct { + ID string `json:"id"` + Enclave string `json:"enclave"` + } + var children []childInfo + if s.treeManager != nil { + for id, conn := range s.treeManager.Children() { + children = append(children, childInfo{ID: id, Enclave: conn.RemoteEnclave()}) + } + } + + resp := map[string]interface{}{ "node_id": s.nodeID, "enclave": s.clusterNode.Enclave(), "peers": peerList, - }) + } + if s.treeManager != nil { + resp["role"] = string(s.treeManager.Role()) + resp["children"] = children + if parent := s.treeManager.Parent(); parent != nil { + resp["parent_id"] = parent.RemoteNodeID() + } + } + + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) } func (s *HTTPServer) putHandler(w http.ResponseWriter, r *http.Request) { @@ -786,3 +908,86 @@ func (s *HTTPServer) bootstrapHandler(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") json.NewEncoder(w).Encode(resp) } + +// clusterPeerer adapts *cluster.ClusterNode to the tree.Peerer interface. +// ClusterNode.Topology already returns the full peer list with enclave +// metadata; this just renames the method to match what tree wants. +type clusterPeerer struct{ cn *cluster.ClusterNode } + +func (c clusterPeerer) GetPeers() []*gossip.Node { return c.cn.Topology() } + +// wsHandler accepts an incoming substrate-transient WebSocket attachment. +// The first non-control frame must be a hello; subsequent gossip-typed +// frames are dispatched to clusterNode.HandleGossipMessage, with PUTs +// recording an ACK route so the substrate can forward the enclave-peer +// ACKs back through the WS pipe to the originating child. +// +// All routing decisions are driven by treeManager.HandleHello — if the +// substrate is at capacity or attachments are disabled, the manager sends +// a goodbye-with-alternatives and closes the connection itself. +func (s *HTTPServer) wsHandler(w http.ResponseWriter, r *http.Request) { + if s.treeManager == nil || !s.treeManager.IsInboundCapable() { + // Transient nodes don't accept inbound; refuse with 404 to + // avoid leaking the role to scanners. + http.NotFound(w, r) + return + } + wsHandler := ws.Handler(s.clusterNode.ClusterSecret(), nil, func(conn *ws.Connection) { + s.bindWSConnection(conn) + }) + wsHandler.ServeHTTP(w, r) +} + +// bindWSConnection sets up the gossip dispatch + ACK-route recording on a +// freshly accepted child connection. Called from ws.Handler's onAccept. +func (s *HTTPServer) bindWSConnection(conn *ws.Connection) { + // One-shot hello handler: install the gossip dispatch only after a + // valid hello arrives. Until then ignore everything (matches the TS + // reference's gate at handleUpgrade attachment handler). + helloDone := make(chan struct{}) + var helloOnce sync.Once + + removeHello := conn.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + if msg.Type != ws.AttachmentTypeHello { + return + } + var h ws.HelloPayload + if err := json.Unmarshal(msg.Payload, &h); err != nil { + logging.Warn("WS attach: hello decode failed: %v", err) + conn.Close(1003, "bad hello") + return + } + helloOnce.Do(func() { close(helloDone) }) + if !s.treeManager.HandleHello(conn, &h) { + // HandleHello already sent goodbye-with-alternatives and is + // scheduling close. Bail. + return + } + // Dispatch any subsequent gossip frame into the cluster handler. + // Recording the ACK route happens for PUTs so the substrate can + // reverse-route ACKs back through this pipe. + conn.OnMessage(func(gmsg *gossip.Message) { + if gmsg.Type == gossip.MessageTypePut { + s.treeManager.RecordAckRoute(gmsg.MessageID, conn, s.clusterNode.WriteTimeout()) + } + if err := s.clusterNode.HandleGossipMessage(gmsg); err != nil { + logging.Debug("WS gossip handler: %v", err) + } + }) + }) + + // If hello never arrives within 30s, close. Mirrors the TS reference's + // silent-attachment ceiling. + go func() { + select { + case <-helloDone: + removeHello() + case <-time.After(30 * time.Second): + if !conn.IsClosed() { + logging.Warn("WS attach: no hello within 30s, closing") + conn.Close(1002, "no hello") + } + removeHello() + } + }() +} diff --git a/cmd/repram/ws_integration_test.go b/cmd/repram/ws_integration_test.go new file mode 100644 index 0000000..01f73fe --- /dev/null +++ b/cmd/repram/ws_integration_test.go @@ -0,0 +1,341 @@ +package main + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + gws "github.com/gorilla/websocket" + + "repram/internal/cluster" + "repram/internal/gossip" + "repram/internal/node" + "repram/internal/transport/ws" + "repram/internal/tree" +) + +// newSubstrateTestServer spins up an HTTPServer running as a substrate +// (inbound=true) with a single-node cluster behind it. Returns the test +// httptest.Server, the cluster node (for direct store inspection), the +// tree manager (for state assertions), and a cleanup func. +func newSubstrateTestServer(t *testing.T) (*httptest.Server, *cluster.ClusterNode, *tree.Manager, func()) { + t.Helper() + cn := cluster.NewClusterNode( + "substrate-1", "127.0.0.1", 0, 0, + 1, 0, 5*time.Second, "", "default", + ) + ctx, cancel := context.WithCancel(context.Background()) + if err := cn.Start(ctx, nil); err != nil { + t.Fatalf("cluster start: %v", err) + } + tm := tree.NewManager( + &gossip.Node{ID: "substrate-1", Address: "127.0.0.1", Port: 0, HTTPPort: 0, Enclave: "default"}, + clusterPeerer{cn: cn}, + tree.Options{Inbound: tree.InboundTrue, MaxChildren: tree.DefaultMaxChildren}, + ) + cn.SetAckRouter(tm) + cn.SetChildBroadcaster(tm) + + securityMW := node.NewSecurityMiddleware(1000, 2000, 10*1024*1024, false) + server := &HTTPServer{ + clusterNode: cn, + treeManager: tm, + nodeID: "substrate-1", + network: "private", + minTTL: 300, + maxTTL: 86400, + startTime: time.Now(), + securityMW: securityMW, + } + + outerMux := http.NewServeMux() + outerMux.HandleFunc("/v1/ws", server.wsHandler) + outerMux.Handle("/", server.Router()) + srv := httptest.NewServer(outerMux) + + cleanup := func() { + srv.Close() + securityMW.Close() + tm.Stop() + cn.Stop() + cancel() + } + return srv, cn, tm, cleanup +} + +// dialWS opens a raw WS to the given httptest server's /v1/ws. +func dialWS(t *testing.T, srv *httptest.Server) *ws.Connection { + t.Helper() + u, _ := url.Parse(srv.URL) + wsURL := "ws://" + u.Host + "/v1/ws" + d := &gws.Dialer{HandshakeTimeout: 2 * time.Second} + raw, resp, err := d.Dial(wsURL, nil) + if err != nil { + if resp != nil { + _ = resp.Body.Close() + } + t.Fatalf("dial /v1/ws: %v", err) + } + return ws.NewConnection(raw, "") +} + +// TestWSAttachHandshake exercises the hello/welcome handshake against the +// real HTTPServer wiring — verifies that /v1/ws is routed, the tree +// manager handles hello, and welcome reaches the client. +func TestWSAttachHandshake(t *testing.T) { + srv, _, tm, cleanup := newSubstrateTestServer(t) + defer cleanup() + + client := dialWS(t, srv) + defer client.Close(gws.CloseNormalClosure, "") + + welcomeCh := make(chan *ws.WelcomePayload, 1) + client.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + if msg.Type != ws.AttachmentTypeWelcome { + return + } + var w ws.WelcomePayload + if err := json.Unmarshal(msg.Payload, &w); err == nil { + welcomeCh <- &w + } + }) + + hello := ws.HelloPayload{ + NodeID: "transient-A", Enclave: "default", + Address: "10.0.0.1", HTTPPort: 0, + Capabilities: ws.Capabilities{Inbound: "false"}, + } + if err := client.SendAttachment(ws.AttachmentTypeHello, hello); err != nil { + t.Fatalf("send hello: %v", err) + } + + select { + case w := <-welcomeCh: + if w.YourPosition.ParentID != "substrate-1" { + t.Errorf("parent_id: %q", w.YourPosition.ParentID) + } + case <-time.After(2 * time.Second): + t.Fatal("no welcome received") + } + + // Wait briefly for child registration (HandleHello may complete after + // the welcome was sent — both happen in the same goroutine; in practice + // child is registered before SendAttachment returns, but be defensive). + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if tm.HasChild("transient-A") { + break + } + time.Sleep(10 * time.Millisecond) + } + if !tm.HasChild("transient-A") { + t.Errorf("tm.HasChild(transient-A) = false; child not registered") + } +} + +// TestWSPutStoresLocallyAndAcks exercises the relay round-trip end-to-end: +// - Client sends PUT over WS +// - Substrate stores it (relayed-PUT enclave-peer fanout is exercised +// elsewhere via the existing HTTP gossip tests; we verify the local +// store side here) +// - Substrate sends an ACK over WS to the originator +// +// quorum=1 because the substrate is a single-node cluster — its local +// store IS the quorum. +func TestWSPutStoresLocallyAndAcks(t *testing.T) { + srv, cn, _, cleanup := newSubstrateTestServer(t) + defer cleanup() + + client := dialWS(t, srv) + defer client.Close(gws.CloseNormalClosure, "") + + ackCh := make(chan *gossip.Message, 1) + client.OnMessage(func(m *gossip.Message) { + if m.Type == gossip.MessageTypeAck { + ackCh <- m + } + }) + + // Hello + wait for welcome. + helloDone := make(chan struct{}, 1) + client.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + if msg.Type == ws.AttachmentTypeWelcome { + select { + case helloDone <- struct{}{}: + default: + } + } + }) + if err := client.SendAttachment(ws.AttachmentTypeHello, ws.HelloPayload{ + NodeID: "transient-A", Enclave: "default", + Capabilities: ws.Capabilities{Inbound: "false"}, + }); err != nil { + t.Fatalf("hello: %v", err) + } + <-helloDone + + // PUT over WS. + put := &gossip.Message{ + Type: gossip.MessageTypePut, + From: "transient-A", + Key: "relayed-key", + Data: []byte("relayed-value"), + TTL: 300, + Timestamp: time.Now(), + MessageID: "ws-put-1", + } + if err := client.SendGossip(put); err != nil { + t.Fatalf("send put: %v", err) + } + + // Substrate should ACK back over WS. + select { + case ack := <-ackCh: + if ack.MessageID != "ws-put-1" { + t.Errorf("ack messageId: %q want ws-put-1", ack.MessageID) + } + if ack.Key != "relayed-key" { + t.Errorf("ack key: %q", ack.Key) + } + case <-time.After(2 * time.Second): + t.Fatal("no ACK received over WS") + } + + // And the substrate's local store must contain the data — confirming + // the relayed write hit the existing PUT handler path. + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) { + if data, ok := cn.Get("relayed-key"); ok && string(data) == "relayed-value" { + return + } + time.Sleep(10 * time.Millisecond) + } + if data, ok := cn.Get("relayed-key"); !ok || string(data) != "relayed-value" { + t.Errorf("substrate store: data=%q ok=%v", data, ok) + } +} + +// TestWSReceivePathFanout exercises phase 5: when the substrate stores a +// PUT (e.g., received via HTTP gossip from an enclave peer), it fans the +// replica out over WS to attached transients. Here we drive +// HandleGossipMessage directly (simulating an HTTP arrival) and verify +// the attached WS client receives the PUT. +func TestWSReceivePathFanout(t *testing.T) { + srv, cn, _, cleanup := newSubstrateTestServer(t) + defer cleanup() + + client := dialWS(t, srv) + defer client.Close(gws.CloseNormalClosure, "") + + gotPut := make(chan *gossip.Message, 1) + client.OnMessage(func(m *gossip.Message) { + if m.Type == gossip.MessageTypePut { + gotPut <- m + } + }) + + welcomeCh := make(chan struct{}, 1) + client.AddAttachmentHandler(func(msg *ws.AttachmentMessage) { + if msg.Type == ws.AttachmentTypeWelcome { + select { + case welcomeCh <- struct{}{}: + default: + } + } + }) + if err := client.SendAttachment(ws.AttachmentTypeHello, ws.HelloPayload{ + NodeID: "transient-B", Enclave: "default", + Capabilities: ws.Capabilities{Inbound: "false"}, + }); err != nil { + t.Fatalf("hello: %v", err) + } + <-welcomeCh + + // Simulate a PUT arriving from another enclave peer via HTTP gossip. + msg := &gossip.Message{ + Type: gossip.MessageTypePut, + From: "some-other-peer", + Key: "other-agents-key", + Data: []byte("other-agents-value"), + TTL: 300, + Timestamp: time.Now(), + MessageID: "http-put-1", + } + if err := cn.HandleGossipMessage(msg); err != nil { + t.Fatalf("HandleGossipMessage: %v", err) + } + + select { + case r := <-gotPut: + if r.MessageID != "http-put-1" { + t.Errorf("MessageID: %q", r.MessageID) + } + if r.Key != "other-agents-key" { + t.Errorf("Key: %q", r.Key) + } + if string(r.Data) != "other-agents-value" { + t.Errorf("Data: %q", r.Data) + } + case <-time.After(2 * time.Second): + t.Fatal("attached transient never received fan-out") + } +} + +// TestWSRejectIfTransient verifies that nodes with inbound=false return +// 404 on /v1/ws — transients shouldn't expose the endpoint at all. +func TestWSRejectIfTransient(t *testing.T) { + cn := cluster.NewClusterNode("transient-1", "127.0.0.1", 0, 0, 1, 0, 5*time.Second, "", "default") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if err := cn.Start(ctx, nil); err != nil { + t.Fatalf("cluster: %v", err) + } + defer cn.Stop() + + tm := tree.NewManager( + &gossip.Node{ID: "transient-1", Address: "127.0.0.1", Enclave: "default"}, + clusterPeerer{cn: cn}, + tree.Options{Inbound: tree.InboundFalse, MaxChildren: 0}, + ) + defer tm.Stop() + + securityMW := node.NewSecurityMiddleware(1000, 2000, 10*1024*1024, false) + defer securityMW.Close() + + server := &HTTPServer{ + clusterNode: cn, + treeManager: tm, + nodeID: "transient-1", + network: "private", + minTTL: 300, + maxTTL: 86400, + startTime: time.Now(), + securityMW: securityMW, + } + + outerMux := http.NewServeMux() + outerMux.HandleFunc("/v1/ws", server.wsHandler) + outerMux.Handle("/", server.Router()) + srv := httptest.NewServer(outerMux) + defer srv.Close() + + u, _ := url.Parse(srv.URL) + wsURL := "ws://" + u.Host + "/v1/ws" + _, resp, err := (&gws.Dialer{HandshakeTimeout: 2 * time.Second}).Dial(wsURL, nil) + if err == nil { + t.Fatal("expected dial to fail on transient /v1/ws") + } + if resp != nil { + defer resp.Body.Close() + if resp.StatusCode != http.StatusNotFound { + t.Errorf("status: %d want 404", resp.StatusCode) + } + } else if !strings.Contains(err.Error(), "404") { + t.Errorf("err: %v want 404-ish", err) + } +} diff --git a/internal/cluster/node.go b/internal/cluster/node.go index dbecd20..9322b9c 100644 --- a/internal/cluster/node.go +++ b/internal/cluster/node.go @@ -41,6 +41,31 @@ type ClusterNode struct { // list. Nil disables recovery (used for tests and for cases where // re-bootstrap doesn't make sense). seedProvider func() []string + + // ackRouter, when set, receives ACKs whose To= field doesn't match a + // known HTTP peer. Substrate nodes use this to forward ACKs back to + // transient children that originated a relayed PUT (#135). + ackRouter AckRouter + + // childBroadcaster, when set, receives every successfully-stored PUT + // so it can fan the replica out to attached transients (#135). Nil + // outside substrate mode. + childBroadcaster ChildBroadcaster +} + +// AckRouter forwards an ACK to a non-peer originator. The substrate's +// tree manager satisfies this interface — its routing table maps each +// relayed messageId to the originating child connection. Returns true if +// the ACK was routed (caller skips other dispatch paths). +type AckRouter interface { + RouteAck(ack *gossip.Message) bool +} + +// ChildBroadcaster fans a gossip message out to attached transient +// children whose enclave matches the local node. The substrate's tree +// manager satisfies this — see tree.Manager.BroadcastToChildren. +type ChildBroadcaster interface { + BroadcastToChildren(msg *gossip.Message) } // IsolationRecoveryInterval is how often the recovery loop checks for @@ -123,6 +148,22 @@ func (cn *ClusterNode) SetSeedProvider(p func() []string) { cn.seedProvider = p } +// SetAckRouter installs the AckRouter used to deliver ACKs back to +// non-peer originators (i.e., transient children attached via WebSocket). +// Substrate nodes wire their tree.Manager here; transient nodes leave it +// nil since their originator is themselves and ACKs come through the +// pendingWrites path. +func (cn *ClusterNode) SetAckRouter(r AckRouter) { + cn.ackRouter = r +} + +// SetChildBroadcaster installs the ChildBroadcaster used to fan stored +// PUTs out to attached transient children. Nil disables WS receive-path +// fan-out (#135 phase 5). +func (cn *ClusterNode) SetChildBroadcaster(b ChildBroadcaster) { + cn.childBroadcaster = b +} + // runIsolationRecovery polls peer count on IsolationRecoveryInterval // and triggers re-bootstrap when the node is fully isolated (#85, F5). // @@ -324,36 +365,62 @@ func (cn *ClusterNode) handlePutMessage(msg *gossip.Message) error { peers := cn.protocol.GetPeers() cn.writesMutex.RUnlock() + delivered := false for _, peer := range peers { if peer.ID == msg.From { logging.Debug("[%s] Sending ACK for key %s to %s", cn.localNode.ID, msg.Key, peer.ID) cn.protocol.Send(context.Background(), peer, ack) + delivered = true break } } + // If the originator isn't in the HTTP peer list, it may be a transient + // child attached via WS. The substrate's AckRouter looks up the + // originating connection by messageId and writes the ACK to that pipe. + if !delivered && cn.ackRouter != nil { + if cn.ackRouter.RouteAck(ack) { + logging.Debug("[%s] Routed ACK for key %s back through WS attachment", cn.localNode.ID, msg.Key) + } + } // Continue epidemic forwarding to other enclave peers cn.protocol.ForwardToEnclave(context.Background(), msg) + // Fan the stored replica out to attached transient children so they + // see other agents' writes in their local store. The broadcaster + // gates on enclave; cross-enclave traffic is dropped. + if cn.childBroadcaster != nil { + cn.childBroadcaster.BroadcastToChildren(msg) + } + return nil } func (cn *ClusterNode) handleAckMessage(msg *gossip.Message) error { cn.writesMutex.Lock() - defer cn.writesMutex.Unlock() - writeOp, exists := cn.pendingWrites[msg.MessageID] - if !exists { + cn.writesMutex.Unlock() + + if exists { + cn.writesMutex.Lock() + writeOp.Confirmations++ + reached := writeOp.Confirmations >= cn.quorumSize() + cn.writesMutex.Unlock() + if reached { + select { + case writeOp.Complete <- true: + default: + } + } return nil } - writeOp.Confirmations++ - - if writeOp.Confirmations >= cn.quorumSize() { - select { - case writeOp.Complete <- true: - default: - } + // Not our pending write — possibly an ACK for a PUT we relayed on + // behalf of a transient child. The substrate's AckRouter table maps + // messageId → originating WS pipe; if we relayed this message, forward + // the ACK upstream so the transient's quorum tally advances. + if cn.ackRouter != nil { + cn.ackRouter.RouteAck(msg) } return nil @@ -393,6 +460,13 @@ func (cn *ClusterNode) Enclave() string { return cn.localNode.Enclave } +// WriteTimeout returns the configured quorum write timeout. Used by +// WS-attached substrates to size the ACK-route eviction window so the +// route is preserved at least as long as the originator waits for ACKs. +func (cn *ClusterNode) WriteTimeout() time.Duration { + return cn.writeTimeout +} + // Topology returns the full peer list with enclave membership. func (cn *ClusterNode) Topology() []*gossip.Node { return cn.protocol.GetPeers() diff --git a/internal/tree/manager.go b/internal/tree/manager.go index 9314cb7..1c3f673 100644 --- a/internal/tree/manager.go +++ b/internal/tree/manager.go @@ -590,6 +590,26 @@ func (m *Manager) LookupAckRoute(messageID string) *ws.Connection { return m.ackRoutes[messageID] } +// RouteAck satisfies cluster.AckRouter. If ack.MessageID is in the routing +// table, the ack is written to that child's WS pipe and the route is cleared +// (an originator gets at most one ACK from this substrate per write). Returns +// true if the ACK was routed. +func (m *Manager) RouteAck(ack *gossip.Message) bool { + conn := m.LookupAckRoute(ack.MessageID) + if conn == nil || conn.IsClosed() { + return false + } + if err := conn.SendGossip(ack); err != nil { + logging.Debug("RouteAck: send to %s failed: %v", conn.RemoteNodeID(), err) + return false + } + // Don't ClearAckRoute on success — the substrate's own ACK is the first + // of potentially several, and the originating transient may still need + // later ACKs routed (e.g., for quorum > 1). The route auto-evicts after + // RecordAckRoute's TTL. + return true +} + // ClearAckRoute removes the ACK route for messageID and stops its eviction // timer. Called by the cluster layer when the originator has accumulated // enough ACKs for quorum, or when the message times out from the write side. From 80181ff6edbd9c5392be463467784314615706fb Mon Sep 17 00:00:00 2001 From: TickTockBent Date: Tue, 12 May 2026 09:19:25 -0400 Subject: [PATCH 5/5] review (#135 / PR #136): address cold-review findings Addresses the four "important" items from the cold sonnet review: 1. handleAckMessage closed-channel panic (internal/cluster/node.go:399-426). Splitting writesMutex to avoid deadlocking on RouteAck's WS write opened a window where two ACKs could both observe exists=true, one closes Complete, the second tries to send on a closed channel, panic. Fix: WriteOperation grows a sync.Once-guarded markComplete() helper; both the local-quorum-met path and the gossip-ACK path call markComplete() instead of close()/buffered-send. The receive side already tolerated either signal style. 2. Post-welcome dispatch race (internal/tree/manager.go, cmd/repram/main.go). The reattach callback wired OnMessage AFTER Attach returned, so a substrate's first PUT after welcome could land on a not-yet-wired onMessage and be silently dropped. Same class of race the earlier attach()-handler-install fix closed, just one level up. Fix: tree.Manager grows SetParentDispatch which is installed inside attach() BEFORE SendAttachment(hello), so the WS readLoop's serial processing guarantees OnMessage is ready for any post-welcome gossip frame. The SetReattachCallback hook stays around for non-dispatch wiring (heartbeat start, metrics scopes); main.go drops the redundant bindParentDispatch closure. 3. Enclave bypass via empty hello.Enclave (internal/tree/manager.go). BroadcastToChildren's filter skipped a child only when its enclave was non-empty AND mismatched. A hello with Enclave="" slipped through and received cross-enclave traffic. Fix: HandleHello normalizes empty to "default" (matching the gossip layer's normalization of peer enclaves), and the BroadcastToChildren filter tightens to strict inequality. New TestEmptyEnclaveNormalizedOnHello regression test exercises a non-default substrate + empty-enclave hello and asserts the broadcast is dropped. 4. Missing CHANGELOG entry (spec DoD requirement). Adds a "Restored" section under [Unreleased] describing the recovery from #125 and the four review fixes folded into this PR. Full repo passes go test -race ./... -count=5 with no new flakes. // ticktockbent --- CHANGELOG.md | 12 ++++++++ cmd/repram/main.go | 23 +++++++-------- internal/cluster/node.go | 26 +++++++++++------ internal/tree/manager.go | 54 +++++++++++++++++++++++++++++++---- internal/tree/manager_test.go | 49 +++++++++++++++++++++++++++++++ 5 files changed, 138 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9005b2d..119a04b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,18 @@ All notable changes to REPRAM are documented here. ## [Unreleased] +### Restored — Substrate/transient WS tree in Go ([#135](https://github.com/TickTockBent/repram/issues/135), recovers the gap from [#133](https://github.com/TickTockBent/repram/issues/133)) +The substrate/transient tree topology that #125 inadvertently removed when it deleted the TypeScript node is now ported into the Go binary. `repram --mcp` (and any node with `REPRAM_INBOUND=false`) is once again a real cluster participant: it attaches to a substrate via persistent outbound WebSocket, sees other agents' writes in its local store, and survives substrate failure via cached alternatives. + +- New `internal/transport/ws` package — WS `Connection` with heartbeat, optional HMAC, multi-subscriber lifecycle handlers; outbound `ConnectToSubstrate` dialer; inbound `Handler` upgrader. Wire format on WS payloads is identical to HTTP gossip — same JSON shape, same handlers process both transports. +- New `internal/tree` package — `Manager` owns substrate (`HandleHello` + child registration + welcome-with-topology) and transient (`Attach` + cached-topology / seed-list reattach loop) lifecycle. Self-skip is enforced literally on address+http\_port to prevent the [#120](https://github.com/TickTockBent/repram/issues/120) regression. Stop-aware context threading and a goroutine WaitGroup ensure clean shutdown. +- `cluster.ClusterNode` gets `AckRouter` and `ChildBroadcaster` interfaces (satisfied by the tree manager). `handlePutMessage` routes the substrate's own ACK back through the WS pipe to the originating transient and broadcasts replicas out to enclave-matched children; `handleAckMessage` forwards enclave-peer ACKs upstream when they're for a relayed write. Quorum-complete signaling is guarded by `sync.Once` so concurrent ACK paths can't panic on close-of-closed-channel. +- `cmd/repram/main.go` mounts `/v1/ws` on an outer `http.ServeMux` that bypasses the data-plane TimeoutHandler + MaxRequestSize wrappers. Substrates accept attachments; transients 404 the endpoint to avoid leaking role to scanners. `--mcp` mode auto-attaches outbound after HTTP bootstrap and falls back to HTTP-only on attach failure. Parent-side gossip dispatch is installed inside `Attach` before the function returns so a fast substrate's first PUT after welcome cannot be silently dropped. `/v1/topology` now exposes `role`, attached `children`, and `parent_id`. +- Enclave isolation hardening: `HandleHello` normalizes empty hello-enclave to `"default"` so the `BroadcastToChildren` filter cannot be bypassed by an underspecified hello. +- 50+ new tests: 25 WS transport, 22 tree manager (including the [#120](https://github.com/TickTockBent/repram/issues/120) self-skip timing assertion and an empty-enclave isolation regression), 4 HTTP-server-level WS integration tests. Full repo passes `go test -race ./...` with no regressions in existing 118+ tests. + +Phase 7 of #135 — a 24h+ burn-in 2.2 on a real multi-substrate + multi-transient cluster — is tracked as a follow-up; it does not gate this PR. + ### Changed — Go-native MCP server ([#123](https://github.com/TickTockBent/repram/issues/123)) The Go binary now serves MCP directly via `repram --mcp`: an embedded node, in-process tool handlers, and JSON-RPC 2.0 on stdin/stdout. The TypeScript node (`repram-mcp/`) has been removed. diff --git a/cmd/repram/main.go b/cmd/repram/main.go index 92d580e..18c2930 100644 --- a/cmd/repram/main.go +++ b/cmd/repram/main.go @@ -257,17 +257,15 @@ func main() { clusterNode.SetAckRouter(treeMgr) clusterNode.SetChildBroadcaster(treeMgr) - // Wire reattach: when the parent connection rolls over to a new - // substrate, re-bind the gossip-message dispatch so incoming PUTs - // and ACKs from the new parent flow into the cluster handler. - bindParentDispatch := func(conn *ws.Connection) { - conn.OnMessage(func(m *gossip.Message) { - if err := clusterNode.HandleGossipMessage(m); err != nil { - logging.Debug("Parent WS dispatch: %v", err) - } - }) - } - treeMgr.SetReattachCallback(bindParentDispatch) + // Wire the parent-side gossip dispatch on the tree manager. Attach() + // and every successful reattach install this handler on the new + // connection BEFORE the function returns — closes the post-welcome + // dispatch race the SetReattachCallback hook used to leave open. + treeMgr.SetParentDispatch(func(m *gossip.Message) { + if err := clusterNode.HandleGossipMessage(m); err != nil { + logging.Debug("Parent WS dispatch: %v", err) + } + }) // Transient bootstrap: if this node accepts no inbound, kick off a // best-effort outbound WS attach to one of the seed substrates after @@ -303,7 +301,8 @@ func main() { conn.Close(1000, "") continue } - bindParentDispatch(conn) + // parent dispatch was installed by treeMgr.Attach via + // SetParentDispatch; nothing extra to wire here. conn.StartHeartbeat() logging.Info("Transient mode: attached to substrate at %s", seed) return diff --git a/internal/cluster/node.go b/internal/cluster/node.go index 9322b9c..1a94657 100644 --- a/internal/cluster/node.go +++ b/internal/cluster/node.go @@ -74,14 +74,27 @@ type ChildBroadcaster interface { const IsolationRecoveryInterval = 30 * time.Second type WriteOperation struct { - Key string - Data []byte - TTL time.Duration + Key string + Data []byte + TTL time.Duration Confirmations int + // signalComplete is guarded by signalOnce so the local-quorum path, + // the gossip-ACK path, and any racing late ACK can all signal once + // without panicking on close-of-closed-channel. Pendant of the + // MessageID dedup that keys pendingWrites: there can still be more + // than one goroutine that observes "quorum reached" within a single + // write's lifetime. Complete chan bool + signalOnce sync.Once Error error } +// markComplete signals the write as quorum-reached at most once. Safe +// to call from any goroutine that observes the quorum threshold. +func (w *WriteOperation) markComplete() { + w.signalOnce.Do(func() { close(w.Complete) }) +} + type Store interface { Put(key string, data []byte, ttl time.Duration) error Get(key string) ([]byte, bool) @@ -272,7 +285,7 @@ func (cn *ClusterNode) Put(ctx context.Context, key string, data []byte, ttl tim cn.writesMutex.Lock() delete(cn.pendingWrites, msg.MessageID) cn.writesMutex.Unlock() - close(writeOp.Complete) + writeOp.markComplete() logging.Debug("Write completed locally (quorum=%d, confirmations=%d)", quorum, writeOp.Confirmations) return nil } @@ -407,10 +420,7 @@ func (cn *ClusterNode) handleAckMessage(msg *gossip.Message) error { reached := writeOp.Confirmations >= cn.quorumSize() cn.writesMutex.Unlock() if reached { - select { - case writeOp.Complete <- true: - default: - } + writeOp.markComplete() } return nil } diff --git a/internal/tree/manager.go b/internal/tree/manager.go index 1c3f673..a669832 100644 --- a/internal/tree/manager.go +++ b/internal/tree/manager.go @@ -118,6 +118,13 @@ type Manager struct { onReattach func(*ws.Connection) seedProvider func() []string + // parentDispatch is the application-level gossip handler installed on + // every parent connection (initial Attach and every successful reattach) + // before Attach returns. Setting this up-front closes the window where + // the substrate could push a PUT between welcome arrival and the caller + // wiring its own OnMessage. + parentDispatch func(*gossip.Message) + // ackRoutes tracks (messageId → child connection) for PUTs relayed // through this substrate. When an enclave peer ACKs a relayed PUT, the // substrate forwards the ACK back through ackRoutes[messageId] so the @@ -213,20 +220,41 @@ func (m *Manager) SetSeedProvider(fn func() []string) { } // SetReattachCallback registers a hook fired after a successful reattach. -// The application uses this to rewire its parent-message router to the new -// connection. +// The application uses this to rewire any per-connection state that isn't +// already handled by SetParentDispatch (e.g., metrics labels, logging +// scopes). Parent message dispatch itself is wired by SetParentDispatch. func (m *Manager) SetReattachCallback(fn func(*ws.Connection)) { m.mu.Lock() m.onReattach = fn m.mu.Unlock() } +// SetParentDispatch registers the gossip-message handler that Attach and +// every successful reattach install on the parent connection — wired +// inside attach() before the function returns, so the substrate's first +// post-welcome PUT can't be silently dropped by a not-yet-wired OnMessage. +// Pass nil to disable parent-side dispatch (transient still receives via +// HTTP gossip; only WS-tree fan-out is muted). +func (m *Manager) SetParentDispatch(fn func(*gossip.Message)) { + m.mu.Lock() + m.parentDispatch = fn + m.mu.Unlock() +} + // HandleHello processes an incoming hello frame on a freshly-accepted child // connection. On accept it sends a welcome with the current peer topology and // registers a close handler that removes the child from the map. On reject // (capacity, attachments disabled) it sends a goodbye-with-alternatives and // closes the connection shortly afterward. func (m *Manager) HandleHello(conn *ws.Connection, hello *ws.HelloPayload) bool { + // Normalize empty enclave to "default" so the BroadcastToChildren + // enclave filter cannot be bypassed by a hello that omits the field + // (the filter skips children whose RemoteEnclave doesn't match the + // local one; "" used to slip through as "any"). Matches the same + // normalization the gossip layer applies to peer enclaves. + if hello.Enclave == "" { + hello.Enclave = "default" + } m.mu.Lock() if m.opts.MaxChildren == 0 { m.mu.Unlock() @@ -340,6 +368,18 @@ func (m *Manager) attach(ctx context.Context, conn *ws.Connection, timeout time. } }) + // Install the parent-dispatch handler before sending hello. The WS + // readLoop is single-threaded and processes frames in order, so once + // welcome is delivered every subsequent gossip frame is guaranteed + // to find OnMessage already wired. Without this, a fast substrate can + // push a PUT into a not-yet-wired handler and the frame is dropped. + m.mu.Lock() + dispatch := m.parentDispatch + m.mu.Unlock() + if dispatch != nil { + conn.OnMessage(dispatch) + } + if err := conn.SendAttachment(ws.AttachmentTypeHello, hello); err != nil { removeAttach() removeClose() @@ -365,8 +405,8 @@ func (m *Manager) attach(ctx context.Context, conn *ws.Connection, timeout time. removeAttach() removeClose() - // Promote to active parent. The role is "transient" by definition when - // we reach here — inbound-capable nodes never call Attach. + // parentDispatch was installed before SendAttachment to close the + // post-welcome dispatch race; nothing to wire here. m.mu.Lock() m.parent = conn m.role = RoleTransient @@ -643,8 +683,10 @@ func (m *Manager) BroadcastToChildren(msg *gossip.Message) { for _, c := range conns { // Enclave gating: substrate must not leak cross-enclave writes to // transients. The child's enclave is recorded on the Connection by - // SetRemote during HandleHello. - if c.RemoteEnclave() != "" && c.RemoteEnclave() != m.local.Enclave { + // SetRemote during HandleHello, where empty hello.Enclave is + // normalized to "default" — so strict inequality is safe (no + // "" → bypass) and any child whose enclave differs is dropped. + if c.RemoteEnclave() != m.local.Enclave { continue } if err := c.SendGossip(msg); err != nil { diff --git a/internal/tree/manager_test.go b/internal/tree/manager_test.go index 5a1370e..12d77ab 100644 --- a/internal/tree/manager_test.go +++ b/internal/tree/manager_test.go @@ -841,6 +841,55 @@ func TestBroadcastToChildren(t *testing.T) { } } +// TestEmptyEnclaveNormalizedOnHello — review finding #4: a hello with an +// empty Enclave field used to slip past the BroadcastToChildren filter +// (which only skipped when enclave was non-empty AND mismatched). After +// normalization the substrate runs in "default" and the child registers +// as "default", so the filter still admits in-enclave traffic but a +// substrate in a NON-default enclave will correctly skip the empty-enclave +// child. +func TestEmptyEnclaveNormalizedOnHello(t *testing.T) { + p := newPair(t) + defer p.cleanupFn() + srv := p.serverConn(t) + + // Substrate in a non-default enclave. + m := NewManager( + makeNode("substrate-1", func(n *gossip.Node) { n.Enclave = "alpha" }), + &fakePeerer{}, + Options{Inbound: InboundTrue, MaxChildren: DefaultMaxChildren}, + ) + defer m.Stop() + + // Hello with empty enclave (the bypass case). + if !m.HandleHello(srv, &ws.HelloPayload{NodeID: "t1", Enclave: ""}) { + t.Fatal("HandleHello rejected hello with empty enclave") + } + if got := srv.RemoteEnclave(); got != "default" { + t.Errorf("RemoteEnclave: got %q want default (normalization)", got) + } + + // Substrate enclave is "alpha"; the child normalized to "default" — + // the filter must drop the broadcast. + var calls atomic.Int32 + p.clientWS.OnMessage(func(*gossip.Message) { calls.Add(1) }) + + msg := &gossip.Message{ + Type: gossip.MessageTypePut, + From: "alpha-peer", + Key: "k", + Data: []byte("v"), + TTL: 60, + Timestamp: time.Now(), + MessageID: "empty-enclave", + } + m.BroadcastToChildren(msg) + time.Sleep(100 * time.Millisecond) + if n := calls.Load(); n != 0 { + t.Errorf("empty-enclave child received broadcast in non-default substrate: %d times", n) + } +} + func TestBroadcastToChildrenSkipsOtherEnclave(t *testing.T) { p := newPair(t) defer p.cleanupFn()