diff --git a/agent/server.go b/agent/server.go index 054163c09..344063cf2 100644 --- a/agent/server.go +++ b/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) diff --git a/agent/server_test.go b/agent/server_test.go index 9cd4fbdd8..b003dd11e 100644 --- a/agent/server_test.go +++ b/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()) +}