jnd/beszel-ipv6
1
0
Fork 0
mirror of https://github.com/henrygd/beszel.git synced 2025-12-17 10:46:16 +01:00
Code Issues Packages Projects Releases 2 Wiki Activity

Add CBOR and agent initiated WebSocket connections (#51, #490, #646, #845, etc)

Browse Source 
- Add version exchange between hub and agent.
- Introduce ConnectionManager for managing WebSocket and SSH connections.
- Implement fingerprint generation and storage in agent.
- Create expiry map package to store universal tokens.
- Update config.yml configuration to include tokens.
- Enhance system management with new methods for handling system states and alerts.
- Update front-end components to support token / fingerprint management features.
- Introduce utility functions for token generation and hub URL retrieval.

Co-authored-by: nhas <jordanatararimu@gmail.com>
This commit is contained in:
henrygd
2025-07-08 18:41:36 -04:00
parent 99d61a0193
commit 402a1584d7
41 changed files with 5567 additions and 989 deletions
Download Patch File Download Diff File Expand all files Collapse all files

99
beszel/internal/agent/agent.go
View File

@@ -4,34 +4,55 @@ package agent
import (
"beszel"
"beszel/internal/entities/system"
"crypto/sha256"
"encoding/hex"
"log/slog"
"os"
"path/filepath"
"strings"
"sync"
"time"
"github.com/gliderlabs/ssh"
"github.com/shirou/gopsutil/v4/host"
gossh "golang.org/x/crypto/ssh"
)
type Agent struct {
sync.Mutex // Used to lock agent while collecting data
debug bool // true if LOG_LEVEL is set to debug
zfs bool // true if system has arcstats
memCalc string // Memory calculation formula
fsNames []string // List of filesystem device names being monitored
fsStats map[string]*system.FsStats // Keeps track of disk stats for each filesystem
netInterfaces map[string]struct{} // Stores all valid network interfaces
netIoStats system.NetIoStats // Keeps track of bandwidth usage
dockerManager *dockerManager // Manages Docker API requests
sensorConfig *SensorConfig // Sensors config
systemInfo system.Info // Host system info
gpuManager *GPUManager // Manages GPU data
cache *SessionCache // Cache for system stats based on primary session ID
sync.Mutex // Used to lock agent while collecting data
debug bool // true if LOG_LEVEL is set to debug
zfs bool // true if system has arcstats
memCalc string // Memory calculation formula
fsNames []string // List of filesystem device names being monitored
fsStats map[string]*system.FsStats // Keeps track of disk stats for each filesystem
netInterfaces map[string]struct{} // Stores all valid network interfaces
netIoStats system.NetIoStats // Keeps track of bandwidth usage
dockerManager *dockerManager // Manages Docker API requests
sensorConfig *SensorConfig // Sensors config
systemInfo system.Info // Host system info
gpuManager *GPUManager // Manages GPU data
cache *SessionCache // Cache for system stats based on primary session ID
connectionManager *ConnectionManager // Channel to signal connection events
server *ssh.Server // SSH server
dataDir string // Directory for persisting data
keys []gossh.PublicKey // SSH public keys
}
func NewAgent() *Agent {
agent := &Agent{
// NewAgent creates a new agent with the given data directory for persisting data.
// If the data directory is not set, it will attempt to find the optimal directory.
func NewAgent(dataDir string) (agent *Agent, err error) {
agent = &Agent{
fsStats: make(map[string]*system.FsStats),
cache: NewSessionCache(69 * time.Second),
}
agent.dataDir, err = getDataDir(dataDir)
if err != nil {
slog.Warn("Data directory not found")
} else {
slog.Info("Data directory", "path", agent.dataDir)
}
agent.memCalc, _ = GetEnv("MEM_CALC")
agent.sensorConfig = agent.newSensorConfig()
// Set up slog with a log level determined by the LOG_LEVEL env var
@@ -49,10 +70,19 @@ func NewAgent() *Agent {
slog.Debug(beszel.Version)
// initialize system info / docker manager
// initialize system info
agent.initializeSystemInfo()
// initialize connection manager
agent.connectionManager = newConnectionManager(agent)
// initialize disk info
agent.initializeDiskInfo()
// initialize net io stats
agent.initializeNetIoStats()
// initialize docker manager
agent.dockerManager = newDockerManager(agent)
// initialize GPU manager
@@ -67,7 +97,7 @@ func NewAgent() *Agent {
slog.Debug("Stats", "data", agent.gatherStats(""))
}
return agent
return agent, nil
}
// GetEnv retrieves an environment variable with a "BESZEL_AGENT_" prefix, or falls back to the unprefixed key.
@@ -115,3 +145,38 @@ func (a *Agent) gatherStats(sessionID string) *system.CombinedData {
a.cache.Set(sessionID, cachedData)
return cachedData
}
// StartAgent initializes and starts the agent with optional WebSocket connection
func (a *Agent) Start(serverOptions ServerOptions) error {
a.keys = serverOptions.Keys
return a.connectionManager.Start(serverOptions)
}
func (a *Agent) getFingerprint() string {
// first look for a fingerprint in the data directory
if a.dataDir != "" {
if fp, err := os.ReadFile(filepath.Join(a.dataDir, "fingerprint")); err == nil {
return string(fp)
}
}
// if no fingerprint is found, generate one
fingerprint, err := host.HostID()
if err != nil || fingerprint == "" {
fingerprint = a.systemInfo.Hostname + a.systemInfo.CpuModel
}
// hash fingerprint
sum := sha256.Sum256([]byte(fingerprint))
fingerprint = hex.EncodeToString(sum[:24])
// save fingerprint to data directory
if a.dataDir != "" {
err = os.WriteFile(filepath.Join(a.dataDir, "fingerprint"), []byte(fingerprint), 0644)
if err != nil {
slog.Warn("Failed to save fingerprint", "err", err)
}
}
return fingerprint
}

9
beszel/internal/agent/agent_test_helpers.go Normal file
View File

@@ -0,0 +1,9 @@
//go:build testing
// +build testing
package agent
// TESTING ONLY: GetConnectionManager is a helper function to get the connection manager for testing.
func (a *Agent) GetConnectionManager() *ConnectionManager {
return a.connectionManager
}

243
beszel/internal/agent/client.go Normal file
View File

@@ -0,0 +1,243 @@
package agent
import (
"beszel"
"beszel/internal/common"
"crypto/tls"
"errors"
"fmt"
"log/slog"
"net"
"net/http"
"net/url"
"path"
"strings"
"time"
"github.com/fxamacker/cbor/v2"
"github.com/lxzan/gws"
"golang.org/x/crypto/ssh"
)
const (
wsDeadline = 70 * time.Second
)
// WebSocketClient manages the WebSocket connection between the agent and hub.
// It handles authentication, message routing, and connection lifecycle management.
type WebSocketClient struct {
gws.BuiltinEventHandler
options *gws.ClientOption // WebSocket client configuration options
agent *Agent // Reference to the parent agent
Conn *gws.Conn // Active WebSocket connection
hubURL *url.URL // Parsed hub URL for connection
token string // Authentication token for hub registration
fingerprint string // System fingerprint for identification
hubRequest *common.HubRequest[cbor.RawMessage] // Reusable request structure for message parsing
lastConnectAttempt time.Time // Timestamp of last connection attempt
hubVerified bool // Whether the hub has been cryptographically verified
}
// newWebSocketClient creates a new WebSocket client for the given agent.
// It reads configuration from environment variables and validates the hub URL.
func newWebSocketClient(agent *Agent) (client *WebSocketClient, err error) {
hubURLStr, exists := GetEnv("HUB_URL")
if !exists {
return nil, errors.New("HUB_URL environment variable not set")
}
client = &WebSocketClient{}
client.hubURL, err = url.Parse(hubURLStr)
if err != nil {
return nil, errors.New("invalid hub URL")
}
// get registration token
client.token, _ = GetEnv("TOKEN")
if client.token == "" {
return nil, errors.New("TOKEN environment variable not set")
}
client.agent = agent
client.hubRequest = &common.HubRequest[cbor.RawMessage]{}
client.fingerprint = agent.getFingerprint()
return client, nil
}
// getOptions returns the WebSocket client options, creating them if necessary.
// It configures the connection URL, TLS settings, and authentication headers.
func (client *WebSocketClient) getOptions() *gws.ClientOption {
if client.options != nil {
return client.options
}
// update the hub url to use websocket scheme and api path
if client.hubURL.Scheme == "https" {
client.hubURL.Scheme = "wss"
} else {
client.hubURL.Scheme = "ws"
}
client.hubURL.Path = path.Join(client.hubURL.Path, "api/beszel/agent-connect")
client.options = &gws.ClientOption{
Addr: client.hubURL.String(),
TlsConfig: &tls.Config{InsecureSkipVerify: true},
RequestHeader: http.Header{
"User-Agent": []string{getUserAgent()},
"X-Token": []string{client.token},
"X-Beszel": []string{beszel.Version},
},
}
return client.options
}
// Connect establishes a WebSocket connection to the hub.
// It closes any existing connection before attempting to reconnect.
func (client *WebSocketClient) Connect() (err error) {
client.lastConnectAttempt = time.Now()
// make sure previous connection is closed
client.Close()
client.Conn, _, err = gws.NewClient(client, client.getOptions())
if err != nil {
return err
}
go client.Conn.ReadLoop()
return nil
}
// OnOpen handles WebSocket connection establishment.
// It sets a deadline for the connection to prevent hanging.
func (client *WebSocketClient) OnOpen(conn *gws.Conn) {
conn.SetDeadline(time.Now().Add(wsDeadline))
}
// OnClose handles WebSocket connection closure.
// It logs the closure reason and notifies the connection manager.
func (client *WebSocketClient) OnClose(conn *gws.Conn, err error) {
slog.Warn("Connection closed", "err", strings.TrimPrefix(err.Error(), "gws: "))
client.agent.connectionManager.eventChan <- WebSocketDisconnect
}
// OnMessage handles incoming WebSocket messages from the hub.
// It decodes CBOR messages and routes them to appropriate handlers.
func (client *WebSocketClient) OnMessage(conn *gws.Conn, message *gws.Message) {
defer message.Close()
conn.SetDeadline(time.Now().Add(wsDeadline))
if message.Opcode != gws.OpcodeBinary {
return
}
if err := cbor.NewDecoder(message.Data).Decode(client.hubRequest); err != nil {
slog.Error("Error parsing message", "err", err)
return
}
if err := client.handleHubRequest(client.hubRequest); err != nil {
slog.Error("Error handling message", "err", err)
}
}
// OnPing handles WebSocket ping frames.
// It responds with a pong and updates the connection deadline.
func (client *WebSocketClient) OnPing(conn *gws.Conn, message []byte) {
conn.SetDeadline(time.Now().Add(wsDeadline))
conn.WritePong(message)
}
// handleAuthChallenge verifies the authenticity of the hub and returns the system's fingerprint.
func (client *WebSocketClient) handleAuthChallenge(msg *common.HubRequest[cbor.RawMessage]) (err error) {
var authRequest common.FingerprintRequest
if err := cbor.Unmarshal(msg.Data, &authRequest); err != nil {
return err
}
if err := client.verifySignature(authRequest.Signature); err != nil {
return err
}
client.hubVerified = true
client.agent.connectionManager.eventChan <- WebSocketConnect
response := &common.FingerprintResponse{
Fingerprint: client.fingerprint,
}
if authRequest.NeedSysInfo {
response.Hostname = client.agent.systemInfo.Hostname
serverAddr := client.agent.connectionManager.serverOptions.Addr
_, response.Port, _ = net.SplitHostPort(serverAddr)
}
return client.sendMessage(response)
}
// verifySignature verifies the signature of the token using the public keys.
func (client *WebSocketClient) verifySignature(signature []byte) (err error) {
for _, pubKey := range client.agent.keys {
sig := ssh.Signature{
Format: pubKey.Type(),
Blob: signature,
}
if err = pubKey.Verify([]byte(client.token), &sig); err == nil {
return nil
}
}
return errors.New("invalid signature - check KEY value")
}
// Close closes the WebSocket connection gracefully.
// This method is safe to call multiple times.
func (client *WebSocketClient) Close() {
if client.Conn != nil {
_ = client.Conn.WriteClose(1000, nil)
}
}
// handleHubRequest routes the request to the appropriate handler.
// It ensures the hub is verified before processing most requests.
func (client *WebSocketClient) handleHubRequest(msg *common.HubRequest[cbor.RawMessage]) error {
if !client.hubVerified && msg.Action != common.CheckFingerprint {
return errors.New("hub not verified")
}
switch msg.Action {
case common.GetData:
return client.sendSystemData()
case common.CheckFingerprint:
return client.handleAuthChallenge(msg)
}
return nil
}
// sendSystemData gathers and sends current system statistics to the hub.
func (client *WebSocketClient) sendSystemData() error {
sysStats := client.agent.gatherStats(client.token)
return client.sendMessage(sysStats)
}
// sendMessage encodes the given data to CBOR and sends it as a binary message over the WebSocket connection to the hub.
func (client *WebSocketClient) sendMessage(data any) error {
bytes, err := cbor.Marshal(data)
if err != nil {
return err
}
return client.Conn.WriteMessage(gws.OpcodeBinary, bytes)
}
// getUserAgent returns one of two User-Agent strings based on current time.
// This is used to avoid being blocked by Cloudflare or other anti-bot measures.
func getUserAgent() string {
const (
uaBase = "Mozilla/5.0 (%s) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36"
uaWindows = "Windows NT 11.0; Win64; x64"
uaMac = "Macintosh; Intel Mac OS X 14_0_0"
)
if time.Now().UnixNano()%2 == 0 {
return fmt.Sprintf(uaBase, uaWindows)
}
return fmt.Sprintf(uaBase, uaMac)
}

220
beszel/internal/agent/connection_manager.go Normal file
View File

@@ -0,0 +1,220 @@
package agent
import (
"beszel/internal/agent/health"
"errors"
"log/slog"
"os"
"os/signal"
"syscall"
"time"
)
// ConnectionManager manages the connection state and events for the agent.
// It handles both WebSocket and SSH connections, automatically switching between
// them based on availability and managing reconnection attempts.
type ConnectionManager struct {
agent *Agent // Reference to the parent agent
State ConnectionState // Current connection state
eventChan chan ConnectionEvent // Channel for connection events
wsClient *WebSocketClient // WebSocket client for hub communication
serverOptions ServerOptions // Configuration for SSH server
wsTicker *time.Ticker // Ticker for WebSocket connection attempts
isConnecting bool // Prevents multiple simultaneous reconnection attempts
}
// ConnectionState represents the current connection state of the agent.
type ConnectionState uint8
// ConnectionEvent represents connection-related events that can occur.
type ConnectionEvent uint8
// Connection states
const (
Disconnected ConnectionState = iota // No active connection
WebSocketConnected // Connected via WebSocket
SSHConnected // Connected via SSH
)
// Connection events
const (
WebSocketConnect ConnectionEvent = iota // WebSocket connection established
WebSocketDisconnect // WebSocket connection lost
SSHConnect // SSH connection established
SSHDisconnect // SSH connection lost
)
const wsTickerInterval = 10 * time.Second
// newConnectionManager creates a new connection manager for the given agent.
func newConnectionManager(agent *Agent) *ConnectionManager {
cm := &ConnectionManager{
agent: agent,
State: Disconnected,
}
return cm
}
// startWsTicker starts or resets the WebSocket connection attempt ticker.
func (c *ConnectionManager) startWsTicker() {
if c.wsTicker == nil {
c.wsTicker = time.NewTicker(wsTickerInterval)
} else {
c.wsTicker.Reset(wsTickerInterval)
}
}
// stopWsTicker stops the WebSocket connection attempt ticker.
func (c *ConnectionManager) stopWsTicker() {
if c.wsTicker != nil {
c.wsTicker.Stop()
}
}
// Start begins connection attempts and enters the main event loop.
// It handles connection events, periodic health updates, and graceful shutdown.
func (c *ConnectionManager) Start(serverOptions ServerOptions) error {
if c.eventChan != nil {
return errors.New("already started")
}
wsClient, err := newWebSocketClient(c.agent)
if err != nil {
slog.Warn("Error creating WebSocket client", "err", err)
}
c.wsClient = wsClient
c.serverOptions = serverOptions
c.eventChan = make(chan ConnectionEvent, 1)
// signal handling for shutdown
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
c.startWsTicker()
c.connect()
// update health status immediately and every 90 seconds
_ = health.Update()
healthTicker := time.Tick(90 * time.Second)
for {
select {
case connectionEvent := <-c.eventChan:
c.handleEvent(connectionEvent)
case <-c.wsTicker.C:
_ = c.startWebSocketConnection()
case <-healthTicker:
_ = health.Update()
case <-sigChan:
slog.Info("Shutting down")
_ = c.agent.StopServer()
c.closeWebSocket()
return health.CleanUp()
}
}
}
// handleEvent processes connection events and updates the connection state accordingly.
func (c *ConnectionManager) handleEvent(event ConnectionEvent) {
switch event {
case WebSocketConnect:
c.handleStateChange(WebSocketConnected)
case SSHConnect:
c.handleStateChange(SSHConnected)
case WebSocketDisconnect:
if c.State == WebSocketConnected {
c.handleStateChange(Disconnected)
}
case SSHDisconnect:
if c.State == SSHConnected {
c.handleStateChange(Disconnected)
}
}
}
// handleStateChange updates the connection state and performs necessary actions
// based on the new state, including stopping services and initiating reconnections.
func (c *ConnectionManager) handleStateChange(newState ConnectionState) {
if c.State == newState {
return
}
c.State = newState
switch newState {
case WebSocketConnected:
slog.Info("WebSocket connected", "host", c.wsClient.hubURL.Host)
c.stopWsTicker()
_ = c.agent.StopServer()
c.isConnecting = false
case SSHConnected:
// stop new ws connection attempts
slog.Info("SSH connection established")
c.stopWsTicker()
c.isConnecting = false
case Disconnected:
if c.isConnecting {
// Already handling reconnection, avoid duplicate attempts
return
}
c.isConnecting = true
slog.Warn("Disconnected from hub")
// make sure old ws connection is closed
c.closeWebSocket()
// reconnect
go c.connect()
}
}
// connect handles the connection logic with proper delays and priority.
// It attempts WebSocket connection first, falling back to SSH server if needed.
func (c *ConnectionManager) connect() {
c.isConnecting = true
defer func() {
c.isConnecting = false
}()
if c.wsClient != nil && time.Since(c.wsClient.lastConnectAttempt) < 5*time.Second {
time.Sleep(5 * time.Second)
}
// Try WebSocket first, if it fails, start SSH server
err := c.startWebSocketConnection()
if err != nil && c.State == Disconnected {
c.startSSHServer()
c.startWsTicker()
}
}
// startWebSocketConnection attempts to establish a WebSocket connection to the hub.
func (c *ConnectionManager) startWebSocketConnection() error {
if c.State != Disconnected {
return errors.New("already connected")
}
if c.wsClient == nil {
return errors.New("WebSocket client not initialized")
}
if time.Since(c.wsClient.lastConnectAttempt) < 5*time.Second {
return errors.New("already connecting")
}
err := c.wsClient.Connect()
if err != nil {
slog.Warn("WebSocket connection failed", "err", err)
c.closeWebSocket()
}
return err
}
// startSSHServer starts the SSH server if the agent is currently disconnected.
func (c *ConnectionManager) startSSHServer() {
if c.State == Disconnected {
go c.agent.StartServer(c.serverOptions)
}
}
// closeWebSocket closes the WebSocket connection if it exists.
func (c *ConnectionManager) closeWebSocket() {
if c.wsClient != nil {
c.wsClient.Close()
}
}

315
beszel/internal/agent/connection_manager_test.go Normal file
View File

@@ -0,0 +1,315 @@
//go:build testing
// +build testing
package agent
import (
"crypto/ed25519"
"fmt"
"net"
"net/url"
"os"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/crypto/ssh"
)
func createTestAgent(t *testing.T) *Agent {
dataDir := t.TempDir()
agent, err := NewAgent(dataDir)
require.NoError(t, err)
return agent
}
func createTestServerOptions(t *testing.T) ServerOptions {
// Generate test key pair
_, privKey, err := ed25519.GenerateKey(nil)
require.NoError(t, err)
sshPubKey, err := ssh.NewPublicKey(privKey.Public().(ed25519.PublicKey))
require.NoError(t, err)
// Find available port
listener, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
port := listener.Addr().(*net.TCPAddr).Port
listener.Close()
return ServerOptions{
Network: "tcp",
Addr: fmt.Sprintf("127.0.0.1:%d", port),
Keys: []ssh.PublicKey{sshPubKey},
}
}
// TestConnectionManager_NewConnectionManager tests connection manager creation
func TestConnectionManager_NewConnectionManager(t *testing.T) {
agent := createTestAgent(t)
cm := newConnectionManager(agent)
assert.NotNil(t, cm, "Connection manager should not be nil")
assert.Equal(t, agent, cm.agent, "Agent reference should be set")
assert.Equal(t, Disconnected, cm.State, "Initial state should be Disconnected")
assert.Nil(t, cm.eventChan, "Event channel should be nil initially")
assert.Nil(t, cm.wsClient, "WebSocket client should be nil initially")
assert.Nil(t, cm.wsTicker, "WebSocket ticker should be nil initially")
assert.False(t, cm.isConnecting, "isConnecting should be false initially")
}
// TestConnectionManager_StateTransitions tests basic state transitions
func TestConnectionManager_StateTransitions(t *testing.T) {
agent := createTestAgent(t)
cm := agent.connectionManager
initialState := cm.State
cm.wsClient = &WebSocketClient{
hubURL: &url.URL{
Host: "localhost:8080",
},
}
assert.NotNil(t, cm, "Connection manager should not be nil")
assert.Equal(t, Disconnected, initialState, "Initial state should be Disconnected")
// Test state transitions
cm.handleStateChange(WebSocketConnected)
assert.Equal(t, WebSocketConnected, cm.State, "State should change to WebSocketConnected")
cm.handleStateChange(SSHConnected)
assert.Equal(t, SSHConnected, cm.State, "State should change to SSHConnected")
cm.handleStateChange(Disconnected)
assert.Equal(t, Disconnected, cm.State, "State should change to Disconnected")
// Test that same state doesn't trigger changes
cm.State = WebSocketConnected
cm.handleStateChange(WebSocketConnected)
assert.Equal(t, WebSocketConnected, cm.State, "Same state should not trigger change")
}
// TestConnectionManager_EventHandling tests event handling logic
func TestConnectionManager_EventHandling(t *testing.T) {
agent := createTestAgent(t)
cm := agent.connectionManager
cm.wsClient = &WebSocketClient{
hubURL: &url.URL{
Host: "localhost:8080",
},
}
testCases := []struct {
name string
initialState ConnectionState
event ConnectionEvent
expectedState ConnectionState
}{
{
name: "WebSocket connect from disconnected",
initialState: Disconnected,
event: WebSocketConnect,
expectedState: WebSocketConnected,
},
{
name: "SSH connect from disconnected",
initialState: Disconnected,
event: SSHConnect,
expectedState: SSHConnected,
},
{
name: "WebSocket disconnect from connected",
initialState: WebSocketConnected,
event: WebSocketDisconnect,
expectedState: Disconnected,
},
{
name: "SSH disconnect from connected",
initialState: SSHConnected,
event: SSHDisconnect,
expectedState: Disconnected,
},
{
name: "WebSocket disconnect from SSH connected (no change)",
initialState: SSHConnected,
event: WebSocketDisconnect,
expectedState: SSHConnected,
},
{
name: "SSH disconnect from WebSocket connected (no change)",
initialState: WebSocketConnected,
event: SSHDisconnect,
expectedState: WebSocketConnected,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
cm.State = tc.initialState
cm.handleEvent(tc.event)
assert.Equal(t, tc.expectedState, cm.State, "State should match expected after event")
})
}
}
// TestConnectionManager_TickerManagement tests WebSocket ticker management
func TestConnectionManager_TickerManagement(t *testing.T) {
agent := createTestAgent(t)
cm := agent.connectionManager
// Test starting ticker
cm.startWsTicker()
assert.NotNil(t, cm.wsTicker, "Ticker should be created")
// Test stopping ticker (should not panic)
assert.NotPanics(t, func() {
cm.stopWsTicker()
}, "Stopping ticker should not panic")
// Test stopping nil ticker (should not panic)
cm.wsTicker = nil
assert.NotPanics(t, func() {
cm.stopWsTicker()
}, "Stopping nil ticker should not panic")
// Test restarting ticker
cm.startWsTicker()
assert.NotNil(t, cm.wsTicker, "Ticker should be recreated")
// Test resetting existing ticker
firstTicker := cm.wsTicker
cm.startWsTicker()
assert.Equal(t, firstTicker, cm.wsTicker, "Same ticker instance should be reused")
cm.stopWsTicker()
}
// TestConnectionManager_WebSocketConnectionFlow tests WebSocket connection logic
func TestConnectionManager_WebSocketConnectionFlow(t *testing.T) {
if testing.Short() {
t.Skip("Skipping WebSocket connection test in short mode")
}
agent := createTestAgent(t)
cm := agent.connectionManager
// Test WebSocket connection without proper environment
err := cm.startWebSocketConnection()
assert.Error(t, err, "WebSocket connection should fail without proper environment")
assert.Equal(t, Disconnected, cm.State, "State should remain Disconnected after failed connection")
// Test with invalid URL
os.Setenv("BESZEL_AGENT_HUB_URL", "invalid-url")
os.Setenv("BESZEL_AGENT_TOKEN", "test-token")
defer func() {
os.Unsetenv("BESZEL_AGENT_HUB_URL")
os.Unsetenv("BESZEL_AGENT_TOKEN")
}()
// Test with missing token
os.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080")
os.Unsetenv("BESZEL_AGENT_TOKEN")
_, err2 := newWebSocketClient(agent)
assert.Error(t, err2, "WebSocket client creation should fail without token")
}
// TestConnectionManager_ReconnectionLogic tests reconnection prevention logic
func TestConnectionManager_ReconnectionLogic(t *testing.T) {
agent := createTestAgent(t)
cm := agent.connectionManager
cm.eventChan = make(chan ConnectionEvent, 1)
// Test that isConnecting flag prevents duplicate reconnection attempts
// Start from connected state, then simulate disconnect
cm.State = WebSocketConnected
cm.isConnecting = false
// First disconnect should trigger reconnection logic
cm.handleStateChange(Disconnected)
assert.Equal(t, Disconnected, cm.State, "Should change to disconnected")
assert.True(t, cm.isConnecting, "Should set isConnecting flag")
}
// TestConnectionManager_ConnectWithRateLimit tests connection rate limiting
func TestConnectionManager_ConnectWithRateLimit(t *testing.T) {
agent := createTestAgent(t)
cm := agent.connectionManager
// Set up environment for WebSocket client creation
os.Setenv("BESZEL_AGENT_HUB_URL", "ws://localhost:8080")
os.Setenv("BESZEL_AGENT_TOKEN", "test-token")
defer func() {
os.Unsetenv("BESZEL_AGENT_HUB_URL")
os.Unsetenv("BESZEL_AGENT_TOKEN")
}()
// Create WebSocket client
wsClient, err := newWebSocketClient(agent)
require.NoError(t, err)
cm.wsClient = wsClient
// Set recent connection attempt
cm.wsClient.lastConnectAttempt = time.Now()
// Test that connection is rate limited
err = cm.startWebSocketConnection()
assert.Error(t, err, "Should error due to rate limiting")
assert.Contains(t, err.Error(), "already connecting", "Error should indicate rate limiting")
// Test connection after rate limit expires
cm.wsClient.lastConnectAttempt = time.Now().Add(-10 * time.Second)
err = cm.startWebSocketConnection()
// This will fail due to no actual server, but should not be rate limited
assert.Error(t, err, "Connection should fail but not due to rate limiting")
assert.NotContains(t, err.Error(), "already connecting", "Error should not indicate rate limiting")
}
// TestConnectionManager_StartWithInvalidConfig tests starting with invalid configuration
func TestConnectionManager_StartWithInvalidConfig(t *testing.T) {
agent := createTestAgent(t)
cm := agent.connectionManager
serverOptions := createTestServerOptions(t)
// Test starting when already started
cm.eventChan = make(chan ConnectionEvent, 5)
err := cm.Start(serverOptions)
assert.Error(t, err, "Should error when starting already started connection manager")
}
// TestConnectionManager_CloseWebSocket tests WebSocket closing
func TestConnectionManager_CloseWebSocket(t *testing.T) {
agent := createTestAgent(t)
cm := agent.connectionManager
// Test closing when no WebSocket client exists
assert.NotPanics(t, func() {
cm.closeWebSocket()
}, "Should not panic when closing nil WebSocket client")
// Set up environment and create WebSocket client
os.Setenv("BESZEL_AGENT_HUB_URL", "ws://localhost:8080")
os.Setenv("BESZEL_AGENT_TOKEN", "test-token")
defer func() {
os.Unsetenv("BESZEL_AGENT_HUB_URL")
os.Unsetenv("BESZEL_AGENT_TOKEN")
}()
wsClient, err := newWebSocketClient(agent)
require.NoError(t, err)
cm.wsClient = wsClient
// Test closing when WebSocket client exists
assert.NotPanics(t, func() {
cm.closeWebSocket()
}, "Should not panic when closing WebSocket client")
}
// TestConnectionManager_ConnectFlow tests the connect method
func TestConnectionManager_ConnectFlow(t *testing.T) {
agent := createTestAgent(t)
cm := agent.connectionManager
// Test connect without WebSocket client
assert.NotPanics(t, func() {
cm.connect()
}, "Connect should not panic without WebSocket client")
}

123
beszel/internal/agent/server.go
View File

@@ -1,25 +1,44 @@
package agent
import (
"beszel"
"beszel/internal/common"
"beszel/internal/entities/system"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net"
"os"
"strings"
"time"
"github.com/blang/semver"
"github.com/fxamacker/cbor/v2"
"github.com/gliderlabs/ssh"
gossh "golang.org/x/crypto/ssh"
)
// ServerOptions contains configuration options for starting the SSH server.
type ServerOptions struct {
Addr string
Network string
Keys []gossh.PublicKey
Addr string // Network address to listen on (e.g., ":45876" or "/path/to/socket")
Network string // Network type ("tcp" or "unix")
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
// is already running or if there's an issue starting the server.
func (a *Agent) StartServer(opts ServerOptions) error {
if a.server != nil {
return errors.New("server already started")
}
slog.Info("Starting SSH server", "addr", opts.Addr, "network", opts.Network)
if opts.Network == "unix" {
@@ -37,7 +56,9 @@ func (a *Agent) StartServer(opts ServerOptions) error {
defer ln.Close()
// base config (limit to allowed algorithms)
config := &gossh.ServerConfig{}
config := &gossh.ServerConfig{
ServerVersion: fmt.Sprintf("SSH-2.0-%s_%s", beszel.AppName, beszel.Version),
}
config.KeyExchanges = common.DefaultKeyExchanges
config.MACs = common.DefaultMACs
config.Ciphers = common.DefaultCiphers
@@ -45,42 +66,92 @@ func (a *Agent) StartServer(opts ServerOptions) error {
// set default handler
ssh.Handle(a.handleSession)
server := ssh.Server{
a.server = &ssh.Server{
ServerConfigCallback: func(ctx ssh.Context) *gossh.ServerConfig {
return config
},
// check public key(s)
PublicKeyHandler: func(ctx ssh.Context, key ssh.PublicKey) bool {
remoteAddr := ctx.RemoteAddr()
for _, pubKey := range opts.Keys {
if ssh.KeysEqual(key, pubKey) {
slog.Info("SSH connected", "addr", remoteAddr)
return true
}
}
slog.Warn("Invalid SSH key", "addr", remoteAddr)
return false
},
// disable pty
PtyCallback: func(ctx ssh.Context, pty ssh.Pty) bool {
return false
},
// log failed connections
ConnectionFailedCallback: func(conn net.Conn, err error) {
slog.Warn("Failed connection attempt", "addr", conn.RemoteAddr().String(), "err", err)
},
// close idle connections after 70 seconds
IdleTimeout: 70 * time.Second,
}
// Start SSH server on the listener
return server.Serve(ln)
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
clientVersion := sessionCtx.Value(ssh.ContextKeyClientVersion)
if versionStr, ok := clientVersion.(string); ok {
hubVersion, _ = extractHubVersion(versionStr)
}
hubVersions[sessionId] = hubVersion
return hubVersion
}
// handleSession handles an incoming SSH session by gathering system statistics
// and sending them to the hub. It signals connection events, determines the
// appropriate encoding format based on hub version, and exits with appropriate
// status codes.
func (a *Agent) handleSession(s ssh.Session) {
slog.Debug("New session", "client", s.RemoteAddr())
stats := a.gatherStats(s.Context().SessionID())
if err := json.NewEncoder(s).Encode(stats); err != nil {
a.connectionManager.eventChan <- SSHConnect
sessionCtx := s.Context()
sessionID := sessionCtx.SessionID()
hubVersion := a.getHubVersion(sessionID, sessionCtx)
stats := a.gatherStats(sessionID)
err := a.writeToSession(s, stats, hubVersion)
if err != nil {
slog.Error("Error encoding stats", "err", err, "stats", stats)
s.Exit(1)
return
} else {
s.Exit(0)
}
s.Exit(0)
}
// writeToSession encodes and writes system statistics to the session.
// It chooses between CBOR and JSON encoding based on the hub version,
// using CBOR for newer versions and JSON for legacy compatibility.
func (a *Agent) writeToSession(w io.Writer, stats *system.CombinedData, hubVersion semver.Version) error {
if hubVersion.GTE(beszel.MinVersionCbor) {
return cbor.NewEncoder(w).Encode(stats)
}
return json.NewEncoder(w).Encode(stats)
}
// extractHubVersion extracts the beszel version from SSH client version string.
// Expected format: "SSH-2.0-beszel_X.Y.Z" or "beszel_X.Y.Z"
func extractHubVersion(versionString string) (semver.Version, error) {
_, after, _ := strings.Cut(versionString, "_")
return semver.Parse(after)
}
// ParseKeys parses a string containing SSH public keys in authorized_keys format.
@@ -103,7 +174,9 @@ func ParseKeys(input string) ([]gossh.PublicKey, error) {
return parsedKeys, nil
}
// GetAddress gets the address to listen on or connect to from environment variables or default value.
// GetAddress determines the network address to listen on from various sources.
// It checks the provided address, then environment variables (LISTEN, PORT),
// and finally defaults to ":45876".
func GetAddress(addr string) string {
if addr == "" {
addr, _ = GetEnv("LISTEN")
@@ -122,7 +195,9 @@ func GetAddress(addr string) string {
return addr
}
// GetNetwork returns the network type to use based on the address
// GetNetwork determines the network type based on the address format.
// It checks the NETWORK environment variable first, then infers from
// the address format: addresses starting with "/" are "unix", others are "tcp".
func GetNetwork(addr string) string {
if network, ok := GetEnv("NETWORK"); ok && network != "" {
return network
@@ -132,3 +207,17 @@ func GetNetwork(addr string) string {
}
return "tcp"
}
// StopServer stops the SSH server if it's running.
// It returns an error if the server is not running or if there's an error stopping it.
func (a *Agent) StopServer() error {
if a.server == nil {
return errors.New("SSH server not running")
}
slog.Info("Stopping SSH server")
_ = a.server.Close()
a.server = nil
a.connectionManager.eventChan <- SSHDisconnect
return nil
}

356
beszel/internal/agent/server_test.go
View File

@@ -1,34 +1,43 @@
package agent
import (
"beszel/internal/entities/container"
"beszel/internal/entities/system"
"context"
"crypto/ed25519"
"encoding/json"
"fmt"
"net"
"os"
"path/filepath"
"strings"
"sync"
"testing"
"time"
"github.com/blang/semver"
"github.com/fxamacker/cbor/v2"
"github.com/gliderlabs/ssh"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/crypto/ssh"
gossh "golang.org/x/crypto/ssh"
)
func TestStartServer(t *testing.T) {
// Generate a test key pair
pubKey, privKey, err := ed25519.GenerateKey(nil)
require.NoError(t, err)
signer, err := ssh.NewSignerFromKey(privKey)
signer, err := gossh.NewSignerFromKey(privKey)
require.NoError(t, err)
sshPubKey, err := ssh.NewPublicKey(pubKey)
sshPubKey, err := gossh.NewPublicKey(pubKey)
require.NoError(t, err)
// Generate a different key pair for bad key test
badPubKey, badPrivKey, err := ed25519.GenerateKey(nil)
require.NoError(t, err)
badSigner, err := ssh.NewSignerFromKey(badPrivKey)
badSigner, err := gossh.NewSignerFromKey(badPrivKey)
require.NoError(t, err)
sshBadPubKey, err := ssh.NewPublicKey(badPubKey)
sshBadPubKey, err := gossh.NewPublicKey(badPubKey)
require.NoError(t, err)
socketFile := filepath.Join(t.TempDir(), "beszel-test.sock")
@@ -46,7 +55,7 @@ func TestStartServer(t *testing.T) {
config: ServerOptions{
Network: "tcp",
Addr: ":45987",
Keys: []ssh.PublicKey{sshPubKey},
Keys: []gossh.PublicKey{sshPubKey},
},
},
{
@@ -54,7 +63,7 @@ func TestStartServer(t *testing.T) {
config: ServerOptions{
Network: "tcp4",
Addr: "127.0.0.1:45988",
Keys: []ssh.PublicKey{sshPubKey},
Keys: []gossh.PublicKey{sshPubKey},
},
},
{
@@ -62,7 +71,7 @@ func TestStartServer(t *testing.T) {
config: ServerOptions{
Network: "tcp6",
Addr: "[::1]:45989",
Keys: []ssh.PublicKey{sshPubKey},
Keys: []gossh.PublicKey{sshPubKey},
},
},
{
@@ -70,7 +79,7 @@ func TestStartServer(t *testing.T) {
config: ServerOptions{
Network: "unix",
Addr: socketFile,
Keys: []ssh.PublicKey{sshPubKey},
Keys: []gossh.PublicKey{sshPubKey},
},
setup: func() error {
// Create a socket file that should be removed
@@ -89,7 +98,7 @@ func TestStartServer(t *testing.T) {
config: ServerOptions{
Network: "tcp",
Addr: ":45987",
Keys: []ssh.PublicKey{sshBadPubKey},
Keys: []gossh.PublicKey{sshBadPubKey},
},
wantErr: true,
errContains: "ssh: handshake failed",
@@ -99,7 +108,7 @@ func TestStartServer(t *testing.T) {
config: ServerOptions{
Network: "tcp",
Addr: ":45987",
Keys: []ssh.PublicKey{sshPubKey},
Keys: []gossh.PublicKey{sshPubKey},
},
},
}
@@ -115,7 +124,8 @@ func TestStartServer(t *testing.T) {
defer tt.cleanup()
}
agent := NewAgent()
agent, err := NewAgent("")
require.NoError(t, err)
// Start server in a goroutine since it blocks
errChan := make(chan error, 1)
@@ -127,8 +137,7 @@ func TestStartServer(t *testing.T) {
time.Sleep(100 * time.Millisecond)
// Try to connect to verify server is running
var client *ssh.Client
var err error
var client *gossh.Client
// Choose the appropriate signer based on the test case
testSigner := signer
@@ -136,23 +145,23 @@ func TestStartServer(t *testing.T) {
testSigner = badSigner
}
sshClientConfig := &ssh.ClientConfig{
sshClientConfig := &gossh.ClientConfig{
User: "a",
Auth: []ssh.AuthMethod{
ssh.PublicKeys(testSigner),
Auth: []gossh.AuthMethod{
gossh.PublicKeys(testSigner),
},
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
HostKeyCallback: gossh.InsecureIgnoreHostKey(),
Timeout: 4 * time.Second,
}
switch tt.config.Network {
case "unix":
client, err = ssh.Dial("unix", tt.config.Addr, sshClientConfig)
client, err = gossh.Dial("unix", tt.config.Addr, sshClientConfig)
default:
if !strings.Contains(tt.config.Addr, ":") {
tt.config.Addr = ":" + tt.config.Addr
}
client, err = ssh.Dial("tcp", tt.config.Addr, sshClientConfig)
client, err = gossh.Dial("tcp", tt.config.Addr, sshClientConfig)
}
if tt.wantErr {
@@ -287,3 +296,310 @@ func TestParseInvalidKey(t *testing.T) {
t.Fatalf("Expected error message to contain '%s', got: %v", expectedErrMsg, err)
}
}
/////////////////////////////////////////////////////////////////
//////////////////// Hub Version Tests //////////////////////////
/////////////////////////////////////////////////////////////////
func TestExtractHubVersion(t *testing.T) {
tests := []struct {
name string
clientVersion string
expectedVersion string
expectError bool
}{
{
name: "valid beszel client version with underscore",
clientVersion: "SSH-2.0-beszel_0.11.1",
expectedVersion: "0.11.1",
expectError: false,
},
{
name: "valid beszel client version with beta",
clientVersion: "SSH-2.0-beszel_1.0.0-beta",
expectedVersion: "1.0.0-beta",
expectError: false,
},
{
name: "valid beszel client version with rc",
clientVersion: "SSH-2.0-beszel_0.12.0-rc1",
expectedVersion: "0.12.0-rc1",
expectError: false,
},
{
name: "different SSH client",
clientVersion: "SSH-2.0-OpenSSH_8.0",
expectedVersion: "8.0",
expectError: true,
},
{
name: "malformed version string without underscore",
clientVersion: "SSH-2.0-beszel",
expectError: true,
},
{
name: "empty version string",
clientVersion: "",
expectError: true,
},
{
name: "version string with underscore but no version",
clientVersion: "beszel_",
expectedVersion: "",
expectError: true,
},
{
name: "version with patch and build metadata",
clientVersion: "SSH-2.0-beszel_1.2.3+build.123",
expectedVersion: "1.2.3+build.123",
expectError: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result, err := extractHubVersion(tt.clientVersion)
if tt.expectError {
assert.Error(t, err)
return
}
require.NoError(t, err)
assert.Equal(t, tt.expectedVersion, result.String())
})
}
}
/////////////////////////////////////////////////////////////////
/////////////// Hub Version Detection Tests ////////////////////
/////////////////////////////////////////////////////////////////
func TestGetHubVersion(t *testing.T) {
agent, err := NewAgent("")
require.NoError(t, err)
// Mock SSH context that implements the ssh.Context interface
mockCtx := &mockSSHContext{
sessionID: "test-session-123",
clientVersion: "SSH-2.0-beszel_0.12.0",
}
// Test first call - should extract and cache version
version := agent.getHubVersion("test-session-123", 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)
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)
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)
assert.True(t, version.EQ(semver.Version{})) // Should be empty version
}
// mockSSHContext implements ssh.Context for testing
type mockSSHContext struct {
context.Context
sync.Mutex
sessionID string
clientVersion string
}
func (m *mockSSHContext) SessionID() string {
return m.sessionID
}
func (m *mockSSHContext) ClientVersion() string {
return m.clientVersion
}
func (m *mockSSHContext) ServerVersion() string {
return "SSH-2.0-beszel_test"
}
func (m *mockSSHContext) Value(key interface{}) interface{} {
if key == ssh.ContextKeyClientVersion {
return m.clientVersion
}
return nil
}
func (m *mockSSHContext) User() string { return "test-user" }
func (m *mockSSHContext) RemoteAddr() net.Addr { return nil }
func (m *mockSSHContext) LocalAddr() net.Addr { return nil }
func (m *mockSSHContext) Permissions() *ssh.Permissions { return nil }
func (m *mockSSHContext) SetValue(key, value interface{}) {}
/////////////////////////////////////////////////////////////////
/////////////// CBOR vs JSON Encoding Tests ////////////////////
/////////////////////////////////////////////////////////////////
// TestWriteToSessionEncoding tests that writeToSession actually encodes data in the correct format
func TestWriteToSessionEncoding(t *testing.T) {
tests := []struct {
name string
hubVersion string
expectedUsesCbor bool
}{
{
name: "old hub version should use JSON",
hubVersion: "0.11.1",
expectedUsesCbor: false,
},
{
name: "non-beta release should use CBOR",
hubVersion: "0.12.0",
expectedUsesCbor: true,
},
{
name: "even newer hub version should use CBOR",
hubVersion: "0.16.4",
expectedUsesCbor: true,
},
{
name: "beta version below release threshold should use JSON",
hubVersion: "0.12.0-beta0",
expectedUsesCbor: false,
},
{
name: "matching beta version should use CBOR",
hubVersion: "0.12.0-beta1",
expectedUsesCbor: true,
},
}
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)
// Parse the test version
version, err := semver.Parse(tt.hubVersion)
require.NoError(t, err)
// Create test data to encode
testData := createTestCombinedData()
var buf strings.Builder
err = agent.writeToSession(&buf, testData, version)
require.NoError(t, err)
encodedData := buf.String()
require.NotEmpty(t, encodedData)
// Verify the encoding format by attempting to decode
if tt.expectedUsesCbor {
var decodedCbor system.CombinedData
err = cbor.Unmarshal([]byte(encodedData), &decodedCbor)
assert.NoError(t, err, "Should be valid CBOR data")
var decodedJson system.CombinedData
err = json.Unmarshal([]byte(encodedData), &decodedJson)
assert.Error(t, err, "Should not be valid JSON data")
assert.Equal(t, testData.Info.Hostname, decodedCbor.Info.Hostname)
assert.Equal(t, testData.Stats.Cpu, decodedCbor.Stats.Cpu)
} else {
// Should be JSON - try to decode as JSON
var decodedJson system.CombinedData
err = json.Unmarshal([]byte(encodedData), &decodedJson)
assert.NoError(t, err, "Should be valid JSON data")
var decodedCbor system.CombinedData
err = cbor.Unmarshal([]byte(encodedData), &decodedCbor)
assert.Error(t, err, "Should not be valid CBOR data")
// Verify the decoded JSON data matches our test data
assert.Equal(t, testData.Info.Hostname, decodedJson.Info.Hostname)
assert.Equal(t, testData.Stats.Cpu, decodedJson.Stats.Cpu)
// Verify it looks like JSON (starts with '{' and contains readable field names)
assert.True(t, strings.HasPrefix(encodedData, "{"), "JSON should start with '{'")
assert.Contains(t, encodedData, `"info"`, "JSON should contain readable field names")
assert.Contains(t, encodedData, `"stats"`, "JSON should contain readable field names")
}
})
}
}
// Helper function to create test data for encoding tests
func createTestCombinedData() *system.CombinedData {
return &system.CombinedData{
Stats: system.Stats{
Cpu: 25.5,
Mem: 8589934592, // 8GB
MemUsed: 4294967296, // 4GB
MemPct: 50.0,
DiskTotal: 1099511627776, // 1TB
DiskUsed: 549755813888, // 512GB
DiskPct: 50.0,
},
Info: system.Info{
Hostname: "test-host",
Cores: 8,
CpuModel: "Test CPU Model",
Uptime: 3600,
AgentVersion: "0.12.0",
Os: system.Linux,
},
Containers: []*container.Stats{
{
Name: "test-container",
Cpu: 10.5,
Mem: 1073741824, // 1GB
},
},
}
}
func TestHubVersionCaching(t *testing.T) {
// Reset the global hubVersions map to ensure clean state
hubVersions = nil
agent, err := NewAgent("")
require.NoError(t, err)
ctx1 := &mockSSHContext{
sessionID: "session1",
clientVersion: "SSH-2.0-beszel_0.12.0",
}
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())
}