Don't let a stub drive handler know non-stub drives

Before this change, a stub drive handler knew non-stub drives, including
the root drive. However patching non-stub drives is not supported and
must be prevented.

Related: #230

Signed-off-by: Kazuyoshi Kato <[email protected]>

Author: kzys

runtime/service.go

@@ -588,6 +588,25 @@ func (s *service) SetVMMetadata(requestCtx context.Context, request *proto.SetVM
 	return &empty.Empty{}, nil
 }
 
+func (s *service) createStubDrives(stubDriveCount int) ([]models.Drive, error) {
+	paths, err := s.stubDriveHandler.StubDrivePaths(stubDriveCount)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to retrieve stub drive paths")
+	}
+
+	stubDrives := make([]models.Drive, 0)
+	for i, path := range paths {
+		stubDrives = append(stubDrives, models.Drive{
+			DriveID:      firecracker.String(fmt.Sprintf("stub%d", i)),
+			IsReadOnly:   firecracker.Bool(false),
+			PathOnHost:   firecracker.String(path),
+			IsRootDevice: firecracker.Bool(false),
+		})
+	}
+
+	return stubDrives, nil
+}
+
 func (s *service) buildVMConfiguration(req *proto.CreateVMRequest) (*firecracker.Config, error) {
 	logger := s.logger.WithField("cid", s.machineCID)
 
@@ -618,15 +637,6 @@ func (s *service) buildVMConfiguration(req *proto.CreateVMRequest) (*firecracker
 	}
 
 	// Drives configuration
-
-	var driveBuilder firecracker.DrivesBuilder
-	if root := req.RootDrive; root != nil {
-		driveBuilder = firecracker.NewDrivesBuilder(root.PathOnHost)
-	} else {
-		driveBuilder = firecracker.NewDrivesBuilder(s.config.RootDrive)
-	}
-
-	stubDriveIndex := int64(len(driveBuilder.Build()) - 1)
 	containerCount := int(req.ContainerCount)
 	if containerCount < 1 {
 		// containerCount should always be positive so that at least one container
@@ -635,26 +645,27 @@ func (s *service) buildVMConfiguration(req *proto.CreateVMRequest) (*firecracker
 		containerCount = 1
 	}
 
-	paths, err := s.stubDriveHandler.StubDrivePaths(containerCount)
+	// Create stub drives first and let stub driver handler manage the drives
+	stubDrives, err := s.createStubDrives(containerCount)
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to retrieve stub drive paths")
+		return nil, errors.Wrap(err, "failed to create stub drives")
 	}
+	s.stubDriveHandler.SetDrives(0, stubDrives)
 
-	for i, path := range paths {
-		driveID := fmt.Sprintf("stub%d", i)
-		driveBuilder = driveBuilder.AddDrive(path, false, func(drive *models.Drive) {
-			drive.DriveID = firecracker.String(driveID)
-		})
+	var driveBuilder firecracker.DrivesBuilder
+	// Create non-stub drives
+	if root := req.RootDrive; root != nil {
+		driveBuilder = firecracker.NewDrivesBuilder(root.PathOnHost)
+	} else {
+		driveBuilder = firecracker.NewDrivesBuilder(s.config.RootDrive)
 	}
 
-	cfg.Drives = driveBuilder.Build()
-	s.stubDriveHandler.SetDrives(stubDriveIndex, cfg.Drives)
-
 	for _, drive := range req.AdditionalDrives {
 		driveBuilder = addDriveFromProto(driveBuilder, drive)
 	}
 
-	cfg.Drives = driveBuilder.Build()
+	// a micro VM must know all drives
+	cfg.Drives = append(stubDrives, driveBuilder.Build()...)
 
 	// Setup network interfaces
 

runtime/service_test.go

@@ -62,17 +62,12 @@ func TestFindNextAvailableVsockCID_Isolated(t *testing.T) {
 
 func TestBuildVMConfiguration(t *testing.T) {
 	namespace := "TestBuildVMConfiguration"
-	rootfsDrive := models.Drive{
-		DriveID:      firecracker.String("stub0"),
-		PathOnHost:   nil, // will be populated in the for loop
-		IsReadOnly:   firecracker.Bool(false),
-		IsRootDevice: firecracker.Bool(false),
-	}
 	testcases := []struct {
-		name        string
-		request     *proto.CreateVMRequest
-		config      *Config
-		expectedCfg *firecracker.Config
+		name                   string
+		request                *proto.CreateVMRequest
+		config                 *Config
+		expectedCfg            *firecracker.Config
+		expectedStubDriveCount int
 	}{
 		{
 			name:    "ConfigFile",
@@ -88,7 +83,6 @@ func TestBuildVMConfiguration(t *testing.T) {
 				KernelArgs:      "KERNEL ARGS",
 				KernelImagePath: "KERNEL IMAGE",
 				Drives: []models.Drive{
-					rootfsDrive,
 					{
 						DriveID:      firecracker.String("root_drive"),
 						PathOnHost:   firecracker.String("ROOT DRIVE"),
@@ -103,6 +97,7 @@ func TestBuildVMConfiguration(t *testing.T) {
 					HtEnabled:   firecracker.Bool(false),
 				},
 			},
+			expectedStubDriveCount: 1,
 		},
 		{
 			name: "Input",
@@ -122,7 +117,6 @@ func TestBuildVMConfiguration(t *testing.T) {
 				KernelArgs:      "REQUEST KERNEL ARGS",
 				KernelImagePath: "REQUEST KERNEL IMAGE",
 				Drives: []models.Drive{
-					rootfsDrive,
 					{
 						DriveID:      firecracker.String("root_drive"),
 						PathOnHost:   firecracker.String("REQUEST ROOT DRIVE"),
@@ -137,6 +131,7 @@ func TestBuildVMConfiguration(t *testing.T) {
 					HtEnabled:   firecracker.Bool(false),
 				},
 			},
+			expectedStubDriveCount: 1,
 		},
 		{
 			name: "Priority",
@@ -160,7 +155,6 @@ func TestBuildVMConfiguration(t *testing.T) {
 				KernelArgs:      "REQUEST KERNEL ARGS",
 				KernelImagePath: "KERNEL IMAGE",
 				Drives: []models.Drive{
-					rootfsDrive,
 					{
 						DriveID:      firecracker.String("root_drive"),
 						PathOnHost:   firecracker.String("REQUEST ROOT DRIVE"),
@@ -175,6 +169,37 @@ func TestBuildVMConfiguration(t *testing.T) {
 					HtEnabled:   firecracker.Bool(false),
 				},
 			},
+			expectedStubDriveCount: 1,
+		},
+		{
+			name:    "Container Count",
+			request: &proto.CreateVMRequest{ContainerCount: 2},
+			config: &Config{
+				KernelArgs:      "KERNEL ARGS",
+				KernelImagePath: "KERNEL IMAGE",
+				RootDrive:       "ROOT DRIVE",
+				CPUTemplate:     "C3",
+				CPUCount:        2,
+			},
+			expectedCfg: &firecracker.Config{
+				KernelArgs:      "KERNEL ARGS",
+				KernelImagePath: "KERNEL IMAGE",
+				Drives: []models.Drive{
+					{
+						DriveID:      firecracker.String("root_drive"),
+						PathOnHost:   firecracker.String("ROOT DRIVE"),
+						IsReadOnly:   firecracker.Bool(false),
+						IsRootDevice: firecracker.Bool(true),
+					},
+				},
+				MachineCfg: models.MachineConfiguration{
+					CPUTemplate: models.CPUTemplateC3,
+					VcpuCount:   firecracker.Int64(2),
+					MemSizeMib:  firecracker.Int64(defaultMemSizeMb),
+					HtEnabled:   firecracker.Bool(false),
+				},
+			},
+			expectedStubDriveCount: 2,
 		},
 	}
 
@@ -197,11 +222,21 @@ func TestBuildVMConfiguration(t *testing.T) {
 			tc.expectedCfg.VsockDevices = []firecracker.VsockDevice{{Path: "root", CID: svc.machineCID}}
 			tc.expectedCfg.LogFifo = svc.shimDir.FirecrackerLogFifoPath()
 			tc.expectedCfg.MetricsFifo = svc.shimDir.FirecrackerMetricsFifoPath()
-			tc.expectedCfg.Drives[0].PathOnHost = firecracker.String(filepath.Join(tempDir, "stub0"))
+
+			drives := make([]models.Drive, tc.expectedStubDriveCount)
+			for i := 0; i < tc.expectedStubDriveCount; i++ {
+				drives[i].PathOnHost = firecracker.String(filepath.Join(tempDir, fmt.Sprintf("stub%d", i)))
+				drives[i].DriveID = firecracker.String(fmt.Sprintf("stub%d", i))
+				drives[i].IsReadOnly = firecracker.Bool(false)
+				drives[i].IsRootDevice = firecracker.Bool(false)
+			}
+			tc.expectedCfg.Drives = append(drives, tc.expectedCfg.Drives...)
 
 			actualCfg, err := svc.buildVMConfiguration(tc.request)
 			assert.NoError(t, err)
 			require.Equal(t, tc.expectedCfg, actualCfg)
+
+			require.Equal(t, tc.expectedStubDriveCount, len(svc.stubDriveHandler.drives), "The stub driver only knows stub drives")
 		})
 	}
 }