From b4c221dfb7dd12b369b7524a829f6cc52956080a Mon Sep 17 00:00:00 2001 From: Caleb Brose <5447118+cmbrose@users.noreply.github.com> Date: Tue, 5 Nov 2024 22:30:15 +0000 Subject: [PATCH 1/4] Create the automatic key when specified with -i --- pkg/cmd/codespace/ssh.go | 27 ++++++++++++++++----------- pkg/cmd/codespace/ssh_test.go | 35 +++++++++++++++++++++++++++++------ pkg/ssh/ssh_keys.go | 6 +++--- 3 files changed, 48 insertions(+), 20 deletions(-) diff --git a/pkg/cmd/codespace/ssh.go b/pkg/cmd/codespace/ssh.go index c6c3943e2..44ed30eb0 100644 --- a/pkg/cmd/codespace/ssh.go +++ b/pkg/cmd/codespace/ssh.go @@ -20,7 +20,6 @@ import ( "github.com/cli/cli/v2/internal/codespaces/api" "github.com/cli/cli/v2/internal/codespaces/portforwarder" "github.com/cli/cli/v2/internal/codespaces/rpc" - "github.com/cli/cli/v2/internal/config" "github.com/cli/cli/v2/pkg/cmdutil" "github.com/cli/cli/v2/pkg/ssh" "github.com/cli/safeexec" @@ -336,10 +335,20 @@ func selectSSHKeys( return nil, false, errors.New("missing value to -i argument") } + privateKeyPath := args[i+1] + + // The --config setup will set the automatic key with -i, but it might not actually be created, so we need to ensure that here + if automaticPrivateKeyPath, _ := automaticPrivateKeyPath(sshContext); automaticPrivateKeyPath == privateKeyPath { + _, err := generateAutomaticSSHKeys(sshContext) + if err != nil { + return nil, false, fmt.Errorf("generating automatic keypair: %w", err) + } + } + // User manually specified an identity file so just trust it is correct return &ssh.KeyPair{ - PrivateKeyPath: args[i+1], - PublicKeyPath: args[i+1] + ".pub", + PrivateKeyPath: privateKeyPath, + PublicKeyPath: privateKeyPath + ".pub", }, false, nil } @@ -636,7 +645,8 @@ func (a *App) printOpenSSHConfig(ctx context.Context, opts sshOptions) (err erro return fmt.Errorf("error formatting template: %w", err) } - automaticIdentityFilePath, err := automaticPrivateKeyPath() + sshContext := ssh.Context{} + automaticIdentityFilePath, err := automaticPrivateKeyPath(sshContext) if err != nil { return fmt.Errorf("error finding .ssh directory: %w", err) } @@ -683,13 +693,8 @@ func (a *App) printOpenSSHConfig(ctx context.Context, opts sshOptions) (err erro return status } -func automaticPrivateKeyPath() (string, error) { - sshDir, err := config.HomeDirPath(".ssh") - if err != nil { - return "", err - } - - return path.Join(sshDir, automaticPrivateKeyName), nil +func automaticPrivateKeyPath(sshContext ssh.Context) (string, error) { + return path.Join(sshContext.ConfigDir, automaticPrivateKeyName), nil } type cpOptions struct { diff --git a/pkg/cmd/codespace/ssh_test.go b/pkg/cmd/codespace/ssh_test.go index b76d304d7..d06bcda8e 100644 --- a/pkg/cmd/codespace/ssh_test.go +++ b/pkg/cmd/codespace/ssh_test.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "os" + "path" "path/filepath" "strings" "testing" @@ -125,6 +126,10 @@ func TestGenerateAutomaticSSHKeys(t *testing.T) { } func TestSelectSSHKeys(t *testing.T) { + // This string will be subsituted in sshArgs for test cases + // This is to work around the temp test ssh dir not being known until the test is executing + substituteSSHDir := "SUB_SSH_DIR" + tests := []struct { sshDirFiles []string sshConfigKeys []string @@ -139,7 +144,7 @@ func TestSelectSSHKeys(t *testing.T) { wantKeyPair: &ssh.KeyPair{PrivateKeyPath: "custom-private-key", PublicKeyPath: "custom-private-key.pub"}, }, { - sshArgs: []string{"-i", automaticPrivateKeyName}, + sshArgs: []string{"-i", path.Join(substituteSSHDir, automaticPrivateKeyName)}, wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, }, { @@ -226,7 +231,12 @@ func TestSelectSSHKeys(t *testing.T) { t.Fatalf("could not write test config %v", err) } - tt.sshArgs = append([]string{"-F", configPath}, tt.sshArgs...) + var subbedSSHArgs []string + for _, arg := range tt.sshArgs { + subbedSSHArgs = append(subbedSSHArgs, strings.Replace(arg, substituteSSHDir, sshDir, -1)) + } + + tt.sshArgs = append([]string{"-F", configPath}, subbedSSHArgs...) gotKeyPair, gotShouldAddArg, err := selectSSHKeys(context.Background(), sshContext, tt.sshArgs, sshOptions{profile: tt.profileOpt}) @@ -254,11 +264,24 @@ func TestSelectSSHKeys(t *testing.T) { } // Strip the dir (sshDir) from the gotKeyPair paths so that they match wantKeyPair (which doesn't know the directory) - gotKeyPair.PrivateKeyPath = filepath.Base(gotKeyPair.PrivateKeyPath) - gotKeyPair.PublicKeyPath = filepath.Base(gotKeyPair.PublicKeyPath) + gotKeyPairJustFileNames := &ssh.KeyPair{ + PrivateKeyPath: filepath.Base(gotKeyPair.PrivateKeyPath), + PublicKeyPath: filepath.Base(gotKeyPair.PublicKeyPath), + } - if fmt.Sprintf("%v", gotKeyPair) != fmt.Sprintf("%v", tt.wantKeyPair) { - t.Errorf("Want selectSSHKeys result to be %v, got %v", tt.wantKeyPair, gotKeyPair) + if fmt.Sprintf("%v", gotKeyPairJustFileNames) != fmt.Sprintf("%v", tt.wantKeyPair) { + t.Errorf("Want selectSSHKeys result to be %v, got %v", tt.wantKeyPair, gotKeyPairJustFileNames) + } + + // If the automatic key pair is selected, it needs to exist no matter what + if strings.Contains(tt.wantKeyPair.PrivateKeyPath, automaticPrivateKeyName) { + if _, err := os.Stat(gotKeyPair.PrivateKeyPath); err != nil { + t.Errorf("Expected automatic key pair private key to exist, but it did not") + } + + if _, err := os.Stat(gotKeyPair.PublicKeyPath); err != nil { + t.Errorf("Expected automatic key pair public key to exist, but it did not") + } } } } diff --git a/pkg/ssh/ssh_keys.go b/pkg/ssh/ssh_keys.go index c750b608a..0b9e12f1f 100644 --- a/pkg/ssh/ssh_keys.go +++ b/pkg/ssh/ssh_keys.go @@ -26,7 +26,7 @@ type KeyPair struct { var ErrKeyAlreadyExists = errors.New("SSH key already exists") func (c *Context) LocalPublicKeys() ([]string, error) { - sshDir, err := c.sshDir() + sshDir, err := c.SshDir() if err != nil { return nil, err } @@ -45,7 +45,7 @@ func (c *Context) GenerateSSHKey(keyName string, passphrase string) (*KeyPair, e return nil, err } - sshDir, err := c.sshDir() + sshDir, err := c.SshDir() if err != nil { return nil, err } @@ -79,7 +79,7 @@ func (c *Context) GenerateSSHKey(keyName string, passphrase string) (*KeyPair, e return &keyPair, nil } -func (c *Context) sshDir() (string, error) { +func (c *Context) SshDir() (string, error) { if c.ConfigDir != "" { return c.ConfigDir, nil } From 940560acf2721a0f09b6795d813d84a1192986d5 Mon Sep 17 00:00:00 2001 From: Caleb Brose <5447118+cmbrose@users.noreply.github.com> Date: Wed, 6 Nov 2024 16:13:29 +0000 Subject: [PATCH 2/4] Fix ssh directory --- pkg/cmd/auth/shared/login_flow_test.go | 24 ++-- pkg/cmd/codespace/ssh.go | 7 +- pkg/cmd/codespace/ssh_test.go | 158 ++++++++++++++++++++++++- pkg/ssh/ssh_keys.go | 25 ++-- 4 files changed, 191 insertions(+), 23 deletions(-) diff --git a/pkg/cmd/auth/shared/login_flow_test.go b/pkg/cmd/auth/shared/login_flow_test.go index 8c8ba5d72..31cf2c107 100644 --- a/pkg/cmd/auth/shared/login_flow_test.go +++ b/pkg/cmd/auth/shared/login_flow_test.go @@ -101,10 +101,7 @@ func TestLogin(t *testing.T) { // simulate that the public key file has been generated _ = os.WriteFile(keyFile+".pub", []byte("PUBKEY asdf"), 0600) }) - opts.sshContext = ssh.Context{ - ConfigDir: dir, - KeygenExe: "ssh-keygen", - } + opts.sshContext = ssh.NewContextForTests(dir, "ssh-keygen") }, wantsConfig: map[string]string{ "example.com:user": "monalisa", @@ -112,6 +109,11 @@ func TestLogin(t *testing.T) { "example.com:git_protocol": "ssh", }, stderrAssert: func(t *testing.T, opts *LoginOptions, stderr string) { + sshDir, err := opts.sshContext.SshDir() + if err != nil { + t.Errorf("Could not load ssh config dir: %v", err) + } + assert.Equal(t, heredoc.Docf(` Tip: you can generate a Personal Access Token here https://example.com/settings/tokens The minimum required scopes are 'repo', 'read:org', 'admin:public_key'. @@ -119,7 +121,7 @@ func TestLogin(t *testing.T) { ✓ Configured git protocol ✓ Uploaded the SSH key to your GitHub account: %s ✓ Logged in as monalisa - `, filepath.Join(opts.sshContext.ConfigDir, "id_ed25519.pub")), stderr) + `, filepath.Join(sshDir, "id_ed25519.pub")), stderr) }, }, { @@ -179,10 +181,7 @@ func TestLogin(t *testing.T) { // simulate that the public key file has been generated _ = os.WriteFile(keyFile+".pub", []byte("PUBKEY asdf"), 0600) }) - opts.sshContext = ssh.Context{ - ConfigDir: dir, - KeygenExe: "ssh-keygen", - } + opts.sshContext = ssh.NewContextForTests(dir, "ssh-keygen") }, wantsConfig: map[string]string{ "example.com:user": "monalisa", @@ -190,6 +189,11 @@ func TestLogin(t *testing.T) { "example.com:git_protocol": "ssh", }, stderrAssert: func(t *testing.T, opts *LoginOptions, stderr string) { + sshDir, err := opts.sshContext.SshDir() + if err != nil { + t.Errorf("Could not load ssh config dir: %v", err) + } + assert.Equal(t, heredoc.Docf(` Tip: you can generate a Personal Access Token here https://example.com/settings/tokens The minimum required scopes are 'repo', 'read:org', 'admin:public_key'. @@ -197,7 +201,7 @@ func TestLogin(t *testing.T) { ✓ Configured git protocol ✓ Uploaded the SSH key to your GitHub account: %s ✓ Logged in as monalisa - `, filepath.Join(opts.sshContext.ConfigDir, "id_ed25519.pub")), stderr) + `, filepath.Join(sshDir, "id_ed25519.pub")), stderr) }, }, { diff --git a/pkg/cmd/codespace/ssh.go b/pkg/cmd/codespace/ssh.go index 44ed30eb0..b79ec68c9 100644 --- a/pkg/cmd/codespace/ssh.go +++ b/pkg/cmd/codespace/ssh.go @@ -694,7 +694,12 @@ func (a *App) printOpenSSHConfig(ctx context.Context, opts sshOptions) (err erro } func automaticPrivateKeyPath(sshContext ssh.Context) (string, error) { - return path.Join(sshContext.ConfigDir, automaticPrivateKeyName), nil + sshDir, err := sshContext.SshDir() + if err != nil { + return "", err + } + + return path.Join(sshDir, automaticPrivateKeyName), nil } type cpOptions struct { diff --git a/pkg/cmd/codespace/ssh_test.go b/pkg/cmd/codespace/ssh_test.go index d06bcda8e..e656e0ecd 100644 --- a/pkg/cmd/codespace/ssh_test.go +++ b/pkg/cmd/codespace/ssh_test.go @@ -69,9 +69,7 @@ func TestGenerateAutomaticSSHKeys(t *testing.T) { for _, tt := range tests { dir := t.TempDir() - sshContext := ssh.Context{ - ConfigDir: dir, - } + sshContext := ssh.NewContextForTests(dir, "") for _, file := range tt.existingFiles { f, err := os.Create(filepath.Join(dir, file)) @@ -207,7 +205,159 @@ func TestSelectSSHKeys(t *testing.T) { for _, tt := range tests { sshDir := t.TempDir() - sshContext := ssh.Context{ConfigDir: sshDir} + sshContext := ssh.NewContextForTests(sshDir, "") + + for _, file := range tt.sshDirFiles { + f, err := os.Create(filepath.Join(sshDir, file)) + if err != nil { + t.Errorf("Failed to create test ssh dir file %q: %v", file, err) + } + f.Close() + } + + configPath := filepath.Join(sshDir, "test-config") + + // Seed the config with a non-existent key so that the default config won't apply + configContent := "IdentityFile dummy\n" + + for _, key := range tt.sshConfigKeys { + configContent += fmt.Sprintf("IdentityFile %s\n", filepath.Join(sshDir, key)) + } + + err := os.WriteFile(configPath, []byte(configContent), 0666) + if err != nil { + t.Fatalf("could not write test config %v", err) + } + + var subbedSSHArgs []string + for _, arg := range tt.sshArgs { + subbedSSHArgs = append(subbedSSHArgs, strings.Replace(arg, substituteSSHDir, sshDir, -1)) + } + + tt.sshArgs = append([]string{"-F", configPath}, subbedSSHArgs...) + + gotKeyPair, gotShouldAddArg, err := selectSSHKeys(context.Background(), sshContext, tt.sshArgs, sshOptions{profile: tt.profileOpt}) + + if tt.wantKeyPair == nil { + if err == nil { + t.Errorf("Expected error from selectSSHKeys but got nil") + } + + continue + } + + if err != nil { + t.Errorf("Unexpected error from selectSSHKeys: %v", err) + continue + } + + if gotKeyPair == nil { + t.Errorf("Expected non-nil result from selectSSHKeys but got nil") + continue + } + + if gotShouldAddArg != tt.wantShouldAddArg { + t.Errorf("Got wrong shouldAddArg value from selectSSHKeys, wanted %v got %v", tt.wantShouldAddArg, gotShouldAddArg) + continue + } + + // Strip the dir (sshDir) from the gotKeyPair paths so that they match wantKeyPair (which doesn't know the directory) + gotKeyPairJustFileNames := &ssh.KeyPair{ + PrivateKeyPath: filepath.Base(gotKeyPair.PrivateKeyPath), + PublicKeyPath: filepath.Base(gotKeyPair.PublicKeyPath), + } + + if fmt.Sprintf("%v", gotKeyPairJustFileNames) != fmt.Sprintf("%v", tt.wantKeyPair) { + t.Errorf("Want selectSSHKeys result to be %v, got %v", tt.wantKeyPair, gotKeyPairJustFileNames) + } + + // If the automatic key pair is selected, it needs to exist no matter what + if strings.Contains(tt.wantKeyPair.PrivateKeyPath, automaticPrivateKeyName) { + if _, err := os.Stat(gotKeyPair.PrivateKeyPath); err != nil { + t.Errorf("Expected automatic key pair private key to exist, but it did not") + } + + if _, err := os.Stat(gotKeyPair.PublicKeyPath); err != nil { + t.Errorf("Expected automatic key pair public key to exist, but it did not") + } + } + } +} + +func TestConfigGeneration(t *testing.T) { + tests := []struct { + selector *CodespaceSelector + }{ + // -i tests + { + sshArgs: []string{"-i", "custom-private-key"}, + wantKeyPair: &ssh.KeyPair{PrivateKeyPath: "custom-private-key", PublicKeyPath: "custom-private-key.pub"}, + }, + { + sshArgs: []string{"-i", path.Join(substituteSSHDir, automaticPrivateKeyName)}, + wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, + }, + { + // Edge case check for missing arg value + sshArgs: []string{"-i"}, + }, + + // Auto key exists tests + { + sshDirFiles: []string{automaticPrivateKeyName, automaticPrivateKeyName + ".pub"}, + wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, + wantShouldAddArg: true, + }, + { + sshDirFiles: []string{automaticPrivateKeyName, automaticPrivateKeyName + ".pub", "custom-private-key", "custom-private-key.pub"}, + wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, + wantShouldAddArg: true, + }, + + // SSH config tests + { + sshDirFiles: []string{"custom-private-key", "custom-private-key.pub"}, + sshConfigKeys: []string{"custom-private-key"}, + wantKeyPair: &ssh.KeyPair{PrivateKeyPath: "custom-private-key", PublicKeyPath: "custom-private-key.pub"}, + wantShouldAddArg: true, + }, + { + // 2 pairs, but only 1 is configured + sshDirFiles: []string{"custom-private-key", "custom-private-key.pub", "custom-private-key-2", "custom-private-key-2.pub"}, + sshConfigKeys: []string{"custom-private-key-2"}, + wantKeyPair: &ssh.KeyPair{PrivateKeyPath: "custom-private-key-2", PublicKeyPath: "custom-private-key-2.pub"}, + wantShouldAddArg: true, + }, + { + // 2 pairs, but only 1 has both public and private + sshDirFiles: []string{"custom-private-key", "custom-private-key-2", "custom-private-key-2.pub"}, + sshConfigKeys: []string{"custom-private-key", "custom-private-key-2"}, + wantKeyPair: &ssh.KeyPair{PrivateKeyPath: "custom-private-key-2", PublicKeyPath: "custom-private-key-2.pub"}, + wantShouldAddArg: true, + }, + + // Automatic key tests + { + wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, + wantShouldAddArg: true, + }, + { + // Renames old key pair to new + sshDirFiles: []string{automaticPrivateKeyNameOld, automaticPrivateKeyNameOld + ".pub"}, + wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, + wantShouldAddArg: true, + }, + { + // Other key is configured, but doesn't exist + sshConfigKeys: []string{"custom-private-key"}, + wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, + wantShouldAddArg: true, + }, + } + + for _, tt := range tests { + sshDir := t.TempDir() + sshContext := ssh.NewContextForTests(sshDir, "") for _, file := range tt.sshDirFiles { f, err := os.Create(filepath.Join(sshDir, file)) diff --git a/pkg/ssh/ssh_keys.go b/pkg/ssh/ssh_keys.go index 0b9e12f1f..309b498c5 100644 --- a/pkg/ssh/ssh_keys.go +++ b/pkg/ssh/ssh_keys.go @@ -14,8 +14,17 @@ import ( ) type Context struct { - ConfigDir string - KeygenExe string + configDir string + keygenExe string +} + +// NewContextForTests creates a new `ssh.Context` with internal properties set to the +// specified values. It should only be used to testing to inject test-specific setup. +func NewContextForTests(configDir, keygenExe string) Context { + return Context{ + configDir, + keygenExe, + } } type KeyPair struct { @@ -80,19 +89,19 @@ func (c *Context) GenerateSSHKey(keyName string, passphrase string) (*KeyPair, e } func (c *Context) SshDir() (string, error) { - if c.ConfigDir != "" { - return c.ConfigDir, nil + if c.configDir != "" { + return c.configDir, nil } dir, err := config.HomeDirPath(".ssh") if err == nil { - c.ConfigDir = dir + c.configDir = dir } return dir, err } func (c *Context) findKeygen() (string, error) { - if c.KeygenExe != "" { - return c.KeygenExe, nil + if c.keygenExe != "" { + return c.keygenExe, nil } keygenExe, err := safeexec.LookPath("ssh-keygen") @@ -107,7 +116,7 @@ func (c *Context) findKeygen() (string, error) { } if err == nil { - c.KeygenExe = keygenExe + c.keygenExe = keygenExe } return keygenExe, err } From 509a181d798fece7768e503db32d9eb0fe050587 Mon Sep 17 00:00:00 2001 From: Caleb Brose <5447118+cmbrose@users.noreply.github.com> Date: Wed, 6 Nov 2024 19:10:28 +0000 Subject: [PATCH 3/4] Remove unimplemented tests --- pkg/cmd/codespace/ssh_test.go | 152 ---------------------------------- 1 file changed, 152 deletions(-) diff --git a/pkg/cmd/codespace/ssh_test.go b/pkg/cmd/codespace/ssh_test.go index e656e0ecd..cd04f05fa 100644 --- a/pkg/cmd/codespace/ssh_test.go +++ b/pkg/cmd/codespace/ssh_test.go @@ -284,158 +284,6 @@ func TestSelectSSHKeys(t *testing.T) { } } -func TestConfigGeneration(t *testing.T) { - tests := []struct { - selector *CodespaceSelector - }{ - // -i tests - { - sshArgs: []string{"-i", "custom-private-key"}, - wantKeyPair: &ssh.KeyPair{PrivateKeyPath: "custom-private-key", PublicKeyPath: "custom-private-key.pub"}, - }, - { - sshArgs: []string{"-i", path.Join(substituteSSHDir, automaticPrivateKeyName)}, - wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, - }, - { - // Edge case check for missing arg value - sshArgs: []string{"-i"}, - }, - - // Auto key exists tests - { - sshDirFiles: []string{automaticPrivateKeyName, automaticPrivateKeyName + ".pub"}, - wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, - wantShouldAddArg: true, - }, - { - sshDirFiles: []string{automaticPrivateKeyName, automaticPrivateKeyName + ".pub", "custom-private-key", "custom-private-key.pub"}, - wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, - wantShouldAddArg: true, - }, - - // SSH config tests - { - sshDirFiles: []string{"custom-private-key", "custom-private-key.pub"}, - sshConfigKeys: []string{"custom-private-key"}, - wantKeyPair: &ssh.KeyPair{PrivateKeyPath: "custom-private-key", PublicKeyPath: "custom-private-key.pub"}, - wantShouldAddArg: true, - }, - { - // 2 pairs, but only 1 is configured - sshDirFiles: []string{"custom-private-key", "custom-private-key.pub", "custom-private-key-2", "custom-private-key-2.pub"}, - sshConfigKeys: []string{"custom-private-key-2"}, - wantKeyPair: &ssh.KeyPair{PrivateKeyPath: "custom-private-key-2", PublicKeyPath: "custom-private-key-2.pub"}, - wantShouldAddArg: true, - }, - { - // 2 pairs, but only 1 has both public and private - sshDirFiles: []string{"custom-private-key", "custom-private-key-2", "custom-private-key-2.pub"}, - sshConfigKeys: []string{"custom-private-key", "custom-private-key-2"}, - wantKeyPair: &ssh.KeyPair{PrivateKeyPath: "custom-private-key-2", PublicKeyPath: "custom-private-key-2.pub"}, - wantShouldAddArg: true, - }, - - // Automatic key tests - { - wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, - wantShouldAddArg: true, - }, - { - // Renames old key pair to new - sshDirFiles: []string{automaticPrivateKeyNameOld, automaticPrivateKeyNameOld + ".pub"}, - wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, - wantShouldAddArg: true, - }, - { - // Other key is configured, but doesn't exist - sshConfigKeys: []string{"custom-private-key"}, - wantKeyPair: &ssh.KeyPair{PrivateKeyPath: automaticPrivateKeyName, PublicKeyPath: automaticPrivateKeyName + ".pub"}, - wantShouldAddArg: true, - }, - } - - for _, tt := range tests { - sshDir := t.TempDir() - sshContext := ssh.NewContextForTests(sshDir, "") - - for _, file := range tt.sshDirFiles { - f, err := os.Create(filepath.Join(sshDir, file)) - if err != nil { - t.Errorf("Failed to create test ssh dir file %q: %v", file, err) - } - f.Close() - } - - configPath := filepath.Join(sshDir, "test-config") - - // Seed the config with a non-existent key so that the default config won't apply - configContent := "IdentityFile dummy\n" - - for _, key := range tt.sshConfigKeys { - configContent += fmt.Sprintf("IdentityFile %s\n", filepath.Join(sshDir, key)) - } - - err := os.WriteFile(configPath, []byte(configContent), 0666) - if err != nil { - t.Fatalf("could not write test config %v", err) - } - - var subbedSSHArgs []string - for _, arg := range tt.sshArgs { - subbedSSHArgs = append(subbedSSHArgs, strings.Replace(arg, substituteSSHDir, sshDir, -1)) - } - - tt.sshArgs = append([]string{"-F", configPath}, subbedSSHArgs...) - - gotKeyPair, gotShouldAddArg, err := selectSSHKeys(context.Background(), sshContext, tt.sshArgs, sshOptions{profile: tt.profileOpt}) - - if tt.wantKeyPair == nil { - if err == nil { - t.Errorf("Expected error from selectSSHKeys but got nil") - } - - continue - } - - if err != nil { - t.Errorf("Unexpected error from selectSSHKeys: %v", err) - continue - } - - if gotKeyPair == nil { - t.Errorf("Expected non-nil result from selectSSHKeys but got nil") - continue - } - - if gotShouldAddArg != tt.wantShouldAddArg { - t.Errorf("Got wrong shouldAddArg value from selectSSHKeys, wanted %v got %v", tt.wantShouldAddArg, gotShouldAddArg) - continue - } - - // Strip the dir (sshDir) from the gotKeyPair paths so that they match wantKeyPair (which doesn't know the directory) - gotKeyPairJustFileNames := &ssh.KeyPair{ - PrivateKeyPath: filepath.Base(gotKeyPair.PrivateKeyPath), - PublicKeyPath: filepath.Base(gotKeyPair.PublicKeyPath), - } - - if fmt.Sprintf("%v", gotKeyPairJustFileNames) != fmt.Sprintf("%v", tt.wantKeyPair) { - t.Errorf("Want selectSSHKeys result to be %v, got %v", tt.wantKeyPair, gotKeyPairJustFileNames) - } - - // If the automatic key pair is selected, it needs to exist no matter what - if strings.Contains(tt.wantKeyPair.PrivateKeyPath, automaticPrivateKeyName) { - if _, err := os.Stat(gotKeyPair.PrivateKeyPath); err != nil { - t.Errorf("Expected automatic key pair private key to exist, but it did not") - } - - if _, err := os.Stat(gotKeyPair.PublicKeyPath); err != nil { - t.Errorf("Expected automatic key pair public key to exist, but it did not") - } - } - } -} - func testingSSHApp() *App { disabledCodespace := &api.Codespace{ Name: "disabledCodespace", From 2435f6915b3abf4134aae74d00ffa923bebc443e Mon Sep 17 00:00:00 2001 From: Andy Feller Date: Wed, 6 Nov 2024 16:12:55 -0500 Subject: [PATCH 4/4] Minor nit suggestion --- pkg/ssh/ssh_keys.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/ssh/ssh_keys.go b/pkg/ssh/ssh_keys.go index 309b498c5..83d4f1b34 100644 --- a/pkg/ssh/ssh_keys.go +++ b/pkg/ssh/ssh_keys.go @@ -19,7 +19,7 @@ type Context struct { } // NewContextForTests creates a new `ssh.Context` with internal properties set to the -// specified values. It should only be used to testing to inject test-specific setup. +// specified values. It should only be used to inject test-specific setup. func NewContextForTests(configDir, keygenExe string) Context { return Context{ configDir,