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 76625d6..18c2930 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,89 @@ 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 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 + // 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 + } + // 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 + } + 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 +358,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 +407,7 @@ func main() { } securityMW.Close() + treeMgr.Stop() clusterNode.Stop() cancel() } @@ -442,6 +533,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 +570,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 +646,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 +907,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/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/cluster/node.go b/internal/cluster/node.go index dbecd20..1a94657 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 @@ -49,14 +74,27 @@ type ClusterNode struct { 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) @@ -123,6 +161,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). // @@ -231,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 } @@ -324,36 +378,59 @@ 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 { + writeOp.markComplete() + } 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 +470,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/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..daa3f2d --- /dev/null +++ b/internal/transport/ws/connection.go @@ -0,0 +1,401 @@ +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) + onError func(error) + attachmentHandlers []handlerEntry[func(*AttachmentMessage)] + closeHandlers []handlerEntry[func(int, string)] + handlerSeq atomic.Uint64 + + 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 +} + +// 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() +} + +// 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.onError = fn + c.handlersMu.Unlock() +} + +// 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.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 + } + } + } +} + +// 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.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 +// 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() + handlers := append([]handlerEntry[func(*AttachmentMessage)](nil), c.attachmentHandlers...) + c.handlersMu.RUnlock() + for _, h := range handlers { + h.fn(msg) + } +} + +func (c *Connection) fireClose(code int, reason string) { + c.handlersMu.RLock() + handlers := append([]handlerEntry[func(int, string)](nil), c.closeHandlers...) + c.handlersMu.RUnlock() + for _, h := range handlers { + h.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..7cfd1e9 --- /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.AddAttachmentHandler(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.AddAttachmentHandler(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.AddAttachmentHandler(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.AddAttachmentHandler(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.AddCloseHandler(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.AddAttachmentHandler(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.AddAttachmentHandler(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.AddAttachmentHandler(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.AddCloseHandler(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.AddCloseHandler(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) + } + } +} diff --git a/internal/tree/manager.go b/internal/tree/manager.go new file mode 100644 index 0000000..a669832 --- /dev/null +++ b/internal/tree/manager.go @@ -0,0 +1,870 @@ +// 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 + + // 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 + // 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 +// 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 + } + stopCtx, stopCancel := context.WithCancel(context.Background()) + m := &Manager{ + 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 + 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 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() + 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)}, + } + // 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) { + 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) + select { + case rejectedCh <- struct{}{}: + default: + } + return + } + select { + case welcomeCh <- &w: + default: + } + case ws.AttachmentTypeGoodbye: + select { + case rejectedCh <- struct{}{}: + default: + } + } + }) + removeClose := conn.AddCloseHandler(func(int, string) { + select { + case rejectedCh <- struct{}{}: + default: + } + }) + + // 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() + return nil, err + } + + 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() + + // 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 + 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() + 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() + 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. 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 + } + m.mu.Lock() + if m.reattachInFlight { + m.mu.Unlock() + return + } + m.reattachInFlight = true + m.wg.Add(1) + m.mu.Unlock() + + go func() { + defer m.wg.Done() + 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(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(m.stopCtx, 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 + } +} + +// 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] +} + +// 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. +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, 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 { + 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. +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.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() + }) +} + +// 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..12d77ab --- /dev/null +++ b/internal/tree/manager_test.go @@ -0,0 +1,943 @@ +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()") + } +} + +// ---- 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) + } + } +} + +// 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() + 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) { + 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 +}