refactor(ssh): construct host key callback only when dialing

knownhosts.New() caches the known hosts database, which means that it
will not be picked up across reconnects.

As such, this commit delays the instantiation of the host key callback
so that this caching does not happen and the user isn't reprompted to
add a key they've already added during the ACK process before magic
rollback.

Author: water-sucks

cmd/apply/apply.go

@@ -646,7 +646,7 @@ func applyMain(cmd *cobra.Command, opts *cmdOpts.ApplyOpts) error {
 		log.Printf("\n")
 
 		var confirm bool
-		confirm, err = cmdUtils.ConfirmationInput("Activate this configuration?", cmdUtils.ConfirmationPromptOptions{
+		confirm, err = cmdUtils.ConfirmationInput(stopCtx, "Activate this configuration?", cmdUtils.ConfirmationPromptOptions{
 			InvalidBehavior: cfg.Confirmation.Invalid,
 			EmptyBehavior:   cfg.Confirmation.Empty,
 		})

cmd/generation/delete/delete.go

@@ -1,6 +1,7 @@
 package delete
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -160,7 +161,7 @@ func generationDeleteMain(cmd *cobra.Command, genOpts *cmdOpts.GenerationOpts, o
 
 	if !opts.AlwaysConfirm && !cfg.Confirmation.Always {
 		var confirm bool
-		confirm, err = cmdUtils.ConfirmationInput("Proceed?", cmdUtils.ConfirmationPromptOptions{
+		confirm, err = cmdUtils.ConfirmationInput(context.Background(), "Proceed?", cmdUtils.ConfirmationPromptOptions{
 			InvalidBehavior: cfg.Confirmation.Invalid,
 			EmptyBehavior:   cfg.Confirmation.Empty,
 		})

cmd/generation/rollback/rollback.go

@@ -1,6 +1,7 @@
 package rollback
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -98,7 +99,7 @@ func generationRollbackMain(cmd *cobra.Command, genOpts *cmdOpts.GenerationOpts,
 		log.Printf("\n")
 
 		var confirm bool
-		confirm, err = cmdUtils.ConfirmationInput("Activate the previous generation?", cmdUtils.ConfirmationPromptOptions{
+		confirm, err = cmdUtils.ConfirmationInput(context.Background(), "Activate the previous generation?", cmdUtils.ConfirmationPromptOptions{
 			InvalidBehavior: cfg.Confirmation.Invalid,
 			EmptyBehavior:   cfg.Confirmation.Empty,
 		})

cmd/generation/switch/switch.go

@@ -1,6 +1,7 @@
 package switch_cmd
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -150,7 +151,7 @@ func generationSwitchMain(cmd *cobra.Command, genOpts *cmdOpts.GenerationOpts, o
 		log.Printf("\n")
 
 		var confirm bool
-		confirm, err = cmdUtils.ConfirmationInput("Activate this generation?", cmdUtils.ConfirmationPromptOptions{
+		confirm, err = cmdUtils.ConfirmationInput(context.Background(), "Activate this generation?", cmdUtils.ConfirmationPromptOptions{
 			InvalidBehavior: cfg.Confirmation.Invalid,
 			EmptyBehavior:   cfg.Confirmation.Empty,
 		})

internal/cmd/utils/confirmation.go

@@ -2,6 +2,7 @@ package cmdUtils
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"os"
 	"strings"
@@ -15,9 +16,26 @@ type ConfirmationPromptOptions struct {
 	EmptyBehavior   settings.ConfirmationPromptBehavior
 }
 
-func ConfirmationInput(msg string, opts ConfirmationPromptOptions) (bool, error) {
+// This operation can be cancelled using the provided context,
+// but any errors returned from here MAY have the potential
+// to keep consuming stdin until another character is typed
+// if a duplicate instance of stdin cannot be opened, so
+// any errors here will result in potentiall undefined behavior
+// for stdin input.
+func ConfirmationInput(ctx context.Context, msg string, opts ConfirmationPromptOptions) (bool, error) {
+	var stdin *os.File
+	if dupStdin, err := os.OpenFile("/dev/stdin", os.O_RDONLY, 0); err == nil {
+		defer dupStdin.Close()
+		stdin = dupStdin
+	} else {
+		// NOTE: falling back to stdin will make context
+		// cancellation behavior a bit unclear, as mentioned
+		// in the doc comment.
+		stdin = os.Stdin
+	}
+
 	var input string
-	scanner := bufio.NewScanner(os.Stdin)
+	scanner := bufio.NewScanner(stdin)
 
 	for {
 		fmt.Fprintf(os.Stderr, "%s\n[y/n]: ", color.GreenString("|> %s", msg))

internal/system/ssh.go

@@ -41,8 +41,9 @@ type SSHConfig struct {
 	Address string
 	Port    int
 
-	AuthMethods     []ssh.AuthMethod
-	HostKeyCallback ssh.HostKeyCallback
+	AuthMethods         []ssh.AuthMethod
+	HostKeyVerification settings.HostKeyVerificationType
+	KnownHostsFiles     []string
 
 	password []byte
 
@@ -167,13 +168,9 @@ func NewSSHConfig(ctx context.Context, host string, log logger.Logger, options S
 	})
 	auth = append(auth, passwordCallback)
 
-	hostKeyCallback, err := knownHostsCallback(log, options.HostKeyVerification, options.KnownHostsFiles)
-	if err != nil {
-		return nil, err
-	}
-
 	cfg.AuthMethods = auth
-	cfg.HostKeyCallback = hostKeyCallback
+	cfg.HostKeyVerification = options.HostKeyVerification
+	cfg.KnownHostsFiles = options.KnownHostsFiles
 
 	return cfg, nil
 }
@@ -229,7 +226,7 @@ func knownHostsCallback(
 }
 
 func NewSSHSystem(cfg *SSHConfig, log logger.Logger) (*SSHSystem, error) {
-	client, sftpClient, err := dialClient(cfg)
+	client, sftpClient, err := dialClient(log, cfg)
 	if err != nil {
 		return nil, err
 	}
@@ -244,11 +241,16 @@ func NewSSHSystem(cfg *SSHConfig, log logger.Logger) (*SSHSystem, error) {
 	}, nil
 }
 
-func dialClient(cfg *SSHConfig) (*ssh.Client, *sftp.Client, error) {
+func dialClient(log logger.Logger, cfg *SSHConfig) (*ssh.Client, *sftp.Client, error) {
+	hostKeyCallback, err := knownHostsCallback(log, cfg.HostKeyVerification, cfg.KnownHostsFiles)
+	if err != nil {
+		return nil, nil, err
+	}
+
 	client, err := ssh.Dial("tcp", net.JoinHostPort(cfg.Address, strconv.Itoa(cfg.Port)), &ssh.ClientConfig{
 		User:            cfg.User,
 		Auth:            cfg.AuthMethods,
-		HostKeyCallback: cfg.HostKeyCallback,
+		HostKeyCallback: hostKeyCallback,
 		Timeout:         30 * time.Second,
 	})
 	if err != nil {
@@ -268,7 +270,7 @@ func (s *SSHSystem) Reconnect() error {
 	_ = s.sftp.Close()
 	_ = s.client.Close()
 
-	client, sftpClient, err := dialClient(s.cfg)
+	client, sftpClient, err := dialClient(s.logger, s.cfg)
 	if err != nil {
 		return fmt.Errorf("failed to reconnect to %s: %w", s.Address(), err)
 	}
@@ -284,7 +286,7 @@ func (s *SSHSystem) Reconnect() error {
 //
 // Caller must keep the `cfg` field alive until ALL clones are closed.
 func (s *SSHSystem) Clone() (*SSHSystem, error) {
-	client, sftpClient, err := dialClient(s.cfg)
+	client, sftpClient, err := dialClient(s.logger, s.cfg)
 	if err != nil {
 		return nil, fmt.Errorf("failed to clone connection to %s: %w", s.Address(), err)
 	}
@@ -379,7 +381,7 @@ func addKeyToKnownHostsCallback(
 			log.Infof("SHA256 fingerprint: %s", fingerprint)
 
 			var confirm bool
-			confirm, err = cmdUtils.ConfirmationInput("Are you sure you want to continue connecting?", cmdUtils.ConfirmationPromptOptions{
+			confirm, err = cmdUtils.ConfirmationInput(context.Background(), "Are you sure you want to continue connecting?", cmdUtils.ConfirmationPromptOptions{
 				// Copy the default SSH behavior of retrying for invalid input.
 				// Disregard user configuration in this case, since this is mimicking
 				// OpenSSH's behavior.