From b4a3cb9ce6009c411a0d36fe4f6732f3af36bca6 Mon Sep 17 00:00:00 2001
From: Henry Dollman <hank@henrygd.me>
Date: Fri, 30 Aug 2024 19:06:16 -0400
Subject: [PATCH] add timeout to ssh session creation to avoid hanging

---
 beszel/internal/hub/hub.go | 35 ++++++++++++++++++++++++++++++++---
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/beszel/internal/hub/hub.go b/beszel/internal/hub/hub.go
index e06b91e..74756a7 100644
--- a/beszel/internal/hub/hub.go
+++ b/beszel/internal/hub/hub.go
@@ -6,6 +6,7 @@ import (
 	"beszel/internal/entities/system"
 	"beszel/internal/records"
 	"beszel/site"
+	"context"
 	"crypto/ed25519"
 	"encoding/pem"
 	"fmt"
@@ -246,8 +247,10 @@ func (h *Hub) updateSystem(record *models.Record) {
 		// create system connection
 		client, err = h.createSystemConnection(record)
 		if err != nil {
-			h.app.Logger().Error("Failed to connect:", "err", err.Error(), "system", record.GetString("host"), "port", record.GetString("port"))
-			h.updateSystemStatus(record, "down")
+			if record.GetString("status") != "down" {
+				h.app.Logger().Error("Failed to connect:", "err", err.Error(), "system", record.GetString("host"), "port", record.GetString("port"))
+				h.updateSystemStatus(record, "down")
+			}
 			return
 		}
 		h.connectionLock.Lock()
@@ -350,7 +353,7 @@ func (h *Hub) createSSHClientConfig() error {
 }
 
 func requestJsonFromAgent(client *ssh.Client, systemData *system.CombinedData) error {
-	session, err := client.NewSession()
+	session, err := newSessionWithTimeout(client, 5*time.Second)
 	if err != nil {
 		return fmt.Errorf("bad client")
 	}
@@ -377,6 +380,32 @@ func requestJsonFromAgent(client *ssh.Client, systemData *system.CombinedData) e
 	return nil
 }
 
+// Adds timeout to SSH session creation to avoid hanging in case of network issues
+func newSessionWithTimeout(client *ssh.Client, timeout time.Duration) (*ssh.Session, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+
+	// use goroutine to create the session
+	sessionChan := make(chan *ssh.Session, 1)
+	errChan := make(chan error, 1)
+	go func() {
+		if session, err := client.NewSession(); err != nil {
+			errChan <- err
+		} else {
+			sessionChan <- session
+		}
+	}()
+
+	select {
+	case session := <-sessionChan:
+		return session, nil
+	case err := <-errChan:
+		return nil, err
+	case <-ctx.Done():
+		return nil, fmt.Errorf("session creation timed out")
+	}
+}
+
 func (h *Hub) getSSHKey() ([]byte, error) {
 	dataDir := h.app.DataDir()
 	// check if the key pair already exists