fix(agent): stop sending legacy hubs a duplicate stats payload (#2529)

agent/server.go

@@ -151,11 +151,16 @@ func (a *Agent) handleSession(s ssh.Session) {
 			s.Exit(1)
 			return
 		}
+		// A legacy hub never sends a request, so the session is done once the payload is out.
+		_ = s.Exit(0)
+		return
 	}
 
 	var req common.HubRequest[cbor.RawMessage]
 	if err := cbor.NewDecoder(s).Decode(&req); err != nil {
-		// Fallback to legacy one-shot if the first decode fails
+		// Only reachable for hubs that report >= MinVersionAgentResponse: a
+		// malformed or truncated request falls back to the one-shot payload
+		// so the peer still receives something parseable.
 		if err2 := a.handleLegacyStats(s, hubVersion); err2 != nil {
 			slog.Error("Error encoding stats (fallback)", "err", err2)
 			s.Exit(1)

agent/server_test.go

@@ -3,6 +3,7 @@
 package agent
 
 import (
+	"bytes"
 	"context"
 	"crypto/ed25519"
 	"encoding/json"
@@ -19,6 +20,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/henrygd/beszel/internal/common"
 	"github.com/henrygd/beszel/internal/entities/container"
 	"github.com/henrygd/beszel/internal/entities/system"
 
@@ -972,3 +974,184 @@ func TestGetHubVersionConcurrent(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+/////////////////////////////////////////////////////////////////
+/////////////// Session Protocol Tests //////////////////////////
+/////////////////////////////////////////////////////////////////
+
+// dialTestAgentClient starts an agent SSH server on a private unix socket and
+// connects a real SSH client whose ClientVersion mimics a hub of the given
+// version, so session-protocol behavior can be exercised end to end.
+func dialTestAgentClient(t *testing.T, clientVersion string) *gossh.Client {
+	t.Helper()
+
+	pubKey, privKey, err := ed25519.GenerateKey(nil)
+	require.NoError(t, err)
+	signer, err := gossh.NewSignerFromKey(privKey)
+	require.NoError(t, err)
+	sshPubKey, err := gossh.NewPublicKey(pubKey)
+	require.NoError(t, err)
+
+	socketFile := filepath.Join(t.TempDir(), "beszel-test.sock")
+	agent, err := NewAgent("")
+	require.NoError(t, err)
+
+	errChan := make(chan error, 1)
+	go func() {
+		errChan <- agent.StartServer(ServerOptions{
+			Network: "unix",
+			Addr:    socketFile,
+			Keys:    []gossh.PublicKey{sshPubKey},
+		})
+	}()
+	t.Cleanup(func() {
+		_ = agent.StopServer()
+		select {
+		case <-errChan:
+		default:
+		}
+	})
+
+	// Wait for the listener instead of a fixed sleep.
+	require.Eventually(t, func() bool {
+		conn, err := net.Dial("unix", socketFile)
+		if err != nil {
+			return false
+		}
+		_ = conn.Close()
+		return true
+	}, 5*time.Second, 25*time.Millisecond)
+
+	client, err := gossh.Dial("unix", socketFile, &gossh.ClientConfig{
+		User:            "u",
+		Auth:            []gossh.AuthMethod{gossh.PublicKeys(signer)},
+		HostKeyCallback: gossh.InsecureIgnoreHostKey(),
+		Timeout:         4 * time.Second,
+		ClientVersion:   clientVersion,
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = client.Close() })
+	return client
+}
+
+// waitSessionOutput runs readFn on a session stream and fails the test if it
+// does not finish quickly. A healthy session ends in milliseconds; the guard
+// exists so a handler that blocks (the bug this suite guards against) fails
+// the test instead of hanging it until the 70s idle timeout.
+func waitSessionOutput(t *testing.T, readFn func() ([]byte, error)) []byte {
+	t.Helper()
+	type result struct {
+		output []byte
+		err    error
+	}
+	resCh := make(chan result, 1)
+	go func() {
+		output, err := readFn()
+		resCh <- result{output, err}
+	}()
+	select {
+	case res := <-resCh:
+		require.NoError(t, res.err)
+		return res.output
+	case <-time.After(10 * time.Second):
+		t.Fatal("session did not complete within 10s; handler appears to be blocked")
+		return nil
+	}
+}
+
+// TestLegacyHubGetsSinglePayload pins the legacy one-shot contract: a hub
+// older than MinVersionAgentResponse never sends a request, so the agent must
+// write exactly one stats payload, exit with status 0, and end the session
+// promptly. Regression target: the handler used to fall through into the
+// request decode after the legacy write, read an immediate EOF from the hub's
+// empty stdin, and send a second payload via the decode-failure fallback.
+// Covers all three routes into the legacy path: a pre-0.13 CBOR hub, a
+// pre-0.12 JSON hub, and a client whose version string cannot be parsed.
+func TestLegacyHubGetsSinglePayload(t *testing.T) {
+	tests := []struct {
+		name          string
+		clientVersion string
+		cborFormat    bool
+	}{
+		{name: "legacy cbor hub 0.12.0", clientVersion: "SSH-2.0-beszel_0.12.0", cborFormat: true},
+		{name: "legacy json hub 0.11.1", clientVersion: "SSH-2.0-beszel_0.11.1", cborFormat: false},
+		{name: "unparsable client version", clientVersion: "SSH-2.0-OpenSSH_8.0", cborFormat: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			client := dialTestAgentClient(t, tt.clientVersion)
+
+			session, err := client.NewSession()
+			require.NoError(t, err)
+			defer func() { _ = session.Close() }()
+
+			output := waitSessionOutput(t, func() ([]byte, error) {
+				return session.Output("")
+			})
+			require.NotEmpty(t, output, "legacy hub must receive a stats payload")
+
+			// Exactly one payload: decode the first object, then require the
+			// stream to be exhausted. A second write would leave bytes here.
+			var data system.CombinedData
+			if tt.cborFormat {
+				dec := cbor.NewDecoder(bytes.NewReader(output))
+				require.NoError(t, dec.Decode(&data), "payload must be valid CBOR")
+				var extra any
+				require.ErrorIs(t, dec.Decode(&extra), io.EOF,
+					"legacy hub must receive exactly one payload, got trailing bytes")
+			} else {
+				dec := json.NewDecoder(bytes.NewReader(output))
+				require.NoError(t, dec.Decode(&data), "payload must be valid JSON")
+				var extra map[string]any
+				require.ErrorIs(t, dec.Decode(&extra), io.EOF,
+					"legacy hub must receive exactly one payload, got trailing bytes")
+			}
+			// ConnectionType is stamped by the SSH path itself, so it is a
+			// deterministic (environment-independent) proof that the payload
+			// came from the SSH handler.
+			assert.Equal(t, system.ConnectionTypeSSH, data.Info.ConnectionType,
+				"payload must be stamped as SSH-sourced")
+
+			// Output() returns an *ssh.ExitError when the handler exits
+			// non-zero; reaching this point with no error means exit status 0.
+			// (An s.Exit(1) from a failed legacy write would surface here.)
+		})
+	}
+}
+
+// TestModernHubRequestResponse guards the new-protocol path against
+// regressions from the legacy-path changes: a hub at or above
+// MinVersionAgentResponse sends a CBOR request and must get exactly one
+// AgentResponse with populated system data and a clean exit.
+func TestModernHubRequestResponse(t *testing.T) {
+	client := dialTestAgentClient(t, "SSH-2.0-beszel_0.21.0")
+
+	session, err := client.NewSession()
+	require.NoError(t, err)
+	defer func() { _ = session.Close() }()
+
+	stdout, err := session.StdoutPipe()
+	require.NoError(t, err)
+	stdin, err := session.StdinPipe()
+	require.NoError(t, err)
+	require.NoError(t, session.Shell())
+
+	// Same request shape the hub's SSH transport sends.
+	require.NoError(t, cbor.NewEncoder(stdin).Encode(common.HubRequest[any]{Action: common.GetData}))
+	require.NoError(t, stdin.Close())
+
+	output := waitSessionOutput(t, func() ([]byte, error) {
+		return io.ReadAll(stdout)
+	})
+
+	var resp common.AgentResponse
+	require.NoError(t, cbor.Unmarshal(output, &resp), "response must be a single valid CBOR AgentResponse")
+	assert.Empty(t, resp.Error)
+	require.NotNil(t, resp.SystemData, "GetData response must carry system data")
+	assert.Equal(t, system.ConnectionTypeSSH, resp.SystemData.Info.ConnectionType,
+		"GetData response must be stamped as SSH-sourced")
+
+	// Wait() returns an *ssh.ExitError on non-zero exit; nil means status 0.
+	require.NoError(t, session.Wait())
+}