diff --git a/agent/server.go b/agent/server.go index 0b584d75..4a86fd3a 100644 --- a/agent/server.go +++ b/agent/server.go @@ -29,9 +29,6 @@ type ServerOptions struct { Keys []gossh.PublicKey // SSH public keys for authentication } -// hubVersions caches hub versions by session ID to avoid repeated parsing. -var hubVersions map[string]semver.Version - // StartServer starts the SSH server with the provided options. // It configures the server with secure defaults, sets up authentication, // and begins listening for connections. Returns an error if the server @@ -99,24 +96,15 @@ func (a *Agent) StartServer(opts ServerOptions) error { return a.server.Serve(ln) } -// getHubVersion retrieves and caches the hub version for a given session. -// It extracts the version from the SSH client version string and caches -// it to avoid repeated parsing. Returns a zero version if parsing fails. -func (a *Agent) getHubVersion(sessionId string, sessionCtx ssh.Context) semver.Version { - if hubVersions == nil { - hubVersions = make(map[string]semver.Version, 1) - } - hubVersion, ok := hubVersions[sessionId] - if ok { - return hubVersion - } - // Extract hub version from SSH client version +// getHubVersion extracts the hub version from the SSH client version string +// for a given session. Returns a zero version if parsing fails. +func (a *Agent) getHubVersion(sessionCtx ssh.Context) semver.Version { clientVersion := sessionCtx.Value(ssh.ContextKeyClientVersion) if versionStr, ok := clientVersion.(string); ok { - hubVersion, _ = extractHubVersion(versionStr) + hubVersion, _ := extractHubVersion(versionStr) + return hubVersion } - hubVersions[sessionId] = hubVersion - return hubVersion + return semver.Version{} } // handleSession handles an incoming SSH session by gathering system statistics @@ -127,9 +115,8 @@ func (a *Agent) handleSession(s ssh.Session) { a.connectionManager.eventChan <- SSHConnect sessionCtx := s.Context() - sessionID := sessionCtx.SessionID() - hubVersion := a.getHubVersion(sessionID, sessionCtx) + hubVersion := a.getHubVersion(sessionCtx) // Legacy one-shot behavior for older hubs if hubVersion.LT(beszel.MinVersionAgentResponse) { diff --git a/agent/server_test.go b/agent/server_test.go index 1cd52626..274cb45c 100644 --- a/agent/server_test.go +++ b/agent/server_test.go @@ -404,27 +404,23 @@ func TestGetHubVersion(t *testing.T) { clientVersion: "SSH-2.0-beszel_0.12.0", } - // Test first call - should extract and cache version - version := agent.getHubVersion("test-session-123", mockCtx) + // Test first call - should extract version + version := agent.getHubVersion(mockCtx) assert.Equal(t, "0.12.0", version.String()) - // Test second call - should return cached version - mockCtx.clientVersion = "SSH-2.0-beszel_0.11.0" // Change version but should still return cached - version = agent.getHubVersion("test-session-123", mockCtx) - assert.Equal(t, "0.12.0", version.String()) // Should still be cached version - - // Test different session - should extract new version - version = agent.getHubVersion("different-session", mockCtx) + // Test that version reflects the current client version (no stale caching) + mockCtx.clientVersion = "SSH-2.0-beszel_0.11.0" + version = agent.getHubVersion(mockCtx) assert.Equal(t, "0.11.0", version.String()) // Test with invalid version string (non-beszel client) mockCtx.clientVersion = "SSH-2.0-OpenSSH_8.0" - version = agent.getHubVersion("invalid-session", mockCtx) + version = agent.getHubVersion(mockCtx) assert.Equal(t, "0.0.0", version.String()) // Should be empty version for non-beszel clients // Test with no client version mockCtx.clientVersion = "" - version = agent.getHubVersion("no-version-session", mockCtx) + version = agent.getHubVersion(mockCtx) assert.True(t, version.EQ(semver.Version{})) // Should be empty version } @@ -501,9 +497,6 @@ func TestWriteToSessionEncoding(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - // Reset the global hubVersions map to ensure clean state for each test - hubVersions = nil - agent, err := NewAgent("") require.NoError(t, err) @@ -585,39 +578,28 @@ func createTestCombinedData() *system.CombinedData { } } -func TestHubVersionCaching(t *testing.T) { - // Reset the global hubVersions map to ensure clean state - hubVersions = nil - +// TestGetHubVersionConcurrent guards against a regression of the +// "concurrent map writes" panic previously caused by a shared, unsynchronized +// hubVersions cache (see https://github.com/henrygd/beszel/issues/2128). +// getHubVersion no longer shares mutable state between sessions, so calling +// it concurrently from many goroutines must be safe under `go test -race`. +func TestGetHubVersionConcurrent(t *testing.T) { agent, err := NewAgent("") require.NoError(t, err) - ctx1 := &mockSSHContext{ - sessionID: "session1", - clientVersion: "SSH-2.0-beszel_0.12.0", + const goroutines = 50 + var wg sync.WaitGroup + wg.Add(goroutines) + for i := 0; i < goroutines; i++ { + go func(i int) { + defer wg.Done() + ctx := &mockSSHContext{ + sessionID: fmt.Sprintf("session-%d", i), + clientVersion: "SSH-2.0-beszel_0.12.0", + } + version := agent.getHubVersion(ctx) + assert.Equal(t, "0.12.0", version.String()) + }(i) } - ctx2 := &mockSSHContext{ - sessionID: "session2", - clientVersion: "SSH-2.0-beszel_0.11.0", - } - - // First calls should cache the versions - v1 := agent.getHubVersion("session1", ctx1) - v2 := agent.getHubVersion("session2", ctx2) - - assert.Equal(t, "0.12.0", v1.String()) - assert.Equal(t, "0.11.0", v2.String()) - - // Verify caching by changing context but keeping same session ID - ctx1.clientVersion = "SSH-2.0-beszel_0.10.0" - v1Cached := agent.getHubVersion("session1", ctx1) - assert.Equal(t, "0.12.0", v1Cached.String()) // Should still be cached version - - // New session should get new version - ctx3 := &mockSSHContext{ - sessionID: "session3", - clientVersion: "SSH-2.0-beszel_0.13.0", - } - v3 := agent.getHubVersion("session3", ctx3) - assert.Equal(t, "0.13.0", v3.String()) + wg.Wait() }