Allow drive paths already in the jailer rootfs

jailer.go

@@ -413,10 +413,19 @@ func LinkFilesHandler(kernelImageFileName string) Handler {
 			}
 
 			// copy all drives to the root fs
+			rootfsPrefix := rootfs + string(os.PathSeparator)
 			for i, drive := range m.Cfg.Drives {
 				hostPath := StringValue(drive.PathOnHost)
-				driveFileName := filepath.Base(hostPath)
 
+				// If the provided host path is already within the rootfs then just
+				// update the drive path to be rootfs-relative.
+				if strings.HasPrefix(hostPath, rootfsPrefix) {
+					rootfsRelativePath := strings.TrimPrefix(hostPath, rootfsPrefix)
+					m.Cfg.Drives[i].PathOnHost = String(rootfsRelativePath)
+					continue
+				}
+
+				driveFileName := filepath.Base(hostPath)
 				if err := os.Link(
 					hostPath,
 					filepath.Join(rootfs, driveFileName),

jailer_test.go

@@ -15,9 +15,12 @@ package firecracker
 import (
 	"context"
 	"fmt"
+	"os"
 	"path/filepath"
 	"reflect"
 	"testing"
+
+	"github.com/firecracker-microvm/firecracker-go-sdk/client/models"
 )
 
 func TestJailerBuilder(t *testing.T) {
@@ -174,9 +177,21 @@ func TestJailerBuilder(t *testing.T) {
 }
 
 func TestJail(t *testing.T) {
+	testTempDir, err := os.MkdirTemp("", "")
+	if err != nil {
+		t.Fatalf("failed to create temp dir for test: %s", err)
+	}
+	t.Cleanup(func() {
+		if err := os.RemoveAll(testTempDir); err != nil {
+			t.Errorf("failed to clean up test temp dir: %s", err)
+		}
+	})
+
 	var testCases = []struct {
 		name             string
 		jailerCfg        JailerConfig
+		drives           []models.Drive
+		testLinkFiles    bool
 		expectedArgs     []string
 		netns            string
 		socketPath       string
@@ -341,41 +356,138 @@ func TestJail(t *testing.T) {
 				rootfsFolderName,
 				"api.sock"),
 		},
+		{
+			name: "files already in jailer root",
+			jailerCfg: JailerConfig{
+				ID:             "my-test-id",
+				UID:            Int(123),
+				GID:            Int(100),
+				NumaNode:       Int(0),
+				ChrootBaseDir:  testTempDir,
+				ChrootStrategy: NewNaiveChrootStrategy("kernel-image-path"),
+				ExecFile:       "/path/to/firecracker",
+			},
+			drives: []models.Drive{
+				{
+					DriveID:      String("test-drive-id"),
+					IsReadOnly:   Bool(true),
+					IsRootDevice: Bool(true),
+					// Test a host path that is already inside the rootfs
+					PathOnHost: String(filepath.Join(
+						testTempDir,
+						"firecracker",
+						"my-test-id",
+						rootfsFolderName,
+						"image.ext4")),
+				},
+			},
+			expectedArgs: []string{
+				defaultJailerBin,
+				"--id",
+				"my-test-id",
+				"--uid",
+				"123",
+				"--gid",
+				"100",
+				"--exec-file",
+				"/path/to/firecracker",
+				"--cgroup",
+				"cpuset.mems=0",
+				"--cgroup",
+				fmt.Sprintf("cpuset.cpus=%s", getNumaCpuset(0)),
+				"--chroot-base-dir",
+				testTempDir,
+				"--",
+				"--no-seccomp",
+				"--api-sock",
+				"/run/firecracker.socket",
+			},
+			expectedSockPath: filepath.Join(
+				testTempDir,
+				"firecracker",
+				"my-test-id",
+				rootfsFolderName,
+				"run",
+				"firecracker.socket"),
+			testLinkFiles: true,
+		},
 	}
 	for _, c := range testCases {
 		t.Run(c.name, func(t *testing.T) {
+			// Clear the temp dir between tests.
+			if err := os.RemoveAll(testTempDir); err != nil {
+				t.Fatalf("failed to clear temp dir: %s", err)
+			}
+			if err := os.MkdirAll(testTempDir, 0755); err != nil {
+				t.Fatalf("failed to create temp dir: %s", err)
+			}
+
+			ctx := context.Background()
 			m := &Machine{
 				Handlers: Handlers{
 					FcInit: defaultFcInitHandlerList,
 				},
 			}
-			cfg := &Config{
-				VMID:       "vmid",
-				JailerCfg:  &c.jailerCfg,
-				NetNS:      c.netns,
-				SocketPath: c.socketPath,
+			m.Cfg = Config{
+				VMID:            "vmid",
+				JailerCfg:       &c.jailerCfg,
+				NetNS:           c.netns,
+				SocketPath:      c.socketPath,
+				Drives:          append([]models.Drive{}, c.drives...), // copy
+				KernelImagePath: createEmptyTempFile(t, testTempDir, "kernel-*"),
 			}
-			jail(context.Background(), m, cfg)
+			jail(ctx, m, &m.Cfg)
 
 			if e, a := c.expectedArgs, m.cmd.Args; !reflect.DeepEqual(e, a) {
 				t.Errorf("expected args %v, but received %v", e, a)
 			}
 
-			if e, a := c.expectedSockPath, cfg.SocketPath; e != a {
+			if e, a := c.expectedSockPath, m.Cfg.SocketPath; e != a {
 				t.Errorf("expected socket path %q, but received %q", e, a)
 			}
 
-			foundJailerHandler := false
+			var jailerHandler *Handler
 			for _, handler := range m.Handlers.FcInit.list {
 				if handler.Name == LinkFilesToRootFSHandlerName {
-					foundJailerHandler = true
+					jailerHandler = &handler
 					break
 				}
 			}
 
-			if !foundJailerHandler {
+			if jailerHandler == nil {
 				t.Errorf("did not find link files handler")
 			}
+
+			if !c.testLinkFiles {
+				return
+			}
+
+			rootfsPath := filepath.Join(
+				testTempDir, "firecracker", "my-test-id", rootfsFolderName,
+			)
+			if err := os.MkdirAll(rootfsPath, 0755); err != nil {
+				t.Fatalf("failed to create rootfs dir: %s", err)
+			}
+			if err := jailerHandler.Fn(ctx, m); err != nil {
+				t.Errorf("failed to run handler: %s", err)
+			}
+
+			// Drive paths should be updated to be rootfs-relative.
+			for i, d := range m.Cfg.Drives {
+				newPath := filepath.Join(rootfsPath, *d.PathOnHost)
+				if e, a := *c.drives[i].PathOnHost, newPath; e != a {
+					t.Errorf("expected drive host path %q, but received %q", e, a)
+				}
+			}
 		})
 	}
 }
+
+func createEmptyTempFile(t *testing.T, testTempDir, pattern string) string {
+	f, err := os.CreateTemp(testTempDir, pattern)
+	if err != nil {
+		t.Fatalf("Failed to create temp kernel image: %s", err)
+	}
+	defer f.Close()
+	return f.Name()
+}