COCOS-492: Cache VCEK on aTLS verification (#524)

* initial commit

* made changes based on errors

* remove unnecessary log

---------

Co-authored-by: Jovan Djukic <jovan.djukic@ultraviolet.com>

.github/workflows/checkproto.yaml

@@ -30,12 +30,12 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: 1.23.x
+          go-version: 1.25.x
 
       - name: Set up protoc
         run: |
           PROTOC_VERSION=29.0
-          PROTOC_GEN_VERSION=v1.36.5
+          PROTOC_GEN_VERSION=v1.36.8
           PROTOC_GRPC_VERSION=v1.5.1
 
           # Download and install protoc

.github/workflows/hal.yml

@@ -42,7 +42,7 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: 1.23.x
+          go-version: 1.25.x
           cache-dependency-path: "go.sum"
 
       - name: Checkout cocos
@@ -56,7 +56,7 @@ jobs:
         with:
           repository: "buildroot/buildroot"
           path: buildroot
-          ref: 2025.05-rc1
+          ref: 2025.08-rc3
 
       - name: Build hal
         run: |

.github/workflows/main.yaml

@@ -9,9 +9,8 @@ on:
       - main
 
 jobs:
-  ci:
+  lint:
     runs-on: ubuntu-latest
-
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -19,39 +18,66 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: 1.23.x
+          go-version: 1.25.x
 
       - name: golangci-lint
-        uses: golangci/golangci-lint-action@v7
+        uses: golangci/golangci-lint-action@v8
         with:
-          version: v2.0.2
+          version: v2.4.0
 
-        
       - name: Build
-        run: |
-          make
+        run: make
+
+  test:
+    runs-on: ubuntu-latest
+    needs: lint
+    strategy:
+      matrix:
+        module: [agent, cli, cmd, internal, pkg, manager]
+        include:
+          - module: manager
+            sudo: true
+    
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Install Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: 1.25.x
 
       - name: Create coverage directory
+        run: mkdir -p coverage
+
+      - name: Run tests for ${{ matrix.module }}
         run: |
-          mkdir coverage
+          if [[ "${{ matrix.module }}" == "manager" ]]; then
+            sudo GOTOOLCHAIN=go1.25.0+auto go test -v --race -covermode=atomic -coverprofile coverage/${{ matrix.module }}.out ./${{ matrix.module }}/...
+          else
+            GOTOOLCHAIN=go1.25.0+auto go test -v --race -covermode=atomic -coverprofile coverage/${{ matrix.module }}.out ./${{ matrix.module }}/...
+          fi
 
-      - name: Run Agent tests
-        run: go test -v --race -covermode=atomic -coverprofile coverage/agent.out ./agent/...
+      - name: Upload coverage artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-${{ matrix.module }}
+          path: coverage/${{ matrix.module }}.out
+          retention-days: 1
 
-      - name: Run cli tests
-        run: go test -v --race -covermode=atomic -coverprofile coverage/cli.out ./cli/...
+  upload-coverage:
+    runs-on: ubuntu-latest
+    needs: test
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
 
-      - name: Run cmd tests
-        run: go test -v --race -covermode=atomic -coverprofile coverage/cmd.out ./cmd/...
-
-      - name: Run internal tests
-        run: go test -v --race -covermode=atomic -coverprofile coverage/internal.out ./internal/...
-
-      - name: Run pkg tests
-        run: go test -v --race -covermode=atomic -coverprofile coverage/pkg.out ./pkg/...
-
-      - name: Run manager tests
-        run: sudo go test -v --race -covermode=atomic -coverprofile coverage/manager.out ./manager/...
+      - name: Download all coverage artifacts
+        uses: actions/download-artifact@v4
+        with:
+          pattern: coverage-*
+          path: coverage/
+          merge-multiple: true
 
       - name: Upload results to Codecov
         uses: codecov/codecov-action@v4

.golangci.yaml

@@ -55,12 +55,6 @@ linters:
       template: |-
         Copyright (c) Ultraviolet
         SPDX-License-Identifier: Apache-2.0
-    importas:
-      alias:
-        - pkg: github.com/absmach/magistrala/logger
-          alias: mglog
-      no-unaliased: true
-      no-extra-aliases: false
     staticcheck:
       checks:
         - -ST1000

.mockery.yml

@@ -0,0 +1,148 @@
+pkgname: mocks
+template: testify
+template-data:
+  boilerplate-file: ./boilerplate.txt
+  unroll-variadic: true
+packages:
+  github.com/ultravioletrs/cocos/agent:
+    interfaces:
+      AgentService_AlgoClient:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+      AgentService_DataClient:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+      AgentService_IMAMeasurementsClient:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+      Service:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/agent/algorithm:
+    interfaces:
+      Algorithm:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/agent/auth:
+    interfaces:
+      Authenticator:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage:
+    interfaces:
+      Storage:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/agent/cvms/server:
+    interfaces:
+      AgentServer:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/agent/events:
+    interfaces:
+      Service:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/agent/statemachine:
+    interfaces:
+      StateMachine:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/pkg/server:
+    interfaces:
+      Server:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/manager:
+    interfaces:
+      ManagerServiceClient:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+      Service:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/manager/qemu:
+    interfaces:
+      Persistence:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/manager/vm:
+    interfaces:
+      StateMachine:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+      VM:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/pkg/attestation:
+    interfaces:
+      Provider:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+      Verifier:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/pkg/attestation/cmdconfig:
+    interfaces:
+      MeasurementProvider:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/pkg/clients/grpc:
+    interfaces:
+      Client:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/pkg/sdk:
+    interfaces:
+      SDK:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"
+  github.com/ultravioletrs/cocos/pkg/atls:
+    interfaces:
+      CertificateProvider:
+        config:
+          dir: '{{.InterfaceDir}}/mocks'
+          structname: '{{.InterfaceName}}'
+          filename: "{{.InterfaceName | lower}}.go"

Makefile

@@ -17,23 +17,23 @@ IGVM_BUILD_SCRIPT := ./scripts/igvmmeasure/igvm.sh
 define compile_service
 	CGO_ENABLED=$(CGO_ENABLED) GOOS=$(GOOS) GOARCH=$(GOARCH) GOARM=$(GOARM) \
 	go build -ldflags "-s -w \
-	-X 'github.com/absmach/magistrala.BuildTime=$(TIME)' \
-	-X 'github.com/absmach/magistrala.Version=$(VERSION)' \
-	-X 'github.com/absmach/magistrala.Commit=$(COMMIT)'" \
+	-X 'github.com/absmach/supermq.BuildTime=$(TIME)' \
+	-X 'github.com/absmach/supermq.Version=$(VERSION)' \
+	-X 'github.com/absmach/supermq.Commit=$(COMMIT)'" \
 	$(if $(filter 1,$(EMBED_ENABLED)),-tags "embed",) \
 	-o ${BUILD_DIR}/cocos-$(1) cmd/$(1)/main.go
 endef
 
 .PHONY: all $(SERVICES) $(ATTESTATION_POLICY) install clean
 
-all: $(SERVICES)
+all: $(SERVICES) $(ATTESTATION_POLICY)
 
 $(SERVICES): 
 	$(call compile_service,$@)
 	@if [ "$@" = "cli" ] || [ "$@" = "manager" ]; then $(MAKE) build-igvm; fi
 
 $(ATTESTATION_POLICY):
-	$(MAKE) -C ./scripts/attestation_policy
+	$(MAKE) -C ./scripts/attestation_policy OUTPUT_DIR=../../$(BUILD_DIR)
 
 protoc:
 	protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/agent.proto
@@ -42,17 +42,19 @@ protoc:
 	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
+	mockery --config ./.mockery.yml
 
-install: $(SERVICES)
+install: $(SERVICES) $(ATTESTATION_POLICY)
 	install -d $(INSTALL_DIR)
 	install $(BUILD_DIR)/cocos-cli $(INSTALL_DIR)/cocos-cli
 	install $(BUILD_DIR)/cocos-manager $(INSTALL_DIR)/cocos-manager
+	install $(BUILD_DIR)/attestation_policy $(INSTALL_DIR)/attestation_policy
 	install -d $(CONFIG_DIR)
 	install cocos-manager.env $(CONFIG_DIR)/cocos-manager.env
 
 clean:
 	rm -rf $(BUILD_DIR)
+	$(MAKE) -C ./scripts/attestation_policy OUTPUT_DIR=../../$(BUILD_DIR) clean
 
 run: install_service
 	sudo systemctl start $(SERVICE_NAME).service

agent/README.md

@@ -6,16 +6,24 @@ 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)                                                        | 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                                                                            | ""                             |
-| AGENT_CVM_CA_URL               | URL for CA service, if provided it will be used for certificate generation, used only with aTLS at the moment | ""                             |
+| Variable                       | Description                                                                                                   | Default                                         |
+| ------------------------------ | ------------------------------------------------------------------------------------------------------------- | ----------------------------------------------- |
+| AGENT_LOG_LEVEL                | Log level for agent service (debug, info, warn, error)                                                        | debug                                           |
+| AGENT_VMPL                     | VMPL (Virtual Machine Privilege Level) for AMD SEV-SNP attestation (0-3)                                      | 2                                               |
+| AGENT_GRPC_HOST                | Agent service gRPC host address                                                                               | 0.0.0.0                                         |
+| 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                                                                            | ""                                              |
+| AGENT_CVM_CA_URL               | URL for CA service, if provided it will be used for certificate generation, used only with aTLS at the moment | ""                                              |
+| AGENT_CVM_ID                   | Unique identifier for the CVM (Confidential Virtual Machine)                                                  | ""                                              |
+| AGENT_CERTS_TOKEN              | Authentication token for certificate service access                                                           | ""                                              |
+| AGENT_MAA_URL                  | Microsoft Azure Attestation service URL for Azure attestation                                                 | https://sharedeus2.eus2.attest.azure.net        |
+| AGENT_OS_BUILD                 | Operating system build information for attestation                                                            | UVC                                             |
+| AGENT_OS_DISTRO                | Operating system distribution information for attestation                                                     | UVC                                             |
+| AGENT_OS_TYPE                  | Operating system type information for attestation                                                             | UVC                                             |
 
 ## Deployment
 

agent/agent.pb.go

@@ -3,7 +3,7 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.36.5
+// 	protoc-gen-go v1.36.8
 // 	protoc        v5.29.0
 // source: agent/agent.proto
 
@@ -472,7 +472,7 @@ func (x *IMAMeasurementsResponse) GetPcr10() []byte {
 	return nil
 }
 
-type AttestationResultRequest struct {
+type AttestationTokenRequest struct {
 	state         protoimpl.MessageState `protogen:"open.v1"`
 	TokenNonce    []byte                 `protobuf:"bytes,1,opt,name=tokenNonce,proto3" json:"tokenNonce,omitempty"` // Should be less or equal 32 bytes
 	Type          int32                  `protobuf:"varint,3,opt,name=type,proto3" json:"type,omitempty"`
@@ -480,20 +480,20 @@ type AttestationResultRequest struct {
 	sizeCache     protoimpl.SizeCache
 }
 
-func (x *AttestationResultRequest) Reset() {
-	*x = AttestationResultRequest{}
+func (x *AttestationTokenRequest) Reset() {
+	*x = AttestationTokenRequest{}
 	mi := &file_agent_agent_proto_msgTypes[10]
 	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 	ms.StoreMessageInfo(mi)
 }
 
-func (x *AttestationResultRequest) String() string {
+func (x *AttestationTokenRequest) String() string {
 	return protoimpl.X.MessageStringOf(x)
 }
 
-func (*AttestationResultRequest) ProtoMessage() {}
+func (*AttestationTokenRequest) ProtoMessage() {}
 
-func (x *AttestationResultRequest) ProtoReflect() protoreflect.Message {
+func (x *AttestationTokenRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[10]
 	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
@@ -505,46 +505,46 @@ func (x *AttestationResultRequest) ProtoReflect() protoreflect.Message {
 	return mi.MessageOf(x)
 }
 
-// Deprecated: Use AttestationResultRequest.ProtoReflect.Descriptor instead.
-func (*AttestationResultRequest) Descriptor() ([]byte, []int) {
+// Deprecated: Use AttestationTokenRequest.ProtoReflect.Descriptor instead.
+func (*AttestationTokenRequest) Descriptor() ([]byte, []int) {
 	return file_agent_agent_proto_rawDescGZIP(), []int{10}
 }
 
-func (x *AttestationResultRequest) GetTokenNonce() []byte {
+func (x *AttestationTokenRequest) GetTokenNonce() []byte {
 	if x != nil {
 		return x.TokenNonce
 	}
 	return nil
 }
 
-func (x *AttestationResultRequest) GetType() int32 {
+func (x *AttestationTokenRequest) GetType() int32 {
 	if x != nil {
 		return x.Type
 	}
 	return 0
 }
 
-type AttestationResultResponse struct {
+type AttestationTokenResponse 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 *AttestationResultResponse) Reset() {
-	*x = AttestationResultResponse{}
+func (x *AttestationTokenResponse) Reset() {
+	*x = AttestationTokenResponse{}
 	mi := &file_agent_agent_proto_msgTypes[11]
 	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 	ms.StoreMessageInfo(mi)
 }
 
-func (x *AttestationResultResponse) String() string {
+func (x *AttestationTokenResponse) String() string {
 	return protoimpl.X.MessageStringOf(x)
 }
 
-func (*AttestationResultResponse) ProtoMessage() {}
+func (*AttestationTokenResponse) ProtoMessage() {}
 
-func (x *AttestationResultResponse) ProtoReflect() protoreflect.Message {
+func (x *AttestationTokenResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[11]
 	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
@@ -556,12 +556,12 @@ func (x *AttestationResultResponse) ProtoReflect() protoreflect.Message {
 	return mi.MessageOf(x)
 }
 
-// Deprecated: Use AttestationResultResponse.ProtoReflect.Descriptor instead.
-func (*AttestationResultResponse) Descriptor() ([]byte, []int) {
+// Deprecated: Use AttestationTokenResponse.ProtoReflect.Descriptor instead.
+func (*AttestationTokenResponse) Descriptor() ([]byte, []int) {
 	return file_agent_agent_proto_rawDescGZIP(), []int{11}
 }
 
-func (x *AttestationResultResponse) GetFile() []byte {
+func (x *AttestationTokenResponse) GetFile() []byte {
 	if x != nil {
 		return x.File
 	}
@@ -570,76 +570,44 @@ func (x *AttestationResultResponse) GetFile() []byte {
 
 var File_agent_agent_proto protoreflect.FileDescriptor
 
-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,
-	0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x09, 0x61, 0x6c,
-	0x67, 0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x12, 0x22, 0x0a, 0x0c, 0x72, 0x65, 0x71, 0x75, 0x69,
-	0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x0c, 0x72,
-	0x65, 0x71, 0x75, 0x69, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x22, 0x0e, 0x0a, 0x0c, 0x41,
-	0x6c, 0x67, 0x6f, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x43, 0x0a, 0x0b, 0x44,
-	0x61, 0x74, 0x61, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x61,
-	0x74, 0x61, 0x73, 0x65, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x61, 0x74,
-	0x61, 0x73, 0x65, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x66, 0x69, 0x6c, 0x65, 0x6e, 0x61, 0x6d, 0x65,
-	0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x66, 0x69, 0x6c, 0x65, 0x6e, 0x61, 0x6d, 0x65,
-	0x22, 0x0e, 0x0a, 0x0c, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
-	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, 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, 0x22, 0x18, 0x0a, 0x16, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61,
-	0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74,
-	0x22, 0x43, 0x0a, 0x17, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61, 0x73, 0x75, 0x72, 0x65, 0x6d, 0x65,
-	0x6e, 0x74, 0x73, 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, 0x12,
-	0x14, 0x0a, 0x05, 0x70, 0x63, 0x72, 0x31, 0x30, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x05,
-	0x70, 0x63, 0x72, 0x31, 0x30, 0x22, 0x4e, 0x0a, 0x18, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61,
-	0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73,
-	0x74, 0x12, 0x1e, 0x0a, 0x0a, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 0x4e, 0x6f, 0x6e, 0x63, 0x65, 0x18,
-	0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x0a, 0x74, 0x6f, 0x6b, 0x65, 0x6e, 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, 0x2f, 0x0a, 0x19, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61,
-	0x74, 0x69, 0x6f, 0x6e, 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, 0x32, 0xad, 0x03, 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, 0x12, 0x54, 0x0a, 0x0f, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61,
-	0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73, 0x12, 0x1d, 0x2e, 0x61, 0x67, 0x65, 0x6e,
-	0x74, 0x2e, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61, 0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74,
-	0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74,
-	0x2e, 0x49, 0x4d, 0x41, 0x4d, 0x65, 0x61, 0x73, 0x75, 0x72, 0x65, 0x6d, 0x65, 0x6e, 0x74, 0x73,
-	0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x12, 0x58, 0x0a, 0x11,
-	0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c,
-	0x74, 0x12, 0x1f, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74,
-	0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65,
-	0x73, 0x74, 0x1a, 0x20, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73,
-	0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x73, 0x70,
-	0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x09, 0x5a, 0x07, 0x2e, 0x2f, 0x61, 0x67, 0x65, 0x6e,
-	0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
-})
+const file_agent_agent_proto_rawDesc = "" +
+	"\n" +
+	"\x11agent/agent.proto\x12\x05agent\"O\n" +
+	"\vAlgoRequest\x12\x1c\n" +
+	"\talgorithm\x18\x01 \x01(\fR\talgorithm\x12\"\n" +
+	"\frequirements\x18\x02 \x01(\fR\frequirements\"\x0e\n" +
+	"\fAlgoResponse\"C\n" +
+	"\vDataRequest\x12\x18\n" +
+	"\adataset\x18\x01 \x01(\fR\adataset\x12\x1a\n" +
+	"\bfilename\x18\x02 \x01(\tR\bfilename\"\x0e\n" +
+	"\fDataResponse\"\x0f\n" +
+	"\rResultRequest\"$\n" +
+	"\x0eResultResponse\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file\"b\n" +
+	"\x12AttestationRequest\x12\x1a\n" +
+	"\bteeNonce\x18\x01 \x01(\fR\bteeNonce\x12\x1c\n" +
+	"\tvtpmNonce\x18\x02 \x01(\fR\tvtpmNonce\x12\x12\n" +
+	"\x04type\x18\x03 \x01(\x05R\x04type\")\n" +
+	"\x13AttestationResponse\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file\"\x18\n" +
+	"\x16IMAMeasurementsRequest\"C\n" +
+	"\x17IMAMeasurementsResponse\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file\x12\x14\n" +
+	"\x05pcr10\x18\x02 \x01(\fR\x05pcr10\"M\n" +
+	"\x17AttestationTokenRequest\x12\x1e\n" +
+	"\n" +
+	"tokenNonce\x18\x01 \x01(\fR\n" +
+	"tokenNonce\x12\x12\n" +
+	"\x04type\x18\x03 \x01(\x05R\x04type\".\n" +
+	"\x18AttestationTokenResponse\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file2\xaf\x03\n" +
+	"\fAgentService\x123\n" +
+	"\x04Algo\x12\x12.agent.AlgoRequest\x1a\x13.agent.AlgoResponse\"\x00(\x01\x123\n" +
+	"\x04Data\x12\x12.agent.DataRequest\x1a\x13.agent.DataResponse\"\x00(\x01\x129\n" +
+	"\x06Result\x12\x14.agent.ResultRequest\x1a\x15.agent.ResultResponse\"\x000\x01\x12H\n" +
+	"\vAttestation\x12\x19.agent.AttestationRequest\x1a\x1a.agent.AttestationResponse\"\x000\x01\x12T\n" +
+	"\x0fIMAMeasurements\x12\x1d.agent.IMAMeasurementsRequest\x1a\x1e.agent.IMAMeasurementsResponse\"\x000\x01\x12Z\n" +
+	"\x15AzureAttestationToken\x12\x1e.agent.AttestationTokenRequest\x1a\x1f.agent.AttestationTokenResponse\"\x00B\tZ\a./agentb\x06proto3"
 
 var (
 	file_agent_agent_proto_rawDescOnce sync.Once
@@ -655,18 +623,18 @@ func file_agent_agent_proto_rawDescGZIP() []byte {
 
 var file_agent_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 12)
 var file_agent_agent_proto_goTypes = []any{
-	(*AlgoRequest)(nil),               // 0: agent.AlgoRequest
-	(*AlgoResponse)(nil),              // 1: agent.AlgoResponse
-	(*DataRequest)(nil),               // 2: agent.DataRequest
-	(*DataResponse)(nil),              // 3: agent.DataResponse
-	(*ResultRequest)(nil),             // 4: agent.ResultRequest
-	(*ResultResponse)(nil),            // 5: agent.ResultResponse
-	(*AttestationRequest)(nil),        // 6: agent.AttestationRequest
-	(*AttestationResponse)(nil),       // 7: agent.AttestationResponse
-	(*IMAMeasurementsRequest)(nil),    // 8: agent.IMAMeasurementsRequest
-	(*IMAMeasurementsResponse)(nil),   // 9: agent.IMAMeasurementsResponse
-	(*AttestationResultRequest)(nil),  // 10: agent.AttestationResultRequest
-	(*AttestationResultResponse)(nil), // 11: agent.AttestationResultResponse
+	(*AlgoRequest)(nil),              // 0: agent.AlgoRequest
+	(*AlgoResponse)(nil),             // 1: agent.AlgoResponse
+	(*DataRequest)(nil),              // 2: agent.DataRequest
+	(*DataResponse)(nil),             // 3: agent.DataResponse
+	(*ResultRequest)(nil),            // 4: agent.ResultRequest
+	(*ResultResponse)(nil),           // 5: agent.ResultResponse
+	(*AttestationRequest)(nil),       // 6: agent.AttestationRequest
+	(*AttestationResponse)(nil),      // 7: agent.AttestationResponse
+	(*IMAMeasurementsRequest)(nil),   // 8: agent.IMAMeasurementsRequest
+	(*IMAMeasurementsResponse)(nil),  // 9: agent.IMAMeasurementsResponse
+	(*AttestationTokenRequest)(nil),  // 10: agent.AttestationTokenRequest
+	(*AttestationTokenResponse)(nil), // 11: agent.AttestationTokenResponse
 }
 var file_agent_agent_proto_depIdxs = []int32{
 	0,  // 0: agent.AgentService.Algo:input_type -> agent.AlgoRequest
@@ -674,13 +642,13 @@ var file_agent_agent_proto_depIdxs = []int32{
 	4,  // 2: agent.AgentService.Result:input_type -> agent.ResultRequest
 	6,  // 3: agent.AgentService.Attestation:input_type -> agent.AttestationRequest
 	8,  // 4: agent.AgentService.IMAMeasurements:input_type -> agent.IMAMeasurementsRequest
-	10, // 5: agent.AgentService.AttestationResult:input_type -> agent.AttestationResultRequest
+	10, // 5: agent.AgentService.AzureAttestationToken:input_type -> agent.AttestationTokenRequest
 	1,  // 6: agent.AgentService.Algo:output_type -> agent.AlgoResponse
 	3,  // 7: agent.AgentService.Data:output_type -> agent.DataResponse
 	5,  // 8: agent.AgentService.Result:output_type -> agent.ResultResponse
 	7,  // 9: agent.AgentService.Attestation:output_type -> agent.AttestationResponse
 	9,  // 10: agent.AgentService.IMAMeasurements:output_type -> agent.IMAMeasurementsResponse
-	11, // 11: agent.AgentService.AttestationResult:output_type -> agent.AttestationResultResponse
+	11, // 11: agent.AgentService.AzureAttestationToken:output_type -> agent.AttestationTokenResponse
 	6,  // [6:12] is the sub-list for method output_type
 	0,  // [0:6] is the sub-list for method input_type
 	0,  // [0:0] is the sub-list for extension type_name

agent/agent.proto

@@ -13,7 +13,7 @@ service AgentService {
   rpc Result(ResultRequest) returns (stream ResultResponse) {}
   rpc Attestation(AttestationRequest) returns (stream AttestationResponse) {}
   rpc IMAMeasurements(IMAMeasurementsRequest) returns (stream IMAMeasurementsResponse) {}
-  rpc AttestationResult(AttestationResultRequest) returns (AttestationResultResponse) {}
+  rpc AzureAttestationToken(AttestationTokenRequest) returns (AttestationTokenResponse) {}
 }
 
 message AlgoRequest {
@@ -55,10 +55,10 @@ message IMAMeasurementsResponse {
   bytes pcr10 = 2;
 }
 
-message AttestationResultRequest{
+message AttestationTokenRequest{
   bytes tokenNonce = 1; // Should be less or equal 32 bytes
   int32 type = 3;
 }
-message AttestationResultResponse{
+message AttestationTokenResponse{
   bytes file = 1;
 }

agent/agent_grpc.pb.go

@@ -22,12 +22,12 @@ import (
 const _ = grpc.SupportPackageIsVersion9
 
 const (
-	AgentService_Algo_FullMethodName              = "/agent.AgentService/Algo"
-	AgentService_Data_FullMethodName              = "/agent.AgentService/Data"
-	AgentService_Result_FullMethodName            = "/agent.AgentService/Result"
-	AgentService_Attestation_FullMethodName       = "/agent.AgentService/Attestation"
-	AgentService_IMAMeasurements_FullMethodName   = "/agent.AgentService/IMAMeasurements"
-	AgentService_AttestationResult_FullMethodName = "/agent.AgentService/AttestationResult"
+	AgentService_Algo_FullMethodName                  = "/agent.AgentService/Algo"
+	AgentService_Data_FullMethodName                  = "/agent.AgentService/Data"
+	AgentService_Result_FullMethodName                = "/agent.AgentService/Result"
+	AgentService_Attestation_FullMethodName           = "/agent.AgentService/Attestation"
+	AgentService_IMAMeasurements_FullMethodName       = "/agent.AgentService/IMAMeasurements"
+	AgentService_AzureAttestationToken_FullMethodName = "/agent.AgentService/AzureAttestationToken"
 )
 
 // AgentServiceClient is the client API for AgentService service.
@@ -39,7 +39,7 @@ type AgentServiceClient interface {
 	Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ResultResponse], error)
 	Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AttestationResponse], error)
 	IMAMeasurements(ctx context.Context, in *IMAMeasurementsRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[IMAMeasurementsResponse], error)
-	AttestationResult(ctx context.Context, in *AttestationResultRequest, opts ...grpc.CallOption) (*AttestationResultResponse, error)
+	AzureAttestationToken(ctx context.Context, in *AttestationTokenRequest, opts ...grpc.CallOption) (*AttestationTokenResponse, error)
 }
 
 type agentServiceClient struct {
@@ -133,10 +133,10 @@ func (c *agentServiceClient) IMAMeasurements(ctx context.Context, in *IMAMeasure
 // This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
 type AgentService_IMAMeasurementsClient = grpc.ServerStreamingClient[IMAMeasurementsResponse]
 
-func (c *agentServiceClient) AttestationResult(ctx context.Context, in *AttestationResultRequest, opts ...grpc.CallOption) (*AttestationResultResponse, error) {
+func (c *agentServiceClient) AzureAttestationToken(ctx context.Context, in *AttestationTokenRequest, opts ...grpc.CallOption) (*AttestationTokenResponse, error) {
 	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
-	out := new(AttestationResultResponse)
-	err := c.cc.Invoke(ctx, AgentService_AttestationResult_FullMethodName, in, out, cOpts...)
+	out := new(AttestationTokenResponse)
+	err := c.cc.Invoke(ctx, AgentService_AzureAttestationToken_FullMethodName, in, out, cOpts...)
 	if err != nil {
 		return nil, err
 	}
@@ -152,7 +152,7 @@ type AgentServiceServer interface {
 	Result(*ResultRequest, grpc.ServerStreamingServer[ResultResponse]) error
 	Attestation(*AttestationRequest, grpc.ServerStreamingServer[AttestationResponse]) error
 	IMAMeasurements(*IMAMeasurementsRequest, grpc.ServerStreamingServer[IMAMeasurementsResponse]) error
-	AttestationResult(context.Context, *AttestationResultRequest) (*AttestationResultResponse, error)
+	AzureAttestationToken(context.Context, *AttestationTokenRequest) (*AttestationTokenResponse, error)
 	mustEmbedUnimplementedAgentServiceServer()
 }
 
@@ -178,8 +178,8 @@ func (UnimplementedAgentServiceServer) Attestation(*AttestationRequest, grpc.Ser
 func (UnimplementedAgentServiceServer) IMAMeasurements(*IMAMeasurementsRequest, grpc.ServerStreamingServer[IMAMeasurementsResponse]) error {
 	return status.Errorf(codes.Unimplemented, "method IMAMeasurements not implemented")
 }
-func (UnimplementedAgentServiceServer) AttestationResult(context.Context, *AttestationResultRequest) (*AttestationResultResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method AttestationResult not implemented")
+func (UnimplementedAgentServiceServer) AzureAttestationToken(context.Context, *AttestationTokenRequest) (*AttestationTokenResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "method AzureAttestationToken not implemented")
 }
 func (UnimplementedAgentServiceServer) mustEmbedUnimplementedAgentServiceServer() {}
 func (UnimplementedAgentServiceServer) testEmbeddedByValue()                      {}
@@ -249,20 +249,20 @@ func _AgentService_IMAMeasurements_Handler(srv interface{}, stream grpc.ServerSt
 // This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
 type AgentService_IMAMeasurementsServer = grpc.ServerStreamingServer[IMAMeasurementsResponse]
 
-func _AgentService_AttestationResult_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(AttestationResultRequest)
+func _AgentService_AzureAttestationToken_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(AttestationTokenRequest)
 	if err := dec(in); err != nil {
 		return nil, err
 	}
 	if interceptor == nil {
-		return srv.(AgentServiceServer).AttestationResult(ctx, in)
+		return srv.(AgentServiceServer).AzureAttestationToken(ctx, in)
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: AgentService_AttestationResult_FullMethodName,
+		FullMethod: AgentService_AzureAttestationToken_FullMethodName,
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(AgentServiceServer).AttestationResult(ctx, req.(*AttestationResultRequest))
+		return srv.(AgentServiceServer).AzureAttestationToken(ctx, req.(*AttestationTokenRequest))
 	}
 	return interceptor(ctx, in, info, handler)
 }
@@ -275,8 +275,8 @@ var AgentService_ServiceDesc = grpc.ServiceDesc{
 	HandlerType: (*AgentServiceServer)(nil),
 	Methods: []grpc.MethodDesc{
 		{
-			MethodName: "AttestationResult",
-			Handler:    _AgentService_AttestationResult_Handler,
+			MethodName: "AzureAttestationToken",
+			Handler:    _AgentService_AzureAttestationToken_Handler,
 		},
 	},
 	Streams: []grpc.StreamDesc{

agent/algorithm/logging/logging_test.go

@@ -6,7 +6,7 @@ import (
 	"strings"
 	"testing"
 
-	mglog "github.com/absmach/magistrala/logger"
+	mglog "github.com/absmach/supermq/logger"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/ultravioletrs/cocos/agent/events/mocks"

agent/algorithm/mocks/algorithm.go

@@ -1,11 +1,29 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
-import mock "github.com/stretchr/testify/mock"
+import (
+	mock "github.com/stretchr/testify/mock"
+)
+
+// NewAlgorithm creates a new instance of Algorithm. 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 NewAlgorithm(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Algorithm {
+	mock := &Algorithm{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
 
 // Algorithm is an autogenerated mock type for the Algorithm type
 type Algorithm struct {
@@ -20,21 +38,20 @@ func (_m *Algorithm) EXPECT() *Algorithm_Expecter {
 	return &Algorithm_Expecter{mock: &_m.Mock}
 }
 
-// Run provides a mock function with no fields
-func (_m *Algorithm) Run() error {
-	ret := _m.Called()
+// Run provides a mock function for the type Algorithm
+func (_mock *Algorithm) Run() error {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for Run")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -55,8 +72,8 @@ func (_c *Algorithm_Run_Call) Run(run func()) *Algorithm_Run_Call {
 	return _c
 }
 
-func (_c *Algorithm_Run_Call) Return(_a0 error) *Algorithm_Run_Call {
-	_c.Call.Return(_a0)
+func (_c *Algorithm_Run_Call) Return(err error) *Algorithm_Run_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -65,21 +82,20 @@ func (_c *Algorithm_Run_Call) RunAndReturn(run func() error) *Algorithm_Run_Call
 	return _c
 }
 
-// Stop provides a mock function with no fields
-func (_m *Algorithm) Stop() error {
-	ret := _m.Called()
+// Stop provides a mock function for the type Algorithm
+func (_mock *Algorithm) Stop() error {
+	ret := _mock.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()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -100,8 +116,8 @@ func (_c *Algorithm_Stop_Call) Run(run func()) *Algorithm_Stop_Call {
 	return _c
 }
 
-func (_c *Algorithm_Stop_Call) Return(_a0 error) *Algorithm_Stop_Call {
-	_c.Call.Return(_a0)
+func (_c *Algorithm_Stop_Call) Return(err error) *Algorithm_Stop_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -109,17 +125,3 @@ func (_c *Algorithm_Stop_Call) RunAndReturn(run func() error) *Algorithm_Stop_Ca
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewAlgorithm creates a new instance of Algorithm. 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 NewAlgorithm(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *Algorithm {
-	mock := &Algorithm{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/api/grpc/endpoint.go

@@ -11,7 +11,7 @@ import (
 )
 
 func algoEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(algoReq)
 
 		if err := req.validate(); err != nil {
@@ -30,7 +30,7 @@ func algoEndpoint(svc agent.Service) endpoint.Endpoint {
 }
 
 func dataEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(dataReq)
 
 		if err := req.validate(); err != nil {
@@ -49,7 +49,7 @@ func dataEndpoint(svc agent.Service) endpoint.Endpoint {
 }
 
 func resultEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(resultReq)
 
 		if err := req.validate(); err != nil {
@@ -65,7 +65,7 @@ func resultEndpoint(svc agent.Service) endpoint.Endpoint {
 }
 
 func attestationEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(attestationReq)
 
 		if err := req.validate(); err != nil {
@@ -81,7 +81,7 @@ func attestationEndpoint(svc agent.Service) endpoint.Endpoint {
 }
 
 func imaMeasurementsEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
+	return func(ctx context.Context, request any) (any, error) {
 		req := request.(imaMeasurementsReq)
 
 		if err := req.validate(); err != nil {
@@ -96,16 +96,16 @@ func imaMeasurementsEndpoint(svc agent.Service) endpoint.Endpoint {
 	}
 }
 
-func attestationResultEndpoint(svc agent.Service) endpoint.Endpoint {
-	return func(ctx context.Context, request interface{}) (interface{}, error) {
-		req := request.(FetchAttestationResultReq)
+func azureAttestationTokenEndpoint(svc agent.Service) endpoint.Endpoint {
+	return func(ctx context.Context, request any) (any, error) {
+		req := request.(azureAttestationTokenReq)
 		if err := req.validate(); err != nil {
-			return fetchAttestationResultRes{}, err
+			return fetchAttestationTokenRes{}, err
 		}
-		file, err := svc.AttestationResult(ctx, req.tokenNonce, attestation.PlatformType(req.AttType))
+		file, err := svc.AzureAttestationToken(ctx, req.tokenNonce)
 		if err != nil {
-			return fetchAttestationResultRes{}, err
+			return fetchAttestationTokenRes{}, err
 		}
-		return fetchAttestationResultRes{File: file}, nil
+		return fetchAttestationTokenRes{File: file}, nil
 	}
 }

agent/api/grpc/endpoint_test.go

@@ -174,23 +174,23 @@ func TestAttestationEndpoint(t *testing.T) {
 	}
 }
 
-func TestAttestationResultEndpoint(t *testing.T) {
+func TestAttestationTokenEndpoint(t *testing.T) {
 	svc := new(mocks.Service)
 	tests := []struct {
 		name        string
-		req         FetchAttestationResultReq
+		req         azureAttestationTokenReq
 		mockErr     error
 		expectedErr bool
 	}{
 		{
 			name:        "Success",
-			req:         FetchAttestationResultReq{tokenNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: attestation.AzureToken},
+			req:         azureAttestationTokenReq{tokenNonce: sha3.Sum256([]byte("vtpm nonce"))},
 			mockErr:     nil,
 			expectedErr: false,
 		},
 		{
 			name:        "Service Error",
-			req:         FetchAttestationResultReq{tokenNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: attestation.AzureToken},
+			req:         azureAttestationTokenReq{tokenNonce: sha3.Sum256([]byte("vtpm nonce"))},
 			mockErr:     errors.New("mock failure"),
 			expectedErr: true,
 		},
@@ -200,21 +200,21 @@ func TestAttestationResultEndpoint(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			// Only call service mock if validation is expected to pass
 			if err := tt.req.validate(); err == nil {
-				svc.On("AttestationResult", mock.Anything, tt.req.tokenNonce, attestation.PlatformType(tt.req.AttType)).
+				svc.On("AzureAttestationToken", mock.Anything, tt.req.tokenNonce).
 					Return([]byte("mock file"), tt.mockErr).Once()
 			}
 
-			endpoint := attestationResultEndpoint(svc)
+			endpoint := azureAttestationTokenEndpoint(svc)
 			res, err := endpoint(context.Background(), tt.req)
 
 			if (err != nil) != tt.expectedErr {
-				t.Errorf("attestationResultEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
+				t.Errorf("attestationTokenEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
 			}
 
 			if !tt.expectedErr {
-				r, ok := res.(fetchAttestationResultRes)
+				r, ok := res.(fetchAttestationTokenRes)
 				if !ok {
-					t.Errorf("attestationResultEndpoint() returned unexpected type %T", res)
+					t.Errorf("attestationTokenEndpoint() returned unexpected type %T", res)
 				}
 				if string(r.File) != "mock file" {
 					t.Errorf("expected file content 'mock file', got %s", r.File)

agent/api/grpc/interceptors.go

@@ -31,7 +31,7 @@ func NewAuthInterceptor(authSvc auth.Authenticator) (grpc.UnaryServerInterceptor
 }
 
 func (s *authInterceptor) AuthStreamInterceptor() grpc.StreamServerInterceptor {
-	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+	return func(srv any, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
 		switch info.FullMethod {
 		case agent.AgentService_Algo_FullMethodName:
 			if _, err := s.auth.AuthenticateUser(stream.Context(), auth.AlgorithmProviderRole); err != nil {
@@ -59,7 +59,7 @@ func (s *authInterceptor) AuthStreamInterceptor() grpc.StreamServerInterceptor {
 }
 
 func (s *authInterceptor) AuthUnaryInterceptor() grpc.UnaryServerInterceptor {
-	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
 		switch info.FullMethod {
 		case agent.AgentService_Result_FullMethodName:
 			ctx, err := s.auth.AuthenticateUser(ctx, auth.ConsumerRole)

agent/api/grpc/interceptors_test.go

@@ -58,7 +58,7 @@ func TestAuthUnaryInterceptor(t *testing.T) {
 			}
 			unaryInt, _ := NewAuthInterceptor(authmock)
 
-			_, err := unaryInt(context.Background(), nil, &grpc.UnaryServerInfo{FullMethod: tt.method}, func(ctx context.Context, req interface{}) (interface{}, error) {
+			_, err := unaryInt(context.Background(), nil, &grpc.UnaryServerInfo{FullMethod: tt.method}, func(ctx context.Context, req any) (any, error) {
 				return nil, nil
 			})
 
@@ -129,7 +129,7 @@ func TestAuthStreamInterceptor(t *testing.T) {
 			}
 			_, streamInt := NewAuthInterceptor(authmock)
 
-			err := streamInt(nil, &mockServerStream{ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs())}, &grpc.StreamServerInfo{FullMethod: tt.method}, func(srv interface{}, stream grpc.ServerStream) error {
+			err := streamInt(nil, &mockServerStream{ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs())}, &grpc.StreamServerInfo{FullMethod: tt.method}, func(srv any, stream grpc.ServerStream) error {
 				return nil
 			})
 

agent/api/grpc/requests.go

@@ -29,7 +29,7 @@ type dataReq struct {
 
 func (req dataReq) validate() error {
 	if len(req.Dataset) == 0 {
-		return errors.New("dataset CSV file is required")
+		return errors.New("dataset is required")
 	}
 	return nil
 }
@@ -47,22 +47,21 @@ type attestationReq struct {
 	AttType   attestation.PlatformType
 }
 
-type FetchAttestationResultReq struct {
+type azureAttestationTokenReq struct {
 	tokenNonce [vtpm.Nonce]byte
-	AttType    attestation.PlatformType
 }
 
 func (req attestationReq) validate() error {
 	return validateAttestationType(req.AttType)
 }
 
-func (req FetchAttestationResultReq) validate() error {
-	return validateAttestationType(req.AttType)
+func (req azureAttestationTokenReq) validate() error {
+	return nil
 }
 
 func validateAttestationType(attType attestation.PlatformType) error {
 	switch attType {
-	case attestation.SNP, attestation.VTPM, attestation.SNPvTPM, attestation.TDX, attestation.AzureToken:
+	case attestation.SNP, attestation.VTPM, attestation.SNPvTPM, attestation.TDX:
 		return nil
 	default:
 		return errors.New("invalid attestation type")

agent/api/grpc/responses.go

@@ -7,7 +7,7 @@ type algoRes struct{}
 type dataRes struct{}
 
 type resultRes struct {
-	File []byte `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"`
+	File []byte
 }
 
 type attestationRes struct {
@@ -19,6 +19,6 @@ type imaMeasurementsRes struct {
 	PCR10 []byte
 }
 
-type fetchAttestationResultRes struct {
-	File []byte `protobuf:"bytes,1,opt,name=AttestationResult,proto3" json:"AttestationResult,omitempty"`
+type fetchAttestationTokenRes struct {
+	File []byte
 }

agent/api/grpc/server.go

@@ -74,10 +74,10 @@ func NewServer(svc agent.Service) agent.AgentServiceServer {
 			decodeRequest:  decodeIMAMeasurementsRequest,
 			encodeResponse: encodeIMAMeasurementsResponse,
 		},
-		"attestationResult": {
-			endpoint:       attestationResultEndpoint,
-			decodeRequest:  decodeAttestationResultRequest,
-			encodeResponse: encodeAttestationResultResponse,
+		"azureAttestationToken": {
+			endpoint:       azureAttestationTokenEndpoint,
+			decodeRequest:  decodeAttestationTokenRequest,
+			encodeResponse: encodeAttestationTokenResponse,
 		},
 	}
 
@@ -96,7 +96,7 @@ func NewServer(svc agent.Service) agent.AgentServiceServer {
 	}
 }
 
-func decodeAlgoRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeAlgoRequest(_ context.Context, grpcReq any) (any, error) {
 	req := grpcReq.(*agent.AlgoRequest)
 	return algoReq{
 		Algorithm:    req.Algorithm,
@@ -104,11 +104,11 @@ func decodeAlgoRequest(_ context.Context, grpcReq interface{}) (interface{}, err
 	}, nil
 }
 
-func encodeAlgoResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeAlgoResponse(_ context.Context, response any) (any, error) {
 	return &agent.AlgoResponse{}, nil
 }
 
-func decodeDataRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeDataRequest(_ context.Context, grpcReq any) (any, error) {
 	req := grpcReq.(*agent.DataRequest)
 	return dataReq{
 		Dataset:  req.Dataset,
@@ -116,22 +116,22 @@ func decodeDataRequest(_ context.Context, grpcReq interface{}) (interface{}, err
 	}, nil
 }
 
-func encodeDataResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeDataResponse(_ context.Context, response any) (any, error) {
 	return &agent.DataResponse{}, nil
 }
 
-func decodeResultRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeResultRequest(_ context.Context, grpcReq any) (any, error) {
 	return resultReq{}, nil
 }
 
-func encodeResultResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeResultResponse(_ context.Context, response any) (any, error) {
 	res := response.(resultRes)
 	return &agent.ResultResponse{
 		File: res.File,
 	}, nil
 }
 
-func validateNonce(nonce []byte, maxLen int, target interface{}) error {
+func validateNonce(nonce []byte, maxLen int, target any) error {
 	if len(nonce) > maxLen {
 		switch maxLen {
 		case quoteprovider.Nonce:
@@ -154,7 +154,7 @@ func validateNonce(nonce []byte, maxLen int, target interface{}) error {
 	return nil
 }
 
-func decodeAttestationRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeAttestationRequest(_ context.Context, grpcReq any) (any, error) {
 	req := grpcReq.(*agent.AttestationRequest)
 	var reportData [quoteprovider.Nonce]byte
 	var nonce [vtpm.Nonce]byte
@@ -174,39 +174,37 @@ func decodeAttestationRequest(_ context.Context, grpcReq interface{}) (interface
 	}, nil
 }
 
-func encodeAttestationResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeAttestationResponse(_ context.Context, response any) (any, error) {
 	res := response.(attestationRes)
 	return &agent.AttestationResponse{
 		File: res.File,
 	}, nil
 }
 
-func decodeAttestationResultRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
-	req := grpcReq.(*agent.AttestationResultRequest)
+func decodeAttestationTokenRequest(_ context.Context, grpcReq any) (any, error) {
+	req := grpcReq.(*agent.AttestationTokenRequest)
 	var nonce [vtpm.Nonce]byte
 
 	if err := validateNonce(req.TokenNonce, vtpm.Nonce, &nonce); err != nil {
 		return nil, err
 	}
-
-	return FetchAttestationResultReq{
+	return azureAttestationTokenReq{
 		tokenNonce: nonce,
-		AttType:    attestation.PlatformType(req.Type),
 	}, nil
 }
 
-func encodeAttestationResultResponse(_ context.Context, response interface{}) (interface{}, error) {
-	res := response.(fetchAttestationResultRes)
-	return &agent.AttestationResultResponse{
+func encodeAttestationTokenResponse(_ context.Context, response any) (any, error) {
+	res := response.(fetchAttestationTokenRes)
+	return &agent.AttestationTokenResponse{
 		File: res.File,
 	}, nil
 }
 
-func decodeIMAMeasurementsRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
+func decodeIMAMeasurementsRequest(_ context.Context, grpcReq any) (any, error) {
 	return imaMeasurementsReq{}, nil
 }
 
-func encodeIMAMeasurementsResponse(_ context.Context, response interface{}) (interface{}, error) {
+func encodeIMAMeasurementsResponse(_ context.Context, response any) (any, error) {
 	res := response.(imaMeasurementsRes)
 	return &agent.IMAMeasurementsResponse{
 		File:  res.File,
@@ -217,10 +215,10 @@ func encodeIMAMeasurementsResponse(_ context.Context, response interface{}) (int
 func (s *grpcServer) streamingHandler(
 	ctx context.Context,
 	handlerName string,
-	req interface{},
-	stream interface{},
+	req any,
+	stream any,
 	sendFn func([]byte) error,
-	getFileData func(interface{}) []byte,
+	getFileData func(any) []byte,
 ) error {
 	handler, ok := s.handlers[handlerName]
 	if !ok {
@@ -352,7 +350,7 @@ func (s *grpcServer) Result(req *agent.ResultRequest, stream agent.AgentService_
 		func(data []byte) error {
 			return stream.Send(&agent.ResultResponse{File: data})
 		},
-		func(res interface{}) []byte {
+		func(res any) []byte {
 			return res.(*agent.ResultResponse).File
 		},
 	)
@@ -367,7 +365,7 @@ func (s *grpcServer) Attestation(req *agent.AttestationRequest, stream agent.Age
 		func(data []byte) error {
 			return stream.Send(&agent.AttestationResponse{File: data})
 		},
-		func(res interface{}) []byte {
+		func(res any) []byte {
 			return res.(*agent.AttestationResponse).File
 		},
 	)
@@ -398,6 +396,20 @@ func (s *grpcServer) IMAMeasurements(req *agent.IMAMeasurementsRequest, stream a
 	)
 }
 
+func (s *grpcServer) AzureAttestationToken(ctx context.Context, req *agent.AttestationTokenRequest) (*agent.AttestationTokenResponse, error) {
+	_, res, err := s.handlers["azureAttestationToken"].ServeGRPC(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+
+	rr, ok := res.(*agent.AttestationTokenResponse)
+	if !ok {
+		return nil, status.Error(codes.Internal, "failed to cast response to AttestationTokenResponse")
+	}
+
+	return rr, nil
+}
+
 func (s *grpcServer) streamDualBuffers(
 	buf1, buf2 *bytes.Buffer,
 	sendFn func([]byte, []byte) error,
@@ -426,17 +438,3 @@ func (s *grpcServer) streamDualBuffers(
 	}
 	return nil
 }
-
-func (s *grpcServer) AttestationResult(ctx context.Context, req *agent.AttestationResultRequest) (*agent.AttestationResultResponse, error) {
-	_, res, err := s.handlers["attestationResult"].ServeGRPC(ctx, req)
-	if err != nil {
-		return nil, err
-	}
-
-	rr, ok := res.(*agent.AttestationResultResponse)
-	if !ok {
-		return nil, status.Error(codes.Internal, "failed to cast response to AttestationResultResponse")
-	}
-
-	return rr, nil
-}

agent/api/grpc/server_test.go

@@ -128,7 +128,7 @@ func TestNewServer(t *testing.T) {
 	assert.Len(t, grpcServer.handlers, 6) // Should have 6 handlers
 
 	// Check that all expected handlers are present
-	expectedHandlers := []string{"algo", "data", "result", "attestation", "imaMeasurements", "attestationResult"}
+	expectedHandlers := []string{"algo", "data", "result", "attestation", "imaMeasurements", "azureAttestationToken"}
 	for _, handler := range expectedHandlers {
 		assert.Contains(t, grpcServer.handlers, handler)
 		assert.NotNil(t, grpcServer.handlers[handler])
@@ -267,17 +267,17 @@ func TestIMAMeasurements(t *testing.T) {
 	mockStream.AssertExpectations(t)
 }
 
-func TestAttestationResult(t *testing.T) {
+func TestAttestationToken(t *testing.T) {
 	mockService := new(mocks.Service)
 	server := NewServer(mockService)
 
-	attestationData := []byte("attestation result data")
+	attestationData := []byte("attestation token data")
 	vtpmNonce := [vtpm.Nonce]byte{}
 	attestationType := attestation.SNP
 
-	mockService.On("AttestationResult", mock.Anything, vtpmNonce, attestationType).Return(attestationData, nil)
+	mockService.On("AzureAttestationToken", mock.Anything, vtpmNonce).Return(attestationData, nil)
 
-	resp, err := server.AttestationResult(context.Background(), &agent.AttestationResultRequest{
+	resp, err := server.AzureAttestationToken(context.Background(), &agent.AttestationTokenRequest{
 		TokenNonce: vtpmNonce[:],
 		Type:       int32(attestationType),
 	})
@@ -428,34 +428,31 @@ func TestEncodeAttestationResponse(t *testing.T) {
 	assert.Equal(t, &agent.AttestationResponse{File: []byte("attestation")}, encoded)
 }
 
-func TestDecodeAttestationResultRequest(t *testing.T) {
+func TestDecodeAttestationTokenRequest(t *testing.T) {
 	tokenNonce := make([]byte, vtpm.Nonce)
-	req := &agent.AttestationResultRequest{
+	req := &agent.AttestationTokenRequest{
 		TokenNonce: tokenNonce,
 		Type:       int32(attestation.SNP),
 	}
 
-	decoded, err := decodeAttestationResultRequest(context.Background(), req)
+	_, err := decodeAttestationTokenRequest(context.Background(), req)
 	assert.NoError(t, err)
-
-	decodedReq := decoded.(FetchAttestationResultReq)
-	assert.Equal(t, attestation.SNP, decodedReq.AttType)
 }
 
-func TestDecodeAttestationResultRequestWithInvalidNonce(t *testing.T) {
+func TestDecodeAttestationTokenRequestWithInvalidNonce(t *testing.T) {
 	// Test with token nonce too long
 	tokenNonce := make([]byte, vtpm.Nonce+1)
-	req := &agent.AttestationResultRequest{TokenNonce: tokenNonce}
+	req := &agent.AttestationTokenRequest{TokenNonce: tokenNonce}
 
-	_, err := decodeAttestationResultRequest(context.Background(), req)
+	_, err := decodeAttestationTokenRequest(context.Background(), req)
 	assert.Error(t, err)
 	assert.Equal(t, ErrVTPMNonceLength, err)
 }
 
-func TestEncodeAttestationResultResponse(t *testing.T) {
-	encoded, err := encodeAttestationResultResponse(context.Background(), fetchAttestationResultRes{File: []byte("attestation")})
+func TestEncodeAttestationTokenResponse(t *testing.T) {
+	encoded, err := encodeAttestationTokenResponse(context.Background(), fetchAttestationTokenRes{File: []byte("attestation")})
 	assert.NoError(t, err)
-	assert.Equal(t, &agent.AttestationResultResponse{File: []byte("attestation")}, encoded)
+	assert.Equal(t, &agent.AttestationTokenResponse{File: []byte("attestation")}, encoded)
 }
 
 func TestDecodeIMAMeasurementsRequest(t *testing.T) {

agent/api/logging.go

@@ -132,9 +132,9 @@ func (lm *loggingMiddleware) IMAMeasurements(ctx context.Context) (file []byte,
 	return lm.svc.IMAMeasurements(ctx)
 }
 
-func (lm *loggingMiddleware) AttestationResult(ctx context.Context, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) (response []byte, err error) {
+func (lm *loggingMiddleware) AzureAttestationToken(ctx context.Context, nonce [vtpm.Nonce]byte) (response []byte, err error) {
 	defer func(begin time.Time) {
-		message := fmt.Sprintf("Method AttestationResult took %s to complete", time.Since(begin))
+		message := fmt.Sprintf("Method AzureAttestationToken took %s to complete", time.Since(begin))
 		if err != nil {
 			lm.logger.Warn(fmt.Sprintf("%s with error: %s", message, err))
 			return
@@ -142,5 +142,5 @@ func (lm *loggingMiddleware) AttestationResult(ctx context.Context, nonce [vtpm.
 		lm.logger.Info(fmt.Sprintf("%s without errors", message))
 	}(time.Now())
 
-	return lm.svc.AttestationResult(ctx, nonce, attType)
+	return lm.svc.AzureAttestationToken(ctx, nonce)
 }

agent/api/metrics.go

@@ -101,13 +101,13 @@ func (ms *metricsMiddleware) Attestation(ctx context.Context, reportData [quotep
 	return ms.svc.Attestation(ctx, reportData, nonce, attType)
 }
 
-func (ms *metricsMiddleware) AttestationResult(ctx context.Context, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error) {
+func (ms *metricsMiddleware) AzureAttestationToken(ctx context.Context, nonce [vtpm.Nonce]byte) ([]byte, error) {
 	defer func(begin time.Time) {
-		ms.counter.With("method", "attestation_result").Add(1)
-		ms.latency.With("method", "attestation_result").Observe(time.Since(begin).Seconds())
+		ms.counter.With("method", "attestation_token").Add(1)
+		ms.latency.With("method", "attestation_token").Observe(time.Since(begin).Seconds())
 	}(time.Now())
 
-	return ms.svc.AttestationResult(ctx, nonce, attType)
+	return ms.svc.AzureAttestationToken(ctx, nonce)
 }
 
 func (ms *metricsMiddleware) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {

agent/auth/auth.go

@@ -13,7 +13,7 @@ import (
 	"crypto/x509"
 	"encoding/base64"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/ultravioletrs/cocos/agent"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -41,9 +41,9 @@ type Authenticator interface {
 }
 
 type service struct {
-	resultConsumers   []interface{}
-	datasetProviders  []interface{}
-	algorithmProvider interface{}
+	resultConsumers   []any
+	datasetProviders  []any
+	algorithmProvider any
 }
 
 func New(manifest agent.Computation) (Authenticator, error) {

agent/auth/auth_test.go

@@ -15,7 +15,7 @@ import (
 	"encoding/base64"
 	"testing"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/ultravioletrs/cocos/agent"

agent/auth/mocks/authenticator.go

@@ -1,18 +1,33 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
-	context "context"
-
-	auth "github.com/ultravioletrs/cocos/agent/auth"
+	"context"
 
 	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent/auth"
 )
 
+// NewAuthenticator creates a new instance of Authenticator. 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 NewAuthenticator(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Authenticator {
+	mock := &Authenticator{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // Authenticator is an autogenerated mock type for the Authenticator type
 type Authenticator struct {
 	mock.Mock
@@ -26,9 +41,9 @@ func (_m *Authenticator) EXPECT() *Authenticator_Expecter {
 	return &Authenticator_Expecter{mock: &_m.Mock}
 }
 
-// AuthenticateUser provides a mock function with given fields: ctx, role
-func (_m *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRole) (context.Context, error) {
-	ret := _m.Called(ctx, role)
+// AuthenticateUser provides a mock function for the type Authenticator
+func (_mock *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRole) (context.Context, error) {
+	ret := _mock.Called(ctx, role)
 
 	if len(ret) == 0 {
 		panic("no return value specified for AuthenticateUser")
@@ -36,23 +51,21 @@ func (_m *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRol
 
 	var r0 context.Context
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, auth.UserRole) (context.Context, error)); ok {
-		return rf(ctx, role)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, auth.UserRole) (context.Context, error)); ok {
+		return returnFunc(ctx, role)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, auth.UserRole) context.Context); ok {
-		r0 = rf(ctx, role)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, auth.UserRole) context.Context); ok {
+		r0 = returnFunc(ctx, role)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(context.Context)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context, auth.UserRole) error); ok {
-		r1 = rf(ctx, role)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, auth.UserRole) error); ok {
+		r1 = returnFunc(ctx, role)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -70,31 +83,28 @@ func (_e *Authenticator_Expecter) AuthenticateUser(ctx interface{}, role interfa
 
 func (_c *Authenticator_AuthenticateUser_Call) Run(run func(ctx context.Context, role auth.UserRole)) *Authenticator_AuthenticateUser_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(auth.UserRole))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 auth.UserRole
+		if args[1] != nil {
+			arg1 = args[1].(auth.UserRole)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Authenticator_AuthenticateUser_Call) Return(_a0 context.Context, _a1 error) *Authenticator_AuthenticateUser_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Authenticator_AuthenticateUser_Call) Return(context1 context.Context, err error) *Authenticator_AuthenticateUser_Call {
+	_c.Call.Return(context1, err)
 	return _c
 }
 
-func (_c *Authenticator_AuthenticateUser_Call) RunAndReturn(run func(context.Context, auth.UserRole) (context.Context, error)) *Authenticator_AuthenticateUser_Call {
+func (_c *Authenticator_AuthenticateUser_Call) RunAndReturn(run func(ctx context.Context, role auth.UserRole) (context.Context, error)) *Authenticator_AuthenticateUser_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewAuthenticator creates a new instance of Authenticator. 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 NewAuthenticator(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *Authenticator {
-	mock := &Authenticator{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/cvms/api/grpc/client.go

@@ -9,14 +9,14 @@ import (
 	"sync"
 	"time"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/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"
 	"github.com/ultravioletrs/cocos/pkg/attestation"
 	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
-	pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
+	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/protobuf/proto"
 )
@@ -24,12 +24,11 @@ import (
 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")
+	errUnknownMessageType = errors.New("unknown message type")
 )
 
 type PendingMessage struct {
@@ -46,12 +45,12 @@ type CVMSClient struct {
 	runReqManager *runRequestManager
 	sp            server.AgentServer
 	storage       storage.Storage
-	reconnectFn   func(context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error)
-	grpcClient    pkggrpc.Client
+	reconnectFn   func(context.Context) (grpc.Client, cvms.Service_ProcessClient, error)
+	grpcClient    grpc.Client
 }
 
 // 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) (pkggrpc.Client, cvms.Service_ProcessClient, error), grpcClient pkggrpc.Client) (*CVMSClient, error) {
+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) (grpc.Client, cvms.Service_ProcessClient, error), grpcClient grpc.Client) (*CVMSClient, error) {
 	store, err := storage.NewFileStorage(storageDir)
 	if err != nil {
 		return nil, err
@@ -187,7 +186,7 @@ func (client *CVMSClient) processIncomingMessage(ctx context.Context, req *cvms.
 		}
 		client.mu.Unlock()
 	default:
-		return errUnknonwMessageType
+		return errUnknownMessageType
 	}
 	return nil
 }

agent/cvms/api/grpc/client_test.go

@@ -7,7 +7,7 @@ import (
 	"testing"
 	"time"
 
-	mglog "github.com/absmach/magistrala/logger"
+	mglog "github.com/absmach/supermq/logger"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/ultravioletrs/cocos/agent/cvms"

agent/cvms/api/grpc/server_test.go

@@ -7,7 +7,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/ultravioletrs/cocos/agent/cvms"

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

@@ -1,17 +1,32 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 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"
+	"github.com/ultravioletrs/cocos/agent/cvms"
+	"github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage"
 )
 
+// 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
+}
+
 // Storage is an autogenerated mock type for the Storage type
 type Storage struct {
 	mock.Mock
@@ -25,21 +40,20 @@ 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)
+// Add provides a mock function for the type Storage
+func (_mock *Storage) Add(msg *cvms.ClientStreamMessage) error {
+	ret := _mock.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)
+	if returnFunc, ok := ret.Get(0).(func(*cvms.ClientStreamMessage) error); ok {
+		r0 = returnFunc(msg)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -56,36 +70,41 @@ func (_e *Storage_Expecter) Add(msg interface{}) *Storage_Add_Call {
 
 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))
+		var arg0 *cvms.ClientStreamMessage
+		if args[0] != nil {
+			arg0 = args[0].(*cvms.ClientStreamMessage)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Storage_Add_Call) Return(_a0 error) *Storage_Add_Call {
-	_c.Call.Return(_a0)
+func (_c *Storage_Add_Call) Return(err error) *Storage_Add_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Storage_Add_Call) RunAndReturn(run func(*cvms.ClientStreamMessage) error) *Storage_Add_Call {
+func (_c *Storage_Add_Call) RunAndReturn(run func(msg *cvms.ClientStreamMessage) error) *Storage_Add_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Clear provides a mock function with no fields
-func (_m *Storage) Clear() error {
-	ret := _m.Called()
+// Clear provides a mock function for the type Storage
+func (_mock *Storage) Clear() error {
+	ret := _mock.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()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -106,8 +125,8 @@ func (_c *Storage_Clear_Call) Run(run func()) *Storage_Clear_Call {
 	return _c
 }
 
-func (_c *Storage_Clear_Call) Return(_a0 error) *Storage_Clear_Call {
-	_c.Call.Return(_a0)
+func (_c *Storage_Clear_Call) Return(err error) *Storage_Clear_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -116,9 +135,9 @@ func (_c *Storage_Clear_Call) RunAndReturn(run func() error) *Storage_Clear_Call
 	return _c
 }
 
-// Load provides a mock function with no fields
-func (_m *Storage) Load() ([]storage.Message, error) {
-	ret := _m.Called()
+// Load provides a mock function for the type Storage
+func (_mock *Storage) Load() ([]storage.Message, error) {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for Load")
@@ -126,23 +145,21 @@ func (_m *Storage) Load() ([]storage.Message, error) {
 
 	var r0 []storage.Message
 	var r1 error
-	if rf, ok := ret.Get(0).(func() ([]storage.Message, error)); ok {
-		return rf()
+	if returnFunc, ok := ret.Get(0).(func() ([]storage.Message, error)); ok {
+		return returnFunc()
 	}
-	if rf, ok := ret.Get(0).(func() []storage.Message); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() []storage.Message); ok {
+		r0 = returnFunc()
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]storage.Message)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -163,8 +180,8 @@ func (_c *Storage_Load_Call) Run(run func()) *Storage_Load_Call {
 	return _c
 }
 
-func (_c *Storage_Load_Call) Return(_a0 []storage.Message, _a1 error) *Storage_Load_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Storage_Load_Call) Return(messages []storage.Message, err error) *Storage_Load_Call {
+	_c.Call.Return(messages, err)
 	return _c
 }
 
@@ -173,21 +190,20 @@ func (_c *Storage_Load_Call) RunAndReturn(run func() ([]storage.Message, error))
 	return _c
 }
 
-// Save provides a mock function with given fields: messages
-func (_m *Storage) Save(messages []storage.Message) error {
-	ret := _m.Called(messages)
+// Save provides a mock function for the type Storage
+func (_mock *Storage) Save(messages []storage.Message) error {
+	ret := _mock.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)
+	if returnFunc, ok := ret.Get(0).(func([]storage.Message) error); ok {
+		r0 = returnFunc(messages)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -204,31 +220,23 @@ func (_e *Storage_Expecter) Save(messages interface{}) *Storage_Save_Call {
 
 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))
+		var arg0 []storage.Message
+		if args[0] != nil {
+			arg0 = args[0].([]storage.Message)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Storage_Save_Call) Return(_a0 error) *Storage_Save_Call {
-	_c.Call.Return(_a0)
+func (_c *Storage_Save_Call) Return(err error) *Storage_Save_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Storage_Save_Call) RunAndReturn(run func([]storage.Message) error) *Storage_Save_Call {
+func (_c *Storage_Save_Call) RunAndReturn(run func(messages []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/cvms.pb.go

@@ -3,7 +3,7 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.36.5
+// 	protoc-gen-go v1.36.8
 // 	protoc        v5.29.0
 // source: agent/cvms/cvms.proto
 
@@ -431,7 +431,7 @@ type ClientStreamMessage struct {
 	//	*ClientStreamMessage_StopComputationRes
 	//	*ClientStreamMessage_AgentStateRes
 	//	*ClientStreamMessage_VTPMattestationReport
-	//	*ClientStreamMessage_AzureAttestationResult
+	//	*ClientStreamMessage_AzureAttestationToken
 	Message       isClientStreamMessage_Message `protobuf_oneof:"message"`
 	unknownFields protoimpl.UnknownFields
 	sizeCache     protoimpl.SizeCache
@@ -528,10 +528,10 @@ func (x *ClientStreamMessage) GetVTPMattestationReport() *AttestationResponse {
 	return nil
 }
 
-func (x *ClientStreamMessage) GetAzureAttestationResult() *AzureAttestationResponse {
+func (x *ClientStreamMessage) GetAzureAttestationToken() *AzureAttestationToken {
 	if x != nil {
-		if x, ok := x.Message.(*ClientStreamMessage_AzureAttestationResult); ok {
-			return x.AzureAttestationResult
+		if x, ok := x.Message.(*ClientStreamMessage_AzureAttestationToken); ok {
+			return x.AzureAttestationToken
 		}
 	}
 	return nil
@@ -565,8 +565,8 @@ type ClientStreamMessage_VTPMattestationReport struct {
 	VTPMattestationReport *AttestationResponse `protobuf:"bytes,6,opt,name=vTPMattestationReport,proto3,oneof"`
 }
 
-type ClientStreamMessage_AzureAttestationResult struct {
-	AzureAttestationResult *AzureAttestationResponse `protobuf:"bytes,7,opt,name=azureAttestationResult,proto3,oneof"`
+type ClientStreamMessage_AzureAttestationToken struct {
+	AzureAttestationToken *AzureAttestationToken `protobuf:"bytes,7,opt,name=azureAttestationToken,proto3,oneof"`
 }
 
 func (*ClientStreamMessage_AgentLog) isClientStreamMessage_Message() {}
@@ -581,7 +581,7 @@ func (*ClientStreamMessage_AgentStateRes) isClientStreamMessage_Message() {}
 
 func (*ClientStreamMessage_VTPMattestationReport) isClientStreamMessage_Message() {}
 
-func (*ClientStreamMessage_AzureAttestationResult) isClientStreamMessage_Message() {}
+func (*ClientStreamMessage_AzureAttestationToken) isClientStreamMessage_Message() {}
 
 type ServerStreamMessage struct {
 	state protoimpl.MessageState `protogen:"open.v1"`
@@ -1209,7 +1209,7 @@ func (x *AttestationResponse) GetCertSerialNumber() string {
 	return ""
 }
 
-type AzureAttestationResponse struct {
+type AzureAttestationToken struct {
 	state            protoimpl.MessageState `protogen:"open.v1"`
 	File             []byte                 `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
 	CertSerialNumber string                 `protobuf:"bytes,2,opt,name=certSerialNumber,proto3" json:"certSerialNumber,omitempty"`
@@ -1217,20 +1217,20 @@ type AzureAttestationResponse struct {
 	sizeCache        protoimpl.SizeCache
 }
 
-func (x *AzureAttestationResponse) Reset() {
-	*x = AzureAttestationResponse{}
+func (x *AzureAttestationToken) Reset() {
+	*x = AzureAttestationToken{}
 	mi := &file_agent_cvms_cvms_proto_msgTypes[17]
 	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 	ms.StoreMessageInfo(mi)
 }
 
-func (x *AzureAttestationResponse) String() string {
+func (x *AzureAttestationToken) String() string {
 	return protoimpl.X.MessageStringOf(x)
 }
 
-func (*AzureAttestationResponse) ProtoMessage() {}
+func (*AzureAttestationToken) ProtoMessage() {}
 
-func (x *AzureAttestationResponse) ProtoReflect() protoreflect.Message {
+func (x *AzureAttestationToken) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_cvms_cvms_proto_msgTypes[17]
 	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
@@ -1242,19 +1242,19 @@ func (x *AzureAttestationResponse) ProtoReflect() protoreflect.Message {
 	return mi.MessageOf(x)
 }
 
-// Deprecated: Use AzureAttestationResponse.ProtoReflect.Descriptor instead.
-func (*AzureAttestationResponse) Descriptor() ([]byte, []int) {
+// Deprecated: Use AzureAttestationToken.ProtoReflect.Descriptor instead.
+func (*AzureAttestationToken) Descriptor() ([]byte, []int) {
 	return file_agent_cvms_cvms_proto_rawDescGZIP(), []int{17}
 }
 
-func (x *AzureAttestationResponse) GetFile() []byte {
+func (x *AzureAttestationToken) GetFile() []byte {
 	if x != nil {
 		return x.File
 	}
 	return nil
 }
 
-func (x *AzureAttestationResponse) GetCertSerialNumber() string {
+func (x *AzureAttestationToken) GetCertSerialNumber() string {
 	if x != nil {
 		return x.CertSerialNumber
 	}
@@ -1263,177 +1263,94 @@ func (x *AzureAttestationResponse) GetCertSerialNumber() string {
 
 var File_agent_cvms_cvms_proto protoreflect.FileDescriptor
 
-var file_agent_cvms_cvms_proto_rawDesc = string([]byte{
-	0x0a, 0x15, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2f, 0x63, 0x76, 0x6d, 0x73, 0x2f, 0x63, 0x76, 0x6d,
-	0x73, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x04, 0x63, 0x76, 0x6d, 0x73, 0x1a, 0x1f, 0x67,
-	0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x74,
-	0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x1f,
-	0x0a, 0x0d, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x12,
-	0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x22,
-	0x35, 0x0a, 0x0d, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73,
-	0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64,
-	0x12, 0x14, 0x0a, 0x05, 0x73, 0x74, 0x61, 0x74, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52,
-	0x05, 0x73, 0x74, 0x61, 0x74, 0x65, 0x22, 0x38, 0x0a, 0x0f, 0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f,
-	0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x25, 0x0a, 0x0e, 0x63, 0x6f, 0x6d,
-	0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64,
-	0x22, 0x5a, 0x0a, 0x17, 0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74,
-	0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x25, 0x0a, 0x0e, 0x63,
-	0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e,
-	0x49, 0x64, 0x12, 0x18, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x18, 0x02, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x4a, 0x0a, 0x0b,
-	0x52, 0x75, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x25, 0x0a, 0x0e, 0x63,
-	0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e,
-	0x49, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x05, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x22, 0xde, 0x01, 0x0a, 0x0a, 0x41, 0x67, 0x65,
-	0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x65, 0x76, 0x65, 0x6e, 0x74,
-	0x5f, 0x74, 0x79, 0x70, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x65, 0x76, 0x65,
-	0x6e, 0x74, 0x54, 0x79, 0x70, 0x65, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74,
-	0x61, 0x6d, 0x70, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67,
-	0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65,
-	0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70,
-	0x12, 0x25, 0x0a, 0x0e, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f,
-	0x69, 0x64, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74,
-	0x61, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69,
-	0x6c, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c,
-	0x73, 0x12, 0x1e, 0x0a, 0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x18,
-	0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f,
-	0x72, 0x12, 0x16, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x06, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x9b, 0x01, 0x0a, 0x08, 0x41, 0x67,
-	0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x12, 0x18, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67,
-	0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65,
-	0x12, 0x25, 0x0a, 0x0e, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f,
-	0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74,
-	0x61, 0x74, 0x69, 0x6f, 0x6e, 0x49, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c,
-	0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x38, 0x0a,
-	0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b,
-	0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62,
-	0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x09, 0x74, 0x69,
-	0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x22, 0xed, 0x03, 0x0a, 0x13, 0x43, 0x6c, 0x69, 0x65,
-	0x6e, 0x74, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12,
-	0x2d, 0x0a, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x6f, 0x67, 0x18, 0x01, 0x20, 0x01,
-	0x28, 0x0b, 0x32, 0x0e, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x4c,
-	0x6f, 0x67, 0x48, 0x00, 0x52, 0x08, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x12, 0x33,
-	0x0a, 0x0b, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x18, 0x02, 0x20,
-	0x01, 0x28, 0x0b, 0x32, 0x10, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74,
-	0x45, 0x76, 0x65, 0x6e, 0x74, 0x48, 0x00, 0x52, 0x0a, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76,
-	0x65, 0x6e, 0x74, 0x12, 0x2c, 0x0a, 0x07, 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x65, 0x73, 0x18, 0x03,
-	0x20, 0x01, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x52, 0x75, 0x6e, 0x52,
-	0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x48, 0x00, 0x52, 0x06, 0x72, 0x75, 0x6e, 0x52, 0x65,
-	0x73, 0x12, 0x4f, 0x0a, 0x12, 0x73, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61,
-	0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1d, 0x2e,
-	0x63, 0x76, 0x6d, 0x73, 0x2e, 0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61,
-	0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x48, 0x00, 0x52, 0x12,
-	0x73, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52,
-	0x65, 0x73, 0x12, 0x3b, 0x0a, 0x0d, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65,
-	0x52, 0x65, 0x73, 0x18, 0x05, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x13, 0x2e, 0x63, 0x76, 0x6d, 0x73,
-	0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x48, 0x00,
-	0x52, 0x0d, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x12,
-	0x51, 0x0a, 0x15, 0x76, 0x54, 0x50, 0x4d, 0x61, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69,
-	0x6f, 0x6e, 0x52, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x06, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x19,
-	0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f,
-	0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x48, 0x00, 0x52, 0x15, 0x76, 0x54, 0x50,
-	0x4d, 0x61, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x70, 0x6f,
-	0x72, 0x74, 0x12, 0x58, 0x0a, 0x16, 0x61, 0x7a, 0x75, 0x72, 0x65, 0x41, 0x74, 0x74, 0x65, 0x73,
-	0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x18, 0x07, 0x20, 0x01,
-	0x28, 0x0b, 0x32, 0x1e, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x61, 0x7a, 0x75, 0x72, 0x65, 0x41,
-	0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e,
-	0x73, 0x65, 0x48, 0x00, 0x52, 0x16, 0x61, 0x7a, 0x75, 0x72, 0x65, 0x41, 0x74, 0x74, 0x65, 0x73,
-	0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x42, 0x09, 0x0a, 0x07,
-	0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0xca, 0x02, 0x0a, 0x13, 0x53, 0x65, 0x72, 0x76,
-	0x65, 0x72, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12,
-	0x38, 0x0a, 0x0c, 0x72, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x43, 0x68, 0x75, 0x6e, 0x6b, 0x73, 0x18,
-	0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x12, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x52, 0x75, 0x6e,
-	0x52, 0x65, 0x71, 0x43, 0x68, 0x75, 0x6e, 0x6b, 0x73, 0x48, 0x00, 0x52, 0x0c, 0x72, 0x75, 0x6e,
-	0x52, 0x65, 0x71, 0x43, 0x68, 0x75, 0x6e, 0x6b, 0x73, 0x12, 0x31, 0x0a, 0x06, 0x72, 0x75, 0x6e,
-	0x52, 0x65, 0x71, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x17, 0x2e, 0x63, 0x76, 0x6d, 0x73,
-	0x2e, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x75, 0x6e, 0x52,
-	0x65, 0x71, 0x48, 0x00, 0x52, 0x06, 0x72, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x12, 0x41, 0x0a, 0x0f,
-	0x73, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x18,
-	0x03, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x15, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x53, 0x74, 0x6f,
-	0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x48, 0x00, 0x52, 0x0f,
-	0x73, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12,
-	0x3b, 0x0a, 0x0d, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71,
-	0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x13, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x41, 0x67,
-	0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x48, 0x00, 0x52, 0x0d, 0x61,
-	0x67, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x12, 0x3b, 0x0a, 0x0d,
-	0x64, 0x69, 0x73, 0x63, 0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x52, 0x65, 0x71, 0x18, 0x05, 0x20,
-	0x01, 0x28, 0x0b, 0x32, 0x13, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x44, 0x69, 0x73, 0x63, 0x6f,
-	0x6e, 0x6e, 0x65, 0x63, 0x74, 0x52, 0x65, 0x71, 0x48, 0x00, 0x52, 0x0d, 0x64, 0x69, 0x73, 0x63,
-	0x6f, 0x6e, 0x6e, 0x65, 0x63, 0x74, 0x52, 0x65, 0x71, 0x42, 0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73,
-	0x73, 0x61, 0x67, 0x65, 0x22, 0x1f, 0x0a, 0x0d, 0x44, 0x69, 0x73, 0x63, 0x6f, 0x6e, 0x6e, 0x65,
-	0x63, 0x74, 0x52, 0x65, 0x71, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x02, 0x69, 0x64, 0x22, 0x4b, 0x0a, 0x0c, 0x52, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x43,
-	0x68, 0x75, 0x6e, 0x6b, 0x73, 0x12, 0x12, 0x0a, 0x04, 0x64, 0x61, 0x74, 0x61, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x0c, 0x52, 0x04, 0x64, 0x61, 0x74, 0x61, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18,
-	0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x17, 0x0a, 0x07, 0x69, 0x73, 0x5f,
-	0x6c, 0x61, 0x73, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x06, 0x69, 0x73, 0x4c, 0x61,
-	0x73, 0x74, 0x22, 0xaa, 0x02, 0x0a, 0x11, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69,
-	0x6f, 0x6e, 0x52, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65,
-	0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x12, 0x20, 0x0a, 0x0b,
-	0x64, 0x65, 0x73, 0x63, 0x72, 0x69, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x18, 0x03, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x0b, 0x64, 0x65, 0x73, 0x63, 0x72, 0x69, 0x70, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x29,
-	0x0a, 0x08, 0x64, 0x61, 0x74, 0x61, 0x73, 0x65, 0x74, 0x73, 0x18, 0x04, 0x20, 0x03, 0x28, 0x0b,
-	0x32, 0x0d, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x44, 0x61, 0x74, 0x61, 0x73, 0x65, 0x74, 0x52,
-	0x08, 0x64, 0x61, 0x74, 0x61, 0x73, 0x65, 0x74, 0x73, 0x12, 0x2d, 0x0a, 0x09, 0x61, 0x6c, 0x67,
-	0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x18, 0x05, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x0f, 0x2e, 0x63,
-	0x76, 0x6d, 0x73, 0x2e, 0x41, 0x6c, 0x67, 0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x52, 0x09, 0x61,
-	0x6c, 0x67, 0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x12, 0x3f, 0x0a, 0x10, 0x72, 0x65, 0x73, 0x75,
-	0x6c, 0x74, 0x5f, 0x63, 0x6f, 0x6e, 0x73, 0x75, 0x6d, 0x65, 0x72, 0x73, 0x18, 0x06, 0x20, 0x03,
-	0x28, 0x0b, 0x32, 0x14, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74,
-	0x43, 0x6f, 0x6e, 0x73, 0x75, 0x6d, 0x65, 0x72, 0x52, 0x0f, 0x72, 0x65, 0x73, 0x75, 0x6c, 0x74,
-	0x43, 0x6f, 0x6e, 0x73, 0x75, 0x6d, 0x65, 0x72, 0x73, 0x12, 0x34, 0x0a, 0x0c, 0x61, 0x67, 0x65,
-	0x6e, 0x74, 0x5f, 0x63, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x18, 0x07, 0x20, 0x01, 0x28, 0x0b, 0x32,
-	0x11, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x66,
-	0x69, 0x67, 0x52, 0x0b, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x22,
-	0x2a, 0x0a, 0x0e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x43, 0x6f, 0x6e, 0x73, 0x75, 0x6d, 0x65,
-	0x72, 0x12, 0x18, 0x0a, 0x07, 0x75, 0x73, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x18, 0x01, 0x20, 0x01,
-	0x28, 0x0c, 0x52, 0x07, 0x75, 0x73, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x22, 0x53, 0x0a, 0x07, 0x44,
-	0x61, 0x74, 0x61, 0x73, 0x65, 0x74, 0x12, 0x12, 0x0a, 0x04, 0x68, 0x61, 0x73, 0x68, 0x18, 0x01,
-	0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x68, 0x61, 0x73, 0x68, 0x12, 0x18, 0x0a, 0x07, 0x75, 0x73,
-	0x65, 0x72, 0x4b, 0x65, 0x79, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x75, 0x73, 0x65,
-	0x72, 0x4b, 0x65, 0x79, 0x12, 0x1a, 0x0a, 0x08, 0x66, 0x69, 0x6c, 0x65, 0x6e, 0x61, 0x6d, 0x65,
-	0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x66, 0x69, 0x6c, 0x65, 0x6e, 0x61, 0x6d, 0x65,
-	0x22, 0x39, 0x0a, 0x09, 0x41, 0x6c, 0x67, 0x6f, 0x72, 0x69, 0x74, 0x68, 0x6d, 0x12, 0x12, 0x0a,
-	0x04, 0x68, 0x61, 0x73, 0x68, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x04, 0x68, 0x61, 0x73,
-	0x68, 0x12, 0x18, 0x0a, 0x07, 0x75, 0x73, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x18, 0x02, 0x20, 0x01,
-	0x28, 0x0c, 0x52, 0x07, 0x75, 0x73, 0x65, 0x72, 0x4b, 0x65, 0x79, 0x22, 0xe5, 0x01, 0x0a, 0x0b,
-	0x41, 0x67, 0x65, 0x6e, 0x74, 0x43, 0x6f, 0x6e, 0x66, 0x69, 0x67, 0x12, 0x12, 0x0a, 0x04, 0x70,
-	0x6f, 0x72, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x70, 0x6f, 0x72, 0x74, 0x12,
-	0x1b, 0x0a, 0x09, 0x63, 0x65, 0x72, 0x74, 0x5f, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x02, 0x20, 0x01,
-	0x28, 0x09, 0x52, 0x08, 0x63, 0x65, 0x72, 0x74, 0x46, 0x69, 0x6c, 0x65, 0x12, 0x19, 0x0a, 0x08,
-	0x6b, 0x65, 0x79, 0x5f, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x07,
-	0x6b, 0x65, 0x79, 0x46, 0x69, 0x6c, 0x65, 0x12, 0x24, 0x0a, 0x0e, 0x63, 0x6c, 0x69, 0x65, 0x6e,
-	0x74, 0x5f, 0x63, 0x61, 0x5f, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52,
-	0x0c, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x43, 0x61, 0x46, 0x69, 0x6c, 0x65, 0x12, 0x24, 0x0a,
-	0x0e, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x5f, 0x63, 0x61, 0x5f, 0x66, 0x69, 0x6c, 0x65, 0x18,
-	0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0c, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x61, 0x46,
-	0x69, 0x6c, 0x65, 0x12, 0x1b, 0x0a, 0x09, 0x6c, 0x6f, 0x67, 0x5f, 0x6c, 0x65, 0x76, 0x65, 0x6c,
-	0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x6c, 0x6f, 0x67, 0x4c, 0x65, 0x76, 0x65, 0x6c,
-	0x12, 0x21, 0x0a, 0x0c, 0x61, 0x74, 0x74, 0x65, 0x73, 0x74, 0x65, 0x64, 0x5f, 0x74, 0x6c, 0x73,
-	0x18, 0x07, 0x20, 0x01, 0x28, 0x08, 0x52, 0x0b, 0x61, 0x74, 0x74, 0x65, 0x73, 0x74, 0x65, 0x64,
-	0x54, 0x6c, 0x73, 0x22, 0x55, 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, 0x12, 0x2a,
-	0x0a, 0x10, 0x63, 0x65, 0x72, 0x74, 0x53, 0x65, 0x72, 0x69, 0x61, 0x6c, 0x4e, 0x75, 0x6d, 0x62,
-	0x65, 0x72, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x10, 0x63, 0x65, 0x72, 0x74, 0x53, 0x65,
-	0x72, 0x69, 0x61, 0x6c, 0x4e, 0x75, 0x6d, 0x62, 0x65, 0x72, 0x22, 0x5a, 0x0a, 0x18, 0x61, 0x7a,
-	0x75, 0x72, 0x65, 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, 0x12, 0x2a, 0x0a, 0x10, 0x63, 0x65,
-	0x72, 0x74, 0x53, 0x65, 0x72, 0x69, 0x61, 0x6c, 0x4e, 0x75, 0x6d, 0x62, 0x65, 0x72, 0x18, 0x02,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x10, 0x63, 0x65, 0x72, 0x74, 0x53, 0x65, 0x72, 0x69, 0x61, 0x6c,
-	0x4e, 0x75, 0x6d, 0x62, 0x65, 0x72, 0x32, 0x50, 0x0a, 0x07, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63,
-	0x65, 0x12, 0x45, 0x0a, 0x07, 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x12, 0x19, 0x2e, 0x63,
-	0x76, 0x6d, 0x73, 0x2e, 0x43, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d,
-	0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x1a, 0x19, 0x2e, 0x63, 0x76, 0x6d, 0x73, 0x2e, 0x53,
-	0x65, 0x72, 0x76, 0x65, 0x72, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73, 0x73, 0x61,
-	0x67, 0x65, 0x22, 0x00, 0x28, 0x01, 0x30, 0x01, 0x42, 0x08, 0x5a, 0x06, 0x2e, 0x2f, 0x63, 0x76,
-	0x6d, 0x73, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
-})
+const file_agent_cvms_cvms_proto_rawDesc = "" +
+	"\n" +
+	"\x15agent/cvms/cvms.proto\x12\x04cvms\x1a\x1fgoogle/protobuf/timestamp.proto\"\x1f\n" +
+	"\rAgentStateReq\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id\"5\n" +
+	"\rAgentStateRes\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id\x12\x14\n" +
+	"\x05state\x18\x02 \x01(\tR\x05state\"8\n" +
+	"\x0fStopComputation\x12%\n" +
+	"\x0ecomputation_id\x18\x01 \x01(\tR\rcomputationId\"Z\n" +
+	"\x17StopComputationResponse\x12%\n" +
+	"\x0ecomputation_id\x18\x01 \x01(\tR\rcomputationId\x12\x18\n" +
+	"\amessage\x18\x02 \x01(\tR\amessage\"J\n" +
+	"\vRunResponse\x12%\n" +
+	"\x0ecomputation_id\x18\x01 \x01(\tR\rcomputationId\x12\x14\n" +
+	"\x05error\x18\x02 \x01(\tR\x05error\"\xde\x01\n" +
+	"\n" +
+	"AgentEvent\x12\x1d\n" +
+	"\n" +
+	"event_type\x18\x01 \x01(\tR\teventType\x128\n" +
+	"\ttimestamp\x18\x02 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12%\n" +
+	"\x0ecomputation_id\x18\x03 \x01(\tR\rcomputationId\x12\x18\n" +
+	"\adetails\x18\x04 \x01(\fR\adetails\x12\x1e\n" +
+	"\n" +
+	"originator\x18\x05 \x01(\tR\n" +
+	"originator\x12\x16\n" +
+	"\x06status\x18\x06 \x01(\tR\x06status\"\x9b\x01\n" +
+	"\bAgentLog\x12\x18\n" +
+	"\amessage\x18\x01 \x01(\tR\amessage\x12%\n" +
+	"\x0ecomputation_id\x18\x02 \x01(\tR\rcomputationId\x12\x14\n" +
+	"\x05level\x18\x03 \x01(\tR\x05level\x128\n" +
+	"\ttimestamp\x18\x04 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\"\xe8\x03\n" +
+	"\x13ClientStreamMessage\x12-\n" +
+	"\tagent_log\x18\x01 \x01(\v2\x0e.cvms.AgentLogH\x00R\bagentLog\x123\n" +
+	"\vagent_event\x18\x02 \x01(\v2\x10.cvms.AgentEventH\x00R\n" +
+	"agentEvent\x12,\n" +
+	"\arun_res\x18\x03 \x01(\v2\x11.cvms.RunResponseH\x00R\x06runRes\x12O\n" +
+	"\x12stopComputationRes\x18\x04 \x01(\v2\x1d.cvms.StopComputationResponseH\x00R\x12stopComputationRes\x12;\n" +
+	"\ragentStateRes\x18\x05 \x01(\v2\x13.cvms.AgentStateResH\x00R\ragentStateRes\x12Q\n" +
+	"\x15vTPMattestationReport\x18\x06 \x01(\v2\x19.cvms.AttestationResponseH\x00R\x15vTPMattestationReport\x12S\n" +
+	"\x15azureAttestationToken\x18\a \x01(\v2\x1b.cvms.azureAttestationTokenH\x00R\x15azureAttestationTokenB\t\n" +
+	"\amessage\"\xca\x02\n" +
+	"\x13ServerStreamMessage\x128\n" +
+	"\frunReqChunks\x18\x01 \x01(\v2\x12.cvms.RunReqChunksH\x00R\frunReqChunks\x121\n" +
+	"\x06runReq\x18\x02 \x01(\v2\x17.cvms.ComputationRunReqH\x00R\x06runReq\x12A\n" +
+	"\x0fstopComputation\x18\x03 \x01(\v2\x15.cvms.StopComputationH\x00R\x0fstopComputation\x12;\n" +
+	"\ragentStateReq\x18\x04 \x01(\v2\x13.cvms.AgentStateReqH\x00R\ragentStateReq\x12;\n" +
+	"\rdisconnectReq\x18\x05 \x01(\v2\x13.cvms.DisconnectReqH\x00R\rdisconnectReqB\t\n" +
+	"\amessage\"\x1f\n" +
+	"\rDisconnectReq\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id\"K\n" +
+	"\fRunReqChunks\x12\x12\n" +
+	"\x04data\x18\x01 \x01(\fR\x04data\x12\x0e\n" +
+	"\x02id\x18\x02 \x01(\tR\x02id\x12\x17\n" +
+	"\ais_last\x18\x03 \x01(\bR\x06isLast\"\xaa\x02\n" +
+	"\x11ComputationRunReq\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id\x12\x12\n" +
+	"\x04name\x18\x02 \x01(\tR\x04name\x12 \n" +
+	"\vdescription\x18\x03 \x01(\tR\vdescription\x12)\n" +
+	"\bdatasets\x18\x04 \x03(\v2\r.cvms.DatasetR\bdatasets\x12-\n" +
+	"\talgorithm\x18\x05 \x01(\v2\x0f.cvms.AlgorithmR\talgorithm\x12?\n" +
+	"\x10result_consumers\x18\x06 \x03(\v2\x14.cvms.ResultConsumerR\x0fresultConsumers\x124\n" +
+	"\fagent_config\x18\a \x01(\v2\x11.cvms.AgentConfigR\vagentConfig\"*\n" +
+	"\x0eResultConsumer\x12\x18\n" +
+	"\auserKey\x18\x01 \x01(\fR\auserKey\"S\n" +
+	"\aDataset\x12\x12\n" +
+	"\x04hash\x18\x01 \x01(\fR\x04hash\x12\x18\n" +
+	"\auserKey\x18\x02 \x01(\fR\auserKey\x12\x1a\n" +
+	"\bfilename\x18\x03 \x01(\tR\bfilename\"9\n" +
+	"\tAlgorithm\x12\x12\n" +
+	"\x04hash\x18\x01 \x01(\fR\x04hash\x12\x18\n" +
+	"\auserKey\x18\x02 \x01(\fR\auserKey\"\xe5\x01\n" +
+	"\vAgentConfig\x12\x12\n" +
+	"\x04port\x18\x01 \x01(\tR\x04port\x12\x1b\n" +
+	"\tcert_file\x18\x02 \x01(\tR\bcertFile\x12\x19\n" +
+	"\bkey_file\x18\x03 \x01(\tR\akeyFile\x12$\n" +
+	"\x0eclient_ca_file\x18\x04 \x01(\tR\fclientCaFile\x12$\n" +
+	"\x0eserver_ca_file\x18\x05 \x01(\tR\fserverCaFile\x12\x1b\n" +
+	"\tlog_level\x18\x06 \x01(\tR\blogLevel\x12!\n" +
+	"\fattested_tls\x18\a \x01(\bR\vattestedTls\"U\n" +
+	"\x13AttestationResponse\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file\x12*\n" +
+	"\x10certSerialNumber\x18\x02 \x01(\tR\x10certSerialNumber\"W\n" +
+	"\x15azureAttestationToken\x12\x12\n" +
+	"\x04file\x18\x01 \x01(\fR\x04file\x12*\n" +
+	"\x10certSerialNumber\x18\x02 \x01(\tR\x10certSerialNumber2P\n" +
+	"\aService\x12E\n" +
+	"\aProcess\x12\x19.cvms.ClientStreamMessage\x1a\x19.cvms.ServerStreamMessage\"\x00(\x010\x01B\bZ\x06./cvmsb\x06proto3"
 
 var (
 	file_agent_cvms_cvms_proto_rawDescOnce sync.Once
@@ -1449,25 +1366,25 @@ func file_agent_cvms_cvms_proto_rawDescGZIP() []byte {
 
 var file_agent_cvms_cvms_proto_msgTypes = make([]protoimpl.MessageInfo, 18)
 var file_agent_cvms_cvms_proto_goTypes = []any{
-	(*AgentStateReq)(nil),            // 0: cvms.AgentStateReq
-	(*AgentStateRes)(nil),            // 1: cvms.AgentStateRes
-	(*StopComputation)(nil),          // 2: cvms.StopComputation
-	(*StopComputationResponse)(nil),  // 3: cvms.StopComputationResponse
-	(*RunResponse)(nil),              // 4: cvms.RunResponse
-	(*AgentEvent)(nil),               // 5: cvms.AgentEvent
-	(*AgentLog)(nil),                 // 6: cvms.AgentLog
-	(*ClientStreamMessage)(nil),      // 7: cvms.ClientStreamMessage
-	(*ServerStreamMessage)(nil),      // 8: cvms.ServerStreamMessage
-	(*DisconnectReq)(nil),            // 9: cvms.DisconnectReq
-	(*RunReqChunks)(nil),             // 10: cvms.RunReqChunks
-	(*ComputationRunReq)(nil),        // 11: cvms.ComputationRunReq
-	(*ResultConsumer)(nil),           // 12: cvms.ResultConsumer
-	(*Dataset)(nil),                  // 13: cvms.Dataset
-	(*Algorithm)(nil),                // 14: cvms.Algorithm
-	(*AgentConfig)(nil),              // 15: cvms.AgentConfig
-	(*AttestationResponse)(nil),      // 16: cvms.AttestationResponse
-	(*AzureAttestationResponse)(nil), // 17: cvms.azureAttestationResponse
-	(*timestamppb.Timestamp)(nil),    // 18: google.protobuf.Timestamp
+	(*AgentStateReq)(nil),           // 0: cvms.AgentStateReq
+	(*AgentStateRes)(nil),           // 1: cvms.AgentStateRes
+	(*StopComputation)(nil),         // 2: cvms.StopComputation
+	(*StopComputationResponse)(nil), // 3: cvms.StopComputationResponse
+	(*RunResponse)(nil),             // 4: cvms.RunResponse
+	(*AgentEvent)(nil),              // 5: cvms.AgentEvent
+	(*AgentLog)(nil),                // 6: cvms.AgentLog
+	(*ClientStreamMessage)(nil),     // 7: cvms.ClientStreamMessage
+	(*ServerStreamMessage)(nil),     // 8: cvms.ServerStreamMessage
+	(*DisconnectReq)(nil),           // 9: cvms.DisconnectReq
+	(*RunReqChunks)(nil),            // 10: cvms.RunReqChunks
+	(*ComputationRunReq)(nil),       // 11: cvms.ComputationRunReq
+	(*ResultConsumer)(nil),          // 12: cvms.ResultConsumer
+	(*Dataset)(nil),                 // 13: cvms.Dataset
+	(*Algorithm)(nil),               // 14: cvms.Algorithm
+	(*AgentConfig)(nil),             // 15: cvms.AgentConfig
+	(*AttestationResponse)(nil),     // 16: cvms.AttestationResponse
+	(*AzureAttestationToken)(nil),   // 17: cvms.azureAttestationToken
+	(*timestamppb.Timestamp)(nil),   // 18: google.protobuf.Timestamp
 }
 var file_agent_cvms_cvms_proto_depIdxs = []int32{
 	18, // 0: cvms.AgentEvent.timestamp:type_name -> google.protobuf.Timestamp
@@ -1478,7 +1395,7 @@ var file_agent_cvms_cvms_proto_depIdxs = []int32{
 	3,  // 5: cvms.ClientStreamMessage.stopComputationRes:type_name -> cvms.StopComputationResponse
 	1,  // 6: cvms.ClientStreamMessage.agentStateRes:type_name -> cvms.AgentStateRes
 	16, // 7: cvms.ClientStreamMessage.vTPMattestationReport:type_name -> cvms.AttestationResponse
-	17, // 8: cvms.ClientStreamMessage.azureAttestationResult:type_name -> cvms.azureAttestationResponse
+	17, // 8: cvms.ClientStreamMessage.azureAttestationToken:type_name -> cvms.azureAttestationToken
 	10, // 9: cvms.ServerStreamMessage.runReqChunks:type_name -> cvms.RunReqChunks
 	11, // 10: cvms.ServerStreamMessage.runReq:type_name -> cvms.ComputationRunReq
 	2,  // 11: cvms.ServerStreamMessage.stopComputation:type_name -> cvms.StopComputation
@@ -1509,7 +1426,7 @@ func file_agent_cvms_cvms_proto_init() {
 		(*ClientStreamMessage_StopComputationRes)(nil),
 		(*ClientStreamMessage_AgentStateRes)(nil),
 		(*ClientStreamMessage_VTPMattestationReport)(nil),
-		(*ClientStreamMessage_AzureAttestationResult)(nil),
+		(*ClientStreamMessage_AzureAttestationToken)(nil),
 	}
 	file_agent_cvms_cvms_proto_msgTypes[8].OneofWrappers = []any{
 		(*ServerStreamMessage_RunReqChunks)(nil),

agent/cvms/cvms.proto

@@ -60,7 +60,7 @@ message ClientStreamMessage {
     StopComputationResponse stopComputationRes = 4;
     AgentStateRes agentStateRes = 5;
     AttestationResponse vTPMattestationReport = 6;
-    azureAttestationResponse azureAttestationResult = 7;
+    azureAttestationToken azureAttestationToken = 7;
   }
 }
 
@@ -124,7 +124,7 @@ message AttestationResponse {
   string certSerialNumber = 2;
 }
 
-message azureAttestationResponse {
+message azureAttestationToken {
   bytes file = 1;
   string certSerialNumber = 2;
 }

agent/cvms/server/cvm.go

@@ -11,8 +11,9 @@ import (
 	"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/atls"
+	"github.com/ultravioletrs/cocos/pkg/server"
+	grpcserver "github.com/ultravioletrs/cocos/pkg/server/grpc"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
 )
@@ -28,21 +29,19 @@ type AgentServer interface {
 }
 
 type agentServer struct {
-	gs     server.Server
-	logger *slog.Logger
-	svc    agent.Service
-	host   string
-	caUrl  string
-	cvmId  string
+	gs           server.Server
+	logger       *slog.Logger
+	svc          agent.Service
+	host         string
+	certProvider atls.CertificateProvider
 }
 
-func NewServer(logger *slog.Logger, svc agent.Service, host string, caUrl string, cvmId string) AgentServer {
+func NewServer(logger *slog.Logger, svc agent.Service, host string, certProvider atls.CertificateProvider) AgentServer {
 	return &agentServer{
-		logger: logger,
-		svc:    svc,
-		host:   host,
-		caUrl:  caUrl,
-		cvmId:  cvmId,
+		logger:       logger,
+		svc:          svc,
+		host:         host,
+		certProvider: certProvider,
 	}
 }
 
@@ -53,7 +52,7 @@ func (as *agentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error
 
 	agentGrpcServerConfig := server.AgentConfig{
 		ServerConfig: server.ServerConfig{
-			BaseConfig: server.BaseConfig{
+			Config: server.Config{
 				Host:         as.host,
 				Port:         cfg.Port,
 				CertFile:     cfg.CertFile,
@@ -78,7 +77,7 @@ func (as *agentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error
 
 	ctx, cancel := context.WithCancel(context.Background())
 
-	as.gs = grpcserver.New(ctx, cancel, svcName, agentGrpcServerConfig, registerAgentServiceServer, as.logger, authSvc, as.caUrl, as.cvmId)
+	as.gs = grpcserver.New(ctx, cancel, svcName, agentGrpcServerConfig, registerAgentServiceServer, as.logger, authSvc, as.certProvider)
 
 	go func() {
 		err := as.gs.Start()

agent/cvms/server/cvm_test.go

@@ -18,12 +18,10 @@ import (
 	"github.com/ultravioletrs/cocos/agent/mocks"
 )
 
-func setupTest(t *testing.T) (*slog.Logger, *mocks.Service, string, string, string, []byte) {
+func setupTest(t *testing.T) (*slog.Logger, *mocks.Service, string, []byte) {
 	logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelError}))
 	mockSvc := new(mocks.Service)
 	host := "localhost"
-	caUrl := "https://ca.example.com"
-	cvmId := "test-cvm-id"
 
 	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 	assert.NoError(t, err, "Failed to generate ECDSA key")
@@ -31,19 +29,17 @@ func setupTest(t *testing.T) (*slog.Logger, *mocks.Service, string, string, stri
 	pubkey, err := x509.MarshalPKIXPublicKey(privateKey.Public())
 	assert.NoError(t, err, "Failed to marshal public key")
 
-	return logger, mockSvc, host, caUrl, cvmId, pubkey
+	return logger, mockSvc, host, pubkey
 }
 
 func TestNewServer(t *testing.T) {
-	logger, svc, host, caUrl, cvmId, _ := setupTest(t)
+	logger, svc, host, _ := setupTest(t)
 
 	tests := []struct {
 		name     string
 		logger   *slog.Logger
 		svc      agent.Service
 		host     string
-		caUrl    string
-		cvmId    string
 		expected AgentServer
 	}{
 		{
@@ -51,38 +47,30 @@ func TestNewServer(t *testing.T) {
 			logger: logger,
 			svc:    svc,
 			host:   host,
-			caUrl:  caUrl,
-			cvmId:  cvmId,
 		},
 		{
 			name:   "server with empty host",
 			logger: logger,
 			svc:    svc,
 			host:   "",
-			caUrl:  caUrl,
-			cvmId:  cvmId,
 		},
 		{
 			name:   "server with empty caUrl",
 			logger: logger,
 			svc:    svc,
 			host:   host,
-			caUrl:  "",
-			cvmId:  cvmId,
 		},
 		{
 			name:   "server with empty cvmId",
 			logger: logger,
 			svc:    svc,
 			host:   host,
-			caUrl:  caUrl,
-			cvmId:  "",
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			server := NewServer(tt.logger, tt.svc, tt.host, tt.caUrl, tt.cvmId)
+			server := NewServer(tt.logger, tt.svc, tt.host, nil)
 
 			assert.NotNil(t, server)
 
@@ -91,14 +79,12 @@ func TestNewServer(t *testing.T) {
 			assert.Equal(t, tt.logger, agentSrv.logger)
 			assert.Equal(t, tt.svc, agentSrv.svc)
 			assert.Equal(t, tt.host, agentSrv.host)
-			assert.Equal(t, tt.caUrl, agentSrv.caUrl)
-			assert.Equal(t, tt.cvmId, agentSrv.cvmId)
 		})
 	}
 }
 
 func TestAgentServer_Start(t *testing.T) {
-	logger, svc, host, caUrl, cvmId, pubKey := setupTest(t)
+	logger, svc, host, pubKey := setupTest(t)
 
 	tests := []struct {
 		name          string
@@ -211,7 +197,7 @@ func TestAgentServer_Start(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			tt.setupMocks(svc)
 
-			server := NewServer(logger, svc, host, caUrl, cvmId)
+			server := NewServer(logger, svc, host, nil)
 
 			err := server.Start(tt.cfg, tt.cmp)
 
@@ -238,7 +224,7 @@ func TestAgentServer_Start(t *testing.T) {
 }
 
 func TestAgentServer_Stop(t *testing.T) {
-	logger, svc, host, caUrl, cvmId, pubKey := setupTest(t)
+	logger, svc, host, pubKey := setupTest(t)
 
 	tests := []struct {
 		name          string
@@ -287,7 +273,7 @@ func TestAgentServer_Stop(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			server := NewServer(logger, svc, host, caUrl, cvmId)
+			server := NewServer(logger, svc, host, nil)
 
 			err := tt.setupServer(server)
 			if err != nil {
@@ -314,8 +300,8 @@ func TestAgentServer_Stop(t *testing.T) {
 }
 
 func TestAgentServer_StopMultipleTimes(t *testing.T) {
-	logger, svc, host, caUrl, cvmId, pubKey := setupTest(t)
-	server := NewServer(logger, svc, host, caUrl, cvmId)
+	logger, svc, host, pubKey := setupTest(t)
+	server := NewServer(logger, svc, host, nil)
 
 	// Start the server
 	cfg := agent.AgentConfig{Port: "7005"}
@@ -358,8 +344,8 @@ func TestAgentServer_StopMultipleTimes(t *testing.T) {
 }
 
 func TestAgentServer_StartAfterStop(t *testing.T) {
-	logger, svc, host, caUrl, cvmId, pubKey := setupTest(t)
-	server := NewServer(logger, svc, host, caUrl, cvmId)
+	logger, svc, host, pubKey := setupTest(t)
+	server := NewServer(logger, svc, host, nil)
 
 	cfg := agent.AgentConfig{Port: "7006"}
 	cmp := agent.Computation{
@@ -425,7 +411,7 @@ func TestAgentServer_StartAfterStop(t *testing.T) {
 }
 
 func TestAgentServer_ConfigValidation(t *testing.T) {
-	logger, svc, host, caUrl, cvmId, pubKey := setupTest(t)
+	logger, svc, host, pubKey := setupTest(t)
 
 	tests := []struct {
 		name   string
@@ -512,7 +498,7 @@ func TestAgentServer_ConfigValidation(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			server := NewServer(logger, svc, host, caUrl, cvmId)
+			server := NewServer(logger, svc, host, nil)
 
 			err := server.Start(tt.config, tt.cmp)
 

agent/cvms/server/mocks/agentserver.go

@@ -1,15 +1,31 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
 	mock "github.com/stretchr/testify/mock"
-	agent "github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/agent"
 )
 
+// 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
+}
+
 // AgentServer is an autogenerated mock type for the AgentServer type
 type AgentServer struct {
 	mock.Mock
@@ -23,21 +39,20 @@ 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)
+// Start provides a mock function for the type AgentServer
+func (_mock *AgentServer) Start(cfg agent.AgentConfig, cmp agent.Computation) error {
+	ret := _mock.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)
+	if returnFunc, ok := ret.Get(0).(func(agent.AgentConfig, agent.Computation) error); ok {
+		r0 = returnFunc(cfg, cmp)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -55,36 +70,46 @@ func (_e *AgentServer_Expecter) Start(cfg interface{}, cmp interface{}) *AgentSe
 
 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))
+		var arg0 agent.AgentConfig
+		if args[0] != nil {
+			arg0 = args[0].(agent.AgentConfig)
+		}
+		var arg1 agent.Computation
+		if args[1] != nil {
+			arg1 = args[1].(agent.Computation)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *AgentServer_Start_Call) Return(_a0 error) *AgentServer_Start_Call {
-	_c.Call.Return(_a0)
+func (_c *AgentServer_Start_Call) Return(err error) *AgentServer_Start_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *AgentServer_Start_Call) RunAndReturn(run func(agent.AgentConfig, agent.Computation) error) *AgentServer_Start_Call {
+func (_c *AgentServer_Start_Call) RunAndReturn(run func(cfg agent.AgentConfig, cmp agent.Computation) error) *AgentServer_Start_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Stop provides a mock function with no fields
-func (_m *AgentServer) Stop() error {
-	ret := _m.Called()
+// Stop provides a mock function for the type AgentServer
+func (_mock *AgentServer) Stop() error {
+	ret := _mock.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()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -105,8 +130,8 @@ func (_c *AgentServer_Stop_Call) Run(run func()) *AgentServer_Stop_Call {
 	return _c
 }
 
-func (_c *AgentServer_Stop_Call) Return(_a0 error) *AgentServer_Stop_Call {
-	_c.Call.Return(_a0)
+func (_c *AgentServer_Stop_Call) Return(err error) *AgentServer_Stop_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -114,17 +139,3 @@ func (_c *AgentServer_Stop_Call) RunAndReturn(run func() error) *AgentServer_Sto
 	_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.pb.go

@@ -3,7 +3,7 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.36.5
+// 	protoc-gen-go v1.36.8
 // 	protoc        v5.29.0
 // source: agent/events/events.proto
 
@@ -261,46 +261,32 @@ func (*EventsLogs_AgentEvent) isEventsLogs_Message() {}
 
 var File_agent_events_events_proto protoreflect.FileDescriptor
 
-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,
-	0x6f, 0x62, 0x75, 0x66, 0x2f, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x2e, 0x70,
-	0x72, 0x6f, 0x74, 0x6f, 0x22, 0xde, 0x01, 0x0a, 0x0a, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76,
-	0x65, 0x6e, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x5f, 0x74, 0x79, 0x70,
-	0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x54, 0x79,
-	0x70, 0x65, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18,
-	0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70,
-	0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d,
-	0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x12, 0x25, 0x0a, 0x0e,
-	0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x03,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f,
-	0x6e, 0x49, 0x64, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, 0x73, 0x18, 0x04,
-	0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, 0x73, 0x12, 0x1e, 0x0a,
-	0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x18, 0x05, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x16, 0x0a,
-	0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x73,
-	0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x9b, 0x01, 0x0a, 0x08, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x4c,
-	0x6f, 0x67, 0x12, 0x18, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x25, 0x0a, 0x0e,
-	0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x02,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f,
-	0x6e, 0x49, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x18, 0x03, 0x20, 0x01,
-	0x28, 0x09, 0x52, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d,
-	0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67,
-	0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54,
-	0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74,
-	0x61, 0x6d, 0x70, 0x22, 0x7f, 0x0a, 0x0a, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x4c, 0x6f, 0x67,
-	0x73, 0x12, 0x2f, 0x0a, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x6f, 0x67, 0x18, 0x01,
-	0x20, 0x01, 0x28, 0x0b, 0x32, 0x10, 0x2e, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x2e, 0x41, 0x67,
-	0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x48, 0x00, 0x52, 0x08, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x4c,
-	0x6f, 0x67, 0x12, 0x35, 0x0a, 0x0b, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x65, 0x76, 0x65, 0x6e,
-	0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x12, 0x2e, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73,
-	0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x48, 0x00, 0x52, 0x0a, 0x61,
-	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,
-})
+const file_agent_events_events_proto_rawDesc = "" +
+	"\n" +
+	"\x19agent/events/events.proto\x12\x06events\x1a\x1fgoogle/protobuf/timestamp.proto\"\xde\x01\n" +
+	"\n" +
+	"AgentEvent\x12\x1d\n" +
+	"\n" +
+	"event_type\x18\x01 \x01(\tR\teventType\x128\n" +
+	"\ttimestamp\x18\x02 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12%\n" +
+	"\x0ecomputation_id\x18\x03 \x01(\tR\rcomputationId\x12\x18\n" +
+	"\adetails\x18\x04 \x01(\fR\adetails\x12\x1e\n" +
+	"\n" +
+	"originator\x18\x05 \x01(\tR\n" +
+	"originator\x12\x16\n" +
+	"\x06status\x18\x06 \x01(\tR\x06status\"\x9b\x01\n" +
+	"\bAgentLog\x12\x18\n" +
+	"\amessage\x18\x01 \x01(\tR\amessage\x12%\n" +
+	"\x0ecomputation_id\x18\x02 \x01(\tR\rcomputationId\x12\x14\n" +
+	"\x05level\x18\x03 \x01(\tR\x05level\x128\n" +
+	"\ttimestamp\x18\x04 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\"\x7f\n" +
+	"\n" +
+	"EventsLogs\x12/\n" +
+	"\tagent_log\x18\x01 \x01(\v2\x10.events.AgentLogH\x00R\bagentLog\x125\n" +
+	"\vagent_event\x18\x02 \x01(\v2\x12.events.AgentEventH\x00R\n" +
+	"agentEventB\t\n" +
+	"\amessageB\n" +
+	"Z\b./eventsb\x06proto3"
 
 var (
 	file_agent_events_events_proto_rawDescOnce sync.Once

agent/events/mocks/service.go

@@ -1,16 +1,32 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
-	json "encoding/json"
+	"encoding/json"
 
 	mock "github.com/stretchr/testify/mock"
 )
 
+// 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 {
+	mock.TestingT
+	Cleanup(func())
+}) *Service {
+	mock := &Service{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // Service is an autogenerated mock type for the Service type
 type Service struct {
 	mock.Mock
@@ -24,9 +40,10 @@ func (_m *Service) EXPECT() *Service_Expecter {
 	return &Service_Expecter{mock: &_m.Mock}
 }
 
-// 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)
+// SendEvent provides a mock function for the type Service
+func (_mock *Service) SendEvent(cmpID string, event string, status string, details json.RawMessage) {
+	_mock.Called(cmpID, event, status, details)
+	return
 }
 
 // Service_SendEvent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendEvent'
@@ -45,7 +62,28 @@ func (_e *Service_Expecter) SendEvent(cmpID interface{}, event interface{}, stat
 
 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].(string), args[3].(json.RawMessage))
+		var arg0 string
+		if args[0] != nil {
+			arg0 = args[0].(string)
+		}
+		var arg1 string
+		if args[1] != nil {
+			arg1 = args[1].(string)
+		}
+		var arg2 string
+		if args[2] != nil {
+			arg2 = args[2].(string)
+		}
+		var arg3 json.RawMessage
+		if args[3] != nil {
+			arg3 = args[3].(json.RawMessage)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+			arg3,
+		)
 	})
 	return _c
 }
@@ -55,21 +93,7 @@ func (_c *Service_SendEvent_Call) Return() *Service_SendEvent_Call {
 	return _c
 }
 
-func (_c *Service_SendEvent_Call) RunAndReturn(run func(string, string, string, json.RawMessage)) *Service_SendEvent_Call {
+func (_c *Service_SendEvent_Call) RunAndReturn(run func(cmpID string, event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
 	_c.Run(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 {
-	mock.TestingT
-	Cleanup(func())
-}) *Service {
-	mock := &Service{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/mocks/agent_grpc_algo.go

@@ -1,434 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-// Code generated by mockery v2.53.3. DO NOT EDIT.
-
-package mocks
-
-import (
-	context "context"
-
-	agent "github.com/ultravioletrs/cocos/agent"
-
-	metadata "google.golang.org/grpc/metadata"
-
-	mock "github.com/stretchr/testify/mock"
-)
-
-// AgentService_AlgoClient is an autogenerated mock type for the AgentService_AlgoClient type
-type AgentService_AlgoClient[Req interface{}, Res interface{}] struct {
-	mock.Mock
-}
-
-type AgentService_AlgoClient_Expecter[Req interface{}, Res interface{}] struct {
-	mock *mock.Mock
-}
-
-func (_m *AgentService_AlgoClient[Req, Res]) EXPECT() *AgentService_AlgoClient_Expecter[Req, Res] {
-	return &AgentService_AlgoClient_Expecter[Req, Res]{mock: &_m.Mock}
-}
-
-// CloseAndRecv provides a mock function with no fields
-func (_m *AgentService_AlgoClient[Req, Res]) CloseAndRecv() (*agent.AlgoResponse, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for CloseAndRecv")
-	}
-
-	var r0 *agent.AlgoResponse
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (*agent.AlgoResponse, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() *agent.AlgoResponse); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*agent.AlgoResponse)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_AlgoClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
-type AgentService_AlgoClient_CloseAndRecv_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// CloseAndRecv is a helper method to define mock.On call
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) CloseAndRecv() *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
-	return &AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]{Call: _e.mock.On("CloseAndRecv")}
-}
-
-func (_c *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]) Return(_a0 *agent.AlgoResponse, _a1 error) *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res]) RunAndReturn(run func() (*agent.AlgoResponse, error)) *AgentService_AlgoClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// CloseSend provides a mock function with no fields
-func (_m *AgentService_AlgoClient[Req, Res]) CloseSend() error {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for CloseSend")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
-type AgentService_AlgoClient_CloseSend_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// CloseSend is a helper method to define mock.On call
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) CloseSend() *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
-	return &AgentService_AlgoClient_CloseSend_Call[Req, Res]{Call: _e.mock.On("CloseSend")}
-}
-
-func (_c *AgentService_AlgoClient_CloseSend_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_CloseSend_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_CloseSend_Call[Req, Res]) RunAndReturn(run func() error) *AgentService_AlgoClient_CloseSend_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Context provides a mock function with no fields
-func (_m *AgentService_AlgoClient[Req, Res]) Context() context.Context {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Context")
-	}
-
-	var r0 context.Context
-	if rf, ok := ret.Get(0).(func() context.Context); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(context.Context)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
-type AgentService_AlgoClient_Context_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Context is a helper method to define mock.On call
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Context() *AgentService_AlgoClient_Context_Call[Req, Res] {
-	return &AgentService_AlgoClient_Context_Call[Req, Res]{Call: _e.mock.On("Context")}
-}
-
-func (_c *AgentService_AlgoClient_Context_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_Context_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Context_Call[Req, Res]) Return(_a0 context.Context) *AgentService_AlgoClient_Context_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Context_Call[Req, Res]) RunAndReturn(run func() context.Context) *AgentService_AlgoClient_Context_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Header provides a mock function with no fields
-func (_m *AgentService_AlgoClient[Req, Res]) Header() (metadata.MD, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Header")
-	}
-
-	var r0 metadata.MD
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_AlgoClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
-type AgentService_AlgoClient_Header_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Header is a helper method to define mock.On call
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Header() *AgentService_AlgoClient_Header_Call[Req, Res] {
-	return &AgentService_AlgoClient_Header_Call[Req, Res]{Call: _e.mock.On("Header")}
-}
-
-func (_c *AgentService_AlgoClient_Header_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_Header_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Header_Call[Req, Res]) Return(_a0 metadata.MD, _a1 error) *AgentService_AlgoClient_Header_Call[Req, Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Header_Call[Req, Res]) RunAndReturn(run func() (metadata.MD, error)) *AgentService_AlgoClient_Header_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// RecvMsg provides a mock function with given fields: m
-func (_m *AgentService_AlgoClient[Req, Res]) RecvMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for RecvMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
-type AgentService_AlgoClient_RecvMsg_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// RecvMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) RecvMsg(m interface{}) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
-	return &AgentService_AlgoClient_RecvMsg_Call[Req, Res]{Call: _e.mock.On("RecvMsg", m)}
-}
-
-func (_c *AgentService_AlgoClient_RecvMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_RecvMsg_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_RecvMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_AlgoClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Send provides a mock function with given fields: _a0
-func (_m *AgentService_AlgoClient[Req, Res]) Send(_a0 *agent.AlgoRequest) error {
-	ret := _m.Called(_a0)
-
-	if len(ret) == 0 {
-		panic("no return value specified for Send")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(*agent.AlgoRequest) error); ok {
-		r0 = rf(_a0)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
-type AgentService_AlgoClient_Send_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Send is a helper method to define mock.On call
-//   - _a0 *agent.AlgoRequest
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Send(_a0 interface{}) *AgentService_AlgoClient_Send_Call[Req, Res] {
-	return &AgentService_AlgoClient_Send_Call[Req, Res]{Call: _e.mock.On("Send", _a0)}
-}
-
-func (_c *AgentService_AlgoClient_Send_Call[Req, Res]) Run(run func(_a0 *agent.AlgoRequest)) *AgentService_AlgoClient_Send_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(*agent.AlgoRequest))
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Send_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_Send_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Send_Call[Req, Res]) RunAndReturn(run func(*agent.AlgoRequest) error) *AgentService_AlgoClient_Send_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// SendMsg provides a mock function with given fields: m
-func (_m *AgentService_AlgoClient[Req, Res]) SendMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for SendMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
-type AgentService_AlgoClient_SendMsg_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// SendMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) SendMsg(m interface{}) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
-	return &AgentService_AlgoClient_SendMsg_Call[Req, Res]{Call: _e.mock.On("SendMsg", m)}
-}
-
-func (_c *AgentService_AlgoClient_SendMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_SendMsg_Call[Req, Res]) Return(_a0 error) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_SendMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_AlgoClient_SendMsg_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Trailer provides a mock function with no fields
-func (_m *AgentService_AlgoClient[Req, Res]) Trailer() metadata.MD {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Trailer")
-	}
-
-	var r0 metadata.MD
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_AlgoClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
-type AgentService_AlgoClient_Trailer_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Trailer is a helper method to define mock.On call
-func (_e *AgentService_AlgoClient_Expecter[Req, Res]) Trailer() *AgentService_AlgoClient_Trailer_Call[Req, Res] {
-	return &AgentService_AlgoClient_Trailer_Call[Req, Res]{Call: _e.mock.On("Trailer")}
-}
-
-func (_c *AgentService_AlgoClient_Trailer_Call[Req, Res]) Run(run func()) *AgentService_AlgoClient_Trailer_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Trailer_Call[Req, Res]) Return(_a0 metadata.MD) *AgentService_AlgoClient_Trailer_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_AlgoClient_Trailer_Call[Req, Res]) RunAndReturn(run func() metadata.MD) *AgentService_AlgoClient_Trailer_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// NewAgentService_AlgoClient creates a new instance of AgentService_AlgoClient. 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 NewAgentService_AlgoClient[Req interface{}, Res interface{}](t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *AgentService_AlgoClient[Req, Res] {
-	mock := &AgentService_AlgoClient[Req, Res]{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/mocks/agent_grpc_data.go

@@ -1,434 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-// Code generated by mockery v2.53.3. DO NOT EDIT.
-
-package mocks
-
-import (
-	context "context"
-
-	agent "github.com/ultravioletrs/cocos/agent"
-
-	metadata "google.golang.org/grpc/metadata"
-
-	mock "github.com/stretchr/testify/mock"
-)
-
-// AgentService_DataClient is an autogenerated mock type for the AgentService_DataClient type
-type AgentService_DataClient[Req interface{}, Res interface{}] struct {
-	mock.Mock
-}
-
-type AgentService_DataClient_Expecter[Req interface{}, Res interface{}] struct {
-	mock *mock.Mock
-}
-
-func (_m *AgentService_DataClient[Req, Res]) EXPECT() *AgentService_DataClient_Expecter[Req, Res] {
-	return &AgentService_DataClient_Expecter[Req, Res]{mock: &_m.Mock}
-}
-
-// CloseAndRecv provides a mock function with no fields
-func (_m *AgentService_DataClient[Req, Res]) CloseAndRecv() (*agent.DataResponse, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for CloseAndRecv")
-	}
-
-	var r0 *agent.DataResponse
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (*agent.DataResponse, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() *agent.DataResponse); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*agent.DataResponse)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_DataClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
-type AgentService_DataClient_CloseAndRecv_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// CloseAndRecv is a helper method to define mock.On call
-func (_e *AgentService_DataClient_Expecter[Req, Res]) CloseAndRecv() *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
-	return &AgentService_DataClient_CloseAndRecv_Call[Req, Res]{Call: _e.mock.On("CloseAndRecv")}
-}
-
-func (_c *AgentService_DataClient_CloseAndRecv_Call[Req, Res]) Run(run func()) *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_CloseAndRecv_Call[Req, Res]) Return(_a0 *agent.DataResponse, _a1 error) *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_DataClient_CloseAndRecv_Call[Req, Res]) RunAndReturn(run func() (*agent.DataResponse, error)) *AgentService_DataClient_CloseAndRecv_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// CloseSend provides a mock function with no fields
-func (_m *AgentService_DataClient[Req, Res]) CloseSend() error {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for CloseSend")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
-type AgentService_DataClient_CloseSend_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// CloseSend is a helper method to define mock.On call
-func (_e *AgentService_DataClient_Expecter[Req, Res]) CloseSend() *AgentService_DataClient_CloseSend_Call[Req, Res] {
-	return &AgentService_DataClient_CloseSend_Call[Req, Res]{Call: _e.mock.On("CloseSend")}
-}
-
-func (_c *AgentService_DataClient_CloseSend_Call[Req, Res]) Run(run func()) *AgentService_DataClient_CloseSend_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_CloseSend_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_CloseSend_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_CloseSend_Call[Req, Res]) RunAndReturn(run func() error) *AgentService_DataClient_CloseSend_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Context provides a mock function with no fields
-func (_m *AgentService_DataClient[Req, Res]) Context() context.Context {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Context")
-	}
-
-	var r0 context.Context
-	if rf, ok := ret.Get(0).(func() context.Context); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(context.Context)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
-type AgentService_DataClient_Context_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Context is a helper method to define mock.On call
-func (_e *AgentService_DataClient_Expecter[Req, Res]) Context() *AgentService_DataClient_Context_Call[Req, Res] {
-	return &AgentService_DataClient_Context_Call[Req, Res]{Call: _e.mock.On("Context")}
-}
-
-func (_c *AgentService_DataClient_Context_Call[Req, Res]) Run(run func()) *AgentService_DataClient_Context_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_Context_Call[Req, Res]) Return(_a0 context.Context) *AgentService_DataClient_Context_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_Context_Call[Req, Res]) RunAndReturn(run func() context.Context) *AgentService_DataClient_Context_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Header provides a mock function with no fields
-func (_m *AgentService_DataClient[Req, Res]) Header() (metadata.MD, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Header")
-	}
-
-	var r0 metadata.MD
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_DataClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
-type AgentService_DataClient_Header_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Header is a helper method to define mock.On call
-func (_e *AgentService_DataClient_Expecter[Req, Res]) Header() *AgentService_DataClient_Header_Call[Req, Res] {
-	return &AgentService_DataClient_Header_Call[Req, Res]{Call: _e.mock.On("Header")}
-}
-
-func (_c *AgentService_DataClient_Header_Call[Req, Res]) Run(run func()) *AgentService_DataClient_Header_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_Header_Call[Req, Res]) Return(_a0 metadata.MD, _a1 error) *AgentService_DataClient_Header_Call[Req, Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_DataClient_Header_Call[Req, Res]) RunAndReturn(run func() (metadata.MD, error)) *AgentService_DataClient_Header_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// RecvMsg provides a mock function with given fields: m
-func (_m *AgentService_DataClient[Req, Res]) RecvMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for RecvMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
-type AgentService_DataClient_RecvMsg_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// RecvMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_DataClient_Expecter[Req, Res]) RecvMsg(m interface{}) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
-	return &AgentService_DataClient_RecvMsg_Call[Req, Res]{Call: _e.mock.On("RecvMsg", m)}
-}
-
-func (_c *AgentService_DataClient_RecvMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_RecvMsg_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_RecvMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_DataClient_RecvMsg_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Send provides a mock function with given fields: _a0
-func (_m *AgentService_DataClient[Req, Res]) Send(_a0 *agent.DataRequest) error {
-	ret := _m.Called(_a0)
-
-	if len(ret) == 0 {
-		panic("no return value specified for Send")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(*agent.DataRequest) error); ok {
-		r0 = rf(_a0)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
-type AgentService_DataClient_Send_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Send is a helper method to define mock.On call
-//   - _a0 *agent.DataRequest
-func (_e *AgentService_DataClient_Expecter[Req, Res]) Send(_a0 interface{}) *AgentService_DataClient_Send_Call[Req, Res] {
-	return &AgentService_DataClient_Send_Call[Req, Res]{Call: _e.mock.On("Send", _a0)}
-}
-
-func (_c *AgentService_DataClient_Send_Call[Req, Res]) Run(run func(_a0 *agent.DataRequest)) *AgentService_DataClient_Send_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(*agent.DataRequest))
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_Send_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_Send_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_Send_Call[Req, Res]) RunAndReturn(run func(*agent.DataRequest) error) *AgentService_DataClient_Send_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// SendMsg provides a mock function with given fields: m
-func (_m *AgentService_DataClient[Req, Res]) SendMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for SendMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
-type AgentService_DataClient_SendMsg_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// SendMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_DataClient_Expecter[Req, Res]) SendMsg(m interface{}) *AgentService_DataClient_SendMsg_Call[Req, Res] {
-	return &AgentService_DataClient_SendMsg_Call[Req, Res]{Call: _e.mock.On("SendMsg", m)}
-}
-
-func (_c *AgentService_DataClient_SendMsg_Call[Req, Res]) Run(run func(m interface{})) *AgentService_DataClient_SendMsg_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_SendMsg_Call[Req, Res]) Return(_a0 error) *AgentService_DataClient_SendMsg_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_SendMsg_Call[Req, Res]) RunAndReturn(run func(interface{}) error) *AgentService_DataClient_SendMsg_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Trailer provides a mock function with no fields
-func (_m *AgentService_DataClient[Req, Res]) Trailer() metadata.MD {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Trailer")
-	}
-
-	var r0 metadata.MD
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_DataClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
-type AgentService_DataClient_Trailer_Call[Req interface{}, Res interface{}] struct {
-	*mock.Call
-}
-
-// Trailer is a helper method to define mock.On call
-func (_e *AgentService_DataClient_Expecter[Req, Res]) Trailer() *AgentService_DataClient_Trailer_Call[Req, Res] {
-	return &AgentService_DataClient_Trailer_Call[Req, Res]{Call: _e.mock.On("Trailer")}
-}
-
-func (_c *AgentService_DataClient_Trailer_Call[Req, Res]) Run(run func()) *AgentService_DataClient_Trailer_Call[Req, Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_DataClient_Trailer_Call[Req, Res]) Return(_a0 metadata.MD) *AgentService_DataClient_Trailer_Call[Req, Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_DataClient_Trailer_Call[Req, Res]) RunAndReturn(run func() metadata.MD) *AgentService_DataClient_Trailer_Call[Req, Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// NewAgentService_DataClient creates a new instance of AgentService_DataClient. 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 NewAgentService_DataClient[Req interface{}, Res interface{}](t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *AgentService_DataClient[Req, Res] {
-	mock := &AgentService_DataClient[Req, Res]{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/mocks/agent_grpc_ima.go

@@ -1,388 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-// Code generated by mockery v2.53.3. DO NOT EDIT.
-
-package mocks
-
-import (
-	context "context"
-
-	agent "github.com/ultravioletrs/cocos/agent"
-
-	metadata "google.golang.org/grpc/metadata"
-
-	mock "github.com/stretchr/testify/mock"
-)
-
-// AgentService_IMAMeasurementsClient is an autogenerated mock type for the AgentService_IMAMeasurementsClient type
-type AgentService_IMAMeasurementsClient[Res interface{}] struct {
-	mock.Mock
-}
-
-type AgentService_IMAMeasurementsClient_Expecter[Res interface{}] struct {
-	mock *mock.Mock
-}
-
-func (_m *AgentService_IMAMeasurementsClient[Res]) EXPECT() *AgentService_IMAMeasurementsClient_Expecter[Res] {
-	return &AgentService_IMAMeasurementsClient_Expecter[Res]{mock: &_m.Mock}
-}
-
-// CloseSend provides a mock function with no fields
-func (_m *AgentService_IMAMeasurementsClient[Res]) CloseSend() error {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for CloseSend")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_IMAMeasurementsClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
-type AgentService_IMAMeasurementsClient_CloseSend_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// CloseSend is a helper method to define mock.On call
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) CloseSend() *AgentService_IMAMeasurementsClient_CloseSend_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_CloseSend_Call[Res]{Call: _e.mock.On("CloseSend")}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call[Res]) Run(run func()) *AgentService_IMAMeasurementsClient_CloseSend_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call[Res]) Return(_a0 error) *AgentService_IMAMeasurementsClient_CloseSend_Call[Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call[Res]) RunAndReturn(run func() error) *AgentService_IMAMeasurementsClient_CloseSend_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Context provides a mock function with no fields
-func (_m *AgentService_IMAMeasurementsClient[Res]) Context() context.Context {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Context")
-	}
-
-	var r0 context.Context
-	if rf, ok := ret.Get(0).(func() context.Context); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(context.Context)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_IMAMeasurementsClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
-type AgentService_IMAMeasurementsClient_Context_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// Context is a helper method to define mock.On call
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) Context() *AgentService_IMAMeasurementsClient_Context_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_Context_Call[Res]{Call: _e.mock.On("Context")}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Context_Call[Res]) Run(run func()) *AgentService_IMAMeasurementsClient_Context_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Context_Call[Res]) Return(_a0 context.Context) *AgentService_IMAMeasurementsClient_Context_Call[Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Context_Call[Res]) RunAndReturn(run func() context.Context) *AgentService_IMAMeasurementsClient_Context_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Header provides a mock function with no fields
-func (_m *AgentService_IMAMeasurementsClient[Res]) Header() (metadata.MD, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Header")
-	}
-
-	var r0 metadata.MD
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_IMAMeasurementsClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
-type AgentService_IMAMeasurementsClient_Header_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// Header is a helper method to define mock.On call
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) Header() *AgentService_IMAMeasurementsClient_Header_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_Header_Call[Res]{Call: _e.mock.On("Header")}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Header_Call[Res]) Run(run func()) *AgentService_IMAMeasurementsClient_Header_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Header_Call[Res]) Return(_a0 metadata.MD, _a1 error) *AgentService_IMAMeasurementsClient_Header_Call[Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Header_Call[Res]) RunAndReturn(run func() (metadata.MD, error)) *AgentService_IMAMeasurementsClient_Header_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Recv provides a mock function with no fields
-func (_m *AgentService_IMAMeasurementsClient[Res]) Recv() (*agent.IMAMeasurementsResponse, error) {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Recv")
-	}
-
-	var r0 *agent.IMAMeasurementsResponse
-	var r1 error
-	if rf, ok := ret.Get(0).(func() (*agent.IMAMeasurementsResponse, error)); ok {
-		return rf()
-	}
-	if rf, ok := ret.Get(0).(func() *agent.IMAMeasurementsResponse); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(*agent.IMAMeasurementsResponse)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// AgentService_IMAMeasurementsClient_Recv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Recv'
-type AgentService_IMAMeasurementsClient_Recv_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// Recv is a helper method to define mock.On call
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) Recv() *AgentService_IMAMeasurementsClient_Recv_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_Recv_Call[Res]{Call: _e.mock.On("Recv")}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Recv_Call[Res]) Run(run func()) *AgentService_IMAMeasurementsClient_Recv_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Recv_Call[Res]) Return(_a0 *agent.IMAMeasurementsResponse, _a1 error) *AgentService_IMAMeasurementsClient_Recv_Call[Res] {
-	_c.Call.Return(_a0, _a1)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Recv_Call[Res]) RunAndReturn(run func() (*agent.IMAMeasurementsResponse, error)) *AgentService_IMAMeasurementsClient_Recv_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// RecvMsg provides a mock function with given fields: m
-func (_m *AgentService_IMAMeasurementsClient[Res]) RecvMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for RecvMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_IMAMeasurementsClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
-type AgentService_IMAMeasurementsClient_RecvMsg_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// RecvMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) RecvMsg(m interface{}) *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_RecvMsg_Call[Res]{Call: _e.mock.On("RecvMsg", m)}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res]) Run(run func(m interface{})) *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res]) Return(_a0 error) *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res]) RunAndReturn(run func(interface{}) error) *AgentService_IMAMeasurementsClient_RecvMsg_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// SendMsg provides a mock function with given fields: m
-func (_m *AgentService_IMAMeasurementsClient[Res]) SendMsg(m interface{}) error {
-	ret := _m.Called(m)
-
-	if len(ret) == 0 {
-		panic("no return value specified for SendMsg")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
-		r0 = rf(m)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// AgentService_IMAMeasurementsClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
-type AgentService_IMAMeasurementsClient_SendMsg_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// SendMsg is a helper method to define mock.On call
-//   - m interface{}
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) SendMsg(m interface{}) *AgentService_IMAMeasurementsClient_SendMsg_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_SendMsg_Call[Res]{Call: _e.mock.On("SendMsg", m)}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call[Res]) Run(run func(m interface{})) *AgentService_IMAMeasurementsClient_SendMsg_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}))
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call[Res]) Return(_a0 error) *AgentService_IMAMeasurementsClient_SendMsg_Call[Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call[Res]) RunAndReturn(run func(interface{}) error) *AgentService_IMAMeasurementsClient_SendMsg_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// Trailer provides a mock function with no fields
-func (_m *AgentService_IMAMeasurementsClient[Res]) Trailer() metadata.MD {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Trailer")
-	}
-
-	var r0 metadata.MD
-	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
-		r0 = rf()
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(metadata.MD)
-		}
-	}
-
-	return r0
-}
-
-// AgentService_IMAMeasurementsClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
-type AgentService_IMAMeasurementsClient_Trailer_Call[Res interface{}] struct {
-	*mock.Call
-}
-
-// Trailer is a helper method to define mock.On call
-func (_e *AgentService_IMAMeasurementsClient_Expecter[Res]) Trailer() *AgentService_IMAMeasurementsClient_Trailer_Call[Res] {
-	return &AgentService_IMAMeasurementsClient_Trailer_Call[Res]{Call: _e.mock.On("Trailer")}
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Trailer_Call[Res]) Run(run func()) *AgentService_IMAMeasurementsClient_Trailer_Call[Res] {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Trailer_Call[Res]) Return(_a0 metadata.MD) *AgentService_IMAMeasurementsClient_Trailer_Call[Res] {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *AgentService_IMAMeasurementsClient_Trailer_Call[Res]) RunAndReturn(run func() metadata.MD) *AgentService_IMAMeasurementsClient_Trailer_Call[Res] {
-	_c.Call.Return(run)
-	return _c
-}
-
-// NewAgentService_IMAMeasurementsClient creates a new instance of AgentService_IMAMeasurementsClient. 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 NewAgentService_IMAMeasurementsClient[Res interface{}](t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *AgentService_IMAMeasurementsClient[Res] {
-	mock := &AgentService_IMAMeasurementsClient[Res]{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/mocks/agentservice_algoclient.go

@@ -0,0 +1,442 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
+
+package mocks
+
+import (
+	"context"
+
+	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"google.golang.org/grpc/metadata"
+)
+
+// NewAgentService_AlgoClient creates a new instance of AgentService_AlgoClient. 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 NewAgentService_AlgoClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *AgentService_AlgoClient {
+	mock := &AgentService_AlgoClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// AgentService_AlgoClient is an autogenerated mock type for the AgentService_AlgoClient type
+type AgentService_AlgoClient struct {
+	mock.Mock
+}
+
+type AgentService_AlgoClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *AgentService_AlgoClient) EXPECT() *AgentService_AlgoClient_Expecter {
+	return &AgentService_AlgoClient_Expecter{mock: &_m.Mock}
+}
+
+// CloseAndRecv provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) CloseAndRecv() (*agent.AlgoResponse, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseAndRecv")
+	}
+
+	var r0 *agent.AlgoResponse
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (*agent.AlgoResponse, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() *agent.AlgoResponse); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*agent.AlgoResponse)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_AlgoClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
+type AgentService_AlgoClient_CloseAndRecv_Call struct {
+	*mock.Call
+}
+
+// CloseAndRecv is a helper method to define mock.On call
+func (_e *AgentService_AlgoClient_Expecter) CloseAndRecv() *AgentService_AlgoClient_CloseAndRecv_Call {
+	return &AgentService_AlgoClient_CloseAndRecv_Call{Call: _e.mock.On("CloseAndRecv")}
+}
+
+func (_c *AgentService_AlgoClient_CloseAndRecv_Call) Run(run func()) *AgentService_AlgoClient_CloseAndRecv_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_CloseAndRecv_Call) Return(algoResponse *agent.AlgoResponse, err error) *AgentService_AlgoClient_CloseAndRecv_Call {
+	_c.Call.Return(algoResponse, err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_CloseAndRecv_Call) RunAndReturn(run func() (*agent.AlgoResponse, error)) *AgentService_AlgoClient_CloseAndRecv_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// CloseSend provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) CloseSend() error {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseSend")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
+type AgentService_AlgoClient_CloseSend_Call struct {
+	*mock.Call
+}
+
+// CloseSend is a helper method to define mock.On call
+func (_e *AgentService_AlgoClient_Expecter) CloseSend() *AgentService_AlgoClient_CloseSend_Call {
+	return &AgentService_AlgoClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
+}
+
+func (_c *AgentService_AlgoClient_CloseSend_Call) Run(run func()) *AgentService_AlgoClient_CloseSend_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_CloseSend_Call) Return(err error) *AgentService_AlgoClient_CloseSend_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_CloseSend_Call) RunAndReturn(run func() error) *AgentService_AlgoClient_CloseSend_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Context provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) Context() context.Context {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Context")
+	}
+
+	var r0 context.Context
+	if returnFunc, ok := ret.Get(0).(func() context.Context); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
+type AgentService_AlgoClient_Context_Call struct {
+	*mock.Call
+}
+
+// Context is a helper method to define mock.On call
+func (_e *AgentService_AlgoClient_Expecter) Context() *AgentService_AlgoClient_Context_Call {
+	return &AgentService_AlgoClient_Context_Call{Call: _e.mock.On("Context")}
+}
+
+func (_c *AgentService_AlgoClient_Context_Call) Run(run func()) *AgentService_AlgoClient_Context_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Context_Call) Return(context1 context.Context) *AgentService_AlgoClient_Context_Call {
+	_c.Call.Return(context1)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Context_Call) RunAndReturn(run func() context.Context) *AgentService_AlgoClient_Context_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Header provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) Header() (metadata.MD, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Header")
+	}
+
+	var r0 metadata.MD
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_AlgoClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
+type AgentService_AlgoClient_Header_Call struct {
+	*mock.Call
+}
+
+// Header is a helper method to define mock.On call
+func (_e *AgentService_AlgoClient_Expecter) Header() *AgentService_AlgoClient_Header_Call {
+	return &AgentService_AlgoClient_Header_Call{Call: _e.mock.On("Header")}
+}
+
+func (_c *AgentService_AlgoClient_Header_Call) Run(run func()) *AgentService_AlgoClient_Header_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Header_Call) Return(mD metadata.MD, err error) *AgentService_AlgoClient_Header_Call {
+	_c.Call.Return(mD, err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *AgentService_AlgoClient_Header_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// RecvMsg provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) RecvMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RecvMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
+type AgentService_AlgoClient_RecvMsg_Call struct {
+	*mock.Call
+}
+
+// RecvMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_AlgoClient_Expecter) RecvMsg(m interface{}) *AgentService_AlgoClient_RecvMsg_Call {
+	return &AgentService_AlgoClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
+}
+
+func (_c *AgentService_AlgoClient_RecvMsg_Call) Run(run func(m any)) *AgentService_AlgoClient_RecvMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_RecvMsg_Call) Return(err error) *AgentService_AlgoClient_RecvMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_RecvMsg_Call) RunAndReturn(run func(m any) error) *AgentService_AlgoClient_RecvMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Send provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) Send(algoRequest *agent.AlgoRequest) error {
+	ret := _mock.Called(algoRequest)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Send")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(*agent.AlgoRequest) error); ok {
+		r0 = returnFunc(algoRequest)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
+type AgentService_AlgoClient_Send_Call struct {
+	*mock.Call
+}
+
+// Send is a helper method to define mock.On call
+//   - algoRequest *agent.AlgoRequest
+func (_e *AgentService_AlgoClient_Expecter) Send(algoRequest interface{}) *AgentService_AlgoClient_Send_Call {
+	return &AgentService_AlgoClient_Send_Call{Call: _e.mock.On("Send", algoRequest)}
+}
+
+func (_c *AgentService_AlgoClient_Send_Call) Run(run func(algoRequest *agent.AlgoRequest)) *AgentService_AlgoClient_Send_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 *agent.AlgoRequest
+		if args[0] != nil {
+			arg0 = args[0].(*agent.AlgoRequest)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Send_Call) Return(err error) *AgentService_AlgoClient_Send_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Send_Call) RunAndReturn(run func(algoRequest *agent.AlgoRequest) error) *AgentService_AlgoClient_Send_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SendMsg provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) SendMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SendMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
+type AgentService_AlgoClient_SendMsg_Call struct {
+	*mock.Call
+}
+
+// SendMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_AlgoClient_Expecter) SendMsg(m interface{}) *AgentService_AlgoClient_SendMsg_Call {
+	return &AgentService_AlgoClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
+}
+
+func (_c *AgentService_AlgoClient_SendMsg_Call) Run(run func(m any)) *AgentService_AlgoClient_SendMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_SendMsg_Call) Return(err error) *AgentService_AlgoClient_SendMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_SendMsg_Call) RunAndReturn(run func(m any) error) *AgentService_AlgoClient_SendMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Trailer provides a mock function for the type AgentService_AlgoClient
+func (_mock *AgentService_AlgoClient) Trailer() metadata.MD {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Trailer")
+	}
+
+	var r0 metadata.MD
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	return r0
+}
+
+// AgentService_AlgoClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
+type AgentService_AlgoClient_Trailer_Call struct {
+	*mock.Call
+}
+
+// Trailer is a helper method to define mock.On call
+func (_e *AgentService_AlgoClient_Expecter) Trailer() *AgentService_AlgoClient_Trailer_Call {
+	return &AgentService_AlgoClient_Trailer_Call{Call: _e.mock.On("Trailer")}
+}
+
+func (_c *AgentService_AlgoClient_Trailer_Call) Run(run func()) *AgentService_AlgoClient_Trailer_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Trailer_Call) Return(mD metadata.MD) *AgentService_AlgoClient_Trailer_Call {
+	_c.Call.Return(mD)
+	return _c
+}
+
+func (_c *AgentService_AlgoClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *AgentService_AlgoClient_Trailer_Call {
+	_c.Call.Return(run)
+	return _c
+}

agent/mocks/agentservice_dataclient.go

@@ -0,0 +1,442 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
+
+package mocks
+
+import (
+	"context"
+
+	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"google.golang.org/grpc/metadata"
+)
+
+// NewAgentService_DataClient creates a new instance of AgentService_DataClient. 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 NewAgentService_DataClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *AgentService_DataClient {
+	mock := &AgentService_DataClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// AgentService_DataClient is an autogenerated mock type for the AgentService_DataClient type
+type AgentService_DataClient struct {
+	mock.Mock
+}
+
+type AgentService_DataClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *AgentService_DataClient) EXPECT() *AgentService_DataClient_Expecter {
+	return &AgentService_DataClient_Expecter{mock: &_m.Mock}
+}
+
+// CloseAndRecv provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) CloseAndRecv() (*agent.DataResponse, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseAndRecv")
+	}
+
+	var r0 *agent.DataResponse
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (*agent.DataResponse, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() *agent.DataResponse); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*agent.DataResponse)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_DataClient_CloseAndRecv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseAndRecv'
+type AgentService_DataClient_CloseAndRecv_Call struct {
+	*mock.Call
+}
+
+// CloseAndRecv is a helper method to define mock.On call
+func (_e *AgentService_DataClient_Expecter) CloseAndRecv() *AgentService_DataClient_CloseAndRecv_Call {
+	return &AgentService_DataClient_CloseAndRecv_Call{Call: _e.mock.On("CloseAndRecv")}
+}
+
+func (_c *AgentService_DataClient_CloseAndRecv_Call) Run(run func()) *AgentService_DataClient_CloseAndRecv_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_CloseAndRecv_Call) Return(dataResponse *agent.DataResponse, err error) *AgentService_DataClient_CloseAndRecv_Call {
+	_c.Call.Return(dataResponse, err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_CloseAndRecv_Call) RunAndReturn(run func() (*agent.DataResponse, error)) *AgentService_DataClient_CloseAndRecv_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// CloseSend provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) CloseSend() error {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseSend")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_DataClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
+type AgentService_DataClient_CloseSend_Call struct {
+	*mock.Call
+}
+
+// CloseSend is a helper method to define mock.On call
+func (_e *AgentService_DataClient_Expecter) CloseSend() *AgentService_DataClient_CloseSend_Call {
+	return &AgentService_DataClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
+}
+
+func (_c *AgentService_DataClient_CloseSend_Call) Run(run func()) *AgentService_DataClient_CloseSend_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_CloseSend_Call) Return(err error) *AgentService_DataClient_CloseSend_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_CloseSend_Call) RunAndReturn(run func() error) *AgentService_DataClient_CloseSend_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Context provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) Context() context.Context {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Context")
+	}
+
+	var r0 context.Context
+	if returnFunc, ok := ret.Get(0).(func() context.Context); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+	return r0
+}
+
+// AgentService_DataClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
+type AgentService_DataClient_Context_Call struct {
+	*mock.Call
+}
+
+// Context is a helper method to define mock.On call
+func (_e *AgentService_DataClient_Expecter) Context() *AgentService_DataClient_Context_Call {
+	return &AgentService_DataClient_Context_Call{Call: _e.mock.On("Context")}
+}
+
+func (_c *AgentService_DataClient_Context_Call) Run(run func()) *AgentService_DataClient_Context_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_Context_Call) Return(context1 context.Context) *AgentService_DataClient_Context_Call {
+	_c.Call.Return(context1)
+	return _c
+}
+
+func (_c *AgentService_DataClient_Context_Call) RunAndReturn(run func() context.Context) *AgentService_DataClient_Context_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Header provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) Header() (metadata.MD, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Header")
+	}
+
+	var r0 metadata.MD
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_DataClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
+type AgentService_DataClient_Header_Call struct {
+	*mock.Call
+}
+
+// Header is a helper method to define mock.On call
+func (_e *AgentService_DataClient_Expecter) Header() *AgentService_DataClient_Header_Call {
+	return &AgentService_DataClient_Header_Call{Call: _e.mock.On("Header")}
+}
+
+func (_c *AgentService_DataClient_Header_Call) Run(run func()) *AgentService_DataClient_Header_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_Header_Call) Return(mD metadata.MD, err error) *AgentService_DataClient_Header_Call {
+	_c.Call.Return(mD, err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *AgentService_DataClient_Header_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// RecvMsg provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) RecvMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RecvMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_DataClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
+type AgentService_DataClient_RecvMsg_Call struct {
+	*mock.Call
+}
+
+// RecvMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_DataClient_Expecter) RecvMsg(m interface{}) *AgentService_DataClient_RecvMsg_Call {
+	return &AgentService_DataClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
+}
+
+func (_c *AgentService_DataClient_RecvMsg_Call) Run(run func(m any)) *AgentService_DataClient_RecvMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_RecvMsg_Call) Return(err error) *AgentService_DataClient_RecvMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_RecvMsg_Call) RunAndReturn(run func(m any) error) *AgentService_DataClient_RecvMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Send provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) Send(dataRequest *agent.DataRequest) error {
+	ret := _mock.Called(dataRequest)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Send")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(*agent.DataRequest) error); ok {
+		r0 = returnFunc(dataRequest)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_DataClient_Send_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Send'
+type AgentService_DataClient_Send_Call struct {
+	*mock.Call
+}
+
+// Send is a helper method to define mock.On call
+//   - dataRequest *agent.DataRequest
+func (_e *AgentService_DataClient_Expecter) Send(dataRequest interface{}) *AgentService_DataClient_Send_Call {
+	return &AgentService_DataClient_Send_Call{Call: _e.mock.On("Send", dataRequest)}
+}
+
+func (_c *AgentService_DataClient_Send_Call) Run(run func(dataRequest *agent.DataRequest)) *AgentService_DataClient_Send_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 *agent.DataRequest
+		if args[0] != nil {
+			arg0 = args[0].(*agent.DataRequest)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_Send_Call) Return(err error) *AgentService_DataClient_Send_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_Send_Call) RunAndReturn(run func(dataRequest *agent.DataRequest) error) *AgentService_DataClient_Send_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SendMsg provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) SendMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SendMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_DataClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
+type AgentService_DataClient_SendMsg_Call struct {
+	*mock.Call
+}
+
+// SendMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_DataClient_Expecter) SendMsg(m interface{}) *AgentService_DataClient_SendMsg_Call {
+	return &AgentService_DataClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
+}
+
+func (_c *AgentService_DataClient_SendMsg_Call) Run(run func(m any)) *AgentService_DataClient_SendMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_SendMsg_Call) Return(err error) *AgentService_DataClient_SendMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_DataClient_SendMsg_Call) RunAndReturn(run func(m any) error) *AgentService_DataClient_SendMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Trailer provides a mock function for the type AgentService_DataClient
+func (_mock *AgentService_DataClient) Trailer() metadata.MD {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Trailer")
+	}
+
+	var r0 metadata.MD
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	return r0
+}
+
+// AgentService_DataClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
+type AgentService_DataClient_Trailer_Call struct {
+	*mock.Call
+}
+
+// Trailer is a helper method to define mock.On call
+func (_e *AgentService_DataClient_Expecter) Trailer() *AgentService_DataClient_Trailer_Call {
+	return &AgentService_DataClient_Trailer_Call{Call: _e.mock.On("Trailer")}
+}
+
+func (_c *AgentService_DataClient_Trailer_Call) Run(run func()) *AgentService_DataClient_Trailer_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_DataClient_Trailer_Call) Return(mD metadata.MD) *AgentService_DataClient_Trailer_Call {
+	_c.Call.Return(mD)
+	return _c
+}
+
+func (_c *AgentService_DataClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *AgentService_DataClient_Trailer_Call {
+	_c.Call.Return(run)
+	return _c
+}

agent/mocks/agentservice_imameasurementsclient.go

@@ -0,0 +1,391 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
+
+package mocks
+
+import (
+	"context"
+
+	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"google.golang.org/grpc/metadata"
+)
+
+// NewAgentService_IMAMeasurementsClient creates a new instance of AgentService_IMAMeasurementsClient. 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 NewAgentService_IMAMeasurementsClient(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *AgentService_IMAMeasurementsClient {
+	mock := &AgentService_IMAMeasurementsClient{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// AgentService_IMAMeasurementsClient is an autogenerated mock type for the AgentService_IMAMeasurementsClient type
+type AgentService_IMAMeasurementsClient struct {
+	mock.Mock
+}
+
+type AgentService_IMAMeasurementsClient_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *AgentService_IMAMeasurementsClient) EXPECT() *AgentService_IMAMeasurementsClient_Expecter {
+	return &AgentService_IMAMeasurementsClient_Expecter{mock: &_m.Mock}
+}
+
+// CloseSend provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) CloseSend() error {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseSend")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_IMAMeasurementsClient_CloseSend_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'CloseSend'
+type AgentService_IMAMeasurementsClient_CloseSend_Call struct {
+	*mock.Call
+}
+
+// CloseSend is a helper method to define mock.On call
+func (_e *AgentService_IMAMeasurementsClient_Expecter) CloseSend() *AgentService_IMAMeasurementsClient_CloseSend_Call {
+	return &AgentService_IMAMeasurementsClient_CloseSend_Call{Call: _e.mock.On("CloseSend")}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call) Run(run func()) *AgentService_IMAMeasurementsClient_CloseSend_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call) Return(err error) *AgentService_IMAMeasurementsClient_CloseSend_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_CloseSend_Call) RunAndReturn(run func() error) *AgentService_IMAMeasurementsClient_CloseSend_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Context provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) Context() context.Context {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Context")
+	}
+
+	var r0 context.Context
+	if returnFunc, ok := ret.Get(0).(func() context.Context); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+	return r0
+}
+
+// AgentService_IMAMeasurementsClient_Context_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Context'
+type AgentService_IMAMeasurementsClient_Context_Call struct {
+	*mock.Call
+}
+
+// Context is a helper method to define mock.On call
+func (_e *AgentService_IMAMeasurementsClient_Expecter) Context() *AgentService_IMAMeasurementsClient_Context_Call {
+	return &AgentService_IMAMeasurementsClient_Context_Call{Call: _e.mock.On("Context")}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Context_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Context_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Context_Call) Return(context1 context.Context) *AgentService_IMAMeasurementsClient_Context_Call {
+	_c.Call.Return(context1)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Context_Call) RunAndReturn(run func() context.Context) *AgentService_IMAMeasurementsClient_Context_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Header provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) Header() (metadata.MD, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Header")
+	}
+
+	var r0 metadata.MD
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_IMAMeasurementsClient_Header_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Header'
+type AgentService_IMAMeasurementsClient_Header_Call struct {
+	*mock.Call
+}
+
+// Header is a helper method to define mock.On call
+func (_e *AgentService_IMAMeasurementsClient_Expecter) Header() *AgentService_IMAMeasurementsClient_Header_Call {
+	return &AgentService_IMAMeasurementsClient_Header_Call{Call: _e.mock.On("Header")}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Header_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Header_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Header_Call) Return(mD metadata.MD, err error) *AgentService_IMAMeasurementsClient_Header_Call {
+	_c.Call.Return(mD, err)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Header_Call) RunAndReturn(run func() (metadata.MD, error)) *AgentService_IMAMeasurementsClient_Header_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Recv provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) Recv() (*agent.IMAMeasurementsResponse, error) {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Recv")
+	}
+
+	var r0 *agent.IMAMeasurementsResponse
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() (*agent.IMAMeasurementsResponse, error)); ok {
+		return returnFunc()
+	}
+	if returnFunc, ok := ret.Get(0).(func() *agent.IMAMeasurementsResponse); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*agent.IMAMeasurementsResponse)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// AgentService_IMAMeasurementsClient_Recv_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Recv'
+type AgentService_IMAMeasurementsClient_Recv_Call struct {
+	*mock.Call
+}
+
+// Recv is a helper method to define mock.On call
+func (_e *AgentService_IMAMeasurementsClient_Expecter) Recv() *AgentService_IMAMeasurementsClient_Recv_Call {
+	return &AgentService_IMAMeasurementsClient_Recv_Call{Call: _e.mock.On("Recv")}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Recv_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Recv_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Recv_Call) Return(iMAMeasurementsResponse *agent.IMAMeasurementsResponse, err error) *AgentService_IMAMeasurementsClient_Recv_Call {
+	_c.Call.Return(iMAMeasurementsResponse, err)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Recv_Call) RunAndReturn(run func() (*agent.IMAMeasurementsResponse, error)) *AgentService_IMAMeasurementsClient_Recv_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// RecvMsg provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) RecvMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RecvMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_IMAMeasurementsClient_RecvMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RecvMsg'
+type AgentService_IMAMeasurementsClient_RecvMsg_Call struct {
+	*mock.Call
+}
+
+// RecvMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_IMAMeasurementsClient_Expecter) RecvMsg(m interface{}) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
+	return &AgentService_IMAMeasurementsClient_RecvMsg_Call{Call: _e.mock.On("RecvMsg", m)}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call) Run(run func(m any)) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call) Return(err error) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_RecvMsg_Call) RunAndReturn(run func(m any) error) *AgentService_IMAMeasurementsClient_RecvMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SendMsg provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) SendMsg(m any) error {
+	ret := _mock.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SendMsg")
+	}
+
+	var r0 error
+	if returnFunc, ok := ret.Get(0).(func(any) error); ok {
+		r0 = returnFunc(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentService_IMAMeasurementsClient_SendMsg_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendMsg'
+type AgentService_IMAMeasurementsClient_SendMsg_Call struct {
+	*mock.Call
+}
+
+// SendMsg is a helper method to define mock.On call
+//   - m any
+func (_e *AgentService_IMAMeasurementsClient_Expecter) SendMsg(m interface{}) *AgentService_IMAMeasurementsClient_SendMsg_Call {
+	return &AgentService_IMAMeasurementsClient_SendMsg_Call{Call: _e.mock.On("SendMsg", m)}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call) Run(run func(m any)) *AgentService_IMAMeasurementsClient_SendMsg_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 any
+		if args[0] != nil {
+			arg0 = args[0].(any)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call) Return(err error) *AgentService_IMAMeasurementsClient_SendMsg_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_SendMsg_Call) RunAndReturn(run func(m any) error) *AgentService_IMAMeasurementsClient_SendMsg_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Trailer provides a mock function for the type AgentService_IMAMeasurementsClient
+func (_mock *AgentService_IMAMeasurementsClient) Trailer() metadata.MD {
+	ret := _mock.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Trailer")
+	}
+
+	var r0 metadata.MD
+	if returnFunc, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = returnFunc()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+	return r0
+}
+
+// AgentService_IMAMeasurementsClient_Trailer_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Trailer'
+type AgentService_IMAMeasurementsClient_Trailer_Call struct {
+	*mock.Call
+}
+
+// Trailer is a helper method to define mock.On call
+func (_e *AgentService_IMAMeasurementsClient_Expecter) Trailer() *AgentService_IMAMeasurementsClient_Trailer_Call {
+	return &AgentService_IMAMeasurementsClient_Trailer_Call{Call: _e.mock.On("Trailer")}
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Trailer_Call) Run(run func()) *AgentService_IMAMeasurementsClient_Trailer_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Trailer_Call) Return(mD metadata.MD) *AgentService_IMAMeasurementsClient_Trailer_Call {
+	_c.Call.Return(mD)
+	return _c
+}
+
+func (_c *AgentService_IMAMeasurementsClient_Trailer_Call) RunAndReturn(run func() metadata.MD) *AgentService_IMAMeasurementsClient_Trailer_Call {
+	_c.Call.Return(run)
+	return _c
+}

agent/mocks/service.go

@@ -1,19 +1,34 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
-	agent "github.com/ultravioletrs/cocos/agent"
-	attestation "github.com/ultravioletrs/cocos/pkg/attestation"
-
-	context "context"
+	"context"
 
 	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
 )
 
+// 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 {
+	mock.TestingT
+	Cleanup(func())
+}) *Service {
+	mock := &Service{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // Service is an autogenerated mock type for the Service type
 type Service struct {
 	mock.Mock
@@ -27,21 +42,20 @@ func (_m *Service) EXPECT() *Service_Expecter {
 	return &Service_Expecter{mock: &_m.Mock}
 }
 
-// Algo provides a mock function with given fields: ctx, algorithm
-func (_m *Service) Algo(ctx context.Context, algorithm agent.Algorithm) error {
-	ret := _m.Called(ctx, algorithm)
+// Algo provides a mock function for the type Service
+func (_mock *Service) Algo(ctx context.Context, algorithm agent.Algorithm) error {
+	ret := _mock.Called(ctx, algorithm)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Algo")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, agent.Algorithm) error); ok {
-		r0 = rf(ctx, algorithm)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, agent.Algorithm) error); ok {
+		r0 = returnFunc(ctx, algorithm)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -59,24 +73,35 @@ func (_e *Service_Expecter) Algo(ctx interface{}, algorithm interface{}) *Servic
 
 func (_c *Service_Algo_Call) Run(run func(ctx context.Context, algorithm agent.Algorithm)) *Service_Algo_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(agent.Algorithm))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 agent.Algorithm
+		if args[1] != nil {
+			arg1 = args[1].(agent.Algorithm)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_Algo_Call) Return(_a0 error) *Service_Algo_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_Algo_Call) Return(err error) *Service_Algo_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Service_Algo_Call) RunAndReturn(run func(context.Context, agent.Algorithm) error) *Service_Algo_Call {
+func (_c *Service_Algo_Call) RunAndReturn(run func(ctx context.Context, algorithm agent.Algorithm) error) *Service_Algo_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// 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 attestation.PlatformType) ([]byte, error) {
-	ret := _m.Called(ctx, reportData, nonce, attType)
+// Attestation provides a mock function for the type Service
+func (_mock *Service) Attestation(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error) {
+	ret := _mock.Called(ctx, reportData, nonce, attType)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Attestation")
@@ -84,23 +109,21 @@ func (_m *Service) Attestation(ctx context.Context, reportData [64]byte, nonce [
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) ([]byte, error)); ok {
-		return rf(ctx, reportData, nonce, attType)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) ([]byte, error)); ok {
+		return returnFunc(ctx, reportData, nonce, attType)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) []byte); ok {
-		r0 = rf(ctx, reportData, nonce, attType)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) []byte); ok {
+		r0 = returnFunc(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, [32]byte, attestation.PlatformType) error); ok {
-		r1 = rf(ctx, reportData, nonce, attType)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, [64]byte, [32]byte, attestation.PlatformType) error); ok {
+		r1 = returnFunc(ctx, reportData, nonce, attType)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -120,96 +143,124 @@ func (_e *Service_Expecter) Attestation(ctx interface{}, reportData interface{},
 
 func (_c *Service_Attestation_Call) Run(run func(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType)) *Service_Attestation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].([64]byte), args[2].([32]byte), args[3].(attestation.PlatformType))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 [64]byte
+		if args[1] != nil {
+			arg1 = args[1].([64]byte)
+		}
+		var arg2 [32]byte
+		if args[2] != nil {
+			arg2 = args[2].([32]byte)
+		}
+		var arg3 attestation.PlatformType
+		if args[3] != nil {
+			arg3 = args[3].(attestation.PlatformType)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+			arg3,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_Attestation_Call) Return(_a0 []byte, _a1 error) *Service_Attestation_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Service_Attestation_Call) Return(bytes []byte, err error) *Service_Attestation_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Service_Attestation_Call) RunAndReturn(run func(context.Context, [64]byte, [32]byte, attestation.PlatformType) ([]byte, error)) *Service_Attestation_Call {
+func (_c *Service_Attestation_Call) RunAndReturn(run func(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error)) *Service_Attestation_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// AttestationResult provides a mock function with given fields: ctx, nonce, attType
-func (_m *Service) AttestationResult(ctx context.Context, nonce [32]byte, attType attestation.PlatformType) ([]byte, error) {
-	ret := _m.Called(ctx, nonce, attType)
+// AzureAttestationToken provides a mock function for the type Service
+func (_mock *Service) AzureAttestationToken(ctx context.Context, nonce [32]byte) ([]byte, error) {
+	ret := _mock.Called(ctx, nonce)
 
 	if len(ret) == 0 {
-		panic("no return value specified for AttestationResult")
+		panic("no return value specified for AzureAttestationToken")
 	}
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, [32]byte, attestation.PlatformType) ([]byte, error)); ok {
-		return rf(ctx, nonce, attType)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, [32]byte) ([]byte, error)); ok {
+		return returnFunc(ctx, nonce)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, [32]byte, attestation.PlatformType) []byte); ok {
-		r0 = rf(ctx, nonce, attType)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, [32]byte) []byte); ok {
+		r0 = returnFunc(ctx, nonce)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context, [32]byte, attestation.PlatformType) error); ok {
-		r1 = rf(ctx, nonce, attType)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, [32]byte) error); ok {
+		r1 = returnFunc(ctx, nonce)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
-// Service_AttestationResult_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AttestationResult'
-type Service_AttestationResult_Call struct {
+// Service_AzureAttestationToken_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AzureAttestationToken'
+type Service_AzureAttestationToken_Call struct {
 	*mock.Call
 }
 
-// AttestationResult is a helper method to define mock.On call
+// AzureAttestationToken is a helper method to define mock.On call
 //   - ctx context.Context
 //   - nonce [32]byte
-//   - attType attestation.PlatformType
-func (_e *Service_Expecter) AttestationResult(ctx interface{}, nonce interface{}, attType interface{}) *Service_AttestationResult_Call {
-	return &Service_AttestationResult_Call{Call: _e.mock.On("AttestationResult", ctx, nonce, attType)}
+func (_e *Service_Expecter) AzureAttestationToken(ctx interface{}, nonce interface{}) *Service_AzureAttestationToken_Call {
+	return &Service_AzureAttestationToken_Call{Call: _e.mock.On("AzureAttestationToken", ctx, nonce)}
 }
 
-func (_c *Service_AttestationResult_Call) Run(run func(ctx context.Context, nonce [32]byte, attType attestation.PlatformType)) *Service_AttestationResult_Call {
+func (_c *Service_AzureAttestationToken_Call) Run(run func(ctx context.Context, nonce [32]byte)) *Service_AzureAttestationToken_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].([32]byte), args[2].(attestation.PlatformType))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 [32]byte
+		if args[1] != nil {
+			arg1 = args[1].([32]byte)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_AttestationResult_Call) Return(_a0 []byte, _a1 error) *Service_AttestationResult_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Service_AzureAttestationToken_Call) Return(bytes []byte, err error) *Service_AzureAttestationToken_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Service_AttestationResult_Call) RunAndReturn(run func(context.Context, [32]byte, attestation.PlatformType) ([]byte, error)) *Service_AttestationResult_Call {
+func (_c *Service_AzureAttestationToken_Call) RunAndReturn(run func(ctx context.Context, nonce [32]byte) ([]byte, error)) *Service_AzureAttestationToken_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Data provides a mock function with given fields: ctx, dataset
-func (_m *Service) Data(ctx context.Context, dataset agent.Dataset) error {
-	ret := _m.Called(ctx, dataset)
+// Data provides a mock function for the type Service
+func (_mock *Service) Data(ctx context.Context, dataset agent.Dataset) error {
+	ret := _mock.Called(ctx, dataset)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Data")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, agent.Dataset) error); ok {
-		r0 = rf(ctx, dataset)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, agent.Dataset) error); ok {
+		r0 = returnFunc(ctx, dataset)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -227,24 +278,35 @@ func (_e *Service_Expecter) Data(ctx interface{}, dataset interface{}) *Service_
 
 func (_c *Service_Data_Call) Run(run func(ctx context.Context, dataset agent.Dataset)) *Service_Data_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(agent.Dataset))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 agent.Dataset
+		if args[1] != nil {
+			arg1 = args[1].(agent.Dataset)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_Data_Call) Return(_a0 error) *Service_Data_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_Data_Call) Return(err error) *Service_Data_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Service_Data_Call) RunAndReturn(run func(context.Context, agent.Dataset) error) *Service_Data_Call {
+func (_c *Service_Data_Call) RunAndReturn(run func(ctx context.Context, dataset agent.Dataset) error) *Service_Data_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// IMAMeasurements provides a mock function with given fields: ctx
-func (_m *Service) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {
-	ret := _m.Called(ctx)
+// IMAMeasurements provides a mock function for the type Service
+func (_mock *Service) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {
+	ret := _mock.Called(ctx)
 
 	if len(ret) == 0 {
 		panic("no return value specified for IMAMeasurements")
@@ -253,31 +315,28 @@ func (_m *Service) IMAMeasurements(ctx context.Context) ([]byte, []byte, error)
 	var r0 []byte
 	var r1 []byte
 	var r2 error
-	if rf, ok := ret.Get(0).(func(context.Context) ([]byte, []byte, error)); ok {
-		return rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) ([]byte, []byte, error)); ok {
+		return returnFunc(ctx)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context) []byte); ok {
-		r0 = rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) []byte); ok {
+		r0 = returnFunc(ctx)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context) []byte); ok {
-		r1 = rf(ctx)
+	if returnFunc, ok := ret.Get(1).(func(context.Context) []byte); ok {
+		r1 = returnFunc(ctx)
 	} else {
 		if ret.Get(1) != nil {
 			r1 = ret.Get(1).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(2).(func(context.Context) error); ok {
-		r2 = rf(ctx)
+	if returnFunc, ok := ret.Get(2).(func(context.Context) error); ok {
+		r2 = returnFunc(ctx)
 	} else {
 		r2 = ret.Error(2)
 	}
-
 	return r0, r1, r2
 }
 
@@ -294,36 +353,41 @@ func (_e *Service_Expecter) IMAMeasurements(ctx interface{}) *Service_IMAMeasure
 
 func (_c *Service_IMAMeasurements_Call) Run(run func(ctx context.Context)) *Service_IMAMeasurements_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_IMAMeasurements_Call) Return(_a0 []byte, _a1 []byte, _a2 error) *Service_IMAMeasurements_Call {
-	_c.Call.Return(_a0, _a1, _a2)
+func (_c *Service_IMAMeasurements_Call) Return(bytes []byte, bytes1 []byte, err error) *Service_IMAMeasurements_Call {
+	_c.Call.Return(bytes, bytes1, err)
 	return _c
 }
 
-func (_c *Service_IMAMeasurements_Call) RunAndReturn(run func(context.Context) ([]byte, []byte, error)) *Service_IMAMeasurements_Call {
+func (_c *Service_IMAMeasurements_Call) RunAndReturn(run func(ctx context.Context) ([]byte, []byte, error)) *Service_IMAMeasurements_Call {
 	_c.Call.Return(run)
 	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)
+// InitComputation provides a mock function for the type Service
+func (_mock *Service) InitComputation(ctx context.Context, cmp agent.Computation) error {
+	ret := _mock.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)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, agent.Computation) error); ok {
+		r0 = returnFunc(ctx, cmp)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -341,24 +405,35 @@ func (_e *Service_Expecter) InitComputation(ctx interface{}, cmp interface{}) *S
 
 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))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 agent.Computation
+		if args[1] != nil {
+			arg1 = args[1].(agent.Computation)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_InitComputation_Call) Return(_a0 error) *Service_InitComputation_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_InitComputation_Call) Return(err error) *Service_InitComputation_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Service_InitComputation_Call) RunAndReturn(run func(context.Context, agent.Computation) error) *Service_InitComputation_Call {
+func (_c *Service_InitComputation_Call) RunAndReturn(run func(ctx context.Context, cmp 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)
+// Result provides a mock function for the type Service
+func (_mock *Service) Result(ctx context.Context) ([]byte, error) {
+	ret := _mock.Called(ctx)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Result")
@@ -366,23 +441,21 @@ func (_m *Service) Result(ctx context.Context) ([]byte, error) {
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context) ([]byte, error)); ok {
-		return rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) ([]byte, error)); ok {
+		return returnFunc(ctx)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context) []byte); ok {
-		r0 = rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) []byte); ok {
+		r0 = returnFunc(ctx)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
-		r1 = rf(ctx)
+	if returnFunc, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = returnFunc(ctx)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -399,36 +472,41 @@ func (_e *Service_Expecter) Result(ctx interface{}) *Service_Result_Call {
 
 func (_c *Service_Result_Call) Run(run func(ctx context.Context)) *Service_Result_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_Result_Call) Return(_a0 []byte, _a1 error) *Service_Result_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Service_Result_Call) Return(bytes []byte, err error) *Service_Result_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Service_Result_Call) RunAndReturn(run func(context.Context) ([]byte, error)) *Service_Result_Call {
+func (_c *Service_Result_Call) RunAndReturn(run func(ctx context.Context) ([]byte, error)) *Service_Result_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// State provides a mock function with no fields
-func (_m *Service) State() string {
-	ret := _m.Called()
+// State provides a mock function for the type Service
+func (_mock *Service) State() string {
+	ret := _mock.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()
+	if returnFunc, ok := ret.Get(0).(func() string); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Get(0).(string)
 	}
-
 	return r0
 }
 
@@ -449,8 +527,8 @@ func (_c *Service_State_Call) Run(run func()) *Service_State_Call {
 	return _c
 }
 
-func (_c *Service_State_Call) Return(_a0 string) *Service_State_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_State_Call) Return(s string) *Service_State_Call {
+	_c.Call.Return(s)
 	return _c
 }
 
@@ -459,21 +537,20 @@ func (_c *Service_State_Call) RunAndReturn(run func() string) *Service_State_Cal
 	return _c
 }
 
-// StopComputation provides a mock function with given fields: ctx
-func (_m *Service) StopComputation(ctx context.Context) error {
-	ret := _m.Called(ctx)
+// StopComputation provides a mock function for the type Service
+func (_mock *Service) StopComputation(ctx context.Context) error {
+	ret := _mock.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)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) error); ok {
+		r0 = returnFunc(ctx)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -490,31 +567,23 @@ func (_e *Service_Expecter) StopComputation(ctx interface{}) *Service_StopComput
 
 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))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_StopComputation_Call) Return(_a0 error) *Service_StopComputation_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_StopComputation_Call) Return(err error) *Service_StopComputation_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Service_StopComputation_Call) RunAndReturn(run func(context.Context) error) *Service_StopComputation_Call {
+func (_c *Service_StopComputation_Call) RunAndReturn(run func(ctx 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 {
-	mock.TestingT
-	Cleanup(func())
-}) *Service {
-	mock := &Service{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

agent/service.go

@@ -14,7 +14,7 @@ import (
 	sync "sync"
 	"time"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/ultravioletrs/cocos/agent/algorithm"
 	"github.com/ultravioletrs/cocos/agent/algorithm/binary"
 	"github.com/ultravioletrs/cocos/agent/algorithm/docker"
@@ -123,7 +123,7 @@ type Service interface {
 	Result(ctx context.Context) ([]byte, error)
 	Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error)
 	IMAMeasurements(ctx context.Context) ([]byte, []byte, error)
-	AttestationResult(ctx context.Context, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error)
+	AzureAttestationToken(ctx context.Context, nonce [vtpm.Nonce]byte) ([]byte, error)
 	State() string
 }
 
@@ -459,17 +459,16 @@ func (as *agentService) Attestation(ctx context.Context, reportData [quoteprovid
 	}
 }
 
-func (as *agentService) AttestationResult(ctx context.Context, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error) {
-	switch attType {
-	case attestation.AzureToken:
-		token, err := as.provider.AzureAttestationToken(nonce[:])
-		if err != nil {
-			return []byte{}, err
-		}
-		return token, nil
-	default:
+func (as *agentService) AzureAttestationToken(ctx context.Context, nonce [vtpm.Nonce]byte) ([]byte, error) {
+	if attestation.CCPlatform() != attestation.Azure {
 		return []byte{}, ErrAttestationType
 	}
+
+	token, err := as.provider.AzureAttestationToken(nonce[:])
+	if err != nil {
+		return []byte{}, err
+	}
+	return token, nil
 }
 
 func (as *agentService) runComputation(state statemachine.State) {

agent/service_test.go

@@ -12,8 +12,8 @@ import (
 	"testing"
 	"time"
 
-	mglog "github.com/absmach/magistrala/logger"
-	"github.com/absmach/magistrala/pkg/errors"
+	mglog "github.com/absmach/supermq/logger"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -411,35 +411,31 @@ func TestAttestation(t *testing.T) {
 	}
 }
 
-func TestAttestationResult(t *testing.T) {
+func TestAzureAttestationToken(t *testing.T) {
 	provider := new(mocks2.Provider)
 	cases := []struct {
-		name     string
-		nonce    [vtpm.Nonce]byte
-		platform attestation.PlatformType
-		token    []byte
-		err      error
+		name  string
+		nonce [vtpm.Nonce]byte
+		token []byte
+		err   error
 	}{
 		{
-			name:     "Azure token fetch successful",
-			nonce:    [32]byte{1, 2, 3}, // any test nonce
-			platform: attestation.AzureToken,
-			token:    []byte("mockToken"),
-			err:      nil,
+			name:  "Azure token fetch successful",
+			nonce: [32]byte{1, 2, 3}, // any test nonce
+			token: []byte("mockToken"),
+			err:   ErrAttestationType,
 		},
 		{
-			name:     "Azure token fetch failed",
-			nonce:    [32]byte{4, 5, 6},
-			platform: attestation.AzureToken,
-			token:    []byte{},
-			err:      ErrFetchAzureToken,
+			name:  "Azure token fetch failed",
+			nonce: [32]byte{4, 5, 6},
+			token: []byte{},
+			err:   ErrAttestationType,
 		},
 		{
-			name:     "Invalid attestation type",
-			nonce:    [32]byte{7, 8, 9},
-			platform: attestation.SNP,
-			token:    []byte{},
-			err:      ErrAttestationType,
+			name:  "Invalid attestation type",
+			nonce: [32]byte{7, 8, 9},
+			token: []byte{},
+			err:   ErrAttestationType,
 		},
 	}
 
@@ -448,17 +444,14 @@ func TestAttestationResult(t *testing.T) {
 			events := new(mocks.Service)
 			events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
 
-			if tc.platform == attestation.AzureToken {
-				provider.On("AzureAttestationToken", tc.nonce[:]).Return(tc.token, tc.err)
-			}
+			provider.On("AzureAttestationToken", tc.nonce[:]).Return(tc.token, tc.err)
 
 			ctx := context.Background()
 
 			svc := New(ctx, mglog.NewMock(), events, provider, 0)
 
-			result, err := svc.AttestationResult(ctx, tc.nonce, tc.platform)
+			_, err := svc.AzureAttestationToken(ctx, tc.nonce)
 			assert.True(t, errors.Contains(err, tc.err), "expected error %v, got %v", tc.err, err)
-			assert.Equal(t, tc.token, result)
 		})
 	}
 }

agent/statemachine/mocks/statemachine.go

@@ -1,17 +1,33 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
-	context "context"
+	"context"
 
 	mock "github.com/stretchr/testify/mock"
-	statemachine "github.com/ultravioletrs/cocos/agent/statemachine"
+	"github.com/ultravioletrs/cocos/agent/statemachine"
 )
 
+// NewStateMachine creates a new instance of StateMachine. 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 NewStateMachine(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *StateMachine {
+	mock := &StateMachine{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // StateMachine is an autogenerated mock type for the StateMachine type
 type StateMachine struct {
 	mock.Mock
@@ -25,9 +41,10 @@ func (_m *StateMachine) EXPECT() *StateMachine_Expecter {
 	return &StateMachine_Expecter{mock: &_m.Mock}
 }
 
-// AddTransition provides a mock function with given fields: t
-func (_m *StateMachine) AddTransition(t statemachine.Transition) {
-	_m.Called(t)
+// AddTransition provides a mock function for the type StateMachine
+func (_mock *StateMachine) AddTransition(t statemachine.Transition) {
+	_mock.Called(t)
+	return
 }
 
 // StateMachine_AddTransition_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AddTransition'
@@ -43,7 +60,13 @@ func (_e *StateMachine_Expecter) AddTransition(t interface{}) *StateMachine_AddT
 
 func (_c *StateMachine_AddTransition_Call) Run(run func(t statemachine.Transition)) *StateMachine_AddTransition_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(statemachine.Transition))
+		var arg0 statemachine.Transition
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.Transition)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
@@ -53,28 +76,27 @@ func (_c *StateMachine_AddTransition_Call) Return() *StateMachine_AddTransition_
 	return _c
 }
 
-func (_c *StateMachine_AddTransition_Call) RunAndReturn(run func(statemachine.Transition)) *StateMachine_AddTransition_Call {
+func (_c *StateMachine_AddTransition_Call) RunAndReturn(run func(t statemachine.Transition)) *StateMachine_AddTransition_Call {
 	_c.Run(run)
 	return _c
 }
 
-// GetState provides a mock function with no fields
-func (_m *StateMachine) GetState() statemachine.State {
-	ret := _m.Called()
+// GetState provides a mock function for the type StateMachine
+func (_mock *StateMachine) GetState() statemachine.State {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for GetState")
 	}
 
 	var r0 statemachine.State
-	if rf, ok := ret.Get(0).(func() statemachine.State); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() statemachine.State); ok {
+		r0 = returnFunc()
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(statemachine.State)
 		}
 	}
-
 	return r0
 }
 
@@ -95,8 +117,8 @@ func (_c *StateMachine_GetState_Call) Run(run func()) *StateMachine_GetState_Cal
 	return _c
 }
 
-func (_c *StateMachine_GetState_Call) Return(_a0 statemachine.State) *StateMachine_GetState_Call {
-	_c.Call.Return(_a0)
+func (_c *StateMachine_GetState_Call) Return(state statemachine.State) *StateMachine_GetState_Call {
+	_c.Call.Return(state)
 	return _c
 }
 
@@ -105,9 +127,10 @@ func (_c *StateMachine_GetState_Call) RunAndReturn(run func() statemachine.State
 	return _c
 }
 
-// Reset provides a mock function with given fields: initialState
-func (_m *StateMachine) Reset(initialState statemachine.State) {
-	_m.Called(initialState)
+// Reset provides a mock function for the type StateMachine
+func (_mock *StateMachine) Reset(initialState statemachine.State) {
+	_mock.Called(initialState)
+	return
 }
 
 // StateMachine_Reset_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Reset'
@@ -123,7 +146,13 @@ func (_e *StateMachine_Expecter) Reset(initialState interface{}) *StateMachine_R
 
 func (_c *StateMachine_Reset_Call) Run(run func(initialState statemachine.State)) *StateMachine_Reset_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(statemachine.State))
+		var arg0 statemachine.State
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.State)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
@@ -133,14 +162,15 @@ func (_c *StateMachine_Reset_Call) Return() *StateMachine_Reset_Call {
 	return _c
 }
 
-func (_c *StateMachine_Reset_Call) RunAndReturn(run func(statemachine.State)) *StateMachine_Reset_Call {
+func (_c *StateMachine_Reset_Call) RunAndReturn(run func(initialState statemachine.State)) *StateMachine_Reset_Call {
 	_c.Run(run)
 	return _c
 }
 
-// SendEvent provides a mock function with given fields: event
-func (_m *StateMachine) SendEvent(event statemachine.Event) {
-	_m.Called(event)
+// SendEvent provides a mock function for the type StateMachine
+func (_mock *StateMachine) SendEvent(event statemachine.Event) {
+	_mock.Called(event)
+	return
 }
 
 // StateMachine_SendEvent_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SendEvent'
@@ -156,7 +186,13 @@ func (_e *StateMachine_Expecter) SendEvent(event interface{}) *StateMachine_Send
 
 func (_c *StateMachine_SendEvent_Call) Run(run func(event statemachine.Event)) *StateMachine_SendEvent_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(statemachine.Event))
+		var arg0 statemachine.Event
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.Event)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
@@ -166,14 +202,15 @@ func (_c *StateMachine_SendEvent_Call) Return() *StateMachine_SendEvent_Call {
 	return _c
 }
 
-func (_c *StateMachine_SendEvent_Call) RunAndReturn(run func(statemachine.Event)) *StateMachine_SendEvent_Call {
+func (_c *StateMachine_SendEvent_Call) RunAndReturn(run func(event statemachine.Event)) *StateMachine_SendEvent_Call {
 	_c.Run(run)
 	return _c
 }
 
-// SetAction provides a mock function with given fields: state, action
-func (_m *StateMachine) SetAction(state statemachine.State, action statemachine.Action) {
-	_m.Called(state, action)
+// SetAction provides a mock function for the type StateMachine
+func (_mock *StateMachine) SetAction(state statemachine.State, action statemachine.Action) {
+	_mock.Called(state, action)
+	return
 }
 
 // StateMachine_SetAction_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'SetAction'
@@ -190,7 +227,18 @@ func (_e *StateMachine_Expecter) SetAction(state interface{}, action interface{}
 
 func (_c *StateMachine_SetAction_Call) Run(run func(state statemachine.State, action statemachine.Action)) *StateMachine_SetAction_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(statemachine.State), args[1].(statemachine.Action))
+		var arg0 statemachine.State
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.State)
+		}
+		var arg1 statemachine.Action
+		if args[1] != nil {
+			arg1 = args[1].(statemachine.Action)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
@@ -200,26 +248,25 @@ func (_c *StateMachine_SetAction_Call) Return() *StateMachine_SetAction_Call {
 	return _c
 }
 
-func (_c *StateMachine_SetAction_Call) RunAndReturn(run func(statemachine.State, statemachine.Action)) *StateMachine_SetAction_Call {
+func (_c *StateMachine_SetAction_Call) RunAndReturn(run func(state statemachine.State, action statemachine.Action)) *StateMachine_SetAction_Call {
 	_c.Run(run)
 	return _c
 }
 
-// Start provides a mock function with given fields: ctx
-func (_m *StateMachine) Start(ctx context.Context) error {
-	ret := _m.Called(ctx)
+// Start provides a mock function for the type StateMachine
+func (_mock *StateMachine) Start(ctx context.Context) error {
+	ret := _mock.Called(ctx)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Start")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context) error); ok {
-		r0 = rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) error); ok {
+		r0 = returnFunc(ctx)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -236,31 +283,23 @@ func (_e *StateMachine_Expecter) Start(ctx interface{}) *StateMachine_Start_Call
 
 func (_c *StateMachine_Start_Call) Run(run func(ctx context.Context)) *StateMachine_Start_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *StateMachine_Start_Call) Return(_a0 error) *StateMachine_Start_Call {
-	_c.Call.Return(_a0)
+func (_c *StateMachine_Start_Call) Return(err error) *StateMachine_Start_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *StateMachine_Start_Call) RunAndReturn(run func(context.Context) error) *StateMachine_Start_Call {
+func (_c *StateMachine_Start_Call) RunAndReturn(run func(ctx context.Context) error) *StateMachine_Start_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewStateMachine creates a new instance of StateMachine. 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 NewStateMachine(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *StateMachine {
-	mock := &StateMachine{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

cli/attestation.go

@@ -11,7 +11,7 @@ import (
 	"strings"
 	"time"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/fatih/color"
 	"github.com/google/go-sev-guest/abi"
 	"github.com/google/go-sev-guest/proto/sevsnp"
@@ -148,25 +148,24 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 				attType = attestation.SNPvTPM
 			case AzureToken:
 				cmd.Println("Fetching Azure token")
-				attType = attestation.AzureToken
 			case TDX:
 				cmd.Println("Fetching TDX attestation report")
 				attType = attestation.TDX
 			}
 
-			if (attType == attestation.VTPM || attType == attestation.SNPvTPM) && len(nonce) == 0 {
+			if (attestationType == VTPM || attestationType == SNPvTPM) && len(nonce) == 0 {
 				msg := color.New(color.FgRed).Sprint("vTPM nonce must be defined for vTPM attestation ❌ ")
 				cmd.Println(msg)
 				return
 			}
 
-			if (attType == attestation.SNP || attType == attestation.SNPvTPM) && len(teeNonce) == 0 {
+			if (attestationType == SNP || attestationType == SNPvTPM) && len(teeNonce) == 0 {
 				msg := color.New(color.FgRed).Sprint("TEE nonce must be defined for SEV-SNP attestation ❌ ")
 				cmd.Println(msg)
 				return
 			}
 
-			if (attType == attestation.AzureToken) && len(tokenNonce) == 0 {
+			if (attestationType == AzureToken) && len(tokenNonce) == 0 {
 				msg := color.New(color.FgRed).Sprint("Token nonce must be defined for Azure attestation ❌ ")
 				cmd.Println(msg)
 				return
@@ -184,13 +183,13 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 			}
 
 			var fixedVtpmNonceByte [vtpm.Nonce]byte
-			if attType != attestation.SNP {
+			if attType != attestation.SNP || attestationType == AzureToken {
 				if (len(nonce) > vtpm.Nonce) || (len(tokenNonce) > 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
 				}
-				if attType == attestation.AzureToken {
+				if attestationType == AzureToken {
 					copy(fixedVtpmNonceByte[:], tokenNonce)
 				} else {
 					copy(fixedVtpmNonceByte[:], nonce)
@@ -199,7 +198,7 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 
 			filename := attestationFilePath
 
-			if attType == attestation.AzureToken {
+			if attestationType == AzureToken {
 				filename = azureAttestResultFilePath
 			}
 
@@ -217,10 +216,10 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 
 			var returnJsonAzureToken bool
 
-			if attType == attestation.AzureToken {
-				err := cli.agentSDK.AttestationResult(cmd.Context(), fixedVtpmNonceByte, int(attType), attestationFile)
+			if attestationType == AzureToken {
+				err := cli.agentSDK.AttestationToken(cmd.Context(), fixedVtpmNonceByte, int(attType), attestationFile)
 				if err != nil {
-					printError(cmd, "Failed to get attestation result due to error: %v ❌", err)
+					printError(cmd, "Failed to get attestation token due to error: %v ❌", err)
 					return
 				}
 				returnJsonAzureToken = !getAzureTokenJWT
@@ -246,7 +245,7 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 
 				switch attestationType {
 				case SNP:
-					result, err = attesationToJSON(result)
+					result, err = attestationToJSON(result)
 					if err != nil {
 						printError(cmd, "Error converting SNP attestation to JSON: %v ❌", err)
 						return
@@ -279,7 +278,7 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 				}
 			}
 
-			cmd.Println("Attestation result retrieved and saved successfully!")
+			cmd.Println("Attestation retrieved and saved successfully!")
 		},
 	}
 
@@ -292,7 +291,7 @@ func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
 	return cmd
 }
 
-func attesationToJSON(report []byte) ([]byte, error) {
+func attestationToJSON(report []byte) ([]byte, error) {
 	if len(report) < abi.ReportSize {
 		return nil, errors.Wrap(errReportSize, fmt.Errorf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes", len(report), abi.ReportSize))
 	}
@@ -304,7 +303,7 @@ func attesationToJSON(report []byte) ([]byte, error) {
 	return json.MarshalIndent(attestationPB, "", "	")
 }
 
-func attesationFromJSON(reportFile []byte) ([]byte, error) {
+func attestationFromJSON(reportFile []byte) ([]byte, error) {
 	var attestationPB sevsnp.Attestation
 	if err := json.Unmarshal(reportFile, &attestationPB); err != nil {
 		return nil, err
@@ -552,7 +551,7 @@ func decodeJWTToJSON(tokenBytes []byte) ([]byte, error) {
 		return nil, fmt.Errorf("invalid JWT: must have at least 2 parts")
 	}
 
-	decode := func(seg string) (map[string]interface{}, error) {
+	decode := func(seg string) (map[string]any, error) {
 		// Add padding if missing
 		if m := len(seg) % 4; m != 0 {
 			seg += strings.Repeat("=", 4-m)
@@ -563,7 +562,7 @@ func decodeJWTToJSON(tokenBytes []byte) ([]byte, error) {
 			return nil, err
 		}
 
-		var result map[string]interface{}
+		var result map[string]any
 		if err := json.Unmarshal(data, &result); err != nil {
 			return nil, err
 		}
@@ -581,7 +580,7 @@ func decodeJWTToJSON(tokenBytes []byte) ([]byte, error) {
 		return nil, fmt.Errorf("failed to decode payload: %v", err)
 	}
 
-	combined := map[string]interface{}{
+	combined := map[string]any{
 		"header":  header,
 		"payload": payload,
 	}

cli/attestation_policy.go

@@ -6,12 +6,14 @@ import (
 	"bytes"
 	"crypto/sha512"
 	"encoding/base64"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"os"
 	"strconv"
+	"strings"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/google/go-sev-guest/proto/check"
 	"github.com/google/go-tpm-tools/proto/attest"
 	"github.com/spf13/cobra"
@@ -20,6 +22,7 @@ import (
 	"github.com/ultravioletrs/cocos/pkg/attestation/azure"
 	"github.com/ultravioletrs/cocos/pkg/attestation/gcp"
 	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
+	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
 )
 
@@ -46,7 +49,10 @@ var (
 	errMarshalJSON                         = errors.New("failed to marshal json")
 	errWriteFile                           = errors.New("failed to write to file")
 	errAttestationPolicyField              = errors.New("the specified field type does not exist in the attestation policy")
+	errReadingManifestFile                 = errors.New("error while reading manifest file")
+	errDecodeHex                           = errors.New("error decoding hex string")
 	policy                          uint64 = 196639
+	isJsonAttestation               bool
 )
 
 func (cli *CLI) NewAttestationPolicyCmd() *cobra.Command {
@@ -110,7 +116,7 @@ func (cli *CLI) NewAddHostDataCmd() *cobra.Command {
 }
 
 func (cli *CLI) NewGCPAttestationPolicy() *cobra.Command {
-	return &cobra.Command{
+	cmd := &cobra.Command{
 		Use:     "gcp",
 		Short:   "Get attestation policy for GCP CVM",
 		Example: `gcp <bin_vtmp_attestation_report_file> <vcpu_count>`,
@@ -130,9 +136,16 @@ func (cli *CLI) NewGCPAttestationPolicy() *cobra.Command {
 
 			attestation := &attest.Attestation{}
 
-			if err := proto.Unmarshal(attestationBin, attestation); err != nil {
-				printError(cmd, "Error unmarshaling attestation report: %v ❌ ", err)
-				return
+			if isJsonAttestation {
+				if err := protojson.Unmarshal(attestationBin, attestation); err != nil {
+					printError(cmd, "Error converting JSON attestation to binary: %v ❌", err)
+					return
+				}
+			} else {
+				if err := proto.Unmarshal(attestationBin, attestation); err != nil {
+					printError(cmd, "Error unmarshaling attestation report: %v ❌ ", err)
+					return
+				}
 			}
 
 			attestationPB := attestation.GetSevSnpAttestation()
@@ -169,10 +182,13 @@ func (cli *CLI) NewGCPAttestationPolicy() *cobra.Command {
 			cmd.Println("Attestation policy file generated successfully ✅")
 		},
 	}
+
+	cmd.Flags().BoolVarP(&isJsonAttestation, "json", "j", false, "Use JSON attestation report instead of binary")
+	return cmd
 }
 
 func (cli *CLI) NewDownloadGCPOvmfFile() *cobra.Command {
-	return &cobra.Command{
+	cmd := &cobra.Command{
 		Use:     "download",
 		Short:   "Download GCP OVMF file",
 		Example: `download <bin_vtmp_attestation_report_file>`,
@@ -186,9 +202,16 @@ func (cli *CLI) NewDownloadGCPOvmfFile() *cobra.Command {
 
 			attestation := &attest.Attestation{}
 
-			if err := proto.Unmarshal(attestationBin, attestation); err != nil {
-				printError(cmd, "Error unmarshaling attestation report: %v ❌ ", err)
-				return
+			if isJsonAttestation {
+				if err := protojson.Unmarshal(attestationBin, attestation); err != nil {
+					printError(cmd, "Error converting JSON attestation to binary: %v ❌", err)
+					return
+				}
+			} else {
+				if err := proto.Unmarshal(attestationBin, attestation); err != nil {
+					printError(cmd, "Error unmarshaling attestation report: %v ❌ ", err)
+					return
+				}
 			}
 
 			attestationPB := attestation.GetSevSnpAttestation()
@@ -227,6 +250,9 @@ func (cli *CLI) NewDownloadGCPOvmfFile() *cobra.Command {
 			cmd.Println("OVMF file downloaded successfully ✅")
 		},
 	}
+
+	cmd.Flags().BoolVarP(&isJsonAttestation, "json", "j", false, "Use JSON attestation report instead of binary")
+	return cmd
 }
 
 func (cli *CLI) NewAzureAttestationPolicy() *cobra.Command {
@@ -275,6 +301,23 @@ func (cli *CLI) NewAzureAttestationPolicy() *cobra.Command {
 	return cmd
 }
 
+func (cli *CLI) NewExtendWithManifestCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:     "extend",
+		Short:   "Extends PCR16 with computation manifests. The first parameter is path to attestation policy file. The rest of the parameters are paths to computation manifest files.",
+		Example: "extend <attestation_policy_file_path> <computation_manifest_file_path> [<computation_manifest_file_path> ...]",
+		Args:    cobra.MinimumNArgs(2),
+		Run: func(cmd *cobra.Command, args []string) {
+			attestationPolicyFilePath := args[0]
+			manifestPaths := args[1:]
+			if err := extendWithManifest(attestationPolicyFilePath, manifestPaths); err != nil {
+				printError(cmd, "Error could not extend PCR16: %v ❌ ", err)
+				return
+			}
+		},
+	}
+}
+
 func changeAttestationConfiguration(fileName, base64Data string, expectedLength int, field fieldType) error {
 	data, err := base64.StdEncoding.DecodeString(base64Data)
 	if err != nil {
@@ -318,3 +361,65 @@ func changeAttestationConfiguration(fileName, base64Data string, expectedLength
 	}
 	return nil
 }
+
+func extendWithManifest(attestationPolicyPath string, manifestPaths []string) error {
+	attestationConfig := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
+
+	attestationPolicyFileData, err := os.ReadFile(attestationPolicyPath)
+	if err != nil {
+		return errors.Wrap(errReadingAttestationPolicyFile, err)
+	}
+
+	if err = vtpm.ReadPolicyFromByte(attestationPolicyFileData, &attestationConfig); err != nil {
+		return errors.Wrap(errUnmarshalJSON, err)
+	}
+
+	for _, manifestPath := range manifestPaths {
+		manifest, err := os.ReadFile(manifestPath)
+		if err != nil {
+			return errors.Wrap(errReadingManifestFile, err)
+		}
+
+		manifestSha256 := sha512.Sum512_256(manifest)
+		manifestSha384 := sha512.Sum384(manifest)
+
+		data256, exists256 := attestationConfig.PCRValues.Sha256["16"]
+
+		if !exists256 {
+			data256 = strings.Repeat("0", 64) // 32 bytes in hex
+		}
+
+		byteData256, err := hex.DecodeString(data256)
+		if err != nil {
+			return errors.Wrap(errDecodeHex, err)
+		}
+
+		newByteData256 := sha512.Sum512_256(append(byteData256, manifestSha256[:]...))
+
+		data384, exists384 := attestationConfig.PCRValues.Sha384["16"]
+
+		if !exists384 {
+			data384 = strings.Repeat("0", 96) // 48 bytes in hex
+		}
+
+		byteData384, err := hex.DecodeString(data384)
+		if err != nil {
+			return errors.Wrap(errDecodeHex, err)
+		}
+
+		newByteData384 := sha512.Sum384(append(byteData384, manifestSha384[:]...))
+
+		attestationConfig.PCRValues.Sha256["16"] = hex.EncodeToString(newByteData256[:])
+		attestationConfig.PCRValues.Sha384["16"] = hex.EncodeToString(newByteData384[:])
+	}
+
+	attestationPolicyJSON, err := vtpm.ConvertPolicyToJSON(&attestationConfig)
+	if err != nil {
+		return errors.Wrap(errMarshalJSON, err)
+	}
+	if err = os.WriteFile(attestationPolicyPath, attestationPolicyJSON, filePermission); err != nil {
+		return errors.Wrap(errWriteFile, err)
+	}
+
+	return nil
+}

cli/attestation_policy_test.go

@@ -5,7 +5,6 @@ package cli
 import (
 	"bytes"
 	"encoding/base64"
-	"encoding/json"
 	"os"
 	"testing"
 
@@ -23,7 +22,7 @@ func TestChangeAttestationConfiguration(t *testing.T) {
 
 	initialConfig := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
 
-	initialJSON, err := json.Marshal(initialConfig)
+	initialJSON, err := vtpm.ConvertPolicyToJSON(&initialConfig)
 	require.NoError(t, err)
 	err = os.WriteFile(tmpfile.Name(), initialJSON, 0o644)
 	require.NoError(t, err)
@@ -376,3 +375,96 @@ func TestCommandErrorHandling(t *testing.T) {
 		assert.Contains(t, output, "❌")
 	})
 }
+
+func TestExtendWithManifestHandling(t *testing.T) {
+	cli := &CLI{}
+
+	t.Run("Invalid policy file", func(t *testing.T) {
+		cmd := cli.NewExtendWithManifestCmd()
+		cmd.SetArgs([]string{"nonexistent.policy.json", "nonexistent.manifest.json"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "error while reading the attestation policy file")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Invalid manifest file", func(t *testing.T) {
+		cmd := cli.NewExtendWithManifestCmd()
+		cmd.SetArgs([]string{"../scripts/attestation_policy/attestation_policy.json", "nonexistent.manifest.json"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "error while reading manifest file")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Valid file paths", func(t *testing.T) {
+		fileContent := `{
+		  "id": "1",
+		  "name": "sample computation",
+		  "description": "sample description",
+		  "datasets": [
+		    {
+		      "hash": "<sha3_encoded string>",
+		      "userKey": "<pem_encoded public key string>"
+		    }
+		  ],
+		  "algorithm": {
+		    "hash": "<sha3_encoded string>",
+		    "userKey": "<pem_encoded public key string>"
+		  },
+		  "result_consumers": [
+		    {
+		      "userKey": "<pem_encoded public key string>"
+		    }
+		  ],
+		  "agent_config": {
+		    "port": "7002",
+		    "cert_file": "<pem encoded cert string>",
+		    "key_file": "<pem encoded private key string>",
+		    "server_ca_file": "<pem encoded cert string>",
+		    "client_ca_file": "<pem encoded cert string>",
+		    "attested_tls": true
+		  }
+		}`
+
+		dir, err := os.Getwd()
+		if err != nil {
+			t.Fatalf("Error getting current working directory: %v", err)
+		}
+
+		manifestFile, err := os.CreateTemp(dir, "manifest.json")
+		if err != nil {
+			t.Fatalf("Error creating temp file: %v", err)
+		}
+		defer os.Remove(manifestFile.Name())
+
+		err = os.WriteFile(manifestFile.Name(), []byte(fileContent), 0o644)
+		if err != nil {
+			t.Fatalf("Error writing temp file: %v", err)
+		}
+
+		cmd := cli.NewExtendWithManifestCmd()
+		cmd.SetArgs([]string{"../scripts/attestation_policy/attestation_policy.json", manifestFile.Name()})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err = cmd.Execute()
+		assert.NoError(t, err)
+	})
+}

cli/attestation_snp.go

@@ -473,7 +473,7 @@ func parseAttestationFile() error {
 	}
 	attestationRaw = file
 	if isFileJSON(attestationFile) {
-		attestationRaw, err = attesationFromJSON(attestationRaw)
+		attestationRaw, err = attestationFromJSON(attestationRaw)
 		if err != nil {
 			return err
 		}

cli/attestation_snp_test.go

@@ -369,15 +369,15 @@ func TestGetBase(t *testing.T) {
 func TestParseConfig(t *testing.T) {
 	tempDir := t.TempDir()
 
-	validConfig := map[string]interface{}{
-		"rootOfTrust": map[string]interface{}{
+	validConfig := map[string]any{
+		"rootOfTrust": map[string]any{
 			"product":         "test_product",
 			"cabundlePaths":   []string{"test_path"},
 			"cabundles":       []string{"test_bundle"},
 			"checkCrl":        true,
 			"disallowNetwork": true,
 		},
-		"policy": map[string]interface{}{
+		"policy": map[string]any{
 			"minimumGuestSvn":  1,
 			"policy":           "1",
 			"minimumBuild":     1,

cli/attestation_tdx.go

@@ -9,7 +9,7 @@ import (
 	"os"
 	"strings"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	ccpb "github.com/google/go-tdx-guest/proto/checkconfig"
 	"github.com/spf13/cobra"
 	"github.com/ultravioletrs/cocos/pkg/attestation"

cli/attestation_test.go

@@ -12,7 +12,7 @@ import (
 	"path/filepath"
 	"testing"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/google/go-sev-guest/abi"
 	"github.com/google/go-sev-guest/proto/check"
 	"github.com/google/go-sev-guest/proto/sevsnp"
@@ -67,21 +67,21 @@ func TestNewGetAttestationCmd(t *testing.T) {
 			args:         []string{"snp", "--tee", teeNonce},
 			mockResponse: []byte("mock attestation"),
 			mockError:    nil,
-			expectedOut:  "Attestation result retrieved and saved successfully!",
+			expectedOut:  "Attestation retrieved and saved successfully!",
 		},
 		{
 			name:         "successful vTPM attestation retrieval",
 			args:         []string{"vtpm", "--vtpm", vtpmNonce},
 			mockResponse: []byte("mock attestation"),
 			mockError:    nil,
-			expectedOut:  "Attestation result retrieved and saved successfully!",
+			expectedOut:  "Attestation 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!",
+			expectedOut:  "Attestation retrieved and saved successfully!",
 		},
 		{
 			name:         "missing vTPM nonce",
@@ -137,7 +137,7 @@ func TestNewGetAttestationCmd(t *testing.T) {
 			args:         []string{"snp", "--tee", teeNonce, "--reporttextproto"},
 			mockResponse: validattestation,
 			mockError:    nil,
-			expectedOut:  "Attestation result retrieved and saved successfully!",
+			expectedOut:  "Attestation retrieved and saved successfully!",
 		},
 		{
 			name:         "connection error",
@@ -151,14 +151,14 @@ func TestNewGetAttestationCmd(t *testing.T) {
 			args:         []string{"azure-token", "--token", tokenNonce},
 			mockResponse: []byte("eyJhbGciOiAiUlMyNTYifQ.eyJzdWIiOiAidGVzdC11c2VyIn0.signature"),
 			mockError:    nil,
-			expectedOut:  "Fetching Azure token\nAttestation result retrieved and saved successfully!\n",
+			expectedOut:  "Fetching Azure token\nAttestation retrieved and saved successfully!\n",
 		},
 		{
 			name:         "failed to retrieve Azure token",
 			args:         []string{"azure-token", "--token", tokenNonce},
 			mockResponse: nil,
 			mockError:    errors.New("error"),
-			expectedErr:  "Fetching Azure token\nFailed to get attestation result due to error: error ❌\n",
+			expectedErr:  "Fetching Azure token\nFailed to get attestation token due to error: error ❌\n",
 		},
 		{
 			name:         "invalid token nonce size",
@@ -189,7 +189,7 @@ func TestNewGetAttestationCmd(t *testing.T) {
 				require.NoError(t, err)
 			})
 
-			mockSDK.On("AttestationResult", mock.Anything, [vtpm.Nonce]byte(bytes.Repeat([]byte{0x00}, vtpm.Nonce)), mock.Anything, mock.Anything).Return(tc.mockError).Run(func(args mock.Arguments) {
+			mockSDK.On("AttestationToken", mock.Anything, [vtpm.Nonce]byte(bytes.Repeat([]byte{0x00}, vtpm.Nonce)), mock.Anything, mock.Anything).Return(tc.mockError).Run(func(args mock.Arguments) {
 				_, err := args.Get(3).(*os.File).Write(tc.mockResponse)
 				require.NoError(t, err)
 			})
@@ -517,7 +517,7 @@ func TestAttestationToJSON(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := attesationToJSON(tt.input)
+			got, err := attestationToJSON(tt.input)
 			assert.True(t, errors.Contains(err, tt.err))
 			if tt.err != nil {
 				assert.Nil(t, got)
@@ -526,7 +526,7 @@ func TestAttestationToJSON(t *testing.T) {
 
 			require.NotNil(t, got)
 
-			var js map[string]interface{}
+			var js map[string]any
 			err = json.Unmarshal(got, &js)
 			assert.NoError(t, err)
 		})
@@ -588,7 +588,7 @@ func TestAttestationFromJSON(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got, err := attesationFromJSON(tt.input)
+			got, err := attestationFromJSON(tt.input)
 			assert.True(t, errors.Contains(err, tt.err))
 			tt.validate(t, got)
 		})
@@ -644,11 +644,11 @@ func TestIsFileJSON(t *testing.T) {
 func TestRoundTrip(t *testing.T) {
 	originalReport, err := os.ReadFile("../attestation.bin")
 	require.NoError(t, err)
-	jsonData, err := attesationToJSON(originalReport)
+	jsonData, err := attestationToJSON(originalReport)
 	require.NoError(t, err)
 	require.NotNil(t, jsonData)
 
-	roundTripReport, err := attesationFromJSON(jsonData)
+	roundTripReport, err := attestationFromJSON(jsonData)
 	require.NoError(t, err)
 	require.NotNil(t, roundTripReport)
 }

cli/datasets.go

@@ -9,7 +9,7 @@ import (
 	"os"
 	"path"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 	"github.com/ultravioletrs/cocos/agent"
@@ -44,6 +44,7 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
 			var dataset *os.File
 
 			if f.IsDir() {
+				cmd.Println("Detected directory, zipping dataset...")
 				dataset, err = internal.ZipDirectoryToTempFile(datasetPath)
 				if err != nil {
 					printError(cmd, "Error zipping dataset directory: %v ❌ ", err)
@@ -88,7 +89,7 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
 	return cmd
 }
 
-func decodeKey(b *pem.Block) (interface{}, error) {
+func decodeKey(b *pem.Block) (any, error) {
 	if b == nil {
 		return nil, errors.New("error decoding key")
 	}

cli/errors.go

@@ -3,7 +3,7 @@
 package cli
 
 import (
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 	"github.com/ultravioletrs/cocos/agent/auth"

cli/errors_test.go

@@ -7,7 +7,7 @@ import (
 	"errors"
 	"testing"
 
-	mgerrors "github.com/absmach/magistrala/pkg/errors"
+	mgerrors "github.com/absmach/supermq/pkg/errors"
 	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 	"github.com/ultravioletrs/cocos/agent/auth"

cli/keys.go

@@ -96,7 +96,7 @@ func (cli *CLI) NewKeysCmd() *cobra.Command {
 	}
 }
 
-func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType string) error {
+func generateAndWriteKeys(privKey any, pubKeyBytes []byte, keyType string) error {
 	privFile, err := os.Create(privateKeyFile)
 	if err != nil {
 		return err

cli/keys_test.go

@@ -57,7 +57,7 @@ func TestGenerateAndWriteKeys(t *testing.T) {
 				t.Fatalf("Failed to decode private key PEM")
 			}
 
-			var privKey interface{}
+			var privKey any
 			switch tt.keyType {
 			case "rsa":
 				privKey, err = x509.ParsePKCS1PrivateKey(privPem.Bytes)

cli/manager.go

@@ -95,15 +95,12 @@ func (c *CLI) NewRemoveVMCmd() *cobra.Command {
 		Args:    cobra.ExactArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
 			if c.managerClient == nil || c.connectErr != nil {
-				if err := c.InitializeManagerClient(cmd); err == nil {
-					defer c.Close()
+				if err := c.InitializeManagerClient(cmd); err != nil {
+					printError(cmd, "Failed to connect to manager: %v ❌ ", err)
+					return
 				}
 			}
-
-			if c.connectErr != nil {
-				printError(cmd, "Failed to connect to manager: %v ❌ ", c.connectErr)
-				return
-			}
+			defer c.Close()
 
 			cmd.Println("🔗 Removing virtual machine")
 

cli/manager_test.go

@@ -113,7 +113,7 @@ func TestCLI_NewCreateVMCmd(t *testing.T) {
 			flags: map[string]string{
 				"server-url": "https://server.com",
 			},
-			expectedError: "Failed to connect to manager: failed to connect to grpc server : failed to exit idle mode: passthrough: received empty target in Build() ❌",
+			expectedError: "failed to exit idle mode: dns resolver: missing address ❌",
 			expectError:   true,
 		},
 		{
@@ -252,7 +252,7 @@ func TestCLI_NewRemoveVMCmd(t *testing.T) {
 				cli.connectErr = errors.New("connection failed")
 			},
 			args:          []string{"vm-123"},
-			expectedError: "Failed to connect to manager: failed to connect to grpc server : failed to exit idle mode: passthrough: received empty target in Build() ❌",
+			expectedError: "failed to exit idle mode: dns resolver: missing address ❌",
 			expectError:   true,
 		},
 		{

cli/sdk.go

@@ -8,6 +8,7 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/ultravioletrs/cocos/manager"
 	"github.com/ultravioletrs/cocos/pkg/attestation/cmdconfig"
+	"github.com/ultravioletrs/cocos/pkg/clients"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
 	managergrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/manager"
@@ -18,15 +19,15 @@ var Verbose bool
 
 type CLI struct {
 	agentSDK      sdk.SDK
-	agentConfig   grpc.AgentClientConfig
-	managerConfig grpc.ManagerClientConfig
+	agentConfig   clients.AttestedClientConfig
+	managerConfig clients.StandardClientConfig
 	client        grpc.Client
 	managerClient manager.ManagerServiceClient
 	connectErr    error
 	measurement   cmdconfig.MeasurementProvider
 }
 
-func New(agentConfig grpc.AgentClientConfig, managerConfig grpc.ManagerClientConfig, measurement cmdconfig.MeasurementProvider) *CLI {
+func New(agentConfig clients.AttestedClientConfig, managerConfig clients.StandardClientConfig, measurement cmdconfig.MeasurementProvider) *CLI {
 	return &CLI{
 		agentConfig:   agentConfig,
 		managerConfig: managerConfig,

cmd/agent/main.go

@@ -14,10 +14,10 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
-	"time"
 
-	mglog "github.com/absmach/magistrala/logger"
-	"github.com/absmach/magistrala/pkg/prometheus"
+	"github.com/absmach/certs/sdk"
+	mglog "github.com/absmach/supermq/logger"
+	"github.com/absmach/supermq/pkg/prometheus"
 	"github.com/caarlos0/env/v11"
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/agent/api"
@@ -26,10 +26,13 @@ import (
 	"github.com/ultravioletrs/cocos/agent/cvms/server"
 	"github.com/ultravioletrs/cocos/agent/events"
 	agentlogger "github.com/ultravioletrs/cocos/internal/logger"
+	"github.com/ultravioletrs/cocos/pkg/atls"
 	"github.com/ultravioletrs/cocos/pkg/attestation"
 	"github.com/ultravioletrs/cocos/pkg/attestation/azure"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
 	"github.com/ultravioletrs/cocos/pkg/attestation/tdx"
 	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
+	"github.com/ultravioletrs/cocos/pkg/clients"
 	pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	cvmsgrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/cvm"
 	"golang.org/x/sync/errgroup"
@@ -37,22 +40,21 @@ import (
 
 const (
 	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"`
-	CAUrl         string `env:"AGENT_CVM_CA_URL" envDefault:""`
-	CVMId         string `env:"AGENT_CVM_ID"     envDefault:""`
-	AgentMaaURL   string `env:"AGENT_MAA_URL"    envDefault:"https://sharedeus2.eus2.attest.azure.net"`
-	AgentOSBuild  string `env:"AGENT_OS_BUILD"   envDefault:"UVC"`
-	AgentOSDistro string `env:"AGENT_OS_DISTRO"  envDefault:"UVC"`
-	AgentOSType   string `env:"AGENT_OS_TYPE"    envDefault:"UVC"`
+	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"`
+	CAUrl         string `env:"AGENT_CVM_CA_URL"  envDefault:""`
+	CVMId         string `env:"AGENT_CVM_ID"      envDefault:""`
+	CertsToken    string `env:"AGENT_CERTS_TOKEN" envDefault:""`
+	AgentMaaURL   string `env:"AGENT_MAA_URL"     envDefault:"https://sharedeus2.eus2.attest.azure.net"`
+	AgentOSBuild  string `env:"AGENT_OS_BUILD"    envDefault:"UVC"`
+	AgentOSDistro string `env:"AGENT_OS_DISTRO"   envDefault:"UVC"`
+	AgentOSType   string `env:"AGENT_OS_TYPE"     envDefault:"UVC"`
 }
 
 func main() {
@@ -111,7 +113,7 @@ func main() {
 		provider = &attestation.EmptyProvider{}
 	}
 
-	cvmGrpcConfig := pkggrpc.CVMClientConfig{}
+	cvmGrpcConfig := clients.StandardClientConfig{}
 	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
@@ -154,6 +156,13 @@ func main() {
 		return
 	}
 
+	err = quoteprovider.FetchCertificates(uint(cfg.Vmpl))
+	if err != nil {
+		logger.Error(fmt.Sprintf("failed to fetch certificates: %s", err))
+		exitCode = 1
+		return
+	}
+
 	svc := newService(ctx, logger, eventSvc, provider, cfg.Vmpl)
 
 	if err := os.MkdirAll(storageDir, 0o755); err != nil {
@@ -162,7 +171,24 @@ func main() {
 		return
 	}
 
-	mc, err := cvmsapi.NewClient(pc, svc, eventsLogsQueue, logger, server.NewServer(logger, svc, cfg.AgentGrpcHost, cfg.CAUrl, cfg.CVMId), storageDir, reconnectFn, cvmGRPCClient)
+	var certProvider atls.CertificateProvider
+
+	if ccPlatform != attestation.NoCC {
+		var certsSDK sdk.SDK
+		if cfg.CAUrl != "" {
+			certsSDK = sdk.NewSDK(sdk.Config{
+				CertsURL: cfg.CAUrl,
+			})
+		}
+		certProvider, err = atls.NewProvider(provider, ccPlatform, cfg.CertsToken, cfg.CVMId, certsSDK)
+		if err != nil {
+			logger.Error(fmt.Sprintf("failed to create certificate provider: %s", err))
+			exitCode = 1
+			return
+		}
+	}
+
+	mc, err := cvmsapi.NewClient(pc, svc, eventsLogsQueue, logger, server.NewServer(logger, svc, cfg.AgentGrpcHost, certProvider), storageDir, reconnectFn, cvmGRPCClient)
 	if err != nil {
 		logger.Error(err.Error())
 		exitCode = 1
@@ -196,16 +222,16 @@ func main() {
 	}
 
 	if ccPlatform == attestation.Azure {
-		azureAttestationResult, azureCertSerialNumber, err := azureAttestationFromCert(ctx, cvmGrpcConfig.ClientCert, svc)
+		azureAttestationToken, azureCertSerialNumber, err := azureAttestationFromCert(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_AzureAttestationResult{
-				AzureAttestationResult: &cvms.AzureAttestationResponse{
-					File:             azureAttestationResult,
+			Message: &cvms.ClientStreamMessage_AzureAttestationToken{
+				AzureAttestationToken: &cvms.AzureAttestationToken{
+					File:             azureAttestationToken,
 					CertSerialNumber: azureCertSerialNumber,
 				},
 			},
@@ -279,7 +305,7 @@ func azureAttestationFromCert(ctx context.Context, certFilePath string, svc agen
 	}
 
 	nonceAzure := sha256.Sum256(certFile)
-	attestation, err := svc.AttestationResult(ctx, nonceAzure, attestation.AzureToken)
+	attestation, err := svc.AzureAttestationToken(ctx, nonceAzure)
 	if err != nil {
 		return nil, "", err
 	}

cmd/cli/main.go

@@ -15,7 +15,7 @@ import (
 	"github.com/spf13/pflag"
 	"github.com/ultravioletrs/cocos/cli"
 	"github.com/ultravioletrs/cocos/pkg/attestation/cmdconfig"
-	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
+	"github.com/ultravioletrs/cocos/pkg/clients"
 	cmd "github.com/virtee/sev-snp-measure-go/sevsnpmeasure/cmd"
 )
 
@@ -94,14 +94,14 @@ func main() {
 		return
 	}
 
-	agentGRPCConfig := grpc.AgentClientConfig{}
+	agentGRPCConfig := clients.AttestedClientConfig{}
 	if err := env.ParseWithOptions(&agentGRPCConfig, env.Options{Prefix: envPrefixAgentGRPC}); err != nil {
 		message := color.New(color.FgRed).Sprintf("failed to load %s gRPC client configuration : %s", svcName, err)
 		rootCmd.Println(message)
 		return
 	}
 
-	managerGRPCConfig := grpc.ManagerClientConfig{}
+	managerGRPCConfig := clients.StandardClientConfig{}
 	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)
@@ -164,6 +164,7 @@ func main() {
 	attestationPolicyCmd.AddCommand(cliSVC.NewGCPAttestationPolicy())
 	attestationPolicyCmd.AddCommand(cliSVC.NewDownloadGCPOvmfFile())
 	attestationPolicyCmd.AddCommand(cliSVC.NewAzureAttestationPolicy())
+	attestationPolicyCmd.AddCommand(cliSVC.NewExtendWithManifestCmd())
 
 	if err := rootCmd.Execute(); err != nil {
 		logErrorCmd(*rootCmd, err)

cmd/manager/main.go

@@ -12,18 +12,22 @@ import (
 	"os"
 	"strings"
 
-	mglog "github.com/absmach/magistrala/logger"
-	"github.com/absmach/magistrala/pkg/jaeger"
-	"github.com/absmach/magistrala/pkg/prometheus"
-	"github.com/absmach/magistrala/pkg/uuid"
+	mglog "github.com/absmach/supermq/logger"
+	"github.com/absmach/supermq/pkg/jaeger"
+	"github.com/absmach/supermq/pkg/prometheus"
+	smqserver "github.com/absmach/supermq/pkg/server"
+	httpserver "github.com/absmach/supermq/pkg/server/http"
+	"github.com/absmach/supermq/pkg/uuid"
 	"github.com/caarlos0/env/v11"
-	"github.com/ultravioletrs/cocos/internal/server"
-	grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
+	"github.com/go-chi/chi/v5"
 	"github.com/ultravioletrs/cocos/manager"
 	"github.com/ultravioletrs/cocos/manager/api"
 	managergrpc "github.com/ultravioletrs/cocos/manager/api/grpc"
+	"github.com/ultravioletrs/cocos/manager/api/http"
 	"github.com/ultravioletrs/cocos/manager/qemu"
 	"github.com/ultravioletrs/cocos/manager/tracing"
+	"github.com/ultravioletrs/cocos/pkg/server"
+	grpcserver "github.com/ultravioletrs/cocos/pkg/server/grpc"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
@@ -31,10 +35,11 @@ import (
 )
 
 const (
-	svcName          = "manager"
-	envPrefixGRPC    = "MANAGER_GRPC_"
-	envPrefixQemu    = "MANAGER_QEMU_"
-	clientBufferSize = 100
+	svcName        = "manager"
+	envPrefixGRPC  = "MANAGER_GRPC_"
+	envPrefixHTTP  = "MANAGER_HTTP_"
+	envPrefixQemu  = "MANAGER_QEMU_"
+	defSvcHTTPPort = "7003"
 )
 
 type config struct {
@@ -46,6 +51,7 @@ type config struct {
 	IgvmMeasureBinary       string  `env:"MANAGER_IGVMMEASURE_BINARY"         envDefault:"../../build/igvmmeasure"`
 	PcrValues               string  `env:"MANAGER_PCR_VALUES"                 envDefault:""`
 	EosVersion              string  `env:"MANAGER_EOS_VERSION"                envDefault:""`
+	MaxVMs                  int     `env:"MANAGER_MAX_VMS"                    envDefault:"10"`
 }
 
 func main() {
@@ -114,7 +120,12 @@ func main() {
 		return
 	}
 
-	svc, err := newService(logger, tracer, *qemuCfg, cfg.AttestationPolicyBinary, cfg.IgvmMeasureBinary, cfg.PcrValues, cfg.EosVersion)
+	httpServerConfig := smqserver.Config{Port: defSvcHTTPPort}
+	if err := env.ParseWithOptions(&httpServerConfig, env.Options{Prefix: envPrefixHTTP}); err != nil {
+		logger.Error(fmt.Sprintf("failed to load %s gRPC server configuration : %s", svcName, err))
+	}
+
+	svc, err := newService(logger, tracer, *qemuCfg, cfg.AttestationPolicyBinary, cfg.IgvmMeasureBinary, cfg.PcrValues, cfg.EosVersion, cfg.MaxVMs)
 	if err != nil {
 		logger.Error(err.Error())
 		exitCode = 1
@@ -134,14 +145,20 @@ func main() {
 		manager.RegisterManagerServiceServer(srv, managergrpc.NewServer(svc))
 	}
 
-	gs := grpcserver.New(ctx, cancel, svcName, managerGRPCConfig, registerManagerServiceServer, logger, nil, "", "")
+	gs := grpcserver.New(ctx, cancel, svcName, managerGRPCConfig, registerManagerServiceServer, logger, nil, nil)
+
+	hs := httpserver.NewServer(ctx, cancel, svcName, httpServerConfig, http.MakeHandler(chi.NewMux(), svcName, cfg.InstanceID), logger)
 
 	g.Go(func() error {
 		return gs.Start()
 	})
 
 	g.Go(func() error {
-		return server.StopHandler(ctx, cancel, logger, svcName, gs)
+		return hs.Start()
+	})
+
+	g.Go(func() error {
+		return server.StopHandler(ctx, cancel, logger, svcName, gs, hs)
 	})
 
 	if err := g.Wait(); err != nil {
@@ -149,8 +166,8 @@ func main() {
 	}
 }
 
-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)
+func newService(logger *slog.Logger, tracer trace.Tracer, qemuCfg qemu.Config, attestationPolicyPath string, igvmMeasurementBinaryPath string, pcrValuesFilePath string, eosVersion string, maxVMs int) (manager.Service, error) {
+	svc, err := manager.New(qemuCfg, attestationPolicyPath, igvmMeasurementBinaryPath, pcrValuesFilePath, logger, qemu.NewVM, eosVersion, maxVMs)
 	if err != nil {
 		return nil, err
 	}

cocos-manager.env

@@ -9,13 +9,21 @@ 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_CERT=
+MANAGER_GRPC_SERVER_KEY=
 MANAGER_GRPC_SERVER_CA_CERTS=
+MANAGER_GRPC_CLIENT_CA_CERTS=
 MANAGER_GRPC_PORT=6101
 MANAGER_GRPC_HOST=0.0.0.0
+MANAGER_HTTP_SERVER_CERT=
+MANAGER_HTTP_SERVER_KEY=
+MANAGER_HTTP_SERVER_CA_CERTS=
+MANAGER_HTTP_CLIENT_CA_CERTS=
+MANAGER_HTTP_PORT=6102
+MANAGER_HTTP_HOST=0.0.0.0
 MANAGER_GRPC_TIMEOUT=60s
 MANAGER_EOS_VERSION=""
+MANAGER_MAX_VMS=10
 
 # QEMU Configuration
 MANAGER_QEMU_MEMORY_SIZE=25G
@@ -47,9 +55,6 @@ MANAGER_QEMU_HOST_DATA=
 MANAGER_QEMU_TDX_ID=tdx0
 MANAGER_QEMU_QUOTE_GENERATION_PORT=4050
 MANAGER_QEMU_OVMF_FILE=/usr/share/ovmf/OVMF.fd
-MANAGER_QEMU_VSOCK_ID=vhost-vsock-pci0
-MANAGER_QEMU_VSOCK_GUEST_CID=3
-MANAGER_QEMU_VSOCK_VNC=0
 MANAGER_QEMU_BIN_PATH=qemu-system-x86_64
 MANAGER_QEMU_USE_SUDO=true
 MANAGER_QEMU_ENABLE_SEV_SNP=false

go.mod

@@ -1,50 +1,49 @@
 module github.com/ultravioletrs/cocos
 
-go 1.23.2
-
-toolchain go1.24.1
+go 1.25.0
 
 require (
-	github.com/absmach/magistrala v0.15.1
 	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.13.0
 	github.com/google/go-tdx-guest v0.3.2-0.20241009005452-097ee70d0843
-	github.com/spf13/cobra v1.9.1
-	github.com/spf13/pflag v1.0.6
-	github.com/stretchr/testify v1.10.0
+	github.com/spf13/cobra v1.10.1
+	github.com/spf13/pflag v1.0.10
+	github.com/stretchr/testify v1.11.1
 	github.com/virtee/sev-snp-measure-go v0.0.0-20240530153610-e6e8dc9b6877
-	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.62.0
-	go.opentelemetry.io/otel/trace v1.37.0
-	golang.org/x/crypto v0.39.0
-	golang.org/x/sync v0.15.0
-	google.golang.org/grpc v1.73.0
-	google.golang.org/protobuf v1.36.6
+	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.63.0
+	go.opentelemetry.io/otel/trace v1.38.0
+	golang.org/x/crypto v0.42.0
+	golang.org/x/sync v0.17.0
+	google.golang.org/grpc v1.75.1
+	google.golang.org/protobuf v1.36.10
 )
 
 require (
-	cloud.google.com/go/storage v1.55.0
+	cloud.google.com/go/storage v1.57.0
+	github.com/absmach/supermq v0.18.1
 	github.com/caarlos0/env/v10 v10.0.0
-	github.com/golang-jwt/jwt/v5 v5.2.2
+	github.com/go-chi/chi/v5 v5.2.3
+	github.com/golang-jwt/jwt/v5 v5.3.0
 	github.com/google/gce-tcb-verifier v0.3.1
 )
 
 require (
-	cel.dev/expr v0.23.0 // indirect
-	cloud.google.com/go v0.121.1 // indirect
-	cloud.google.com/go/auth v0.16.1 // indirect
+	cel.dev/expr v0.24.0 // indirect
+	cloud.google.com/go v0.121.6 // indirect
+	cloud.google.com/go/auth v0.16.5 // indirect
 	cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect
-	cloud.google.com/go/compute/metadata v0.7.0 // indirect
+	cloud.google.com/go/compute/metadata v0.8.0 // indirect
 	cloud.google.com/go/iam v1.5.2 // indirect
 	cloud.google.com/go/monitoring v1.24.2 // indirect
-	github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.27.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/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.29.0 // indirect
+	github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.53.0 // indirect
+	github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.53.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-20250326154945-ae57f3c0d45f // indirect
+	github.com/cenkalti/backoff/v5 v5.0.3 // indirect
+	github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 // indirect
 	github.com/containerd/errdefs v1.0.0 // indirect
 	github.com/containerd/errdefs/pkg v0.3.0 // indirect
 	github.com/containerd/log v0.1.0 // indirect
@@ -54,19 +53,17 @@ require (
 	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/go-jose/go-jose/v3 v3.0.3 // indirect
-	github.com/go-jose/go-jose/v4 v4.0.5 // indirect
-	github.com/gofrs/uuid/v5 v5.3.0 // indirect
-	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/go-jose/go-jose/v3 v3.0.4 // indirect
+	github.com/go-jose/go-jose/v4 v4.1.1 // indirect
+	github.com/gofrs/uuid/v5 v5.3.2 // indirect
 	github.com/google/certificate-transparency-go v1.1.8 // 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.6 // indirect
-	github.com/googleapis/gax-go/v2 v2.14.2 // indirect
-	github.com/hokaccha/go-prettyjson v0.0.0-20211117102719-0474bc63780f // indirect
-	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/googleapis/gax-go/v2 v2.15.0 // indirect
+	github.com/mattn/go-colorable v0.1.14 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/moby/docker-image-spec v1.3.1 // indirect
 	github.com/moby/sys/atomicwriter v0.1.0 // indirect
@@ -74,62 +71,61 @@ 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
-	github.com/pelletier/go-toml v1.9.5 // indirect
-	github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
+	github.com/planetscale/vtprotobuf v0.6.1-0.20240917153116-6f2963f01587 // indirect
 	github.com/spiffe/go-spiffe/v2 v2.5.0 // indirect
 	github.com/zeebo/errs v1.4.0 // indirect
 	go.opentelemetry.io/auto/sdk v1.1.0 // indirect
 	go.opentelemetry.io/contrib/detectors/gcp v1.36.0 // indirect
-	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 // indirect
-	go.opentelemetry.io/otel v1.37.0 // indirect
-	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.36.0 // indirect
-	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0 // indirect
-	go.opentelemetry.io/otel/sdk v1.37.0 // indirect
-	go.opentelemetry.io/otel/sdk/metric v1.37.0 // indirect
-	golang.org/x/exp v0.0.0-20240909161429-701f63a606c0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 // indirect
+	go.opentelemetry.io/otel v1.38.0 // indirect
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0 // indirect
+	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0 // indirect
+	go.opentelemetry.io/otel/sdk v1.38.0 // indirect
+	go.opentelemetry.io/otel/sdk/metric v1.38.0 // indirect
+	go.yaml.in/yaml/v2 v2.4.2 // indirect
+	golang.org/x/exp v0.0.0-20250606033433-dcc06ee1d476 // indirect
 	golang.org/x/oauth2 v0.30.0 // indirect
-	golang.org/x/time v0.11.0 // indirect
-	google.golang.org/api v0.235.0 // indirect
-	google.golang.org/genproto v0.0.0-20250505200425-f936aa4a68b2 // indirect
-	gopkg.in/yaml.v2 v2.4.0 // indirect
+	golang.org/x/time v0.12.0 // indirect
+	google.golang.org/api v0.247.0 // indirect
+	google.golang.org/genproto v0.0.0-20250603155806-513f23925822 // indirect
 	gotest.tools/v3 v3.5.1 // indirect
 	moul.io/http2curl v1.0.0 // indirect
 )
 
 require (
-	github.com/absmach/certs v0.0.0-20250707105817-4c41124bba54
+	github.com/absmach/certs v0.18.0
 	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 v28.2.2+incompatible
+	github.com/docker/docker v28.5.0+incompatible
 	github.com/edgelesssys/go-azguestattestation v0.0.0-20250408071817-8c4457b235ff
 	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.3 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/google/go-configfs-tsm v0.3.3-0.20240919001351-b4b5b84fdcbc // indirect
-	github.com/google/go-tpm v0.9.5
+	github.com/google/go-tpm v0.9.6
 	github.com/google/go-tpm-tools v0.4.4
 	github.com/google/logger v1.1.1
 	github.com/google/uuid v1.6.0
-	github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.3 // indirect
+	github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // 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.22.0 // indirect
-	github.com/prometheus/client_model v0.6.1 // indirect
-	github.com/prometheus/common v0.62.0 // indirect
-	github.com/prometheus/procfs v0.15.1 // indirect
+	github.com/prometheus/client_golang v1.23.2
+	github.com/prometheus/client_model v0.6.2 // indirect
+	github.com/prometheus/common v0.66.1 // indirect
+	github.com/prometheus/procfs v0.17.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
-	go.opentelemetry.io/otel/metric v1.37.0 // indirect
-	go.opentelemetry.io/proto/otlp v1.6.0 // indirect
+	go.opentelemetry.io/otel/metric v1.38.0 // indirect
+	go.opentelemetry.io/proto/otlp v1.7.1 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	golang.org/x/net v0.41.0 // indirect
-	golang.org/x/sys v0.33.0 // indirect
-	golang.org/x/term v0.32.0
-	golang.org/x/text v0.26.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20250512202823-5a2f75b736a9 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20250603155806-513f23925822 // indirect
+	golang.org/x/net v0.43.0 // indirect
+	golang.org/x/sys v0.36.0 // indirect
+	golang.org/x/term v0.35.0
+	golang.org/x/text v0.29.0 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20250825161204-c5933d9347a5 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20250825161204-c5933d9347a5 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 

go.sum

@@ -1,13 +1,13 @@
-cel.dev/expr v0.23.0 h1:wUb94w6OYQS4uXraxo9U+wUAs9jT47Xvl4iPgAwM2ss=
-cel.dev/expr v0.23.0/go.mod h1:hLPLo1W4QUmuYdA72RBX06QTs6MXw941piREPl3Yfiw=
-cloud.google.com/go v0.121.1 h1:S3kTQSydxmu1JfLRLpKtxRPA7rSrYPRPEUmL/PavVUw=
-cloud.google.com/go v0.121.1/go.mod h1:nRFlrHq39MNVWu+zESP2PosMWA0ryJw8KUBZ2iZpxbw=
-cloud.google.com/go/auth v0.16.1 h1:XrXauHMd30LhQYVRHLGvJiYeczweKQXZxsTbV9TiguU=
-cloud.google.com/go/auth v0.16.1/go.mod h1:1howDHJ5IETh/LwYs3ZxvlkXF48aSqqJUM+5o02dNOI=
+cel.dev/expr v0.24.0 h1:56OvJKSH3hDGL0ml5uSxZmz3/3Pq4tJ+fb1unVLAFcY=
+cel.dev/expr v0.24.0/go.mod h1:hLPLo1W4QUmuYdA72RBX06QTs6MXw941piREPl3Yfiw=
+cloud.google.com/go v0.121.6 h1:waZiuajrI28iAf40cWgycWNgaXPO06dupuS+sgibK6c=
+cloud.google.com/go v0.121.6/go.mod h1:coChdst4Ea5vUpiALcYKXEpR1S9ZgXbhEzzMcMR66vI=
+cloud.google.com/go/auth v0.16.5 h1:mFWNQ2FEVWAliEQWpAdH80omXFokmrnbDhUS9cBywsI=
+cloud.google.com/go/auth v0.16.5/go.mod h1:utzRfHMP+Vv0mpOkTRQoWD2q3BatTOoWbA7gCc2dUhQ=
 cloud.google.com/go/auth/oauth2adapt v0.2.8 h1:keo8NaayQZ6wimpNSmW5OPc283g65QNIiLpZnkHRbnc=
 cloud.google.com/go/auth/oauth2adapt v0.2.8/go.mod h1:XQ9y31RkqZCcwJWNSx2Xvric3RrU88hAYYbjDWYDL+c=
-cloud.google.com/go/compute/metadata v0.7.0 h1:PBWF+iiAerVNe8UCHxdOt6eHLVc3ydFeOCw78U8ytSU=
-cloud.google.com/go/compute/metadata v0.7.0/go.mod h1:j5MvL9PprKL39t166CoB1uVHfQMs4tFQZZcKwksXUjo=
+cloud.google.com/go/compute/metadata v0.8.0 h1:HxMRIbao8w17ZX6wBnjhcDkW6lTFpgcaobyVfZWqRLA=
+cloud.google.com/go/compute/metadata v0.8.0/go.mod h1:sYOGTp851OV9bOFJ9CH7elVvyzopvWQFNNghtDQ/Biw=
 cloud.google.com/go/iam v1.5.2 h1:qgFRAGEmd8z6dJ/qyEchAuL9jpswyODjA2lS+w234g8=
 cloud.google.com/go/iam v1.5.2/go.mod h1:SE1vg0N81zQqLzQEwxL2WI6yhetBdbNQuTvIKCSkUHE=
 cloud.google.com/go/logging v1.13.0 h1:7j0HgAp0B94o1YRDqiqm26w4q1rDMH7XNRU34lJXHYc=
@@ -16,40 +16,42 @@ cloud.google.com/go/longrunning v0.6.7 h1:IGtfDWHhQCgCjwQjV9iiLnUta9LBCo8R9QmAFs
 cloud.google.com/go/longrunning v0.6.7/go.mod h1:EAFV3IZAKmM56TyiE6VAP3VoTzhZzySwI/YI1s/nRsY=
 cloud.google.com/go/monitoring v1.24.2 h1:5OTsoJ1dXYIiMiuL+sYscLc9BumrL3CarVLL7dd7lHM=
 cloud.google.com/go/monitoring v1.24.2/go.mod h1:x7yzPWcgDRnPEv3sI+jJGBkwl5qINf+6qY4eq0I9B4U=
-cloud.google.com/go/storage v1.55.0 h1:NESjdAToN9u1tmhVqhXCaCwYBuvEhZLLv0gBr+2znf0=
-cloud.google.com/go/storage v1.55.0/go.mod h1:ztSmTTwzsdXe5syLVS0YsbFxXuvEmEyZj7v7zChEmuY=
+cloud.google.com/go/storage v1.57.0 h1:4g7NB7Ta7KetVbOMpCqy89C+Vg5VE8scqlSHUPm7Rds=
+cloud.google.com/go/storage v1.57.0/go.mod h1:329cwlpzALLgJuu8beyJ/uvQznDHpa2U5lGjWednkzg=
 cloud.google.com/go/trace v1.11.6 h1:2O2zjPzqPYAHrn3OKl029qlqG6W8ZdYaOWRyr8NgMT4=
 cloud.google.com/go/trace v1.11.6/go.mod h1:GA855OeDEBiBMzcckLPE2kDunIpC72N+Pq8WFieFjnI=
 github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 h1:L/gRVlceqvL25UVaW/CKtUDjefjrs0SPonmDGUVOYP0=
 github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
-github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.27.0 h1:ErKg/3iS1AKcTkf3yixlZ54f9U1rljCkQyEXWUnIUxc=
-github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.27.0/go.mod h1:yAZHSGnqScoU556rBOVkwLze6WP5N+U11RHuWaGVxwY=
-github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.51.0 h1:fYE9p3esPxA/C0rQ0AHhP0drtPXDRhaWiwg1DPqO7IU=
-github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.51.0/go.mod h1:BnBReJLvVYx2CS/UHOgVz2BXKXD9wsQPxZug20nZhd0=
-github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.51.0 h1:OqVGm6Ei3x5+yZmSJG1Mh2NwHvpVmZ08CB5qJhT9Nuk=
-github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.51.0/go.mod h1:SZiPHWGOOk3bl8tkevxkoiwPgsIl6CwrWcbwjfHZpdM=
-github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.51.0 h1:6/0iUd0xrnX7qt+mLNRwg5c0PGv8wpE8K90ryANQwMI=
-github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.51.0/go.mod h1:otE2jQekW/PqXk1Awf5lmfokJx4uwuqcj1ab5SpGeW0=
+github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.29.0 h1:UQUsRi8WTzhZntp5313l+CHIAT95ojUI2lpP/ExlZa4=
+github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.29.0/go.mod h1:Cz6ft6Dkn3Et6l2v2a9/RpN7epQ1GtDlO6lj8bEcOvw=
+github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.53.0 h1:owcC2UnmsZycprQ5RfRgjydWhuoxg71LUfyiQdijZuM=
+github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.53.0/go.mod h1:ZPpqegjbE99EPKsu3iUWV22A04wzGPcAY/ziSIQEEgs=
+github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.53.0 h1:4LP6hvB4I5ouTbGgWtixJhgED6xdf67twf9PoY96Tbg=
+github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.53.0/go.mod h1:jUZ5LYlw40WMd07qxcQJD5M40aUxrfwqQX1g7zxYnrQ=
+github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.53.0 h1:Ron4zCA/yk6U7WOBXhTJcDpsUBG9npumK6xw2auFltQ=
+github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.53.0/go.mod h1:cSgYe11MCNYunTnRXrKiR/tHc0eoKjICUuWpNZoVCOo=
 github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
 github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
 github.com/VividCortex/gohistogram v1.0.0 h1:6+hBz+qvs0JOrrNhhmR7lFxo5sINxBCGXrdtl/UvroE=
 github.com/VividCortex/gohistogram v1.0.0/go.mod h1:Pf5mBqqDxYaXu3hDrrU+w6nw50o/4+TcAqDqk/vUH7g=
-github.com/absmach/certs v0.0.0-20250707105817-4c41124bba54 h1:5pTDktMVvtJrQMY0d2vwBRFOgH0LVEoArwcaqrbuGQ4=
-github.com/absmach/certs v0.0.0-20250707105817-4c41124bba54/go.mod h1:MGDO7iEXbtJCRKOwL34+HR8MQI2mo7l3wcz93e2Km68=
-github.com/absmach/magistrala v0.15.1 h1:3Bk2hlyWcV591LxPYwlvRcyCXTfuZ1g/EkNmU+o3NNQ=
-github.com/absmach/magistrala v0.15.1/go.mod h1:9pto6xuBt/IuCtZRdEha0iDQKNQ5tyNOjLXJgUiikYk=
+github.com/absmach/certs v0.18.0 h1:4m1uFhFOmIoNyKIVyknaynwnME4OaMcrdYb9zeZzZ20=
+github.com/absmach/certs v0.18.0/go.mod h1:m6CWmAio930laR1TLZ3HhQna4f9KmSyyUGfD8EFZHwE=
+github.com/absmach/senml v1.0.8 h1:+opem/r4g6c6eA/JLyCIuksyEhj7eBdysY3pEmy1mqo=
+github.com/absmach/senml v1.0.8/go.mod h1:DRhzHLgvQoIUHroBgpFrSWso+bJZO9E96RlHAHy+VRI=
+github.com/absmach/supermq v0.18.1 h1:JRLP6rfSzZoHgRGPfwNSmzJ7a4K4b4Dvz2nCmR32rxI=
+github.com/absmach/supermq v0.18.1/go.mod h1:dYnFOIcGQzZ1WpYt1qNv1g609WmOWYzWCBBRjPQV7Uk=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/caarlos0/env/v10 v10.0.0 h1:yIHUBZGsyqCnpTkbjk8asUlx6RFhhEs+h7TOBdgdzXA=
 github.com/caarlos0/env/v10 v10.0.0/go.mod h1:ZfulV76NvVPw3tm591U4SwL3Xx9ldzBP9aGxzeN7G18=
 github.com/caarlos0/env/v11 v11.3.1 h1:cArPWC15hWmEt+gWk7YBi7lEXTXCvpaSdCiZE2X5mCA=
 github.com/caarlos0/env/v11 v11.3.1/go.mod h1:qupehSf/Y0TUTsxKywqRt/vJjN5nz6vauiYEUUr8P4U=
-github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
-github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
+github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM=
+github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
-github.com/cncf/xds/go v0.0.0-20250326154945-ae57f3c0d45f h1:C5bqEmzEPLsHm9Mv73lSE9e9bKV23aB1vxOsmZrkl3k=
-github.com/cncf/xds/go v0.0.0-20250326154945-ae57f3c0d45f/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8=
+github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 h1:aQ3y1lwWyqYPiWZThqv1aFbZMiM9vblcSArJRf2Irls=
+github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8=
 github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI=
 github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M=
 github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE=
@@ -62,14 +64,20 @@ github.com/danko-miladinovic/go-tpm-tools v0.0.0-20250228160324-1ebcfd79567c/go.
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/decred/dcrd/dcrec/secp256k1/v4 v4.4.0 h1:NMZiJj8QnKe1LgsbDayM4UoHwbvwDRwnI3hwNaAHRnc=
+github.com/decred/dcrd/dcrec/secp256k1/v4 v4.4.0/go.mod h1:ZXNYxsqcloTdSy/rNShjYzMhyjf0LaoftYK0p+A3h40=
+github.com/dgraph-io/ristretto/v2 v2.3.0 h1:qTQ38m7oIyd4GAed/QkUZyPFNMnvVWyazGXRwvOt5zk=
+github.com/dgraph-io/ristretto/v2 v2.3.0/go.mod h1:gpoRV3VzrEY1a9dWAYV6T1U7YzfgttXdd/ZzL1s9OZM=
 github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk=
 github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E=
-github.com/docker/docker v28.2.2+incompatible h1:CjwRSksz8Yo4+RmQ339Dp/D2tGO5JxwYeqtMOEe0LDw=
-github.com/docker/docker v28.2.2+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
+github.com/docker/docker v28.5.0+incompatible h1:ZdSQoRUE9XxhFI/B8YLvhnEFMmYN9Pp8Egd2qcaFk1E=
+github.com/docker/docker v28.5.0+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/go-connections v0.5.0 h1:USnMq7hx7gwdVZq1L49hLXaFtUdTADjXGp+uj1Br63c=
 github.com/docker/go-connections v0.5.0/go.mod h1:ov60Kzw0kKElRwhNs9UlUHAE/F9Fe6GLaXnqyDdmEXc=
 github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4=
 github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
+github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
+github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
 github.com/edgelesssys/go-azguestattestation v0.0.0-20250408071817-8c4457b235ff h1:V6A5kD0+c1Qg4X72Lg+zxhCZk+par436sQdgLvMCBBc=
 github.com/edgelesssys/go-azguestattestation v0.0.0-20250408071817-8c4457b235ff/go.mod h1:Lz4QaomI4wU2YbatD4/W7vatW2Q35tnkoJezB1clscc=
 github.com/envoyproxy/go-control-plane v0.13.4 h1:zEqyPVyku6IvWCFwux4x9RxkLOMUL+1vC9xUFv5l2/M=
@@ -84,12 +92,16 @@ github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
 github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
 github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
 github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
-github.com/go-chi/chi/v5 v5.2.1 h1:KOIHODQj58PmL80G2Eak4WdvUzjSJSm0vG72crDCqb8=
-github.com/go-chi/chi/v5 v5.2.1/go.mod h1:L2yAIGWB3H+phAw1NxKwWM+7eUH/lU8pOMm5hHcoops=
-github.com/go-jose/go-jose/v3 v3.0.3 h1:fFKWeig/irsp7XD2zBxvnmA/XaRWp5V3CBsZXJF7G7k=
-github.com/go-jose/go-jose/v3 v3.0.3/go.mod h1:5b+7YgP7ZICgJDBdfjZaIt+H/9L9T/YQrVfLAMboGkQ=
-github.com/go-jose/go-jose/v4 v4.0.5 h1:M6T8+mKZl/+fNNuFHvGIzDz7BTLQPIounk/b9dw3AaE=
-github.com/go-jose/go-jose/v4 v4.0.5/go.mod h1:s3P1lRrkT8igV8D9OjyL4WRyHvjB6a4JSllnOrmmBOA=
+github.com/fxamacker/cbor/v2 v2.9.0 h1:NpKPmjDBgUfBms6tr6JZkTHtfFGcMKsw3eGcmD/sapM=
+github.com/fxamacker/cbor/v2 v2.9.0/go.mod h1:vM4b+DJCtHn+zz7h3FFp/hDAI9WNWCsZj23V5ytsSxQ=
+github.com/go-chi/chi/v5 v5.2.3 h1:WQIt9uxdsAbgIYgid+BpYc+liqQZGMHRaUwp0JUcvdE=
+github.com/go-chi/chi/v5 v5.2.3/go.mod h1:L2yAIGWB3H+phAw1NxKwWM+7eUH/lU8pOMm5hHcoops=
+github.com/go-gorp/gorp/v3 v3.1.0 h1:ItKF/Vbuj31dmV4jxA1qblpSwkl9g1typ24xoe70IGs=
+github.com/go-gorp/gorp/v3 v3.1.0/go.mod h1:dLEjIyyRNiXvNZ8PSmzpt1GsWAUK8kjVhEpjH8TixEw=
+github.com/go-jose/go-jose/v3 v3.0.4 h1:Wp5HA7bLQcKnf6YYao/4kpRpVMp/yf6+pJKV8WFSaNY=
+github.com/go-jose/go-jose/v3 v3.0.4/go.mod h1:5b+7YgP7ZICgJDBdfjZaIt+H/9L9T/YQrVfLAMboGkQ=
+github.com/go-jose/go-jose/v4 v4.1.1 h1:JYhSgy4mXXzAdF3nUx3ygx347LRXJRrpgyU3adRmkAI=
+github.com/go-jose/go-jose/v4 v4.1.1/go.mod h1:BdsZGqgdO3b6tTc6LSE56wcDbMMLuPsw5d4ZD5f94kA=
 github.com/go-kit/kit v0.13.0 h1:OoneCcHKHQ03LfBpoQCUfCluwd2Vt3ohz+kvbJneZAU=
 github.com/go-kit/kit v0.13.0/go.mod h1:phqEHMMUbyrCFCTgH48JueqrM3md2HcAZ8N3XE4FKDg=
 github.com/go-kit/log v0.2.1 h1:MRVx0/zhvdseW+Gza6N9rVzU/IVzaeE1SFI4raAhmBU=
@@ -101,14 +113,14 @@ github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
 github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
 github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
 github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
+github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4=
+github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M=
 github.com/gofrs/uuid v4.4.0+incompatible h1:3qXRTX8/NbyulANqlc0lchS1gqAVxRgsuW1YrTJupqA=
 github.com/gofrs/uuid v4.4.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
-github.com/gofrs/uuid/v5 v5.3.0 h1:m0mUMr+oVYUdxpMLgSYCZiXe7PuVPnI94+OMeVBNedk=
-github.com/gofrs/uuid/v5 v5.3.0/go.mod h1:CDOjlDMVAtN56jqyRUZh58JT31Tiw7/oQyEXZV+9bD8=
-github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
-github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
-github.com/golang-jwt/jwt/v5 v5.2.2 h1:Rl4B7itRWVtYIHFrSNd7vhTiz9UpLdi6gZhZ3wEeDy8=
-github.com/golang-jwt/jwt/v5 v5.2.2/go.mod h1:pqrtFR0X4osieyHYxtmOUWsAWrfe1Q5UVIyoH402zdk=
+github.com/gofrs/uuid/v5 v5.3.2 h1:2jfO8j3XgSwlz/wHqemAEugfnTlikAYHhnqQ8Xh4fE0=
+github.com/gofrs/uuid/v5 v5.3.2/go.mod h1:CDOjlDMVAtN56jqyRUZh58JT31Tiw7/oQyEXZV+9bD8=
+github.com/golang-jwt/jwt/v5 v5.3.0 h1:pv4AsKCKKZuqlgs5sUmn4x8UlGa0kEVt/puTpKx9vvo=
+github.com/golang-jwt/jwt/v5 v5.3.0/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
 github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
 github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
 github.com/google/certificate-transparency-go v1.0.21/go.mod h1:QeJfpSbVSfYc7RgB3gJFj9cbuQMMchQxrWXz8Ruopmg=
@@ -129,8 +141,8 @@ github.com/google/go-sev-guest v0.13.0 h1:DJB6ACdykyweMU0HGOp/TQ7cjsnbV2ecbYunu2
 github.com/google/go-sev-guest v0.13.0/go.mod h1:SK9vW+uyfuzYdVN0m8BShL3OQCtXZe/JPF7ZkpD3760=
 github.com/google/go-tdx-guest v0.3.2-0.20241009005452-097ee70d0843 h1:+MoPobRN9HrDhGyn6HnF5NYo4uMBKaiFqAtf/D/OB4A=
 github.com/google/go-tdx-guest v0.3.2-0.20241009005452-097ee70d0843/go.mod h1:g/n8sKITIT9xRivBUbizo34DTsUm2nN2uU3A662h09g=
-github.com/google/go-tpm v0.9.5 h1:ocUmnDebX54dnW+MQWGQRbdaAcJELsa6PqZhJ48KwVU=
-github.com/google/go-tpm v0.9.5/go.mod h1:h9jEsEECg7gtLis0upRBQU+GhYVH6jMjrFxI8u6bVUY=
+github.com/google/go-tpm v0.9.6 h1:Ku42PT4LmjDu1H5C5ISWLlpI1mj+Zq7sPGKoRw2XROA=
+github.com/google/go-tpm v0.9.6/go.mod h1:h9jEsEECg7gtLis0upRBQU+GhYVH6jMjrFxI8u6bVUY=
 github.com/google/go-tspi v0.3.0 h1:ADtq8RKfP+jrTyIWIZDIYcKOMecRqNJFOew2IT0Inus=
 github.com/google/go-tspi v0.3.0/go.mod h1:xfMGI3G0PhxCdNVcYr1C4C+EizojDg/TXuX5by8CiHI=
 github.com/google/logger v1.1.1 h1:+6Z2geNxc9G+4D4oDO9njjjn2d0wN5d7uOo0vOIW1NQ=
@@ -143,27 +155,48 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/googleapis/enterprise-certificate-proxy v0.3.6 h1:GW/XbdyBFQ8Qe+YAmFU9uHLo7OnF5tL52HFAgMmyrf4=
 github.com/googleapis/enterprise-certificate-proxy v0.3.6/go.mod h1:MkHOF77EYAE7qfSuSS9PU6g4Nt4e11cnsDUowfwewLA=
-github.com/googleapis/gax-go/v2 v2.14.2 h1:eBLnkZ9635krYIPD+ag1USrOAI0Nr0QYF3+/3GqO0k0=
-github.com/googleapis/gax-go/v2 v2.14.2/go.mod h1:ON64QhlJkhVtSqp4v1uaK92VyZ2gmvDQsweuyLV+8+w=
+github.com/googleapis/gax-go/v2 v2.15.0 h1:SyjDc1mGgZU5LncH8gimWo9lW1DtIfPibOG81vgd/bo=
+github.com/googleapis/gax-go/v2 v2.15.0/go.mod h1:zVVkkxAQHa1RQpg9z2AUCMnKhi0Qld9rcmyfL1OZhoc=
 github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g=
 github.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k=
-github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.3 h1:5ZPtiqj0JL5oKWmcsq4VMaAW5ukBEgSGXEN89zeH1Jo=
-github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.3/go.mod h1:ndYquD05frm2vACXE1nsccT4oJzjhw2arTS2cpUD1PI=
-github.com/hokaccha/go-prettyjson v0.0.0-20211117102719-0474bc63780f h1:7LYC+Yfkj3CTRcShK0KOL/w6iTiKyqqBA9a41Wnggw8=
-github.com/hokaccha/go-prettyjson v0.0.0-20211117102719-0474bc63780f/go.mod h1:pFlLw2CfqZiIBOx6BuCeRLCrfxBJipTY0nIOF/VbGcI=
+github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2 h1:8Tjv8EJ+pM1xP8mK6egEbD1OgnVTyacbefKhmbLhIhU=
+github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2/go.mod h1:pkJQ2tZHJ0aFOVEEot6oZmaVEZcRme73eIFmhiVuRWs=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
+github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
+github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
+github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
+github.com/jackc/pgx/v5 v5.7.6 h1:rWQc5FwZSPX58r1OQmkuaNicxdmExaEz5A2DO2hUuTk=
+github.com/jackc/pgx/v5 v5.7.6/go.mod h1:aruU7o91Tc2q2cFp5h4uP3f6ztExVpyVv88Xl/8Vl8M=
+github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
+github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
+github.com/jmoiron/sqlx v1.4.0 h1:1PLqN7S1UYp5t4SrVVnt4nUVNemrDAtxlulVe+Qgm3o=
+github.com/jmoiron/sqlx v1.4.0/go.mod h1:ZrZ7UsYB/weZdl2Bxg6jCRO9c3YHl8r3ahlKmRT4JLY=
 github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo=
 github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
-github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
-github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
+github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
+github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
-github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
-github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
-github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
+github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
+github.com/lestrrat-go/blackmagic v1.0.3 h1:94HXkVLxkZO9vJI/w2u1T0DAoprShFd13xtnSINtDWs=
+github.com/lestrrat-go/blackmagic v1.0.3/go.mod h1:6AWFyKNNj0zEXQYfTMPfZrAXUWUfTIZ5ECEUEJaijtw=
+github.com/lestrrat-go/httpcc v1.0.1 h1:ydWCStUeJLkpYyjLDHihupbn2tYmZ7m22BGkcvZZrIE=
+github.com/lestrrat-go/httpcc v1.0.1/go.mod h1:qiltp3Mt56+55GPVCbTdM9MlqhvzyuL6W/NMDA8vA5E=
+github.com/lestrrat-go/httprc v1.0.6 h1:qgmgIRhpvBqexMJjA/PmwSvhNk679oqD1RbovdCGW8k=
+github.com/lestrrat-go/httprc v1.0.6/go.mod h1:mwwz3JMTPBjHUkkDv/IGJ39aALInZLrhBp0X7KGUZlo=
+github.com/lestrrat-go/iter v1.0.2 h1:gMXo1q4c2pHmC3dn8LzRhJfP1ceCbgSiT9lUydIzltI=
+github.com/lestrrat-go/iter v1.0.2/go.mod h1:Momfcq3AnRlRjI5b5O8/G5/BvpzrhoFTZcn06fEOPt4=
+github.com/lestrrat-go/jwx/v2 v2.1.6 h1:hxM1gfDILk/l5ylers6BX/Eq1m/pnxe9NBwW6lVfecA=
+github.com/lestrrat-go/jwx/v2 v2.1.6/go.mod h1:Y722kU5r/8mV7fYDifjug0r8FK8mZdw0K0GpJw/l8pU=
+github.com/lestrrat-go/option v1.0.1 h1:oAzP2fvZGQKWkvHa1/SAcFolBEca1oN+mQ7eooNBEYU=
+github.com/lestrrat-go/option v1.0.1/go.mod h1:5ZHFbivi4xwXxhxY9XHDe2FHo6/Z7WWmtT7T5nBBp3I=
+github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
+github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0=
@@ -182,48 +215,51 @@ github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/image-spec v1.1.0 h1:8SG7/vwALn54lVB/0yZ/MMwhFrPYtpEHQb2IpWsCzug=
 github.com/opencontainers/image-spec v1.1.0/go.mod h1:W4s4sFTMaBeK1BQLXbG4AdM2szdn85PY75RI83NrTrM=
-github.com/pelletier/go-toml v1.9.5 h1:4yBQzkHv+7BHq2PQUZF3Mx0IYxG7LsP222s7Agd3ve8=
-github.com/pelletier/go-toml v1.9.5/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
-github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo=
-github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8=
+github.com/planetscale/vtprotobuf v0.6.1-0.20240917153116-6f2963f01587 h1:xzZOeCMQLA/W198ZkdVdt4EKFKJtS26B773zNU377ZY=
+github.com/planetscale/vtprotobuf v0.6.1-0.20240917153116-6f2963f01587/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/prometheus/client_golang v1.22.0 h1:rb93p9lokFEsctTys46VnV1kLCDpVZ0a/Y92Vm0Zc6Q=
-github.com/prometheus/client_golang v1.22.0/go.mod h1:R7ljNsLXhuQXYZYtw6GAE9AZg8Y7vEW5scdCXrWRXC0=
-github.com/prometheus/client_model v0.6.1 h1:ZKSh/rekM+n3CeS952MLRAdFwIKqeY8b62p8ais2e9E=
-github.com/prometheus/client_model v0.6.1/go.mod h1:OrxVMOVHjw3lKMa8+x6HeMGkHMQyHDk9E3jmP2AmGiY=
-github.com/prometheus/common v0.62.0 h1:xasJaQlnWAeyHdUBeGjXmutelfJHWMRr+Fg4QszZ2Io=
-github.com/prometheus/common v0.62.0/go.mod h1:vyBcEuLSvWos9B1+CyL7JZ2up+uFzXhkqml0W5zIY1I=
-github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc=
-github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk=
+github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o=
+github.com/prometheus/client_golang v1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg=
+github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk=
+github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE=
+github.com/prometheus/common v0.66.1 h1:h5E0h5/Y8niHc5DlaLlWLArTQI7tMrsfQjHV+d9ZoGs=
+github.com/prometheus/common v0.66.1/go.mod h1:gcaUsgf3KfRSwHY4dIMXLPV0K/Wg1oZ8+SbZk/HH/dA=
+github.com/prometheus/procfs v0.17.0 h1:FuLQ+05u4ZI+SS/w9+BWEM2TXiHKsUQ9TADiRH7DuK0=
+github.com/prometheus/procfs v0.17.0/go.mod h1:oPQLaDAMRbA+u8H5Pbfq+dl3VDAvHxMUOVhe0wYB2zw=
 github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
 github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
+github.com/rubenv/sql-migrate v1.8.0 h1:dXnYiJk9k3wetp7GfQbKJcPHjVJL6YK19tKj8t2Ns0o=
+github.com/rubenv/sql-migrate v1.8.0/go.mod h1:F2bGFBwCU+pnmbtNYDeKvSuvL6lBVtXDXUUv5t+u1qw=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/sammyoina/sev-snp-measure-go v0.0.0-20241202151803-ef189f0ff825 h1:SqNaL9udBIc026SGNEuEuiVL0/hw9fXxM5qrFhWGkdE=
 github.com/sammyoina/sev-snp-measure-go v0.0.0-20241202151803-ef189f0ff825/go.mod h1:dEkBe8JnxU5itNjZDEQINFd7f7l4DtjfqRuzPQcit4w=
+github.com/segmentio/asm v1.2.0 h1:9BQrFxC+YOHJlTlHGkTrFWf59nbL3XnCoFLTwDCI7ys=
+github.com/segmentio/asm v1.2.0/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
-github.com/smarty/assertions v1.15.0 h1:cR//PqUBUiQRakZWqBiFFQ9wb8emQGDb0HeGdqGByCY=
-github.com/smarty/assertions v1.15.0/go.mod h1:yABtdzeQs6l1brC900WlRNwj6ZR55d7B+E8C6HtKdec=
+github.com/smarty/assertions v1.16.0 h1:EvHNkdRA4QHMrn75NZSoUQ/mAUXAYWfatfB01yTCzfY=
+github.com/smarty/assertions v1.16.0/go.mod h1:duaaFdCS0K9dnoM50iyek/eYINOZ64gbh1Xlf6LG7AI=
 github.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY=
 github.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60=
-github.com/spf13/cobra v1.9.1 h1:CXSaggrXdbHK9CF+8ywj8Amf7PBRmPCOJugH954Nnlo=
-github.com/spf13/cobra v1.9.1/go.mod h1:nDyEzZ8ogv936Cinf6g1RU9MRY64Ir93oCnqb9wxYW0=
-github.com/spf13/pflag v1.0.6 h1:jFzHGLGAlb3ruxLB8MhbI6A8+AQX/2eW4qeyNZXNp2o=
-github.com/spf13/pflag v1.0.6/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/spf13/cobra v1.10.1 h1:lJeBwCfmrnXthfAupyUTzJ/J4Nc1RsHC/mSRU2dll/s=
+github.com/spf13/cobra v1.10.1/go.mod h1:7SmJGaTHFVBY0jW4NXGluQoLvhqFQM+6XSKD+P4XaB0=
+github.com/spf13/pflag v1.0.9/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
+github.com/spf13/pflag v1.0.10 h1:4EBh2KAYBwaONj6b2Ye1GiHfwjqyROoF4RwYO+vPwFk=
+github.com/spf13/pflag v1.0.10/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/spiffe/go-spiffe/v2 v2.5.0 h1:N2I01KCUkv1FAjZXJMwh95KK1ZIQLYbPfhaxw8WS0hE=
 github.com/spiffe/go-spiffe/v2 v2.5.0/go.mod h1:P+NxobPc6wXhVtINNtFjNWGBTreew1GBUCwT2wPmb7g=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
 github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
-github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
-github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
+github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
+github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
+github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM=
+github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
 github.com/zeebo/errs v1.4.0 h1:XNdoD/RRMKP7HD0UhJnIzUy74ISdGGxURlYG8HSWSfM=
 github.com/zeebo/errs v1.4.0/go.mod h1:sgbWHsvVuTPHcqJJGQ1WhI5KbWlHYz+2+2C/LSEtCw4=
@@ -231,123 +267,108 @@ go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJyS
 go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
 go.opentelemetry.io/contrib/detectors/gcp v1.36.0 h1:F7q2tNlCaHY9nMKHR6XH9/qkp8FktLnIcy6jJNyOCQw=
 go.opentelemetry.io/contrib/detectors/gcp v1.36.0/go.mod h1:IbBN8uAIIx734PTonTPxAxnjc2pQTxWNkwfstZ+6H2k=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.62.0 h1:rbRJ8BBoVMsQShESYZ0FkvcITu8X8QNwJogcLUmDNNw=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.62.0/go.mod h1:ru6KHrNtNHxM4nD/vd6QrLVWgKhxPYgblq4VAtNawTQ=
-go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 h1:sbiXRNDSWJOTobXh5HyQKjq6wUC5tNybqjIqDpAY4CU=
-go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0/go.mod h1:69uWxva0WgAA/4bu2Yy70SLDBwZXuQ6PbBpbsa5iZrQ=
-go.opentelemetry.io/otel v1.37.0 h1:9zhNfelUvx0KBfu/gb+ZgeAfAgtWrfHJZcAqFC228wQ=
-go.opentelemetry.io/otel v1.37.0/go.mod h1:ehE/umFRLnuLa/vSccNq9oS1ErUlkkK71gMcN34UG8I=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.36.0 h1:dNzwXjZKpMpE2JhmO+9HsPl42NIXFIFSUSSs0fiqra0=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.36.0/go.mod h1:90PoxvaEB5n6AOdZvi+yWJQoE95U8Dhhw2bSyRqnTD0=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0 h1:xJ2qHD0C1BeYVTLLR9sX12+Qb95kfeD/byKj6Ky1pXg=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.35.0/go.mod h1:u5BF1xyjstDowA1R5QAO9JHzqK+ublenEW/dyqTjBVk=
+go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.63.0 h1:YH4g8lQroajqUwWbq/tr2QX1JFmEXaDLgG+ew9bLMWo=
+go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.63.0/go.mod h1:fvPi2qXDqFs8M4B4fmJhE92TyQs9Ydjlg3RvfUp+NbQ=
+go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0 h1:RbKq8BG0FI8OiXhBfcRtqqHcZcka+gU3cskNuf05R18=
+go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.63.0/go.mod h1:h06DGIukJOevXaj/xrNjhi/2098RZzcLTbc0jDAUbsg=
+go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8=
+go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM=
+go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0 h1:GqRJVj7UmLjCVyVJ3ZFLdPRmhDUp2zFmQe3RHIOsw24=
+go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.38.0/go.mod h1:ri3aaHSmCTVYu2AWv44YMauwAQc0aqI9gHKIcSbI1pU=
+go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0 h1:aTL7F04bJHUlztTsNGJ2l+6he8c+y/b//eR0jjjemT4=
+go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.38.0/go.mod h1:kldtb7jDTeol0l3ewcmd8SDvx3EmIE7lyvqbasU3QC4=
 go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.36.0 h1:rixTyDGXFxRy1xzhKrotaHy3/KXdPhlWARrCgK+eqUY=
 go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.36.0/go.mod h1:dowW6UsM9MKbJq5JTz2AMVp3/5iW5I/TStsk8S+CfHw=
-go.opentelemetry.io/otel/metric v1.37.0 h1:mvwbQS5m0tbmqML4NqK+e3aDiO02vsf/WgbsdpcPoZE=
-go.opentelemetry.io/otel/metric v1.37.0/go.mod h1:04wGrZurHYKOc+RKeye86GwKiTb9FKm1WHtO+4EVr2E=
-go.opentelemetry.io/otel/sdk v1.37.0 h1:ItB0QUqnjesGRvNcmAcU0LyvkVyGJ2xftD29bWdDvKI=
-go.opentelemetry.io/otel/sdk v1.37.0/go.mod h1:VredYzxUvuo2q3WRcDnKDjbdvmO0sCzOvVAiY+yUkAg=
-go.opentelemetry.io/otel/sdk/metric v1.37.0 h1:90lI228XrB9jCMuSdA0673aubgRobVZFhbjxHHspCPc=
-go.opentelemetry.io/otel/sdk/metric v1.37.0/go.mod h1:cNen4ZWfiD37l5NhS+Keb5RXVWZWpRE+9WyVCpbo5ps=
-go.opentelemetry.io/otel/trace v1.37.0 h1:HLdcFNbRQBE2imdSEgm/kwqmQj1Or1l/7bW6mxVK7z4=
-go.opentelemetry.io/otel/trace v1.37.0/go.mod h1:TlgrlQ+PtQO5XFerSPUYG0JSgGyryXewPGyayAWSBS0=
-go.opentelemetry.io/proto/otlp v1.6.0 h1:jQjP+AQyTf+Fe7OKj/MfkDrmK4MNVtw2NpXsf9fefDI=
-go.opentelemetry.io/proto/otlp v1.6.0/go.mod h1:cicgGehlFuNdgZkcALOCh3VE6K/u2tAjzlRhDwmVpZc=
+go.opentelemetry.io/otel/metric v1.38.0 h1:Kl6lzIYGAh5M159u9NgiRkmoMKjvbsKtYRwgfrA6WpA=
+go.opentelemetry.io/otel/metric v1.38.0/go.mod h1:kB5n/QoRM8YwmUahxvI3bO34eVtQf2i4utNVLr9gEmI=
+go.opentelemetry.io/otel/sdk v1.38.0 h1:l48sr5YbNf2hpCUj/FoGhW9yDkl+Ma+LrVl8qaM5b+E=
+go.opentelemetry.io/otel/sdk v1.38.0/go.mod h1:ghmNdGlVemJI3+ZB5iDEuk4bWA3GkTpW+DOoZMYBVVg=
+go.opentelemetry.io/otel/sdk/metric v1.38.0 h1:aSH66iL0aZqo//xXzQLYozmWrXxyFkBJ6qT5wthqPoM=
+go.opentelemetry.io/otel/sdk/metric v1.38.0/go.mod h1:dg9PBnW9XdQ1Hd6ZnRz689CbtrUp0wMMs9iPcgT9EZA=
+go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE=
+go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs=
+go.opentelemetry.io/proto/otlp v1.7.1 h1:gTOMpGDb0WTBOP8JaO72iL3auEZhVmAQg4ipjOVAtj4=
+go.opentelemetry.io/proto/otlp v1.7.1/go.mod h1:b2rVh6rfI/s2pHWNlB7ILJcRALpcNDzKhACevjI+ZnE=
 go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
 go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
 go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
 go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
+go.yaml.in/yaml/v2 v2.4.2 h1:DzmwEr2rDGHl7lsFgAHxmNz/1NlQ7xLIrlN2h5d1eGI=
+go.yaml.in/yaml/v2 v2.4.2/go.mod h1:081UH+NErpNdqlCXm3TtEran0rJZGxAYx9hb/ELlsPU=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
-golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
-golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
-golang.org/x/crypto v0.39.0 h1:SHs+kF4LP+f+p14esP5jAoDpHU8Gu/v9lFRK6IT5imM=
-golang.org/x/crypto v0.39.0/go.mod h1:L+Xg3Wf6HoL4Bn4238Z6ft6KfEpN0tJGo53AAPC632U=
-golang.org/x/exp v0.0.0-20240909161429-701f63a606c0 h1:e66Fs6Z+fZTbFBAxKfP3PALWBtpfqks2bwGcexMxgtk=
-golang.org/x/exp v0.0.0-20240909161429-701f63a606c0/go.mod h1:2TbTHSBQa924w8M6Xs1QcRcFwyucIwBGpK1p2f1YFFY=
-golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
-golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
+golang.org/x/crypto v0.42.0 h1:chiH31gIWm57EkTXpwnqf8qeuMUi0yekh6mT2AvFlqI=
+golang.org/x/crypto v0.42.0/go.mod h1:4+rDnOTJhQCx2q7/j6rAN5XDw8kPjeaXEUR2eL94ix8=
+golang.org/x/exp v0.0.0-20250606033433-dcc06ee1d476 h1:bsqhLWFR6G6xiQcb+JoGqdKdRU6WzPWmK8E0jxTjzo4=
+golang.org/x/exp v0.0.0-20250606033433-dcc06ee1d476/go.mod h1:3//PLf8L/X+8b4vuAfHzxeRUl04Adcb341+IGKfnqS8=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
-golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
-golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
-golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
 golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
 golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
-golang.org/x/net v0.41.0 h1:vBTly1HeNPEn3wtREYfy4GZ/NECgw2Cnl+nK6Nz3uvw=
-golang.org/x/net v0.41.0/go.mod h1:B/K4NNqkfmg07DQYrbwvSluqCJOOXwUjeb/5lOisjbA=
+golang.org/x/net v0.43.0 h1:lat02VYK2j4aLzMzecihNvTlJNQUq316m2Mr9rnM6YE=
+golang.org/x/net v0.43.0/go.mod h1:vhO1fvI4dGsIjh73sWfUVjj3N7CA9WkKJNQm2svM6Jg=
 golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI=
 golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.15.0 h1:KWH3jNZsfyT6xfAfKiz6MRNmd46ByHDYaZ7KSkCtdW8=
-golang.org/x/sync v0.15.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA=
+golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug=
+golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
-golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210426230700-d19ff857e887/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw=
-golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+golang.org/x/sys v0.36.0 h1:KVRy2GtZBrk1cBYA7MKu5bEZFxQk4NIDV6RLVcC8o0k=
+golang.org/x/sys v0.36.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
 golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo=
 golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
-golang.org/x/term v0.32.0 h1:DR4lr0TjUs3epypdhTOkMmuF5CDFJ/8pOnbzMZPQ7bg=
-golang.org/x/term v0.32.0/go.mod h1:uZG1FhGx848Sqfsq4/DlJr3xGGsYMu/L5GW4abiaEPQ=
+golang.org/x/term v0.35.0 h1:bZBVKBudEyhRcajGcNc3jIfWPqV4y/Kt2XcoigOWtDQ=
+golang.org/x/term v0.35.0/go.mod h1:TPGtkTLesOwf2DE8CgVYiZinHAOuy5AYUYT1lENIZnA=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
 golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
-golang.org/x/text v0.26.0 h1:P42AVeLghgTYr4+xUnTRKDMqpar+PtX7KWuNQL21L8M=
-golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA=
-golang.org/x/time v0.11.0 h1:/bpjEDfN9tkoN/ryeYHnv5hcMlc8ncjMcM4XBk5NWV0=
-golang.org/x/time v0.11.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
+golang.org/x/text v0.29.0 h1:1neNs90w9YzJ9BocxfsQNHKuAT4pkghyXc4nhZ6sJvk=
+golang.org/x/text v0.29.0/go.mod h1:7MhJOA9CD2qZyOKYazxdYMF85OwPdEr9jTtBpO7ydH4=
+golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE=
+golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
-golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
-golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
 golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-google.golang.org/api v0.235.0 h1:C3MkpQSRxS1Jy6AkzTGKKrpSCOd2WOGrezZ+icKSkKo=
-google.golang.org/api v0.235.0/go.mod h1:QpeJkemzkFKe5VCE/PMv7GsUfn9ZF+u+q1Q7w6ckxTg=
-google.golang.org/genproto v0.0.0-20250505200425-f936aa4a68b2 h1:1tXaIXCracvtsRxSBsYDiSBN0cuJvM7QYW+MrpIRY78=
-google.golang.org/genproto v0.0.0-20250505200425-f936aa4a68b2/go.mod h1:49MsLSx0oWMOZqcpB3uL8ZOkAh1+TndpJ8ONoCBWiZk=
-google.golang.org/genproto/googleapis/api v0.0.0-20250512202823-5a2f75b736a9 h1:WvBuA5rjZx9SNIzgcU53OohgZy6lKSus++uY4xLaWKc=
-google.golang.org/genproto/googleapis/api v0.0.0-20250512202823-5a2f75b736a9/go.mod h1:W3S/3np0/dPWsWLi1h/UymYctGXaGBM2StwzD0y140U=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20250603155806-513f23925822 h1:fc6jSaCT0vBduLYZHYrBBNY4dsWuvgyff9noRNDdBeE=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20250603155806-513f23925822/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A=
-google.golang.org/grpc v1.73.0 h1:VIWSmpI2MegBtTuFt5/JWy2oXxtjJ/e89Z70ImfD2ok=
-google.golang.org/grpc v1.73.0/go.mod h1:50sbHOUqWoCQGI8V2HQLJM0B+LMlIUjNSZmow7EVBQc=
-google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=
-google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY=
+gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
+gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
+google.golang.org/api v0.247.0 h1:tSd/e0QrUlLsrwMKmkbQhYVa109qIintOls2Wh6bngc=
+google.golang.org/api v0.247.0/go.mod h1:r1qZOPmxXffXg6xS5uhx16Fa/UFY8QU/K4bfKrnvovM=
+google.golang.org/genproto v0.0.0-20250603155806-513f23925822 h1:rHWScKit0gvAPuOnu87KpaYtjK5zBMLcULh7gxkCXu4=
+google.golang.org/genproto v0.0.0-20250603155806-513f23925822/go.mod h1:HubltRL7rMh0LfnQPkMH4NPDFEWp0jw3vixw7jEM53s=
+google.golang.org/genproto/googleapis/api v0.0.0-20250825161204-c5933d9347a5 h1:BIRfGDEjiHRrk0QKZe3Xv2ieMhtgRGeLcZQ0mIVn4EY=
+google.golang.org/genproto/googleapis/api v0.0.0-20250825161204-c5933d9347a5/go.mod h1:j3QtIyytwqGr1JUDtYXwtMXWPKsEa5LtzIFN1Wn5WvE=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20250825161204-c5933d9347a5 h1:eaY8u2EuxbRv7c3NiGK0/NedzVsCcV6hDuU5qPX5EGE=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20250825161204-c5933d9347a5/go.mod h1:M4/wBTSeyLxupu3W3tJtOgB14jILAS/XWPSSa3TAlJc=
+google.golang.org/grpc v1.75.1 h1:/ODCNEuf9VghjgO3rqLcfg8fiOP0nSluljWFlDxELLI=
+google.golang.org/grpc v1.75.1/go.mod h1:JtPAzKiq4v1xcAB2hydNlWI2RnF85XXcV0mhKXr2ecQ=
+google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE=
+google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
-gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
-gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

hal/cloud/.env

@@ -41,10 +41,6 @@ SEV_SNP_CBIT_POS=51
 SEV_SNP_REDUCED_PHYS_BITS=1
 SEV_SNP_HOST_DATA=
 
-#### VSock Config
-VSOCK_ID=vhost-vsock-pci0
-VSOCK_GUEST_CID=3
-
 BIN_PATH=qemu-system-x86_64
 USE_SUDO=false
 ENABLE_SEV_SNP=false

hal/cloud/qemu.sh

@@ -72,8 +72,7 @@ construct_qemu_args() {
     # 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_SNP (if enabled)
     if [ "$ENABLE_SEV_SNP" == "true" ]; then
         kernel_hash=""

hal/linux/README.md

@@ -10,7 +10,7 @@ HAL uses [Buildroot](https://buildroot.org/)'s [_External Tree_ mechanism](https
 git clone git@github.com:ultravioletrs/cocos.git
 git clone git@github.com:buildroot/buildroot.git
 cd buildroot
-git checkout 2025.05-rc1
+git checkout 2025.08-rc3
 make BR2_EXTERNAL=../cocos/hal/linux cocos_defconfig
 # Execute 'make menuconfig' only if you want to make additional configuration changes to Buildroot.
 make menuconfig

hal/linux/package/agent/agent.mk

@@ -24,4 +24,4 @@ define AGENT_INSTALL_INIT_SYSTEMD
 	$(INSTALL) -D -m 0750 $(@D)/init/systemd/agent_start_script.sh $(TARGET_DIR)/cocos_init/agent_start_script.sh
 endef
 
-$(eval $(golang-package))
+$(eval $(generic-package))

hal/linux/package/wasmedge/wasmedge.mk

@@ -1,7 +1,7 @@
 WASMEDGE_DOWNLOAD_URL = https://raw.githubusercontent.com/WasmEdge/WasmEdge/master/utils/install.sh
 
 define WASMEDGE_INSTALL_TARGET_CMDS
-    curl -sSf $(WASMEDGE_DOWNLOAD_URL) | bash -s -- -p $(TARGET_DIR)/usr
+    curl -sSf $(WASMEDGE_DOWNLOAD_URL) | bash -s -- -p $(TARGET_DIR)/usr -v 0.14.1
     echo "source /usr/env" >> $(TARGET_DIR)/etc/profile
 endef
 

internal/logger/protohandler_test.go

@@ -9,7 +9,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/ultravioletrs/cocos/agent/cvms"
 )

internal/server/grpc/grpc.go

@@ -1,307 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-package grpc
-
-import (
-	"context"
-	"crypto/tls"
-	"crypto/x509"
-	"fmt"
-	"log/slog"
-	"net"
-	"os"
-	"strings"
-	"sync"
-	"time"
-
-	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
-	"github.com/ultravioletrs/cocos/agent/auth"
-	"github.com/ultravioletrs/cocos/internal/server"
-	"github.com/ultravioletrs/cocos/pkg/atls"
-	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
-	"google.golang.org/grpc"
-	"google.golang.org/grpc/credentials"
-	"google.golang.org/grpc/credentials/insecure"
-	"google.golang.org/grpc/health"
-	grpchealth "google.golang.org/grpc/health/grpc_health_v1"
-)
-
-const (
-	stopWaitTime  = 5 * time.Second
-	organization  = "Ultraviolet"
-	country       = "Serbia"
-	province      = ""
-	locality      = "Belgrade"
-	streetAddress = "Bulevar Arsenija Carnojevica 103"
-	postalCode    = "11000"
-	notAfterYear  = 1
-	notAfterMonth = 0
-	notAfterDay   = 0
-	nonceSize     = 32
-)
-
-type Server struct {
-	server.BaseServer
-	mu              sync.RWMutex
-	server          *grpc.Server
-	health          *health.Server
-	registerService serviceRegister
-	authSvc         auth.Authenticator
-	caUrl           string
-	cvmId           string
-	started         bool
-	stopped         bool
-}
-
-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, authSvc auth.Authenticator, caUrl string, cvmId string) server.Server {
-	base := config.GetBaseConfig()
-	listenFullAddress := fmt.Sprintf("%s:%s", base.Host, base.Port)
-	return &Server{
-		BaseServer: server.BaseServer{
-			Ctx:     ctx,
-			Cancel:  cancel,
-			Name:    name,
-			Address: listenFullAddress,
-			Config:  config,
-			Logger:  logger,
-		},
-		registerService: registerService,
-		authSvc:         authSvc,
-		caUrl:           caUrl,
-		cvmId:           cvmId,
-	}
-}
-
-func (s *Server) Start() error {
-	s.mu.Lock()
-	if s.started {
-		s.mu.Unlock()
-		return fmt.Errorf("server already started")
-	}
-	if s.stopped {
-		s.mu.Unlock()
-		return fmt.Errorf("server already stopped")
-	}
-	s.started = true
-	s.mu.Unlock()
-
-	errCh := make(chan error)
-	grpcServerOptions := []grpc.ServerOption{
-		grpc.StatsHandler(otelgrpc.NewServerHandler()),
-	}
-
-	if s.authSvc != nil {
-		unary, stream := agentgrpc.NewAuthInterceptor(s.authSvc)
-		grpcServerOptions = append(grpcServerOptions, grpc.UnaryInterceptor(unary))
-		grpcServerOptions = append(grpcServerOptions, grpc.StreamInterceptor(stream))
-	}
-
-	creds := grpc.Creds(insecure.NewCredentials())
-
-	c := s.Config.GetBaseConfig()
-	if agCfg, ok := s.Config.(server.AgentConfig); ok && agCfg.AttestedTLS {
-		tlsConfig := &tls.Config{
-			ClientAuth:     tls.NoClientCert,
-			GetCertificate: atls.GetCertificate(s.caUrl, s.cvmId),
-		}
-
-		var mtls bool
-		mtls = false
-
-		// Loading Server CA file
-		rootCA, err := loadCertFile(c.ServerCAFile)
-		if err != nil {
-			return fmt.Errorf("failed to load server ca file: %w", err)
-		}
-		if len(rootCA) > 0 {
-			if tlsConfig.RootCAs == nil {
-				tlsConfig.RootCAs = x509.NewCertPool()
-			}
-			if !tlsConfig.RootCAs.AppendCertsFromPEM(rootCA) {
-				return fmt.Errorf("failed to append server ca to tls.Config")
-			}
-			mtls = true
-		}
-
-		// Loading Client CA File
-		clientCA, err := loadCertFile(c.ClientCAFile)
-		if err != nil {
-			return fmt.Errorf("failed to load client ca file: %w", err)
-		}
-		if len(clientCA) > 0 {
-			if tlsConfig.ClientCAs == nil {
-				tlsConfig.ClientCAs = x509.NewCertPool()
-			}
-			if !tlsConfig.ClientCAs.AppendCertsFromPEM(clientCA) {
-				return fmt.Errorf("failed to append client ca to tls.Config")
-			}
-
-			tlsConfig.ClientAuth = tls.RequireAndVerifyClientCert
-			mtls = true
-		}
-
-		creds = grpc.Creds(credentials.NewTLS(tlsConfig))
-
-		if mtls {
-			s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s with Attested mTLS", s.Name, s.Address))
-		} else {
-			s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s with Attested TLS", s.Name, s.Address))
-		}
-	} else {
-		switch {
-		case c.CertFile != "" || c.KeyFile != "":
-			certificate, err := loadX509KeyPair(c.CertFile, c.KeyFile)
-			if err != nil {
-				return fmt.Errorf("failed to load auth certificates: %w", err)
-			}
-			tlsConfig := &tls.Config{
-				ClientAuth:   tls.NoClientCert,
-				Certificates: []tls.Certificate{certificate},
-			}
-
-			var mtlsCA string
-			// Loading Server CA file
-			rootCA, err := loadCertFile(c.ServerCAFile)
-			if err != nil {
-				return fmt.Errorf("failed to load root ca file: %w", err)
-			}
-			if len(rootCA) > 0 {
-				if tlsConfig.RootCAs == nil {
-					tlsConfig.RootCAs = x509.NewCertPool()
-				}
-				if !tlsConfig.RootCAs.AppendCertsFromPEM(rootCA) {
-					return fmt.Errorf("failed to append root ca to tls.Config")
-				}
-				mtlsCA = fmt.Sprintf("root ca %s", c.ServerCAFile)
-			}
-
-			// Loading Client CA File
-			clientCA, err := loadCertFile(c.ClientCAFile)
-			if err != nil {
-				return fmt.Errorf("failed to load client ca file: %w", err)
-			}
-			if len(clientCA) > 0 {
-				if tlsConfig.ClientCAs == nil {
-					tlsConfig.ClientCAs = x509.NewCertPool()
-				}
-				if !tlsConfig.ClientCAs.AppendCertsFromPEM(clientCA) {
-					return fmt.Errorf("failed to append client ca to tls.Config")
-				}
-				mtlsCA = fmt.Sprintf("%s client ca %s", mtlsCA, c.ClientCAFile)
-			}
-			creds = grpc.Creds(credentials.NewTLS(tlsConfig))
-			switch {
-			case mtlsCA != "":
-				tlsConfig.ClientAuth = tls.RequireAndVerifyClientCert
-				creds = grpc.Creds(credentials.NewTLS(tlsConfig))
-				s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s with TLS/mTLS cert %s , key %s and %s", s.Name, s.Address, c.CertFile, c.KeyFile, mtlsCA))
-			default:
-				s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s with TLS cert %s and key %s", s.Name, s.Address, c.CertFile, c.KeyFile))
-			}
-		default:
-			s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s without TLS", s.Name, s.Address))
-		}
-	}
-
-	listener, err := net.Listen("tcp", s.Address)
-	if err != nil {
-		return fmt.Errorf("failed to listen on port %s: %w", s.Address, err)
-	}
-
-	grpcServerOptions = append(grpcServerOptions, creds)
-
-	s.mu.Lock()
-	s.server = grpc.NewServer(grpcServerOptions...)
-	s.health = health.NewServer()
-	grpchealth.RegisterHealthServer(s.server, s.health)
-	s.registerService(s.server)
-	s.health.SetServingStatus(s.Name, grpchealth.HealthCheckResponse_SERVING)
-	s.mu.Unlock()
-
-	go func() {
-		s.mu.RLock()
-		server := s.server
-		s.mu.RUnlock()
-
-		if server != nil {
-			errCh <- server.Serve(listener)
-		}
-	}()
-
-	select {
-	case <-s.Ctx.Done():
-		return s.Stop()
-	case err := <-errCh:
-		s.Cancel()
-		return err
-	}
-}
-
-func (s *Server) Stop() error {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
-	if s.stopped {
-		return nil
-	}
-	s.stopped = true
-
-	defer s.Cancel()
-
-	c := make(chan bool)
-	go func() {
-		defer close(c)
-		if s.health != nil {
-			s.health.Shutdown()
-		}
-		if s.server != nil {
-			s.server.GracefulStop()
-		}
-	}()
-
-	select {
-	case <-c:
-	case <-time.After(stopWaitTime):
-	}
-
-	s.Logger.Info(fmt.Sprintf("%s gRPC service shutdown at %s", s.Name, s.Address))
-	return nil
-}
-
-func loadCertFile(certFile string) ([]byte, error) {
-	if certFile != "" {
-		return readFileOrData(certFile)
-	}
-	return []byte{}, nil
-}
-
-func readFileOrData(input string) ([]byte, error) {
-	if len(input) < 1000 && !strings.Contains(input, "\n") {
-		data, err := os.ReadFile(input)
-		if err == nil {
-			return data, nil
-		} else {
-			return nil, err
-		}
-	}
-	return []byte(input), nil
-}
-
-func loadX509KeyPair(certfile, keyfile string) (tls.Certificate, error) {
-	cert, err := readFileOrData(certfile)
-	if err != nil {
-		return tls.Certificate{}, fmt.Errorf("failed to read cert: %v", err)
-	}
-
-	key, err := readFileOrData(keyfile)
-	if err != nil {
-		return tls.Certificate{}, fmt.Errorf("failed to read key: %v", err)
-	}
-
-	return tls.X509KeyPair(cert, key)
-}

internal/zip.go

@@ -111,6 +111,10 @@ func ZipDirectoryToTempFile(sourceDir string) (*os.File, error) {
 		return nil, err
 	}
 
+	if _, err := tmpFile.Seek(0, 0); err != nil {
+		return nil, err
+	}
+
 	return tmpFile, nil
 }
 

manager/README.md

@@ -14,11 +14,18 @@ The service is configured using the environment variables from the following tab
 | 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_HTTP_HOST                          | Manager service HTTP host                                                                                        | ""                             |
+| MANAGER_HTTP_PORT                          | Manager service HTTP port                                                                                        | 7003                           |
+| MANAGER_HTTP_SERVER_CERT                   | Manager to HTTP server certificate in pem format                                                                 | ""                             |
+| MANAGER_HTTP_SERVER_KEY                    | Path to HTTP server key in pem format                                                                            | ""                             |
+| MANAGER_HTTP_SERVER_CA_CERTS               | Path to HTTP server CA certificate                                                                               | ""                             |
+| MANAGER_HTTP_CLIENT_CA_CERTS               | Path to HTTP client CA certificate                                                                               | ""                             |
+| MANAGER_GRPC_HOST                          | Manager service gRPC host                                                                                        | ""                             |
+| MANAGER_GRPC_PORT                          | Manager service gRPC port                                                                                        | 7001                           |
+| MANAGER_GRPC_SERVER_CERT                   | Path to gRPC server certificate in pem format                                                                    | ""                             |
+| MANAGER_GRPC_SERVER_KEY                    | Path to gRPC server key in pem format                                                                            | ""                             |
+| MANAGER_GRPC_SERVER_CA_CERTS               | Path to gRPC server CA certificate                                                                               | ""                             |
+| MANAGER_GRPC_CLIENT_CA_CERTS               | Path to gRPC client CA certificate                                                                               | ""                             |
 | MANAGER_EOS_VERSION                        | The EOS version used for booting CVMs.                                                                           |                                |
 | 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                          |
@@ -53,9 +60,6 @@ The service is configured using the environment variables from the following tab
 | MANAGER_QEMU_OVMF_FILE                     | The file path for the OVMF file (combined OVMF_CODE and OVMF_VARS file).                                         | /usr/share/ovmf/OVMF.fd        |
 | 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_SNP                | Whether to enable Secure Nested Paging (SEV-SNP).                                                                | true                           |
@@ -69,6 +73,7 @@ The service is configured using the environment variables from the following tab
 | 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                      |
+| MANAGER_MAX_VMS                            | The maximum number of vms running concurrently on manager.                                                       | 10                             |
 
 ## Setup
 

manager/api/http/transport.go

@@ -0,0 +1,20 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package http
+
+import (
+	"net/http"
+
+	"github.com/absmach/supermq"
+	"github.com/go-chi/chi/v5"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+)
+
+// MakeHandler returns a HTTP handler for API endpoints.
+func MakeHandler(r *chi.Mux, svcName, instanceID string) http.Handler {
+	r.Get("/health", supermq.Health(svcName, instanceID))
+	r.Handle("/metrics", promhttp.Handler())
+
+	return r
+}

manager/api/http/transport_test.go

@@ -0,0 +1,379 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package http
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/go-chi/chi/v5"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMakeHandler(t *testing.T) {
+	const (
+		testServiceName = "test-service"
+		testInstanceID  = "test-instance-123"
+	)
+
+	tests := []struct {
+		name           string
+		serviceName    string
+		instanceID     string
+		expectedRoutes int
+	}{
+		{
+			name:           "valid handler creation",
+			serviceName:    testServiceName,
+			instanceID:     testInstanceID,
+			expectedRoutes: 2, // /health and /metrics
+		},
+		{
+			name:           "empty service name",
+			serviceName:    "",
+			instanceID:     testInstanceID,
+			expectedRoutes: 2,
+		},
+		{
+			name:           "empty instance ID",
+			serviceName:    testServiceName,
+			instanceID:     "",
+			expectedRoutes: 2,
+		},
+		{
+			name:           "both empty",
+			serviceName:    "",
+			instanceID:     "",
+			expectedRoutes: 2,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := chi.NewRouter()
+			handler := MakeHandler(r, tt.serviceName, tt.instanceID)
+
+			require.NotNil(t, handler)
+			assert.Implements(t, (*http.Handler)(nil), handler)
+
+			// Verify that the handler is actually the chi router
+			assert.Equal(t, r, handler)
+		})
+	}
+}
+
+func TestHealthEndpoint(t *testing.T) {
+	const (
+		testServiceName = "test-service"
+		testInstanceID  = "test-instance-123"
+	)
+
+	tests := []struct {
+		name           string
+		serviceName    string
+		instanceID     string
+		method         string
+		path           string
+		expectedStatus int
+	}{
+		{
+			name:           "GET health endpoint success",
+			serviceName:    testServiceName,
+			instanceID:     testInstanceID,
+			method:         http.MethodGet,
+			path:           "/health",
+			expectedStatus: http.StatusOK,
+		},
+		{
+			name:           "POST health endpoint not allowed",
+			serviceName:    testServiceName,
+			instanceID:     testInstanceID,
+			method:         http.MethodPost,
+			path:           "/health",
+			expectedStatus: http.StatusMethodNotAllowed,
+		},
+		{
+			name:           "PUT health endpoint not allowed",
+			serviceName:    testServiceName,
+			instanceID:     testInstanceID,
+			method:         http.MethodPut,
+			path:           "/health",
+			expectedStatus: http.StatusMethodNotAllowed,
+		},
+		{
+			name:           "DELETE health endpoint not allowed",
+			serviceName:    testServiceName,
+			instanceID:     testInstanceID,
+			method:         http.MethodDelete,
+			path:           "/health",
+			expectedStatus: http.StatusMethodNotAllowed,
+		},
+		{
+			name:           "health with empty service name",
+			serviceName:    "",
+			instanceID:     testInstanceID,
+			method:         http.MethodGet,
+			path:           "/health",
+			expectedStatus: http.StatusOK,
+		},
+		{
+			name:           "health with empty instance ID",
+			serviceName:    testServiceName,
+			instanceID:     "",
+			method:         http.MethodGet,
+			path:           "/health",
+			expectedStatus: http.StatusOK,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := chi.NewRouter()
+			handler := MakeHandler(r, tt.serviceName, tt.instanceID)
+
+			req, err := http.NewRequest(tt.method, tt.path, nil)
+			require.NoError(t, err)
+
+			rr := httptest.NewRecorder()
+			handler.ServeHTTP(rr, req)
+
+			assert.Equal(t, tt.expectedStatus, rr.Code)
+
+			if tt.expectedStatus == http.StatusOK {
+				// Verify content type for successful health checks
+				contentType := rr.Header().Get("Content-Type")
+				assert.Contains(t, contentType, "application/health+json")
+
+				// Verify response body contains service info
+				body := rr.Body.String()
+				if tt.serviceName != "" {
+					assert.Contains(t, body, tt.serviceName)
+				}
+				if tt.instanceID != "" {
+					assert.Contains(t, body, tt.instanceID)
+				}
+			}
+		})
+	}
+}
+
+func TestMetricsEndpoint(t *testing.T) {
+	const (
+		testServiceName = "test-service"
+		testInstanceID  = "test-instance-123"
+	)
+
+	tests := []struct {
+		name           string
+		method         string
+		path           string
+		expectedStatus int
+	}{
+		{
+			name:           "GET metrics endpoint success",
+			method:         http.MethodGet,
+			path:           "/metrics",
+			expectedStatus: http.StatusOK,
+		},
+		{
+			name:           "POST metrics endpoint not allowed",
+			method:         http.MethodPost,
+			path:           "/metrics",
+			expectedStatus: http.StatusOK,
+		},
+		{
+			name:           "PUT metrics endpoint not allowed",
+			method:         http.MethodPut,
+			path:           "/metrics",
+			expectedStatus: http.StatusOK,
+		},
+		{
+			name:           "DELETE metrics endpoint not allowed",
+			method:         http.MethodDelete,
+			path:           "/metrics",
+			expectedStatus: http.StatusOK,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := chi.NewRouter()
+			handler := MakeHandler(r, testServiceName, testInstanceID)
+
+			req, err := http.NewRequest(tt.method, tt.path, nil)
+			require.NoError(t, err)
+
+			rr := httptest.NewRecorder()
+			handler.ServeHTTP(rr, req)
+
+			assert.Equal(t, tt.expectedStatus, rr.Code)
+
+			if tt.expectedStatus == http.StatusOK {
+				// Verify content type for Prometheus metrics
+				contentType := rr.Header().Get("Content-Type")
+				assert.Contains(t, contentType, "text/plain")
+
+				// Verify response contains Prometheus metrics format
+				body := rr.Body.String()
+				assert.Contains(t, body, "# HELP")
+				assert.Contains(t, body, "# TYPE")
+			}
+		})
+	}
+}
+
+func TestNotFoundEndpoint(t *testing.T) {
+	const (
+		testServiceName = "test-service"
+		testInstanceID  = "test-instance-123"
+	)
+
+	tests := []struct {
+		name           string
+		path           string
+		expectedStatus int
+	}{
+		{
+			name:           "root path not found",
+			path:           "/",
+			expectedStatus: http.StatusNotFound,
+		},
+		{
+			name:           "random path not found",
+			path:           "/random-path",
+			expectedStatus: http.StatusNotFound,
+		},
+		{
+			name:           "health typo not found",
+			path:           "/helth",
+			expectedStatus: http.StatusNotFound,
+		},
+		{
+			name:           "metrics typo not found",
+			path:           "/metric",
+			expectedStatus: http.StatusNotFound,
+		},
+		{
+			name:           "health with trailing slash",
+			path:           "/health/",
+			expectedStatus: http.StatusNotFound,
+		},
+		{
+			name:           "metrics with trailing slash",
+			path:           "/metrics/",
+			expectedStatus: http.StatusNotFound,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			r := chi.NewRouter()
+			handler := MakeHandler(r, testServiceName, testInstanceID)
+
+			req, err := http.NewRequest(http.MethodGet, tt.path, nil)
+			require.NoError(t, err)
+
+			rr := httptest.NewRecorder()
+			handler.ServeHTTP(rr, req)
+
+			assert.Equal(t, tt.expectedStatus, rr.Code)
+		})
+	}
+}
+
+func TestConcurrentRequests(t *testing.T) {
+	const (
+		testServiceName = "test-service"
+		testInstanceID  = "test-instance-123"
+		numRequests     = 100
+	)
+
+	r := chi.NewRouter()
+	handler := MakeHandler(r, testServiceName, testInstanceID)
+
+	// Test concurrent health requests
+	t.Run("concurrent health requests", func(t *testing.T) {
+		results := make(chan int, numRequests)
+
+		for i := 0; i < numRequests; i++ {
+			go func() {
+				req, err := http.NewRequest(http.MethodGet, "/health", nil)
+				require.NoError(t, err)
+
+				rr := httptest.NewRecorder()
+				handler.ServeHTTP(rr, req)
+
+				results <- rr.Code
+			}()
+		}
+
+		// Collect all results
+		for i := 0; i < numRequests; i++ {
+			status := <-results
+			assert.Equal(t, http.StatusOK, status)
+		}
+	})
+
+	// Test concurrent metrics requests
+	t.Run("concurrent metrics requests", func(t *testing.T) {
+		results := make(chan int, numRequests)
+
+		for i := 0; i < numRequests; i++ {
+			go func() {
+				req, err := http.NewRequest(http.MethodGet, "/metrics", nil)
+				require.NoError(t, err)
+
+				rr := httptest.NewRecorder()
+				handler.ServeHTTP(rr, req)
+
+				results <- rr.Code
+			}()
+		}
+
+		// Collect all results
+		for i := 0; i < numRequests; i++ {
+			status := <-results
+			assert.Equal(t, http.StatusOK, status)
+		}
+	})
+}
+
+func TestHandlerWithCustomRouter(t *testing.T) {
+	const (
+		testServiceName = "test-service"
+		testInstanceID  = "test-instance-123"
+	)
+
+	// Test with a router that already has some routes
+	r := chi.NewRouter()
+	r.Get("/existing", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		if _, err := w.Write([]byte("existing")); err != nil {
+			http.Error(w, "Failed to write response", http.StatusInternalServerError)
+		}
+	})
+
+	handler := MakeHandler(r, testServiceName, testInstanceID)
+
+	// Test that existing route still works
+	req, err := http.NewRequest(http.MethodGet, "/existing", nil)
+	require.NoError(t, err)
+
+	rr := httptest.NewRecorder()
+	handler.ServeHTTP(rr, req)
+
+	assert.Equal(t, http.StatusOK, rr.Code)
+	assert.Equal(t, "existing", rr.Body.String())
+
+	// Test that new routes work
+	req, err = http.NewRequest(http.MethodGet, "/health", nil)
+	require.NoError(t, err)
+
+	rr = httptest.NewRecorder()
+	handler.ServeHTTP(rr, req)
+
+	assert.Equal(t, http.StatusOK, rr.Code)
+}

manager/attestation_policy_test.go

@@ -47,10 +47,10 @@ func TestFetchAttestationPolicy(t *testing.T) {
 	testCases := []struct {
 		name           string
 		computationId  string
-		vmConfig       interface{}
+		vmConfig       any
 		binaryBehavior string
 		expectedError  string
-		expectedResult map[string]interface{}
+		expectedResult map[string]any
 	}{
 		{
 			name:           "Valid SEV-SNP configuration",
@@ -136,7 +136,7 @@ func TestFetchAttestationPolicy(t *testing.T) {
 				assert.NoError(t, err)
 				assert.NotNil(t, result)
 
-				var attestationPolicy map[string]interface{}
+				var attestationPolicy map[string]any
 				err = json.Unmarshal(result, &attestationPolicy)
 				assert.NoError(t, err)
 

manager/manager.pb.go

@@ -3,7 +3,7 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.36.5
+// 	protoc-gen-go v1.36.8
 // 	protoc        v5.29.0
 // source: manager/manager.proto
 
@@ -34,6 +34,7 @@ type CreateReq struct {
 	AgentCvmServerUrl    string                 `protobuf:"bytes,5,opt,name=agent_cvm_server_url,json=agentCvmServerUrl,proto3" json:"agent_cvm_server_url,omitempty"`
 	AgentCvmCaUrl        string                 `protobuf:"bytes,6,opt,name=agent_cvm_ca_url,json=agentCvmCaUrl,proto3" json:"agent_cvm_ca_url,omitempty"`
 	Ttl                  string                 `protobuf:"bytes,7,opt,name=ttl,proto3" json:"ttl,omitempty"`
+	AgentCertsToken      string                 `protobuf:"bytes,8,opt,name=agent_certs_token,json=agentCertsToken,proto3" json:"agent_certs_token,omitempty"`
 	unknownFields        protoimpl.UnknownFields
 	sizeCache            protoimpl.SizeCache
 }
@@ -117,6 +118,13 @@ func (x *CreateReq) GetTtl() string {
 	return ""
 }
 
+func (x *CreateReq) GetAgentCertsToken() string {
+	if x != nil {
+		return x.AgentCertsToken
+	}
+	return ""
+}
+
 type CreateRes struct {
 	state         protoimpl.MessageState `protogen:"open.v1"`
 	ForwardedPort string                 `protobuf:"bytes,1,opt,name=forwarded_port,json=forwardedPort,proto3" json:"forwarded_port,omitempty"`
@@ -439,77 +447,46 @@ func (x *CVMInfoReq) GetId() string {
 
 var File_manager_manager_proto protoreflect.FileDescriptor
 
-var file_manager_manager_proto_rawDesc = string([]byte{
-	0x0a, 0x15, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2f, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65,
-	0x72, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x07, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72,
-	0x1a, 0x1b, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75,
-	0x66, 0x2f, 0x65, 0x6d, 0x70, 0x74, 0x79, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0xbb, 0x02,
-	0x0a, 0x09, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x12, 0x26, 0x0a, 0x0f, 0x61,
-	0x67, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x6f, 0x67, 0x5f, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x18, 0x01,
-	0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x4c, 0x65,
-	0x76, 0x65, 0x6c, 0x12, 0x36, 0x0a, 0x18, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x63, 0x76, 0x6d,
-	0x5f, 0x73, 0x65, 0x72, 0x76, 0x65, 0x72, 0x5f, 0x63, 0x61, 0x5f, 0x63, 0x65, 0x72, 0x74, 0x18,
-	0x02, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x14, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x43, 0x76, 0x6d, 0x53,
-	0x65, 0x72, 0x76, 0x65, 0x72, 0x43, 0x61, 0x43, 0x65, 0x72, 0x74, 0x12, 0x2f, 0x0a, 0x14, 0x61,
-	0x67, 0x65, 0x6e, 0x74, 0x5f, 0x63, 0x76, 0x6d, 0x5f, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x5f,
-	0x6b, 0x65, 0x79, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x11, 0x61, 0x67, 0x65, 0x6e, 0x74,
-	0x43, 0x76, 0x6d, 0x43, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x4b, 0x65, 0x79, 0x12, 0x31, 0x0a, 0x15,
-	0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x63, 0x76, 0x6d, 0x5f, 0x63, 0x6c, 0x69, 0x65, 0x6e, 0x74,
-	0x5f, 0x63, 0x65, 0x72, 0x74, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0c, 0x52, 0x12, 0x61, 0x67, 0x65,
-	0x6e, 0x74, 0x43, 0x76, 0x6d, 0x43, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x43, 0x65, 0x72, 0x74, 0x12,
-	0x2f, 0x0a, 0x14, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x63, 0x76, 0x6d, 0x5f, 0x73, 0x65, 0x72,
-	0x76, 0x65, 0x72, 0x5f, 0x75, 0x72, 0x6c, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x11, 0x61,
-	0x67, 0x65, 0x6e, 0x74, 0x43, 0x76, 0x6d, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x55, 0x72, 0x6c,
-	0x12, 0x27, 0x0a, 0x10, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x63, 0x76, 0x6d, 0x5f, 0x63, 0x61,
-	0x5f, 0x75, 0x72, 0x6c, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x61, 0x67, 0x65, 0x6e,
-	0x74, 0x43, 0x76, 0x6d, 0x43, 0x61, 0x55, 0x72, 0x6c, 0x12, 0x10, 0x0a, 0x03, 0x74, 0x74, 0x6c,
-	0x18, 0x07, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x74, 0x74, 0x6c, 0x22, 0x49, 0x0a, 0x09, 0x43,
-	0x72, 0x65, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x12, 0x25, 0x0a, 0x0e, 0x66, 0x6f, 0x72, 0x77,
-	0x61, 0x72, 0x64, 0x65, 0x64, 0x5f, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09,
-	0x52, 0x0d, 0x66, 0x6f, 0x72, 0x77, 0x61, 0x72, 0x64, 0x65, 0x64, 0x50, 0x6f, 0x72, 0x74, 0x12,
-	0x15, 0x0a, 0x06, 0x63, 0x76, 0x6d, 0x5f, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52,
-	0x05, 0x63, 0x76, 0x6d, 0x49, 0x64, 0x22, 0x22, 0x0a, 0x09, 0x52, 0x65, 0x6d, 0x6f, 0x76, 0x65,
-	0x52, 0x65, 0x71, 0x12, 0x15, 0x0a, 0x06, 0x63, 0x76, 0x6d, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x05, 0x63, 0x76, 0x6d, 0x49, 0x64, 0x22, 0x3a, 0x0a, 0x14, 0x41, 0x74,
-	0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x52,
-	0x65, 0x73, 0x12, 0x12, 0x0a, 0x04, 0x69, 0x6e, 0x66, 0x6f, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0c,
-	0x52, 0x04, 0x69, 0x6e, 0x66, 0x6f, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x02, 0x20, 0x01,
-	0x28, 0x09, 0x52, 0x02, 0x69, 0x64, 0x22, 0xb3, 0x01, 0x0a, 0x0a, 0x43, 0x56, 0x4d, 0x49, 0x6e,
-	0x66, 0x6f, 0x52, 0x65, 0x73, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28,
-	0x09, 0x52, 0x02, 0x69, 0x64, 0x12, 0x21, 0x0a, 0x0c, 0x6f, 0x76, 0x6d, 0x66, 0x5f, 0x76, 0x65,
-	0x72, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0b, 0x6f, 0x76, 0x6d,
-	0x66, 0x56, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x12, 0x17, 0x0a, 0x07, 0x63, 0x70, 0x75, 0x5f,
-	0x6e, 0x75, 0x6d, 0x18, 0x03, 0x20, 0x01, 0x28, 0x05, 0x52, 0x06, 0x63, 0x70, 0x75, 0x4e, 0x75,
-	0x6d, 0x12, 0x19, 0x0a, 0x08, 0x63, 0x70, 0x75, 0x5f, 0x74, 0x79, 0x70, 0x65, 0x18, 0x04, 0x20,
-	0x01, 0x28, 0x09, 0x52, 0x07, 0x63, 0x70, 0x75, 0x54, 0x79, 0x70, 0x65, 0x12, 0x1d, 0x0a, 0x0a,
-	0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x5f, 0x63, 0x6d, 0x64, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09,
-	0x52, 0x09, 0x6b, 0x65, 0x72, 0x6e, 0x65, 0x6c, 0x43, 0x6d, 0x64, 0x12, 0x1f, 0x0a, 0x0b, 0x65,
-	0x6f, 0x73, 0x5f, 0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09,
-	0x52, 0x0a, 0x65, 0x6f, 0x73, 0x56, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x22, 0x26, 0x0a, 0x14,
-	0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x50, 0x6f, 0x6c, 0x69, 0x63,
-	0x79, 0x52, 0x65, 0x71, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09,
-	0x52, 0x02, 0x69, 0x64, 0x22, 0x1c, 0x0a, 0x0a, 0x43, 0x56, 0x4d, 0x49, 0x6e, 0x66, 0x6f, 0x52,
-	0x65, 0x71, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x02,
-	0x69, 0x64, 0x32, 0x8c, 0x02, 0x0a, 0x0e, 0x4d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x53, 0x65,
-	0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x34, 0x0a, 0x08, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x56,
-	0x6d, 0x12, 0x12, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x43, 0x72, 0x65, 0x61,
-	0x74, 0x65, 0x52, 0x65, 0x71, 0x1a, 0x12, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e,
-	0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x52, 0x65, 0x73, 0x22, 0x00, 0x12, 0x38, 0x0a, 0x08, 0x52,
-	0x65, 0x6d, 0x6f, 0x76, 0x65, 0x56, 0x6d, 0x12, 0x12, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65,
-	0x72, 0x2e, 0x52, 0x65, 0x6d, 0x6f, 0x76, 0x65, 0x52, 0x65, 0x71, 0x1a, 0x16, 0x2e, 0x67, 0x6f,
-	0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d,
-	0x70, 0x74, 0x79, 0x22, 0x00, 0x12, 0x35, 0x0a, 0x07, 0x43, 0x56, 0x4d, 0x49, 0x6e, 0x66, 0x6f,
-	0x12, 0x13, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x43, 0x56, 0x4d, 0x49, 0x6e,
-	0x66, 0x6f, 0x52, 0x65, 0x71, 0x1a, 0x13, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e,
-	0x43, 0x56, 0x4d, 0x49, 0x6e, 0x66, 0x6f, 0x52, 0x65, 0x73, 0x22, 0x00, 0x12, 0x53, 0x0a, 0x11,
-	0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x50, 0x6f, 0x6c, 0x69, 0x63,
-	0x79, 0x12, 0x1d, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x41, 0x74, 0x74, 0x65,
-	0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x52, 0x65, 0x71,
-	0x1a, 0x1d, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73,
-	0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x50, 0x6f, 0x6c, 0x69, 0x63, 0x79, 0x52, 0x65, 0x73, 0x22,
-	0x00, 0x42, 0x0b, 0x5a, 0x09, 0x2e, 0x2f, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x62, 0x06,
-	0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
-})
+const file_manager_manager_proto_rawDesc = "" +
+	"\n" +
+	"\x15manager/manager.proto\x12\amanager\x1a\x1bgoogle/protobuf/empty.proto\"\xe7\x02\n" +
+	"\tCreateReq\x12&\n" +
+	"\x0fagent_log_level\x18\x01 \x01(\tR\ragentLogLevel\x126\n" +
+	"\x18agent_cvm_server_ca_cert\x18\x02 \x01(\fR\x14agentCvmServerCaCert\x12/\n" +
+	"\x14agent_cvm_client_key\x18\x03 \x01(\fR\x11agentCvmClientKey\x121\n" +
+	"\x15agent_cvm_client_cert\x18\x04 \x01(\fR\x12agentCvmClientCert\x12/\n" +
+	"\x14agent_cvm_server_url\x18\x05 \x01(\tR\x11agentCvmServerUrl\x12'\n" +
+	"\x10agent_cvm_ca_url\x18\x06 \x01(\tR\ragentCvmCaUrl\x12\x10\n" +
+	"\x03ttl\x18\a \x01(\tR\x03ttl\x12*\n" +
+	"\x11agent_certs_token\x18\b \x01(\tR\x0fagentCertsToken\"I\n" +
+	"\tCreateRes\x12%\n" +
+	"\x0eforwarded_port\x18\x01 \x01(\tR\rforwardedPort\x12\x15\n" +
+	"\x06cvm_id\x18\x02 \x01(\tR\x05cvmId\"\"\n" +
+	"\tRemoveReq\x12\x15\n" +
+	"\x06cvm_id\x18\x01 \x01(\tR\x05cvmId\":\n" +
+	"\x14AttestationPolicyRes\x12\x12\n" +
+	"\x04info\x18\x01 \x01(\fR\x04info\x12\x0e\n" +
+	"\x02id\x18\x02 \x01(\tR\x02id\"\xb3\x01\n" +
+	"\n" +
+	"CVMInfoRes\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id\x12!\n" +
+	"\fovmf_version\x18\x02 \x01(\tR\vovmfVersion\x12\x17\n" +
+	"\acpu_num\x18\x03 \x01(\x05R\x06cpuNum\x12\x19\n" +
+	"\bcpu_type\x18\x04 \x01(\tR\acpuType\x12\x1d\n" +
+	"\n" +
+	"kernel_cmd\x18\x05 \x01(\tR\tkernelCmd\x12\x1f\n" +
+	"\veos_version\x18\x06 \x01(\tR\n" +
+	"eosVersion\"&\n" +
+	"\x14AttestationPolicyReq\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id\"\x1c\n" +
+	"\n" +
+	"CVMInfoReq\x12\x0e\n" +
+	"\x02id\x18\x01 \x01(\tR\x02id2\x8c\x02\n" +
+	"\x0eManagerService\x124\n" +
+	"\bCreateVm\x12\x12.manager.CreateReq\x1a\x12.manager.CreateRes\"\x00\x128\n" +
+	"\bRemoveVm\x12\x12.manager.RemoveReq\x1a\x16.google.protobuf.Empty\"\x00\x125\n" +
+	"\aCVMInfo\x12\x13.manager.CVMInfoReq\x1a\x13.manager.CVMInfoRes\"\x00\x12S\n" +
+	"\x11AttestationPolicy\x12\x1d.manager.AttestationPolicyReq\x1a\x1d.manager.AttestationPolicyRes\"\x00B\vZ\t./managerb\x06proto3"
 
 var (
 	file_manager_manager_proto_rawDescOnce sync.Once

manager/manager.proto

@@ -24,6 +24,7 @@ message CreateReq{
   string agent_cvm_server_url = 5;
   string agent_cvm_ca_url = 6;
   string ttl = 7;
+  string agent_certs_token = 8;
 }
 
 message CreateRes{

manager/mocks/managerserviceclient.go

@@ -1,21 +1,35 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
-	context "context"
-
-	grpc "google.golang.org/grpc"
-	emptypb "google.golang.org/protobuf/types/known/emptypb"
-
-	manager "github.com/ultravioletrs/cocos/manager"
+	"context"
 
 	mock "github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/manager"
+	"google.golang.org/grpc"
+	"google.golang.org/protobuf/types/known/emptypb"
 )
 
+// 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
+}
+
 // ManagerServiceClient is an autogenerated mock type for the ManagerServiceClient type
 type ManagerServiceClient struct {
 	mock.Mock
@@ -29,8 +43,9 @@ 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) {
+// AttestationPolicy provides a mock function for the type ManagerServiceClient
+func (_mock *ManagerServiceClient) AttestationPolicy(ctx context.Context, in *manager.AttestationPolicyReq, opts ...grpc.CallOption) (*manager.AttestationPolicyRes, error) {
+	// grpc.CallOption
 	_va := make([]interface{}, len(opts))
 	for _i := range opts {
 		_va[_i] = opts[_i]
@@ -38,7 +53,7 @@ func (_m *ManagerServiceClient) AttestationPolicy(ctx context.Context, in *manag
 	var _ca []interface{}
 	_ca = append(_ca, ctx, in)
 	_ca = append(_ca, _va...)
-	ret := _m.Called(_ca...)
+	ret := _mock.Called(_ca...)
 
 	if len(ret) == 0 {
 		panic("no return value specified for AttestationPolicy")
@@ -46,23 +61,21 @@ func (_m *ManagerServiceClient) AttestationPolicy(ctx context.Context, in *manag
 
 	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 returnFunc, ok := ret.Get(0).(func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) (*manager.AttestationPolicyRes, error)); ok {
+		return returnFunc(ctx, in, opts...)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) *manager.AttestationPolicyRes); ok {
-		r0 = rf(ctx, in, opts...)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) *manager.AttestationPolicyRes); ok {
+		r0 = returnFunc(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...)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) error); ok {
+		r1 = returnFunc(ctx, in, opts...)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -82,29 +95,44 @@ func (_e *ManagerServiceClient_Expecter) AttestationPolicy(ctx interface{}, in i
 
 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) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 *manager.AttestationPolicyReq
+		if args[1] != nil {
+			arg1 = args[1].(*manager.AttestationPolicyReq)
+		}
+		var arg2 []grpc.CallOption
 		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...)
+		arg2 = variadicArgs
+		run(
+			arg0,
+			arg1,
+			arg2...,
+		)
 	})
 	return _c
 }
 
-func (_c *ManagerServiceClient_AttestationPolicy_Call) Return(_a0 *manager.AttestationPolicyRes, _a1 error) *ManagerServiceClient_AttestationPolicy_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *ManagerServiceClient_AttestationPolicy_Call) Return(attestationPolicyRes *manager.AttestationPolicyRes, err error) *ManagerServiceClient_AttestationPolicy_Call {
+	_c.Call.Return(attestationPolicyRes, err)
 	return _c
 }
 
-func (_c *ManagerServiceClient_AttestationPolicy_Call) RunAndReturn(run func(context.Context, *manager.AttestationPolicyReq, ...grpc.CallOption) (*manager.AttestationPolicyRes, error)) *ManagerServiceClient_AttestationPolicy_Call {
+func (_c *ManagerServiceClient_AttestationPolicy_Call) RunAndReturn(run func(ctx context.Context, in *manager.AttestationPolicyReq, opts ...grpc.CallOption) (*manager.AttestationPolicyRes, error)) *ManagerServiceClient_AttestationPolicy_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// CVMInfo provides a mock function with given fields: ctx, in, opts
-func (_m *ManagerServiceClient) CVMInfo(ctx context.Context, in *manager.CVMInfoReq, opts ...grpc.CallOption) (*manager.CVMInfoRes, error) {
+// CVMInfo provides a mock function for the type ManagerServiceClient
+func (_mock *ManagerServiceClient) CVMInfo(ctx context.Context, in *manager.CVMInfoReq, opts ...grpc.CallOption) (*manager.CVMInfoRes, error) {
+	// grpc.CallOption
 	_va := make([]interface{}, len(opts))
 	for _i := range opts {
 		_va[_i] = opts[_i]
@@ -112,7 +140,7 @@ func (_m *ManagerServiceClient) CVMInfo(ctx context.Context, in *manager.CVMInfo
 	var _ca []interface{}
 	_ca = append(_ca, ctx, in)
 	_ca = append(_ca, _va...)
-	ret := _m.Called(_ca...)
+	ret := _mock.Called(_ca...)
 
 	if len(ret) == 0 {
 		panic("no return value specified for CVMInfo")
@@ -120,23 +148,21 @@ func (_m *ManagerServiceClient) CVMInfo(ctx context.Context, in *manager.CVMInfo
 
 	var r0 *manager.CVMInfoRes
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, *manager.CVMInfoReq, ...grpc.CallOption) (*manager.CVMInfoRes, error)); ok {
-		return rf(ctx, in, opts...)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *manager.CVMInfoReq, ...grpc.CallOption) (*manager.CVMInfoRes, error)); ok {
+		return returnFunc(ctx, in, opts...)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, *manager.CVMInfoReq, ...grpc.CallOption) *manager.CVMInfoRes); ok {
-		r0 = rf(ctx, in, opts...)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *manager.CVMInfoReq, ...grpc.CallOption) *manager.CVMInfoRes); ok {
+		r0 = returnFunc(ctx, in, opts...)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).(*manager.CVMInfoRes)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context, *manager.CVMInfoReq, ...grpc.CallOption) error); ok {
-		r1 = rf(ctx, in, opts...)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, *manager.CVMInfoReq, ...grpc.CallOption) error); ok {
+		r1 = returnFunc(ctx, in, opts...)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -156,29 +182,44 @@ func (_e *ManagerServiceClient_Expecter) CVMInfo(ctx interface{}, in interface{}
 
 func (_c *ManagerServiceClient_CVMInfo_Call) Run(run func(ctx context.Context, in *manager.CVMInfoReq, opts ...grpc.CallOption)) *ManagerServiceClient_CVMInfo_Call {
 	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 *manager.CVMInfoReq
+		if args[1] != nil {
+			arg1 = args[1].(*manager.CVMInfoReq)
+		}
+		var arg2 []grpc.CallOption
 		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.CVMInfoReq), variadicArgs...)
+		arg2 = variadicArgs
+		run(
+			arg0,
+			arg1,
+			arg2...,
+		)
 	})
 	return _c
 }
 
-func (_c *ManagerServiceClient_CVMInfo_Call) Return(_a0 *manager.CVMInfoRes, _a1 error) *ManagerServiceClient_CVMInfo_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *ManagerServiceClient_CVMInfo_Call) Return(cVMInfoRes *manager.CVMInfoRes, err error) *ManagerServiceClient_CVMInfo_Call {
+	_c.Call.Return(cVMInfoRes, err)
 	return _c
 }
 
-func (_c *ManagerServiceClient_CVMInfo_Call) RunAndReturn(run func(context.Context, *manager.CVMInfoReq, ...grpc.CallOption) (*manager.CVMInfoRes, error)) *ManagerServiceClient_CVMInfo_Call {
+func (_c *ManagerServiceClient_CVMInfo_Call) RunAndReturn(run func(ctx context.Context, in *manager.CVMInfoReq, opts ...grpc.CallOption) (*manager.CVMInfoRes, error)) *ManagerServiceClient_CVMInfo_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) {
+// CreateVm provides a mock function for the type ManagerServiceClient
+func (_mock *ManagerServiceClient) CreateVm(ctx context.Context, in *manager.CreateReq, opts ...grpc.CallOption) (*manager.CreateRes, error) {
+	// grpc.CallOption
 	_va := make([]interface{}, len(opts))
 	for _i := range opts {
 		_va[_i] = opts[_i]
@@ -186,7 +227,7 @@ func (_m *ManagerServiceClient) CreateVm(ctx context.Context, in *manager.Create
 	var _ca []interface{}
 	_ca = append(_ca, ctx, in)
 	_ca = append(_ca, _va...)
-	ret := _m.Called(_ca...)
+	ret := _mock.Called(_ca...)
 
 	if len(ret) == 0 {
 		panic("no return value specified for CreateVm")
@@ -194,23 +235,21 @@ func (_m *ManagerServiceClient) CreateVm(ctx context.Context, in *manager.Create
 
 	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 returnFunc, ok := ret.Get(0).(func(context.Context, *manager.CreateReq, ...grpc.CallOption) (*manager.CreateRes, error)); ok {
+		return returnFunc(ctx, in, opts...)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, *manager.CreateReq, ...grpc.CallOption) *manager.CreateRes); ok {
-		r0 = rf(ctx, in, opts...)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *manager.CreateReq, ...grpc.CallOption) *manager.CreateRes); ok {
+		r0 = returnFunc(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...)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, *manager.CreateReq, ...grpc.CallOption) error); ok {
+		r1 = returnFunc(ctx, in, opts...)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -230,29 +269,44 @@ func (_e *ManagerServiceClient_Expecter) CreateVm(ctx interface{}, in interface{
 
 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) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 *manager.CreateReq
+		if args[1] != nil {
+			arg1 = args[1].(*manager.CreateReq)
+		}
+		var arg2 []grpc.CallOption
 		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...)
+		arg2 = variadicArgs
+		run(
+			arg0,
+			arg1,
+			arg2...,
+		)
 	})
 	return _c
 }
 
-func (_c *ManagerServiceClient_CreateVm_Call) Return(_a0 *manager.CreateRes, _a1 error) *ManagerServiceClient_CreateVm_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *ManagerServiceClient_CreateVm_Call) Return(createRes *manager.CreateRes, err error) *ManagerServiceClient_CreateVm_Call {
+	_c.Call.Return(createRes, err)
 	return _c
 }
 
-func (_c *ManagerServiceClient_CreateVm_Call) RunAndReturn(run func(context.Context, *manager.CreateReq, ...grpc.CallOption) (*manager.CreateRes, error)) *ManagerServiceClient_CreateVm_Call {
+func (_c *ManagerServiceClient_CreateVm_Call) RunAndReturn(run func(ctx context.Context, in *manager.CreateReq, opts ...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) {
+// RemoveVm provides a mock function for the type ManagerServiceClient
+func (_mock *ManagerServiceClient) RemoveVm(ctx context.Context, in *manager.RemoveReq, opts ...grpc.CallOption) (*emptypb.Empty, error) {
+	// grpc.CallOption
 	_va := make([]interface{}, len(opts))
 	for _i := range opts {
 		_va[_i] = opts[_i]
@@ -260,7 +314,7 @@ func (_m *ManagerServiceClient) RemoveVm(ctx context.Context, in *manager.Remove
 	var _ca []interface{}
 	_ca = append(_ca, ctx, in)
 	_ca = append(_ca, _va...)
-	ret := _m.Called(_ca...)
+	ret := _mock.Called(_ca...)
 
 	if len(ret) == 0 {
 		panic("no return value specified for RemoveVm")
@@ -268,23 +322,21 @@ func (_m *ManagerServiceClient) RemoveVm(ctx context.Context, in *manager.Remove
 
 	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 returnFunc, ok := ret.Get(0).(func(context.Context, *manager.RemoveReq, ...grpc.CallOption) (*emptypb.Empty, error)); ok {
+		return returnFunc(ctx, in, opts...)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, *manager.RemoveReq, ...grpc.CallOption) *emptypb.Empty); ok {
-		r0 = rf(ctx, in, opts...)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *manager.RemoveReq, ...grpc.CallOption) *emptypb.Empty); ok {
+		r0 = returnFunc(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...)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, *manager.RemoveReq, ...grpc.CallOption) error); ok {
+		r1 = returnFunc(ctx, in, opts...)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -304,37 +356,37 @@ func (_e *ManagerServiceClient_Expecter) RemoveVm(ctx interface{}, in interface{
 
 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) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 *manager.RemoveReq
+		if args[1] != nil {
+			arg1 = args[1].(*manager.RemoveReq)
+		}
+		var arg2 []grpc.CallOption
 		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...)
+		arg2 = variadicArgs
+		run(
+			arg0,
+			arg1,
+			arg2...,
+		)
 	})
 	return _c
 }
 
-func (_c *ManagerServiceClient_RemoveVm_Call) Return(_a0 *emptypb.Empty, _a1 error) *ManagerServiceClient_RemoveVm_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *ManagerServiceClient_RemoveVm_Call) Return(empty *emptypb.Empty, err error) *ManagerServiceClient_RemoveVm_Call {
+	_c.Call.Return(empty, err)
 	return _c
 }
 
-func (_c *ManagerServiceClient_RemoveVm_Call) RunAndReturn(run func(context.Context, *manager.RemoveReq, ...grpc.CallOption) (*emptypb.Empty, error)) *ManagerServiceClient_RemoveVm_Call {
+func (_c *ManagerServiceClient_RemoveVm_Call) RunAndReturn(run func(ctx context.Context, in *manager.RemoveReq, opts ...grpc.CallOption) (*emptypb.Empty, error)) *ManagerServiceClient_RemoveVm_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
-}

manager/mocks/service.go

@@ -1,17 +1,33 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
-	context "context"
+	"context"
 
 	mock "github.com/stretchr/testify/mock"
-	manager "github.com/ultravioletrs/cocos/manager"
+	"github.com/ultravioletrs/cocos/manager"
 )
 
+// 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 {
+	mock.TestingT
+	Cleanup(func())
+}) *Service {
+	mock := &Service{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // Service is an autogenerated mock type for the Service type
 type Service struct {
 	mock.Mock
@@ -25,9 +41,9 @@ func (_m *Service) EXPECT() *Service_Expecter {
 	return &Service_Expecter{mock: &_m.Mock}
 }
 
-// CreateVM provides a mock function with given fields: ctx, req
-func (_m *Service) CreateVM(ctx context.Context, req *manager.CreateReq) (string, string, error) {
-	ret := _m.Called(ctx, req)
+// CreateVM provides a mock function for the type Service
+func (_mock *Service) CreateVM(ctx context.Context, req *manager.CreateReq) (string, string, error) {
+	ret := _mock.Called(ctx, req)
 
 	if len(ret) == 0 {
 		panic("no return value specified for CreateVM")
@@ -36,27 +52,24 @@ func (_m *Service) CreateVM(ctx context.Context, req *manager.CreateReq) (string
 	var r0 string
 	var r1 string
 	var r2 error
-	if rf, ok := ret.Get(0).(func(context.Context, *manager.CreateReq) (string, string, error)); ok {
-		return rf(ctx, req)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *manager.CreateReq) (string, string, error)); ok {
+		return returnFunc(ctx, req)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, *manager.CreateReq) string); ok {
-		r0 = rf(ctx, req)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, *manager.CreateReq) string); ok {
+		r0 = returnFunc(ctx, req)
 	} else {
 		r0 = ret.Get(0).(string)
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context, *manager.CreateReq) string); ok {
-		r1 = rf(ctx, req)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, *manager.CreateReq) string); ok {
+		r1 = returnFunc(ctx, req)
 	} else {
 		r1 = ret.Get(1).(string)
 	}
-
-	if rf, ok := ret.Get(2).(func(context.Context, *manager.CreateReq) error); ok {
-		r2 = rf(ctx, req)
+	if returnFunc, ok := ret.Get(2).(func(context.Context, *manager.CreateReq) error); ok {
+		r2 = returnFunc(ctx, req)
 	} else {
 		r2 = ret.Error(2)
 	}
-
 	return r0, r1, r2
 }
 
@@ -74,24 +87,35 @@ func (_e *Service_Expecter) CreateVM(ctx interface{}, req interface{}) *Service_
 
 func (_c *Service_CreateVM_Call) Run(run func(ctx context.Context, req *manager.CreateReq)) *Service_CreateVM_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(*manager.CreateReq))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 *manager.CreateReq
+		if args[1] != nil {
+			arg1 = args[1].(*manager.CreateReq)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_CreateVM_Call) Return(_a0 string, _a1 string, _a2 error) *Service_CreateVM_Call {
-	_c.Call.Return(_a0, _a1, _a2)
+func (_c *Service_CreateVM_Call) Return(s string, s1 string, err error) *Service_CreateVM_Call {
+	_c.Call.Return(s, s1, err)
 	return _c
 }
 
-func (_c *Service_CreateVM_Call) RunAndReturn(run func(context.Context, *manager.CreateReq) (string, string, error)) *Service_CreateVM_Call {
+func (_c *Service_CreateVM_Call) RunAndReturn(run func(ctx context.Context, req *manager.CreateReq) (string, string, error)) *Service_CreateVM_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// FetchAttestationPolicy provides a mock function with given fields: ctx, computationID
-func (_m *Service) FetchAttestationPolicy(ctx context.Context, computationID string) ([]byte, error) {
-	ret := _m.Called(ctx, computationID)
+// FetchAttestationPolicy provides a mock function for the type Service
+func (_mock *Service) FetchAttestationPolicy(ctx context.Context, computationID string) ([]byte, error) {
+	ret := _mock.Called(ctx, computationID)
 
 	if len(ret) == 0 {
 		panic("no return value specified for FetchAttestationPolicy")
@@ -99,23 +123,21 @@ func (_m *Service) FetchAttestationPolicy(ctx context.Context, computationID str
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, string) ([]byte, error)); ok {
-		return rf(ctx, computationID)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, string) ([]byte, error)); ok {
+		return returnFunc(ctx, computationID)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context, string) []byte); ok {
-		r0 = rf(ctx, computationID)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, string) []byte); ok {
+		r0 = returnFunc(ctx, computationID)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
-		r1 = rf(ctx, computationID)
+	if returnFunc, ok := ret.Get(1).(func(context.Context, string) error); ok {
+		r1 = returnFunc(ctx, computationID)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -133,36 +155,46 @@ func (_e *Service_Expecter) FetchAttestationPolicy(ctx interface{}, computationI
 
 func (_c *Service_FetchAttestationPolicy_Call) Run(run func(ctx context.Context, computationID string)) *Service_FetchAttestationPolicy_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(string))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 string
+		if args[1] != nil {
+			arg1 = args[1].(string)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_FetchAttestationPolicy_Call) Return(_a0 []byte, _a1 error) *Service_FetchAttestationPolicy_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Service_FetchAttestationPolicy_Call) Return(bytes []byte, err error) *Service_FetchAttestationPolicy_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Service_FetchAttestationPolicy_Call) RunAndReturn(run func(context.Context, string) ([]byte, error)) *Service_FetchAttestationPolicy_Call {
+func (_c *Service_FetchAttestationPolicy_Call) RunAndReturn(run func(ctx context.Context, computationID string) ([]byte, error)) *Service_FetchAttestationPolicy_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// RemoveVM provides a mock function with given fields: ctx, computationID
-func (_m *Service) RemoveVM(ctx context.Context, computationID string) error {
-	ret := _m.Called(ctx, computationID)
+// RemoveVM provides a mock function for the type Service
+func (_mock *Service) RemoveVM(ctx context.Context, computationID string) error {
+	ret := _mock.Called(ctx, computationID)
 
 	if len(ret) == 0 {
 		panic("no return value specified for RemoveVM")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(context.Context, string) error); ok {
-		r0 = rf(ctx, computationID)
+	if returnFunc, ok := ret.Get(0).(func(context.Context, string) error); ok {
+		r0 = returnFunc(ctx, computationID)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -180,24 +212,35 @@ func (_e *Service_Expecter) RemoveVM(ctx interface{}, computationID interface{})
 
 func (_c *Service_RemoveVM_Call) Run(run func(ctx context.Context, computationID string)) *Service_RemoveVM_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context), args[1].(string))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		var arg1 string
+		if args[1] != nil {
+			arg1 = args[1].(string)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_RemoveVM_Call) Return(_a0 error) *Service_RemoveVM_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_RemoveVM_Call) Return(err error) *Service_RemoveVM_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Service_RemoveVM_Call) RunAndReturn(run func(context.Context, string) error) *Service_RemoveVM_Call {
+func (_c *Service_RemoveVM_Call) RunAndReturn(run func(ctx context.Context, computationID string) error) *Service_RemoveVM_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// ReturnCVMInfo provides a mock function with given fields: ctx
-func (_m *Service) ReturnCVMInfo(ctx context.Context) (string, int, string, string) {
-	ret := _m.Called(ctx)
+// ReturnCVMInfo provides a mock function for the type Service
+func (_mock *Service) ReturnCVMInfo(ctx context.Context) (string, int, string, string) {
+	ret := _mock.Called(ctx)
 
 	if len(ret) == 0 {
 		panic("no return value specified for ReturnCVMInfo")
@@ -207,33 +250,29 @@ func (_m *Service) ReturnCVMInfo(ctx context.Context) (string, int, string, stri
 	var r1 int
 	var r2 string
 	var r3 string
-	if rf, ok := ret.Get(0).(func(context.Context) (string, int, string, string)); ok {
-		return rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) (string, int, string, string)); ok {
+		return returnFunc(ctx)
 	}
-	if rf, ok := ret.Get(0).(func(context.Context) string); ok {
-		r0 = rf(ctx)
+	if returnFunc, ok := ret.Get(0).(func(context.Context) string); ok {
+		r0 = returnFunc(ctx)
 	} else {
 		r0 = ret.Get(0).(string)
 	}
-
-	if rf, ok := ret.Get(1).(func(context.Context) int); ok {
-		r1 = rf(ctx)
+	if returnFunc, ok := ret.Get(1).(func(context.Context) int); ok {
+		r1 = returnFunc(ctx)
 	} else {
 		r1 = ret.Get(1).(int)
 	}
-
-	if rf, ok := ret.Get(2).(func(context.Context) string); ok {
-		r2 = rf(ctx)
+	if returnFunc, ok := ret.Get(2).(func(context.Context) string); ok {
+		r2 = returnFunc(ctx)
 	} else {
 		r2 = ret.Get(2).(string)
 	}
-
-	if rf, ok := ret.Get(3).(func(context.Context) string); ok {
-		r3 = rf(ctx)
+	if returnFunc, ok := ret.Get(3).(func(context.Context) string); ok {
+		r3 = returnFunc(ctx)
 	} else {
 		r3 = ret.Get(3).(string)
 	}
-
 	return r0, r1, r2, r3
 }
 
@@ -250,36 +289,41 @@ func (_e *Service_Expecter) ReturnCVMInfo(ctx interface{}) *Service_ReturnCVMInf
 
 func (_c *Service_ReturnCVMInfo_Call) Run(run func(ctx context.Context)) *Service_ReturnCVMInfo_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(context.Context))
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Service_ReturnCVMInfo_Call) Return(_a0 string, _a1 int, _a2 string, _a3 string) *Service_ReturnCVMInfo_Call {
-	_c.Call.Return(_a0, _a1, _a2, _a3)
+func (_c *Service_ReturnCVMInfo_Call) Return(s string, n int, s1 string, s2 string) *Service_ReturnCVMInfo_Call {
+	_c.Call.Return(s, n, s1, s2)
 	return _c
 }
 
-func (_c *Service_ReturnCVMInfo_Call) RunAndReturn(run func(context.Context) (string, int, string, string)) *Service_ReturnCVMInfo_Call {
+func (_c *Service_ReturnCVMInfo_Call) RunAndReturn(run func(ctx context.Context) (string, int, string, string)) *Service_ReturnCVMInfo_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Shutdown provides a mock function with no fields
-func (_m *Service) Shutdown() error {
-	ret := _m.Called()
+// Shutdown provides a mock function for the type Service
+func (_mock *Service) Shutdown() error {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for Shutdown")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -300,8 +344,8 @@ func (_c *Service_Shutdown_Call) Run(run func()) *Service_Shutdown_Call {
 	return _c
 }
 
-func (_c *Service_Shutdown_Call) Return(_a0 error) *Service_Shutdown_Call {
-	_c.Call.Return(_a0)
+func (_c *Service_Shutdown_Call) Return(err error) *Service_Shutdown_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -309,17 +353,3 @@ func (_c *Service_Shutdown_Call) RunAndReturn(run func() error) *Service_Shutdow
 	_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 {
-	mock.TestingT
-	Cleanup(func())
-}) *Service {
-	mock := &Service{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

manager/monitoring/alert_rules.yml

@@ -0,0 +1,57 @@
+groups:
+  - name: cocos-manager-alerts
+    rules:
+      # Service Down Alert - Primary method using up metric
+      - alert: CocosManagerDown
+        expr: up{job="cocos-manager"} == 0
+        for: 1m
+        labels:
+          severity: critical
+          service: cocos-manager
+        annotations:
+          summary: "Cocos Manager service is down"
+          description: "Cocos Manager service has been down for more than 1 minute. Instance: {{ $labels.instance }}"
+
+      # Alternative: No metrics received (for cases where up metric might not be reliable)
+      - alert: CocosManagerNoMetrics
+        expr: absent(up{job="cocos-manager"})
+        for: 2m
+        labels:
+          severity: critical
+          service: cocos-manager
+        annotations:
+          summary: "No metrics received from Cocos Manager"
+          description: "No metrics have been received from Cocos Manager for more than 2 minutes"
+
+      # Health endpoint specific alert (if you expose health as a metric)
+      - alert: CocosManagerUnhealthy
+        expr: health_check{job="cocos-manager"} == 0
+        for: 30s
+        labels:
+          severity: warning
+          service: cocos-manager
+        annotations:
+          summary: "Cocos Manager health check failing"
+          description: "Cocos Manager health check has been failing for more than 30 seconds. Instance: {{ $labels.instance }}"
+
+      # High latency alert
+      - alert: CocosManagerHighLatency
+        expr: histogram_quantile(0.95, rate(cocos_manager_latency_bucket[5m])) > 2
+        for: 5m
+        labels:
+          severity: warning
+          service: cocos-manager
+        annotations:
+          summary: "Cocos Manager high latency"
+          description: "Cocos Manager 95th percentile latency is above 2 seconds for more than 5 minutes. Current value: {{ $value }}s"
+
+      # High error rate alert
+      - alert: CocosManagerHighErrorRate
+        expr: rate(cocos_manager_errors_total[5m]) > 0.1
+        for: 2m
+        labels:
+          severity: warning
+          service: cocos-manager
+        annotations:
+          summary: "Cocos Manager high error rate"
+          description: "Cocos Manager error rate is above 10% for more than 2 minutes. Current rate: {{ $value }}"

manager/monitoring/alertmanager.yml

@@ -0,0 +1,56 @@
+global:
+  smtp_smarthost: 'localhost:587'
+  smtp_from: 'alerts@yourcompany.com'
+
+route:
+  group_by: ['alertname', 'service']
+  group_wait: 10s
+  group_interval: 10s
+  repeat_interval: 1h
+  receiver: 'web.hook'
+  routes:
+    - match:
+        severity: critical
+      receiver: 'critical-alerts'
+    - match:
+        severity: warning
+      receiver: 'warning-alerts'
+
+receivers:
+  - name: 'web.hook'
+    webhook_configs:
+      - url: 'http://127.0.0.1:5001/'
+
+  - name: 'critical-alerts'
+    email_configs:
+      - to: 'oncall@yourcompany.com'
+        subject: 'CRITICAL: {{ .GroupLabels.service }} Alert'
+        body: |
+          {{ range .Alerts }}
+          Alert: {{ .Annotations.summary }}
+          Description: {{ .Annotations.description }}
+          Labels: {{ range .Labels.SortedPairs }}{{ .Name }}: {{ .Value }} {{ end }}
+          {{ end }}
+    slack_configs:
+      - api_url: 'YOUR_SLACK_WEBHOOK_URL'
+        channel: '#alerts'
+        title: 'CRITICAL Alert: {{ .GroupLabels.service }}'
+        text: '{{ range .Alerts }}{{ .Annotations.summary }}: {{ .Annotations.description }}{{ end }}'
+
+  - name: 'warning-alerts'
+    email_configs:
+      - to: 'team@yourcompany.com'
+        subject: 'WARNING: {{ .GroupLabels.service }} Alert'
+        body: |
+          {{ range .Alerts }}
+          Alert: {{ .Annotations.summary }}
+          Description: {{ .Annotations.description }}
+          Labels: {{ range .Labels.SortedPairs }}{{ .Name }}: {{ .Value }} {{ end }}
+          {{ end }}
+
+inhibit_rules:
+  - source_match:
+      severity: 'critical'
+    target_match:
+      severity: 'warning'
+    equal: ['alertname', 'service']

manager/monitoring/docker-compose.yml

@@ -0,0 +1,58 @@
+version: '3.8'
+
+services:
+  prometheus:
+    image: prom/prometheus:latest
+    container_name: prometheus
+    ports:
+      - "9090:9090"
+    volumes:
+      - ./prometheus.yml:/etc/prometheus/prometheus.yml
+      - ./alert_rules.yml:/etc/prometheus/alert_rules.yml
+      - prometheus_data:/prometheus
+    command:
+      - '--config.file=/etc/prometheus/prometheus.yml'
+      - '--storage.tsdb.path=/prometheus'
+      - '--web.console.libraries=/etc/prometheus/console_libraries'
+      - '--web.console.templates=/etc/prometheus/consoles'
+      - '--storage.tsdb.retention.time=200h'
+      - '--web.enable-lifecycle'
+      - '--web.enable-admin-api'
+    networks:
+      - monitoring
+
+  alertmanager:
+    image: prom/alertmanager:latest
+    container_name: alertmanager
+    ports:
+      - "9093:9093"
+    volumes:
+      - ./alertmanager.yml:/etc/alertmanager/alertmanager.yml
+      - alertmanager_data:/alertmanager
+    command:
+      - '--config.file=/etc/alertmanager/alertmanager.yml'
+      - '--storage.path=/alertmanager'
+      - '--web.external-url=http://localhost:9093'
+    networks:
+      - monitoring
+
+  grafana:
+    image: grafana/grafana:latest
+    container_name: grafana
+    ports:
+      - "3000:3000"
+    environment:
+      - GF_SECURITY_ADMIN_PASSWORD=admin
+    volumes:
+      - grafana_data:/var/lib/grafana
+    networks:
+      - monitoring
+
+volumes:
+  prometheus_data:
+  alertmanager_data:
+  grafana_data:
+
+networks:
+  monitoring:
+    driver: bridge

manager/monitoring/prometheus.yml

@@ -0,0 +1,20 @@
+global:
+  scrape_interval: 15s
+  evaluation_interval: 15s
+
+rule_files:
+  - "alert_rules.yml"
+
+alerting:
+  alertmanagers:
+    - static_configs:
+        - targets:
+          - alertmanager:9093
+
+scrape_configs:
+  - job_name: 'cocos-manager'
+    static_configs:
+      - targets: ['192.168.100.4:7003']  # Adjust to your manager http port
+    metrics_path: /metrics
+    scrape_interval: 10s
+    scrape_timeout: 5s

manager/qemu/config.go

@@ -10,7 +10,6 @@ import (
 )
 
 const (
-	BaseGuestCID      = 3
 	KernelCommandLine = "quiet console=null"
 	TDXObject         = "{\"qom-type\":\"tdx-guest\",\"id\":\"%s\",\"quote-generation-socket\":{\"type\": \"vsock\", \"cid\":\"2\",\"port\":\"%d\"}}"
 )
@@ -74,11 +73,6 @@ type IGVMConfig struct {
 	File string `env:"IGVM_FILE"      envDefault:"/root/coconut-qemu.igvm"`
 }
 
-type VSockConfig struct {
-	ID       string `env:"VSOCK_ID"        envDefault:"vhost-vsock-pci0"`
-	GuestCID int    `env:"VSOCK_GUEST_CID" envDefault:"3"`
-}
-
 type Config struct {
 	EnableSEVSNP bool
 	EnableTDX    bool
@@ -103,9 +97,6 @@ type Config struct {
 	NetDevConfig
 	VirtioNetPciConfig
 
-	// Vsock
-	VSockConfig
-
 	// disk
 	DiskImgConfig
 
@@ -186,8 +177,6 @@ func (config Config) ConstructQemuArgs() []string {
 			config.VirtioNetPciConfig.Addr,
 			config.VirtioNetPciConfig.ROMFile))
 
-	args = append(args, "-device", fmt.Sprintf("vhost-vsock-pci,id=%s,guest-cid=%d", config.VSockConfig.ID, config.VSockConfig.GuestCID))
-
 	// SEV-SNP
 	if config.EnableSEVSNP {
 		sevSnpType := "sev-snp-guest"

manager/qemu/config_test.go

@@ -51,10 +51,6 @@ func TestConstructQemuArgs(t *testing.T) {
 					IOMMUPlatform: true,
 					Addr:          "0x2",
 				},
-				VSockConfig: VSockConfig{
-					ID:       "vhost-vsock-pci0",
-					GuestCID: 3,
-				},
 				DiskImgConfig: DiskImgConfig{
 					KernelFile: "img/bzImage",
 					RootFsFile: "img/rootfs.cpio.gz",
@@ -72,7 +68,6 @@ func TestConstructQemuArgs(t *testing.T) {
 				"-drive", "if=pflash,format=raw,unit=1,file=/usr/share/OVMF/OVMF_VARS.fd",
 				"-netdev", "user,id=vmnic,hostfwd=tcp::7020-:7002",
 				"-device", "virtio-net-pci,disable-legacy=on,iommu_platform=true,netdev=vmnic,addr=0x2,romfile=",
-				"-device", "vhost-vsock-pci,id=vhost-vsock-pci0,guest-cid=3",
 				"-kernel", "img/bzImage",
 				"-append", "\"quiet console=null\"",
 				"-initrd", "img/rootfs.cpio.gz",
@@ -119,10 +114,6 @@ func TestConstructQemuArgs(t *testing.T) {
 					IOMMUPlatform: true,
 					Addr:          "0x2",
 				},
-				VSockConfig: VSockConfig{
-					ID:       "vhost-vsock-pci0",
-					GuestCID: 3,
-				},
 				DiskImgConfig: DiskImgConfig{
 					KernelFile: "img/bzImage",
 					RootFsFile: "img/rootfs.cpio.gz",
@@ -147,7 +138,6 @@ func TestConstructQemuArgs(t *testing.T) {
 				"-m", "2048M,slots=5,maxmem=30G",
 				"-netdev", "user,id=vmnic,hostfwd=tcp::7020-:7002",
 				"-device", "virtio-net-pci,disable-legacy=on,iommu_platform=true,netdev=vmnic,addr=0x2,romfile=",
-				"-device", "vhost-vsock-pci,id=vhost-vsock-pci0,guest-cid=3",
 				"-machine", "confidential-guest-support=sev0,memory-backend=ram1,igvm-cfg=igvm0",
 				"-object", "memory-backend-memfd,id=ram1,size=2048M,share=true,prealloc=false",
 				"-object", "sev-snp-guest,id=sev0,cbitpos=51,reduced-phys-bits=1",

manager/qemu/mocks/persistence.go

@@ -1,15 +1,31 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
 	mock "github.com/stretchr/testify/mock"
-	qemu "github.com/ultravioletrs/cocos/manager/qemu"
+	"github.com/ultravioletrs/cocos/manager/qemu"
 )
 
+// NewPersistence creates a new instance of Persistence. 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 NewPersistence(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Persistence {
+	mock := &Persistence{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // Persistence is an autogenerated mock type for the Persistence type
 type Persistence struct {
 	mock.Mock
@@ -23,21 +39,20 @@ func (_m *Persistence) EXPECT() *Persistence_Expecter {
 	return &Persistence_Expecter{mock: &_m.Mock}
 }
 
-// DeleteVM provides a mock function with given fields: id
-func (_m *Persistence) DeleteVM(id string) error {
-	ret := _m.Called(id)
+// DeleteVM provides a mock function for the type Persistence
+func (_mock *Persistence) DeleteVM(id string) error {
+	ret := _mock.Called(id)
 
 	if len(ret) == 0 {
 		panic("no return value specified for DeleteVM")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
-		r0 = rf(id)
+	if returnFunc, ok := ret.Get(0).(func(string) error); ok {
+		r0 = returnFunc(id)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -54,24 +69,30 @@ func (_e *Persistence_Expecter) DeleteVM(id interface{}) *Persistence_DeleteVM_C
 
 func (_c *Persistence_DeleteVM_Call) Run(run func(id string)) *Persistence_DeleteVM_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(string))
+		var arg0 string
+		if args[0] != nil {
+			arg0 = args[0].(string)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Persistence_DeleteVM_Call) Return(_a0 error) *Persistence_DeleteVM_Call {
-	_c.Call.Return(_a0)
+func (_c *Persistence_DeleteVM_Call) Return(err error) *Persistence_DeleteVM_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Persistence_DeleteVM_Call) RunAndReturn(run func(string) error) *Persistence_DeleteVM_Call {
+func (_c *Persistence_DeleteVM_Call) RunAndReturn(run func(id string) error) *Persistence_DeleteVM_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// LoadVMs provides a mock function with no fields
-func (_m *Persistence) LoadVMs() ([]qemu.VMState, error) {
-	ret := _m.Called()
+// LoadVMs provides a mock function for the type Persistence
+func (_mock *Persistence) LoadVMs() ([]qemu.VMState, error) {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for LoadVMs")
@@ -79,23 +100,21 @@ func (_m *Persistence) LoadVMs() ([]qemu.VMState, error) {
 
 	var r0 []qemu.VMState
 	var r1 error
-	if rf, ok := ret.Get(0).(func() ([]qemu.VMState, error)); ok {
-		return rf()
+	if returnFunc, ok := ret.Get(0).(func() ([]qemu.VMState, error)); ok {
+		return returnFunc()
 	}
-	if rf, ok := ret.Get(0).(func() []qemu.VMState); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() []qemu.VMState); ok {
+		r0 = returnFunc()
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]qemu.VMState)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func() error); ok {
-		r1 = rf()
+	if returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -116,8 +135,8 @@ func (_c *Persistence_LoadVMs_Call) Run(run func()) *Persistence_LoadVMs_Call {
 	return _c
 }
 
-func (_c *Persistence_LoadVMs_Call) Return(_a0 []qemu.VMState, _a1 error) *Persistence_LoadVMs_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Persistence_LoadVMs_Call) Return(vMStates []qemu.VMState, err error) *Persistence_LoadVMs_Call {
+	_c.Call.Return(vMStates, err)
 	return _c
 }
 
@@ -126,21 +145,20 @@ func (_c *Persistence_LoadVMs_Call) RunAndReturn(run func() ([]qemu.VMState, err
 	return _c
 }
 
-// SaveVM provides a mock function with given fields: state
-func (_m *Persistence) SaveVM(state qemu.VMState) error {
-	ret := _m.Called(state)
+// SaveVM provides a mock function for the type Persistence
+func (_mock *Persistence) SaveVM(state qemu.VMState) error {
+	ret := _mock.Called(state)
 
 	if len(ret) == 0 {
 		panic("no return value specified for SaveVM")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(qemu.VMState) error); ok {
-		r0 = rf(state)
+	if returnFunc, ok := ret.Get(0).(func(qemu.VMState) error); ok {
+		r0 = returnFunc(state)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -157,31 +175,23 @@ func (_e *Persistence_Expecter) SaveVM(state interface{}) *Persistence_SaveVM_Ca
 
 func (_c *Persistence_SaveVM_Call) Run(run func(state qemu.VMState)) *Persistence_SaveVM_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(qemu.VMState))
+		var arg0 qemu.VMState
+		if args[0] != nil {
+			arg0 = args[0].(qemu.VMState)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Persistence_SaveVM_Call) Return(_a0 error) *Persistence_SaveVM_Call {
-	_c.Call.Return(_a0)
+func (_c *Persistence_SaveVM_Call) Return(err error) *Persistence_SaveVM_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Persistence_SaveVM_Call) RunAndReturn(run func(qemu.VMState) error) *Persistence_SaveVM_Call {
+func (_c *Persistence_SaveVM_Call) RunAndReturn(run func(state qemu.VMState) error) *Persistence_SaveVM_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewPersistence creates a new instance of Persistence. 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 NewPersistence(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *Persistence {
-	mock := &Persistence{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

manager/qemu/vm.go

@@ -39,7 +39,7 @@ type qemuVM struct {
 	vm.StateMachine
 }
 
-func NewVM(config interface{}, cvmId string, logger *slog.Logger) vm.VM {
+func NewVM(config any, cvmId string, logger *slog.Logger) vm.VM {
 	return &qemuVM{
 		vmi:          config.(VMInfo),
 		cvmId:        cvmId,
@@ -192,11 +192,7 @@ func processExists(pid int) bool {
 	return false
 }
 
-func (v *qemuVM) GetCID() int {
-	return v.vmi.Config.GuestCID
-}
-
-func (v *qemuVM) GetConfig() interface{} {
+func (v *qemuVM) GetConfig() any {
 	return v.vmi
 }
 

manager/qemu/vm_test.go

@@ -132,22 +132,6 @@ func TestGetProcess(t *testing.T) {
 	assert.Equal(t, expectedPid, pid)
 }
 
-func TestGetCID(t *testing.T) {
-	expectedCID := 42
-	vm := &qemuVM{
-		vmi: VMInfo{
-			Config: Config{
-				VSockConfig: VSockConfig{
-					GuestCID: expectedCID,
-				},
-			},
-		},
-	}
-
-	cid := vm.GetCID()
-	assert.Equal(t, expectedCID, cid)
-}
-
 func TestGetConfig(t *testing.T) {
 	expectedConfig := VMInfo{
 		Config: Config{

manager/service.go

@@ -15,7 +15,7 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/google/go-sev-guest/proto/check"
 	"github.com/google/uuid"
 	"github.com/ultravioletrs/cocos/manager/qemu"
@@ -27,7 +27,6 @@ import (
 )
 
 const (
-	hashLength              = 32
 	persistenceDir          = "/tmp/cocos"
 	agentLogLevelKey        = "AGENT_LOG_LEVEL"
 	agentCvmGrpcUrlKey      = "AGENT_CVM_GRPC_URL"
@@ -35,6 +34,7 @@ const (
 	agentCvmClientKey       = "AGENT_CVM_GRPC_CLIENT_KEY"
 	agentCvmServerCaCertKey = "AGENT_CVM_GRPC_SERVER_CA_CERTS"
 	agentCvmId              = "AGENT_CVM_ID"
+	agentCaToken            = "AGENT_CERTS_TOKEN"
 	agentCvmCaUrl           = "AGENT_CVM_CA_URL"
 	defClientCertPath       = "/etc/certs/cert.pem"
 	defClientKeyPath        = "/etc/certs/key.pem"
@@ -68,6 +68,9 @@ var (
 
 	// ErrUnmarshalFailed indicates that the file for the attestation policy could not be unmarshaled.
 	ErrUnmarshalFailed = errors.New("error while unmarshaling the attestation policy")
+
+	// ErrMaxVMsExceeded indicates that the maximum number of VMs has been reached.
+	ErrMaxVMsExceeded = errors.New("maximum number of VMs exceeded")
 )
 
 // Service specifies an API that must be fulfilled by the domain service
@@ -100,12 +103,13 @@ type managerService struct {
 	persistence                 qemu.Persistence
 	eosVersion                  string
 	ttlManager                  *TTLManager
+	maxVMs                      int
 }
 
 var _ Service = (*managerService)(nil)
 
 // New instantiates the manager service implementation.
-func New(cfg qemu.Config, attestationPolicyBinPath string, igvmMeasurementBinaryPath string, pcrValuesFilePath string, logger *slog.Logger, vmFactory vm.Provider, eosVersion string) (Service, error) {
+func New(cfg qemu.Config, attestationPolicyBinPath string, igvmMeasurementBinaryPath string, pcrValuesFilePath string, logger *slog.Logger, vmFactory vm.Provider, eosVersion string, maxVMs int) (Service, error) {
 	start, end, err := decodeRange(cfg.HostFwdRange)
 	if err != nil {
 		return nil, err
@@ -129,6 +133,7 @@ func New(cfg qemu.Config, attestationPolicyBinPath string, igvmMeasurementBinary
 		persistence:                 persistence,
 		eosVersion:                  eosVersion,
 		ttlManager:                  NewTTLManager(),
+		maxVMs:                      maxVMs,
 	}
 
 	if err := ms.restoreVMs(); err != nil {
@@ -140,7 +145,13 @@ func New(cfg qemu.Config, attestationPolicyBinPath string, igvmMeasurementBinary
 
 func (ms *managerService) CreateVM(ctx context.Context, req *CreateReq) (string, string, error) {
 	id := uuid.New().String()
+
 	ms.mu.Lock()
+	if ms.maxVMs > 0 && len(ms.vms) >= ms.maxVMs {
+		ms.mu.Unlock()
+		return "", id, ErrMaxVMsExceeded
+	}
+
 	cfg := qemu.VMInfo{
 		Config:    ms.qemuCfg,
 		LaunchTCB: 0,
@@ -190,22 +201,6 @@ func (ms *managerService) CreateVM(ctx context.Context, req *CreateReq) (string,
 	}
 	cfg.Config.HostFwdAgent = agentPort
 
-	var cid int = qemu.BaseGuestCID
-	for {
-		available := true
-		for _, vm := range ms.vms {
-			if vm.GetCID() == cid {
-				available = false
-				break
-			}
-		}
-		if available {
-			break
-		}
-		cid++
-	}
-	cfg.Config.VSockConfig.GuestCID = cid
-
 	if cfg.Config.EnableSEVSNP {
 		todo := sha3.Sum256([]byte("TODO"))
 		// Define host-data value of QEMU for SEV-SNP, with a base64 encoding of the computation hash.
@@ -216,7 +211,15 @@ func (ms *managerService) CreateVM(ctx context.Context, req *CreateReq) (string,
 	if err = cvm.Start(); err != nil {
 		return "", id, err
 	}
+
 	ms.mu.Lock()
+	if ms.maxVMs > 0 && len(ms.vms) >= ms.maxVMs {
+		ms.mu.Unlock()
+		if stopErr := cvm.Stop(); stopErr != nil {
+			ms.logger.Error("Failed to stop VM after exceeding max limit", "vmID", id, "error", stopErr)
+		}
+		return "", id, ErrMaxVMsExceeded
+	}
 	ms.vms[id] = cvm
 	ms.mu.Unlock()
 
@@ -443,6 +446,7 @@ func tmpEnvironment(id string, req *CreateReq) (string, error) {
 		agentLogLevelKey:   req.AgentLogLevel,
 		agentCvmGrpcUrlKey: req.AgentCvmServerUrl,
 		agentCvmId:         id,
+		agentCaToken:       req.AgentCertsToken,
 		agentCvmCaUrl:      req.AgentCvmCaUrl,
 	}
 

manager/service_test.go

@@ -12,8 +12,8 @@ import (
 	"path"
 	"testing"
 
-	mglog "github.com/absmach/magistrala/logger"
-	"github.com/absmach/magistrala/pkg/errors"
+	mglog "github.com/absmach/supermq/logger"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -30,7 +30,7 @@ func TestNew(t *testing.T) {
 	logger := slog.Default()
 	vmf := new(mocks.Provider)
 
-	service, err := New(cfg, "", "", "", logger, vmf.Execute, "")
+	service, err := New(cfg, "", "", "", logger, vmf.Execute, "", 10)
 	require.NoError(t, err)
 
 	assert.NotNil(t, service)
@@ -77,6 +77,13 @@ func TestRun(t *testing.T) {
 			expectedError:  nil,
 			ttl:            "10s",
 		},
+		{
+			name:           "with exceeded max vms",
+			binaryBehavior: "success",
+			vmStartError:   nil,
+			expectedError:  errors.New("maximum number of VMs exceeded"),
+			ttl:            "",
+		},
 	}
 
 	for _, tt := range tests {
@@ -95,9 +102,6 @@ func TestRun(t *testing.T) {
 
 			qemuCfg := qemu.Config{
 				EnableSEVSNP: true,
-				VSockConfig: qemu.VSockConfig{
-					GuestCID: 3,
-				},
 			}
 			logger := slog.Default()
 
@@ -115,6 +119,11 @@ func TestRun(t *testing.T) {
 				ttlManager:                  NewTTLManager(),
 			}
 
+			if tt.name == "with exceeded max vms" {
+				ms.maxVMs = 1
+				ms.vms["existing-vm"] = vmMock // Simulate an existing VM
+			}
+
 			ctx := context.Background()
 
 			port, _, err := ms.CreateVM(ctx, &CreateReq{Ttl: tt.ttl})

manager/vm/mocks/provider.go

@@ -27,7 +27,7 @@ func (_m *Provider) EXPECT() *Provider_Expecter {
 }
 
 // Execute provides a mock function with given fields: config, computationId, logger
-func (_m *Provider) Execute(config interface{}, computationId string, logger *slog.Logger) vm.VM {
+func (_m *Provider) Execute(config any, computationId string, logger *slog.Logger) vm.VM {
 	ret := _m.Called(config, computationId, logger)
 
 	if len(ret) == 0 {
@@ -35,7 +35,7 @@ func (_m *Provider) Execute(config interface{}, computationId string, logger *sl
 	}
 
 	var r0 vm.VM
-	if rf, ok := ret.Get(0).(func(interface{}, string, *slog.Logger) vm.VM); ok {
+	if rf, ok := ret.Get(0).(func(any, string, *slog.Logger) vm.VM); ok {
 		r0 = rf(config, computationId, logger)
 	} else {
 		if ret.Get(0) != nil {
@@ -52,16 +52,16 @@ type Provider_Execute_Call struct {
 }
 
 // Execute is a helper method to define mock.On call
-//   - config interface{}
+//   - config any
 //   - computationId string
 //   - logger *slog.Logger
-func (_e *Provider_Expecter) Execute(config interface{}, computationId interface{}, logger interface{}) *Provider_Execute_Call {
+func (_e *Provider_Expecter) Execute(config any, computationId any, logger any) *Provider_Execute_Call {
 	return &Provider_Execute_Call{Call: _e.mock.On("Execute", config, computationId, logger)}
 }
 
-func (_c *Provider_Execute_Call) Run(run func(config interface{}, computationId string, logger *slog.Logger)) *Provider_Execute_Call {
+func (_c *Provider_Execute_Call) Run(run func(config any, computationId string, logger *slog.Logger)) *Provider_Execute_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(interface{}), args[1].(string), args[2].(*slog.Logger))
+		run(args[0].(any), args[1].(string), args[2].(*slog.Logger))
 	})
 	return _c
 }
@@ -71,7 +71,7 @@ func (_c *Provider_Execute_Call) Return(_a0 vm.VM) *Provider_Execute_Call {
 	return _c
 }
 
-func (_c *Provider_Execute_Call) RunAndReturn(run func(interface{}, string, *slog.Logger) vm.VM) *Provider_Execute_Call {
+func (_c *Provider_Execute_Call) RunAndReturn(run func(any, string, *slog.Logger) vm.VM) *Provider_Execute_Call {
 	_c.Call.Return(run)
 	return _c
 }

manager/vm/mocks/statemachine.go

@@ -1,15 +1,31 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
 	mock "github.com/stretchr/testify/mock"
-	manager "github.com/ultravioletrs/cocos/pkg/manager"
+	"github.com/ultravioletrs/cocos/pkg/manager"
 )
 
+// NewStateMachine creates a new instance of StateMachine. 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 NewStateMachine(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *StateMachine {
+	mock := &StateMachine{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // StateMachine is an autogenerated mock type for the StateMachine type
 type StateMachine struct {
 	mock.Mock
@@ -23,21 +39,20 @@ func (_m *StateMachine) EXPECT() *StateMachine_Expecter {
 	return &StateMachine_Expecter{mock: &_m.Mock}
 }
 
-// State provides a mock function with no fields
-func (_m *StateMachine) State() string {
-	ret := _m.Called()
+// State provides a mock function for the type StateMachine
+func (_mock *StateMachine) State() string {
+	ret := _mock.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()
+	if returnFunc, ok := ret.Get(0).(func() string); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Get(0).(string)
 	}
-
 	return r0
 }
 
@@ -58,8 +73,8 @@ func (_c *StateMachine_State_Call) Run(run func()) *StateMachine_State_Call {
 	return _c
 }
 
-func (_c *StateMachine_State_Call) Return(_a0 string) *StateMachine_State_Call {
-	_c.Call.Return(_a0)
+func (_c *StateMachine_State_Call) Return(s string) *StateMachine_State_Call {
+	_c.Call.Return(s)
 	return _c
 }
 
@@ -68,21 +83,20 @@ func (_c *StateMachine_State_Call) RunAndReturn(run func() string) *StateMachine
 	return _c
 }
 
-// Transition provides a mock function with given fields: newState
-func (_m *StateMachine) Transition(newState manager.ManagerState) error {
-	ret := _m.Called(newState)
+// Transition provides a mock function for the type StateMachine
+func (_mock *StateMachine) Transition(newState manager.ManagerState) error {
+	ret := _mock.Called(newState)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Transition")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(manager.ManagerState) error); ok {
-		r0 = rf(newState)
+	if returnFunc, ok := ret.Get(0).(func(manager.ManagerState) error); ok {
+		r0 = returnFunc(newState)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -99,31 +113,23 @@ func (_e *StateMachine_Expecter) Transition(newState interface{}) *StateMachine_
 
 func (_c *StateMachine_Transition_Call) Run(run func(newState manager.ManagerState)) *StateMachine_Transition_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(manager.ManagerState))
+		var arg0 manager.ManagerState
+		if args[0] != nil {
+			arg0 = args[0].(manager.ManagerState)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *StateMachine_Transition_Call) Return(_a0 error) *StateMachine_Transition_Call {
-	_c.Call.Return(_a0)
+func (_c *StateMachine_Transition_Call) Return(err error) *StateMachine_Transition_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *StateMachine_Transition_Call) RunAndReturn(run func(manager.ManagerState) error) *StateMachine_Transition_Call {
+func (_c *StateMachine_Transition_Call) RunAndReturn(run func(newState manager.ManagerState) error) *StateMachine_Transition_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewStateMachine creates a new instance of StateMachine. 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 NewStateMachine(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *StateMachine {
-	mock := &StateMachine{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

manager/vm/mocks/vm.go

@@ -1,15 +1,31 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Code generated by mockery v2.53.3. DO NOT EDIT.
+// Code generated by mockery; DO NOT EDIT.
+// github.com/vektra/mockery
+// template: testify
 
 package mocks
 
 import (
 	mock "github.com/stretchr/testify/mock"
-	manager "github.com/ultravioletrs/cocos/pkg/manager"
+	"github.com/ultravioletrs/cocos/pkg/manager"
 )
 
+// NewVM creates a new instance of VM. 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 NewVM(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *VM {
+	mock := &VM{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
 // VM is an autogenerated mock type for the VM type
 type VM struct {
 	mock.Mock
@@ -23,68 +39,22 @@ func (_m *VM) EXPECT() *VM_Expecter {
 	return &VM_Expecter{mock: &_m.Mock}
 }
 
-// GetCID provides a mock function with no fields
-func (_m *VM) GetCID() int {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for GetCID")
-	}
-
-	var r0 int
-	if rf, ok := ret.Get(0).(func() int); ok {
-		r0 = rf()
-	} else {
-		r0 = ret.Get(0).(int)
-	}
-
-	return r0
-}
-
-// VM_GetCID_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetCID'
-type VM_GetCID_Call struct {
-	*mock.Call
-}
-
-// GetCID is a helper method to define mock.On call
-func (_e *VM_Expecter) GetCID() *VM_GetCID_Call {
-	return &VM_GetCID_Call{Call: _e.mock.On("GetCID")}
-}
-
-func (_c *VM_GetCID_Call) Run(run func()) *VM_GetCID_Call {
-	_c.Call.Run(func(args mock.Arguments) {
-		run()
-	})
-	return _c
-}
-
-func (_c *VM_GetCID_Call) Return(_a0 int) *VM_GetCID_Call {
-	_c.Call.Return(_a0)
-	return _c
-}
-
-func (_c *VM_GetCID_Call) RunAndReturn(run func() int) *VM_GetCID_Call {
-	_c.Call.Return(run)
-	return _c
-}
-
-// GetConfig provides a mock function with no fields
-func (_m *VM) GetConfig() interface{} {
-	ret := _m.Called()
+// GetConfig provides a mock function for the type VM
+func (_mock *VM) GetConfig() any {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for GetConfig")
 	}
 
-	var r0 interface{}
-	if rf, ok := ret.Get(0).(func() interface{}); ok {
-		r0 = rf()
+	var r0 any
+	if returnFunc, ok := ret.Get(0).(func() any); ok {
+		r0 = returnFunc()
 	} else {
 		if ret.Get(0) != nil {
-			r0 = ret.Get(0).(interface{})
+			r0 = ret.Get(0).(any)
 		}
 	}
-
 	return r0
 }
 
@@ -105,31 +75,30 @@ func (_c *VM_GetConfig_Call) Run(run func()) *VM_GetConfig_Call {
 	return _c
 }
 
-func (_c *VM_GetConfig_Call) Return(_a0 interface{}) *VM_GetConfig_Call {
-	_c.Call.Return(_a0)
+func (_c *VM_GetConfig_Call) Return(v any) *VM_GetConfig_Call {
+	_c.Call.Return(v)
 	return _c
 }
 
-func (_c *VM_GetConfig_Call) RunAndReturn(run func() interface{}) *VM_GetConfig_Call {
+func (_c *VM_GetConfig_Call) RunAndReturn(run func() any) *VM_GetConfig_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// GetProcess provides a mock function with no fields
-func (_m *VM) GetProcess() int {
-	ret := _m.Called()
+// GetProcess provides a mock function for the type VM
+func (_mock *VM) GetProcess() int {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for GetProcess")
 	}
 
 	var r0 int
-	if rf, ok := ret.Get(0).(func() int); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() int); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Get(0).(int)
 	}
-
 	return r0
 }
 
@@ -150,8 +119,8 @@ func (_c *VM_GetProcess_Call) Run(run func()) *VM_GetProcess_Call {
 	return _c
 }
 
-func (_c *VM_GetProcess_Call) Return(_a0 int) *VM_GetProcess_Call {
-	_c.Call.Return(_a0)
+func (_c *VM_GetProcess_Call) Return(n int) *VM_GetProcess_Call {
+	_c.Call.Return(n)
 	return _c
 }
 
@@ -160,21 +129,20 @@ func (_c *VM_GetProcess_Call) RunAndReturn(run func() int) *VM_GetProcess_Call {
 	return _c
 }
 
-// SetProcess provides a mock function with given fields: pid
-func (_m *VM) SetProcess(pid int) error {
-	ret := _m.Called(pid)
+// SetProcess provides a mock function for the type VM
+func (_mock *VM) SetProcess(pid int) error {
+	ret := _mock.Called(pid)
 
 	if len(ret) == 0 {
 		panic("no return value specified for SetProcess")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(int) error); ok {
-		r0 = rf(pid)
+	if returnFunc, ok := ret.Get(0).(func(int) error); ok {
+		r0 = returnFunc(pid)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -191,36 +159,41 @@ func (_e *VM_Expecter) SetProcess(pid interface{}) *VM_SetProcess_Call {
 
 func (_c *VM_SetProcess_Call) Run(run func(pid int)) *VM_SetProcess_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(int))
+		var arg0 int
+		if args[0] != nil {
+			arg0 = args[0].(int)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *VM_SetProcess_Call) Return(_a0 error) *VM_SetProcess_Call {
-	_c.Call.Return(_a0)
+func (_c *VM_SetProcess_Call) Return(err error) *VM_SetProcess_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *VM_SetProcess_Call) RunAndReturn(run func(int) error) *VM_SetProcess_Call {
+func (_c *VM_SetProcess_Call) RunAndReturn(run func(pid int) error) *VM_SetProcess_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Start provides a mock function with no fields
-func (_m *VM) Start() error {
-	ret := _m.Called()
+// Start provides a mock function for the type VM
+func (_mock *VM) Start() error {
+	ret := _mock.Called()
 
 	if len(ret) == 0 {
 		panic("no return value specified for Start")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -241,8 +214,8 @@ func (_c *VM_Start_Call) Run(run func()) *VM_Start_Call {
 	return _c
 }
 
-func (_c *VM_Start_Call) Return(_a0 error) *VM_Start_Call {
-	_c.Call.Return(_a0)
+func (_c *VM_Start_Call) Return(err error) *VM_Start_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -251,21 +224,20 @@ func (_c *VM_Start_Call) RunAndReturn(run func() error) *VM_Start_Call {
 	return _c
 }
 
-// State provides a mock function with no fields
-func (_m *VM) State() string {
-	ret := _m.Called()
+// State provides a mock function for the type VM
+func (_mock *VM) State() string {
+	ret := _mock.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()
+	if returnFunc, ok := ret.Get(0).(func() string); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Get(0).(string)
 	}
-
 	return r0
 }
 
@@ -286,8 +258,8 @@ func (_c *VM_State_Call) Run(run func()) *VM_State_Call {
 	return _c
 }
 
-func (_c *VM_State_Call) Return(_a0 string) *VM_State_Call {
-	_c.Call.Return(_a0)
+func (_c *VM_State_Call) Return(s string) *VM_State_Call {
+	_c.Call.Return(s)
 	return _c
 }
 
@@ -296,21 +268,20 @@ func (_c *VM_State_Call) RunAndReturn(run func() string) *VM_State_Call {
 	return _c
 }
 
-// Stop provides a mock function with no fields
-func (_m *VM) Stop() error {
-	ret := _m.Called()
+// Stop provides a mock function for the type VM
+func (_mock *VM) Stop() error {
+	ret := _mock.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()
+	if returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -331,8 +302,8 @@ func (_c *VM_Stop_Call) Run(run func()) *VM_Stop_Call {
 	return _c
 }
 
-func (_c *VM_Stop_Call) Return(_a0 error) *VM_Stop_Call {
-	_c.Call.Return(_a0)
+func (_c *VM_Stop_Call) Return(err error) *VM_Stop_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -341,21 +312,20 @@ func (_c *VM_Stop_Call) RunAndReturn(run func() error) *VM_Stop_Call {
 	return _c
 }
 
-// Transition provides a mock function with given fields: newState
-func (_m *VM) Transition(newState manager.ManagerState) error {
-	ret := _m.Called(newState)
+// Transition provides a mock function for the type VM
+func (_mock *VM) Transition(newState manager.ManagerState) error {
+	ret := _mock.Called(newState)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Transition")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(manager.ManagerState) error); ok {
-		r0 = rf(newState)
+	if returnFunc, ok := ret.Get(0).(func(manager.ManagerState) error); ok {
+		r0 = returnFunc(newState)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -372,31 +342,23 @@ func (_e *VM_Expecter) Transition(newState interface{}) *VM_Transition_Call {
 
 func (_c *VM_Transition_Call) Run(run func(newState manager.ManagerState)) *VM_Transition_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].(manager.ManagerState))
+		var arg0 manager.ManagerState
+		if args[0] != nil {
+			arg0 = args[0].(manager.ManagerState)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *VM_Transition_Call) Return(_a0 error) *VM_Transition_Call {
-	_c.Call.Return(_a0)
+func (_c *VM_Transition_Call) Return(err error) *VM_Transition_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *VM_Transition_Call) RunAndReturn(run func(manager.ManagerState) error) *VM_Transition_Call {
+func (_c *VM_Transition_Call) RunAndReturn(run func(newState manager.ManagerState) error) *VM_Transition_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewVM creates a new instance of VM. 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 NewVM(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *VM {
-	mock := &VM{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

manager/vm/vm.go

@@ -15,13 +15,12 @@ type VM interface {
 	Stop() error
 	SetProcess(pid int) error
 	GetProcess() int
-	GetCID() int
 	Transition(newState pkgmanager.ManagerState) error
 	State() string
-	GetConfig() interface{}
+	GetConfig() any
 }
 
-type Provider func(config interface{}, computationId string, logger *slog.Logger) VM
+type Provider func(config any, computationId string, logger *slog.Logger) VM
 
 type Event struct {
 	EventType     string

mockery.yml

@@ -1,145 +0,0 @@
-with-expecter: true
-filename: "{{.InterfaceName}}.go"
-outpkg: "mocks"
-boilerplate-file: "./boilerplate.txt"
-packages:
-  github.com/ultravioletrs/cocos/agent:
-    interfaces:
-      Service:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "agent.go"
-          mockname: "{{.InterfaceName}}"
-      AgentService_DataClient:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "agent_grpc_data.go"
-          mockname: "{{.InterfaceName}}"
-      AgentService_AlgoClient:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "agent_grpc_algo.go"
-          mockname: "{{.InterfaceName}}"
-      AgentService_IMAMeasurementsClient:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "agent_grpc_ima.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/agent/auth:
-    interfaces:
-      Authenticator:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "auth.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/agent/events:
-    interfaces:
-      Service:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "events.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/agent/statemachine:
-    interfaces:
-      StateMachine:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "state.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/internal/server:
-    interfaces:
-      Server:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "server.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/manager:
-    interfaces:
-      Service:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "service.go"
-          mockname: "{{.InterfaceName}}"
-      ManagerServiceClient:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "manager_service_client.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/manager/qemu:
-    interfaces:
-      Persistence:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "persistence.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/manager/vm:
-    interfaces:
-      StateMachine:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "state_machine.go"
-          mockname: "{{.InterfaceName}}"
-      VM:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "vm.go"
-          mockname: "{{.InterfaceName}}"
-      Provider:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "provider.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/pkg/sdk:
-    interfaces:
-      SDK:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "sdk.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/agent/cvms/server:
-    interfaces:
-      AgentServer:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "server.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage:
-    interfaces:
-      Storage:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "storage.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/pkg/clients/grpc:
-    interfaces:
-      Client:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "client.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/pkg/attestation:
-    interfaces:
-      Provider:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "attestation.go"
-          mockname: "{{.InterfaceName}}"
-      Verifier:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "verifier.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/agent/algorithm:
-    interfaces:
-      Algorithm:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "algorithm.go"
-          mockname: "{{.InterfaceName}}"
-  github.com/ultravioletrs/cocos/pkg/attestation/cmdconfig:
-    interfaces:
-      MeasurementProvider:
-        config:
-          dir: "{{.InterfaceDir}}/mocks"
-          filename: "measurement_provider.go"
-          mockname: "{{.InterfaceName}}"

pkg/atls/atls.go

@@ -3,331 +3,67 @@
 package atls
 
 import (
-	"bytes"
-	"crypto/ecdsa"
-	"crypto/elliptic"
-	"crypto/rand"
-	"crypto/tls"
-	"crypto/x509"
-	"crypto/x509/pkix"
 	"encoding/asn1"
 	"encoding/hex"
-	"encoding/json"
-	"encoding/pem"
 	"fmt"
-	"io"
-	"math/big"
-	"net/http"
-	"net/url"
-	"strings"
-	"time"
-
-	"github.com/absmach/certs"
-	certscli "github.com/absmach/certs/cli"
-	"github.com/absmach/certs/errors"
-	certssdk "github.com/absmach/certs/sdk"
-	"github.com/ultravioletrs/cocos/pkg/attestation"
-	"github.com/ultravioletrs/cocos/pkg/attestation/azure"
-	"github.com/ultravioletrs/cocos/pkg/attestation/tdx"
-	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
-	"golang.org/x/crypto/sha3"
 )
 
 const (
-	vmpl2         = 2
-	organization  = "Ultraviolet"
-	country       = "Serbia"
-	province      = ""
-	locality      = "Belgrade"
-	streetAddress = "Bulevar Arsenija Carnojevica 103"
-	postalCode    = "11000"
-	notAfterYear  = 1
-	notAfterMonth = 0
-	notAfterDay   = 0
+	defaultNotAfterYears = 1
+	nonceLength          = 64
+	nonceSuffix          = ".nonce"
 )
 
+// Platform-specific OIDs for certificate extensions.
 var (
 	SNPvTPMOID = asn1.ObjectIdentifier{2, 99999, 1, 0}
 	AzureOID   = asn1.ObjectIdentifier{2, 99999, 1, 1}
 	TDXOID     = asn1.ObjectIdentifier{2, 99999, 1, 2}
 )
 
-type csrReq struct {
-	CSR string `json:"csr,omitempty"`
+// CertificateSubject contains certificate subject information.
+type CertificateSubject struct {
+	Organization  string
+	CommonName    string
+	Country       string
+	Province      string
+	Locality      string
+	StreetAddress string
+	PostalCode    string
 }
 
-func getPlatformProvider(platformType attestation.PlatformType) (attestation.Provider, error) {
-	switch platformType {
-	case attestation.SNPvTPM:
-		return vtpm.NewProvider(true, vmpl2), nil
-	case attestation.Azure:
-		return azure.NewProvider(), nil
-	case attestation.TDX:
-		return tdx.NewProvider(), nil
-	default:
-		return nil, fmt.Errorf("unsupported platform type: %d", platformType)
+// DefaultCertificateSubject returns the default certificate subject for Ultraviolet.
+func DefaultCertificateSubject() CertificateSubject {
+	return CertificateSubject{
+		Organization:  "Ultraviolet",
+		CommonName:    "Ultraviolet",
+		Country:       "Serbia",
+		Province:      "",
+		Locality:      "Belgrade",
+		StreetAddress: "Bulevar Arsenija Carnojevica 103",
+		PostalCode:    "11000",
 	}
 }
 
-func getPlatformVerifier(platformType attestation.PlatformType) (attestation.Verifier, error) {
-	var verifier attestation.Verifier
-
-	switch platformType {
-	case attestation.SNPvTPM:
-		verifier = vtpm.NewVerifier(nil)
-	case attestation.Azure:
-		verifier = azure.NewVerifier(nil)
-	case attestation.TDX:
-		verifier = tdx.NewVerifier()
-	default:
-		return nil, fmt.Errorf("unsupported platform type: %d", platformType)
+func extractNonceFromSNI(serverName string) ([]byte, error) {
+	if len(serverName) < len(nonceSuffix) || !hasNonceSuffix(serverName) {
+		return nil, fmt.Errorf("invalid server name: %s", serverName)
 	}
 
-	err := verifier.JSONToPolicy(attestation.AttestationPolicyPath)
+	nonceStr := serverName[:len(serverName)-len(nonceSuffix)]
+	nonce, err := hex.DecodeString(nonceStr)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to decode nonce: %w", err)
 	}
-	return verifier, nil
+
+	if len(nonce) != nonceLength {
+		return nil, fmt.Errorf("invalid nonce length: expected %d bytes, got %d bytes", nonceLength, len(nonce))
+	}
+
+	return nonce, nil
 }
 
-func getOID(platformType attestation.PlatformType) (asn1.ObjectIdentifier, error) {
-	switch platformType {
-	case attestation.SNPvTPM:
-		return SNPvTPMOID, nil
-	case attestation.Azure:
-		return AzureOID, nil
-	case attestation.TDX:
-		return TDXOID, nil
-	default:
-		return nil, fmt.Errorf("unsupported platform type: %d", platformType)
-	}
-}
-
-func GetPlatformTypeFromOID(oid asn1.ObjectIdentifier) (attestation.PlatformType, error) {
-	switch {
-	case oid.Equal(SNPvTPMOID):
-		return attestation.SNPvTPM, nil
-	case oid.Equal(AzureOID):
-		return attestation.Azure, nil
-	case oid.Equal(TDXOID):
-		return attestation.TDX, nil
-	default:
-		return 0, fmt.Errorf("unsupported OID: %v", oid)
-	}
-}
-
-func VerifyCertificateExtension(extension []byte, pubKey []byte, nonce []byte, pType attestation.PlatformType) error {
-	teeNonce := append(pubKey, nonce...)
-	hashNonce := sha3.Sum512(teeNonce)
-
-	verifier, err := getPlatformVerifier(pType)
-	if err != nil {
-		return fmt.Errorf("failed to get platform verifier: %w", err)
-	}
-
-	if err = verifier.VerifyAttestation(extension, hashNonce[:], hashNonce[:vtpm.Nonce]); err != nil {
-		fmt.Printf("failed to verify attestation: %v\n", err)
-		return err
-	}
-
-	return nil
-}
-
-func GetCertificate(caUrl string, cvmId string) func(*tls.ClientHelloInfo) (*tls.Certificate, error) {
-	pType := attestation.CCPlatform()
-
-	provider, err := getPlatformProvider(pType)
-	if err != nil {
-		return func(*tls.ClientHelloInfo) (*tls.Certificate, error) {
-			return nil, fmt.Errorf("failed to get platform provider: %w", err)
-		}
-	}
-
-	teeOid, err := getOID(pType)
-	if err != nil {
-		return func(*tls.ClientHelloInfo) (*tls.Certificate, error) {
-			return nil, fmt.Errorf("failed to get OID for platform type: %w", err)
-		}
-	}
-
-	return func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {
-		curve := elliptic.P256()
-
-		privateKey, err := ecdsa.GenerateKey(curve, rand.Reader)
-		if err != nil {
-			return nil, fmt.Errorf("failed to generate private/public key: %w", err)
-		}
-
-		pubKeyDER, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
-		if err != nil {
-			return nil, fmt.Errorf("failed to marshal public key to DER format: %w", err)
-		}
-
-		sniLength := len(clientHello.ServerName)
-		if sniLength < 7 || clientHello.ServerName[sniLength-6:] != ".nonce" {
-			return nil, fmt.Errorf("invalid server name: %s", clientHello.ServerName)
-		}
-
-		nonceStr := clientHello.ServerName[:sniLength-6]
-		nonce, err := hex.DecodeString(nonceStr)
-		if err != nil {
-			return nil, fmt.Errorf("failed to decode nonce from server name: %w", err)
-		}
-
-		if len(nonce) != 64 {
-			return nil, fmt.Errorf("invalid nonce length: expected 64 bytes, got %d bytes", len(nonce))
-		}
-
-		attestExtension, err := getCertificateExtension(provider, pubKeyDER, nonce, teeOid)
-		if err != nil {
-			return nil, fmt.Errorf("failed to get certificate extension: %w", err)
-		}
-
-		var certDERBytes []byte
-
-		if caUrl == "" && cvmId == "" {
-			certTemplate := &x509.Certificate{
-				SerialNumber: big.NewInt(202403311),
-				Subject: pkix.Name{
-					Organization:  []string{organization},
-					Country:       []string{country},
-					Province:      []string{province},
-					Locality:      []string{locality},
-					StreetAddress: []string{streetAddress},
-					PostalCode:    []string{postalCode},
-				},
-				NotBefore:             time.Now(),
-				NotAfter:              time.Now().AddDate(notAfterYear, notAfterMonth, notAfterDay),
-				KeyUsage:              x509.KeyUsageDigitalSignature,
-				ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
-				BasicConstraintsValid: true,
-				ExtraExtensions:       []pkix.Extension{attestExtension},
-			}
-
-			DERBytes, err := x509.CreateCertificate(rand.Reader, certTemplate, certTemplate, &privateKey.PublicKey, privateKey)
-			if err != nil {
-				return nil, fmt.Errorf("failed to create certificate: %w", err)
-			}
-
-			certDERBytes = DERBytes
-		} else {
-			csrmd := certs.CSRMetadata{
-				Organization:    []string{organization},
-				Country:         []string{country},
-				Province:        []string{province},
-				Locality:        []string{locality},
-				StreetAddress:   []string{streetAddress},
-				PostalCode:      []string{postalCode},
-				ExtraExtensions: []pkix.Extension{attestExtension},
-			}
-
-			csr, err := certscli.CreateCSR(csrmd, privateKey)
-			if err != nil {
-				return nil, fmt.Errorf("failed to create CSR: %w", err)
-			}
-
-			csrData := string(csr.CSR)
-
-			r := csrReq{
-				CSR: csrData,
-			}
-
-			data, sdkErr := json.Marshal(r)
-			if sdkErr != nil {
-				return nil, fmt.Errorf("failed to marshal CSR request: %w", sdkErr)
-			}
-
-			notBefore := time.Now()
-			notAfter := time.Now().AddDate(notAfterYear, notAfterMonth, notAfterDay)
-			ttlString := notAfter.Sub(notBefore).String()
-
-			query := url.Values{}
-			query.Add("ttl", ttlString)
-			query_string := query.Encode()
-
-			certsEndpoint := "certs"
-			csrEndpoint := "csrs"
-			endpoint := fmt.Sprintf("%s/%s/%s", certsEndpoint, csrEndpoint, cvmId)
-
-			url := fmt.Sprintf("%s/%s?%s", caUrl, endpoint, query_string)
-
-			_, body, err := processRequest(http.MethodPost, url, data, nil, http.StatusOK)
-			if err != nil {
-				return nil, fmt.Errorf("failed to process request: %w", err)
-			}
-
-			var cert certssdk.Certificate
-			if err := json.Unmarshal(body, &cert); err != nil {
-				return nil, fmt.Errorf("failed to unmarshal certificate response: %w", err)
-			}
-
-			cleanCertificateString := strings.ReplaceAll(cert.Certificate, "\\n", "\n")
-
-			block, rest := pem.Decode([]byte(cleanCertificateString))
-
-			if len(rest) != 0 {
-				return nil, fmt.Errorf("failed to convert generated certificate to DER format: %s", cleanCertificateString)
-			}
-
-			certDERBytes = block.Bytes
-		}
-
-		return &tls.Certificate{
-			Certificate: [][]byte{certDERBytes},
-			PrivateKey:  privateKey,
-		}, nil
-	}
-}
-
-func getCertificateExtension(provider attestation.Provider, pubKey []byte, nonce []byte, teeOid asn1.ObjectIdentifier) (pkix.Extension, error) {
-	teeNonce := append(pubKey, nonce...)
-	hashNonce := sha3.Sum512(teeNonce)
-
-	rawAttestation, err := provider.Attestation(hashNonce[:], hashNonce[:vtpm.Nonce])
-	if err != nil {
-		return pkix.Extension{}, fmt.Errorf("failed to get attestation: %w", err)
-	}
-
-	return pkix.Extension{
-		Id:    teeOid,
-		Value: rawAttestation,
-	}, nil
-}
-
-func processRequest(method, reqUrl string, data []byte, headers map[string]string, expectedRespCodes ...int) (http.Header, []byte, errors.SDKError) {
-	req, err := http.NewRequest(method, reqUrl, bytes.NewReader(data))
-	if err != nil {
-		return make(http.Header), []byte{}, errors.NewSDKError(err)
-	}
-
-	// Sets a default value for the Content-Type.
-	// Overridden if Content-Type is passed in the headers arguments.
-	req.Header.Add("Content-Type", "application/json")
-
-	for key, value := range headers {
-		req.Header.Add(key, value)
-	}
-
-	client := &http.Client{
-		Transport: &http.Transport{
-			TLSClientConfig: &tls.Config{
-				InsecureSkipVerify: true,
-			},
-		},
-	}
-
-	resp, err := client.Do(req)
-	if err != nil {
-		return make(http.Header), []byte{}, errors.NewSDKError(err)
-	}
-	defer resp.Body.Close()
-	sdkerr := errors.CheckError(resp, expectedRespCodes...)
-	if sdkerr != nil {
-		return make(http.Header), []byte{}, sdkerr
-	}
-	body, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return make(http.Header), []byte{}, errors.NewSDKError(err)
-	}
-	return resp.Header, body, nil
+func hasNonceSuffix(serverName string) bool {
+	return len(serverName) >= len(nonceSuffix) &&
+		serverName[len(serverName)-len(nonceSuffix):] == nonceSuffix
 }