beszel-ipv6/agent/connection_manager.go

package agent

import (
	"context"
	"errors"
	"log/slog"
	"net"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/henrygd/beszel/agent/health"
	"github.com/henrygd/beszel/agent/utils"
	"github.com/henrygd/beszel/internal/entities/system"
)

// 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
	// mu guards State and isConnecting, which are read and written from both
	// the main event loop and the goroutine spawned by connect().
	mu             sync.Mutex
	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
	ConnectionType system.ConnectionType
}

// 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()
	}
}

// getState returns the current connection state.
func (c *ConnectionManager) getState() ConnectionState {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.State
}

// setConnecting sets the isConnecting flag and reports its previous value.
func (c *ConnectionManager) setConnecting(v bool) (previous bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	previous = c.isConnecting
	c.isConnecting = v
	return previous
}

// isConnectingNow reports whether a reconnection attempt is currently in flight.
func (c *ConnectionManager) isConnectingNow() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.isConnecting
}

// 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 {
		var caCertErr *caCertFileError
		if errors.As(err, &caCertErr) {
			return err
		}
		disableSSH, _ := utils.GetEnv("DISABLE_SSH")
		if errors.Is(err, errNoHubURL) && disableSSH != "true" {
			// SSH-only mode: the hub dials the agent, so there is nothing to warn
			// about. With SSH also disabled there is no connection method at all,
			// so that case still warns.
			slog.Debug("WebSocket client not configured", "err", err)
		} else {
			slog.Warn("Error creating WebSocket client", "err", err)
		}
	}
	c.wsClient = wsClient

	c.serverOptions = serverOptions
	c.eventChan = make(chan ConnectionEvent, 1)

	// signal handling for shutdown
	sigCtx, stopSignals := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stopSignals()

	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:
			// skip if connect() is still running its own attempt
			if !c.isConnectingNow() {
				_ = c.startWebSocketConnection()
			}
		case <-healthTicker:
			_ = health.Update()
		case <-sigCtx.Done():
			slog.Info("Shutting down", "cause", context.Cause(sigCtx))
			return c.stop()
		}
	}
}

// stop does not stop the connection manager itself, just any active connections. The manager will attempt to reconnect after stopping, so this should only be called immediately before shutting down the entire agent.
//
// If we need or want to expose a graceful Stop method in the future, do something like this to actually stop the manager:
//
//	func (c *ConnectionManager) Start(serverOptions ServerOptions) error {
//		ctx, cancel := context.WithCancel(context.Background())
//		c.cancel = cancel
//
//		for {
//			select {
//			case <-ctx.Done():
//				return c.stop()
//			}
//		}
//	}
//
//	func (c *ConnectionManager) Stop() {
//		c.cancel()
//	}
func (c *ConnectionManager) stop() error {
	_ = c.agent.StopServer()
	c.agent.monitorManager.Stop()
	c.closeWebSocket()
	c.agent.cleanupSensorShadow()
	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:
		if c.getState() == Disconnected {
			c.handleStateChange(SSHConnected)
		}
	case WebSocketDisconnect:
		if c.getState() == WebSocketConnected {
			c.handleStateChange(Disconnected)
		}
	case SSHDisconnect:
		if c.getState() == 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) {
	c.mu.Lock()
	if c.State == newState {
		c.mu.Unlock()
		return
	}
	c.State = newState
	c.mu.Unlock()

	switch newState {
	case WebSocketConnected:
		slog.Info("WebSocket connected", "host", c.wsClient.hubURL.Host)
		c.ConnectionType = system.ConnectionTypeWebSocket
		c.stopWsTicker()
		_ = c.agent.StopServer()
		c.setConnecting(false)
	case SSHConnected:
		// stop new ws connection attempts
		slog.Info("SSH connection established")
		c.ConnectionType = system.ConnectionTypeSSH
		c.stopWsTicker()
		c.setConnecting(false)
	case Disconnected:
		c.ConnectionType = system.ConnectionTypeNone
		// Always keep the ticker running while disconnected. A pending WebSocket
		// handshake started by connect() can fail asynchronously (e.g. the hub
		// closes the socket, or the deadline set in OnOpen expires) after
		// connect() has already returned with a nil error, in which case the
		// ticker would otherwise never get re-armed and the agent would stop
		// retrying entirely (#2326).
		c.startWsTicker()
		if c.setConnecting(true) {
			// Already handling reconnection, avoid duplicate attempts
			return
		}
		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.setConnecting(true)
	defer c.setConnecting(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 {
		if shouldExitOnErr(err) {
			time.Sleep(2 * time.Second) // prevent tight restart loop
			_ = c.stop()
			os.Exit(1)
		}
		if c.getState() == Disconnected {
			c.startSSHServer()
			c.startWsTicker()
		}
	}
}

// startWebSocketConnection attempts to establish a WebSocket connection to the hub.
func (c *ConnectionManager) startWebSocketConnection() error {
	if c.getState() != 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.getState() == 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()
	}
}

// shouldExitOnErr checks if the error is a DNS resolution failure and if the
// EXIT_ON_DNS_ERROR env var is set. https://github.com/henrygd/beszel/issues/1924.
func shouldExitOnErr(err error) bool {
	if val, _ := utils.GetEnv("EXIT_ON_DNS_ERROR"); val == "true" {
		if opErr, ok := errors.AsType[*net.OpError](err); ok {
			return strings.Contains(opErr.Err.Error(), "lookup")
		}
	}
	return false
}