fix: Phase 3 — Long-term improvements (SSH TOFU, EC2Client split, context propagation) + E2E tests

pkg/provisioner/provisioner.go

@@ -20,6 +20,7 @@ import (
 	"bytes"
 	"fmt"
 	"io"
+	"net"
 	"os"
 	"path/filepath"
 	"strings"
@@ -441,7 +442,7 @@ func connectOrDie(keyPath, userName, hostUrl string) (*ssh.Client, error) {
 		Auth: []ssh.AuthMethod{
 			ssh.PublicKeys(signer),
 		},
-		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // nolint:gosec
+		HostKeyCallback: tofuHostKeyCallback(),
 	}
 
 	connectionFailed := false
@@ -462,3 +463,54 @@ func connectOrDie(keyPath, userName, hostUrl string) (*ssh.Client, error) {
 
 	return client, nil
 }
+
+// tofuHostKeyCallback implements a Trust-On-First-Use (TOFU) pattern for SSH
+// host key verification. On first connection to a host, the key is recorded in
+// $HOME/.cache/holodeck/known_hosts (or os.UserCacheDir fallback). On subsequent
+// connections the stored key is compared and a mismatch (potential MITM) is rejected.
+func tofuHostKeyCallback() ssh.HostKeyCallback {
+	return func(hostname string, remote net.Addr, key ssh.PublicKey) error {
+		cacheBase, err := os.UserCacheDir()
+		if err != nil {
+			return fmt.Errorf("cannot determine cache directory for TOFU host keys: %w", err)
+		}
+		knownHostsPath := filepath.Join(cacheBase, "holodeck", "known_hosts")
+
+		if err := os.MkdirAll(filepath.Dir(knownHostsPath), 0700); err != nil {
+			return fmt.Errorf("failed to create known_hosts directory: %w", err)
+		}
+
+		keyStr := strings.TrimSpace(string(ssh.MarshalAuthorizedKey(key)))
+		host := hostname
+
+		// Try to read existing known hosts file
+		data, err := os.ReadFile(knownHostsPath) // nolint:gosec // path from UserCacheDir + static components
+		if err != nil && !os.IsNotExist(err) {
+			return fmt.Errorf("failed to read known_hosts: %w", err)
+		}
+		if err == nil {
+			for _, line := range strings.Split(string(data), "\n") {
+				parts := strings.SplitN(line, " ", 2)
+				if len(parts) == 2 && parts[0] == host {
+					if strings.TrimSpace(parts[1]) == keyStr {
+						return nil // Key matches
+					}
+					return fmt.Errorf("host key mismatch for %s: stored key differs from presented key (possible MITM)", host)
+				}
+			}
+		}
+
+		// First connection to this host: record the key (TOFU)
+		f, err := os.OpenFile(knownHostsPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600) // nolint:gosec
+		if err != nil {
+			return fmt.Errorf("failed to open known_hosts for writing: %w", err)
+		}
+		defer func() { _ = f.Close() }()
+
+		if _, err := fmt.Fprintf(f, "%s %s\n", host, keyStr); err != nil {
+			return fmt.Errorf("failed to write known host: %w", err)
+		}
+
+		return nil
+	}
+}

pkg/provisioner/tofu_test.go

@@ -0,0 +1,132 @@
+/*
+ * Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package provisioner
+
+import (
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"net"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"golang.org/x/crypto/ssh"
+)
+
+func generateTestKey(t *testing.T) ssh.PublicKey {
+	t.Helper()
+	privKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	if err != nil {
+		t.Fatalf("failed to generate test key: %v", err)
+	}
+	pubKey, err := ssh.NewPublicKey(&privKey.PublicKey)
+	if err != nil {
+		t.Fatalf("failed to create SSH public key: %v", err)
+	}
+	return pubKey
+}
+
+// setupTOFUTest isolates the TOFU cache in a temp directory by overriding HOME.
+// Returns the expected known_hosts path.
+func setupTOFUTest(t *testing.T) string {
+	t.Helper()
+	tmpDir := t.TempDir()
+	t.Setenv("HOME", tmpDir)
+
+	cacheDir, err := os.UserCacheDir()
+	if err != nil {
+		t.Fatalf("os.UserCacheDir failed: %v", err)
+	}
+	return filepath.Join(cacheDir, "holodeck", "known_hosts")
+}
+
+func TestTOFU_FirstConnection_RecordsKey(t *testing.T) {
+	knownHostsPath := setupTOFUTest(t)
+
+	key := generateTestKey(t)
+	addr := &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 22}
+
+	cb := tofuHostKeyCallback()
+	if err := cb("testhost:22", addr, key); err != nil {
+		t.Fatalf("first connection should succeed: %v", err)
+	}
+
+	data, err := os.ReadFile(knownHostsPath) // nolint:gosec // test helper with controlled tmpdir path
+	if err != nil {
+		t.Fatalf("known_hosts should exist at %s: %v", knownHostsPath, err)
+	}
+	if len(data) == 0 {
+		t.Fatal("known_hosts should not be empty")
+	}
+}
+
+func TestTOFU_SubsequentConnection_SameKey_Accepted(t *testing.T) {
+	_ = setupTOFUTest(t)
+
+	key := generateTestKey(t)
+	addr := &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 22}
+
+	cb := tofuHostKeyCallback()
+
+	if err := cb("testhost:22", addr, key); err != nil {
+		t.Fatalf("first connection should succeed: %v", err)
+	}
+	if err := cb("testhost:22", addr, key); err != nil {
+		t.Fatalf("second connection with same key should succeed: %v", err)
+	}
+}
+
+func TestTOFU_SubsequentConnection_DifferentKey_Rejected(t *testing.T) {
+	_ = setupTOFUTest(t)
+
+	key1 := generateTestKey(t)
+	key2 := generateTestKey(t)
+	addr := &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 22}
+
+	cb := tofuHostKeyCallback()
+
+	if err := cb("testhost:22", addr, key1); err != nil {
+		t.Fatalf("first connection should succeed: %v", err)
+	}
+	if err := cb("testhost:22", addr, key2); err == nil {
+		t.Fatal("connection with different key should be rejected")
+	}
+}
+
+func TestTOFU_UnreadableFile_ReturnsError(t *testing.T) {
+	knownHostsPath := setupTOFUTest(t)
+
+	if err := os.MkdirAll(filepath.Dir(knownHostsPath), 0700); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(knownHostsPath, []byte("some data"), 0600); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.Chmod(knownHostsPath, 0200); err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = os.Chmod(knownHostsPath, 0600) })
+
+	key := generateTestKey(t)
+	addr := &net.TCPAddr{IP: net.ParseIP("10.0.0.1"), Port: 22}
+
+	cb := tofuHostKeyCallback()
+	if err := cb("testhost:22", addr, key); err == nil {
+		t.Fatal("should return error when known_hosts is not readable")
+	}
+}