From 44bd0ff0893cf6f38ce5c6d1064da5ed540499c5 Mon Sep 17 00:00:00 2001
From: Harivansh Rathi <rathiharivansh@gmail.com>
Date: Sat, 11 Apr 2026 15:38:05 +0000
Subject: [PATCH] fix: seed mmds data via --metadata flag at boot, eliminates
 race with guest init

---
 internal/firecracker/configfile.go  | 18 ++++++++++++++++++
 internal/firecracker/launch.go      |  5 ++++-
 internal/firecracker/launch_test.go |  4 ++--
 internal/firecracker/runtime.go     | 22 ++++++++--------------
 4 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/internal/firecracker/configfile.go b/internal/firecracker/configfile.go
index f6a7805..9904889 100644
--- a/internal/firecracker/configfile.go
+++ b/internal/firecracker/configfile.go
@@ -141,3 +141,21 @@ func writeConfigFile(chrootRootDir string, spec MachineSpec, paths machinePaths,
 
 	return "/vm_config.json", nil
 }
+
+func writeMetadataFile(chrootRootDir string, spec MachineSpec) (string, error) {
+	if spec.MMDS == nil || spec.MMDS.Data == nil {
+		return "", nil
+	}
+
+	data, err := json.Marshal(spec.MMDS.Data)
+	if err != nil {
+		return "", fmt.Errorf("marshal mmds data: %w", err)
+	}
+
+	metadataPath := filepath.Join(chrootRootDir, "mmds.json")
+	if err := os.WriteFile(metadataPath, data, 0o644); err != nil {
+		return "", fmt.Errorf("write mmds data: %w", err)
+	}
+
+	return "/mmds.json", nil
+}
diff --git a/internal/firecracker/launch.go b/internal/firecracker/launch.go
index a011f21..2fd0bbe 100644
--- a/internal/firecracker/launch.go
+++ b/internal/firecracker/launch.go
@@ -66,7 +66,7 @@ func configureMachine(ctx context.Context, client *apiClient, paths machinePaths
 	return nil
 }
 
-func launchJailedFirecracker(paths machinePaths, machineID MachineID, firecrackerBinaryPath string, jailerBinaryPath string, enablePCI bool, configFilePath string) (*exec.Cmd, error) {
+func launchJailedFirecracker(paths machinePaths, machineID MachineID, firecrackerBinaryPath string, jailerBinaryPath string, enablePCI bool, configFilePath string, metadataFilePath string) (*exec.Cmd, error) {
 	args := []string{
 		"--id", string(machineID),
 		"--uid", strconv.Itoa(os.Getuid()),
@@ -81,6 +81,9 @@ func launchJailedFirecracker(paths machinePaths, machineID MachineID, firecracke
 	}
 	if configFilePath != "" {
 		args = append(args, "--config-file", configFilePath)
+		if metadataFilePath != "" {
+			args = append(args, "--metadata", metadataFilePath)
+		}
 	} else {
 		args = append(args,
 			"--log-path", paths.JailedFirecrackerLogPath,
diff --git a/internal/firecracker/launch_test.go b/internal/firecracker/launch_test.go
index 9501fe2..01aa86f 100644
--- a/internal/firecracker/launch_test.go
+++ b/internal/firecracker/launch_test.go
@@ -26,7 +26,7 @@ func TestLaunchJailedFirecrackerPassesDaemonAndLoggingFlags(t *testing.T) {
 		t.Fatalf("create log dir: %v", err)
 	}
 
-	if _, err := launchJailedFirecracker(paths, "vm-1", "/usr/bin/firecracker", jailerPath, false, ""); err != nil {
+	if _, err := launchJailedFirecracker(paths, "vm-1", "/usr/bin/firecracker", jailerPath, false, "", ""); err != nil {
 		t.Fatalf("launch jailed firecracker: %v", err)
 	}
 
@@ -62,7 +62,7 @@ func TestLaunchJailedFirecrackerPassesEnablePCIWhenConfigured(t *testing.T) {
 		t.Fatalf("create log dir: %v", err)
 	}
 
-	if _, err := launchJailedFirecracker(paths, "vm-1", "/usr/bin/firecracker", jailerPath, true, ""); err != nil {
+	if _, err := launchJailedFirecracker(paths, "vm-1", "/usr/bin/firecracker", jailerPath, true, "", ""); err != nil {
 		t.Fatalf("launch jailed firecracker: %v", err)
 	}
 
diff --git a/internal/firecracker/runtime.go b/internal/firecracker/runtime.go
index d9698c9..af1598d 100644
--- a/internal/firecracker/runtime.go
+++ b/internal/firecracker/runtime.go
@@ -125,7 +125,13 @@ func (r *Runtime) Boot(ctx context.Context, spec MachineSpec, usedNetworks []Net
 		return nil, fmt.Errorf("write config file: %w", err)
 	}
 
-	command, err := launchJailedFirecracker(paths, spec.ID, r.firecrackerBinaryPath, r.jailerBinaryPath, r.enablePCI, configFilePath)
+	metadataFilePath, err := writeMetadataFile(paths.ChrootRootDir, spec)
+	if err != nil {
+		cleanup(network, paths, nil, 0)
+		return nil, fmt.Errorf("write metadata file: %w", err)
+	}
+
+	command, err := launchJailedFirecracker(paths, spec.ID, r.firecrackerBinaryPath, r.jailerBinaryPath, r.enablePCI, configFilePath, metadataFilePath)
 	if err != nil {
 		cleanup(network, paths, nil, 0)
 		return nil, err
@@ -137,18 +143,6 @@ func (r *Runtime) Boot(ctx context.Context, spec MachineSpec, usedNetworks []Net
 	}
 	socketPath := procSocketPath(firecrackerPID)
 
-	if spec.MMDS != nil && spec.MMDS.Data != nil {
-		client := newAPIClient(socketPath)
-		if err := waitForSocket(ctx, client, socketPath); err != nil {
-			cleanup(network, paths, command, firecrackerPID)
-			return nil, fmt.Errorf("wait for firecracker socket: %w", err)
-		}
-		if err := client.PutMMDS(ctx, spec.MMDS.Data); err != nil {
-			cleanup(network, paths, command, firecrackerPID)
-			return nil, fmt.Errorf("put mmds data: %w", err)
-		}
-	}
-
 	now := time.Now().UTC()
 	state := MachineState{
 		ID:          spec.ID,
@@ -288,7 +282,7 @@ func (r *Runtime) RestoreBoot(ctx context.Context, loadSpec SnapshotLoadSpec, us
 		return nil, err
 	}
 
-	command, err := launchJailedFirecracker(paths, loadSpec.ID, r.firecrackerBinaryPath, r.jailerBinaryPath, r.enablePCI, "")
+	command, err := launchJailedFirecracker(paths, loadSpec.ID, r.firecrackerBinaryPath, r.jailerBinaryPath, r.enablePCI, "", "")
 	if err != nil {
 		cleanup(network, paths, nil, 0)
 		return nil, err