Add tests for GetSSHKey and handle read errors on key file

beszel/internal/hub/hub.go

@@ -58,7 +58,6 @@ func GetEnv(key string) (value string, exists bool) {
 }
 
 func (h *Hub) StartHub() error {
-
 	h.App.OnServe().BindFunc(func(e *core.ServeEvent) error {
 		// initialize settings / collections
 		if err := h.initialize(e); err != nil {
@@ -158,7 +157,7 @@ func (h *Hub) initialize(e *core.ServeEvent) error {
 	return nil
 }
 
-// startServer starts the server for the Beszel (not PocketBase)
+// startServer sets up the server for Beszel
 func (h *Hub) startServer(se *core.ServeEvent) error {
 	// TODO: exclude dev server from production binary
 	switch h.IsDev() {
@@ -242,8 +241,8 @@ func (h *Hub) registerApiRoutes(se *core.ServeEvent) error {
 }
 
 // generates key pair if it doesn't exist and returns signer
-func (h *Hub) GetSSHKey() (ssh.Signer, error) {
+func (h *Hub) GetSSHKey(dataDir string) (ssh.Signer, error) {
-	privateKeyPath := path.Join(h.DataDir(), "id_ed25519")
+	privateKeyPath := path.Join(dataDir, "id_ed25519")
 
 	// check if the key pair already exists
 	existingKey, err := os.ReadFile(privateKeyPath)
@@ -255,6 +254,9 @@ func (h *Hub) GetSSHKey() (ssh.Signer, error) {
 		pubKeyBytes := ssh.MarshalAuthorizedKey(private.PublicKey())
 		h.pubKey = strings.TrimSuffix(string(pubKeyBytes), "\n")
 		return private, nil
+	} else if !os.IsNotExist(err) {
+		// File exists but couldn't be read for some other reason
+		return nil, fmt.Errorf("failed to read %s: %w", privateKeyPath, err)
 	}
 
 	// Generate the Ed25519 key pair

beszel/internal/hub/hub_test.go

@@ -6,17 +6,26 @@ package hub
 import (
 	"testing"
 
+	"crypto/ed25519"
+	"encoding/pem"
+	"os"
+	"path/filepath"
+	"strings"
+
 	"github.com/pocketbase/pocketbase"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/ssh"
 )
 
-func TestMakeLink(t *testing.T) {
+func getTestHub() *Hub {
-	// The Hub's MakeLink method uses h.Settings().Meta.AppURL.
-	// h.Settings() is a method on h.App (of type core.App).
-	// We use a pocketbase.PocketBase instance as core.App and set its Meta.AppURL
-	// directly for isolated testing of MakeLink.
 	app := pocketbase.New()
-	h := NewHub(app)
+	return NewHub(app)
+}
+
+func TestMakeLink(t *testing.T) {
+	hub := getTestHub()
 
 	tests := []struct {
 		name     string
@@ -94,13 +103,155 @@ func TestMakeLink(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			// Store original and defer restoration if app instance is reused across test functions (good practice)
+			// Store original app URL and restore it after the test
-			originalAppURL := app.Settings().Meta.AppURL
+			originalAppURL := hub.Settings().Meta.AppURL
-			app.Settings().Meta.AppURL = tt.appURL
+			hub.Settings().Meta.AppURL = tt.appURL
-			defer func() { app.Settings().Meta.AppURL = originalAppURL }()
+			defer func() { hub.Settings().Meta.AppURL = originalAppURL }()
 
-			got := h.MakeLink(tt.parts...)
+			got := hub.MakeLink(tt.parts...)
 			assert.Equal(t, tt.expected, got, "MakeLink generated URL does not match expected")
 		})
 	}
 }
+
+func TestGetSSHKey(t *testing.T) {
+	hub := getTestHub()
+
+	// Test Case 1: Key generation (no existing key)
+	t.Run("KeyGeneration", func(t *testing.T) {
+		tempDir := t.TempDir()
+
+		// Ensure pubKey is initially empty or different to ensure GetSSHKey sets it
+		hub.pubKey = ""
+
+		signer, err := hub.GetSSHKey(tempDir)
+		assert.NoError(t, err, "GetSSHKey should not error when generating a new key")
+		assert.NotNil(t, signer, "GetSSHKey should return a non-nil signer")
+
+		// Check if private key file was created
+		privateKeyPath := filepath.Join(tempDir, "id_ed25519")
+		info, err := os.Stat(privateKeyPath)
+		assert.NoError(t, err, "Private key file should be created")
+		assert.False(t, info.IsDir(), "Private key path should be a file, not a directory")
+
+		// Check if h.pubKey was set
+		assert.NotEmpty(t, hub.pubKey, "h.pubKey should be set after key generation")
+		assert.True(t, strings.HasPrefix(hub.pubKey, "ssh-ed25519 "), "h.pubKey should start with 'ssh-ed25519 '")
+
+		// Verify the generated private key is parsable
+		keyData, err := os.ReadFile(privateKeyPath)
+		require.NoError(t, err)
+		_, err = ssh.ParsePrivateKey(keyData)
+		assert.NoError(t, err, "Generated private key should be parsable by ssh.ParsePrivateKey")
+	})
+
+	// Test Case 2: Existing key
+	t.Run("ExistingKey", func(t *testing.T) {
+		tempDir := t.TempDir()
+
+		// Manually create a valid key pair for the test
+		rawPubKey, rawPrivKey, err := ed25519.GenerateKey(nil)
+		require.NoError(t, err, "Failed to generate raw ed25519 key pair for pre-existing key test")
+
+		// Marshal the private key into OpenSSH PEM format
+		pemBlock, err := ssh.MarshalPrivateKey(rawPrivKey, "")
+		require.NoError(t, err, "Failed to marshal private key to PEM block for pre-existing key test")
+
+		privateKeyBytes := pem.EncodeToMemory(pemBlock)
+		require.NotNil(t, privateKeyBytes, "PEM encoded private key bytes should not be nil")
+
+		privateKeyPath := filepath.Join(tempDir, "id_ed25519")
+		err = os.WriteFile(privateKeyPath, privateKeyBytes, 0600)
+		require.NoError(t, err, "Failed to write pre-existing private key")
+
+		// Determine the expected public key string
+		sshPubKey, err := ssh.NewPublicKey(rawPubKey)
+		require.NoError(t, err)
+		expectedPubKeyStr := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(sshPubKey)))
+
+		// Reset h.pubKey to ensure it's set by GetSSHKey from the file
+		hub.pubKey = ""
+
+		signer, err := hub.GetSSHKey(tempDir)
+		assert.NoError(t, err, "GetSSHKey should not error when reading an existing key")
+		assert.NotNil(t, signer, "GetSSHKey should return a non-nil signer for an existing key")
+
+		// Check if h.pubKey was set correctly to the public key from the file
+		assert.Equal(t, expectedPubKeyStr, hub.pubKey, "h.pubKey should match the existing public key")
+
+		// Verify the signer's public key matches the original public key
+		signerPubKey := signer.PublicKey()
+		marshaledSignerPubKey := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(signerPubKey)))
+		assert.Equal(t, expectedPubKeyStr, marshaledSignerPubKey, "Signer's public key should match the existing public key")
+	})
+
+	// Test Case 3: Error cases
+	t.Run("ErrorCases", func(t *testing.T) {
+		tests := []struct {
+			name       string
+			setupFunc  func(dir string) error
+			errorCheck func(t *testing.T, err error)
+		}{
+			{
+				name: "CorruptedKey",
+				setupFunc: func(dir string) error {
+					return os.WriteFile(filepath.Join(dir, "id_ed25519"), []byte("this is not a valid SSH key"), 0600)
+				},
+				errorCheck: func(t *testing.T, err error) {
+					assert.Error(t, err)
+					assert.Contains(t, err.Error(), "ssh: no key found")
+				},
+			},
+			{
+				name: "PermissionDenied",
+				setupFunc: func(dir string) error {
+					// Create the key file
+					keyPath := filepath.Join(dir, "id_ed25519")
+					if err := os.WriteFile(keyPath, []byte("dummy content"), 0600); err != nil {
+						return err
+					}
+					// Make it read-only (can't be opened for writing in case a new key needs to be written)
+					return os.Chmod(keyPath, 0400)
+				},
+				errorCheck: func(t *testing.T, err error) {
+					// On read-only key, the parser will attempt to parse it and fail with "ssh: no key found"
+					assert.Error(t, err)
+				},
+			},
+			{
+				name: "EmptyFile",
+				setupFunc: func(dir string) error {
+					// Create an empty file
+					return os.WriteFile(filepath.Join(dir, "id_ed25519"), []byte{}, 0600)
+				},
+				errorCheck: func(t *testing.T, err error) {
+					assert.Error(t, err)
+					// The error from attempting to parse an empty file
+					assert.Contains(t, err.Error(), "ssh: no key found")
+				},
+			},
+		}
+
+		for _, tc := range tests {
+			t.Run(tc.name, func(t *testing.T) {
+				tempDir := t.TempDir()
+
+				// Setup the test case
+				err := tc.setupFunc(tempDir)
+				require.NoError(t, err, "Setup failed")
+
+				// Reset h.pubKey before each test case
+				hub.pubKey = ""
+
+				// Attempt to get SSH key
+				_, err = hub.GetSSHKey(tempDir)
+
+				// Verify the error
+				tc.errorCheck(t, err)
+
+				// Check that pubKey was not set in error cases
+				assert.Empty(t, hub.pubKey, "h.pubKey should not be set if there was an error")
+			})
+		}
+	})
+}

beszel/internal/hub/systems/systems.go

@@ -46,7 +46,7 @@ type System struct {
 
 type hubLike interface {
 	core.App
-	GetSSHKey() (ssh.Signer, error)
+	GetSSHKey(dataDir string) (ssh.Signer, error)
 	HandleSystemAlerts(systemRecord *core.Record, data *system.CombinedData) error
 	HandleStatusAlerts(status string, systemRecord *core.Record) error
 }
@@ -362,7 +362,7 @@ func (sys *System) fetchDataFromAgent() (*system.CombinedData, error) {
 
 // createSSHClientConfig initializes the ssh config for the system manager
 func (sm *SystemManager) createSSHClientConfig() error {
-	privateKey, err := sm.hub.GetSSHKey()
+	privateKey, err := sm.hub.GetSSHKey(sm.hub.DataDir())
 	if err != nil {
 		return err
 	}