ultravioletrs · drasko · Dec 11, 2024 · Dec 11, 2024 · Dec 11, 2024 · Dec 11, 2024
diff --git a/manager/agentEventsLogs_test.go b/manager/agentEventsLogs_test.go
@@ -14,8 +14,8 @@ import (
 func TestComputationIDFromAddress(t *testing.T) {
 	ms := &managerService{
 		vms: map[string]vm.VM{
-			"comp1": qemu.NewVM(qemu.Config{VSockConfig: qemu.VSockConfig{GuestCID: 3}}, func(event interface{}) error { return nil }, "comp1"),
-			"comp2": qemu.NewVM(qemu.Config{VSockConfig: qemu.VSockConfig{GuestCID: 5}}, func(event interface{}) error { return nil }, "comp2"),
+			"comp1": qemu.NewVM(qemu.VMInfo{Config: qemu.Config{VSockConfig: qemu.VSockConfig{GuestCID: 3}}}, func(event interface{}) error { return nil }, "comp1"),
+			"comp2": qemu.NewVM(qemu.VMInfo{Config: qemu.Config{VSockConfig: qemu.VSockConfig{GuestCID: 5}}}, func(event interface{}) error { return nil }, "comp2"),
 		},
 	}
 
@@ -47,7 +47,7 @@ func TestReportBrokenConnection(t *testing.T) {
 	ms := &managerService{
 		eventsChan: make(chan *ClientStreamMessage, 1),
 		vms: map[string]vm.VM{
-			"comp1": qemu.NewVM(qemu.Config{VSockConfig: qemu.VSockConfig{GuestCID: 3}}, func(event interface{}) error { return nil }, "comp1"),
+			"comp1": qemu.NewVM(qemu.VMInfo{Config: qemu.Config{VSockConfig: qemu.VSockConfig{GuestCID: 3}}}, func(event interface{}) error { return nil }, "comp1"),
 		},
 	}
 

diff --git a/manager/attestation_policy.go b/manager/attestation_policy.go
@@ -34,17 +34,21 @@
 		return nil, fmt.Errorf("computationId %s not found", computationId)
 	}
 
-	config, ok := vm.GetConfig().(qemu.Config)
+	vmi, ok := vm.GetConfig().(qemu.VMInfo)
 	if !ok {
-		return nil, fmt.Errorf("failed to cast config to qemu.Config")
+		return nil, fmt.Errorf("failed to cast config to qemu.VMInfo")
 	}
 
+	ms.ap.Lock()
 	_, err := cmd.Output()
+	ms.ap.Unlock()
 	if err != nil {
 		return nil, err
 	}
 
+	ms.ap.Lock()
 	f, err := os.ReadFile("./attestation_policy.json")
+	ms.ap.Unlock()
 	if err != nil {
 		return nil, err
 	}
@@ -57,13 +61,13 @@
 
 	var measurement []byte
 	switch {
-	case config.EnableSEV:
-		measurement, err = guest.CalcLaunchDigest(guest.SEV, config.SMPCount, uint64(cpuid.CpuSigs[ms.qemuCfg.CPU]), config.OVMFCodeConfig.File, config.KernelFile, config.RootFsFile, strconv.Quote(qemu.KernelCommandLine), defGuestFeatures, "", vmmtypes.QEMU, false, "", 0)
+	case vmi.Config.EnableSEV:
+		measurement, err = guest.CalcLaunchDigest(guest.SEV, vmi.Config.SMPCount, uint64(cpuid.CpuSigs[ms.qemuCfg.CPU]), vmi.Config.OVMFCodeConfig.File, vmi.Config.KernelFile, vmi.Config.RootFsFile, strconv.Quote(qemu.KernelCommandLine), defGuestFeatures, "", vmmtypes.QEMU, false, "", 0)
 		if err != nil {
 			return nil, err
 		}
-	case config.EnableSEVSNP:
-		measurement, err = guest.CalcLaunchDigest(guest.SEV_SNP, config.SMPCount, uint64(cpuid.CpuSigs[config.CPU]), config.OVMFCodeConfig.File, config.KernelFile, config.RootFsFile, strconv.Quote(qemu.KernelCommandLine), defGuestFeatures, "", vmmtypes.QEMU, false, "", 0)
+	case vmi.Config.EnableSEVSNP:
+		measurement, err = guest.CalcLaunchDigest(guest.SEV_SNP, vmi.Config.SMPCount, uint64(cpuid.CpuSigs[vmi.Config.CPU]), vmi.Config.OVMFCodeConfig.File, vmi.Config.KernelFile, vmi.Config.RootFsFile, strconv.Quote(qemu.KernelCommandLine), defGuestFeatures, "", vmmtypes.QEMU, false, "", 0)
 		if err != nil {
 			return nil, err
 		}
@@ -72,14 +76,16 @@
 		attestationPolicy.Policy.Measurement = measurement
 	}
 
-	if config.HostData != "" {
-		hostData, err := base64.StdEncoding.DecodeString(config.HostData)
+	if vmi.Config.HostData != "" {
+		hostData, err := base64.StdEncoding.DecodeString(vmi.Config.HostData)
 		if err != nil {
 			return nil, err
 		}
 		attestationPolicy.Policy.HostData = hostData
 	}
 
+	attestationPolicy.Policy.MinimumLaunchTcb = vmi.LaunchTCB
+
 	f, err = protojson.Marshal(&attestationPolicy)
 	if err != nil {
 		return nil, err

diff --git a/manager/attestation_policy_test.go b/manager/attestation_policy_test.go
@@ -15,7 +15,7 @@ import (
 	"github.com/ultravioletrs/cocos/manager/vm/mocks"
 )
 
-func createDummyAttestationPolicyBinary(t *testing.T, behavior string) string {
+func CreateDummyAttestationPolicyBinary(t *testing.T, behavior string) string {
 	var content []byte
 	switch behavior {
 	case "success":
@@ -55,67 +55,79 @@ func TestFetchAttestationPolicy(t *testing.T) {
 			name:           "Valid SEV configuration",
 			computationId:  "sev-computation",
 			binaryBehavior: "success",
-			vmConfig: qemu.Config{
-				EnableSEV: true,
-				SMPCount:  2,
-				CPU:       "EPYC",
-				OVMFCodeConfig: qemu.OVMFCodeConfig{
-					File: "/path/to/OVMF_CODE.fd",
+			vmConfig: qemu.VMInfo{
+				Config: qemu.Config{
+					EnableSEV: true,
+					SMPCount:  2,
+					CPU:       "EPYC",
+					OVMFCodeConfig: qemu.OVMFCodeConfig{
+						File: "/path/to/OVMF_CODE.fd",
+					},
 				},
+				LaunchTCB: 0,
 			},
 			expectedError: "open /path/to/OVMF_CODE.fd: no such file or directory",
 		},
 		{
 			name:           "Valid SEV-SNP configuration",
 			computationId:  "sev-snp-computation",
 			binaryBehavior: "success",
-			vmConfig: qemu.Config{
-				EnableSEVSNP: true,
-				SMPCount:     4,
-				CPU:          "EPYC-v2",
-				OVMFCodeConfig: qemu.OVMFCodeConfig{
-					File: "/path/to/OVMF_CODE_SNP.fd",
+			vmConfig: qemu.VMInfo{
+				Config: qemu.Config{
+					EnableSEVSNP: true,
+					SMPCount:     4,
+					CPU:          "EPYC-v2",
+					OVMFCodeConfig: qemu.OVMFCodeConfig{
+						File: "/path/to/OVMF_CODE_SNP.fd",
+					},
 				},
+				LaunchTCB: 0,
 			},
 			expectedError: "open /path/to/OVMF_CODE_SNP.fd: no such file or director",
 		},
 		{
 			name:           "Invalid computation ID",
 			computationId:  "non-existent",
 			binaryBehavior: "success",
-			vmConfig:       qemu.Config{},
+			vmConfig:       qemu.VMInfo{Config: qemu.Config{}, LaunchTCB: 0},
 			expectedError:  "computationId non-existent not found",
 		},
 		{
 			name:           "Invalid config type",
 			computationId:  "invalid-config",
 			binaryBehavior: "success",
 			vmConfig:       struct{}{},
-			expectedError:  "failed to cast config to qemu.Config",
+			expectedError:  "failed to cast config to qemu.VMInfo",
 		},
 		{
 			name:           "Binary execution failure",
 			computationId:  "binary-fail",
 			binaryBehavior: "fail",
-			vmConfig: qemu.Config{
-				EnableSEV: true,
+			vmConfig: qemu.VMInfo{
+				Config: qemu.Config{
+					EnableSEV: true,
+				},
+				LaunchTCB: 0,
 			},
 			expectedError: "exit status 1",
 		},
 		{
 			name:           "JSON file not created",
 			computationId:  "no-json",
 			binaryBehavior: "no_json",
-			vmConfig: qemu.Config{
-				EnableSEV: true,
+			vmConfig: qemu.VMInfo{
+				Config: qemu.Config{
+					EnableSEV: true,
+				},
+				LaunchTCB: 0,
 			},
 			expectedError: "no such file or directory",
 		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			tempDir := createDummyAttestationPolicyBinary(t, tc.binaryBehavior)
+			tempDir := CreateDummyAttestationPolicyBinary(t, tc.binaryBehavior)
 			defer os.RemoveAll(tempDir)
 
 			ms := &managerService{

diff --git a/manager/qemu/persistence.go b/manager/qemu/persistence.go
@@ -13,7 +13,7 @@ const jsonExt = ".json"
 
 type VMState struct {
 	ID     string
-	Config Config
+	VMinfo VMInfo
 	PID    int
 }
 

diff --git a/manager/qemu/persistence_test.go b/manager/qemu/persistence_test.go
@@ -29,7 +29,7 @@ func TestSaveVM(t *testing.T) {
 
 	state := VMState{
 		ID:     "test-vm",
-		Config: Config{},
+		VMinfo: VMInfo{Config: Config{}},
 		PID:    1234,
 	}
 
@@ -50,8 +50,8 @@ func TestLoadVMs(t *testing.T) {
 
 	// Save two VMs
 	states := []VMState{
-		{ID: "vm1", Config: Config{}, PID: 1234},
-		{ID: "vm2", Config: Config{}, PID: 5678},
+		{ID: "vm1", VMinfo: VMInfo{Config: Config{}}, PID: 1234},
+		{ID: "vm2", VMinfo: VMInfo{Config: Config{}}, PID: 5678},
 	}
 
 	for _, state := range states {
@@ -82,7 +82,7 @@ func TestDeleteVM(t *testing.T) {
 	tempDir := t.TempDir()
 	fp, _ := NewFilePersistence(tempDir)
 
-	state := VMState{ID: "test-vm", Config: Config{}, PID: 1234}
+	state := VMState{ID: "test-vm", VMinfo: VMInfo{Config: Config{}}, PID: 1234}
 
 	// Save VM
 	if err := fp.SaveVM(state); err != nil {
@@ -126,7 +126,7 @@ func TestConcurrentAccess(t *testing.T) {
 	for i := 0; i < numGoroutines; i++ {
 		go func(id int) {
 			defer wg.Done()
-			state := VMState{ID: fmt.Sprintf("vm-%d", id), Config: Config{}, PID: id}
+			state := VMState{ID: fmt.Sprintf("vm-%d", id), VMinfo: VMInfo{Config: Config{}}, PID: id}
 			if err := fp.SaveVM(state); err != nil {
 				t.Errorf("Concurrent SaveVM failed: %v", err)
 			}

diff --git a/manager/qemu/vm.go b/manager/qemu/vm.go
@@ -25,8 +25,13 @@ const (
 	shutdownTimeout = 30 * time.Second
 )
 
+type VMInfo struct {
+	Config    Config
+	LaunchTCB uint64 `env:"LAUNCH_TCB" envDefault:"0"`
+}
+
 type qemuVM struct {
-	config           Config
+	vmi              VMInfo
 	cmd              *exec.Cmd
 	eventsLogsSender vm.EventSender
 	computationId    string
@@ -35,7 +40,7 @@ type qemuVM struct {
 
 func NewVM(config interface{}, eventsLogsSender vm.EventSender, computationId string) vm.VM {
 	return &qemuVM{
-		config:           config.(Config),
+		vmi:              config.(VMInfo),
 		eventsLogsSender: eventsLogsSender,
 		computationId:    computationId,
 		StateMachine:     vm.NewStateMachine(),
@@ -54,18 +59,18 @@ func (v *qemuVM) Start() (err error) {
 		return err
 	}
 
-	v.config.NetDevConfig.ID = fmt.Sprintf("%s-%s", v.config.NetDevConfig.ID, id)
-	v.config.SevConfig.ID = fmt.Sprintf("%s-%s", v.config.SevConfig.ID, id)
+	v.vmi.Config.NetDevConfig.ID = fmt.Sprintf("%s-%s", v.vmi.Config.NetDevConfig.ID, id)
+	v.vmi.Config.SevConfig.ID = fmt.Sprintf("%s-%s", v.vmi.Config.SevConfig.ID, id)
 
-	if !v.config.KernelHash {
+	if !v.vmi.Config.KernelHash {
 		// Copy firmware vars file.
-		srcFile := v.config.OVMFVarsConfig.File
+		srcFile := v.vmi.Config.OVMFVarsConfig.File
 		dstFile := fmt.Sprintf("%s/%s-%s.fd", tmpDir, firmwareVars, id)
 		err = internal.CopyFile(srcFile, dstFile)
 		if err != nil {
 			return err
 		}
-		v.config.OVMFVarsConfig.File = dstFile
+		v.vmi.Config.OVMFVarsConfig.File = dstFile
 	}
 
 	exe, args, err := v.executableAndArgs()
@@ -140,14 +145,14 @@ func (v *qemuVM) GetProcess() int {
 }
 
 func (v *qemuVM) executableAndArgs() (string, []string, error) {
-	exe, err := exec.LookPath(v.config.QemuBinPath)
+	exe, err := exec.LookPath(v.vmi.Config.QemuBinPath)
 	if err != nil {
 		return "", nil, err
 	}
 
-	args := v.config.ConstructQemuArgs()
+	args := v.vmi.Config.ConstructQemuArgs()
 
-	if v.config.UseSudo {
+	if v.vmi.Config.UseSudo {
 		args = append([]string{exe}, args...)
 		exe = "sudo"
 	}
@@ -191,9 +196,9 @@ func processExists(pid int) bool {
 }
 
 func (v *qemuVM) GetCID() int {
-	return v.config.GuestCID
+	return v.vmi.Config.GuestCID
 }
 
 func (v *qemuVM) GetConfig() interface{} {
-	return v.config
+	return v.vmi.Config
 }
diff --git a/manager/qemu/vm_test.go b/manager/qemu/vm_test.go
@@ -17,7 +17,7 @@ import (
 const testComputationID = "test-computation"
 
 func TestNewVM(t *testing.T) {
-	config := Config{}
+	config := VMInfo{Config: Config{}}
 
 	vm := NewVM(config, func(event interface{}) error { return nil }, testComputationID)
 
@@ -31,12 +31,12 @@ func TestStart(t *testing.T) {
 	assert.NoError(t, err)
 	defer os.Remove(tmpFile.Name())
 
-	config := Config{
+	config := VMInfo{Config: Config{
 		OVMFVarsConfig: OVMFVarsConfig{
 			File: tmpFile.Name(),
 		},
 		QemuBinPath: "echo",
-	}
+	}}
 
 	vm := NewVM(config, func(event interface{}) error { return nil }, testComputationID).(*qemuVM)
 
@@ -53,13 +53,13 @@ func TestStartSudo(t *testing.T) {
 	assert.NoError(t, err)
 	defer os.Remove(tmpFile.Name())
 
-	config := Config{
+	config := VMInfo{Config: Config{
 		OVMFVarsConfig: OVMFVarsConfig{
 			File: tmpFile.Name(),
 		},
 		QemuBinPath: "echo",
 		UseSudo:     true,
-	}
+	}}
 
 	vm := NewVM(config, func(event interface{}) error { return nil }, testComputationID).(*qemuVM)
 
@@ -113,8 +113,8 @@ func TestStop(t *testing.T) {
 
 func TestSetProcess(t *testing.T) {
 	vm := &qemuVM{
-		config: Config{
-			QemuBinPath: "echo", // Use 'echo' as a dummy QEMU binary
+		vmi: VMInfo{
+			Config: Config{QemuBinPath: "echo"}, // Use 'echo' as a dummy QEMU binary
 		},
 	}
 
@@ -139,9 +139,11 @@ func TestGetProcess(t *testing.T) {
 func TestGetCID(t *testing.T) {
 	expectedCID := 42
 	vm := &qemuVM{
-		config: Config{
-			VSockConfig: VSockConfig{
-				GuestCID: expectedCID,
+		vmi: VMInfo{
+			Config: Config{
+				VSockConfig: VSockConfig{
+					GuestCID: expectedCID,
+				},
 			},
 		},
 	}
@@ -155,7 +157,9 @@ func TestGetConfig(t *testing.T) {
 		QemuBinPath: "echo",
 	}
 	vm := &qemuVM{
-		config: expectedConfig,
+		vmi: VMInfo{
+			Config: expectedConfig,
+		},
 	}
 
 	config := vm.GetConfig()

diff --git a/manager/qemu/vsock.go b/manager/qemu/vsock.go
@@ -12,7 +12,7 @@
 const VsockConfigPort uint32 = 9999
 
 func (v *qemuVM) SendAgentConfig(ac agent.Computation) error {
-	conn, err := vsock.Dial(uint32(v.config.GuestCID), VsockConfigPort, nil)
+	conn, err := vsock.Dial(uint32(v.vmi.Config.GuestCID), VsockConfigPort, nil)
 	if err != nil {
 		return err
 	}