diff --git a/agent/connection_manager.go b/agent/connection_manager.go index fc91a9b1..30b5c543 100644 --- a/agent/connection_manager.go +++ b/agent/connection_manager.go @@ -4,11 +4,15 @@ import ( "context" "errors" "log/slog" + "net" + "os" "os/signal" + "strings" "syscall" "time" "github.com/henrygd/beszel/agent/health" + "github.com/henrygd/beszel/agent/utils" "github.com/henrygd/beszel/internal/entities/system" ) @@ -111,13 +115,36 @@ func (c *ConnectionManager) Start(serverOptions ServerOptions) error { _ = health.Update() case <-sigCtx.Done(): slog.Info("Shutting down", "cause", context.Cause(sigCtx)) - _ = c.agent.StopServer() - c.closeWebSocket() - return health.CleanUp() + 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.closeWebSocket() + return health.CleanUp() +} + // handleEvent processes connection events and updates the connection state accordingly. func (c *ConnectionManager) handleEvent(event ConnectionEvent) { switch event { @@ -185,9 +212,16 @@ func (c *ConnectionManager) connect() { // Try WebSocket first, if it fails, start SSH server err := c.startWebSocketConnection() - if err != nil && c.State == Disconnected { - c.startSSHServer() - c.startWsTicker() + if err != nil { + if shouldExitOnErr(err) { + time.Sleep(2 * time.Second) // prevent tight restart loop + _ = c.stop() + os.Exit(1) + } + if c.State == Disconnected { + c.startSSHServer() + c.startWsTicker() + } } } @@ -224,3 +258,14 @@ func (c *ConnectionManager) closeWebSocket() { 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 +} diff --git a/agent/connection_manager_test.go b/agent/connection_manager_test.go index 7a9ca4a0..8aba1579 100644 --- a/agent/connection_manager_test.go +++ b/agent/connection_manager_test.go @@ -4,6 +4,7 @@ package agent import ( "crypto/ed25519" + "errors" "fmt" "net" "net/url" @@ -298,3 +299,65 @@ func TestConnectionManager_ConnectFlow(t *testing.T) { cm.connect() }, "Connect should not panic without WebSocket client") } + +func TestShouldExitOnErr(t *testing.T) { + createDialErr := func(msg string) error { + return &net.OpError{ + Op: "dial", + Net: "tcp", + Err: errors.New(msg), + } + } + + tests := []struct { + name string + err error + envValue string + expected bool + }{ + { + name: "no env var", + err: createDialErr("lookup lkahsdfasdf: no such host"), + envValue: "", + expected: false, + }, + { + name: "env var false", + err: createDialErr("lookup lkahsdfasdf: no such host"), + envValue: "false", + expected: false, + }, + { + name: "env var true, matching error", + err: createDialErr("lookup lkahsdfasdf: no such host"), + envValue: "true", + expected: true, + }, + { + name: "env var true, matching error with extra context", + err: createDialErr("lookup beszel.server.lan on [::1]:53: read udp [::1]:44557->[::1]:53: read: connection refused"), + envValue: "true", + expected: true, + }, + { + name: "env var true, non-matching error", + err: errors.New("connection refused"), + envValue: "true", + expected: false, + }, + { + name: "env var true, dial but not lookup", + err: createDialErr("connection timeout"), + envValue: "true", + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("EXIT_ON_DNS_ERROR", tt.envValue) + result := shouldExitOnErr(tt.err) + assert.Equal(t, tt.expected, result) + }) + } +} diff --git a/internal/cmd/agent/agent.go b/internal/cmd/agent/agent.go index a27908db..ea2dbd69 100644 --- a/internal/cmd/agent/agent.go +++ b/internal/cmd/agent/agent.go @@ -195,6 +195,6 @@ func main() { } if err := a.Start(serverConfig); err != nil { - log.Fatal("Failed to start server: ", err) + log.Fatal("Failed to start: ", err) } }