COCOS-391- GCP Attestation policy (#405)

* Add AgentGrpcHost configuration to agent server

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Add SHA1 support to PcrValues and implement GCP attestation functions

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Add GCP attestation policy and OVMF download commands

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Add vTPM attestation support and update protobuf versions

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Remove Host field from AgentConfig and update related references

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Update GCP attestation policy to accept vCPU count as an argument

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Add SHA512 digest verification for OVMF file in GCP download command

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Update OVMF object name format in GCP attestation package

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

* Refactor attestation policy structure to use nested Config field

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

---------

Signed-off-by: Sammy Oina <sammyoina@gmail.com>

.github/dependabot.yaml

@@ -16,7 +16,3 @@ updates:
     schedule:
       interval: "weekly"
       day: "monday"
-    groups:
-      go-dependency:
-        patterns:
-          - "*"

.github/workflows/checkproto.yaml

@@ -33,8 +33,8 @@ jobs:
 
       - name: Set up protoc
         run: |
-          PROTOC_VERSION=28.1
-          PROTOC_GEN_VERSION=v1.34.2
+          PROTOC_VERSION=29.0
+          PROTOC_GEN_VERSION=v1.36.5
           PROTOC_GRPC_VERSION=v1.5.1
 
           # Download and install protoc
@@ -55,7 +55,7 @@ jobs:
       - name: Set up Cocos-AI
         run: |
           # Rename .pb.go files to .pb.go.tmp to prevent conflicts
-          for p in $(ls pkg/manager/*.pb.go); do
+          for p in $(ls manager/*.pb.go); do
             mv $p $p.tmp
           done
 
@@ -67,7 +67,7 @@ jobs:
           make protoc
 
           # Compare generated Go files with the original ones
-          for p in $(ls pkg/manager/*.pb.go); do
+          for p in $(ls manager/*.pb.go); do
             if ! cmp -s $p $p.tmp; then
               echo "Proto file and generated Go file $p are out of sync!"
               exit 1

.github/workflows/hal.yml

@@ -1,9 +1,9 @@
-name: Build and Release
+name: Build and Release Hal
 
 on:
   push:
     tags:
-      - '*'
+      - "*"
 
 jobs:
   build:
@@ -32,8 +32,8 @@ jobs:
         with:
           root-reserve-mb: 35000
           swap-size-mb: 1024
-          remove-dotnet: 'true'
-          remove-android: 'true'
+          remove-dotnet: "true"
+          remove-android: "true"
       - name: Check free space
         run: |
           echo "Free space:"
@@ -48,26 +48,33 @@ jobs:
       - name: Checkout cocos
         uses: actions/checkout@v4
         with:
-          repository: 'ultravioletrs/cocos'
+          repository: "ultravioletrs/cocos"
           path: cocos
 
       - name: Checkout buildroot
         uses: actions/checkout@v4
         with:
-          repository: 'buildroot/buildroot'
+          repository: "buildroot/buildroot"
           path: buildroot
           ref: 2024.11-rc2
 
-      - name: Build
+      - name: Build hal
         run: |
           cd buildroot
           make BR2_EXTERNAL=../cocos/hal/linux cocos_defconfig
           make
 
+      - name: Build cocos
+        run: |
+          cd cocos
+          make
+
       - name: Release
         uses: softprops/action-gh-release@v2
         with:
           files: |
             buildroot/output/images/bzImage
             buildroot/output/images/rootfs.cpio.gz
-
+            cocos/build/cocos-agent  
+            cocos/build/cocos-cli  
+            cocos/build/cocos-manager

.github/workflows/main.yaml

@@ -35,7 +35,7 @@ jobs:
           mkdir coverage
 
       - name: Run Agent tests
-        run: go test -v --race -covermode=atomic -coverprofile coverage/agent.out ./agent/...
+        run: go test --tags embed -v --race -covermode=atomic -coverprofile coverage/agent.out ./agent/...
 
       - name: Run cli tests
         run: go test -v --race -covermode=atomic -coverprofile coverage/cli.out ./cli/...

.golangci.yaml

@@ -4,6 +4,10 @@ run:
 issues:
   max-issues-per-linter: 10
   max-same-issues: 10
+  exclude-rules:
+    - linters:
+        - makezero
+      text: "with non-zero initialized length"
 
 linters-settings:
   importas:
@@ -67,7 +71,6 @@ linters:
     - dogsled
     - errchkjson
     - errname
-    - execinquery
     - copyloopvar
     - ginkgolinter
     - gocheckcompilerdirectives
@@ -78,3 +81,4 @@ linters:
     - mirror
     - nakedret
     - dupword
+

Makefile

@@ -12,6 +12,7 @@ CONFIG_DIR ?= /etc/cocos
 SERVICE_NAME ?= cocos-manager
 SERVICE_DIR ?= /etc/systemd/system
 SERVICE_FILE = init/systemd/$(SERVICE_NAME).service
+IGVM_BUILD_SCRIPT := ./scripts/igvmmeasure/igvm.sh
 
 define compile_service
 	CGO_ENABLED=$(CGO_ENABLED) GOOS=$(GOOS) GOARCH=$(GOARCH) GOARM=$(GOARM) \
@@ -27,8 +28,9 @@ endef
 
 all: $(SERVICES)
 
-$(SERVICES):
+$(SERVICES): 
 	$(call compile_service,$@)
+	@if [ "$@" = "cli" ] || [ "$@" = "manager" ]; then $(MAKE) build-igvm; fi
 
 $(ATTESTATION_POLICY):
 	$(MAKE) -C ./scripts/attestation_policy
@@ -37,6 +39,7 @@ protoc:
 	protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/agent.proto
 	protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative manager/manager.proto
 	protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/events/events.proto
+	protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/cvms/cvms.proto
 
 mocks:
 	mockery --config ./mockery.yml
@@ -60,3 +63,7 @@ stop:
 install_service:
 	sudo install -m 644 $(SERVICE_FILE) $(SERVICE_DIR)/$(SERVICE_NAME).service
 	sudo systemctl daemon-reload
+
+build-igvm:
+	@echo "Running build script for igvmmeasure..."
+	@$(IGVM_BUILD_SCRIPT)

README.md

@@ -1,65 +1,78 @@
-# Cocos AI
+<div align="center">
+
+# Cocos AI 🥥
+
+**Confidential Computing System for AI**
+
+**Made with ❤️ by [Ultraviolet](https://ultraviolet.rs/)**
 
 [![codecov](https://codecov.io/gh/ultravioletrs/cocos/graph/badge.svg?token=HX01LR01K9)](https://codecov.io/gh/ultravioletrs/cocos)
-![Go report card](https://goreportcard.com/badge/github.com/ultravioletrs/cocos)
+[![Go report card](https://goreportcard.com/badge/github.com/ultravioletrs/cocos)](https://goreportcard.com/report/github.com/ultravioletrs/cocos)
+[![License](https://img.shields.io/badge/license-Apache--2.0-blue)](LICENSE)
 
-[Cocos AI (Confdential Computing System for AI/ML)][cocos] is a platform for secure multiparty computation (SMPC)
-based on the [Confidential Computing][cc] and [Trusted Execution Environments (TEEs)][tee].
+### [Guide](https://docs.cocos.ultraviolet.rs) | [Contributing](CONTRIBUTING.md) | [Website](https://cocos.ai/)
+
+</div>
+
+## Introduction 🚀
+
+Cocos AI is a **cutting-edge platform** designed to enable secure multiparty computation (SMPC) using **Confidential Computing** and **Trusted Execution Environments (TEEs)**.
+
+It empowers organizations to collaboratively process sensitive data for AI/ML workloads while ensuring:
+
+- 🔒 **Data Privacy**: Your data stays encrypted and secure throughout the computation.
+- 🛡️ **Trust and Integrity**: Protected by hardware enclaves with robust remote attestation protocols.
+- 🤝 **Seamless Collaboration**: Multiple organizations can work together without exposing sensitive information.
 
 <p align="center">
-  <img src="https://cocos.ai/images/Collaborative%20AI.drawio.svg" width="500" height="500">
+  <img src="https://cocos.ai/images/Collaborative%20AI.drawio.svg" alt="Cocos AI Illustration" width="400" height="400">
 </p>
 
-With Cocos AI it becomes possible to run AI/ML workloads on combined datasets from multiple organizations
-while guaranteeing the privacy and security of the data and the algorithm.
-Data is always encrypted, protected by hardware secure enclaves (Trusted Execution Environments),
-attested via secure remote attestation protocols, and invisible to cloud processors or any other
-3rd party to which computation is offloaded.
+## Features 🛠️
 
-## Features
+Cocos AI provides essential features for secure and efficient collaborative AI/ML:
 
-Cocos AI is implementing the following features:
+- 🖥️ **TEE Enablement and Monitoring**: Secure VM management for deploying and monitoring workloads.
+- 🛡️ **Hardware Abstraction Layer (HAL)**: Built on a hardened Linux kernel, secure bootloader, and minimal root filesystem (minimal TCB).
+- 🕵️ **In-Enclave Agent and Networking Controller**: Essential system software for managing secure workloads.
+- 🔒 **Encrypted Data Transfer**: Asynchronous data transfer and secure result delivery.
+- 🛠️ **API for Platform Manipulation**: Programmatic control for managing workloads.
+- ✅ **Attestation and Verification Tools**: Hardware- and software-supported attestation for integrity assurance.
+- 🖱️ **Command-Line Interface (CLI)**: A user-friendly CLI for system interaction.
 
-- TEE enablement, deployment and monitoring (secure VM manager)
-- HAL for TEEs based on hardened Linux kernel, secure bootloader and custom-tailored embedded rootfs for minimal TCB
-- In-enclave agent, netowrking controller and other system software
-- Encrypted asynchronous data transfer and result delivery
-- API for programmable platform manipulation
-- HW and SW supported attestation with verification tools
-- CLI for system interaction
-
-## Usage
-
-Clone the repo and create binaries:
+## 🚀 Quick Start
 
+### Clone the Repository and Build Binaries
 ```bash
 git clone git@github.com:ultravioletrs/cocos.git
 make
 ```
 
-This will create 3 binaries:
+This will generate three binaries:
 ```bash
 ls build/
 # cocos-agent  cocos-cli  cocos-manager
 ```
 
-- Manager can be deployed on the AMD SEV-SNP host
-- Agent can be built into [EOS][eos]-based HAL
-- CLI can be used to communicate to remote Agent.
+### Deployment Overview:
+- **Manager**: Deploy on the AMD SEV-SNP host to orchestrate workloads.
+- **Agent**: Build into the [EOS](https://github.com/ultravioletrs/eos)-based HAL for secure enclave management.
+- **CLI**: Interact with remote agents to control operations.
 
-## Documentation
+## 📚 Documentation
 
-Project documentation is hosted at [Cocos AI official docs page][docs].
+Comprehensive documentation is available at the [official documentation page](https://docs.cocos.ultraviolet.rs).  
+For CLI usage details, visit the [CLI Documentation](https://docs.cocos.ultraviolet.rs/cli).
 
-Documentation is generated from the [docs repository](https://github.com/ultravioletrs/docs).
+Documentation is automatically generated from the [docs repository](https://github.com/ultravioletrs/docs). Contributions to documentation are welcome!
 
-## License
-Cocos AI is published under permissive open-source [Apache-2.0](LICENSE) license.
+## 🛡️ License
 
-[cc]: https://confidentialcomputing.io/white-papers-reports/
-[cocos]: https://cocos.ai/
-[rel]: https://github.com/ultravioletrs/cocos/releases
-[tee]: https://en.wikipedia.org/wiki/Trusted_execution_environment
-[docs]: https://docs.cocos.ultraviolet.rs
-[cli]: https://docs.cocos.ultraviolet.rs/cli
-[eos]: https://github.com/ultravioletrs/eos
+Cocos AI is published under the permissive open-source [Apache-2.0](LICENSE) license. Contributions are encouraged and appreciated!
+
+## 🌐 Links and Resources
+
+- [Cocos AI Website](https://cocos.ai/)
+- [Official Releases](https://github.com/ultravioletrs/cocos/releases)
+- [Confidential Computing Overview](https://confidentialcomputing.io/white-papers-reports/)
+- [Trusted Execution Environments (TEEs)](https://en.wikipedia.org/wiki/Trusted_execution_environment)

agent/README.md

@@ -6,15 +6,15 @@ Agent service provides a barebones HTTP and gRPC API and Service interface imple
 
 The service is configured using the environment variables from the following table. Note that any unset variables will be replaced with their default values.
 
-| Variable                      | Description                                            | Default                        |
-| ----------------------------- | ------------------------------------------------------ | ------------------------------ |
-| AGENT_LOG_LEVEL               | Log level for agent service (debug, info, warn, error) | info                           |
-| AGENT_GRPC_HOST               | Agent service gRPC host                                | ""                             |
-| AGENT_GRPC_PORT               | Agent service gRPC port                                | 7002                           |
-| AGENT_GRPC_SERVER_CERT        | Path to gRPC server certificate in pem format          | ""                             |
-| AGENT_GRPC_SERVER_KEY         | Path to gRPC server key in pem format                  | ""                             |
-| AGENT_GRPC_SERVER_CA_CERTS    | Path to gRPC server CA certificate                     | ""                             |
-| AGENT_GRPC_CLIENT_CA_CERTS    | Path to gRPC client CA certificate                     | ""                             |
+| Variable                       | Description                                            | Default                        |
+| ------------------------------ | ------------------------------------------------------ | ------------------------------ |
+| AGENT_LOG_LEVEL                | Log level for agent service (debug, info, warn, error) | debug                          |
+| AGENT_CVM_GRPC_HOST            | Agent service gRPC host                                | ""                             |
+| AGENT_CVM_GRPC_PORT            | Agent service gRPC port                                | 7001                           |
+| AGENT_CVM_GRPC_SERVER_CERT     | Path to gRPC server certificate in pem format          | ""                             |
+| AGENT_CVM_GRPC_SERVER_KEY      | Path to gRPC server key in pem format                  | ""                             |
+| AGENT_CVM_GRPC_SERVER_CA_CERTS | Path to gRPC server CA certificate                     | ""                             |
+| AGENT_CVM_GRPC_CLIENT_CA_CERTS | Path to gRPC client CA certificate                     | ""                             |
 
 
 ## Deployment

agent/agent.pb.go

@@ -3,8 +3,8 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.34.2
-// 	protoc        v5.28.1
+// 	protoc-gen-go v1.36.5
+// 	protoc        v5.29.0
 // source: agent/agent.proto
 
 package agent
@@ -14,6 +14,7 @@ import (
 	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 	reflect "reflect"
 	sync "sync"
+	unsafe "unsafe"
 )
 
 const (
@@ -24,21 +25,18 @@ const (
 )
 
 type AlgoRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	Algorithm     []byte                 `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
+	Requirements  []byte                 `protobuf:"bytes,2,opt,name=requirements,proto3" json:"requirements,omitempty"`
 	unknownFields protoimpl.UnknownFields
-
-	Algorithm    []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
-	Requirements []byte `protobuf:"bytes,2,opt,name=requirements,proto3" json:"requirements,omitempty"`
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *AlgoRequest) Reset() {
 	*x = AlgoRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[0]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[0]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *AlgoRequest) String() string {
@@ -49,7 +47,7 @@ func (*AlgoRequest) ProtoMessage() {}
 
 func (x *AlgoRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[0]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -79,18 +77,16 @@ func (x *AlgoRequest) GetRequirements() []byte {
 }
 
 type AlgoResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
 	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *AlgoResponse) Reset() {
 	*x = AlgoResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[1]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[1]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *AlgoResponse) String() string {
@@ -101,7 +97,7 @@ func (*AlgoResponse) ProtoMessage() {}
 
 func (x *AlgoResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[1]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -117,21 +113,18 @@ func (*AlgoResponse) Descriptor() ([]byte, []int) {
 }
 
 type DataRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	Dataset       []byte                 `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
+	Filename      string                 `protobuf:"bytes,2,opt,name=filename,proto3" json:"filename,omitempty"`
 	unknownFields protoimpl.UnknownFields
-
-	Dataset  []byte `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
-	Filename string `protobuf:"bytes,2,opt,name=filename,proto3" json:"filename,omitempty"`
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *DataRequest) Reset() {
 	*x = DataRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[2]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[2]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *DataRequest) String() string {
@@ -142,7 +135,7 @@ func (*DataRequest) ProtoMessage() {}
 
 func (x *DataRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[2]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -172,18 +165,16 @@ func (x *DataRequest) GetFilename() string {
 }
 
 type DataResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
 	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *DataResponse) Reset() {
 	*x = DataResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[3]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[3]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *DataResponse) String() string {
@@ -194,7 +185,7 @@ func (*DataResponse) ProtoMessage() {}
 
 func (x *DataResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[3]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -210,18 +201,16 @@ func (*DataResponse) Descriptor() ([]byte, []int) {
 }
 
 type ResultRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
 	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *ResultRequest) Reset() {
 	*x = ResultRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[4]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[4]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *ResultRequest) String() string {
@@ -232,7 +221,7 @@ func (*ResultRequest) ProtoMessage() {}
 
 func (x *ResultRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[4]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -248,20 +237,17 @@ func (*ResultRequest) Descriptor() ([]byte, []int) {
 }
 
 type ResultResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	File          []byte                 `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
 	unknownFields protoimpl.UnknownFields
-
-	File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *ResultResponse) Reset() {
 	*x = ResultResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[5]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[5]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *ResultResponse) String() string {
@@ -272,7 +258,7 @@ func (*ResultResponse) ProtoMessage() {}
 
 func (x *ResultResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[5]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -295,20 +281,19 @@ func (x *ResultResponse) GetFile() []byte {
 }
 
 type AttestationRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	TeeNonce      []byte                 `protobuf:"bytes,1,opt,name=teeNonce,proto3" json:"teeNonce,omitempty"`   // Should be less or equal 64 bytes.
+	VtpmNonce     []byte                 `protobuf:"bytes,2,opt,name=vtpmNonce,proto3" json:"vtpmNonce,omitempty"` // Should be less or equal 32 bytes.
+	Type          int32                  `protobuf:"varint,3,opt,name=type,proto3" json:"type,omitempty"`
 	unknownFields protoimpl.UnknownFields
-
-	ReportData []byte `protobuf:"bytes,1,opt,name=report_data,json=reportData,proto3" json:"report_data,omitempty"` // Should be of length 64.
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *AttestationRequest) Reset() {
 	*x = AttestationRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[6]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[6]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *AttestationRequest) String() string {
@@ -319,7 +304,7 @@ func (*AttestationRequest) ProtoMessage() {}
 
 func (x *AttestationRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[6]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -334,28 +319,39 @@ func (*AttestationRequest) Descriptor() ([]byte, []int) {
 	return file_agent_agent_proto_rawDescGZIP(), []int{6}
 }
 
-func (x *AttestationRequest) GetReportData() []byte {
+func (x *AttestationRequest) GetTeeNonce() []byte {
 	if x != nil {
-		return x.ReportData
+		return x.TeeNonce
 	}
 	return nil
 }
 
-type AttestationResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
+func (x *AttestationRequest) GetVtpmNonce() []byte {
+	if x != nil {
+		return x.VtpmNonce
+	}
+	return nil
+}
 
-	File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
+func (x *AttestationRequest) GetType() int32 {
+	if x != nil {
+		return x.Type
+	}
+	return 0
+}
+
+type AttestationResponse struct {
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	File          []byte                 `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *AttestationResponse) Reset() {
 	*x = AttestationResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[7]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[7]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *AttestationResponse) String() string {
@@ -366,7 +362,7 @@ func (*AttestationResponse) ProtoMessage() {}
 
 func (x *AttestationResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[7]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -390,7 +386,7 @@ func (x *AttestationResponse) GetFile() []byte {
 
 var File_agent_agent_proto protoreflect.FileDescriptor
 
-var file_agent_agent_proto_rawDesc = []byte{
+var file_agent_agent_proto_rawDesc = string([]byte{
 	0x0a, 0x11, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2f, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x70, 0x72,
 	0x6f, 0x74, 0x6f, 0x12, 0x05, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x22, 0x4f, 0x0a, 0x0b, 0x41, 0x6c,
 	0x67, 0x6f, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1c, 0x0a, 0x09, 0x61, 0x6c, 0x67,
@@ -407,40 +403,43 @@ var file_agent_agent_proto_rawDesc = []byte{
 	0x22, 0x0f, 0x0a, 0x0d, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
 	0x74, 0x22, 0x24, 0x0a, 0x0e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f,
 	0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28,
-	0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x22, 0x35, 0x0a, 0x12, 0x41, 0x74, 0x74, 0x65, 0x73,
-	0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1f, 0x0a,
-	0x0b, 0x72, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x5f, 0x64, 0x61, 0x74, 0x61, 0x18, 0x01, 0x20, 0x01,
-	0x28, 0x0c, 0x52, 0x0a, 0x72, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x44, 0x61, 0x74, 0x61, 0x22, 0x29,
-	0x0a, 0x13, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73,
-	0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x32, 0xfd, 0x01, 0x0a, 0x0c, 0x41, 0x67,
-	0x65, 0x6e, 0x74, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x33, 0x0a, 0x04, 0x41, 0x6c,
-	0x67, 0x6f, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f, 0x52,
-	0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41,
-	0x6c, 0x67, 0x6f, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x28, 0x01, 0x12,
-	0x33, 0x0a, 0x04, 0x44, 0x61, 0x74, 0x61, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e,
-	0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67,
-	0x65, 0x6e, 0x74, 0x2e, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
-	0x22, 0x00, 0x28, 0x01, 0x12, 0x39, 0x0a, 0x06, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x12, 0x14,
-	0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71,
-	0x75, 0x65, 0x73, 0x74, 0x1a, 0x15, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73,
-	0x75, 0x6c, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x12,
-	0x48, 0x0a, 0x0b, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x19,
-	0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69,
-	0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x61, 0x67, 0x65, 0x6e,
-	0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73,
-	0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x42, 0x09, 0x5a, 0x07, 0x2e, 0x2f, 0x61,
-	0x67, 0x65, 0x6e, 0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
-}
+	0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x22, 0x62, 0x0a, 0x12, 0x41, 0x74, 0x74, 0x65, 0x73,
+	0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1a, 0x0a,
+	0x08, 0x74, 0x65, 0x65, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52,
+	0x08, 0x74, 0x65, 0x65, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x12, 0x1c, 0x0a, 0x09, 0x76, 0x74, 0x70,
+	0x6d, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x09, 0x76, 0x74,
+	0x70, 0x6d, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, 0x65, 0x18,
+	0x03, 0x20, 0x01, 0x28, 0x05, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x22, 0x29, 0x0a, 0x13, 0x41,
+	0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e,
+	0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c,
+	0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x32, 0xfd, 0x01, 0x0a, 0x0c, 0x41, 0x67, 0x65, 0x6e, 0x74,
+	0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x33, 0x0a, 0x04, 0x41, 0x6c, 0x67, 0x6f, 0x12,
+	0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f, 0x52, 0x65, 0x71, 0x75,
+	0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f,
+	0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x28, 0x01, 0x12, 0x33, 0x0a, 0x04,
+	0x44, 0x61, 0x74, 0x61, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x44, 0x61, 0x74,
+	0x61, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74,
+	0x2e, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x28,
+	0x01, 0x12, 0x39, 0x0a, 0x06, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x12, 0x14, 0x2e, 0x61, 0x67,
+	0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
+	0x74, 0x1a, 0x15, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74,
+	0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x12, 0x48, 0x0a, 0x0b,
+	0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x19, 0x2e, 0x61, 0x67,
+	0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52,
+	0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41,
+	0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e,
+	0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x42, 0x09, 0x5a, 0x07, 0x2e, 0x2f, 0x61, 0x67, 0x65, 0x6e,
+	0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
+})
 
 var (
 	file_agent_agent_proto_rawDescOnce sync.Once
-	file_agent_agent_proto_rawDescData = file_agent_agent_proto_rawDesc
+	file_agent_agent_proto_rawDescData []byte
 )
 
 func file_agent_agent_proto_rawDescGZIP() []byte {
 	file_agent_agent_proto_rawDescOnce.Do(func() {
-		file_agent_agent_proto_rawDescData = protoimpl.X.CompressGZIP(file_agent_agent_proto_rawDescData)
+		file_agent_agent_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_agent_agent_proto_rawDesc), len(file_agent_agent_proto_rawDesc)))
 	})
 	return file_agent_agent_proto_rawDescData
 }
@@ -477,109 +476,11 @@ func file_agent_agent_proto_init() {
 	if File_agent_agent_proto != nil {
 		return
 	}
-	if !protoimpl.UnsafeEnabled {
-		file_agent_agent_proto_msgTypes[0].Exporter = func(v any, i int) any {
-			switch v := v.(*AlgoRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[1].Exporter = func(v any, i int) any {
-			switch v := v.(*AlgoResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[2].Exporter = func(v any, i int) any {
-			switch v := v.(*DataRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[3].Exporter = func(v any, i int) any {
-			switch v := v.(*DataResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[4].Exporter = func(v any, i int) any {
-			switch v := v.(*ResultRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[5].Exporter = func(v any, i int) any {
-			switch v := v.(*ResultResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[6].Exporter = func(v any, i int) any {
-			switch v := v.(*AttestationRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[7].Exporter = func(v any, i int) any {
-			switch v := v.(*AttestationResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-	}
 	type x struct{}
 	out := protoimpl.TypeBuilder{
 		File: protoimpl.DescBuilder{
 			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
-			RawDescriptor: file_agent_agent_proto_rawDesc,
+			RawDescriptor: unsafe.Slice(unsafe.StringData(file_agent_agent_proto_rawDesc), len(file_agent_agent_proto_rawDesc)),
 			NumEnums:      0,
 			NumMessages:   8,
 			NumExtensions: 0,
@@ -590,7 +491,6 @@ func file_agent_agent_proto_init() {
 		MessageInfos:      file_agent_agent_proto_msgTypes,
 	}.Build()
 	File_agent_agent_proto = out.File
-	file_agent_agent_proto_rawDesc = nil
 	file_agent_agent_proto_goTypes = nil
 	file_agent_agent_proto_depIdxs = nil
 }

agent/agent.proto

@@ -36,7 +36,9 @@ message ResultResponse {
 }
 
 message AttestationRequest {
-  bytes report_data = 1; // Should be of length 64.
+  bytes teeNonce = 1; // Should be less or equal 64 bytes.
+  bytes vtpmNonce = 2; // Should be less or equal 32 bytes.
+  int32 type = 3;
 }
 
 message AttestationResponse {

agent/agent_grpc.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
 // - protoc-gen-go-grpc v1.5.1
-// - protoc             v5.28.1
+// - protoc             v5.29.0
 // source: agent/agent.proto
 
 package agent

agent/algorithm/algorithm.go

@@ -46,4 +46,7 @@ func AlgorithmArgsFromContext(ctx context.Context) []string {
 type Algorithm interface {
 	// Run executes the algorithm and returns the result.
 	Run() error
+
+	// Stop stops the algorithm.
+	Stop() error
 }

agent/algorithm/binary/binary.go

@@ -20,29 +20,46 @@ type binary struct {
 	stderr   io.Writer
 	stdout   io.Writer
 	args     []string
+	cmd      *exec.Cmd
 }
 
-func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string, args []string) algorithm.Algorithm {
+func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string, args []string, cmpID string) algorithm.Algorithm {
 	return &binary{
 		algoFile: algoFile,
-		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
+		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
 		stdout:   &logging.Stdout{Logger: logger},
 		args:     args,
 	}
 }
 
 func (b *binary) Run() error {
-	cmd := exec.Command(b.algoFile, b.args...)
-	cmd.Stderr = b.stderr
-	cmd.Stdout = b.stdout
+	b.cmd = exec.Command(b.algoFile, b.args...)
+	b.cmd.Stderr = b.stderr
+	b.cmd.Stdout = b.stdout
 
-	if err := cmd.Start(); err != nil {
+	if err := b.cmd.Start(); err != nil {
 		return fmt.Errorf("error starting algorithm: %v", err)
 	}
 
-	if err := cmd.Wait(); err != nil {
+	if err := b.cmd.Wait(); err != nil {
 		return fmt.Errorf("algorithm execution error: %v", err)
 	}
 
 	return nil
 }
+
+func (b *binary) Stop() error {
+	if b.cmd == nil {
+		return nil
+	}
+
+	if b.cmd.ProcessState != nil && b.cmd.ProcessState.Exited() {
+		return nil
+	}
+
+	if err := b.cmd.Process.Kill(); err != nil {
+		return fmt.Errorf("error stopping algorithm: %v", err)
+	}
+
+	return nil
+}

agent/algorithm/binary/binary_test.go

@@ -18,7 +18,7 @@ func TestNewAlgorithm(t *testing.T) {
 	algoFile := "/path/to/algo"
 	args := []string{"arg1", "arg2"}
 
-	algo := NewAlgorithm(logger, eventsSvc, algoFile, args)
+	algo := NewAlgorithm(logger, eventsSvc, algoFile, args, "")
 
 	b, ok := algo.(*binary)
 	if !ok {
@@ -74,7 +74,7 @@ func TestBinaryRun(t *testing.T) {
 			logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 			eventsSvc := new(mocks.Service)
 
-			b := NewAlgorithm(logger, eventsSvc, tt.algoFile, tt.args).(*binary)
+			b := NewAlgorithm(logger, eventsSvc, tt.algoFile, tt.args, "").(*binary)
 
 			var stdout, stderr bytes.Buffer
 			b.stdout = &stdout

agent/algorithm/docker/docker.go

@@ -35,11 +35,11 @@ type docker struct {
 	stdout   io.Writer
 }
 
-func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string) algorithm.Algorithm {
+func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile, cmpID string) algorithm.Algorithm {
 	d := &docker{
 		algoFile: algoFile,
 		logger:   logger,
-		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
+		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
 		stdout:   &logging.Stdout{Logger: logger},
 	}
 
@@ -47,8 +47,6 @@ func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string
 }
 
 func (d *docker) Run() error {
-	ctx := context.Background()
-
 	// Create a new Docker client.
 	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
 	if err != nil {
@@ -62,8 +60,9 @@ func (d *docker) Run() error {
 	}
 	defer imageFile.Close()
 
+	ctx := context.Background()
 	// Load the Docker image from the tar file.
-	resp, err := cli.ImageLoad(ctx, imageFile, true)
+	resp, err := cli.ImageLoad(ctx, imageFile, client.ImageLoadWithQuiet(true))
 	if err != nil {
 		return fmt.Errorf("could not load Docker image from file: %v", err)
 	}
@@ -176,3 +175,8 @@ func writeToOut(readCloser io.ReadCloser, ioWriter io.Writer) error {
 
 	return nil
 }
+
+func (d *docker) Stop() error {
+	// To be supported later.
+	return nil
+}

agent/algorithm/docker/docker_test.go

@@ -18,7 +18,7 @@ func TestNewAlgorithm(t *testing.T) {
 	eventsSvc := new(mocks.Service)
 	algoFile := "/path/to/algo.tar"
 
-	algo := NewAlgorithm(logger, eventsSvc, algoFile)
+	algo := NewAlgorithm(logger, eventsSvc, algoFile, "")
 
 	d, ok := algo.(*docker)
 	assert.True(t, ok, "NewAlgorithm should return a *docker")

agent/algorithm/logging/logging.go

@@ -50,6 +50,7 @@ func (s *Stdout) Write(p []byte) (n int, err error) {
 type Stderr struct {
 	Logger   *slog.Logger
 	EventSvc events.Service
+	CmpID    string
 }
 
 // Write implements io.Writer.
@@ -70,9 +71,7 @@ func (s *Stderr) Write(p []byte) (n int, err error) {
 		s.Logger.Error(string(buf[:n]))
 	}
 
-	if err := s.EventSvc.SendEvent(algorithmRun, warningStatus, json.RawMessage{}); err != nil {
-		return len(p), err
-	}
+	s.EventSvc.SendEvent(s.CmpID, algorithmRun, warningStatus, json.RawMessage{})
 
 	return len(p), nil
 }

agent/algorithm/logging/logging_test.go

@@ -73,7 +73,7 @@ func TestStderrWrite(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			mockEventService := mocks.NewService(t)
-			mockEventService.On("SendEvent", "AlgorithmRun", manager.Warning.String(), mock.Anything).Return(nil)
+			mockEventService.On("SendEvent", mock.Anything, "AlgorithmRun", manager.Warning.String(), mock.Anything).Return(nil)
 
 			stderr := &Stderr{Logger: mglog.NewMock(), EventSvc: mockEventService}
 			n, err := stderr.Write([]byte(tt.input))

agent/algorithm/python/python.go

@@ -39,12 +39,13 @@ type python struct {
 	runtime          string
 	requirementsFile string
 	args             []string
+	cmd              *exec.Cmd
 }
 
-func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, runtime, requirementsFile, algoFile string, args []string) algorithm.Algorithm {
+func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, runtime, requirementsFile, algoFile string, args []string, cmpID string) algorithm.Algorithm {
 	p := &python{
 		algoFile:         algoFile,
-		stderr:           &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
+		stderr:           &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
 		stdout:           &logging.Stdout{Logger: logger},
 		requirementsFile: requirementsFile,
 		args:             args,
@@ -85,15 +86,15 @@ func (p *python) Run() error {
 	}
 
 	args := append([]string{p.algoFile}, p.args...)
-	cmd := exec.Command(pythonPath, args...)
-	cmd.Stderr = p.stderr
-	cmd.Stdout = p.stdout
+	p.cmd = exec.Command(pythonPath, args...)
+	p.cmd.Stderr = p.stderr
+	p.cmd.Stdout = p.stdout
 
-	if err := cmd.Start(); err != nil {
+	if err := p.cmd.Start(); err != nil {
 		return fmt.Errorf("error starting algorithm: %v", err)
 	}
 
-	if err := cmd.Wait(); err != nil {
+	if err := p.cmd.Wait(); err != nil {
 		return fmt.Errorf("algorithm execution error: %v", err)
 	}
 
@@ -103,3 +104,19 @@ func (p *python) Run() error {
 
 	return nil
 }
+
+func (p *python) Stop() error {
+	if p.cmd == nil {
+		return nil
+	}
+
+	if p.cmd.ProcessState != nil && p.cmd.ProcessState.Exited() {
+		return nil
+	}
+
+	if err := p.cmd.Process.Kill(); err != nil {
+		return fmt.Errorf("error stopping algorithm: %v", err)
+	}
+
+	return nil
+}

agent/algorithm/python/python_test.go

@@ -50,7 +50,7 @@ func TestNewAlgorithm(t *testing.T) {
 	algoFile := "algorithm.py"
 	args := []string{"--arg1", "value1"}
 
-	algo := NewAlgorithm(logger, eventsSvc, runtime, requirementsFile, algoFile, args)
+	algo := NewAlgorithm(logger, eventsSvc, runtime, requirementsFile, algoFile, args, "")
 
 	p, ok := algo.(*python)
 	if !ok {

agent/algorithm/wasm/wasm.go

@@ -24,12 +24,13 @@ type wasm struct {
 	stderr   io.Writer
 	stdout   io.Writer
 	args     []string
+	cmd      *exec.Cmd
 }
 
-func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string, args []string) algorithm.Algorithm {
+func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, args []string, algoFile, cmpID string) algorithm.Algorithm {
 	return &wasm{
 		algoFile: algoFile,
-		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
+		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
 		stdout:   &logging.Stdout{Logger: logger},
 		args:     args,
 	}
@@ -38,17 +39,33 @@ func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string
 func (w *wasm) Run() error {
 	args := append(mapDirOption, w.algoFile)
 	args = append(args, w.args...)
-	cmd := exec.Command(wasmRuntime, args...)
-	cmd.Stderr = w.stderr
-	cmd.Stdout = w.stdout
+	w.cmd = exec.Command(wasmRuntime, args...)
+	w.cmd.Stderr = w.stderr
+	w.cmd.Stdout = w.stdout
 
-	if err := cmd.Start(); err != nil {
+	if err := w.cmd.Start(); err != nil {
 		return fmt.Errorf("error starting algorithm: %v", err)
 	}
 
-	if err := cmd.Wait(); err != nil {
+	if err := w.cmd.Wait(); err != nil {
 		return fmt.Errorf("algorithm execution error: %v", err)
 	}
 
 	return nil
 }
+
+func (w *wasm) Stop() error {
+	if w.cmd == nil {
+		return nil
+	}
+
+	if w.cmd.ProcessState != nil && w.cmd.ProcessState.Exited() {
+		return nil
+	}
+
+	if err := w.cmd.Process.Kill(); err != nil {
+		return fmt.Errorf("error stopping algorithm: %v", err)
+	}
+
+	return nil
+}

agent/algorithm/wasm/wasm_test.go

@@ -18,7 +18,7 @@ func TestNewAlgorithm(t *testing.T) {
 	algoFile := "test.wasm"
 	args := []string{"arg1", "arg2"}
 
-	algo := NewAlgorithm(logger, eventsSvc, algoFile, args)
+	algo := NewAlgorithm(logger, eventsSvc, args, algoFile, "")
 
 	w, ok := algo.(*wasm)
 	if !ok {
@@ -54,7 +54,7 @@ func TestRunError(t *testing.T) {
 	algoFile := "test.wasm"
 	args := []string{"arg1", "arg2"}
 
-	w := NewAlgorithm(logger, eventsSvc, algoFile, args).(*wasm)
+	w := NewAlgorithm(logger, eventsSvc, args, algoFile, "").(*wasm)
 
 	err := w.Run()
 	if err == nil {

agent/api/grpc/endpoint.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/go-kit/kit/endpoint"
 	"github.com/ultravioletrs/cocos/agent"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
 )
 
 func algoEndpoint(svc agent.Service) endpoint.Endpoint {
@@ -70,7 +71,7 @@ func attestationEndpoint(svc agent.Service) endpoint.Endpoint {
 		if err := req.validate(); err != nil {
 			return attestationRes{}, err
 		}
-		file, err := svc.Attestation(ctx, req.ReportData)
+		file, err := svc.Attestation(ctx, req.TeeNonce, req.VtpmNonce, config.AttestationType(req.AttType))
 		if err != nil {
 			return attestationRes{}, err
 		}

agent/api/grpc/endpoint_test.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/agent/mocks"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
 	"golang.org/x/crypto/sha3"
 )
 
@@ -141,11 +142,11 @@ func TestAttestationEndpoint(t *testing.T) {
 	}{
 		{
 			name: "Success",
-			req:  attestationReq{ReportData: sha3.Sum512([]byte("report data"))},
+			req:  attestationReq{TeeNonce: sha3.Sum512([]byte("report data")), VtpmNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: config.SNP},
 		},
 		{
 			name:        "Service Error",
-			req:         attestationReq{ReportData: sha3.Sum512([]byte("report data"))},
+			req:         attestationReq{TeeNonce: sha3.Sum512([]byte("report data")), VtpmNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: config.SNP},
 			expectedErr: true,
 		},
 	}
@@ -153,9 +154,9 @@ func TestAttestationEndpoint(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			if tt.name == svcErr {
-				svc.On("Attestation", context.Background(), tt.req.ReportData).Return([]byte{}, errors.New("")).Once()
+				svc.On("Attestation", context.Background(), tt.req.TeeNonce, tt.req.VtpmNonce, tt.req.AttType).Return([]byte{}, errors.New("")).Once()
 			} else {
-				svc.On("Attestation", context.Background(), tt.req.ReportData).Return([]byte{}, nil).Once()
+				svc.On("Attestation", context.Background(), tt.req.TeeNonce, tt.req.VtpmNonce, tt.req.AttType).Return([]byte{}, nil).Once()
 			}
 			endpoint := attestationEndpoint(svc)
 			res, err := endpoint(context.Background(), tt.req)

agent/api/grpc/requests.go

@@ -4,6 +4,10 @@ package grpc
 
 import (
 	"errors"
+
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 )
 
 type algoReq struct {
@@ -38,9 +42,16 @@ func (req resultReq) validate() error {
 }
 
 type attestationReq struct {
-	ReportData [64]byte
+	TeeNonce  [quoteprovider.Nonce]byte
+	VtpmNonce [vtpm.Nonce]byte
+	AttType   config.AttestationType
 }
 
 func (req attestationReq) validate() error {
-	return nil
+	switch req.AttType {
+	case config.SNP, config.VTPM, config.SNPvTPM:
+		return nil
+	default:
+		return errors.New("invalid attestation type in attestation request")
+	}
 }

agent/api/grpc/server.go

@@ -11,6 +11,9 @@ import (
 
 	"github.com/go-kit/kit/transport/grpc"
 	"github.com/ultravioletrs/cocos/agent"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
@@ -21,6 +24,11 @@ const (
 	FileSizeKey = "file-size"
 )
 
+var (
+	ErrTEENonceLength  = errors.New("malformed report data, expect less or equal to 64 bytes")
+	ErrVTpmNonceLength = errors.New("malformed vTPM nonce, expect less or equal to 32 bytes")
+)
+
 var _ agent.AgentServiceServer = (*grpcServer)(nil)
 
 type grpcServer struct {
@@ -96,10 +104,20 @@ func encodeResultResponse(_ context.Context, response interface{}) (interface{},
 
 func decodeAttestationRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
 	req := grpcReq.(*agent.AttestationRequest)
-	if len(req.ReportData) != agent.ReportDataSize {
-		return nil, errors.New("malformed report data, expect 64 bytes")
+	var reportData [quoteprovider.Nonce]byte
+	var nonce [vtpm.Nonce]byte
+
+	if len(req.TeeNonce) > quoteprovider.Nonce {
+		return nil, ErrTEENonceLength
 	}
-	return attestationReq{ReportData: [agent.ReportDataSize]byte(req.ReportData)}, nil
+
+	if len(req.VtpmNonce) > vtpm.Nonce {
+		return nil, ErrVTpmNonceLength
+	}
+
+	copy(reportData[:], req.TeeNonce)
+	copy(nonce[:], req.VtpmNonce)
+	return attestationReq{TeeNonce: reportData, VtpmNonce: nonce, AttType: config.AttestationType(req.Type)}, nil
 }
 
 func encodeAttestationResponse(_ context.Context, response interface{}) (interface{}, error) {

agent/api/grpc/server_test.go

@@ -11,6 +11,9 @@ import (
 	"github.com/stretchr/testify/mock"
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/agent/mocks"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
 )
@@ -151,10 +154,12 @@ func TestAttestation(t *testing.T) {
 	mockStream := &MockAgentService_AttestationServer{ctx: context.Background()}
 	mockStream.On("Send", mock.AnythingOfType("*agent.AttestationResponse")).Return(nil)
 
-	reportData := [agent.ReportDataSize]byte{}
-	mockService.On("Attestation", mock.Anything, reportData).Return([]byte("attestation data"), nil)
+	reportData := [quoteprovider.Nonce]byte{}
+	vtpmNonce := [vtpm.Nonce]byte{}
+	attestationType := config.SNP
+	mockService.On("Attestation", mock.Anything, reportData, vtpmNonce, attestationType).Return([]byte("attestation data"), nil)
 
-	err := server.Attestation(&agent.AttestationRequest{ReportData: reportData[:]}, mockStream)
+	err := server.Attestation(&agent.AttestationRequest{TeeNonce: reportData[:]}, mockStream)
 	assert.NoError(t, err)
 
 	mockService.AssertExpectations(t)
@@ -199,11 +204,11 @@ func TestEncodeResultResponse(t *testing.T) {
 }
 
 func TestDecodeAttestationRequest(t *testing.T) {
-	reportData := [agent.ReportDataSize]byte{}
-	req := &agent.AttestationRequest{ReportData: reportData[:]}
+	nonce := [quoteprovider.Nonce]byte{}
+	req := &agent.AttestationRequest{TeeNonce: nonce[:]}
 	decoded, err := decodeAttestationRequest(context.Background(), req)
 	assert.NoError(t, err)
-	assert.Equal(t, attestationReq{ReportData: reportData}, decoded)
+	assert.Equal(t, attestationReq{TeeNonce: nonce}, decoded)
 }
 
 func TestEncodeAttestationResponse(t *testing.T) {

agent/api/logging.go

@@ -13,6 +13,9 @@ import (
 	"time"
 
 	"github.com/ultravioletrs/cocos/agent"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 )
 
 var _ agent.Service = (*loggingMiddleware)(nil)
@@ -27,6 +30,43 @@ func LoggingMiddleware(svc agent.Service, logger *slog.Logger) agent.Service {
 	return &loggingMiddleware{logger, svc}
 }
 
+// State implements agent.Service.
+func (lm *loggingMiddleware) State() (state string) {
+	defer func(begin time.Time) {
+		message := fmt.Sprintf("Method State took %s to complete with state %s", time.Since(begin), state)
+		lm.logger.Info(message)
+	}(time.Now())
+	return lm.svc.State()
+}
+
+// InitComputation implements agent.Service.
+func (lm *loggingMiddleware) InitComputation(ctx context.Context, cmp agent.Computation) (err error) {
+	defer func(begin time.Time) {
+		message := fmt.Sprintf("Method InitComputation for computation id %s took %s to complete", cmp.ID, time.Since(begin))
+		if err != nil {
+			lm.logger.WithGroup(cmp.ID).Warn(fmt.Sprintf("%s with error: %s", message, err))
+			return
+		}
+		lm.logger.WithGroup(cmp.ID).Info(fmt.Sprintf("%s without errors", message))
+	}(time.Now())
+
+	return lm.svc.InitComputation(ctx, cmp)
+}
+
+// StopComputation implements agent.Service.
+func (lm *loggingMiddleware) StopComputation(ctx context.Context) (err error) {
+	defer func(begin time.Time) {
+		message := fmt.Sprintf("Method StopComputation took %s to complete", time.Since(begin))
+		if err != nil {
+			lm.logger.Warn(fmt.Sprintf("%s with error: %s", message, err))
+			return
+		}
+		lm.logger.Info(fmt.Sprintf("%s without errors", message))
+	}(time.Now())
+
+	return lm.svc.StopComputation(ctx)
+}
+
 func (lm *loggingMiddleware) Algo(ctx context.Context, algorithm agent.Algorithm) (err error) {
 	defer func(begin time.Time) {
 		message := fmt.Sprintf("Method Algo took %s to complete", time.Since(begin))
@@ -66,7 +106,7 @@ func (lm *loggingMiddleware) Result(ctx context.Context) (response []byte, err e
 	return lm.svc.Result(ctx)
 }
 
-func (lm *loggingMiddleware) Attestation(ctx context.Context, reportData [agent.ReportDataSize]byte) (response []byte, err error) {
+func (lm *loggingMiddleware) Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType config.AttestationType) (response []byte, err error) {
 	defer func(begin time.Time) {
 		message := fmt.Sprintf("Method Attestation took %s to complete", time.Since(begin))
 		if err != nil {
@@ -76,5 +116,5 @@ func (lm *loggingMiddleware) Attestation(ctx context.Context, reportData [agent.
 		lm.logger.Info(fmt.Sprintf("%s without errors", message))
 	}(time.Now())
 
-	return lm.svc.Attestation(ctx, reportData)
+	return lm.svc.Attestation(ctx, reportData, nonce, attType)
 }

agent/api/metrics.go

@@ -12,6 +12,9 @@ import (
 
 	"github.com/go-kit/kit/metrics"
 	"github.com/ultravioletrs/cocos/agent"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 )
 
 var _ agent.Service = (*metricsMiddleware)(nil)
@@ -32,6 +35,36 @@ func MetricsMiddleware(svc agent.Service, counter metrics.Counter, latency metri
 	}
 }
 
+// State implements agent.Service.
+func (ms *metricsMiddleware) State() string {
+	defer func(begin time.Time) {
+		ms.counter.With("method", "state").Add(1)
+		ms.latency.With("method", "state").Observe(time.Since(begin).Seconds())
+	}(time.Now())
+
+	return ms.svc.State()
+}
+
+// InitComputation implements agent.Service.
+func (ms *metricsMiddleware) InitComputation(ctx context.Context, cmp agent.Computation) error {
+	defer func(begin time.Time) {
+		ms.counter.With("method", "init_computation").Add(1)
+		ms.latency.With("method", "init_computation").Observe(time.Since(begin).Seconds())
+	}(time.Now())
+
+	return ms.svc.InitComputation(ctx, cmp)
+}
+
+// StopComputation implements agent.Service.
+func (ms *metricsMiddleware) StopComputation(ctx context.Context) error {
+	defer func(begin time.Time) {
+		ms.counter.With("method", "stop_computation").Add(1)
+		ms.latency.With("method", "stop_computation").Observe(time.Since(begin).Seconds())
+	}(time.Now())
+
+	return ms.svc.StopComputation(ctx)
+}
+
 func (ms *metricsMiddleware) Algo(ctx context.Context, algorithm agent.Algorithm) error {
 	defer func(begin time.Time) {
 		ms.counter.With("method", "algo").Add(1)
@@ -59,11 +92,11 @@ func (ms *metricsMiddleware) Result(ctx context.Context) ([]byte, error) {
 	return ms.svc.Result(ctx)
 }
 
-func (ms *metricsMiddleware) Attestation(ctx context.Context, reportData [agent.ReportDataSize]byte) ([]byte, error) {
+func (ms *metricsMiddleware) Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType config.AttestationType) ([]byte, error) {
 	defer func(begin time.Time) {
 		ms.counter.With("method", "attestation").Add(1)
 		ms.latency.With("method", "attestation").Observe(time.Since(begin).Seconds())
 	}(time.Now())
 
-	return ms.svc.Attestation(ctx, reportData)
+	return ms.svc.Attestation(ctx, reportData, nonce, attType)
 }

agent/computations.go

@@ -13,8 +13,6 @@ import (
 var _ fmt.Stringer = (*Datasets)(nil)
 
 type AgentConfig struct {
-	LogLevel     string `json:"log_level,omitempty"`
-	Host         string `json:"host,omitempty"`
 	Port         string `json:"port,omitempty"`
 	CertFile     string `json:"cert_file,omitempty"`
 	KeyFile      string `json:"server_key,omitempty"`
@@ -30,7 +28,6 @@ type Computation struct {
 	Datasets        Datasets         `json:"datasets,omitempty"`
 	Algorithm       Algorithm        `json:"algorithm,omitempty"`
 	ResultConsumers []ResultConsumer `json:"result_consumers,omitempty"`
-	AgentConfig     AgentConfig      `json:"agent_config,omitempty"`
 }
 
 type ResultConsumer struct {

agent/computations_test.go

@@ -106,8 +106,6 @@ func TestDecompressToContext(t *testing.T) {
 
 func TestAgentConfigJSON(t *testing.T) {
 	config := AgentConfig{
-		LogLevel:     "info",
-		Host:         "localhost",
 		Port:         "8080",
 		CertFile:     "cert.pem",
 		KeyFile:      "key.pem",

agent/cvms/api/grpc/client.go

@@ -0,0 +1,349 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package grpc
+
+import (
+	"context"
+	"log/slog"
+	"sync"
+	"time"
+
+	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/agent/cvms"
+	"github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage"
+	"github.com/ultravioletrs/cocos/agent/cvms/server"
+	"golang.org/x/sync/errgroup"
+	"google.golang.org/protobuf/proto"
+)
+
+const (
+	reconnectInterval = 5 * time.Second
+	sendTimeout       = 5 * time.Second
+	pendingMsgFile    = "pending_messages.json"
+)
+
+var (
+	errCorruptedManifest  = errors.New("received manifest may be corrupted")
+	errUnknonwMessageType = errors.New("unknown message type")
+)
+
+type PendingMessage struct {
+	Message *cvms.ClientStreamMessage
+	Time    time.Time
+}
+
+type CVMSClient struct {
+	mu            sync.Mutex
+	stream        cvms.Service_ProcessClient
+	svc           agent.Service
+	messageQueue  chan *cvms.ClientStreamMessage
+	logger        *slog.Logger
+	runReqManager *runRequestManager
+	sp            server.AgentServer
+	storage       storage.Storage
+	reconnectFn   func(context.Context) (cvms.Service_ProcessClient, error)
+}
+
+// NewClient returns new gRPC client instance.
+func NewClient(stream cvms.Service_ProcessClient, svc agent.Service, messageQueue chan *cvms.ClientStreamMessage, logger *slog.Logger, sp server.AgentServer, storageDir string, reconnectFn func(context.Context) (cvms.Service_ProcessClient, error)) (*CVMSClient, error) {
+	store, err := storage.NewFileStorage(storageDir)
+	if err != nil {
+		return nil, err
+	}
+
+	return &CVMSClient{
+		stream:        stream,
+		svc:           svc,
+		messageQueue:  messageQueue,
+		logger:        logger,
+		runReqManager: newRunRequestManager(),
+		sp:            sp,
+		storage:       store,
+		reconnectFn:   reconnectFn,
+	}, nil
+}
+
+func (client *CVMSClient) Process(ctx context.Context, cancel context.CancelFunc) error {
+	for {
+		err := client.processWithRetry(ctx)
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
+		client.logger.Info("Connection lost, attempting to reconnect...", "error", err)
+		time.Sleep(reconnectInterval)
+
+		stream, err := client.reconnectFn(ctx)
+		if err != nil {
+			client.logger.Error("Failed to reconnect", "error", err)
+			continue
+		}
+
+		client.mu.Lock()
+		client.stream = stream
+		client.mu.Unlock()
+	}
+}
+
+func (client *CVMSClient) processWithRetry(ctx context.Context) error {
+	eg, ctx := errgroup.WithContext(ctx)
+
+	eg.Go(func() error {
+		return client.handleIncomingMessages(ctx)
+	})
+
+	eg.Go(func() error {
+		return client.handleOutgoingMessages(ctx)
+	})
+
+	return eg.Wait()
+}
+
+func (client *CVMSClient) handleIncomingMessages(ctx context.Context) error {
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			req, err := client.stream.Recv()
+			if err != nil {
+				return err
+			}
+			if err := client.processIncomingMessage(ctx, req); err != nil {
+				return err
+			}
+		}
+	}
+}
+
+func (client *CVMSClient) handleOutgoingMessages(ctx context.Context) error {
+	pendingMsgs, err := client.storage.Load()
+	if err != nil {
+		client.logger.Error("Failed to load pending messages", "error", err)
+	} else {
+		client.sendPendingMessages(pendingMsgs)
+	}
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case msg := <-client.messageQueue:
+			if err := client.sendStreamMessage(msg); err != nil {
+				if err := client.storage.Add(msg); err != nil {
+					client.logger.Error("Failed to store pending message", "error", err)
+				}
+				client.logger.Error("Failed to send message, stored for retry", "error", err)
+			}
+		}
+	}
+}
+
+func (client *CVMSClient) sendStreamMessage(msg *cvms.ClientStreamMessage) error {
+	client.mu.Lock()
+	defer client.mu.Unlock()
+
+	return client.stream.Send(msg)
+}
+
+func (client *CVMSClient) sendPendingMessages(pending []storage.Message) {
+	for _, pm := range pending {
+		if err := client.sendStreamMessage(pm.Message); err != nil {
+			if err := client.storage.Add(pm.Message); err != nil {
+				client.logger.Error("Failed to store pending message", "error", err)
+			}
+			client.logger.Error("Failed to resend pending message", "error", err)
+		} else {
+			client.logger.Info("Successfully resent pending message")
+		}
+	}
+
+	if err := client.storage.Clear(); err != nil {
+		client.logger.Error("Failed to clear pending messages", "error", err)
+	}
+}
+
+func (client *CVMSClient) processIncomingMessage(ctx context.Context, req *cvms.ServerStreamMessage) error {
+	switch mes := req.Message.(type) {
+	case *cvms.ServerStreamMessage_RunReqChunks:
+		return client.handleRunReqChunks(ctx, mes)
+	case *cvms.ServerStreamMessage_StopComputation:
+		go client.handleStopComputation(ctx, mes)
+	case *cvms.ServerStreamMessage_AgentStateReq:
+		client.handleAgentStateReq(mes)
+	default:
+		return errUnknonwMessageType
+	}
+	return nil
+}
+
+func (client *CVMSClient) handleAgentStateReq(mes *cvms.ServerStreamMessage_AgentStateReq) {
+	state := client.svc.State()
+
+	msg := &cvms.ClientStreamMessage_AgentStateRes{
+		AgentStateRes: &cvms.AgentStateRes{
+			State: state,
+			Id:    mes.AgentStateReq.Id,
+		},
+	}
+
+	client.sendMessage(&cvms.ClientStreamMessage{Message: msg})
+}
+
+func (client *CVMSClient) handleRunReqChunks(ctx context.Context, msg *cvms.ServerStreamMessage_RunReqChunks) error {
+	buffer, complete := client.runReqManager.addChunk(msg.RunReqChunks.Id, msg.RunReqChunks.Data, msg.RunReqChunks.IsLast)
+
+	if complete {
+		var runReq cvms.ComputationRunReq
+		if err := proto.Unmarshal(buffer, &runReq); err != nil {
+			return errors.Wrap(err, errCorruptedManifest)
+		}
+
+		go client.executeRun(ctx, &runReq)
+	}
+
+	return nil
+}
+
+func (client *CVMSClient) executeRun(ctx context.Context, runReq *cvms.ComputationRunReq) {
+	ac := agent.Computation{
+		ID:          runReq.Id,
+		Name:        runReq.Name,
+		Description: runReq.Description,
+	}
+
+	if runReq.Algorithm != nil {
+		ac.Algorithm = agent.Algorithm{
+			Hash:    [32]byte(runReq.Algorithm.Hash),
+			UserKey: runReq.Algorithm.UserKey,
+		}
+	}
+
+	for _, ds := range runReq.Datasets {
+		ac.Datasets = append(ac.Datasets, agent.Dataset{
+			Hash:    [32]byte(ds.Hash),
+			UserKey: ds.UserKey,
+		})
+	}
+
+	for _, rc := range runReq.ResultConsumers {
+		ac.ResultConsumers = append(ac.ResultConsumers, agent.ResultConsumer{
+			UserKey: rc.UserKey,
+		})
+	}
+
+	if err := client.svc.InitComputation(ctx, ac); err != nil {
+		client.logger.Warn(err.Error())
+		return
+	}
+
+	client.mu.Lock()
+	defer client.mu.Unlock()
+
+	if runReq.AgentConfig == nil {
+		runReq.AgentConfig = &cvms.AgentConfig{}
+	}
+
+	runRes := &cvms.ClientStreamMessage_RunRes{
+		RunRes: &cvms.RunResponse{
+			ComputationId: runReq.Id,
+		},
+	}
+
+	if err := client.sp.Start(agent.AgentConfig{
+		Port:         runReq.AgentConfig.Port,
+		CertFile:     runReq.AgentConfig.CertFile,
+		KeyFile:      runReq.AgentConfig.KeyFile,
+		ServerCAFile: runReq.AgentConfig.ServerCaFile,
+		ClientCAFile: runReq.AgentConfig.ClientCaFile,
+		AttestedTls:  runReq.AgentConfig.AttestedTls,
+	}, ac); err != nil {
+		client.logger.Warn(err.Error())
+		runRes.RunRes.Error = err.Error()
+	}
+
+	client.sendMessage(&cvms.ClientStreamMessage{Message: runRes})
+}
+
+func (client *CVMSClient) handleStopComputation(ctx context.Context, mes *cvms.ServerStreamMessage_StopComputation) {
+	msg := &cvms.ClientStreamMessage_StopComputationRes{
+		StopComputationRes: &cvms.StopComputationResponse{
+			ComputationId: mes.StopComputation.ComputationId,
+		},
+	}
+	if err := client.svc.StopComputation(ctx); err != nil {
+		msg.StopComputationRes.Message = err.Error()
+	}
+
+	client.mu.Lock()
+	if err := client.sp.Stop(); err != nil {
+		msg.StopComputationRes.Message = err.Error()
+	}
+	client.mu.Unlock()
+
+	client.sendMessage(&cvms.ClientStreamMessage{Message: msg})
+}
+
+func (client *CVMSClient) sendMessage(mes *cvms.ClientStreamMessage) {
+	ctx, cancel := context.WithTimeout(context.Background(), sendTimeout)
+	defer cancel()
+
+	select {
+	case client.messageQueue <- mes:
+	case <-ctx.Done():
+		client.logger.Warn("Failed to send message: timeout exceeded")
+	}
+}
+
+type runRequestManager struct {
+	requests map[string]*runRequest
+	mu       sync.Mutex
+}
+
+type runRequest struct {
+	buffer    []byte
+	lastChunk time.Time
+	timer     *time.Timer
+}
+
+func newRunRequestManager() *runRequestManager {
+	return &runRequestManager{
+		requests: make(map[string]*runRequest),
+	}
+}
+
+func (m *runRequestManager) addChunk(id string, chunk []byte, isLast bool) ([]byte, bool) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	req, exists := m.requests[id]
+	if !exists {
+		req = &runRequest{
+			buffer:    make([]byte, 0),
+			lastChunk: time.Now(),
+			timer:     time.AfterFunc(runReqTimeout, func() { m.timeoutRequest(id) }),
+		}
+		m.requests[id] = req
+	}
+
+	req.buffer = append(req.buffer, chunk...)
+	req.lastChunk = time.Now()
+	req.timer.Reset(runReqTimeout)
+
+	if isLast {
+		delete(m.requests, id)
+		req.timer.Stop()
+		return req.buffer, true
+	}
+
+	return nil, false
+}
+
+func (m *runRequestManager) timeoutRequest(id string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	delete(m.requests, id)
+	// Log timeout or handle it as needed
+}

agent/cvms/api/grpc/client_test.go

@@ -0,0 +1,209 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package grpc
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	mglog "github.com/absmach/magistrala/logger"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent/cvms"
+	servermocks "github.com/ultravioletrs/cocos/agent/cvms/server/mocks"
+	"github.com/ultravioletrs/cocos/agent/mocks"
+	"google.golang.org/grpc"
+	"google.golang.org/protobuf/proto"
+)
+
+type mockStream struct {
+	mock.Mock
+	grpc.ClientStream
+}
+
+func (m *mockStream) Recv() (*cvms.ServerStreamMessage, error) {
+	args := m.Called()
+	return args.Get(0).(*cvms.ServerStreamMessage), args.Error(1)
+}
+
+func (m *mockStream) Send(msg *cvms.ClientStreamMessage) error {
+	args := m.Called(msg)
+	return args.Error(0)
+}
+
+func TestManagerClient_Process(t *testing.T) {
+	tests := []struct {
+		name        string
+		setupMocks  func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer)
+		expectError bool
+		errorMsg    string
+	}{
+		{
+			name: "Stop computation",
+			setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer) {
+				mockStream.On("Recv").Return(&cvms.ServerStreamMessage{
+					Message: &cvms.ServerStreamMessage_StopComputation{
+						StopComputation: &cvms.StopComputation{},
+					},
+				}, nil)
+				mockStream.On("Send", mock.Anything).Return(nil)
+				mockSvc.On("StopComputation", mock.Anything).Return(nil)
+				mockServerSvc.On("Stop").Return(nil)
+			},
+			expectError: true,
+			errorMsg:    "context deadline exceeded",
+		},
+		{
+			name: "Run request chunks",
+			setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer) {
+				mockStream.On("Recv").Return(&cvms.ServerStreamMessage{
+					Message: &cvms.ServerStreamMessage_RunReqChunks{
+						RunReqChunks: &cvms.RunReqChunks{},
+					},
+				}, nil)
+				mockStream.On("Send", mock.Anything).Return(nil).Once()
+				mockSvc.On("Run", mock.Anything, mock.Anything).Return("", assert.AnError).Once()
+			},
+			expectError: true,
+		},
+		{
+			name: "Receive error",
+			setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer) {
+				mockStream.On("Recv").Return(&cvms.ServerStreamMessage{}, assert.AnError)
+			},
+			expectError: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			mockStream := new(mockStream)
+			mockSvc := new(mocks.Service)
+			mockServerSvc := new(servermocks.AgentServer)
+			messageQueue := make(chan *cvms.ClientStreamMessage)
+			logger := mglog.NewMock()
+
+			go func() {
+				<-messageQueue
+			}()
+
+			client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, t.TempDir(), func(ctx context.Context) (cvms.Service_ProcessClient, error) { return nil, nil })
+			assert.NoError(t, err)
+
+			ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+			defer cancel()
+
+			tc.setupMocks(mockStream, mockSvc, mockServerSvc)
+
+			err = client.Process(ctx, cancel)
+
+			if tc.expectError {
+				assert.Error(t, err)
+				if tc.errorMsg != "" {
+					assert.Contains(t, err.Error(), tc.errorMsg)
+				}
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
+
+func TestManagerClient_handleRunReqChunks(t *testing.T) {
+	mockStream := new(mockStream)
+	mockSvc := new(mocks.Service)
+	mockServerSvc := new(servermocks.AgentServer)
+	messageQueue := make(chan *cvms.ClientStreamMessage, 10)
+	logger := mglog.NewMock()
+
+	client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, t.TempDir(), func(ctx context.Context) (cvms.Service_ProcessClient, error) { return nil, nil })
+	assert.NoError(t, err)
+
+	runReq := &cvms.ComputationRunReq{
+		Id: "test-id",
+	}
+	runReqBytes, _ := proto.Marshal(runReq)
+
+	chunk1 := &cvms.ServerStreamMessage_RunReqChunks{
+		RunReqChunks: &cvms.RunReqChunks{
+			Id:     "chunk-1",
+			Data:   runReqBytes[:len(runReqBytes)/2],
+			IsLast: false,
+		},
+	}
+	chunk2 := &cvms.ServerStreamMessage_RunReqChunks{
+		RunReqChunks: &cvms.RunReqChunks{
+			Id:     "chunk-1",
+			Data:   runReqBytes[len(runReqBytes)/2:],
+			IsLast: true,
+		},
+	}
+
+	mockSvc.On("InitComputation", mock.Anything, mock.Anything).Return(nil)
+	mockServerSvc.On("Start", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+
+	err = client.handleRunReqChunks(context.Background(), chunk1)
+	assert.NoError(t, err)
+
+	err = client.handleRunReqChunks(context.Background(), chunk2)
+	assert.NoError(t, err)
+
+	// Wait for the goroutine to finish
+	time.Sleep(50 * time.Millisecond)
+
+	mockSvc.AssertExpectations(t)
+	assert.Len(t, messageQueue, 1)
+
+	msg := <-messageQueue
+	runRes, ok := msg.Message.(*cvms.ClientStreamMessage_RunRes)
+	assert.True(t, ok)
+	assert.Equal(t, "test-id", runRes.RunRes.ComputationId)
+}
+
+func TestManagerClient_handleStopComputation(t *testing.T) {
+	mockStream := new(mockStream)
+	mockSvc := new(mocks.Service)
+	mockServerSvc := new(servermocks.AgentServer)
+	messageQueue := make(chan *cvms.ClientStreamMessage, 10)
+	logger := mglog.NewMock()
+
+	client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, t.TempDir(), func(ctx context.Context) (cvms.Service_ProcessClient, error) { return nil, nil })
+	assert.NoError(t, err)
+
+	stopReq := &cvms.ServerStreamMessage_StopComputation{
+		StopComputation: &cvms.StopComputation{
+			ComputationId: "test-comp-id",
+		},
+	}
+
+	mockSvc.On("StopComputation", mock.Anything).Return(nil)
+	mockServerSvc.On("Stop").Return(nil)
+
+	client.handleStopComputation(context.Background(), stopReq)
+
+	// Wait for the goroutine to finish
+	time.Sleep(50 * time.Millisecond)
+
+	mockSvc.AssertExpectations(t)
+	assert.Len(t, messageQueue, 1)
+
+	msg := <-messageQueue
+	stopRes, ok := msg.Message.(*cvms.ClientStreamMessage_StopComputationRes)
+	assert.True(t, ok)
+	assert.Equal(t, "test-comp-id", stopRes.StopComputationRes.ComputationId)
+	assert.Empty(t, stopRes.StopComputationRes.Message)
+}
+
+func TestManagerClient_timeoutRequest(t *testing.T) {
+	rm := newRunRequestManager()
+	rm.requests["test-id"] = &runRequest{
+		timer:     time.NewTimer(100 * time.Millisecond),
+		buffer:    []byte("test-data"),
+		lastChunk: time.Now(),
+	}
+
+	rm.timeoutRequest("test-id")
+
+	assert.Len(t, rm.requests, 0)
+}

agent/cvms/api/grpc/doc.go

@@ -0,0 +1,5 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Package grpc contains implementation of kit service gRPC API.
+package grpc

agent/cvms/api/grpc/server.go

@@ -0,0 +1,133 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package grpc
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"io"
+	"time"
+
+	"github.com/ultravioletrs/cocos/agent/cvms"
+	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/protobuf/proto"
+)
+
+var (
+	_                cvms.ServiceServer = (*grpcServer)(nil)
+	ErrUnexpectedMsg                    = errors.New("unknown message type")
+)
+
+const (
+	bufferSize    = 1024 * 1024 // 1 MB
+	runReqTimeout = 30 * time.Second
+)
+
+type SendFunc func(*cvms.ServerStreamMessage) error
+
+type grpcServer struct {
+	cvms.UnimplementedServiceServer
+	incoming chan *cvms.ClientStreamMessage
+	svc      Service
+}
+
+type Service interface {
+	Run(ctx context.Context, ipAddress string, sendMessage SendFunc, authInfo credentials.AuthInfo)
+}
+
+// NewServer returns new AuthServiceServer instance.
+func NewServer(incoming chan *cvms.ClientStreamMessage, svc Service) cvms.ServiceServer {
+	return &grpcServer{
+		incoming: incoming,
+		svc:      svc,
+	}
+}
+
+func (s *grpcServer) Process(stream cvms.Service_ProcessServer) error {
+	client, ok := peer.FromContext(stream.Context())
+	if !ok {
+		return errors.New("failed to get peer info")
+	}
+
+	eg, ctx := errgroup.WithContext(stream.Context())
+
+	eg.Go(func() error {
+		for {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
+				req, err := stream.Recv()
+				if err != nil {
+					return err
+				}
+				s.incoming <- req
+			}
+		}
+	})
+
+	eg.Go(func() error {
+		sendMessage := func(msg *cvms.ServerStreamMessage) error {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
+				switch m := msg.Message.(type) {
+				case *cvms.ServerStreamMessage_RunReq:
+					return s.sendRunReqInChunks(stream, m.RunReq)
+				default:
+					return stream.Send(msg)
+				}
+			}
+		}
+
+		s.svc.Run(ctx, client.Addr.String(), sendMessage, client.AuthInfo)
+		return nil
+	})
+
+	return eg.Wait()
+}
+
+func (s *grpcServer) sendRunReqInChunks(stream cvms.Service_ProcessServer, runReq *cvms.ComputationRunReq) error {
+	data, err := proto.Marshal(runReq)
+	if err != nil {
+		return err
+	}
+
+	dataBuffer := bytes.NewBuffer(data)
+	buf := make([]byte, bufferSize)
+
+	for {
+		n, err := dataBuffer.Read(buf)
+		isLast := false
+
+		if err == io.EOF {
+			isLast = true
+		} else if err != nil {
+			return err
+		}
+
+		chunk := &cvms.ServerStreamMessage{
+			Message: &cvms.ServerStreamMessage_RunReqChunks{
+				RunReqChunks: &cvms.RunReqChunks{
+					Id:     runReq.Id,
+					Data:   buf[:n],
+					IsLast: isLast,
+				},
+			},
+		}
+
+		if err := stream.Send(chunk); err != nil {
+			return err
+		}
+
+		if isLast {
+			break
+		}
+	}
+
+	return nil
+}

agent/cvms/api/grpc/server_test.go

@@ -10,24 +10,24 @@ import (
 	"github.com/absmach/magistrala/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
-	"github.com/ultravioletrs/cocos/manager"
+	"github.com/ultravioletrs/cocos/agent/cvms"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/peer"
 )
 
 type mockServerStream struct {
 	mock.Mock
-	manager.ManagerService_ProcessServer
+	cvms.Service_ProcessServer
 }
 
-func (m *mockServerStream) Send(msg *manager.ServerStreamMessage) error {
+func (m *mockServerStream) Send(msg *cvms.ServerStreamMessage) error {
 	args := m.Called(msg)
 	return args.Error(0)
 }
 
-func (m *mockServerStream) Recv() (*manager.ClientStreamMessage, error) {
+func (m *mockServerStream) Recv() (*cvms.ClientStreamMessage, error) {
 	args := m.Called()
-	return args.Get(0).(*manager.ClientStreamMessage), args.Error(1)
+	return args.Get(0).(*cvms.ClientStreamMessage), args.Error(1)
 }
 
 func (m *mockServerStream) Context() context.Context {
@@ -44,7 +44,7 @@ func (m *mockService) Run(ctx context.Context, ipAddress string, sendMessage Sen
 }
 
 func TestNewServer(t *testing.T) {
-	incoming := make(chan *manager.ClientStreamMessage)
+	incoming := make(chan *cvms.ClientStreamMessage)
 	mockSvc := new(mockService)
 
 	server := NewServer(incoming, mockSvc)
@@ -56,19 +56,19 @@ func TestNewServer(t *testing.T) {
 func TestGrpcServer_Process(t *testing.T) {
 	tests := []struct {
 		name          string
-		recvReturn    *manager.ClientStreamMessage
+		recvReturn    *cvms.ClientStreamMessage
 		recvError     error
 		expectedError string
 	}{
 		{
 			name:          "Process with context deadline exceeded",
-			recvReturn:    &manager.ClientStreamMessage{},
+			recvReturn:    &cvms.ClientStreamMessage{},
 			recvError:     nil,
 			expectedError: "context deadline exceeded",
 		},
 		{
 			name:          "Process with Recv error",
-			recvReturn:    &manager.ClientStreamMessage{},
+			recvReturn:    &cvms.ClientStreamMessage{},
 			recvError:     errors.New("recv error"),
 			expectedError: "recv error",
 		},
@@ -76,7 +76,7 @@ func TestGrpcServer_Process(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			incoming := make(chan *manager.ClientStreamMessage, 1)
+			incoming := make(chan *cvms.ClientStreamMessage, 1)
 			mockSvc := new(mockService)
 			server := NewServer(incoming, mockSvc).(*grpcServer)
 
@@ -111,13 +111,13 @@ func TestGrpcServer_Process(t *testing.T) {
 }
 
 func TestGrpcServer_sendRunReqInChunks(t *testing.T) {
-	incoming := make(chan *manager.ClientStreamMessage)
+	incoming := make(chan *cvms.ClientStreamMessage)
 	mockSvc := new(mockService)
 	server := NewServer(incoming, mockSvc).(*grpcServer)
 
 	mockStream := new(mockServerStream)
 
-	runReq := &manager.ComputationRunReq{
+	runReq := &cvms.ComputationRunReq{
 		Id: "test-id",
 	}
 
@@ -125,10 +125,10 @@ func TestGrpcServer_sendRunReqInChunks(t *testing.T) {
 	for i := range largePayload {
 		largePayload[i] = byte(i % 256)
 	}
-	runReq.Algorithm = &manager.Algorithm{}
+	runReq.Algorithm = &cvms.Algorithm{}
 	runReq.Algorithm.UserKey = largePayload
 
-	mockStream.On("Send", mock.AnythingOfType("*manager.ServerStreamMessage")).Return(nil).Times(4)
+	mockStream.On("Send", mock.AnythingOfType("*cvms.ServerStreamMessage")).Return(nil).Times(4)
 
 	err := server.sendRunReqInChunks(mockStream, runReq)
 
@@ -139,7 +139,7 @@ func TestGrpcServer_sendRunReqInChunks(t *testing.T) {
 	assert.Equal(t, 4, len(calls))
 
 	for i, call := range calls {
-		msg := call.Arguments[0].(*manager.ServerStreamMessage)
+		msg := call.Arguments[0].(*cvms.ServerStreamMessage)
 		chunk := msg.GetRunReqChunks()
 
 		assert.NotNil(t, chunk)
@@ -174,9 +174,9 @@ func TestGrpcServer_ProcessWithMockService(t *testing.T) {
 				mockSvc.On("Run", mock.Anything, "test", mock.Anything, mock.AnythingOfType("mockAuthInfo")).
 					Run(func(args mock.Arguments) {
 						sendFunc := args.Get(2).(SendFunc)
-						runReq := &manager.ComputationRunReq{Id: "test-run-id"}
-						err := sendFunc(&manager.ServerStreamMessage{
-							Message: &manager.ServerStreamMessage_RunReq{
+						runReq := &cvms.ComputationRunReq{Id: "test-run-id"}
+						err := sendFunc(&cvms.ServerStreamMessage{
+							Message: &cvms.ServerStreamMessage_RunReq{
 								RunReq: runReq,
 							},
 						})
@@ -184,34 +184,17 @@ func TestGrpcServer_ProcessWithMockService(t *testing.T) {
 					}).
 					Return()
 
-				mockStream.On("Send", mock.MatchedBy(func(msg *manager.ServerStreamMessage) bool {
+				mockStream.On("Send", mock.MatchedBy(func(msg *cvms.ServerStreamMessage) bool {
 					chunks := msg.GetRunReqChunks()
 					return chunks != nil && chunks.Id == "test-run-id"
 				})).Return(nil)
 			},
 		},
-		{
-			name: "Terminate Request Test",
-			setupMockFn: func(mockSvc *mockService, mockStream *mockServerStream) {
-				mockSvc.On("Run", mock.Anything, "test", mock.Anything, mock.AnythingOfType("mockAuthInfo")).
-					Run(func(args mock.Arguments) {
-						sendFunc := args.Get(2).(SendFunc)
-						err := sendFunc(&manager.ServerStreamMessage{
-							Message: &manager.ServerStreamMessage_TerminateReq{
-								TerminateReq: &manager.Terminate{},
-							},
-						})
-						assert.NoError(t, err)
-					}).Return()
-
-				mockStream.On("Send", mock.AnythingOfType("*manager.ServerStreamMessage")).Return(nil)
-			},
-		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			incoming := make(chan *manager.ClientStreamMessage, 10)
+			incoming := make(chan *cvms.ClientStreamMessage, 10)
 			mockSvc := new(mockService)
 			server := NewServer(incoming, mockSvc).(*grpcServer)
 
@@ -231,7 +214,7 @@ func TestGrpcServer_ProcessWithMockService(t *testing.T) {
 			})
 
 			mockStream.On("Context").Return(peerCtx)
-			mockStream.On("Recv").Return(&manager.ClientStreamMessage{}, nil).Maybe()
+			mockStream.On("Recv").Return(&cvms.ClientStreamMessage{}, nil).Maybe()
 
 			tt.setupMockFn(mockSvc, mockStream)
 
@@ -251,18 +234,18 @@ func TestGrpcServer_ProcessWithMockService(t *testing.T) {
 }
 
 func TestGrpcServer_sendRunReqInChunksError(t *testing.T) {
-	incoming := make(chan *manager.ClientStreamMessage)
+	incoming := make(chan *cvms.ClientStreamMessage)
 	mockSvc := new(mockService)
 	server := NewServer(incoming, mockSvc).(*grpcServer)
 
 	mockStream := new(mockServerStream)
 
-	runReq := &manager.ComputationRunReq{
+	runReq := &cvms.ComputationRunReq{
 		Id: "test-id",
 	}
 
 	// Simulate an error when sending
-	mockStream.On("Send", mock.AnythingOfType("*manager.ServerStreamMessage")).Return(errors.New("send error")).Once()
+	mockStream.On("Send", mock.AnythingOfType("*cvms.ServerStreamMessage")).Return(errors.New("send error")).Once()
 
 	err := server.sendRunReqInChunks(mockStream, runReq)
 
@@ -272,7 +255,7 @@ func TestGrpcServer_sendRunReqInChunksError(t *testing.T) {
 }
 
 func TestGrpcServer_ProcessMissingPeerInfo(t *testing.T) {
-	incoming := make(chan *manager.ClientStreamMessage)
+	incoming := make(chan *cvms.ClientStreamMessage)
 	mockSvc := new(mockService)
 	server := NewServer(incoming, mockSvc).(*grpcServer)
 

agent/cvms/api/grpc/storage/mocks/storage.go

@@ -0,0 +1,234 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+package mocks
+
+import (
+	mock "github.com/stretchr/testify/mock"
+	cvms "github.com/ultravioletrs/cocos/agent/cvms"
+
+	storage "github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage"
+)
+
+// Storage is an autogenerated mock type for the Storage type
+type Storage struct {
+	mock.Mock
+}
+
+type Storage_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *Storage) EXPECT() *Storage_Expecter {
+	return &Storage_Expecter{mock: &_m.Mock}
+}
+
+// Add provides a mock function with given fields: msg
+func (_m *Storage) Add(msg *cvms.ClientStreamMessage) error {
+	ret := _m.Called(msg)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Add")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(*cvms.ClientStreamMessage) error); ok {
+		r0 = rf(msg)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Storage_Add_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Add'
+type Storage_Add_Call struct {
+	*mock.Call
+}
+
+// Add is a helper method to define mock.On call
+//   - msg *cvms.ClientStreamMessage
+func (_e *Storage_Expecter) Add(msg interface{}) *Storage_Add_Call {
+	return &Storage_Add_Call{Call: _e.mock.On("Add", msg)}
+}
+
+func (_c *Storage_Add_Call) Run(run func(msg *cvms.ClientStreamMessage)) *Storage_Add_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(*cvms.ClientStreamMessage))
+	})
+	return _c
+}
+
+func (_c *Storage_Add_Call) Return(_a0 error) *Storage_Add_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Storage_Add_Call) RunAndReturn(run func(*cvms.ClientStreamMessage) error) *Storage_Add_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Clear provides a mock function with given fields:
+func (_m *Storage) Clear() error {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Clear")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Storage_Clear_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Clear'
+type Storage_Clear_Call struct {
+	*mock.Call
+}
+
+// Clear is a helper method to define mock.On call
+func (_e *Storage_Expecter) Clear() *Storage_Clear_Call {
+	return &Storage_Clear_Call{Call: _e.mock.On("Clear")}
+}
+
+func (_c *Storage_Clear_Call) Run(run func()) *Storage_Clear_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *Storage_Clear_Call) Return(_a0 error) *Storage_Clear_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Storage_Clear_Call) RunAndReturn(run func() error) *Storage_Clear_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Load provides a mock function with given fields:
+func (_m *Storage) Load() ([]storage.Message, error) {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Load")
+	}
+
+	var r0 []storage.Message
+	var r1 error
+	if rf, ok := ret.Get(0).(func() ([]storage.Message, error)); ok {
+		return rf()
+	}
+	if rf, ok := ret.Get(0).(func() []storage.Message); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]storage.Message)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// Storage_Load_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Load'
+type Storage_Load_Call struct {
+	*mock.Call
+}
+
+// Load is a helper method to define mock.On call
+func (_e *Storage_Expecter) Load() *Storage_Load_Call {
+	return &Storage_Load_Call{Call: _e.mock.On("Load")}
+}
+
+func (_c *Storage_Load_Call) Run(run func()) *Storage_Load_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *Storage_Load_Call) Return(_a0 []storage.Message, _a1 error) *Storage_Load_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *Storage_Load_Call) RunAndReturn(run func() ([]storage.Message, error)) *Storage_Load_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Save provides a mock function with given fields: messages
+func (_m *Storage) Save(messages []storage.Message) error {
+	ret := _m.Called(messages)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Save")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func([]storage.Message) error); ok {
+		r0 = rf(messages)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Storage_Save_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Save'
+type Storage_Save_Call struct {
+	*mock.Call
+}
+
+// Save is a helper method to define mock.On call
+//   - messages []storage.Message
+func (_e *Storage_Expecter) Save(messages interface{}) *Storage_Save_Call {
+	return &Storage_Save_Call{Call: _e.mock.On("Save", messages)}
+}
+
+func (_c *Storage_Save_Call) Run(run func(messages []storage.Message)) *Storage_Save_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].([]storage.Message))
+	})
+	return _c
+}
+
+func (_c *Storage_Save_Call) Return(_a0 error) *Storage_Save_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Storage_Save_Call) RunAndReturn(run func([]storage.Message) error) *Storage_Save_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewStorage creates a new instance of Storage. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewStorage(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Storage {
+	mock := &Storage{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

agent/cvms/api/grpc/storage/storage.go

@@ -0,0 +1,111 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package storage
+
+import (
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"sync"
+	"time"
+
+	"github.com/ultravioletrs/cocos/agent/cvms"
+)
+
+// Message represents a pending message with its timestamp.
+type Message struct {
+	Message *cvms.ClientStreamMessage
+	Time    time.Time
+}
+
+// Storage defines the interface for message persistence operations.
+type Storage interface {
+	// Load retrieves all pending messages from storage.
+	Load() ([]Message, error)
+
+	// Save persists the given messages to storage.
+	Save(messages []Message) error
+
+	// Add appends a new message to storage.
+	Add(msg *cvms.ClientStreamMessage) error
+
+	// Clear removes all messages from storage.
+	Clear() error
+}
+
+// FileStorage implements Storage interface using file-based persistence.
+type FileStorage struct {
+	mu   sync.Mutex
+	path string
+	msgs []Message
+}
+
+// NewFileStorage creates a new file-based storage instance.
+func NewFileStorage(storageDir string) (*FileStorage, error) {
+	if err := os.MkdirAll(storageDir, 0o755); err != nil {
+		return nil, err
+	}
+
+	return &FileStorage{
+		path: filepath.Join(storageDir, "pending_messages.json"),
+		msgs: make([]Message, 0),
+	}, nil
+}
+
+func (fs *FileStorage) Load() ([]Message, error) {
+	fs.mu.Lock()
+	defer fs.mu.Unlock()
+
+	data, err := os.ReadFile(fs.path)
+	if os.IsNotExist(err) {
+		return nil, nil
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	if err := json.Unmarshal(data, &fs.msgs); err != nil {
+		return nil, err
+	}
+
+	return fs.msgs, nil
+}
+
+func (fs *FileStorage) Save(messages []Message) error {
+	fs.mu.Lock()
+	defer fs.mu.Unlock()
+
+	fs.msgs = messages
+
+	data, err := json.Marshal(messages)
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(fs.path, data, 0o644)
+}
+
+func (fs *FileStorage) Add(msg *cvms.ClientStreamMessage) error {
+	fs.mu.Lock()
+	defer fs.mu.Unlock()
+
+	fs.msgs = append(fs.msgs, Message{
+		Message: msg,
+		Time:    time.Now(),
+	})
+
+	data, err := json.Marshal(fs.msgs)
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(fs.path, data, 0o644)
+}
+
+func (fs *FileStorage) Clear() error {
+	fs.mu.Lock()
+	defer fs.mu.Unlock()
+
+	fs.msgs = make([]Message, 0)
+	return os.WriteFile(fs.path, []byte("[]"), 0o644)
+}

agent/cvms/cvms.proto

@@ -0,0 +1,119 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+syntax = "proto3";
+
+import "google/protobuf/timestamp.proto";
+
+package cvms;
+
+option go_package = "./cvms";
+
+service Service {
+  rpc Process(stream ClientStreamMessage) returns (stream ServerStreamMessage) {}
+}
+
+message AgentStateReq {
+  string id = 1;
+}
+
+message AgentStateRes {
+  string id = 1;
+  string state = 2;
+}
+
+message StopComputation {
+  string computation_id = 1;
+}
+
+message StopComputationResponse {
+  string computation_id = 1;
+  string message = 2;
+}
+
+message RunResponse{
+  string computation_id = 1;
+  string error = 2;
+}
+
+message AgentEvent {
+  string event_type = 1;
+	google.protobuf.Timestamp timestamp = 2;
+	string	computation_id = 3;
+	bytes	details = 4;
+	string	originator = 5;
+	string	status = 6;
+}
+
+message AgentLog {
+  string message = 1;
+  string computation_id = 2;
+  string level = 3;
+  google.protobuf.Timestamp timestamp = 4;
+}
+
+message ClientStreamMessage {
+  oneof message {
+    AgentLog agent_log = 1;
+    AgentEvent agent_event = 2;
+    RunResponse run_res = 3;
+    StopComputationResponse stopComputationRes = 4;
+    AgentStateRes agentStateRes = 5;
+    AttestationResponse vTPMattestationReport = 6;
+  }
+}
+
+message ServerStreamMessage {
+  oneof message {
+    RunReqChunks runReqChunks = 1;
+    ComputationRunReq runReq = 2;
+    StopComputation stopComputation = 3;
+    AgentStateReq agentStateReq = 4;
+  }
+}
+
+message RunReqChunks {
+  bytes data = 1;
+  string id = 2;
+  bool is_last = 3;
+}
+
+message ComputationRunReq {
+  string id = 1;
+  string name = 2;
+  string description = 3;
+  repeated Dataset datasets = 4;
+  Algorithm algorithm = 5;
+  repeated ResultConsumer result_consumers = 6;
+  AgentConfig agent_config = 7;
+}
+
+message ResultConsumer {
+  bytes userKey = 1;
+}
+
+message Dataset {
+  bytes hash = 1; // should be sha3.Sum256, 32 byte length.
+  bytes userKey = 2;
+  string filename = 3;
+}
+
+message Algorithm {
+  bytes hash = 1; // should be sha3.Sum256, 32 byte length.
+  bytes userKey = 2;
+}
+
+message AgentConfig {
+  string port = 1;
+  string cert_file = 2;
+  string key_file = 3;
+  string client_ca_file = 4;
+  string server_ca_file = 5;
+  string log_level = 6;
+  bool   attested_tls = 7;
+}
+
+message AttestationResponse {
+  bytes file = 1;
+  string certSerialNumber = 2;
+}

agent/cvms/cvms_grpc.pb.go

@@ -0,0 +1,118 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
+// versions:
+// - protoc-gen-go-grpc v1.5.1
+// - protoc             v5.29.0
+// source: agent/cvms/cvms.proto
+
+package cvms
+
+import (
+	context "context"
+	grpc "google.golang.org/grpc"
+	codes "google.golang.org/grpc/codes"
+	status "google.golang.org/grpc/status"
+)
+
+// This is a compile-time assertion to ensure that this generated file
+// is compatible with the grpc package it is being compiled against.
+// Requires gRPC-Go v1.64.0 or later.
+const _ = grpc.SupportPackageIsVersion9
+
+const (
+	Service_Process_FullMethodName = "/cvms.Service/Process"
+)
+
+// ServiceClient is the client API for Service service.
+//
+// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
+type ServiceClient interface {
+	Process(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage], error)
+}
+
+type serviceClient struct {
+	cc grpc.ClientConnInterface
+}
+
+func NewServiceClient(cc grpc.ClientConnInterface) ServiceClient {
+	return &serviceClient{cc}
+}
+
+func (c *serviceClient) Process(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage], error) {
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
+	stream, err := c.cc.NewStream(ctx, &Service_ServiceDesc.Streams[0], Service_Process_FullMethodName, cOpts...)
+	if err != nil {
+		return nil, err
+	}
+	x := &grpc.GenericClientStream[ClientStreamMessage, ServerStreamMessage]{ClientStream: stream}
+	return x, nil
+}
+
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type Service_ProcessClient = grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage]
+
+// ServiceServer is the server API for Service service.
+// All implementations must embed UnimplementedServiceServer
+// for forward compatibility.
+type ServiceServer interface {
+	Process(grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]) error
+	mustEmbedUnimplementedServiceServer()
+}
+
+// UnimplementedServiceServer must be embedded to have
+// forward compatible implementations.
+//
+// NOTE: this should be embedded by value instead of pointer to avoid a nil
+// pointer dereference when methods are called.
+type UnimplementedServiceServer struct{}
+
+func (UnimplementedServiceServer) Process(grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]) error {
+	return status.Errorf(codes.Unimplemented, "method Process not implemented")
+}
+func (UnimplementedServiceServer) mustEmbedUnimplementedServiceServer() {}
+func (UnimplementedServiceServer) testEmbeddedByValue()                 {}
+
+// UnsafeServiceServer may be embedded to opt out of forward compatibility for this service.
+// Use of this interface is not recommended, as added methods to ServiceServer will
+// result in compilation errors.
+type UnsafeServiceServer interface {
+	mustEmbedUnimplementedServiceServer()
+}
+
+func RegisterServiceServer(s grpc.ServiceRegistrar, srv ServiceServer) {
+	// If the following call pancis, it indicates UnimplementedServiceServer was
+	// embedded by pointer and is nil.  This will cause panics if an
+	// unimplemented method is ever invoked, so we test this at initialization
+	// time to prevent it from happening at runtime later due to I/O.
+	if t, ok := srv.(interface{ testEmbeddedByValue() }); ok {
+		t.testEmbeddedByValue()
+	}
+	s.RegisterService(&Service_ServiceDesc, srv)
+}
+
+func _Service_Process_Handler(srv interface{}, stream grpc.ServerStream) error {
+	return srv.(ServiceServer).Process(&grpc.GenericServerStream[ClientStreamMessage, ServerStreamMessage]{ServerStream: stream})
+}
+
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type Service_ProcessServer = grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]
+
+// Service_ServiceDesc is the grpc.ServiceDesc for Service service.
+// It's only intended for direct use with grpc.RegisterService,
+// and not to be introspected or modified (even as a copy)
+var Service_ServiceDesc = grpc.ServiceDesc{
+	ServiceName: "cvms.Service",
+	HandlerType: (*ServiceServer)(nil),
+	Methods:     []grpc.MethodDesc{},
+	Streams: []grpc.StreamDesc{
+		{
+			StreamName:    "Process",
+			Handler:       _Service_Process_Handler,
+			ServerStreams: true,
+			ClientStreams: true,
+		},
+	},
+	Metadata: "agent/cvms/cvms.proto",
+}

agent/cvms/server/cvm.go

@@ -0,0 +1,98 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package server
+
+import (
+	context "context"
+	"fmt"
+	"log/slog"
+
+	"github.com/ultravioletrs/cocos/agent"
+	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
+	"github.com/ultravioletrs/cocos/agent/auth"
+	"github.com/ultravioletrs/cocos/internal/server"
+	grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/reflection"
+)
+
+const (
+	svcName        = "agent"
+	defSvcGRPCPort = "7002"
+)
+
+type AgentServer interface {
+	Start(cfg agent.AgentConfig, cmp agent.Computation) error
+	Stop() error
+}
+
+type agentServer struct {
+	gs     server.Server
+	logger *slog.Logger
+	svc    agent.Service
+	host   string
+}
+
+func NewServer(logger *slog.Logger, svc agent.Service, host string) AgentServer {
+	return &agentServer{
+		logger: logger,
+		svc:    svc,
+		host:   host,
+	}
+}
+
+func (as *agentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error {
+	if cfg.Port == "" {
+		cfg.Port = defSvcGRPCPort
+	}
+
+	agentGrpcServerConfig := server.AgentConfig{
+		ServerConfig: server.ServerConfig{
+			BaseConfig: server.BaseConfig{
+				Host:         as.host,
+				Port:         cfg.Port,
+				CertFile:     cfg.CertFile,
+				KeyFile:      cfg.KeyFile,
+				ServerCAFile: cfg.ServerCAFile,
+				ClientCAFile: cfg.ClientCAFile,
+			},
+		},
+		AttestedTLS: cfg.AttestedTls,
+	}
+
+	registerAgentServiceServer := func(srv *grpc.Server) {
+		reflection.Register(srv)
+		agent.RegisterAgentServiceServer(srv, agentgrpc.NewServer(as.svc))
+	}
+
+	authSvc, err := auth.New(cmp)
+	if err != nil {
+		as.logger.WithGroup(cmp.ID).Error(fmt.Sprintf("failed to create auth service %s", err.Error()))
+		return err
+	}
+
+	qp, err := quoteprovider.GetLeveledQuoteProvider()
+	if err != nil {
+		as.logger.Error(fmt.Sprintf("failed to create quote provider %s", err.Error()))
+		return err
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	as.gs = grpcserver.New(ctx, cancel, svcName, agentGrpcServerConfig, registerAgentServiceServer, as.logger, qp, authSvc)
+
+	go func() {
+		err := as.gs.Start()
+		if err != nil {
+			as.logger.Error(fmt.Sprintf("failed to start grpc server %s", err.Error()))
+		}
+	}()
+
+	return nil
+}
+
+func (as *agentServer) Stop() error {
+	return as.gs.Stop()
+}

agent/cvms/server/mocks/server.go

@@ -0,0 +1,130 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+package mocks
+
+import (
+	mock "github.com/stretchr/testify/mock"
+	agent "github.com/ultravioletrs/cocos/agent"
+)
+
+// AgentServer is an autogenerated mock type for the AgentServer type
+type AgentServer struct {
+	mock.Mock
+}
+
+type AgentServer_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *AgentServer) EXPECT() *AgentServer_Expecter {
+	return &AgentServer_Expecter{mock: &_m.Mock}
+}
+
+// Start provides a mock function with given fields: cfg, cmp
+func (_m *AgentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error {
+	ret := _m.Called(cfg, cmp)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Start")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(agent.AgentConfig, agent.Computation) error); ok {
+		r0 = rf(cfg, cmp)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// AgentServer_Start_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Start'
+type AgentServer_Start_Call struct {
+	*mock.Call
+}
+
+// Start is a helper method to define mock.On call
+//   - cfg agent.AgentConfig
+//   - cmp agent.Computation
+func (_e *AgentServer_Expecter) Start(cfg interface{}, cmp interface{}) *AgentServer_Start_Call {
+	return &AgentServer_Start_Call{Call: _e.mock.On("Start", cfg, cmp)}
+}
+
+func (_c *AgentServer_Start_Call) Run(run func(cfg agent.AgentConfig, cmp agent.Computation)) *AgentServer_Start_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(agent.AgentConfig), args[1].(agent.Computation))
+	})
+	return _c
+}
+
+func (_c *AgentServer_Start_Call) Return(_a0 error) *AgentServer_Start_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *AgentServer_Start_Call) RunAndReturn(run func(agent.AgentConfig, agent.Computation) error) *AgentServer_Start_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Stop provides a mock function with given fields:
+func (_m *AgentServer) Stop() error {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Stop")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// AgentServer_Stop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Stop'
+type AgentServer_Stop_Call struct {
+	*mock.Call
+}
+
+// Stop is a helper method to define mock.On call
+func (_e *AgentServer_Expecter) Stop() *AgentServer_Stop_Call {
+	return &AgentServer_Stop_Call{Call: _e.mock.On("Stop")}
+}
+
+func (_c *AgentServer_Stop_Call) Run(run func()) *AgentServer_Stop_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentServer_Stop_Call) Return(_a0 error) *AgentServer_Stop_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *AgentServer_Stop_Call) RunAndReturn(run func() error) *AgentServer_Stop_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewAgentServer creates a new instance of AgentServer. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewAgentServer(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *AgentServer {
+	mock := &AgentServer{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

agent/events/events.go

@@ -4,43 +4,38 @@ package events
 
 import (
 	"encoding/json"
-	"io"
 
-	"google.golang.org/protobuf/proto"
+	"github.com/ultravioletrs/cocos/agent/cvms"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 type service struct {
-	service       string
-	computationID string
-	conn          io.Writer
+	service string
+	queue   chan *cvms.ClientStreamMessage
 }
 
 type Service interface {
-	SendEvent(event, status string, details json.RawMessage) error
+	SendEvent(cmpID, event, status string, details json.RawMessage)
 }
 
-func New(svc, computationID string, conn io.Writer) (Service, error) {
+func New(svc string, queue chan *cvms.ClientStreamMessage) (Service, error) {
 	return &service{
-		service:       svc,
-		computationID: computationID,
-		conn:          conn,
+		service: svc,
+		queue:   queue,
 	}, nil
 }
 
-func (s *service) SendEvent(event, status string, details json.RawMessage) error {
-	body := EventsLogs{Message: &EventsLogs_AgentEvent{AgentEvent: &AgentEvent{
-		EventType:     event,
-		Timestamp:     timestamppb.Now(),
-		ComputationId: s.computationID,
-		Originator:    s.service,
-		Status:        status,
-		Details:       details,
-	}}}
-	protoBody, err := proto.Marshal(&body)
-	if err != nil {
-		return err
+func (s *service) SendEvent(cmpID, event, status string, details json.RawMessage) {
+	s.queue <- &cvms.ClientStreamMessage{
+		Message: &cvms.ClientStreamMessage_AgentEvent{
+			AgentEvent: &cvms.AgentEvent{
+				EventType:     event,
+				Timestamp:     timestamppb.Now(),
+				ComputationId: cmpID,
+				Originator:    s.service,
+				Status:        status,
+				Details:       details,
+			},
+		},
 	}
-	_, err = s.conn.Write(protoBody)
-	return err
 }

agent/events/events.pb.go

@@ -3,8 +3,8 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.34.2
-// 	protoc        v5.28.1
+// 	protoc-gen-go v1.36.5
+// 	protoc        v5.29.0
 // source: agent/events/events.proto
 
 package events
@@ -15,6 +15,7 @@ import (
 	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
 	reflect "reflect"
 	sync "sync"
+	unsafe "unsafe"
 )
 
 const (
@@ -25,25 +26,22 @@ const (
 )
 
 type AgentEvent struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
+	state         protoimpl.MessageState `protogen:"open.v1"`
 	EventType     string                 `protobuf:"bytes,1,opt,name=event_type,json=eventType,proto3" json:"event_type,omitempty"`
 	Timestamp     *timestamppb.Timestamp `protobuf:"bytes,2,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
 	ComputationId string                 `protobuf:"bytes,3,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
 	Details       []byte                 `protobuf:"bytes,4,opt,name=details,proto3" json:"details,omitempty"`
 	Originator    string                 `protobuf:"bytes,5,opt,name=originator,proto3" json:"originator,omitempty"`
 	Status        string                 `protobuf:"bytes,6,opt,name=status,proto3" json:"status,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *AgentEvent) Reset() {
 	*x = AgentEvent{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_events_events_proto_msgTypes[0]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_events_events_proto_msgTypes[0]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *AgentEvent) String() string {
@@ -54,7 +52,7 @@ func (*AgentEvent) ProtoMessage() {}
 
 func (x *AgentEvent) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_events_events_proto_msgTypes[0]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -112,23 +110,20 @@ func (x *AgentEvent) GetStatus() string {
 }
 
 type AgentLog struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
+	state         protoimpl.MessageState `protogen:"open.v1"`
 	Message       string                 `protobuf:"bytes,1,opt,name=message,proto3" json:"message,omitempty"`
 	ComputationId string                 `protobuf:"bytes,2,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
 	Level         string                 `protobuf:"bytes,3,opt,name=level,proto3" json:"level,omitempty"`
 	Timestamp     *timestamppb.Timestamp `protobuf:"bytes,4,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *AgentLog) Reset() {
 	*x = AgentLog{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_events_events_proto_msgTypes[1]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_events_events_proto_msgTypes[1]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *AgentLog) String() string {
@@ -139,7 +134,7 @@ func (*AgentLog) ProtoMessage() {}
 
 func (x *AgentLog) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_events_events_proto_msgTypes[1]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -183,24 +178,21 @@ func (x *AgentLog) GetTimestamp() *timestamppb.Timestamp {
 }
 
 type EventsLogs struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
-
-	// Types that are assignable to Message:
+	state protoimpl.MessageState `protogen:"open.v1"`
+	// Types that are valid to be assigned to Message:
 	//
 	//	*EventsLogs_AgentLog
 	//	*EventsLogs_AgentEvent
-	Message isEventsLogs_Message `protobuf_oneof:"message"`
+	Message       isEventsLogs_Message `protobuf_oneof:"message"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *EventsLogs) Reset() {
 	*x = EventsLogs{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_events_events_proto_msgTypes[2]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_events_events_proto_msgTypes[2]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *EventsLogs) String() string {
@@ -211,7 +203,7 @@ func (*EventsLogs) ProtoMessage() {}
 
 func (x *EventsLogs) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_events_events_proto_msgTypes[2]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -226,23 +218,27 @@ func (*EventsLogs) Descriptor() ([]byte, []int) {
 	return file_agent_events_events_proto_rawDescGZIP(), []int{2}
 }
 
-func (m *EventsLogs) GetMessage() isEventsLogs_Message {
-	if m != nil {
-		return m.Message
+func (x *EventsLogs) GetMessage() isEventsLogs_Message {
+	if x != nil {
+		return x.Message
 	}
 	return nil
 }
 
 func (x *EventsLogs) GetAgentLog() *AgentLog {
-	if x, ok := x.GetMessage().(*EventsLogs_AgentLog); ok {
-		return x.AgentLog
+	if x != nil {
+		if x, ok := x.Message.(*EventsLogs_AgentLog); ok {
+			return x.AgentLog
+		}
 	}
 	return nil
 }
 
 func (x *EventsLogs) GetAgentEvent() *AgentEvent {
-	if x, ok := x.GetMessage().(*EventsLogs_AgentEvent); ok {
-		return x.AgentEvent
+	if x != nil {
+		if x, ok := x.Message.(*EventsLogs_AgentEvent); ok {
+			return x.AgentEvent
+		}
 	}
 	return nil
 }
@@ -265,7 +261,7 @@ func (*EventsLogs_AgentEvent) isEventsLogs_Message() {}
 
 var File_agent_events_events_proto protoreflect.FileDescriptor
 
-var file_agent_events_events_proto_rawDesc = []byte{
+var file_agent_events_events_proto_rawDesc = string([]byte{
 	0x0a, 0x19, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x2f, 0x65,
 	0x76, 0x65, 0x6e, 0x74, 0x73, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x06, 0x65, 0x76, 0x65,
 	0x6e, 0x74, 0x73, 0x1a, 0x1f, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74,
@@ -304,16 +300,16 @@ var file_agent_events_events_proto_rawDesc = []byte{
 	0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x42, 0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73,
 	0x73, 0x61, 0x67, 0x65, 0x42, 0x0a, 0x5a, 0x08, 0x2e, 0x2f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73,
 	0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
-}
+})
 
 var (
 	file_agent_events_events_proto_rawDescOnce sync.Once
-	file_agent_events_events_proto_rawDescData = file_agent_events_events_proto_rawDesc
+	file_agent_events_events_proto_rawDescData []byte
 )
 
 func file_agent_events_events_proto_rawDescGZIP() []byte {
 	file_agent_events_events_proto_rawDescOnce.Do(func() {
-		file_agent_events_events_proto_rawDescData = protoimpl.X.CompressGZIP(file_agent_events_events_proto_rawDescData)
+		file_agent_events_events_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_agent_events_events_proto_rawDesc), len(file_agent_events_events_proto_rawDesc)))
 	})
 	return file_agent_events_events_proto_rawDescData
 }
@@ -342,44 +338,6 @@ func file_agent_events_events_proto_init() {
 	if File_agent_events_events_proto != nil {
 		return
 	}
-	if !protoimpl.UnsafeEnabled {
-		file_agent_events_events_proto_msgTypes[0].Exporter = func(v any, i int) any {
-			switch v := v.(*AgentEvent); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_events_events_proto_msgTypes[1].Exporter = func(v any, i int) any {
-			switch v := v.(*AgentLog); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_events_events_proto_msgTypes[2].Exporter = func(v any, i int) any {
-			switch v := v.(*EventsLogs); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-	}
 	file_agent_events_events_proto_msgTypes[2].OneofWrappers = []any{
 		(*EventsLogs_AgentLog)(nil),
 		(*EventsLogs_AgentEvent)(nil),
@@ -388,7 +346,7 @@ func file_agent_events_events_proto_init() {
 	out := protoimpl.TypeBuilder{
 		File: protoimpl.DescBuilder{
 			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
-			RawDescriptor: file_agent_events_events_proto_rawDesc,
+			RawDescriptor: unsafe.Slice(unsafe.StringData(file_agent_events_events_proto_rawDesc), len(file_agent_events_events_proto_rawDesc)),
 			NumEnums:      0,
 			NumMessages:   3,
 			NumExtensions: 0,
@@ -399,7 +357,6 @@ func file_agent_events_events_proto_init() {
 		MessageInfos:      file_agent_events_events_proto_msgTypes,
 	}.Build()
 	File_agent_events_events_proto = out.File
-	file_agent_events_events_proto_rawDesc = nil
 	file_agent_events_events_proto_goTypes = nil
 	file_agent_events_events_proto_depIdxs = nil
 }

agent/events/events_test.go

@@ -3,62 +3,36 @@
 package events
 
 import (
-	"bytes"
 	"encoding/json"
-	"errors"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
-	"google.golang.org/protobuf/proto"
+	"github.com/ultravioletrs/cocos/agent/cvms"
 )
 
-type mockConn struct {
-	writeErr error
-	buf      bytes.Buffer
-}
-
-func (m *mockConn) Write(p []byte) (n int, err error) {
-	if m.writeErr != nil {
-		return 0, m.writeErr
-	}
-	return m.buf.Write(p)
-}
-
 func TestSendEventSuccess(t *testing.T) {
-	mockConnection := &mockConn{}
-
-	svc, err := New("test_service", "12345", mockConnection)
+	queue := make(chan *cvms.ClientStreamMessage, 1)
+	svc, err := New("test_service", queue)
 	assert.NoError(t, err)
 
 	details := json.RawMessage(`{"key": "value"}`)
 
-	err = svc.SendEvent("test_event", "success", details)
-	assert.NoError(t, err)
+	go func() {
+		msg := <-queue
+		assert.NotNil(t, msg)
+		assert.NotNil(t, msg.GetAgentEvent())
+		assert.Equal(t, "test_event", msg.GetAgentEvent().EventType)
+		assert.Equal(t, "testid", msg.GetAgentEvent().ComputationId)
+		assert.Equal(t, "test_service", msg.GetAgentEvent().Originator)
+		assert.Equal(t, "success", msg.GetAgentEvent().Status)
 
-	var writtenMessage EventsLogs
-	err = proto.Unmarshal(mockConnection.buf.Bytes(), &writtenMessage)
-	assert.NoError(t, err)
+		now := time.Now()
+		eventTimestamp := msg.GetAgentEvent().GetTimestamp().AsTime()
+		assert.WithinDuration(t, now, eventTimestamp, 1*time.Second)
+	}()
 
-	assert.Equal(t, "test_event", writtenMessage.GetAgentEvent().EventType)
-	assert.Equal(t, "12345", writtenMessage.GetAgentEvent().ComputationId)
-	assert.Equal(t, "test_service", writtenMessage.GetAgentEvent().Originator)
-	assert.Equal(t, "success", writtenMessage.GetAgentEvent().Status)
+	svc.SendEvent("testid", "test_event", "success", details)
 
-	now := time.Now()
-	eventTimestamp := writtenMessage.GetAgentEvent().GetTimestamp().AsTime()
-	assert.WithinDuration(t, now, eventTimestamp, 1*time.Second)
-}
-
-func TestSendEventFailure(t *testing.T) {
-	mockConnection := &mockConn{writeErr: errors.New("write error")}
-
-	svc, err := New("test_service", "12345", mockConnection)
-	assert.NoError(t, err)
-
-	details := json.RawMessage(`{"key": "value"}`)
-
-	err = svc.SendEvent("test_event", "failure", details)
-	assert.Error(t, err)
-	assert.Equal(t, "write error", err.Error())
+	time.Sleep(1 * time.Second)
 }

agent/events/mocks/events.go

@@ -24,22 +24,9 @@ func (_m *Service) EXPECT() *Service_Expecter {
 	return &Service_Expecter{mock: &_m.Mock}
 }
 
-// SendEvent provides a mock function with given fields: event, status, details
-func (_m *Service) SendEvent(event string, status string, details json.RawMessage) error {
-	ret := _m.Called(event, status, details)
-
-	if len(ret) == 0 {
-		panic("no return value specified for SendEvent")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(string, string, json.RawMessage) error); ok {
-		r0 = rf(event, status, details)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
+// SendEvent provides a mock function with given fields: cmpID, event, status, details
+func (_m *Service) SendEvent(cmpID string, event string, status string, details json.RawMessage) {
+	_m.Called(cmpID, event, status, details)
 }
 
 // Service_SendEvent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendEvent'
@@ -48,26 +35,27 @@ type Service_SendEvent_Call struct {
 }
 
 // SendEvent is a helper method to define mock.On call
+//   - cmpID string
 //   - event string
 //   - status string
 //   - details json.RawMessage
-func (_e *Service_Expecter) SendEvent(event interface{}, status interface{}, details interface{}) *Service_SendEvent_Call {
-	return &Service_SendEvent_Call{Call: _e.mock.On("SendEvent", event, status, details)}
+func (_e *Service_Expecter) SendEvent(cmpID interface{}, event interface{}, status interface{}, details interface{}) *Service_SendEvent_Call {
+	return &Service_SendEvent_Call{Call: _e.mock.On("SendEvent", cmpID, event, status, details)}
 }
 
-func (_c *Service_SendEvent_Call) Run(run func(event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
+func (_c *Service_SendEvent_Call) Run(run func(cmpID string, event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string), args[1].(string), args[2].(json.RawMessage))
+		run(args[0].(string), args[1].(string), args[2].(string), args[3].(json.RawMessage))
 	})
 	return _c
 }
 
-func (_c *Service_SendEvent_Call) Return(_a0 error) *Service_SendEvent_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_SendEvent_Call) Return() *Service_SendEvent_Call {
+	_c.Call.Return()
 	return _c
 }
 
-func (_c *Service_SendEvent_Call) RunAndReturn(run func(string, string, json.RawMessage) error) *Service_SendEvent_Call {
+func (_c *Service_SendEvent_Call) RunAndReturn(run func(string, string, string, json.RawMessage)) *Service_SendEvent_Call {
 	_c.Call.Return(run)
 	return _c
 }

agent/mocks/agent.go

@@ -6,9 +6,10 @@
 package mocks
 
 import (
-	context "context"
-
 	agent "github.com/ultravioletrs/cocos/agent"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
+
+	context "context"
 
 	mock "github.com/stretchr/testify/mock"
 )
@@ -73,9 +74,9 @@ func (_c *Service_Algo_Call) RunAndReturn(run func(context.Context, agent.Algori
 	return _c
 }
 
-// Attestation provides a mock function with given fields: ctx, reportData
-func (_m *Service) Attestation(ctx context.Context, reportData [64]byte) ([]byte, error) {
-	ret := _m.Called(ctx, reportData)
+// Attestation provides a mock function with given fields: ctx, reportData, nonce, attType
+func (_m *Service) Attestation(ctx context.Context, reportData [64]byte, nonce [32]byte, attType config.AttestationType) ([]byte, error) {
+	ret := _m.Called(ctx, reportData, nonce, attType)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Attestation")
@@ -83,19 +84,19 @@ func (_m *Service) Attestation(ctx context.Context, reportData [64]byte) ([]byte
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, [64]byte) ([]byte, error)); ok {
-		return rf(ctx, reportData)
+	if rf, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, config.AttestationType) ([]byte, error)); ok {
+		return rf(ctx, reportData, nonce, attType)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, [64]byte) []byte); ok {
-		r0 = rf(ctx, reportData)
+	if rf, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, config.AttestationType) []byte); ok {
+		r0 = rf(ctx, reportData, nonce, attType)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
 
-	if rf, ok := ret.Get(1).(func(context.Context, [64]byte) error); ok {
-		r1 = rf(ctx, reportData)
+	if rf, ok := ret.Get(1).(func(context.Context, [64]byte, [32]byte, config.AttestationType) error); ok {
+		r1 = rf(ctx, reportData, nonce, attType)
 	} else {
 		r1 = ret.Error(1)
 	}
@@ -111,13 +112,15 @@ type Service_Attestation_Call struct {
 // Attestation is a helper method to define mock.On call
 //   - ctx context.Context
 //   - reportData [64]byte
-func (_e *Service_Expecter) Attestation(ctx interface{}, reportData interface{}) *Service_Attestation_Call {
-	return &Service_Attestation_Call{Call: _e.mock.On("Attestation", ctx, reportData)}
+//   - nonce [32]byte
+//   - attType config.AttestationType
+func (_e *Service_Expecter) Attestation(ctx interface{}, reportData interface{}, nonce interface{}, attType interface{}) *Service_Attestation_Call {
+	return &Service_Attestation_Call{Call: _e.mock.On("Attestation", ctx, reportData, nonce, attType)}
 }
 
-func (_c *Service_Attestation_Call) Run(run func(ctx context.Context, reportData [64]byte)) *Service_Attestation_Call {
+func (_c *Service_Attestation_Call) Run(run func(ctx context.Context, reportData [64]byte, nonce [32]byte, attType config.AttestationType)) *Service_Attestation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].([64]byte))
+		run(args[0].(context.Context), args[1].([64]byte), args[2].([32]byte), args[3].(config.AttestationType))
 	})
 	return _c
 }
@@ -127,7 +130,7 @@ func (_c *Service_Attestation_Call) Return(_a0 []byte, _a1 error) *Service_Attes
 	return _c
 }
 
-func (_c *Service_Attestation_Call) RunAndReturn(run func(context.Context, [64]byte) ([]byte, error)) *Service_Attestation_Call {
+func (_c *Service_Attestation_Call) RunAndReturn(run func(context.Context, [64]byte, [32]byte, config.AttestationType) ([]byte, error)) *Service_Attestation_Call {
 	_c.Call.Return(run)
 	return _c
 }
@@ -179,6 +182,53 @@ func (_c *Service_Data_Call) RunAndReturn(run func(context.Context, agent.Datase
 	return _c
 }
 
+// InitComputation provides a mock function with given fields: ctx, cmp
+func (_m *Service) InitComputation(ctx context.Context, cmp agent.Computation) error {
+	ret := _m.Called(ctx, cmp)
+
+	if len(ret) == 0 {
+		panic("no return value specified for InitComputation")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, agent.Computation) error); ok {
+		r0 = rf(ctx, cmp)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Service_InitComputation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'InitComputation'
+type Service_InitComputation_Call struct {
+	*mock.Call
+}
+
+// InitComputation is a helper method to define mock.On call
+//   - ctx context.Context
+//   - cmp agent.Computation
+func (_e *Service_Expecter) InitComputation(ctx interface{}, cmp interface{}) *Service_InitComputation_Call {
+	return &Service_InitComputation_Call{Call: _e.mock.On("InitComputation", ctx, cmp)}
+}
+
+func (_c *Service_InitComputation_Call) Run(run func(ctx context.Context, cmp agent.Computation)) *Service_InitComputation_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(agent.Computation))
+	})
+	return _c
+}
+
+func (_c *Service_InitComputation_Call) Return(_a0 error) *Service_InitComputation_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Service_InitComputation_Call) RunAndReturn(run func(context.Context, agent.Computation) error) *Service_InitComputation_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Result provides a mock function with given fields: ctx
 func (_m *Service) Result(ctx context.Context) ([]byte, error) {
 	ret := _m.Called(ctx)
@@ -237,6 +287,97 @@ func (_c *Service_Result_Call) RunAndReturn(run func(context.Context) ([]byte, e
 	return _c
 }
 
+// State provides a mock function with given fields:
+func (_m *Service) State() string {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for State")
+	}
+
+	var r0 string
+	if rf, ok := ret.Get(0).(func() string); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Get(0).(string)
+	}
+
+	return r0
+}
+
+// Service_State_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'State'
+type Service_State_Call struct {
+	*mock.Call
+}
+
+// State is a helper method to define mock.On call
+func (_e *Service_Expecter) State() *Service_State_Call {
+	return &Service_State_Call{Call: _e.mock.On("State")}
+}
+
+func (_c *Service_State_Call) Run(run func()) *Service_State_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *Service_State_Call) Return(_a0 string) *Service_State_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Service_State_Call) RunAndReturn(run func() string) *Service_State_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// StopComputation provides a mock function with given fields: ctx
+func (_m *Service) StopComputation(ctx context.Context) error {
+	ret := _m.Called(ctx)
+
+	if len(ret) == 0 {
+		panic("no return value specified for StopComputation")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context) error); ok {
+		r0 = rf(ctx)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Service_StopComputation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'StopComputation'
+type Service_StopComputation_Call struct {
+	*mock.Call
+}
+
+// StopComputation is a helper method to define mock.On call
+//   - ctx context.Context
+func (_e *Service_Expecter) StopComputation(ctx interface{}) *Service_StopComputation_Call {
+	return &Service_StopComputation_Call{Call: _e.mock.On("StopComputation", ctx)}
+}
+
+func (_c *Service_StopComputation_Call) Run(run func(ctx context.Context)) *Service_StopComputation_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context))
+	})
+	return _c
+}
+
+func (_c *Service_StopComputation_Call) Return(_a0 error) *Service_StopComputation_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Service_StopComputation_Call) RunAndReturn(run func(context.Context) error) *Service_StopComputation_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
 // The first argument is typically a *testing.T value.
 func NewService(t interface {

agent/service.go

@@ -23,6 +23,9 @@ import (
 	"github.com/ultravioletrs/cocos/agent/events"
 	"github.com/ultravioletrs/cocos/agent/statemachine"
 	"github.com/ultravioletrs/cocos/internal"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	"golang.org/x/crypto/sha3"
 )
 
@@ -69,8 +72,6 @@ const (
 )
 
 const (
-	// ReportDataSize is the size of the report data expected by the attestation service.
-	ReportDataSize     = 64
 	algoFilePermission = 0o700
 )
 
@@ -99,41 +100,50 @@ var (
 	ErrAllResultsConsumed = errors.New("all results have been consumed by declared consumers")
 	// ErrAttestationFailed attestation failed.
 	ErrAttestationFailed = errors.New("failed to get raw quote")
+	// ErrAttType indicates that the attestation type that is requested does not exist or is not supported.
+	ErrAttestationType = errors.New("attestation type does not exist or is not supported")
 )
 
 // Service specifies an API that must be fullfiled by the domain service
 // implementation, and all of its decorators (e.g. logging & metrics).
 type Service interface {
+	InitComputation(ctx context.Context, cmp Computation) error
+	StopComputation(ctx context.Context) error
 	Algo(ctx context.Context, algorithm Algorithm) error
 	Data(ctx context.Context, dataset Dataset) error
 	Result(ctx context.Context) ([]byte, error)
-	Attestation(ctx context.Context, reportData [ReportDataSize]byte) ([]byte, error)
+	Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType config.AttestationType) ([]byte, error)
+	State() string
 }
 
 type agentService struct {
 	mu              sync.Mutex
-	computation     Computation               // Holds the current computation request details.
-	algorithm       algorithm.Algorithm       // Filepath to the algorithm received for the computation.
-	result          []byte                    // Stores the result of the computation.
-	sm              statemachine.StateMachine // Manages the state transitions of the agent service.
-	runError        error                     // Stores any error encountered during the computation run.
-	eventSvc        events.Service            // Service for publishing events related to computation.
-	quoteProvider   client.QuoteProvider      // Provider for generating attestation quotes.
-	logger          *slog.Logger              // Logger for the agent service.
-	resultsConsumed bool                      // Indicates if the results have been consumed.
+	computation     Computation                 // Holds the current computation request details.
+	algorithm       algorithm.Algorithm         // Filepath to the algorithm received for the computation.
+	result          []byte                      // Stores the result of the computation.
+	sm              statemachine.StateMachine   // Manages the state transitions of the agent service.
+	runError        error                       // Stores any error encountered during the computation run.
+	eventSvc        events.Service              // Service for publishing events related to computation.
+	quoteProvider   client.LeveledQuoteProvider // Provider for generating attestation quotes.
+	logger          *slog.Logger                // Logger for the agent service.
+	resultsConsumed bool                        // Indicates if the results have been consumed.
+	cancel          context.CancelFunc          // Cancels the computation context.
+	vmpl            int                         // VMPL at which the Agent is running.
 }
 
 var _ Service = (*agentService)(nil)
 
 // New instantiates the agent service implementation.
-func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp Computation, quoteProvider client.QuoteProvider) Service {
+func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, quoteProvider client.LeveledQuoteProvider, vmlp int) Service {
 	sm := statemachine.NewStateMachine(Idle)
+	ctx, cancel := context.WithCancel(ctx)
 	svc := &agentService{
 		sm:            sm,
 		eventSvc:      eventSvc,
 		quoteProvider: quoteProvider,
 		logger:        logger,
-		computation:   cmp,
+		cancel:        cancel,
+		vmpl:          vmlp,
 	}
 
 	transitions := []statemachine.Transition{
@@ -141,13 +151,6 @@ func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp
 		{From: ReceivingManifest, Event: ManifestReceived, To: ReceivingAlgorithm},
 	}
 
-	if len(cmp.Datasets) == 0 {
-		transitions = append(transitions, statemachine.Transition{From: ReceivingAlgorithm, Event: AlgorithmReceived, To: Running})
-	} else {
-		transitions = append(transitions, statemachine.Transition{From: ReceivingAlgorithm, Event: AlgorithmReceived, To: ReceivingData})
-		transitions = append(transitions, statemachine.Transition{From: ReceivingData, Event: DataReceived, To: Running})
-	}
-
 	transitions = append(transitions, []statemachine.Transition{
 		{From: Running, Event: RunComplete, To: ConsumingResults},
 		{From: Running, Event: RunFailed, To: Failed},
@@ -158,8 +161,6 @@ func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp
 		sm.AddTransition(t)
 	}
 
-	sm.SetAction(Idle, svc.publishEvent(IdleState.String()))
-	sm.SetAction(ReceivingManifest, svc.publishEvent(InProgress.String()))
 	sm.SetAction(ReceivingAlgorithm, svc.publishEvent(InProgress.String()))
 	sm.SetAction(ReceivingData, svc.publishEvent(InProgress.String()))
 	sm.SetAction(Running, svc.runComputation)
@@ -173,11 +174,71 @@ func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp
 		}
 	}()
 	sm.SendEvent(Start)
-	defer sm.SendEvent(ManifestReceived)
 
 	return svc
 }
 
+func (as *agentService) State() string {
+	return as.sm.GetState().String()
+}
+
+func (as *agentService) InitComputation(ctx context.Context, cmp Computation) error {
+	if as.sm.GetState() != ReceivingManifest {
+		return ErrStateNotReady
+	}
+	defer as.sm.SendEvent(ManifestReceived)
+
+	as.mu.Lock()
+	defer as.mu.Unlock()
+
+	as.computation = cmp
+
+	transitions := []statemachine.Transition{}
+
+	if len(cmp.Datasets) == 0 {
+		transitions = append(transitions, statemachine.Transition{From: ReceivingAlgorithm, Event: AlgorithmReceived, To: Running})
+	} else {
+		transitions = append(transitions, statemachine.Transition{From: ReceivingAlgorithm, Event: AlgorithmReceived, To: ReceivingData})
+		transitions = append(transitions, statemachine.Transition{From: ReceivingData, Event: DataReceived, To: Running})
+	}
+
+	for _, t := range transitions {
+		as.sm.AddTransition(t)
+	}
+
+	return nil
+}
+
+func (as *agentService) StopComputation(ctx context.Context) error {
+	as.mu.Lock()
+	defer as.mu.Unlock()
+
+	as.cancel()
+
+	if err := as.algorithm.Stop(); err != nil {
+		return fmt.Errorf("error stopping computation: %v", err)
+	}
+
+	sm := statemachine.NewStateMachine(Idle)
+
+	if err := os.RemoveAll(algorithm.DatasetsDir); err != nil {
+		return fmt.Errorf("error removing datasets directory: %v", err)
+	}
+
+	if err := os.RemoveAll(algorithm.ResultsDir); err != nil {
+		return fmt.Errorf("error removing results directory: %v", err)
+	}
+
+	as.sm = sm
+	as.computation = Computation{}
+	as.algorithm = nil
+	as.result = nil
+	as.runError = nil
+	as.resultsConsumed = false
+
+	return nil
+}
+
 func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 	if as.sm.GetState() != ReceivingAlgorithm {
 		return ErrStateNotReady
@@ -225,7 +286,7 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 
 	switch algoType {
 	case string(algorithm.AlgoTypeBin):
-		as.algorithm = binary.NewAlgorithm(as.logger, as.eventSvc, f.Name(), args)
+		as.algorithm = binary.NewAlgorithm(as.logger, as.eventSvc, f.Name(), args, as.computation.ID)
 	case string(algorithm.AlgoTypePython):
 		var requirementsFile string
 		if len(algo.Requirements) > 0 {
@@ -243,11 +304,11 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 			requirementsFile = fr.Name()
 		}
 		runtime := python.PythonRunTimeFromContext(ctx)
-		as.algorithm = python.NewAlgorithm(as.logger, as.eventSvc, runtime, requirementsFile, f.Name(), args)
+		as.algorithm = python.NewAlgorithm(as.logger, as.eventSvc, runtime, requirementsFile, f.Name(), args, as.computation.ID)
 	case string(algorithm.AlgoTypeWasm):
-		as.algorithm = wasm.NewAlgorithm(as.logger, as.eventSvc, f.Name(), args)
+		as.algorithm = wasm.NewAlgorithm(as.logger, as.eventSvc, args, f.Name(), as.computation.ID)
 	case string(algorithm.AlgoTypeDocker):
-		as.algorithm = docker.NewAlgorithm(as.logger, as.eventSvc, f.Name())
+		as.algorithm = docker.NewAlgorithm(as.logger, as.eventSvc, f.Name(), as.computation.ID)
 	}
 
 	if err := os.Mkdir(algorithm.DatasetsDir, 0o755); err != nil {
@@ -341,13 +402,29 @@ func (as *agentService) Result(ctx context.Context) ([]byte, error) {
 	return as.result, as.runError
 }
 
-func (as *agentService) Attestation(ctx context.Context, reportData [ReportDataSize]byte) ([]byte, error) {
-	rawQuote, err := as.quoteProvider.GetRawQuote(reportData)
-	if err != nil {
-		return []byte{}, err
+func (as *agentService) Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType config.AttestationType) ([]byte, error) {
+	switch attType {
+	case config.SNP:
+		rawQuote, err := as.quoteProvider.GetRawQuoteAtLevel(reportData, uint(as.vmpl))
+		if err != nil {
+			return []byte{}, err
+		}
+		return rawQuote, nil
+	case config.VTPM:
+		vTPMQuote, err := vtpm.Attest(reportData[:], nonce[:], false)
+		if err != nil {
+			return []byte{}, err
+		}
+		return vTPMQuote, nil
+	case config.SNPvTPM:
+		vTPMQuote, err := vtpm.Attest(reportData[:], nonce[:], true)
+		if err != nil {
+			return []byte{}, err
+		}
+		return vTPMQuote, nil
+	default:
+		return []byte{}, ErrAttestationType
 	}
-
-	return rawQuote, nil
 }
 
 func (as *agentService) runComputation(state statemachine.State) {
@@ -400,8 +477,6 @@ func (as *agentService) runComputation(state statemachine.State) {
 
 func (as *agentService) publishEvent(status string) statemachine.Action {
 	return func(state statemachine.State) {
-		if err := as.eventSvc.SendEvent(state.String(), status, json.RawMessage{}); err != nil {
-			as.logger.Warn(err.Error())
-		}
+		as.eventSvc.SendEvent(as.computation.ID, state.String(), status, json.RawMessage{})
 	}
 }

agent/service_test.go

@@ -22,6 +22,7 @@ import (
 	smmocks "github.com/ultravioletrs/cocos/agent/statemachine/mocks"
 	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
 	mocks2 "github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider/mocks"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	"golang.org/x/crypto/sha3"
 	"google.golang.org/grpc/metadata"
 )
@@ -35,12 +36,7 @@ var (
 const datasetFile = "iris.csv"
 
 func TestAlgo(t *testing.T) {
-	events := new(mocks.Service)
-
-	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	defer evCall.Unset()
-
-	qp, err := quoteprovider.GetQuoteProvider()
+	qp, err := quoteprovider.GetLeveledQuoteProvider()
 	require.NoError(t, err)
 
 	algo, err := os.ReadFile(algoPath)
@@ -120,9 +116,15 @@ func TestAlgo(t *testing.T) {
 				metadata.Pairs(algorithm.AlgoTypeKey, tc.algoType, python.PyRuntimeKey, python.PyRuntime),
 			)
 
+			events := new(mocks.Service)
+			events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
 			ctx, cancel := context.WithCancel(ctx)
 			defer cancel()
-			svc := New(ctx, mglog.NewMock(), events, testComputation(t), qp)
+			svc := New(ctx, mglog.NewMock(), events, qp, 0)
+
+			err := svc.InitComputation(ctx, testComputation(t))
+			require.NoError(t, err)
 
 			time.Sleep(300 * time.Millisecond)
 
@@ -138,12 +140,7 @@ func TestAlgo(t *testing.T) {
 }
 
 func TestData(t *testing.T) {
-	events := new(mocks.Service)
-
-	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	defer evCall.Unset()
-
-	qp, err := quoteprovider.GetQuoteProvider()
+	qp, err := quoteprovider.GetLeveledQuoteProvider()
 	require.NoError(t, err)
 
 	algo, err := os.ReadFile(algoPath)
@@ -209,6 +206,9 @@ func TestData(t *testing.T) {
 					python.PyRuntimeKey, python.PyRuntime),
 			)
 
+			events := new(mocks.Service)
+			events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
 			if tc.err != ErrUndeclaredDataset {
 				ctx = IndexToContext(ctx, 0)
 			}
@@ -216,13 +216,16 @@ func TestData(t *testing.T) {
 			ctx, cancel := context.WithCancel(ctx)
 			defer cancel()
 
-			comp := testComputation(t)
+			svc := New(ctx, mglog.NewMock(), events, qp, 0)
+
+			err := svc.InitComputation(ctx, testComputation(t))
+			require.NoError(t, err)
 
-			svc := New(ctx, mglog.NewMock(), events, comp, qp)
 			time.Sleep(300 * time.Millisecond)
 
 			if tc.err != ErrStateNotReady {
-				_ = svc.Algo(ctx, alg)
+				err = svc.Algo(ctx, alg)
+				require.NoError(t, err)
 				time.Sleep(300 * time.Millisecond)
 			}
 			err = svc.Data(ctx, tc.data)
@@ -238,12 +241,7 @@ func TestData(t *testing.T) {
 }
 
 func TestResult(t *testing.T) {
-	events := new(mocks.Service)
-
-	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	defer evCall.Unset()
-
-	qp, err := quoteprovider.GetQuoteProvider()
+	qp, err := quoteprovider.GetLeveledQuoteProvider()
 	require.NoError(t, err)
 
 	cases := []struct {
@@ -285,6 +283,9 @@ func TestResult(t *testing.T) {
 	}
 
 	for _, tc := range cases {
+		events := new(mocks.Service)
+		events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
 		t.Run(tc.name, func(t *testing.T) {
 			ctx := metadata.NewIncomingContext(context.Background(),
 				metadata.Pairs(algorithm.AlgoTypeKey, "python", python.PyRuntimeKey, python.PyRuntime),
@@ -323,55 +324,57 @@ func TestResult(t *testing.T) {
 }
 
 func TestAttestation(t *testing.T) {
-	events := new(mocks.Service)
-	qp := new(mocks2.QuoteProvider)
-
-	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	defer evCall.Unset()
+	qp := new(mocks2.LeveledQuoteProvider)
 
 	cases := []struct {
 		name       string
-		reportData [ReportDataSize]byte
+		reportData [quoteprovider.Nonce]byte
+		nonce      [vtpm.Nonce]byte
 		rawQuote   []uint8
 		err        error
 	}{
 		{
 			name:       "Test attestation successful",
 			reportData: generateReportData(),
+			nonce:      [32]byte{},
 			rawQuote:   make([]uint8, 0),
 			err:        nil,
 		},
 		{
 			name:       "Test attestation failed",
 			reportData: generateReportData(),
+			nonce:      [32]byte{},
 			rawQuote:   nil,
 			err:        ErrAttestationFailed,
 		},
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
+			events := new(mocks.Service)
+			events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
 			ctx := metadata.NewIncomingContext(context.Background(),
 				metadata.Pairs(algorithm.AlgoTypeKey, "python", python.PyRuntimeKey, python.PyRuntime),
 			)
 			ctx, cancel := context.WithCancel(ctx)
 			defer cancel()
 
-			getQuote := qp.On("GetRawQuote", mock.Anything).Return(tc.rawQuote, tc.err)
+			getQuote := qp.On("GetRawQuoteAtLevel", mock.Anything, mock.Anything).Return(tc.rawQuote, tc.err)
 			if tc.err != ErrAttestationFailed {
-				getQuote = qp.On("GetRawQuote", mock.Anything).Return(tc.reportData, nil)
+				getQuote = qp.On("GetRawQuoteAtLevel", mock.Anything, mock.Anything).Return(tc.nonce, nil)
 			}
 			defer getQuote.Unset()
 
-			svc := New(ctx, mglog.NewMock(), events, testComputation(t), qp)
+			svc := New(ctx, mglog.NewMock(), events, qp, 0)
 			time.Sleep(300 * time.Millisecond)
-			_, err := svc.Attestation(ctx, tc.reportData)
+			_, err := svc.Attestation(ctx, tc.reportData, tc.nonce, 0)
 			assert.True(t, errors.Contains(err, tc.err), "expected %v, got %v", tc.err, err)
 		})
 	}
 }
 
-func generateReportData() [ReportDataSize]byte {
-	bytes := make([]byte, ReportDataSize)
+func generateReportData() [quoteprovider.Nonce]byte {
+	bytes := make([]byte, quoteprovider.Nonce)
 	_, err := rand.Read(bytes)
 	if err != nil {
 		log.Fatalf("Failed to generate random bytes: %v", err)
@@ -397,10 +400,5 @@ func testComputation(t *testing.T) Computation {
 		Datasets:        []Dataset{{Hash: dataHash, UserKey: []byte("key"), Dataset: data, Filename: datasetFile}},
 		Algorithm:       Algorithm{Hash: algoHash, UserKey: []byte("key"), Algorithm: algo},
 		ResultConsumers: []ResultConsumer{{UserKey: []byte("key")}},
-		AgentConfig: AgentConfig{
-			Port:        "7002",
-			LogLevel:    "debug",
-			AttestedTls: false,
-		},
 	}
 }

cli/README.md

@@ -100,3 +100,22 @@ When defining the manifest dataset and algorithm checksums are required. This ca
 ```bash
 ./build/cocos-cli checksum <path_to_dataset_or_algorithm>
 ```
+
+#### Measure IGVM file
+We assume that our current working directory is the root of the cocos repository, both on the host machine and in the VM.
+
+`igvmmeasure` calculates the launch measurement for an IGVM file and can generate a signed version. It ensures integrity by precomputing the expected launch digest, which can be verified against the attestation report. The tool parses IGVM directives, outputs the measurement as a hex string, or creates a signed file for verification at guest launch.
+
+##### Example
+We measure an IGVM file using our measure command, run:
+
+```bash
+./build/cocos-cli igvmmeasure /path/to/igvm/file
+```
+
+The tool will parse the directives in the IGVM file, calculate the launch measurement, and output the computed digest. If successful, it prints the measurement to standard output.
+
+Here is a sample output
+```
+91c4929bec2d0ecf11a708e09f0a57d7d82208bcba2451564444a4b01c22d047995ca27f9053f86de4e8063e9f810548
+```

cli/attestation.go

@@ -3,9 +3,11 @@
 package cli
 
 import (
+	"crypto"
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
+	"io"
 	"os"
 	"strconv"
 	"strings"
@@ -17,11 +19,18 @@ import (
 	"github.com/google/go-sev-guest/proto/check"
 	"github.com/google/go-sev-guest/proto/sevsnp"
 	"github.com/google/go-sev-guest/tools/lib/report"
+	sevVerify "github.com/google/go-sev-guest/verify"
+	tpmAttest "github.com/google/go-tpm-tools/proto/attest"
+	"github.com/google/go-tpm-tools/server"
+	"github.com/google/go-tpm/legacy/tpm2"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
-	"github.com/ultravioletrs/cocos/agent"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
 	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/encoding/prototext"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
@@ -42,9 +51,12 @@ const (
 	size48                  = 48
 	size64                  = 64
 	attestationFilePath     = "attestation.bin"
+	vtpmFilePath            = "../quote.dat"
 	attestationJson         = "attestation.json"
 	sevProductNameMilan     = "Milan"
 	sevProductNameGenoa     = "Genoa"
+	FormatBinaryPB          = "binarypb"
+	FormatTextProto         = "textproto"
 	exampleJSONConfig       = `
 	{
 		"rootOfTrust":{
@@ -98,29 +110,47 @@ const (
 		}
 	}
 	`
+	SNP     = "snp"
+	VTPM    = "vtpm"
+	SNPvTPM = "snp-vtpm"
 )
 
 var (
-	cfg                 = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
-	cfgString           string
-	timeout             time.Duration
-	maxRetryDelay       time.Duration
-	platformInfo        string
-	stepping            string
-	trustedAuthorKeys   []string
-	trustedAuthorHashes []string
-	trustedIdKeys       []string
-	trustedIdKeyHashes  []string
-	attestationFile     string
-	attestation         []byte
-	empty16             = [size16]byte{}
-	empty32             = [size32]byte{}
-	empty64             = [size64]byte{}
-	defaultReportIdMa   = []byte{255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255}
-	getJsonAttestation  bool
-	errReportSize       = errors.New("attestation contents too small")
+	mode                    string
+	cfg                     = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+	cfgString               string
+	timeout                 time.Duration
+	maxRetryDelay           time.Duration
+	platformInfo            string
+	stepping                string
+	trustedAuthorKeys       []string
+	trustedAuthorHashes     []string
+	trustedIdKeys           []string
+	trustedIdKeyHashes      []string
+	attestationFile         string
+	tpmAttestationFile      string
+	attestation             []byte
+	empty16                 = [size16]byte{}
+	empty32                 = [size32]byte{}
+	empty64                 = [size64]byte{}
+	defaultReportIdMa       = []byte{255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255}
+	errReportSize           = errors.New("attestation contents too small")
+	ErrBadAttestation       = errors.New("attestation file is corrupted or in wrong format")
+	output                  string
+	nonce                   []byte
+	format                  string
+	teeNonce                []byte
+	getTextProtoAttestation bool
 )
 
+var errEmptyFile = errors.New("input file is empty")
+
+var marshalOptions = prototext.MarshalOptions{
+	Multiline: true,
+	EmitASCII: true,
+}
+var unmarshalOptions = prototext.UnmarshalOptions{}
+
 func (cli *CLI) NewAttestationCmd() *cobra.Command {
 	return &cobra.Command{
 		Use:   "attestation [command]",
@@ -153,31 +183,75 @@ func (cli *CLI) NewAttestationCmd() *cobra.Command {
 
 func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "get",
-		Short:   "Retrieve attestation information from agent. Report data expected in hex enoded string of length 64 bytes.",
-		Example: "get <report_data>",
-		Args:    cobra.ExactArgs(1),
+		Use:       "get",
+		Short:     "Retrieve attestation information from agent. The argument of the command must be the type of the report (snp or vtpm or snp-vtpm).",
+		ValidArgs: []cobra.Completion{SNP, VTPM, SNPvTPM},
+		Example: fmt.Sprintf(`Based on attestation report type:
+		get %s --tee <512 bit hex value>
+		get %s --vtpm <256 bit hex value>
+		get %s --tee <512 bit hex value> --vtpm <256 bit hex value>`, SNP, VTPM, SNPvTPM),
+		Args: cobra.ExactArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
 			if cli.connectErr != nil {
 				printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
 				return
 			}
 
-			cmd.Println("Getting attestation")
-
-			reportData, err := hex.DecodeString(args[0])
-			if err != nil {
-				printError(cmd, "Error decoding report data: %v ❌ ", err)
+			if err := cobra.OnlyValidArgs(cmd, args); err != nil {
+				printError(cmd, "Bad attestation type: %v ❌ ", err)
 				return
 			}
-			if len(reportData) != agent.ReportDataSize {
-				msg := color.New(color.FgRed).Sprintf("report data must be a hex encoded string of length %d bytes ❌ ", agent.ReportDataSize)
+
+			attestationType := args[0]
+
+			attType := config.SNP
+			switch attestationType {
+			case SNP:
+				cmd.Println("Fetching SEV-SNP attestation report")
+			case VTPM:
+				cmd.Println("Fetching vTPM report")
+				attType = config.VTPM
+			case SNPvTPM:
+				cmd.Println("Fetching SEV-SNP and vTPM report")
+				attType = config.SNPvTPM
+			}
+
+			if (attType == config.VTPM || attType == config.SNPvTPM) && len(nonce) == 0 {
+				msg := color.New(color.FgRed).Sprint("vTPM nonce must be defined for vTPM attestation ❌ ")
 				cmd.Println(msg)
 				return
 			}
 
+			if (attType == config.SNP || attType == config.SNPvTPM) && len(teeNonce) == 0 {
+				msg := color.New(color.FgRed).Sprint("TEE nonce must be defined for SEV-SNP attestation ❌ ")
+				cmd.Println(msg)
+				return
+			}
+
+			var fixedReportData [quoteprovider.Nonce]byte
+			if attType != config.VTPM {
+				if len(teeNonce) > quoteprovider.Nonce {
+					msg := color.New(color.FgRed).Sprintf("nonce must be a hex encoded string of length lesser or equal %d bytes ❌ ", quoteprovider.Nonce)
+					cmd.Println(msg)
+					return
+				}
+
+				copy(fixedReportData[:], teeNonce)
+			}
+
+			var fixedVtpmNonceByte [vtpm.Nonce]byte
+			if attType != config.SNP {
+				if len(nonce) > vtpm.Nonce {
+					msg := color.New(color.FgRed).Sprintf("vTPM nonce must be a hex encoded string of length lesser or equal %d bytes ❌ ", vtpm.Nonce)
+					cmd.Println(msg)
+					return
+				}
+
+				copy(fixedVtpmNonceByte[:], nonce)
+			}
+
 			filename := attestationFilePath
-			if getJsonAttestation {
+			if getTextProtoAttestation {
 				filename = attestationJson
 			}
 
@@ -187,7 +261,7 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 				return
 			}
 
-			if err := cli.agentSDK.Attestation(cmd.Context(), [agent.ReportDataSize]byte(reportData), attestationFile); err != nil {
+			if err := cli.agentSDK.Attestation(cmd.Context(), fixedReportData, fixedVtpmNonceByte, int(attType), attestationFile); err != nil {
 				printError(cmd, "Failed to get attestation due to error: %v ❌ ", err)
 				return
 			}
@@ -197,16 +271,32 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 				return
 			}
 
-			if getJsonAttestation {
+			if getTextProtoAttestation {
 				result, err := os.ReadFile(filename)
 				if err != nil {
 					printError(cmd, "Error reading attestation file: %v ❌ ", err)
 					return
 				}
 
-				result, err = attesationToJSON(result)
+				switch attestationType {
+				case SNP:
+					result, err = attesationToJSON(result)
+				case VTPM, SNPvTPM:
+					marshalOptions := prototext.MarshalOptions{
+						Multiline: true,
+						EmitASCII: true,
+					}
+					var attvTPM tpmAttest.Attestation
+					err = proto.Unmarshal(result, &attvTPM)
+					if err != nil {
+						printError(cmd, "failed to unmarshal the attestation report: %v ❌ ", ErrBadAttestation)
+					}
+
+					result = []byte(marshalOptions.Format(&attvTPM))
+				}
+
 				if err != nil {
-					printError(cmd, "Error converting attestation to json: %v ❌ ", err)
+					printError(cmd, "Error converting attestation to textproto: %v ❌ ", err)
 					return
 				}
 
@@ -220,7 +310,9 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 		},
 	}
 
-	cmd.Flags().BoolVarP(&getJsonAttestation, "json", "j", false, "Get attestation in json format")
+	cmd.Flags().BoolVarP(&getTextProtoAttestation, "textproto", "p", false, "Get attestation in textproto format")
+	cmd.Flags().BytesHexVarP(&teeNonce, "tee", "e", []byte{}, "Define the nonce for the SNP attestation report (must be used with attestation type snp and snp-vtpm)")
+	cmd.Flags().BytesHexVarP(&nonce, "vtpm", "t", []byte{}, "Define the nonce for the vTPM attestation report (must be used with attestation type vtpm and snp-vtpm)")
 
 	return cmd
 }
@@ -252,52 +344,108 @@ func isFileJSON(filename string) bool {
 
 func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:     "validate",
-		Short:   "Validate and verify attestation information. The report is provided as a file path.",
-		Example: "validate <attestation_report_file_path>",
-		Args:    cobra.ExactArgs(1),
-		Run: func(cmd *cobra.Command, args []string) {
-			cmd.Println("Checking attestation")
+		Use:   "validate",
+		Short: "Validate and verify attestation information. You can choose from 3 modes: snp,vtpm and snp-vtpm.Default mode is snp.",
+		Example: `Based on mode:
+		validate <attestationreportfilepath> --report_data <reportdata> --product <product data> //default
+		validate --mode snp <attestationreportfilepath> --report_data <reportdata> --product <product data>
+		validate --mode vtpm <attestationreportfilepath> --nonce <noncevalue> --format <formatvalue>  --output <outputvalue>
+		validate --mode snp-vtpm <attestationreportfilepath> --nonce <noncevalue> --format <formatvalue>  --output <outputvalue>`,
 
-			attestationFile = string(args[0])
-
-			if err := parseConfig(); err != nil {
-				printError(cmd, "Error parsing config: %v ❌ ", err)
-				return
-			}
-			if err := parseHashes(); err != nil {
-				printError(cmd, "Error parsing hashes: %v ❌ ", err)
-				return
-			}
-			if err := parseFiles(); err != nil {
-				printError(cmd, "Error parsing files: %v ❌ ", err)
-				return
-			}
-			// This format is the attestation report in AMD's specified ABI format, immediately
-			// followed by the certificate table bytes.
-			if len(attestation) < abi.ReportSize {
-				msg := color.New(color.FgRed).Sprintf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes ❌ ", len(attestation), abi.ReportSize)
-				cmd.Println(msg)
-				return
-			}
-			if err := parseUints(); err != nil {
-				printError(cmd, "Error parsing uints: %v ❌ ", err)
-				return
-			}
-			cfg.Policy.Vmpl = wrapperspb.UInt32(0)
-
-			if err := validateInput(); err != nil {
-				printError(cmd, "Error validating input: %v ❌ ", err)
-				return
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			mode, _ := cmd.Flags().GetString("mode")
+			if len(args) != 1 {
+				return fmt.Errorf("please pass the attestation report file path")
 			}
 
-			if err := quoteprovider.VerifyAndValidate(attestation, &cfg); err != nil {
-				printError(cmd, "Attestation validation and verification failed with error: %v ❌ ", err)
-				return
+			// Validate flags based on the mode
+			switch mode {
+			case "snp":
+				if err := cmd.MarkFlagRequired("report_data"); err != nil {
+					return fmt.Errorf("failed to mark 'report_data' as required for SEV-SNP mode: %v", err)
+				}
+				if err := cmd.MarkFlagRequired("product"); err != nil {
+					return fmt.Errorf("failed to mark flag as required: %v ❌ ", err)
+				}
+			case "snp-vtpm":
+				if err := cmd.MarkFlagRequired("nonce"); err != nil {
+					return fmt.Errorf("failed to mark 'nonce' as required for vTPM mode: %v", err)
+				}
+				if err := cmd.MarkFlagRequired("format"); err != nil {
+					return fmt.Errorf("failed to mark 'format' as required for vTPM mode: %v", err)
+				}
+				if err := cmd.MarkFlagRequired("output"); err != nil {
+					return fmt.Errorf("failed to mark 'output' as required for vTPM mode: %v", err)
+				}
+
+			case "vtpm":
+				if err := cmd.MarkFlagRequired("nonce"); err != nil {
+					return fmt.Errorf("failed to mark 'nonce' as required for vTPM mode: %v", err)
+				}
+				if err := cmd.MarkFlagRequired("format"); err != nil {
+					return fmt.Errorf("failed to mark 'format' as required for vTPM mode: %v", err)
+				}
+				if err := cmd.MarkFlagRequired("output"); err != nil {
+					return fmt.Errorf("failed to mark 'output' as required for vTPM mode: %v", err)
+				}
+			default:
+				return fmt.Errorf("unknown mode: %s", mode)
 			}
-			cmd.Println("Attestation validation and verification is successful!")
+			return nil
 		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			mode, _ := cmd.Flags().GetString("mode")
+			switch mode {
+			case "snp":
+				return sevsnpverify(cmd, args)
+			case "snp-vtpm":
+				return vtpmSevSnpverify(args)
+			case "vtpm":
+				return vtpmverify(args)
+			default:
+				return fmt.Errorf("unknown mode: %s", mode)
+			}
+		},
+		SilenceUsage:  true,
+		SilenceErrors: true,
 	}
+	cmd.Flags().StringVar(
+		&mode,
+		"mode",
+		"snp", // default mode
+		"The attestation validation mode. Example: sevsnp",
+	)
+
+	// VTPM FLAGS
+	cmd.Flags().BytesHexVar(
+		&nonce,
+		"nonce",
+		[]byte{},
+		"hex encoded nonce for vTPM attestation, cannot be empty",
+	)
+
+	cmd.Flags().StringVar(
+		&format,
+		"format",
+		"binarypb", // default value
+		"type of output file where attestation report stored <binarypb|textproto>",
+	)
+
+	cmd.Flags().StringVar(
+		&output,
+		"output",
+		"",
+		"output file",
+	)
+
+	cmd.Flags().BytesHexVar(
+		&teeNonce,
+		"tee-nonce",
+		[]byte{},
+		"hex encoded teenonce for hardware attestation, can be empty",
+	)
+
+	// SEV-SNP FLAGS
 	cmd.Flags().StringVar(
 		&cfgString,
 		"config",
@@ -473,19 +621,233 @@ func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
 		"PEM format CA bundles for the AMD product. Combined with contents of cabundle_paths.",
 	)
 
-	if err := cmd.MarkFlagRequired("report_data"); err != nil {
-		printError(cmd, "Failed to mark flag as required: %v ❌ ", err)
-		return nil
-	}
-
-	if err := cmd.MarkFlagRequired("product"); err != nil {
-		printError(cmd, "Failed to mark flag as required: %v ❌ ", err)
-		return nil
-	}
-
 	return cmd
 }
 
+func (cli *CLI) NewMeasureCmd(igvmBinaryPath string) *cobra.Command {
+	igvmmeasureCmd := &cobra.Command{
+		Use:   "igvmmeasure <INPUT>",
+		Short: "Measure an IGVM file",
+		Long: `igvmmeasure measures an IGVM file and outputs the calculated measurement.
+			It ensures integrity verification for the IGVM file.`,
+
+		Args: cobra.MinimumNArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if len(args) == 0 {
+				return fmt.Errorf("error: No input file provided")
+			}
+
+			inputFile := args[0]
+
+			measurement, err := cli.measurement.Run(inputFile)
+			if err != nil {
+				return err
+			}
+
+			outputString := string(measurement)
+			lines := strings.Split(strings.TrimSpace(outputString), "\n")
+
+			if len(lines) == 1 {
+				outputString = strings.ToLower(outputString)
+			} else {
+				return fmt.Errorf("error: %s", outputString)
+			}
+
+			cmd.Print(outputString)
+
+			return nil
+		},
+	}
+
+	return igvmmeasureCmd
+}
+
+func sevsnpverify(cmd *cobra.Command, args []string) error {
+	cmd.Println("Checking attestation")
+
+	attestationFile = string(args[0])
+
+	if err := parseConfig(); err != nil {
+		return fmt.Errorf("error parsing config: %v ❌ ", err)
+	}
+	if err := parseHashes(); err != nil {
+		return fmt.Errorf("error parsing hashes: %v ❌ ", err)
+	}
+	if err := parseFiles(); err != nil {
+		return fmt.Errorf("error parsing files: %v ❌ ", err)
+	}
+	// This format is the attestation report in AMD's specified ABI format, immediately
+	// followed by the certificate table bytes.
+	if len(attestation) < abi.ReportSize {
+		return fmt.Errorf("attestation too small: got 0x%x bytes, need at least 0x%x bytes", len(attestation), abi.ReportSize)
+	}
+	if err := parseUints(); err != nil {
+		return fmt.Errorf("error parsing uints: %v ❌ ", err)
+	}
+	cfg.Policy.Vmpl = wrapperspb.UInt32(0)
+
+	if err := validateInput(); err != nil {
+		return fmt.Errorf("error validating input: %v ❌ ", err)
+	}
+
+	attestationPB, err := abi.ReportCertsToProto(attestation)
+	if err != nil {
+		return fmt.Errorf("failed to convert attestation bytes to struct %v ❌ ", err)
+	}
+
+	if err := quoteprovider.VerifyAndValidate(attestationPB, &cfg); err != nil {
+		return fmt.Errorf("attestation validation and verification failed with error: %v ❌ ", err)
+	}
+	cmd.Println("Attestation validation and verification is successful!")
+	return nil
+}
+
+func vtpmSevSnpverify(args []string) error {
+	tpmAttestationFile = string(args[0])
+	input, err := openInputFile()
+	if err != nil {
+		return err
+	}
+	if closer, ok := input.(*os.File); ok {
+		defer closer.Close()
+	}
+	attestationBytes, err := io.ReadAll(input)
+	if err != nil {
+		return err
+	}
+	attestation := &tpmAttest.Attestation{}
+
+	if format == FormatBinaryPB {
+		err = proto.Unmarshal(attestationBytes, attestation)
+	} else if format == FormatTextProto {
+		err = unmarshalOptions.Unmarshal(attestationBytes, attestation)
+	} else {
+		return fmt.Errorf("format should be either binarypb or textproto")
+	}
+	if err != nil {
+		return fmt.Errorf("fail to unmarshal attestation report: %v", err)
+	}
+
+	pub, err := tpm2.DecodePublic(attestation.GetAkPub())
+	if err != nil {
+		return err
+	}
+	cryptoPub, err := pub.Key()
+	if err != nil {
+		return err
+	}
+
+	var validateOpts interface{}
+	switch attestation.GetTeeAttestation().(type) {
+	case *tpmAttest.Attestation_SevSnpAttestation:
+		if len(teeNonce) != 0 {
+			validateOpts = &server.VerifySnpOpts{
+				Validation:   server.SevSnpDefaultValidateOpts(teeNonce),
+				Verification: &sevVerify.Options{},
+			}
+		} else {
+			validateOpts = &server.VerifySnpOpts{
+				Validation:   server.SevSnpDefaultValidateOpts(nonce),
+				Verification: &sevVerify.Options{},
+			}
+		}
+	default:
+		validateOpts = nil
+	}
+
+	ms, err := server.VerifyAttestation(attestation, server.VerifyOpts{Nonce: nonce, TrustedAKs: []crypto.PublicKey{cryptoPub}, TEEOpts: validateOpts})
+	if err != nil {
+		return fmt.Errorf("verifying attestation: %w", err)
+	}
+	out, err := marshalOptions.Marshal(ms)
+	if err != nil {
+		return nil
+	}
+	output, err := createOutputFile()
+	if err != nil {
+		return err
+	}
+	if closer, ok := output.(*os.File); ok {
+		defer closer.Close()
+	}
+	if _, err := output.Write(out); err != nil {
+		return fmt.Errorf("failed to write verified attestation report: %v", err)
+	}
+
+	return nil
+}
+
+func vtpmverify(args []string) error {
+	tpmAttestationFile = string(args[0])
+	input, err := openInputFile()
+	if err != nil {
+		return err
+	}
+	if closer, ok := input.(*os.File); ok {
+		defer closer.Close()
+	}
+	attestationBytes, err := io.ReadAll(input)
+	if err != nil {
+		return err
+	}
+	attestation := &tpmAttest.Attestation{}
+
+	if format == FormatBinaryPB {
+		err = proto.Unmarshal(attestationBytes, attestation)
+	} else if format == FormatTextProto {
+		err = unmarshalOptions.Unmarshal(attestationBytes, attestation)
+	} else {
+		return fmt.Errorf("format should be either binarypb or textproto")
+	}
+	if err != nil {
+		return fmt.Errorf("fail to unmarshal attestation report: %v", err)
+	}
+
+	pub, err := tpm2.DecodePublic(attestation.GetAkPub())
+	if err != nil {
+		return err
+	}
+	cryptoPub, err := pub.Key()
+	if err != nil {
+		return err
+	}
+
+	ms, err := server.VerifyAttestation(attestation, server.VerifyOpts{Nonce: nonce, TrustedAKs: []crypto.PublicKey{cryptoPub}, TEEOpts: nil})
+	if err != nil {
+		return nil
+	}
+	out, err := marshalOptions.Marshal(ms)
+	if err != nil {
+		return nil
+	}
+	output, err := createOutputFile()
+	if err != nil {
+		return err
+	}
+	if closer, ok := output.(*os.File); ok {
+		defer closer.Close()
+	}
+	if _, err := output.Write(out); err != nil {
+		return fmt.Errorf("failed to write verified attestation report: %v", err)
+	}
+
+	return nil
+}
+
+func openInputFile() (io.Reader, error) {
+	if tpmAttestationFile == "" {
+		return nil, errEmptyFile
+	}
+	return os.Open(tpmAttestationFile)
+}
+
+func createOutputFile() (io.Writer, error) {
+	if output == "" {
+		return os.Stdout, nil
+	}
+	return os.Create(output)
+}
+
 // parseConfig decodes config passed as json for check.Config struct.
 // example
 /* {

cli/attestation_policy.go

@@ -3,15 +3,22 @@
 package cli
 
 import (
+	"bytes"
+	"crypto/sha512"
 	"encoding/base64"
+	"encoding/json"
 	"fmt"
 	"os"
+	"strconv"
 
 	"github.com/absmach/magistrala/pkg/errors"
 	"github.com/google/go-sev-guest/proto/check"
+	"github.com/google/go-tpm-tools/proto/attest"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
-	"google.golang.org/protobuf/encoding/protojson"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/gcp"
+	"google.golang.org/protobuf/proto"
 )
 
 type fieldType int
@@ -99,6 +106,126 @@ func (cli *CLI) NewAddHostDataCmd() *cobra.Command {
 	}
 }
 
+func (cli *CLI) NewGCPAttestationPolicy() *cobra.Command {
+	return &cobra.Command{
+		Use:     "gcp",
+		Short:   "Get attestation policy for GCP CVM",
+		Example: `gcp <bin_vtmp_attestation_report_file> <vcpu_count>`,
+		Args:    cobra.ExactArgs(2),
+		Run: func(cmd *cobra.Command, args []string) {
+			attestationBin, err := os.ReadFile(args[0])
+			if err != nil {
+				printError(cmd, "Error reading attestation report file: %v ❌ ", err)
+				return
+			}
+
+			vcpuCount, err := strconv.Atoi(args[1])
+			if err != nil {
+				printError(cmd, "Error converting vCPU count to integer: %v ❌ ", err)
+				return
+			}
+
+			attestation := &attest.Attestation{}
+
+			if err := proto.Unmarshal(attestationBin, attestation); err != nil {
+				printError(cmd, "Error unmarshaling attestation report: %v ❌ ", err)
+				return
+			}
+
+			attestationPB := attestation.GetSevSnpAttestation()
+
+			measurement, err := gcp.Extract384BitMeasurement(attestationPB)
+			if err != nil {
+				printError(cmd, "Error extracting 384-bit measurement: %v ❌ ", err)
+				return
+			}
+
+			launchEndorsement, err := gcp.GetLaunchEndorsement(cmd.Context(), measurement)
+			if err != nil {
+				printError(cmd, "Error getting launch endorsement: %v ❌ ", err)
+				return
+			}
+
+			attestationPolicy, err := gcp.GenerateAttestationPolicy(launchEndorsement, uint32(vcpuCount))
+			if err != nil {
+				printError(cmd, "Error generating attestation policy: %v ❌ ", err)
+				return
+			}
+
+			attestationPolicyJson, err := json.MarshalIndent(attestationPolicy, "", "  ")
+			if err != nil {
+				printError(cmd, "Error marshaling attestation policy: %v ❌ ", err)
+				return
+			}
+
+			if err := os.WriteFile("attestation_policy.json", attestationPolicyJson, filePermission); err != nil {
+				printError(cmd, "Error writing attestation policy file: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println("Attestation policy file generated successfully ✅")
+		},
+	}
+}
+
+func (cli *CLI) NewDownloadGCPOvmfFile() *cobra.Command {
+	return &cobra.Command{
+		Use:     "download",
+		Short:   "Download GCP OVMF file",
+		Example: `download <bin_vtmp_attestation_report_file>`,
+		Args:    cobra.ExactArgs(1),
+		Run: func(cmd *cobra.Command, args []string) {
+			attestationBin, err := os.ReadFile(args[0])
+			if err != nil {
+				printError(cmd, "Error reading attestation report file: %v ❌ ", err)
+				return
+			}
+
+			attestation := &attest.Attestation{}
+
+			if err := proto.Unmarshal(attestationBin, attestation); err != nil {
+				printError(cmd, "Error unmarshaling attestation report: %v ❌ ", err)
+				return
+			}
+
+			attestationPB := attestation.GetSevSnpAttestation()
+
+			measurement, err := gcp.Extract384BitMeasurement(attestationPB)
+			if err != nil {
+				printError(cmd, "Error extracting 384-bit measurement: %v ❌ ", err)
+				return
+			}
+
+			launchEndorsement, err := gcp.GetLaunchEndorsement(cmd.Context(), measurement)
+			if err != nil {
+				printError(cmd, "Error getting launch endorsement: %v ❌ ", err)
+				return
+			}
+
+			ovmf, err := gcp.DownloadOvmfFile(cmd.Context(), fmt.Sprintf("%x", launchEndorsement.Digest))
+			if err != nil {
+				printError(cmd, "Error downloading OVMF file: %v ❌ ", err)
+				return
+			}
+
+			sum384 := sha512.Sum384(ovmf)
+
+			if !bytes.Equal(sum384[:], launchEndorsement.Digest) {
+				printError(cmd, "Error OVMF file does not match the measurement: %v ❌ ", fmt.Errorf("digest mismatch"))
+			} else {
+				cmd.Println("OVMF firmware in vm is unmodified ✅")
+			}
+
+			if err := os.WriteFile("ovmf.fd", ovmf, filePermission); err != nil {
+				printError(cmd, "Error writing OVMF file: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println("OVMF file downloaded successfully ✅")
+		},
+	}
+}
+
 func changeAttestationConfiguration(fileName, base64Data string, expectedLength int, field fieldType) error {
 	data, err := base64.StdEncoding.DecodeString(base64Data)
 	if err != nil {
@@ -109,27 +236,31 @@ func changeAttestationConfiguration(fileName, base64Data string, expectedLength
 		return errDataLength
 	}
 
-	ac := check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+	ac := config.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &config.PcrConfig{}}
 
-	attestationPolicy, err := os.ReadFile(fileName)
+	f, err := os.ReadFile(fileName)
 	if err != nil {
 		return errors.Wrap(errReadingAttestationPolicyFile, err)
 	}
 
-	if err = protojson.Unmarshal(attestationPolicy, &ac); err != nil {
+	if err = config.ReadAttestationPolicyFromByte(f, &ac); err != nil {
 		return errors.Wrap(errUnmarshalJSON, err)
 	}
 
+	if ac.Config.Policy == nil {
+		ac.Config.Policy = &check.Policy{}
+	}
+
 	switch field {
 	case measurementField:
-		ac.Policy.Measurement = data
+		ac.Config.Policy.Measurement = data
 	case hostDataField:
-		ac.Policy.HostData = data
+		ac.Config.Policy.HostData = data
 	default:
 		return errAttestationPolicyField
 	}
 
-	fileJson, err := protojson.Marshal(&ac)
+	fileJson, err := json.MarshalIndent(&ac, "", " ")
 	if err != nil {
 		return errors.Wrap(errMarshalJSON, err)
 	}

cli/attestation_policy_test.go

@@ -4,13 +4,14 @@ package cli
 
 import (
 	"encoding/base64"
+	"encoding/json"
 	"os"
 	"testing"
 
 	"github.com/google/go-sev-guest/proto/check"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"google.golang.org/protobuf/encoding/protojson"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
 )
 
 func TestChangeAttestationConfiguration(t *testing.T) {
@@ -18,9 +19,9 @@ func TestChangeAttestationConfiguration(t *testing.T) {
 	require.NoError(t, err)
 	defer os.Remove(tmpfile.Name())
 
-	initialConfig := check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+	initialConfig := config.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &config.PcrConfig{}}
 
-	initialJSON, err := protojson.Marshal(&initialConfig)
+	initialJSON, err := json.Marshal(initialConfig)
 	require.NoError(t, err)
 	err = os.WriteFile(tmpfile.Name(), initialJSON, 0o644)
 	require.NoError(t, err)
@@ -86,15 +87,15 @@ func TestChangeAttestationConfiguration(t *testing.T) {
 				content, err := os.ReadFile(tmpfile.Name())
 				require.NoError(t, err)
 
-				config := check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
-				err = protojson.Unmarshal(content, &config)
+				ap := config.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &config.PcrConfig{}}
+				err = config.ReadAttestationPolicyFromByte(content, &ap)
 				require.NoError(t, err)
 
 				decodedData, _ := base64.StdEncoding.DecodeString(tt.base64Data)
 				if tt.field == measurementField {
-					assert.Equal(t, decodedData, config.Policy.Measurement)
+					assert.Equal(t, decodedData, ap.Config.Policy.Measurement)
 				} else if tt.field == hostDataField {
-					assert.Equal(t, decodedData, config.Policy.HostData)
+					assert.Equal(t, decodedData, ap.Config.Policy.HostData)
 				}
 			}
 		})

cli/attestation_test.go

@@ -14,10 +14,12 @@ import (
 	"github.com/google/go-sev-guest/abi"
 	"github.com/google/go-sev-guest/proto/check"
 	"github.com/google/go-sev-guest/proto/sevsnp"
+	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
-	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
 )
 
@@ -32,10 +34,8 @@ func TestNewAttestationCmd(t *testing.T) {
 	var buf bytes.Buffer
 	cmd.SetOut(&buf)
 
-	cmd.SetOutput(&buf)
-
-	reportData := bytes.Repeat([]byte{0x01}, agent.ReportDataSize)
-	mockSDK.On("Attestation", mock.Anything, [agent.ReportDataSize]byte(reportData), mock.Anything).Return(nil)
+	reportData := bytes.Repeat([]byte{0x01}, quoteprovider.Nonce)
+	mockSDK.On("Attestation", mock.Anything, [quoteprovider.Nonce]byte(reportData), mock.Anything).Return(nil)
 
 	cmd.SetArgs([]string{hex.EncodeToString(reportData)})
 	err := cmd.Execute()
@@ -46,6 +46,10 @@ func TestNewAttestationCmd(t *testing.T) {
 func TestNewGetAttestationCmd(t *testing.T) {
 	validattestation, err := os.ReadFile("../attestation.bin")
 	require.NoError(t, err)
+
+	teeNonce := hex.EncodeToString(bytes.Repeat([]byte{0x00}, quoteprovider.Nonce))
+	vtpmNonce := hex.EncodeToString(bytes.Repeat([]byte{0x00}, vtpm.Nonce))
+
 	testCases := []struct {
 		name         string
 		args         []string
@@ -55,57 +59,85 @@ func TestNewGetAttestationCmd(t *testing.T) {
 		expectedOut  string
 	}{
 		{
-			name:         "successful attestation retrieval",
-			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
+			name:         "successful SNP attestation retrieval",
+			args:         []string{"snp", "--tee", teeNonce},
 			mockResponse: []byte("mock attestation"),
 			mockError:    nil,
 			expectedOut:  "Attestation result retrieved and saved successfully!",
 		},
 		{
-			name:         "invalid report data (decoding error)",
-			args:         []string{"invalid"},
-			mockResponse: nil,
-			mockError:    errors.New("error"),
-			expectedErr:  "Error decoding report data",
+			name:         "successful vTPM attestation retrieval",
+			args:         []string{"vtpm", "--vtpm", vtpmNonce},
+			mockResponse: []byte("mock attestation"),
+			mockError:    nil,
+			expectedOut:  "Attestation result retrieved and saved successfully!",
+		},
+		{
+			name:         "successful SNP-vTPM attestation retrieval",
+			args:         []string{"snp-vtpm", "--tee", teeNonce, "--vtpm", vtpmNonce},
+			mockResponse: []byte("mock attestation"),
+			mockError:    nil,
+			expectedOut:  "Attestation result retrieved and saved successfully!",
+		},
+		{
+			name:         "missing vTPM nonce",
+			args:         []string{"snp-vtpm", "--tee", teeNonce},
+			mockResponse: []byte("mock attestation"),
+			mockError:    nil,
+			expectedOut:  "vTPM nonce must be defined for vTPM attestation",
+		},
+		{
+			name:         "missing TEE nonce",
+			args:         []string{"snp-vtpm", "--vtpm", vtpmNonce},
+			mockResponse: []byte("mock attestation"),
+			mockError:    nil,
+			expectedOut:  "TEE nonce must be defined for SEV-SNP attestation",
 		},
 		{
 			name:         "invalid report data size",
-			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, 32))},
+			args:         []string{"snp", "--tee", hex.EncodeToString(bytes.Repeat([]byte{0x00}, 65))},
 			mockResponse: nil,
 			mockError:    errors.New("error"),
-			expectedErr:  "report data must be a hex encoded string of length 64 bytes",
+			expectedErr:  "nonce must be a hex encoded string of length lesser or equal 64 bytes",
 		},
 		{
-			name:         "invalid report data hex",
+			name:         "invalid vTPM data size",
+			args:         []string{"vtpm", "-t", hex.EncodeToString(bytes.Repeat([]byte{0x00}, 33))},
+			mockResponse: nil,
+			mockError:    errors.New("error"),
+			expectedErr:  "vTPM nonce must be a hex encoded string of length lesser or equal 32 bytes",
+		},
+		{
+			name:         "invalid arguments",
 			args:         []string{"invalid"},
 			mockResponse: nil,
 			mockError:    errors.New("error"),
-			expectedErr:  "Error decoding report data",
+			expectedErr:  "Bad attestation type: invalid argument ",
 		},
 		{
 			name:         "failed to get attestation",
-			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
+			args:         []string{"snp", "-e", teeNonce},
 			mockResponse: nil,
 			mockError:    errors.New("error"),
 			expectedErr:  "Failed to get attestation due to error",
 		},
 		{
-			name:         "JSON report error",
-			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), "--json"},
+			name:         "Textproto report error",
+			args:         []string{"snp", "-e", teeNonce, "--textproto"},
 			mockResponse: []byte("mock attestation"),
 			mockError:    nil,
-			expectedErr:  "Error converting attestation to json",
+			expectedErr:  "Error converting attestation to textproto",
 		},
 		{
-			name:         "successful JSON report",
-			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), "--json"},
+			name:         "successful Textproto report",
+			args:         []string{"snp", "-e", teeNonce, "--textproto"},
 			mockResponse: validattestation,
 			mockError:    nil,
 			expectedOut:  "Attestation result retrieved and saved successfully!",
 		},
 		{
 			name:         "connection error",
-			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
+			args:         []string{"snp", "-e", teeNonce},
 			mockResponse: nil,
 			mockError:    errors.New("failed to connect to agent"),
 			expectedErr:  "Failed to connect to agent",
@@ -125,10 +157,10 @@ func TestNewGetAttestationCmd(t *testing.T) {
 			}
 			cmd := cli.NewGetAttestationCmd()
 			var buf bytes.Buffer
-			cmd.SetOutput(&buf)
+			cmd.SetOut(&buf)
 
-			mockSDK.On("Attestation", mock.Anything, [agent.ReportDataSize]byte(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), mock.Anything).Return(tc.mockError).Run(func(args mock.Arguments) {
-				_, err := args.Get(2).(*os.File).Write(tc.mockResponse)
+			mockSDK.On("Attestation", mock.Anything, [quoteprovider.Nonce]byte(bytes.Repeat([]byte{0x00}, quoteprovider.Nonce)), [vtpm.Nonce]byte(bytes.Repeat([]byte{0x00}, vtpm.Nonce)), mock.Anything, mock.Anything).Return(tc.mockError).Run(func(args mock.Arguments) {
+				_, err := args.Get(4).(*os.File).Write(tc.mockResponse)
 				require.NoError(t, err)
 			})
 
@@ -145,12 +177,12 @@ func TestNewGetAttestationCmd(t *testing.T) {
 	}
 }
 
-func TestNewValidateAttestationValidationCmd(t *testing.T) {
+func TestNewValidateAttestationValidationCmdDefaults(t *testing.T) {
 	cli := &CLI{}
 	cmd := cli.NewValidateAttestationValidationCmd()
 
 	assert.Equal(t, "validate", cmd.Use)
-	assert.Equal(t, "Validate and verify attestation information. The report is provided as a file path.", cmd.Short)
+	assert.Equal(t, "Validate and verify attestation information. You can choose from 3 modes: snp,vtpm and snp-vtpm.Default mode is snp.", cmd.Short)
 
 	assert.Equal(t, fmt.Sprint(defaultMinimumTcb), cmd.Flag("minimum_tcb").Value.String())
 	assert.Equal(t, fmt.Sprint(defaultMinimumLaunchTcb), cmd.Flag("minimum_lauch_tcb").Value.String())
@@ -162,6 +194,150 @@ func TestNewValidateAttestationValidationCmd(t *testing.T) {
 	assert.Equal(t, fmt.Sprint(defaultMaxRetryDelay), cmd.Flag("max_retry_delay").Value.String())
 }
 
+func TestNewValidateAttestationValidationCmd(t *testing.T) {
+	cli := &CLI{}
+	cmd := cli.NewValidateAttestationValidationCmd()
+
+	t.Run("missing attestation report file path", func(t *testing.T) {
+		err := cmd.Execute()
+		assert.Error(t, err)
+		assert.Equal(t, "please pass the attestation report file path", err.Error())
+	})
+
+	t.Run("unknown mode", func(t *testing.T) {
+		cmd.SetArgs([]string{attestationFilePath, "--mode=invalid"})
+		err := cmd.Execute()
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "unknown mode")
+	})
+
+	t.Run("snp mode with missing flags", func(t *testing.T) {
+		cmd.SetArgs([]string{attestationFilePath, "--mode=snp"})
+		err := cmd.Execute()
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "required flag(s) \"product\", \"report_data\" not set")
+	})
+
+	t.Run("vtpm mode with missing flags", func(t *testing.T) {
+		cmd.SetArgs([]string{vtpmFilePath, "--mode=vtpm"})
+		err := cmd.Execute()
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "required flag(s) \"format\", \"nonce\", \"output\", \"product\", \"report_data\" not set")
+	})
+
+	t.Run("snp-vtpm mode with missing flags", func(t *testing.T) {
+		cmd.SetArgs([]string{vtpmFilePath, "--mode=snp-vtpm"})
+		err := cmd.Execute()
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "required flag(s) \"format\", \"nonce\", \"output\", \"product\", \"report_data\" not set")
+	})
+
+	t.Run("valid snp mode execution", func(t *testing.T) {
+		cli := CLI{}
+		cmd := cli.NewValidateAttestationValidationCmd()
+
+		cmd.RunE = func(_ *cobra.Command, _ []string) error {
+			t.Log("Mock RunE executed instead of sevsnpverify")
+			return nil
+		}
+
+		cmd.SetArgs([]string{
+			"../attestation.bin",
+			"--mode=snp",
+			"--report_data=" +
+				"11223344556677889900aabbccddeeff11223344556677889900aabbccddeeff" +
+				"11223344556677889900aabbccddeeff11223344556677889900aabbccddeeff",
+			"--product=Milan",
+		})
+		err := cmd.PreRunE(cmd, []string{"../attestation.bin"})
+		assert.NoError(t, err)
+	})
+
+	t.Run("valid vtpm mode execution", func(t *testing.T) {
+		cli := CLI{}
+		cmd := cli.NewValidateAttestationValidationCmd()
+
+		cmd.RunE = func(_ *cobra.Command, _ []string) error {
+			t.Log("Mock RunE executed instead of vtpmverify")
+			return nil
+		}
+
+		cmd.SetArgs([]string{vtpmFilePath, "--mode=vtpm", "--nonce=123abc", "--format=binarypb", "--output=some_output"})
+
+		err := cmd.PreRunE(cmd, []string{"../quote.dat"})
+		assert.NoError(t, err)
+	})
+
+	t.Run("valid snp-vtpm mode execution", func(t *testing.T) {
+		cli := CLI{}
+		cmd := cli.NewValidateAttestationValidationCmd()
+
+		cmd.RunE = func(_ *cobra.Command, _ []string) error {
+			t.Log("Mock RunE executed instead of vtpmSevSnpverify")
+			return nil
+		}
+
+		cmd.SetArgs([]string{vtpmFilePath, "--mode=snp-vtpm", "--nonce=123abc", "--format=textproto", "--output=some_output"})
+		err := cmd.PreRunE(cmd, []string{"../quote.dat"})
+		assert.NoError(t, err)
+	})
+}
+
+type MockMeasurement struct {
+	mock.Mock
+}
+
+func (m *MockMeasurement) Run(igvmBinaryPath string) ([]byte, error) {
+	args := m.Called(igvmBinaryPath)
+	return nil, args.Error(0)
+}
+
+func (m *MockMeasurement) Stop() error {
+	args := m.Called()
+	return args.Error(0)
+}
+
+func TestNewMeasureCmd_RunSuccess(t *testing.T) {
+	cliInstance := &CLI{}
+	mockMeasurement := new(MockMeasurement)
+	cliInstance.measurement = mockMeasurement
+
+	mockMeasurement.On("Run", "testfile.igvm").Return(nil)
+
+	cmd := cliInstance.NewMeasureCmd("fake_binary_path")
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetErr(buf)
+	cmd.SetArgs([]string{"testfile.igvm"})
+
+	err := cmd.Execute()
+
+	assert.NoError(t, err)
+	mockMeasurement.AssertExpectations(t)
+}
+
+func TestNewMeasureCmd_RunError(t *testing.T) {
+	cliInstance := &CLI{}
+	mockMeasurement := new(MockMeasurement)
+	cliInstance.measurement = mockMeasurement
+	expectedError := errors.New("mocked measurement error")
+
+	mockMeasurement.On("Run", "testfile.igvm").Return(expectedError)
+
+	cmd := cliInstance.NewMeasureCmd("fake_binary_path")
+
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetErr(buf)
+	cmd.SetArgs([]string{"testfile.igvm"})
+
+	err := cmd.Execute()
+
+	assert.Error(t, err)
+	assert.Equal(t, expectedError.Error(), err.Error())
+	mockMeasurement.AssertExpectations(t)
+}
+
 func TestParseConfig(t *testing.T) {
 	cfgString = ""
 	err := parseConfig()

cli/cache.go

@@ -12,7 +12,7 @@ import (
 	"github.com/google/go-sev-guest/proto/check"
 	"github.com/google/go-sev-guest/verify/trust"
 	"github.com/spf13/cobra"
-	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
 )
 
 const (
@@ -27,14 +27,14 @@ func (cli *CLI) NewCABundleCmd(fileSavePath string) *cobra.Command {
 		Example: "ca-bundle <path_to_platform_info_json>",
 		Args:    cobra.ExactArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
-			attestationConfiguration := check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
-			err := grpc.ReadAttestationPolicy(args[0], &attestationConfiguration)
+			attestationConfiguration := config.Config{Config: &check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}, PcrConfig: &config.PcrConfig{}}
+			err := config.ReadAttestationPolicy(args[0], &attestationConfiguration)
 			if err != nil {
 				printError(cmd, "Error while reading manifest: %v ❌ ", err)
 				return
 			}
 
-			product := attestationConfiguration.RootOfTrust.ProductLine
+			product := attestationConfiguration.Config.RootOfTrust.ProductLine
 
 			getter := trust.DefaultHTTPSGetter()
 			caURL := kds.ProductCertChainURL(abi.VcekReportSigner, product)

cli/checksum_test.go

@@ -131,7 +131,7 @@ func TestManifestChecksum(t *testing.T) {
 				"name": "Example Computation",
 				"description": "This is an example computation"
 			}`,
-			expectedSum: "868825367c32c4b6d621d5d95e2890f233d8554df2348ab743aac2663a936f08",
+			expectedSum: "a99683e4d22ba54cefa51aa49fb2e97a92b828c088395992ddff16a6236f3299",
 		},
 		{
 			name:        "Invalid JSON",

cli/manager.go

@@ -0,0 +1,130 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"os"
+
+	"github.com/fatih/color"
+	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/manager"
+)
+
+const (
+	serverURL = "server-url"
+	serverCA  = "server-ca"
+	clientKey = "client-key"
+	clientCrt = "client-crt"
+	logLevel  = "log-level"
+)
+
+var (
+	agentCVMServerUrl string
+	agentCVMServerCA  string
+	agentCVMClientKey string
+	agentCVMClientCrt string
+	agentLogLevel     string
+)
+
+func (c *CLI) NewCreateVMCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "create-vm",
+		Short:   "Create a new virtual machine",
+		Example: `create-vm`,
+		Args:    cobra.ExactArgs(0),
+		Run: func(cmd *cobra.Command, args []string) {
+			if err := c.InitializeManagerClient(cmd); err != nil {
+				printError(cmd, "Failed to connect to manager: %v ❌ ", c.connectErr)
+				return
+			}
+			defer c.Close()
+
+			createReq, err := loadCerts()
+			if err != nil {
+				printError(cmd, "Error loading certs: %v ❌ ", err)
+				return
+			}
+
+			createReq.AgentCvmServerUrl = agentCVMServerUrl
+			createReq.AgentLogLevel = agentLogLevel
+
+			cmd.Println("🔗 Creating a new virtual machine")
+
+			res, err := c.managerClient.CreateVm(cmd.Context(), createReq)
+			if err != nil {
+				printError(cmd, "Error creating virtual machine: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println(color.New(color.FgGreen).Sprintf("✅ Virtual machine created successfully with id %s and port %s", res.SvmId, res.ForwardedPort))
+		},
+	}
+
+	cmd.Flags().StringVar(&agentCVMServerUrl, serverURL, "", "CVM server URL")
+	cmd.Flags().StringVar(&agentCVMServerCA, serverCA, "", "CVM server CA")
+	cmd.Flags().StringVar(&agentCVMClientKey, clientKey, "", "CVM client key")
+	cmd.Flags().StringVar(&agentCVMClientCrt, clientCrt, "", "CVM client crt")
+	cmd.Flags().StringVar(&agentLogLevel, logLevel, "", "Agent Log level")
+
+	return cmd
+}
+
+func (c *CLI) NewRemoveVMCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:     "remove-vm",
+		Short:   "Remove a virtual machine",
+		Example: `remove-vm <svm_id>`,
+		Args:    cobra.ExactArgs(1),
+		Run: func(cmd *cobra.Command, args []string) {
+			if err := c.InitializeManagerClient(cmd); err == nil {
+				defer c.Close()
+			}
+
+			if c.connectErr != nil {
+				printError(cmd, "Failed to connect to manager: %v ❌ ", c.connectErr)
+				return
+			}
+
+			cmd.Println("🔗 Removing virtual machine")
+
+			_, err := c.managerClient.RemoveVm(cmd.Context(), &manager.RemoveReq{SvmId: args[0]})
+			if err != nil {
+				printError(cmd, "Error removing virtual machine: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println(color.New(color.FgGreen).Sprintf("✅ Virtual machine removed successfully"))
+		},
+	}
+}
+
+func fileReader(path string) ([]byte, error) {
+	if path == "" {
+		return nil, nil
+	}
+
+	return os.ReadFile(path)
+}
+
+func loadCerts() (*manager.CreateReq, error) {
+	clientKey, err := fileReader(agentCVMClientKey)
+	if err != nil {
+		return nil, err
+	}
+
+	clientCrt, err := fileReader(agentCVMClientCrt)
+	if err != nil {
+		return nil, err
+	}
+
+	serverCA, err := fileReader(agentCVMServerCA)
+	if err != nil {
+		return nil, err
+	}
+
+	return &manager.CreateReq{
+		AgentCvmServerCaCert: serverCA,
+		AgentCvmClientKey:    clientKey,
+		AgentCvmClientCert:   clientCrt,
+	}, nil
+}

cli/sdk.go

@@ -6,28 +6,36 @@ import (
 	"context"
 
 	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/manager"
+	"github.com/ultravioletrs/cocos/pkg/attestation/cmdconfig"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
+	managergrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/manager"
 	"github.com/ultravioletrs/cocos/pkg/sdk"
 )
 
 var Verbose bool
 
 type CLI struct {
-	agentSDK   sdk.SDK
-	config     grpc.AgentClientConfig
-	client     grpc.Client
-	connectErr error
+	agentSDK      sdk.SDK
+	agentConfig   grpc.AgentClientConfig
+	managerConfig grpc.ManagerClientConfig
+	client        grpc.Client
+	managerClient manager.ManagerServiceClient
+	connectErr    error
+	measurement   cmdconfig.MeasurementProvider
 }
 
-func New(config grpc.AgentClientConfig) *CLI {
+func New(agentConfig grpc.AgentClientConfig, managerConfig grpc.ManagerClientConfig, measurement cmdconfig.MeasurementProvider) *CLI {
 	return &CLI{
-		config: config,
+		agentConfig:   agentConfig,
+		managerConfig: managerConfig,
+		measurement:   measurement,
 	}
 }
 
-func (c *CLI) InitializeSDK(cmd *cobra.Command) error {
-	agentGRPCClient, agentClient, err := agent.NewAgentClient(context.Background(), c.config)
+func (c *CLI) InitializeAgentSDK(cmd *cobra.Command) error {
+	agentGRPCClient, agentClient, err := agent.NewAgentClient(context.Background(), c.agentConfig)
 	if err != nil {
 		c.connectErr = err
 		return err
@@ -39,6 +47,20 @@ func (c *CLI) InitializeSDK(cmd *cobra.Command) error {
 	return nil
 }
 
+func (c *CLI) InitializeManagerClient(cmd *cobra.Command) error {
+	managerGRPCClient, managerClient, err := managergrpc.NewManagerClient(c.managerConfig)
+	if err != nil {
+		c.connectErr = err
+		return err
+	}
+
+	cmd.Println("🔗 Connected to manager using ", managerGRPCClient.Secure())
+	c.client = managerGRPCClient
+
+	c.managerClient = managerClient
+	return nil
+}
+
 func (c *CLI) Close() {
 	c.client.Close()
 }

cmd/agent/main.go

@@ -3,157 +3,196 @@
 package main
 
 import (
-	"bytes"
 	"context"
-	"encoding/json"
+	"crypto/sha256"
+	"crypto/sha512"
+	"crypto/x509"
+	"encoding/pem"
+	"errors"
 	"fmt"
-	"io"
 	"log"
 	"log/slog"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	mglog "github.com/absmach/magistrala/logger"
 	"github.com/absmach/magistrala/pkg/prometheus"
-	"github.com/cenkalti/backoff/v4"
-	"github.com/google/go-sev-guest/abi"
+	"github.com/caarlos0/env/v11"
 	"github.com/google/go-sev-guest/client"
-	"github.com/mdlayher/vsock"
+	"github.com/stretchr/testify/mock"
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/agent/api"
-	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
-	"github.com/ultravioletrs/cocos/agent/auth"
+	"github.com/ultravioletrs/cocos/agent/cvms"
+	cvmsapi "github.com/ultravioletrs/cocos/agent/cvms/api/grpc"
+	"github.com/ultravioletrs/cocos/agent/cvms/server"
 	"github.com/ultravioletrs/cocos/agent/events"
 	agentlogger "github.com/ultravioletrs/cocos/internal/logger"
-	"github.com/ultravioletrs/cocos/internal/server"
-	grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
-	ackvsock "github.com/ultravioletrs/cocos/internal/vsock"
-	managerevents "github.com/ultravioletrs/cocos/manager/events"
-	"github.com/ultravioletrs/cocos/manager/qemu"
+	attestationconfig "github.com/ultravioletrs/cocos/pkg/attestation"
 	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
-	"golang.org/x/crypto/sha3"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider/mocks"
+	pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
+	cvmsgrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/cvm"
 	"golang.org/x/sync/errgroup"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/reflection"
 )
 
 const (
-	svcName        = "agent"
-	defSvcGRPCPort = "7002"
-	retryInterval  = 5 * time.Second
+	svcName          = "agent"
+	defSvcGRPCPort   = "7002"
+	retryInterval    = 5 * time.Second
+	envPrefixCVMGRPC = "AGENT_CVM_GRPC_"
+	storageDir       = "/var/lib/cocos/agent"
 )
 
+type config struct {
+	LogLevel      string `env:"AGENT_LOG_LEVEL" envDefault:"debug"`
+	Vmpl          int    `env:"AGENT_VMPL"      envDefault:"2"`
+	AgentGrpcHost string `env:"AGENT_GRPC_HOST" envDefault:"0.0.0.0"`
+}
+
 func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	g, ctx := errgroup.WithContext(ctx)
 
-	cfg, err := readConfig()
-	if err != nil {
-		log.Fatalf("failed to read agent configuration from vsock %s", err.Error())
+	var cfg config
+	if err := env.Parse(&cfg); err != nil {
+		log.Fatalf("failed to load %s configuration : %s", svcName, err)
 	}
 
-	conn, err := dialVsock()
-	if err != nil {
-		log.Fatal(err)
-	}
-	defer conn.Close()
-
-	ackConn := ackvsock.NewAckWriter(conn)
-
 	var exitCode int
 	defer mglog.ExitWithError(&exitCode)
 
 	var level slog.Level
-	if err := level.UnmarshalText([]byte(cfg.AgentConfig.LogLevel)); err != nil {
+	if err := level.UnmarshalText([]byte(cfg.LogLevel)); err != nil {
 		log.Println(err)
 		exitCode = 1
 		return
 	}
 
-	handler := agentlogger.NewProtoHandler(ackConn, &slog.HandlerOptions{Level: level}, cfg.ID)
+	eventsLogsQueue := make(chan *cvms.ClientStreamMessage, 1000)
+
+	handler := agentlogger.NewProtoHandler(os.Stdout, &slog.HandlerOptions{Level: level}, eventsLogsQueue)
 	logger := slog.New(handler)
 
-	eventSvc, err := events.New(svcName, cfg.ID, ackConn)
+	eventSvc, err := events.New(svcName, eventsLogsQueue)
 	if err != nil {
 		logger.Error(fmt.Sprintf("failed to create events service %s", err.Error()))
 		exitCode = 1
 		return
 	}
 
-	qp, err := quoteprovider.GetQuoteProvider()
-	if err != nil {
-		logger.Error(fmt.Sprintf("failed to create quote provider %s", err.Error()))
+	var qp client.LeveledQuoteProvider
+
+	if !sevGuesDeviceExists() {
+		logger.Info("SEV-SNP device not found")
+		qpMock := new(mocks.LeveledQuoteProvider)
+		qpMock.On("GetRawQuoteAtLevel", mock.Anything, mock.Anything).Return([]uint8{}, errors.New("SEV-SNP device not found"))
+		qp = qpMock
+	} else {
+		qp, err = quoteprovider.GetLeveledQuoteProvider()
+		if err != nil {
+			logger.Error(fmt.Sprintf("failed to create quote provider %s", err.Error()))
+			exitCode = 1
+			return
+		}
+	}
+
+	cvmGrpcConfig := pkggrpc.CVMClientConfig{}
+	if err := env.ParseWithOptions(&cvmGrpcConfig, env.Options{Prefix: envPrefixCVMGRPC}); err != nil {
+		logger.Error(fmt.Sprintf("failed to load %s gRPC client configuration : %s", svcName, err))
 		exitCode = 1
 		return
 	}
 
-	if err := verifyManifest(cfg, qp); err != nil {
+	cvmGRPCClient, cvmsClient, err := cvmsgrpc.NewCVMClient(cvmGrpcConfig)
+	if err != nil {
+		logger.Error(err.Error())
+		exitCode = 1
+		return
+	}
+	defer cvmGRPCClient.Close()
+
+	reconnectFn := func(ctx context.Context) (cvms.Service_ProcessClient, error) {
+		_, newClient, err := cvmsgrpc.NewCVMClient(cvmGrpcConfig)
+		if err != nil {
+			return nil, err
+		}
+		// Don't defer close here as we want to keep the connection open
+
+		return newClient.Process(ctx)
+	}
+
+	pc, err := cvmsClient.Process(ctx)
+	if err != nil {
 		logger.Error(err.Error())
 		exitCode = 1
 		return
 	}
 
-	setDefaultValues(&cfg)
-
-	svc := newService(ctx, logger, eventSvc, cfg, qp)
-
-	agentGrpcServerConfig := server.AgentConfig{
-		ServerConfig: server.ServerConfig{
-			BaseConfig: server.BaseConfig{
-				Host:         cfg.AgentConfig.Host,
-				Port:         cfg.AgentConfig.Port,
-				CertFile:     cfg.AgentConfig.CertFile,
-				KeyFile:      cfg.AgentConfig.KeyFile,
-				ServerCAFile: cfg.AgentConfig.ServerCAFile,
-				ClientCAFile: cfg.AgentConfig.ClientCAFile,
-			},
-		},
-		AttestedTLS: cfg.AgentConfig.AttestedTls,
-	}
-
-	registerAgentServiceServer := func(srv *grpc.Server) {
-		reflection.Register(srv)
-		agent.RegisterAgentServiceServer(srv, agentgrpc.NewServer(svc))
-	}
-
-	authSvc, err := auth.New(cfg)
-	if err != nil {
-		logger.Error(fmt.Sprintf("failed to create auth service %s", err.Error()))
+	if cfg.Vmpl < 0 || cfg.Vmpl > 3 {
+		logger.Error("vmpl level must be in a range [0, 3]")
 		exitCode = 1
 		return
 	}
 
-	gs := grpcserver.New(ctx, cancel, svcName, agentGrpcServerConfig, registerAgentServiceServer, logger, qp, authSvc)
+	svc := newService(ctx, logger, eventSvc, qp, cfg.Vmpl)
+
+	if err := os.MkdirAll(storageDir, 0o755); err != nil {
+		logger.Error(fmt.Sprintf("failed to create storage directory: %s", err))
+		exitCode = 1
+		return
+	}
+
+	mc, err := cvmsapi.NewClient(pc, svc, eventsLogsQueue, logger, server.NewServer(logger, svc, cfg.AgentGrpcHost), storageDir, reconnectFn)
+	if err != nil {
+		logger.Error(err.Error())
+		exitCode = 1
+		return
+	}
 
 	g.Go(func() error {
-		for {
-			if _, err := io.Copy(io.Discard, conn); err != nil {
-				log.Printf("vsock connection lost: %v, reconnecting...", err)
-				conn.Close()
-				conn, err = dialVsock()
-				if err != nil {
-					log.Fatal("failed to reconnect: ", err)
-				}
-			}
-			time.Sleep(retryInterval)
+		ch := make(chan os.Signal, 1)
+		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
+		defer signal.Stop(ch)
+
+		select {
+		case <-ch:
+			logger.Info("Received signal, shutting down...")
+			cancel()
+			return nil
+		case <-ctx.Done():
+			return ctx.Err()
 		}
 	})
 
 	g.Go(func() error {
-		return gs.Start()
+		return mc.Process(ctx, cancel)
 	})
 
-	g.Go(func() error {
-		return server.StopHandler(ctx, cancel, logger, svcName, gs)
-	})
+	attestation, certSerialNumber, err := attestationFromCert(ctx, cvmGrpcConfig.ClientCert, svc)
+	if err != nil {
+		logger.Error(fmt.Sprintf("failed to get attestation: %s", err))
+		exitCode = 1
+		return
+	}
+
+	eventsLogsQueue <- &cvms.ClientStreamMessage{
+		Message: &cvms.ClientStreamMessage_VTPMattestationReport{
+			VTPMattestationReport: &cvms.AttestationResponse{
+				File:             attestation,
+				CertSerialNumber: certSerialNumber,
+			},
+		},
+	}
 
 	if err := g.Wait(); err != nil {
 		logger.Error(fmt.Sprintf("%s service terminated: %s", svcName, err))
 	}
 }
 
-func newService(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp agent.Computation, qp client.QuoteProvider) agent.Service {
-	svc := agent.New(ctx, logger, eventSvc, cmp, qp)
+func newService(ctx context.Context, logger *slog.Logger, eventSvc events.Service, qp client.LeveledQuoteProvider, vmpl int) agent.Service {
+	svc := agent.New(ctx, logger, eventSvc, qp, vmpl)
 
 	svc = api.LoggingMiddleware(svc, logger)
 	counter, latency := prometheus.MakeMetrics(svcName, "api")
@@ -162,103 +201,37 @@ func newService(ctx context.Context, logger *slog.Logger, eventSvc events.Servic
 	return svc
 }
 
-func readConfig() (agent.Computation, error) {
-	l, err := vsock.Listen(qemu.VsockConfigPort, nil)
+func sevGuesDeviceExists() bool {
+	d, err := client.OpenDevice()
 	if err != nil {
-		return agent.Computation{}, err
+		return false
 	}
-	defer l.Close()
-
-	conn, err := l.Accept()
-	if err != nil {
-		return agent.Computation{}, err
-	}
-	defer conn.Close()
-
-	var buffer []byte
-	for {
-		chunk := make([]byte, 1024)
-		n, err := conn.Read(chunk)
-		if err != nil {
-			if err == io.EOF {
-				break
-			}
-			return agent.Computation{}, err
-		}
-		buffer = append(buffer, chunk[:n]...)
-	}
-
-	ac := agent.Computation{
-		AgentConfig: agent.AgentConfig{},
-	}
-	if err := json.Unmarshal(buffer, &ac); err != nil {
-		return agent.Computation{}, err
-	}
-	return ac, nil
+	d.Close()
+	return true
 }
 
-func setDefaultValues(cfg *agent.Computation) {
-	if cfg.AgentConfig.LogLevel == "" {
-		cfg.AgentConfig.LogLevel = "info"
+func attestationFromCert(ctx context.Context, certFilePath string, svc agent.Service) ([]byte, string, error) {
+	if certFilePath == "" {
+		return nil, "", nil
 	}
-	if cfg.AgentConfig.Port == "" {
-		cfg.AgentConfig.Port = defSvcGRPCPort
-	}
-}
 
-func isTEE() bool {
-	_, err := os.Stat("/dev/sev-guest")
-	return !os.IsNotExist(err)
-}
-
-func dialVsock() (*vsock.Conn, error) {
-	var conn *vsock.Conn
-	var err error
-
-	err = backoff.Retry(func() error {
-		conn, err = vsock.Dial(vsock.Host, managerevents.ManagerVsockPort, nil)
-		if err == nil {
-			log.Println("vsock connection established")
-			return nil
-		}
-		log.Printf("vsock connection failed, retrying in %s... Error: %v", retryInterval, err)
-		return err
-	}, backoff.NewExponentialBackOff())
+	certFile, err := os.ReadFile(certFilePath)
 	if err != nil {
-		return nil, err
+		return nil, "", err
 	}
 
-	return conn, nil
-}
-
-func verifyManifest(cfg agent.Computation, qp client.QuoteProvider) error {
-	if !isTEE() {
-		return nil
-	}
-
-	ar, err := qp.GetRawQuote(sha3.Sum512([]byte(cfg.ID)))
-	if err != nil {
-		return err
-	}
-
-	arProto, err := abi.ReportCertsToProto(ar[:abi.ReportSize])
-	if err != nil {
-		return err
-	}
-
-	cfgBytes, err := json.Marshal(cfg)
-	if err != nil {
-		return err
-	}
-
-	mcHash := sha3.Sum256(cfgBytes)
-
-	if arProto.Report.HostData == nil {
-		return fmt.Errorf("manifest verification failed: HostData is nil")
-	}
-	if !bytes.Equal(arProto.Report.HostData, mcHash[:]) {
-		return fmt.Errorf("manifest verification failed")
-	}
-
-	return nil
+	certPem, _ := pem.Decode(certFile)
+	certx509, err := x509.ParseCertificate(certPem.Bytes)
+	if err != nil {
+		return nil, "", err
+	}
+
+	nonceSNP := sha512.Sum512(certFile)
+	nonceVTPM := sha256.Sum256(certFile)
+	attestation, err := svc.Attestation(ctx, nonceSNP, nonceVTPM, attestationconfig.SNPvTPM)
+	if err != nil {
+		return nil, "", err
+	}
+
+	return attestation, certx509.SerialNumber.String(), nil
 }

cmd/agent/main_test.go

@@ -1,94 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package main
-
-import (
-	"context"
-	"log/slog"
-	"os"
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
-	"github.com/ultravioletrs/cocos/agent"
-	"github.com/ultravioletrs/cocos/agent/events/mocks"
-	qpmocks "github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider/mocks"
-)
-
-func TestSetDefaultValues(t *testing.T) {
-	tests := []struct {
-		name     string
-		input    agent.Computation
-		expected agent.Computation
-	}{
-		{
-			name: "Empty config",
-			input: agent.Computation{
-				AgentConfig: agent.AgentConfig{},
-			},
-			expected: agent.Computation{
-				AgentConfig: agent.AgentConfig{
-					LogLevel: "info",
-					Port:     "7002",
-				},
-			},
-		},
-		{
-			name: "Partial config",
-			input: agent.Computation{
-				AgentConfig: agent.AgentConfig{
-					LogLevel: "debug",
-				},
-			},
-			expected: agent.Computation{
-				AgentConfig: agent.AgentConfig{
-					LogLevel: "debug",
-					Port:     "7002",
-				},
-			},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			setDefaultValues(&tt.input)
-			assert.Equal(t, tt.expected, tt.input)
-		})
-	}
-}
-
-func TestNewService(t *testing.T) {
-	ctx := context.Background()
-	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
-	eventSvc := new(mocks.Service)
-	eventSvc.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
-	cmp := agent.Computation{
-		ID: "test-computation",
-		AgentConfig: agent.AgentConfig{
-			LogLevel: "info",
-			Port:     "7002",
-		},
-	}
-	qp := new(qpmocks.QuoteProvider)
-
-	svc := newService(ctx, logger, eventSvc, cmp, qp)
-
-	assert.NotNil(t, svc)
-}
-
-func TestVerifyManifest(t *testing.T) {
-	cfg := agent.Computation{
-		ID: "test-computation",
-		AgentConfig: agent.AgentConfig{
-			LogLevel: "info",
-			Port:     "7002",
-		},
-	}
-
-	mockQP := new(qpmocks.QuoteProvider)
-	mockQP.On("GetRawQuote", mock.Anything).Return([]byte{}, nil)
-
-	err := verifyManifest(cfg, mockQP)
-
-	assert.NoError(t, err)
-}

cmd/cli/main.go

@@ -14,20 +14,23 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"github.com/ultravioletrs/cocos/cli"
+	"github.com/ultravioletrs/cocos/pkg/attestation/cmdconfig"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	cmd "github.com/virtee/sev-snp-measure-go/sevsnpmeasure/cmd"
 )
 
 const (
-	svcName            = "cli"
-	envPrefixAgentGRPC = "AGENT_GRPC_"
-	completion         = "completion"
-	filePermision      = 0o755
-	cocosDirectory     = ".cocos"
+	svcName              = "cli"
+	envPrefixAgentGRPC   = "AGENT_GRPC_"
+	envPrefixManagerGRPC = "MANAGER_GRPC_"
+	completion           = "completion"
+	filePermision        = 0o755
+	cocosDirectory       = ".cocos"
 )
 
 type config struct {
-	LogLevel string `env:"AGENT_LOG_LEVEL" envDefault:"info"`
+	LogLevel       string `env:"AGENT_LOG_LEVEL" envDefault:"info"`
+	IgvmBinaryPath string `env:"IGVM_BINARY_PATH" envDefault:"./build/igvmmeasure"`
 }
 
 func main() {
@@ -98,9 +101,24 @@ func main() {
 		return
 	}
 
-	cliSVC := cli.New(agentGRPCConfig)
+	managerGRPCConfig := grpc.ManagerClientConfig{}
+	if err := env.ParseWithOptions(&managerGRPCConfig, env.Options{Prefix: envPrefixManagerGRPC}); err != nil {
+		message := color.New(color.FgRed).Sprintf("failed to load %s gRPC client configuration : %s", svcName, err)
+		rootCmd.Println(message)
+		return
+	}
 
-	if err := cliSVC.InitializeSDK(rootCmd); err == nil {
+	options := cmdconfig.IgvmMeasureOptions
+	measurement, err := cmdconfig.NewCmdConfig(cfg.IgvmBinaryPath, options, os.Stderr)
+	if err != nil {
+		message := color.New(color.FgRed).Sprintf("failed to initialize measurement: %s", err) // Use %s instead of %w
+		rootCmd.Println(message)
+		return
+	}
+
+	cliSVC := cli.New(agentGRPCConfig, managerGRPCConfig, measurement)
+
+	if err := cliSVC.InitializeAgentSDK(rootCmd); err == nil {
 		defer cliSVC.Close()
 	}
 
@@ -119,6 +137,8 @@ func main() {
 	rootCmd.AddCommand(attestationPolicyCmd)
 	rootCmd.AddCommand(keysCmd)
 	rootCmd.AddCommand(cliSVC.NewCABundleCmd(directoryCachePath))
+	rootCmd.AddCommand(cliSVC.NewCreateVMCmd())
+	rootCmd.AddCommand(cliSVC.NewRemoveVMCmd())
 
 	// Attestation commands
 	attestationCmd.AddCommand(cliSVC.NewGetAttestationCmd())
@@ -126,6 +146,7 @@ func main() {
 
 	// measure.
 	rootCmd.AddCommand(cmd.NewRootCmd())
+	rootCmd.AddCommand(cliSVC.NewMeasureCmd(cfg.IgvmBinaryPath))
 
 	// Flags
 	keysCmd.PersistentFlags().StringVarP(
@@ -139,6 +160,8 @@ func main() {
 	// Attestation Policy commands
 	attestationPolicyCmd.AddCommand(cliSVC.NewAddMeasurementCmd())
 	attestationPolicyCmd.AddCommand(cliSVC.NewAddHostDataCmd())
+	attestationPolicyCmd.AddCommand(cliSVC.NewGCPAttestationPolicy())
+	attestationPolicyCmd.AddCommand(cliSVC.NewDownloadGCPOvmfFile())
 
 	if err := rootCmd.Execute(); err != nil {
 		logErrorCmd(*rootCmd, err)

cmd/manager/main.go

@@ -10,25 +10,24 @@ import (
 	"log/slog"
 	"net/url"
 	"os"
-	"os/signal"
 	"strings"
-	"syscall"
 
 	mglog "github.com/absmach/magistrala/logger"
 	"github.com/absmach/magistrala/pkg/jaeger"
 	"github.com/absmach/magistrala/pkg/prometheus"
 	"github.com/absmach/magistrala/pkg/uuid"
 	"github.com/caarlos0/env/v11"
+	"github.com/ultravioletrs/cocos/internal/server"
+	grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
 	"github.com/ultravioletrs/cocos/manager"
 	"github.com/ultravioletrs/cocos/manager/api"
-	managerapi "github.com/ultravioletrs/cocos/manager/api/grpc"
-	"github.com/ultravioletrs/cocos/manager/events"
+	managergrpc "github.com/ultravioletrs/cocos/manager/api/grpc"
 	"github.com/ultravioletrs/cocos/manager/qemu"
 	"github.com/ultravioletrs/cocos/manager/tracing"
-	pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
-	managergrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/manager"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/reflection"
 )
 
 const (
@@ -43,7 +42,9 @@ type config struct {
 	JaegerURL               url.URL `env:"COCOS_JAEGER_URL"                   envDefault:"http://localhost:4318"`
 	TraceRatio              float64 `env:"COCOS_JAEGER_TRACE_RATIO"           envDefault:"1.0"`
 	InstanceID              string  `env:"MANAGER_INSTANCE_ID"                envDefault:""`
-	AttestationPolicyBinary string  `env:"MANAGER_ATTESTATION_POLICY_BINARY"  envDefault:"../../build"`
+	AttestationPolicyBinary string  `env:"MANAGER_ATTESTATION_POLICY_BINARY"  envDefault:"../../build/attestation_policy"`
+	IgvmMeasureBinary       string  `env:"MANAGER_IGVMMEASURE_BINARY"         envDefault:"../../build/igvmmeasure"`
+	PcrValues               string  `env:"MANAGER_PCR_VALUES"                 envDefault:""`
 	EosVersion              string  `env:"MANAGER_EOS_VERSION"                envDefault:""`
 }
 
@@ -92,64 +93,33 @@ func main() {
 	args := qemuCfg.ConstructQemuArgs()
 	logger.Info(strings.Join(args, " "))
 
-	managerGRPCConfig := pkggrpc.ManagerClientConfig{}
+	managerGRPCConfig := server.ServerConfig{}
 	if err := env.ParseWithOptions(&managerGRPCConfig, env.Options{Prefix: envPrefixGRPC}); err != nil {
 		logger.Error(fmt.Sprintf("failed to load %s gRPC client configuration : %s", svcName, err))
 		exitCode = 1
 		return
 	}
 
-	managerGRPCClient, managerClient, err := managergrpc.NewManagerClient(managerGRPCConfig)
-	if err != nil {
-		logger.Error(err.Error())
-		exitCode = 1
-		return
-	}
-	defer managerGRPCClient.Close()
-
-	pc, err := managerClient.Process(ctx)
+	svc, err := newService(logger, tracer, qemuCfg, cfg.AttestationPolicyBinary, cfg.IgvmMeasureBinary, cfg.PcrValues, cfg.EosVersion)
 	if err != nil {
 		logger.Error(err.Error())
 		exitCode = 1
 		return
 	}
 
-	eventsChan := make(chan *manager.ClientStreamMessage, clientBufferSize)
-	svc, err := newService(logger, tracer, qemuCfg, eventsChan, cfg.AttestationPolicyBinary, cfg.EosVersion)
-	if err != nil {
-		logger.Error(err.Error())
-		exitCode = 1
-		return
+	registerManagerServiceServer := func(srv *grpc.Server) {
+		reflection.Register(srv)
+		manager.RegisterManagerServiceServer(srv, managergrpc.NewServer(svc))
 	}
 
-	eventsSvc, err := events.New(logger, svc.ReportBrokenConnection, eventsChan)
-	if err != nil {
-		logger.Error(err.Error())
-		exitCode = 1
-		return
-	}
-
-	go eventsSvc.Listen(ctx)
-
-	mc := managerapi.NewClient(pc, svc, eventsChan, logger)
+	gs := grpcserver.New(ctx, cancel, svcName, managerGRPCConfig, registerManagerServiceServer, logger, nil, nil)
 
 	g.Go(func() error {
-		ch := make(chan os.Signal, 1)
-		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
-		defer signal.Stop(ch)
-
-		select {
-		case <-ch:
-			logger.Info("Received signal, shutting down...")
-			cancel()
-			return nil
-		case <-ctx.Done():
-			return ctx.Err()
-		}
+		return gs.Start()
 	})
 
 	g.Go(func() error {
-		return mc.Process(ctx, cancel)
+		return server.StopHandler(ctx, cancel, logger, svcName, gs)
 	})
 
 	if err := g.Wait(); err != nil {
@@ -157,8 +127,8 @@ func main() {
 	}
 }
 
-func newService(logger *slog.Logger, tracer trace.Tracer, qemuCfg qemu.Config, eventsChan chan *manager.ClientStreamMessage, attestationPolicyPath string, eosVersion string) (manager.Service, error) {
-	svc, err := manager.New(qemuCfg, attestationPolicyPath, logger, eventsChan, qemu.NewVM, eosVersion)
+func newService(logger *slog.Logger, tracer trace.Tracer, qemuCfg qemu.Config, attestationPolicyPath string, igvmMeasurementBinaryPath string, pcrValuesFilePath string, eosVersion string) (manager.Service, error) {
+	svc, err := manager.New(qemuCfg, attestationPolicyPath, igvmMeasurementBinaryPath, pcrValuesFilePath, logger, qemu.NewVM, eosVersion)
 	if err != nil {
 		return nil, err
 	}

cocos-manager.env

@@ -7,10 +7,13 @@ COCOS_JAEGER_TRACE_RATIO=1.0
 # Manager Service Configuration
 MANAGER_INSTANCE_ID=
 MANAGER_ATTESTATION_POLICY_BINARY=../../build
+MANAGER_IGVMMEASURE_BINARY=../../build
+MANAGER_PCR_VALUES=/etc/cocos/pcr_values.json
 MANAGER_GRPC_CLIENT_CERT=
 MANAGER_GRPC_CLIENT_KEY=
 MANAGER_GRPC_SERVER_CA_CERTS=
-MANAGER_GRPC_URL=localhost:7001
+MANAGER_GRPC_PORT=6101
+MANAGER_GRPC_HOST=0.0.0.0
 MANAGER_GRPC_TIMEOUT=60s
 MANAGER_EOS_VERSION=""
 
@@ -21,13 +24,13 @@ MANAGER_QEMU_MAX_MEMORY=30G
 MANAGER_QEMU_OVMF_CODE_IF=pflash
 MANAGER_QEMU_OVMF_CODE_FORMAT=raw
 MANAGER_QEMU_OVMF_CODE_UNIT=0
-MANAGER_QEMU_OVMF_CODE_FILE=/usr/share/OVMF/x64/OVMF_CODE.fd
+MANAGER_QEMU_OVMF_CODE_FILE=/usr/share/edk2/x64/OVMF_CODE.fd
 MANAGER_QEMU_OVMF_VERSION=edk2-stable202408
 MANAGER_QEMU_OVMF_CODE_READONLY=on
 MANAGER_QEMU_OVMF_VARS_IF=pflash
 MANAGER_QEMU_OVMF_VARS_FORMAT=raw
 MANAGER_QEMU_OVMF_VARS_UNIT=1
-MANAGER_QEMU_OVMF_VARS_FILE=/usr/share/OVMF/x64/OVMF_VARS.fd
+MANAGER_QEMU_OVMF_VARS_FILE=/usr/share/edk2/x64/OVMF_VARS.fd
 MANAGER_QEMU_NETDEV_ID=vmnic
 MANAGER_QEMU_HOST_FWD_AGENT=7020
 MANAGER_QEMU_GUEST_FWD_AGENT=7002
@@ -35,8 +38,8 @@ MANAGER_QEMU_VIRTIO_NET_PCI_DISABLE_LEGACY=on
 MANAGER_QEMU_VIRTIO_NET_PCI_IOMMU_PLATFORM=true
 MANAGER_QEMU_VIRTIO_NET_PCI_ADDR=0x2
 MANAGER_QEMU_VIRTIO_NET_PCI_ROMFILE=
-MANAGER_QEMU_DISK_IMG_KERNEL_FILE=/home/sammyk/Documents/cocos-ai/cmd/manager/img/bzImage
-MANAGER_QEMU_DISK_IMG_ROOTFS_FILE=/home/sammyk/Documents/cocos-ai/cmd/manager/img/rootfs.cpio.gz
+MANAGER_QEMU_DISK_IMG_KERNEL_FILE=/etc/cocos/bzImage
+MANAGER_QEMU_DISK_IMG_ROOTFS_FILE=/etc/cocos/rootfs.cpio.gz
 MANAGER_QEMU_SEV_ID=sev0
 MANAGER_QEMU_SEV_CBITPOS=51
 MANAGER_QEMU_SEV_REDUCED_PHYS_BITS=1
@@ -48,13 +51,13 @@ MANAGER_QEMU_BIN_PATH=qemu-system-x86_64
 MANAGER_QEMU_USE_SUDO=true
 MANAGER_QEMU_ENABLE_SEV=false
 MANAGER_QEMU_ENABLE_SEV_SNP=false
+MANAGER_QEMU_IGVM_FILE=/etc/cocos/coconut-qemu.igvm
 MANAGER_QEMU_ENABLE_KVM=true
 MANAGER_QEMU_MACHINE=q35
 MANAGER_QEMU_CPU=EPYC
 MANAGER_QEMU_SMP_COUNT=4
 MANAGER_QEMU_SMP_MAXCPUS=16
 MANAGER_QEMU_MEM_ID=ram1
-MANAGER_QEMU_KERNEL_HASH=false
 MANAGER_QEMU_NO_GRAPHIC=true
 MANAGER_QEMU_MONITOR=pty
 MANAGER_QEMU_HOST_FWD_RANGE=6100-6200

go.mod

@@ -1,37 +1,61 @@
 module github.com/ultravioletrs/cocos
 
-go 1.23.0
+go 1.23.2
 
 require (
 	github.com/absmach/magistrala v0.15.1
-	github.com/caarlos0/env/v11 v11.2.2
-	github.com/cenkalti/backoff/v4 v4.3.0
+	github.com/caarlos0/env/v11 v11.3.1
 	github.com/fatih/color v1.18.0
 	github.com/go-kit/kit v0.13.0
 	github.com/gofrs/uuid v4.4.0+incompatible
-	github.com/google/go-sev-guest v0.11.1
+	github.com/google/go-sev-guest v0.13.0
+	github.com/google/go-tdx-guest v0.3.2-0.20241009005452-097ee70d0843 // indirect
 	github.com/mdlayher/vsock v1.2.1
-	github.com/spf13/cobra v1.8.1
-	github.com/spf13/pflag v1.0.5
+	github.com/spf13/cobra v1.9.1
+	github.com/spf13/pflag v1.0.6
 	github.com/stretchr/testify v1.10.0
 	github.com/virtee/sev-snp-measure-go v0.0.0-20240530153610-e6e8dc9b6877
-	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.57.0
-	go.opentelemetry.io/otel/trace v1.32.0
-	golang.org/x/crypto v0.30.0
-	golang.org/x/sync v0.10.0
-	google.golang.org/grpc v1.68.1
-	google.golang.org/protobuf v1.35.2
+	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.59.0
+	go.opentelemetry.io/otel/trace v1.34.0
+	golang.org/x/crypto v0.35.0
+	golang.org/x/sync v0.12.0
+	google.golang.org/grpc v1.71.0
+	google.golang.org/protobuf v1.36.5
 )
 
 require (
+	cel.dev/expr v0.19.2 // indirect
+	cloud.google.com/go v0.118.3 // indirect
+	cloud.google.com/go/auth v0.15.0 // indirect
+	cloud.google.com/go/auth/oauth2adapt v0.2.7 // indirect
+	cloud.google.com/go/compute/metadata v0.6.0 // indirect
+	cloud.google.com/go/iam v1.4.1 // indirect
+	cloud.google.com/go/monitoring v1.24.0 // indirect
+	cloud.google.com/go/storage v1.51.0 // indirect
+	github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.25.0 // indirect
+	github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.51.0 // indirect
+	github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.51.0 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
+	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
+	github.com/cncf/xds/go v0.0.0-20250121191232-2f005788dc42 // indirect
 	github.com/containerd/log v0.1.0 // indirect
 	github.com/distribution/reference v0.6.0 // indirect
 	github.com/docker/go-connections v0.5.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
+	github.com/envoyproxy/go-control-plane/envoy v1.32.4 // indirect
+	github.com/envoyproxy/protoc-gen-validate v1.2.1 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/gofrs/uuid/v5 v5.3.0 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang/protobuf v1.5.4 // indirect
+	github.com/google/certificate-transparency-go v1.1.2 // indirect
+	github.com/google/gce-tcb-verifier v0.3.1 // indirect
+	github.com/google/go-attestation v0.5.1 // indirect
+	github.com/google/go-eventlog v0.0.2-0.20241003021507-01bb555f7cba // indirect
+	github.com/google/go-tspi v0.3.0 // indirect
+	github.com/google/s2a-go v0.1.9 // indirect
+	github.com/googleapis/enterprise-certificate-proxy v0.3.5 // indirect
+	github.com/googleapis/gax-go/v2 v2.14.1 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/moby/docker-image-spec v1.3.1 // indirect
@@ -39,11 +63,20 @@ require (
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/opencontainers/go-digest v1.0.0 // indirect
 	github.com/opencontainers/image-spec v1.1.0 // indirect
-	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.57.0 // indirect
-	go.opentelemetry.io/otel v1.32.0 // indirect
+	github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
+	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
+	go.opentelemetry.io/contrib/detectors/gcp v1.34.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.59.0 // indirect
+	go.opentelemetry.io/otel v1.34.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.32.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.32.0 // indirect
-	go.opentelemetry.io/otel/sdk v1.32.0 // indirect
+	go.opentelemetry.io/otel/sdk v1.34.0 // indirect
+	go.opentelemetry.io/otel/sdk/metric v1.34.0 // indirect
+	golang.org/x/exp v0.0.0-20240909161429-701f63a606c0 // indirect
+	golang.org/x/oauth2 v0.28.0 // indirect
+	golang.org/x/time v0.10.0 // indirect
+	google.golang.org/api v0.224.0 // indirect
+	google.golang.org/genproto v0.0.0-20250303144028-a0af3efb3deb // indirect
 	gotest.tools/v3 v3.5.1 // indirect
 )
 
@@ -51,19 +84,19 @@ require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
-	github.com/docker/docker v27.4.0+incompatible
+	github.com/docker/docker v28.0.1+incompatible
 	github.com/go-kit/log v0.2.1 // indirect
 	github.com/go-logfmt/logfmt v0.6.0 // indirect
 	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
-	github.com/golang/protobuf v1.5.4 // indirect
-	github.com/google/go-configfs-tsm v0.2.2 // indirect
+	github.com/google/go-configfs-tsm v0.3.3-0.20240919001351-b4b5b84fdcbc // indirect
+	github.com/google/go-tpm v0.9.3
+	github.com/google/go-tpm-tools v0.4.4
 	github.com/google/logger v1.1.1
-	github.com/google/uuid v1.6.0 // indirect
-	github.com/grpc-ecosystem/grpc-gateway/v2 v2.23.0 // indirect
+	github.com/google/uuid v1.6.0
+	github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.1 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
 	github.com/mdlayher/socket v0.4.1 // indirect
-	github.com/pborman/uuid v1.2.1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/prometheus/client_golang v1.20.5 // indirect
@@ -71,16 +104,18 @@ require (
 	github.com/prometheus/common v0.59.1 // indirect
 	github.com/prometheus/procfs v0.15.1 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
-	go.opentelemetry.io/otel/metric v1.32.0 // indirect
+	go.opentelemetry.io/otel/metric v1.34.0 // indirect
 	go.opentelemetry.io/proto/otlp v1.3.1 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	golang.org/x/net v0.31.0 // indirect
-	golang.org/x/sys v0.28.0 // indirect
-	golang.org/x/term v0.27.0
-	golang.org/x/text v0.21.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20241104194629-dd2ea8efbc28 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20241104194629-dd2ea8efbc28 // indirect
+	golang.org/x/net v0.35.0 // indirect
+	golang.org/x/sys v0.30.0 // indirect
+	golang.org/x/term v0.29.0
+	golang.org/x/text v0.22.0 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20250303144028-a0af3efb3deb // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20250303144028-a0af3efb3deb // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
-replace github.com/virtee/sev-snp-measure-go => github.com/sammyoina/sev-snp-measure-go v0.0.0-20241107163739-38915ab517c7
+replace github.com/virtee/sev-snp-measure-go => github.com/sammyoina/sev-snp-measure-go v0.0.0-20241202151803-ef189f0ff825
+
+replace github.com/google/go-tpm-tools => github.com/danko-miladinovic/go-tpm-tools v0.0.0-20250228160324-1ebcfd79567c

hal/cloud/.env

@@ -0,0 +1,89 @@
+####  memory config
+MEMORY_SIZE=2048M
+MEMORY_SLOTS=5
+MAX_MEMORY=30G
+
+
+####  ovmf code config
+OVMF_CODE_IF=pflash
+OVMF_CODE_FORMAT=raw
+OVMF_CODE_UNIT=0
+OVMF_CODE_FILE=/usr/share/OVMF/OVMF_CODE.fd
+OVMF_CODE_READONLY=on
+OVMF_VERSION=
+
+####  ovmf vars config
+OVMF_VARS_IF=pflash
+OVMF_VARS_FORMAT=raw
+OVMF_VARS_UNIT=1
+OVMF_VARS_FILE=/usr/share/OVMF/OVMF_VARS.fd
+
+####  net dev config
+NET_DEV_ID=vmnic
+NET_DEV_HOST_FWD_AGENT=7020
+NET_DEV_GUEST_FWD_AGENT=7002
+
+####  Virtio Net Pci Config
+VIRTIO_NET_PCI_DISABLE_LEGACY=on
+VIRTIO_NET_PCI_IOMMU_PLATFORM=true
+VIRTIO_NET_PCI_ADDR=0x2
+VIRTIO_NET_PCI_ROMFILE=
+
+####  Disk image config
+DISK_IMG_KERNEL_FILE=
+DISK_IMG_ROOTFS_FILE=
+
+KERNEL_COMMAND_LINE="quiet console=null"
+
+####  Sev Config
+SEV_ID=sev0
+SEV_CBIT_POS=51
+SEV_REDUCED_PHYS_BITS=1
+SEV_HOST_DATA=
+
+#### VSock Config
+VSOCK_ID=vhost-vsock-pci0
+VSOCK_GUEST_CID=3
+
+BIN_PATH=qemu-system-x86_64
+USE_SUDO=false
+ENABLE_SEV=false
+ENABLE_SEV_SNP=false
+
+ENABLE_KVM=true
+
+MACHINE=q35
+CPU=EPYC
+SMP_COUNT=8
+SMP_MAXCPUS=64
+MEM_ID=ram1
+
+KERNEL_HASH=false
+
+NO_GRAPHIC=true
+MONITOR=pty
+HOST_FWD_RANGE=6100-6200
+CERTS_MOUNT=/etc/cocos/certs
+ENV_MOUNT=/etc/cocos/environment
+
+COCOS_AGENT_VERSION=v0.3.1
+
+#### Base image URL and names
+BASE_IMAGE_URL=https://cloud-images.ubuntu.com/noble/current/noble-server-cloudimg-amd64.img
+BASE_IMAGE=ubuntu-base.qcow2
+CUSTOM_IMAGE=ubuntu-custom.qcow2
+
+#### Paths for OVMF firmware
+OVMF_CODE=/usr/share/ovmf/x64/OVMF_CODE.4m.fd
+OVMF_VARS=/usr/share/ovmf/x64/OVMF_VARS.4m.fd
+
+#### VM parameters
+VM_NAME=cocos-vm
+RAM=16G
+DISK_SIZE=10G # Size for root filesystem
+QEMU_BINARY=qemu-system-x86_64
+
+AGENT_GRPC_SERVER_CERT=/etc/cocos/certs/server.pem
+AGENT_GRPC_SERVER_KEY=/etc/cocos/certs/key.pem
+AGENT_GRPC_SERVER_CA_CERTS=/etc/cocos/ca.pem
+AGENT_GRPC_CLIENT_CA_CERTS=/etc/cocos/ca.pem

hal/cloud/README.md

@@ -0,0 +1,114 @@
+# Agent Cloud Init Setup
+
+## Overview
+
+The `hal/cloud` directory contains essential files required for setting up a virtual machine (VM) with cloud-init. This setup ensures the automated installation of dependencies, configuration of the environment, and deployment of the Cocos agent as a systemd service.
+
+### Directory Contents
+
+- **`config.yaml`**: This YAML file provides configuration instructions for the cloud image.
+- **`meta-data`**: Contains VM metadata, such as instance-specific details and identifiers.
+- **`qemu.sh`**: A Bash script for downloading and configuring a cloud image, running QEMU to simulate a VM with the cloud-init configuration.
+- **`.env`**: Contains environment variables for starting the VM in different modes, configuring disk space, memory allocation, and other parameters.
+
+## Configuration
+
+### Preparing the Cloud-Config File
+
+The `config.yaml` file defines system configurations, including user creation, package installations, file management, and command execution.
+Ensure that the cloud-config file is set up with the following configurations:
+
+- **User Credentials**: Specify the default username and password.
+- **Certificates and Keys**: Certificate files for agent for secure communication.
+- **Environment Variables**: Configuration parameters required by the system.
+
+The `config.yaml` file is divided into multiple sections, each addressing a specific aspect of the setup process.
+
+### 1. User Configuration
+
+This section creates a default user with specific permissions and configurations:
+
+- Creates a user named **`cocos_user`**.
+- Adds `cocos_user` to the `sudo` and `docker` groups.
+- Sets a default password (should be changed for production use).
+- Configures the user’s shell as `/bin/bash`.
+
+### 2. Package Installation
+
+Installs essential system packages required for various operations:
+
+- **`curl`**: For downloading files from the web.
+- **`make`**: A utility for building software.
+- **`git`**: Version control system for managing code repositories.
+- **`python3` and `python3-dev`**: Required for running Python-based tools.
+- **`net-tools`**: Provides networking utilities such as `ifconfig` and `route`.
+
+### 3. File Management (write_files)
+
+Creates and configures critical files required for the setup:
+
+- **Certificates**: Cert files (`cert.pem`, `ca.pem`, `key.pem`) located at `/etc/cocos/certs/`.
+- **Environment Variables**: An env file stored at `/etc/cocos/environment`.
+- **Systemd Service File**: Cocos agent service configuration file at `/etc/systemd/system/cocos-agent.service` for managing the Cocos agent.
+- **Agent Scripts**:
+  - `agent_setup.sh`: Configures network interfaces and resizes the root filesystem.
+  - `agent_start_script.sh`: Sets up Docker and starts the Cocos agent.
+
+### 4. Execution of Commands (runcmd)
+
+A sequence of commands is executed to finalize the setup:
+
+- Creates necessary directories: `/cocos`, `/cocos_init`, `/var/log/cocos`, `/etc/cocos`.
+- Downloads and installs the Cocos agent binary.
+- Installs **Wasmtime** and configures its environment variables.
+- Installs **Docker** and adds `cocos_user` to the Docker group.
+- Reloads systemd and enables the Cocos agent service.
+
+## Running the Agent
+
+To test the cloud-init configuration, execute the `qemu.sh` script to bring up a VM using QEMU:
+
+```bash
+sudo ./qemu.sh
+```
+
+**Important:** The script must be executed as root.
+
+Once the QEMU boots the VM, the Cocos agent will run as a systemd service. The service is configured to start automatically on boot and restart in case of failure.
+
+## Debugging and Monitoring
+
+For troubleshooting and monitoring the Cocos agent service, use the following commands within the VM:
+
+### Manually Start the Service
+
+To manually start the agent service, execute:
+
+```bash
+sudo systemctl start cocos-agent.service
+```
+
+### Verify Service Status
+
+To check if the service is running properly, use:
+
+```bash
+sudo systemctl status cocos-agent.service
+```
+
+### View Service Logs
+
+To inspect logs generated by the agent service, execute:
+
+```bash
+journalctl -u cocos-agent.service
+```
+
+### Check Standard Output and Error Logs
+
+To check logs stored in the system, use:
+
+```bash
+cat /var/log/cocos/agent.stdout.log
+cat /var/log/cocos/agent.stderr.log
+```

hal/cloud/config.yaml

@@ -0,0 +1,174 @@
+#cloud-config
+package_update: true
+package_upgrade: false
+
+users:
+  - default
+  - name: cocos_user
+    gecos: Default User
+    groups:
+      - sudo
+      - docker # Add cocos user to the docker group
+    sudo:
+      - ALL=(ALL:ALL) ALL
+    shell: /bin/bash
+
+chpasswd:
+  list: |
+    cocos_user:password
+  expire: False
+
+ssh_pwauth: True
+
+packages:
+  - curl
+  - make
+  - git
+  - python3
+  - python3-dev
+  - net-tools # Add net-tools to install the 'route' command
+
+write_files:
+  - path: /etc/cocos/certs/cert.pem
+    content: |
+      # Add certificate content here
+    permissions: "0644"
+
+  - path: /etc/cocos/certs/ca.pem
+    content: |
+      # Add CA certificate content here
+    permissions: "0644"
+
+  - path: /etc/cocos/certs/key.pem
+    content: |
+      # Add private key content here
+    permissions: "0600"
+
+  - path: /etc/cocos/environment
+    content: |
+      # Add environment variables here
+    permissions: "0644"
+
+  - path: /etc/systemd/system/cocos-agent.service
+    content: |
+      [Unit]
+      Description=Cocos AI agent
+      After=network.target
+      Before=docker.service
+
+      [Service]
+      WorkingDirectory=/cocos
+      StandardOutput=file:/var/log/cocos/agent.stdout
+      StandardError=file:/var/log/cocos/agent.stderr
+      EnvironmentFile=/etc/cocos/environment
+      ExecStartPre=/cocos_init/agent_setup.sh
+      ExecStart=/cocos_init/agent_start_script.sh
+      Restart=always
+
+      [Install]
+      WantedBy=default.target
+    permissions: "0644"
+
+  # Agent setup script
+  - path: /cocos_init/agent_setup.sh
+    content: |
+      #!/bin/sh
+
+      WORK_DIR="/cocos"
+
+      # IFACES are all network interfaces excluding lo (LOOPBACK) and sit interfaces
+      IFACES=$(ip link show | grep -vE 'LOOPBACK|sit*' | awk -F': ' '{print $2}')
+
+      # This for loop brings up all network interfaces in IFACES and dhclient obtains an IP address for the every interface
+      for IFACE in $IFACES; do
+          STATE=$(ip link show $IFACE | grep DOWN)
+          if [ -n "$STATE" ]; then
+              ip link set $IFACE up
+          fi
+
+          IP_ADDR=$(ip addr show $IFACE | grep 'inet ')
+          if [ -z "$IP_ADDR" ]; then
+              dhclient $IFACE
+          fi
+      done
+
+      if [ ! -d "$WORK_DIR" ]; then
+          mkdir -p $WORK_DIR
+      fi
+
+      # Resize the root filesystem to 100% of available space
+      ROOT_DEV=$(findmnt / -o SOURCE -n)  # Get the root filesystem device
+      resize2fs "$ROOT_DEV" && echo "Root filesystem resized successfully" || echo "Failed to resize root filesystem"
+    permissions: "0755"
+
+  # Agent start script
+  - path: /cocos_init/agent_start_script.sh
+    content: |
+      #!/bin/sh
+
+      # Change the docker.service file to allow Docker to run in RAM
+      mkdir -p /etc/systemd/system/docker.service.d
+
+      # Create or overwrite the override.conf file with the new Environment variable
+      tee /etc/systemd/system/docker.service.d/override.conf > /dev/null <<EOF
+      [Service]
+      Environment=DOCKER_RAMDISK=true
+      EOF
+
+      systemctl daemon-reload
+
+      NUM_OF_PERMITED_IFACE=1
+
+      NUM_OF_IFACE=$(ip route | grep -Eo 'dev [a-z0-9]+' | awk '{ print $2 }' | grep -v '^docker' | sort | uniq | wc -l)
+
+      if [ $NUM_OF_IFACE -gt $NUM_OF_PERMITED_IFACE ]; then
+          echo "More than one network interface in the VM"
+          exit 1
+      fi
+
+      DEFAULT_IFACE=$(route | grep '^default' | grep -o '[^ ]*$')
+      AGENT_GRPC_HOST=$(ip -4 addr show $DEFAULT_IFACE | grep inet | awk '{print $2}' | cut -d/ -f1)
+
+      export AGENT_GRPC_HOST
+
+      exec /bin/cocos-agent
+    permissions: "0755"
+
+runcmd:
+  # Create necessary directories
+  - mkdir -p /cocos
+  - mkdir -p /cocos_init
+  - mkdir -p /var/log/cocos
+  - mkdir -p /etc/cocos
+
+  # Download the cocos-agent binary
+  - echo "[ COCOS AGENT SETUP ] Downloading the cocos-agent binary..."
+  - curl -L -O -J https://github.com/smithjilks/cocos/releases/download/v1.0.0/cocos-agent --progress-bar && echo "[ COCOS AGENT SETUP ] cocos-agent binary downloaded successfully" || echo "Failed to download cocos-agent binary"
+
+  # Install the agent binary
+  - echo "[ COCOS AGENT SETUP ] Installing cocos-agent binary..."
+  - install -D -m 0755 cocos-agent /bin/cocos-agent && echo "[ COCOS AGENT SETUP ] cocos-agent binary installed successfully" || echo "[ COCOS AGENT SETUP ] Failed to install cocos-agent binary"
+
+  # Install Wasmtime
+  - echo "Installing Wasmtime runtime..."
+  - curl https://wasmtime.dev/install.sh -sSf | bash && echo "Wasmtime installed successfully" || echo "Failed to install Wasmtime"
+  - echo "Configuring Wasmtime environment variables..."
+  - echo "export WASMTIME_HOME=$HOME/.wasmtime" >> /etc/profile.d/wasm_env.sh
+  - echo "export PATH=\$WASMTIME_HOME/bin:\$PATH" >> /etc/profile.d/wasm_env.sh
+  - . /etc/profile.d/wasm_env.sh && echo "Wasmtime environment variables configured successfully" || echo "Failed to configure Wasmtime environment variables"
+
+  # Install Docker
+  - echo "Starting Docker installation..."
+  - curl -fsSL https://get.docker.com -o get-docker.sh && echo "Docker install script downloaded successfully" || echo "Failed to download Docker install script"
+  - sh ./get-docker.sh && echo "Docker installed successfully" || echo "Failed to install Docker"
+  - usermod -aG docker cocos_user && echo "Added cocos_user to the docker group" || echo "Failed to add cocos_user to the docker group"
+
+  # Reload systemd and enable the service
+  - echo "[ COCOS AGENT SETUP ] Reloading systemd daemon..."
+  - systemctl daemon-reload && echo "[ COCOS AGENT SETUP ] Systemd daemon reloaded successfully" || echo "[ COCOS AGENT SETUP ] Failed to reload systemd daemon"
+  - echo "[ COCOS AGENT SETUP ] Enabling cocos-agent.service..."
+  - systemctl enable cocos-agent.service && echo "[ COCOS AGENT SETUP ] cocos-agent.service enabled successfully" || echo "[ COCOS AGENT SETUP ] Failed to enable cocos-agent.service"
+  - echo "[ COCOS AGENT SETUP ] Starting cocos-agent.service..."
+  - systemctl start cocos-agent.service && echo "[ COCOS AGENT SETUP ] cocos-agent.service started successfully" || echo "[ COCOS AGENT SETUP ] Failed to start cocos-agent.service"
+
+final_message: "Cocos agent setup complete. Verify logs to confirm successful service startup."

hal/cloud/meta-data

@@ -0,0 +1,2 @@
+instance-id: iid-cocos-vm
+local-hostname: cocos-vm

hal/cloud/qemu.sh

@@ -0,0 +1,124 @@
+#!/bin/bash
+
+# Source environment variables
+source ./.env
+
+# Required commands
+REQUIRED_CMDS=("wget" "cloud-localds" "$QEMU_BINARY" "qemu-img")
+
+# Check for required commands
+for cmd in "${REQUIRED_CMDS[@]}"; do
+    if ! command -v "$cmd" &> /dev/null; then
+        echo "Error: $cmd is not installed. Please install it and try again."
+        exit 1
+    fi
+done
+
+# Ensure script is run as root
+if [[ $EUID -ne 0 ]]; then
+    echo "Error: This script must be run as root."
+    exit 1
+fi
+
+# Create the root filesystem image if it doesn't exist
+if [ ! -f "$BASE_IMAGE" ]; then
+  echo "Downloading base Ubuntu image..."
+  wget -q "$BASE_IMAGE_URL" -O "$BASE_IMAGE" --show-progress
+fi
+
+# Create custom image
+echo "Creating custom QEMU image..."
+qemu-img create -f qcow2 -b "$BASE_IMAGE" -F qcow2 "$CUSTOM_IMAGE" "$DISK_SIZE"
+
+# Cloud-init configuration files
+CLOUD_CONFIG="config.yaml"
+META_DATA="meta-data"
+SEED_IMAGE="seed.img"
+
+# Create seed image for cloud-init
+echo "Creating seed image..."
+cloud-localds "$SEED_IMAGE" "$CLOUD_CONFIG" "$META_DATA"
+
+# Construct QEMU arguments from environment variables
+construct_qemu_args() {
+    args=()
+
+    args+=("-name" "$VM_NAME")
+
+    # Virtualization (Enable KVM)
+    if [ "$ENABLE_KVM" == "true" ]; then
+        args+=("-enable-kvm")
+    fi
+
+    # Machine, CPU, RAM
+    if [ -n "$MACHINE" ]; then
+        args+=("-machine" "$MACHINE")
+    fi
+
+    if [ -n "$CPU" ]; then
+        args+=("-cpu" "$CPU")
+    fi
+
+    args+=("-boot" "d")
+    args+=("-smp" "$SMP_COUNT,maxcpus=$SMP_MAXCPUS")
+    args+=("-m" "$MEMORY_SIZE,slots=$MEMORY_SLOTS,maxmem=$MAX_MEMORY")
+
+    # OVMF (if applicable)
+    if [ "$ENABLE_SEV_SNP" != "true" ]; then
+        args+=("-drive" "if=$OVMF_CODE_IF,format=$OVMF_CODE_FORMAT,unit=$OVMF_CODE_UNIT,file=$OVMF_CODE,readonly=$OVMF_CODE_READONLY")
+        args+=("-drive" "if=$OVMF_VARS_IF,format=$OVMF_VARS_FORMAT,unit=$OVMF_VARS_UNIT,file=$OVMF_VARS")
+    fi
+
+    # Network configuration
+    args+=("-netdev" "user,id=$NET_DEV_ID,hostfwd=tcp::$NET_DEV_HOST_FWD_AGENT-:$NET_DEV_GUEST_FWD_AGENT")
+    args+=("-device" "virtio-net-pci,disable-legacy=$VIRTIO_NET_PCI_DISABLE_LEGACY,iommu_platform=$VIRTIO_NET_PCI_IOMMU_PLATFORM,netdev=$NET_DEV_ID,addr=$VIRTIO_NET_PCI_ADDR,romfile=$VIRTIO_NET_PCI_ROMFILE")
+    args+=("-device" "vhost-vsock-pci,id=$VSOCK_ID,guest-cid=$VSOCK_GUEST_CID")
+
+    # SEV (if enabled)
+    if [ "$ENABLE_SEV" == "true" ] || [ "$ENABLE_SEV_SNP" == "true" ]; then
+        sev_type="sev-guest"
+        kernel_hash=""
+        host_data=""
+
+        args+=("-machine" "confidential-guest-support=$SEV_ID,memory-backend=$MEM_ID")
+
+        if [ "$ENABLE_SEV_SNP" == "true" ]; then
+            args+=("-bios" "$OVMF_CODE_FILE")
+            sev_type="sev-snp-guest"
+
+            if [ -n "$SEV_HOST_DATA" ]; then
+                host_data=",host-data=$SEV_HOST_DATA"
+            fi
+        fi
+
+        if [ "$ENABLE_KERNEL_HASH" == "true" ]; then
+            kernel_hash=",kernel-hashes=on"
+        fi
+
+        args+=("-object" "memory-backend-memfd,id=$MEM_ID,size=$MEMORY_SIZE,share=true,prealloc=false")
+        args+=("-object" "$sev_type,id=$SEV_ID,cbitpos=$SEV_CBIT_POS,reduced-phys-bits=$SEV_REDUCED_PHYS_BITS$kernel_hash$host_data")
+    fi
+
+    # Disk image configuration
+    args+=("-drive" "file=$SEED_IMAGE,media=cdrom")
+    args+=("-drive" "file=$CUSTOM_IMAGE,if=none,id=disk0,format=qcow2")
+    args+=("-device" "virtio-scsi-pci,id=scsi,disable-legacy=on,iommu_platform=true")
+    args+=("-device" "scsi-hd,drive=disk0")
+
+    # Display options
+    if [ "$NO_GRAPHIC" == "true" ]; then
+        args+=("-nographic")
+    fi
+
+    args+=("-monitor" "$MONITOR")
+    args+=("-no-reboot")
+    args+=("-vnc" ":9")
+
+    echo "${args[@]}"
+}
+
+qemu_args=$(construct_qemu_args)
+
+echo "Running QEMU with the following arguments: $qemu_args"
+echo "Starting QEMU VM..."
+$QEMU_BINARY $qemu_args

hal/linux/board/cocos/linux.config

@@ -65,3 +65,14 @@ CONFIG_PREEMPT_DYNAMIC=n
 CONFIG_DEBUG_PREEMPT=n
 CONFIG_CGROUP_MISC=y
 CONFIG_X86_CPUID=y
+
+CONFIG_NET_9P=y
+CONFIG_NET_9P_VIRTIO=y
+CONFIG_9P_FS=y
+CONFIG_9P_FS_POSIX_ACL=y
+CONFIG_9P_FS_SECURITY=y
+
+# TCG TPM
+CONFIG_TCG_TPM=y
+CONFIG_TCG_TPM2_HMAC=y
+CONFIG_TCG_PLATFORM=y

hal/linux/board/cocos/post-build.sh

@@ -6,6 +6,23 @@ set -e
 # Add a console on tty1
 if [ -e ${TARGET_DIR}/etc/inittab ]; then
     grep -qE '^tty1::' ${TARGET_DIR}/etc/inittab || \
-	sed -i '/GENERIC_SERIAL/a\
+    sed -i '/GENERIC_SERIAL/a\
 tty1::respawn:/sbin/getty -L  tty1 0 vt100 # QEMU graphical window' ${TARGET_DIR}/etc/inittab
 fi
+
+# Create the mount points
+# Create the mount points
+mkdir -p ${TARGET_DIR}/etc/certs
+mkdir -p ${TARGET_DIR}/etc/cocos
+
+# Ensure /etc/fstab exists
+if [ ! -f "${TARGET_DIR}/etc/fstab" ]; then
+    touch "${TARGET_DIR}/etc/fstab"
+fi
+
+# Add the 9p entries to /etc/fstab
+grep -q "certs_share /etc/certs" ${TARGET_DIR}/etc/fstab || \
+echo "certs_share /etc/certs 9p trans=virtio,version=9p2000.L,cache=mmap 0 0" >> "${TARGET_DIR}/etc/fstab"
+
+grep -q "env_share /etc/cocos" ${TARGET_DIR}/etc/fstab || \
+echo "env_share /etc/cocos 9p trans=virtio,version=9p2000.L,cache=mmap 0 0" >> "${TARGET_DIR}/etc/fstab"

hal/linux/configs/cocos_defconfig

@@ -14,6 +14,8 @@ BR2_SYSTEM_BIN_SH_BASH=y
 BR2_TARGET_ROOTFS_CPIO=y
 BR2_TARGET_ROOTFS_CPIO_FULL=y
 BR2_TARGET_ROOTFS_CPIO_GZIP=y
+BR2_TARGET_ROOTFS_OVERLAY="overlay"
+BR2_PACKAGE_9PFS=y
 
 # Image
 BR2_ROOTFS_POST_BUILD_SCRIPT="$(BR2_EXTERNAL_COCOS_PATH)/board/cocos/post-build.sh"
@@ -25,14 +27,14 @@ BR2_ROOTFS_POST_SCRIPT_ARGS="$(BR2_DEFCONFIG)"
 # Linux headers same as kernel
 BR2_PACKAGE_HOST_LINUX_HEADERS_CUSTOM_6_11=y
 BR2_TOOLCHAIN_HEADERS_LATEST=y
-BR2_TOOLCHAIN_HEADERS_AT_LEAST="6.12-rc6"
+BR2_TOOLCHAIN_HEADERS_AT_LEAST="6.11-rc7"
 
 # Kernel
 BR2_LINUX_KERNEL=y
 BR2_LINUX_KERNEL_CUSTOM_GIT=y
-BR2_LINUX_KERNEL_CUSTOM_REPO_URL="https://github.com/torvalds/linux.git"
-BR2_LINUX_KERNEL_CUSTOM_REPO_VERSION="v6.12-rc6"
-BR2_LINUX_KERNEL_VERSION="v6.12-rc6"
+BR2_LINUX_KERNEL_CUSTOM_REPO_URL="https://github.com/coconut-svsm/linux.git"
+BR2_LINUX_KERNEL_CUSTOM_REPO_VERSION="svsm"
+BR2_LINUX_KERNEL_VERSION="svsm"
 BR2_LINUX_KERNEL_PATCH=""
 BR2_LINUX_KERNEL_USE_CUSTOM_CONFIG=y
 BR2_LINUX_KERNEL_CUSTOM_CONFIG_FILE="$(BR2_EXTERNAL_COCOS_PATH)/board/cocos/linux.config"

init/systemd/cocos-agent.service

@@ -8,8 +8,7 @@ WorkingDirectory=/cocos
 StandardOutput=file:/var/log/cocos/agent.stdout
 StandardError=file:/var/log/cocos/agent.stderr
 
-Environment=AGENT_GRPC_PORT=7002
-Environment=AGENT_LOG_LEVEL=info
+EnvironmentFile=/etc/cocos/environment
 
 ExecStartPre=/cocos_init/agent_setup.sh
 ExecStart=/cocos_init/agent_start_script.sh

internal/logger/protohandler.go

@@ -7,8 +7,9 @@ import (
 	"io"
 	"log/slog"
 
+	"github.com/ultravioletrs/cocos/agent/cvms"
 	"github.com/ultravioletrs/cocos/agent/events"
-	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
@@ -18,16 +19,17 @@ type handler struct {
 	opts  slog.HandlerOptions
 	w     io.Writer
 	cmpID string
+	queue chan *cvms.ClientStreamMessage
 }
 
-func NewProtoHandler(conn io.Writer, opts *slog.HandlerOptions, cmpID string) slog.Handler {
+func NewProtoHandler(conn io.Writer, opts *slog.HandlerOptions, queue chan *cvms.ClientStreamMessage) slog.Handler {
 	if opts == nil {
 		opts = &slog.HandlerOptions{}
 	}
 	h := &handler{
 		opts:  *opts,
 		w:     conn,
-		cmpID: cmpID,
+		queue: queue,
 	}
 
 	return h
@@ -69,7 +71,18 @@ func (h *handler) Handle(_ context.Context, r slog.Record) error {
 			},
 		}
 
-		b, err := proto.Marshal(&agentLog)
+		h.queue <- &cvms.ClientStreamMessage{
+			Message: &cvms.ClientStreamMessage_AgentLog{
+				AgentLog: &cvms.AgentLog{
+					Timestamp:     timestamp,
+					Message:       chunk,
+					Level:         level,
+					ComputationId: h.cmpID,
+				},
+			},
+		}
+
+		b, err := protojson.Marshal(&agentLog)
 		if err != nil {
 			return err
 		}
@@ -78,6 +91,11 @@ func (h *handler) Handle(_ context.Context, r slog.Record) error {
 		if err != nil {
 			return err
 		}
+
+		_, err = h.w.Write([]byte("\n"))
+		if err != nil {
+			return err
+		}
 	}
 
 	return nil
@@ -88,7 +106,8 @@ func (h *handler) WithAttrs(attrs []slog.Attr) slog.Handler {
 }
 
 func (h *handler) WithGroup(name string) slog.Handler {
-	panic("unimplemented")
+	h.cmpID = name
+	return h
 }
 
 func (h *handler) Close() error {

internal/logger/protohandler_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/absmach/magistrala/pkg/errors"
 	"github.com/stretchr/testify/assert"
+	"github.com/ultravioletrs/cocos/agent/cvms"
 )
 
 type failedWriter struct{}
@@ -21,14 +22,14 @@ func (f *failedWriter) Write(p []byte) (n int, err error) {
 
 // TestNewProtoHandler tests the initialization of the ProtoHandler.
 func TestNewProtoHandler(t *testing.T) {
-	handler := NewProtoHandler(io.Discard, nil, "testCmpID")
+	handler := NewProtoHandler(io.Discard, nil, make(chan *cvms.ClientStreamMessage))
 
 	assert.NotNil(t, handler, "Handler should not be nil")
 }
 
 // TestHandleMessageSuccess tests the handling of a message when the write succeeds.
 func TestHandleMessageSuccess(t *testing.T) {
-	handler := NewProtoHandler(io.Discard, nil, "testCmpID")
+	handler := NewProtoHandler(io.Discard, nil, make(chan *cvms.ClientStreamMessage, 1))
 	record := slog.Record{
 		Time:    time.Now(),
 		Message: "Test message",
@@ -42,7 +43,7 @@ func TestHandleMessageSuccess(t *testing.T) {
 
 // TestHandleMessageFailure tests the caching mechanism when the write fails.
 func TestHandleMessageFailure(t *testing.T) {
-	protohandler := NewProtoHandler(&failedWriter{}, nil, "testCmpID")
+	protohandler := NewProtoHandler(&failedWriter{}, nil, make(chan *cvms.ClientStreamMessage, 1))
 	record := slog.Record{
 		Time:    time.Now(),
 		Message: "Test message",
@@ -56,7 +57,7 @@ func TestHandleMessageFailure(t *testing.T) {
 
 // TestEnabled tests that the handler enables logging based on level.
 func TestEnabled(t *testing.T) {
-	handler := NewProtoHandler(io.Discard, nil, "testCmpID")
+	handler := NewProtoHandler(io.Discard, nil, make(chan *cvms.ClientStreamMessage, 1))
 
 	assert.True(t, handler.Enabled(context.Background(), slog.LevelInfo), "Logging should be enabled for LevelInfo")
 	assert.False(t, handler.Enabled(context.Background(), slog.LevelDebug), "Logging should be disabled for LevelDebug by default")
@@ -66,7 +67,7 @@ func TestEnabled(t *testing.T) {
 func TestCloseStopsRetry(t *testing.T) {
 	mockWriter := io.Discard
 
-	handler := NewProtoHandler(mockWriter, nil, "testCmpID").(*handler)
+	handler := NewProtoHandler(mockWriter, nil, make(chan *cvms.ClientStreamMessage, 1)).(*handler)
 
 	time.Sleep(2 * time.Second)
 	err := handler.Close()

internal/server/grpc/grpc.go

@@ -25,6 +25,7 @@ import (
 	"github.com/ultravioletrs/cocos/agent/auth"
 	"github.com/ultravioletrs/cocos/internal/server"
 	"github.com/ultravioletrs/cocos/pkg/atls"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
@@ -51,7 +52,7 @@ type Server struct {
 	server.BaseServer
 	server          *grpc.Server
 	registerService serviceRegister
-	quoteProvider   client.QuoteProvider
+	quoteProvider   client.LeveledQuoteProvider
 	authSvc         auth.Authenticator
 	health          *health.Server
 }
@@ -60,7 +61,7 @@ type serviceRegister func(srv *grpc.Server)
 
 var _ server.Server = (*Server)(nil)
 
-func New(ctx context.Context, cancel context.CancelFunc, name string, config server.ServerConfiguration, registerService serviceRegister, logger *slog.Logger, qp client.QuoteProvider, authSvc auth.Authenticator) server.Server {
+func New(ctx context.Context, cancel context.CancelFunc, name string, config server.ServerConfiguration, registerService serviceRegister, logger *slog.Logger, qp client.LeveledQuoteProvider, authSvc auth.Authenticator) server.Server {
 	base := config.GetBaseConfig()
 	listenFullAddress := fmt.Sprintf("%s:%s", base.Host, base.Port)
 	return &Server{
@@ -301,5 +302,19 @@ func generateCertificatesForATLS() ([]byte, []byte, error) {
 		Bytes: privateKeyBytes,
 	})
 
+	cert, err := x509.ParseCertificate(certDERBytes)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to parse certificate: %w", err)
+	}
+
+	pubKeyDER, err := x509.MarshalPKIXPublicKey(cert.PublicKey)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to marshal public key to DER format: %w", err)
+	}
+
+	if err := vtpm.ExtendPCR(vtpm.PCR15, pubKeyDER); err != nil {
+		return nil, nil, fmt.Errorf("failed to extend vTPM PCR with public key: %w", err)
+	}
+
 	return certBytes, keyBytes, nil
 }

internal/server/grpc/grpc_test.go

@@ -22,6 +22,7 @@ import (
 	authmocks "github.com/ultravioletrs/cocos/agent/auth/mocks"
 	"github.com/ultravioletrs/cocos/internal/server"
 	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider/mocks"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/test/bufconn"
 )
@@ -30,6 +31,28 @@ const bufSize = 1024 * 1024
 
 var lis *bufconn.Listener
 
+type DummyRWC struct{}
+
+// Read fills p with byte(len(p)) and returns len(p).
+func (l *DummyRWC) Read(p []byte) (int, error) {
+	n := len(p)
+	// Fill each byte in p with the value of n as a byte.
+	for i := range p {
+		p[i] = byte(n)
+	}
+	return n, nil
+}
+
+// Write simply returns len(p) indicating that all bytes were written.
+func (l *DummyRWC) Write(p []byte) (int, error) {
+	// In this simple implementation, we ignore the data.
+	return len(p), nil
+}
+
+func (l *DummyRWC) Close() error {
+	return nil
+}
+
 func init() {
 	lis = bufconn.Listen(bufSize)
 }
@@ -47,7 +70,7 @@ func TestNew(t *testing.T) {
 		},
 	}
 	logger := slog.Default()
-	qp := new(mocks.QuoteProvider)
+	qp := new(mocks.LeveledQuoteProvider)
 	authSvc := new(authmocks.Authenticator)
 
 	srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
@@ -97,7 +120,7 @@ func TestServerStartWithTLSFile(t *testing.T) {
 
 	logBuffer := &ThreadSafeBuffer{}
 	logger := slog.New(slog.NewTextHandler(logBuffer, &slog.HandlerOptions{Level: slog.LevelDebug}))
-	qp := new(mocks.QuoteProvider)
+	qp := new(mocks.LeveledQuoteProvider)
 	authSvc := new(authmocks.Authenticator)
 
 	srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
@@ -144,7 +167,7 @@ func TestServerStartWithmTLSFile(t *testing.T) {
 
 	logBuffer := &ThreadSafeBuffer{}
 	logger := slog.New(slog.NewTextHandler(logBuffer, &slog.HandlerOptions{Level: slog.LevelDebug}))
-	qp := new(mocks.QuoteProvider)
+	qp := new(mocks.LeveledQuoteProvider)
 	authSvc := new(authmocks.Authenticator)
 
 	srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
@@ -184,7 +207,7 @@ func TestServerStop(t *testing.T) {
 	}
 	buf := &ThreadSafeBuffer{}
 	logger := slog.New(slog.NewTextHandler(buf, &slog.HandlerOptions{Level: slog.LevelDebug}))
-	qp := new(mocks.QuoteProvider)
+	qp := new(mocks.LeveledQuoteProvider)
 	authSvc := new(authmocks.Authenticator)
 
 	srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
@@ -259,6 +282,8 @@ func (b *ThreadSafeBuffer) String() string {
 }
 
 func TestServerInitializationAndStartup(t *testing.T) {
+	vtpm.ExternalTPM = &DummyRWC{}
+
 	testCases := []struct {
 		name          string
 		config        server.AgentConfig
@@ -374,7 +399,7 @@ func TestServerInitializationAndStartup(t *testing.T) {
 
 			logBuffer := &ThreadSafeBuffer{}
 			logger := slog.New(slog.NewTextHandler(logBuffer, &slog.HandlerOptions{Level: slog.LevelDebug}))
-			qp := new(mocks.QuoteProvider)
+			qp := new(mocks.LeveledQuoteProvider)
 			authSvc := new(authmocks.Authenticator)
 
 			srv := New(ctx, cancel, "TestServer", tc.config, func(srv *grpc.Server) {}, logger, qp, authSvc)

internal/vsock/client.go

@@ -1,247 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-package vsock
-
-import (
-	"context"
-	"encoding/binary"
-	"fmt"
-	"io"
-	"log"
-	"net"
-	"sync"
-	"sync/atomic"
-	"time"
-
-	"google.golang.org/protobuf/proto"
-)
-
-const (
-	maxRetries     = 3
-	retryDelay     = time.Second
-	maxMessageSize = 1 << 20 // 1 MB
-	ackTimeout     = 5 * time.Second
-	maxConcurrent  = 100
-)
-
-type MessageStatus int
-
-const (
-	StatusPending MessageStatus = iota
-	StatusSent
-	StatusAcknowledged
-	StatusFailed
-)
-
-type Message struct {
-	ID      uint32
-	Content []byte
-	Status  MessageStatus
-	Retries int
-}
-
-type AckWriter struct {
-	conn            net.Conn
-	pendingMessages chan *Message
-	messageStore    sync.Map // map[uint32]*Message
-	nextID          uint32
-	ctx             context.Context
-	cancel          context.CancelFunc
-	wg              sync.WaitGroup
-}
-
-func NewAckWriter(conn net.Conn) io.WriteCloser {
-	ctx, cancel := context.WithCancel(context.Background())
-	aw := &AckWriter{
-		conn:            conn,
-		pendingMessages: make(chan *Message, maxConcurrent),
-		nextID:          1,
-		ctx:             ctx,
-		cancel:          cancel,
-	}
-	aw.wg.Add(2)
-	go aw.sendMessages()
-	go aw.handleAcknowledgments()
-	return aw
-}
-
-func (aw *AckWriter) Write(p []byte) (int, error) {
-	if len(p) > maxMessageSize {
-		return 0, fmt.Errorf("message size exceeds maximum allowed size of %d bytes", maxMessageSize)
-	}
-
-	messageID := atomic.AddUint32(&aw.nextID, 1)
-	message := &Message{
-		ID:      messageID,
-		Content: make([]byte, len(p)),
-		Status:  StatusPending,
-	}
-	copy(message.Content, p)
-
-	aw.messageStore.Store(messageID, message)
-	select {
-	case aw.pendingMessages <- message:
-		return len(p), nil
-	case <-aw.ctx.Done():
-		return 0, fmt.Errorf("writer is closed")
-	}
-}
-
-func (aw *AckWriter) sendMessages() {
-	defer aw.wg.Done()
-
-	for {
-		select {
-		case <-aw.ctx.Done():
-			return
-		case msg := <-aw.pendingMessages:
-			if err := aw.sendWithRetry(msg); err != nil {
-				log.Printf("Failed to send message %d after all retries: %v", msg.ID, err)
-				msg.Status = StatusFailed
-				aw.messageStore.Store(msg.ID, msg)
-			}
-		}
-	}
-}
-
-func (aw *AckWriter) sendWithRetry(msg *Message) error {
-	for msg.Retries < maxRetries {
-		if err := aw.writeMessage(msg.ID, msg.Content); err != nil {
-			msg.Retries++
-			msg.Status = StatusPending
-			log.Printf("Error writing message %d (attempt %d): %v", msg.ID, msg.Retries, err)
-			time.Sleep(retryDelay)
-			continue
-		}
-		msg.Status = StatusSent
-		aw.messageStore.Store(msg.ID, msg)
-		return nil
-	}
-	return fmt.Errorf("max retries reached")
-}
-
-func (aw *AckWriter) writeMessage(messageID uint32, p []byte) error {
-	if err := binary.Write(aw.conn, binary.LittleEndian, messageID); err != nil {
-		return fmt.Errorf("failed to write message ID: %w", err)
-	}
-
-	messageLen := uint32(len(p))
-	if err := binary.Write(aw.conn, binary.LittleEndian, messageLen); err != nil {
-		return fmt.Errorf("failed to write message length: %w", err)
-	}
-
-	if _, err := aw.conn.Write(p); err != nil {
-		return fmt.Errorf("failed to write message content: %w", err)
-	}
-
-	return nil
-}
-
-func (aw *AckWriter) handleAcknowledgments() {
-	defer aw.wg.Done()
-
-	for {
-		select {
-		case <-aw.ctx.Done():
-			return
-		default:
-			var ackID uint32
-			if err := binary.Read(aw.conn, binary.LittleEndian, &ackID); err != nil {
-				if err == io.EOF {
-					log.Println("Connection closed, stopping acknowledgment handler")
-					return
-				}
-				log.Printf("Error reading ACK: %v", err)
-				time.Sleep(retryDelay)
-				continue
-			}
-
-			if msg, ok := aw.messageStore.Load(ackID); ok {
-				m := msg.(*Message)
-				m.Status = StatusAcknowledged
-				aw.messageStore.Store(ackID, m)
-
-				// Clean up old messages periodically
-				go aw.cleanupOldMessages(ackID)
-			} else {
-				log.Printf("Received ACK for unknown message ID: %d", ackID)
-			}
-		}
-	}
-}
-
-func (aw *AckWriter) cleanupOldMessages(currentID uint32) {
-	aw.messageStore.Range(func(key, value interface{}) bool {
-		msgID := key.(uint32)
-		msg := value.(*Message)
-
-		// Clean up acknowledged messages that are old
-		if msg.Status == StatusAcknowledged && msgID < currentID-maxConcurrent {
-			aw.messageStore.Delete(msgID)
-		}
-		return true
-	})
-}
-
-func (aw *AckWriter) Close() error {
-	aw.cancel()
-	aw.wg.Wait()
-	return aw.conn.Close()
-}
-
-type Reader interface {
-	Read() ([]byte, error)
-	ReadProto(msg proto.Message) error
-}
-
-type AckReader struct {
-	conn net.Conn
-	ctx  context.Context
-}
-
-func NewAckReader(conn net.Conn) Reader {
-	return &AckReader{
-		conn: conn,
-		ctx:  context.Background(),
-	}
-}
-
-func (ar *AckReader) ReadProto(msg proto.Message) error {
-	data, err := ar.Read()
-	if err != nil {
-		return fmt.Errorf("failed to read proto message: %w", err)
-	}
-	return proto.Unmarshal(data, msg)
-}
-
-func (ar *AckReader) Read() ([]byte, error) {
-	var messageID uint32
-	if err := binary.Read(ar.conn, binary.LittleEndian, &messageID); err != nil {
-		return nil, fmt.Errorf("error reading message ID: %w", err)
-	}
-
-	var messageLen uint32
-	if err := binary.Read(ar.conn, binary.LittleEndian, &messageLen); err != nil {
-		return nil, fmt.Errorf("error reading message length: %w", err)
-	}
-
-	if messageLen > maxMessageSize {
-		return nil, fmt.Errorf("message size %d exceeds maximum allowed size of %d bytes", messageLen, maxMessageSize)
-	}
-
-	data := make([]byte, messageLen)
-	if _, err := io.ReadFull(ar.conn, data); err != nil {
-		return nil, fmt.Errorf("error reading message content: %w", err)
-	}
-
-	if err := ar.sendAck(messageID); err != nil {
-		return nil, fmt.Errorf("error sending ACK: %w", err)
-	}
-
-	return data, nil
-}
-
-func (ar *AckReader) sendAck(messageID uint32) error {
-	return binary.Write(ar.conn, binary.LittleEndian, messageID)
-}

internal/vsock/client_test.go

@@ -1,337 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package vsock
-
-import (
-	"bytes"
-	"encoding/binary"
-	"errors"
-	"fmt"
-	"io"
-	"net"
-	"sync"
-	"testing"
-	"time"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/ultravioletrs/cocos/manager"
-	"google.golang.org/protobuf/proto"
-)
-
-// MockConn implements net.Conn for testing purposes.
-type MockConn struct {
-	ReadData    []byte
-	WrittenData []byte
-	ReadErr     error
-	WriteErr    error
-	closed      bool
-	mu          sync.Mutex
-}
-
-func (m *MockConn) Read(b []byte) (n int, err error) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-	if m.closed {
-		return 0, io.EOF
-	}
-	if len(m.ReadData) == 0 {
-		return 0, io.EOF // Ensure we handle this case more predictably
-	}
-	if m.ReadErr != nil {
-		return 0, m.ReadErr
-	}
-	n = copy(b, m.ReadData)
-	m.ReadData = m.ReadData[n:]
-	return n, nil
-}
-
-func (m *MockConn) Write(b []byte) (n int, err error) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-	if m.closed {
-		return 0, errors.New("connection closed")
-	}
-	if m.WriteErr != nil {
-		return 0, m.WriteErr
-	}
-	m.WrittenData = append(m.WrittenData, b...)
-	return len(b), nil
-}
-
-func (m *MockConn) Close() error {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-	m.closed = true
-	return nil
-}
-
-// Implement other net.Conn methods with empty implementations.
-func (m *MockConn) LocalAddr() net.Addr                { return nil }
-func (m *MockConn) RemoteAddr() net.Addr               { return nil }
-func (m *MockConn) SetDeadline(t time.Time) error      { return nil }
-func (m *MockConn) SetReadDeadline(t time.Time) error  { return nil }
-func (m *MockConn) SetWriteDeadline(t time.Time) error { return nil }
-
-func TestAckReader_Read(t *testing.T) {
-	tests := []struct {
-		name    string
-		data    []byte
-		wantErr bool
-	}{
-		{"Valid message", []byte("Hello, World!"), false},
-		{"Empty message", []byte{}, false},
-		{"Message at max size", make([]byte, maxMessageSize), false},
-		{"Message exceeds max size", make([]byte, maxMessageSize+1), true},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			mockConn := &MockConn{}
-			ar := NewAckReader(mockConn)
-
-			// Prepare mock data
-			messageID := uint32(1)
-			messageLen := uint32(len(tt.data))
-			mockData := make([]byte, 8+len(tt.data))
-			binary.LittleEndian.PutUint32(mockData[:4], messageID)
-			binary.LittleEndian.PutUint32(mockData[4:8], messageLen)
-			copy(mockData[8:], tt.data)
-			mockConn.ReadData = mockData
-
-			data, err := ar.Read()
-
-			if (err != nil) != tt.wantErr {
-				t.Errorf("AckReader.Read() error = %v, wantErr %v", err, tt.wantErr)
-				return
-			}
-
-			if !tt.wantErr {
-				if !bytes.Equal(data, tt.data) {
-					t.Errorf("AckReader.Read() got = %v, want %v", data, tt.data)
-				}
-
-				// Check if ACK was sent
-				if len(mockConn.WrittenData) != 4 {
-					t.Errorf("AckReader.Read() did not send ACK")
-				} else {
-					ackID := binary.LittleEndian.Uint32(mockConn.WrittenData)
-					if ackID != messageID {
-						t.Errorf("AckReader.Read() sent wrong ACK ID, got %d, want %d", ackID, messageID)
-					}
-				}
-			}
-		})
-	}
-}
-
-func TestAckReader_ReadProto(t *testing.T) {
-	tests := []struct {
-		name    string
-		msg     *manager.ClientStreamMessage
-		wantErr bool
-	}{
-		{"Valid proto message", &manager.ClientStreamMessage{}, false},
-		{"Empty proto message", &manager.ClientStreamMessage{}, false},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			mockConn := &MockConn{}
-			ar := NewAckReader(mockConn)
-
-			// Prepare mock data
-			protoData, _ := proto.Marshal(tt.msg)
-			messageID := uint32(1)
-			messageLen := uint32(len(protoData))
-			mockData := make([]byte, 8+len(protoData))
-			binary.LittleEndian.PutUint32(mockData[:4], messageID)
-			binary.LittleEndian.PutUint32(mockData[4:8], messageLen)
-			copy(mockData[8:], protoData)
-			mockConn.ReadData = mockData
-
-			receivedMsg := &manager.ClientStreamMessage{}
-			err := ar.ReadProto(receivedMsg)
-
-			if (err != nil) != tt.wantErr {
-				t.Errorf("AckReader.ReadProto() error = %v, wantErr %v", err, tt.wantErr)
-				return
-			}
-
-			if !tt.wantErr {
-				if receivedMsg.Message != tt.msg.Message {
-					t.Errorf("AckReader.ReadProto() got = %v, want %v", receivedMsg, tt.msg)
-				}
-
-				// Check if ACK was sent
-				if len(mockConn.WrittenData) != 4 {
-					t.Errorf("AckReader.ReadProto() did not send ACK")
-				} else {
-					ackID := binary.LittleEndian.Uint32(mockConn.WrittenData)
-					if ackID != messageID {
-						t.Errorf("AckReader.ReadProto() sent wrong ACK ID, got %d, want %d", ackID, messageID)
-					}
-				}
-			}
-		})
-	}
-}
-
-func TestNewAckWriter(t *testing.T) {
-	mockConn := &MockConn{}
-	writer := NewAckWriter(mockConn)
-
-	if _, ok := writer.(io.Writer); !ok {
-		t.Errorf("NewAckWriter() did not return an io.Writer")
-	}
-}
-
-func TestNewAckReader(t *testing.T) {
-	mockConn := &MockConn{}
-	reader := NewAckReader(mockConn)
-
-	assert.NotNil(t, reader)
-}
-
-func TestAckWriter_Close(t *testing.T) {
-	mockConn := &MockConn{}
-	aw := NewAckWriter(mockConn)
-
-	err := aw.Close()
-	if err != nil {
-		t.Errorf("AckWriter.Close() error = %v, wantErr %v", err, nil)
-	}
-
-	if !mockConn.closed {
-		t.Errorf("AckWriter.Close() did not close the connection")
-	}
-}
-
-func TestAckWriter_Write(t *testing.T) {
-	tests := []struct {
-		name          string
-		input         []byte
-		expectErr     bool
-		expectedError string
-	}{
-		{
-			name:          "Message exceeds max size",
-			input:         make([]byte, maxMessageSize+1),
-			expectErr:     true,
-			expectedError: "message size exceeds maximum allowed size",
-		},
-		{
-			name:      "Write succeeds",
-			input:     []byte("Hello, world!"),
-			expectErr: false,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			mockConn := &MockConn{
-				mu: sync.Mutex{},
-			}
-
-			writer := NewAckWriter(mockConn)
-			defer writer.Close()
-
-			if tt.expectErr {
-				writer.(*AckWriter).ctx.Done()
-			}
-
-			time.Sleep(100 * time.Millisecond)
-
-			n, err := writer.Write(tt.input)
-
-			if tt.expectErr {
-				assert.Error(t, err)
-				if tt.expectedError != "" {
-					assert.Contains(t, err.Error(), tt.expectedError)
-				}
-				assert.Zero(t, n)
-			} else {
-				assert.NoError(t, err)
-				assert.Equal(t, len(tt.input), n)
-			}
-		})
-	}
-}
-
-func TestAckWriter_CleanupOldMessages(t *testing.T) {
-	mockConn := &MockConn{}
-	writer := NewAckWriter(mockConn).(*AckWriter)
-	defer writer.Close()
-
-	for i := uint32(1); i <= maxConcurrent+10; i++ {
-		msg := &Message{
-			ID:      i,
-			Content: []byte("test"),
-			Status:  StatusAcknowledged,
-		}
-		writer.messageStore.Store(i, msg)
-	}
-
-	writer.cleanupOldMessages(maxConcurrent + 11)
-
-	var count int
-	writer.messageStore.Range(func(key, value interface{}) bool {
-		count++
-		return true
-	})
-
-	assert.LessOrEqual(t, count, maxConcurrent)
-}
-
-func TestAckReader_LargeMessage(t *testing.T) {
-	mockConn := &MockConn{}
-	reader := NewAckReader(mockConn)
-
-	largeMessage := make([]byte, maxMessageSize-1)
-	for i := range largeMessage {
-		largeMessage[i] = byte(i % 256)
-	}
-
-	messageID := uint32(1)
-	messageLen := uint32(len(largeMessage))
-	mockData := make([]byte, 8+len(largeMessage))
-	binary.LittleEndian.PutUint32(mockData[:4], messageID)
-	binary.LittleEndian.PutUint32(mockData[4:8], messageLen)
-	copy(mockData[8:], largeMessage)
-	mockConn.ReadData = mockData
-
-	data, err := reader.Read()
-	assert.NoError(t, err)
-	assert.Equal(t, largeMessage, data)
-
-	assert.Equal(t, 4, len(mockConn.WrittenData))
-	ackID := binary.LittleEndian.Uint32(mockConn.WrittenData)
-	assert.Equal(t, messageID, ackID)
-}
-
-func TestAckWriter_FailedSends(t *testing.T) {
-	mockConn := &MockConn{
-		WriteErr: errors.New("write error"),
-	}
-	writer := NewAckWriter(mockConn).(*AckWriter)
-	defer writer.Close()
-
-	// Add some messages to the channel
-	for i := 0; i < 5; i++ {
-		msg := &Message{
-			ID:      uint32(i + 1),
-			Content: []byte(fmt.Sprintf("Message %d", i+1)),
-			Status:  StatusPending,
-		}
-		writer.pendingMessages <- msg
-	}
-
-	// Wait for the messages to be sent
-	time.Sleep(100 * time.Millisecond)
-
-	// Check that the messages were marked as failed
-	writer.messageStore.Range(func(key, value interface{}) bool {
-		msg := value.(*Message)
-		assert.Equal(t, StatusFailed, msg.Status)
-		return true
-	})
-}

manager/README.md

@@ -6,62 +6,66 @@ Manager service provides a barebones gRPC API and Service interface implementati
 
 The service is configured using the environment variables from the following table. Note that any unset variables will be replaced with their default values.
 
-| Variable                                  | Description                                                                                                      | Default                      |
-| ----------------------------------------- | ---------------------------------------------------------------------------------------------------------------- | ---------------------------- |
-| COCOS_JAEGER_URL                          | The URL for the Jaeger tracing endpoint.                                                                         | http://localhost:4318        |
-| COCOS_JAEGER_TRACE_RATIO                  | The ratio of traces to sample.                                                                                   | 1.0                          |
-| MANAGER_INSTANCE_ID                       | The instance ID for the manager service.                                                                         |                              |
-| MANAGER_ATTESTATION_POLICY_BINARY         | The file path for the attestation policy binary.                                                                 | ../../build                  |
-| MANAGER_GRPC_CLIENT_CERT                  | The file path for the client certificate.                                                                        |                              |
-| MANAGER_GRPC_CLIENT_KEY                   | The file path for the client private key.                                                                        |                              |
-| MANAGER_GRPC_SERVER_CA_CERTS              | The file path for the server CA certificate(s).                                                                  |                              |
-| MANAGER_GRPC_URL                          | The URL for the gRPC endpoint.                                                                                   | localhost:7001               |
-| MANAGER_GRPC_TIMEOUT                      | The timeout for gRPC requests.                                                                                   | 60s                          |
-| MANAGER_EOS_VERSION                       | The EOS version used for booting SVMs.                                                                           |                              |
-| MANAGER_INSTANCE_ID                       | Manager service instance ID                                                                                      |                              |
-| MANAGER_QEMU_MEMORY_SIZE                  | The total memory size for the virtual machine. Can be specified in a human-readable format like "2048M" or "4G". | 2048M                        |
-| MANAGER_QEMU_MEMORY_SLOTS                 | The number of memory slots for the virtual machine.                                                              | 5                            |
-| MANAGER_QEMU_MAX_MEMORY                   | The maximum memory size for the virtual machine. Can be specified in a human-readable format like "30G".         | 30G                          |
-| MANAGER_QEMU_OVMF_CODE_IF                 | The interface type for the OVMF code.                                                                            | pflash                       |
-| MANAGER_QEMU_OVMF_CODE_FORMAT             | The format of the OVMF code file.                                                                                | raw                          |
-| MANAGER_QEMU_OVMF_CODE_UNIT               | The unit number for the OVMF code.                                                                               | 0                            |
-| MANAGER_QEMU_OVMF_CODE_FILE               | The file path for the OVMF code.                                                                                 | /usr/share/OVMF/OVMF_CODE.fd |
-| MANAGER_QEMU_OVMF_VERSION                 | The version number of EDKII from which OVMF was built                                                            | edk2-stable202408            |
-| MANAGER_QEMU_OVMF_CODE_READONLY           | Whether the OVMF code should be read-only.                                                                       | on                           |
-| MANAGER_QEMU_OVMF_VARS_IF                 | The interface type for the OVMF variables.                                                                       | pflash                       |
-| MANAGER_QEMU_OVMF_VARS_FORMAT             | The format of the OVMF variables file.                                                                           | raw                          |
-| MANAGER_QEMU_OVMF_VARS_UNIT               | The unit number for the OVMF variables.                                                                          | 1                            |
-| MANAGER_QEMU_OVMF_VARS_FILE               | The file path for the OVMF variables.                                                                            | /usr/share/OVMF/OVMF_VARS.fd |
-| MANAGER_QEMU_NETDEV_ID                    | The ID for the network device.                                                                                   | vmnic                        |
-| MANAGER_QEMU_HOST_FWD_AGENT               | The port number for the host forward agent.                                                                      | 7020                         |
-| MANAGER_QEMU_GUEST_FWD_AGENT              | The port number for the guest forward agent.                                                                     | 7002                         |
-| MANAGER_QEMU_VIRTIO_NET_PCI_DISABLE_LEGACY | Whether to disable the legacy PCI device.                                                                       | on                           |
-| MANAGER_QEMU_VIRTIO_NET_PCI_IOMMU_PLATFORM | Whether to enable the IOMMU platform for the virtio-net PCI device.                                             | true                         |
-| MANAGER_QEMU_VIRTIO_NET_PCI_ADDR           | The PCI address for the virtio-net PCI device.                                                                  | 0x2                          |
-| MANAGER_QEMU_VIRTIO_NET_PCI_ROMFILE        | The file path for the ROM image for the virtio-net PCI device.                                                  |                              |
-| MANAGER_QEMU_DISK_IMG_KERNEL_FILE          | The file path for the kernel image.                                                                             | img/bzImage                  |
-| MANAGER_QEMU_DISK_IMG_ROOTFS_FILE          | The file path for the root filesystem image.                                                                    | img/rootfs.cpio.gz           |
-| MANAGER_QEMU_SEV_ID                        | The ID for the Secure Encrypted Virtualization (SEV) device.                                                    | sev0                         |
-| MANAGER_QEMU_SEV_CBITPOS                   | The position of the C-bit in the physical address.                                                              | 51                           |
-| MANAGER_QEMU_SEV_REDUCED_PHYS_BITS         | The number of reduced physical address bits for SEV.                                                            | 1                            |
-| MANAGER_QEMU_HOST_DATA                     | Additional data for the SEV host.                                                                               |                              |
-| MANAGER_QEMU_VSOCK_ID                      | The ID for the virtual socket device.                                                                           | vhost-vsock-pci0             |
-| MANAGER_QEMU_VSOCK_GUEST_CID               | The guest-side CID (Context ID) for the virtual socket device.                                                  | 3                            |
-| MANAGER_QEMU_VSOCK_VNC                     | Whether to enable the virtual socket device for VNC.                                                            | 0                            |
-| MANAGER_QEMU_BIN_PATH                      | The file path for the QEMU binary.                                                                              | qemu-system-x86_64           |
-| MANAGER_QEMU_USE_SUDO                      | Whether to use sudo to run QEMU.                                                                                | false                        |
-| MANAGER_QEMU_ENABLE_SEV                    | Whether to enable Secure Encrypted Virtualization (SEV).                                                        | false                        |
-| MANAGER_QEMU_ENABLE_SEV_SNP                | Whether to enable Secure Nested Paging (SEV-SNP).                                                               | true                         |
-| MANAGER_QEMU_ENABLE_KVM                    | Whether to enable the Kernel-based Virtual Machine (KVM) acceleration.                                          | true                         |
-| MANAGER_QEMU_MACHINE                       | The machine type for QEMU.                                                                                      | q35                          |
-| MANAGER_QEMU_CPU                           | The CPU model for QEMU.                                                                                         | EPYC                         |
-| MANAGER_QEMU_SMP_COUNT                     | The number of virtual CPUs.                                                                                     | 4                            |
-| MANAGER_QEMU_SMP_MAXCPUS                   | The maximum number of virtual CPUs.                                                                             | 64                           |
-| MANAGER_QEMU_MEM_ID                        | The ID for the memory device.                                                                                   | ram1                         |
-| MANAGER_QEMU_KERNEL_HASH                   | Whether to enable kernel hash verification.                                                                     | false                        |
-| MANAGER_QEMU_NO_GRAPHIC                    | Whether to disable the graphical display.                                                                       | true                         |
-| MANAGER_QEMU_MONITOR                       | The type of monitor to use.                                                                                     | pty                          |
-| MANAGER_QEMU_HOST_FWD_RANGE                | The range of host ports to forward.                                                                             | 6100-6200                    |
+| Variable                                   | Description                                                                                                      | Default                        |
+| ------------------------------------------ | ---------------------------------------------------------------------------------------------------------------- | ------------------------------ |
+| COCOS_JAEGER_URL                           | The URL for the Jaeger tracing endpoint.                                                                         | http://localhost:4318          |
+| COCOS_JAEGER_TRACE_RATIO                   | The ratio of traces to sample.                                                                                   | 1.0                            |
+| MANAGER_INSTANCE_ID                        | The instance ID for the manager service.                                                                         |                                |
+| MANAGER_ATTESTATION_POLICY_BINARY          | The file path for the attestation policy binarie.                                                                | ../../build/attestation_policy |
+| MANAGER_IGVMMEASURE_BINARY                 | The file path for the igvmmeasure binarie.                                                                       | ../../build/igvmmeasure        |
+| MANAGER_PCR_VALUES                         | The file path for the file with the expected PCR values.                                                         |                                |
+| MANAGER_GRPC_CLIENT_CERT                   | The file path for the client certificate.                                                                        |                                |
+| MANAGER_GRPC_CLIENT_KEY                    | The file path for the client private key.                                                                        |                                |
+| MANAGER_GRPC_SERVER_CA_CERTS               | The file path for the server CA certificate(s).                                                                  |                                |
+| MANAGER_GRPC_URL                           | The URL for the gRPC endpoint.                                                                                   | localhost:7001                 |
+| MANAGER_GRPC_TIMEOUT                       | The timeout for gRPC requests.                                                                                   | 60s                            |
+| MANAGER_EOS_VERSION                        | The EOS version used for booting SVMs.                                                                           |                                |
+| MANAGER_INSTANCE_ID                        | Manager service instance ID                                                                                      |                                |
+| MANAGER_QEMU_MEMORY_SIZE                   | The total memory size for the virtual machine. Can be specified in a human-readable format like "2048M" or "4G". | 2048M                          |
+| MANAGER_QEMU_MEMORY_SLOTS                  | The number of memory slots for the virtual machine.                                                              | 5                              |
+| MANAGER_QEMU_MAX_MEMORY                    | The maximum memory size for the virtual machine. Can be specified in a human-readable format like "30G".         | 30G                            |
+| MANAGER_QEMU_OVMF_CODE_IF                  | The interface type for the OVMF code.                                                                            | pflash                         |
+| MANAGER_QEMU_OVMF_CODE_FORMAT              | The format of the OVMF code file.                                                                                | raw                            |
+| MANAGER_QEMU_OVMF_CODE_UNIT                | The unit number for the OVMF code.                                                                               | 0                              |
+| MANAGER_QEMU_OVMF_CODE_FILE                | The file path for the OVMF code.                                                                                 | /usr/share/OVMF/OVMF_CODE.fd   |
+| MANAGER_QEMU_OVMF_VERSION                  | The version number of EDKII from which OVMF was built                                                            | edk2-stable202408              |
+| MANAGER_QEMU_OVMF_CODE_READONLY            | Whether the OVMF code should be read-only.                                                                       | on                             |
+| MANAGER_QEMU_OVMF_VARS_IF                  | The interface type for the OVMF variables.                                                                       | pflash                         |
+| MANAGER_QEMU_OVMF_VARS_FORMAT              | The format of the OVMF variables file.                                                                           | raw                            |
+| MANAGER_QEMU_OVMF_VARS_UNIT                | The unit number for the OVMF variables.                                                                          | 1                              |
+| MANAGER_QEMU_OVMF_VARS_FILE                | The file path for the OVMF variables.                                                                            | /usr/share/OVMF/OVMF_VARS.fd   |
+| MANAGER_QEMU_NETDEV_ID                     | The ID for the network device.                                                                                   | vmnic                          |
+| MANAGER_QEMU_HOST_FWD_AGENT                | The port number for the host forward agent.                                                                      | 7020                           |
+| MANAGER_QEMU_GUEST_FWD_AGENT               | The port number for the guest forward agent.                                                                     | 7002                           |
+| MANAGER_QEMU_VIRTIO_NET_PCI_DISABLE_LEGACY | Whether to disable the legacy PCI device.                                                                        | on                             |
+| MANAGER_QEMU_VIRTIO_NET_PCI_IOMMU_PLATFORM | Whether to enable the IOMMU platform for the virtio-net PCI device.                                              | true                           |
+| MANAGER_QEMU_VIRTIO_NET_PCI_ADDR           | The PCI address for the virtio-net PCI device.                                                                   | 0x2                            |
+| MANAGER_QEMU_VIRTIO_NET_PCI_ROMFILE        | The file path for the ROM image for the virtio-net PCI device.                                                   |                                |
+| MANAGER_QEMU_DISK_IMG_KERNEL_FILE          | The file path for the kernel image.                                                                              | img/bzImage                    |
+| MANAGER_QEMU_DISK_IMG_ROOTFS_FILE          | The file path for the root filesystem image.                                                                     | img/rootfs.cpio.gz             |
+| MANAGER_QEMU_SEV_ID                        | The ID for the Secure Encrypted Virtualization (SEV) device.                                                     | sev0                           |
+| MANAGER_QEMU_SEV_CBITPOS                   | The position of the C-bit in the physical address.                                                               | 51                             |
+| MANAGER_QEMU_SEV_REDUCED_PHYS_BITS         | The number of reduced physical address bits for SEV.                                                             | 1                              |
+| MANAGER_QEMU_ENABLE_HOST_DATA              | Enable additional data for the SEV host.                                                                         | false                          |
+| MANAGER_QEMU_HOST_DATA                     | Additional data for the SEV host.                                                                                |                                |
+| MANAGER_QEMU_IGVM_ID                       | The ID of the IGVM file.                                                                                         | igvm0                          |
+| MANAGER_QEMU_IGVM_FILE                     | The file path to the IGVM file.                                                                                  | /root/coconut-qemu.igvm        |
+| MANAGER_QEMU_VSOCK_ID                      | The ID for the virtual socket device.                                                                            | vhost-vsock-pci0               |
+| MANAGER_QEMU_VSOCK_GUEST_CID               | The guest-side CID (Context ID) for the virtual socket device.                                                   | 3                              |
+| MANAGER_QEMU_VSOCK_VNC                     | Whether to enable the virtual socket device for VNC.                                                             | 0                              |
+| MANAGER_QEMU_BIN_PATH                      | The file path for the QEMU binary.                                                                               | qemu-system-x86_64             |
+| MANAGER_QEMU_USE_SUDO                      | Whether to use sudo to run QEMU.                                                                                 | false                          |
+| MANAGER_QEMU_ENABLE_SEV                    | Whether to enable Secure Encrypted Virtualization (SEV).                                                         | false                          |
+| MANAGER_QEMU_ENABLE_SEV_SNP                | Whether to enable Secure Nested Paging (SEV-SNP).                                                                | true                           |
+| MANAGER_QEMU_ENABLE_KVM                    | Whether to enable the Kernel-based Virtual Machine (KVM) acceleration.                                           | true                           |
+| MANAGER_QEMU_MACHINE                       | The machine type for QEMU.                                                                                       | q35                            |
+| MANAGER_QEMU_CPU                           | The CPU model for QEMU.                                                                                          | EPYC                           |
+| MANAGER_QEMU_SMP_COUNT                     | The number of virtual CPUs.                                                                                      | 4                              |
+| MANAGER_QEMU_SMP_MAXCPUS                   | The maximum number of virtual CPUs.                                                                              | 64                             |
+| MANAGER_QEMU_MEM_ID                        | The ID for the memory device.                                                                                    | ram1                           |
+| MANAGER_QEMU_NO_GRAPHIC                    | Whether to disable the graphical display.                                                                        | true                           |
+| MANAGER_QEMU_MONITOR                       | The type of monitor to use.                                                                                      | pty                            |
+| MANAGER_QEMU_HOST_FWD_RANGE                | The range of host ports to forward.                                                                              | 6100-6200                      |
 
 ## Setup
 
@@ -88,21 +92,11 @@ sudo apt install qemu-kvm
 
 Create `img` directory in `cmd/manager`.
 
-#### Add Vsock
-The necessary kernel modules must be loaded on the hypervisor. To check if `vhost_vsock` is loaded run:
-```shell
-lsmod | grep vhost_vsock
-```
+#### Virtual filesystem
 
-If `vhost_vsock` is not loaded run the following commands:
+9P (or Plan 9 Filesystem) in QEMU is a lightweight, network-based file-sharing protocol. In Cocos, the 9P is used to transfer environment variables and TLS certificates for cloud communication from the Manager to the Agent.
 
-```shell
-sudo modprobe vhost_vsock
-ls -l /dev/vhost-vsock
-# crw-rw-rw- 1 root kvm 10, 241 Jan 16 12:05 /dev/vhost-vsock
-ls -l /dev/vsock
-# crw-rw-rw- 1 root root 10, 121 Jan 16 12:05 /dev/vsock
-```
+You should define the environment variables in a file called environment. For the number and meaning of the environment variables, please refer to the Agent [Readme](https://github.com/ultravioletrs/cocos/blob/main/agent/README.md).
 
 ### Prepare Cocos HAL
 
@@ -122,10 +116,32 @@ OVMF_CODE=/usr/share/OVMF/OVMF_CODE.fd
 
 sudo find / -name OVMF_VARS.fd
 # => /usr/share/OVMF/OVMF_VARS.fd
-OVMF_VARS=/usr/share/OVMF/OVMF_VARS.fd
 
+# Create a local copy of OVMF_VARS.
+cp /usr/share/OVMF/OVMF_VARS.fd .
+
+# Create a directory for the environment file and the certificates for cloud certificates.
+mkdir env
+mkdir certs
+
+# Enter the env directory and create the environemnt file.
+cd env
+touch environment
+
+# Define Computations endpoint URL for agent.
+# Make sure the Computation endpoint is running (like Cocos Prism).
+echo AGENT_CVM_GRPC_URL=localhost:7001 >> ./environment
+# Define log level for the agent.
+echo AGENT_LOG_LEVEL=debug >> ./environment
+
+# Return to cmd/manager
+cd ..
+
+OVMF_VARS=./OVMF_VARS.fd
 KERNEL="img/bzImage"
 INITRD="img/rootfs.cpio.gz"
+ENV_PATH=./env
+CERTH_PATH=./certs
 
 qemu-system-x86_64 \
     -enable-kvm \
@@ -135,15 +151,19 @@ qemu-system-x86_64 \
     -m 2048M,slots=5,maxmem=10240M \
     -no-reboot \
     -drive if=pflash,format=raw,unit=0,file=$OVMF_CODE,readonly=on \
+    -drive if=pflash,format=raw,unit=1,file=$OVMF_VARS \
     -netdev user,id=vmnic,hostfwd=tcp::7020-:7002 \
     -device virtio-net-pci,disable-legacy=on,iommu_platform=true,netdev=vmnic,romfile= \
-    -device vhost-vsock-pci,id=vhost-vsock-pci0,guest-cid=3 -vnc :0 \
     -kernel $KERNEL \
     -append "earlyprintk=serial console=ttyS0" \
     -initrd $INITRD \
     -nographic \
     -monitor pty \
-    -monitor unix:monitor,server,nowait
+    -monitor unix:monitor,server,nowait \
+    -fsdev local,id=env_fs,path=$ENV_PATH,security_model=mapped \
+    -device virtio-9p-pci,fsdev=env_fs,mount_tag=env_share \
+    -fsdev local,id=cert_fs,path=$CERTH_PATH,security_model=mapped \
+    -device virtio-9p-pci,fsdev=cert_fs,mount_tag=certs_share
 ```
 Once the VM is booted press enter and on the login use username `root`.
 
@@ -188,10 +208,21 @@ MANAGER_QEMU_OVMF_VARS_FILE=/usr/share/OVMF/OVMF_VARS.fd
 
 NB: we set environment variables that we will use in the shell process where we run `manager`.
 
+### Trusted Platform Module (TPM)
+
+The Trusted Platform Module (TPM) plays a fundamental role in this process by providing a tamper-resistant foundation for cryptographic operations, securing sensitive artifacts, measuring system state, and enabling attestation mechanisms.
+
+### IGVM
+
+An IGVM file contains all the necessary information to launch a virtual machine on different virtualization platforms. It includes setup commands for the guest system and verification data to ensure the VM is loaded securely and correctly.
+
+Cocos uses the [COCONUT-SVSM](https://github.com/coconut-svsm/svsm/blob/main/Documentation/docs/installation/INSTALL.md) for the vTPM. The IGVM file contains the OVMF file and the vTPM.
 
 ## Deployment
 
-To start the service, execute the following shell script (note a server needs to be running see  [here](../test/computations/README.md)):
+To start the service, execute the following shell script (note a server needs to be running see [here](../test/cvms/README.md)):
+
+The manager can be started as a *systemd* service or a standalone executable. To start the manager as a systemd service, look at the systemd service script [here](https://github.com/ultravioletrs/cocos/blob/main/init/systemd/cocos-manager.service). The environment variables are defined in the `cocos-manager.env` file. Below are examples of how to start the manager.
 
 ```bash
 # Download the latest version of the service
@@ -203,13 +234,14 @@ cd cocos
 make manager
 
 # Set the environment variables and run the service
-MANAGER_GRPC_URL=localhost:7001
+MANAGER_GRPC_URL=localhost:7001 \
 MANAGER_LOG_LEVEL=debug \
 MANAGER_QEMU_USE_SUDO=false \
 MANAGER_QEMU_ENABLE_SEV=false \
 ./build/cocos-manager
 ```
 
+
 To enable [AMD SEV](https://www.amd.com/en/developer/sev.html) support, start manager like this 
 
 ```sh
@@ -221,7 +253,7 @@ MANAGER_QEMU_SEV_CBITPOS=51 \
 ./build/cocos-manager
 ```
 
-To build the OVMF with the kernel hash capability, we must build the AmdSev package of OVMF. The result of the build should be a single `OVMF.fd` file (unlike the regular two OVFM files). The OVMF package is located at `OvmfPkg/AmdSev/AmdSevX64.dsc`.
+To start SEV-SNP, define the IGVM file that contains the vTPM and the OVMF (combined OVMF_CODE and OVMF_VARS) of the CVM.
 
 To enable [AMD SEV-SNP](https://www.amd.com/en/developer/sev.html) support, start manager like this 
 
@@ -232,60 +264,10 @@ MANAGER_QEMU_ENABLE_SEV=false \
 MANAGER_QEMU_ENABLE_SEV_SNP=true \
 MANAGER_QEMU_SEV_CBITPOS=51 \
 MANAGER_QEMU_BIN_PATH=<path to QEMU binary> \
-MANAGER_QEMU_QEMU_OVMF_CODE_FILE=<path to OVMF.fd Amd Sev built package> \
+MANAGER_QEMU_IGVM_FILE=<path to IGVM file> \
 ./build/cocos-manager
 ```
 
-To include the kernel hash into the measurement of the attestation report (SEV or SEV-SNP), start manager like this
-
-```sh
-MANAGER_GRPC_URL=localhost:7001 \
-MANAGER_LOG_LEVEL=debug \
-MANAGER_QEMU_ENABLE_SEV=false \
-MANAGER_QEMU_ENABLE_SEV_SNP=true \
-MANAGER_QEMU_SEV_CBITPOS=51 \
-MANAGER_QEMU_KERNEL_HASH=true \
-MANAGER_QEMU_BIN_PATH=<path to QEMU binary> \
-MANAGER_QEMU_QEMU_OVMF_CODE_FILE=<path to OVMF.fd Amd Sev built package> \
-./build/cocos-manager
-```
-
-### Verifying VM launch
-
-NB: To verify that the manager successfully launched the VM, you need to open three terminals on the same machine. In one terminal, you need to launch the computations server by executing (with the environment variables of choice):
-
-```bash
-go run ./test/computations/main.go <dataset path> <algo path>
-```
-
-and in the second the manager by executing (with the environment variables of choice):
-
-```bash
-go run ./cmd/manager/main.go
-```
-
-Ensure that the Manager can connect to the Manager test server by setting the MANAGER_GRPC_PORT with the port value of the Manager test server. In the last terminal, you can run the verification commands.
-
-To verify that the manager launched the VM successfully, run the following command:
-
-```sh
-ps aux | grep qemu-system-x86_64
-```
-
-You should get something similar to this
-```
-darko     324763 95.3  6.0 6398136 981044 ?      Sl   16:17   0:15 /usr/bin/qemu-system-x86_64 -enable-kvm -machine q35 -cpu EPYC -smp 4,maxcpus=64 -m 4096M,slots=5,maxmem=30G -drive if=pflash,format=raw,unit=0,file=/usr/share/OVMF/OVMF_CODE.fd,readonly=on -drive if=pflash,format=raw,unit=1,file=img/OVMF_VARS.fd -device virtio-scsi-pci,id=scsi,disable-legacy=on,iommu_platform=true -drive file=img/focal-server-cloudimg-amd64.img,if=none,id=disk0,format=qcow2 -device scsi-hd,drive=disk0 -netdev user,id=vmnic,hostfwd=tcp::2222-:22,hostfwd=tcp::9301-:9031,hostfwd=tcp::7020-:7002 -device virtio-net-pci,disable-legacy=on,iommu_platform=true,netdev=vmnic,romfile= -nographic -monitor pty
-```
-
-If you run a command as `sudo`, you should get the output similar to this one
-
-```
-root       37982  0.0  0.0   9444  4572 pts/0    S+   16:18   0:00 sudo /usr/local/bin/qemu-system-x86_64 -enable-kvm -machine q35 -cpu EPYC -smp 4,maxcpus=64 -m 4096M,slots=5,maxmem=30G -drive if=pflash,format=raw,unit=0,file=/usr/share/OVMF/OVMF_CODE.fd,readonly=on -drive if=pflash,format=raw,unit=1,file=img/OVMF_VARS.fd -device virtio-scsi-pci,id=scsi,disable-legacy=on,iommu_platform=true -drive file=img/focal-server-cloudimg-amd64.img,if=none,id=disk0,format=qcow2 -device scsi-hd,drive=disk0 -netdev user,id=vmnic,hostfwd=tcp::2222-:22,hostfwd=tcp::9301-:9031,hostfwd=tcp::7020-:7002 -device virtio-net-pci,disable-legacy=on,iommu_platform=true,netdev=vmnic,romfile= -object sev-guest,id=sev0,cbitpos=51,reduced-phys-bits=1 -machine memory-encryption=sev0 -nographic -monitor pty
-root       37989  122 13.1 5345816 4252312 pts/0 Sl+  16:19   0:04 /usr/local/bin/qemu-system-x86_64 -enable-kvm -machine q35 -cpu EPYC -smp 4,maxcpus=64 -m 4096M,slots=5,maxmem=30G -drive if=pflash,format=raw,unit=0,file=/usr/share/OVMF/OVMF_CODE.fd,readonly=on -drive if=pflash,format=raw,unit=1,file=img/OVMF_VARS.fd -device virtio-scsi-pci,id=scsi,disable-legacy=on,iommu_platform=true -drive file=img/focal-server-cloudimg-amd64.img,if=none,id=disk0,format=qcow2 -device scsi-hd,drive=disk0 -netdev user,id=vmnic,hostfwd=tcp::2222-:22,hostfwd=tcp::9301-:9031,hostfwd=tcp::7020-:7002 -device virtio-net-pci,disable-legacy=on,iommu_platform=true,netdev=vmnic,romfile= -object sev-guest,id=sev0,cbitpos=51,reduced-phys-bits=1 -machine memory-encryption=sev0 -nographic -monitor pty
-```
-
-The two processes are due to the fact that we run the command `/usr/bin/qemu-system-x86_64` as `sudo`, so there is one process for `sudo` command and the other for `/usr/bin/qemu-system-x86_64`.
-
 ### Troubleshooting
 
 If the `ps aux | grep qemu-system-x86_64` give you something like this

manager/agentEventsLogs.go

@@ -1,66 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package manager
-
-import (
-	"fmt"
-	"regexp"
-	"strconv"
-
-	"github.com/ultravioletrs/cocos/pkg/manager"
-	"google.golang.org/protobuf/types/known/timestamppb"
-)
-
-var (
-	errFailedToParseCID    = fmt.Errorf("failed to parse computation ID")
-	errComputationNotFound = fmt.Errorf("computation not found")
-)
-
-func (ms *managerService) computationIDFromAddress(address string) (string, error) {
-	re := regexp.MustCompile(`vm\((\d+)\)`)
-	matches := re.FindStringSubmatch(address)
-
-	if len(matches) > 1 {
-		cid, err := strconv.Atoi(matches[1])
-		if err != nil {
-			return "", err
-		}
-		return ms.findComputationID(cid)
-	}
-	return "", errFailedToParseCID
-}
-
-func (ms *managerService) findComputationID(cid int) (string, error) {
-	ms.mu.Lock()
-	defer ms.mu.Unlock()
-	for cmpID, vm := range ms.vms {
-		if vm.GetCID() == cid {
-			return cmpID, nil
-		}
-	}
-
-	return "", errComputationNotFound
-}
-
-func (ms *managerService) reportBrokenConnection(cmpID string) {
-	ms.eventsChan <- &ClientStreamMessage{
-		Message: &ClientStreamMessage_AgentEvent{
-			AgentEvent: &AgentEvent{
-				EventType:     ms.vms[cmpID].State(),
-				ComputationId: cmpID,
-				Status:        manager.Disconnected.String(),
-				Timestamp:     timestamppb.Now(),
-				Originator:    "manager",
-			},
-		},
-	}
-}
-
-func (ms *managerService) ReportBrokenConnection(addr string) {
-	cmpID, err := ms.computationIDFromAddress(addr)
-	if err != nil {
-		ms.logger.Warn(err.Error())
-		return
-	}
-	ms.reportBrokenConnection(cmpID)
-}

manager/agentEventsLogs_test.go

@@ -1,64 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package manager
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/assert"
-	"github.com/ultravioletrs/cocos/manager/qemu"
-	"github.com/ultravioletrs/cocos/manager/vm"
-	"github.com/ultravioletrs/cocos/pkg/manager"
-)
-
-func TestComputationIDFromAddress(t *testing.T) {
-	ms := &managerService{
-		vms: map[string]vm.VM{
-			"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"),
-		},
-	}
-
-	tests := []struct {
-		name    string
-		address string
-		want    string
-		wantErr bool
-	}{
-		{"Valid address", "vm(3)", "comp1", false},
-		{"Invalid address", "invalid", "", true},
-		{"Non-existent CID", "vm(10)", "", true},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got, err := ms.computationIDFromAddress(tt.address)
-			if tt.wantErr {
-				assert.Error(t, err)
-			} else {
-				assert.NoError(t, err)
-				assert.Equal(t, tt.want, got)
-			}
-		})
-	}
-}
-
-func TestReportBrokenConnection(t *testing.T) {
-	ms := &managerService{
-		eventsChan: make(chan *ClientStreamMessage, 1),
-		vms: map[string]vm.VM{
-			"comp1": qemu.NewVM(qemu.VMInfo{Config: qemu.Config{VSockConfig: qemu.VSockConfig{GuestCID: 3}}}, func(event interface{}) error { return nil }, "comp1"),
-		},
-	}
-
-	ms.reportBrokenConnection("comp1")
-
-	select {
-	case msg := <-ms.eventsChan:
-		assert.Equal(t, "comp1", msg.GetAgentEvent().ComputationId)
-		assert.Equal(t, manager.Disconnected.String(), msg.GetAgentEvent().Status)
-		assert.Equal(t, "manager", msg.GetAgentEvent().Originator)
-	default:
-		t.Error("Expected message in eventsChan, but none received")
-	}
-}

manager/api/grpc/client.go

@@ -1,242 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package grpc
-
-import (
-	"context"
-	"log/slog"
-	"sync"
-	"time"
-
-	"github.com/absmach/magistrala/pkg/errors"
-	"github.com/ultravioletrs/cocos/manager"
-	"github.com/ultravioletrs/cocos/manager/qemu"
-	"golang.org/x/sync/errgroup"
-	"google.golang.org/protobuf/proto"
-)
-
-var (
-	errTerminationFromServer = errors.New("server requested client termination")
-	errCorruptedManifest     = errors.New("received manifest may be corrupted")
-	sendTimeout              = 5 * time.Second
-)
-
-type ManagerClient struct {
-	stream        manager.ManagerService_ProcessClient
-	svc           manager.Service
-	messageQueue  chan *manager.ClientStreamMessage
-	logger        *slog.Logger
-	runReqManager *runRequestManager
-}
-
-// NewClient returns new gRPC client instance.
-func NewClient(stream manager.ManagerService_ProcessClient, svc manager.Service, messageQueue chan *manager.ClientStreamMessage, logger *slog.Logger) ManagerClient {
-	return ManagerClient{
-		stream:        stream,
-		svc:           svc,
-		messageQueue:  messageQueue,
-		logger:        logger,
-		runReqManager: newRunRequestManager(),
-	}
-}
-
-func (client ManagerClient) Process(ctx context.Context, cancel context.CancelFunc) error {
-	eg, ctx := errgroup.WithContext(ctx)
-
-	eg.Go(func() error {
-		return client.handleIncomingMessages(ctx)
-	})
-
-	eg.Go(func() error {
-		return client.handleOutgoingMessages(ctx)
-	})
-
-	return eg.Wait()
-}
-
-func (client ManagerClient) handleIncomingMessages(ctx context.Context) error {
-	for {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		default:
-			req, err := client.stream.Recv()
-			if err != nil {
-				return err
-			}
-			if err := client.processIncomingMessage(ctx, req); err != nil {
-				return err
-			}
-		}
-	}
-}
-
-func (client ManagerClient) processIncomingMessage(ctx context.Context, req *manager.ServerStreamMessage) error {
-	switch mes := req.Message.(type) {
-	case *manager.ServerStreamMessage_RunReqChunks:
-		return client.handleRunReqChunks(ctx, mes)
-	case *manager.ServerStreamMessage_TerminateReq:
-		return client.handleTerminateReq(mes)
-	case *manager.ServerStreamMessage_StopComputation:
-		go client.handleStopComputation(ctx, mes)
-	case *manager.ServerStreamMessage_AttestationPolicyReq:
-		go client.handleAttestationPolicyReq(ctx, mes)
-	case *manager.ServerStreamMessage_SvmInfoReq:
-		go client.handleSVMInfoReq(ctx, mes)
-	default:
-		return errors.New("unknown message type")
-	}
-	return nil
-}
-
-func (client *ManagerClient) handleRunReqChunks(ctx context.Context, mes *manager.ServerStreamMessage_RunReqChunks) error {
-	buffer, complete := client.runReqManager.addChunk(mes.RunReqChunks.Id, mes.RunReqChunks.Data, mes.RunReqChunks.IsLast)
-
-	if complete {
-		var runReq manager.ComputationRunReq
-		if err := proto.Unmarshal(buffer, &runReq); err != nil {
-			return errors.Wrap(err, errCorruptedManifest)
-		}
-
-		go client.executeRun(ctx, &runReq)
-	}
-
-	return nil
-}
-
-func (client ManagerClient) executeRun(ctx context.Context, runReq *manager.ComputationRunReq) {
-	port, err := client.svc.Run(ctx, runReq)
-	if err != nil {
-		client.logger.Warn(err.Error())
-		return
-	}
-	runRes := &manager.ClientStreamMessage_RunRes{
-		RunRes: &manager.RunResponse{
-			AgentPort:     port,
-			ComputationId: runReq.Id,
-		},
-	}
-	client.sendMessage(&manager.ClientStreamMessage{Message: runRes})
-}
-
-func (client ManagerClient) handleTerminateReq(mes *manager.ServerStreamMessage_TerminateReq) error {
-	return errors.Wrap(errTerminationFromServer, errors.New(mes.TerminateReq.Message))
-}
-
-func (client ManagerClient) handleStopComputation(ctx context.Context, mes *manager.ServerStreamMessage_StopComputation) {
-	msg := &manager.ClientStreamMessage_StopComputationRes{
-		StopComputationRes: &manager.StopComputationResponse{
-			ComputationId: mes.StopComputation.ComputationId,
-		},
-	}
-	if err := client.svc.Stop(ctx, mes.StopComputation.ComputationId); err != nil {
-		msg.StopComputationRes.Message = err.Error()
-	}
-	client.sendMessage(&manager.ClientStreamMessage{Message: msg})
-}
-
-func (client ManagerClient) handleAttestationPolicyReq(ctx context.Context, mes *manager.ServerStreamMessage_AttestationPolicyReq) {
-	res, err := client.svc.FetchAttestationPolicy(ctx, mes.AttestationPolicyReq.Id)
-	if err != nil {
-		client.logger.Warn(err.Error())
-		return
-	}
-	info := &manager.ClientStreamMessage_AttestationPolicy{
-		AttestationPolicy: &manager.AttestationPolicy{
-			Info: res,
-			Id:   mes.AttestationPolicyReq.Id,
-		},
-	}
-	client.sendMessage(&manager.ClientStreamMessage{Message: info})
-}
-
-func (client ManagerClient) handleSVMInfoReq(ctx context.Context, mes *manager.ServerStreamMessage_SvmInfoReq) {
-	ovmfVersion, cpuNum, cpuType, eosVersion := client.svc.ReturnSVMInfo(ctx)
-	info := &manager.ClientStreamMessage_SvmInfo{
-		SvmInfo: &manager.SVMInfo{
-			OvmfVersion: ovmfVersion,
-			CpuNum:      int32(cpuNum),
-			CpuType:     cpuType,
-			KernelCmd:   qemu.KernelCommandLine,
-			EosVersion:  eosVersion,
-			Id:          mes.SvmInfoReq.Id,
-		},
-	}
-	client.sendMessage(&manager.ClientStreamMessage{Message: info})
-}
-
-func (client ManagerClient) handleOutgoingMessages(ctx context.Context) error {
-	for {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case mes := <-client.messageQueue:
-			if err := client.stream.Send(mes); err != nil {
-				return err
-			}
-		}
-	}
-}
-
-func (client ManagerClient) sendMessage(mes *manager.ClientStreamMessage) {
-	ctx, cancel := context.WithTimeout(context.Background(), sendTimeout)
-	defer cancel()
-
-	select {
-	case client.messageQueue <- mes:
-	case <-ctx.Done():
-		client.logger.Warn("Failed to send message: timeout exceeded")
-	}
-}
-
-type runRequestManager struct {
-	requests map[string]*runRequest
-	mu       sync.Mutex
-}
-
-type runRequest struct {
-	buffer    []byte
-	lastChunk time.Time
-	timer     *time.Timer
-}
-
-func newRunRequestManager() *runRequestManager {
-	return &runRequestManager{
-		requests: make(map[string]*runRequest),
-	}
-}
-
-func (m *runRequestManager) addChunk(id string, chunk []byte, isLast bool) ([]byte, bool) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	req, exists := m.requests[id]
-	if !exists {
-		req = &runRequest{
-			buffer:    make([]byte, 0),
-			lastChunk: time.Now(),
-			timer:     time.AfterFunc(runReqTimeout, func() { m.timeoutRequest(id) }),
-		}
-		m.requests[id] = req
-	}
-
-	req.buffer = append(req.buffer, chunk...)
-	req.lastChunk = time.Now()
-	req.timer.Reset(runReqTimeout)
-
-	if isLast {
-		delete(m.requests, id)
-		req.timer.Stop()
-		return req.buffer, true
-	}
-
-	return nil, false
-}
-
-func (m *runRequestManager) timeoutRequest(id string) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-
-	delete(m.requests, id)
-	// Log timeout or handle it as needed
-}

manager/api/grpc/client_test.go

@@ -1,322 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package grpc
-
-import (
-	"context"
-	"testing"
-	"time"
-
-	mglog "github.com/absmach/magistrala/logger"
-	"github.com/absmach/magistrala/pkg/errors"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
-	"github.com/ultravioletrs/cocos/manager"
-	"github.com/ultravioletrs/cocos/manager/mocks"
-	"github.com/ultravioletrs/cocos/manager/qemu"
-	"google.golang.org/grpc"
-	"google.golang.org/protobuf/proto"
-)
-
-type mockStream struct {
-	mock.Mock
-	grpc.ClientStream
-}
-
-func (m *mockStream) Recv() (*manager.ServerStreamMessage, error) {
-	args := m.Called()
-	return args.Get(0).(*manager.ServerStreamMessage), args.Error(1)
-}
-
-func (m *mockStream) Send(msg *manager.ClientStreamMessage) error {
-	args := m.Called(msg)
-	return args.Error(0)
-}
-
-func TestManagerClient_Process1(t *testing.T) {
-	tests := []struct {
-		name        string
-		setupMocks  func(mockStream *mockStream, mockSvc *mocks.Service)
-		expectError bool
-		errorMsg    string
-	}{
-		{
-			name: "Stop computation",
-			setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service) {
-				mockStream.On("Recv").Return(&manager.ServerStreamMessage{
-					Message: &manager.ServerStreamMessage_StopComputation{
-						StopComputation: &manager.StopComputation{},
-					},
-				}, nil)
-				mockStream.On("Send", mock.Anything).Return(nil)
-				mockSvc.On("Stop", mock.Anything, mock.Anything).Return(nil)
-			},
-			expectError: true,
-			errorMsg:    "context deadline exceeded",
-		},
-		{
-			name: "Terminate request",
-			setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service) {
-				mockStream.On("Recv").Return(&manager.ServerStreamMessage{
-					Message: &manager.ServerStreamMessage_TerminateReq{
-						TerminateReq: &manager.Terminate{},
-					},
-				}, nil)
-			},
-			expectError: true,
-			errorMsg:    errTerminationFromServer.Error(),
-		},
-		{
-			name: "Attestation Policy request",
-			setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service) {
-				mockStream.On("Recv").Return(&manager.ServerStreamMessage{
-					Message: &manager.ServerStreamMessage_AttestationPolicyReq{
-						AttestationPolicyReq: &manager.AttestationPolicyReq{},
-					},
-				}, nil)
-				mockStream.On("Send", mock.Anything).Return(nil).Once()
-				mockSvc.On("FetchAttestationPolicy", mock.Anything, mock.Anything).Return(nil, assert.AnError)
-			},
-			expectError: true,
-		},
-		{
-			name: "Run request chunks",
-			setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service) {
-				mockStream.On("Recv").Return(&manager.ServerStreamMessage{
-					Message: &manager.ServerStreamMessage_RunReqChunks{
-						RunReqChunks: &manager.RunReqChunks{},
-					},
-				}, nil)
-				mockStream.On("Send", mock.Anything).Return(nil).Once()
-				mockSvc.On("Run", mock.Anything, mock.Anything).Return("", assert.AnError).Once()
-			},
-			expectError: true,
-		},
-		{
-			name: "Receive error",
-			setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service) {
-				mockStream.On("Recv").Return(&manager.ServerStreamMessage{}, assert.AnError)
-			},
-			expectError: true,
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			mockStream := new(mockStream)
-			mockSvc := new(mocks.Service)
-			messageQueue := make(chan *manager.ClientStreamMessage, 10)
-			logger := mglog.NewMock()
-
-			client := NewClient(mockStream, mockSvc, messageQueue, logger)
-
-			ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
-			defer cancel()
-
-			tc.setupMocks(mockStream, mockSvc)
-
-			err := client.Process(ctx, cancel)
-
-			if tc.expectError {
-				assert.Error(t, err)
-				if tc.errorMsg != "" {
-					assert.Contains(t, err.Error(), tc.errorMsg)
-				}
-			} else {
-				assert.NoError(t, err)
-			}
-		})
-	}
-}
-
-func TestManagerClient_handleRunReqChunks(t *testing.T) {
-	mockStream := new(mockStream)
-	mockSvc := new(mocks.Service)
-	messageQueue := make(chan *manager.ClientStreamMessage, 10)
-	logger := mglog.NewMock()
-
-	client := NewClient(mockStream, mockSvc, messageQueue, logger)
-
-	runReq := &manager.ComputationRunReq{
-		Id: "test-id",
-	}
-	runReqBytes, _ := proto.Marshal(runReq)
-
-	chunk1 := &manager.ServerStreamMessage_RunReqChunks{
-		RunReqChunks: &manager.RunReqChunks{
-			Id:     "chunk-1",
-			Data:   runReqBytes[:len(runReqBytes)/2],
-			IsLast: false,
-		},
-	}
-	chunk2 := &manager.ServerStreamMessage_RunReqChunks{
-		RunReqChunks: &manager.RunReqChunks{
-			Id:     "chunk-1",
-			Data:   runReqBytes[len(runReqBytes)/2:],
-			IsLast: true,
-		},
-	}
-
-	mockSvc.On("Run", mock.Anything, mock.AnythingOfType("*manager.ComputationRunReq")).Return("8080", nil)
-
-	err := client.handleRunReqChunks(context.Background(), chunk1)
-	assert.NoError(t, err)
-
-	err = client.handleRunReqChunks(context.Background(), chunk2)
-	assert.NoError(t, err)
-
-	// Wait for the goroutine to finish
-	time.Sleep(50 * time.Millisecond)
-
-	mockSvc.AssertExpectations(t)
-	assert.Len(t, messageQueue, 1)
-
-	msg := <-messageQueue
-	runRes, ok := msg.Message.(*manager.ClientStreamMessage_RunRes)
-	assert.True(t, ok)
-	assert.Equal(t, "8080", runRes.RunRes.AgentPort)
-	assert.Equal(t, "test-id", runRes.RunRes.ComputationId)
-}
-
-func TestManagerClient_handleTerminateReq(t *testing.T) {
-	client := ManagerClient{}
-
-	terminateReq := &manager.ServerStreamMessage_TerminateReq{
-		TerminateReq: &manager.Terminate{
-			Message: "Test termination",
-		},
-	}
-
-	err := client.handleTerminateReq(terminateReq)
-	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "Test termination")
-	assert.True(t, errors.Contains(err, errTerminationFromServer))
-}
-
-func TestManagerClient_handleStopComputation(t *testing.T) {
-	mockStream := new(mockStream)
-	mockSvc := new(mocks.Service)
-	messageQueue := make(chan *manager.ClientStreamMessage, 10)
-	logger := mglog.NewMock()
-
-	client := NewClient(mockStream, mockSvc, messageQueue, logger)
-
-	stopReq := &manager.ServerStreamMessage_StopComputation{
-		StopComputation: &manager.StopComputation{
-			ComputationId: "test-comp-id",
-		},
-	}
-
-	mockSvc.On("Stop", mock.Anything, "test-comp-id").Return(nil)
-
-	client.handleStopComputation(context.Background(), stopReq)
-
-	// Wait for the goroutine to finish
-	time.Sleep(50 * time.Millisecond)
-
-	mockSvc.AssertExpectations(t)
-	assert.Len(t, messageQueue, 1)
-
-	msg := <-messageQueue
-	stopRes, ok := msg.Message.(*manager.ClientStreamMessage_StopComputationRes)
-	assert.True(t, ok)
-	assert.Equal(t, "test-comp-id", stopRes.StopComputationRes.ComputationId)
-	assert.Empty(t, stopRes.StopComputationRes.Message)
-}
-
-func TestManagerClient_handleAttestationPolicyReq(t *testing.T) {
-	t.Run("success", func(t *testing.T) {
-		mockStream := new(mockStream)
-		mockSvc := new(mocks.Service)
-		messageQueue := make(chan *manager.ClientStreamMessage, 10)
-		logger := mglog.NewMock()
-
-		client := NewClient(mockStream, mockSvc, messageQueue, logger)
-
-		infoReq := &manager.ServerStreamMessage_AttestationPolicyReq{
-			AttestationPolicyReq: &manager.AttestationPolicyReq{
-				Id: "test-info-id",
-			},
-		}
-
-		mockSvc.On("FetchAttestationPolicy", context.Background(), infoReq.AttestationPolicyReq.Id).Return([]byte("test-attestation-policy"), nil)
-
-		client.handleAttestationPolicyReq(context.Background(), infoReq)
-
-		// Wait for the goroutine to finish
-		time.Sleep(50 * time.Millisecond)
-
-		mockSvc.AssertExpectations(t)
-		assert.Len(t, messageQueue, 1)
-
-		msg := <-messageQueue
-		infoRes, ok := msg.Message.(*manager.ClientStreamMessage_AttestationPolicy)
-		assert.True(t, ok)
-		assert.Equal(t, "test-info-id", infoRes.AttestationPolicy.Id)
-		assert.Equal(t, []byte("test-attestation-policy"), infoRes.AttestationPolicy.Info)
-	})
-	t.Run("error", func(t *testing.T) {
-		mockStream := new(mockStream)
-		mockSvc := new(mocks.Service)
-		messageQueue := make(chan *manager.ClientStreamMessage, 10)
-		logger := mglog.NewMock()
-
-		client := NewClient(mockStream, mockSvc, messageQueue, logger)
-
-		infoReq := &manager.ServerStreamMessage_AttestationPolicyReq{
-			AttestationPolicyReq: &manager.AttestationPolicyReq{
-				Id: "test-info-id",
-			},
-		}
-
-		mockSvc.On("FetchAttestationPolicy", context.Background(), infoReq.AttestationPolicyReq.Id).Return(nil, assert.AnError)
-
-		client.handleAttestationPolicyReq(context.Background(), infoReq)
-
-		time.Sleep(50 * time.Millisecond)
-
-		mockSvc.AssertExpectations(t)
-		assert.Len(t, messageQueue, 0)
-	})
-}
-
-func TestManagerClient_handleSVMInfoReq(t *testing.T) {
-	mockStream := new(mockStream)
-	mockSvc := new(mocks.Service)
-	messageQueue := make(chan *manager.ClientStreamMessage, 10)
-	logger := mglog.NewMock()
-
-	client := NewClient(mockStream, mockSvc, messageQueue, logger)
-
-	mockSvc.On("ReturnSVMInfo", context.Background()).Return("edk2-stable202408", 4, "EPYC", "")
-
-	client.handleSVMInfoReq(context.Background(), &manager.ServerStreamMessage_SvmInfoReq{SvmInfoReq: &manager.SVMInfoReq{Id: "test-svm-info-id"}})
-
-	// Wait for the goroutine to finish
-	time.Sleep(50 * time.Millisecond)
-
-	mockSvc.AssertExpectations(t)
-	assert.Len(t, messageQueue, 1)
-
-	msg := <-messageQueue
-	infoRes, ok := msg.Message.(*manager.ClientStreamMessage_SvmInfo)
-	assert.True(t, ok)
-	assert.Equal(t, "edk2-stable202408", infoRes.SvmInfo.OvmfVersion)
-	assert.Equal(t, int32(4), infoRes.SvmInfo.CpuNum)
-	assert.Equal(t, "EPYC", infoRes.SvmInfo.CpuType)
-	assert.Equal(t, "", infoRes.SvmInfo.EosVersion)
-	assert.Equal(t, qemu.KernelCommandLine, infoRes.SvmInfo.KernelCmd)
-}
-
-func TestManagerClient_timeoutRequest(t *testing.T) {
-	rm := newRunRequestManager()
-	rm.requests["test-id"] = &runRequest{
-		timer:     time.NewTimer(100 * time.Millisecond),
-		buffer:    []byte("test-data"),
-		lastChunk: time.Now(),
-	}
-
-	rm.timeoutRequest("test-id")
-
-	assert.Len(t, rm.requests, 0)
-}

manager/api/grpc/server.go

@@ -3,17 +3,11 @@
 package grpc
 
 import (
-	"bytes"
 	"context"
 	"errors"
-	"io"
-	"time"
 
 	"github.com/ultravioletrs/cocos/manager"
-	"golang.org/x/sync/errgroup"
-	"google.golang.org/grpc/credentials"
-	"google.golang.org/grpc/peer"
-	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/emptypb"
 )
 
 var (
@@ -21,113 +15,58 @@ var (
 	ErrUnexpectedMsg                              = errors.New("unknown message type")
 )
 
-const (
-	bufferSize    = 1024 * 1024 // 1 MB
-	runReqTimeout = 30 * time.Second
-)
-
-type SendFunc func(*manager.ServerStreamMessage) error
-
 type grpcServer struct {
 	manager.UnimplementedManagerServiceServer
-	incoming chan *manager.ClientStreamMessage
-	svc      Service
-}
-
-type Service interface {
-	Run(ctx context.Context, ipAddress string, sendMessage SendFunc, authInfo credentials.AuthInfo)
+	svc manager.Service
 }
 
 // NewServer returns new AuthServiceServer instance.
-func NewServer(incoming chan *manager.ClientStreamMessage, svc Service) manager.ManagerServiceServer {
+func NewServer(svc manager.Service) manager.ManagerServiceServer {
 	return &grpcServer{
-		incoming: incoming,
-		svc:      svc,
+		svc: svc,
 	}
 }
 
-func (s *grpcServer) Process(stream manager.ManagerService_ProcessServer) error {
-	client, ok := peer.FromContext(stream.Context())
-	if !ok {
-		return errors.New("failed to get peer info")
-	}
-
-	eg, ctx := errgroup.WithContext(stream.Context())
-
-	eg.Go(func() error {
-		for {
-			select {
-			case <-ctx.Done():
-				return ctx.Err()
-			default:
-				req, err := stream.Recv()
-				if err != nil {
-					return err
-				}
-				s.incoming <- req
-			}
-		}
-	})
-
-	eg.Go(func() error {
-		sendMessage := func(msg *manager.ServerStreamMessage) error {
-			select {
-			case <-ctx.Done():
-				return ctx.Err()
-			default:
-				switch m := msg.Message.(type) {
-				case *manager.ServerStreamMessage_RunReq:
-					return s.sendRunReqInChunks(stream, m.RunReq)
-				default:
-					return stream.Send(msg)
-				}
-			}
-		}
-
-		s.svc.Run(ctx, client.Addr.String(), sendMessage, client.AuthInfo)
-		return nil
-	})
-
-	return eg.Wait()
-}
-
-func (s *grpcServer) sendRunReqInChunks(stream manager.ManagerService_ProcessServer, runReq *manager.ComputationRunReq) error {
-	data, err := proto.Marshal(runReq)
+func (s *grpcServer) CreateVm(ctx context.Context, req *manager.CreateReq) (*manager.CreateRes, error) {
+	port, id, err := s.svc.CreateVM(ctx, req)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	dataBuffer := bytes.NewBuffer(data)
-	buf := make([]byte, bufferSize)
-
-	for {
-		n, err := dataBuffer.Read(buf)
-		isLast := false
-
-		if err == io.EOF {
-			isLast = true
-		} else if err != nil {
-			return err
-		}
-
-		chunk := &manager.ServerStreamMessage{
-			Message: &manager.ServerStreamMessage_RunReqChunks{
-				RunReqChunks: &manager.RunReqChunks{
-					Id:     runReq.Id,
-					Data:   buf[:n],
-					IsLast: isLast,
-				},
-			},
-		}
-
-		if err := stream.Send(chunk); err != nil {
-			return err
-		}
-
-		if isLast {
-			break
-		}
-	}
-
-	return nil
+	return &manager.CreateRes{
+		ForwardedPort: port,
+		SvmId:         id,
+	}, nil
+}
+
+func (s *grpcServer) RemoveVm(ctx context.Context, req *manager.RemoveReq) (*emptypb.Empty, error) {
+	if err := s.svc.RemoveVM(ctx, req.SvmId); err != nil {
+		return nil, err
+	}
+
+	return &emptypb.Empty{}, nil
+}
+
+func (s *grpcServer) SVMInfo(ctx context.Context, req *manager.SVMInfoReq) (*manager.SVMInfoRes, error) {
+	ovmf, cpunum, cputype, eosversion := s.svc.ReturnSVMInfo(ctx)
+
+	return &manager.SVMInfoRes{
+		OvmfVersion: ovmf,
+		CpuNum:      int32(cpunum),
+		CpuType:     cputype,
+		EosVersion:  eosversion,
+		Id:          req.Id,
+	}, nil
+}
+
+func (s *grpcServer) AttestationPolicy(ctx context.Context, req *manager.AttestationPolicyReq) (*manager.AttestationPolicyRes, error) {
+	policy, err := s.svc.FetchAttestationPolicy(ctx, req.Id)
+	if err != nil {
+		return nil, err
+	}
+
+	return &manager.AttestationPolicyRes{
+		Info: policy,
+		Id:   req.Id,
+	}, nil
 }

manager/api/logging.go

@@ -27,9 +27,9 @@ func LoggingMiddleware(svc manager.Service, logger *slog.Logger) manager.Service
 	return &loggingMiddleware{logger, svc}
 }
 
-func (lm *loggingMiddleware) Run(ctx context.Context, mc *manager.ComputationRunReq) (agentAddr string, err error) {
+func (lm *loggingMiddleware) CreateVM(ctx context.Context, req *manager.CreateReq) (agentAddr string, id string, err error) {
 	defer func(begin time.Time) {
-		message := fmt.Sprintf("Method Run for computation took %s to complete", time.Since(begin))
+		message := fmt.Sprintf("Method CreateVM for id %s on port %s took %s to complete", id, agentAddr, time.Since(begin))
 		if err != nil {
 			lm.logger.Warn(fmt.Sprintf("%s with error: %s.", message, err))
 			return
@@ -37,12 +37,12 @@ func (lm *loggingMiddleware) Run(ctx context.Context, mc *manager.ComputationRun
 		lm.logger.Info(message)
 	}(time.Now())
 
-	return lm.svc.Run(ctx, mc)
+	return lm.svc.CreateVM(ctx, req)
 }
 
-func (lm *loggingMiddleware) Stop(ctx context.Context, computationID string) (err error) {
+func (lm *loggingMiddleware) RemoveVM(ctx context.Context, id string) (err error) {
 	defer func(begin time.Time) {
-		message := fmt.Sprintf("Method Stop for computation took %s to complete", time.Since(begin))
+		message := fmt.Sprintf("Method RemoveVM for vm %s took %s to complete", id, time.Since(begin))
 		if err != nil {
 			lm.logger.Warn(fmt.Sprintf("%s with error: %s.", message, err))
 			return
@@ -50,7 +50,7 @@ func (lm *loggingMiddleware) Stop(ctx context.Context, computationID string) (er
 		lm.logger.Info(message)
 	}(time.Now())
 
-	return lm.svc.Stop(ctx, computationID)
+	return lm.svc.RemoveVM(ctx, id)
 }
 
 func (lm *loggingMiddleware) FetchAttestationPolicy(ctx context.Context, cmpId string) (body []byte, err error) {
@@ -67,10 +67,6 @@ func (lm *loggingMiddleware) FetchAttestationPolicy(ctx context.Context, cmpId s
 	return lm.svc.FetchAttestationPolicy(ctx, cmpId)
 }
 
-func (lm *loggingMiddleware) ReportBrokenConnection(addr string) {
-	lm.svc.ReportBrokenConnection(addr)
-}
-
 func (lm *loggingMiddleware) ReturnSVMInfo(ctx context.Context) (string, int, string, string) {
 	defer func(begin time.Time) {
 		message := fmt.Sprintf("Method ReturnSVMInfo for computation took %s to complete", time.Since(begin))

manager/api/metrics.go

@@ -32,22 +32,22 @@ func MetricsMiddleware(svc manager.Service, counter metrics.Counter, latency met
 	}
 }
 
-func (ms *metricsMiddleware) Run(ctx context.Context, mc *manager.ComputationRunReq) (string, error) {
+func (ms *metricsMiddleware) CreateVM(ctx context.Context, req *manager.CreateReq) (string, string, error) {
 	defer func(begin time.Time) {
 		ms.counter.With("method", "Run").Add(1)
 		ms.latency.With("method", "Run").Observe(time.Since(begin).Seconds())
 	}(time.Now())
 
-	return ms.svc.Run(ctx, mc)
+	return ms.svc.CreateVM(ctx, req)
 }
 
-func (ms *metricsMiddleware) Stop(ctx context.Context, computationID string) error {
+func (ms *metricsMiddleware) RemoveVM(ctx context.Context, computationID string) error {
 	defer func(begin time.Time) {
 		ms.counter.With("method", "Stop").Add(1)
 		ms.latency.With("method", "Stop").Observe(time.Since(begin).Seconds())
 	}(time.Now())
 
-	return ms.svc.Stop(ctx, computationID)
+	return ms.svc.RemoveVM(ctx, computationID)
 }
 
 func (ms *metricsMiddleware) FetchAttestationPolicy(ctx context.Context, cmpId string) ([]byte, error) {
@@ -59,10 +59,6 @@ func (ms *metricsMiddleware) FetchAttestationPolicy(ctx context.Context, cmpId s
 	return ms.svc.FetchAttestationPolicy(ctx, cmpId)
 }
 
-func (ms *metricsMiddleware) ReportBrokenConnection(addr string) {
-	ms.svc.ReportBrokenConnection(addr)
-}
-
 func (ms *metricsMiddleware) ReturnSVMInfo(ctx context.Context) (string, int, string, string) {
 	defer func(begin time.Time) {
 		ms.counter.With("method", "ReturnSVMInfo").Add(1)

manager/attestation_policy.go

@@ -7,25 +7,42 @@
 package manager
 
 import (
+	"bufio"
+	"bytes"
 	"context"
 	"encoding/base64"
+	"encoding/hex"
+	"encoding/json"
 	"fmt"
-	"os"
-	"os/exec"
 	"strconv"
+	"strings"
 
 	"github.com/google/go-sev-guest/proto/check"
 	"github.com/ultravioletrs/cocos/manager/qemu"
+	config "github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/cmdconfig"
 	"github.com/virtee/sev-snp-measure-go/cpuid"
 	"github.com/virtee/sev-snp-measure-go/guest"
 	"github.com/virtee/sev-snp-measure-go/vmmtypes"
-	"google.golang.org/protobuf/encoding/protojson"
 )
 
 const defGuestFeatures = 0x1
 
 func (ms *managerService) FetchAttestationPolicy(_ context.Context, computationId string) ([]byte, error) {
-	cmd := exec.Command("sudo", fmt.Sprintf("%s/attestation_policy", ms.attestationPolicyBinaryPath), "--policy", "196608")
+	var stderrBuffer bytes.Buffer
+	options := []string{"--policy", "196608"}
+
+	if ms.pcrValuesFilePath != "" {
+		pcrValues := []string{"--pcr", ms.pcrValuesFilePath}
+		options = append(options, pcrValues...)
+	}
+
+	stderr := bufio.NewWriter(&stderrBuffer)
+
+	attestPolicyCmd, err := cmdconfig.NewCmdConfig("sudo", options, stderr)
+	if err != nil {
+		return nil, err
+	}
 
 	ms.mu.Lock()
 	vm, exists := ms.vms[computationId]
@@ -40,22 +57,15 @@ func (ms *managerService) FetchAttestationPolicy(_ context.Context, computationI
 	}
 
 	ms.ap.Lock()
-	_, err := cmd.Output()
+	stdOutByte, err := attestPolicyCmd.Run(ms.attestationPolicyBinaryPath)
 	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
-	}
+	attestationPolicy := config.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &config.PcrConfig{}}
 
-	var attestationPolicy check.Config
-
-	if err = protojson.Unmarshal(f, &attestationPolicy); err != nil {
+	if err = config.ReadAttestationPolicyFromByte(stdOutByte, &attestationPolicy); err != nil {
 		return nil, err
 	}
 
@@ -67,26 +77,50 @@ func (ms *managerService) FetchAttestationPolicy(_ context.Context, computationI
 			return nil, err
 		}
 	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)
+		stderr := bufio.NewWriter(&stderrBuffer)
+		options := cmdconfig.IgvmMeasureOptions
+
+		igvmMeasurement, err := cmdconfig.NewCmdConfig(ms.igvmMeasurementBinaryPath, options, stderr)
+		if err != nil {
+			return nil, err
+		}
+
+		outputByte, err := igvmMeasurement.Run(ms.qemuCfg.IGVMConfig.File)
+		if err != nil {
+			return nil, err
+		}
+
+		outputString := string(outputByte)
+		lines := strings.Split(strings.TrimSpace(outputString), "\n")
+
+		if len(lines) == 1 {
+			outputString = strings.TrimSpace(outputString)
+			outputString = strings.ToLower(outputString)
+		} else {
+			return nil, fmt.Errorf("error: %s", outputString)
+		}
+
+		measurement, err = hex.DecodeString(outputString)
 		if err != nil {
 			return nil, err
 		}
 	}
+
 	if measurement != nil {
-		attestationPolicy.Policy.Measurement = measurement
+		attestationPolicy.Config.Policy.Measurement = measurement
 	}
 
-	if vmi.Config.HostData != "" {
-		hostData, err := base64.StdEncoding.DecodeString(vmi.Config.HostData)
+	if vmi.Config.SevConfig.EnableHostData {
+		hostData, err := base64.StdEncoding.DecodeString(vmi.Config.SevConfig.HostData)
 		if err != nil {
 			return nil, err
 		}
-		attestationPolicy.Policy.HostData = hostData
+		attestationPolicy.Config.Policy.HostData = hostData
 	}
 
-	attestationPolicy.Policy.MinimumLaunchTcb = vmi.LaunchTCB
+	attestationPolicy.Config.Policy.MinimumLaunchTcb = vmi.LaunchTCB
 
-	f, err = protojson.Marshal(&attestationPolicy)
+	f, err := json.MarshalIndent(attestationPolicy, "", " ")
 	if err != nil {
 		return nil, err
 	}

manager/attestation_policy_test.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"encoding/json"
 	"os"
+	"path"
 	"path/filepath"
 	"testing"
 
@@ -20,7 +21,7 @@ func CreateDummyAttestationPolicyBinary(t *testing.T, behavior string) string {
 	switch behavior {
 	case "success":
 		content = []byte(`#!/bin/sh
-echo '{"policy": {"measurement": null, "host_data": null}}' > attestation_policy.json
+echo '{"pcr_values": {"sha256": null, "sha384": null}, "policy": {"measurement": null, "host_data": null}}'
 `)
 	case "fail":
 		content = []byte(`#!/bin/sh
@@ -57,9 +58,10 @@ func TestFetchAttestationPolicy(t *testing.T) {
 			binaryBehavior: "success",
 			vmConfig: qemu.VMInfo{
 				Config: qemu.Config{
-					EnableSEV: true,
-					SMPCount:  2,
-					CPU:       "EPYC",
+					EnableSEV:    true,
+					EnableSEVSNP: false,
+					SMPCount:     2,
+					CPU:          "EPYC",
 					OVMFCodeConfig: qemu.OVMFCodeConfig{
 						File: "/path/to/OVMF_CODE.fd",
 					},
@@ -68,23 +70,6 @@ func TestFetchAttestationPolicy(t *testing.T) {
 			},
 			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.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",
@@ -121,7 +106,7 @@ func TestFetchAttestationPolicy(t *testing.T) {
 				},
 				LaunchTCB: 0,
 			},
-			expectedError: "no such file or directory",
+			expectedError: "failed to decode Attestation Policy file",
 		},
 	}
 
@@ -132,7 +117,8 @@ func TestFetchAttestationPolicy(t *testing.T) {
 
 			ms := &managerService{
 				vms:                         make(map[string]vm.VM),
-				attestationPolicyBinaryPath: tempDir,
+				attestationPolicyBinaryPath: path.Join(tempDir, "attestation_policy"),
+				pcrValuesFilePath:           tempDir,
 				qemuCfg: qemu.Config{
 					CPU: "EPYC",
 				},

manager/events/events.go

@@ -1,9 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package events
-
-import "context"
-
-type Listener interface {
-	Listen(ctx context.Context)
-}

manager/events/vsock.go

@@ -1,125 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package events
-
-import (
-	"context"
-	"log/slog"
-	"net"
-
-	"github.com/mdlayher/vsock"
-	agentevents "github.com/ultravioletrs/cocos/agent/events"
-	internalvsock "github.com/ultravioletrs/cocos/internal/vsock"
-	"github.com/ultravioletrs/cocos/manager"
-	"google.golang.org/protobuf/proto"
-)
-
-const ManagerVsockPort = 9997
-
-type ReportBrokenConnectionFunc func(address string)
-
-type events struct {
-	reportBrokenConnection ReportBrokenConnectionFunc
-	lis                    net.Listener
-	logger                 *slog.Logger
-	eventsChan             chan *manager.ClientStreamMessage
-}
-
-func New(logger *slog.Logger, reportBrokenConnection ReportBrokenConnectionFunc, eventsChan chan *manager.ClientStreamMessage) (Listener, error) {
-	l, err := vsock.Listen(ManagerVsockPort, nil)
-	if err != nil {
-		return nil, err
-	}
-	return &events{
-		lis:                    l,
-		reportBrokenConnection: reportBrokenConnection,
-		logger:                 logger,
-		eventsChan:             eventsChan,
-	}, nil
-}
-
-func (e *events) Listen(ctx context.Context) {
-	for {
-		select {
-		case <-ctx.Done():
-			e.logger.Info("Listener shutting down")
-			return
-		default:
-			conn, err := e.lis.Accept()
-			if err != nil {
-				e.logger.Warn(err.Error())
-				continue
-			}
-
-			go e.handleConnection(conn)
-		}
-	}
-}
-
-func (e *events) handleConnection(conn net.Conn) {
-	defer conn.Close()
-
-	ackReader := internalvsock.NewAckReader(conn)
-
-	for {
-		var message agentevents.EventsLogs
-		data, err := ackReader.Read()
-		if err != nil {
-			go e.reportBrokenConnection(conn.RemoteAddr().String())
-			e.logger.Warn(err.Error())
-			return
-		}
-
-		if err := proto.Unmarshal(data, &message); err != nil {
-			e.logger.Warn(err.Error())
-			continue
-		}
-
-		var mes manager.ClientStreamMessage
-
-		args := []any{}
-
-		switch message.Message.(type) {
-		case *agentevents.EventsLogs_AgentEvent:
-			args = append(args, slog.Group("agent-event",
-				slog.String("event-type", message.GetAgentEvent().GetEventType()),
-				slog.String("computation-id", message.GetAgentEvent().GetComputationId()),
-				slog.String("status", message.GetAgentEvent().GetStatus()),
-				slog.String("originator", message.GetAgentEvent().GetOriginator()),
-				slog.String("timestamp", message.GetAgentEvent().GetTimestamp().String()),
-				slog.String("details", string(message.GetAgentEvent().GetDetails()))))
-			mes = manager.ClientStreamMessage{
-				Message: &manager.ClientStreamMessage_AgentEvent{
-					AgentEvent: &manager.AgentEvent{
-						EventType:     message.GetAgentEvent().GetEventType(),
-						ComputationId: message.GetAgentEvent().GetComputationId(),
-						Status:        message.GetAgentEvent().GetStatus(),
-						Originator:    message.GetAgentEvent().GetOriginator(),
-						Timestamp:     message.GetAgentEvent().GetTimestamp(),
-						Details:       message.GetAgentEvent().GetDetails(),
-					},
-				},
-			}
-		case *agentevents.EventsLogs_AgentLog:
-			args = append(args, slog.Group("agent-log",
-				slog.String("computation-id", message.GetAgentLog().GetComputationId()),
-				slog.String("level", message.GetAgentLog().GetLevel()),
-				slog.String("timestamp", message.GetAgentLog().GetTimestamp().String()),
-				slog.String("message", message.GetAgentLog().GetMessage())))
-			mes = manager.ClientStreamMessage{
-				Message: &manager.ClientStreamMessage_AgentLog{
-					AgentLog: &manager.AgentLog{
-						ComputationId: message.GetAgentLog().GetComputationId(),
-						Level:         message.GetAgentLog().GetLevel(),
-						Timestamp:     message.GetAgentLog().GetTimestamp(),
-						Message:       message.GetAgentLog().GetMessage(),
-					},
-				},
-			}
-		}
-
-		e.eventsChan <- &mes
-
-		e.logger.Info("", args...)
-	}
-}

manager/events/vsock_test.go

@@ -1,295 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package events
-
-import (
-	"bytes"
-	"context"
-	"encoding/binary"
-	"fmt"
-	"log/slog"
-	"net"
-	"os"
-	"testing"
-	"time"
-
-	mglog "github.com/absmach/magistrala/logger"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
-	"github.com/ultravioletrs/cocos/manager"
-	"google.golang.org/protobuf/proto"
-	"google.golang.org/protobuf/types/known/timestamppb"
-)
-
-type MockVsockListener struct {
-	mock.Mock
-}
-
-func (m *MockVsockListener) Accept() (net.Conn, error) {
-	args := m.Called()
-	return args.Get(0).(net.Conn), args.Error(1)
-}
-
-func (m *MockVsockListener) Close() error {
-	args := m.Called()
-	return args.Error(0)
-}
-
-func (m *MockVsockListener) Addr() net.Addr {
-	args := m.Called()
-	return args.Get(0).(net.Addr)
-}
-
-var _ net.Conn = (*MockConn)(nil)
-
-type MockConn struct {
-	mock.Mock
-}
-
-func (m *MockConn) Read(b []byte) (n int, err error) {
-	args := m.Called(b)
-	return args.Int(0), args.Error(1)
-}
-
-func (m *MockConn) Write(b []byte) (n int, err error) {
-	args := m.Called(b)
-	return args.Int(0), args.Error(1)
-}
-
-func (m *MockConn) Close() error {
-	args := m.Called()
-	return args.Error(0)
-}
-
-func (m *MockConn) LocalAddr() net.Addr {
-	args := m.Called()
-	return args.Get(0).(net.Addr)
-}
-
-func (m *MockConn) RemoteAddr() net.Addr {
-	args := m.Called()
-	return args.Get(0).(net.Addr)
-}
-
-func (m *MockConn) SetDeadline(t time.Time) error {
-	args := m.Called(t)
-	return args.Error(0)
-}
-
-func (m *MockConn) SetReadDeadline(t time.Time) error {
-	args := m.Called(t)
-	return args.Error(0)
-}
-
-func (m *MockConn) SetWriteDeadline(t time.Time) error {
-	args := m.Called(t)
-	return args.Error(0)
-}
-
-func TestNew(t *testing.T) {
-	logger := &slog.Logger{}
-	reportBrokenConnection := func(address string) {}
-	eventsChan := make(chan *manager.ClientStreamMessage)
-
-	e, err := New(logger, reportBrokenConnection, eventsChan)
-
-	if vsockDeviceExists() {
-		assert.NoError(t, err)
-
-		assert.NotNil(t, e)
-		assert.IsType(t, &events{}, e)
-	} else {
-		assert.Error(t, err)
-	}
-}
-
-func TestListen(t *testing.T) {
-	mockListener := new(MockVsockListener)
-	mockConn := new(MockConn)
-
-	e := &events{
-		lis:    mockListener,
-		logger: mglog.NewMock(),
-	}
-
-	mockListener.On("Accept").Return(mockConn, fmt.Errorf("mock error")).Once()
-	mockListener.On("Accept").Return(mockConn, nil)
-	mockConn.On("Close").Return(nil)
-	mockConn.On("Read", mock.Anything).Return(0, nil)
-
-	go e.Listen(context.Background())
-
-	time.Sleep(100 * time.Millisecond)
-
-	mockListener.AssertExpectations(t)
-}
-
-func TestListenContextDone(t *testing.T) {
-	mockListener := new(MockVsockListener)
-	mockConn := new(MockConn)
-
-	e := &events{
-		lis:    mockListener,
-		logger: mglog.NewMock(),
-	}
-
-	ctx, cancel := context.WithCancel(context.Background())
-	cancel()
-
-	mockListener.On("Accept").Return(mockConn, nil)
-
-	e.Listen(ctx)
-
-	time.Sleep(100 * time.Millisecond)
-}
-
-func vsockDeviceExists() bool {
-	fs, err := os.Stat("/dev/vsock")
-	if err != nil {
-		return false
-	}
-	if fs.Mode()&os.ModeDevice == 0 {
-		return false
-	}
-	return true
-}
-
-type MockConnWithBuffer struct {
-	mock.Mock
-	readBuf  *bytes.Buffer
-	writeBuf *bytes.Buffer
-}
-
-func NewMockConnWithBuffer() *MockConnWithBuffer {
-	return &MockConnWithBuffer{
-		readBuf:  new(bytes.Buffer),
-		writeBuf: new(bytes.Buffer),
-	}
-}
-
-func (m *MockConnWithBuffer) Read(b []byte) (n int, err error) {
-	return m.readBuf.Read(b)
-}
-
-func (m *MockConnWithBuffer) Write(b []byte) (n int, err error) {
-	return m.writeBuf.Write(b)
-}
-
-func (m *MockConnWithBuffer) Close() error {
-	return nil
-}
-
-func (m *MockConnWithBuffer) LocalAddr() net.Addr {
-	return nil
-}
-
-func (m *MockConnWithBuffer) RemoteAddr() net.Addr {
-	return &net.IPAddr{IP: net.ParseIP("localhost")}
-}
-
-func (m *MockConnWithBuffer) SetDeadline(t time.Time) error {
-	return nil
-}
-
-func (m *MockConnWithBuffer) SetReadDeadline(t time.Time) error {
-	return nil
-}
-
-func (m *MockConnWithBuffer) SetWriteDeadline(t time.Time) error {
-	return nil
-}
-
-func TestHandleConnection(t *testing.T) {
-	tests := []struct {
-		name    string
-		message *manager.ClientStreamMessage
-	}{
-		{
-			name: "handle agent event",
-			message: &manager.ClientStreamMessage{
-				Message: &manager.ClientStreamMessage_AgentEvent{
-					AgentEvent: &manager.AgentEvent{
-						EventType:     "test_event",
-						ComputationId: "test_computation",
-						Status:        "test_status",
-						Originator:    "test_originator",
-						Timestamp:     timestamppb.Now(),
-						Details:       []byte("test_details"),
-					},
-				},
-			},
-		},
-		{
-			name: "handle agent log",
-			message: &manager.ClientStreamMessage{
-				Message: &manager.ClientStreamMessage_AgentLog{
-					AgentLog: &manager.AgentLog{
-						ComputationId: "test_computation",
-						Timestamp:     timestamppb.Now(),
-					},
-				},
-			},
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			mockConn := NewMockConnWithBuffer()
-			eventsChan := make(chan *manager.ClientStreamMessage, 1)
-
-			e := &events{
-				logger:                 mglog.NewMock(),
-				eventsChan:             eventsChan,
-				reportBrokenConnection: func(address string) {},
-			}
-
-			data, err := proto.Marshal(tt.message)
-			assert.NoError(t, err)
-
-			messageID := uint32(1)
-			err = binary.Write(mockConn.readBuf, binary.LittleEndian, messageID)
-			assert.NoError(t, err)
-			err = binary.Write(mockConn.readBuf, binary.LittleEndian, uint32(len(data)))
-			assert.NoError(t, err)
-			_, err = mockConn.readBuf.Write(data)
-			assert.NoError(t, err)
-
-			// Add EOF to signal end of stream
-			err = binary.Write(mockConn.readBuf, binary.LittleEndian, uint32(0))
-			assert.NoError(t, err)
-			err = binary.Write(mockConn.readBuf, binary.LittleEndian, uint32(0))
-			assert.NoError(t, err)
-
-			done := make(chan struct{})
-			go func() {
-				e.handleConnection(mockConn)
-				close(done)
-			}()
-
-			var receivedMessage *manager.ClientStreamMessage
-			select {
-			case receivedMessage = <-eventsChan:
-			case <-time.After(2 * time.Second):
-				t.Fatal("Timeout waiting for message in eventsChan")
-			}
-
-			assert.NotNil(t, receivedMessage)
-
-			select {
-			case <-done:
-				// handleConnection has exited
-			case <-time.After(2 * time.Second):
-				t.Fatal("Timeout waiting for handleConnection to exit")
-			}
-
-			// Check if ack was written
-			var receivedAck uint32
-			err = binary.Read(mockConn.writeBuf, binary.LittleEndian, &receivedAck)
-			assert.NoError(t, err)
-			assert.Equal(t, messageID, receivedAck)
-
-			// Ensure no unexpected calls were made on the mock
-			mockConn.AssertExpectations(t)
-		})
-	}
-}

manager/manager.proto

@@ -3,40 +3,42 @@
 
 syntax = "proto3";
 
-import "google/protobuf/timestamp.proto";
+import "google/protobuf/empty.proto";
 
 package manager;
 
 option go_package = "./manager";
 
 service ManagerService {
-  rpc Process(stream ClientStreamMessage) returns (stream ServerStreamMessage) {}
+  rpc CreateVm(CreateReq) returns (CreateRes) {}
+  rpc RemoveVm(RemoveReq) returns (google.protobuf.Empty) {}
+  rpc SVMInfo(SVMInfoReq) returns (SVMInfoRes) {}
+  rpc AttestationPolicy(AttestationPolicyReq) returns (AttestationPolicyRes) {}
 }
 
-message Terminate {
-  string message = 1;
+message CreateReq{
+  string agent_log_level = 1;
+  bytes agent_cvm_server_ca_cert = 2;
+  bytes agent_cvm_client_key = 3;
+  bytes agent_cvm_client_cert = 4;
+  string agent_cvm_server_url = 5;
 }
 
-message StopComputation {
-  string computation_id = 1;
+message CreateRes{
+  string forwarded_port = 1;
+  string svm_id = 2;
 }
 
-message StopComputationResponse {
-  string computation_id = 1;
-  string message = 2;
+message RemoveReq{
+  string svm_id = 1;
 }
 
-message RunResponse{
-  string agent_port = 1;
-  string computation_id = 2;
-}
-
-message AttestationPolicy{
+message AttestationPolicyRes{
   bytes info = 1;
   string id = 2;
 }
 
-message SVMInfo{
+message SVMInfoRes{
   string id = 1;
   string ovmf_version = 2;
   int32 cpu_num = 3;
@@ -45,60 +47,6 @@ message SVMInfo{
   string eos_version = 6;
 }
 
-message AgentEvent {
-  string event_type = 1;
-	google.protobuf.Timestamp timestamp = 2;
-	string	computation_id = 3;
-	bytes	details = 4;
-	string	originator = 5;
-	string	status = 6;
-}
-
-message AgentLog {
-  string message = 1;
-  string computation_id = 2;
-  string level = 3;
-  google.protobuf.Timestamp timestamp = 4;
-}
-
-message ClientStreamMessage {
-  oneof message {
-    AgentLog agent_log = 1;
-    AgentEvent agent_event = 2;
-    RunResponse run_res = 3;
-    AttestationPolicy attestationPolicy = 4;
-    StopComputationResponse stopComputationRes = 5;
-    SVMInfo svm_info = 6;
-  }
-}
-
-message ServerStreamMessage {
-  oneof message {
-    RunReqChunks runReqChunks = 1;
-    ComputationRunReq runReq = 2;
-    Terminate terminateReq = 3;
-    StopComputation stopComputation = 4;
-    AttestationPolicyReq attestationPolicyReq = 5;
-    SVMInfoReq svmInfoReq = 6;
-  }
-}
-
-message RunReqChunks {
-  bytes data = 1;
-  string id = 2;
-  bool is_last = 3;
-}
-
-message ComputationRunReq {
-  string id = 1;
-  string name = 2;
-  string description = 3;
-  repeated Dataset datasets = 4;
-  Algorithm algorithm = 5;
-  repeated ResultConsumer result_consumers = 6;
-  AgentConfig agent_config = 7;
-}
-
 message AttestationPolicyReq {
   string id = 1;
 }
@@ -107,29 +55,3 @@ message SVMInfoReq {
   string id = 1;
 }
 
-message ResultConsumer {
-  bytes userKey = 1;
-}
-
-message Dataset {
-  bytes hash = 1; // should be sha3.Sum256, 32 byte length.
-  bytes userKey = 2;
-  string filename = 3;
-}
-
-message Algorithm {
-  bytes hash = 1; // should be sha3.Sum256, 32 byte length.
-  bytes userKey = 2;
-}
-
-message AgentConfig {
-  string port = 1;
-  string host = 2;
-  string cert_file = 3;
-  string key_file = 4;
-  string client_ca_file = 5;
-  string server_ca_file = 6;
-  string log_level = 7;
-  bool   attested_tls = 8;
-}
-

manager/manager_grpc.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
 // - protoc-gen-go-grpc v1.5.1
-// - protoc             v5.28.1
+// - protoc             v5.29.0
 // source: manager/manager.proto
 
 package manager
@@ -14,6 +14,7 @@ import (
 	grpc "google.golang.org/grpc"
 	codes "google.golang.org/grpc/codes"
 	status "google.golang.org/grpc/status"
+	emptypb "google.golang.org/protobuf/types/known/emptypb"
 )
 
 // This is a compile-time assertion to ensure that this generated file
@@ -22,14 +23,20 @@ import (
 const _ = grpc.SupportPackageIsVersion9
 
 const (
-	ManagerService_Process_FullMethodName = "/manager.ManagerService/Process"
+	ManagerService_CreateVm_FullMethodName          = "/manager.ManagerService/CreateVm"
+	ManagerService_RemoveVm_FullMethodName          = "/manager.ManagerService/RemoveVm"
+	ManagerService_SVMInfo_FullMethodName           = "/manager.ManagerService/SVMInfo"
+	ManagerService_AttestationPolicy_FullMethodName = "/manager.ManagerService/AttestationPolicy"
 )
 
 // ManagerServiceClient is the client API for ManagerService service.
 //
 // For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
 type ManagerServiceClient interface {
-	Process(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage], error)
+	CreateVm(ctx context.Context, in *CreateReq, opts ...grpc.CallOption) (*CreateRes, error)
+	RemoveVm(ctx context.Context, in *RemoveReq, opts ...grpc.CallOption) (*emptypb.Empty, error)
+	SVMInfo(ctx context.Context, in *SVMInfoReq, opts ...grpc.CallOption) (*SVMInfoRes, error)
+	AttestationPolicy(ctx context.Context, in *AttestationPolicyReq, opts ...grpc.CallOption) (*AttestationPolicyRes, error)
 }
 
 type managerServiceClient struct {
@@ -40,24 +47,54 @@ func NewManagerServiceClient(cc grpc.ClientConnInterface) ManagerServiceClient {
 	return &managerServiceClient{cc}
 }
 
-func (c *managerServiceClient) Process(ctx context.Context, opts ...grpc.CallOption) (grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage], error) {
+func (c *managerServiceClient) CreateVm(ctx context.Context, in *CreateReq, opts ...grpc.CallOption) (*CreateRes, error) {
 	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
-	stream, err := c.cc.NewStream(ctx, &ManagerService_ServiceDesc.Streams[0], ManagerService_Process_FullMethodName, cOpts...)
+	out := new(CreateRes)
+	err := c.cc.Invoke(ctx, ManagerService_CreateVm_FullMethodName, in, out, cOpts...)
 	if err != nil {
 		return nil, err
 	}
-	x := &grpc.GenericClientStream[ClientStreamMessage, ServerStreamMessage]{ClientStream: stream}
-	return x, nil
+	return out, nil
 }
 
-// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
-type ManagerService_ProcessClient = grpc.BidiStreamingClient[ClientStreamMessage, ServerStreamMessage]
+func (c *managerServiceClient) RemoveVm(ctx context.Context, in *RemoveReq, opts ...grpc.CallOption) (*emptypb.Empty, error) {
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
+	out := new(emptypb.Empty)
+	err := c.cc.Invoke(ctx, ManagerService_RemoveVm_FullMethodName, in, out, cOpts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *managerServiceClient) SVMInfo(ctx context.Context, in *SVMInfoReq, opts ...grpc.CallOption) (*SVMInfoRes, error) {
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
+	out := new(SVMInfoRes)
+	err := c.cc.Invoke(ctx, ManagerService_SVMInfo_FullMethodName, in, out, cOpts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *managerServiceClient) AttestationPolicy(ctx context.Context, in *AttestationPolicyReq, opts ...grpc.CallOption) (*AttestationPolicyRes, error) {
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
+	out := new(AttestationPolicyRes)
+	err := c.cc.Invoke(ctx, ManagerService_AttestationPolicy_FullMethodName, in, out, cOpts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
 
 // ManagerServiceServer is the server API for ManagerService service.
 // All implementations must embed UnimplementedManagerServiceServer
 // for forward compatibility.
 type ManagerServiceServer interface {
-	Process(grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]) error
+	CreateVm(context.Context, *CreateReq) (*CreateRes, error)
+	RemoveVm(context.Context, *RemoveReq) (*emptypb.Empty, error)
+	SVMInfo(context.Context, *SVMInfoReq) (*SVMInfoRes, error)
+	AttestationPolicy(context.Context, *AttestationPolicyReq) (*AttestationPolicyRes, error)
 	mustEmbedUnimplementedManagerServiceServer()
 }
 
@@ -68,8 +105,17 @@ type ManagerServiceServer interface {
 // pointer dereference when methods are called.
 type UnimplementedManagerServiceServer struct{}
 
-func (UnimplementedManagerServiceServer) Process(grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]) error {
-	return status.Errorf(codes.Unimplemented, "method Process not implemented")
+func (UnimplementedManagerServiceServer) CreateVm(context.Context, *CreateReq) (*CreateRes, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method CreateVm not implemented")
+}
+func (UnimplementedManagerServiceServer) RemoveVm(context.Context, *RemoveReq) (*emptypb.Empty, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method RemoveVm not implemented")
+}
+func (UnimplementedManagerServiceServer) SVMInfo(context.Context, *SVMInfoReq) (*SVMInfoRes, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method SVMInfo not implemented")
+}
+func (UnimplementedManagerServiceServer) AttestationPolicy(context.Context, *AttestationPolicyReq) (*AttestationPolicyRes, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method AttestationPolicy not implemented")
 }
 func (UnimplementedManagerServiceServer) mustEmbedUnimplementedManagerServiceServer() {}
 func (UnimplementedManagerServiceServer) testEmbeddedByValue()                        {}
@@ -92,12 +138,77 @@ func RegisterManagerServiceServer(s grpc.ServiceRegistrar, srv ManagerServiceSer
 	s.RegisterService(&ManagerService_ServiceDesc, srv)
 }
 
-func _ManagerService_Process_Handler(srv interface{}, stream grpc.ServerStream) error {
-	return srv.(ManagerServiceServer).Process(&grpc.GenericServerStream[ClientStreamMessage, ServerStreamMessage]{ServerStream: stream})
+func _ManagerService_CreateVm_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(CreateReq)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(ManagerServiceServer).CreateVm(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: ManagerService_CreateVm_FullMethodName,
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(ManagerServiceServer).CreateVm(ctx, req.(*CreateReq))
+	}
+	return interceptor(ctx, in, info, handler)
 }
 
-// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
-type ManagerService_ProcessServer = grpc.BidiStreamingServer[ClientStreamMessage, ServerStreamMessage]
+func _ManagerService_RemoveVm_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(RemoveReq)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(ManagerServiceServer).RemoveVm(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: ManagerService_RemoveVm_FullMethodName,
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(ManagerServiceServer).RemoveVm(ctx, req.(*RemoveReq))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+func _ManagerService_SVMInfo_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(SVMInfoReq)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(ManagerServiceServer).SVMInfo(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: ManagerService_SVMInfo_FullMethodName,
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(ManagerServiceServer).SVMInfo(ctx, req.(*SVMInfoReq))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
+func _ManagerService_AttestationPolicy_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(AttestationPolicyReq)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(ManagerServiceServer).AttestationPolicy(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: ManagerService_AttestationPolicy_FullMethodName,
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(ManagerServiceServer).AttestationPolicy(ctx, req.(*AttestationPolicyReq))
+	}
+	return interceptor(ctx, in, info, handler)
+}
 
 // ManagerService_ServiceDesc is the grpc.ServiceDesc for ManagerService service.
 // It's only intended for direct use with grpc.RegisterService,
@@ -105,14 +216,24 @@ type ManagerService_ProcessServer = grpc.BidiStreamingServer[ClientStreamMessage
 var ManagerService_ServiceDesc = grpc.ServiceDesc{
 	ServiceName: "manager.ManagerService",
 	HandlerType: (*ManagerServiceServer)(nil),
-	Methods:     []grpc.MethodDesc{},
-	Streams: []grpc.StreamDesc{
+	Methods: []grpc.MethodDesc{
 		{
-			StreamName:    "Process",
-			Handler:       _ManagerService_Process_Handler,
-			ServerStreams: true,
-			ClientStreams: true,
+			MethodName: "CreateVm",
+			Handler:    _ManagerService_CreateVm_Handler,
+		},
+		{
+			MethodName: "RemoveVm",
+			Handler:    _ManagerService_RemoveVm_Handler,
+		},
+		{
+			MethodName: "SVMInfo",
+			Handler:    _ManagerService_SVMInfo_Handler,
+		},
+		{
+			MethodName: "AttestationPolicy",
+			Handler:    _ManagerService_AttestationPolicy_Handler,
 		},
 	},
+	Streams:  []grpc.StreamDesc{},
 	Metadata: "manager/manager.proto",
 }

manager/manager_test.go

@@ -1,68 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package manager_test
-
-import (
-	"bytes"
-	"context"
-	"testing"
-
-	"github.com/ultravioletrs/cocos/manager"
-	"google.golang.org/grpc"
-	"google.golang.org/protobuf/proto"
-)
-
-func TestProcess(t *testing.T) {
-	ctx := context.Background()
-	conn, err := grpc.DialContext(ctx, "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
-	if err != nil {
-		t.Fatalf("Failed to dial bufnet: %v", err)
-	}
-	defer conn.Close()
-
-	client := manager.NewManagerServiceClient(conn)
-	stream, err := client.Process(ctx)
-	if err != nil {
-		t.Fatalf("Process failed: %v", err)
-	}
-
-	var data bytes.Buffer
-	for {
-		msg, err := stream.Recv()
-		if err != nil {
-			t.Fatalf("Failed to receive ServerStreamMessage: %v", err)
-		}
-
-		switch m := msg.Message.(type) {
-		case *manager.ServerStreamMessage_TerminateReq:
-			if m.TerminateReq.Message != "test terminate" {
-				t.Fatalf("Unexpected terminate message: %v", m.TerminateReq.Message)
-			}
-		case *manager.ServerStreamMessage_RunReqChunks:
-			if len(m.RunReqChunks.Data) == 0 {
-				var runReq manager.ComputationRunReq
-				if err = proto.Unmarshal(data.Bytes(), &runReq); err != nil {
-					t.Fatalf("Failed to create run request: %v", err)
-				}
-
-				runRes := &manager.ClientStreamMessage_AgentLog{
-					AgentLog: &manager.AgentLog{
-						Message:       "test log",
-						ComputationId: "comp1",
-						Level:         "DEBUG",
-					},
-				}
-				if runReq.Id != "1" || runReq.Name != "sample computation" || runReq.Description != "sample description" {
-					t.Fatalf("Unexpected run request message: %v", &runReq)
-				}
-				if err := stream.Send(&manager.ClientStreamMessage{Message: runRes}); err != nil {
-					t.Fatalf("Failed to send ClientStreamMessage: %v", err)
-				}
-				return
-			}
-			data.Write(m.RunReqChunks.Data)
-		default:
-			t.Fatalf("Unexpected message type: %T", m)
-		}
-	}
-}

manager/mocks/manager_service_client.go

@@ -0,0 +1,340 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+package mocks
+
+import (
+	context "context"
+
+	grpc "google.golang.org/grpc"
+	emptypb "google.golang.org/protobuf/types/known/emptypb"
+
+	manager "github.com/ultravioletrs/cocos/manager"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// ManagerServiceClient is an autogenerated mock type for the ManagerServiceClient type
+type ManagerServiceClient struct {
+	mock.Mock
+}
+
+type ManagerServiceClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *ManagerServiceClient) EXPECT() *ManagerServiceClient_Expecter {
+	return &ManagerServiceClient_Expecter{mock: &_m.Mock}
+}
+
+// AttestationPolicy provides a mock function with given fields: ctx, in, opts
+func (_m *ManagerServiceClient) AttestationPolicy(ctx context.Context, in *manager.AttestationPolicyReq, opts ...grpc.CallOption) (*manager.AttestationPolicyRes, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	if len(ret) == 0 {
+		panic("no return value specified for AttestationPolicy")
+	}
+
+	var r0 *manager.AttestationPolicyRes
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) (*manager.AttestationPolicyRes, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) *manager.AttestationPolicyRes); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*manager.AttestationPolicyRes)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ManagerServiceClient_AttestationPolicy_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AttestationPolicy'
+type ManagerServiceClient_AttestationPolicy_Call struct {
+	*mock.Call
+}
+
+// AttestationPolicy is a helper method to define mock.On call
+//   - ctx context.Context
+//   - in *manager.AttestationPolicyReq
+//   - opts ...grpc.CallOption
+func (_e *ManagerServiceClient_Expecter) AttestationPolicy(ctx interface{}, in interface{}, opts ...interface{}) *ManagerServiceClient_AttestationPolicy_Call {
+	return &ManagerServiceClient_AttestationPolicy_Call{Call: _e.mock.On("AttestationPolicy",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *ManagerServiceClient_AttestationPolicy_Call) Run(run func(ctx context.Context, in *manager.AttestationPolicyReq, opts ...grpc.CallOption)) *ManagerServiceClient_AttestationPolicy_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*manager.AttestationPolicyReq), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *ManagerServiceClient_AttestationPolicy_Call) Return(_a0 *manager.AttestationPolicyRes, _a1 error) *ManagerServiceClient_AttestationPolicy_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *ManagerServiceClient_AttestationPolicy_Call) RunAndReturn(run func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) (*manager.AttestationPolicyRes, error)) *ManagerServiceClient_AttestationPolicy_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// CreateVm provides a mock function with given fields: ctx, in, opts
+func (_m *ManagerServiceClient) CreateVm(ctx context.Context, in *manager.CreateReq, opts ...grpc.CallOption) (*manager.CreateRes, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	if len(ret) == 0 {
+		panic("no return value specified for CreateVm")
+	}
+
+	var r0 *manager.CreateRes
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *manager.CreateReq, ...grpc.CallOption) (*manager.CreateRes, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *manager.CreateReq, ...grpc.CallOption) *manager.CreateRes); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*manager.CreateRes)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *manager.CreateReq, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ManagerServiceClient_CreateVm_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CreateVm'
+type ManagerServiceClient_CreateVm_Call struct {
+	*mock.Call
+}
+
+// CreateVm is a helper method to define mock.On call
+//   - ctx context.Context
+//   - in *manager.CreateReq
+//   - opts ...grpc.CallOption
+func (_e *ManagerServiceClient_Expecter) CreateVm(ctx interface{}, in interface{}, opts ...interface{}) *ManagerServiceClient_CreateVm_Call {
+	return &ManagerServiceClient_CreateVm_Call{Call: _e.mock.On("CreateVm",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *ManagerServiceClient_CreateVm_Call) Run(run func(ctx context.Context, in *manager.CreateReq, opts ...grpc.CallOption)) *ManagerServiceClient_CreateVm_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*manager.CreateReq), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *ManagerServiceClient_CreateVm_Call) Return(_a0 *manager.CreateRes, _a1 error) *ManagerServiceClient_CreateVm_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *ManagerServiceClient_CreateVm_Call) RunAndReturn(run func(context.Context, *manager.CreateReq, ...grpc.CallOption) (*manager.CreateRes, error)) *ManagerServiceClient_CreateVm_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// RemoveVm provides a mock function with given fields: ctx, in, opts
+func (_m *ManagerServiceClient) RemoveVm(ctx context.Context, in *manager.RemoveReq, opts ...grpc.CallOption) (*emptypb.Empty, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RemoveVm")
+	}
+
+	var r0 *emptypb.Empty
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *manager.RemoveReq, ...grpc.CallOption) (*emptypb.Empty, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *manager.RemoveReq, ...grpc.CallOption) *emptypb.Empty); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*emptypb.Empty)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *manager.RemoveReq, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ManagerServiceClient_RemoveVm_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RemoveVm'
+type ManagerServiceClient_RemoveVm_Call struct {
+	*mock.Call
+}
+
+// RemoveVm is a helper method to define mock.On call
+//   - ctx context.Context
+//   - in *manager.RemoveReq
+//   - opts ...grpc.CallOption
+func (_e *ManagerServiceClient_Expecter) RemoveVm(ctx interface{}, in interface{}, opts ...interface{}) *ManagerServiceClient_RemoveVm_Call {
+	return &ManagerServiceClient_RemoveVm_Call{Call: _e.mock.On("RemoveVm",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *ManagerServiceClient_RemoveVm_Call) Run(run func(ctx context.Context, in *manager.RemoveReq, opts ...grpc.CallOption)) *ManagerServiceClient_RemoveVm_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*manager.RemoveReq), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *ManagerServiceClient_RemoveVm_Call) Return(_a0 *emptypb.Empty, _a1 error) *ManagerServiceClient_RemoveVm_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *ManagerServiceClient_RemoveVm_Call) RunAndReturn(run func(context.Context, *manager.RemoveReq, ...grpc.CallOption) (*emptypb.Empty, error)) *ManagerServiceClient_RemoveVm_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SVMInfo provides a mock function with given fields: ctx, in, opts
+func (_m *ManagerServiceClient) SVMInfo(ctx context.Context, in *manager.SVMInfoReq, opts ...grpc.CallOption) (*manager.SVMInfoRes, error) {
+	_va := make([]interface{}, len(opts))
+	for _i := range opts {
+		_va[_i] = opts[_i]
+	}
+	var _ca []interface{}
+	_ca = append(_ca, ctx, in)
+	_ca = append(_ca, _va...)
+	ret := _m.Called(_ca...)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SVMInfo")
+	}
+
+	var r0 *manager.SVMInfoRes
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *manager.SVMInfoReq, ...grpc.CallOption) (*manager.SVMInfoRes, error)); ok {
+		return rf(ctx, in, opts...)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *manager.SVMInfoReq, ...grpc.CallOption) *manager.SVMInfoRes); ok {
+		r0 = rf(ctx, in, opts...)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*manager.SVMInfoRes)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *manager.SVMInfoReq, ...grpc.CallOption) error); ok {
+		r1 = rf(ctx, in, opts...)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// ManagerServiceClient_SVMInfo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SVMInfo'
+type ManagerServiceClient_SVMInfo_Call struct {
+	*mock.Call
+}
+
+// SVMInfo is a helper method to define mock.On call
+//   - ctx context.Context
+//   - in *manager.SVMInfoReq
+//   - opts ...grpc.CallOption
+func (_e *ManagerServiceClient_Expecter) SVMInfo(ctx interface{}, in interface{}, opts ...interface{}) *ManagerServiceClient_SVMInfo_Call {
+	return &ManagerServiceClient_SVMInfo_Call{Call: _e.mock.On("SVMInfo",
+		append([]interface{}{ctx, in}, opts...)...)}
+}
+
+func (_c *ManagerServiceClient_SVMInfo_Call) Run(run func(ctx context.Context, in *manager.SVMInfoReq, opts ...grpc.CallOption)) *ManagerServiceClient_SVMInfo_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		variadicArgs := make([]grpc.CallOption, len(args)-2)
+		for i, a := range args[2:] {
+			if a != nil {
+				variadicArgs[i] = a.(grpc.CallOption)
+			}
+		}
+		run(args[0].(context.Context), args[1].(*manager.SVMInfoReq), variadicArgs...)
+	})
+	return _c
+}
+
+func (_c *ManagerServiceClient_SVMInfo_Call) Return(_a0 *manager.SVMInfoRes, _a1 error) *ManagerServiceClient_SVMInfo_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *ManagerServiceClient_SVMInfo_Call) RunAndReturn(run func(context.Context, *manager.SVMInfoReq, ...grpc.CallOption) (*manager.SVMInfoRes, error)) *ManagerServiceClient_SVMInfo_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewManagerServiceClient creates a new instance of ManagerServiceClient. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewManagerServiceClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *ManagerServiceClient {
+	mock := &ManagerServiceClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}