feat: prepare runtime control sockets (#31)

InderdeepBajwa · web-flow · commit ad4ab0ad97bd · 2026-04-26T00:48:15.000-04:00
diff --git a/schedune-control-plane/internal/runtime/cloudhypervisor.go b/schedune-control-plane/internal/runtime/cloudhypervisor.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"path/filepath"
 
 	"github.com/TechnologyTailors/Schedune/schedune-control-plane/pkg/schema/launch"
 )
@@ -24,20 +25,28 @@ func (k *CloudHypervisorExecutor) Prepare(spec launch.LaunchSpec) (launch.Prepar
 		return launch.PreparedLaunch{}, fmt.Errorf("artifact missing at host path: %s", artifactPath)
 	}
 
+	controlSocket, err := GetControlSocketPath(spec.WorkloadID, "cloudhypervisor")
+	if err != nil {
+		return launch.PreparedLaunch{}, fmt.Errorf("failed to resolve control socket: %w", err)
+	}
+
 	args := []string{
 		"--memory", fmt.Sprintf("size=%dM", spec.MemoryMB),
 		"--cpus", fmt.Sprintf("boot=%d", spec.Vcpu),
 		"--disk", fmt.Sprintf("path=%s", artifactPath),
+		"--api-socket", controlSocket,
 	}
 
 	return launch.PreparedLaunch{
-		RuntimeBackend: "cloud_hypervisor",
-		MemoryMB:       spec.MemoryMB,
-		Vcpu:           spec.Vcpu,
+		RuntimeBackend:  "cloud_hypervisor",
+		MemoryMB:        spec.MemoryMB,
+		Vcpu:            spec.Vcpu,
+		StartupGraceSec: 3,
 		CloudHypervisor: &launch.PreparedCloudHypervisorLaunch{
-			BinaryPath:   binPath,
-			ArtifactPath: artifactPath,
-			CommandArgs:  args,
+			BinaryPath:        binPath,
+			ArtifactPath:      artifactPath,
+			CommandArgs:       args,
+			ControlSocketPath: controlSocket,
 		},
 	}, nil
 }
@@ -46,6 +55,13 @@ func (k *CloudHypervisorExecutor) Execute(prepared launch.PreparedLaunch) (int,
 	if prepared.CloudHypervisor == nil {
 		return 0, fmt.Errorf("missing cloud_hypervisor prepared state")
 	}
+
+	if prepared.CloudHypervisor.ControlSocketPath != "" {
+		if err := os.MkdirAll(filepath.Dir(prepared.CloudHypervisor.ControlSocketPath), 0755); err != nil {
+			return 0, fmt.Errorf("failed to create runtime directory: %w", err)
+		}
+	}
+
 	cmd := exec.Command(prepared.CloudHypervisor.BinaryPath, prepared.CloudHypervisor.CommandArgs...)
 	if err := cmd.Start(); err != nil {
 		return 0, fmt.Errorf("executable failed to start: %w", err)
diff --git a/schedune-control-plane/internal/runtime/firecracker_validate.go b/schedune-control-plane/internal/runtime/firecracker_validate.go
@@ -2,6 +2,7 @@ package runtime
 
 import (
 	"fmt"
+	"path/filepath"
 
 	"github.com/TechnologyTailors/Schedune/schedune-control-plane/pkg/schema/launch"
 )
@@ -18,22 +19,31 @@ func (k *FirecrackerExecutor) Prepare(spec launch.LaunchSpec) (launch.PreparedLa
 		return launch.PreparedLaunch{}, fmt.Errorf("missing kernel or rootfs path for firecracker artifact model")
 	}
 
+	controlSocket, err := GetControlSocketPath(spec.WorkloadID, "firecracker")
+	if err != nil {
+		return launch.PreparedLaunch{}, fmt.Errorf("failed to resolve control socket: %w", err)
+	}
+
+	runtimeDir := filepath.Dir(controlSocket)
+
 	args := []string{
-		"--api-sock", "/tmp/firecracker.socket",
+		"--api-sock", controlSocket,
 		// In a real execution, we'd write a config JSON and pass it, or call the API.
 		// For V0 dry-run, we just mock the arguments.
-		"--config-file", "/tmp/fc-config.json",
+		"--config-file", filepath.Join(runtimeDir, "fc-config.json"),
 	}
 
 	return launch.PreparedLaunch{
-		RuntimeBackend: "firecracker",
-		MemoryMB:       spec.MemoryMB,
-		Vcpu:           spec.Vcpu,
+		RuntimeBackend:  "firecracker",
+		MemoryMB:        spec.MemoryMB,
+		Vcpu:            spec.Vcpu,
+		StartupGraceSec: 2,
 		Firecracker: &launch.PreparedFirecrackerLaunch{
-			BinaryPath:      binPath,
-			KernelImagePath: kernel,
-			RootfsPath:      rootfs,
-			CommandArgs:     args,
+			BinaryPath:        binPath,
+			KernelImagePath:   kernel,
+			RootfsPath:        rootfs,
+			CommandArgs:       args,
+			ControlSocketPath: controlSocket,
 		},
 	}, nil
 }
diff --git a/schedune-control-plane/internal/runtime/kvm.go b/schedune-control-plane/internal/runtime/kvm.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"path/filepath"
 
 	"github.com/TechnologyTailors/Schedune/schedune-control-plane/pkg/schema/launch"
 )
@@ -32,21 +33,29 @@ func (k *KvmExecutor) Prepare(spec launch.LaunchSpec) (launch.PreparedLaunch, er
 		return launch.PreparedLaunch{}, fmt.Errorf("artifact missing at host path: %s", artifactPath)
 	}
 
+	controlSocket, err := GetControlSocketPath(spec.WorkloadID, "qemu")
+	if err != nil {
+		return launch.PreparedLaunch{}, fmt.Errorf("failed to resolve control socket: %w", err)
+	}
+
 	args := []string{
 		"-m", fmt.Sprintf("%d", spec.MemoryMB),
 		"-smp", fmt.Sprintf("%d", spec.Vcpu),
 		"-drive", fmt.Sprintf("file=%s,format=%s", artifactPath, format),
 		"-nographic",
+		"-qmp", fmt.Sprintf("unix:%s,server,nowait", controlSocket),
 	}
 
 	return launch.PreparedLaunch{
-		RuntimeBackend: "kvm_qemu",
-		MemoryMB:       spec.MemoryMB,
-		Vcpu:           spec.Vcpu,
+		RuntimeBackend:  "kvm_qemu",
+		MemoryMB:        spec.MemoryMB,
+		Vcpu:            spec.Vcpu,
+		StartupGraceSec: 5,
 		KvmQemu: &launch.PreparedQemuLaunch{
-			BinaryPath:   binPath,
-			ArtifactPath: artifactPath,
-			CommandArgs:  args,
+			BinaryPath:        binPath,
+			ArtifactPath:      artifactPath,
+			CommandArgs:       args,
+			ControlSocketPath: controlSocket,
 		},
 	}, nil
 }
@@ -55,6 +64,14 @@ func (k *KvmExecutor) Execute(prepared launch.PreparedLaunch) (int, error) {
 	if prepared.KvmQemu == nil {
 		return 0, fmt.Errorf("missing kvm_qemu prepared state")
 	}
+
+	// Ensure the runtime directory exists before starting the binary
+	if prepared.KvmQemu.ControlSocketPath != "" {
+		if err := os.MkdirAll(filepath.Dir(prepared.KvmQemu.ControlSocketPath), 0755); err != nil {
+			return 0, fmt.Errorf("failed to create runtime directory: %w", err)
+		}
+	}
+
 	cmd := exec.Command(prepared.KvmQemu.BinaryPath, prepared.KvmQemu.CommandArgs...)
 	if err := cmd.Start(); err != nil {
 		return 0, fmt.Errorf("executable failed to start: %w", err)
diff --git a/schedune-control-plane/internal/runtime/kvm_test.go b/schedune-control-plane/internal/runtime/kvm_test.go
@@ -1,6 +1,7 @@
 package runtime
 
 import (
+	"fmt"
 	"github.com/TechnologyTailors/Schedune/schedune-control-plane/pkg/schema/launch"
 	"os"
 	"strings"
@@ -11,6 +12,7 @@ func TestKvmExecutor_PrepareMissingImage(t *testing.T) {
 	exec := &KvmExecutor{}
 
 	spec := launch.LaunchSpec{
+		WorkloadID:   "test-missing-id",
 		Architecture: "aarch64",
 		Storage: []launch.StorageAttachmentSpec{
 			{HostPath: "/tmp/non_existent_image_12345.qcow2", Format: "qcow2"},
@@ -35,6 +37,7 @@ func TestKvmExecutor_PrepareValidImageLegacy(t *testing.T) {
 	defer os.Remove(f.Name())
 
 	spec := launch.LaunchSpec{
+		WorkloadID:     "test-legacy-id",
 		Architecture:   "aarch64",
 		ImageReference: f.Name(),
 		Vcpu:           2,
@@ -43,7 +46,7 @@ func TestKvmExecutor_PrepareValidImageLegacy(t *testing.T) {
 
 	prep, err := exec.Prepare(spec)
 	if err != nil {
-		t.Errorf("expected Prepare to succeed, got %v", err)
+		t.Fatalf("expected Prepare to succeed, got %v", err)
 	}
 
 	if prep.KvmQemu == nil {
@@ -54,15 +57,26 @@ func TestKvmExecutor_PrepareValidImageLegacy(t *testing.T) {
 		t.Errorf("expected qemu-system-aarch64, got %s", prep.KvmQemu.BinaryPath)
 	}
 
+	if prep.KvmQemu.ControlSocketPath == "" || !strings.Contains(prep.KvmQemu.ControlSocketPath, "test-legacy-id/qemu.sock") {
+		t.Errorf("expected qemu control socket path, got %v", prep.KvmQemu.ControlSocketPath)
+	}
+
 	foundDrive := false
-	for _, arg := range prep.KvmQemu.CommandArgs {
+	foundQmp := false
+	for i, arg := range prep.KvmQemu.CommandArgs {
 		if strings.Contains(arg, "format=qcow2") {
 			foundDrive = true
 		}
+		if arg == "-qmp" && i+1 < len(prep.KvmQemu.CommandArgs) && prep.KvmQemu.CommandArgs[i+1] == fmt.Sprintf("unix:%s,server,nowait", prep.KvmQemu.ControlSocketPath) {
+			foundQmp = true
+		}
 	}
 	if !foundDrive {
 		t.Errorf("expected format=qcow2 in args, got %v", prep.KvmQemu.CommandArgs)
 	}
+	if !foundQmp {
+		t.Errorf("expected exact -qmp socket in args, got %v", prep.KvmQemu.CommandArgs)
+	}
 }
 
 func TestKvmExecutor_PrepareValidImageTyped(t *testing.T) {
@@ -75,6 +89,7 @@ func TestKvmExecutor_PrepareValidImageTyped(t *testing.T) {
 	defer os.Remove(f.Name())
 
 	spec := launch.LaunchSpec{
+		WorkloadID:   "test-typed-id",
 		Architecture: "x86_64",
 		Storage: []launch.StorageAttachmentSpec{
 			{HostPath: f.Name(), Format: "raw"},
@@ -85,7 +100,7 @@ func TestKvmExecutor_PrepareValidImageTyped(t *testing.T) {
 
 	prep, err := exec.Prepare(spec)
 	if err != nil {
-		t.Errorf("expected Prepare to succeed, got %v", err)
+		t.Fatalf("expected Prepare to succeed, got %v", err)
 	}
 
 	if prep.KvmQemu == nil {
@@ -113,6 +128,7 @@ func TestCloudHypervisorExecutor_PrepareValidImageTyped(t *testing.T) {
 	defer os.Remove(f.Name())
 
 	spec := launch.LaunchSpec{
+		WorkloadID:   "test-ch-id",
 		Architecture: "x86_64",
 		Storage: []launch.StorageAttachmentSpec{
 			{HostPath: f.Name(), Format: "raw"},
@@ -123,28 +139,40 @@ func TestCloudHypervisorExecutor_PrepareValidImageTyped(t *testing.T) {
 
 	prep, err := exec.Prepare(spec)
 	if err != nil {
-		t.Errorf("expected Prepare to succeed, got %v", err)
+		t.Fatalf("expected Prepare to succeed, got %v", err)
 	}
 
 	if prep.CloudHypervisor == nil {
 		t.Fatalf("expected CloudHypervisor prepared state, got nil")
 	}
 
+	if prep.CloudHypervisor.ControlSocketPath == "" || !strings.Contains(prep.CloudHypervisor.ControlSocketPath, "test-ch-id/cloudhypervisor.sock") {
+		t.Errorf("expected cloudhypervisor control socket path, got %v", prep.CloudHypervisor.ControlSocketPath)
+	}
+
 	foundDrive := false
-	for _, arg := range prep.CloudHypervisor.CommandArgs {
+	foundApi := false
+	for i, arg := range prep.CloudHypervisor.CommandArgs {
 		if strings.Contains(arg, "path="+f.Name()) {
 			foundDrive = true
 		}
+		if arg == "--api-socket" && i+1 < len(prep.CloudHypervisor.CommandArgs) && prep.CloudHypervisor.CommandArgs[i+1] == prep.CloudHypervisor.ControlSocketPath {
+			foundApi = true
+		}
 	}
 	if !foundDrive {
 		t.Errorf("expected path in args, got %v", prep.CloudHypervisor.CommandArgs)
 	}
+	if !foundApi {
+		t.Errorf("expected api-socket in args matching control socket, got %v", prep.CloudHypervisor.CommandArgs)
+	}
 }
 
 func TestFirecrackerExecutor_PrepareValidImageTyped(t *testing.T) {
 	exec := &FirecrackerExecutor{}
 
 	spec := launch.LaunchSpec{
+		WorkloadID:   "test-fc-id",
 		Architecture: "x86_64",
 		Storage: []launch.StorageAttachmentSpec{
 			{HostPath: "/tmp/rootfs.ext4", Format: "ext4", MountPoint: "/"},
@@ -156,7 +184,7 @@ func TestFirecrackerExecutor_PrepareValidImageTyped(t *testing.T) {
 
 	prep, err := exec.Prepare(spec)
 	if err != nil {
-		t.Errorf("expected Prepare to succeed, got %v", err)
+		t.Fatalf("expected Prepare to succeed, got %v", err)
 	}
 
 	if prep.Firecracker == nil {
@@ -169,6 +197,28 @@ func TestFirecrackerExecutor_PrepareValidImageTyped(t *testing.T) {
 	if prep.Firecracker.KernelImagePath != "/tmp/vmlinux" {
 		t.Errorf("expected kernel image path to be set from typed storage")
 	}
+
+	if prep.Firecracker.ControlSocketPath == "" || !strings.Contains(prep.Firecracker.ControlSocketPath, "test-fc-id/firecracker.sock") {
+		t.Errorf("expected firecracker control socket path, got %v", prep.Firecracker.ControlSocketPath)
+	}
+
+	foundApi := false
+	foundConfig := false
+	for i, arg := range prep.Firecracker.CommandArgs {
+		if arg == "--api-sock" && i+1 < len(prep.Firecracker.CommandArgs) && prep.Firecracker.CommandArgs[i+1] == prep.Firecracker.ControlSocketPath {
+			foundApi = true
+		}
+		if arg == "--config-file" && i+1 < len(prep.Firecracker.CommandArgs) && strings.Contains(prep.Firecracker.CommandArgs[i+1], "test-fc-id/fc-config.json") {
+			foundConfig = true
+		}
+	}
+
+	if !foundApi {
+		t.Errorf("expected --api-sock argument with control socket path, got %v", prep.Firecracker.CommandArgs)
+	}
+	if !foundConfig {
+		t.Errorf("expected --config-file argument with config path, got %v", prep.Firecracker.CommandArgs)
+	}
 }
 
 func TestKvmExecutor_ExecuteSpawnFails(t *testing.T) {
@@ -179,6 +229,7 @@ func TestKvmExecutor_ExecuteSpawnFails(t *testing.T) {
 		KvmQemu: &launch.PreparedQemuLaunch{
 			BinaryPath:  "qemu-system-non-existent-binary-12345",
 			CommandArgs: []string{"-m", "1024"},
+			// Omit ControlSocketPath to intentionally bypass MkdirAll during this test
 		},
 	}
 
diff --git a/schedune-control-plane/internal/runtime/paths.go b/schedune-control-plane/internal/runtime/paths.go
@@ -0,0 +1,39 @@
+package runtime
+
+import (
+	"fmt"
+	"path/filepath"
+	"regexp"
+)
+
+var validIDPattern = regexp.MustCompile(`^[a-zA-Z0-9_-]+$`)
+
+// SanitizeWorkloadID returns the ID if it strictly matches safe characters, preventing path traversal.
+func SanitizeWorkloadID(id string) (string, error) {
+	if !validIDPattern.MatchString(id) {
+		return "", fmt.Errorf("invalid characters in workload ID: %s", id)
+	}
+	return id, nil
+}
+
+// GetRuntimeDir returns the deterministic host directory for a specific workload's runtime artifacts.
+// The resulting path is isolated under var/run/schedune to ensure a predictable and secure workspace.
+func GetRuntimeDir(workloadID string) (string, error) {
+	safeID, err := SanitizeWorkloadID(workloadID)
+	if err != nil {
+		return "", err
+	}
+	return filepath.Join("var", "run", "schedune", safeID), nil
+}
+
+// GetControlSocketPath returns the deterministic path for the runtime control socket.
+func GetControlSocketPath(workloadID, backend string) (string, error) {
+	dir, err := GetRuntimeDir(workloadID)
+	if err != nil {
+		return "", err
+	}
+	if !validIDPattern.MatchString(backend) {
+		return "", fmt.Errorf("invalid characters in backend token: %s", backend)
+	}
+	return filepath.Join(dir, fmt.Sprintf("%s.sock", backend)), nil
+}
diff --git a/schedune-control-plane/internal/runtime/paths_test.go b/schedune-control-plane/internal/runtime/paths_test.go
