From f4fb7a89e57ea1a634804b7da598b3203c4b20a6 Mon Sep 17 00:00:00 2001 From: henrygd Date: Thu, 8 May 2025 18:54:14 -0400 Subject: [PATCH] Add tests for GetSSHKey and handle read errors on key file --- beszel/internal/hub/hub.go | 10 +- beszel/internal/hub/hub_test.go | 173 +++++++++++++++++++++++-- beszel/internal/hub/systems/systems.go | 4 +- 3 files changed, 170 insertions(+), 17 deletions(-) diff --git a/beszel/internal/hub/hub.go b/beszel/internal/hub/hub.go index a20e550..54eb98e 100644 --- a/beszel/internal/hub/hub.go +++ b/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) { - privateKeyPath := path.Join(h.DataDir(), "id_ed25519") +func (h *Hub) GetSSHKey(dataDir string) (ssh.Signer, error) { + 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 diff --git a/beszel/internal/hub/hub_test.go b/beszel/internal/hub/hub_test.go index 1031bc3..fc1f4ea 100644 --- a/beszel/internal/hub/hub_test.go +++ b/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) { - // 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. +func getTestHub() *Hub { 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) - originalAppURL := app.Settings().Meta.AppURL - app.Settings().Meta.AppURL = tt.appURL - defer func() { app.Settings().Meta.AppURL = originalAppURL }() + // Store original app URL and restore it after the test + originalAppURL := hub.Settings().Meta.AppURL + hub.Settings().Meta.AppURL = tt.appURL + 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") + }) + } + }) +} diff --git a/beszel/internal/hub/systems/systems.go b/beszel/internal/hub/systems/systems.go index d69fc6f..20e7b35 100644 --- a/beszel/internal/hub/systems/systems.go +++ b/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 }