NOISSUE - Update certs and smq versions (#544)

* Update dependencies and refactor certificate generation to include context

- Updated `cloud.google.com/go/compute/metadata` from v0.8.0 to v0.9.0.
- Updated `github.com/absmach/certs` from v0.18.0 to v0.18.2.
- Updated `github.com/absmach/supermq` from v0.18.1 to v0.18.2.
- Updated `github.com/go-logfmt/logfmt` from v0.6.0 to v0.6.1.
- Updated `github.com/grpc-ecosystem/grpc-gateway/v2` from v2.27.2 to v2.27.3.
- Updated `github.com/prometheus/common` from v0.66.1 to v0.67.1.
- Updated `github.com/rogpeppe/go-internal` from v1.13.1 to v1.14.1.
- Updated `github.com/segmentio/asm` from v1.2.0 to v1.2.1.
- Updated `go.opentelemetry.io/auto/sdk` from v1.1.0 to v1.2.1.
- Updated `go.opentelemetry.io/proto/otlp` from v1.7.1 to v1.8.0.
- Updated `golang.org/x/net` from v0.45.0 to v0.46.0.
- Updated `golang.org/x/oauth2` from v0.30.0 to v0.32.0.
- Updated `google.golang.org/genproto/googleapis/api` and `google.golang.org/genproto/googleapis/rpc` to the latest versions.

- Refactored `generateCASignedCertificate` method in `certificate_provider.go` to accept a context parameter.
- Updated calls to `generateCASignedCertificate` in `GetCertificate` and `TestCASignedCertificateErrors` to pass the context.

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

* Update mockSDK method signatures in certificate error tests to include additional parameters

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

---------

Signed-off-by: Sammy Oina <sammyoina@gmail.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

.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

@@ -26,14 +26,14 @@ 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/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

@@ -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/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

@@ -16,7 +16,7 @@ import (
 	"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/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

@@ -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

@@ -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

@@ -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

@@ -22,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"
 )
 
@@ -51,6 +52,7 @@ var (
 	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 {
@@ -114,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>`,
@@ -134,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()
@@ -173,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>`,
@@ -190,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()
@@ -231,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 {

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)

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_test.go

@@ -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

@@ -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/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_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,8 +14,8 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
-	"time"
 
+	"github.com/absmach/certs/sdk"
 	mglog "github.com/absmach/supermq/logger"
 	"github.com/absmach/supermq/pkg/prometheus"
 	"github.com/caarlos0/env/v11"
@@ -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,15 @@ func main() {
 		return
 	}
 
+	if ccPlatform == attestation.SNP || ccPlatform == attestation.SNPvTPM {
+		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 +173,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 +224,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 +307,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)

cmd/manager/main.go

@@ -20,14 +20,14 @@ import (
 	"github.com/absmach/supermq/pkg/uuid"
 	"github.com/caarlos0/env/v11"
 	"github.com/go-chi/chi/v5"
-	"github.com/ultravioletrs/cocos/internal/server"
-	grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
 	"github.com/ultravioletrs/cocos/manager"
 	"github.com/ultravioletrs/cocos/manager/api"
 	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"
@@ -35,12 +35,11 @@ import (
 )
 
 const (
-	svcName          = "manager"
-	envPrefixGRPC    = "MANAGER_GRPC_"
-	envPrefixHTTP    = "MANAGER_HTTP_"
-	envPrefixQemu    = "MANAGER_QEMU_"
-	clientBufferSize = 100
-	defSvcHTTPPort   = "7003"
+	svcName        = "manager"
+	envPrefixGRPC  = "MANAGER_GRPC_"
+	envPrefixHTTP  = "MANAGER_HTTP_"
+	envPrefixQemu  = "MANAGER_QEMU_"
+	defSvcHTTPPort = "7003"
 )
 
 type config struct {
@@ -146,7 +145,7 @@ 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)
 

go.mod

@@ -1,6 +1,6 @@
 module github.com/ultravioletrs/cocos
 
-go 1.24.5
+go 1.25.3
 
 require (
 	github.com/caarlos0/env/v11 v11.3.1
@@ -9,40 +9,40 @@ require (
 	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.7
-	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.40.0
-	golang.org/x/sync v0.16.0
-	google.golang.org/grpc v1.74.2
-	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.43.0
+	golang.org/x/sync v0.17.0
+	google.golang.org/grpc v1.76.0
+	google.golang.org/protobuf v1.36.10
 )
 
 require (
-	cloud.google.com/go/storage v1.55.0
-	github.com/absmach/supermq v0.17.0
+	cloud.google.com/go/storage v1.57.0
+	github.com/absmach/supermq v0.18.2
 	github.com/caarlos0/env/v10 v10.0.0
-	github.com/go-chi/chi/v5 v5.2.2
-	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.24.0 // indirect
-	cloud.google.com/go v0.121.1 // indirect
-	cloud.google.com/go/auth v0.16.1 // 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.9.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/v5 v5.0.2 // 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
@@ -54,17 +54,15 @@ require (
 	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.4 // indirect
-	github.com/go-jose/go-jose/v4 v4.0.5 // indirect
+	github.com/go-jose/go-jose/v4 v4.1.2 // indirect
 	github.com/gofrs/uuid/v5 v5.3.2 // indirect
-	github.com/gogo/protobuf v1.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/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
@@ -73,65 +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.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/auto/sdk v1.2.1 // indirect
 	go.opentelemetry.io/contrib/detectors/gcp v1.36.0 // indirect
-	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.62.0 // indirect
-	go.opentelemetry.io/otel v1.37.0 // indirect
-	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.37.0 // indirect
-	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.37.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-20250606033433-dcc06ee1d476 // indirect
-	golang.org/x/oauth2 v0.30.0 // 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.3 // indirect
+	golang.org/x/oauth2 v0.32.0 // indirect
 	golang.org/x/time v0.12.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
+	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.2
 	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.3.2+incompatible
+	github.com/docker/docker v28.5.1+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-logfmt/logfmt v0.6.1 // 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-tools v0.4.4
+	github.com/google/go-tpm v0.9.6
+	github.com/google/go-tpm-tools v0.4.6
 	github.com/google/logger v1.1.1
 	github.com/google/uuid v1.6.0
-	github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.1 // indirect
+	github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3 // 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.23.0
+	github.com/prometheus/client_golang v1.23.2
 	github.com/prometheus/client_model v0.6.2 // indirect
-	github.com/prometheus/common v0.65.0 // indirect
-	github.com/prometheus/procfs v0.16.1 // indirect
-	github.com/stretchr/objx v0.5.2 // indirect
-	go.opentelemetry.io/otel/metric v1.37.0 // indirect
-	go.opentelemetry.io/proto/otlp v1.7.0 // indirect
+	github.com/prometheus/common v0.67.1 // indirect
+	github.com/prometheus/procfs v0.17.0 // indirect
+	github.com/stretchr/objx v0.5.3 // indirect
+	go.opentelemetry.io/otel/metric v1.38.0 // indirect
+	go.opentelemetry.io/proto/otlp v1.8.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	golang.org/x/net v0.41.0 // indirect
-	golang.org/x/sys v0.34.0 // indirect
-	golang.org/x/term v0.33.0
-	golang.org/x/text v0.27.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20250603155806-513f23925822 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20250603155806-513f23925822 // indirect
+	golang.org/x/net v0.46.0 // indirect
+	golang.org/x/sys v0.37.0 // indirect
+	golang.org/x/term v0.36.0
+	golang.org/x/text v0.30.0 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20251014184007-4626949a642f // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20251014184007-4626949a642f // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 
 replace github.com/virtee/sev-snp-measure-go => github.com/sammyoina/sev-snp-measure-go v0.0.0-20241202151803-ef189f0ff825
-
-replace github.com/google/go-tpm-tools => github.com/danko-miladinovic/go-tpm-tools v0.0.0-20250228160324-1ebcfd79567c

go.sum

@@ -1,13 +1,13 @@
 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.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=
+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.9.0 h1:pDUj4QMoPejqq20dK0Pg2N4yG9zIkYGdBtwLoEkH9Zs=
+cloud.google.com/go/compute/metadata v0.9.0/go.mod h1:E0bWwX5wTnLPedCKqk3pJmVgCBSM6qQI1yTBdEb3C10=
 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,36 +16,36 @@ 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/supermq v0.17.0 h1:l8XkoPJ2GIs3od4FBUu52DJRuzUIcVfScj7UkXr4/3w=
-github.com/absmach/supermq v0.17.0/go.mod h1:BTpQiHHevhP5ahA34vhyeLK8YIJsx+7gYUFpqgMGSAI=
+github.com/absmach/certs v0.18.2 h1:s6KKL3/KfDZ6z0IxvNCksIOUwRnEgQyCpeAonuR15No=
+github.com/absmach/certs v0.18.2/go.mod h1:scqVZsmW2xPScnpMTtE70oN6cn0LLjFcJVPi4JKZ4+E=
+github.com/absmach/supermq v0.18.2 h1:oGiuMyjiFJhAmysmApOsYKtN4Kf/WXpLWElEUqmcZec=
+github.com/absmach/supermq v0.18.2/go.mod h1:ehhfnpFhHZLF9EAdxqDSUCWLJxYhu0mqFd7a4CKMY0g=
 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/v5 v5.0.2 h1:rIfFVxEf1QsI7E1ZHfp/B4DF/6QBAUhmgkxc0H7Zss8=
-github.com/cenkalti/backoff/v5 v5.0.2/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw=
+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-20250501225837-2ac532fd4443 h1:aQ3y1lwWyqYPiWZThqv1aFbZMiM9vblcSArJRf2Irls=
@@ -57,15 +57,15 @@ github.com/containerd/errdefs/pkg v0.3.0/go.mod h1:NJw6s9HwNuRhnjJhM7pylWwMyAkmC
 github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I=
 github.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo=
 github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
-github.com/danko-miladinovic/go-tpm-tools v0.0.0-20250228160324-1ebcfd79567c h1:gFo8kqRXFoM6ttqMrK+M3xffxco+Yj80kUo3NoMe8LU=
-github.com/danko-miladinovic/go-tpm-tools v0.0.0-20250228160324-1ebcfd79567c/go.mod h1:ktjTNq8yZFD6TzdBFefUfen96rF3NpYwpSb2d8bc+Y8=
 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/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.3.2+incompatible h1:wn66NJ6pWB1vBZIilP8G3qQPqHy5XymfYn5vsqeA5oA=
-github.com/docker/docker v28.3.2+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
+github.com/docker/docker v28.5.1+incompatible h1:Bm8DchhSD2J6PsFzxC35TZo4TLGR2PdW/E69rU45NhM=
+github.com/docker/docker v28.5.1+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=
@@ -84,31 +84,33 @@ 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.2 h1:CMwsvRVTbXVytCk1Wd72Zy1LAsAh9GxMmSNWLHCG618=
-github.com/go-chi/chi/v5 v5.2.2/go.mod h1:L2yAIGWB3H+phAw1NxKwWM+7eUH/lU8pOMm5hHcoops=
+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.0.5 h1:M6T8+mKZl/+fNNuFHvGIzDz7BTLQPIounk/b9dw3AaE=
-github.com/go-jose/go-jose/v4 v4.0.5/go.mod h1:s3P1lRrkT8igV8D9OjyL4WRyHvjB6a4JSllnOrmmBOA=
+github.com/go-jose/go-jose/v4 v4.1.2 h1:TK/7NqRQZfgAh+Td8AlsrvtPoUyiHh0LqVvokh+1vHI=
+github.com/go-jose/go-jose/v4 v4.1.2/go.mod h1:22cg9HWM1pOlnRiY+9cQYJ9XHmya1bYW8OeDM6Ku6Oo=
 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=
 github.com/go-kit/log v0.2.1/go.mod h1:NwTd00d/i8cPZ3xOwwiv2PO5MOcx78fFErGNcVmBjv0=
-github.com/go-logfmt/logfmt v0.6.0 h1:wGYYu3uicYdqXVgoYbvnkrPVXkuLM1p1ifugDMEdRi4=
-github.com/go-logfmt/logfmt v0.6.0/go.mod h1:WYhtIu8zTZfxdn5+rREduYbwxfcBr/Vr6KEVveWlfTs=
+github.com/go-logfmt/logfmt v0.6.1 h1:4hvbpePJKnIzH1B+8OR/JPbTx37NktoI9LE2QZBBkvE=
+github.com/go-logfmt/logfmt v0.6.1/go.mod h1:EV2pOAQoZaT1ZXZbqDl5hrymndi4SY9ED9/z6CO0XAk=
 github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
 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.2 h1:2jfO8j3XgSwlz/wHqemAEugfnTlikAYHhnqQ8Xh4fE0=
 github.com/gofrs/uuid/v5 v5.3.2/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/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 +131,10 @@ 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-tpm-tools v0.4.6 h1:hwIwPG7w4z5eQEBq11gYw8YYr9xXLfBQ/0JsKyq5AJM=
+github.com/google/go-tpm-tools v0.4.6/go.mod h1:MsVQbJnRhKDfWwf5zgr3cDGpj13P1uLAFF0wMEP/n5w=
 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,28 +147,46 @@ 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.27.1 h1:X5VWvz21y3gzm9Nw/kaUeku/1+uBhcekkmy4IkffJww=
-github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.1/go.mod h1:Zanoh4+gvIgluNqcfMVTJueD4wSS5hT7zTt4Mrutd90=
-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.3 h1:NmZ1PKzSTQbuGHw9DGPFomqkkLWMC+vZCkfs+FHv1Vg=
+github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3/go.mod h1:zQrxl1YP88HQlA6i9c63DSVPFklWpGX4OWAc9bFuaH4=
 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/klauspost/compress v1.18.1 h1:bcSGx7UbpBqMChDtsF28Lw6v/G94LPrrbMbdC3JH2co=
+github.com/klauspost/compress v1.18.1/go.mod h1:ZQFFVG+MdnR0P+l6wpXgIL4NTtwiKIdBnrBd8Nrxr+0=
 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/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.4 h1:IwQibdnf8l2KoO+qC3uT4OaTWsW7tuRQXy9TRN9QanA=
+github.com/lestrrat-go/blackmagic v1.0.4/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=
@@ -185,8 +207,6 @@ 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.20240917153116-6f2963f01587 h1:xzZOeCMQLA/W198ZkdVdt4EKFKJtS26B773zNU377ZY=
@@ -194,108 +214,99 @@ github.com/planetscale/vtprotobuf v0.6.1-0.20240917153116-6f2963f01587/go.mod h1
 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.23.0 h1:ust4zpdl9r4trLY/gSjlm07PuiBq2ynaXXlptpfy8Uc=
-github.com/prometheus/client_golang v1.23.0/go.mod h1:i/o0R9ByOnHX0McrTMTyhYvKE4haaf2mW08I+jGAjEE=
+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.65.0 h1:QDwzd+G1twt//Kwj/Ww6E9FQq1iVMmODnILtW1t2VzE=
-github.com/prometheus/common v0.65.0/go.mod h1:0gZns+BLRQ3V6NdaerOhMbwwRbNh9hkGINtQAsP5GS8=
-github.com/prometheus/procfs v0.16.1 h1:hZ15bTNuirocR6u0JZ6BAHHmwS1p8B4P6MRqxtzMyRg=
-github.com/prometheus/procfs v0.16.1/go.mod h1:teAbpZRB1iIAJYREa1LsoWUXykVXA1KlTmWl8x/U+Is=
-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/prometheus/common v0.67.1 h1:OTSON1P4DNxzTg4hmKCc37o4ZAZDv0cfXLkOt0oEowI=
+github.com/prometheus/common v0.67.1/go.mod h1:RpmT9v35q2Y+lsieQsdOh5sXZ6ajUGC8NjZAmr8vb0Q=
+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.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
+github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
+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.1 h1:DTNbBqs57ioxAD4PrArqftgypG4/qNpXoJx8TVXxPR0=
+github.com/segmentio/asm v1.2.1/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.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/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
-github.com/spf13/pflag v1.0.7 h1:vN6T9TfwStFPFM5XzjsvmzZkLuaLX+HS+0SeFLRgU6M=
-github.com/spf13/pflag v1.0.7/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/objx v0.5.3 h1:jmXUvGomnU1o3W/V5h2VEradbpJDwGrzugQQvL0POH4=
+github.com/stretchr/objx v0.5.3/go.mod h1:rDQraq+vQZU7Fde9LOZLr8Tax6zZvy4kuNKF+QYS+U0=
 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/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=
-go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA=
-go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A=
+go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64=
+go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y=
 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.62.0 h1:Hf9xI/XLML9ElpiHVDNwvqI0hIFlzV8dgIr35kV1kRU=
-go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.62.0/go.mod h1:NfchwuyNoMcZ5MLHwPrODwUF1HWCXWrL31s8gSAdIKY=
-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.37.0 h1:Ahq7pZmv87yiyn3jeFz/LekZmPLLdKejuO3NcK9MssM=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.37.0/go.mod h1:MJTqhM0im3mRLw1i8uGHnCvUEeS7VwRyxlLC78PA18M=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.37.0 h1:bDMKF3RUSxshZ5OjOTi8rsHGaPKsAt76FaqgvIUySLc=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.37.0/go.mod h1:dDT67G/IkA46Mr2l9Uj7HsQVwsjASyV9SjGofsiUZDA=
+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.7.0 h1:jX1VolD6nHuFzOYso2E73H85i92Mv8JQYk0K9vz09os=
-go.opentelemetry.io/proto/otlp v1.7.0/go.mod h1:fSKjH6YJ7HDlwzltzyMj036AJ3ejJLCgCSHGj4efDDo=
+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.8.0 h1:fRAZQDcAFHySxpJ1TwlA1cJ4tvcrw7nXl9xWWC8N5CE=
+go.opentelemetry.io/proto/otlp v1.8.0/go.mod h1:tIeYOeNBU4cvmPqpaji1P+KbB4Oloai8wN4rWzRrFF0=
 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.3 h1:6gvOSjQoTB3vt1l+CU+tSyi/HOjfOjRLJ4YwYZGwRO0=
+go.yaml.in/yaml/v2 v2.4.3/go.mod h1:zSxWcmIDjOzPXpjlTTbAsKokqkDNAVtZO0WOMiT90s8=
 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.40.0 h1:r4x+VvoG5Fm+eJcxMaY8CQM7Lb0l1lsmjGBQ6s8BfKM=
-golang.org/x/crypto v0.40.0/go.mod h1:Qr1vMER5WyS2dfPHAlsOj01wgLbsyWtFn/aY+5+ZdxY=
-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.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
-golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
+golang.org/x/crypto v0.43.0 h1:dduJYIi3A3KOfdGOHX8AVZ/jGiyPa3IbBozJ5kNuE04=
+golang.org/x/crypto v0.43.0/go.mod h1:BFbav4mRNlXJL4wNeejLpWxB7wMbc79PdRGhWKncxR0=
 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/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI=
-golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU=
+golang.org/x/net v0.46.0 h1:giFlY12I07fugqwPuWJi68oOnpfqFnJIJzaIIm2JVV4=
+golang.org/x/net v0.46.0/go.mod h1:Q9BGdFy1y4nkUwiLvT5qtyhAnEHgnQ/zd8PfU6nc210=
+golang.org/x/oauth2 v0.32.0 h1:jsCblLleRMDrxMN29H3z/k1KliIvpLgCkE6R8FXXNgY=
+golang.org/x/oauth2 v0.32.0/go.mod h1:lzm5WQJQwKZ3nwavOZ3IS5Aulzxi68dUSgRHujetwEA=
 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.16.0 h1:ycBJEhp9p4vXvUZNszeOq0kGTPghopOL8q0fq3vstxw=
-golang.org/x/sync v0.16.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=
@@ -305,52 +316,47 @@ 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.34.0 h1:H5Y5sJ2L2JRdyv7ROF1he/lPdvFsd0mJHFw2ThKHxLA=
-golang.org/x/sys v0.34.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ=
+golang.org/x/sys v0.37.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.33.0 h1:NuFncQrRcaRvVmgRkvM3j/F00gWIAlcmlB8ACEKmGIg=
-golang.org/x/term v0.33.0/go.mod h1:s18+ql9tYWp1IfpV9DmCtQDDSRBUjKaw9M1eAv5UeF0=
+golang.org/x/term v0.36.0 h1:zMPR+aF8gfksFprF/Nc/rd1wRS1EI6nDBGyWAvDzx2Q=
+golang.org/x/term v0.36.0/go.mod h1:Qu394IJq6V6dCBRgwqshf3mPF85AqzYEzofzRdZkWss=
 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.27.0 h1:4fGWRpyh641NLlecmyl4LOe6yDdfaYNrGb2zdfo4JV4=
-golang.org/x/text v0.27.0/go.mod h1:1D28KMCvyooCX9hBiosv5Tz/+YLxj0j7XhWjpSUF7CU=
+golang.org/x/text v0.30.0 h1:yznKA/E9zq54KzlzBEAWn1NXSQ8DIp/NYMy88xJjl4k=
+golang.org/x/text v0.30.0/go.mod h1:yDdHFIX9t+tORqspjENWgzaCVXgk0yYnYuSZ8UzzBVM=
 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-20250603155806-513f23925822 h1:oWVWY3NzT7KJppx2UKhKmzPq4SRe0LdCijVRwvGeikY=
-google.golang.org/genproto/googleapis/api v0.0.0-20250603155806-513f23925822/go.mod h1:h3c4v36UTKzUiuaOKQ6gr3S+0hovBtUrXzTG/i3+XEc=
-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.74.2 h1:WoosgB65DlWVC9FqI82dGsZhWFNBSLjQ84bjROOpMu4=
-google.golang.org/grpc v1.74.2/go.mod h1:CtQ+BGjaAIXHs/5YS3i473GqwBBa1zGQNevxdeBEXrM=
-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-20251014184007-4626949a642f h1:OiFuztEyBivVKDvguQJYWq1yDcfAHIID/FVrPR4oiI0=
+google.golang.org/genproto/googleapis/api v0.0.0-20251014184007-4626949a642f/go.mod h1:kprOiu9Tr0JYyD6DORrc4Hfyk3RFXqkQ3ctHEum3ZbM=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251014184007-4626949a642f h1:1FTH6cpXFsENbPR5Bu8NQddPSaUUE6NA2XdZdDSAJK4=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20251014184007-4626949a642f/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk=
+google.golang.org/grpc v1.76.0 h1:UnVkv1+uMLYXoIz6o7chp59WfQUYA2ex/BXQ9rHZu7A=
+google.golang.org/grpc v1.76.0/go.mod h1:Ju12QI8M6iQJtbcsV+awF5a4hfJMLi4X0JLo94ULZ6c=
+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/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/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/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/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,7 +192,7 @@ func processExists(pid int) bool {
 	return false
 }
 
-func (v *qemuVM) GetConfig() interface{} {
+func (v *qemuVM) GetConfig() any {
 	return v.vmi
 }
 

manager/service.go

@@ -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"
@@ -446,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/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,23 +39,22 @@ func (_m *VM) EXPECT() *VM_Expecter {
 	return &VM_Expecter{mock: &_m.Mock}
 }
 
-// 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
 }
 
@@ -60,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
 }
 
@@ -105,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
 }
 
@@ -115,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
 }
 
@@ -146,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
 }
 
@@ -196,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
 }
 
@@ -206,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
 }
 
@@ -241,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
 }
 
@@ -251,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
 }
 
@@ -286,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
 }
 
@@ -296,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
 }
 
@@ -327,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

@@ -17,10 +17,10 @@ type VM interface {
 	GetProcess() 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
 }

pkg/atls/attestation_provider.go

@@ -0,0 +1,66 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package atls
+
+import (
+	"encoding/asn1"
+	"fmt"
+
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"golang.org/x/crypto/sha3"
+)
+
+// AttestationProvider defines the interface for platform attestation operations.
+type AttestationProvider interface {
+	Attest(pubKey []byte, nonce []byte) ([]byte, error)
+	OID() asn1.ObjectIdentifier
+	PlatformType() attestation.PlatformType
+}
+
+// PlatformAttestationProvider handles platform attestation operations.
+type platformAttestationProvider struct {
+	provider     attestation.Provider
+	oid          asn1.ObjectIdentifier
+	platformType attestation.PlatformType
+}
+
+// NewAttestationProvider creates a new attestation provider for the given platform type.
+func NewAttestationProvider(provider attestation.Provider, platformType attestation.PlatformType) (AttestationProvider, error) {
+	oid, err := OID(platformType)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get OID: %w", err)
+	}
+
+	return &platformAttestationProvider{
+		provider:     provider,
+		oid:          oid,
+		platformType: platformType,
+	}, nil
+}
+
+func (p *platformAttestationProvider) Attest(pubKey []byte, nonce []byte) ([]byte, error) {
+	teeNonce := append(pubKey, nonce...)
+	hashNonce := sha3.Sum512(teeNonce)
+	return p.provider.Attestation(hashNonce[:], hashNonce[:32])
+}
+
+func (p *platformAttestationProvider) OID() asn1.ObjectIdentifier {
+	return p.oid
+}
+
+func (p *platformAttestationProvider) PlatformType() attestation.PlatformType {
+	return p.platformType
+}
+
+func OID(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)
+	}
+}

pkg/atls/certificate_provider.go

@@ -0,0 +1,189 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package atls
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/pem"
+	"fmt"
+	"math/big"
+	"strings"
+	"time"
+
+	"github.com/absmach/certs"
+	sdk "github.com/absmach/certs/sdk"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+)
+
+// CertificateProvider defines the interface for providing TLS certificates.
+type CertificateProvider interface {
+	GetCertificate(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error)
+}
+
+// AttestedCertificateProvider provides attested TLS certificates.
+type attestedCertificateProvider struct {
+	attestationProvider AttestationProvider
+	certsSDK            sdk.SDK
+	agentToken          string
+	subject             CertificateSubject
+	useCA               bool
+	cvmID               string
+	ttl                 time.Duration
+	notAfterYears       int
+}
+
+// NewAttestedProvider creates a new attested certificate provider for self-signed certificates.
+func NewAttestedProvider(
+	attestationProvider AttestationProvider,
+	subject CertificateSubject,
+) CertificateProvider {
+	return &attestedCertificateProvider{
+		attestationProvider: attestationProvider,
+		subject:             subject,
+		useCA:               false,
+		notAfterYears:       defaultNotAfterYears,
+	}
+}
+
+// NewAttestedCAProvider creates a new attested certificate provider for CA-signed certificates.
+func NewAttestedCAProvider(
+	attestationProvider AttestationProvider,
+	subject CertificateSubject,
+	certsSDK sdk.SDK, cvmID, agentToken string,
+) CertificateProvider {
+	return &attestedCertificateProvider{
+		attestationProvider: attestationProvider,
+		subject:             subject,
+		certsSDK:            certsSDK,
+		agentToken:          agentToken,
+		useCA:               true,
+		cvmID:               cvmID,
+		ttl:                 time.Hour * 24 * 365, // Default 1 year
+	}
+}
+
+// SetTTL sets the certificate TTL for CA-signed certificates.
+func (p *attestedCertificateProvider) SetTTL(ttl time.Duration) {
+	p.ttl = ttl
+}
+
+func (p *attestedCertificateProvider) GetCertificate(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {
+	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate private key: %w", err)
+	}
+
+	pubKeyDER, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal public key: %w", err)
+	}
+
+	nonce, err := extractNonceFromSNI(clientHello.ServerName)
+	if err != nil {
+		return nil, fmt.Errorf("failed to extract nonce: %w", err)
+	}
+
+	attestationData, err := p.attestationProvider.Attest(pubKeyDER, nonce)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get attestation: %w", err)
+	}
+
+	extension := pkix.Extension{
+		Id:    p.attestationProvider.OID(),
+		Value: attestationData,
+	}
+
+	var certDERBytes []byte
+	if p.useCA {
+		certDERBytes, err = p.generateCASignedCertificate(clientHello.Context(), privateKey, extension)
+	} else {
+		certDERBytes, err = p.generateSelfSignedCertificate(privateKey, extension)
+	}
+
+	if err != nil {
+		return nil, fmt.Errorf("failed to generate certificate: %w", err)
+	}
+
+	return &tls.Certificate{
+		Certificate: [][]byte{certDERBytes},
+		PrivateKey:  privateKey,
+	}, nil
+}
+
+func (p *attestedCertificateProvider) generateSelfSignedCertificate(privateKey *ecdsa.PrivateKey, extension pkix.Extension) ([]byte, error) {
+	certTemplate := &x509.Certificate{
+		SerialNumber: big.NewInt(time.Now().Unix()),
+		Subject: pkix.Name{
+			Organization:  []string{p.subject.Organization},
+			Country:       []string{p.subject.Country},
+			Province:      []string{p.subject.Province},
+			Locality:      []string{p.subject.Locality},
+			StreetAddress: []string{p.subject.StreetAddress},
+			PostalCode:    []string{p.subject.PostalCode},
+		},
+		NotBefore:             time.Now(),
+		NotAfter:              time.Now().AddDate(p.notAfterYears, 0, 0),
+		KeyUsage:              x509.KeyUsageDigitalSignature,
+		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
+		BasicConstraintsValid: true,
+		ExtraExtensions:       []pkix.Extension{extension},
+	}
+
+	return x509.CreateCertificate(rand.Reader, certTemplate, certTemplate, &privateKey.PublicKey, privateKey)
+}
+
+func (p *attestedCertificateProvider) generateCASignedCertificate(ctx context.Context, privateKey *ecdsa.PrivateKey, extension pkix.Extension) ([]byte, error) {
+	csrMetadata := certs.CSRMetadata{
+		Organization:    []string{p.subject.Organization},
+		Country:         []string{p.subject.Country},
+		CommonName:      p.subject.CommonName,
+		Province:        []string{p.subject.Province},
+		Locality:        []string{p.subject.Locality},
+		StreetAddress:   []string{p.subject.StreetAddress},
+		PostalCode:      []string{p.subject.PostalCode},
+		ExtraExtensions: []pkix.Extension{extension},
+	}
+
+	csr, sdkerr := p.certsSDK.CreateCSR(ctx, csrMetadata, privateKey)
+	if sdkerr != nil {
+		return nil, fmt.Errorf("failed to create CSR: %w", sdkerr)
+	}
+
+	cert, err := p.certsSDK.IssueFromCSRInternal(ctx, p.cvmID, p.ttl.String(), string(csr.CSR), p.agentToken)
+	if err != nil {
+		return nil, err
+	}
+
+	cleanCertificateString := strings.ReplaceAll(cert.Certificate, "\\n", "\n")
+	block, rest := pem.Decode([]byte(cleanCertificateString))
+
+	if len(rest) != 0 {
+		return nil, fmt.Errorf("failed to decode certificate PEM: unexpected remaining data")
+	}
+	if block == nil {
+		return nil, fmt.Errorf("failed to decode certificate PEM: no PEM block found")
+	}
+
+	return block.Bytes, nil
+}
+
+func NewProvider(provider attestation.Provider, platformType attestation.PlatformType, agentToken, cvmID string, certsSDK sdk.SDK) (CertificateProvider, error) {
+	attestationProvider, err := NewAttestationProvider(provider, platformType)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create attestation provider: %w", err)
+	}
+
+	subject := DefaultCertificateSubject()
+
+	if certsSDK != nil {
+		return NewAttestedCAProvider(attestationProvider, subject, certsSDK, cvmID, agentToken), nil
+	}
+
+	return NewAttestedProvider(attestationProvider, subject), nil
+}

pkg/atls/certificate_verifier.go

@@ -0,0 +1,125 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package atls
+
+import (
+	"crypto/x509"
+	"encoding/asn1"
+	"fmt"
+	"time"
+
+	"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"
+)
+
+type CertificateVerifier interface {
+	VerifyPeerCertificate(rawCerts [][]byte, verifiedChains [][]*x509.Certificate, nonce []byte) error
+}
+
+// CertificateVerifier handles certificate verification operations.
+type certificateVerifier struct {
+	rootCAs *x509.CertPool
+}
+
+func NewCertificateVerifier(rootCAs *x509.CertPool) CertificateVerifier {
+	return &certificateVerifier{rootCAs: rootCAs}
+}
+
+func (v *certificateVerifier) VerifyPeerCertificate(rawCerts [][]byte, _ [][]*x509.Certificate, nonce []byte) error {
+	if len(rawCerts) == 0 {
+		return fmt.Errorf("no certificates provided")
+	}
+
+	cert, err := x509.ParseCertificate(rawCerts[0])
+	if err != nil {
+		return fmt.Errorf("failed to parse x509 certificate: %w", err)
+	}
+
+	if err := v.verifyCertificateSignature(cert); err != nil {
+		return fmt.Errorf("certificate signature verification failed: %w", err)
+	}
+
+	return v.verifyAttestationExtension(cert, nonce)
+}
+
+func (v *certificateVerifier) verifyCertificateSignature(cert *x509.Certificate) error {
+	rootCAs := v.rootCAs
+	if rootCAs == nil {
+		rootCAs = x509.NewCertPool()
+		rootCAs.AddCert(cert)
+	}
+
+	opts := x509.VerifyOptions{
+		Roots:       rootCAs,
+		CurrentTime: time.Now(),
+	}
+
+	_, err := cert.Verify(opts)
+	return err
+}
+
+func (v *certificateVerifier) verifyAttestationExtension(cert *x509.Certificate, nonce []byte) error {
+	for _, ext := range cert.Extensions {
+		if platformType, err := platformTypeFromOID(ext.Id); err == nil {
+			pubKeyDER, err := x509.MarshalPKIXPublicKey(cert.PublicKey)
+			if err != nil {
+				return fmt.Errorf("failed to marshal public key: %w", err)
+			}
+			return v.verifyCertificateExtension(ext.Value, pubKeyDER, nonce, platformType)
+		}
+	}
+	return fmt.Errorf("attestation extension not found in certificate")
+}
+
+func (v *certificateVerifier) verifyCertificateExtension(extension []byte, pubKey []byte, nonce []byte, platformType attestation.PlatformType) error {
+	verifier, err := platformVerifier(platformType)
+	if err != nil {
+		return fmt.Errorf("failed to get platform verifier: %w", err)
+	}
+
+	teeNonce := append(pubKey, nonce...)
+	hashNonce := sha3.Sum512(teeNonce)
+
+	if err = verifier.VerifyAttestation(extension, hashNonce[:], hashNonce[:32]); err != nil {
+		return fmt.Errorf("failed to verify attestation: %w", err)
+	}
+
+	return nil
+}
+
+func platformTypeFromOID(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 platformVerifier(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)
+	}
+
+	err := verifier.JSONToPolicy(attestation.AttestationPolicyPath)
+	if err != nil {
+		return nil, err
+	}
+	return verifier, nil
+}

pkg/atls/mocks/certificateprovider.go

@@ -0,0 +1,103 @@
+// 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 (
+	"crypto/tls"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// NewCertificateProvider creates a new instance of CertificateProvider. 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 NewCertificateProvider(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *CertificateProvider {
+	mock := &CertificateProvider{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
+
+// CertificateProvider is an autogenerated mock type for the CertificateProvider type
+type CertificateProvider struct {
+	mock.Mock
+}
+
+type CertificateProvider_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *CertificateProvider) EXPECT() *CertificateProvider_Expecter {
+	return &CertificateProvider_Expecter{mock: &_m.Mock}
+}
+
+// GetCertificate provides a mock function for the type CertificateProvider
+func (_mock *CertificateProvider) GetCertificate(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error) {
+	ret := _mock.Called(clientHello)
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetCertificate")
+	}
+
+	var r0 *tls.Certificate
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(*tls.ClientHelloInfo) (*tls.Certificate, error)); ok {
+		return returnFunc(clientHello)
+	}
+	if returnFunc, ok := ret.Get(0).(func(*tls.ClientHelloInfo) *tls.Certificate); ok {
+		r0 = returnFunc(clientHello)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*tls.Certificate)
+		}
+	}
+	if returnFunc, ok := ret.Get(1).(func(*tls.ClientHelloInfo) error); ok {
+		r1 = returnFunc(clientHello)
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// CertificateProvider_GetCertificate_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetCertificate'
+type CertificateProvider_GetCertificate_Call struct {
+	*mock.Call
+}
+
+// GetCertificate is a helper method to define mock.On call
+//   - clientHello *tls.ClientHelloInfo
+func (_e *CertificateProvider_Expecter) GetCertificate(clientHello interface{}) *CertificateProvider_GetCertificate_Call {
+	return &CertificateProvider_GetCertificate_Call{Call: _e.mock.On("GetCertificate", clientHello)}
+}
+
+func (_c *CertificateProvider_GetCertificate_Call) Run(run func(clientHello *tls.ClientHelloInfo)) *CertificateProvider_GetCertificate_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 *tls.ClientHelloInfo
+		if args[0] != nil {
+			arg0 = args[0].(*tls.ClientHelloInfo)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *CertificateProvider_GetCertificate_Call) Return(certificate *tls.Certificate, err error) *CertificateProvider_GetCertificate_Call {
+	_c.Call.Return(certificate, err)
+	return _c
+}
+
+func (_c *CertificateProvider_GetCertificate_Call) RunAndReturn(run func(clientHello *tls.ClientHelloInfo) (*tls.Certificate, error)) *CertificateProvider_GetCertificate_Call {
+	_c.Call.Return(run)
+	return _c
+}

pkg/attestation/attestation.go

@@ -20,7 +20,6 @@ const (
 	SNP PlatformType = iota
 	VTPM
 	SNPvTPM
-	AzureToken
 	Azure
 	TDX
 	NoCC

pkg/attestation/azure/snp.go

@@ -63,7 +63,6 @@ func (a provider) Attestation(teeNonce []byte, vTpmNonce []byte) ([]byte, error)
 	quote.TeeAttestation = &attest.Attestation_SevSnpAttestation{
 		SevSnpAttestation: snpReport,
 	}
-
 	return proto.Marshal(quote)
 }
 
@@ -165,7 +164,7 @@ func GenerateAttestationPolicy(token, product string, policy uint64) (*attestati
 		return nil, fmt.Errorf("failed to validate token: %w", err)
 	}
 
-	tee, ok := claims["x-ms-isolation-tee"].(map[string]interface{})
+	tee, ok := claims["x-ms-isolation-tee"].(map[string]any)
 	if !ok {
 		return nil, fmt.Errorf("failed to get tee from claims")
 	}
@@ -284,7 +283,7 @@ func FetchAzureAttestationToken(tokenNonce []byte, maaURL string) ([]byte, error
 	return []byte(token), nil
 }
 
-func validateToken(token string) (map[string]interface{}, error) {
+func validateToken(token string) (map[string]any, error) {
 	unverifiedToken, _, err := new(jwt.Parser).ParseUnverified(token, jwt.MapClaims{})
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse token: %w", err)

pkg/attestation/azure/snp_test.go

@@ -452,7 +452,7 @@ func TestIntegration_FullAttestationFlow(t *testing.T) {
 		maaServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			switch r.URL.Path {
 			case "/attest":
-				response := map[string]interface{}{
+				response := map[string]any{
 					"token": createMockJWT(),
 				}
 				w.Header().Set("Content-Type", "application/json")
@@ -460,7 +460,7 @@ func TestIntegration_FullAttestationFlow(t *testing.T) {
 					t.Fatalf("Failed to encode response: %v", err)
 				}
 			case "/.well-known/openid_configuration":
-				config := map[string]interface{}{
+				config := map[string]any{
 					"jwks_uri": "maaServer.URL" + "/certs",
 				}
 				w.Header().Set("Content-Type", "application/json")
@@ -468,8 +468,8 @@ func TestIntegration_FullAttestationFlow(t *testing.T) {
 					t.Fatalf("Failed to encode OpenID configuration: %v", err)
 				}
 			case "/certs":
-				jwks := map[string]interface{}{
-					"keys": []map[string]interface{}{
+				jwks := map[string]any{
+					"keys": []map[string]any{
 						{
 							"kid": "test-kid",
 							"kty": "RSA",
@@ -555,7 +555,7 @@ func createMockJWT() string {
 		"aud": "test-audience",
 		"exp": time.Now().Add(time.Hour).Unix(),
 		"iat": time.Now().Unix(),
-		"x-ms-isolation-tee": map[string]interface{}{
+		"x-ms-isolation-tee": map[string]any{
 			"x-ms-sevsnpvm-familyId":          "1234567890abcdef",
 			"x-ms-sevsnpvm-imageId":           "fedcba0987654321",
 			"x-ms-sevsnpvm-launchmeasurement": "abcdef1234567890",

pkg/attestation/cmdconfig/mocks/measurementprovider.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"
+)
+
+// NewMeasurementProvider creates a new instance of MeasurementProvider. 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 NewMeasurementProvider(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *MeasurementProvider {
+	mock := &MeasurementProvider{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
 
 // MeasurementProvider is an autogenerated mock type for the MeasurementProvider type
 type MeasurementProvider struct {
@@ -20,9 +38,9 @@ func (_m *MeasurementProvider) EXPECT() *MeasurementProvider_Expecter {
 	return &MeasurementProvider_Expecter{mock: &_m.Mock}
 }
 
-// Run provides a mock function with given fields: binaryPath
-func (_m *MeasurementProvider) Run(binaryPath string) ([]byte, error) {
-	ret := _m.Called(binaryPath)
+// Run provides a mock function for the type MeasurementProvider
+func (_mock *MeasurementProvider) Run(binaryPath string) ([]byte, error) {
+	ret := _mock.Called(binaryPath)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Run")
@@ -30,23 +48,21 @@ func (_m *MeasurementProvider) Run(binaryPath string) ([]byte, error) {
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func(string) ([]byte, error)); ok {
-		return rf(binaryPath)
+	if returnFunc, ok := ret.Get(0).(func(string) ([]byte, error)); ok {
+		return returnFunc(binaryPath)
 	}
-	if rf, ok := ret.Get(0).(func(string) []byte); ok {
-		r0 = rf(binaryPath)
+	if returnFunc, ok := ret.Get(0).(func(string) []byte); ok {
+		r0 = returnFunc(binaryPath)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func(string) error); ok {
-		r1 = rf(binaryPath)
+	if returnFunc, ok := ret.Get(1).(func(string) error); ok {
+		r1 = returnFunc(binaryPath)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -63,36 +79,41 @@ func (_e *MeasurementProvider_Expecter) Run(binaryPath interface{}) *Measurement
 
 func (_c *MeasurementProvider_Run_Call) Run(run func(binaryPath string)) *MeasurementProvider_Run_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 *MeasurementProvider_Run_Call) Return(_a0 []byte, _a1 error) *MeasurementProvider_Run_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *MeasurementProvider_Run_Call) Return(bytes []byte, err error) *MeasurementProvider_Run_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *MeasurementProvider_Run_Call) RunAndReturn(run func(string) ([]byte, error)) *MeasurementProvider_Run_Call {
+func (_c *MeasurementProvider_Run_Call) RunAndReturn(run func(binaryPath string) ([]byte, error)) *MeasurementProvider_Run_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// Stop provides a mock function with no fields
-func (_m *MeasurementProvider) Stop() error {
-	ret := _m.Called()
+// Stop provides a mock function for the type MeasurementProvider
+func (_mock *MeasurementProvider) 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
 }
 
@@ -113,8 +134,8 @@ func (_c *MeasurementProvider_Stop_Call) Run(run func()) *MeasurementProvider_St
 	return _c
 }
 
-func (_c *MeasurementProvider_Stop_Call) Return(_a0 error) *MeasurementProvider_Stop_Call {
-	_c.Call.Return(_a0)
+func (_c *MeasurementProvider_Stop_Call) Return(err error) *MeasurementProvider_Stop_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
@@ -122,17 +143,3 @@ func (_c *MeasurementProvider_Stop_Call) RunAndReturn(run func() error) *Measure
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewMeasurementProvider creates a new instance of MeasurementProvider. 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 NewMeasurementProvider(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *MeasurementProvider {
-	mock := &MeasurementProvider{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

pkg/attestation/mocks/provider.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"
+)
+
+// NewProvider creates a new instance of Provider. 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 NewProvider(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Provider {
+	mock := &Provider{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
 
 // Provider is an autogenerated mock type for the Provider type
 type Provider struct {
@@ -20,9 +38,9 @@ func (_m *Provider) EXPECT() *Provider_Expecter {
 	return &Provider_Expecter{mock: &_m.Mock}
 }
 
-// Attestation provides a mock function with given fields: teeNonce, vTpmNonce
-func (_m *Provider) Attestation(teeNonce []byte, vTpmNonce []byte) ([]byte, error) {
-	ret := _m.Called(teeNonce, vTpmNonce)
+// Attestation provides a mock function for the type Provider
+func (_mock *Provider) Attestation(teeNonce []byte, vTpmNonce []byte) ([]byte, error) {
+	ret := _mock.Called(teeNonce, vTpmNonce)
 
 	if len(ret) == 0 {
 		panic("no return value specified for Attestation")
@@ -30,23 +48,21 @@ func (_m *Provider) Attestation(teeNonce []byte, vTpmNonce []byte) ([]byte, erro
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func([]byte, []byte) ([]byte, error)); ok {
-		return rf(teeNonce, vTpmNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte, []byte) ([]byte, error)); ok {
+		return returnFunc(teeNonce, vTpmNonce)
 	}
-	if rf, ok := ret.Get(0).(func([]byte, []byte) []byte); ok {
-		r0 = rf(teeNonce, vTpmNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte, []byte) []byte); ok {
+		r0 = returnFunc(teeNonce, vTpmNonce)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func([]byte, []byte) error); ok {
-		r1 = rf(teeNonce, vTpmNonce)
+	if returnFunc, ok := ret.Get(1).(func([]byte, []byte) error); ok {
+		r1 = returnFunc(teeNonce, vTpmNonce)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -64,24 +80,35 @@ func (_e *Provider_Expecter) Attestation(teeNonce interface{}, vTpmNonce interfa
 
 func (_c *Provider_Attestation_Call) Run(run func(teeNonce []byte, vTpmNonce []byte)) *Provider_Attestation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]byte), args[1].([]byte))
+		var arg0 []byte
+		if args[0] != nil {
+			arg0 = args[0].([]byte)
+		}
+		var arg1 []byte
+		if args[1] != nil {
+			arg1 = args[1].([]byte)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Provider_Attestation_Call) Return(_a0 []byte, _a1 error) *Provider_Attestation_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Provider_Attestation_Call) Return(bytes []byte, err error) *Provider_Attestation_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Provider_Attestation_Call) RunAndReturn(run func([]byte, []byte) ([]byte, error)) *Provider_Attestation_Call {
+func (_c *Provider_Attestation_Call) RunAndReturn(run func(teeNonce []byte, vTpmNonce []byte) ([]byte, error)) *Provider_Attestation_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// AzureAttestationToken provides a mock function with given fields: tokenNonce
-func (_m *Provider) AzureAttestationToken(tokenNonce []byte) ([]byte, error) {
-	ret := _m.Called(tokenNonce)
+// AzureAttestationToken provides a mock function for the type Provider
+func (_mock *Provider) AzureAttestationToken(tokenNonce []byte) ([]byte, error) {
+	ret := _mock.Called(tokenNonce)
 
 	if len(ret) == 0 {
 		panic("no return value specified for AzureAttestationToken")
@@ -89,23 +116,21 @@ func (_m *Provider) AzureAttestationToken(tokenNonce []byte) ([]byte, error) {
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func([]byte) ([]byte, error)); ok {
-		return rf(tokenNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte) ([]byte, error)); ok {
+		return returnFunc(tokenNonce)
 	}
-	if rf, ok := ret.Get(0).(func([]byte) []byte); ok {
-		r0 = rf(tokenNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte) []byte); ok {
+		r0 = returnFunc(tokenNonce)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func([]byte) error); ok {
-		r1 = rf(tokenNonce)
+	if returnFunc, ok := ret.Get(1).(func([]byte) error); ok {
+		r1 = returnFunc(tokenNonce)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -122,24 +147,30 @@ func (_e *Provider_Expecter) AzureAttestationToken(tokenNonce interface{}) *Prov
 
 func (_c *Provider_AzureAttestationToken_Call) Run(run func(tokenNonce []byte)) *Provider_AzureAttestationToken_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]byte))
+		var arg0 []byte
+		if args[0] != nil {
+			arg0 = args[0].([]byte)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Provider_AzureAttestationToken_Call) Return(_a0 []byte, _a1 error) *Provider_AzureAttestationToken_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Provider_AzureAttestationToken_Call) Return(bytes []byte, err error) *Provider_AzureAttestationToken_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Provider_AzureAttestationToken_Call) RunAndReturn(run func([]byte) ([]byte, error)) *Provider_AzureAttestationToken_Call {
+func (_c *Provider_AzureAttestationToken_Call) RunAndReturn(run func(tokenNonce []byte) ([]byte, error)) *Provider_AzureAttestationToken_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// TeeAttestation provides a mock function with given fields: teeNonce
-func (_m *Provider) TeeAttestation(teeNonce []byte) ([]byte, error) {
-	ret := _m.Called(teeNonce)
+// TeeAttestation provides a mock function for the type Provider
+func (_mock *Provider) TeeAttestation(teeNonce []byte) ([]byte, error) {
+	ret := _mock.Called(teeNonce)
 
 	if len(ret) == 0 {
 		panic("no return value specified for TeeAttestation")
@@ -147,23 +178,21 @@ func (_m *Provider) TeeAttestation(teeNonce []byte) ([]byte, error) {
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func([]byte) ([]byte, error)); ok {
-		return rf(teeNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte) ([]byte, error)); ok {
+		return returnFunc(teeNonce)
 	}
-	if rf, ok := ret.Get(0).(func([]byte) []byte); ok {
-		r0 = rf(teeNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte) []byte); ok {
+		r0 = returnFunc(teeNonce)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func([]byte) error); ok {
-		r1 = rf(teeNonce)
+	if returnFunc, ok := ret.Get(1).(func([]byte) error); ok {
+		r1 = returnFunc(teeNonce)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -180,24 +209,30 @@ func (_e *Provider_Expecter) TeeAttestation(teeNonce interface{}) *Provider_TeeA
 
 func (_c *Provider_TeeAttestation_Call) Run(run func(teeNonce []byte)) *Provider_TeeAttestation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]byte))
+		var arg0 []byte
+		if args[0] != nil {
+			arg0 = args[0].([]byte)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Provider_TeeAttestation_Call) Return(_a0 []byte, _a1 error) *Provider_TeeAttestation_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Provider_TeeAttestation_Call) Return(bytes []byte, err error) *Provider_TeeAttestation_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Provider_TeeAttestation_Call) RunAndReturn(run func([]byte) ([]byte, error)) *Provider_TeeAttestation_Call {
+func (_c *Provider_TeeAttestation_Call) RunAndReturn(run func(teeNonce []byte) ([]byte, error)) *Provider_TeeAttestation_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// VTpmAttestation provides a mock function with given fields: vTpmNonce
-func (_m *Provider) VTpmAttestation(vTpmNonce []byte) ([]byte, error) {
-	ret := _m.Called(vTpmNonce)
+// VTpmAttestation provides a mock function for the type Provider
+func (_mock *Provider) VTpmAttestation(vTpmNonce []byte) ([]byte, error) {
+	ret := _mock.Called(vTpmNonce)
 
 	if len(ret) == 0 {
 		panic("no return value specified for VTpmAttestation")
@@ -205,23 +240,21 @@ func (_m *Provider) VTpmAttestation(vTpmNonce []byte) ([]byte, error) {
 
 	var r0 []byte
 	var r1 error
-	if rf, ok := ret.Get(0).(func([]byte) ([]byte, error)); ok {
-		return rf(vTpmNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte) ([]byte, error)); ok {
+		return returnFunc(vTpmNonce)
 	}
-	if rf, ok := ret.Get(0).(func([]byte) []byte); ok {
-		r0 = rf(vTpmNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte) []byte); ok {
+		r0 = returnFunc(vTpmNonce)
 	} else {
 		if ret.Get(0) != nil {
 			r0 = ret.Get(0).([]byte)
 		}
 	}
-
-	if rf, ok := ret.Get(1).(func([]byte) error); ok {
-		r1 = rf(vTpmNonce)
+	if returnFunc, ok := ret.Get(1).(func([]byte) error); ok {
+		r1 = returnFunc(vTpmNonce)
 	} else {
 		r1 = ret.Error(1)
 	}
-
 	return r0, r1
 }
 
@@ -238,31 +271,23 @@ func (_e *Provider_Expecter) VTpmAttestation(vTpmNonce interface{}) *Provider_VT
 
 func (_c *Provider_VTpmAttestation_Call) Run(run func(vTpmNonce []byte)) *Provider_VTpmAttestation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]byte))
+		var arg0 []byte
+		if args[0] != nil {
+			arg0 = args[0].([]byte)
+		}
+		run(
+			arg0,
+		)
 	})
 	return _c
 }
 
-func (_c *Provider_VTpmAttestation_Call) Return(_a0 []byte, _a1 error) *Provider_VTpmAttestation_Call {
-	_c.Call.Return(_a0, _a1)
+func (_c *Provider_VTpmAttestation_Call) Return(bytes []byte, err error) *Provider_VTpmAttestation_Call {
+	_c.Call.Return(bytes, err)
 	return _c
 }
 
-func (_c *Provider_VTpmAttestation_Call) RunAndReturn(run func([]byte) ([]byte, error)) *Provider_VTpmAttestation_Call {
+func (_c *Provider_VTpmAttestation_Call) RunAndReturn(run func(vTpmNonce []byte) ([]byte, error)) *Provider_VTpmAttestation_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewProvider creates a new instance of Provider. 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 NewProvider(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *Provider {
-	mock := &Provider{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

pkg/attestation/mocks/verifier.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"
+)
+
+// NewVerifier creates a new instance of Verifier. 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 NewVerifier(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Verifier {
+	mock := &Verifier{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}
 
 // Verifier is an autogenerated mock type for the Verifier type
 type Verifier struct {
@@ -20,21 +38,20 @@ func (_m *Verifier) EXPECT() *Verifier_Expecter {
 	return &Verifier_Expecter{mock: &_m.Mock}
 }
 
-// JSONToPolicy provides a mock function with given fields: path
-func (_m *Verifier) JSONToPolicy(path string) error {
-	ret := _m.Called(path)
+// JSONToPolicy provides a mock function for the type Verifier
+func (_mock *Verifier) JSONToPolicy(path string) error {
+	ret := _mock.Called(path)
 
 	if len(ret) == 0 {
 		panic("no return value specified for JSONToPolicy")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func(string) error); ok {
-		r0 = rf(path)
+	if returnFunc, ok := ret.Get(0).(func(string) error); ok {
+		r0 = returnFunc(path)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -51,36 +68,41 @@ func (_e *Verifier_Expecter) JSONToPolicy(path interface{}) *Verifier_JSONToPoli
 
 func (_c *Verifier_JSONToPolicy_Call) Run(run func(path string)) *Verifier_JSONToPolicy_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 *Verifier_JSONToPolicy_Call) Return(_a0 error) *Verifier_JSONToPolicy_Call {
-	_c.Call.Return(_a0)
+func (_c *Verifier_JSONToPolicy_Call) Return(err error) *Verifier_JSONToPolicy_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Verifier_JSONToPolicy_Call) RunAndReturn(run func(string) error) *Verifier_JSONToPolicy_Call {
+func (_c *Verifier_JSONToPolicy_Call) RunAndReturn(run func(path string) error) *Verifier_JSONToPolicy_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// VerifTeeAttestation provides a mock function with given fields: report, teeNonce
-func (_m *Verifier) VerifTeeAttestation(report []byte, teeNonce []byte) error {
-	ret := _m.Called(report, teeNonce)
+// VerifTeeAttestation provides a mock function for the type Verifier
+func (_mock *Verifier) VerifTeeAttestation(report []byte, teeNonce []byte) error {
+	ret := _mock.Called(report, teeNonce)
 
 	if len(ret) == 0 {
 		panic("no return value specified for VerifTeeAttestation")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func([]byte, []byte) error); ok {
-		r0 = rf(report, teeNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte, []byte) error); ok {
+		r0 = returnFunc(report, teeNonce)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -98,36 +120,46 @@ func (_e *Verifier_Expecter) VerifTeeAttestation(report interface{}, teeNonce in
 
 func (_c *Verifier_VerifTeeAttestation_Call) Run(run func(report []byte, teeNonce []byte)) *Verifier_VerifTeeAttestation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]byte), args[1].([]byte))
+		var arg0 []byte
+		if args[0] != nil {
+			arg0 = args[0].([]byte)
+		}
+		var arg1 []byte
+		if args[1] != nil {
+			arg1 = args[1].([]byte)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Verifier_VerifTeeAttestation_Call) Return(_a0 error) *Verifier_VerifTeeAttestation_Call {
-	_c.Call.Return(_a0)
+func (_c *Verifier_VerifTeeAttestation_Call) Return(err error) *Verifier_VerifTeeAttestation_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Verifier_VerifTeeAttestation_Call) RunAndReturn(run func([]byte, []byte) error) *Verifier_VerifTeeAttestation_Call {
+func (_c *Verifier_VerifTeeAttestation_Call) RunAndReturn(run func(report []byte, teeNonce []byte) error) *Verifier_VerifTeeAttestation_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// VerifVTpmAttestation provides a mock function with given fields: report, vTpmNonce
-func (_m *Verifier) VerifVTpmAttestation(report []byte, vTpmNonce []byte) error {
-	ret := _m.Called(report, vTpmNonce)
+// VerifVTpmAttestation provides a mock function for the type Verifier
+func (_mock *Verifier) VerifVTpmAttestation(report []byte, vTpmNonce []byte) error {
+	ret := _mock.Called(report, vTpmNonce)
 
 	if len(ret) == 0 {
 		panic("no return value specified for VerifVTpmAttestation")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func([]byte, []byte) error); ok {
-		r0 = rf(report, vTpmNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte, []byte) error); ok {
+		r0 = returnFunc(report, vTpmNonce)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -145,36 +177,46 @@ func (_e *Verifier_Expecter) VerifVTpmAttestation(report interface{}, vTpmNonce
 
 func (_c *Verifier_VerifVTpmAttestation_Call) Run(run func(report []byte, vTpmNonce []byte)) *Verifier_VerifVTpmAttestation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]byte), args[1].([]byte))
+		var arg0 []byte
+		if args[0] != nil {
+			arg0 = args[0].([]byte)
+		}
+		var arg1 []byte
+		if args[1] != nil {
+			arg1 = args[1].([]byte)
+		}
+		run(
+			arg0,
+			arg1,
+		)
 	})
 	return _c
 }
 
-func (_c *Verifier_VerifVTpmAttestation_Call) Return(_a0 error) *Verifier_VerifVTpmAttestation_Call {
-	_c.Call.Return(_a0)
+func (_c *Verifier_VerifVTpmAttestation_Call) Return(err error) *Verifier_VerifVTpmAttestation_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Verifier_VerifVTpmAttestation_Call) RunAndReturn(run func([]byte, []byte) error) *Verifier_VerifVTpmAttestation_Call {
+func (_c *Verifier_VerifVTpmAttestation_Call) RunAndReturn(run func(report []byte, vTpmNonce []byte) error) *Verifier_VerifVTpmAttestation_Call {
 	_c.Call.Return(run)
 	return _c
 }
 
-// VerifyAttestation provides a mock function with given fields: report, teeNonce, vTpmNonce
-func (_m *Verifier) VerifyAttestation(report []byte, teeNonce []byte, vTpmNonce []byte) error {
-	ret := _m.Called(report, teeNonce, vTpmNonce)
+// VerifyAttestation provides a mock function for the type Verifier
+func (_mock *Verifier) VerifyAttestation(report []byte, teeNonce []byte, vTpmNonce []byte) error {
+	ret := _mock.Called(report, teeNonce, vTpmNonce)
 
 	if len(ret) == 0 {
 		panic("no return value specified for VerifyAttestation")
 	}
 
 	var r0 error
-	if rf, ok := ret.Get(0).(func([]byte, []byte, []byte) error); ok {
-		r0 = rf(report, teeNonce, vTpmNonce)
+	if returnFunc, ok := ret.Get(0).(func([]byte, []byte, []byte) error); ok {
+		r0 = returnFunc(report, teeNonce, vTpmNonce)
 	} else {
 		r0 = ret.Error(0)
 	}
-
 	return r0
 }
 
@@ -193,31 +235,33 @@ func (_e *Verifier_Expecter) VerifyAttestation(report interface{}, teeNonce inte
 
 func (_c *Verifier_VerifyAttestation_Call) Run(run func(report []byte, teeNonce []byte, vTpmNonce []byte)) *Verifier_VerifyAttestation_Call {
 	_c.Call.Run(func(args mock.Arguments) {
-		run(args[0].([]byte), args[1].([]byte), args[2].([]byte))
+		var arg0 []byte
+		if args[0] != nil {
+			arg0 = args[0].([]byte)
+		}
+		var arg1 []byte
+		if args[1] != nil {
+			arg1 = args[1].([]byte)
+		}
+		var arg2 []byte
+		if args[2] != nil {
+			arg2 = args[2].([]byte)
+		}
+		run(
+			arg0,
+			arg1,
+			arg2,
+		)
 	})
 	return _c
 }
 
-func (_c *Verifier_VerifyAttestation_Call) Return(_a0 error) *Verifier_VerifyAttestation_Call {
-	_c.Call.Return(_a0)
+func (_c *Verifier_VerifyAttestation_Call) Return(err error) *Verifier_VerifyAttestation_Call {
+	_c.Call.Return(err)
 	return _c
 }
 
-func (_c *Verifier_VerifyAttestation_Call) RunAndReturn(run func([]byte, []byte, []byte) error) *Verifier_VerifyAttestation_Call {
+func (_c *Verifier_VerifyAttestation_Call) RunAndReturn(run func(report []byte, teeNonce []byte, vTpmNonce []byte) error) *Verifier_VerifyAttestation_Call {
 	_c.Call.Return(run)
 	return _c
 }
-
-// NewVerifier creates a new instance of Verifier. 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 NewVerifier(t interface {
-	mock.TestingT
-	Cleanup(func())
-}) *Verifier {
-	mock := &Verifier{}
-	mock.Mock.Test(t)
-
-	t.Cleanup(func() { mock.AssertExpectations(t) })
-
-	return mock
-}

pkg/attestation/quoteprovider/sev.go

@@ -7,10 +7,14 @@
 package quoteprovider
 
 import (
+	"crypto/rand"
+	"crypto/x509"
+	"encoding/pem"
 	"fmt"
 	"io"
 	"os"
 	"path"
+	"path/filepath"
 	"time"
 
 	"github.com/absmach/supermq/pkg/errors"
@@ -22,11 +26,13 @@ import (
 	"github.com/google/go-sev-guest/verify/trust"
 	"github.com/google/logger"
 	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"google.golang.org/protobuf/proto"
 )
 
 const (
-	cocosDirectory     = ".cocos"
-	caBundleName       = "ask_ark.pem"
+	cocosDirectory     = "/cocos"
+	arkAskBundleName   = "ask_ark.pem"
+	vcekName           = "vcek.pem"
 	Nonce              = 64
 	sevSnpProductMilan = "Milan"
 	sevSnpProductGenoa = "Genoa"
@@ -57,7 +63,7 @@ func fillInAttestationLocal(attestation *sevsnp.Attestation, cfg *check.Config)
 			return err
 		}
 
-		bundlePath := path.Join(homePath, cocosDirectory, product, caBundleName)
+		bundlePath := path.Join(homePath, cocosDirectory, product, arkAskBundleName)
 		if _, err := os.Stat(bundlePath); err == nil {
 			amdRootCerts := trust.AMDRootCerts{}
 			if err := amdRootCerts.FromKDSCert(bundlePath); err != nil {
@@ -166,12 +172,39 @@ func FetchAttestation(reportDataSlice []byte, vmpl uint) ([]byte, error) {
 	}
 	copy(reportData[:], reportDataSlice)
 
-	rawQuote, err := qp.GetRawQuoteAtLevel(reportData, vmpl)
+	quoteProto, err := client.GetQuoteProtoAtLevel(qp, reportData, vmpl)
 	if err != nil {
-		return []byte{}, fmt.Errorf("failed to get raw quote")
+		return []byte{}, fmt.Errorf("failed to get quote proto")
 	}
 
-	return rawQuote, nil
+	homePath, _ := os.UserHomeDir()
+	vcekPath := path.Join(homePath, cocosDirectory, fmt.Sprintf("%d", quoteProto.Product.Name), vcekName)
+	arkAskBundlePath := path.Join(homePath, cocosDirectory, fmt.Sprintf("%d", quoteProto.Product.Name), arkAskBundleName)
+
+	vcekBytes, err := os.ReadFile(vcekPath)
+	if err != nil {
+		return []byte{}, fmt.Errorf("could not read VCEK file: %v", err)
+	}
+
+	arkAskBundleBytes, err := os.ReadFile(arkAskBundlePath)
+	if err != nil {
+		return []byte{}, fmt.Errorf("could not read ask/ark bundle file: %v", err)
+	}
+
+	vcekPem, _ := pem.Decode(vcekBytes)
+	arkPem, rest := pem.Decode(arkAskBundleBytes)
+	askPem, _ := pem.Decode(rest)
+
+	quoteProto.CertificateChain.VcekCert = vcekPem.Bytes
+	quoteProto.CertificateChain.AskCert = askPem.Bytes
+	quoteProto.CertificateChain.ArkCert = arkPem.Bytes
+
+	result, err := proto.Marshal(quoteProto)
+	if err != nil {
+		return []byte{}, fmt.Errorf("failed to marshal quote proto: %v", err)
+	}
+
+	return result, nil
 }
 
 func GetProductName(product string) sevsnp.SevProduct_SevProductName {
@@ -184,3 +217,83 @@ func GetProductName(product string) sevsnp.SevProduct_SevProductName {
 		return sevsnp.SevProduct_SEV_PRODUCT_UNKNOWN
 	}
 }
+
+func derToPem(der []byte) []byte {
+	// Try to parse to make sure it's a certificate
+	if _, err := x509.ParseCertificate(der); err != nil {
+		// cert_chain endpoint already returns PEM; just pass through
+		return der
+	}
+	return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
+}
+
+func FetchCertificates(vmpl uint) error {
+	var reportData [Nonce]byte
+
+	qp, err := GetLeveledQuoteProvider()
+	if err != nil {
+		return fmt.Errorf("could not get quote provider")
+	}
+
+	if len(reportData) > Nonce {
+		return fmt.Errorf("attestation report size mismatch")
+	}
+
+	_, err = rand.Read(reportData[:])
+	if err != nil {
+		return fmt.Errorf("failed to read random data: %v", err)
+	}
+
+	quoteProto, err := client.GetQuoteProtoAtLevel(qp, reportData, vmpl) // for coverage
+	if err != nil {
+		return fmt.Errorf("failed to get quote proto")
+	}
+
+	options := &verify.Options{
+		CheckRevocations:    true,
+		DisableCertFetching: false,
+		Getter:              trust.DefaultHTTPSGetter(),
+		Now:                 time.Now(),
+		TrustedRoots:        nil,
+		Product:             quoteProto.Product,
+	}
+
+	result, err := verify.GetAttestationFromReport(quoteProto.Report, options)
+	if err != nil {
+		return fmt.Errorf("could not get fetch certificates: %v", err)
+	}
+
+	homePath, _ := os.UserHomeDir()
+
+	vcekPath := path.Join(homePath, cocosDirectory, fmt.Sprintf("%d", quoteProto.Product.Name), vcekName)
+	arkAskBundlePath := path.Join(homePath, cocosDirectory, fmt.Sprintf("%d", quoteProto.Product.Name), arkAskBundleName)
+
+	vcekPem := derToPem(result.CertificateChain.VcekCert)
+	askPem := derToPem(result.CertificateChain.AskCert)
+	arkPem := derToPem(result.CertificateChain.ArkCert)
+
+	arkAskBundlePem := append(askPem, arkPem...)
+
+	vcekDir := filepath.Dir(vcekPath)
+	err = os.MkdirAll(vcekDir, 0o755)
+	if err != nil {
+		return fmt.Errorf("could not create VCEK directory: %v", err)
+	}
+	askArkBundleDir := filepath.Dir(arkAskBundlePath)
+	err = os.MkdirAll(askArkBundleDir, 0o755)
+	if err != nil {
+		return fmt.Errorf("could not create ask/ark bundle directory: %v", err)
+	}
+
+	err = os.WriteFile(vcekPath, vcekPem, 0o644)
+	if err != nil {
+		return fmt.Errorf("could not write VCEK file: %v", err)
+	}
+
+	err = os.WriteFile(arkAskBundlePath, arkAskBundlePem, 0o644)
+	if err != nil {
+		return fmt.Errorf("could not write ark/ask bundle file: %v", err)
+	}
+
+	return nil
+}

pkg/attestation/quoteprovider/sev_test.go

@@ -34,7 +34,7 @@ func TestFillInAttestationLocal(t *testing.T) {
 	require.NoError(t, err)
 
 	bundleContent := []byte("mock ASK ARK bundle")
-	bundlePath := path.Join(cocosDir, caBundleName)
+	bundlePath := path.Join(cocosDir, arkAskBundleName)
 	err = os.WriteFile(bundlePath, bundleContent, 0o644)
 	require.NoError(t, err)
 

pkg/attestation/vtpm/vtpm.go

@@ -213,7 +213,15 @@ func VTPMVerify(quote []byte, teeNonce []byte, vtpmNonce []byte, writer io.Write
 		return errors.Wrap(fmt.Errorf("failed to unmarshal quote"), err)
 	}
 
-	if err := quoteprovider.VerifyAttestationReportTLS(attestation.GetSevSnpAttestation(), teeNonce, policy); err != nil {
+	akPub := attestation.GetAkPub()
+
+	nonce := make([]byte, 0, len(teeNonce)+len(akPub))
+	nonce = append(nonce, teeNonce...)
+	nonce = append(nonce, akPub...)
+
+	attestData := sha3.Sum512(nonce)
+
+	if err := quoteprovider.VerifyAttestationReportTLS(attestation.GetSevSnpAttestation(), attestData[:], policy); err != nil {
 		return fmt.Errorf("failed to verify TEE attestation report: %v", err)
 	}
 
@@ -239,7 +247,9 @@ func VerifyQuote(quote []byte, vtpmNonce []byte, writer io.Writer, policy *attes
 		return err
 	}
 
-	ms, err := server.VerifyAttestation(attestation, server.VerifyOpts{Nonce: vtpmNonce, TrustedAKs: []crypto.PublicKey{cryptoPub}})
+	verifyOpts := server.VerifyOpts{Nonce: vtpmNonce, TrustedAKs: []crypto.PublicKey{cryptoPub}, AllowEFIAppBeforeCallingEvent: true}
+
+	ms, err := server.VerifyAttestation(attestation, verifyOpts)
 	if err != nil {
 		return errors.Wrap(fmt.Errorf("failed to verify attestation"), err)
 	}
@@ -305,14 +315,23 @@ func FetchQuote(nonce []byte) (*attest.Attestation, error) {
 }
 
 func addTEEAttestation(attestation *attest.Attestation, nonce []byte, vmpl uint) error {
-	rawTeeAttestation, err := quoteprovider.FetchAttestation(nonce, vmpl)
+	akPub := attestation.GetAkPub()
+
+	teeNonce := make([]byte, 0, len(nonce)+len(akPub))
+	teeNonce = append(teeNonce, nonce...)
+	teeNonce = append(teeNonce, akPub...)
+
+	attestData := sha3.Sum512(teeNonce)
+
+	rawTeeAttestation, err := quoteprovider.FetchAttestation(attestData[:], vmpl)
 	if err != nil {
 		return fmt.Errorf("failed to fetch TEE attestation report: %v", err)
 	}
 
-	extReport, err := abi.ReportCertsToProto(rawTeeAttestation)
+	extReport := &sevsnp.Attestation{}
+	err = proto.Unmarshal(rawTeeAttestation, extReport)
 	if err != nil {
-		return errors.Wrap(fmt.Errorf("failed to convert TEE report to proto"), err)
+		return errors.Wrap(fmt.Errorf("failed to unmarshal TEE report proto"), err)
 	}
 	attestation.TeeAttestation = &attest.Attestation_SevSnpAttestation{
 		SevSnpAttestation: extReport,
@@ -415,7 +434,7 @@ func ConvertPolicyToJSON(attestationConfiguration *attestation.Config) ([]byte,
 		return nil, errors.Wrap(ErrProtoMarshalFailed, err)
 	}
 
-	var pbMap map[string]interface{}
+	var pbMap map[string]any
 	if err := json.Unmarshal(pbJson, &pbMap); err != nil {
 		return nil, errors.Wrap(ErrJsonUnarshalFailed, err)
 	}
@@ -425,7 +444,7 @@ func ConvertPolicyToJSON(attestationConfiguration *attestation.Config) ([]byte,
 		return nil, errors.Wrap(ErrJsonMarshalFailed, err)
 	}
 
-	var pcrMap map[string]interface{}
+	var pcrMap map[string]any
 	if err := json.Unmarshal(pcrJson, &pcrMap); err != nil {
 		return nil, errors.Wrap(ErrJsonUnarshalFailed, err)
 	}

pkg/attestation/vtpm/vtpm_test.go

@@ -371,17 +371,17 @@ func TestReadPolicy(t *testing.T) {
 	require.NoError(t, err)
 	defer os.RemoveAll(tempDir)
 
-	validPolicy := map[string]interface{}{
-		"policy": map[string]interface{}{
-			"product": map[string]interface{}{
+	validPolicy := map[string]any{
+		"policy": map[string]any{
+			"product": map[string]any{
 				"name": "test-product",
 			},
 		},
-		"rootOfTrust": map[string]interface{}{
+		"rootOfTrust": map[string]any{
 			"productLine": "test-line",
 		},
-		"pcrConfig": map[string]interface{}{
-			"pcrValues": map[string]interface{}{
+		"pcrConfig": map[string]any{
+			"pcrValues": map[string]any{
 				"sha256": map[string]string{
 					"0": "0000000000000000000000000000000000000000000000000000000000000000",
 				},
@@ -557,7 +557,7 @@ func TestConvertPolicyToJSON(t *testing.T) {
 				assert.NoError(t, err)
 				assert.NotNil(t, jsonData)
 
-				var result map[string]interface{}
+				var result map[string]any
 				err = json.Unmarshal(jsonData, &result)
 				assert.NoError(t, err)
 			}

pkg/clients/clients.go

@@ -0,0 +1,40 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package clients
+
+import "time"
+
+var (
+	_ ClientConfiguration = (*AttestedClientConfig)(nil)
+	_ ClientConfiguration = (*StandardClientConfig)(nil)
+)
+
+type ClientConfiguration interface {
+	Config() StandardClientConfig
+}
+
+// StandardClientConfig represents a basic client configuration without attested TLS.
+type StandardClientConfig struct {
+	URL          string        `env:"URL"             envDefault:"localhost:7001"`
+	Timeout      time.Duration `env:"TIMEOUT"         envDefault:"60s"`
+	ClientCert   string        `env:"CLIENT_CERT"     envDefault:""`
+	ClientKey    string        `env:"CLIENT_KEY"      envDefault:""`
+	ServerCAFile string        `env:"SERVER_CA_CERTS" envDefault:""`
+}
+
+// AttestedClientConfig represents a client configuration with attested TLS capabilities.
+type AttestedClientConfig struct {
+	StandardClientConfig
+	AttestationPolicy string `env:"ATTESTATION_POLICY" envDefault:""`
+	AttestedTLS       bool   `env:"ATTESTED_TLS"       envDefault:"false"`
+	ProductName       string `env:"PRODUCT_NAME"       envDefault:"Milan"`
+}
+
+func (c AttestedClientConfig) Config() StandardClientConfig {
+	return c.StandardClientConfig
+}
+
+func (c StandardClientConfig) Config() StandardClientConfig {
+	return c
+}

pkg/clients/doc.go

@@ -2,5 +2,5 @@
 // SPDX-License-Identifier: Apache-2.0
 
 // Package clients contains the domain concept definitions needed to support
-// Agent Client functionality.
+// HTTP/gRPC Client functionality.
 package clients

pkg/clients/grpc/agent/agent.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/absmach/supermq/pkg/errors"
 	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/pkg/clients"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	grpchealth "google.golang.org/grpc/health/grpc_health_v1"
 )
@@ -14,13 +15,13 @@ import (
 var ErrAgentServiceUnavailable = errors.New("agent service is unavailable")
 
 // NewAgentClient creates new agent gRPC client instance.
-func NewAgentClient(ctx context.Context, cfg grpc.AgentClientConfig) (grpc.Client, agent.AgentServiceClient, error) {
+func NewAgentClient(ctx context.Context, cfg clients.AttestedClientConfig) (grpc.Client, agent.AgentServiceClient, error) {
 	client, err := grpc.NewClient(cfg)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	if client.Secure() != grpc.WithMATLS && client.Secure() != grpc.WithATLS && client.Secure() != grpc.WithTLS {
+	if client.Secure() != clients.WithMATLS.String() && client.Secure() != clients.WithATLS.String() && client.Secure() != clients.WithTLS.String() {
 		health := grpchealth.NewHealthClient(client.Connection())
 		resp, err := health.Check(ctx, &grpchealth.HealthCheckRequest{
 			Service: "agent",

pkg/clients/grpc/agent/agent_test.go

@@ -15,7 +15,7 @@ import (
 	"github.com/ultravioletrs/cocos/agent"
 	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
 	"github.com/ultravioletrs/cocos/agent/mocks"
-	pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
+	"github.com/ultravioletrs/cocos/pkg/clients"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health"
 	grpchealth "google.golang.org/grpc/health/grpc_health_v1"
@@ -78,14 +78,14 @@ func TestAgentClientIntegration(t *testing.T) {
 	tests := []struct {
 		name          string
 		serverRunning bool
-		config        pkggrpc.AgentClientConfig
+		config        clients.AttestedClientConfig
 		err           error
 	}{
 		{
 			name:          "successful connection",
 			serverRunning: true,
-			config: pkggrpc.AgentClientConfig{
-				BaseConfig: pkggrpc.BaseConfig{
+			config: clients.AttestedClientConfig{
+				StandardClientConfig: clients.StandardClientConfig{
 					URL:     testServer.listenAddr,
 					Timeout: 1,
 				},
@@ -95,18 +95,18 @@ func TestAgentClientIntegration(t *testing.T) {
 		{
 			name:          "server not healthy",
 			serverRunning: false,
-			config: pkggrpc.AgentClientConfig{
-				BaseConfig: pkggrpc.BaseConfig{
+			config: clients.AttestedClientConfig{
+				StandardClientConfig: clients.StandardClientConfig{
 					URL:     "",
 					Timeout: 1,
 				},
 			},
-			err: errors.New("failed to connect to grpc server"),
+			err: errors.New("agent service is unavailable"),
 		},
 		{
 			name: "invalid config, missing AttestationPolicy with aTLS",
-			config: pkggrpc.AgentClientConfig{
-				BaseConfig: pkggrpc.BaseConfig{
+			config: clients.AttestedClientConfig{
+				StandardClientConfig: clients.StandardClientConfig{
 					URL:     testServer.listenAddr,
 					Timeout: 1,
 				},
@@ -127,7 +127,7 @@ func TestAgentClientIntegration(t *testing.T) {
 			}
 
 			client, agentClient, err := NewAgentClient(ctx, tt.config)
-			assert.True(t, errors.Contains(err, tt.err))
+			assert.True(t, errors.Contains(err, tt.err), fmt.Sprintf("expected error to contain: %v, got: %v", tt.err, err))
 			if err != nil {
 				assert.Nil(t, client)
 				assert.Nil(t, agentClient)

pkg/clients/grpc/atls.go

@@ -1,133 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-package grpc
-
-import (
-	"crypto/rand"
-	"crypto/tls"
-	"crypto/x509"
-	"encoding/hex"
-	"encoding/pem"
-	"fmt"
-	"os"
-	"time"
-
-	"github.com/absmach/supermq/pkg/errors"
-	"github.com/ultravioletrs/cocos/pkg/atls"
-	"github.com/ultravioletrs/cocos/pkg/attestation"
-	"google.golang.org/grpc/credentials"
-)
-
-func setupATLS(cfg AgentClientConfig) (credentials.TransportCredentials, security, error) {
-	security := withaTLS
-
-	info, err := os.Stat(cfg.AttestationPolicy)
-	if err != nil {
-		return nil, withoutTLS, errors.Wrap(fmt.Errorf("failed to stat attestation policy file"), err)
-	}
-
-	if !info.Mode().IsRegular() {
-		return nil, withoutTLS, fmt.Errorf("attestation policy file is not a regular file: %s", cfg.AttestationPolicy)
-	}
-
-	attestation.AttestationPolicyPath = cfg.AttestationPolicy
-
-	var rootCAs *x509.CertPool = nil
-
-	if len(cfg.ServerCAFile) > 0 {
-		// Read the certificate file
-		certPEM, err := os.ReadFile(cfg.ServerCAFile)
-		if err != nil {
-			return nil, withoutTLS, errors.Wrap(fmt.Errorf("failed to read certificate file"), err)
-		}
-
-		// Decode the PEM block
-		block, _ := pem.Decode(certPEM)
-		if block == nil {
-			return nil, withoutTLS, fmt.Errorf("failed to decode PEM block")
-		}
-
-		// Parse the certificate
-		cert, err := x509.ParseCertificate(block.Bytes)
-		if err != nil {
-			return nil, withoutTLS, errors.Wrap(fmt.Errorf("failed to parse certificate"), err)
-		}
-
-		rootCAs = x509.NewCertPool()
-		rootCAs.AddCert(cert)
-
-		security = withmaTLS
-	}
-
-	nonce := make([]byte, 64)
-	if _, err := rand.Read(nonce); err != nil {
-		return nil, withoutTLS, errors.Wrap(fmt.Errorf("failed to generate nonce"), err)
-	}
-
-	encoded := hex.EncodeToString(nonce)
-	sni := fmt.Sprintf("%s.nonce", encoded)
-
-	tlsConfig := &tls.Config{
-		InsecureSkipVerify: true,
-		RootCAs:            rootCAs,
-		ServerName:         sni,
-		VerifyPeerCertificate: func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error {
-			return verifyPeerCertificateATLS(rawCerts, verifiedChains, nonce, rootCAs)
-		},
-	}
-
-	if cfg.ClientCert != "" || cfg.ClientKey != "" {
-		certificate, err := tls.LoadX509KeyPair(cfg.ClientCert, cfg.ClientKey)
-		if err != nil {
-			return nil, withoutTLS, errors.Wrap(errFailedToLoadClientCertKey, err)
-		}
-		tlsConfig.Certificates = []tls.Certificate{certificate}
-	}
-
-	return credentials.NewTLS(tlsConfig), security, nil
-}
-
-func verifyPeerCertificateATLS(rawCerts [][]byte, verifiedChains [][]*x509.Certificate, nonce []byte, rootCAs *x509.CertPool) error {
-	cert, err := x509.ParseCertificate(rawCerts[0])
-	if err != nil {
-		return errors.Wrap(errCertificateParse, err)
-	}
-
-	err = checkIfCertificateSigned(cert, rootCAs)
-	if err != nil {
-		return errors.Wrap(errAttVerification, err)
-	}
-
-	for _, ext := range cert.Extensions {
-		pType, err := atls.GetPlatformTypeFromOID(ext.Id)
-		if err == nil {
-			pubKeyDER, err := x509.MarshalPKIXPublicKey(cert.PublicKey)
-			if err != nil {
-				return fmt.Errorf("failed to marshal public key to DER format: %w", err)
-			}
-
-			return atls.VerifyCertificateExtension(ext.Value, pubKeyDER, nonce, pType)
-		}
-	}
-
-	return fmt.Errorf("attestation extension not found in certificate")
-}
-
-func checkIfCertificateSigned(cert *x509.Certificate, rootCAs *x509.CertPool) error {
-	if rootCAs == nil {
-		rootCAs = x509.NewCertPool()
-		rootCAs.AddCert(cert)
-	}
-
-	opts := x509.VerifyOptions{
-		Roots:       rootCAs,
-		CurrentTime: time.Now(),
-	}
-
-	if _, err := cert.Verify(opts); err != nil {
-		return err
-	}
-
-	return nil
-}

pkg/clients/grpc/connect_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"github.com/ultravioletrs/cocos/pkg/attestation"
 	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
+	"github.com/ultravioletrs/cocos/pkg/clients"
 )
 
 func TestNewClient(t *testing.T) {
@@ -35,14 +36,14 @@ func TestNewClient(t *testing.T) {
 
 	tests := []struct {
 		name     string
-		cfg      BaseConfig
-		agentCfg AgentClientConfig
+		cfg      clients.StandardClientConfig
+		agentCfg clients.AttestedClientConfig
 		wantErr  bool
 		err      error
 	}{
 		{
 			name: "Success without TLS",
-			cfg: BaseConfig{
+			cfg: clients.StandardClientConfig{
 				URL: "localhost:7001",
 			},
 			wantErr: false,
@@ -50,7 +51,7 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Success with TLS",
-			cfg: BaseConfig{
+			cfg: clients.StandardClientConfig{
 				URL:          "localhost:7001",
 				ServerCAFile: caCertFile,
 			},
@@ -59,7 +60,7 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Success with mTLS",
-			cfg: BaseConfig{
+			cfg: clients.StandardClientConfig{
 				URL:          "localhost:7001",
 				ServerCAFile: caCertFile,
 				ClientCert:   clientCertFile,
@@ -70,8 +71,8 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Success agent client with mTLS",
-			agentCfg: AgentClientConfig{
-				BaseConfig: BaseConfig{
+			agentCfg: clients.AttestedClientConfig{
+				StandardClientConfig: clients.StandardClientConfig{
 					URL:          "localhost:7001",
 					ServerCAFile: caCertFile,
 					ClientCert:   clientCertFile,
@@ -83,8 +84,8 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Success agent client with aTLS",
-			agentCfg: AgentClientConfig{
-				BaseConfig: BaseConfig{
+			agentCfg: clients.AttestedClientConfig{
+				StandardClientConfig: clients.StandardClientConfig{
 					URL:          "localhost:7001",
 					ServerCAFile: caCertFile,
 					ClientCert:   clientCertFile,
@@ -98,8 +99,8 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Failed agent client with aTLS",
-			agentCfg: AgentClientConfig{
-				BaseConfig: BaseConfig{
+			agentCfg: clients.AttestedClientConfig{
+				StandardClientConfig: clients.StandardClientConfig{
 					URL:          "localhost:7001",
 					ServerCAFile: caCertFile,
 					ClientCert:   clientCertFile,
@@ -113,34 +114,34 @@ func TestNewClient(t *testing.T) {
 		},
 		{
 			name: "Fail with invalid ServerCAFile",
-			cfg: BaseConfig{
+			cfg: clients.StandardClientConfig{
 				URL:          "localhost:7001",
 				ServerCAFile: "nonexistent.pem",
 			},
 			wantErr: true,
-			err:     errFailedToLoadRootCA,
+			err:     clients.ErrFailedToLoadRootCA,
 		},
 		{
 			name: "Fail with invalid ClientCert",
-			cfg: BaseConfig{
+			cfg: clients.StandardClientConfig{
 				URL:          "localhost:7001",
 				ServerCAFile: caCertFile,
 				ClientCert:   "nonexistent.pem",
 				ClientKey:    clientKeyFile,
 			},
 			wantErr: true,
-			err:     errFailedToLoadClientCertKey,
+			err:     clients.ErrFailedToLoadClientCertKey,
 		},
 		{
 			name: "Fail with invalid ClientKey",
-			cfg: BaseConfig{
+			cfg: clients.StandardClientConfig{
 				URL:          "localhost:7001",
 				ServerCAFile: caCertFile,
 				ClientCert:   clientCertFile,
 				ClientKey:    "nonexistent.pem",
 			},
 			wantErr: true,
-			err:     errFailedToLoadClientCertKey,
+			err:     clients.ErrFailedToLoadClientCertKey,
 		},
 	}
 
@@ -168,39 +169,39 @@ func TestNewClient(t *testing.T) {
 func TestClientSecure(t *testing.T) {
 	tests := []struct {
 		name     string
-		secure   security
+		secure   clients.Security
 		expected string
 	}{
 		{
 			name:     "Without TLS",
-			secure:   withoutTLS,
+			secure:   clients.WithoutTLS,
 			expected: "without TLS",
 		},
 		{
 			name:     "With TLS",
-			secure:   withTLS,
+			secure:   clients.WithTLS,
 			expected: "with TLS",
 		},
 		{
 			name:     "With mTLS",
-			secure:   withmTLS,
+			secure:   clients.WithMTLS,
 			expected: "with mTLS",
 		},
 		{
 			name:     "With aTLS",
-			secure:   withaTLS,
+			secure:   clients.WithATLS,
 			expected: "with aTLS",
 		},
 		{
 			name:     "With maTLS",
-			secure:   withmaTLS,
-			expected: WithMATLS,
+			secure:   clients.WithMATLS,
+			expected: "with maTLS",
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			c := &client{secure: tt.secure}
+			c := &client{security: tt.secure}
 			assert.Equal(t, tt.expected, c.Secure())
 		})
 	}
@@ -354,56 +355,3 @@ func createTempFile(data []byte) (string, error) {
 func createTempFileHandle() (*os.File, error) {
 	return os.CreateTemp("", "test")
 }
-
-func TestCheckIfCertificateSelfSigned(t *testing.T) {
-	selfSignedCert := createSelfSignedCert(t)
-
-	tests := []struct {
-		name string
-		cert *x509.Certificate
-		err  error
-	}{
-		{
-			name: "Self-signed certificate",
-			cert: selfSignedCert,
-			err:  nil,
-		},
-		{
-			name: "missing certificate contents",
-			cert: &x509.Certificate{},
-			err:  errors.New("x509: missing ASN.1 contents; use ParseCertificate"),
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			err := checkIfCertificateSigned(tt.cert, nil)
-			assert.True(t, errors.Contains(err, tt.err), fmt.Sprintf("expected error %v, got %v", tt.err, err))
-		})
-	}
-}
-
-func createSelfSignedCert(t *testing.T) *x509.Certificate {
-	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
-	require.NoError(t, err)
-
-	template := x509.Certificate{
-		SerialNumber: big.NewInt(1),
-		Subject: pkix.Name{
-			Organization: []string{"Test Org"},
-		},
-		NotBefore:             time.Now(),
-		NotAfter:              time.Now().Add(time.Hour * 24),
-		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
-		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
-		BasicConstraintsValid: true,
-	}
-
-	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &privateKey.PublicKey, privateKey)
-	require.NoError(t, err)
-
-	cert, err := x509.ParseCertificate(certDER)
-	require.NoError(t, err)
-
-	return cert
-}

pkg/clients/grpc/cvm/cvm.go

@@ -4,11 +4,12 @@ package cvm
 
 import (
 	"github.com/ultravioletrs/cocos/agent/cvms"
+	"github.com/ultravioletrs/cocos/pkg/clients"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
 )
 
 // NewManagerClient creates new manager gRPC client instance.
-func NewCVMClient(cfg grpc.CVMClientConfig) (grpc.Client, cvms.ServiceClient, error) {
+func NewCVMClient(cfg clients.StandardClientConfig) (grpc.Client, cvms.ServiceClient, error) {
 	client, err := grpc.NewClient(cfg)
 	if err != nil {
 		return nil, nil, err

pkg/clients/grpc/cvm/cvm_test.go

@@ -14,7 +14,7 @@ import (
 	"github.com/ultravioletrs/cocos/agent"
 	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
 	"github.com/ultravioletrs/cocos/agent/mocks"
-	pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
+	"github.com/ultravioletrs/cocos/pkg/clients"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/health"
 	grpchealth "google.golang.org/grpc/health/grpc_health_v1"
@@ -77,31 +77,18 @@ func TestAgentClientIntegration(t *testing.T) {
 	tests := []struct {
 		name          string
 		serverRunning bool
-		config        pkggrpc.CVMClientConfig
+		config        clients.StandardClientConfig
 		err           error
 	}{
 		{
 			name:          "successful connection",
 			serverRunning: true,
-			config: pkggrpc.CVMClientConfig{
-				BaseConfig: pkggrpc.BaseConfig{
-					URL:     testServer.listenAddr,
-					Timeout: 1,
-				},
+			config: clients.StandardClientConfig{
+				URL:     testServer.listenAddr,
+				Timeout: 1,
 			},
 			err: nil,
 		},
-		{
-			name:          "server not healthy",
-			serverRunning: false,
-			config: pkggrpc.CVMClientConfig{
-				BaseConfig: pkggrpc.BaseConfig{
-					URL:     "",
-					Timeout: 1,
-				},
-			},
-			err: errors.New("failed to connect to grpc server"),
-		},
 	}
 
 	for _, tt := range tests {
@@ -113,7 +100,7 @@ func TestAgentClientIntegration(t *testing.T) {
 			}
 
 			client, agentClient, err := NewCVMClient(tt.config)
-			assert.True(t, errors.Contains(err, tt.err))
+			assert.True(t, errors.Contains(err, tt.err), fmt.Sprintf("expected error to contain: %v, got: %v", tt.err, err))
 			if err != nil {
 				assert.Nil(t, client)
 				assert.Nil(t, agentClient)

pkg/clients/grpc/grpc.go

@@ -4,84 +4,19 @@
 package grpc
 
 import (
-	"crypto/tls"
-	"crypto/x509"
-	"fmt"
-	"os"
-	"time"
-
 	"github.com/absmach/supermq/pkg/errors"
+	"github.com/ultravioletrs/cocos/pkg/clients"
 	"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"
 )
 
-type security int
-
-const (
-	withoutTLS security = iota
-	withTLS
-	withmTLS
-	withaTLS
-	withmaTLS
-)
-
-const (
-	AttestationReportSize = 0x4A0
-	WithMATLS             = "with maTLS"
-	WithATLS              = "with aTLS"
-	WithTLS               = "with TLS"
-)
-
 var (
-	errGrpcConnect               = errors.New("failed to connect to grpc server")
-	errGrpcClose                 = errors.New("failed to close grpc connection")
-	errCertificateParse          = errors.New("failed to parse x509 certificate")
-	errAttVerification           = errors.New("certificat is not self signed")
-	errFailedToLoadClientCertKey = errors.New("failed to load client certificate and key")
-	errFailedToLoadRootCA        = errors.New("failed to load root ca file")
+	errGrpcConnect = errors.New("failed to connect to grpc server")
+	errGrpcClose   = errors.New("failed to close grpc connection")
 )
 
-type ClientConfiguration interface {
-	GetBaseConfig() BaseConfig
-}
-
-type BaseConfig struct {
-	URL          string        `env:"URL"             envDefault:"localhost:7001"`
-	Timeout      time.Duration `env:"TIMEOUT"         envDefault:"60s"`
-	ClientCert   string        `env:"CLIENT_CERT"     envDefault:""`
-	ClientKey    string        `env:"CLIENT_KEY"      envDefault:""`
-	ServerCAFile string        `env:"SERVER_CA_CERTS" envDefault:""`
-}
-
-type AgentClientConfig struct {
-	BaseConfig
-	AttestationPolicy string `env:"ATTESTATION_POLICY" envDefault:""`
-	AttestedTLS       bool   `env:"ATTESTED_TLS"       envDefault:"false"`
-	ProductName       string `env:"PRODUCT_NAME"       envDefault:"Milan"`
-}
-
-type ManagerClientConfig struct {
-	BaseConfig
-}
-
-type CVMClientConfig struct {
-	BaseConfig
-}
-
-func (a BaseConfig) GetBaseConfig() BaseConfig {
-	return a
-}
-
-func (a AgentClientConfig) GetBaseConfig() BaseConfig {
-	return a.BaseConfig
-}
-
-func (a CVMClientConfig) GetBaseConfig() BaseConfig {
-	return a.BaseConfig
-}
-
 type Client interface {
 	Close() error
 	Secure() string
@@ -90,14 +25,14 @@ type Client interface {
 
 type client struct {
 	*grpc.ClientConn
-	cfg    ClientConfiguration
-	secure security
+	cfg      clients.ClientConfiguration
+	security clients.Security
 }
 
 var _ Client = (*client)(nil)
 
-func NewClient(cfg ClientConfiguration) (Client, error) {
-	conn, secure, err := connect(cfg)
+func NewClient(cfg clients.ClientConfiguration) (Client, error) {
+	conn, security, err := connect(cfg)
 	if err != nil {
 		return nil, err
 	}
@@ -105,7 +40,7 @@ func NewClient(cfg ClientConfiguration) (Client, error) {
 	return &client{
 		ClientConn: conn,
 		cfg:        cfg,
-		secure:     secure,
+		security:   security,
 	}, nil
 }
 
@@ -117,86 +52,53 @@ func (c *client) Close() error {
 }
 
 func (c *client) Secure() string {
-	switch c.secure {
-	case withTLS:
-		return WithTLS
-	case withmTLS:
-		return "with mTLS"
-	case withaTLS:
-		return "with aTLS"
-	case withmaTLS:
-		return WithMATLS
-	default:
-		return "without TLS"
-	}
+	return c.security.String()
 }
 
 func (c *client) Connection() *grpc.ClientConn {
 	return c.ClientConn
 }
 
-func connect(cfg ClientConfiguration) (*grpc.ClientConn, security, error) {
+func connect(cfg clients.ClientConfiguration) (*grpc.ClientConn, clients.Security, error) {
 	opts := []grpc.DialOption{
 		grpc.WithStatsHandler(otelgrpc.NewClientHandler()),
 	}
-	secure := withoutTLS
+	security := clients.WithoutTLS
 
-	if agcfg, ok := cfg.(AgentClientConfig); ok && agcfg.AttestedTLS {
-		tc, sec, err := setupATLS(agcfg)
+	if agcfg, ok := cfg.(clients.AttestedClientConfig); ok && agcfg.AttestedTLS {
+		result, err := clients.LoadATLSConfig(agcfg)
 		if err != nil {
-			return nil, secure, err
+			return nil, security, err
 		}
 
-		opts = append(opts, grpc.WithTransportCredentials(tc))
-
-		secure = sec
+		opts = append(opts, grpc.WithTransportCredentials(credentials.NewTLS(result.Config)))
+		security = result.Security
 	} else {
-		conf := cfg.GetBaseConfig()
+		conf := cfg.Config()
 		transportCreds, sec, err := loadTLSConfig(conf.ServerCAFile, conf.ClientCert, conf.ClientKey)
 		if err != nil {
-			return nil, secure, err
+			return nil, security, err
 		}
 		opts = append(opts, grpc.WithTransportCredentials(transportCreds))
-		secure = sec
+		security = sec
 	}
 
-	conn, err := grpc.Dial(cfg.GetBaseConfig().URL, opts...)
+	conn, err := grpc.NewClient(cfg.Config().URL, opts...)
 	if err != nil {
-		return nil, secure, errors.Wrap(errGrpcConnect, err)
+		return nil, security, errors.Wrap(errGrpcConnect, err)
 	}
-	return conn, secure, nil
+	return conn, security, nil
 }
 
-func loadTLSConfig(serverCAFile, clientCert, clientKey string) (credentials.TransportCredentials, security, error) {
-	tlsConfig := &tls.Config{}
-	secure := withoutTLS
-	tc := insecure.NewCredentials()
-
-	if serverCAFile != "" {
-		rootCA, err := os.ReadFile(serverCAFile)
-		if err != nil {
-			return nil, secure, errors.Wrap(errFailedToLoadRootCA, err)
-		}
-		if len(rootCA) > 0 {
-			capool := x509.NewCertPool()
-			if !capool.AppendCertsFromPEM(rootCA) {
-				return nil, secure, fmt.Errorf("failed to append root ca to tls.Config")
-			}
-			tlsConfig.RootCAs = capool
-			secure = withTLS
-			tc = credentials.NewTLS(tlsConfig)
-		}
+func loadTLSConfig(serverCAFile, clientCert, clientKey string) (credentials.TransportCredentials, clients.Security, error) {
+	result, err := clients.LoadBasicTLSConfig(serverCAFile, clientCert, clientKey)
+	if err != nil {
+		return nil, clients.WithoutTLS, err
 	}
 
-	if clientCert != "" || clientKey != "" {
-		certificate, err := tls.LoadX509KeyPair(clientCert, clientKey)
-		if err != nil {
-			return nil, secure, errors.Wrap(errFailedToLoadClientCertKey, err)
-		}
-		tlsConfig.Certificates = []tls.Certificate{certificate}
-		secure = withmTLS
-		tc = credentials.NewTLS(tlsConfig)
+	if result.Security == clients.WithoutTLS || result.Config == nil {
+		return insecure.NewCredentials(), result.Security, nil
 	}
 
-	return tc, secure, nil
+	return credentials.NewTLS(result.Config), result.Security, nil
 }