From 9e4c2d830875ae40dc1f6bef5f2e34eb6f3ec9d7 Mon Sep 17 00:00:00 2001 From: Shawn Burke Date: Thu, 24 Sep 2026 08:43:33 +1000 Subject: [PATCH 1/4] Don't run the relay idle watchdog in serve mode A serve-mode (handler) agent's relay tunnel only carries what relay-dispatcher sends it, mostly the periodic /__axon/info liveness poll. The broker server delivers a token's request to the newest client on each server, so when several replicas share a token only one of them is polled. The others see no traffic, hit RELAY_IDLE_TIMEOUT, restart, become the newest client, and the rotation repeats forever. Serve mode now defaults RelayIdleTimeout to 0. Setting RELAY_IDLE_TIMEOUT explicitly still enables the watchdog. Co-Authored-By: Claude Opus 5.5 (1M context) --- agent/cmd/serve.go | 15 +++++++++++++++ agent/cmd/serve_test.go | 13 +++++++++++++ 2 files changed, 28 insertions(+) diff --git a/agent/cmd/serve.go b/agent/cmd/serve.go index f3e6563..9920f78 100644 --- a/agent/cmd/serve.go +++ b/agent/cmd/serve.go @@ -4,6 +4,7 @@ import ( _ "embed" "fmt" "os" + "time" "github.com/cortexapps/axon/common" "github.com/cortexapps/axon/config" @@ -35,6 +36,8 @@ var serveCmd = &cobra.Command{ config.IntegrationAlias = id } + config.RelayIdleTimeout = serveRelayIdleTimeout(config.RelayIdleTimeout) + config.Print() info := common.IntegrationInfo{ @@ -59,6 +62,18 @@ var serveCmd = &cobra.Command{ }, } +// serveRelayIdleTimeout disables the relay idle watchdog for serve mode +// unless RELAY_IDLE_TIMEOUT is set explicitly. A handler agent's tunnel +// carries only what relay-dispatcher sends it, and the broker server +// delivers that to just the newest client per token, so the other +// replicas would look idle and be restarted in rotation forever. +func serveRelayIdleTimeout(configured time.Duration) time.Duration { + if _, set := os.LookupEnv("RELAY_IDLE_TIMEOUT"); set { + return configured + } + return 0 +} + func init() { serveCmd.Flags().Bool("dry-run", false, "Dry run mode") serveCmd.Flags().StringP("alias", "a", "customer-agent", "Alias (identifier) for this agent type") diff --git a/agent/cmd/serve_test.go b/agent/cmd/serve_test.go index 9b8b595..19ed21a 100644 --- a/agent/cmd/serve_test.go +++ b/agent/cmd/serve_test.go @@ -52,6 +52,19 @@ func TestBuildServeStack(t *testing.T) { app.Stop(context.Background()) } +func TestServeRelayIdleTimeout(t *testing.T) { + oldEnv := util.SaveEnv(false) + defer util.RestoreEnv(oldEnv) + + os.Unsetenv("RELAY_IDLE_TIMEOUT") + require.Equal(t, time.Duration(0), serveRelayIdleTimeout(10*time.Minute), + "serve mode must not run the idle watchdog by default") + + os.Setenv("RELAY_IDLE_TIMEOUT", "3m") + require.Equal(t, 3*time.Minute, serveRelayIdleTimeout(3*time.Minute), + "an explicit RELAY_IDLE_TIMEOUT still wins") +} + func TestBuildServeStackLive(t *testing.T) { // create a fake server that serves http://localhost:xxx/relay/register From 26d204558de81b73ab8600d3d431acb6c657aa76 Mon Sep 17 00:00:00 2001 From: Shawn Burke Date: Thu, 24 Sep 2026 09:11:52 +1000 Subject: [PATCH 2/4] Count frames from the broker server as tunnel activity for the idle watchdog The idle watchdog restarted the broker when no request had been relayed for RELAY_IDLE_TIMEOUT. The broker server hands a token's requests to only its newest client on each server, so every other replica sharing a token looked idle and was restarted in rotation, in relay and serve mode alike. The reflector's websocket proxy now reports bytes that arrive from the broker server, which heartbeats every tunnel, and the watchdog treats those as activity alongside relayed requests and the last restart. Bytes we send don't count: a write into a dead tunnel can still succeed. A tunnel that goes quiet is still restarted. This replaces the earlier change that turned the watchdog off in serve mode, which would also have stopped it recovering a dead tunnel there. Co-Authored-By: Claude Opus 5.5 (1M context) --- agent/cmd/serve.go | 15 ----- agent/cmd/serve_test.go | 13 ---- agent/server/snykbroker/reflector.go | 15 +++++ .../snykbroker/relay_instance_manager.go | 14 +++- .../snykbroker/relay_instance_manager_test.go | 31 +++++++++ agent/server/snykbroker/ws_proxy.go | 18 +++-- agent/server/snykbroker/ws_proxy_test.go | 65 +++++++++++++++++++ 7 files changed, 136 insertions(+), 35 deletions(-) diff --git a/agent/cmd/serve.go b/agent/cmd/serve.go index 9920f78..f3e6563 100644 --- a/agent/cmd/serve.go +++ b/agent/cmd/serve.go @@ -4,7 +4,6 @@ import ( _ "embed" "fmt" "os" - "time" "github.com/cortexapps/axon/common" "github.com/cortexapps/axon/config" @@ -36,8 +35,6 @@ var serveCmd = &cobra.Command{ config.IntegrationAlias = id } - config.RelayIdleTimeout = serveRelayIdleTimeout(config.RelayIdleTimeout) - config.Print() info := common.IntegrationInfo{ @@ -62,18 +59,6 @@ var serveCmd = &cobra.Command{ }, } -// serveRelayIdleTimeout disables the relay idle watchdog for serve mode -// unless RELAY_IDLE_TIMEOUT is set explicitly. A handler agent's tunnel -// carries only what relay-dispatcher sends it, and the broker server -// delivers that to just the newest client per token, so the other -// replicas would look idle and be restarted in rotation forever. -func serveRelayIdleTimeout(configured time.Duration) time.Duration { - if _, set := os.LookupEnv("RELAY_IDLE_TIMEOUT"); set { - return configured - } - return 0 -} - func init() { serveCmd.Flags().Bool("dry-run", false, "Dry run mode") serveCmd.Flags().StringP("alias", "a", "customer-agent", "Alias (identifier) for this agent type") diff --git a/agent/cmd/serve_test.go b/agent/cmd/serve_test.go index 19ed21a..9b8b595 100644 --- a/agent/cmd/serve_test.go +++ b/agent/cmd/serve_test.go @@ -52,19 +52,6 @@ func TestBuildServeStack(t *testing.T) { app.Stop(context.Background()) } -func TestServeRelayIdleTimeout(t *testing.T) { - oldEnv := util.SaveEnv(false) - defer util.RestoreEnv(oldEnv) - - os.Unsetenv("RELAY_IDLE_TIMEOUT") - require.Equal(t, time.Duration(0), serveRelayIdleTimeout(10*time.Minute), - "serve mode must not run the idle watchdog by default") - - os.Setenv("RELAY_IDLE_TIMEOUT", "3m") - require.Equal(t, 3*time.Minute, serveRelayIdleTimeout(3*time.Minute), - "an explicit RELAY_IDLE_TIMEOUT still wins") -} - func TestBuildServeStackLive(t *testing.T) { // create a fake server that serves http://localhost:xxx/relay/register diff --git a/agent/server/snykbroker/reflector.go b/agent/server/snykbroker/reflector.go index f49eb5a..36ca5b2 100644 --- a/agent/server/snykbroker/reflector.go +++ b/agent/server/snykbroker/reflector.go @@ -33,6 +33,7 @@ type RegistrationReflector struct { config config.AgentConfig lastTrafficTime atomic.Int64 lastStartupTime atomic.Int64 + lastTunnelTime atomic.Int64 wsProxy *WebSocketProxy } @@ -77,6 +78,7 @@ func NewRegistrationReflector(p RegistrationReflectorParams) *RegistrationReflec rr.wsProxy.OnTunnelEstablished = func(target string) { rr.logger.Info("WebSocket tunnel established", zap.String("target", target)) } + rr.wsProxy.OnActivity = rr.RecordTunnelActivity rr.RecordStartup() @@ -137,6 +139,19 @@ func (rr *RegistrationReflector) LastStartupTime() time.Time { return time.UnixMilli(rr.lastStartupTime.Load()) } +// RecordTunnelActivity marks that the broker server just sent something down +// the websocket tunnel. It is kept apart from lastTrafficTime: a heartbeat +// proves the tunnel is alive, not that anything was relayed. +func (rr *RegistrationReflector) RecordTunnelActivity() { + rr.lastTunnelTime.Store(time.Now().UnixMilli()) +} + +// LastTunnelActivityTime returns the time the broker server last sent +// anything down the tunnel. +func (rr *RegistrationReflector) LastTunnelActivityTime() time.Time { + return time.UnixMilli(rr.lastTunnelTime.Load()) +} + // SetOnWSTunnelClose sets a callback invoked when a WebSocket tunnel closes. // Used by the relay instance manager to trigger a broker restart. func (rr *RegistrationReflector) SetOnWSTunnelClose(fn func()) { diff --git a/agent/server/snykbroker/relay_instance_manager.go b/agent/server/snykbroker/relay_instance_manager.go index 65dfaa2..1e75e96 100644 --- a/agent/server/snykbroker/relay_instance_manager.go +++ b/agent/server/snykbroker/relay_instance_manager.go @@ -311,9 +311,19 @@ func (r *relayInstanceManager) shouldRestart() (bool, string) { if !r.config.HttpRelayReflectorMode.ReflectsTraffic() { return false, "" } + // Relayed requests alone can't tell a dead tunnel from a quiet one: the + // broker server hands a token's requests to only its newest client, so + // every other replica sharing the token looks idle and gets restarted in + // rotation. Frames from the server (its heartbeat included) show the + // tunnel is alive, so they count too. lastActivity := r.reflector.LastTrafficTime() - if startup := r.reflector.LastStartupTime(); startup.After(lastActivity) { - lastActivity = startup + for _, t := range []time.Time{ + r.reflector.LastStartupTime(), + r.reflector.LastTunnelActivityTime(), + } { + if t.After(lastActivity) { + lastActivity = t + } } if time.Since(lastActivity) >= r.config.RelayIdleTimeout { return true, "idle_timeout" diff --git a/agent/server/snykbroker/relay_instance_manager_test.go b/agent/server/snykbroker/relay_instance_manager_test.go index 3dda95c..3213df4 100644 --- a/agent/server/snykbroker/relay_instance_manager_test.go +++ b/agent/server/snykbroker/relay_instance_manager_test.go @@ -486,6 +486,37 @@ func TestIdleTimeoutDetectsIdleReflector(t *testing.T) { require.NoError(t, err) } +// Only the newest client for a token gets relayed requests, so a healthy +// replica can go a long time without traffic. Frames from the broker server +// keep it out of the idle watchdog; a tunnel that has gone quiet does not. +func TestShouldRestartCountsTunnelActivity(t *testing.T) { + cfg := config.NewAgentEnvConfig() + cfg.HttpRelayReflectorMode = config.RelayReflectorAllTraffic + cfg.RelayIdleTimeout = time.Minute + + rr := NewRegistrationReflector(RegistrationReflectorParams{ + Logger: zap.NewNop(), + Config: cfg, + }) + r := &relayInstanceManager{config: cfg, reflector: rr} + + stale := time.Now().Add(-2 * time.Minute).UnixMilli() + rr.lastStartupTime.Store(stale) + rr.lastTrafficTime.Store(stale) + + restart, reason := r.shouldRestart() + require.True(t, restart, "no traffic, no startup and no tunnel frames within the window is idle") + require.Equal(t, "idle_timeout", reason) + + rr.RecordTunnelActivity() + restart, _ = r.shouldRestart() + require.False(t, restart, "a frame from the broker server shows the tunnel is alive") + + rr.lastTunnelTime.Store(stale) + restart, _ = r.shouldRestart() + require.True(t, restart, "a tunnel that has gone quiet is idle again") +} + // A broker with nothing to relay looks identical, idle-wise, to one whose // tunnel died silently: shouldRestart() can only see "no activity recorded." // A restart must reset the idle clock (LastStartupTime) so a healthy-but-idle diff --git a/agent/server/snykbroker/ws_proxy.go b/agent/server/snykbroker/ws_proxy.go index cd5430b..1b33803 100644 --- a/agent/server/snykbroker/ws_proxy.go +++ b/agent/server/snykbroker/ws_proxy.go @@ -31,6 +31,11 @@ type WebSocketProxy struct { // Callbacks OnTunnelEstablished func(target string) OnTunnelClosed func(target string, duration time.Duration) + // OnActivity is called whenever bytes arrive from the target (the broker + // server). The server heartbeats every tunnel, so this proves the far end + // is alive even when no request is being relayed. Bytes we send don't + // count: writing into a dead tunnel can still succeed. + OnActivity func() // State tracking - use Int32 to properly track multiple concurrent tunnels activeConnections atomic.Int32 @@ -256,13 +261,13 @@ func (wp *WebSocketProxy) runTunnel(clientConn, targetConn net.Conn, targetAddr done := make(chan struct{}, 2) - copy := func(dst, src net.Conn, direction string) { + copy := func(dst, src net.Conn, direction string, onRead func()) { defer func() { done <- struct{}{} }() - wp.copyWithIdleTimeout(dst, src, direction) + wp.copyWithIdleTimeout(dst, src, direction, onRead) } - go copy(clientConn, targetConn, "target->client") - go copy(targetConn, clientConn, "client->target") + go copy(clientConn, targetConn, "target->client", wp.OnActivity) + go copy(targetConn, clientConn, "client->target", nil) <-done clientConn.Close() @@ -270,7 +275,7 @@ func (wp *WebSocketProxy) runTunnel(clientConn, targetConn net.Conn, targetAddr <-done } -func (wp *WebSocketProxy) copyWithIdleTimeout(dst, src net.Conn, direction string) { +func (wp *WebSocketProxy) copyWithIdleTimeout(dst, src net.Conn, direction string, onRead func()) { buf := make([]byte, 32*1024) isFirstRead := true for { @@ -298,6 +303,9 @@ func (wp *WebSocketProxy) copyWithIdleTimeout(dst, src net.Conn, direction strin } if n > 0 { + if onRead != nil { + onRead() + } dst.SetWriteDeadline(time.Now().Add(wp.HandshakeTimeout)) if _, writeErr := dst.Write(buf[:n]); writeErr != nil { return diff --git a/agent/server/snykbroker/ws_proxy_test.go b/agent/server/snykbroker/ws_proxy_test.go index 71de8df..6b13cb2 100644 --- a/agent/server/snykbroker/ws_proxy_test.go +++ b/agent/server/snykbroker/ws_proxy_test.go @@ -479,6 +479,71 @@ func TestWebSocketProxyCallbacks(t *testing.T) { assert.True(t, closedDuration > 0, "duration should be positive") } +// OnActivity must fire for bytes from the broker server and not for bytes we +// send: a write into a dead tunnel can still succeed, so it proves nothing. +func TestWebSocketProxyOnActivityCountsTargetFramesOnly(t *testing.T) { + logger := newTestLogger(t) + + serverSend := make(chan struct{}) + var upgrader = websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { return true }, + } + router := mux.NewRouter() + router.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return + } + defer conn.Close() + go func() { + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + }() + <-serverSend + _ = conn.WriteMessage(websocket.TextMessage, []byte("primus::ping::1")) + time.Sleep(200 * time.Millisecond) + }) + targetServer := httptest.NewServer(router) + defer targetServer.Close() + + var activity atomic.Int32 + rr := NewRegistrationReflector(RegistrationReflectorParams{ + Logger: logger, + Config: config.AgentConfig{ + ReflectorWebSocketUpgrade: true, + }, + }) + rr.wsProxy.OnActivity = func() { activity.Add(1) } + + reflectorRouter := mux.NewRouter() + rr.RegisterRoutes(reflectorRouter) + reflectorServer := httptest.NewServer(reflectorRouter) + defer reflectorServer.Close() + + wsURL := "ws" + rr.ProxyURI(targetServer.URL)[4:] + "/ws" + conn, _, err := (&websocket.Dialer{}).Dial(wsURL, nil) + require.NoError(t, err) + defer conn.Close() + + // The upgrade response is the first thing back from the target; only + // frames after it are under test. + time.Sleep(50 * time.Millisecond) + afterHandshake := activity.Load() + + require.NoError(t, conn.WriteMessage(websocket.TextMessage, []byte("primus::pong::1"))) + time.Sleep(50 * time.Millisecond) + require.Equal(t, afterHandshake, activity.Load(), "our own writes must not count as activity") + + close(serverSend) + _, msg, err := conn.ReadMessage() + require.NoError(t, err) + require.Equal(t, "primus::ping::1", string(msg)) + require.Greater(t, activity.Load(), afterHandshake, "a frame from the server must count as activity") +} + func TestWebSocketProxyIsConnected(t *testing.T) { logger := newTestLogger(t) From f1c6951753fca6dd7055881e14214d672dd3f7b9 Mon Sep 17 00:00:00 2001 From: Shawn Burke Date: Thu, 24 Sep 2026 09:24:05 +1000 Subject: [PATCH 3/4] Take the idle watchdog's clock as a single max over the three watermarks Review feedback on #151. The reflector now exposes LastActivityTime, the latest of the last relayed request, the last frame from the broker server, and the last (re)start, and shouldRestart compares only that. LastTunnelActivityTime had no other caller and is gone. Co-Authored-By: Claude Opus 5.5 (1M context) --- agent/server/snykbroker/reflector.go | 18 ++++++++++++++---- .../snykbroker/relay_instance_manager.go | 16 +--------------- 2 files changed, 15 insertions(+), 19 deletions(-) diff --git a/agent/server/snykbroker/reflector.go b/agent/server/snykbroker/reflector.go index 36ca5b2..872a84f 100644 --- a/agent/server/snykbroker/reflector.go +++ b/agent/server/snykbroker/reflector.go @@ -146,10 +146,20 @@ func (rr *RegistrationReflector) RecordTunnelActivity() { rr.lastTunnelTime.Store(time.Now().UnixMilli()) } -// LastTunnelActivityTime returns the time the broker server last sent -// anything down the tunnel. -func (rr *RegistrationReflector) LastTunnelActivityTime() time.Time { - return time.UnixMilli(rr.lastTunnelTime.Load()) +// LastActivityTime is the latest of the last relayed request, the last frame +// from the broker server, and the last (re)start: the idle watchdog's clock. +// +// Relayed requests alone can't tell a dead tunnel from a quiet one: the +// broker server hands a token's requests to only its newest client, so every +// other replica sharing the token looks idle. Frames cover that when the +// tunnel runs through this reflector; in "traffic" mode it doesn't, and +// relayed requests are the only signal. +func (rr *RegistrationReflector) LastActivityTime() time.Time { + return time.UnixMilli(max( + rr.lastTrafficTime.Load(), + rr.lastTunnelTime.Load(), + rr.lastStartupTime.Load(), + )) } // SetOnWSTunnelClose sets a callback invoked when a WebSocket tunnel closes. diff --git a/agent/server/snykbroker/relay_instance_manager.go b/agent/server/snykbroker/relay_instance_manager.go index 1e75e96..aadcb03 100644 --- a/agent/server/snykbroker/relay_instance_manager.go +++ b/agent/server/snykbroker/relay_instance_manager.go @@ -311,21 +311,7 @@ func (r *relayInstanceManager) shouldRestart() (bool, string) { if !r.config.HttpRelayReflectorMode.ReflectsTraffic() { return false, "" } - // Relayed requests alone can't tell a dead tunnel from a quiet one: the - // broker server hands a token's requests to only its newest client, so - // every other replica sharing the token looks idle and gets restarted in - // rotation. Frames from the server (its heartbeat included) show the - // tunnel is alive, so they count too. - lastActivity := r.reflector.LastTrafficTime() - for _, t := range []time.Time{ - r.reflector.LastStartupTime(), - r.reflector.LastTunnelActivityTime(), - } { - if t.After(lastActivity) { - lastActivity = t - } - } - if time.Since(lastActivity) >= r.config.RelayIdleTimeout { + if time.Since(r.reflector.LastActivityTime()) >= r.config.RelayIdleTimeout { return true, "idle_timeout" } return false, "" From 2bdbfdf15e3922dfffb6efc5de56bdd44a184803 Mon Sep 17 00:00:00 2001 From: Shawn Burke Date: Thu, 24 Sep 2026 09:32:44 +1000 Subject: [PATCH 4/4] Scope tunnel activity to the broker's own tunnel and watch it in registration mode Review feedback on #151. The reflector proxies every websocket upgrade, including websockets to customer origins in accept-file rules, so a busy one could keep the tunnel watermark fresh while the broker's tunnel to the server was dead. Proxy now takes the activity callback per call, and the reflector passes it only for the default entry, which is the broker server. In "registration" mode the broker's tunnel runs through the reflector but relayed requests don't, and shouldRestart returned before looking. It now runs in every reflecting mode and stays off only when the reflector is disabled, where the agent can see neither signal. Co-Authored-By: Claude Opus 5.5 (1M context) --- agent/server/snykbroker/reflector.go | 9 +- .../snykbroker/relay_instance_manager.go | 6 +- .../snykbroker/relay_instance_manager_test.go | 15 +++ agent/server/snykbroker/ws_proxy.go | 16 ++- agent/server/snykbroker/ws_proxy_test.go | 113 ++++++++++-------- 5 files changed, 99 insertions(+), 60 deletions(-) diff --git a/agent/server/snykbroker/reflector.go b/agent/server/snykbroker/reflector.go index 872a84f..d9a720f 100644 --- a/agent/server/snykbroker/reflector.go +++ b/agent/server/snykbroker/reflector.go @@ -78,7 +78,6 @@ func NewRegistrationReflector(p RegistrationReflectorParams) *RegistrationReflec rr.wsProxy.OnTunnelEstablished = func(target string) { rr.logger.Info("WebSocket tunnel established", zap.String("target", target)) } - rr.wsProxy.OnActivity = rr.RecordTunnelActivity rr.RecordStartup() @@ -421,7 +420,13 @@ func (rr *RegistrationReflector) ServeHTTP(w http.ResponseWriter, r *http.Reques // Check if this is a WebSocket upgrade request if rr.config.ReflectorWebSocketUpgrade && IsWebSocketUpgrade(r) { rr.logger.Debug("Detected WebSocket upgrade request, using WebSocket proxy") - if err := rr.wsProxy.Proxy(w, r, entry.TargetURI); err != nil { + // Only the default entry is the broker's own tunnel to the server; a + // websocket to a customer origin says nothing about that tunnel. + var onActivity func() + if entry.isDefault { + onActivity = rr.RecordTunnelActivity + } + if err := rr.wsProxy.Proxy(w, r, entry.TargetURI, onActivity); err != nil { rr.logger.Error("WebSocket proxy failed", zap.Error(err)) } return diff --git a/agent/server/snykbroker/relay_instance_manager.go b/agent/server/snykbroker/relay_instance_manager.go index aadcb03..b07786e 100644 --- a/agent/server/snykbroker/relay_instance_manager.go +++ b/agent/server/snykbroker/relay_instance_manager.go @@ -308,7 +308,11 @@ func (r *relayInstanceManager) shouldRestart() (bool, string) { if r.config.RelayIdleTimeout == 0 || r.reflector == nil { return false, "" } - if !r.config.HttpRelayReflectorMode.ReflectsTraffic() { + // Each reflecting mode gives the watchdog a signal: relayed requests in + // "traffic", frames on the broker's tunnel in "registration", both in + // "all". Only "disabled" leaves it blind. + mode := r.config.HttpRelayReflectorMode + if !mode.ReflectsTraffic() && !mode.ReflectsRegistration() { return false, "" } if time.Since(r.reflector.LastActivityTime()) >= r.config.RelayIdleTimeout { diff --git a/agent/server/snykbroker/relay_instance_manager_test.go b/agent/server/snykbroker/relay_instance_manager_test.go index 3213df4..d5c7043 100644 --- a/agent/server/snykbroker/relay_instance_manager_test.go +++ b/agent/server/snykbroker/relay_instance_manager_test.go @@ -515,6 +515,21 @@ func TestShouldRestartCountsTunnelActivity(t *testing.T) { rr.lastTunnelTime.Store(stale) restart, _ = r.shouldRestart() require.True(t, restart, "a tunnel that has gone quiet is idle again") + + // "registration" routes the broker's tunnel through the reflector but no + // relayed requests, so tunnel frames are its only signal. + r.config.HttpRelayReflectorMode = config.RelayReflectorRegistrationOnly + restart, _ = r.shouldRestart() + require.True(t, restart, "registration mode watches the tunnel too") + rr.RecordTunnelActivity() + restart, _ = r.shouldRestart() + require.False(t, restart, "registration mode counts tunnel frames") + + // With the reflector off the agent sees neither, so it must not guess. + rr.lastTunnelTime.Store(stale) + r.config.HttpRelayReflectorMode = config.RelayReflectorDisabled + restart, _ = r.shouldRestart() + require.False(t, restart, "disabled mode has no signal and never restarts") } // A broker with nothing to relay looks identical, idle-wise, to one whose diff --git a/agent/server/snykbroker/ws_proxy.go b/agent/server/snykbroker/ws_proxy.go index 1b33803..f853fa5 100644 --- a/agent/server/snykbroker/ws_proxy.go +++ b/agent/server/snykbroker/ws_proxy.go @@ -31,11 +31,6 @@ type WebSocketProxy struct { // Callbacks OnTunnelEstablished func(target string) OnTunnelClosed func(target string, duration time.Duration) - // OnActivity is called whenever bytes arrive from the target (the broker - // server). The server heartbeats every tunnel, so this proves the far end - // is alive even when no request is being relayed. Bytes we send don't - // count: writing into a dead tunnel can still succeed. - OnActivity func() // State tracking - use Int32 to properly track multiple concurrent tunnels activeConnections atomic.Int32 @@ -65,7 +60,10 @@ func (wp *WebSocketProxy) ActiveConnections() int32 { // Proxy handles a WebSocket upgrade request by establishing a tunnel to the target. // It hijacks the client connection and proxies bidirectionally. -func (wp *WebSocketProxy) Proxy(w http.ResponseWriter, r *http.Request, targetURI string) error { +// +// onActivity, if set, is called whenever bytes arrive from the target. Bytes +// sent to the target don't count: writing into a dead tunnel can still succeed. +func (wp *WebSocketProxy) Proxy(w http.ResponseWriter, r *http.Request, targetURI string, onActivity func()) error { targetURL, err := url.Parse(targetURI) if err != nil { return fmt.Errorf("invalid target URI: %w", err) @@ -96,7 +94,7 @@ func (wp *WebSocketProxy) Proxy(w http.ResponseWriter, r *http.Request, targetUR } // Run the bidirectional tunnel - wp.runTunnel(clientConn, targetConn, targetAddr) + wp.runTunnel(clientConn, targetConn, targetAddr, onActivity) return nil } @@ -244,7 +242,7 @@ func (wp *WebSocketProxy) hijackConnection(w http.ResponseWriter) (net.Conn, err return conn, nil } -func (wp *WebSocketProxy) runTunnel(clientConn, targetConn net.Conn, targetAddr string) { +func (wp *WebSocketProxy) runTunnel(clientConn, targetConn net.Conn, targetAddr string, onActivity func()) { start := time.Now() wp.activeConnections.Add(1) @@ -266,7 +264,7 @@ func (wp *WebSocketProxy) runTunnel(clientConn, targetConn net.Conn, targetAddr wp.copyWithIdleTimeout(dst, src, direction, onRead) } - go copy(clientConn, targetConn, "target->client", wp.OnActivity) + go copy(clientConn, targetConn, "target->client", onActivity) go copy(targetConn, clientConn, "client->target", nil) <-done diff --git a/agent/server/snykbroker/ws_proxy_test.go b/agent/server/snykbroker/ws_proxy_test.go index 6b13cb2..7024224 100644 --- a/agent/server/snykbroker/ws_proxy_test.go +++ b/agent/server/snykbroker/ws_proxy_test.go @@ -479,69 +479,86 @@ func TestWebSocketProxyCallbacks(t *testing.T) { assert.True(t, closedDuration > 0, "duration should be positive") } -// OnActivity must fire for bytes from the broker server and not for bytes we -// send: a write into a dead tunnel can still succeed, so it proves nothing. -func TestWebSocketProxyOnActivityCountsTargetFramesOnly(t *testing.T) { +// Frames from the broker server on the default entry (the broker's own +// tunnel) count as tunnel activity. Our own writes don't, since writing into a +// dead tunnel can still succeed, and neither does a websocket to a customer +// origin, which says nothing about the broker's tunnel. +func TestReflectorRecordsTunnelActivityForBrokerTunnelOnly(t *testing.T) { logger := newTestLogger(t) - serverSend := make(chan struct{}) - var upgrader = websocket.Upgrader{ - CheckOrigin: func(r *http.Request) bool { return true }, - } - router := mux.NewRouter() - router.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { - conn, err := upgrader.Upgrade(w, r, nil) - if err != nil { - return - } - defer conn.Close() - go func() { - for { - if _, _, err := conn.ReadMessage(); err != nil { - return - } + // A websocket server that sends one frame when told to. + newTarget := func(t *testing.T) (*httptest.Server, chan struct{}) { + send := make(chan struct{}) + upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }} + router := mux.NewRouter() + router.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + return } - }() - <-serverSend - _ = conn.WriteMessage(websocket.TextMessage, []byte("primus::ping::1")) - time.Sleep(200 * time.Millisecond) - }) - targetServer := httptest.NewServer(router) - defer targetServer.Close() + defer conn.Close() + go func() { + for { + if _, _, err := conn.ReadMessage(); err != nil { + return + } + } + }() + <-send + _ = conn.WriteMessage(websocket.TextMessage, []byte("primus::ping::1")) + time.Sleep(200 * time.Millisecond) + }) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + return server, send + } - var activity atomic.Int32 rr := NewRegistrationReflector(RegistrationReflectorParams{ Logger: logger, - Config: config.AgentConfig{ - ReflectorWebSocketUpgrade: true, - }, + Config: config.AgentConfig{ReflectorWebSocketUpgrade: true}, }) - rr.wsProxy.OnActivity = func() { activity.Add(1) } - reflectorRouter := mux.NewRouter() rr.RegisterRoutes(reflectorRouter) reflectorServer := httptest.NewServer(reflectorRouter) defer reflectorServer.Close() - wsURL := "ws" + rr.ProxyURI(targetServer.URL)[4:] + "/ws" - conn, _, err := (&websocket.Dialer{}).Dial(wsURL, nil) - require.NoError(t, err) - defer conn.Close() + // Dial through the reflector, write one frame of our own, then have the + // target send one. Returns the tunnel watermark after each step. + exchange := func(t *testing.T, proxyURI string, send chan struct{}) (afterOwnWrite, afterTargetFrame int64) { + conn, _, err := (&websocket.Dialer{}).Dial("ws"+proxyURI[4:]+"/ws", nil) + require.NoError(t, err) + defer conn.Close() - // The upgrade response is the first thing back from the target; only - // frames after it are under test. - time.Sleep(50 * time.Millisecond) - afterHandshake := activity.Load() + // The upgrade response is the first read from the target; only frames + // after it are under test. + time.Sleep(50 * time.Millisecond) + before := rr.lastTunnelTime.Load() - require.NoError(t, conn.WriteMessage(websocket.TextMessage, []byte("primus::pong::1"))) - time.Sleep(50 * time.Millisecond) - require.Equal(t, afterHandshake, activity.Load(), "our own writes must not count as activity") + require.NoError(t, conn.WriteMessage(websocket.TextMessage, []byte("primus::pong::1"))) + time.Sleep(50 * time.Millisecond) + afterOwnWrite = rr.lastTunnelTime.Load() + require.Equal(t, before, afterOwnWrite, "our own writes must not count as tunnel activity") - close(serverSend) - _, msg, err := conn.ReadMessage() - require.NoError(t, err) - require.Equal(t, "primus::ping::1", string(msg)) - require.Greater(t, activity.Load(), afterHandshake, "a frame from the server must count as activity") + // Millisecond watermark: make sure a new stamp can differ from the old. + time.Sleep(5 * time.Millisecond) + close(send) + _, msg, err := conn.ReadMessage() + require.NoError(t, err) + require.Equal(t, "primus::ping::1", string(msg)) + return afterOwnWrite, rr.lastTunnelTime.Load() + } + + t.Run("customer origin", func(t *testing.T) { + target, send := newTarget(t) + before, after := exchange(t, rr.ProxyURI(target.URL), send) + require.Equal(t, before, after, "a websocket to a customer origin must not count as tunnel activity") + }) + + t.Run("broker tunnel", func(t *testing.T) { + target, send := newTarget(t) + before, after := exchange(t, rr.ProxyURI(target.URL, WithDefault(true)), send) + require.Greater(t, after, before, "a frame from the broker server must count as tunnel activity") + }) } func TestWebSocketProxyIsConnected(t *testing.T) {