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>

.env

@@ -1,33 +0,0 @@
-## Jaeger
-COCOS_JAEGER_PORT=6831
-COCOS_JAEGER_FRONTEND=16686
-COCOS_JAEGER_COLLECTOR=14268
-COCOS_JAEGER_CONFIGS=5778
-COCOS_JAEGER_URL=http://jaeger:14268/api/traces
-COCOS_JAEGER_TRACE_RATIO=1.0
-
-## Core Services
-
-### Manager
-MANAGER_HTTP_HOST="cocos-manager"
-MANAGER_HTTP_PORT=9021
-MANAGER_HTTP_SERVER_CERT=""
-MANAGER_HTTP_SERVER_KEY=""
-MANAGER_GRPC_HOST="cocos-manager"
-MANAGER_GRPC_PORT=7003
-MANAGER_GRPC_SERVER_CERT=""
-MANAGER_GRPC_SERVER_KEY=""
-AGENT_GRPC_URL="192.168.100.4:7002"
-AGENT_GRPC_TIMEOUT=""
-AGENT_GRPC_CA_CERTS=""
-AGENT_GRPC_CLIENT_TLS=""
-MANAGER_INSTANCE_ID=""
-MANAGER_LOG_LEVEL=debug
-MANAGER_QEMU_USE_SUDO=false
-MANAGER_QEMU_ENABLE_SEV=false
-MANAGER_QEMU_SEV_CBITPOS=51
-MANAGER_QEMU_OVMF_CODE_FILE=/usr/share/OVMF/OVMF_CODE.fd
-MANAGER_QEMU_OVMF_VARS_FILE=/usr/share/OVMF/OVMF_VARS.fd
-
-# Docker image tag
-COCOS_RELEASE_TAG=latest

.github/dependabot.yaml

@@ -0,0 +1,18 @@
+version: 2
+updates:
+  - package-ecosystem: "cargo"
+    directory: "/scripts/attestation_policy"
+    schedule:
+      interval: "weekly"
+      day: "monday"
+    groups:
+      rs-dependencies:
+        patterns:
+          - "*"
+
+  - package-ecosystem: "gomod"
+    directories:
+      - "/"
+    schedule:
+      interval: "weekly"
+      day: "monday"

.github/workflows/checkproto.yaml

@@ -9,6 +9,7 @@ on:
       - "pkg/manager/*.pb.go"
       - "agent/agent.proto"
       - "agent/*.pb.go"
+
   pull_request:
     branches:
       - main
@@ -29,13 +30,13 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: 1.22.x
+          go-version: 1.25.x
 
       - name: Set up protoc
         run: |
-          PROTOC_VERSION=27.2
-          PROTOC_GEN_VERSION=v1.34.2
-          PROTOC_GRPC_VERSION=v1.4.0
+          PROTOC_VERSION=29.0
+          PROTOC_GEN_VERSION=v1.36.8
+          PROTOC_GRPC_VERSION=v1.5.1
 
           # Download and install protoc
           PROTOC_ZIP=protoc-$PROTOC_VERSION-linux-x86_64.zip
@@ -55,7 +56,7 @@ jobs:
       - name: Set up Cocos-AI
         run: |
           # Rename .pb.go files to .pb.go.tmp to prevent conflicts
-          for p in $(ls pkg/manager/*.pb.go); do
+          for p in $(ls manager/*.pb.go); do
             mv $p $p.tmp
           done
 
@@ -67,7 +68,7 @@ jobs:
           make protoc
 
           # Compare generated Go files with the original ones
-          for p in $(ls pkg/manager/*.pb.go); do
+          for p in $(ls manager/*.pb.go); do
             if ! cmp -s $p $p.tmp; then
               echo "Proto file and generated Go file $p are out of sync!"
               exit 1

.github/workflows/hal.yml

@@ -1,73 +1,80 @@
-name: Build and Release
+name: Build and Release Hal
 
 on:
   push:
     tags:
-      - '*'
+      - "*"
 
 jobs:
   build:
     runs-on: ubuntu-latest
-
+    timeout-minutes: 120
     steps:
-    - name: Update Ubuntu
-      run: |
-        sudo apt-get update
-        sudo apt-get upgrade -y
-    
-    - name: Install Go
-      uses: actions/setup-go@v5
-      with:
-        go-version: 1.22.x
-        cache-dependency-path: "go.sum"
+      - name: Free Disk Space
+        run: |
+          sudo apt-get clean
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf "/usr/local/share/boost"
+          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+          df -h
 
-    - name: Checkout cocos
-      uses: actions/checkout@v4
-      with:
-        repository: 'ultravioletrs/cocos'
-        path: cocos
+      - name: Update Ubuntu
+        run: |
+          sudo apt-get update
+          sudo apt-get upgrade -y
+          sudo apt-get clean
+          sudo apt-get autoremove -y
+          df -h
 
-    - name: Checkout buildroot
-      uses: actions/checkout@v4
-      with:
-        repository: 'buildroot/buildroot'
-        path: buildroot
+      - name: Maximize build space
+        uses: easimon/maximize-build-space@master
+        with:
+          root-reserve-mb: 35000
+          swap-size-mb: 1024
+          remove-dotnet: "true"
+          remove-android: "true"
+      - name: Check free space
+        run: |
+          echo "Free space:"
+          df -h
 
-    - name: Build
-      run: |
-        cd buildroot
-        make BR2_EXTERNAL=../cocos/hal/linux cocos_defconfig
-        make
+      - name: Install Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: 1.25.x
+          cache-dependency-path: "go.sum"
 
-    - name: Create Release
-      id: create_release
-      uses: actions/create-release@latest
-      env:
-        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-      with:
-        tag_name: ${{ github.ref }}
-        release_name: Release ${{ github.ref }}
-        draft: false
-        prerelease: false
+      - name: Checkout cocos
+        uses: actions/checkout@v4
+        with:
+          repository: "ultravioletrs/cocos"
+          path: cocos
 
-    - name: Upload Release Asset
-      id: upload-release-kernel
-      uses: actions/upload-release-asset@latest
-      env:
-        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-      with:
-        upload_url: ${{ steps.create_release.outputs.upload_url }}
-        asset_path: ./buildroot/output/images/bzImage
-        asset_name: bzImage
-        asset_content_type: application/octet-stream
+      - name: Checkout buildroot
+        uses: actions/checkout@v4
+        with:
+          repository: "buildroot/buildroot"
+          path: buildroot
+          ref: 2025.08-rc3
 
-    - name: Upload Release Asset
-      id: upload-release-rootfs
-      uses: actions/upload-release-asset@latest
-      env:
-        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-      with:
-        upload_url: ${{ steps.create_release.outputs.upload_url }}
-        asset_path: ./buildroot/output/images/rootfs.cpio.gz
-        asset_name: rootfs.cpio.gz
-        asset_content_type: application/gzip
+      - name: Build hal
+        run: |
+          cd buildroot
+          make BR2_EXTERNAL=../cocos/hal/linux cocos_defconfig
+          make
+
+      - name: Build cocos
+        run: |
+          cd cocos
+          make
+
+      - name: Release
+        uses: softprops/action-gh-release@v2
+        with:
+          files: |
+            buildroot/output/images/bzImage
+            buildroot/output/images/rootfs.cpio.gz
+            cocos/build/cocos-agent  
+            cocos/build/cocos-cli  
+            cocos/build/cocos-manager

.github/workflows/main.yaml

@@ -9,9 +9,8 @@ on:
       - main
 
 jobs:
-  ci:
+  lint:
     runs-on: ubuntu-latest
-
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -19,16 +18,71 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: 1.22.x
+          go-version: 1.25.x
 
       - name: golangci-lint
-        uses: golangci/golangci-lint-action@v6
+        uses: golangci/golangci-lint-action@v8
         with:
-          version: v1.59.1
+          version: v2.4.0
 
       - name: Build
-        run: |
-          make
+        run: make
 
-      - name: Run tests
-        run: go test -v --race -covermode=atomic -coverprofile cover.out ./...
+  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: |
+          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: Upload coverage artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-${{ matrix.module }}
+          path: coverage/${{ matrix.module }}.out
+          retention-days: 1
+
+  upload-coverage:
+    runs-on: ubuntu-latest
+    needs: test
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - 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
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          files: ./coverage/*.out
+          codecov_yml_path: codecov.yml
+          verbose: true

.github/workflows/rust.yaml

@@ -0,0 +1,41 @@
+name: Rust CI Pipeline
+
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - "scripts/attestation_policy/**"
+      - ".github/workflows/rust.yaml"
+  pull_request:
+    branches:
+      - main
+    paths:
+      - "scripts/attestation_policy/**"
+      - ".github/workflows/rust.yaml"
+
+env:
+  CARGO_TERM_COLOR: always
+
+jobs:
+  rust-check:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: ./scripts/attestation_policy
+
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+
+      - name: Check cargo
+        run: cargo check --release --all-targets
+
+      - name: Check formatting
+        run: cargo fmt --all -- --check
+
+      - name: Run linter
+        run: cargo clippy -- -D warnings
+
+      - name: Build for all features
+        run: cargo build --release --all-features

.gitignore

@@ -1,10 +1,27 @@
 build
 
-build
 cmd/manager/img
-cmd/manager/iso
-cmd/manager/tmp
 
 .cov
 
 *.pem
+
+dist/
+*.zip
+*.spec
+*.tar
+
+# Generated by Cargo
+# will have compiled files and executables
+debug/
+target/
+
+# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
+# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
+Cargo.lock
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
+
+# MSVC Windows builds of rustc generate these, which store debugging information
+*.pdb

.golangci.yaml

@@ -1,80 +1,94 @@
-run:
-  timeout: 3m
-
-issues:
-  max-issues-per-linter: 10
-  max-same-issues: 10
-
-linters-settings:
-  importas:
-    no-unaliased: true
-    no-extra-aliases: false
-    alias:
-      - pkg: github.com/absmach/magistrala/logger
-        alias: mglog
-
-  gocritic:
-    enabled-checks:
-      - dupImport
-      - importShadow
-      - httpNoBody
-      - paramTypeCombine
-      - emptyStringTest
-      - builtinShadow
-      - exposedSyncMutex
-    disabled-checks:
-      - appendAssign
-    enabled-tags:
-      - diagnostic
-    disabled-tags:
-      - performance
-      - style
-      - experimental
-      - opinionated
-  stylecheck:
-    checks: ["-ST1000", "-ST1003", "-ST1020", "-ST1021", "-ST1022"]
-  goheader:
-    template: |-
-      Copyright (c) Ultraviolet
-      SPDX-License-Identifier: Apache-2.0
-
+version: "2"
 linters:
-  disable-all: true
+  default: none
   enable:
-    - importas
-    - gocritic
-    - gosimple
-    - errcheck
-    - govet
-    - unused
-    - goconst
-    - godot
-    - godox
-    - ineffassign
-    - misspell
-    - stylecheck
-    - whitespace
-    - gci
-    - gofmt
-    - goimports
-    - loggercheck
-    - goheader
     - asasalint
     - asciicheck
     - bidichk
     - contextcheck
+    - copyloopvar
     - decorder
     - dogsled
+    - dupword
+    - errcheck
     - errchkjson
     - errname
-    - execinquery
-    - exportloopref
     - ginkgolinter
     - gocheckcompilerdirectives
-    - gofumpt
+    - goconst
+    - gocritic
+    - godot
+    - godox
+    - goheader
     - goprintffuncname
+    - govet
     - importas
+    - ineffassign
+    - loggercheck
     - makezero
     - mirror
+    - misspell
     - nakedret
-    - dupword
+    - staticcheck
+    - unused
+    - whitespace
+  settings:
+    gocritic:
+      enabled-checks:
+        - dupImport
+        - importShadow
+        - httpNoBody
+        - paramTypeCombine
+        - emptyStringTest
+        - builtinShadow
+        - exposedSyncMutex
+      disabled-checks:
+        - appendAssign
+      enabled-tags:
+        - diagnostic
+      disabled-tags:
+        - performance
+        - style
+        - experimental
+        - opinionated
+    goheader:
+      template: |-
+        Copyright (c) Ultraviolet
+        SPDX-License-Identifier: Apache-2.0
+    staticcheck:
+      checks:
+        - -ST1000
+        - -ST1003
+        - -ST1020
+        - -ST1021
+        - -ST1022
+  exclusions:
+    generated: lax
+    presets:
+      - comments
+      - common-false-positives
+      - legacy
+      - std-error-handling
+    rules:
+      - linters:
+          - makezero
+        text: with non-zero initialized length
+    paths:
+      - third_party$
+      - builtin$
+      - examples$
+issues:
+  max-issues-per-linter: 10
+  max-same-issues: 10
+formatters:
+  enable:
+    - gci
+    - gofmt
+    - gofumpt
+    - goimports
+  exclusions:
+    generated: lax
+    paths:
+      - third_party$
+      - builtin$
+      - examples$

.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

@@ -1,34 +1,71 @@
 BUILD_DIR = build
 SERVICES = manager agent cli
-PLATFORM_INFO = platform_info
+ATTESTATION_POLICY = attestation_policy
 CGO_ENABLED ?= 0
 GOARCH ?= amd64
 VERSION ?= $(shell git describe --abbrev=0 --tags --always)
 COMMIT ?= $(shell git rev-parse HEAD)
 TIME ?= $(shell date +%F_%T)
+EMBED_ENABLED ?= 0
+INSTALL_DIR ?= /usr/local/bin
+CONFIG_DIR ?= /etc/cocos
+SERVICE_NAME ?= cocos-manager
+SERVICE_DIR ?= /etc/systemd/system
+SERVICE_FILE = init/systemd/$(SERVICE_NAME).service
+IGVM_BUILD_SCRIPT := ./scripts/igvmmeasure/igvm.sh
 
 define compile_service
 	CGO_ENABLED=$(CGO_ENABLED) GOOS=$(GOOS) GOARCH=$(GOARCH) GOARM=$(GOARM) \
 	go build -ldflags "-s -w \
-	-X 'github.com/absmach/magistrala.BuildTime=$(TIME)' \
-	-X 'github.com/absmach/magistrala.Version=$(VERSION)' \
-	-X 'github.com/absmach/magistrala.Commit=$(COMMIT)'" \
+	-X 'github.com/absmach/supermq.BuildTime=$(TIME)' \
+	-X 'github.com/absmach/supermq.Version=$(VERSION)' \
+	-X 'github.com/absmach/supermq.Commit=$(COMMIT)'" \
+	$(if $(filter 1,$(EMBED_ENABLED)),-tags "embed",) \
 	-o ${BUILD_DIR}/cocos-$(1) cmd/$(1)/main.go
 endef
 
-.PHONY: all $(SERVICES) $(PLATFORM_INFO)
+.PHONY: all $(SERVICES) $(ATTESTATION_POLICY) install clean
 
-all: $(SERVICES)
+all: $(SERVICES) $(ATTESTATION_POLICY)
 
-$(SERVICES):
-	$(call compile_service,$(@))
+$(SERVICES): 
+	$(call compile_service,$@)
+	@if [ "$@" = "cli" ] || [ "$@" = "manager" ]; then $(MAKE) build-igvm; fi
 
-$(PLATFORM_INFO):
-	$(MAKE) -C ./scripts/platform_info
+$(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
-	protoc -I. --go_out=./pkg --go_opt=paths=source_relative --go-grpc_out=./pkg --go-grpc_opt=paths=source_relative manager/manager.proto
+	protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative manager/manager.proto
+	protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/events/events.proto
+	protoc -I. --go_out=. --go_opt=paths=source_relative --go-grpc_out=. --go-grpc_opt=paths=source_relative agent/cvms/cvms.proto
 
 mocks:
-	go generate ./...
+	mockery --config ./.mockery.yml
+
+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
+
+stop:
+	sudo systemctl stop $(SERVICE_NAME).service
+
+install_service:
+	sudo install -m 644 $(SERVICE_FILE) $(SERVICE_DIR)/$(SERVICE_NAME).service
+	sudo systemctl daemon-reload
+
+build-igvm:
+	@echo "Running build script for igvmmeasure..."
+	@$(IGVM_BUILD_SCRIPT)

README.md

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

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) | info                           |
-| AGENT_GRPC_HOST               | Agent service gRPC host                                | ""                             |
-| AGENT_GRPC_PORT               | Agent service gRPC port                                | 7002                           |
-| AGENT_GRPC_SERVER_CERT        | Path to gRPC server certificate in pem format          | ""                             |
-| AGENT_GRPC_SERVER_KEY         | Path to gRPC server key in pem format                  | ""                             |
-| AGENT_GRPC_SERVER_CA_CERTS    | Path to gRPC server CA certificate                     | ""                             |
-| AGENT_GRPC_CLIENT_CA_CERTS    | Path to gRPC client CA certificate                     | ""                             |
-
+| Variable                       | Description                                                                                                   | Default                                         |
+| ------------------------------ | ------------------------------------------------------------------------------------------------------------- | ----------------------------------------------- |
+| AGENT_LOG_LEVEL                | Log level for agent service (debug, info, warn, error)                                                        | debug                                           |
+| AGENT_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,8 +3,8 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.34.2
-// 	protoc        v5.27.2
+// 	protoc-gen-go v1.36.8
+// 	protoc        v5.29.0
 // source: agent/agent.proto
 
 package agent
@@ -14,6 +14,7 @@ import (
 	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 	reflect "reflect"
 	sync "sync"
+	unsafe "unsafe"
 )
 
 const (
@@ -24,20 +25,18 @@ const (
 )
 
 type AlgoRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	Algorithm     []byte                 `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
+	Requirements  []byte                 `protobuf:"bytes,2,opt,name=requirements,proto3" json:"requirements,omitempty"`
 	unknownFields protoimpl.UnknownFields
-
-	Algorithm []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *AlgoRequest) Reset() {
 	*x = AlgoRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[0]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[0]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *AlgoRequest) String() string {
@@ -48,7 +47,7 @@ func (*AlgoRequest) ProtoMessage() {}
 
 func (x *AlgoRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[0]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -70,19 +69,24 @@ func (x *AlgoRequest) GetAlgorithm() []byte {
 	return nil
 }
 
+func (x *AlgoRequest) GetRequirements() []byte {
+	if x != nil {
+		return x.Requirements
+	}
+	return nil
+}
+
 type AlgoResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
 	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *AlgoResponse) Reset() {
 	*x = AlgoResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[1]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[1]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *AlgoResponse) String() string {
@@ -93,7 +97,7 @@ func (*AlgoResponse) ProtoMessage() {}
 
 func (x *AlgoResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[1]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -109,20 +113,18 @@ func (*AlgoResponse) Descriptor() ([]byte, []int) {
 }
 
 type DataRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	Dataset       []byte                 `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
+	Filename      string                 `protobuf:"bytes,2,opt,name=filename,proto3" json:"filename,omitempty"`
 	unknownFields protoimpl.UnknownFields
-
-	Dataset []byte `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *DataRequest) Reset() {
 	*x = DataRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[2]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[2]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *DataRequest) String() string {
@@ -133,7 +135,7 @@ func (*DataRequest) ProtoMessage() {}
 
 func (x *DataRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[2]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -155,19 +157,24 @@ func (x *DataRequest) GetDataset() []byte {
 	return nil
 }
 
+func (x *DataRequest) GetFilename() string {
+	if x != nil {
+		return x.Filename
+	}
+	return ""
+}
+
 type DataResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
 	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *DataResponse) Reset() {
 	*x = DataResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[3]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[3]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *DataResponse) String() string {
@@ -178,7 +185,7 @@ func (*DataResponse) ProtoMessage() {}
 
 func (x *DataResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[3]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -194,18 +201,16 @@ func (*DataResponse) Descriptor() ([]byte, []int) {
 }
 
 type ResultRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
 	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *ResultRequest) Reset() {
 	*x = ResultRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[4]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[4]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *ResultRequest) String() string {
@@ -216,7 +221,7 @@ func (*ResultRequest) ProtoMessage() {}
 
 func (x *ResultRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[4]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -232,20 +237,17 @@ func (*ResultRequest) Descriptor() ([]byte, []int) {
 }
 
 type ResultResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	File          []byte                 `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
 	unknownFields protoimpl.UnknownFields
-
-	File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *ResultResponse) Reset() {
 	*x = ResultResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[5]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[5]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *ResultResponse) String() string {
@@ -256,7 +258,7 @@ func (*ResultResponse) ProtoMessage() {}
 
 func (x *ResultResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[5]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -279,20 +281,19 @@ func (x *ResultResponse) GetFile() []byte {
 }
 
 type AttestationRequest struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	TeeNonce      []byte                 `protobuf:"bytes,1,opt,name=teeNonce,proto3" json:"teeNonce,omitempty"`   // Should be less or equal 64 bytes.
+	VtpmNonce     []byte                 `protobuf:"bytes,2,opt,name=vtpmNonce,proto3" json:"vtpmNonce,omitempty"` // Should be less or equal 32 bytes.
+	Type          int32                  `protobuf:"varint,3,opt,name=type,proto3" json:"type,omitempty"`
 	unknownFields protoimpl.UnknownFields
-
-	ReportData []byte `protobuf:"bytes,1,opt,name=report_data,json=reportData,proto3" json:"report_data,omitempty"` // Should be of length 64.
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *AttestationRequest) Reset() {
 	*x = AttestationRequest{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[6]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[6]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *AttestationRequest) String() string {
@@ -303,7 +304,7 @@ func (*AttestationRequest) ProtoMessage() {}
 
 func (x *AttestationRequest) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[6]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -318,28 +319,39 @@ func (*AttestationRequest) Descriptor() ([]byte, []int) {
 	return file_agent_agent_proto_rawDescGZIP(), []int{6}
 }
 
-func (x *AttestationRequest) GetReportData() []byte {
+func (x *AttestationRequest) GetTeeNonce() []byte {
 	if x != nil {
-		return x.ReportData
+		return x.TeeNonce
 	}
 	return nil
 }
 
-type AttestationResponse struct {
-	state         protoimpl.MessageState
-	sizeCache     protoimpl.SizeCache
-	unknownFields protoimpl.UnknownFields
+func (x *AttestationRequest) GetVtpmNonce() []byte {
+	if x != nil {
+		return x.VtpmNonce
+	}
+	return nil
+}
 
-	File []byte `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
+func (x *AttestationRequest) GetType() int32 {
+	if x != nil {
+		return x.Type
+	}
+	return 0
+}
+
+type AttestationResponse struct {
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	File          []byte                 `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
 }
 
 func (x *AttestationResponse) Reset() {
 	*x = AttestationResponse{}
-	if protoimpl.UnsafeEnabled {
-		mi := &file_agent_agent_proto_msgTypes[7]
-		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
-		ms.StoreMessageInfo(mi)
-	}
+	mi := &file_agent_agent_proto_msgTypes[7]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
 }
 
 func (x *AttestationResponse) String() string {
@@ -350,7 +362,7 @@ func (*AttestationResponse) ProtoMessage() {}
 
 func (x *AttestationResponse) ProtoReflect() protoreflect.Message {
 	mi := &file_agent_agent_proto_msgTypes[7]
-	if protoimpl.UnsafeEnabled && x != nil {
+	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
 			ms.StoreMessageInfo(mi)
@@ -372,84 +384,276 @@ func (x *AttestationResponse) GetFile() []byte {
 	return nil
 }
 
+type IMAMeasurementsRequest struct {
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *IMAMeasurementsRequest) Reset() {
+	*x = IMAMeasurementsRequest{}
+	mi := &file_agent_agent_proto_msgTypes[8]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *IMAMeasurementsRequest) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*IMAMeasurementsRequest) ProtoMessage() {}
+
+func (x *IMAMeasurementsRequest) ProtoReflect() protoreflect.Message {
+	mi := &file_agent_agent_proto_msgTypes[8]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use IMAMeasurementsRequest.ProtoReflect.Descriptor instead.
+func (*IMAMeasurementsRequest) Descriptor() ([]byte, []int) {
+	return file_agent_agent_proto_rawDescGZIP(), []int{8}
+}
+
+type IMAMeasurementsResponse struct {
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	File          []byte                 `protobuf:"bytes,1,opt,name=file,proto3" json:"file,omitempty"`
+	Pcr10         []byte                 `protobuf:"bytes,2,opt,name=pcr10,proto3" json:"pcr10,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *IMAMeasurementsResponse) Reset() {
+	*x = IMAMeasurementsResponse{}
+	mi := &file_agent_agent_proto_msgTypes[9]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *IMAMeasurementsResponse) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*IMAMeasurementsResponse) ProtoMessage() {}
+
+func (x *IMAMeasurementsResponse) ProtoReflect() protoreflect.Message {
+	mi := &file_agent_agent_proto_msgTypes[9]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use IMAMeasurementsResponse.ProtoReflect.Descriptor instead.
+func (*IMAMeasurementsResponse) Descriptor() ([]byte, []int) {
+	return file_agent_agent_proto_rawDescGZIP(), []int{9}
+}
+
+func (x *IMAMeasurementsResponse) GetFile() []byte {
+	if x != nil {
+		return x.File
+	}
+	return nil
+}
+
+func (x *IMAMeasurementsResponse) GetPcr10() []byte {
+	if x != nil {
+		return x.Pcr10
+	}
+	return nil
+}
+
+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"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+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 *AttestationTokenRequest) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*AttestationTokenRequest) ProtoMessage() {}
+
+func (x *AttestationTokenRequest) ProtoReflect() protoreflect.Message {
+	mi := &file_agent_agent_proto_msgTypes[10]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use AttestationTokenRequest.ProtoReflect.Descriptor instead.
+func (*AttestationTokenRequest) Descriptor() ([]byte, []int) {
+	return file_agent_agent_proto_rawDescGZIP(), []int{10}
+}
+
+func (x *AttestationTokenRequest) GetTokenNonce() []byte {
+	if x != nil {
+		return x.TokenNonce
+	}
+	return nil
+}
+
+func (x *AttestationTokenRequest) GetType() int32 {
+	if x != nil {
+		return x.Type
+	}
+	return 0
+}
+
+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 *AttestationTokenResponse) Reset() {
+	*x = AttestationTokenResponse{}
+	mi := &file_agent_agent_proto_msgTypes[11]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *AttestationTokenResponse) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*AttestationTokenResponse) ProtoMessage() {}
+
+func (x *AttestationTokenResponse) ProtoReflect() protoreflect.Message {
+	mi := &file_agent_agent_proto_msgTypes[11]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use AttestationTokenResponse.ProtoReflect.Descriptor instead.
+func (*AttestationTokenResponse) Descriptor() ([]byte, []int) {
+	return file_agent_agent_proto_rawDescGZIP(), []int{11}
+}
+
+func (x *AttestationTokenResponse) GetFile() []byte {
+	if x != nil {
+		return x.File
+	}
+	return nil
+}
+
 var File_agent_agent_proto protoreflect.FileDescriptor
 
-var file_agent_agent_proto_rawDesc = []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, 0x2b, 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, 0x22, 0x0e, 0x0a, 0x0c, 0x41, 0x6c, 0x67, 0x6f, 0x52,
-	0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x27, 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,
-	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, 0x35, 0x0a, 0x12, 0x41, 0x74, 0x74, 0x65, 0x73,
-	0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1f, 0x0a,
-	0x0b, 0x72, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x5f, 0x64, 0x61, 0x74, 0x61, 0x18, 0x01, 0x20, 0x01,
-	0x28, 0x0c, 0x52, 0x0a, 0x72, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x44, 0x61, 0x74, 0x61, 0x22, 0x29,
-	0x0a, 0x13, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73,
-	0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x12, 0x0a, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x18, 0x01, 0x20,
-	0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x32, 0xf9, 0x01, 0x0a, 0x0c, 0x41, 0x67,
-	0x65, 0x6e, 0x74, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x33, 0x0a, 0x04, 0x41, 0x6c,
-	0x67, 0x6f, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f, 0x52,
-	0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41,
-	0x6c, 0x67, 0x6f, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x28, 0x01, 0x12,
-	0x33, 0x0a, 0x04, 0x44, 0x61, 0x74, 0x61, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e,
-	0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67,
-	0x65, 0x6e, 0x74, 0x2e, 0x44, 0x61, 0x74, 0x61, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65,
-	0x22, 0x00, 0x28, 0x01, 0x12, 0x37, 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, 0x12, 0x46, 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, 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
-	file_agent_agent_proto_rawDescData = file_agent_agent_proto_rawDesc
+	file_agent_agent_proto_rawDescData []byte
 )
 
 func file_agent_agent_proto_rawDescGZIP() []byte {
 	file_agent_agent_proto_rawDescOnce.Do(func() {
-		file_agent_agent_proto_rawDescData = protoimpl.X.CompressGZIP(file_agent_agent_proto_rawDescData)
+		file_agent_agent_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_agent_agent_proto_rawDesc), len(file_agent_agent_proto_rawDesc)))
 	})
 	return file_agent_agent_proto_rawDescData
 }
 
-var file_agent_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 8)
+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
+	(*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
-	2, // 1: agent.AgentService.Data:input_type -> agent.DataRequest
-	4, // 2: agent.AgentService.Result:input_type -> agent.ResultRequest
-	6, // 3: agent.AgentService.Attestation:input_type -> agent.AttestationRequest
-	1, // 4: agent.AgentService.Algo:output_type -> agent.AlgoResponse
-	3, // 5: agent.AgentService.Data:output_type -> agent.DataResponse
-	5, // 6: agent.AgentService.Result:output_type -> agent.ResultResponse
-	7, // 7: agent.AgentService.Attestation:output_type -> agent.AttestationResponse
-	4, // [4:8] is the sub-list for method output_type
-	0, // [0:4] is the sub-list for method input_type
-	0, // [0:0] is the sub-list for extension type_name
-	0, // [0:0] is the sub-list for extension extendee
-	0, // [0:0] is the sub-list for field type_name
+	0,  // 0: agent.AgentService.Algo:input_type -> agent.AlgoRequest
+	2,  // 1: agent.AgentService.Data:input_type -> agent.DataRequest
+	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.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.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
+	0,  // [0:0] is the sub-list for extension extendee
+	0,  // [0:0] is the sub-list for field type_name
 }
 
 func init() { file_agent_agent_proto_init() }
@@ -457,111 +661,13 @@ func file_agent_agent_proto_init() {
 	if File_agent_agent_proto != nil {
 		return
 	}
-	if !protoimpl.UnsafeEnabled {
-		file_agent_agent_proto_msgTypes[0].Exporter = func(v any, i int) any {
-			switch v := v.(*AlgoRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[1].Exporter = func(v any, i int) any {
-			switch v := v.(*AlgoResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[2].Exporter = func(v any, i int) any {
-			switch v := v.(*DataRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[3].Exporter = func(v any, i int) any {
-			switch v := v.(*DataResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[4].Exporter = func(v any, i int) any {
-			switch v := v.(*ResultRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[5].Exporter = func(v any, i int) any {
-			switch v := v.(*ResultResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[6].Exporter = func(v any, i int) any {
-			switch v := v.(*AttestationRequest); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-		file_agent_agent_proto_msgTypes[7].Exporter = func(v any, i int) any {
-			switch v := v.(*AttestationResponse); i {
-			case 0:
-				return &v.state
-			case 1:
-				return &v.sizeCache
-			case 2:
-				return &v.unknownFields
-			default:
-				return nil
-			}
-		}
-	}
 	type x struct{}
 	out := protoimpl.TypeBuilder{
 		File: protoimpl.DescBuilder{
 			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
-			RawDescriptor: file_agent_agent_proto_rawDesc,
+			RawDescriptor: unsafe.Slice(unsafe.StringData(file_agent_agent_proto_rawDesc), len(file_agent_agent_proto_rawDesc)),
 			NumEnums:      0,
-			NumMessages:   8,
+			NumMessages:   12,
 			NumExtensions: 0,
 			NumServices:   1,
 		},
@@ -570,7 +676,6 @@ func file_agent_agent_proto_init() {
 		MessageInfos:      file_agent_agent_proto_msgTypes,
 	}.Build()
 	File_agent_agent_proto = out.File
-	file_agent_agent_proto_rawDesc = nil
 	file_agent_agent_proto_goTypes = nil
 	file_agent_agent_proto_depIdxs = nil
 }

agent/agent.proto

@@ -10,18 +10,22 @@ option go_package = "./agent";
 service AgentService {
   rpc Algo(stream AlgoRequest) returns (AlgoResponse) {}
   rpc Data(stream DataRequest) returns (DataResponse) {}
-  rpc Result(ResultRequest) returns (ResultResponse) {}
-  rpc Attestation(AttestationRequest) returns (AttestationResponse) {}
+  rpc Result(ResultRequest) returns (stream ResultResponse) {}
+  rpc Attestation(AttestationRequest) returns (stream AttestationResponse) {}
+  rpc IMAMeasurements(IMAMeasurementsRequest) returns (stream IMAMeasurementsResponse) {}
+  rpc AzureAttestationToken(AttestationTokenRequest) returns (AttestationTokenResponse) {}
 }
 
 message AlgoRequest {
   bytes algorithm = 1;
+  bytes requirements = 2;
 }
 
 message AlgoResponse {}
 
 message DataRequest {
   bytes dataset = 1;
+  string filename = 2;
 }
 
 message DataResponse {}
@@ -34,9 +38,27 @@ message ResultResponse {
 }
 
 message AttestationRequest {
-  bytes report_data = 1; // Should be of length 64.
+  bytes teeNonce = 1; // Should be less or equal 64 bytes.
+  bytes vtpmNonce = 2; // Should be less or equal 32 bytes.
+  int32 type = 3;
 }
 
 message AttestationResponse {
   bytes file = 1;
 }
+
+message IMAMeasurementsRequest {
+}
+
+message IMAMeasurementsResponse {
+  bytes file = 1;
+  bytes pcr10 = 2;
+}
+
+message AttestationTokenRequest{
+  bytes tokenNonce = 1; // Should be less or equal 32 bytes
+  int32 type = 3;
+}
+message AttestationTokenResponse{
+  bytes file = 1;
+}

agent/agent_grpc.pb.go

@@ -3,8 +3,8 @@
 
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
-// - protoc-gen-go-grpc v1.4.0
-// - protoc             v5.27.2
+// - protoc-gen-go-grpc v1.5.1
+// - protoc             v5.29.0
 // source: agent/agent.proto
 
 package agent
@@ -18,24 +18,28 @@ import (
 
 // This is a compile-time assertion to ensure that this generated file
 // is compatible with the grpc package it is being compiled against.
-// Requires gRPC-Go v1.62.0 or later.
-const _ = grpc.SupportPackageIsVersion8
+// Requires gRPC-Go v1.64.0 or later.
+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_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.
 //
 // For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
 type AgentServiceClient interface {
-	Algo(ctx context.Context, opts ...grpc.CallOption) (AgentService_AlgoClient, error)
-	Data(ctx context.Context, opts ...grpc.CallOption) (AgentService_DataClient, error)
-	Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (*ResultResponse, error)
-	Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (*AttestationResponse, error)
+	Algo(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[AlgoRequest, AlgoResponse], error)
+	Data(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[DataRequest, DataResponse], error)
+	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)
+	AzureAttestationToken(ctx context.Context, in *AttestationTokenRequest, opts ...grpc.CallOption) (*AttestationTokenResponse, error)
 }
 
 type agentServiceClient struct {
@@ -46,90 +50,93 @@ func NewAgentServiceClient(cc grpc.ClientConnInterface) AgentServiceClient {
 	return &agentServiceClient{cc}
 }
 
-func (c *agentServiceClient) Algo(ctx context.Context, opts ...grpc.CallOption) (AgentService_AlgoClient, error) {
+func (c *agentServiceClient) Algo(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[AlgoRequest, AlgoResponse], error) {
 	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[0], AgentService_Algo_FullMethodName, cOpts...)
 	if err != nil {
 		return nil, err
 	}
-	x := &agentServiceAlgoClient{ClientStream: stream}
+	x := &grpc.GenericClientStream[AlgoRequest, AlgoResponse]{ClientStream: stream}
 	return x, nil
 }
 
-type AgentService_AlgoClient interface {
-	Send(*AlgoRequest) error
-	CloseAndRecv() (*AlgoResponse, error)
-	grpc.ClientStream
-}
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type AgentService_AlgoClient = grpc.ClientStreamingClient[AlgoRequest, AlgoResponse]
 
-type agentServiceAlgoClient struct {
-	grpc.ClientStream
-}
-
-func (x *agentServiceAlgoClient) Send(m *AlgoRequest) error {
-	return x.ClientStream.SendMsg(m)
-}
-
-func (x *agentServiceAlgoClient) CloseAndRecv() (*AlgoResponse, error) {
-	if err := x.ClientStream.CloseSend(); err != nil {
-		return nil, err
-	}
-	m := new(AlgoResponse)
-	if err := x.ClientStream.RecvMsg(m); err != nil {
-		return nil, err
-	}
-	return m, nil
-}
-
-func (c *agentServiceClient) Data(ctx context.Context, opts ...grpc.CallOption) (AgentService_DataClient, error) {
+func (c *agentServiceClient) Data(ctx context.Context, opts ...grpc.CallOption) (grpc.ClientStreamingClient[DataRequest, DataResponse], error) {
 	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[1], AgentService_Data_FullMethodName, cOpts...)
 	if err != nil {
 		return nil, err
 	}
-	x := &agentServiceDataClient{ClientStream: stream}
+	x := &grpc.GenericClientStream[DataRequest, DataResponse]{ClientStream: stream}
 	return x, nil
 }
 
-type AgentService_DataClient interface {
-	Send(*DataRequest) error
-	CloseAndRecv() (*DataResponse, error)
-	grpc.ClientStream
-}
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type AgentService_DataClient = grpc.ClientStreamingClient[DataRequest, DataResponse]
 
-type agentServiceDataClient struct {
-	grpc.ClientStream
-}
-
-func (x *agentServiceDataClient) Send(m *DataRequest) error {
-	return x.ClientStream.SendMsg(m)
-}
-
-func (x *agentServiceDataClient) CloseAndRecv() (*DataResponse, error) {
-	if err := x.ClientStream.CloseSend(); err != nil {
-		return nil, err
-	}
-	m := new(DataResponse)
-	if err := x.ClientStream.RecvMsg(m); err != nil {
-		return nil, err
-	}
-	return m, nil
-}
-
-func (c *agentServiceClient) Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (*ResultResponse, error) {
+func (c *agentServiceClient) Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ResultResponse], error) {
 	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
-	out := new(ResultResponse)
-	err := c.cc.Invoke(ctx, AgentService_Result_FullMethodName, in, out, cOpts...)
+	stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[2], AgentService_Result_FullMethodName, cOpts...)
 	if err != nil {
 		return nil, err
 	}
-	return out, nil
+	x := &grpc.GenericClientStream[ResultRequest, ResultResponse]{ClientStream: stream}
+	if err := x.ClientStream.SendMsg(in); err != nil {
+		return nil, err
+	}
+	if err := x.ClientStream.CloseSend(); err != nil {
+		return nil, err
+	}
+	return x, nil
 }
 
-func (c *agentServiceClient) Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (*AttestationResponse, error) {
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type AgentService_ResultClient = grpc.ServerStreamingClient[ResultResponse]
+
+func (c *agentServiceClient) Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[AttestationResponse], error) {
 	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
-	out := new(AttestationResponse)
-	err := c.cc.Invoke(ctx, AgentService_Attestation_FullMethodName, in, out, cOpts...)
+	stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[3], AgentService_Attestation_FullMethodName, cOpts...)
+	if err != nil {
+		return nil, err
+	}
+	x := &grpc.GenericClientStream[AttestationRequest, AttestationResponse]{ClientStream: stream}
+	if err := x.ClientStream.SendMsg(in); err != nil {
+		return nil, err
+	}
+	if err := x.ClientStream.CloseSend(); err != nil {
+		return nil, err
+	}
+	return x, nil
+}
+
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type AgentService_AttestationClient = grpc.ServerStreamingClient[AttestationResponse]
+
+func (c *agentServiceClient) IMAMeasurements(ctx context.Context, in *IMAMeasurementsRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[IMAMeasurementsResponse], error) {
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
+	stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[4], AgentService_IMAMeasurements_FullMethodName, cOpts...)
+	if err != nil {
+		return nil, err
+	}
+	x := &grpc.GenericClientStream[IMAMeasurementsRequest, IMAMeasurementsResponse]{ClientStream: stream}
+	if err := x.ClientStream.SendMsg(in); err != nil {
+		return nil, err
+	}
+	if err := x.ClientStream.CloseSend(); err != nil {
+		return nil, err
+	}
+	return x, nil
+}
+
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type AgentService_IMAMeasurementsClient = grpc.ServerStreamingClient[IMAMeasurementsResponse]
+
+func (c *agentServiceClient) AzureAttestationToken(ctx context.Context, in *AttestationTokenRequest, opts ...grpc.CallOption) (*AttestationTokenResponse, error) {
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
+	out := new(AttestationTokenResponse)
+	err := c.cc.Invoke(ctx, AgentService_AzureAttestationToken_FullMethodName, in, out, cOpts...)
 	if err != nil {
 		return nil, err
 	}
@@ -138,32 +145,44 @@ func (c *agentServiceClient) Attestation(ctx context.Context, in *AttestationReq
 
 // AgentServiceServer is the server API for AgentService service.
 // All implementations must embed UnimplementedAgentServiceServer
-// for forward compatibility
+// for forward compatibility.
 type AgentServiceServer interface {
-	Algo(AgentService_AlgoServer) error
-	Data(AgentService_DataServer) error
-	Result(context.Context, *ResultRequest) (*ResultResponse, error)
-	Attestation(context.Context, *AttestationRequest) (*AttestationResponse, error)
+	Algo(grpc.ClientStreamingServer[AlgoRequest, AlgoResponse]) error
+	Data(grpc.ClientStreamingServer[DataRequest, DataResponse]) error
+	Result(*ResultRequest, grpc.ServerStreamingServer[ResultResponse]) error
+	Attestation(*AttestationRequest, grpc.ServerStreamingServer[AttestationResponse]) error
+	IMAMeasurements(*IMAMeasurementsRequest, grpc.ServerStreamingServer[IMAMeasurementsResponse]) error
+	AzureAttestationToken(context.Context, *AttestationTokenRequest) (*AttestationTokenResponse, error)
 	mustEmbedUnimplementedAgentServiceServer()
 }
 
-// UnimplementedAgentServiceServer must be embedded to have forward compatible implementations.
-type UnimplementedAgentServiceServer struct {
-}
+// UnimplementedAgentServiceServer must be embedded to have
+// forward compatible implementations.
+//
+// NOTE: this should be embedded by value instead of pointer to avoid a nil
+// pointer dereference when methods are called.
+type UnimplementedAgentServiceServer struct{}
 
-func (UnimplementedAgentServiceServer) Algo(AgentService_AlgoServer) error {
+func (UnimplementedAgentServiceServer) Algo(grpc.ClientStreamingServer[AlgoRequest, AlgoResponse]) error {
 	return status.Errorf(codes.Unimplemented, "method Algo not implemented")
 }
-func (UnimplementedAgentServiceServer) Data(AgentService_DataServer) error {
+func (UnimplementedAgentServiceServer) Data(grpc.ClientStreamingServer[DataRequest, DataResponse]) error {
 	return status.Errorf(codes.Unimplemented, "method Data not implemented")
 }
-func (UnimplementedAgentServiceServer) Result(context.Context, *ResultRequest) (*ResultResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method Result not implemented")
+func (UnimplementedAgentServiceServer) Result(*ResultRequest, grpc.ServerStreamingServer[ResultResponse]) error {
+	return status.Errorf(codes.Unimplemented, "method Result not implemented")
 }
-func (UnimplementedAgentServiceServer) Attestation(context.Context, *AttestationRequest) (*AttestationResponse, error) {
-	return nil, status.Errorf(codes.Unimplemented, "method Attestation not implemented")
+func (UnimplementedAgentServiceServer) Attestation(*AttestationRequest, grpc.ServerStreamingServer[AttestationResponse]) error {
+	return status.Errorf(codes.Unimplemented, "method Attestation not implemented")
+}
+func (UnimplementedAgentServiceServer) IMAMeasurements(*IMAMeasurementsRequest, grpc.ServerStreamingServer[IMAMeasurementsResponse]) error {
+	return status.Errorf(codes.Unimplemented, "method IMAMeasurements 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()                      {}
 
 // UnsafeAgentServiceServer may be embedded to opt out of forward compatibility for this service.
 // Use of this interface is not recommended, as added methods to AgentServiceServer will
@@ -173,93 +192,77 @@ type UnsafeAgentServiceServer interface {
 }
 
 func RegisterAgentServiceServer(s grpc.ServiceRegistrar, srv AgentServiceServer) {
+	// If the following call pancis, it indicates UnimplementedAgentServiceServer was
+	// embedded by pointer and is nil.  This will cause panics if an
+	// unimplemented method is ever invoked, so we test this at initialization
+	// time to prevent it from happening at runtime later due to I/O.
+	if t, ok := srv.(interface{ testEmbeddedByValue() }); ok {
+		t.testEmbeddedByValue()
+	}
 	s.RegisterService(&AgentService_ServiceDesc, srv)
 }
 
 func _AgentService_Algo_Handler(srv interface{}, stream grpc.ServerStream) error {
-	return srv.(AgentServiceServer).Algo(&agentServiceAlgoServer{ServerStream: stream})
+	return srv.(AgentServiceServer).Algo(&grpc.GenericServerStream[AlgoRequest, AlgoResponse]{ServerStream: stream})
 }
 
-type AgentService_AlgoServer interface {
-	SendAndClose(*AlgoResponse) error
-	Recv() (*AlgoRequest, error)
-	grpc.ServerStream
-}
-
-type agentServiceAlgoServer struct {
-	grpc.ServerStream
-}
-
-func (x *agentServiceAlgoServer) SendAndClose(m *AlgoResponse) error {
-	return x.ServerStream.SendMsg(m)
-}
-
-func (x *agentServiceAlgoServer) Recv() (*AlgoRequest, error) {
-	m := new(AlgoRequest)
-	if err := x.ServerStream.RecvMsg(m); err != nil {
-		return nil, err
-	}
-	return m, nil
-}
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type AgentService_AlgoServer = grpc.ClientStreamingServer[AlgoRequest, AlgoResponse]
 
 func _AgentService_Data_Handler(srv interface{}, stream grpc.ServerStream) error {
-	return srv.(AgentServiceServer).Data(&agentServiceDataServer{ServerStream: stream})
+	return srv.(AgentServiceServer).Data(&grpc.GenericServerStream[DataRequest, DataResponse]{ServerStream: stream})
 }
 
-type AgentService_DataServer interface {
-	SendAndClose(*DataResponse) error
-	Recv() (*DataRequest, error)
-	grpc.ServerStream
-}
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type AgentService_DataServer = grpc.ClientStreamingServer[DataRequest, DataResponse]
 
-type agentServiceDataServer struct {
-	grpc.ServerStream
-}
-
-func (x *agentServiceDataServer) SendAndClose(m *DataResponse) error {
-	return x.ServerStream.SendMsg(m)
-}
-
-func (x *agentServiceDataServer) Recv() (*DataRequest, error) {
-	m := new(DataRequest)
-	if err := x.ServerStream.RecvMsg(m); err != nil {
-		return nil, err
+func _AgentService_Result_Handler(srv interface{}, stream grpc.ServerStream) error {
+	m := new(ResultRequest)
+	if err := stream.RecvMsg(m); err != nil {
+		return err
 	}
-	return m, nil
+	return srv.(AgentServiceServer).Result(m, &grpc.GenericServerStream[ResultRequest, ResultResponse]{ServerStream: stream})
 }
 
-func _AgentService_Result_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(ResultRequest)
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type AgentService_ResultServer = grpc.ServerStreamingServer[ResultResponse]
+
+func _AgentService_Attestation_Handler(srv interface{}, stream grpc.ServerStream) error {
+	m := new(AttestationRequest)
+	if err := stream.RecvMsg(m); err != nil {
+		return err
+	}
+	return srv.(AgentServiceServer).Attestation(m, &grpc.GenericServerStream[AttestationRequest, AttestationResponse]{ServerStream: stream})
+}
+
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type AgentService_AttestationServer = grpc.ServerStreamingServer[AttestationResponse]
+
+func _AgentService_IMAMeasurements_Handler(srv interface{}, stream grpc.ServerStream) error {
+	m := new(IMAMeasurementsRequest)
+	if err := stream.RecvMsg(m); err != nil {
+		return err
+	}
+	return srv.(AgentServiceServer).IMAMeasurements(m, &grpc.GenericServerStream[IMAMeasurementsRequest, IMAMeasurementsResponse]{ServerStream: stream})
+}
+
+// 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_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).Result(ctx, in)
+		return srv.(AgentServiceServer).AzureAttestationToken(ctx, in)
 	}
 	info := &grpc.UnaryServerInfo{
 		Server:     srv,
-		FullMethod: AgentService_Result_FullMethodName,
+		FullMethod: AgentService_AzureAttestationToken_FullMethodName,
 	}
 	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(AgentServiceServer).Result(ctx, req.(*ResultRequest))
-	}
-	return interceptor(ctx, in, info, handler)
-}
-
-func _AgentService_Attestation_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
-	in := new(AttestationRequest)
-	if err := dec(in); err != nil {
-		return nil, err
-	}
-	if interceptor == nil {
-		return srv.(AgentServiceServer).Attestation(ctx, in)
-	}
-	info := &grpc.UnaryServerInfo{
-		Server:     srv,
-		FullMethod: AgentService_Attestation_FullMethodName,
-	}
-	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
-		return srv.(AgentServiceServer).Attestation(ctx, req.(*AttestationRequest))
+		return srv.(AgentServiceServer).AzureAttestationToken(ctx, req.(*AttestationTokenRequest))
 	}
 	return interceptor(ctx, in, info, handler)
 }
@@ -272,12 +275,8 @@ var AgentService_ServiceDesc = grpc.ServiceDesc{
 	HandlerType: (*AgentServiceServer)(nil),
 	Methods: []grpc.MethodDesc{
 		{
-			MethodName: "Result",
-			Handler:    _AgentService_Result_Handler,
-		},
-		{
-			MethodName: "Attestation",
-			Handler:    _AgentService_Attestation_Handler,
+			MethodName: "AzureAttestationToken",
+			Handler:    _AgentService_AzureAttestationToken_Handler,
 		},
 	},
 	Streams: []grpc.StreamDesc{
@@ -291,6 +290,21 @@ var AgentService_ServiceDesc = grpc.ServiceDesc{
 			Handler:       _AgentService_Data_Handler,
 			ClientStreams: true,
 		},
+		{
+			StreamName:    "Result",
+			Handler:       _AgentService_Result_Handler,
+			ServerStreams: true,
+		},
+		{
+			StreamName:    "Attestation",
+			Handler:       _AgentService_Attestation_Handler,
+			ServerStreams: true,
+		},
+		{
+			StreamName:    "IMAMeasurements",
+			Handler:       _AgentService_IMAMeasurements_Handler,
+			ServerStreams: true,
+		},
 	},
 	Metadata: "agent/agent.proto",
 }

agent/agentevent_string.go

@@ -0,0 +1,29 @@
+// Code generated by "stringer -type=AgentEvent"; DO NOT EDIT.
+
+package agent
+
+import "strconv"
+
+func _() {
+	// An "invalid array index" compiler error signifies that the constant values have changed.
+	// Re-run the stringer command to generate them again.
+	var x [1]struct{}
+	_ = x[Start-0]
+	_ = x[ManifestReceived-1]
+	_ = x[AlgorithmReceived-2]
+	_ = x[DataReceived-3]
+	_ = x[RunComplete-4]
+	_ = x[ResultsConsumed-5]
+	_ = x[RunFailed-6]
+}
+
+const _AgentEvent_name = "StartManifestReceivedAlgorithmReceivedDataReceivedRunCompleteResultsConsumedRunFailed"
+
+var _AgentEvent_index = [...]uint8{0, 5, 21, 38, 50, 61, 76, 85}
+
+func (i AgentEvent) String() string {
+	if i < 0 || i >= AgentEvent(len(_AgentEvent_index)-1) {
+		return "AgentEvent(" + strconv.FormatInt(int64(i), 10) + ")"
+	}
+	return _AgentEvent_name[_AgentEvent_index[i]:_AgentEvent_index[i+1]]
+}

agent/agentstate_string.go

@@ -0,0 +1,30 @@
+// Code generated by "stringer -type=AgentState"; DO NOT EDIT.
+
+package agent
+
+import "strconv"
+
+func _() {
+	// An "invalid array index" compiler error signifies that the constant values have changed.
+	// Re-run the stringer command to generate them again.
+	var x [1]struct{}
+	_ = x[Idle-0]
+	_ = x[ReceivingManifest-1]
+	_ = x[ReceivingAlgorithm-2]
+	_ = x[ReceivingData-3]
+	_ = x[Running-4]
+	_ = x[ConsumingResults-5]
+	_ = x[Complete-6]
+	_ = x[Failed-7]
+}
+
+const _AgentState_name = "IdleReceivingManifestReceivingAlgorithmReceivingDataRunningConsumingResultsCompleteFailed"
+
+var _AgentState_index = [...]uint8{0, 4, 21, 39, 52, 59, 75, 83, 89}
+
+func (i AgentState) String() string {
+	if i < 0 || i >= AgentState(len(_AgentState_index)-1) {
+		return "AgentState(" + strconv.FormatInt(int64(i), 10) + ")"
+	}
+	return _AgentState_name[_AgentState_index[i]:_AgentState_index[i+1]]
+}

agent/algorithm/algorithm.go

@@ -2,8 +2,51 @@
 // SPDX-License-Identifier: Apache-2.0
 package algorithm
 
+import (
+	"context"
+
+	"google.golang.org/grpc/metadata"
+)
+
+type AlgorithType string
+
+const (
+	AlgoTypeBin    AlgorithType = "bin"
+	AlgoTypePython AlgorithType = "python"
+	AlgoTypeWasm   AlgorithType = "wasm"
+	AlgoTypeDocker AlgorithType = "docker"
+	AlgoTypeKey                 = "algo_type"
+	AlgoArgsKey                 = "algo_args"
+
+	ResultsDir     = "results"
+	DatasetsDir    = "datasets"
+	AlgoWorkingDir = "/cocos"
+)
+
+func AlgorithmTypeToContext(ctx context.Context, algoType string) context.Context {
+	return metadata.AppendToOutgoingContext(ctx, AlgoTypeKey, algoType)
+}
+
+func AlgorithmTypeFromContext(ctx context.Context) string {
+	return metadata.ValueFromIncomingContext(ctx, AlgoTypeKey)[0]
+}
+
+func AlgorithmArgsToContext(ctx context.Context, algoArgs []string) context.Context {
+	for _, arg := range algoArgs {
+		ctx = metadata.AppendToOutgoingContext(ctx, AlgoArgsKey, arg)
+	}
+	return ctx
+}
+
+func AlgorithmArgsFromContext(ctx context.Context) []string {
+	return metadata.ValueFromIncomingContext(ctx, AlgoArgsKey)
+}
+
 // Algorithm is an interface that specifies the API for an algorithm.
 type Algorithm interface {
 	// Run executes the algorithm and returns the result.
-	Run() ([]byte, error)
+	Run() error
+
+	// Stop stops the algorithm.
+	Stop() error
 }

agent/algorithm/binary/binary.go

@@ -6,74 +6,64 @@ import (
 	"fmt"
 	"io"
 	"log/slog"
-	"os"
 	"os/exec"
 
 	"github.com/ultravioletrs/cocos/agent/algorithm"
+	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
 	"github.com/ultravioletrs/cocos/agent/events"
-	"github.com/ultravioletrs/cocos/pkg/socket"
 )
 
-const socketPath = "unix_socket"
-
 var _ algorithm.Algorithm = (*binary)(nil)
 
 type binary struct {
 	algoFile string
-	datasets []string
-	logger   *slog.Logger
 	stderr   io.Writer
 	stdout   io.Writer
+	args     []string
+	cmd      *exec.Cmd
 }
 
-func New(logger *slog.Logger, eventsSvc events.Service, algoFile string, datasets ...string) algorithm.Algorithm {
+func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string, args []string, cmpID string) algorithm.Algorithm {
 	return &binary{
 		algoFile: algoFile,
-		datasets: datasets,
-		logger:   logger,
-		stderr:   &algorithm.Stderr{Logger: logger, EventSvc: eventsSvc},
-		stdout:   &algorithm.Stdout{Logger: logger},
+		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
+		stdout:   &logging.Stdout{Logger: logger},
+		args:     args,
 	}
 }
 
-func (b *binary) Run() ([]byte, error) {
-	defer os.Remove(b.algoFile)
-	defer func() {
-		for _, file := range b.datasets {
-			os.Remove(file)
-		}
-	}()
-	listener, err := socket.StartUnixSocketServer(socketPath)
-	if err != nil {
-		return nil, fmt.Errorf("error creating stdout pipe: %v", err)
-	}
-	defer listener.Close()
+func (b *binary) Run() error {
+	b.cmd = exec.Command(b.algoFile, b.args...)
+	b.cmd.Stderr = b.stderr
+	b.cmd.Stdout = b.stdout
 
-	// Create channels for received data and errors
-	dataChannel := make(chan []byte)
-	errorChannel := make(chan error)
-
-	var result []byte
-
-	go socket.AcceptConnection(listener, dataChannel, errorChannel)
-
-	args := append([]string{socketPath}, b.datasets...)
-	cmd := exec.Command(b.algoFile, args...)
-	cmd.Stderr = b.stderr
-	cmd.Stdout = b.stdout
-
-	if err := cmd.Start(); err != nil {
-		return nil, fmt.Errorf("error starting algorithm: %v", err)
+	if err := b.cmd.Start(); err != nil {
+		return fmt.Errorf("error starting algorithm: %v", err)
 	}
 
-	if err := cmd.Wait(); err != nil {
-		return nil, fmt.Errorf("algorithm execution error: %v", err)
+	if err := b.cmd.Wait(); err != nil {
+		return fmt.Errorf("algorithm execution error: %v", err)
 	}
 
-	select {
-	case result = <-dataChannel:
-		return result, nil
-	case err = <-errorChannel:
-		return nil, fmt.Errorf("error receiving data: %v", err)
-	}
+	return nil
+}
+
+func (b *binary) Stop() error {
+	if b.cmd == nil {
+		return nil
+	}
+
+	if b.cmd.ProcessState != nil && b.cmd.ProcessState.Exited() {
+		return nil
+	}
+
+	if b.cmd.Process == nil {
+		return nil
+	}
+
+	if err := b.cmd.Process.Kill(); err != nil {
+		return fmt.Errorf("error stopping algorithm: %v", err)
+	}
+
+	return nil
 }

agent/algorithm/binary/binary_test.go

@@ -0,0 +1,100 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package binary
+
+import (
+	"bytes"
+	"log/slog"
+	"os"
+	"testing"
+
+	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
+	"github.com/ultravioletrs/cocos/agent/events/mocks"
+)
+
+func TestNewAlgorithm(t *testing.T) {
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+	eventsSvc := new(mocks.Service)
+	algoFile := "/path/to/algo"
+	args := []string{"arg1", "arg2"}
+
+	algo := NewAlgorithm(logger, eventsSvc, algoFile, args, "")
+
+	b, ok := algo.(*binary)
+	if !ok {
+		t.Fatalf("NewAlgorithm did not return a *binary")
+	}
+
+	if b.algoFile != algoFile {
+		t.Errorf("Expected algoFile to be %s, got %s", algoFile, b.algoFile)
+	}
+
+	if len(b.args) != len(args) {
+		t.Errorf("Expected %d args, got %d", len(args), len(b.args))
+	}
+
+	for i, arg := range args {
+		if b.args[i] != arg {
+			t.Errorf("Expected arg %d to be %s, got %s", i, arg, b.args[i])
+		}
+	}
+
+	if _, ok := b.stderr.(*logging.Stderr); !ok {
+		t.Errorf("Expected stderr to be *algorithm.Stderr")
+	}
+
+	if _, ok := b.stdout.(*logging.Stdout); !ok {
+		t.Errorf("Expected stdout to be *algorithm.Stdout")
+	}
+}
+
+func TestBinaryRun(t *testing.T) {
+	tests := []struct {
+		name          string
+		algoFile      string
+		args          []string
+		expectedError bool
+	}{
+		{
+			name:          "Successful execution",
+			algoFile:      "echo",
+			args:          []string{"Hello, World!"},
+			expectedError: false,
+		},
+		{
+			name:          "Non-existent binary",
+			algoFile:      "non_existent_binary",
+			args:          []string{},
+			expectedError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+			eventsSvc := new(mocks.Service)
+
+			b := NewAlgorithm(logger, eventsSvc, tt.algoFile, tt.args, "").(*binary)
+
+			var stdout, stderr bytes.Buffer
+			b.stdout = &stdout
+			b.stderr = &stderr
+
+			err := b.Run()
+
+			if tt.expectedError && err == nil {
+				t.Errorf("Expected an error, but got none")
+			}
+
+			if !tt.expectedError && err != nil {
+				t.Errorf("Unexpected error: %v", err)
+			}
+
+			if !tt.expectedError {
+				if stdout.Len() == 0 {
+					t.Errorf("Expected non-empty stdout")
+				}
+			}
+		})
+	}
+}

agent/algorithm/docker/docker.go

@@ -0,0 +1,182 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package docker
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+	"os"
+	"path"
+
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/api/types/image"
+	"github.com/docker/docker/api/types/mount"
+	"github.com/docker/docker/client"
+	"github.com/ultravioletrs/cocos/agent/algorithm"
+	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
+	"github.com/ultravioletrs/cocos/agent/events"
+)
+
+const (
+	containerName     = "agent_container"
+	datasetsMountPath = "/cocos/datasets"
+	resultsMountPath  = "/cocos/results"
+)
+
+var _ algorithm.Algorithm = (*docker)(nil)
+
+type docker struct {
+	algoFile string
+	logger   *slog.Logger
+	stderr   io.Writer
+	stdout   io.Writer
+}
+
+func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile, cmpID string) algorithm.Algorithm {
+	d := &docker{
+		algoFile: algoFile,
+		logger:   logger,
+		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
+		stdout:   &logging.Stdout{Logger: logger},
+	}
+
+	return d
+}
+
+func (d *docker) Run() error {
+	// Create a new Docker client.
+	cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
+	if err != nil {
+		return fmt.Errorf("could not create a new Docker client: %v", err)
+	}
+
+	// Open the Docker image tar file.
+	imageFile, err := os.Open(d.algoFile)
+	if err != nil {
+		return fmt.Errorf("could not open Docker image: %v", err)
+	}
+	defer imageFile.Close()
+
+	ctx := context.Background()
+	// Load the Docker image from the tar file.
+	resp, err := cli.ImageLoad(ctx, imageFile, client.ImageLoadWithQuiet(true))
+	if err != nil {
+		return fmt.Errorf("could not load Docker image from file: %v", err)
+	}
+	defer resp.Body.Close()
+
+	// List the loaded images to get the image ID.
+	images, err := cli.ImageList(ctx, image.ListOptions{})
+	if err != nil {
+		return fmt.Errorf("could not get the Docker image list: %v", err)
+	}
+
+	var imageID string = ""
+	var dockerImageName string = ""
+	for _, image := range images {
+		for _, tag := range image.RepoTags {
+			imageID = image.ID
+			dockerImageName = tag
+			break
+		}
+	}
+
+	if imageID == "" {
+		return fmt.Errorf("could not find image ID")
+	}
+
+	// Create and start the container.
+	respContainer, err := cli.ContainerCreate(ctx, &container.Config{
+		Image:        dockerImageName,
+		Tty:          true,
+		AttachStdout: true,
+		AttachStderr: true,
+	}, &container.HostConfig{
+		Mounts: []mount.Mount{
+			{
+				Type:   mount.TypeBind,
+				Source: path.Join(algorithm.AlgoWorkingDir, algorithm.DatasetsDir),
+				Target: datasetsMountPath,
+			},
+			{
+				Type:   mount.TypeBind,
+				Source: path.Join(algorithm.AlgoWorkingDir, algorithm.ResultsDir),
+				Target: resultsMountPath,
+			},
+		},
+	}, nil, nil, containerName)
+	if err != nil {
+		return fmt.Errorf("could not create a Docker container: %v", err)
+	}
+
+	if err := cli.ContainerStart(ctx, respContainer.ID, container.StartOptions{}); err != nil {
+		return fmt.Errorf("could not start a Docker container: %v", err)
+	}
+
+	stdout, err := cli.ContainerLogs(ctx, respContainer.ID, container.LogsOptions{ShowStdout: true, Follow: true})
+	if err != nil {
+		return fmt.Errorf("could not read stdout from the container: %v", err)
+	}
+	defer stdout.Close()
+
+	go func() {
+		if err := writeToOut(stdout, d.stdout); err != nil {
+			d.logger.Warn(fmt.Sprintf("could not write to stdout: %v", err))
+		}
+	}()
+
+	stderr, err := cli.ContainerLogs(ctx, respContainer.ID, container.LogsOptions{ShowStderr: true, Follow: true})
+	if err != nil {
+		d.logger.Warn(fmt.Sprintf("could not read stderr from the container: %v", err))
+	}
+	defer stderr.Close()
+
+	go func() {
+		if err := writeToOut(stderr, d.stderr); err != nil {
+			d.logger.Warn(fmt.Sprintf("could not write to stderr: %v", err))
+		}
+	}()
+
+	statusCh, errCh := cli.ContainerWait(ctx, respContainer.ID, container.WaitConditionNotRunning)
+	select {
+	case err := <-errCh:
+		if err != nil {
+			return fmt.Errorf("could not wait for a Docker container: %v", err)
+		}
+	case <-statusCh:
+	}
+
+	defer func() {
+		if err = cli.ContainerRemove(ctx, respContainer.ID, container.RemoveOptions{Force: true}); err != nil {
+			d.logger.Warn(fmt.Sprintf("error could not remove container: %v", err))
+		}
+
+		if _, err := cli.ImageRemove(ctx, imageID, image.RemoveOptions{Force: true}); err != nil {
+			d.logger.Warn(fmt.Sprintf("error could not remove image: %v", err))
+		}
+	}()
+
+	return nil
+}
+
+func writeToOut(readCloser io.ReadCloser, ioWriter io.Writer) error {
+	scanner := bufio.NewScanner(readCloser)
+	for scanner.Scan() {
+		if _, err := ioWriter.Write(scanner.Bytes()); err != nil {
+			return fmt.Errorf("error writing to output: %v", err)
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		return fmt.Errorf("error reading container logs error: %v", err)
+	}
+
+	return nil
+}
+
+func (d *docker) Stop() error {
+	// To be supported later.
+	return nil
+}

agent/algorithm/docker/docker_test.go

@@ -0,0 +1,29 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package docker
+
+import (
+	"log/slog"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
+	"github.com/ultravioletrs/cocos/agent/events/mocks"
+)
+
+// TestNewAlgorithm tests the NewAlgorithm function.
+func TestNewAlgorithm(t *testing.T) {
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+	eventsSvc := new(mocks.Service)
+	algoFile := "/path/to/algo.tar"
+
+	algo := NewAlgorithm(logger, eventsSvc, algoFile, "")
+
+	d, ok := algo.(*docker)
+	assert.True(t, ok, "NewAlgorithm should return a *docker")
+	assert.Equal(t, algoFile, d.algoFile, "algoFile should be set correctly")
+	assert.NotNil(t, d.logger, "logger should be set")
+	assert.IsType(t, &logging.Stderr{}, d.stderr, "stderr should be of type *algorithm.Stderr")
+	assert.IsType(t, &logging.Stdout{}, d.stdout, "stdout should be of type *algorithm.Stdout")
+}

agent/algorithm/logging/logging.go

@@ -1,9 +1,10 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
-package algorithm
+package logging
 
 import (
 	"bytes"
+	"encoding/json"
 	"io"
 	"log/slog"
 
@@ -15,7 +16,11 @@ var (
 	_ io.Writer = &Stderr{}
 )
 
-const bufSize = 1024
+const (
+	bufSize       = 1024
+	algorithmRun  = "AlgorithmRun"
+	warningStatus = "Warning"
+)
 
 type Stdout struct {
 	Logger *slog.Logger
@@ -45,6 +50,7 @@ func (s *Stdout) Write(p []byte) (n int, err error) {
 type Stderr struct {
 	Logger   *slog.Logger
 	EventSvc events.Service
+	CmpID    string
 }
 
 // Write implements io.Writer.
@@ -65,9 +71,7 @@ func (s *Stderr) Write(p []byte) (n int, err error) {
 		s.Logger.Error(string(buf[:n]))
 	}
 
-	if err := s.EventSvc.SendEvent("algorithm-run", "failed", nil); err != nil {
-		return len(p), err
-	}
+	s.EventSvc.SendEvent(s.CmpID, algorithmRun, warningStatus, json.RawMessage{})
 
 	return len(p), nil
 }

agent/algorithm/logging/logging_test.go

@@ -1,15 +1,16 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
-package algorithm
+package logging
 
 import (
 	"strings"
 	"testing"
 
-	mglog "github.com/absmach/magistrala/logger"
+	mglog "github.com/absmach/supermq/logger"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/ultravioletrs/cocos/agent/events/mocks"
+	"github.com/ultravioletrs/cocos/pkg/manager"
 )
 
 func TestStdoutWrite(t *testing.T) {
@@ -72,7 +73,7 @@ func TestStderrWrite(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			mockEventService := mocks.NewService(t)
-			mockEventService.On("SendEvent", "algorithm-run", "failed", mock.Anything).Return(nil)
+			mockEventService.On("SendEvent", mock.Anything, "AlgorithmRun", manager.Warning.String(), mock.Anything).Return(nil)
 
 			stderr := &Stderr{Logger: mglog.NewMock(), EventSvc: mockEventService}
 			n, err := stderr.Write([]byte(tt.input))

agent/algorithm/mocks/algorithm.go

@@ -0,0 +1,127 @@
+// 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 (
+	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 {
+	mock.Mock
+}
+
+type Algorithm_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *Algorithm) EXPECT() *Algorithm_Expecter {
+	return &Algorithm_Expecter{mock: &_m.Mock}
+}
+
+// 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 returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// Algorithm_Run_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Run'
+type Algorithm_Run_Call struct {
+	*mock.Call
+}
+
+// Run is a helper method to define mock.On call
+func (_e *Algorithm_Expecter) Run() *Algorithm_Run_Call {
+	return &Algorithm_Run_Call{Call: _e.mock.On("Run")}
+}
+
+func (_c *Algorithm_Run_Call) Run(run func()) *Algorithm_Run_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *Algorithm_Run_Call) Return(err error) *Algorithm_Run_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *Algorithm_Run_Call) RunAndReturn(run func() error) *Algorithm_Run_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// 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 returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// Algorithm_Stop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Stop'
+type Algorithm_Stop_Call struct {
+	*mock.Call
+}
+
+// Stop is a helper method to define mock.On call
+func (_e *Algorithm_Expecter) Stop() *Algorithm_Stop_Call {
+	return &Algorithm_Stop_Call{Call: _e.mock.On("Stop")}
+}
+
+func (_c *Algorithm_Stop_Call) Run(run func()) *Algorithm_Stop_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *Algorithm_Stop_Call) Return(err error) *Algorithm_Stop_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *Algorithm_Stop_Call) RunAndReturn(run func() error) *Algorithm_Stop_Call {
+	_c.Call.Return(run)
+	return _c
+}

agent/algorithm/python/python.go

@@ -0,0 +1,126 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package python
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+	"os"
+	"os/exec"
+	"path/filepath"
+
+	"github.com/ultravioletrs/cocos/agent/algorithm"
+	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
+	"github.com/ultravioletrs/cocos/agent/events"
+	"google.golang.org/grpc/metadata"
+)
+
+const (
+	PyRuntime    = "python3"
+	PyRuntimeKey = "python_runtime"
+)
+
+func PythonRunTimeToContext(ctx context.Context, runtime string) context.Context {
+	return metadata.AppendToOutgoingContext(ctx, PyRuntimeKey, runtime)
+}
+
+func PythonRunTimeFromContext(ctx context.Context) string {
+	return metadata.ValueFromIncomingContext(ctx, PyRuntimeKey)[0]
+}
+
+var _ algorithm.Algorithm = (*python)(nil)
+
+type python struct {
+	algoFile         string
+	stderr           io.Writer
+	stdout           io.Writer
+	runtime          string
+	requirementsFile string
+	args             []string
+	cmd              *exec.Cmd
+}
+
+func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, runtime, requirementsFile, algoFile string, args []string, cmpID string) algorithm.Algorithm {
+	p := &python{
+		algoFile:         algoFile,
+		stderr:           &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
+		stdout:           &logging.Stdout{Logger: logger},
+		requirementsFile: requirementsFile,
+		args:             args,
+	}
+	if runtime != "" {
+		p.runtime = runtime
+	} else {
+		p.runtime = PyRuntime
+	}
+	return p
+}
+
+func (p *python) Run() error {
+	venvPath := "venv"
+	createVenvCmd := exec.Command(p.runtime, "-m", "venv", venvPath)
+	createVenvCmd.Stderr = p.stderr
+	createVenvCmd.Stdout = p.stdout
+	if err := createVenvCmd.Run(); err != nil {
+		return fmt.Errorf("error creating virtual environment: %v", err)
+	}
+
+	pythonPath := filepath.Join(venvPath, "bin", "python")
+
+	updatePipCmd := exec.Command(pythonPath, "-m", "pip", "install", "--upgrade", "pip")
+	updatePipCmd.Stderr = p.stderr
+	updatePipCmd.Stdout = p.stdout
+	if err := updatePipCmd.Run(); err != nil {
+		return fmt.Errorf("error updating pip: %v", err)
+	}
+
+	if p.requirementsFile != "" {
+		rcmd := exec.Command(pythonPath, "-m", "pip", "install", "-r", p.requirementsFile)
+		rcmd.Stderr = p.stderr
+		rcmd.Stdout = p.stdout
+		if err := rcmd.Run(); err != nil {
+			return fmt.Errorf("error installing requirements: %v", err)
+		}
+	}
+
+	args := append([]string{p.algoFile}, p.args...)
+	p.cmd = exec.Command(pythonPath, args...)
+	p.cmd.Stderr = p.stderr
+	p.cmd.Stdout = p.stdout
+
+	if err := p.cmd.Start(); err != nil {
+		return fmt.Errorf("error starting algorithm: %v", err)
+	}
+
+	if err := p.cmd.Wait(); err != nil {
+		return fmt.Errorf("algorithm execution error: %v", err)
+	}
+
+	if err := os.RemoveAll(venvPath); err != nil {
+		return fmt.Errorf("error removing virtual environment: %v", err)
+	}
+
+	return nil
+}
+
+func (p *python) Stop() error {
+	if p.cmd == nil {
+		return nil
+	}
+
+	if p.cmd.ProcessState != nil && p.cmd.ProcessState.Exited() {
+		return nil
+	}
+
+	if p.cmd.Process == nil {
+		return nil
+	}
+
+	if err := p.cmd.Process.Kill(); err != nil {
+		return fmt.Errorf("error stopping algorithm: %v", err)
+	}
+
+	return nil
+}

agent/algorithm/python/python_test.go

@@ -0,0 +1,148 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package python
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"log/slog"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
+	"github.com/ultravioletrs/cocos/agent/events/mocks"
+	"google.golang.org/grpc/metadata"
+)
+
+const runtime = "python3"
+
+func TestPythonRunTimeToContext(t *testing.T) {
+	ctx := context.Background()
+	newCtx := PythonRunTimeToContext(ctx, runtime)
+
+	md, ok := metadata.FromOutgoingContext(newCtx)
+	if !ok {
+		t.Fatal("Expected metadata in context")
+	}
+
+	values := md.Get(PyRuntimeKey)
+	if len(values) != 1 || values[0] != runtime {
+		t.Errorf("Expected runtime %s, got %v", runtime, values)
+	}
+}
+
+func TestPythonRunTimeFromContext(t *testing.T) {
+	ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs(PyRuntimeKey, runtime))
+
+	got := PythonRunTimeFromContext(ctx)
+	if got != runtime {
+		t.Errorf("Expected runtime %s, got %s", runtime, got)
+	}
+}
+
+func TestNewAlgorithm(t *testing.T) {
+	logger := &slog.Logger{}
+	eventsSvc := new(mocks.Service)
+	requirementsFile := "requirements.txt"
+	algoFile := "algorithm.py"
+	args := []string{"--arg1", "value1"}
+
+	algo := NewAlgorithm(logger, eventsSvc, runtime, requirementsFile, algoFile, args, "")
+
+	p, ok := algo.(*python)
+	if !ok {
+		t.Fatal("Expected *python type")
+	}
+
+	if p.runtime != runtime {
+		t.Errorf("Expected runtime %s, got %s", runtime, p.runtime)
+	}
+	if p.requirementsFile != requirementsFile {
+		t.Errorf("Expected requirementsFile %s, got %s", requirementsFile, p.requirementsFile)
+	}
+	if p.algoFile != algoFile {
+		t.Errorf("Expected algoFile %s, got %s", algoFile, p.algoFile)
+	}
+	if len(p.args) != len(args) {
+		t.Errorf("Expected %d args, got %d", len(args), len(p.args))
+	}
+}
+
+func TestRun(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "python-test")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	scriptContent := []byte("print('Hello, World!')")
+	scriptPath := filepath.Join(tmpDir, "test_script.py")
+	if err := os.WriteFile(scriptPath, scriptContent, 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	eventsSvc := new(mocks.Service)
+
+	var stdout, stderr bytes.Buffer
+
+	algo := &python{
+		algoFile: scriptPath,
+		stderr:   io.MultiWriter(&stderr, &logging.Stderr{Logger: slog.Default(), EventSvc: eventsSvc}),
+		stdout:   io.MultiWriter(&stdout, &logging.Stdout{Logger: slog.Default()}),
+		runtime:  "python3",
+	}
+
+	err = algo.Run()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	expectedOutput := "Hello, World!\n"
+	if !strings.Contains(stdout.String(), expectedOutput) {
+		t.Errorf("Expected output to contain %q, got %q", expectedOutput, stdout.String())
+	}
+}
+
+func TestRunWithRequirements(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "python-test")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	scriptContent := []byte("import requests\nprint(requests.__version__)")
+	scriptPath := filepath.Join(tmpDir, "test_script.py")
+	if err := os.WriteFile(scriptPath, scriptContent, 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	requirementsContent := []byte("requests==2.26.0")
+	requirementsPath := filepath.Join(tmpDir, "requirements.txt")
+	if err := os.WriteFile(requirementsPath, requirementsContent, 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	eventsSvc := new(mocks.Service)
+
+	var stdout, stderr bytes.Buffer
+
+	algo := &python{
+		algoFile:         scriptPath,
+		requirementsFile: requirementsPath,
+		stderr:           io.MultiWriter(&stderr, &logging.Stderr{Logger: slog.Default(), EventSvc: eventsSvc}),
+		stdout:           io.MultiWriter(&stdout, &logging.Stdout{Logger: slog.Default()}),
+		runtime:          "python3",
+	}
+
+	err = algo.Run()
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	if !strings.Contains(stdout.String(), "2.26.0") {
+		t.Errorf("Expected output to contain requests version 2.26.0, got %q", stdout.String())
+	}
+}

agent/algorithm/results_test.go

@@ -0,0 +1,82 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package algorithm_test
+
+import (
+	"os"
+	"testing"
+
+	"github.com/ultravioletrs/cocos/agent/algorithm"
+	"github.com/ultravioletrs/cocos/internal"
+)
+
+func TestZipDirectory(t *testing.T) {
+	cases := []struct {
+		name        string
+		directories []string
+		files       []string
+		expected    []string
+	}{
+		{
+			name:        "empty directory",
+			directories: []string{"testdata"},
+		},
+		{
+			name:  "single file",
+			files: []string{"file1.txt"},
+		},
+		{
+			name:        "directory with single file",
+			directories: []string{"testdata"},
+			expected:    []string{"testdata/file1.txt"},
+		},
+		{
+			name:        "directory with multiple files",
+			directories: []string{"testdata"},
+			expected: []string{
+				"testdata/file1.txt",
+				"testdata/file2.txt",
+				"testdata/file3.txt",
+			},
+		},
+		{
+			name:        "nested directories",
+			directories: []string{"testdata", "testdata/nested"},
+			expected: []string{
+				"testdata/nested/file1.txt",
+				"testdata/nested/file2.txt",
+				"testdata/nested/file3.txt",
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			if err := os.Mkdir(algorithm.ResultsDir, 0o755); err != nil {
+				t.Fatalf("error creating results directory: %s", err.Error())
+			}
+			defer func() {
+				if err := os.RemoveAll(algorithm.ResultsDir); err != nil {
+					t.Fatalf("error removing results directory and its contents: %s", err.Error())
+				}
+			}()
+
+			for _, dir := range tc.directories {
+				if dir != "" {
+					if err := os.Mkdir(algorithm.ResultsDir+"/"+dir, 0o755); err != nil {
+						t.Fatalf("error creating test directory: %s", err.Error())
+					}
+				}
+			}
+			for _, file := range tc.files {
+				if _, err := os.Create(algorithm.ResultsDir + "/" + file); err != nil {
+					t.Fatalf("error creating test file: %s", err.Error())
+				}
+			}
+
+			if _, err := internal.ZipDirectoryToMemory(algorithm.ResultsDir); err != nil {
+				t.Errorf("ZipDirectory() error = %v", err)
+			}
+		})
+	}
+}

agent/algorithm/wasm/wasm.go

@@ -0,0 +1,75 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package wasm
+
+import (
+	"fmt"
+	"io"
+	"log/slog"
+	"os/exec"
+
+	"github.com/ultravioletrs/cocos/agent/algorithm"
+	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
+	"github.com/ultravioletrs/cocos/agent/events"
+)
+
+const wasmRuntime = "wasmedge"
+
+var mapDirOption = []string{"--dir", ".:" + algorithm.ResultsDir}
+
+var _ algorithm.Algorithm = (*wasm)(nil)
+
+type wasm struct {
+	algoFile string
+	stderr   io.Writer
+	stdout   io.Writer
+	args     []string
+	cmd      *exec.Cmd
+}
+
+func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, args []string, algoFile, cmpID string) algorithm.Algorithm {
+	return &wasm{
+		algoFile: algoFile,
+		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc, CmpID: cmpID},
+		stdout:   &logging.Stdout{Logger: logger},
+		args:     args,
+	}
+}
+
+func (w *wasm) Run() error {
+	args := append(mapDirOption, w.algoFile)
+	args = append(args, w.args...)
+	w.cmd = exec.Command(wasmRuntime, args...)
+	w.cmd.Stderr = w.stderr
+	w.cmd.Stdout = w.stdout
+
+	if err := w.cmd.Start(); err != nil {
+		return fmt.Errorf("error starting algorithm: %v", err)
+	}
+
+	if err := w.cmd.Wait(); err != nil {
+		return fmt.Errorf("algorithm execution error: %v", err)
+	}
+
+	return nil
+}
+
+func (w *wasm) Stop() error {
+	if w.cmd == nil {
+		return nil
+	}
+
+	if w.cmd.ProcessState != nil && w.cmd.ProcessState.Exited() {
+		return nil
+	}
+
+	if w.cmd.Process == nil {
+		return nil
+	}
+
+	if err := w.cmd.Process.Kill(); err != nil {
+		return fmt.Errorf("error stopping algorithm: %v", err)
+	}
+
+	return nil
+}

agent/algorithm/wasm/wasm_test.go

@@ -0,0 +1,89 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package wasm
+
+import (
+	"log/slog"
+	"os"
+	"os/exec"
+	"testing"
+
+	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
+	"github.com/ultravioletrs/cocos/agent/events/mocks"
+)
+
+func TestNewAlgorithm(t *testing.T) {
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+	eventsSvc := new(mocks.Service)
+	algoFile := "test.wasm"
+	args := []string{"arg1", "arg2"}
+
+	algo := NewAlgorithm(logger, eventsSvc, args, algoFile, "")
+
+	w, ok := algo.(*wasm)
+	if !ok {
+		t.Fatalf("NewAlgorithm did not return a *wasm")
+	}
+
+	if w.algoFile != algoFile {
+		t.Errorf("Expected algoFile to be %s, got %s", algoFile, w.algoFile)
+	}
+
+	if len(w.args) != len(args) {
+		t.Errorf("Expected %d args, got %d", len(args), len(w.args))
+	}
+
+	_, ok = w.stderr.(*logging.Stderr)
+	if !ok {
+		t.Errorf("Expected stderr to be *algorithm.Stderr")
+	}
+
+	_, ok = w.stdout.(*logging.Stdout)
+	if !ok {
+		t.Errorf("Expected stdout to be *algorithm.Stdout")
+	}
+}
+
+func TestRunError(t *testing.T) {
+	// Mock exec.Command to return an error
+	execCommand = mockExecCommandError
+	defer func() { execCommand = exec.Command }()
+
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+	eventsSvc := new(mocks.Service)
+	algoFile := "test.wasm"
+	args := []string{"arg1", "arg2"}
+
+	w := NewAlgorithm(logger, eventsSvc, args, algoFile, "").(*wasm)
+
+	err := w.Run()
+	if err == nil {
+		t.Errorf("Run() should have returned an error")
+	}
+}
+
+func mockExecCommand(command string, args ...string) *exec.Cmd {
+	cs := []string{"-test.run=TestHelperProcess", "--", command}
+	cs = append(cs, args...)
+	cmd := exec.Command(os.Args[0], cs...)
+	cmd.Env = []string{"GO_WANT_HELPER_PROCESS=1"}
+	return cmd
+}
+
+func mockExecCommandError(command string, args ...string) *exec.Cmd {
+	cmd := mockExecCommand(command, args...)
+	cmd.Env = append(cmd.Env, "GO_WANT_HELPER_PROCESS_ERROR=1")
+	return cmd
+}
+
+func TestHelperProcess(t *testing.T) {
+	if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
+		return
+	}
+	if os.Getenv("GO_WANT_HELPER_PROCESS_ERROR") == "1" {
+		os.Exit(1)
+	}
+	os.Exit(0)
+}
+
+var execCommand = exec.Command

agent/api/grpc/endpoint.go

@@ -7,17 +7,18 @@ import (
 
 	"github.com/go-kit/kit/endpoint"
 	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
 )
 
 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 {
 			return algoRes{}, err
 		}
 
-		algo := agent.Algorithm{Algorithm: req.Algorithm}
+		algo := agent.Algorithm{Algorithm: req.Algorithm, Requirements: req.Requirements}
 
 		err := svc.Algo(ctx, algo)
 		if err != nil {
@@ -29,14 +30,14 @@ 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 {
 			return dataRes{}, err
 		}
 
-		dataset := agent.Dataset{Dataset: req.Dataset}
+		dataset := agent.Dataset{Dataset: req.Dataset, Filename: req.Filename}
 
 		err := svc.Data(ctx, dataset)
 		if err != nil {
@@ -48,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 {
@@ -64,13 +65,13 @@ 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 {
 			return attestationRes{}, err
 		}
-		file, err := svc.Attestation(ctx, req.ReportData)
+		file, err := svc.Attestation(ctx, req.TeeNonce, req.VtpmNonce, attestation.PlatformType(req.AttType))
 		if err != nil {
 			return attestationRes{}, err
 		}
@@ -78,3 +79,33 @@ func attestationEndpoint(svc agent.Service) endpoint.Endpoint {
 		return attestationRes{File: file}, nil
 	}
 }
+
+func imaMeasurementsEndpoint(svc agent.Service) endpoint.Endpoint {
+	return func(ctx context.Context, request any) (any, error) {
+		req := request.(imaMeasurementsReq)
+
+		if err := req.validate(); err != nil {
+			return imaMeasurementsRes{}, err
+		}
+		file, pcr10, err := svc.IMAMeasurements(ctx)
+		if err != nil {
+			return imaMeasurementsRes{}, err
+		}
+
+		return imaMeasurementsRes{File: file, PCR10: pcr10}, nil
+	}
+}
+
+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 fetchAttestationTokenRes{}, err
+		}
+		file, err := svc.AzureAttestationToken(ctx, req.tokenNonce)
+		if err != nil {
+			return fetchAttestationTokenRes{}, err
+		}
+		return fetchAttestationTokenRes{File: file}, nil
+	}
+}

agent/api/grpc/endpoint_test.go

@@ -0,0 +1,227 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package grpc
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/agent/mocks"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"golang.org/x/crypto/sha3"
+)
+
+const svcErr = "Service Error"
+
+func TestAlgoEndpoint(t *testing.T) {
+	svc := new(mocks.Service)
+	tests := []struct {
+		name        string
+		req         algoReq
+		expectedErr bool
+	}{
+		{
+			name: "Success",
+			req:  algoReq{Algorithm: []byte("algorithm")},
+		},
+		{
+			name:        "Validation Error",
+			req:         algoReq{},
+			expectedErr: true,
+		},
+		{
+			name:        "Service Error",
+			req:         algoReq{Algorithm: []byte("algorithm")},
+			expectedErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.name == svcErr {
+				svc.On("Algo", context.Background(), agent.Algorithm{Algorithm: tt.req.Algorithm}).Return(errors.New("")).Once()
+			} else {
+				svc.On("Algo", context.Background(), agent.Algorithm{Algorithm: tt.req.Algorithm}).Return(nil).Once()
+			}
+			endpoint := algoEndpoint(svc)
+			_, err := endpoint(context.Background(), tt.req)
+			if (err != nil) != tt.expectedErr {
+				t.Errorf("algoEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
+			}
+		})
+	}
+}
+
+func TestDataEndpoint(t *testing.T) {
+	svc := new(mocks.Service)
+	tests := []struct {
+		name        string
+		req         dataReq
+		expectedErr bool
+	}{
+		{
+			name: "Success",
+			req:  dataReq{Dataset: []byte("dataset")},
+		},
+		{
+			name:        "Validation Error",
+			req:         dataReq{},
+			expectedErr: true,
+		},
+		{
+			name:        "Service Error",
+			req:         dataReq{Dataset: []byte("dataset")},
+			expectedErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.name == svcErr {
+				svc.On("Data", context.Background(), agent.Dataset{Dataset: tt.req.Dataset}).Return(errors.New("")).Once()
+			} else {
+				svc.On("Data", context.Background(), agent.Dataset{Dataset: tt.req.Dataset}).Return(nil).Once()
+			}
+			endpoint := dataEndpoint(svc)
+			_, err := endpoint(context.Background(), tt.req)
+			if (err != nil) != tt.expectedErr {
+				t.Errorf("dataEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
+			}
+		})
+	}
+}
+
+func TestResultEndpoint(t *testing.T) {
+	svc := new(mocks.Service)
+	tests := []struct {
+		name        string
+		req         resultReq
+		expectedErr bool
+	}{
+		{
+			name: "Success",
+			req:  resultReq{},
+		},
+		{
+			name:        "Service Error",
+			req:         resultReq{},
+			expectedErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.name == svcErr {
+				svc.On("Result", context.Background()).Return([]byte{}, errors.New("")).Once()
+			} else {
+				svc.On("Result", context.Background()).Return([]byte{}, nil).Once()
+			}
+			endpoint := resultEndpoint(svc)
+			res, err := endpoint(context.Background(), tt.req)
+			if (err != nil) != tt.expectedErr {
+				t.Errorf("resultEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
+			}
+			if err == nil {
+				_, ok := res.(resultRes)
+				if !ok {
+					t.Errorf("resultEndpoint() returned unexpected type %T", res)
+				}
+			}
+		})
+	}
+}
+
+func TestAttestationEndpoint(t *testing.T) {
+	svc := new(mocks.Service)
+	tests := []struct {
+		name        string
+		req         attestationReq
+		expectedErr bool
+	}{
+		{
+			name: "Success",
+			req:  attestationReq{TeeNonce: sha3.Sum512([]byte("report data")), VtpmNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: attestation.SNP},
+		},
+		{
+			name:        "Service Error",
+			req:         attestationReq{TeeNonce: sha3.Sum512([]byte("report data")), VtpmNonce: sha3.Sum256([]byte("vtpm nonce")), AttType: attestation.SNP},
+			expectedErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.name == svcErr {
+				svc.On("Attestation", context.Background(), tt.req.TeeNonce, tt.req.VtpmNonce, tt.req.AttType).Return([]byte{}, errors.New("")).Once()
+			} else {
+				svc.On("Attestation", context.Background(), tt.req.TeeNonce, tt.req.VtpmNonce, tt.req.AttType).Return([]byte{}, nil).Once()
+			}
+			endpoint := attestationEndpoint(svc)
+			res, err := endpoint(context.Background(), tt.req)
+			if (err != nil) != tt.expectedErr {
+				t.Errorf("attestationEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
+			}
+			if err == nil {
+				_, ok := res.(attestationRes)
+				if !ok {
+					t.Errorf("attestationEndpoint() returned unexpected type %T", res)
+				}
+			}
+		})
+	}
+}
+
+func TestAttestationTokenEndpoint(t *testing.T) {
+	svc := new(mocks.Service)
+	tests := []struct {
+		name        string
+		req         azureAttestationTokenReq
+		mockErr     error
+		expectedErr bool
+	}{
+		{
+			name:        "Success",
+			req:         azureAttestationTokenReq{tokenNonce: sha3.Sum256([]byte("vtpm nonce"))},
+			mockErr:     nil,
+			expectedErr: false,
+		},
+		{
+			name:        "Service Error",
+			req:         azureAttestationTokenReq{tokenNonce: sha3.Sum256([]byte("vtpm nonce"))},
+			mockErr:     errors.New("mock failure"),
+			expectedErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		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("AzureAttestationToken", mock.Anything, tt.req.tokenNonce).
+					Return([]byte("mock file"), tt.mockErr).Once()
+			}
+
+			endpoint := azureAttestationTokenEndpoint(svc)
+			res, err := endpoint(context.Background(), tt.req)
+
+			if (err != nil) != tt.expectedErr {
+				t.Errorf("attestationTokenEndpoint() error = %v, expectedErr %v", err, tt.expectedErr)
+			}
+
+			if !tt.expectedErr {
+				r, ok := res.(fetchAttestationTokenRes)
+				if !ok {
+					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)
+				}
+			}
+
+			svc.AssertExpectations(t)
+		})
+	}
+}

agent/api/grpc/interceptors.go

@@ -31,17 +31,24 @@ 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 {
-				return status.Errorf(codes.Unauthenticated, err.Error())
+				return status.Errorf(codes.Unauthenticated, "%v", err.Error())
 			}
 			return handler(srv, stream)
 		case agent.AgentService_Data_FullMethodName:
 			ctx, err := s.auth.AuthenticateUser(stream.Context(), auth.DataProviderRole)
 			if err != nil {
-				return status.Errorf(codes.Unauthenticated, err.Error())
+				return status.Errorf(codes.Unauthenticated, "%s", err.Error())
+			}
+			wrapped := &wrappedServerStream{ServerStream: stream, ctx: ctx}
+			return handler(srv, wrapped)
+		case agent.AgentService_Result_FullMethodName:
+			ctx, err := s.auth.AuthenticateUser(stream.Context(), auth.ConsumerRole)
+			if err != nil {
+				return status.Errorf(codes.Unauthenticated, "%v", err.Error())
 			}
 			wrapped := &wrappedServerStream{ServerStream: stream, ctx: ctx}
 			return handler(srv, wrapped)
@@ -52,12 +59,12 @@ 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)
 			if err != nil {
-				return nil, status.Errorf(codes.Unauthenticated, err.Error())
+				return nil, status.Errorf(codes.Unauthenticated, "%v", err.Error())
 			}
 			return handler(ctx, req)
 		default:

agent/api/grpc/interceptors_test.go

@@ -9,7 +9,7 @@ import (
 	"github.com/stretchr/testify/mock"
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/agent/auth"
-	"github.com/ultravioletrs/cocos/agent/mocks"
+	"github.com/ultravioletrs/cocos/agent/auth/mocks"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
 )
@@ -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

@@ -4,10 +4,15 @@ package grpc
 
 import (
 	"errors"
+
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 )
 
 type algoReq struct {
-	Algorithm []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
+	Algorithm    []byte `protobuf:"bytes,1,opt,name=algorithm,proto3" json:"algorithm,omitempty"`
+	Requirements []byte
 }
 
 func (req algoReq) validate() error {
@@ -18,12 +23,13 @@ func (req algoReq) validate() error {
 }
 
 type dataReq struct {
-	Dataset []byte `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
+	Dataset  []byte `protobuf:"bytes,1,opt,name=dataset,proto3" json:"dataset,omitempty"`
+	Filename string
 }
 
 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
 }
@@ -36,9 +42,35 @@ func (req resultReq) validate() error {
 }
 
 type attestationReq struct {
-	ReportData [64]byte
+	TeeNonce  [quoteprovider.Nonce]byte
+	VtpmNonce [vtpm.Nonce]byte
+	AttType   attestation.PlatformType
+}
+
+type azureAttestationTokenReq struct {
+	tokenNonce [vtpm.Nonce]byte
 }
 
 func (req attestationReq) 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:
+		return nil
+	default:
+		return errors.New("invalid attestation type")
+	}
+}
+
+type imaMeasurementsReq struct{}
+
+func (req imaMeasurementsReq) validate() error {
+	// No request parameters to validate, so no validation logic needed
 	return nil
 }

agent/api/grpc/responses.go

@@ -7,9 +7,18 @@ 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 {
 	File []byte
 }
+
+type imaMeasurementsRes struct {
+	File  []byte
+	PCR10 []byte
+}
+
+type fetchAttestationTokenRes struct {
+	File []byte
+}

agent/api/grpc/server.go

@@ -3,158 +3,438 @@
 package grpc
 
 import (
+	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"io"
+	"strconv"
 
+	"github.com/go-kit/kit/endpoint"
 	"github.com/go-kit/kit/transport/grpc"
 	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 )
 
+const (
+	bufferSize  = 1024 * 1024
+	FileSizeKey = "file-size"
+)
+
+var (
+	ErrTEENonceLength   = errors.New("malformed report data, expect less or equal to 64 bytes")
+	ErrVTPMNonceLength  = errors.New("malformed vTPM nonce, expect less or equal to 32 bytes")
+	ErrTokenNonceLength = errors.New("malformed token nonce, expect less or equal to 32 bytes")
+)
+
 var _ agent.AgentServiceServer = (*grpcServer)(nil)
 
 type grpcServer struct {
-	algo        grpc.Handler
-	data        grpc.Handler
-	result      grpc.Handler
-	attestation grpc.Handler
+	handlers map[string]grpc.Handler
 	agent.UnimplementedAgentServiceServer
 }
 
+type endpointConfig struct {
+	endpoint       func(agent.Service) endpoint.Endpoint
+	decodeRequest  grpc.DecodeRequestFunc
+	encodeResponse grpc.EncodeResponseFunc
+}
+
 // NewServer returns new AgentServiceServer instance.
 func NewServer(svc agent.Service) agent.AgentServiceServer {
+	// Define endpoint configurations
+	endpoints := map[string]endpointConfig{
+		"algo": {
+			endpoint:       algoEndpoint,
+			decodeRequest:  decodeAlgoRequest,
+			encodeResponse: encodeAlgoResponse,
+		},
+		"data": {
+			endpoint:       dataEndpoint,
+			decodeRequest:  decodeDataRequest,
+			encodeResponse: encodeDataResponse,
+		},
+		"result": {
+			endpoint:       resultEndpoint,
+			decodeRequest:  decodeResultRequest,
+			encodeResponse: encodeResultResponse,
+		},
+		"attestation": {
+			endpoint:       attestationEndpoint,
+			decodeRequest:  decodeAttestationRequest,
+			encodeResponse: encodeAttestationResponse,
+		},
+		"imaMeasurements": {
+			endpoint:       imaMeasurementsEndpoint,
+			decodeRequest:  decodeIMAMeasurementsRequest,
+			encodeResponse: encodeIMAMeasurementsResponse,
+		},
+		"azureAttestationToken": {
+			endpoint:       azureAttestationTokenEndpoint,
+			decodeRequest:  decodeAttestationTokenRequest,
+			encodeResponse: encodeAttestationTokenResponse,
+		},
+	}
+
+	// Create handlers using the configurations
+	handlers := make(map[string]grpc.Handler)
+	for name, config := range endpoints {
+		handlers[name] = grpc.NewServer(
+			config.endpoint(svc),
+			config.decodeRequest,
+			config.encodeResponse,
+		)
+	}
+
 	return &grpcServer{
-		algo: grpc.NewServer(
-			algoEndpoint(svc),
-			decodeAlgoRequest,
-			encodeAlgoResponse,
-		),
-		data: grpc.NewServer(
-			dataEndpoint(svc),
-			decodeDataRequest,
-			encodeDataResponse,
-		),
-		result: grpc.NewServer(
-			resultEndpoint(svc),
-			decodeResultRequest,
-			encodeResultResponse,
-		),
-		attestation: grpc.NewServer(
-			attestationEndpoint(svc),
-			decodeAttestationRequest,
-			encodeAttestationResponse,
-		),
+		handlers: handlers,
 	}
 }
 
-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,
+		Algorithm:    req.Algorithm,
+		Requirements: req.Requirements,
 	}, 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,
+		Dataset:  req.Dataset,
+		Filename: req.Filename,
 	}, 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 decodeAttestationRequest(_ context.Context, grpcReq interface{}) (interface{}, error) {
-	req := grpcReq.(*agent.AttestationRequest)
-	if len(req.ReportData) != agent.ReportDataSize {
-		return nil, errors.New("malformed report data, expect 64 bytes")
+func validateNonce(nonce []byte, maxLen int, target any) error {
+	if len(nonce) > maxLen {
+		switch maxLen {
+		case quoteprovider.Nonce:
+			return ErrTEENonceLength
+		case vtpm.Nonce:
+			return ErrVTPMNonceLength
+		default:
+			return ErrTokenNonceLength
+		}
 	}
-	return attestationReq{ReportData: [agent.ReportDataSize]byte(req.ReportData)}, nil
+
+	switch t := target.(type) {
+	case *[quoteprovider.Nonce]byte:
+		copy(t[:], nonce)
+	case *[vtpm.Nonce]byte:
+		copy(t[:], nonce)
+	default:
+		return fmt.Errorf("unsupported target type for nonce validation: %T", target)
+	}
+	return nil
 }
 
-func encodeAttestationResponse(_ context.Context, response 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
+
+	if err := validateNonce(req.TeeNonce, quoteprovider.Nonce, &reportData); err != nil {
+		return nil, err
+	}
+
+	if err := validateNonce(req.VtpmNonce, vtpm.Nonce, &nonce); err != nil {
+		return nil, err
+	}
+
+	return attestationReq{
+		TeeNonce:  reportData,
+		VtpmNonce: nonce,
+		AttType:   attestation.PlatformType(req.Type),
+	}, nil
+}
+
+func encodeAttestationResponse(_ context.Context, response any) (any, error) {
 	res := response.(attestationRes)
 	return &agent.AttestationResponse{
 		File: res.File,
 	}, nil
 }
 
-// Algo implements agent.AgentServiceServer.
-func (s *grpcServer) Algo(stream agent.AgentService_AlgoServer) error {
-	var algoFile []byte
+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 azureAttestationTokenReq{
+		tokenNonce: nonce,
+	}, nil
+}
+
+func encodeAttestationTokenResponse(_ context.Context, response any) (any, error) {
+	res := response.(fetchAttestationTokenRes)
+	return &agent.AttestationTokenResponse{
+		File: res.File,
+	}, nil
+}
+
+func decodeIMAMeasurementsRequest(_ context.Context, grpcReq any) (any, error) {
+	return imaMeasurementsReq{}, nil
+}
+
+func encodeIMAMeasurementsResponse(_ context.Context, response any) (any, error) {
+	res := response.(imaMeasurementsRes)
+	return &agent.IMAMeasurementsResponse{
+		File:  res.File,
+		Pcr10: res.PCR10,
+	}, nil
+}
+
+func (s *grpcServer) streamingHandler(
+	ctx context.Context,
+	handlerName string,
+	req any,
+	stream any,
+	sendFn func([]byte) error,
+	getFileData func(any) []byte,
+) error {
+	handler, ok := s.handlers[handlerName]
+	if !ok {
+		return status.Errorf(codes.NotFound, "handler %q not found", handlerName)
+	}
+
+	_, res, err := handler.ServeGRPC(ctx, req)
+	if err != nil {
+		return err
+	}
+
+	fileData := getFileData(res)
+
+	// Set file size header
+	if setter, ok := stream.(interface{ SetHeader(metadata.MD) error }); ok {
+		if err := setter.SetHeader(metadata.New(map[string]string{
+			FileSizeKey: fmt.Sprint(len(fileData)),
+		})); err != nil {
+			return status.Error(codes.Internal, err.Error())
+		}
+	}
+
+	// Stream the file data
+	return s.streamFileData(bytes.NewBuffer(fileData), sendFn)
+}
+
+func (s *grpcServer) streamFileData(buffer *bytes.Buffer, sendFn func([]byte) error) error {
+	buf := make([]byte, bufferSize)
 	for {
-		algoChunk, err := stream.Recv()
+		n, err := buffer.Read(buf)
 		if err == io.EOF {
 			break
 		}
 		if err != nil {
 			return status.Error(codes.Internal, err.Error())
 		}
-		algoFile = append(algoFile, algoChunk.Algorithm...)
+
+		if err := sendFn(buf[:n]); err != nil {
+			return status.Error(codes.Internal, err.Error())
+		}
 	}
-	_, res, err := s.algo.ServeGRPC(stream.Context(), &agent.AlgoRequest{Algorithm: algoFile})
+	return nil
+}
+
+func receiveStreamingData(getData func() ([]byte, string, error)) ([]byte, string, error) {
+	var data []byte
+	var filename string
+
+	for {
+		chunk, fname, err := getData()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, "", status.Error(codes.Internal, err.Error())
+		}
+		data = append(data, chunk...)
+		if fname != "" {
+			filename = fname
+		}
+	}
+	return data, filename, nil
+}
+
+// Algo implements agent.AgentServiceServer.
+func (s *grpcServer) Algo(stream agent.AgentService_AlgoServer) error {
+	algoFile, reqFile, err := s.receiveAlgoData(stream)
 	if err != nil {
 		return err
 	}
-	ar := res.(*agent.AlgoResponse)
-	return stream.SendAndClose(ar)
+
+	_, res, err := s.handlers["algo"].ServeGRPC(stream.Context(), &agent.AlgoRequest{
+		Algorithm:    algoFile,
+		Requirements: reqFile,
+	})
+	if err != nil {
+		return err
+	}
+
+	return stream.SendAndClose(res.(*agent.AlgoResponse))
+}
+
+func (s *grpcServer) receiveAlgoData(stream agent.AgentService_AlgoServer) ([]byte, []byte, error) {
+	var algoFile, reqFile []byte
+	for {
+		chunk, err := stream.Recv()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return nil, nil, status.Error(codes.Internal, err.Error())
+		}
+		algoFile = append(algoFile, chunk.Algorithm...)
+		reqFile = append(reqFile, chunk.Requirements...)
+	}
+	return algoFile, reqFile, nil
 }
 
 // Data implements agent.AgentServiceServer.
 func (s *grpcServer) Data(stream agent.AgentService_DataServer) error {
-	var dataFile []byte
-	for {
-		dataChunk, err := stream.Recv()
-		if err == io.EOF {
-			break
-		}
+	dataFile, filename, err := receiveStreamingData(func() ([]byte, string, error) {
+		chunk, err := stream.Recv()
 		if err != nil {
-			return status.Error(codes.Internal, err.Error())
+			return nil, "", err
 		}
-		dataFile = append(dataFile, dataChunk.Dataset...)
-	}
-	_, res, err := s.data.ServeGRPC(stream.Context(), &agent.DataRequest{Dataset: dataFile})
+		return chunk.Dataset, chunk.Filename, nil
+	})
 	if err != nil {
 		return err
 	}
-	ar := res.(*agent.DataResponse)
-	return stream.SendAndClose(ar)
+
+	_, res, err := s.handlers["data"].ServeGRPC(stream.Context(), &agent.DataRequest{
+		Dataset:  dataFile,
+		Filename: filename,
+	})
+	if err != nil {
+		return err
+	}
+
+	return stream.SendAndClose(res.(*agent.DataResponse))
 }
 
-func (s *grpcServer) Result(ctx context.Context, req *agent.ResultRequest) (*agent.ResultResponse, error) {
-	_, res, err := s.result.ServeGRPC(ctx, req)
+func (s *grpcServer) Result(req *agent.ResultRequest, stream agent.AgentService_ResultServer) error {
+	return s.streamingHandler(
+		stream.Context(),
+		"result",
+		req,
+		stream,
+		func(data []byte) error {
+			return stream.Send(&agent.ResultResponse{File: data})
+		},
+		func(res any) []byte {
+			return res.(*agent.ResultResponse).File
+		},
+	)
+}
+
+func (s *grpcServer) Attestation(req *agent.AttestationRequest, stream agent.AgentService_AttestationServer) error {
+	return s.streamingHandler(
+		stream.Context(),
+		"attestation",
+		req,
+		stream,
+		func(data []byte) error {
+			return stream.Send(&agent.AttestationResponse{File: data})
+		},
+		func(res any) []byte {
+			return res.(*agent.AttestationResponse).File
+		},
+	)
+}
+
+func (s *grpcServer) IMAMeasurements(req *agent.IMAMeasurementsRequest, stream agent.AgentService_IMAMeasurementsServer) error {
+	_, res, err := s.handlers["imaMeasurements"].ServeGRPC(stream.Context(), req)
+	if err != nil {
+		return err
+	}
+	rr := res.(*agent.IMAMeasurementsResponse)
+
+	if err := stream.SetHeader(metadata.New(map[string]string{
+		FileSizeKey: strconv.Itoa(len(rr.File)),
+	})); err != nil {
+		return status.Error(codes.Internal, err.Error())
+	}
+
+	return s.streamDualBuffers(
+		bytes.NewBuffer(rr.File),
+		bytes.NewBuffer(rr.Pcr10),
+		func(fileData, pcr10Data []byte) error {
+			return stream.Send(&agent.IMAMeasurementsResponse{
+				File:  fileData,
+				Pcr10: pcr10Data,
+			})
+		},
+	)
+}
+
+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 := res.(*agent.ResultResponse)
+
+	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) Attestation(ctx context.Context, req *agent.AttestationRequest) (*agent.AttestationResponse, error) {
-	_, res, err := s.attestation.ServeGRPC(ctx, req)
-	if err != nil {
-		return nil, err
+func (s *grpcServer) streamDualBuffers(
+	buf1, buf2 *bytes.Buffer,
+	sendFn func([]byte, []byte) error,
+) error {
+	buff1 := make([]byte, bufferSize)
+	buff2 := make([]byte, bufferSize)
+
+	for {
+		n1, err1 := buf1.Read(buff1)
+		if err1 != nil && err1 != io.EOF {
+			return status.Error(codes.Internal, err1.Error())
+		}
+
+		n2, err2 := buf2.Read(buff2)
+		if err2 != nil && err2 != io.EOF {
+			return status.Error(codes.Internal, err2.Error())
+		}
+
+		if n1 == 0 && err1 == io.EOF && n2 == 0 && err2 == io.EOF {
+			break
+		}
+
+		if err := sendFn(buff1[:n1], buff2[:n2]); err != nil {
+			return status.Error(codes.Internal, err.Error())
+		}
 	}
-	rr := res.(*agent.AttestationResponse)
-	return rr, nil
+	return nil
 }

agent/api/grpc/server_test.go

@@ -0,0 +1,500 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package grpc
+
+import (
+	"context"
+	"io"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/agent/mocks"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+type MockAgentService_AlgoServer struct {
+	grpc.ServerStream
+	mock.Mock
+	ctx context.Context
+}
+
+func (m *MockAgentService_AlgoServer) Context() context.Context {
+	return m.ctx
+}
+
+func (m *MockAgentService_AlgoServer) Recv() (*agent.AlgoRequest, error) {
+	args := m.Called()
+	return args.Get(0).(*agent.AlgoRequest), args.Error(1)
+}
+
+func (m *MockAgentService_AlgoServer) SendAndClose(resp *agent.AlgoResponse) error {
+	args := m.Called(resp)
+	return args.Error(0)
+}
+
+type MockAgentService_DataServer struct {
+	grpc.ServerStream
+	mock.Mock
+	ctx context.Context
+}
+
+func (m *MockAgentService_DataServer) Context() context.Context {
+	return m.ctx
+}
+
+func (m *MockAgentService_DataServer) Recv() (*agent.DataRequest, error) {
+	args := m.Called()
+	return args.Get(0).(*agent.DataRequest), args.Error(1)
+}
+
+func (m *MockAgentService_DataServer) SendAndClose(resp *agent.DataResponse) error {
+	args := m.Called(resp)
+	return args.Error(0)
+}
+
+type MockAgentService_ResultServer struct {
+	grpc.ServerStream
+	mock.Mock
+	ctx context.Context
+}
+
+func (m *MockAgentService_ResultServer) Context() context.Context {
+	return m.ctx
+}
+
+func (m *MockAgentService_ResultServer) SetHeader(md metadata.MD) error {
+	args := m.Called(md)
+	return args.Error(0)
+}
+
+func (m *MockAgentService_ResultServer) Send(resp *agent.ResultResponse) error {
+	args := m.Called(resp)
+	return args.Error(0)
+}
+
+type MockAgentService_AttestationServer struct {
+	grpc.ServerStream
+	mock.Mock
+	ctx context.Context
+}
+
+func (m *MockAgentService_AttestationServer) Context() context.Context {
+	return m.ctx
+}
+
+func (m *MockAgentService_AttestationServer) Send(resp *agent.AttestationResponse) error {
+	args := m.Called(resp)
+	return args.Error(0)
+}
+
+func (m *MockAgentService_AttestationServer) SetHeader(md metadata.MD) error {
+	args := m.Called(md)
+	return args.Error(0)
+}
+
+type MockAgentService_IMAMeasurementsServer struct {
+	grpc.ServerStream
+	mock.Mock
+	ctx context.Context
+}
+
+func (m *MockAgentService_IMAMeasurementsServer) Context() context.Context {
+	return m.ctx
+}
+
+func (m *MockAgentService_IMAMeasurementsServer) Send(resp *agent.IMAMeasurementsResponse) error {
+	args := m.Called(resp)
+	return args.Error(0)
+}
+
+func (m *MockAgentService_IMAMeasurementsServer) SetHeader(md metadata.MD) error {
+	args := m.Called(md)
+	return args.Error(0)
+}
+
+func TestNewServer(t *testing.T) {
+	mockService := new(mocks.Service)
+	server := NewServer(mockService)
+
+	grpcServer, ok := server.(*grpcServer)
+	assert.True(t, ok)
+	assert.NotNil(t, grpcServer.handlers)
+	assert.Len(t, grpcServer.handlers, 6) // Should have 6 handlers
+
+	// Check that all expected handlers are present
+	expectedHandlers := []string{"algo", "data", "result", "attestation", "imaMeasurements", "azureAttestationToken"}
+	for _, handler := range expectedHandlers {
+		assert.Contains(t, grpcServer.handlers, handler)
+		assert.NotNil(t, grpcServer.handlers[handler])
+	}
+}
+
+func TestAlgo(t *testing.T) {
+	mockService := new(mocks.Service)
+	server := NewServer(mockService)
+
+	mockStream := &MockAgentService_AlgoServer{ctx: context.Background()}
+	mockStream.On("Recv").Return(&agent.AlgoRequest{Algorithm: []byte("algo"), Requirements: []byte("req")}, nil).Once()
+	mockStream.On("Recv").Return(&agent.AlgoRequest{}, io.EOF).Once()
+	mockStream.On("SendAndClose", &agent.AlgoResponse{}).Return(nil).Once()
+
+	mockService.On("Algo", context.Background(), agent.Algorithm{Algorithm: []byte("algo"), Requirements: []byte("req")}).Return(nil)
+
+	err := server.Algo(mockStream)
+	assert.NoError(t, err)
+
+	mockStream.AssertExpectations(t)
+	mockService.AssertExpectations(t)
+}
+
+func TestAlgoWithMultipleChunks(t *testing.T) {
+	mockService := new(mocks.Service)
+	server := NewServer(mockService)
+
+	mockStream := &MockAgentService_AlgoServer{ctx: context.Background()}
+	mockStream.On("Recv").Return(&agent.AlgoRequest{Algorithm: []byte("algo"), Requirements: []byte("req")}, nil).Once()
+	mockStream.On("Recv").Return(&agent.AlgoRequest{Algorithm: []byte("2"), Requirements: []byte("2")}, nil).Once()
+	mockStream.On("Recv").Return(&agent.AlgoRequest{}, io.EOF).Once()
+	mockStream.On("SendAndClose", &agent.AlgoResponse{}).Return(nil).Once()
+
+	mockService.On("Algo", context.Background(), agent.Algorithm{Algorithm: []byte("algo2"), Requirements: []byte("req2")}).Return(nil)
+
+	err := server.Algo(mockStream)
+	assert.NoError(t, err)
+
+	mockStream.AssertExpectations(t)
+	mockService.AssertExpectations(t)
+}
+
+func TestData(t *testing.T) {
+	mockService := new(mocks.Service)
+	server := NewServer(mockService)
+
+	mockStream := &MockAgentService_DataServer{ctx: context.Background()}
+	mockStream.On("Recv").Return(&agent.DataRequest{Dataset: []byte("data"), Filename: "test.txt"}, nil).Once()
+	mockStream.On("Recv").Return(&agent.DataRequest{}, io.EOF).Once()
+	mockStream.On("SendAndClose", &agent.DataResponse{}).Return(nil).Once()
+
+	mockService.On("Data", context.Background(), agent.Dataset{Dataset: []byte("data"), Filename: "test.txt"}).Return(nil)
+
+	err := server.Data(mockStream)
+	assert.NoError(t, err)
+
+	mockStream.AssertExpectations(t)
+	mockService.AssertExpectations(t)
+}
+
+func TestResult(t *testing.T) {
+	mockService := new(mocks.Service)
+	server := NewServer(mockService)
+
+	resultData := []byte("result data")
+	mockStream := &MockAgentService_ResultServer{ctx: context.Background()}
+
+	// Mock the SetHeader call
+	mockStream.On("SetHeader", mock.AnythingOfType("metadata.MD")).Return(nil).Once()
+
+	// Mock the Send call - it should be called with the result data
+	mockStream.On("Send", mock.MatchedBy(func(resp *agent.ResultResponse) bool {
+		return len(resp.File) > 0
+	})).Return(nil).Once()
+
+	mockService.On("Result", mock.Anything).Return(resultData, nil)
+
+	err := server.Result(&agent.ResultRequest{}, mockStream)
+	assert.NoError(t, err)
+
+	mockStream.AssertExpectations(t)
+	mockService.AssertExpectations(t)
+}
+
+func TestAttestation(t *testing.T) {
+	mockService := new(mocks.Service)
+	server := NewServer(mockService)
+
+	attestationData := []byte("attestation data")
+	mockStream := &MockAgentService_AttestationServer{ctx: context.Background()}
+
+	// Mock the SetHeader call
+	mockStream.On("SetHeader", mock.AnythingOfType("metadata.MD")).Return(nil).Once()
+
+	// Mock the Send call
+	mockStream.On("Send", mock.MatchedBy(func(resp *agent.AttestationResponse) bool {
+		return len(resp.File) > 0
+	})).Return(nil).Once()
+
+	reportData := [quoteprovider.Nonce]byte{}
+	vtpmNonce := [vtpm.Nonce]byte{}
+	attestationType := attestation.SNP
+	mockService.On("Attestation", mock.Anything, reportData, vtpmNonce, attestationType).Return(attestationData, nil)
+
+	err := server.Attestation(&agent.AttestationRequest{TeeNonce: reportData[:], Type: int32(attestationType)}, mockStream)
+	assert.NoError(t, err)
+
+	mockService.AssertExpectations(t)
+	mockStream.AssertExpectations(t)
+}
+
+func TestIMAMeasurements(t *testing.T) {
+	mockService := new(mocks.Service)
+	server := NewServer(mockService)
+
+	imaData := []byte("ima data")
+	pcr10Data := []byte("pcr10 data")
+
+	mockStream := &MockAgentService_IMAMeasurementsServer{ctx: context.Background()}
+
+	// Mock the SetHeader call
+	mockStream.On("SetHeader", mock.AnythingOfType("metadata.MD")).Return(nil).Once()
+
+	// Mock the Send call
+	mockStream.On("Send", mock.MatchedBy(func(resp *agent.IMAMeasurementsResponse) bool {
+		return len(resp.File) > 0 || len(resp.Pcr10) > 0
+	})).Return(nil).Once()
+
+	mockService.On("IMAMeasurements", mock.Anything).Return(imaData, pcr10Data, nil)
+
+	err := server.IMAMeasurements(&agent.IMAMeasurementsRequest{}, mockStream)
+	assert.NoError(t, err)
+
+	mockService.AssertExpectations(t)
+	mockStream.AssertExpectations(t)
+}
+
+func TestAttestationToken(t *testing.T) {
+	mockService := new(mocks.Service)
+	server := NewServer(mockService)
+
+	attestationData := []byte("attestation token data")
+	vtpmNonce := [vtpm.Nonce]byte{}
+	attestationType := attestation.SNP
+
+	mockService.On("AzureAttestationToken", mock.Anything, vtpmNonce).Return(attestationData, nil)
+
+	resp, err := server.AzureAttestationToken(context.Background(), &agent.AttestationTokenRequest{
+		TokenNonce: vtpmNonce[:],
+		Type:       int32(attestationType),
+	})
+
+	assert.NoError(t, err)
+	assert.Equal(t, attestationData, resp.File)
+
+	mockService.AssertExpectations(t)
+}
+
+func TestValidateNonce(t *testing.T) {
+	tests := []struct {
+		name        string
+		nonce       []byte
+		maxLen      int
+		shouldError bool
+		expectedErr error
+	}{
+		{
+			name:        "valid TEE nonce",
+			nonce:       make([]byte, quoteprovider.Nonce),
+			maxLen:      quoteprovider.Nonce,
+			shouldError: false,
+		},
+		{
+			name:        "valid vTPM nonce",
+			nonce:       make([]byte, vtpm.Nonce),
+			maxLen:      vtpm.Nonce,
+			shouldError: false,
+		},
+		{
+			name:        "TEE nonce too long",
+			nonce:       make([]byte, quoteprovider.Nonce+1),
+			maxLen:      quoteprovider.Nonce,
+			shouldError: true,
+			expectedErr: ErrTEENonceLength,
+		},
+		{
+			name:        "vTPM nonce too long",
+			nonce:       make([]byte, vtpm.Nonce+1),
+			maxLen:      vtpm.Nonce,
+			shouldError: true,
+			expectedErr: ErrVTPMNonceLength,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.maxLen == quoteprovider.Nonce {
+				var target [quoteprovider.Nonce]byte
+				err := validateNonce(tt.nonce, tt.maxLen, &target)
+				if tt.shouldError {
+					assert.Error(t, err)
+					assert.Equal(t, tt.expectedErr, err)
+				} else {
+					assert.NoError(t, err)
+				}
+			} else {
+				var target [vtpm.Nonce]byte
+				err := validateNonce(tt.nonce, tt.maxLen, &target)
+				if tt.shouldError {
+					assert.Error(t, err)
+					assert.Equal(t, tt.expectedErr, err)
+				} else {
+					assert.NoError(t, err)
+				}
+			}
+		})
+	}
+}
+
+func TestDecodeAlgoRequest(t *testing.T) {
+	req := &agent.AlgoRequest{Algorithm: []byte("algo"), Requirements: []byte("req")}
+	decoded, err := decodeAlgoRequest(context.Background(), req)
+	assert.NoError(t, err)
+	assert.Equal(t, algoReq{Algorithm: []byte("algo"), Requirements: []byte("req")}, decoded)
+}
+
+func TestEncodeAlgoResponse(t *testing.T) {
+	encoded, err := encodeAlgoResponse(context.Background(), algoRes{})
+	assert.NoError(t, err)
+	assert.Equal(t, &agent.AlgoResponse{}, encoded)
+}
+
+func TestDecodeDataRequest(t *testing.T) {
+	req := &agent.DataRequest{Dataset: []byte("data"), Filename: "test.txt"}
+	decoded, err := decodeDataRequest(context.Background(), req)
+	assert.NoError(t, err)
+	assert.Equal(t, dataReq{Dataset: []byte("data"), Filename: "test.txt"}, decoded)
+}
+
+func TestEncodeDataResponse(t *testing.T) {
+	encoded, err := encodeDataResponse(context.Background(), dataRes{})
+	assert.NoError(t, err)
+	assert.Equal(t, &agent.DataResponse{}, encoded)
+}
+
+func TestDecodeResultRequest(t *testing.T) {
+	decoded, err := decodeResultRequest(context.Background(), &agent.ResultRequest{})
+	assert.NoError(t, err)
+	assert.Equal(t, resultReq{}, decoded)
+}
+
+func TestEncodeResultResponse(t *testing.T) {
+	encoded, err := encodeResultResponse(context.Background(), resultRes{File: []byte("result")})
+	assert.NoError(t, err)
+	assert.Equal(t, &agent.ResultResponse{File: []byte("result")}, encoded)
+}
+
+func TestDecodeAttestationRequest(t *testing.T) {
+	teeNonce := make([]byte, quoteprovider.Nonce)
+	vtpmNonce := make([]byte, vtpm.Nonce)
+
+	req := &agent.AttestationRequest{
+		TeeNonce:  teeNonce,
+		VtpmNonce: vtpmNonce,
+		Type:      int32(attestation.SNP),
+	}
+
+	decoded, err := decodeAttestationRequest(context.Background(), req)
+	assert.NoError(t, err)
+
+	decodedReq := decoded.(attestationReq)
+	assert.Equal(t, attestation.SNP, decodedReq.AttType)
+}
+
+func TestDecodeAttestationRequestWithInvalidNonce(t *testing.T) {
+	// Test with TEE nonce too long
+	teeNonce := make([]byte, quoteprovider.Nonce+1)
+	req := &agent.AttestationRequest{TeeNonce: teeNonce}
+
+	_, err := decodeAttestationRequest(context.Background(), req)
+	assert.Error(t, err)
+	assert.Equal(t, ErrTEENonceLength, err)
+
+	// Test with vTPM nonce too long
+	vtpmNonce := make([]byte, vtpm.Nonce+1)
+	req = &agent.AttestationRequest{VtpmNonce: vtpmNonce}
+
+	_, err = decodeAttestationRequest(context.Background(), req)
+	assert.Error(t, err)
+	assert.Equal(t, ErrVTPMNonceLength, err)
+}
+
+func TestEncodeAttestationResponse(t *testing.T) {
+	encoded, err := encodeAttestationResponse(context.Background(), attestationRes{File: []byte("attestation")})
+	assert.NoError(t, err)
+	assert.Equal(t, &agent.AttestationResponse{File: []byte("attestation")}, encoded)
+}
+
+func TestDecodeAttestationTokenRequest(t *testing.T) {
+	tokenNonce := make([]byte, vtpm.Nonce)
+	req := &agent.AttestationTokenRequest{
+		TokenNonce: tokenNonce,
+		Type:       int32(attestation.SNP),
+	}
+
+	_, err := decodeAttestationTokenRequest(context.Background(), req)
+	assert.NoError(t, err)
+}
+
+func TestDecodeAttestationTokenRequestWithInvalidNonce(t *testing.T) {
+	// Test with token nonce too long
+	tokenNonce := make([]byte, vtpm.Nonce+1)
+	req := &agent.AttestationTokenRequest{TokenNonce: tokenNonce}
+
+	_, err := decodeAttestationTokenRequest(context.Background(), req)
+	assert.Error(t, err)
+	assert.Equal(t, ErrVTPMNonceLength, err)
+}
+
+func TestEncodeAttestationTokenResponse(t *testing.T) {
+	encoded, err := encodeAttestationTokenResponse(context.Background(), fetchAttestationTokenRes{File: []byte("attestation")})
+	assert.NoError(t, err)
+	assert.Equal(t, &agent.AttestationTokenResponse{File: []byte("attestation")}, encoded)
+}
+
+func TestDecodeIMAMeasurementsRequest(t *testing.T) {
+	decoded, err := decodeIMAMeasurementsRequest(context.Background(), &agent.IMAMeasurementsRequest{})
+	assert.NoError(t, err)
+	assert.Equal(t, imaMeasurementsReq{}, decoded)
+}
+
+func TestEncodeIMAMeasurementsResponse(t *testing.T) {
+	encoded, err := encodeIMAMeasurementsResponse(context.Background(), imaMeasurementsRes{
+		File:  []byte("ima"),
+		PCR10: []byte("pcr10"),
+	})
+	assert.NoError(t, err)
+	assert.Equal(t, &agent.IMAMeasurementsResponse{
+		File:  []byte("ima"),
+		Pcr10: []byte("pcr10"),
+	}, encoded)
+}
+
+func TestAlgoWithStreamError(t *testing.T) {
+	mockService := new(mocks.Service)
+	server := NewServer(mockService)
+
+	mockStream := &MockAgentService_AlgoServer{ctx: context.Background()}
+	mockStream.On("Recv").Return(&agent.AlgoRequest{}, assert.AnError).Once()
+
+	err := server.Algo(mockStream)
+	assert.Error(t, err)
+
+	mockStream.AssertExpectations(t)
+}
+
+func TestDataWithStreamError(t *testing.T) {
+	mockService := new(mocks.Service)
+	server := NewServer(mockService)
+
+	mockStream := &MockAgentService_DataServer{ctx: context.Background()}
+	mockStream.On("Recv").Return(&agent.DataRequest{}, assert.AnError).Once()
+
+	err := server.Data(mockStream)
+	assert.Error(t, err)
+
+	mockStream.AssertExpectations(t)
+}

agent/api/logging.go

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

agent/api/metrics.go

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

agent/auth/auth.go

@@ -13,7 +13,7 @@ import (
 	"crypto/x509"
 	"encoding/base64"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/ultravioletrs/cocos/agent"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -36,15 +36,14 @@ var (
 	ErrSignatureVerificationFailed = errors.New("signature verification failed")
 )
 
-//go:generate mockery --name Authenticator --output=../mocks --filename auth.go --quiet --note "Copyright (c) Ultraviolet \n // SPDX-License-Identifier: Apache-2.0"
 type Authenticator interface {
 	AuthenticateUser(ctx context.Context, role UserRole) (context.Context, error)
 }
 
 type service struct {
-	resultConsumers   []interface{}
-	datasetProviders  []interface{}
-	algorithmProvider interface{}
+	resultConsumers   []any
+	datasetProviders  []any
+	algorithmProvider any
 }
 
 func New(manifest agent.Computation) (Authenticator, error) {
@@ -146,9 +145,9 @@ func (s *service) AuthenticateUser(ctx context.Context, role UserRole) (context.
 			}
 		}
 	case DataProviderRole:
-		for i, dp := range s.datasetProviders {
+		for _, dp := range s.datasetProviders {
 			if err := verifySignature(role, signature, dp); err == nil {
-				return agent.IndexToContext(ctx, i), nil
+				return ctx, nil
 			}
 		}
 	case AlgorithmProviderRole:

agent/auth/auth_test.go

@@ -15,7 +15,7 @@ import (
 	"encoding/base64"
 	"testing"
 
-	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/absmach/supermq/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/ultravioletrs/cocos/agent"
@@ -124,7 +124,7 @@ func TestAuthenticateUser(t *testing.T) {
 
 			if err == nil {
 				switch id, ok := agent.IndexFromContext(ctx); {
-				case tc.role == ConsumerRole, tc.role == DataProviderRole:
+				case tc.role == ConsumerRole:
 					assert.True(t, ok, "expected index in context")
 					assert.Equal(t, 0, id, "expected index 0 in context")
 				default:

agent/auth/mocks/authenticator.go

@@ -0,0 +1,110 @@
+// 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/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
+}
+
+type Authenticator_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *Authenticator) EXPECT() *Authenticator_Expecter {
+	return &Authenticator_Expecter{mock: &_m.Mock}
+}
+
+// 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")
+	}
+
+	var r0 context.Context
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, auth.UserRole) (context.Context, error)); ok {
+		return returnFunc(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 returnFunc, ok := ret.Get(1).(func(context.Context, auth.UserRole) error); ok {
+		r1 = returnFunc(ctx, role)
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// Authenticator_AuthenticateUser_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'AuthenticateUser'
+type Authenticator_AuthenticateUser_Call struct {
+	*mock.Call
+}
+
+// AuthenticateUser is a helper method to define mock.On call
+//   - ctx context.Context
+//   - role auth.UserRole
+func (_e *Authenticator_Expecter) AuthenticateUser(ctx interface{}, role interface{}) *Authenticator_AuthenticateUser_Call {
+	return &Authenticator_AuthenticateUser_Call{Call: _e.mock.On("AuthenticateUser", ctx, role)}
+}
+
+func (_c *Authenticator_AuthenticateUser_Call) Run(run func(ctx context.Context, role auth.UserRole)) *Authenticator_AuthenticateUser_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		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(context1 context.Context, err error) *Authenticator_AuthenticateUser_Call {
+	_c.Call.Return(context1, err)
+	return _c
+}
+
+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
+}

agent/computations.go

@@ -6,19 +6,19 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+
+	"google.golang.org/grpc/metadata"
 )
 
 var _ fmt.Stringer = (*Datasets)(nil)
 
 type AgentConfig struct {
-	LogLevel     string `json:"log_level"`
-	Host         string `json:"host"`
-	Port         string `json:"port"`
-	CertFile     string `json:"cert_file"`
-	KeyFile      string `json:"server_key"`
-	ServerCAFile string `json:"server_ca_file"`
-	ClientCAFile string `json:"client_ca_file"`
-	AttestedTls  bool   `json:"attested_tls"`
+	Port         string `json:"port,omitempty"`
+	CertFile     string `json:"cert_file,omitempty"`
+	KeyFile      string `json:"server_key,omitempty"`
+	ServerCAFile string `json:"server_ca_file,omitempty"`
+	ClientCAFile string `json:"client_ca_file,omitempty"`
+	AttestedTls  bool   `json:"attested_tls,omitempty"`
 }
 
 type Computation struct {
@@ -28,7 +28,6 @@ type Computation struct {
 	Datasets        Datasets         `json:"datasets,omitempty"`
 	Algorithm       Algorithm        `json:"algorithm,omitempty"`
 	ResultConsumers []ResultConsumer `json:"result_consumers,omitempty"`
-	AgentConfig     AgentConfig      `json:"agent_config,omitempty"`
 }
 
 type ResultConsumer struct {
@@ -44,17 +43,19 @@ func (d *Datasets) String() string {
 }
 
 type Dataset struct {
-	Dataset []byte   `json:"-"`
-	Hash    [32]byte `json:"hash,omitempty"`
-	UserKey []byte   `json:"user_key,omitempty"`
+	Dataset  []byte   `json:"-"`
+	Hash     [32]byte `json:"hash,omitempty"`
+	UserKey  []byte   `json:"user_key,omitempty"`
+	Filename string   `json:"filename,omitempty"`
 }
 
 type Datasets []Dataset
 
 type Algorithm struct {
-	Algorithm []byte   `json:"-"`
-	Hash      [32]byte `json:"hash,omitempty"`
-	UserKey   []byte   `json:"user_key,omitempty"`
+	Algorithm    []byte   `json:"-"`
+	Hash         [32]byte `json:"hash,omitempty"`
+	UserKey      []byte   `json:"user_key,omitempty"`
+	Requirements []byte   `json:"-"`
 }
 
 type ManifestIndexKey struct{}
@@ -67,3 +68,18 @@ func IndexFromContext(ctx context.Context) (int, bool) {
 	index, ok := ctx.Value(ManifestIndexKey{}).(int)
 	return index, ok
 }
+
+const DecompressKey = "decompress"
+
+func DecompressFromContext(ctx context.Context) bool {
+	vals := metadata.ValueFromIncomingContext(ctx, DecompressKey)
+	if len(vals) == 0 {
+		return false
+	}
+
+	return vals[0] == "true"
+}
+
+func DecompressToContext(ctx context.Context, decompress bool) context.Context {
+	return metadata.AppendToOutgoingContext(ctx, DecompressKey, fmt.Sprintf("%t", decompress))
+}

agent/computations_test.go

@@ -0,0 +1,131 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"reflect"
+	"testing"
+
+	"google.golang.org/grpc/metadata"
+)
+
+func TestDatasetsString(t *testing.T) {
+	datasets := Datasets{
+		{
+			Hash:     [32]byte{1, 2, 3},
+			UserKey:  []byte("user_key"),
+			Filename: "test.dat",
+		},
+	}
+
+	expected := `[{"hash":[1,2,3,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],"user_key":"dXNlcl9rZXk=","filename":"test.dat"}]`
+	result := datasets.String()
+
+	if result != expected {
+		t.Errorf("Datasets.String() = %v, want %v", result, expected)
+	}
+}
+
+func TestIndexToContext(t *testing.T) {
+	ctx := context.Background()
+	index := 5
+
+	newCtx := IndexToContext(ctx, index)
+	result, ok := IndexFromContext(newCtx)
+
+	if !ok {
+		t.Errorf("IndexFromContext() ok = false, want true")
+	}
+
+	if result != index {
+		t.Errorf("IndexFromContext() = %v, want %v", result, index)
+	}
+}
+
+func TestDecompressFromContext(t *testing.T) {
+	tests := []struct {
+		name     string
+		ctx      context.Context
+		expected bool
+	}{
+		{
+			name:     "No decompress metadata",
+			ctx:      context.Background(),
+			expected: false,
+		},
+		{
+			name: "Decompress true",
+			ctx: metadata.NewIncomingContext(
+				context.Background(),
+				metadata.Pairs(DecompressKey, "true"),
+			),
+			expected: true,
+		},
+		{
+			name: "Decompress false",
+			ctx: metadata.NewIncomingContext(
+				context.Background(),
+				metadata.Pairs(DecompressKey, "false"),
+			),
+			expected: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := DecompressFromContext(tt.ctx)
+			if result != tt.expected {
+				t.Errorf("DecompressFromContext() = %v, want %v", result, tt.expected)
+			}
+		})
+	}
+}
+
+func TestDecompressToContext(t *testing.T) {
+	ctx := context.Background()
+	decompress := true
+
+	newCtx := DecompressToContext(ctx, decompress)
+	md, ok := metadata.FromOutgoingContext(newCtx)
+
+	if !ok {
+		t.Errorf("metadata.FromOutgoingContext() ok = false, want true")
+	}
+
+	vals := md.Get(DecompressKey)
+	if len(vals) != 1 {
+		t.Errorf("len(md.Get(DecompressKey)) = %v, want 1", len(vals))
+	}
+
+	if vals[0] != "true" {
+		t.Errorf("md.Get(DecompressKey)[0] = %v, want 'true'", vals[0])
+	}
+}
+
+func TestAgentConfigJSON(t *testing.T) {
+	config := AgentConfig{
+		Port:         "8080",
+		CertFile:     "cert.pem",
+		KeyFile:      "key.pem",
+		ServerCAFile: "server_ca.pem",
+		ClientCAFile: "client_ca.pem",
+		AttestedTls:  true,
+	}
+
+	data, err := json.Marshal(config)
+	if err != nil {
+		t.Fatalf("Failed to marshal AgentConfig: %v", err)
+	}
+
+	var unmarshaledConfig AgentConfig
+	err = json.Unmarshal(data, &unmarshaledConfig)
+	if err != nil {
+		t.Fatalf("Failed to unmarshal AgentConfig: %v", err)
+	}
+
+	if !reflect.DeepEqual(config, unmarshaledConfig) {
+		t.Errorf("Unmarshaled config does not match original. Got %+v, want %+v", unmarshaledConfig, config)
+	}
+}

agent/cvms/api/grpc/client.go

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

agent/cvms/api/grpc/client_test.go

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

agent/cvms/api/grpc/doc.go

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

agent/cvms/api/grpc/server.go

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

agent/cvms/api/grpc/server_test.go

@@ -0,0 +1,273 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package grpc
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/absmach/supermq/pkg/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent/cvms"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/peer"
+)
+
+type mockServerStream struct {
+	mock.Mock
+	cvms.Service_ProcessServer
+}
+
+func (m *mockServerStream) Send(msg *cvms.ServerStreamMessage) error {
+	args := m.Called(msg)
+	return args.Error(0)
+}
+
+func (m *mockServerStream) Recv() (*cvms.ClientStreamMessage, error) {
+	args := m.Called()
+	return args.Get(0).(*cvms.ClientStreamMessage), args.Error(1)
+}
+
+func (m *mockServerStream) Context() context.Context {
+	args := m.Called()
+	return args.Get(0).(context.Context)
+}
+
+type mockService struct {
+	mock.Mock
+}
+
+func (m *mockService) Run(ctx context.Context, ipAddress string, sendMessage SendFunc, authInfo credentials.AuthInfo) {
+	m.Called(ctx, ipAddress, sendMessage, authInfo)
+}
+
+func TestNewServer(t *testing.T) {
+	incoming := make(chan *cvms.ClientStreamMessage)
+	mockSvc := new(mockService)
+
+	server := NewServer(incoming, mockSvc)
+
+	assert.NotNil(t, server)
+	assert.IsType(t, &grpcServer{}, server)
+}
+
+func TestGrpcServer_Process(t *testing.T) {
+	tests := []struct {
+		name          string
+		recvReturn    *cvms.ClientStreamMessage
+		recvError     error
+		expectedError string
+	}{
+		{
+			name:          "Process with context deadline exceeded",
+			recvReturn:    &cvms.ClientStreamMessage{},
+			recvError:     nil,
+			expectedError: "context deadline exceeded",
+		},
+		{
+			name:          "Process with Recv error",
+			recvReturn:    &cvms.ClientStreamMessage{},
+			recvError:     errors.New("recv error"),
+			expectedError: "recv error",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			incoming := make(chan *cvms.ClientStreamMessage, 1)
+			mockSvc := new(mockService)
+			server := NewServer(incoming, mockSvc).(*grpcServer)
+
+			mockStream := new(mockServerStream)
+			ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+			defer cancel()
+
+			mockStream.On("Context").Return(peer.NewContext(ctx, &peer.Peer{
+				Addr:     mockAddr{},
+				AuthInfo: mockAuthInfo{},
+			}))
+
+			if tt.recvError == nil {
+				go func() {
+					for mes := range incoming {
+						assert.NotNil(t, mes)
+					}
+				}()
+			}
+
+			mockStream.On("Recv").Return(tt.recvReturn, tt.recvError)
+			mockSvc.On("Run", mock.Anything, "test", mock.Anything, mock.AnythingOfType("mockAuthInfo")).Return()
+
+			err := server.Process(mockStream)
+
+			assert.Error(t, err)
+			assert.Contains(t, err.Error(), tt.expectedError)
+			mockStream.AssertExpectations(t)
+			mockSvc.AssertExpectations(t)
+		})
+	}
+}
+
+func TestGrpcServer_sendRunReqInChunks(t *testing.T) {
+	incoming := make(chan *cvms.ClientStreamMessage)
+	mockSvc := new(mockService)
+	server := NewServer(incoming, mockSvc).(*grpcServer)
+
+	mockStream := new(mockServerStream)
+
+	runReq := &cvms.ComputationRunReq{
+		Id: "test-id",
+	}
+
+	largePayload := make([]byte, bufferSize*2)
+	for i := range largePayload {
+		largePayload[i] = byte(i % 256)
+	}
+	runReq.Algorithm = &cvms.Algorithm{}
+	runReq.Algorithm.UserKey = largePayload
+
+	mockStream.On("Send", mock.AnythingOfType("*cvms.ServerStreamMessage")).Return(nil).Times(4)
+
+	err := server.sendRunReqInChunks(mockStream, runReq)
+
+	assert.NoError(t, err)
+	mockStream.AssertExpectations(t)
+
+	calls := mockStream.Calls
+	assert.Equal(t, 4, len(calls))
+
+	for i, call := range calls {
+		msg := call.Arguments[0].(*cvms.ServerStreamMessage)
+		chunk := msg.GetRunReqChunks()
+
+		assert.NotNil(t, chunk)
+		assert.Equal(t, "test-id", chunk.Id)
+
+		if i < 3 {
+			assert.False(t, chunk.IsLast)
+		} else {
+			assert.Equal(t, 0, len(chunk.Data))
+			assert.True(t, chunk.IsLast)
+		}
+	}
+}
+
+type mockAddr struct{}
+
+func (mockAddr) Network() string { return "test network" }
+func (mockAddr) String() string  { return "test" }
+
+type mockAuthInfo struct{}
+
+func (mockAuthInfo) AuthType() string { return "test auth" }
+
+func TestGrpcServer_ProcessWithMockService(t *testing.T) {
+	tests := []struct {
+		name        string
+		setupMockFn func(*mockService, *mockServerStream)
+	}{
+		{
+			name: "Run Request Test",
+			setupMockFn: func(mockSvc *mockService, mockStream *mockServerStream) {
+				mockSvc.On("Run", mock.Anything, "test", mock.Anything, mock.AnythingOfType("mockAuthInfo")).
+					Run(func(args mock.Arguments) {
+						sendFunc := args.Get(2).(SendFunc)
+						runReq := &cvms.ComputationRunReq{Id: "test-run-id"}
+						err := sendFunc(&cvms.ServerStreamMessage{
+							Message: &cvms.ServerStreamMessage_RunReq{
+								RunReq: runReq,
+							},
+						})
+						assert.NoError(t, err)
+					}).
+					Return()
+
+				mockStream.On("Send", mock.MatchedBy(func(msg *cvms.ServerStreamMessage) bool {
+					chunks := msg.GetRunReqChunks()
+					return chunks != nil && chunks.Id == "test-run-id"
+				})).Return(nil)
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			incoming := make(chan *cvms.ClientStreamMessage, 10)
+			mockSvc := new(mockService)
+			server := NewServer(incoming, mockSvc).(*grpcServer)
+
+			go func() {
+				for mes := range incoming {
+					assert.NotNil(t, mes)
+				}
+			}()
+
+			mockStream := new(mockServerStream)
+			ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+			defer cancel()
+
+			peerCtx := peer.NewContext(ctx, &peer.Peer{
+				Addr:     mockAddr{},
+				AuthInfo: mockAuthInfo{},
+			})
+
+			mockStream.On("Context").Return(peerCtx)
+			mockStream.On("Recv").Return(&cvms.ClientStreamMessage{}, nil).Maybe()
+
+			tt.setupMockFn(mockSvc, mockStream)
+
+			go func() {
+				time.Sleep(150 * time.Millisecond)
+				cancel()
+			}()
+
+			err := server.Process(mockStream)
+
+			assert.Error(t, err)
+			assert.Contains(t, err.Error(), "context canceled")
+			mockStream.AssertExpectations(t)
+			mockSvc.AssertExpectations(t)
+		})
+	}
+}
+
+func TestGrpcServer_sendRunReqInChunksError(t *testing.T) {
+	incoming := make(chan *cvms.ClientStreamMessage)
+	mockSvc := new(mockService)
+	server := NewServer(incoming, mockSvc).(*grpcServer)
+
+	mockStream := new(mockServerStream)
+
+	runReq := &cvms.ComputationRunReq{
+		Id: "test-id",
+	}
+
+	// Simulate an error when sending
+	mockStream.On("Send", mock.AnythingOfType("*cvms.ServerStreamMessage")).Return(errors.New("send error")).Once()
+
+	err := server.sendRunReqInChunks(mockStream, runReq)
+
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "send error")
+	mockStream.AssertExpectations(t)
+}
+
+func TestGrpcServer_ProcessMissingPeerInfo(t *testing.T) {
+	incoming := make(chan *cvms.ClientStreamMessage)
+	mockSvc := new(mockService)
+	server := NewServer(incoming, mockSvc).(*grpcServer)
+
+	mockStream := new(mockServerStream)
+	ctx := context.Background()
+
+	// Return a context without peer info
+	mockStream.On("Context").Return(ctx)
+
+	err := server.Process(mockStream)
+
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to get peer info")
+	mockStream.AssertExpectations(t)
+}

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

@@ -0,0 +1,242 @@
+// 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 (
+	mock "github.com/stretchr/testify/mock"
+	"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
+}
+
+type Storage_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *Storage) EXPECT() *Storage_Expecter {
+	return &Storage_Expecter{mock: &_m.Mock}
+}
+
+// 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 returnFunc, ok := ret.Get(0).(func(*cvms.ClientStreamMessage) error); ok {
+		r0 = returnFunc(msg)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// Storage_Add_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Add'
+type Storage_Add_Call struct {
+	*mock.Call
+}
+
+// Add is a helper method to define mock.On call
+//   - msg *cvms.ClientStreamMessage
+func (_e *Storage_Expecter) Add(msg interface{}) *Storage_Add_Call {
+	return &Storage_Add_Call{Call: _e.mock.On("Add", msg)}
+}
+
+func (_c *Storage_Add_Call) Run(run func(msg *cvms.ClientStreamMessage)) *Storage_Add_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 *cvms.ClientStreamMessage
+		if args[0] != nil {
+			arg0 = args[0].(*cvms.ClientStreamMessage)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+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(msg *cvms.ClientStreamMessage) error) *Storage_Add_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// 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 returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// Storage_Clear_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Clear'
+type Storage_Clear_Call struct {
+	*mock.Call
+}
+
+// Clear is a helper method to define mock.On call
+func (_e *Storage_Expecter) Clear() *Storage_Clear_Call {
+	return &Storage_Clear_Call{Call: _e.mock.On("Clear")}
+}
+
+func (_c *Storage_Clear_Call) Run(run func()) *Storage_Clear_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *Storage_Clear_Call) Return(err error) *Storage_Clear_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *Storage_Clear_Call) RunAndReturn(run func() error) *Storage_Clear_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Load provides a mock function 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")
+	}
+
+	var r0 []storage.Message
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func() ([]storage.Message, error)); ok {
+		return returnFunc()
+	}
+	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 returnFunc, ok := ret.Get(1).(func() error); ok {
+		r1 = returnFunc()
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// Storage_Load_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Load'
+type Storage_Load_Call struct {
+	*mock.Call
+}
+
+// Load is a helper method to define mock.On call
+func (_e *Storage_Expecter) Load() *Storage_Load_Call {
+	return &Storage_Load_Call{Call: _e.mock.On("Load")}
+}
+
+func (_c *Storage_Load_Call) Run(run func()) *Storage_Load_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *Storage_Load_Call) Return(messages []storage.Message, err error) *Storage_Load_Call {
+	_c.Call.Return(messages, err)
+	return _c
+}
+
+func (_c *Storage_Load_Call) RunAndReturn(run func() ([]storage.Message, error)) *Storage_Load_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Save provides a mock function 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 returnFunc, ok := ret.Get(0).(func([]storage.Message) error); ok {
+		r0 = returnFunc(messages)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// Storage_Save_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Save'
+type Storage_Save_Call struct {
+	*mock.Call
+}
+
+// Save is a helper method to define mock.On call
+//   - messages []storage.Message
+func (_e *Storage_Expecter) Save(messages interface{}) *Storage_Save_Call {
+	return &Storage_Save_Call{Call: _e.mock.On("Save", messages)}
+}
+
+func (_c *Storage_Save_Call) Run(run func(messages []storage.Message)) *Storage_Save_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 []storage.Message
+		if args[0] != nil {
+			arg0 = args[0].([]storage.Message)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+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(messages []storage.Message) error) *Storage_Save_Call {
+	_c.Call.Return(run)
+	return _c
+}

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

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

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

@@ -0,0 +1,450 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package storage
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/agent/cvms"
+)
+
+func createTempDir(t *testing.T) string {
+	tmpDir, err := os.MkdirTemp("", "storage_test_*")
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		os.RemoveAll(tmpDir)
+	})
+	return tmpDir
+}
+
+func createTestMessage(content string) *cvms.ClientStreamMessage {
+	return &cvms.ClientStreamMessage{
+		Message: &cvms.ClientStreamMessage_RunRes{
+			RunRes: &cvms.RunResponse{
+				Error:         "",
+				ComputationId: content,
+			},
+		},
+	}
+}
+
+func TestNewFileStorage(t *testing.T) {
+	tests := []struct {
+		name        string
+		storageDir  string
+		expectError bool
+	}{
+		{
+			name:        "valid directory",
+			storageDir:  createTempDir(t),
+			expectError: false,
+		},
+		{
+			name:        "non-existent directory gets created",
+			storageDir:  filepath.Join(createTempDir(t), "subdir"),
+			expectError: false,
+		},
+		{
+			name:        "invalid directory path",
+			storageDir:  "/invalid/path/that/cannot/be/created",
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			storage, err := NewFileStorage(tt.storageDir)
+
+			if tt.expectError {
+				assert.Error(t, err)
+				assert.Nil(t, storage)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, storage)
+				assert.Equal(t, filepath.Join(tt.storageDir, "pending_messages.json"), storage.path)
+				assert.Empty(t, storage.msgs)
+			}
+		})
+	}
+}
+
+func TestFileStorage_Load(t *testing.T) {
+	tests := []struct {
+		name         string
+		setupFile    func(string) error
+		expectedMsgs int
+		expectError  bool
+	}{
+		{
+			name: "load from non-existent file",
+			setupFile: func(path string) error {
+				// Don't create file
+				return nil
+			},
+			expectedMsgs: 0,
+			expectError:  false,
+		},
+		{
+			name: "load from empty file",
+			setupFile: func(path string) error {
+				return os.WriteFile(path, []byte("[]"), 0o644)
+			},
+			expectedMsgs: 0,
+			expectError:  false,
+		},
+		{
+			name: "load from corrupted file",
+			setupFile: func(path string) error {
+				return os.WriteFile(path, []byte("invalid json"), 0o644)
+			},
+			expectedMsgs: 0,
+			expectError:  true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tmpDir := createTempDir(t)
+			storage, err := NewFileStorage(tmpDir)
+			require.NoError(t, err)
+
+			err = tt.setupFile(storage.path)
+			require.NoError(t, err)
+
+			msgs, err := storage.Load()
+
+			if tt.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Len(t, msgs, tt.expectedMsgs)
+			}
+		})
+	}
+}
+
+func TestFileStorage_Save(t *testing.T) {
+	tests := []struct {
+		name        string
+		messages    []Message
+		expectError bool
+	}{
+		{
+			name:        "save empty messages",
+			messages:    []Message{},
+			expectError: false,
+		},
+		{
+			name: "save single message",
+			messages: []Message{
+				{
+					Message: createTestMessage("test"),
+					Time:    time.Now(),
+				},
+			},
+			expectError: false,
+		},
+		{
+			name: "save multiple messages",
+			messages: []Message{
+				{
+					Message: createTestMessage("test1"),
+					Time:    time.Now(),
+				},
+				{
+					Message: createTestMessage("test2"),
+					Time:    time.Now().Add(time.Second),
+				},
+			},
+			expectError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tmpDir := createTempDir(t)
+			storage, err := NewFileStorage(tmpDir)
+			require.NoError(t, err)
+
+			err = storage.Save(tt.messages)
+
+			if tt.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+
+				// Verify file was written correctly
+				_, err := os.ReadFile(storage.path)
+				assert.NoError(t, err)
+
+				// Verify internal state was updated
+				assert.Equal(t, tt.messages, storage.msgs)
+			}
+		})
+	}
+}
+
+func TestFileStorage_Add(t *testing.T) {
+	tests := []struct {
+		name          string
+		initialMsgs   []Message
+		newMessage    *cvms.ClientStreamMessage
+		expectError   bool
+		expectedCount int
+	}{
+		{
+			name:          "add to empty storage",
+			initialMsgs:   []Message{},
+			newMessage:    createTestMessage("new"),
+			expectError:   false,
+			expectedCount: 1,
+		},
+		{
+			name: "add to existing messages",
+			initialMsgs: []Message{
+				{
+					Message: createTestMessage("existing"),
+					Time:    time.Now(),
+				},
+			},
+			newMessage:    createTestMessage("new"),
+			expectError:   false,
+			expectedCount: 2,
+		},
+		{
+			name:          "add nil message",
+			initialMsgs:   []Message{},
+			newMessage:    nil,
+			expectError:   false,
+			expectedCount: 1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tmpDir := createTempDir(t)
+			storage, err := NewFileStorage(tmpDir)
+			require.NoError(t, err)
+
+			// Setup initial messages
+			if len(tt.initialMsgs) > 0 {
+				err = storage.Save(tt.initialMsgs)
+				require.NoError(t, err)
+			}
+
+			beforeTime := time.Now()
+			err = storage.Add(tt.newMessage)
+			afterTime := time.Now()
+
+			if tt.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+
+				// Verify message was added to internal state
+				assert.Len(t, storage.msgs, tt.expectedCount)
+
+				// Verify timestamp is reasonable
+				if tt.expectedCount > 0 {
+					lastMsg := storage.msgs[len(storage.msgs)-1]
+					assert.True(t, lastMsg.Time.After(beforeTime) || lastMsg.Time.Equal(beforeTime))
+					assert.True(t, lastMsg.Time.Before(afterTime) || lastMsg.Time.Equal(afterTime))
+					assert.Equal(t, tt.newMessage, lastMsg.Message)
+				}
+
+				_, err := os.ReadFile(storage.path)
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
+
+func TestFileStorage_Clear(t *testing.T) {
+	tests := []struct {
+		name        string
+		initialMsgs []Message
+		expectError bool
+	}{
+		{
+			name:        "clear empty storage",
+			initialMsgs: []Message{},
+			expectError: false,
+		},
+		{
+			name: "clear storage with messages",
+			initialMsgs: []Message{
+				{
+					Message: createTestMessage("test1"),
+					Time:    time.Now(),
+				},
+				{
+					Message: createTestMessage("test2"),
+					Time:    time.Now(),
+				},
+			},
+			expectError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tmpDir := createTempDir(t)
+			storage, err := NewFileStorage(tmpDir)
+			require.NoError(t, err)
+
+			// Setup initial messages
+			if len(tt.initialMsgs) > 0 {
+				err = storage.Save(tt.initialMsgs)
+				require.NoError(t, err)
+			}
+
+			err = storage.Clear()
+
+			if tt.expectError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+
+				// Verify internal state is cleared
+				assert.Empty(t, storage.msgs)
+
+				// Verify file contains empty array
+				data, err := os.ReadFile(storage.path)
+				assert.NoError(t, err)
+				assert.Equal(t, "[]", string(data))
+			}
+		})
+	}
+}
+
+func TestFileStorage_ConcurrentAccess(t *testing.T) {
+	tmpDir := createTempDir(t)
+	storage, err := NewFileStorage(tmpDir)
+	require.NoError(t, err)
+
+	// Test concurrent Add operations
+	numGoroutines := 10
+	done := make(chan bool, numGoroutines)
+
+	for i := 0; i < numGoroutines; i++ {
+		go func(id int) {
+			defer func() { done <- true }()
+
+			msg := createTestMessage(string(rune('A' + id)))
+			err := storage.Add(msg)
+			assert.NoError(t, err)
+		}(i)
+	}
+
+	// Wait for all goroutines to complete
+	for i := 0; i < numGoroutines; i++ {
+		<-done
+	}
+
+	// Verify all messages were added
+	msgs, err := storage.Load()
+	assert.NoError(t, err)
+	assert.Len(t, msgs, numGoroutines)
+}
+
+func TestFileStorage_IntegrationFlow(t *testing.T) {
+	tmpDir := createTempDir(t)
+	storage, err := NewFileStorage(tmpDir)
+	require.NoError(t, err)
+
+	// Test full workflow
+
+	// 1. Load from empty storage
+	msgs, err := storage.Load()
+	assert.NoError(t, err)
+	assert.Empty(t, msgs)
+
+	// 2. Add some messages
+	msg1 := createTestMessage("message1")
+	err = storage.Add(msg1)
+	assert.NoError(t, err)
+
+	msg2 := createTestMessage("message2")
+	err = storage.Add(msg2)
+	assert.NoError(t, err)
+
+	// 3. Load and verify
+	msgs, err = storage.Load()
+	assert.NoError(t, err)
+	assert.Len(t, msgs, 2)
+
+	// 4. Save new set of messages
+	newMsgs := []Message{
+		{
+			Message: createTestMessage("new1"),
+			Time:    time.Now(),
+		},
+	}
+	err = storage.Save(newMsgs)
+	assert.NoError(t, err)
+
+	// 5. Load and verify replacement
+	msgs, err = storage.Load()
+	assert.NoError(t, err)
+	assert.Len(t, msgs, 1)
+
+	// 6. Clear storage
+	err = storage.Clear()
+	assert.NoError(t, err)
+
+	// 7. Verify empty
+	msgs, err = storage.Load()
+	assert.NoError(t, err)
+	assert.Empty(t, msgs)
+}
+
+func TestFileStorage_FilePermissions(t *testing.T) {
+	tmpDir := createTempDir(t)
+	storage, err := NewFileStorage(tmpDir)
+	require.NoError(t, err)
+
+	// Add a message to create the file
+	msg := createTestMessage("test")
+	err = storage.Add(msg)
+	assert.NoError(t, err)
+
+	// Check file permissions
+	info, err := os.Stat(storage.path)
+	assert.NoError(t, err)
+	assert.Equal(t, os.FileMode(0o644), info.Mode().Perm())
+}
+
+func TestFileStorage_ErrorHandling(t *testing.T) {
+	tmpDir := createTempDir(t)
+	storage, err := NewFileStorage(tmpDir)
+	require.NoError(t, err)
+
+	// Make directory read-only to trigger write errors
+	err = os.Chmod(tmpDir, 0o555)
+	require.NoError(t, err)
+
+	// Restore permissions for cleanup
+	t.Cleanup(func() {
+		if err := os.Chmod(tmpDir, 0o755); err != nil {
+			t.Errorf("Failed to restore permissions: %v", err)
+		}
+	})
+
+	// Try to add a message - should fail due to write permissions
+	msg := createTestMessage("test")
+	err = storage.Add(msg)
+	assert.Error(t, err)
+
+	// Try to save - should fail due to write permissions
+	err = storage.Save([]Message{})
+	assert.Error(t, err)
+
+	// Try to clear - should fail due to write permissions
+	err = storage.Clear()
+	assert.Error(t, err)
+}

agent/cvms/cvms.proto

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

agent/cvms/cvms_grpc.pb.go

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

agent/cvms/server/cvm.go

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

agent/cvms/server/cvm_test.go

@@ -0,0 +1,530 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package server
+
+import (
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/x509"
+	"log/slog"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/agent/mocks"
+)
+
+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"
+
+	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	assert.NoError(t, err, "Failed to generate ECDSA key")
+
+	pubkey, err := x509.MarshalPKIXPublicKey(privateKey.Public())
+	assert.NoError(t, err, "Failed to marshal public key")
+
+	return logger, mockSvc, host, pubkey
+}
+
+func TestNewServer(t *testing.T) {
+	logger, svc, host, _ := setupTest(t)
+
+	tests := []struct {
+		name     string
+		logger   *slog.Logger
+		svc      agent.Service
+		host     string
+		expected AgentServer
+	}{
+		{
+			name:   "valid server creation",
+			logger: logger,
+			svc:    svc,
+			host:   host,
+		},
+		{
+			name:   "server with empty host",
+			logger: logger,
+			svc:    svc,
+			host:   "",
+		},
+		{
+			name:   "server with empty caUrl",
+			logger: logger,
+			svc:    svc,
+			host:   host,
+		},
+		{
+			name:   "server with empty cvmId",
+			logger: logger,
+			svc:    svc,
+			host:   host,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			server := NewServer(tt.logger, tt.svc, tt.host, nil)
+
+			assert.NotNil(t, server)
+
+			agentSrv, ok := server.(*agentServer)
+			assert.True(t, ok)
+			assert.Equal(t, tt.logger, agentSrv.logger)
+			assert.Equal(t, tt.svc, agentSrv.svc)
+			assert.Equal(t, tt.host, agentSrv.host)
+		})
+	}
+}
+
+func TestAgentServer_Start(t *testing.T) {
+	logger, svc, host, pubKey := setupTest(t)
+
+	tests := []struct {
+		name          string
+		cfg           agent.AgentConfig
+		cmp           agent.Computation
+		setupMocks    func(*mocks.Service)
+		expectedError bool
+		errorContains string
+	}{
+		{
+			name: "successful start with default port",
+			cfg: agent.AgentConfig{
+				Port:         "",
+				CertFile:     "cert.pem",
+				KeyFile:      "key.pem",
+				ServerCAFile: "server-ca.pem",
+				ClientCAFile: "client-ca.pem",
+				AttestedTls:  true,
+			},
+			cmp: agent.Computation{
+				ID:          "test-computation-1",
+				Name:        "Test Computation",
+				Description: "A test computation",
+				Algorithm: agent.Algorithm{
+					Hash:    [32]byte{0x01, 0x02, 0x03},
+					UserKey: pubKey,
+				},
+				Datasets: []agent.Dataset{
+					{
+						Hash:    [32]byte{0x04, 0x05, 0x06},
+						UserKey: pubKey,
+					},
+				},
+				ResultConsumers: []agent.ResultConsumer{
+					{
+						UserKey: pubKey,
+					},
+				},
+			},
+			setupMocks: func(m *mocks.Service) {
+			},
+			expectedError: false,
+		},
+		{
+			name: "successful start with custom port",
+			cfg: agent.AgentConfig{
+				Port:         "8080",
+				CertFile:     "cert.pem",
+				KeyFile:      "key.pem",
+				ServerCAFile: "server-ca.pem",
+				ClientCAFile: "client-ca.pem",
+				AttestedTls:  false,
+			},
+			cmp: agent.Computation{
+				ID:          "test-computation-2",
+				Name:        "Test Computation 2",
+				Description: "Another test computation",
+				Algorithm: agent.Algorithm{
+					Hash:    [32]byte{0x07, 0x08, 0x09},
+					UserKey: pubKey,
+				},
+				Datasets: []agent.Dataset{
+					{
+						Hash:    [32]byte{0x0a, 0x0b, 0x0c},
+						UserKey: pubKey,
+					},
+				},
+				ResultConsumers: []agent.ResultConsumer{
+					{
+						UserKey: pubKey,
+					},
+				},
+			},
+			setupMocks: func(m *mocks.Service) {
+			},
+			expectedError: false,
+		},
+		{
+			name: "start with minimal config",
+			cfg: agent.AgentConfig{
+				Port:        "9090",
+				AttestedTls: false,
+			},
+			cmp: agent.Computation{
+				ID:   "test-computation-3",
+				Name: "Minimal Test",
+				Algorithm: agent.Algorithm{
+					Hash:    [32]byte{0x0d, 0x0e, 0x0f},
+					UserKey: pubKey,
+				},
+				Datasets: []agent.Dataset{
+					{
+						Hash:    [32]byte{0x10, 0x11, 0x12},
+						UserKey: pubKey,
+					},
+				},
+				ResultConsumers: []agent.ResultConsumer{
+					{
+						UserKey: pubKey,
+					},
+				},
+			},
+			setupMocks: func(m *mocks.Service) {
+			},
+			expectedError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.setupMocks(svc)
+
+			server := NewServer(logger, svc, host, nil)
+
+			err := server.Start(tt.cfg, tt.cmp)
+
+			if tt.expectedError {
+				assert.Error(t, err)
+				if tt.errorContains != "" {
+					assert.Contains(t, err.Error(), tt.errorContains)
+				}
+			} else {
+				assert.NoError(t, err)
+
+				// Verify the port was set correctly
+				agentSrv := server.(*agentServer)
+				assert.NotNil(t, agentSrv.gs)
+
+				if err := server.Stop(); err != nil {
+					t.Fatalf("Failed to stop server after start: %v", err)
+				}
+			}
+
+			svc.AssertExpectations(t)
+		})
+	}
+}
+
+func TestAgentServer_Stop(t *testing.T) {
+	logger, svc, host, pubKey := setupTest(t)
+
+	tests := []struct {
+		name          string
+		setupServer   func(AgentServer) error
+		expectedError bool
+		errorContains string
+	}{
+		{
+			name: "stop unstarted server",
+			setupServer: func(server AgentServer) error {
+				// Don't start the server
+				return nil
+			},
+			expectedError: false,
+		},
+		{
+			name: "stop started server",
+			setupServer: func(server AgentServer) error {
+				cfg := agent.AgentConfig{
+					Port: "7004",
+				}
+				cmp := agent.Computation{
+					ID:   "test-stop-computation",
+					Name: "Stop Test",
+					Algorithm: agent.Algorithm{
+						Hash:    [32]byte{0x19, 0x1a, 0x1b},
+						UserKey: pubKey,
+					},
+					Datasets: []agent.Dataset{
+						{
+							Hash:    [32]byte{0x1c, 0x1d, 0x1e},
+							UserKey: pubKey,
+						},
+					},
+					ResultConsumers: []agent.ResultConsumer{
+						{
+							UserKey: pubKey,
+						},
+					},
+				}
+				return server.Start(cfg, cmp)
+			},
+			expectedError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			server := NewServer(logger, svc, host, nil)
+
+			err := tt.setupServer(server)
+			if err != nil {
+				t.Fatalf("Setup failed: %v", err)
+			}
+
+			// Give the server a moment to start if it was started
+			time.Sleep(10 * time.Millisecond)
+
+			err = server.Stop()
+
+			if tt.expectedError {
+				assert.Error(t, err)
+				if tt.errorContains != "" {
+					assert.Contains(t, err.Error(), tt.errorContains)
+				}
+			} else {
+				assert.NoError(t, err)
+			}
+
+			svc.AssertExpectations(t)
+		})
+	}
+}
+
+func TestAgentServer_StopMultipleTimes(t *testing.T) {
+	logger, svc, host, pubKey := setupTest(t)
+	server := NewServer(logger, svc, host, nil)
+
+	// Start the server
+	cfg := agent.AgentConfig{Port: "7005"}
+	cmp := agent.Computation{
+		ID:   "test-multiple-stop",
+		Name: "Multiple Stop Test",
+		Algorithm: agent.Algorithm{
+			Hash:    [32]byte{0x1f, 0x20, 0x21},
+			UserKey: pubKey,
+		},
+		Datasets: []agent.Dataset{
+			{
+				Hash:    [32]byte{0x22, 0x23, 0x24},
+				UserKey: pubKey,
+			},
+		},
+		ResultConsumers: []agent.ResultConsumer{
+			{
+				UserKey: pubKey,
+			},
+		},
+	}
+
+	err := server.Start(cfg, cmp)
+	assert.NoError(t, err)
+
+	// Give the server a moment to start
+	time.Sleep(10 * time.Millisecond)
+
+	// Stop the server multiple times
+	err1 := server.Stop()
+	err2 := server.Stop()
+	err3 := server.Stop()
+
+	assert.NoError(t, err1)
+	assert.NoError(t, err2)
+	assert.NoError(t, err3)
+
+	svc.AssertExpectations(t)
+}
+
+func TestAgentServer_StartAfterStop(t *testing.T) {
+	logger, svc, host, pubKey := setupTest(t)
+	server := NewServer(logger, svc, host, nil)
+
+	cfg := agent.AgentConfig{Port: "7006"}
+	cmp := agent.Computation{
+		ID:   "test-restart",
+		Name: "Restart Test",
+		Algorithm: agent.Algorithm{
+			Hash:    [32]byte{0x25, 0x26, 0x27},
+			UserKey: pubKey,
+		},
+		Datasets: []agent.Dataset{
+			{
+				Hash:    [32]byte{0x28, 0x29, 0x2a},
+				UserKey: pubKey,
+			},
+		},
+		ResultConsumers: []agent.ResultConsumer{
+			{
+				UserKey: pubKey,
+			},
+		},
+	}
+
+	// Start, stop, then start again
+	err := server.Start(cfg, cmp)
+	assert.NoError(t, err)
+
+	time.Sleep(10 * time.Millisecond)
+
+	err = server.Stop()
+	assert.NoError(t, err)
+
+	// Start again with different config
+	cfg2 := agent.AgentConfig{Port: "7007"}
+	cmp2 := agent.Computation{
+		ID:   "test-restart-2",
+		Name: "Restart Test 2",
+		Algorithm: agent.Algorithm{
+			Hash:    [32]byte{0x2b, 0x2c, 0x2d},
+			UserKey: pubKey,
+		},
+		Datasets: []agent.Dataset{
+			{
+				Hash:    [32]byte{0x2e, 0x2f, 0x30},
+				UserKey: pubKey,
+			},
+		},
+		ResultConsumers: []agent.ResultConsumer{
+			{
+				UserKey: pubKey,
+			},
+		},
+	}
+
+	err = server.Start(cfg2, cmp2)
+	assert.NoError(t, err)
+
+	time.Sleep(10 * time.Millisecond)
+
+	err = server.Stop()
+	assert.NoError(t, err)
+
+	svc.AssertExpectations(t)
+}
+
+func TestAgentServer_ConfigValidation(t *testing.T) {
+	logger, svc, host, pubKey := setupTest(t)
+
+	tests := []struct {
+		name   string
+		config agent.AgentConfig
+		cmp    agent.Computation
+		valid  bool
+	}{
+		{
+			name: "valid config with all fields",
+			config: agent.AgentConfig{
+				Port:         "8080",
+				CertFile:     "cert.pem",
+				KeyFile:      "key.pem",
+				ServerCAFile: "server-ca.pem",
+				ClientCAFile: "client-ca.pem",
+				AttestedTls:  true,
+			},
+			cmp: agent.Computation{
+				ID:   "valid-config-test",
+				Name: "Valid Config Test",
+				Algorithm: agent.Algorithm{
+					Hash:    [32]byte{0x31, 0x32, 0x33},
+					UserKey: pubKey,
+				},
+				Datasets: []agent.Dataset{
+					{
+						Hash:    [32]byte{0x34, 0x35, 0x36},
+						UserKey: pubKey,
+					},
+				},
+				ResultConsumers: []agent.ResultConsumer{
+					{
+						UserKey: pubKey,
+					},
+				},
+			},
+			valid: true,
+		},
+		{
+			name: "valid config with minimal fields",
+			config: agent.AgentConfig{
+				Port: "9090",
+			},
+			cmp: agent.Computation{
+				ID:   "minimal-config-test",
+				Name: "Minimal Config Test",
+				Algorithm: agent.Algorithm{
+					Hash:    [32]byte{0x37, 0x38, 0x39},
+					UserKey: pubKey,
+				},
+				Datasets: []agent.Dataset{
+					{
+						Hash:    [32]byte{0x3a, 0x3b, 0x3c},
+						UserKey: pubKey,
+					},
+				},
+				ResultConsumers: []agent.ResultConsumer{
+					{
+						UserKey: pubKey,
+					},
+				},
+			},
+			valid: true,
+		},
+		{
+			name: "config with empty port uses default",
+			config: agent.AgentConfig{
+				Port: "",
+			},
+			cmp: agent.Computation{
+				ID:        "default-port-test",
+				Name:      "Default Port Test",
+				Algorithm: agent.Algorithm{Hash: [32]byte{0x3d, 0x3e, 0x3f}, UserKey: pubKey},
+				Datasets: []agent.Dataset{
+					{Hash: [32]byte{0x40, 0x41, 0x42}, UserKey: pubKey},
+				},
+				ResultConsumers: []agent.ResultConsumer{
+					{UserKey: pubKey},
+				},
+			},
+			valid: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			server := NewServer(logger, svc, host, nil)
+
+			err := server.Start(tt.config, tt.cmp)
+
+			if tt.valid {
+				assert.NoError(t, err)
+
+				// Verify default port is used when empty
+				if tt.config.Port == "" {
+					agentSrv := server.(*agentServer)
+					assert.NotNil(t, agentSrv.gs)
+				}
+
+				time.Sleep(10 * time.Millisecond)
+				if err := server.Stop(); err != nil {
+					t.Fatalf("Failed to stop server after start: %v", err)
+				}
+			} else {
+				assert.Error(t, err)
+			}
+
+			svc.AssertExpectations(t)
+		})
+	}
+}
+
+func TestConstants(t *testing.T) {
+	assert.Equal(t, "agent", svcName)
+	assert.Equal(t, "7002", defSvcGRPCPort)
+}

agent/cvms/server/mocks/agentserver.go

@@ -0,0 +1,141 @@
+// 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 (
+	mock "github.com/stretchr/testify/mock"
+	"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
+}
+
+type AgentServer_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *AgentServer) EXPECT() *AgentServer_Expecter {
+	return &AgentServer_Expecter{mock: &_m.Mock}
+}
+
+// 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 returnFunc, ok := ret.Get(0).(func(agent.AgentConfig, agent.Computation) error); ok {
+		r0 = returnFunc(cfg, cmp)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentServer_Start_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Start'
+type AgentServer_Start_Call struct {
+	*mock.Call
+}
+
+// Start is a helper method to define mock.On call
+//   - cfg agent.AgentConfig
+//   - cmp agent.Computation
+func (_e *AgentServer_Expecter) Start(cfg interface{}, cmp interface{}) *AgentServer_Start_Call {
+	return &AgentServer_Start_Call{Call: _e.mock.On("Start", cfg, cmp)}
+}
+
+func (_c *AgentServer_Start_Call) Run(run func(cfg agent.AgentConfig, cmp agent.Computation)) *AgentServer_Start_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		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(err error) *AgentServer_Start_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+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 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 returnFunc, ok := ret.Get(0).(func() error); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// AgentServer_Stop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Stop'
+type AgentServer_Stop_Call struct {
+	*mock.Call
+}
+
+// Stop is a helper method to define mock.On call
+func (_e *AgentServer_Expecter) Stop() *AgentServer_Stop_Call {
+	return &AgentServer_Stop_Call{Call: _e.mock.On("Stop")}
+}
+
+func (_c *AgentServer_Stop_Call) Run(run func()) *AgentServer_Stop_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *AgentServer_Stop_Call) Return(err error) *AgentServer_Stop_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+func (_c *AgentServer_Stop_Call) RunAndReturn(run func() error) *AgentServer_Stop_Call {
+	_c.Call.Return(run)
+	return _c
+}

agent/events/events.go

@@ -4,66 +4,38 @@ package events
 
 import (
 	"encoding/json"
-	"time"
 
-	"github.com/mdlayher/vsock"
-	"github.com/ultravioletrs/cocos/pkg/manager"
-	"google.golang.org/protobuf/proto"
+	"github.com/ultravioletrs/cocos/agent/cvms"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 type service struct {
-	service       string
-	computationID string
-	conn          *vsock.Conn
+	service string
+	queue   chan *cvms.ClientStreamMessage
 }
 
-type AgentEvent struct {
-	EventType     string          `json:"event_type"`
-	Timestamp     time.Time       `json:"timestamp"`
-	ComputationID string          `json:"computation_id,omitempty"`
-	Details       json.RawMessage `json:"details,omitempty"`
-	Originator    string          `json:"originator"`
-	Status        string          `json:"status,omitempty"`
-}
-
-//go:generate mockery --name Service --output=./mocks --filename events.go --quiet --note "Copyright (c) Ultraviolet \n // SPDX-License-Identifier: Apache-2.0"
 type Service interface {
-	SendEvent(event, status string, details json.RawMessage) error
-	Close() error
+	SendEvent(cmpID, event, status string, details json.RawMessage)
 }
 
-func New(svc, computationID string, sockPort uint32) (Service, error) {
-	conn, err := vsock.Dial(vsock.Host, sockPort, nil)
-	if err != nil {
-		return nil, err
-	}
+func New(svc string, queue chan *cvms.ClientStreamMessage) (Service, error) {
 	return &service{
-		service:       svc,
-		computationID: computationID,
-		conn:          conn,
+		service: svc,
+		queue:   queue,
 	}, nil
 }
 
-func (s *service) SendEvent(event, status string, details json.RawMessage) error {
-	body := manager.ClientStreamMessage{Message: &manager.ClientStreamMessage_AgentEvent{AgentEvent: &manager.AgentEvent{
-		EventType:     event,
-		Timestamp:     timestamppb.Now(),
-		ComputationId: s.computationID,
-		Originator:    s.service,
-		Status:        status,
-		Details:       details,
-	}}}
-	protoBody, err := proto.Marshal(&body)
-	if err != nil {
-		return err
+func (s *service) SendEvent(cmpID, event, status string, details json.RawMessage) {
+	s.queue <- &cvms.ClientStreamMessage{
+		Message: &cvms.ClientStreamMessage_AgentEvent{
+			AgentEvent: &cvms.AgentEvent{
+				EventType:     event,
+				Timestamp:     timestamppb.Now(),
+				ComputationId: cmpID,
+				Originator:    s.service,
+				Status:        status,
+				Details:       details,
+			},
+		},
 	}
-	if _, err := s.conn.Write(protoBody); err != nil {
-		return err
-	}
-	return nil
-}
-
-func (s *service) Close() error {
-	return s.conn.Close()
 }

agent/events/events.pb.go

@@ -0,0 +1,348 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by protoc-gen-go. DO NOT EDIT.
+// versions:
+// 	protoc-gen-go v1.36.8
+// 	protoc        v5.29.0
+// source: agent/events/events.proto
+
+package events
+
+import (
+	protoreflect "google.golang.org/protobuf/reflect/protoreflect"
+	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
+	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
+	reflect "reflect"
+	sync "sync"
+	unsafe "unsafe"
+)
+
+const (
+	// Verify that this generated code is sufficiently up-to-date.
+	_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
+	// Verify that runtime/protoimpl is sufficiently up-to-date.
+	_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
+)
+
+type AgentEvent struct {
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	EventType     string                 `protobuf:"bytes,1,opt,name=event_type,json=eventType,proto3" json:"event_type,omitempty"`
+	Timestamp     *timestamppb.Timestamp `protobuf:"bytes,2,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
+	ComputationId string                 `protobuf:"bytes,3,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
+	Details       []byte                 `protobuf:"bytes,4,opt,name=details,proto3" json:"details,omitempty"`
+	Originator    string                 `protobuf:"bytes,5,opt,name=originator,proto3" json:"originator,omitempty"`
+	Status        string                 `protobuf:"bytes,6,opt,name=status,proto3" json:"status,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *AgentEvent) Reset() {
+	*x = AgentEvent{}
+	mi := &file_agent_events_events_proto_msgTypes[0]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *AgentEvent) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*AgentEvent) ProtoMessage() {}
+
+func (x *AgentEvent) ProtoReflect() protoreflect.Message {
+	mi := &file_agent_events_events_proto_msgTypes[0]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use AgentEvent.ProtoReflect.Descriptor instead.
+func (*AgentEvent) Descriptor() ([]byte, []int) {
+	return file_agent_events_events_proto_rawDescGZIP(), []int{0}
+}
+
+func (x *AgentEvent) GetEventType() string {
+	if x != nil {
+		return x.EventType
+	}
+	return ""
+}
+
+func (x *AgentEvent) GetTimestamp() *timestamppb.Timestamp {
+	if x != nil {
+		return x.Timestamp
+	}
+	return nil
+}
+
+func (x *AgentEvent) GetComputationId() string {
+	if x != nil {
+		return x.ComputationId
+	}
+	return ""
+}
+
+func (x *AgentEvent) GetDetails() []byte {
+	if x != nil {
+		return x.Details
+	}
+	return nil
+}
+
+func (x *AgentEvent) GetOriginator() string {
+	if x != nil {
+		return x.Originator
+	}
+	return ""
+}
+
+func (x *AgentEvent) GetStatus() string {
+	if x != nil {
+		return x.Status
+	}
+	return ""
+}
+
+type AgentLog struct {
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	Message       string                 `protobuf:"bytes,1,opt,name=message,proto3" json:"message,omitempty"`
+	ComputationId string                 `protobuf:"bytes,2,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
+	Level         string                 `protobuf:"bytes,3,opt,name=level,proto3" json:"level,omitempty"`
+	Timestamp     *timestamppb.Timestamp `protobuf:"bytes,4,opt,name=timestamp,proto3" json:"timestamp,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *AgentLog) Reset() {
+	*x = AgentLog{}
+	mi := &file_agent_events_events_proto_msgTypes[1]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *AgentLog) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*AgentLog) ProtoMessage() {}
+
+func (x *AgentLog) ProtoReflect() protoreflect.Message {
+	mi := &file_agent_events_events_proto_msgTypes[1]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use AgentLog.ProtoReflect.Descriptor instead.
+func (*AgentLog) Descriptor() ([]byte, []int) {
+	return file_agent_events_events_proto_rawDescGZIP(), []int{1}
+}
+
+func (x *AgentLog) GetMessage() string {
+	if x != nil {
+		return x.Message
+	}
+	return ""
+}
+
+func (x *AgentLog) GetComputationId() string {
+	if x != nil {
+		return x.ComputationId
+	}
+	return ""
+}
+
+func (x *AgentLog) GetLevel() string {
+	if x != nil {
+		return x.Level
+	}
+	return ""
+}
+
+func (x *AgentLog) GetTimestamp() *timestamppb.Timestamp {
+	if x != nil {
+		return x.Timestamp
+	}
+	return nil
+}
+
+type EventsLogs struct {
+	state protoimpl.MessageState `protogen:"open.v1"`
+	// Types that are valid to be assigned to Message:
+	//
+	//	*EventsLogs_AgentLog
+	//	*EventsLogs_AgentEvent
+	Message       isEventsLogs_Message `protobuf_oneof:"message"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *EventsLogs) Reset() {
+	*x = EventsLogs{}
+	mi := &file_agent_events_events_proto_msgTypes[2]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *EventsLogs) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*EventsLogs) ProtoMessage() {}
+
+func (x *EventsLogs) ProtoReflect() protoreflect.Message {
+	mi := &file_agent_events_events_proto_msgTypes[2]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use EventsLogs.ProtoReflect.Descriptor instead.
+func (*EventsLogs) Descriptor() ([]byte, []int) {
+	return file_agent_events_events_proto_rawDescGZIP(), []int{2}
+}
+
+func (x *EventsLogs) GetMessage() isEventsLogs_Message {
+	if x != nil {
+		return x.Message
+	}
+	return nil
+}
+
+func (x *EventsLogs) GetAgentLog() *AgentLog {
+	if x != nil {
+		if x, ok := x.Message.(*EventsLogs_AgentLog); ok {
+			return x.AgentLog
+		}
+	}
+	return nil
+}
+
+func (x *EventsLogs) GetAgentEvent() *AgentEvent {
+	if x != nil {
+		if x, ok := x.Message.(*EventsLogs_AgentEvent); ok {
+			return x.AgentEvent
+		}
+	}
+	return nil
+}
+
+type isEventsLogs_Message interface {
+	isEventsLogs_Message()
+}
+
+type EventsLogs_AgentLog struct {
+	AgentLog *AgentLog `protobuf:"bytes,1,opt,name=agent_log,json=agentLog,proto3,oneof"`
+}
+
+type EventsLogs_AgentEvent struct {
+	AgentEvent *AgentEvent `protobuf:"bytes,2,opt,name=agent_event,json=agentEvent,proto3,oneof"`
+}
+
+func (*EventsLogs_AgentLog) isEventsLogs_Message() {}
+
+func (*EventsLogs_AgentEvent) isEventsLogs_Message() {}
+
+var File_agent_events_events_proto protoreflect.FileDescriptor
+
+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
+	file_agent_events_events_proto_rawDescData []byte
+)
+
+func file_agent_events_events_proto_rawDescGZIP() []byte {
+	file_agent_events_events_proto_rawDescOnce.Do(func() {
+		file_agent_events_events_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_agent_events_events_proto_rawDesc), len(file_agent_events_events_proto_rawDesc)))
+	})
+	return file_agent_events_events_proto_rawDescData
+}
+
+var file_agent_events_events_proto_msgTypes = make([]protoimpl.MessageInfo, 3)
+var file_agent_events_events_proto_goTypes = []any{
+	(*AgentEvent)(nil),            // 0: events.AgentEvent
+	(*AgentLog)(nil),              // 1: events.AgentLog
+	(*EventsLogs)(nil),            // 2: events.EventsLogs
+	(*timestamppb.Timestamp)(nil), // 3: google.protobuf.Timestamp
+}
+var file_agent_events_events_proto_depIdxs = []int32{
+	3, // 0: events.AgentEvent.timestamp:type_name -> google.protobuf.Timestamp
+	3, // 1: events.AgentLog.timestamp:type_name -> google.protobuf.Timestamp
+	1, // 2: events.EventsLogs.agent_log:type_name -> events.AgentLog
+	0, // 3: events.EventsLogs.agent_event:type_name -> events.AgentEvent
+	4, // [4:4] is the sub-list for method output_type
+	4, // [4:4] is the sub-list for method input_type
+	4, // [4:4] is the sub-list for extension type_name
+	4, // [4:4] is the sub-list for extension extendee
+	0, // [0:4] is the sub-list for field type_name
+}
+
+func init() { file_agent_events_events_proto_init() }
+func file_agent_events_events_proto_init() {
+	if File_agent_events_events_proto != nil {
+		return
+	}
+	file_agent_events_events_proto_msgTypes[2].OneofWrappers = []any{
+		(*EventsLogs_AgentLog)(nil),
+		(*EventsLogs_AgentEvent)(nil),
+	}
+	type x struct{}
+	out := protoimpl.TypeBuilder{
+		File: protoimpl.DescBuilder{
+			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
+			RawDescriptor: unsafe.Slice(unsafe.StringData(file_agent_events_events_proto_rawDesc), len(file_agent_events_events_proto_rawDesc)),
+			NumEnums:      0,
+			NumMessages:   3,
+			NumExtensions: 0,
+			NumServices:   0,
+		},
+		GoTypes:           file_agent_events_events_proto_goTypes,
+		DependencyIndexes: file_agent_events_events_proto_depIdxs,
+		MessageInfos:      file_agent_events_events_proto_msgTypes,
+	}.Build()
+	File_agent_events_events_proto = out.File
+	file_agent_events_events_proto_goTypes = nil
+	file_agent_events_events_proto_depIdxs = nil
+}

agent/events/events.proto

@@ -0,0 +1,33 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+syntax = "proto3";
+
+package events;
+
+import "google/protobuf/timestamp.proto";
+
+option go_package = "./events";
+
+message AgentEvent {
+  string event_type = 1;
+	google.protobuf.Timestamp timestamp = 2;
+	string	computation_id = 3;
+	bytes	details = 4;
+	string	originator = 5;
+	string	status = 6;
+}
+
+message AgentLog {
+  string message = 1;
+  string computation_id = 2;
+  string level = 3;
+  google.protobuf.Timestamp timestamp = 4;
+}
+
+message EventsLogs {
+  oneof message {
+    AgentLog agent_log = 1;
+    AgentEvent agent_event = 2;
+  }
+}

agent/events/events_test.go

@@ -0,0 +1,38 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package events
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/ultravioletrs/cocos/agent/cvms"
+)
+
+func TestSendEventSuccess(t *testing.T) {
+	queue := make(chan *cvms.ClientStreamMessage, 1)
+	svc, err := New("test_service", queue)
+	assert.NoError(t, err)
+
+	details := json.RawMessage(`{"key": "value"}`)
+
+	go func() {
+		msg := <-queue
+		assert.NotNil(t, msg)
+		assert.NotNil(t, msg.GetAgentEvent())
+		assert.Equal(t, "test_event", msg.GetAgentEvent().EventType)
+		assert.Equal(t, "testid", msg.GetAgentEvent().ComputationId)
+		assert.Equal(t, "test_service", msg.GetAgentEvent().Originator)
+		assert.Equal(t, "success", msg.GetAgentEvent().Status)
+
+		now := time.Now()
+		eventTimestamp := msg.GetAgentEvent().GetTimestamp().AsTime()
+		assert.WithinDuration(t, now, eventTimestamp, 1*time.Second)
+	}()
+
+	svc.SendEvent("testid", "test_event", "success", details)
+
+	time.Sleep(1 * time.Second)
+}

agent/events/mocks/events.go

@@ -1,67 +0,0 @@
-// Code generated by mockery v2.43.2. DO NOT EDIT.
-
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-package mocks
-
-import (
-	json "encoding/json"
-
-	mock "github.com/stretchr/testify/mock"
-)
-
-// Service is an autogenerated mock type for the Service type
-type Service struct {
-	mock.Mock
-}
-
-// Close provides a mock function with given fields:
-func (_m *Service) Close() error {
-	ret := _m.Called()
-
-	if len(ret) == 0 {
-		panic("no return value specified for Close")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func() error); ok {
-		r0 = rf()
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// SendEvent provides a mock function with given fields: event, status, details
-func (_m *Service) SendEvent(event string, status string, details json.RawMessage) error {
-	ret := _m.Called(event, status, details)
-
-	if len(ret) == 0 {
-		panic("no return value specified for SendEvent")
-	}
-
-	var r0 error
-	if rf, ok := ret.Get(0).(func(string, string, json.RawMessage) error); ok {
-		r0 = rf(event, status, details)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// 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/events/mocks/service.go

@@ -0,0 +1,99 @@
+// 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 (
+	"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
+}
+
+type Service_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *Service) EXPECT() *Service_Expecter {
+	return &Service_Expecter{mock: &_m.Mock}
+}
+
+// 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'
+type Service_SendEvent_Call struct {
+	*mock.Call
+}
+
+// SendEvent is a helper method to define mock.On call
+//   - cmpID string
+//   - event string
+//   - status string
+//   - details json.RawMessage
+func (_e *Service_Expecter) SendEvent(cmpID interface{}, event interface{}, status interface{}, details interface{}) *Service_SendEvent_Call {
+	return &Service_SendEvent_Call{Call: _e.mock.On("SendEvent", cmpID, event, status, details)}
+}
+
+func (_c *Service_SendEvent_Call) Run(run func(cmpID string, event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		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
+}
+
+func (_c *Service_SendEvent_Call) Return() *Service_SendEvent_Call {
+	_c.Call.Return()
+	return _c
+}
+
+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
+}

agent/mocks/agent.go

@@ -1,129 +0,0 @@
-// Code generated by mockery v2.42.3. DO NOT EDIT.
-
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-package mocks
-
-import (
-	context "context"
-
-	agent "github.com/ultravioletrs/cocos/agent"
-
-	mock "github.com/stretchr/testify/mock"
-)
-
-// Service is an autogenerated mock type for the Service type
-type Service struct {
-	mock.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)
-
-	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)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// Attestation provides a mock function with given fields: ctx, reportData
-func (_m *Service) Attestation(ctx context.Context, reportData [64]byte) ([]byte, error) {
-	ret := _m.Called(ctx, reportData)
-
-	if len(ret) == 0 {
-		panic("no return value specified for Attestation")
-	}
-
-	var r0 []byte
-	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context, [64]byte) ([]byte, error)); ok {
-		return rf(ctx, reportData)
-	}
-	if rf, ok := ret.Get(0).(func(context.Context, [64]byte) []byte); ok {
-		r0 = rf(ctx, reportData)
-	} else {
-		if ret.Get(0) != nil {
-			r0 = ret.Get(0).([]byte)
-		}
-	}
-
-	if rf, ok := ret.Get(1).(func(context.Context, [64]byte) error); ok {
-		r1 = rf(ctx, reportData)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// 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)
-
-	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)
-	} else {
-		r0 = ret.Error(0)
-	}
-
-	return r0
-}
-
-// Result provides a mock function with given fields: ctx
-func (_m *Service) Result(ctx context.Context) ([]byte, error) {
-	ret := _m.Called(ctx)
-
-	if len(ret) == 0 {
-		panic("no return value specified for Result")
-	}
-
-	var r0 []byte
-	var r1 error
-	if rf, ok := ret.Get(0).(func(context.Context) ([]byte, error)); ok {
-		return rf(ctx)
-	}
-	if rf, ok := ret.Get(0).(func(context.Context) []byte); ok {
-		r0 = rf(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)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// 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/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/auth.go

@@ -1,63 +0,0 @@
-// Code generated by mockery v2.42.3. DO NOT EDIT.
-
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-package mocks
-
-import (
-	context "context"
-
-	auth "github.com/ultravioletrs/cocos/agent/auth"
-
-	mock "github.com/stretchr/testify/mock"
-)
-
-// Authenticator is an autogenerated mock type for the Authenticator type
-type Authenticator struct {
-	mock.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)
-
-	if len(ret) == 0 {
-		panic("no return value specified for AuthenticateUser")
-	}
-
-	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 rf, ok := ret.Get(0).(func(context.Context, auth.UserRole) context.Context); ok {
-		r0 = rf(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)
-	} else {
-		r1 = ret.Error(1)
-	}
-
-	return r0, r1
-}
-
-// 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/mocks/service.go

@@ -0,0 +1,589 @@
+// 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"
+	"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
+}
+
+type Service_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *Service) EXPECT() *Service_Expecter {
+	return &Service_Expecter{mock: &_m.Mock}
+}
+
+// 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 returnFunc, ok := ret.Get(0).(func(context.Context, agent.Algorithm) error); ok {
+		r0 = returnFunc(ctx, algorithm)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// Service_Algo_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Algo'
+type Service_Algo_Call struct {
+	*mock.Call
+}
+
+// Algo is a helper method to define mock.On call
+//   - ctx context.Context
+//   - algorithm agent.Algorithm
+func (_e *Service_Expecter) Algo(ctx interface{}, algorithm interface{}) *Service_Algo_Call {
+	return &Service_Algo_Call{Call: _e.mock.On("Algo", ctx, algorithm)}
+}
+
+func (_c *Service_Algo_Call) Run(run func(ctx context.Context, algorithm agent.Algorithm)) *Service_Algo_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		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(err error) *Service_Algo_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+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 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")
+	}
+
+	var r0 []byte
+	var r1 error
+	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 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 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
+}
+
+// Service_Attestation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Attestation'
+type Service_Attestation_Call struct {
+	*mock.Call
+}
+
+// Attestation is a helper method to define mock.On call
+//   - ctx context.Context
+//   - reportData [64]byte
+//   - nonce [32]byte
+//   - attType attestation.PlatformType
+func (_e *Service_Expecter) Attestation(ctx interface{}, reportData interface{}, nonce interface{}, attType interface{}) *Service_Attestation_Call {
+	return &Service_Attestation_Call{Call: _e.mock.On("Attestation", ctx, reportData, nonce, attType)}
+}
+
+func (_c *Service_Attestation_Call) Run(run func(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType)) *Service_Attestation_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		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(bytes []byte, err error) *Service_Attestation_Call {
+	_c.Call.Return(bytes, err)
+	return _c
+}
+
+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
+}
+
+// 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 AzureAttestationToken")
+	}
+
+	var r0 []byte
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context, [32]byte) ([]byte, error)); ok {
+		return returnFunc(ctx, nonce)
+	}
+	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 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_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
+}
+
+// AzureAttestationToken is a helper method to define mock.On call
+//   - ctx context.Context
+//   - nonce [32]byte
+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_AzureAttestationToken_Call) Run(run func(ctx context.Context, nonce [32]byte)) *Service_AzureAttestationToken_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		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_AzureAttestationToken_Call) Return(bytes []byte, err error) *Service_AzureAttestationToken_Call {
+	_c.Call.Return(bytes, err)
+	return _c
+}
+
+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 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 returnFunc, ok := ret.Get(0).(func(context.Context, agent.Dataset) error); ok {
+		r0 = returnFunc(ctx, dataset)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// Service_Data_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Data'
+type Service_Data_Call struct {
+	*mock.Call
+}
+
+// Data is a helper method to define mock.On call
+//   - ctx context.Context
+//   - dataset agent.Dataset
+func (_e *Service_Expecter) Data(ctx interface{}, dataset interface{}) *Service_Data_Call {
+	return &Service_Data_Call{Call: _e.mock.On("Data", ctx, dataset)}
+}
+
+func (_c *Service_Data_Call) Run(run func(ctx context.Context, dataset agent.Dataset)) *Service_Data_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		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(err error) *Service_Data_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+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 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")
+	}
+
+	var r0 []byte
+	var r1 []byte
+	var r2 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context) ([]byte, []byte, error)); ok {
+		return returnFunc(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 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 returnFunc, ok := ret.Get(2).(func(context.Context) error); ok {
+		r2 = returnFunc(ctx)
+	} else {
+		r2 = ret.Error(2)
+	}
+	return r0, r1, r2
+}
+
+// Service_IMAMeasurements_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IMAMeasurements'
+type Service_IMAMeasurements_Call struct {
+	*mock.Call
+}
+
+// IMAMeasurements is a helper method to define mock.On call
+//   - ctx context.Context
+func (_e *Service_Expecter) IMAMeasurements(ctx interface{}) *Service_IMAMeasurements_Call {
+	return &Service_IMAMeasurements_Call{Call: _e.mock.On("IMAMeasurements", ctx)}
+}
+
+func (_c *Service_IMAMeasurements_Call) Run(run func(ctx context.Context)) *Service_IMAMeasurements_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+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(ctx context.Context) ([]byte, []byte, error)) *Service_IMAMeasurements_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// 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 returnFunc, ok := ret.Get(0).(func(context.Context, agent.Computation) error); ok {
+		r0 = returnFunc(ctx, cmp)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// Service_InitComputation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'InitComputation'
+type Service_InitComputation_Call struct {
+	*mock.Call
+}
+
+// InitComputation is a helper method to define mock.On call
+//   - ctx context.Context
+//   - cmp agent.Computation
+func (_e *Service_Expecter) InitComputation(ctx interface{}, cmp interface{}) *Service_InitComputation_Call {
+	return &Service_InitComputation_Call{Call: _e.mock.On("InitComputation", ctx, cmp)}
+}
+
+func (_c *Service_InitComputation_Call) Run(run func(ctx context.Context, cmp agent.Computation)) *Service_InitComputation_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		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(err error) *Service_InitComputation_Call {
+	_c.Call.Return(err)
+	return _c
+}
+
+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 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")
+	}
+
+	var r0 []byte
+	var r1 error
+	if returnFunc, ok := ret.Get(0).(func(context.Context) ([]byte, error)); ok {
+		return returnFunc(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 returnFunc, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = returnFunc(ctx)
+	} else {
+		r1 = ret.Error(1)
+	}
+	return r0, r1
+}
+
+// Service_Result_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Result'
+type Service_Result_Call struct {
+	*mock.Call
+}
+
+// Result is a helper method to define mock.On call
+//   - ctx context.Context
+func (_e *Service_Expecter) Result(ctx interface{}) *Service_Result_Call {
+	return &Service_Result_Call{Call: _e.mock.On("Result", ctx)}
+}
+
+func (_c *Service_Result_Call) Run(run func(ctx context.Context)) *Service_Result_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+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(ctx context.Context) ([]byte, error)) *Service_Result_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// 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 returnFunc, ok := ret.Get(0).(func() string); ok {
+		r0 = returnFunc()
+	} else {
+		r0 = ret.Get(0).(string)
+	}
+	return r0
+}
+
+// Service_State_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'State'
+type Service_State_Call struct {
+	*mock.Call
+}
+
+// State is a helper method to define mock.On call
+func (_e *Service_Expecter) State() *Service_State_Call {
+	return &Service_State_Call{Call: _e.mock.On("State")}
+}
+
+func (_c *Service_State_Call) Run(run func()) *Service_State_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *Service_State_Call) Return(s string) *Service_State_Call {
+	_c.Call.Return(s)
+	return _c
+}
+
+func (_c *Service_State_Call) RunAndReturn(run func() string) *Service_State_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// StopComputation provides a mock function 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 returnFunc, ok := ret.Get(0).(func(context.Context) error); ok {
+		r0 = returnFunc(ctx)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// Service_StopComputation_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'StopComputation'
+type Service_StopComputation_Call struct {
+	*mock.Call
+}
+
+// StopComputation is a helper method to define mock.On call
+//   - ctx context.Context
+func (_e *Service_Expecter) StopComputation(ctx interface{}) *Service_StopComputation_Call {
+	return &Service_StopComputation_Call{Call: _e.mock.On("StopComputation", ctx)}
+}
+
+func (_c *Service_StopComputation_Call) Run(run func(ctx context.Context)) *Service_StopComputation_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+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(ctx context.Context) error) *Service_StopComputation_Call {
+	_c.Call.Return(run)
+	return _c
+}

agent/service.go

@@ -6,26 +6,80 @@ package agent
 import (
 	"context"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"log/slog"
 	"os"
+	"path/filepath"
 	"slices"
+	sync "sync"
+	"time"
 
-	"github.com/google/go-sev-guest/client"
+	"github.com/absmach/supermq/pkg/errors"
+	"github.com/ultravioletrs/cocos/agent/algorithm"
 	"github.com/ultravioletrs/cocos/agent/algorithm/binary"
+	"github.com/ultravioletrs/cocos/agent/algorithm/docker"
+	"github.com/ultravioletrs/cocos/agent/algorithm/python"
+	"github.com/ultravioletrs/cocos/agent/algorithm/wasm"
 	"github.com/ultravioletrs/cocos/agent/events"
+	"github.com/ultravioletrs/cocos/agent/statemachine"
+	"github.com/ultravioletrs/cocos/internal"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	"golang.org/x/crypto/sha3"
 )
 
 var _ Service = (*agentService)(nil)
 
+//go:generate stringer -type=AgentState
+type AgentState int
+
+const (
+	Idle AgentState = iota
+	ReceivingManifest
+	ReceivingAlgorithm
+	ReceivingData
+	Running
+	ConsumingResults
+	Complete
+	Failed
+)
+
+//go:generate stringer -type=AgentEvent
+type AgentEvent int
+
+const (
+	Start AgentEvent = iota
+	ManifestReceived
+	AlgorithmReceived
+	DataReceived
+	RunComplete
+	ResultsConsumed
+	RunFailed
+)
+
+//go:generate stringer -type=Status
+type Status uint8
+
+const (
+	IdleState Status = iota
+	InProgress
+	Ready
+	Completed
+	Terminated
+	Warning
+	Starting
+)
+
 const (
-	// ReportDataSize is the size of the report data expected by the attestation service.
-	ReportDataSize     = 64
 	algoFilePermission = 0o700
 )
 
+const (
+	ImaMeasurementsFilePath = "/sys/kernel/security/integrity/ima/ascii_runtime_measurements"
+	ImaPcrIndex             = 10
+)
+
 var (
 	// ErrMalformedEntity indicates malformed entity specification (e.g.
 	// invalid username or password).
@@ -33,85 +87,217 @@ var (
 	// ErrUnauthorizedAccess indicates missing or invalid credentials provided
 	// when accessing a protected resource.
 	ErrUnauthorizedAccess = errors.New("missing or invalid credentials provided")
-	// errUndeclaredAlgorithm indicates algorithm was not declared in computation manifest.
+	// ErrUndeclaredAlgorithm indicates algorithm was not declared in computation manifest.
 	ErrUndeclaredDataset = errors.New("dataset not declared in computation manifest")
-	// errAllManifestItemsReceived indicates no new computation manifest items expected.
+	// ErrAllManifestItemsReceived indicates no new computation manifest items expected.
 	ErrAllManifestItemsReceived = errors.New("all expected manifest Items have been received")
-	// errUndeclaredConsumer indicates the consumer requesting results in not declared in computation manifest.
+	// ErrUndeclaredConsumer indicates the consumer requesting results in not declared in computation manifest.
 	ErrUndeclaredConsumer = errors.New("result consumer is undeclared in computation manifest")
-	// errResultsNotReady indicates the computation results are not ready.
+	// ErrResultsNotReady indicates the computation results are not ready.
 	ErrResultsNotReady = errors.New("computation results are not yet ready")
-	// errStateNotReady agent received a request in the wrong state.
+	// ErrStateNotReady agent received a request in the wrong state.
 	ErrStateNotReady = errors.New("agent not expecting this operation in the current state")
-	// errHashMismatch provided algorithm/dataset does not match hash in manifest.
+	// ErrHashMismatch provided algorithm/dataset does not match hash in manifest.
 	ErrHashMismatch = errors.New("malformed data, hash does not match manifest")
+	// ErrFileNameMismatch provided dataset filename does not match filename in manifest.
+	ErrFileNameMismatch = errors.New("malformed data, filename does not match manifest")
+	// ErrAllResultsConsumed indicates all results have been consumed.
+	ErrAllResultsConsumed = errors.New("all results have been consumed by declared consumers")
+	// ErrAttestationFailed attestation failed.
+	ErrAttestationFailed = errors.New("failed to get raw quote")
+	// ErrAttestationVTpmFailed vTPM attestation failed.
+	ErrAttestationVTpmFailed = errors.New("failed to get vTPM quote")
+	// ErrFetchAzureToken azure token fetch failed.
+	ErrFetchAzureToken = errors.New("failed to get azure token")
+	// ErrAttType indicates that the attestation type that is requested does not exist or is not supported.
+	ErrAttestationType = errors.New("attestation type does not exist or is not supported")
 )
 
 // Service specifies an API that must be fullfiled by the domain service
 // implementation, and all of its decorators (e.g. logging & metrics).
-//
-//go:generate mockery --name Service --output=mocks --filename agent.go --quiet --note "Copyright (c) Ultraviolet \n // SPDX-License-Identifier: Apache-2.0"
 type Service interface {
+	InitComputation(ctx context.Context, cmp Computation) error
+	StopComputation(ctx context.Context) error
 	Algo(ctx context.Context, algorithm Algorithm) error
 	Data(ctx context.Context, dataset Dataset) error
 	Result(ctx context.Context) ([]byte, error)
-	Attestation(ctx context.Context, reportData [ReportDataSize]byte) ([]byte, error)
+	Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error)
+	IMAMeasurements(ctx context.Context) ([]byte, []byte, error)
+	AzureAttestationToken(ctx context.Context, nonce [vtpm.Nonce]byte) ([]byte, error)
+	State() string
 }
 
 type agentService struct {
-	computation Computation    // Holds the current computation request details.
-	algorithm   string         // Filepath to the algorithm received for the computation.
-	datasets    []string       // Filepath to the datasets received for the computation.
-	result      []byte         // Stores the result of the computation.
-	sm          *StateMachine  // Manages the state transitions of the agent service.
-	runError    error          // Stores any error encountered during the computation run.
-	eventSvc    events.Service // Service for publishing events related to computation.
+	mu              sync.Mutex
+	computation     Computation               // Holds the current computation request details.
+	algorithm       algorithm.Algorithm       // Filepath to the algorithm received for the computation.
+	result          []byte                    // Stores the result of the computation.
+	sm              statemachine.StateMachine // Manages the state transitions of the agent service.
+	runError        error                     // Stores any error encountered during the computation run.
+	eventSvc        events.Service            // Service for publishing events related to computation.
+	provider        attestation.Provider      // Provider for generating attestation quotes.
+	logger          *slog.Logger              // Logger for the agent service.
+	resultsConsumed bool                      // Indicates if the results have been consumed.
+	cancel          context.CancelFunc        // Cancels the computation context.
+	vmpl            int                       // VMPL at which the Agent is running.
 }
 
 var _ Service = (*agentService)(nil)
 
 // New instantiates the agent service implementation.
-func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp Computation) Service {
+func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, provider attestation.Provider, vmlp int) Service {
+	sm := statemachine.NewStateMachine(Idle)
+	ctx, cancel := context.WithCancel(ctx)
 	svc := &agentService{
-		sm:       NewStateMachine(logger),
+		sm:       sm,
 		eventSvc: eventSvc,
+		provider: provider,
+		logger:   logger,
+		cancel:   cancel,
+		vmpl:     vmlp,
 	}
 
-	go svc.sm.Start(ctx)
-	svc.sm.SendEvent(start)
-	svc.sm.StateFunctions[idle] = svc.publishEvent("in-progress", json.RawMessage{})
-	svc.sm.StateFunctions[receivingManifest] = svc.publishEvent("in-progress", json.RawMessage{})
-	svc.sm.StateFunctions[receivingAlgorithm] = svc.publishEvent("in-progress", json.RawMessage{})
-	svc.sm.StateFunctions[receivingData] = svc.publishEvent("in-progress", json.RawMessage{})
-	svc.sm.StateFunctions[resultsReady] = svc.publishEvent("in-progress", json.RawMessage{})
-	svc.sm.StateFunctions[complete] = svc.publishEvent("in-progress", json.RawMessage{})
-	svc.sm.StateFunctions[running] = svc.runComputation
+	transitions := []statemachine.Transition{
+		{From: Idle, Event: Start, To: ReceivingManifest},
+		{From: ReceivingManifest, Event: ManifestReceived, To: ReceivingAlgorithm},
+	}
+
+	transitions = append(transitions, []statemachine.Transition{
+		{From: Running, Event: RunComplete, To: ConsumingResults},
+		{From: Running, Event: RunFailed, To: Failed},
+		{From: ConsumingResults, Event: ResultsConsumed, To: Complete},
+	}...)
+
+	for _, t := range transitions {
+		sm.AddTransition(t)
+	}
+
+	sm.SetAction(ReceivingAlgorithm, svc.publishEvent(InProgress.String()))
+	sm.SetAction(ReceivingData, svc.publishEvent(InProgress.String()))
+	sm.SetAction(Running, svc.runComputation)
+	sm.SetAction(ConsumingResults, svc.publishEvent(Ready.String()))
+	sm.SetAction(Complete, svc.publishEvent(Completed.String()))
+	sm.SetAction(Failed, svc.publishEvent(Failed.String()))
+
+	go func() {
+		if err := sm.Start(ctx); err != nil {
+			logger.Error(err.Error())
+		}
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+	sm.SendEvent(Start)
+
+	time.Sleep(100 * time.Millisecond)
 
-	svc.computation = cmp
-	svc.sm.SendEvent(manifestReceived)
 	return svc
 }
 
-func (as *agentService) Algo(ctx context.Context, algorithm Algorithm) error {
-	if as.sm.GetState() != receivingAlgorithm {
+func (as *agentService) State() string {
+	return as.sm.GetState().String()
+}
+
+func (as *agentService) InitComputation(ctx context.Context, cmp Computation) error {
+	if as.sm.GetState() != ReceivingManifest {
 		return ErrStateNotReady
 	}
-	if as.algorithm != "" {
+	defer as.sm.SendEvent(ManifestReceived)
+
+	as.mu.Lock()
+	defer as.mu.Unlock()
+
+	as.computation = cmp
+
+	transitions := []statemachine.Transition{}
+
+	if len(cmp.Datasets) == 0 {
+		transitions = append(transitions, statemachine.Transition{From: ReceivingAlgorithm, Event: AlgorithmReceived, To: Running})
+	} else {
+		transitions = append(transitions, statemachine.Transition{From: ReceivingAlgorithm, Event: AlgorithmReceived, To: ReceivingData})
+		transitions = append(transitions, statemachine.Transition{From: ReceivingData, Event: DataReceived, To: Running})
+	}
+
+	for _, t := range transitions {
+		as.sm.AddTransition(t)
+	}
+
+	return nil
+}
+
+func (as *agentService) StopComputation(ctx context.Context) error {
+	as.mu.Lock()
+	defer as.mu.Unlock()
+
+	as.eventSvc.SendEvent(as.computation.ID, "Stopped", "Stopped", json.RawMessage{})
+
+	as.cancel()
+
+	if as.algorithm != nil {
+		if err := as.algorithm.Stop(); err != nil {
+			return fmt.Errorf("error stopping computation: %v", err)
+		}
+	}
+
+	if err := os.RemoveAll(algorithm.DatasetsDir); err != nil {
+		return fmt.Errorf("error removing datasets directory: %v", err)
+	}
+
+	if err := os.RemoveAll(algorithm.ResultsDir); err != nil {
+		return fmt.Errorf("error removing results directory: %v", err)
+	}
+
+	as.sm.Reset(Idle)
+
+	as.computation = Computation{}
+	as.algorithm = nil
+	as.result = nil
+	as.runError = nil
+	as.resultsConsumed = false
+
+	ctx, cancel := context.WithCancel(ctx)
+	as.cancel = cancel
+
+	go func() {
+		if err := as.sm.Start(ctx); err != nil {
+			as.logger.Error(err.Error())
+		}
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+	as.sm.SendEvent(Start)
+
+	time.Sleep(100 * time.Millisecond)
+
+	return nil
+}
+
+func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
+	if as.sm.GetState() != ReceivingAlgorithm {
+		return ErrStateNotReady
+	}
+	as.mu.Lock()
+	defer as.mu.Unlock()
+	if as.algorithm != nil {
 		return ErrAllManifestItemsReceived
 	}
 
-	hash := sha3.Sum256(algorithm.Algorithm)
+	hash := sha3.Sum256(algo.Algorithm)
 
 	if hash != as.computation.Algorithm.Hash {
 		return ErrHashMismatch
 	}
 
-	f, err := os.CreateTemp("", "algorithm")
+	currentDir, err := os.Getwd()
+	if err != nil {
+		return fmt.Errorf("error getting current directory: %v", err)
+	}
+
+	f, err := os.Create(filepath.Join(currentDir, "algo"))
 	if err != nil {
 		return fmt.Errorf("error creating algorithm file: %v", err)
 	}
 
-	if _, err := f.Write(algorithm.Algorithm); err != nil {
+	if _, err := f.Write(algo.Algorithm); err != nil {
 		return fmt.Errorf("error writing algorithm to file: %v", err)
 	}
 
@@ -123,110 +309,232 @@ func (as *agentService) Algo(ctx context.Context, algorithm Algorithm) error {
 		return fmt.Errorf("error closing file: %v", err)
 	}
 
-	as.algorithm = f.Name()
+	algoType := algorithm.AlgorithmTypeFromContext(ctx)
+	if algoType == "" {
+		algoType = string(algorithm.AlgoTypeBin)
+	}
 
-	if as.algorithm != "" {
-		as.sm.SendEvent(algorithmReceived)
+	args := algorithm.AlgorithmArgsFromContext(ctx)
+
+	switch algoType {
+	case string(algorithm.AlgoTypeBin):
+		as.algorithm = binary.NewAlgorithm(as.logger, as.eventSvc, f.Name(), args, as.computation.ID)
+	case string(algorithm.AlgoTypePython):
+		var requirementsFile string
+		if len(algo.Requirements) > 0 {
+			fr, err := os.CreateTemp("", "requirements.txt")
+			if err != nil {
+				return fmt.Errorf("error creating requirments file: %v", err)
+			}
+
+			if _, err := fr.Write(algo.Requirements); err != nil {
+				return fmt.Errorf("error writing requirements to file: %v", err)
+			}
+			if err := fr.Close(); err != nil {
+				return fmt.Errorf("error closing file: %v", err)
+			}
+			requirementsFile = fr.Name()
+		}
+		runtime := python.PythonRunTimeFromContext(ctx)
+		as.algorithm = python.NewAlgorithm(as.logger, as.eventSvc, runtime, requirementsFile, f.Name(), args, as.computation.ID)
+	case string(algorithm.AlgoTypeWasm):
+		as.algorithm = wasm.NewAlgorithm(as.logger, as.eventSvc, args, f.Name(), as.computation.ID)
+	case string(algorithm.AlgoTypeDocker):
+		as.algorithm = docker.NewAlgorithm(as.logger, as.eventSvc, f.Name(), as.computation.ID)
+	}
+
+	if err := os.Mkdir(algorithm.DatasetsDir, 0o755); err != nil {
+		return fmt.Errorf("error creating datasets directory: %v", err)
+	}
+
+	if as.algorithm != nil {
+		as.sm.SendEvent(AlgorithmReceived)
 	}
 
 	return nil
 }
 
 func (as *agentService) Data(ctx context.Context, dataset Dataset) error {
-	if as.sm.GetState() != receivingData {
+	if as.sm.GetState() != ReceivingData {
 		return ErrStateNotReady
 	}
+	as.mu.Lock()
+	defer as.mu.Unlock()
 	if len(as.computation.Datasets) == 0 {
 		return ErrAllManifestItemsReceived
 	}
 
 	hash := sha3.Sum256(dataset.Dataset)
 
-	index, ok := IndexFromContext(ctx)
-	if !ok {
+	matched := false
+	for i, d := range as.computation.Datasets {
+		if hash == d.Hash {
+			if d.Filename != "" && d.Filename != dataset.Filename {
+				return ErrFileNameMismatch
+			}
+
+			as.computation.Datasets = slices.Delete(as.computation.Datasets, i, i+1)
+
+			if DecompressFromContext(ctx) {
+				if err := internal.UnzipFromMemory(dataset.Dataset, algorithm.DatasetsDir); err != nil {
+					return fmt.Errorf("error decompressing dataset: %v", err)
+				}
+			} else {
+				f, err := os.Create(fmt.Sprintf("%s/%s", algorithm.DatasetsDir, dataset.Filename))
+				if err != nil {
+					return fmt.Errorf("error creating dataset file: %v", err)
+				}
+
+				if _, err := f.Write(dataset.Dataset); err != nil {
+					return fmt.Errorf("error writing dataset to file: %v", err)
+				}
+				if err := f.Close(); err != nil {
+					return fmt.Errorf("error closing file: %v", err)
+				}
+			}
+
+			matched = true
+			break
+		}
+	}
+
+	if !matched {
 		return ErrUndeclaredDataset
 	}
 
-	if hash != as.computation.Datasets[index].Hash {
-		return ErrHashMismatch
-	}
-	as.computation.Datasets = slices.Delete(as.computation.Datasets, index, index+1)
-
-	f, err := os.CreateTemp("", fmt.Sprintf("dataset-%d", index))
-	if err != nil {
-		return fmt.Errorf("error creating dataset file: %v", err)
-	}
-
-	if _, err := f.Write(dataset.Dataset); err != nil {
-		return fmt.Errorf("error writing dataset to file: %v", err)
-	}
-	if err := f.Close(); err != nil {
-		return fmt.Errorf("error closing file: %v", err)
-	}
-
-	as.datasets = append(as.datasets, f.Name())
-
 	if len(as.computation.Datasets) == 0 {
-		as.sm.SendEvent(dataReceived)
+		defer as.sm.SendEvent(DataReceived)
 	}
 
 	return nil
 }
 
 func (as *agentService) Result(ctx context.Context) ([]byte, error) {
-	if as.sm.GetState() != resultsReady {
+	currentState := as.sm.GetState()
+	if currentState != ConsumingResults && currentState != Complete && currentState != Failed {
 		return []byte{}, ErrResultsNotReady
 	}
-	if len(as.computation.ResultConsumers) == 0 {
-		return []byte{}, ErrAllManifestItemsReceived
-	}
+
 	index, ok := IndexFromContext(ctx)
 	if !ok {
 		return []byte{}, ErrUndeclaredConsumer
 	}
-	as.computation.ResultConsumers = slices.Delete(as.computation.ResultConsumers, index, index+1)
 
-	if len(as.computation.ResultConsumers) == 0 {
-		as.sm.SendEvent(resultsConsumed)
+	as.mu.Lock()
+	defer as.mu.Unlock()
+	if index < 0 || index >= len(as.computation.ResultConsumers) {
+		return []byte{}, ErrUndeclaredConsumer
 	}
-	// Return the result file or an error
+
+	if !as.resultsConsumed && currentState == ConsumingResults {
+		as.resultsConsumed = true
+		defer as.sm.SendEvent(ResultsConsumed)
+	}
+
 	return as.result, as.runError
 }
 
-func (as *agentService) Attestation(ctx context.Context, reportData [ReportDataSize]byte) ([]byte, error) {
-	provider, err := client.GetQuoteProvider()
-	if err != nil {
-		return []byte{}, err
+func (as *agentService) Attestation(ctx context.Context, reportData [quoteprovider.Nonce]byte, nonce [vtpm.Nonce]byte, attType attestation.PlatformType) ([]byte, error) {
+	switch attType {
+	case attestation.SNP, attestation.TDX:
+		rawQuote, err := as.provider.TeeAttestation(reportData[:])
+		if err != nil {
+			return []byte{}, errors.Wrap(ErrAttestationFailed, err)
+		}
+		return rawQuote, nil
+	case attestation.VTPM:
+		vTPMQuote, err := as.provider.VTpmAttestation(nonce[:])
+		if err != nil {
+			return []byte{}, errors.Wrap(ErrAttestationVTpmFailed, err)
+		}
+		return vTPMQuote, nil
+	case attestation.SNPvTPM:
+		vTPMQuote, err := as.provider.Attestation(reportData[:], nonce[:])
+		if err != nil {
+			return []byte{}, errors.Wrap(ErrAttestationVTpmFailed, err)
+		}
+		return vTPMQuote, nil
+	default:
+		return []byte{}, ErrAttestationType
 	}
-	rawQuote, err := provider.GetRawQuote(reportData)
-	if err != nil {
-		return []byte{}, err
-	}
-
-	return rawQuote, nil
 }
 
-func (as *agentService) runComputation() {
-	as.publishEvent("starting", json.RawMessage{})()
-	as.sm.logger.Debug("computation run started")
-	defer as.sm.SendEvent(runComplete)
-	as.publishEvent("in-progress", json.RawMessage{})()
-	algorithm := binary.New(as.sm.logger, as.eventSvc, as.algorithm, as.datasets...)
-	result, err := algorithm.Run()
+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 {
-		as.runError = err
-		as.sm.logger.Warn(fmt.Sprintf("computation failed with error: %s", err.Error()))
-		as.publishEvent("failed", json.RawMessage{})()
+		return []byte{}, err
+	}
+	return token, nil
+}
+
+func (as *agentService) runComputation(state statemachine.State) {
+	as.publishEvent(Starting.String())(state)
+	as.logger.Debug("computation run started")
+	defer func() {
+		if as.runError != nil {
+			as.sm.SendEvent(RunFailed)
+		} else {
+			as.sm.SendEvent(RunComplete)
+		}
+	}()
+
+	if err := os.Mkdir(algorithm.ResultsDir, 0o755); err != nil {
+		as.runError = fmt.Errorf("error creating results directory: %s", err.Error())
+		as.logger.Warn(as.runError.Error())
+		as.publishEvent(Failed.String())(state)
 		return
 	}
-	as.publishEvent("complete", json.RawMessage{})()
-	as.result = result
+
+	defer func() {
+		if err := os.RemoveAll(algorithm.ResultsDir); err != nil {
+			as.logger.Warn(fmt.Sprintf("error removing results directory and its contents: %s", err.Error()))
+		}
+		if err := os.RemoveAll(algorithm.DatasetsDir); err != nil {
+			as.logger.Warn(fmt.Sprintf("error removing datasets directory and its contents: %s", err.Error()))
+		}
+	}()
+
+	as.publishEvent(InProgress.String())(state)
+	if err := as.algorithm.Run(); err != nil {
+		as.runError = err
+		as.logger.Warn(fmt.Sprintf("failed to run computation: %s", err.Error()))
+		as.publishEvent(Failed.String())(state)
+		return
+	}
+
+	results, err := internal.ZipDirectoryToMemory(algorithm.ResultsDir)
+	if err != nil {
+		as.runError = err
+		as.logger.Warn(fmt.Sprintf("failed to zip results: %s", err.Error()))
+		as.publishEvent(Failed.String())(state)
+		return
+	}
+
+	as.publishEvent(Completed.String())(state)
+
+	as.result = results
 }
 
-func (as *agentService) publishEvent(status string, details json.RawMessage) func() {
-	return func() {
-		if err := as.eventSvc.SendEvent(as.sm.State.String(), status, details); err != nil {
-			as.sm.logger.Warn(err.Error())
-		}
+func (as *agentService) publishEvent(status string) statemachine.Action {
+	return func(state statemachine.State) {
+		as.eventSvc.SendEvent(as.computation.ID, state.String(), status, json.RawMessage{})
 	}
 }
+
+func (as *agentService) IMAMeasurements(ctx context.Context) ([]byte, []byte, error) {
+	data, err := os.ReadFile(ImaMeasurementsFilePath)
+	if err != nil {
+		return nil, nil, fmt.Errorf("error reading Linux IMA measurements file: %s", err.Error())
+	}
+
+	pcr10, err := vtpm.GetPCRSHA1Value(ImaPcrIndex)
+	if err != nil {
+		return nil, nil, fmt.Errorf("error reading TPM PCR #10: %s", err.Error())
+	}
+
+	return data, pcr10, nil
+}

agent/service_test.go

@@ -0,0 +1,671 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package agent
+
+import (
+	"context"
+	"crypto/rand"
+	"fmt"
+	"log"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	mglog "github.com/absmach/supermq/logger"
+	"github.com/absmach/supermq/pkg/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/agent/algorithm"
+	algomocks "github.com/ultravioletrs/cocos/agent/algorithm/mocks"
+	"github.com/ultravioletrs/cocos/agent/algorithm/python"
+	"github.com/ultravioletrs/cocos/agent/events/mocks"
+	"github.com/ultravioletrs/cocos/agent/statemachine"
+	smmocks "github.com/ultravioletrs/cocos/agent/statemachine/mocks"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	mocks2 "github.com/ultravioletrs/cocos/pkg/attestation/mocks"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
+	"golang.org/x/crypto/sha3"
+	"google.golang.org/grpc/metadata"
+)
+
+var (
+	algoPath = "../test/manual/algo/lin_reg.py"
+	reqPath  = "../test/manual/algo/requirements.txt"
+	dataPath = "../test/manual/data/iris.csv"
+)
+
+const datasetFile = "iris.csv"
+
+func TestAlgo(t *testing.T) {
+	algo, err := os.ReadFile(algoPath)
+	require.NoError(t, err)
+
+	algoHash := sha3.Sum256(algo)
+	vtpm.ExternalTPM = &vtpm.DummyRWC{}
+
+	reqFile, err := os.ReadFile(reqPath)
+	require.NoError(t, err)
+
+	testCases := []struct {
+		name     string
+		err      error
+		algo     Algorithm
+		algoType string
+	}{
+		{
+			name: "Test Algo successfully",
+			algo: Algorithm{
+				Algorithm: algo,
+				Hash:      algoHash,
+			},
+			algoType: "python",
+			err:      nil,
+		},
+		{
+			name: "Test Algo successfully with requirements file",
+			algo: Algorithm{
+				Algorithm:    algo,
+				Hash:         algoHash,
+				Requirements: reqFile,
+			},
+			algoType: "python",
+			err:      nil,
+		},
+		{
+			name: "Test Algo type binary successfully",
+			algo: Algorithm{
+				Algorithm: algo,
+				Hash:      algoHash,
+			},
+			algoType: "bin",
+			err:      nil,
+		},
+		{
+			name: "Test Algo type wasm successfully",
+			algo: Algorithm{
+				Algorithm: algo,
+				Hash:      algoHash,
+			},
+			algoType: "wasm",
+			err:      nil,
+		},
+		{
+			name: "Test Algo type docker successfully",
+			algo: Algorithm{
+				Algorithm: algo,
+				Hash:      algoHash,
+			},
+			algoType: "docker",
+			err:      nil,
+		},
+		{
+			name:     "Test algo hash mismatch",
+			algo:     Algorithm{},
+			algoType: "python",
+			err:      ErrHashMismatch,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			err = os.RemoveAll("datasets")
+			require.NoError(t, err)
+
+			ctx := metadata.NewIncomingContext(context.Background(),
+				metadata.Pairs(algorithm.AlgoTypeKey, tc.algoType, python.PyRuntimeKey, python.PyRuntime),
+			)
+
+			events := new(mocks.Service)
+			events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
+			ctx, cancel := context.WithCancel(ctx)
+			defer cancel()
+			svc := New(ctx, mglog.NewMock(), events, &attestation.EmptyProvider{}, 0)
+
+			err := svc.InitComputation(ctx, testComputation(t))
+			require.NoError(t, err)
+
+			time.Sleep(300 * time.Millisecond)
+
+			err = svc.Algo(ctx, tc.algo)
+			assert.True(t, errors.Contains(err, tc.err), "expected %v, got %v", tc.err, err)
+			t.Cleanup(func() {
+				err = os.RemoveAll("venv")
+				err = os.RemoveAll("algo")
+				err = os.RemoveAll("datasets")
+			})
+		})
+	}
+}
+
+func TestData(t *testing.T) {
+	algo, err := os.ReadFile(algoPath)
+	require.NoError(t, err)
+
+	algoHash := sha3.Sum256(algo)
+	vtpm.ExternalTPM = &vtpm.DummyRWC{}
+
+	alg := Algorithm{
+		Hash:      algoHash,
+		Algorithm: algo,
+	}
+
+	data, err := os.ReadFile(dataPath)
+	require.NoError(t, err)
+
+	dataHash := sha3.Sum256(data)
+
+	cases := []struct {
+		name string
+		data Dataset
+		err  error
+	}{
+		{
+			name: "Test data successfully",
+			data: Dataset{
+				Hash:     dataHash,
+				Dataset:  data,
+				Filename: datasetFile,
+			},
+		},
+		{
+			name: "Test State not ready",
+			data: Dataset{
+				Dataset:  data,
+				Hash:     dataHash,
+				Filename: datasetFile,
+			},
+			err: ErrStateNotReady,
+		},
+		{
+			name: "Test File name does not match manifest",
+			data: Dataset{
+				Dataset:  data,
+				Hash:     dataHash,
+				Filename: "invalid",
+			},
+			err: ErrFileNameMismatch,
+		},
+		{
+			name: "Test dataset not declared in manifest",
+			data: Dataset{
+				Filename: datasetFile,
+			},
+			err: ErrUndeclaredDataset,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := metadata.NewIncomingContext(context.Background(),
+				metadata.Pairs(
+					algorithm.AlgoTypeKey, "python",
+					python.PyRuntimeKey, python.PyRuntime),
+			)
+
+			events := new(mocks.Service)
+			events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
+			if tc.err != ErrUndeclaredDataset {
+				ctx = IndexToContext(ctx, 0)
+			}
+
+			ctx, cancel := context.WithCancel(ctx)
+			defer cancel()
+
+			svc := New(ctx, mglog.NewMock(), events, &attestation.EmptyProvider{}, 0)
+
+			err := svc.InitComputation(ctx, testComputation(t))
+			require.NoError(t, err)
+
+			time.Sleep(300 * time.Millisecond)
+
+			if tc.err != ErrStateNotReady {
+				err = svc.Algo(ctx, alg)
+				require.NoError(t, err)
+				time.Sleep(300 * time.Millisecond)
+			}
+			err = svc.Data(ctx, tc.data)
+			assert.True(t, errors.Contains(err, tc.err), "expected %v, got %v", tc.err, err)
+			t.Cleanup(func() {
+				_ = os.RemoveAll("datasets")
+				_ = os.RemoveAll("results")
+				err = os.RemoveAll("venv")
+				err = os.RemoveAll("algo")
+			})
+		})
+	}
+}
+
+func TestResult(t *testing.T) {
+	cases := []struct {
+		name     string
+		err      error
+		setup    func(svc *agentService)
+		ctxSetup func(ctx context.Context) context.Context
+		state    statemachine.State
+	}{
+		{
+			name: "Test results not ready",
+			err:  ErrResultsNotReady,
+			setup: func(svc *agentService) {
+			},
+			state: Running,
+		},
+		{
+			name: "Test undeclared consumer",
+			err:  ErrUndeclaredConsumer,
+			setup: func(svc *agentService) {
+				svc.computation.ResultConsumers = []ResultConsumer{{UserKey: []byte("user")}}
+			},
+			ctxSetup: func(ctx context.Context) context.Context {
+				return ctx
+			},
+			state: ConsumingResults,
+		},
+		{
+			name: "Test results consumed and event sent",
+			err:  nil,
+			setup: func(svc *agentService) {
+				svc.computation.ResultConsumers = []ResultConsumer{{UserKey: []byte("key")}}
+			},
+			ctxSetup: func(ctx context.Context) context.Context {
+				return IndexToContext(ctx, 0)
+			},
+			state: ConsumingResults,
+		},
+	}
+
+	for _, tc := range cases {
+		events := new(mocks.Service)
+		events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
+		t.Run(tc.name, func(t *testing.T) {
+			ctx := metadata.NewIncomingContext(context.Background(),
+				metadata.Pairs(algorithm.AlgoTypeKey, "python", python.PyRuntimeKey, python.PyRuntime),
+			)
+
+			if tc.ctxSetup != nil {
+				ctx = tc.ctxSetup(ctx)
+			}
+
+			sm := new(smmocks.StateMachine)
+			sm.On("Start", ctx).Return(nil)
+			sm.On("GetState").Return(tc.state)
+			sm.On("SendEvent", mock.Anything).Return()
+
+			svc := &agentService{
+				sm:          sm,
+				eventSvc:    events,
+				provider:    &attestation.EmptyProvider{},
+				computation: testComputation(t),
+			}
+
+			go func() {
+				if err := svc.sm.Start(ctx); err != nil {
+					t.Errorf("Error starting state machine: %v", err)
+				}
+			}()
+			tc.setup(svc)
+			_, err := svc.Result(ctx)
+			t.Cleanup(func() {
+				_ = os.RemoveAll("datasets")
+				_ = os.RemoveAll("results")
+			})
+			assert.ErrorIs(t, err, tc.err, "expected %v, got %v", tc.err, err)
+		})
+	}
+}
+
+func TestAttestation(t *testing.T) {
+	provider := new(mocks2.Provider)
+
+	cases := []struct {
+		name       string
+		reportData [quoteprovider.Nonce]byte
+		nonce      [vtpm.Nonce]byte
+		rawQuote   []uint8
+		platform   attestation.PlatformType
+		err        error
+	}{
+		{
+			name:       "Test SNP attestation successful",
+			reportData: generateReportData(),
+			nonce:      [32]byte{},
+			rawQuote:   make([]uint8, 0),
+			platform:   attestation.SNP,
+			err:        nil,
+		},
+		{
+			name:       "Test SNP attestation failed",
+			reportData: generateReportData(),
+			nonce:      [32]byte{},
+			rawQuote:   nil,
+			platform:   attestation.SNP,
+			err:        ErrAttestationFailed,
+		},
+		{
+			name:       "Test vTPM attestation successful",
+			reportData: generateReportData(),
+			nonce:      [32]byte{},
+			rawQuote:   make([]uint8, 0),
+			platform:   attestation.VTPM,
+			err:        nil,
+		},
+		{
+			name:       "Test vTPM attestation failed",
+			reportData: generateReportData(),
+			nonce:      [32]byte{},
+			rawQuote:   nil,
+			platform:   attestation.VTPM,
+			err:        ErrAttestationVTpmFailed,
+		},
+		{
+			name:       "Test SNP-vTPM attestation successful",
+			reportData: generateReportData(),
+			nonce:      [32]byte{},
+			rawQuote:   make([]uint8, 0),
+			platform:   attestation.SNPvTPM,
+			err:        nil,
+		},
+		{
+			name:       "Test SNP-vTPM attestation failed",
+			reportData: generateReportData(),
+			nonce:      [32]byte{},
+			rawQuote:   nil,
+			platform:   attestation.SNPvTPM,
+			err:        ErrAttestationVTpmFailed,
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			events := new(mocks.Service)
+			events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
+			ctx := metadata.NewIncomingContext(context.Background(),
+				metadata.Pairs(algorithm.AlgoTypeKey, "python", python.PyRuntimeKey, python.PyRuntime),
+			)
+			ctx, cancel := context.WithCancel(ctx)
+			defer cancel()
+
+			getQuote := provider.On("TeeAttestation", mock.Anything).Return(tc.rawQuote, tc.err)
+			vtpmQuote := provider.On("VTpmAttestation", mock.Anything).Return(tc.rawQuote, tc.err)
+			snpVtpm := provider.On("Attestation", mock.Anything, mock.Anything).Return(tc.rawQuote, tc.err)
+			if tc.err != ErrAttestationFailed && tc.err != ErrAttestationVTpmFailed {
+				getQuote = provider.On("TeeAttestation", mock.Anything).Return(tc.nonce, nil)
+				vtpmQuote = provider.On("VTpmAttestation", mock.Anything).Return(tc.nonce[:], nil)
+				snpVtpm = provider.On("Attestation", mock.Anything, mock.Anything).Return(tc.nonce[:], nil)
+			}
+			defer getQuote.Unset()
+			defer vtpmQuote.Unset()
+			defer snpVtpm.Unset()
+
+			svc := New(ctx, mglog.NewMock(), events, provider, 0)
+			time.Sleep(300 * time.Millisecond)
+			_, err := svc.Attestation(ctx, tc.reportData, tc.nonce, tc.platform)
+			assert.True(t, errors.Contains(err, tc.err), "expected %v, got %v", tc.err, err)
+		})
+	}
+}
+
+func TestAzureAttestationToken(t *testing.T) {
+	provider := new(mocks2.Provider)
+	cases := []struct {
+		name  string
+		nonce [vtpm.Nonce]byte
+		token []byte
+		err   error
+	}{
+		{
+			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},
+			token: []byte{},
+			err:   ErrAttestationType,
+		},
+		{
+			name:  "Invalid attestation type",
+			nonce: [32]byte{7, 8, 9},
+			token: []byte{},
+			err:   ErrAttestationType,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			events := new(mocks.Service)
+			events.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
+			provider.On("AzureAttestationToken", tc.nonce[:]).Return(tc.token, tc.err)
+
+			ctx := context.Background()
+
+			svc := New(ctx, mglog.NewMock(), events, provider, 0)
+
+			_, err := svc.AzureAttestationToken(ctx, tc.nonce)
+			assert.True(t, errors.Contains(err, tc.err), "expected error %v, got %v", tc.err, err)
+		})
+	}
+}
+
+func generateReportData() [quoteprovider.Nonce]byte {
+	bytes := make([]byte, quoteprovider.Nonce)
+	_, err := rand.Read(bytes)
+	if err != nil {
+		log.Fatalf("Failed to generate random bytes: %v", err)
+	}
+	return [64]byte(bytes)
+}
+
+func testComputation(t *testing.T) Computation {
+	algo, err := os.ReadFile(algoPath)
+	require.NoError(t, err)
+
+	algoHash := sha3.Sum256(algo)
+
+	data, err := os.ReadFile(dataPath)
+	require.NoError(t, err)
+
+	dataHash := sha3.Sum256(data)
+
+	return Computation{
+		ID:              "1",
+		Name:            "sample computation",
+		Description:     "sample description",
+		Datasets:        []Dataset{{Hash: dataHash, UserKey: []byte("key"), Dataset: data, Filename: datasetFile}},
+		Algorithm:       Algorithm{Hash: algoHash, UserKey: []byte("key"), Algorithm: algo},
+		ResultConsumers: []ResultConsumer{{UserKey: []byte("key")}},
+	}
+}
+
+func TestStopComputation(t *testing.T) {
+	testDataDir := "test_datasets"
+	testResultsDir := "test_results"
+
+	cases := []struct {
+		name        string
+		setupDirs   bool
+		setupAlgo   bool
+		algoStopErr error
+		expectedErr error
+	}{
+		{
+			name:        "Stop computation successfully",
+			setupDirs:   true,
+			setupAlgo:   true,
+			algoStopErr: nil,
+			expectedErr: nil,
+		},
+		{
+			name:        "Stop computation with algorithm stop error",
+			setupDirs:   true,
+			setupAlgo:   true,
+			algoStopErr: fmt.Errorf("algorithm stop failed"),
+			expectedErr: fmt.Errorf("error stopping computation: algorithm stop failed"),
+		},
+		{
+			name:        "Stop computation without algorithm",
+			setupDirs:   true,
+			setupAlgo:   false,
+			algoStopErr: nil,
+			expectedErr: nil,
+		},
+		{
+			name:        "Stop computation with missing directories",
+			setupDirs:   false,
+			setupAlgo:   false,
+			algoStopErr: nil,
+			expectedErr: nil, // os.RemoveAll doesn't error on non-existing directories
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			events := new(mocks.Service)
+			events.On("SendEvent", mock.Anything, "Stopped", "Stopped", mock.Anything).Return()
+
+			ctx := context.Background()
+			ctx, cancel := context.WithCancel(ctx)
+			defer cancel()
+
+			svc := New(ctx, mglog.NewMock(), events, &attestation.EmptyProvider{}, 0).(*agentService)
+
+			svc.computation = Computation{
+				ID:   "test-computation",
+				Name: "test",
+			}
+
+			if tc.setupDirs {
+				err := os.MkdirAll(testDataDir, 0o755)
+				require.NoError(t, err)
+				err = os.MkdirAll(testResultsDir, 0o755)
+				require.NoError(t, err)
+			}
+
+			if tc.setupAlgo {
+				mockAlgo := new(algomocks.Algorithm)
+				mockAlgo.On("Stop").Return(tc.algoStopErr)
+				svc.algorithm = mockAlgo
+			}
+
+			err := svc.StopComputation(ctx)
+
+			if tc.expectedErr != nil {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), tc.expectedErr.Error())
+			} else {
+				assert.NoError(t, err)
+			}
+
+			assert.Equal(t, ReceivingManifest, svc.sm.GetState())
+			assert.Nil(t, svc.result)
+			assert.Nil(t, svc.runError)
+			assert.False(t, svc.resultsConsumed)
+
+			events.AssertExpectations(t)
+
+			_ = os.RemoveAll(testDataDir)
+			_ = os.RemoveAll(testResultsDir)
+		})
+	}
+}
+
+func TestStopComputationIntegration(t *testing.T) {
+	if testing.Short() {
+		t.Skip("Skipping integration test in short mode")
+	}
+
+	algo := []byte("#!/bin/bash\necho 'test algorithm'")
+	algoHash := sha3.Sum256(algo)
+
+	testDir := "test_integration"
+	err := os.MkdirAll(testDir, 0o755)
+	require.NoError(t, err)
+	defer os.RemoveAll(testDir)
+
+	algoFile := filepath.Join(testDir, "test_algo")
+	err = os.WriteFile(algoFile, algo, 0o755)
+	require.NoError(t, err)
+
+	events := new(mocks.Service)
+	events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
+	ctx := metadata.NewIncomingContext(context.Background(),
+		metadata.Pairs(algorithm.AlgoTypeKey, "bin"),
+	)
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	svc := New(ctx, mglog.NewMock(), events, &attestation.EmptyProvider{}, 0)
+
+	computation := Computation{
+		ID:   "integration-test",
+		Name: "Integration Test",
+		Algorithm: Algorithm{
+			Hash:      algoHash,
+			Algorithm: algo,
+		},
+	}
+
+	err = svc.InitComputation(ctx, computation)
+	require.NoError(t, err)
+
+	time.Sleep(100 * time.Millisecond)
+
+	err = svc.Algo(ctx, Algorithm{
+		Hash:      algoHash,
+		Algorithm: algo,
+	})
+	require.NoError(t, err)
+
+	time.Sleep(100 * time.Millisecond)
+
+	err = svc.StopComputation(ctx)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "ReceivingManifest", svc.State())
+}
+
+func TestStopComputationConcurrent(t *testing.T) {
+	events := new(mocks.Service)
+	events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return()
+
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	svc := New(ctx, mglog.NewMock(), events, &attestation.EmptyProvider{}, 0)
+
+	svc.(*agentService).computation = Computation{
+		ID:   "concurrent-test",
+		Name: "Concurrent Test",
+	}
+
+	const numGoroutines = 10
+	errChan := make(chan error, numGoroutines)
+
+	for i := 0; i < numGoroutines; i++ {
+		go func() {
+			err := svc.StopComputation(ctx)
+			errChan <- err
+		}()
+	}
+
+	var errors []error
+	for i := 0; i < numGoroutines; i++ {
+		err := <-errChan
+		if err != nil {
+			errors = append(errors, err)
+		}
+	}
+
+	assert.True(t, len(errors) < numGoroutines, "All StopComputation calls failed")
+}

agent/state.go

@@ -1,121 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package agent
-
-import (
-	"context"
-	"fmt"
-	"log/slog"
-	"sync"
-)
-
-//go:generate stringer -type=state
-type state int
-
-const (
-	idle state = iota
-	receivingManifest
-	receivingAlgorithm
-	receivingData
-	running
-	resultsReady
-	complete
-)
-
-type event int
-
-const (
-	start event = iota
-	manifestReceived
-	algorithmReceived
-	dataReceived
-	runComplete
-	resultsConsumed
-)
-
-// StateMachine represents the state machine.
-type StateMachine struct {
-	mu             sync.Mutex
-	State          state
-	EventChan      chan event
-	Transitions    map[state]map[event]state
-	StateFunctions map[state]func()
-	logger         *slog.Logger
-	wg             *sync.WaitGroup
-}
-
-// NewStateMachine creates a new StateMachine.
-func NewStateMachine(logger *slog.Logger) *StateMachine {
-	sm := &StateMachine{
-		State:          idle,
-		EventChan:      make(chan event),
-		Transitions:    make(map[state]map[event]state),
-		StateFunctions: make(map[state]func()),
-		logger:         logger,
-		wg:             &sync.WaitGroup{},
-	}
-
-	sm.Transitions[idle] = make(map[event]state)
-	sm.Transitions[idle][start] = receivingManifest
-
-	sm.Transitions[receivingManifest] = make(map[event]state)
-	sm.Transitions[receivingManifest][manifestReceived] = receivingAlgorithm
-
-	sm.Transitions[receivingAlgorithm] = make(map[event]state)
-	sm.Transitions[receivingAlgorithm][algorithmReceived] = receivingData
-
-	sm.Transitions[receivingData] = make(map[event]state)
-	sm.Transitions[receivingData][dataReceived] = running
-
-	sm.Transitions[running] = make(map[event]state)
-	sm.Transitions[running][runComplete] = resultsReady
-
-	sm.Transitions[resultsReady] = make(map[event]state)
-	sm.Transitions[resultsReady][resultsConsumed] = complete
-
-	return sm
-}
-
-// Start the state machine.
-func (sm *StateMachine) Start(ctx context.Context) {
-	sm.wg.Add(1)
-	defer sm.wg.Done()
-	for {
-		select {
-		case event := <-sm.EventChan:
-			nextState, valid := sm.Transitions[sm.GetState()][event]
-			if valid {
-				sm.mu.Lock()
-				sm.State = nextState
-				sm.mu.Unlock()
-				sm.logger.Debug(fmt.Sprintf("Transition: %v -> %v\n", sm.GetState(), nextState))
-			} else {
-				sm.logger.Error(fmt.Sprintf("Invalid transition: %v -> ???\n", sm.GetState()))
-			}
-			stateFunc, exists := sm.StateFunctions[sm.GetState()]
-			if exists {
-				go stateFunc()
-			}
-		case <-ctx.Done():
-			return
-		}
-	}
-}
-
-// SendEvent sends an event to the state machine.
-func (sm *StateMachine) SendEvent(event event) {
-	sm.EventChan <- event
-}
-
-func (sm *StateMachine) GetState() state {
-	sm.mu.Lock()
-	state := sm.State
-	sm.mu.Unlock()
-	return state
-}
-
-func (sm *StateMachine) SetState(state state) {
-	sm.mu.Lock()
-	sm.State = state
-	sm.mu.Unlock()
-}

agent/state_string.go

@@ -1,29 +0,0 @@
-// Code generated by "stringer -type=state"; DO NOT EDIT.
-
-package agent
-
-import "strconv"
-
-func _() {
-	// An "invalid array index" compiler error signifies that the constant values have changed.
-	// Re-run the stringer command to generate them again.
-	var x [1]struct{}
-	_ = x[idle-0]
-	_ = x[receivingManifest-1]
-	_ = x[receivingAlgorithm-2]
-	_ = x[receivingData-3]
-	_ = x[running-4]
-	_ = x[resultsReady-5]
-	_ = x[complete-6]
-}
-
-const _state_name = "idlereceivingManifestreceivingAlgorithmreceivingDatarunningresultsReadycomplete"
-
-var _state_index = [...]uint8{0, 4, 21, 39, 52, 59, 71, 79}
-
-func (i state) String() string {
-	if i < 0 || i >= state(len(_state_index)-1) {
-		return "state(" + strconv.FormatInt(int64(i), 10) + ")"
-	}
-	return _state_name[_state_index[i]:_state_index[i+1]]
-}

agent/state_test.go

@@ -4,58 +4,255 @@ package agent
 
 import (
 	"context"
-	"fmt"
+	sync "sync"
 	"testing"
+	"time"
 
-	mglog "github.com/absmach/magistrala/logger"
+	"github.com/ultravioletrs/cocos/agent/statemachine"
 )
 
-func TestStateMachineTransitions(t *testing.T) {
-	testCases := []struct {
-		fromState state
-		event     event
-		expected  state
+type MockState int
+
+type MockEvent int
+
+func (s MockState) String() string {
+	return []string{"State1", "State2", "State3"}[s]
+}
+
+func (e MockEvent) String() string {
+	return []string{"Event1", "Event2", "Event3"}[e]
+}
+
+const (
+	State1 MockState = iota
+	State2
+	State3
+)
+
+const (
+	Event1 MockEvent = iota
+	Event2
+	Event3
+)
+
+func TestNewStateMachine(t *testing.T) {
+	sm := statemachine.NewStateMachine(State1)
+	if sm == nil {
+		t.Fatal("NewStateMachine returned nil")
+	}
+	if sm.GetState() != State1 {
+		t.Errorf("Initial state not set correctly, got %v, want %v", sm.GetState(), State1)
+	}
+}
+
+func TestAddTransition(t *testing.T) {
+	sm := statemachine.NewStateMachine(State1)
+	sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	go func() {
+		if err := sm.Start(ctx); err != context.Canceled {
+		}
+	}()
+
+	time.Sleep(50 * time.Millisecond)
+
+	sm.SendEvent(Event1)
+
+	time.Sleep(50 * time.Millisecond)
+
+	if sm.GetState() != State2 {
+		t.Errorf("Transition not applied correctly, got state %v, want %v", sm.GetState(), State2)
+	}
+}
+
+func TestSetAction(t *testing.T) {
+	sm := statemachine.NewStateMachine(State1)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	sm.SetAction(State2, func(s statemachine.State) {
+		defer wg.Done()
+	})
+
+	sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
+	defer cancel()
+
+	go func() {
+		if err := sm.Start(ctx); err != context.Canceled {
+			t.Errorf("Start returned error: %v", err)
+		}
+	}()
+
+	time.Sleep(50 * time.Millisecond)
+
+	sm.SendEvent(Event1)
+
+	time.Sleep(50 * time.Millisecond)
+
+	wg.Wait()
+
+	if ctx.Err() != nil {
+		t.Error("Action was not called within the expected time")
+	}
+}
+
+func TestInvalidTransition(t *testing.T) {
+	sm := statemachine.NewStateMachine(State1)
+	sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+
+	errChan := make(chan error)
+	go func() {
+		errChan <- sm.Start(ctx)
+	}()
+
+	sm.SendEvent(Event2)
+
+	select {
+	case err := <-errChan:
+		if err == nil {
+			t.Errorf("Expected invalid transition error, got: %v", err)
+		}
+	case <-time.After(150 * time.Millisecond):
+		t.Error("Timeout waiting for invalid transition error")
+	}
+}
+
+func TestMultipleTransitions(t *testing.T) {
+	sm := statemachine.NewStateMachine(State1)
+	sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
+	sm.AddTransition(statemachine.Transition{From: State2, Event: Event2, To: State3})
+	sm.AddTransition(statemachine.Transition{From: State3, Event: Event3, To: State1})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
+	go func() {
+		if err := sm.Start(ctx); err != context.Canceled {
+		}
+	}()
+
+	time.Sleep(50 * time.Millisecond)
+
+	transitions := []struct {
+		event MockEvent
+		want  MockState
 	}{
-		{idle, start, receivingManifest},
-		{receivingManifest, manifestReceived, receivingAlgorithm},
-		{receivingAlgorithm, algorithmReceived, receivingData},
-		{receivingData, dataReceived, running},
-		{running, runComplete, resultsReady},
-		{resultsReady, resultsConsumed, complete},
+		{Event1, State2},
+		{Event2, State3},
+		{Event3, State1},
 	}
 
-	for _, testCase := range testCases {
-		t.Run(fmt.Sprintf("Transition from %v to %v", testCase.fromState, testCase.expected), func(t *testing.T) {
-			sm := NewStateMachine(mglog.NewMock())
-			ctx, cancel := context.WithCancel(context.Background())
-			go func() {
-				sm.Start(ctx)
-			}()
-			sm.wg.Wait()
-			sm.SetState(testCase.fromState)
+	for _, tt := range transitions {
+		sm.SendEvent(tt.event)
+		time.Sleep(50 * time.Millisecond)
 
-			sm.SendEvent(testCase.event)
-
-			if sm.GetState() != testCase.expected {
-				t.Errorf("Expected state %v after the event, but got %v", testCase.expected, sm.GetState())
-			}
-			close(sm.EventChan)
-			cancel()
-		})
+		if sm.GetState() != tt.want {
+			t.Errorf("After event %v, got state %v, want %v", tt.event, sm.GetState(), tt.want)
+		}
 	}
 }
 
-func TestStateMachineInvalidTransition(t *testing.T) {
-	sm := NewStateMachine(mglog.NewMock())
-	ctx, cancel := context.WithCancel(context.Background())
-	go sm.Start(ctx)
+func TestConcurrency(t *testing.T) {
+	sm := statemachine.NewStateMachine(State1)
+	sm.AddTransition(statemachine.Transition{From: State1, Event: Event1, To: State2})
+	sm.AddTransition(statemachine.Transition{From: State2, Event: Event2, To: State1})
 
-	sm.SetState(idle)
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
 
-	sm.SendEvent(dataReceived)
+	go func() {
+		if err := sm.Start(ctx); err == nil {
+			t.Errorf("Expected context error, got nil")
+		}
+	}()
 
-	if sm.State != idle {
-		t.Errorf("State should not change on an invalid event, but got %v", sm.State)
+	for i := 0; i < 100; i++ {
+		go func() {
+			sm.SendEvent(Event1)
+			sm.SendEvent(Event2)
+		}()
+	}
+
+	time.Sleep(400 * time.Millisecond)
+
+	finalState := sm.GetState()
+	if finalState != State1 && finalState != State2 {
+		t.Errorf("Unexpected final state: %v", finalState)
+	}
+}
+
+func TestAgentEventString(t *testing.T) {
+	tests := []struct {
+		event AgentEvent
+		want  string
+	}{
+		{Start, "Start"},
+		{ManifestReceived, "ManifestReceived"},
+		{AlgorithmReceived, "AlgorithmReceived"},
+		{DataReceived, "DataReceived"},
+		{RunComplete, "RunComplete"},
+		{ResultsConsumed, "ResultsConsumed"},
+		{RunFailed, "RunFailed"},
+		{AgentEvent(-1), "AgentEvent(-1)"},
+	}
+
+	for _, tt := range tests {
+		if got := tt.event.String(); got != tt.want {
+			t.Errorf("AgentEvent.String() = %v, want %v", got, tt.want)
+		}
+	}
+}
+
+func TestAgentStateString(t *testing.T) {
+	tests := []struct {
+		state AgentState
+		want  string
+	}{
+		{Idle, "Idle"},
+		{ReceivingManifest, "ReceivingManifest"},
+		{ReceivingAlgorithm, "ReceivingAlgorithm"},
+		{ReceivingData, "ReceivingData"},
+		{Running, "Running"},
+		{ConsumingResults, "ConsumingResults"},
+		{Complete, "Complete"},
+		{Failed, "Failed"},
+		{AgentState(-1), "AgentState(-1)"},
+	}
+
+	for _, tt := range tests {
+		if got := tt.state.String(); got != tt.want {
+			t.Errorf("AgentState.String() = %v, want %v", got, tt.want)
+		}
+	}
+}
+
+func TestStatusString(t *testing.T) {
+	tests := []struct {
+		status Status
+		want   string
+	}{
+		{IdleState, "IdleState"},
+		{InProgress, "InProgress"},
+		{Ready, "Ready"},
+		{Completed, "Completed"},
+		{Terminated, "Terminated"},
+		{Warning, "Warning"},
+		{Starting, "Starting"},
+		{Status(uint8(8)), "Status(8)"},
+	}
+
+	for _, tt := range tests {
+		if got := tt.status.String(); got != tt.want {
+			t.Errorf("Status.String() = %v, want %v", got, tt.want)
+		}
 	}
-	cancel()
 }

agent/statemachine/mocks/statemachine.go

@@ -0,0 +1,305 @@
+// 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/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
+}
+
+type StateMachine_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *StateMachine) EXPECT() *StateMachine_Expecter {
+	return &StateMachine_Expecter{mock: &_m.Mock}
+}
+
+// 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'
+type StateMachine_AddTransition_Call struct {
+	*mock.Call
+}
+
+// AddTransition is a helper method to define mock.On call
+//   - t statemachine.Transition
+func (_e *StateMachine_Expecter) AddTransition(t interface{}) *StateMachine_AddTransition_Call {
+	return &StateMachine_AddTransition_Call{Call: _e.mock.On("AddTransition", t)}
+}
+
+func (_c *StateMachine_AddTransition_Call) Run(run func(t statemachine.Transition)) *StateMachine_AddTransition_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 statemachine.Transition
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.Transition)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *StateMachine_AddTransition_Call) Return() *StateMachine_AddTransition_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *StateMachine_AddTransition_Call) RunAndReturn(run func(t statemachine.Transition)) *StateMachine_AddTransition_Call {
+	_c.Run(run)
+	return _c
+}
+
+// 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 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
+}
+
+// StateMachine_GetState_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetState'
+type StateMachine_GetState_Call struct {
+	*mock.Call
+}
+
+// GetState is a helper method to define mock.On call
+func (_e *StateMachine_Expecter) GetState() *StateMachine_GetState_Call {
+	return &StateMachine_GetState_Call{Call: _e.mock.On("GetState")}
+}
+
+func (_c *StateMachine_GetState_Call) Run(run func()) *StateMachine_GetState_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *StateMachine_GetState_Call) Return(state statemachine.State) *StateMachine_GetState_Call {
+	_c.Call.Return(state)
+	return _c
+}
+
+func (_c *StateMachine_GetState_Call) RunAndReturn(run func() statemachine.State) *StateMachine_GetState_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// 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'
+type StateMachine_Reset_Call struct {
+	*mock.Call
+}
+
+// Reset is a helper method to define mock.On call
+//   - initialState statemachine.State
+func (_e *StateMachine_Expecter) Reset(initialState interface{}) *StateMachine_Reset_Call {
+	return &StateMachine_Reset_Call{Call: _e.mock.On("Reset", initialState)}
+}
+
+func (_c *StateMachine_Reset_Call) Run(run func(initialState statemachine.State)) *StateMachine_Reset_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 statemachine.State
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.State)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *StateMachine_Reset_Call) Return() *StateMachine_Reset_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *StateMachine_Reset_Call) RunAndReturn(run func(initialState statemachine.State)) *StateMachine_Reset_Call {
+	_c.Run(run)
+	return _c
+}
+
+// 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'
+type StateMachine_SendEvent_Call struct {
+	*mock.Call
+}
+
+// SendEvent is a helper method to define mock.On call
+//   - event statemachine.Event
+func (_e *StateMachine_Expecter) SendEvent(event interface{}) *StateMachine_SendEvent_Call {
+	return &StateMachine_SendEvent_Call{Call: _e.mock.On("SendEvent", event)}
+}
+
+func (_c *StateMachine_SendEvent_Call) Run(run func(event statemachine.Event)) *StateMachine_SendEvent_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 statemachine.Event
+		if args[0] != nil {
+			arg0 = args[0].(statemachine.Event)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+func (_c *StateMachine_SendEvent_Call) Return() *StateMachine_SendEvent_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *StateMachine_SendEvent_Call) RunAndReturn(run func(event statemachine.Event)) *StateMachine_SendEvent_Call {
+	_c.Run(run)
+	return _c
+}
+
+// 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'
+type StateMachine_SetAction_Call struct {
+	*mock.Call
+}
+
+// SetAction is a helper method to define mock.On call
+//   - state statemachine.State
+//   - action statemachine.Action
+func (_e *StateMachine_Expecter) SetAction(state interface{}, action interface{}) *StateMachine_SetAction_Call {
+	return &StateMachine_SetAction_Call{Call: _e.mock.On("SetAction", state, action)}
+}
+
+func (_c *StateMachine_SetAction_Call) Run(run func(state statemachine.State, action statemachine.Action)) *StateMachine_SetAction_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		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
+}
+
+func (_c *StateMachine_SetAction_Call) Return() *StateMachine_SetAction_Call {
+	_c.Call.Return()
+	return _c
+}
+
+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 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 returnFunc, ok := ret.Get(0).(func(context.Context) error); ok {
+		r0 = returnFunc(ctx)
+	} else {
+		r0 = ret.Error(0)
+	}
+	return r0
+}
+
+// StateMachine_Start_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Start'
+type StateMachine_Start_Call struct {
+	*mock.Call
+}
+
+// Start is a helper method to define mock.On call
+//   - ctx context.Context
+func (_e *StateMachine_Expecter) Start(ctx interface{}) *StateMachine_Start_Call {
+	return &StateMachine_Start_Call{Call: _e.mock.On("Start", ctx)}
+}
+
+func (_c *StateMachine_Start_Call) Run(run func(ctx context.Context)) *StateMachine_Start_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		var arg0 context.Context
+		if args[0] != nil {
+			arg0 = args[0].(context.Context)
+		}
+		run(
+			arg0,
+		)
+	})
+	return _c
+}
+
+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(ctx context.Context) error) *StateMachine_Start_Call {
+	_c.Call.Return(run)
+	return _c
+}

agent/statemachine/state.go

@@ -0,0 +1,147 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package statemachine
+
+import (
+	"context"
+	"fmt"
+	"sync"
+)
+
+type State interface {
+	String() string
+}
+
+type Event interface {
+	String() string
+}
+
+type Action func(State)
+
+type Transition struct {
+	From  State
+	Event Event
+	To    State
+}
+
+type StateMachine interface {
+	AddTransition(t Transition)
+	SetAction(state State, action Action)
+	GetState() State
+	SendEvent(event Event)
+	Start(ctx context.Context) error
+	Reset(initialState State)
+}
+
+type stateMachine struct {
+	mu           sync.Mutex
+	currentState State
+	transitions  map[State]map[Event]State
+	actions      map[State]Action
+	eventChan    chan Event
+	resetChan    chan struct{}
+}
+
+func NewStateMachine(initialState State) StateMachine {
+	return &stateMachine{
+		currentState: initialState,
+		transitions:  make(map[State]map[Event]State),
+		actions:      make(map[State]Action),
+		eventChan:    make(chan Event),
+		resetChan:    make(chan struct{}),
+	}
+}
+
+func (sm *stateMachine) AddTransition(t Transition) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	if _, ok := sm.transitions[t.From]; !ok {
+		sm.transitions[t.From] = make(map[Event]State)
+	}
+	sm.transitions[t.From][t.Event] = t.To
+}
+
+func (sm *stateMachine) SetAction(state State, action Action) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	sm.actions[state] = action
+}
+
+func (sm *stateMachine) GetState() State {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+	return sm.currentState
+}
+
+func (sm *stateMachine) SendEvent(event Event) {
+	sm.mu.Lock()
+	eventChan := sm.eventChan
+	sm.mu.Unlock()
+
+	select {
+	case eventChan <- event:
+	default:
+		// Channel might be closed or full, ignore the event
+	}
+}
+
+func (sm *stateMachine) Start(ctx context.Context) error {
+	for {
+		sm.mu.Lock()
+		eventChan := sm.eventChan
+		resetChan := sm.resetChan
+		sm.mu.Unlock()
+
+		select {
+		case event := <-eventChan:
+			if err := sm.handleEvent(event); err != nil {
+				return err
+			}
+		case <-resetChan:
+			continue
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+}
+
+func (sm *stateMachine) Reset(initialState State) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	// Reset current state to initial state
+	sm.currentState = initialState
+
+	// Close the existing event channel to stop processing events
+	close(sm.eventChan)
+
+	// Close the reset channel to signal Start() to restart
+	close(sm.resetChan)
+
+	sm.eventChan = make(chan Event)
+	sm.resetChan = make(chan struct{})
+}
+
+func (sm *stateMachine) handleEvent(event Event) error {
+	sm.mu.Lock()
+	currentState := sm.currentState
+	nextState, valid := sm.transitions[currentState][event]
+	sm.mu.Unlock()
+
+	if !valid {
+		return fmt.Errorf("invalid transition: %v -> %v", currentState, event)
+	}
+
+	sm.mu.Lock()
+	sm.currentState = nextState
+	action := sm.actions[nextState]
+	sm.mu.Unlock()
+
+	if action != nil {
+		go action(nextState)
+	}
+
+	return nil
+}

agent/statemachine/state_test.go

@@ -0,0 +1,607 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package statemachine
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+)
+
+type testState string
+
+func (s testState) String() string {
+	return string(s)
+}
+
+type testEvent string
+
+func (e testEvent) String() string {
+	return string(e)
+}
+
+const (
+	StateIdle    testState = "idle"
+	StateRunning testState = "running"
+	StatePaused  testState = "paused"
+	StateStopped testState = "stopped"
+	StateError   testState = "error"
+)
+
+const (
+	EventStart testEvent = "start"
+	EventPause testEvent = "pause"
+	EventStop  testEvent = "stop"
+	EventReset testEvent = "reset"
+	EventError testEvent = "error"
+)
+
+func TestNewStateMachine(t *testing.T) {
+	tests := []struct {
+		name         string
+		initialState State
+		want         State
+	}{
+		{
+			name:         "create with idle state",
+			initialState: StateIdle,
+			want:         StateIdle,
+		},
+		{
+			name:         "create with running state",
+			initialState: StateRunning,
+			want:         StateRunning,
+		},
+		{
+			name:         "create with custom state",
+			initialState: testState("custom"),
+			want:         testState("custom"),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sm := NewStateMachine(tt.initialState)
+			if got := sm.GetState(); got != tt.want {
+				t.Errorf("NewStateMachine() initial state = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestStateMachine_AddTransition(t *testing.T) {
+	tests := []struct {
+		name        string
+		transitions []Transition
+		from        State
+		event       Event
+		expectTo    State
+		expectValid bool
+	}{
+		{
+			name: "single transition",
+			transitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+			},
+			from:        StateIdle,
+			event:       EventStart,
+			expectTo:    StateRunning,
+			expectValid: true,
+		},
+		{
+			name: "multiple transitions from same state",
+			transitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+				{From: StateIdle, Event: EventError, To: StateError},
+			},
+			from:        StateIdle,
+			event:       EventError,
+			expectTo:    StateError,
+			expectValid: true,
+		},
+		{
+			name: "overwrite existing transition",
+			transitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+				{From: StateIdle, Event: EventStart, To: StatePaused}, // Overwrite
+			},
+			from:        StateIdle,
+			event:       EventStart,
+			expectTo:    StatePaused,
+			expectValid: true,
+		},
+		{
+			name: "transition not found",
+			transitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+			},
+			from:        StateRunning,
+			event:       EventPause,
+			expectValid: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sm := NewStateMachine(StateIdle).(*stateMachine)
+
+			for _, transition := range tt.transitions {
+				sm.AddTransition(transition)
+			}
+
+			sm.mu.Lock()
+			nextState, valid := sm.transitions[tt.from][tt.event]
+			sm.mu.Unlock()
+
+			if valid != tt.expectValid {
+				t.Errorf("Transition validity = %v, want %v", valid, tt.expectValid)
+			}
+
+			if tt.expectValid && nextState != tt.expectTo {
+				t.Errorf("Transition destination = %v, want %v", nextState, tt.expectTo)
+			}
+		})
+	}
+}
+
+func TestStateMachine_SetAction(t *testing.T) {
+	tests := []struct {
+		name         string
+		state        State
+		action       Action
+		expectAction bool
+	}{
+		{
+			name:  "set action for state",
+			state: StateRunning,
+			action: func(s State) {
+			},
+			expectAction: true,
+		},
+		{
+			name:         "set nil action",
+			state:        StatePaused,
+			action:       nil,
+			expectAction: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sm := NewStateMachine(StateIdle).(*stateMachine)
+			sm.SetAction(tt.state, tt.action)
+
+			sm.mu.Lock()
+			action := sm.actions[tt.state]
+			sm.mu.Unlock()
+
+			if tt.expectAction && action == nil {
+				t.Error("Expected action to be set, but it was nil")
+			}
+			if !tt.expectAction && action != nil {
+				t.Error("Expected action to be nil, but it was set")
+			}
+		})
+	}
+}
+
+func TestStateMachine_GetState(t *testing.T) {
+	tests := []struct {
+		name         string
+		initialState State
+		transitions  []Transition
+		events       []Event
+		finalState   State
+	}{
+		{
+			name:         "get initial state",
+			initialState: StateIdle,
+			finalState:   StateIdle,
+		},
+		{
+			name:         "get state after transition",
+			initialState: StateIdle,
+			transitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+			},
+			events:     []Event{EventStart},
+			finalState: StateRunning,
+		},
+		{
+			name:         "get state after multiple transitions",
+			initialState: StateIdle,
+			transitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+				{From: StateRunning, Event: EventPause, To: StatePaused},
+				{From: StatePaused, Event: EventStart, To: StateRunning},
+			},
+			events:     []Event{EventStart, EventPause, EventStart},
+			finalState: StateRunning,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sm := NewStateMachine(tt.initialState)
+
+			for _, transition := range tt.transitions {
+				sm.AddTransition(transition)
+			}
+
+			smImpl := sm.(*stateMachine)
+			for _, event := range tt.events {
+				if err := smImpl.handleEvent(event); err != nil {
+					t.Fatalf("Failed to handle event %v: %v", event, err)
+				}
+			}
+
+			if got := sm.GetState(); got != tt.finalState {
+				t.Errorf("GetState() = %v, want %v", got, tt.finalState)
+			}
+		})
+	}
+}
+
+func TestStateMachine_Start(t *testing.T) {
+	tests := []struct {
+		name           string
+		initialState   State
+		transitions    []Transition
+		events         []Event
+		cancelAfter    time.Duration
+		expectError    bool
+		expectedStates []State
+	}{
+		{
+			name:         "start and cancel immediately",
+			initialState: StateIdle,
+			cancelAfter:  10 * time.Millisecond,
+			expectError:  true, // context.Canceled
+		},
+		{
+			name:         "process events then cancel",
+			initialState: StateIdle,
+			transitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+				{From: StateRunning, Event: EventStop, To: StateStopped},
+			},
+			events:         []Event{EventStart, EventStop},
+			cancelAfter:    100 * time.Millisecond,
+			expectError:    true, // context.Canceled
+			expectedStates: []State{StateRunning, StateStopped},
+		},
+		{
+			name:         "invalid transition error",
+			initialState: StateIdle,
+			transitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+			},
+			events:      []Event{EventPause}, // Invalid from StateIdle
+			cancelAfter: 50 * time.Millisecond,
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sm := NewStateMachine(tt.initialState)
+
+			for _, transition := range tt.transitions {
+				sm.AddTransition(transition)
+			}
+
+			var states []State
+			var mu sync.Mutex
+
+			for _, state := range tt.expectedStates {
+				sm.SetAction(state, func(s State) {
+					mu.Lock()
+					states = append(states, s)
+					mu.Unlock()
+				})
+			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+
+			errChan := make(chan error, 1)
+			go func() {
+				errChan <- sm.Start(ctx)
+			}()
+
+			time.Sleep(5 * time.Millisecond)
+
+			for _, event := range tt.events {
+				sm.SendEvent(event)
+				time.Sleep(5 * time.Millisecond)
+			}
+
+			time.Sleep(tt.cancelAfter)
+			cancel()
+
+			err := <-errChan
+
+			if tt.expectError && err == nil {
+				t.Error("Expected error but got none")
+			}
+			if !tt.expectError && err != nil {
+				t.Errorf("Unexpected error: %v", err)
+			}
+
+			time.Sleep(10 * time.Millisecond)
+
+			mu.Lock()
+			if len(states) != len(tt.expectedStates) {
+				t.Errorf("Expected %d state changes, got %d", len(tt.expectedStates), len(states))
+			}
+			for i, expectedState := range tt.expectedStates {
+				if i < len(states) && states[i] != expectedState {
+					t.Errorf("State change %d = %v, want %v", i, states[i], expectedState)
+				}
+			}
+			mu.Unlock()
+		})
+	}
+}
+
+func TestStateMachine_Reset(t *testing.T) {
+	tests := []struct {
+		name              string
+		initialState      State
+		resetState        State
+		setupTransitions  []Transition
+		eventsBeforeReset []Event
+		eventsAfterReset  []Event
+		expectedState     State
+	}{
+		{
+			name:          "reset to same state",
+			initialState:  StateIdle,
+			resetState:    StateIdle,
+			expectedState: StateIdle,
+		},
+		{
+			name:          "reset to different state",
+			initialState:  StateIdle,
+			resetState:    StateRunning,
+			expectedState: StateRunning,
+		},
+		{
+			name:         "reset after state changes",
+			initialState: StateIdle,
+			resetState:   StateIdle,
+			setupTransitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+			},
+			eventsBeforeReset: []Event{EventStart},
+			expectedState:     StateIdle,
+		},
+		{
+			name:         "reset and send new events",
+			initialState: StateIdle,
+			resetState:   StateIdle,
+			setupTransitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+				{From: StateRunning, Event: EventStop, To: StateStopped},
+			},
+			eventsBeforeReset: []Event{EventStart},
+			eventsAfterReset:  []Event{EventStart},
+			expectedState:     StateIdle,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sm := NewStateMachine(tt.initialState)
+			smImpl := sm.(*stateMachine)
+
+			for _, transition := range tt.setupTransitions {
+				sm.AddTransition(transition)
+			}
+
+			for _, event := range tt.eventsBeforeReset {
+				if err := smImpl.handleEvent(event); err != nil {
+					// Ignore errors for this test
+				}
+			}
+
+			sm.Reset(tt.resetState)
+
+			if got := sm.GetState(); got != tt.expectedState {
+				t.Errorf("State after reset = %v, want %v", got, tt.expectedState)
+			}
+
+			for _, event := range tt.eventsAfterReset {
+				sm.SendEvent(event)
+			}
+
+			// For events after reset, we can't easily check the channel length
+			// due to the synchronization changes, so we just verify the reset worked
+			if len(tt.eventsAfterReset) > 0 {
+				time.Sleep(5 * time.Millisecond)
+			}
+		})
+	}
+}
+
+func TestStateMachine_Reset_WithRunningStateMachine(t *testing.T) {
+	sm := NewStateMachine(StateIdle)
+	sm.AddTransition(Transition{From: StateIdle, Event: EventStart, To: StateRunning})
+	sm.AddTransition(Transition{From: StateRunning, Event: EventStop, To: StateStopped})
+
+	var stateChanges []State
+	var mu sync.Mutex
+
+	sm.SetAction(StateRunning, func(s State) {
+		mu.Lock()
+		stateChanges = append(stateChanges, s)
+		mu.Unlock()
+	})
+
+	sm.SetAction(StateStopped, func(s State) {
+		mu.Lock()
+		stateChanges = append(stateChanges, s)
+		mu.Unlock()
+	})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go func() {
+		if err := sm.Start(ctx); err != nil {
+		}
+	}()
+
+	// Give it time to start
+	time.Sleep(5 * time.Millisecond)
+
+	// Send an event
+	sm.SendEvent(EventStart)
+	time.Sleep(10 * time.Millisecond)
+
+	// Reset while running
+	sm.Reset(StateIdle)
+
+	// Verify state was reset
+	if got := sm.GetState(); got != StateIdle {
+		t.Errorf("State after reset = %v, want %v", got, StateIdle)
+	}
+
+	// Send another event after reset
+	sm.SendEvent(EventStart)
+	time.Sleep(10 * time.Millisecond)
+
+	mu.Lock()
+	changes := len(stateChanges)
+	mu.Unlock()
+
+	// Should have at least processed the first event
+	if changes < 1 {
+		t.Errorf("Expected at least 1 state change, got %d", changes)
+	}
+}
+
+func TestStateMachine_HandleEvent(t *testing.T) {
+	tests := []struct {
+		name             string
+		initialState     State
+		transitions      []Transition
+		event            Event
+		expectedState    State
+		expectError      bool
+		expectActionCall bool
+	}{
+		{
+			name:         "valid transition",
+			initialState: StateIdle,
+			transitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+			},
+			event:            EventStart,
+			expectedState:    StateRunning,
+			expectError:      false,
+			expectActionCall: true,
+		},
+		{
+			name:         "invalid transition",
+			initialState: StateIdle,
+			transitions: []Transition{
+				{From: StateRunning, Event: EventPause, To: StatePaused},
+			},
+			event:            EventStart,
+			expectedState:    StateIdle,
+			expectError:      true,
+			expectActionCall: false,
+		},
+		{
+			name:         "transition with no action",
+			initialState: StateIdle,
+			transitions: []Transition{
+				{From: StateIdle, Event: EventStart, To: StateRunning},
+			},
+			event:            EventStart,
+			expectedState:    StateRunning,
+			expectError:      false,
+			expectActionCall: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sm := NewStateMachine(tt.initialState).(*stateMachine)
+
+			for _, transition := range tt.transitions {
+				sm.AddTransition(transition)
+			}
+
+			var actionCalled bool
+			var mu sync.Mutex
+
+			if tt.expectActionCall {
+				sm.SetAction(tt.expectedState, func(s State) {
+					mu.Lock()
+					actionCalled = true
+					mu.Unlock()
+				})
+			}
+
+			err := sm.handleEvent(tt.event)
+
+			if tt.expectError && err == nil {
+				t.Error("Expected error but got none")
+			}
+			if !tt.expectError && err != nil {
+				t.Errorf("Unexpected error: %v", err)
+			}
+
+			if sm.GetState() != tt.expectedState {
+				t.Errorf("State after handleEvent = %v, want %v", sm.GetState(), tt.expectedState)
+			}
+
+			if tt.expectActionCall {
+				time.Sleep(10 * time.Millisecond)
+				mu.Lock()
+				called := actionCalled
+				mu.Unlock()
+				if !called {
+					t.Error("Expected action to be called but it wasn't")
+				}
+			}
+		})
+	}
+}
+
+func TestStateMachine_SendEvent_ThreadSafety(t *testing.T) {
+	sm := NewStateMachine(StateIdle)
+	sm.AddTransition(Transition{From: StateIdle, Event: EventStart, To: StateRunning})
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	go func() {
+		if err := sm.Start(ctx); err != nil {
+		}
+	}()
+
+	time.Sleep(5 * time.Millisecond)
+
+	var wg sync.WaitGroup
+	numGoroutines := 10
+	eventsPerGoroutine := 100
+
+	// Send events concurrently
+	for i := 0; i < numGoroutines; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < eventsPerGoroutine; j++ {
+				sm.SendEvent(EventStart)
+			}
+		}()
+	}
+
+	wg.Wait()
+	time.Sleep(10 * time.Millisecond)
+
+	// If we reach here without panicking, the test passes
+}

agent/status_string.go

@@ -0,0 +1,29 @@
+// Code generated by "stringer -type=Status"; DO NOT EDIT.
+
+package agent
+
+import "strconv"
+
+func _() {
+	// An "invalid array index" compiler error signifies that the constant values have changed.
+	// Re-run the stringer command to generate them again.
+	var x [1]struct{}
+	_ = x[IdleState-0]
+	_ = x[InProgress-1]
+	_ = x[Ready-2]
+	_ = x[Completed-3]
+	_ = x[Terminated-4]
+	_ = x[Warning-5]
+	_ = x[Starting-6]
+}
+
+const _Status_name = "IdleStateInProgressReadyCompletedTerminatedWarningStarting"
+
+var _Status_index = [...]uint8{0, 9, 19, 24, 33, 43, 50, 58}
+
+func (i Status) String() string {
+	if i >= Status(len(_Status_index)-1) {
+		return "Status(" + strconv.FormatInt(int64(i), 10) + ")"
+	}
+	return _Status_name[_Status_index[i]:_Status_index[i+1]]
+}

attestation.go

@@ -0,0 +1,9 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package cocosai
+
+import _ "embed"
+
+//go:embed attestation.bin
+var EmbeddedAttestation []byte

boilerplate.txt

@@ -0,0 +1,2 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0

cli/README.md

@@ -16,14 +16,14 @@ make cli
 Retrieves attestation information from the SEV guest and saves it to a file.
 To retrieve attestation from agent, use the following command:
 ```bash
-./build/cocos-cli agent attestation get '<report_data>'
+./build/cocos-cli attestation get '<report_data>'
 ```
 
 #### Validate attestation
 Validates the retrieved attestation information against a specified policy and checks its authenticity.
 To validate and verify attestation from agent, use the following command:
 ```bash
-./build/cocos-cli agent attestation validate '<attestation>' --report_data '<report_data>'
+./build/cocos-cli attestation validate '<attestation>' --report_data '<report_data>'
 ```
 ##### Flags
 - --config: Path to a JSON file containing the validation configuration. This can be used to override individual flags.
@@ -62,21 +62,60 @@ To validate and verify attestation from agent, use the following command:
 To upload an algorithm, use the following command:
 
 ```bash
-./build/cocos-cli agent algo /path/to/algorithm <private_key_file_path>
+./build/cocos-cli algo /path/to/algorithm <private_key_file_path>
 ```
 
+##### Flags
+- -a, --algorithm string        Algorithm type to run (default "bin")
+-     --args stringArray        Arguments to pass to the algorithm
+-     --python-runtime string   Python runtime to use (default "python3")
+- -r, --requirements string     Python requirements file
+
 #### Upload Dataset
 
 To upload a dataset, use the following command:
 
 ```bash
-./build/cocos-cli agent data /path/to/dataset.csv <private_key_file_path>
+./build/cocos-cli data /path/to/dataset.csv <private_key_file_path>
 ```
 
+Users can also upload directories which will be compressed on transit. Once received by agent they will be stored as compressed files or decompressed if the user passed the decompression argument.
+
+##### Flags
+- -d, --decompress   Decompress the dataset on agent
+
+
+
 #### Retrieve result
 
 To retrieve the computation result, use the following command:
 
 ```bash
-./build/cocos-cli agent result <private_key_file_path>
+./build/cocos-cli result <private_key_file_path>
+```
+
+#### Checksum
+When defining the manifest dataset and algorithm checksums are required. This can be done as below:
+
+```bash
+./build/cocos-cli checksum <path_to_dataset_or_algorithm>
+```
+
+#### Measure IGVM file
+We assume that our current working directory is the root of the cocos repository, both on the host machine and in the VM.
+
+`igvmmeasure` calculates the launch measurement for an IGVM file and can generate a signed version. It ensures integrity by precomputing the expected launch digest, which can be verified against the attestation report. The tool parses IGVM directives, outputs the measurement as a hex string, or creates a signed file for verification at guest launch.
+
+##### Example
+We measure an IGVM file using our measure command, run:
+
+```bash
+./build/cocos-cli igvmmeasure /path/to/igvm/file
+```
+
+The tool will parse the directives in the IGVM file, calculate the launch measurement, and output the computed digest. If successful, it prints the measurement to standard output.
+
+Here is a sample output
+```
+91c4929bec2d0ecf11a708e09f0a57d7d82208bcba2451564444a4b01c22d047995ca27f9053f86de4e8063e9f810548
 ```

cli/algorithm_test.go

@@ -0,0 +1,180 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"bytes"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/pem"
+	"errors"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
+)
+
+const algorithmFile = "test_algo_file.py"
+
+func generateRSAPrivateKeyFile(fileName string) error {
+	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		return err
+	}
+
+	privKeyFile, err := os.Create(fileName)
+	if err != nil {
+		return err
+	}
+	defer privKeyFile.Close()
+
+	privateKeyPEM := &pem.Block{
+		Type:  rsaKeyType,
+		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
+	}
+
+	err = pem.Encode(privKeyFile, privateKeyPEM)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func TestAlgorithmCmd(t *testing.T) {
+	tests := []struct {
+		name           string
+		setupMock      func(*mocks.SDK)
+		setupFiles     func() error
+		args           []string
+		connectErr     error
+		expectedOutput string
+		cleanup        func()
+	}{
+		{
+			name: "successful upload",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			},
+			setupFiles: func() error {
+				if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
+					return err
+				}
+				return generateRSAPrivateKeyFile(privateKeyFile)
+			},
+			args:           []string{algorithmFile, privateKeyFile},
+			expectedOutput: "Successfully uploaded algorithm",
+			cleanup: func() {
+				os.Remove(privateKeyFile)
+				os.Remove(algorithmFile)
+			},
+		},
+		{
+			name: "missing algorithm file",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			},
+			args:           []string{"non_existent_algo_file.py", privateKeyFile},
+			expectedOutput: "Error reading algorithm file",
+		},
+		{
+			name: "missing private key file",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			},
+			setupFiles: func() error {
+				return os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644)
+			},
+			args:           []string{algorithmFile, "non_existent_private_key.pem"},
+			expectedOutput: "Error reading private key file",
+			cleanup: func() {
+				os.Remove(algorithmFile)
+			},
+		},
+		{
+			name: "upload failure",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
+			},
+			setupFiles: func() error {
+				if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
+					return err
+				}
+				return generateRSAPrivateKeyFile(privateKeyFile)
+			},
+			args:           []string{algorithmFile, privateKeyFile},
+			expectedOutput: "Failed to upload algorithm",
+			cleanup: func() {
+				os.Remove(privateKeyFile)
+				os.Remove(algorithmFile)
+			},
+		},
+		{
+			name: "invalid private key",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Algo", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			},
+			setupFiles: func() error {
+				if err := os.WriteFile(algorithmFile, []byte("test algorithm"), 0o644); err != nil {
+					return err
+				}
+				privKeyFile, err := os.Create(privateKeyFile)
+				if err != nil {
+					return err
+				}
+				defer privKeyFile.Close()
+				_, err = privKeyFile.WriteString("invalid private key")
+				return err
+			},
+			args:           []string{algorithmFile, privateKeyFile},
+			expectedOutput: "Error decoding private key",
+			cleanup: func() {
+				os.Remove(algorithmFile)
+				os.Remove(privateKeyFile)
+			},
+		},
+		{
+			name: "connection error",
+			setupMock: func(m *mocks.SDK) {
+			},
+			connectErr:     errors.New("failed to connect to agent"),
+			args:           []string{algorithmFile, privateKeyFile},
+			expectedOutput: "Failed to connect to agent",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockSDK := new(mocks.SDK)
+			if tt.setupMock != nil {
+				tt.setupMock(mockSDK)
+			}
+
+			testCLI := CLI{
+				agentSDK:   mockSDK,
+				connectErr: tt.connectErr,
+			}
+
+			if tt.setupFiles != nil {
+				err := tt.setupFiles()
+				require.NoError(t, err)
+			}
+
+			cmd := testCLI.NewAlgorithmCmd()
+			buf := new(bytes.Buffer)
+			cmd.SetOut(buf)
+			cmd.SetArgs(tt.args)
+			err := cmd.Execute()
+			require.NoError(t, err)
+
+			require.Contains(t, buf.String(), tt.expectedOutput)
+
+			if tt.cleanup != nil {
+				tt.cleanup()
+			}
+		})
+	}
+}

cli/algorithms.go

@@ -3,48 +3,94 @@
 package cli
 
 import (
+	"context"
 	"encoding/pem"
-	"log"
 	"os"
 
+	"github.com/fatih/color"
 	"github.com/spf13/cobra"
-	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/agent/algorithm"
+	"github.com/ultravioletrs/cocos/agent/algorithm/python"
+	"google.golang.org/grpc/metadata"
+)
+
+var (
+	pythonRuntime    string
+	algoType         string
+	requirementsFile string
+	algoArgs         []string
 )
 
 func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
-	return &cobra.Command{
+	cmd := &cobra.Command{
 		Use:     "algo",
 		Short:   "Upload an algorithm binary",
 		Example: "algo <algo_file> <private_key_file_path>",
 		Args:    cobra.ExactArgs(2),
 		Run: func(cmd *cobra.Command, args []string) {
-			algorithmFile := args[0]
-
-			log.Println("Uploading algorithm binary:", algorithmFile)
-
-			algorithm, err := os.ReadFile(algorithmFile)
-			if err != nil {
-				log.Fatalf("Error reading algorithm file: %v", err)
+			if cli.connectErr != nil {
+				printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
+				return
 			}
 
-			algoReq := agent.Algorithm{
-				Algorithm: algorithm,
+			algorithmFile := args[0]
+
+			cmd.Println("Uploading algorithm file:", algorithmFile)
+
+			algorithm, err := os.Open(algorithmFile)
+			if err != nil {
+				printError(cmd, "Error reading algorithm file: %v ❌ ", err)
+				return
+			}
+
+			defer algorithm.Close()
+
+			var req *os.File
+			if requirementsFile != "" {
+				req, err = os.Open(requirementsFile)
+				if err != nil {
+					printError(cmd, "Error reading requirments file: %v ❌ ", err)
+					return
+				}
+				defer req.Close()
 			}
 
 			privKeyFile, err := os.ReadFile(args[1])
 			if err != nil {
-				log.Fatalf("Error reading private key file: %v", err)
+				printError(cmd, "Error reading private key file: %v ❌ ", err)
+				return
 			}
 
 			pemBlock, _ := pem.Decode(privKeyFile)
 
-			privKey := decodeKey(pemBlock)
-
-			if err := cli.agentSDK.Algo(cmd.Context(), algoReq, privKey); err != nil {
-				log.Fatalf("Error uploading algorithm with error: %v", err)
+			privKey, err := decodeKey(pemBlock)
+			if err != nil {
+				printError(cmd, "Error decoding private key: %v ❌ ", err)
+				return
 			}
 
-			log.Println("Successfully uploaded algorithm")
+			ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
+
+			if err := cli.agentSDK.Algo(addAlgoMetadata(ctx), algorithm, req, privKey); err != nil {
+				printError(cmd, "Failed to upload algorithm due to error: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println(color.New(color.FgGreen).Sprint("Successfully uploaded algorithm! ✔ "))
 		},
 	}
+
+	cmd.Flags().StringVarP(&algoType, "algorithm", "a", string(algorithm.AlgoTypeBin), "Algorithm type to run")
+	cmd.Flags().StringVar(&pythonRuntime, "python-runtime", python.PyRuntime, "Python runtime to use")
+	cmd.Flags().StringVarP(&requirementsFile, "requirements", "r", "", "Python requirements file")
+	cmd.Flags().StringArrayVar(&algoArgs, "args", []string{}, "Arguments to pass to the algorithm")
+
+	return cmd
+}
+
+func addAlgoMetadata(ctx context.Context) context.Context {
+	ctx = algorithm.AlgorithmTypeToContext(ctx, algoType)
+	ctx = algorithm.AlgorithmArgsToContext(ctx, algoArgs)
+	ctx = python.PythonRunTimeToContext(ctx, pythonRuntime)
+	return ctx
 }

cli/attestation_policy.go

@@ -0,0 +1,425 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"bytes"
+	"crypto/sha512"
+	"encoding/base64"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"os"
+	"strconv"
+	"strings"
+
+	"github.com/absmach/supermq/pkg/errors"
+	"github.com/google/go-sev-guest/proto/check"
+	"github.com/google/go-tpm-tools/proto/attest"
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"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"
+)
+
+type fieldType int
+
+const (
+	measurementField fieldType = iota
+	hostDataField
+)
+
+const (
+	// 0o744 file permission gives RWX permission to the user and only the R permission to others.
+	filePermission = 0o744
+	// Length of the expected host data and measurement field in bytes.
+	hostDataLength    = 32
+	measurementLength = 48
+)
+
+var (
+	errDecode                              = errors.New("base64 string could not be decoded")
+	errDataLength                          = errors.New("data does not have an adequate length")
+	errReadingAttestationPolicyFile        = errors.New("error while reading the attestation policy file")
+	errUnmarshalJSON                       = errors.New("failed to unmarshal json")
+	errMarshalJSON                         = errors.New("failed to marshal json")
+	errWriteFile                           = errors.New("failed to write to file")
+	errAttestationPolicyField              = errors.New("the specified field type does not exist in the attestation policy")
+	errReadingManifestFile                 = errors.New("error while reading manifest file")
+	errDecodeHex                           = errors.New("error decoding hex string")
+	policy                          uint64 = 196639
+	isJsonAttestation               bool
+)
+
+func (cli *CLI) NewAttestationPolicyCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:   "policy [command]",
+		Short: "Change attestation policy",
+		Run: func(cmd *cobra.Command, args []string) {
+			fmt.Printf("Change attestation policy\n\n")
+			fmt.Printf("Usage:\n  %s [command]\n\n", cmd.CommandPath())
+			fmt.Printf("Available Commands:\n")
+
+			// Filter out "completion" command
+			availableCommands := make([]*cobra.Command, 0)
+			for _, subCmd := range cmd.Commands() {
+				if subCmd.Name() != "completion" {
+					availableCommands = append(availableCommands, subCmd)
+				}
+			}
+
+			for _, subCmd := range availableCommands {
+				fmt.Printf("  %-15s%s\n", subCmd.Name(), subCmd.Short)
+			}
+
+			fmt.Printf("\nFlags:\n")
+			cmd.Flags().VisitAll(func(flag *pflag.Flag) {
+				fmt.Printf("  -%s, --%s %s\n", flag.Shorthand, flag.Name, flag.Usage)
+			})
+			fmt.Printf("\nUse \"%s [command] --help\" for more information about a command.\n", cmd.CommandPath())
+		},
+	}
+}
+
+func (cli *CLI) NewAddMeasurementCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:     "measurement",
+		Short:   "Add measurement to the attestation policy file. The value should be in base64. The second parameter is attestation_policy.json file",
+		Example: "measurement <measurement> <attestation_policy.json>",
+		Args:    cobra.ExactArgs(2),
+		Run: func(cmd *cobra.Command, args []string) {
+			if err := changeAttestationConfiguration(args[1], args[0], measurementLength, measurementField); err != nil {
+				printError(cmd, "Error could not change measurement data: %v ❌ ", err)
+				return
+			}
+		},
+	}
+}
+
+func (cli *CLI) NewAddHostDataCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:     "hostdata",
+		Short:   "Add host data to the attestation policy file. The value should be in base64. The second parameter is attestation_policy.json file",
+		Example: "hostdata <host-data> <attestation_policy.json>",
+		Args:    cobra.ExactArgs(2),
+		Run: func(cmd *cobra.Command, args []string) {
+			if err := changeAttestationConfiguration(args[1], args[0], hostDataLength, hostDataField); err != nil {
+				printError(cmd, "Error could not change host data: %v ❌ ", err)
+				return
+			}
+		},
+	}
+}
+
+func (cli *CLI) NewGCPAttestationPolicy() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "gcp",
+		Short:   "Get attestation policy for GCP CVM",
+		Example: `gcp <bin_vtmp_attestation_report_file> <vcpu_count>`,
+		Args:    cobra.ExactArgs(2),
+		Run: func(cmd *cobra.Command, args []string) {
+			attestationBin, err := os.ReadFile(args[0])
+			if err != nil {
+				printError(cmd, "Error reading attestation report file: %v ❌ ", err)
+				return
+			}
+
+			vcpuCount, err := strconv.Atoi(args[1])
+			if err != nil {
+				printError(cmd, "Error converting vCPU count to integer: %v ❌ ", err)
+				return
+			}
+
+			attestation := &attest.Attestation{}
+
+			if 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()
+
+			measurement, err := gcp.Extract384BitMeasurement(attestationPB)
+			if err != nil {
+				printError(cmd, "Error extracting 384-bit measurement: %v ❌ ", err)
+				return
+			}
+
+			launchEndorsement, err := gcp.GetLaunchEndorsement(cmd.Context(), measurement)
+			if err != nil {
+				printError(cmd, "Error getting launch endorsement: %v ❌ ", err)
+				return
+			}
+
+			attestationPolicy, err := gcp.GenerateAttestationPolicy(launchEndorsement, uint32(vcpuCount))
+			if err != nil {
+				printError(cmd, "Error generating attestation policy: %v ❌ ", err)
+				return
+			}
+
+			attestationPolicyJson, err := json.MarshalIndent(attestationPolicy, "", "  ")
+			if err != nil {
+				printError(cmd, "Error marshaling attestation policy: %v ❌ ", err)
+				return
+			}
+
+			if err := os.WriteFile("attestation_policy.json", attestationPolicyJson, filePermission); err != nil {
+				printError(cmd, "Error writing attestation policy file: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println("Attestation policy file generated successfully ✅")
+		},
+	}
+
+	cmd.Flags().BoolVarP(&isJsonAttestation, "json", "j", false, "Use JSON attestation report instead of binary")
+	return cmd
+}
+
+func (cli *CLI) NewDownloadGCPOvmfFile() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "download",
+		Short:   "Download GCP OVMF file",
+		Example: `download <bin_vtmp_attestation_report_file>`,
+		Args:    cobra.ExactArgs(1),
+		Run: func(cmd *cobra.Command, args []string) {
+			attestationBin, err := os.ReadFile(args[0])
+			if err != nil {
+				printError(cmd, "Error reading attestation report file: %v ❌ ", err)
+				return
+			}
+
+			attestation := &attest.Attestation{}
+
+			if 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()
+
+			measurement, err := gcp.Extract384BitMeasurement(attestationPB)
+			if err != nil {
+				printError(cmd, "Error extracting 384-bit measurement: %v ❌ ", err)
+				return
+			}
+
+			launchEndorsement, err := gcp.GetLaunchEndorsement(cmd.Context(), measurement)
+			if err != nil {
+				printError(cmd, "Error getting launch endorsement: %v ❌ ", err)
+				return
+			}
+
+			ovmf, err := gcp.DownloadOvmfFile(cmd.Context(), fmt.Sprintf("%x", launchEndorsement.Digest))
+			if err != nil {
+				printError(cmd, "Error downloading OVMF file: %v ❌ ", err)
+				return
+			}
+
+			sum384 := sha512.Sum384(ovmf)
+
+			if !bytes.Equal(sum384[:], launchEndorsement.Digest) {
+				printError(cmd, "Error OVMF file does not match the measurement: %v ❌ ", fmt.Errorf("digest mismatch"))
+			} else {
+				cmd.Println("OVMF firmware in vm is unmodified ✅")
+			}
+
+			if err := os.WriteFile("ovmf.fd", ovmf, filePermission); err != nil {
+				printError(cmd, "Error writing OVMF file: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println("OVMF file downloaded successfully ✅")
+		},
+	}
+
+	cmd.Flags().BoolVarP(&isJsonAttestation, "json", "j", false, "Use JSON attestation report instead of binary")
+	return cmd
+}
+
+func (cli *CLI) NewAzureAttestationPolicy() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "azure",
+		Short:   "Get attestation policy for Azure CVM",
+		Example: `azure <azure_maa_token_file> <product_name>`,
+		Args:    cobra.ExactArgs(2),
+		Run: func(cmd *cobra.Command, args []string) {
+			token, err := os.ReadFile(args[0])
+			if err != nil {
+				printError(cmd, "Error reading attestation report file: %v ❌ ", err)
+				return
+			}
+
+			product := args[1]
+
+			config, err := azure.GenerateAttestationPolicy(string(token), product, policy)
+			if err != nil {
+				printError(cmd, "Error generating attestation policy: %v ❌ ", err)
+				return
+			}
+
+			attestationPolicyJson, err := json.MarshalIndent(&config, "", " ")
+			if err != nil {
+				printError(cmd, "Error marshaling attestation policy: %v ❌ ", err)
+				return
+			}
+
+			if err := os.WriteFile("attestation_policy.json", attestationPolicyJson, filePermission); err != nil {
+				printError(cmd, "Error writing attestation policy file: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println("Attestation policy file generated successfully ✅")
+		},
+	}
+
+	cmd.Flags().Uint64Var(
+		&policy,
+		"policy",
+		policy,
+		"Policy of the guest CVM",
+	)
+
+	return cmd
+}
+
+func (cli *CLI) NewExtendWithManifestCmd() *cobra.Command {
+	return &cobra.Command{
+		Use:     "extend",
+		Short:   "Extends PCR16 with computation manifests. The first parameter is path to attestation policy file. The rest of the parameters are paths to computation manifest files.",
+		Example: "extend <attestation_policy_file_path> <computation_manifest_file_path> [<computation_manifest_file_path> ...]",
+		Args:    cobra.MinimumNArgs(2),
+		Run: func(cmd *cobra.Command, args []string) {
+			attestationPolicyFilePath := args[0]
+			manifestPaths := args[1:]
+			if err := extendWithManifest(attestationPolicyFilePath, manifestPaths); err != nil {
+				printError(cmd, "Error could not extend PCR16: %v ❌ ", err)
+				return
+			}
+		},
+	}
+}
+
+func changeAttestationConfiguration(fileName, base64Data string, expectedLength int, field fieldType) error {
+	data, err := base64.StdEncoding.DecodeString(base64Data)
+	if err != nil {
+		return errDecode
+	}
+
+	if len(data) != expectedLength {
+		return errDataLength
+	}
+
+	ac := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
+
+	f, err := os.ReadFile(fileName)
+	if err != nil {
+		return errors.Wrap(errReadingAttestationPolicyFile, err)
+	}
+
+	if err = vtpm.ReadPolicyFromByte(f, &ac); err != nil {
+		return errors.Wrap(errUnmarshalJSON, err)
+	}
+
+	if ac.Config.Policy == nil {
+		ac.Config.Policy = &check.Policy{}
+	}
+
+	switch field {
+	case measurementField:
+		ac.Config.Policy.Measurement = data
+	case hostDataField:
+		ac.Config.Policy.HostData = data
+	default:
+		return errAttestationPolicyField
+	}
+
+	fileJson, err := vtpm.ConvertPolicyToJSON(&ac)
+	if err != nil {
+		return errors.Wrap(errMarshalJSON, err)
+	}
+	if err = os.WriteFile(fileName, fileJson, filePermission); err != nil {
+		return errors.Wrap(errWriteFile, err)
+	}
+	return nil
+}
+
+func extendWithManifest(attestationPolicyPath string, manifestPaths []string) error {
+	attestationConfig := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
+
+	attestationPolicyFileData, err := os.ReadFile(attestationPolicyPath)
+	if err != nil {
+		return errors.Wrap(errReadingAttestationPolicyFile, err)
+	}
+
+	if err = vtpm.ReadPolicyFromByte(attestationPolicyFileData, &attestationConfig); err != nil {
+		return errors.Wrap(errUnmarshalJSON, err)
+	}
+
+	for _, manifestPath := range manifestPaths {
+		manifest, err := os.ReadFile(manifestPath)
+		if err != nil {
+			return errors.Wrap(errReadingManifestFile, err)
+		}
+
+		manifestSha256 := sha512.Sum512_256(manifest)
+		manifestSha384 := sha512.Sum384(manifest)
+
+		data256, exists256 := attestationConfig.PCRValues.Sha256["16"]
+
+		if !exists256 {
+			data256 = strings.Repeat("0", 64) // 32 bytes in hex
+		}
+
+		byteData256, err := hex.DecodeString(data256)
+		if err != nil {
+			return errors.Wrap(errDecodeHex, err)
+		}
+
+		newByteData256 := sha512.Sum512_256(append(byteData256, manifestSha256[:]...))
+
+		data384, exists384 := attestationConfig.PCRValues.Sha384["16"]
+
+		if !exists384 {
+			data384 = strings.Repeat("0", 96) // 48 bytes in hex
+		}
+
+		byteData384, err := hex.DecodeString(data384)
+		if err != nil {
+			return errors.Wrap(errDecodeHex, err)
+		}
+
+		newByteData384 := sha512.Sum384(append(byteData384, manifestSha384[:]...))
+
+		attestationConfig.PCRValues.Sha256["16"] = hex.EncodeToString(newByteData256[:])
+		attestationConfig.PCRValues.Sha384["16"] = hex.EncodeToString(newByteData384[:])
+	}
+
+	attestationPolicyJSON, err := vtpm.ConvertPolicyToJSON(&attestationConfig)
+	if err != nil {
+		return errors.Wrap(errMarshalJSON, err)
+	}
+	if err = os.WriteFile(attestationPolicyPath, attestationPolicyJSON, filePermission); err != nil {
+		return errors.Wrap(errWriteFile, err)
+	}
+
+	return nil
+}

cli/attestation_policy_test.go

@@ -0,0 +1,470 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"bytes"
+	"encoding/base64"
+	"os"
+	"testing"
+
+	"github.com/google/go-sev-guest/proto/check"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
+)
+
+func TestChangeAttestationConfiguration(t *testing.T) {
+	tmpfile, err := os.CreateTemp("", "attestation_policy.json")
+	require.NoError(t, err)
+	defer os.Remove(tmpfile.Name())
+
+	initialConfig := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
+
+	initialJSON, err := vtpm.ConvertPolicyToJSON(&initialConfig)
+	require.NoError(t, err)
+	err = os.WriteFile(tmpfile.Name(), initialJSON, 0o644)
+	require.NoError(t, err)
+
+	tests := []struct {
+		name           string
+		base64Data     string
+		expectedLength int
+		field          fieldType
+		expectError    bool
+		errorType      error
+	}{
+		{
+			name:           "Valid Measurement",
+			base64Data:     base64.StdEncoding.EncodeToString(make([]byte, measurementLength)),
+			expectedLength: measurementLength,
+			field:          measurementField,
+			expectError:    false,
+		},
+		{
+			name:           "Valid Host Data",
+			base64Data:     base64.StdEncoding.EncodeToString(make([]byte, hostDataLength)),
+			expectedLength: hostDataLength,
+			field:          hostDataField,
+			expectError:    false,
+		},
+		{
+			name:           "Invalid Base64",
+			base64Data:     "Invalid Base64",
+			expectedLength: measurementLength,
+			field:          measurementField,
+			expectError:    true,
+			errorType:      errDecode,
+		},
+		{
+			name:           "Invalid Data Length",
+			base64Data:     base64.StdEncoding.EncodeToString(make([]byte, measurementLength-1)),
+			expectedLength: measurementLength,
+			field:          measurementField,
+			expectError:    true,
+			errorType:      errDataLength,
+		},
+		{
+			name:           "Invalid Field Type",
+			base64Data:     base64.StdEncoding.EncodeToString(make([]byte, measurementLength)),
+			expectedLength: measurementLength,
+			field:          fieldType(999),
+			expectError:    true,
+			errorType:      errAttestationPolicyField,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := changeAttestationConfiguration(tmpfile.Name(), tt.base64Data, tt.expectedLength, tt.field)
+
+			if tt.expectError {
+				assert.Error(t, err)
+				assert.ErrorIs(t, err, tt.errorType)
+			} else {
+				assert.NoError(t, err)
+
+				content, err := os.ReadFile(tmpfile.Name())
+				require.NoError(t, err)
+
+				ap := attestation.Config{Config: &check.Config{RootOfTrust: &check.RootOfTrust{}, Policy: &check.Policy{}}, PcrConfig: &attestation.PcrConfig{}}
+				err = vtpm.ReadPolicyFromByte(content, &ap)
+				require.NoError(t, err)
+
+				decodedData, _ := base64.StdEncoding.DecodeString(tt.base64Data)
+				if tt.field == measurementField {
+					assert.Equal(t, decodedData, ap.Config.Policy.Measurement)
+				} else if tt.field == hostDataField {
+					assert.Equal(t, decodedData, ap.Config.Policy.HostData)
+				}
+			}
+		})
+	}
+}
+
+func TestNewAttestationPolicyCmd(t *testing.T) {
+	cli := &CLI{}
+	cmd := cli.NewAttestationPolicyCmd()
+
+	assert.Equal(t, "policy [command]", cmd.Use)
+	assert.Equal(t, "Change attestation policy", cmd.Short)
+	assert.NotNil(t, cmd.Run)
+}
+
+func TestNewAddMeasurementCmd(t *testing.T) {
+	cli := &CLI{}
+	cmd := cli.NewAddMeasurementCmd()
+
+	assert.Equal(t, "measurement", cmd.Use)
+	assert.Equal(t, "Add measurement to the attestation policy file. The value should be in base64. The second parameter is attestation_policy.json file", cmd.Short)
+	assert.Equal(t, "measurement <measurement> <attestation_policy.json>", cmd.Example)
+	assert.NotNil(t, cmd.Run)
+}
+
+func TestNewAddHostDataCmd(t *testing.T) {
+	cli := &CLI{}
+	cmd := cli.NewAddHostDataCmd()
+
+	assert.Equal(t, "hostdata", cmd.Use)
+	assert.Equal(t, "Add host data to the attestation policy file. The value should be in base64. The second parameter is attestation_policy.json file", cmd.Short)
+	assert.Equal(t, "hostdata <host-data> <attestation_policy.json>", cmd.Example)
+	assert.NotNil(t, cmd.Run)
+}
+
+func TestChangeAttestationConfigurationFileErrors(t *testing.T) {
+	t.Run("File Not Found", func(t *testing.T) {
+		err := changeAttestationConfiguration("nonexistent.json", base64.StdEncoding.EncodeToString(make([]byte, measurementLength)), measurementLength, measurementField)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "error while reading the attestation policy file")
+	})
+
+	t.Run("Invalid JSON Content", func(t *testing.T) {
+		tmpfile, err := os.CreateTemp("", "invalid.json")
+		require.NoError(t, err)
+		defer os.Remove(tmpfile.Name())
+
+		err = os.WriteFile(tmpfile.Name(), []byte("invalid json"), 0o644)
+		require.NoError(t, err)
+
+		err = changeAttestationConfiguration(tmpfile.Name(), base64.StdEncoding.EncodeToString(make([]byte, measurementLength)), measurementLength, measurementField)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "failed to unmarshal json")
+	})
+}
+
+func TestNewGCPAttestationPolicy(t *testing.T) {
+	cli := &CLI{}
+	cmd := cli.NewGCPAttestationPolicy()
+
+	assert.Equal(t, "gcp", cmd.Use)
+	assert.Equal(t, "Get attestation policy for GCP CVM", cmd.Short)
+	assert.Equal(t, "gcp <bin_vtmp_attestation_report_file> <vcpu_count>", cmd.Example)
+	assert.NotNil(t, cmd.Run)
+
+	t.Run("File Not Found", func(t *testing.T) {
+		cmd.SetArgs([]string{"nonexistent.bin", "4"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "Error reading attestation report file")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Invalid vCPU Count", func(t *testing.T) {
+		tmpfile, err := os.CreateTemp("", "attestation.bin")
+		require.NoError(t, err)
+		defer os.Remove(tmpfile.Name())
+
+		err = os.WriteFile(tmpfile.Name(), []byte("dummy content"), 0o644)
+		require.NoError(t, err)
+
+		cmd.SetArgs([]string{tmpfile.Name(), "invalid"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err = cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "Error converting vCPU count to integer")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Invalid Attestation Data", func(t *testing.T) {
+		tmpfile, err := os.CreateTemp("", "attestation.bin")
+		require.NoError(t, err)
+		defer os.Remove(tmpfile.Name())
+
+		err = os.WriteFile(tmpfile.Name(), []byte("invalid protobuf data"), 0o644)
+		require.NoError(t, err)
+
+		cmd.SetArgs([]string{tmpfile.Name(), "4"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err = cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "Error unmarshaling attestation report")
+		assert.Contains(t, output, "❌")
+	})
+}
+
+func TestNewDownloadGCPOvmfFile(t *testing.T) {
+	cli := &CLI{}
+	cmd := cli.NewDownloadGCPOvmfFile()
+
+	assert.Equal(t, "download", cmd.Use)
+	assert.Equal(t, "Download GCP OVMF file", cmd.Short)
+	assert.Equal(t, "download <bin_vtmp_attestation_report_file>", cmd.Example)
+	assert.NotNil(t, cmd.Run)
+
+	t.Run("File Not Found", func(t *testing.T) {
+		cmd.SetArgs([]string{"nonexistent.bin"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "Error reading attestation report file")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Invalid Attestation Data", func(t *testing.T) {
+		tmpfile, err := os.CreateTemp("", "attestation.bin")
+		require.NoError(t, err)
+		defer os.Remove(tmpfile.Name())
+
+		err = os.WriteFile(tmpfile.Name(), []byte("invalid protobuf data"), 0o644)
+		require.NoError(t, err)
+
+		cmd.SetArgs([]string{tmpfile.Name()})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err = cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "Error unmarshaling attestation report")
+		assert.Contains(t, output, "❌")
+	})
+}
+
+func TestNewAzureAttestationPolicy(t *testing.T) {
+	cli := &CLI{}
+	cmd := cli.NewAzureAttestationPolicy()
+
+	assert.Equal(t, "azure", cmd.Use)
+	assert.Equal(t, "Get attestation policy for Azure CVM", cmd.Short)
+	assert.Equal(t, "azure <azure_maa_token_file> <product_name>", cmd.Example)
+	assert.NotNil(t, cmd.Run)
+
+	flag := cmd.Flags().Lookup("policy")
+	assert.NotNil(t, flag)
+	assert.Equal(t, "Policy of the guest CVM", flag.Usage)
+
+	t.Run("File Not Found", func(t *testing.T) {
+		cmd.SetArgs([]string{"nonexistent.token", "test-product"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "Error reading attestation report file")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Valid Token File", func(t *testing.T) {
+		tmpfile, err := os.CreateTemp("", "token.maa")
+		require.NoError(t, err)
+		defer os.Remove(tmpfile.Name())
+
+		err = os.WriteFile(tmpfile.Name(), []byte("dummy.token.content"), 0o644)
+		require.NoError(t, err)
+
+		defer os.Remove("attestation_policy.json")
+
+		cmd.SetArgs([]string{tmpfile.Name(), "test-product"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err = cmd.Execute()
+		assert.NoError(t, err)
+	})
+
+	t.Run("Custom Policy Flag", func(t *testing.T) {
+		tmpfile, err := os.CreateTemp("", "token.maa")
+		require.NoError(t, err)
+		defer os.Remove(tmpfile.Name())
+
+		err = os.WriteFile(tmpfile.Name(), []byte("dummy.token.content"), 0o644)
+		require.NoError(t, err)
+
+		cmd.SetArgs([]string{"--policy", "123456", tmpfile.Name(), "test-product"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err = cmd.Execute()
+		assert.NoError(t, err)
+
+		flag := cmd.Flags().Lookup("policy")
+		assert.NotNil(t, flag)
+		assert.Equal(t, "123456", flag.Value.String())
+	})
+}
+
+func TestCommandErrorHandling(t *testing.T) {
+	cli := &CLI{}
+
+	t.Run("Measurement Command Error", func(t *testing.T) {
+		cmd := cli.NewAddMeasurementCmd()
+		cmd.SetArgs([]string{"invalid-base64", "nonexistent.json"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "Error could not change measurement data")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Host Data Command Error", func(t *testing.T) {
+		cmd := cli.NewAddHostDataCmd()
+		cmd.SetArgs([]string{"invalid-base64", "nonexistent.json"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "Error could not change host data")
+		assert.Contains(t, output, "❌")
+	})
+}
+
+func TestExtendWithManifestHandling(t *testing.T) {
+	cli := &CLI{}
+
+	t.Run("Invalid policy file", func(t *testing.T) {
+		cmd := cli.NewExtendWithManifestCmd()
+		cmd.SetArgs([]string{"nonexistent.policy.json", "nonexistent.manifest.json"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "error while reading the attestation policy file")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Invalid manifest file", func(t *testing.T) {
+		cmd := cli.NewExtendWithManifestCmd()
+		cmd.SetArgs([]string{"../scripts/attestation_policy/attestation_policy.json", "nonexistent.manifest.json"})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err := cmd.Execute()
+		assert.NoError(t, err)
+
+		output := buf.String()
+		assert.Contains(t, output, "error while reading manifest file")
+		assert.Contains(t, output, "❌")
+	})
+
+	t.Run("Valid file paths", func(t *testing.T) {
+		fileContent := `{
+		  "id": "1",
+		  "name": "sample computation",
+		  "description": "sample description",
+		  "datasets": [
+		    {
+		      "hash": "<sha3_encoded string>",
+		      "userKey": "<pem_encoded public key string>"
+		    }
+		  ],
+		  "algorithm": {
+		    "hash": "<sha3_encoded string>",
+		    "userKey": "<pem_encoded public key string>"
+		  },
+		  "result_consumers": [
+		    {
+		      "userKey": "<pem_encoded public key string>"
+		    }
+		  ],
+		  "agent_config": {
+		    "port": "7002",
+		    "cert_file": "<pem encoded cert string>",
+		    "key_file": "<pem encoded private key string>",
+		    "server_ca_file": "<pem encoded cert string>",
+		    "client_ca_file": "<pem encoded cert string>",
+		    "attested_tls": true
+		  }
+		}`
+
+		dir, err := os.Getwd()
+		if err != nil {
+			t.Fatalf("Error getting current working directory: %v", err)
+		}
+
+		manifestFile, err := os.CreateTemp(dir, "manifest.json")
+		if err != nil {
+			t.Fatalf("Error creating temp file: %v", err)
+		}
+		defer os.Remove(manifestFile.Name())
+
+		err = os.WriteFile(manifestFile.Name(), []byte(fileContent), 0o644)
+		if err != nil {
+			t.Fatalf("Error writing temp file: %v", err)
+		}
+
+		cmd := cli.NewExtendWithManifestCmd()
+		cmd.SetArgs([]string{"../scripts/attestation_policy/attestation_policy.json", manifestFile.Name()})
+
+		var buf bytes.Buffer
+		cmd.SetOut(&buf)
+		cmd.SetErr(&buf)
+
+		err = cmd.Execute()
+		assert.NoError(t, err)
+	})
+}

cli/attestation_snp.go

@@ -0,0 +1,597 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"encoding/hex"
+	"fmt"
+	"io"
+	"os"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/google/go-sev-guest/abi"
+	"github.com/google/go-sev-guest/proto/check"
+	tpmAttest "github.com/google/go-tpm-tools/proto/attest"
+	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/encoding/prototext"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/wrapperspb"
+)
+
+const (
+	defaultMinimumTcb       = 0
+	defaultMinimumLaunchTcb = 0
+	defaultMinimumGuestSvn  = 0
+	defaultGuestPolicy      = 0x0000000000030000
+	defaultMinimumBuild     = 0
+	defaultCheckCrl         = false
+	defaultTimeout          = 2 * time.Minute
+	defaultMaxRetryDelay    = 30 * time.Second
+	defaultRequireAuthor    = false
+	defaultRequireIdBlock   = false
+	defaultMinVersion       = "0.0"
+	vtpmFilePath            = "../quote.dat"
+	attestationReportJson   = "attestation.json"
+	sevSnpProductMilan      = "Milan"
+	sevSnpProductGenoa      = "Genoa"
+	FormatBinaryPB          = "binarypb"
+	FormatTextProto         = "textproto"
+	exampleJSONConfig       = `
+	{
+		"rootOfTrust":{
+		   "product":"test_product",
+		   "cabundlePaths":[
+			  "test_cabundlePaths"
+		   ],
+		   "cabundles":[
+			  "test_Cabundles"
+		   ],
+		   "checkCrl":true,
+		   "disallowNetwork":true
+		},
+		"policy":{
+		   "minimumGuestSvn":1,
+		   "policy":"1",
+		   "familyId":"AQIDBAUGBwgJCgsMDQ4PEA==",
+		   "imageId":"AQIDBAUGBwgJCgsMDQ4PEA==",
+		   "vmpl":0,
+		   "minimumTcb":"1",
+		   "minimumLaunchTcb":"1",
+		   "platformInfo":"1",
+		   "requireAuthorKey":true,
+		   "reportData":"J+60aXs8btm8VcGgaJYURGeNCu0FIyWMFXQ7ZUlJDC0FJGJizJsOzDIXgQ75UtPC+Zqe0A3dvnnf5VEeQ61RTg==",
+		   "measurement":"8s78ewoX7Xkfy1qsgVnkZwLDotD768Nqt6qTL5wtQOxHsLczipKM6bhDmWiHLdP4",
+		   "hostData":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
+		   "reportId":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
+		   "reportIdMa":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
+		   "chipId":"J+60aXs8btm8VcGgaJYURGeNCu0FIyWMFXQ7ZUlJDC0FJGJizJsOzDIXgQ75UtPC+Zqe0A3dvnnf5VEeQ61RTg==",
+		   "minimumBuild":1,
+		   "minimumVersion":"0.90",
+		   "permitProvisionalFirmware":true,
+		   "requireIdBlock":true,
+		   "trustedAuthorKeys":[
+			  "GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
+		   ],
+		   "trustedAuthorKeyHashes":[
+			  "GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
+		   ],
+		   "trustedIdKeys":[
+			  "GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
+		   ],
+		   "trustedIdKeyHashes":[
+			  "GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
+		   ],
+		   "product":{
+			  "name":1,
+			  "stepping":1,
+			  "machineStepping":1
+		   }
+		}
+	}
+	`
+)
+
+var cfg = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+
+func addSEVSNPVerificationOptions(cmd *cobra.Command) *cobra.Command {
+	cmd.Flags().BytesHexVar(
+		&cfg.Policy.HostData,
+		"host_data",
+		empty32[:],
+		"The expected HOST_DATA field as a hex string. Must encode 32 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfg.Policy.FamilyId,
+		"family_id",
+		empty16[:],
+		"The expected FAMILY_ID field as a hex string. Must encode 16 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfg.Policy.ImageId,
+		"image_id",
+		empty16[:],
+		"The expected IMAGE_ID field as a hex string. Must encode 16 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfg.Policy.ReportId,
+		"report_id",
+		nil,
+		"The expected REPORT_ID field as a hex string. Must encode 32 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfg.Policy.ReportIdMa,
+		"report_id_ma",
+		defaultReportIdMa,
+		"The expected REPORT_ID_MA field as a hex string. Must encode 32 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfg.Policy.Measurement,
+		"measurement",
+		nil,
+		"The expected MEASUREMENT field as a hex string. Must encode 48 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfg.Policy.ChipId,
+		"chip_id",
+		nil,
+		"The expected MEASUREMENT field as a hex string. Must encode 48 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().Uint64Var(
+		&cfg.Policy.MinimumTcb,
+		"minimum_tcb",
+		defaultMinimumTcb,
+		"The minimum acceptable value for CURRENT_TCB, COMMITTED_TCB, and REPORTED_TCB.",
+	)
+	cmd.Flags().Uint64Var(
+		&cfg.Policy.MinimumLaunchTcb,
+		"minimum_lauch_tcb",
+		defaultMinimumLaunchTcb,
+		"The minimum acceptable value for LAUNCH_TCB.",
+	)
+	cmd.Flags().Uint64Var(
+		&cfg.Policy.Policy,
+		"guest_policy",
+		defaultGuestPolicy,
+		"The most acceptable guest SnpPolicy.",
+	)
+	cmd.Flags().Uint32Var(
+		&cfg.Policy.MinimumGuestSvn,
+		"minimum_guest_svn",
+		defaultMinimumGuestSvn,
+		"The most acceptable GUEST_SVN.",
+	)
+	cmd.Flags().Uint32Var(
+		&cfg.Policy.MinimumBuild,
+		"minimum_build",
+		defaultMinimumBuild,
+		"The 8-bit minimum build number for AMD-SP firmware",
+	)
+	cmd.Flags().BoolVar(
+		&checkCrl,
+		"check_crl",
+		defaultCheckCrl,
+		"Download and check the CRL for revoked certificates.",
+	)
+	cmd.Flags().DurationVar(
+		&timeout,
+		"timeout",
+		defaultTimeout,
+		"Duration to continue to retry failed HTTP requests.",
+	)
+	cmd.Flags().DurationVar(
+		&maxRetryDelay,
+		"max_retry_delay",
+		defaultMaxRetryDelay,
+		"Maximum Duration to wait between HTTP request retries.",
+	)
+	cmd.Flags().BoolVar(
+		&cfg.Policy.RequireAuthorKey,
+		"require_author_key",
+		defaultRequireAuthor,
+		"Require that AUTHOR_KEY_EN is 1.",
+	)
+	cmd.Flags().BoolVar(
+		&cfg.Policy.RequireIdBlock,
+		"require_id_block",
+		defaultRequireIdBlock,
+		"Require that the VM was launch with an ID_BLOCK signed by a trusted id key or author key",
+	)
+	cmd.Flags().StringVar(
+		&platformInfo,
+		"platform_info",
+		"",
+		"The maximum acceptable PLATFORM_INFO field bit-wise. May be empty or a 64-bit unsigned integer",
+	)
+	cmd.Flags().StringVar(
+		&cfg.Policy.MinimumVersion,
+		"minimum_version",
+		defaultMinVersion,
+		"Minimum AMD-SP firmware API version (major.minor). Each number must be 8-bit non-negative.",
+	)
+	cmd.Flags().StringArrayVar(
+		&trustedAuthorKeys,
+		"trusted_author_keys",
+		[]string{},
+		"Paths to x.509 certificates of trusted author keys",
+	)
+	cmd.Flags().StringArrayVar(
+		&trustedAuthorHashes,
+		"trusted_author_key_hashes",
+		[]string{},
+		"Hex-encoded SHA-384 hash values of trusted author keys in AMD public key format",
+	)
+	cmd.Flags().StringArrayVar(
+		&trustedIdKeys,
+		"trusted_id_keys",
+		[]string{},
+		"Paths to x.509 certificates of trusted author keys",
+	)
+	cmd.Flags().StringArrayVar(
+		&trustedIdKeyHashes,
+		"trusted_id_key_hashes",
+		[]string{},
+		"Hex-encoded SHA-384 hash values of trusted identity keys in AMD public key format",
+	)
+	cmd.Flags().StringVar(
+		&cfg.RootOfTrust.ProductLine,
+		"product",
+		"",
+		"The AMD product name for the chip that generated the attestation report.",
+	)
+	cmd.Flags().StringVar(
+		&stepping,
+		"stepping",
+		"",
+		"The machine stepping for the chip that generated the attestation report. Default unchecked.",
+	)
+	cmd.Flags().StringArrayVar(
+		&cfg.RootOfTrust.CabundlePaths,
+		"CA_bundles_paths",
+		[]string{},
+		"Paths to CA bundles for the AMD product. Must be in PEM format, ASK, then ARK certificates. If unset, uses embedded root certificates.",
+	)
+	cmd.Flags().StringArrayVar(
+		&cfg.RootOfTrust.Cabundles,
+		"CA_bundles",
+		[]string{},
+		"PEM format CA bundles for the AMD product. Combined with contents of cabundle_paths.",
+	)
+
+	return cmd
+}
+
+func validateInput() error {
+	if len(cfg.RootOfTrust.CabundlePaths) != 0 || len(cfg.RootOfTrust.Cabundles) != 0 && cfg.RootOfTrust.ProductLine == "" {
+		return fmt.Errorf("product name must be set if CA bundles are provided")
+	}
+
+	if err := validateFieldLength("report_data", cfg.Policy.ReportData, size64); err != nil {
+		return err
+	}
+	if err := validateFieldLength("host_data", cfg.Policy.HostData, size32); err != nil {
+		return err
+	}
+	if err := validateFieldLength("family_id", cfg.Policy.FamilyId, size16); err != nil {
+		return err
+	}
+	if err := validateFieldLength("image_id", cfg.Policy.ImageId, size16); err != nil {
+		return err
+	}
+	if err := validateFieldLength("report_id", cfg.Policy.ReportId, size32); err != nil {
+		return err
+	}
+	if err := validateFieldLength("report_id_ma", cfg.Policy.ReportIdMa, size32); err != nil {
+		return err
+	}
+	if err := validateFieldLength("measurement", cfg.Policy.Measurement, size48); err != nil {
+		return err
+	}
+	if err := validateFieldLength("chip_id", cfg.Policy.ChipId, size64); err != nil {
+		return err
+	}
+	for _, hash := range cfg.Policy.TrustedAuthorKeyHashes {
+		if err := validateFieldLength("trusted_author_key_hash", hash, size48); err != nil {
+			return err
+		}
+	}
+	for _, hash := range cfg.Policy.TrustedIdKeyHashes {
+		if err := validateFieldLength("trusted_id_key_hash", hash, size48); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func parseTrustedKeys() error {
+	for _, path := range trustedAuthorKeys {
+		file, err := os.ReadFile(path)
+		if err != nil {
+			return err
+		}
+		cfg.Policy.TrustedAuthorKeys = append(cfg.Policy.TrustedAuthorKeys, file)
+	}
+	for _, path := range trustedIdKeys {
+		file, err := os.ReadFile(path)
+		if err != nil {
+			return err
+		}
+		cfg.Policy.TrustedIdKeys = append(cfg.Policy.TrustedIdKeys, file)
+	}
+	return nil
+}
+
+func parseUints() error {
+	if stepping != "" {
+		if base := getBase(stepping); base == 10 {
+			num, err := strconv.ParseUint(stepping, getBase(stepping), 8)
+			if err != nil {
+				return err
+			}
+
+			cfg.Policy.Product.MachineStepping = wrapperspb.UInt32(uint32(num))
+		} else {
+			num, err := strconv.ParseUint(stepping[2:], base, 8)
+			if err != nil {
+				return err
+			}
+			cfg.Policy.Product.MachineStepping = wrapperspb.UInt32(uint32(num))
+		}
+	}
+	if platformInfo != "" {
+		if base := getBase(platformInfo); base == 10 {
+			num, err := strconv.ParseUint(platformInfo, getBase(platformInfo), 8)
+			if err != nil {
+				return err
+			}
+			cfg.Policy.PlatformInfo = wrapperspb.UInt64(num)
+		} else {
+			num, err := strconv.ParseUint(platformInfo[2:], base, 8)
+			if err != nil {
+				return err
+			}
+			cfg.Policy.PlatformInfo = wrapperspb.UInt64(num)
+		}
+	}
+	return nil
+}
+
+func getBase(val string) int {
+	switch {
+	case strings.HasPrefix(val, "0x"):
+		return 16
+	case strings.HasPrefix(val, "0o"):
+		return 8
+	case strings.HasPrefix(val, "0b"):
+		return 2
+	default:
+		return 10
+	}
+}
+
+// parseConfig decodes config passed as json for check.Config struct.
+// example
+/* {
+	"rootOfTrust":{
+		"product":"test_product",
+		"cabundlePaths":[
+		   "test_cabundlePaths"
+		],
+		"cabundles":[
+		   "test_Cabundles"
+		],
+		"checkCrl":true,
+		"disallowNetwork":true
+	 },
+	 "policy":{
+		"minimumGuestSvn":1,
+		"policy":"1",
+		"familyId":"AQIDBAUGBwgJCgsMDQ4PEA==",
+		"imageId":"AQIDBAUGBwgJCgsMDQ4PEA==",
+		"vmpl":0,
+		"minimumTcb":"1",
+		"minimumLaunchTcb":"1",
+		"platformInfo":"1",
+		"requireAuthorKey":true,
+		"reportData":"J+60aXs8btm8VcGgaJYURGeNCu0FIyWMFXQ7ZUlJDC0FJGJizJsOzDIXgQ75UtPC+Zqe0A3dvnnf5VEeQ61RTg==",
+		"measurement":"8s78ewoX7Xkfy1qsgVnkZwLDotD768Nqt6qTL5wtQOxHsLczipKM6bhDmWiHLdP4",
+		"hostData":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
+		"reportId":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
+		"reportIdMa":"GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw=",
+		"chipId":"J+60aXs8btm8VcGgaJYURGeNCu0FIyWMFXQ7ZUlJDC0FJGJizJsOzDIXgQ75UtPC+Zqe0A3dvnnf5VEeQ61RTg==",
+		"minimumBuild":1,
+		"minimumVersion":"0.90",
+		"permitProvisionalFirmware":true,
+		"requireIdBlock":true,
+		"trustedAuthorKeys":[
+		   "GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
+		],
+		"trustedAuthorKeyHashes":[
+		   "GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
+		],
+		"trustedIdKeys":[
+		   "GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
+		],
+		"trustedIdKeyHashes":[
+		   "GSvLKpfu59Y9QOF6vhq0vQsOIvb4+5O/UOHLGLBTkdw="
+		],
+		"product":{
+		   "name":"1",
+		   "stepping":1,
+		   "machineStepping":1
+		}
+	 }
+  }*/
+func parseConfig() error {
+	if cfgString == "" {
+		return nil
+	}
+
+	policyByte, err := os.ReadFile(cfgString)
+	if err != nil {
+		return err
+	}
+
+	if err := protojson.Unmarshal(policyByte, &cfg); err != nil {
+		return err
+	}
+	// Populate fields that should not be nil
+	if cfg.RootOfTrust == nil {
+		cfg.RootOfTrust = &check.RootOfTrust{}
+	}
+	if cfg.Policy == nil {
+		cfg.Policy = &check.Policy{}
+	}
+	return nil
+}
+
+func parseHashes() error {
+	for _, hash := range trustedAuthorHashes {
+		hashBytes, err := hex.DecodeString(hash)
+		if err != nil {
+			return err
+		}
+		cfg.Policy.TrustedAuthorKeyHashes = append(cfg.Policy.TrustedAuthorKeyHashes, hashBytes)
+	}
+	for _, hash := range trustedIdKeyHashes {
+		hashBytes, err := hex.DecodeString(hash)
+		if err != nil {
+			return err
+		}
+		cfg.Policy.TrustedIdKeyHashes = append(cfg.Policy.TrustedIdKeyHashes, hashBytes)
+	}
+	return nil
+}
+
+func parseAttestationFile() error {
+	file, err := os.ReadFile(attestationFile)
+	if err != nil {
+		return err
+	}
+	attestationRaw = file
+	if isFileJSON(attestationFile) {
+		attestationRaw, err = attestationFromJSON(attestationRaw)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func sevsnpverify(cmd *cobra.Command, verifier attestation.Verifier, args []string) error {
+	cmd.Println("Checking attestation")
+
+	attestationFile = string(args[0])
+
+	if err := parseAttestationFile(); err != nil {
+		return fmt.Errorf("error parsing config: %v ❌ ", err)
+	}
+
+	// This format is the attestation report in AMD's specified ABI format, immediately
+	// followed by the certificate table bytes.
+	if len(attestationRaw) < abi.ReportSize {
+		return fmt.Errorf("attestation too small: got 0x%x bytes, need at least 0x%x bytes", len(attestationRaw), abi.ReportSize)
+	}
+
+	if err := parseAttestationConfig(); err != nil {
+		return err
+	}
+
+	if err := verifier.VerifTeeAttestation(attestationRaw, cfg.Policy.ReportData); err != nil {
+		return fmt.Errorf("attestation validation and verification failed with error: %v ❌ ", err)
+	}
+
+	cmd.Println("Attestation validation and verification is successful!")
+	return nil
+}
+
+func parseAttestationConfig() error {
+	if err := parseConfig(); err != nil {
+		return fmt.Errorf("error parsing config: %v ❌ ", err)
+	}
+	if err := parseHashes(); err != nil {
+		return fmt.Errorf("error parsing hashes: %v ❌ ", err)
+	}
+	if err := parseTrustedKeys(); err != nil {
+		return fmt.Errorf("error parsing files: %v ❌ ", err)
+	}
+
+	if err := parseUints(); err != nil {
+		return fmt.Errorf("error parsing uints: %v ❌ ", err)
+	}
+
+	if err := validateInput(); err != nil {
+		return fmt.Errorf("error validating input: %v ❌ ", err)
+	}
+
+	return nil
+}
+
+func vtpmSevSnpverify(args []string, verifier attestation.Verifier) error {
+	attest, err := returnvTPMAttestation(args)
+	if err != nil {
+		return err
+	}
+
+	if err := parseAttestationConfig(); err != nil {
+		return err
+	}
+
+	if err := verifier.VerifyAttestation(attest, cfg.Policy.ReportData, nonce); err != nil {
+		return fmt.Errorf("attestation validation and verification failed with error: %v ❌ ", err)
+	}
+
+	return nil
+}
+
+func vtpmverify(args []string, verifier attestation.Verifier) error {
+	attestation, err := returnvTPMAttestation(args)
+	if err != nil {
+		return err
+	}
+
+	if err := verifier.VerifVTpmAttestation(attestation, nonce); err != nil {
+		return fmt.Errorf("attestation validation and verification failed with error: %v ❌ ", err)
+	}
+
+	return nil
+}
+
+func returnvTPMAttestation(args []string) ([]byte, error) {
+	attestationFile = string(args[0])
+	input, err := openInputFile()
+	if err != nil {
+		return nil, err
+	}
+	if closer, ok := input.(*os.File); ok {
+		defer closer.Close()
+	}
+	attestationBytes, err := io.ReadAll(input)
+	if err != nil {
+		return nil, err
+	}
+	attestation := &tpmAttest.Attestation{}
+
+	if format == FormatBinaryPB {
+		return attestationBytes, nil
+	} else if format == FormatTextProto {
+		unmarshalOptions := prototext.UnmarshalOptions{}
+		err = unmarshalOptions.Unmarshal(attestationBytes, attestation)
+	} else {
+		return nil, fmt.Errorf("format should be either binarypb or textproto")
+	}
+	if err != nil {
+		return nil, fmt.Errorf("fail to unmarshal attestation report: %v", err)
+	}
+
+	attestationBytes, err = proto.Marshal(attestation)
+	if err != nil {
+		return nil, fmt.Errorf("fail to marshal vTPM attestation report: %v", err)
+	}
+
+	return attestationBytes, nil
+}

cli/attestation_snp_test.go

@@ -0,0 +1,870 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"bytes"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/google/go-sev-guest/abi"
+	"github.com/google/go-sev-guest/proto/check"
+	"github.com/google/go-sev-guest/proto/sevsnp"
+	tpmAttest "github.com/google/go-tpm-tools/proto/attest"
+	"github.com/google/go-tpm-tools/proto/tpm"
+	"github.com/spf13/cobra"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/pkg/attestation/mocks"
+	"google.golang.org/protobuf/encoding/prototext"
+	"google.golang.org/protobuf/proto"
+)
+
+func TestAddSEVSNPVerificationOptions(t *testing.T) {
+	cmd := &cobra.Command{
+		Use: "test",
+	}
+
+	result := addSEVSNPVerificationOptions(cmd)
+
+	assert.Equal(t, cmd, result)
+
+	// Check that important flags are added
+	flags := []string{
+		"host_data",
+		"family_id",
+		"image_id",
+		"report_id",
+		"report_id_ma",
+		"measurement",
+		"chip_id",
+		"minimum_tcb",
+		"minimum_lauch_tcb",
+		"guest_policy",
+		"minimum_guest_svn",
+		"minimum_build",
+		"check_crl",
+		"timeout",
+		"max_retry_delay",
+		"require_author_key",
+		"require_id_block",
+		"platform_info",
+		"minimum_version",
+		"trusted_author_keys",
+		"trusted_author_key_hashes",
+		"trusted_id_keys",
+		"trusted_id_key_hashes",
+		"product",
+		"stepping",
+		"CA_bundles_paths",
+		"CA_bundles",
+	}
+
+	for _, flagName := range flags {
+		flag := cmd.Flags().Lookup(flagName)
+		assert.NotNil(t, flag, "Flag %s should exist", flagName)
+	}
+}
+
+func TestValidateInput(t *testing.T) {
+	tests := []struct {
+		name      string
+		setupCfg  func()
+		expectErr bool
+		errMsg    string
+	}{
+		{
+			name: "valid empty config",
+			setupCfg: func() {
+				cfg = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+			},
+			expectErr: false,
+		},
+		{
+			name: "CA bundles without product name",
+			setupCfg: func() {
+				cfg = check.Config{
+					Policy: &check.Policy{},
+					RootOfTrust: &check.RootOfTrust{
+						CabundlePaths: []string{"test.pem"},
+						ProductLine:   "",
+					},
+				}
+			},
+			expectErr: true,
+			errMsg:    "product name must be set if CA bundles are provided",
+		},
+		{
+			name: "invalid report_data length",
+			setupCfg: func() {
+				cfg = check.Config{
+					Policy: &check.Policy{
+						ReportData: []byte("invalid"),
+					},
+					RootOfTrust: &check.RootOfTrust{},
+				}
+			},
+			expectErr: true,
+			errMsg:    "report_data",
+		},
+		{
+			name: "invalid host_data length",
+			setupCfg: func() {
+				cfg = check.Config{
+					Policy: &check.Policy{
+						HostData: []byte("invalid"),
+					},
+					RootOfTrust: &check.RootOfTrust{},
+				}
+			},
+			expectErr: true,
+			errMsg:    "host_data",
+		},
+		{
+			name: "invalid family_id length",
+			setupCfg: func() {
+				cfg = check.Config{
+					Policy: &check.Policy{
+						FamilyId: []byte("invalid"),
+					},
+					RootOfTrust: &check.RootOfTrust{},
+				}
+			},
+			expectErr: true,
+			errMsg:    "family_id",
+		},
+		{
+			name: "invalid image_id length",
+			setupCfg: func() {
+				cfg = check.Config{
+					Policy: &check.Policy{
+						ImageId: []byte("invalid"),
+					},
+					RootOfTrust: &check.RootOfTrust{},
+				}
+			},
+			expectErr: true,
+			errMsg:    "image_id",
+		},
+		{
+			name: "invalid trusted author key hash",
+			setupCfg: func() {
+				cfg = check.Config{
+					Policy: &check.Policy{
+						TrustedAuthorKeyHashes: [][]byte{[]byte("invalid")},
+					},
+					RootOfTrust: &check.RootOfTrust{},
+				}
+			},
+			expectErr: true,
+			errMsg:    "trusted_author_key_hash",
+		},
+		{
+			name: "invalid trusted id key hash",
+			setupCfg: func() {
+				cfg = check.Config{
+					Policy: &check.Policy{
+						TrustedIdKeyHashes: [][]byte{[]byte("invalid")},
+					},
+					RootOfTrust: &check.RootOfTrust{},
+				}
+			},
+			expectErr: true,
+			errMsg:    "trusted_id_key_hash",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tt.setupCfg()
+			err := validateInput()
+			if tt.expectErr {
+				assert.Error(t, err)
+				assert.Contains(t, err.Error(), tt.errMsg)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
+
+func TestParseTrustedKeys(t *testing.T) {
+	tempDir := t.TempDir()
+
+	authorKeyFile := filepath.Join(tempDir, "author.pem")
+	idKeyFile := filepath.Join(tempDir, "id.pem")
+	nonExistentFile := filepath.Join(tempDir, "nonexistent.pem")
+
+	authorKeyContent := "-----BEGIN CERTIFICATE-----\nMIIBkTCB+wIJAOI..."
+	idKeyContent := "-----BEGIN CERTIFICATE-----\nMIIBkTCB+wIJAOI..."
+
+	require.NoError(t, os.WriteFile(authorKeyFile, []byte(authorKeyContent), 0o644))
+	require.NoError(t, os.WriteFile(idKeyFile, []byte(idKeyContent), 0o644))
+
+	tests := []struct {
+		name              string
+		trustedAuthorKeys []string
+		trustedIdKeys     []string
+		expectErr         bool
+	}{
+		{
+			name:              "valid files",
+			trustedAuthorKeys: []string{authorKeyFile},
+			trustedIdKeys:     []string{idKeyFile},
+			expectErr:         false,
+		},
+		{
+			name:              "nonexistent author key file",
+			trustedAuthorKeys: []string{nonExistentFile},
+			trustedIdKeys:     []string{},
+			expectErr:         true,
+		},
+		{
+			name:              "nonexistent id key file",
+			trustedAuthorKeys: []string{},
+			trustedIdKeys:     []string{nonExistentFile},
+			expectErr:         true,
+		},
+		{
+			name:              "empty file lists",
+			trustedAuthorKeys: []string{},
+			trustedIdKeys:     []string{},
+			expectErr:         false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+			trustedAuthorKeys = tt.trustedAuthorKeys
+			trustedIdKeys = tt.trustedIdKeys
+
+			err := parseTrustedKeys()
+			if tt.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				if len(tt.trustedAuthorKeys) > 0 {
+					assert.Len(t, cfg.Policy.TrustedAuthorKeys, len(tt.trustedAuthorKeys))
+					assert.Equal(t, []byte(authorKeyContent), cfg.Policy.TrustedAuthorKeys[0])
+				}
+				if len(tt.trustedIdKeys) > 0 {
+					assert.Len(t, cfg.Policy.TrustedIdKeys, len(tt.trustedIdKeys))
+					assert.Equal(t, []byte(idKeyContent), cfg.Policy.TrustedIdKeys[0])
+				}
+			}
+		})
+	}
+}
+
+func TestParseUints(t *testing.T) {
+	tests := []struct {
+		name             string
+		stepping         string
+		platformInfo     string
+		expectErr        bool
+		expectedStep     *uint32
+		expectedPlatform *uint64
+	}{
+		{
+			name:         "empty values",
+			stepping:     "",
+			platformInfo: "",
+			expectErr:    false,
+		},
+		{
+			name:             "decimal values",
+			stepping:         "5",
+			platformInfo:     "10",
+			expectErr:        false,
+			expectedStep:     uint32Ptr(5),
+			expectedPlatform: uint64Ptr(10),
+		},
+		{
+			name:             "hex values",
+			stepping:         "0x5",
+			platformInfo:     "0xa",
+			expectErr:        false,
+			expectedStep:     uint32Ptr(5),
+			expectedPlatform: uint64Ptr(10),
+		},
+		{
+			name:             "octal values",
+			stepping:         "0o7",
+			platformInfo:     "0o12",
+			expectErr:        false,
+			expectedStep:     uint32Ptr(7),
+			expectedPlatform: uint64Ptr(10),
+		},
+		{
+			name:             "binary values",
+			stepping:         "0b101",
+			platformInfo:     "0b1010",
+			expectErr:        false,
+			expectedStep:     uint32Ptr(5),
+			expectedPlatform: uint64Ptr(10),
+		},
+		{
+			name:         "invalid stepping",
+			stepping:     "invalid",
+			platformInfo: "",
+			expectErr:    true,
+		},
+		{
+			name:         "invalid platform info",
+			stepping:     "",
+			platformInfo: "invalid",
+			expectErr:    true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg = check.Config{Policy: &check.Policy{Product: &sevsnp.SevProduct{}}, RootOfTrust: &check.RootOfTrust{}}
+			stepping = tt.stepping
+			platformInfo = tt.platformInfo
+
+			err := parseUints()
+			if tt.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				if tt.expectedStep != nil {
+					assert.Equal(t, *tt.expectedStep, cfg.Policy.Product.MachineStepping.Value)
+				}
+				if tt.expectedPlatform != nil {
+					assert.Equal(t, *tt.expectedPlatform, cfg.Policy.PlatformInfo.Value)
+				}
+			}
+		})
+	}
+}
+
+func TestGetBase(t *testing.T) {
+	tests := []struct {
+		input    string
+		expected int
+	}{
+		{"0x10", 16},
+		{"0o10", 8},
+		{"0b10", 2},
+		{"10", 10},
+		{"", 10},
+		{"abc", 10},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.input, func(t *testing.T) {
+			result := getBase(tt.input)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+func TestParseConfig(t *testing.T) {
+	tempDir := t.TempDir()
+
+	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]any{
+			"minimumGuestSvn":  1,
+			"policy":           "1",
+			"minimumBuild":     1,
+			"minimumVersion":   "0.90",
+			"requireAuthorKey": true,
+			"requireIdBlock":   true,
+		},
+	}
+
+	tests := []struct {
+		name        string
+		setupConfig func() string
+		expectErr   bool
+	}{
+		{
+			name: "empty config string",
+			setupConfig: func() string {
+				return ""
+			},
+			expectErr: false,
+		},
+		{
+			name: "valid config file",
+			setupConfig: func() string {
+				configFile := filepath.Join(tempDir, "valid_config.json")
+				configBytes, err := json.Marshal(validConfig)
+				assert.NoError(t, err)
+				if err := os.WriteFile(configFile, configBytes, 0o644); err != nil {
+					t.Errorf("failed to write config file: %v", err)
+				}
+				return configFile
+			},
+			expectErr: false,
+		},
+		{
+			name: "nonexistent config file",
+			setupConfig: func() string {
+				return filepath.Join(tempDir, "nonexistent.json")
+			},
+			expectErr: true,
+		},
+		{
+			name: "invalid JSON config",
+			setupConfig: func() string {
+				configFile := filepath.Join(tempDir, "invalid_config.json")
+				if err := os.WriteFile(configFile, []byte("invalid json"), 0o644); err != nil {
+					t.Errorf("failed to write invalid config file: %v", err)
+				}
+				return configFile
+			},
+			expectErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+			cfgString = tt.setupConfig()
+
+			err := parseConfig()
+			if tt.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, cfg.Policy)
+				assert.NotNil(t, cfg.RootOfTrust)
+			}
+		})
+	}
+}
+
+func TestParseHashes(t *testing.T) {
+	tests := []struct {
+		name                string
+		trustedAuthorHashes []string
+		trustedIdKeyHashes  []string
+		expectErr           bool
+	}{
+		{
+			name:                "valid hashes",
+			trustedAuthorHashes: []string{"deadbeef", "cafebabe"},
+			trustedIdKeyHashes:  []string{"12345678", "87654321"},
+			expectErr:           false,
+		},
+		{
+			name:                "empty hashes",
+			trustedAuthorHashes: []string{},
+			trustedIdKeyHashes:  []string{},
+			expectErr:           false,
+		},
+		{
+			name:                "invalid author hash",
+			trustedAuthorHashes: []string{"invalid_hex"},
+			trustedIdKeyHashes:  []string{},
+			expectErr:           true,
+		},
+		{
+			name:                "invalid id key hash",
+			trustedAuthorHashes: []string{},
+			trustedIdKeyHashes:  []string{"invalid_hex"},
+			expectErr:           true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+			trustedAuthorHashes = tt.trustedAuthorHashes
+			trustedIdKeyHashes = tt.trustedIdKeyHashes
+
+			err := parseHashes()
+			if tt.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.Len(t, cfg.Policy.TrustedAuthorKeyHashes, len(tt.trustedAuthorHashes))
+				assert.Len(t, cfg.Policy.TrustedIdKeyHashes, len(tt.trustedIdKeyHashes))
+
+				for i, hash := range tt.trustedAuthorHashes {
+					expected, _ := hex.DecodeString(hash)
+					assert.Equal(t, expected, cfg.Policy.TrustedAuthorKeyHashes[i])
+				}
+
+				for i, hash := range tt.trustedIdKeyHashes {
+					expected, _ := hex.DecodeString(hash)
+					assert.Equal(t, expected, cfg.Policy.TrustedIdKeyHashes[i])
+				}
+			}
+		})
+	}
+}
+
+func TestParseAttestationFile(t *testing.T) {
+	tempDir := t.TempDir()
+
+	binaryFile := filepath.Join(tempDir, "attestation.bin")
+	jsonFile := filepath.Join(tempDir, "attestation.json")
+
+	binaryData := make([]byte, 1024)
+	for i := range binaryData {
+		binaryData[i] = byte(i % 256)
+	}
+
+	jsonData := &sevsnp.Attestation{
+		Report: &sevsnp.Report{
+			FamilyId:        make([]byte, 16),
+			ImageId:         make([]byte, 16),
+			ReportData:      make([]byte, 64),
+			Measurement:     make([]byte, 48),
+			HostData:        make([]byte, 32),
+			IdKeyDigest:     make([]byte, 48),
+			AuthorKeyDigest: make([]byte, 48),
+			ReportId:        make([]byte, 32),
+			ReportIdMa:      make([]byte, 32),
+			ChipId:          make([]byte, 64),
+			Signature:       make([]byte, 512),
+		},
+	}
+	jsonBytes, err := json.Marshal(jsonData)
+	require.NoError(t, err)
+
+	require.NoError(t, os.WriteFile(binaryFile, binaryData, 0o644))
+	require.NoError(t, os.WriteFile(jsonFile, jsonBytes, 0o644))
+
+	tests := []struct {
+		name            string
+		attestationFile string
+		expectErr       bool
+	}{
+		{
+			name:            "valid binary file",
+			attestationFile: binaryFile,
+			expectErr:       false,
+		},
+		{
+			name:            "valid JSON file",
+			attestationFile: jsonFile,
+			expectErr:       false,
+		},
+		{
+			name:            "nonexistent file",
+			attestationFile: filepath.Join(tempDir, "nonexistent.bin"),
+			expectErr:       true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			attestationFile = tt.attestationFile
+
+			err := parseAttestationFile()
+			if tt.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, attestationRaw)
+				assert.NotEmpty(t, attestationRaw)
+			}
+		})
+	}
+}
+
+func TestSevsnpverify(t *testing.T) {
+	trustedAuthorHashes = []string{}
+	trustedIdKeyHashes = []string{}
+	stepping = ""
+	platformInfo = ""
+	tempDir := t.TempDir()
+	cfg = check.Config{Policy: &check.Policy{Product: &sevsnp.SevProduct{}}, RootOfTrust: &check.RootOfTrust{}}
+
+	attestationFile := filepath.Join(tempDir, "attestation.bin")
+	attestationData := make([]byte, abi.ReportSize+100)
+	for i := range attestationData {
+		attestationData[i] = byte(i % 256)
+	}
+	require.NoError(t, os.WriteFile(attestationFile, attestationData, 0o644))
+
+	tests := []struct {
+		name        string
+		args        []string
+		setupMock   func(*mocks.Verifier)
+		expectErr   bool
+		expectedMsg string
+	}{
+		{
+			name: "successful verification",
+			args: []string{attestationFile},
+			setupMock: func(m *mocks.Verifier) {
+				m.On("VerifTeeAttestation", mock.Anything, mock.Anything).Return(nil)
+			},
+			expectErr:   false,
+			expectedMsg: "Attestation validation and verification is successful!",
+		},
+		{
+			name: "verification failure",
+			args: []string{attestationFile},
+			setupMock: func(m *mocks.Verifier) {
+				m.On("VerifTeeAttestation", mock.Anything, mock.Anything).Return(fmt.Errorf("verification failed"))
+			},
+			expectErr:   true,
+			expectedMsg: "attestation validation and verification failed",
+		},
+		{
+			name:      "nonexistent file",
+			args:      []string{filepath.Join(tempDir, "nonexistent.bin")},
+			setupMock: func(m *mocks.Verifier) {},
+			expectErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cfgString = ""
+
+			mockVerifier := new(mocks.Verifier)
+			tt.setupMock(mockVerifier)
+
+			var output bytes.Buffer
+			cmd := &cobra.Command{}
+			cmd.SetOut(&output)
+
+			err := sevsnpverify(cmd, mockVerifier, tt.args)
+			fmt.Println("error1", err)
+
+			if tt.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+				if tt.expectedMsg != "" {
+					assert.Contains(t, output.String(), tt.expectedMsg)
+				}
+			}
+
+			mockVerifier.AssertExpectations(t)
+		})
+	}
+}
+
+func TestReturnvTPMAttestation(t *testing.T) {
+	tempDir := t.TempDir()
+
+	attestation := &tpmAttest.Attestation{
+		Quotes: []*tpm.Quote{
+			{
+				Quote:  []byte("test quote"),
+				RawSig: []byte("test signature"),
+			},
+		},
+	}
+
+	binaryData, err := proto.Marshal(attestation)
+	require.NoError(t, err)
+
+	binaryFile := filepath.Join(tempDir, "attestation.pb")
+	require.NoError(t, os.WriteFile(binaryFile, binaryData, 0o644))
+
+	textData, err := prototext.Marshal(attestation)
+	require.NoError(t, err)
+
+	textFile := filepath.Join(tempDir, "attestation.txtpb")
+	require.NoError(t, os.WriteFile(textFile, textData, 0o644))
+
+	tests := []struct {
+		name      string
+		args      []string
+		format    string
+		expectErr bool
+	}{
+		{
+			name:      "binary protobuf format",
+			args:      []string{binaryFile},
+			format:    FormatBinaryPB,
+			expectErr: false,
+		},
+		{
+			name:      "text protobuf format",
+			args:      []string{textFile},
+			format:    FormatTextProto,
+			expectErr: false,
+		},
+		{
+			name:      "invalid format",
+			args:      []string{binaryFile},
+			format:    "invalid",
+			expectErr: true,
+		},
+		{
+			name:      "nonexistent file",
+			args:      []string{filepath.Join(tempDir, "nonexistent.pb")},
+			format:    FormatBinaryPB,
+			expectErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			format = tt.format
+
+			result, err := returnvTPMAttestation(tt.args)
+
+			if tt.expectErr {
+				assert.Error(t, err)
+				assert.Nil(t, result)
+			} else {
+				assert.NoError(t, err)
+				assert.NotNil(t, result)
+				assert.NotEmpty(t, result)
+			}
+		})
+	}
+}
+
+func TestVtpmSevSnpverify(t *testing.T) {
+	stepping = ""
+	platformInfo = ""
+	trustedAuthorHashes = []string{}
+	trustedIdKeyHashes = []string{}
+	cfg = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+	tempDir := t.TempDir()
+
+	attestation := &tpmAttest.Attestation{
+		Quotes: []*tpm.Quote{
+			{
+				Quote:  []byte("test quote"),
+				RawSig: []byte("test signature"),
+			},
+		},
+	}
+
+	binaryData, err := proto.Marshal(attestation)
+	require.NoError(t, err)
+
+	attestationFile := filepath.Join(tempDir, "vtpm_attestation.pb")
+	require.NoError(t, os.WriteFile(attestationFile, binaryData, 0o644))
+
+	tests := []struct {
+		name      string
+		args      []string
+		setupMock func(*mocks.Verifier)
+		expectErr bool
+	}{
+		{
+			name: "successful verification",
+			args: []string{attestationFile},
+			setupMock: func(m *mocks.Verifier) {
+				m.On("VerifyAttestation", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			},
+			expectErr: false,
+		},
+		{
+			name: "verification failure",
+			args: []string{attestationFile},
+			setupMock: func(m *mocks.Verifier) {
+				m.On("VerifyAttestation", mock.Anything, mock.Anything, mock.Anything).Return(fmt.Errorf("verification failed"))
+			},
+			expectErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cfg = check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+			cfgString = ""
+			format = FormatBinaryPB
+
+			mockVerifier := new(mocks.Verifier)
+			tt.setupMock(mockVerifier)
+
+			err := vtpmSevSnpverify(tt.args, mockVerifier)
+
+			if tt.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+
+			mockVerifier.AssertExpectations(t)
+		})
+	}
+}
+
+func TestVtpmverify(t *testing.T) {
+	tempDir := t.TempDir()
+
+	attestation := &tpmAttest.Attestation{
+		Quotes: []*tpm.Quote{
+			{
+				Quote:  []byte("test quote"),
+				RawSig: []byte("test signature"),
+			},
+		},
+	}
+
+	binaryData, err := proto.Marshal(attestation)
+	require.NoError(t, err)
+
+	attestationFile := filepath.Join(tempDir, "vtpm_attestation.pb")
+	require.NoError(t, os.WriteFile(attestationFile, binaryData, 0o644))
+
+	tests := []struct {
+		name      string
+		args      []string
+		setupMock func(*mocks.Verifier)
+		expectErr bool
+	}{
+		{
+			name: "successful verification",
+			args: []string{attestationFile},
+			setupMock: func(m *mocks.Verifier) {
+				m.On("VerifVTpmAttestation", mock.Anything, mock.Anything).Return(nil)
+			},
+			expectErr: false,
+		},
+		{
+			name: "verification failure",
+			args: []string{attestationFile},
+			setupMock: func(m *mocks.Verifier) {
+				m.On("VerifVTpmAttestation", mock.Anything, mock.Anything).Return(fmt.Errorf("verification failed"))
+			},
+			expectErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			format = FormatBinaryPB
+
+			mockVerifier := new(mocks.Verifier)
+			tt.setupMock(mockVerifier)
+
+			err := vtpmverify(tt.args, mockVerifier)
+
+			if tt.expectErr {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+
+			mockVerifier.AssertExpectations(t)
+		})
+	}
+}
+
+func uint32Ptr(v uint32) *uint32 {
+	return &v
+}
+
+func uint64Ptr(v uint64) *uint64 {
+	return &v
+}

cli/attestation_tdx.go

@@ -0,0 +1,258 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"encoding/hex"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/absmach/supermq/pkg/errors"
+	ccpb "github.com/google/go-tdx-guest/proto/checkconfig"
+	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"google.golang.org/protobuf/encoding/protojson"
+)
+
+var (
+	cfgTDX = &ccpb.Config{
+		RootOfTrust: &ccpb.RootOfTrust{},
+		Policy:      &ccpb.Policy{HeaderPolicy: &ccpb.HeaderPolicy{}, TdQuoteBodyPolicy: &ccpb.TDQuoteBodyPolicy{}},
+	}
+	rtmrsS             string
+	trustedRootS       string
+	errNumberRtmrs     = fmt.Errorf("expected 4 RTMRS values")
+	errDecodeRtmrs     = fmt.Errorf("failed to decode RTMRS hex string")
+	errTrustedRootPath = fmt.Errorf("trusted root path must be a file, not a directory")
+	errNotAFile        = fmt.Errorf("trusted root path must be a file")
+)
+
+func addTDXVerificationOptions(cmd *cobra.Command) *cobra.Command {
+	cmd.Flags().BytesHexVar(
+		&cfgTDX.Policy.HeaderPolicy.QeVendorId,
+		"qe_vendor_id",
+		[]byte{},
+		"The expected QE_VENDOR_ID field as a hex string. Must encode 16 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfgTDX.Policy.TdQuoteBodyPolicy.MrSeam,
+		"mr_seam",
+		[]byte{},
+		"The expected MR_SEAM field as a hex string. Must encode 48 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfgTDX.Policy.TdQuoteBodyPolicy.TdAttributes,
+		"td_attributes",
+		[]byte{},
+		"The expected TD_ATTRIBUTES field as a hex string. Must encode 8 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfgTDX.Policy.TdQuoteBodyPolicy.Xfam,
+		"xfam",
+		[]byte{},
+		"The expected XFAM field as a hex string. Must encode 8 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfgTDX.Policy.TdQuoteBodyPolicy.MrTd,
+		"mr_td",
+		[]byte{},
+		"The expected MR_TD field as a hex string. Must encode 48 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfgTDX.Policy.TdQuoteBodyPolicy.MrConfigId,
+		"mr_config_id",
+		[]byte{},
+		"The expected MR_CONFIG_ID field as a hex string. Must encode 48 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfgTDX.Policy.TdQuoteBodyPolicy.MrOwnerConfig,
+		"mr_owner",
+		[]byte{},
+		"The expected MR_OWNER field as a hex string. Must encode 48 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfgTDX.Policy.TdQuoteBodyPolicy.MrOwnerConfig,
+		"mr_config_owner",
+		[]byte{},
+		"The expected MR_OWNER_CONFIG field as a hex string. Must encode 48 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().BytesHexVar(
+		&cfgTDX.Policy.TdQuoteBodyPolicy.MinimumTeeTcbSvn,
+		"minimum_tee_tcb_svn",
+		[]byte{},
+		"The minimum acceptable value for TEE_TCB_SVN field as a hex string. Must encode 16 bytes. Unchecked if unset.",
+	)
+	cmd.Flags().StringVar(
+		&rtmrsS,
+		"rtmrs",
+		"",
+		"Comma-separated hex strings representing expected values of RTMRS field. Expected 4 strings, either empty or each must encode 48 bytes. Unchecked if unset",
+	)
+	cmd.Flags().StringVar(
+		&trustedRootS,
+		"trusted_root",
+		"",
+		"Comma-separated paths to CA bundles for the Intel TDX. Must be in PEM format, Root CA certificate. If unset, uses embedded root certificate.",
+	)
+	cmd.Flags().Uint32Var(
+		&cfgTDX.Policy.HeaderPolicy.MinimumQeSvn,
+		"minimum_qe_svn",
+		0,
+		"The minimum acceptable value for QE_SVN field.",
+	)
+	cmd.Flags().Uint32Var(
+		&cfgTDX.Policy.HeaderPolicy.MinimumPceSvn,
+		"minimum_pce_svn",
+		0,
+		"The minimum acceptable value for PCE_SVN field.",
+	)
+	cmd.Flags().BoolVar(
+		&cfgTDX.RootOfTrust.GetCollateral,
+		"get_collateral",
+		false,
+		"If true, then permitted to download necessary collaterals for additional checks.",
+	)
+
+	return cmd
+}
+
+func parseRtmrs() ([][]byte, error) {
+	if rtmrsS == "" {
+		return nil, nil // No RTMRS provided, return nil
+	}
+
+	hexString := strings.Split(rtmrsS, ",")
+	if len(hexString) != 4 {
+		return nil, errNumberRtmrs
+	}
+
+	var result [][]byte
+	for _, hexStr := range hexString {
+		h, err := hex.DecodeString(strings.TrimSpace(hexStr))
+		if err != nil {
+			return nil, errors.Wrap(errDecodeRtmrs, err)
+		}
+
+		result = append(result, h)
+	}
+
+	return result, nil
+}
+
+func parseTrustedRoot() ([]string, error) {
+	if trustedRootS == "" {
+		return nil, nil // No trusted roots provided, return nil
+	}
+
+	roots := strings.Split(trustedRootS, ",")
+	var result []string
+	for _, root := range roots {
+		p := strings.TrimSpace(root)
+		state, err := os.Stat(p)
+		if err != nil {
+			return nil, errors.Wrap(errTrustedRootPath, err)
+		}
+		if state.IsDir() {
+			return nil, errNotAFile
+		}
+
+		result = append(result, p)
+	}
+
+	return result, nil
+}
+
+func parseTDXConfig() error {
+	if cfgString == "" {
+		return nil // No config provided, return nil
+	}
+
+	policyByte, err := os.ReadFile(cfgString)
+	if err != nil {
+		return err
+	}
+
+	if err := protojson.Unmarshal(policyByte, cfgTDX); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func validateTDXFlags() error {
+	if err := parseTDXConfig(); err != nil {
+		return err
+	}
+
+	rtrms, err := parseRtmrs()
+	if err != nil {
+		return err
+	}
+	if rtrms != nil {
+		cfgTDX.Policy.TdQuoteBodyPolicy.Rtmrs = rtrms
+	}
+	trustedRoots, err := parseTrustedRoot()
+	if err != nil {
+		return err
+	}
+	if trustedRoots != nil {
+		cfgTDX.RootOfTrust.CabundlePaths = trustedRoots
+	}
+
+	if err := validateTDXinput(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func tdxVerify(reportFilePath string, verifier attestation.Verifier) error {
+	attestationFile = reportFilePath
+	input, err := openInputFile()
+	if err != nil {
+		return err
+	}
+	if closer, ok := input.(*os.File); ok {
+		defer closer.Close()
+	}
+	attestationBytes, err := io.ReadAll(input)
+	if err != nil {
+		return err
+	}
+
+	return verifier.VerifyAttestation(attestationBytes, reportData, nil)
+}
+
+func validateTDXinput() error {
+	if err := validateFieldLength("qe_vendor_id", cfgTDX.Policy.HeaderPolicy.QeVendorId, size16); err != nil {
+		return err
+	}
+	if err := validateFieldLength("mr_seam", cfgTDX.Policy.TdQuoteBodyPolicy.MrSeam, size48); err != nil {
+		return err
+	}
+	if err := validateFieldLength("td_attributes", cfgTDX.Policy.TdQuoteBodyPolicy.TdAttributes, size8); err != nil {
+		return err
+	}
+	if err := validateFieldLength("xfam", cfgTDX.Policy.TdQuoteBodyPolicy.Xfam, size8); err != nil {
+		return err
+	}
+	if err := validateFieldLength("mr_td", cfgTDX.Policy.TdQuoteBodyPolicy.MrTd, size48); err != nil {
+		return err
+	}
+	if err := validateFieldLength("mr_config_id", cfgTDX.Policy.TdQuoteBodyPolicy.MrConfigId, size48); err != nil {
+		return err
+	}
+	if err := validateFieldLength("mr_owner", cfgTDX.Policy.TdQuoteBodyPolicy.MrOwnerConfig, size48); err != nil {
+		return err
+	}
+	if err := validateFieldLength("mr_config_owner", cfgTDX.Policy.TdQuoteBodyPolicy.MrOwnerConfig, size48); err != nil {
+		return err
+	}
+	if err := validateFieldLength("minimum_tee_tcb_svn", cfgTDX.Policy.TdQuoteBodyPolicy.MinimumTeeTcbSvn, size16); err != nil {
+		return err
+	}
+
+	return nil
+}

cli/backend_info.go

@@ -1,144 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package cli
-
-import (
-	"encoding/base64"
-	"encoding/json"
-	"fmt"
-	"log"
-	"os"
-
-	"github.com/absmach/magistrala/pkg/errors"
-	"github.com/google/go-sev-guest/proto/check"
-	"github.com/spf13/cobra"
-	"github.com/spf13/pflag"
-)
-
-type fieldType int
-
-const (
-	measurementField fieldType = iota
-	hostDataField
-)
-
-const (
-	// 0o744 file permission gives RWX permission to the user and only the R permission to others.
-	filePermission = 0o744
-	// Length of the expected host data and measurement field in bytes.
-	hostDataLength    = 32
-	measurementLength = 48
-)
-
-var (
-	errDecode                 = errors.New("base64 string could not be decoded")
-	errDataLength             = errors.New("data does not have an adequate length")
-	errReadingBackendInfoFile = errors.New("error while reading the backend information file")
-	errUnmarshalJSON          = errors.New("failed to unmarshal json")
-	errMarshalJSON            = errors.New("failed to marshal json")
-	errWriteFile              = errors.New("failed to write to file")
-	errBackendField           = errors.New("the specified field type does not exist in the backend information")
-)
-
-type AttestationConfiguration struct {
-	SNPPolicy   *check.Policy      `json:"snp_policy,omitempty"`
-	RootOfTrust *check.RootOfTrust `json:"root_of_trust,omitempty"`
-}
-
-func (cli *CLI) NewBackendCmd() *cobra.Command {
-	return &cobra.Command{
-		Use:   "backend [command]",
-		Short: "Change backend information",
-		Run: func(cmd *cobra.Command, args []string) {
-			fmt.Printf("Change backend information\n\n")
-			fmt.Printf("Usage:\n  %s [command]\n\n", cmd.CommandPath())
-			fmt.Printf("Available Commands:\n")
-
-			// Filter out "completion" command
-			availableCommands := make([]*cobra.Command, 0)
-			for _, subCmd := range cmd.Commands() {
-				if subCmd.Name() != "completion" {
-					availableCommands = append(availableCommands, subCmd)
-				}
-			}
-
-			for _, subCmd := range availableCommands {
-				fmt.Printf("  %-15s%s\n", subCmd.Name(), subCmd.Short)
-			}
-
-			fmt.Printf("\nFlags:\n")
-			cmd.Flags().VisitAll(func(flag *pflag.Flag) {
-				fmt.Printf("  -%s, --%s %s\n", flag.Shorthand, flag.Name, flag.Usage)
-			})
-			fmt.Printf("\nUse \"%s [command] --help\" for more information about a command.\n", cmd.CommandPath())
-		},
-	}
-}
-
-func (cli *CLI) NewAddMeasurementCmd() *cobra.Command {
-	return &cobra.Command{
-		Use:     "measurement",
-		Short:   "Add measurement to the backend info file. The value should be in base64. The second parameter is backend_info.json file",
-		Example: "measurement <measurement> <backend_info.json>",
-		Args:    cobra.ExactArgs(2),
-		Run: func(cmd *cobra.Command, args []string) {
-			if err := changeAttestationConfiguration(args[1], args[0], measurementLength, measurementField); err != nil {
-				log.Fatalf("Error could not change measurement data %v", err)
-			}
-		},
-	}
-}
-
-func (cli *CLI) NewAddHostDataCmd() *cobra.Command {
-	return &cobra.Command{
-		Use:     "hostdata",
-		Short:   "Add host data to the backend info file. The value should be in base64. The second parameter is backend_info.json file",
-		Example: "hostdata <host-data> <backend_info.json>",
-		Args:    cobra.ExactArgs(2),
-		Run: func(cmd *cobra.Command, args []string) {
-			if err := changeAttestationConfiguration(args[1], args[0], hostDataLength, hostDataField); err != nil {
-				log.Fatalf("Error could not change host data %v", err)
-			}
-		},
-	}
-}
-
-func changeAttestationConfiguration(fileName, base64Data string, expectedLength int, field fieldType) error {
-	data, err := base64.StdEncoding.DecodeString(base64Data)
-	if err != nil {
-		return errDecode
-	}
-
-	if len(data) != expectedLength {
-		return errDataLength
-	}
-
-	ac := AttestationConfiguration{}
-
-	backendInfo, err := os.ReadFile(fileName)
-	if err != nil {
-		return errors.Wrap(errReadingBackendInfoFile, err)
-	}
-
-	if err = json.Unmarshal(backendInfo, &ac); err != nil {
-		return errors.Wrap(errUnmarshalJSON, err)
-	}
-
-	switch field {
-	case measurementField:
-		ac.SNPPolicy.Measurement = data
-	case hostDataField:
-		ac.SNPPolicy.HostData = data
-	default:
-		return errBackendField
-	}
-
-	fileJson, err := json.MarshalIndent(ac, "", " ")
-	if err != nil {
-		return errors.Wrap(errMarshalJSON, err)
-	}
-	if err = os.WriteFile(fileName, fileJson, filePermission); err != nil {
-		return errors.Wrap(errWriteFile, err)
-	}
-	return nil
-}

cli/cache.go

@@ -3,15 +3,17 @@
 package cli
 
 import (
-	"log"
+	"fmt"
 	"os"
 	"path"
 
 	"github.com/google/go-sev-guest/abi"
 	"github.com/google/go-sev-guest/kds"
+	"github.com/google/go-sev-guest/proto/check"
 	"github.com/google/go-sev-guest/verify/trust"
 	"github.com/spf13/cobra"
-	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
+	"github.com/ultravioletrs/cocos/pkg/attestation"
+	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 )
 
 const (
@@ -26,30 +28,38 @@ func (cli *CLI) NewCABundleCmd(fileSavePath string) *cobra.Command {
 		Example: "ca-bundle <path_to_platform_info_json>",
 		Args:    cobra.ExactArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
-			attestationConfiguration := grpc.AttestationConfiguration{}
-			err := grpc.ReadManifest(args[0], &attestationConfiguration)
+			attestationConfiguration := attestation.Config{Config: &check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}, PcrConfig: &attestation.PcrConfig{}}
+			err := vtpm.ReadPolicy(args[0], &attestationConfiguration)
 			if err != nil {
-				log.Fatalf("Error while reading manifest: %v", err)
+				printError(cmd, "Error while reading manifest: %v ❌ ", err)
+				return
 			}
 
-			product := attestationConfiguration.RootOfTrust.Product
+			product := attestationConfiguration.Config.RootOfTrust.ProductLine
 
 			getter := trust.DefaultHTTPSGetter()
 			caURL := kds.ProductCertChainURL(abi.VcekReportSigner, product)
 
 			bundle, err := getter.Get(caURL)
 			if err != nil {
-				log.Fatalf("Error fetching ARK and ASK from AMD KDS for product: %s, error: %v", product, err)
+				message := fmt.Sprintf("Error fetching ARK and ASK from AMD KDS for product: %s", product)
+				message += ", error: %v ❌ "
+				printError(cmd, message, err)
+				return
 			}
 
 			err = os.MkdirAll(path.Join(fileSavePath, product), filePermisionKeys)
 			if err != nil {
-				log.Fatalf("Error while creating directory for product name %s, error: %v", product, err)
+				message := fmt.Sprintf("Error while creating directory for product name %s", product)
+				message += ", error: %v ❌ "
+				printError(cmd, message, err)
+				return
 			}
 
-			bundleFilePath := path.Join(fileSavePath, product, caBundleName)
-			if err = saveToFile(bundleFilePath, bundle); err != nil {
-				log.Fatalf("Error while saving ARK-ASK to file: %v", err)
+			bundlePath := path.Join(fileSavePath, product, caBundleName)
+			if err = saveToFile(bundlePath, bundle); err != nil {
+				printError(cmd, "Error while saving ARK-ASK to file: %v ❌ ", err)
+				return
 			}
 		},
 	}

cli/cache_test.go

@@ -0,0 +1,59 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"bytes"
+	"os"
+	"path"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewCABundleCmd(t *testing.T) {
+	cli := &CLI{}
+	tempDir, err := os.MkdirTemp("", "ca-bundle-test")
+	assert.NoError(t, err)
+	defer os.RemoveAll(tempDir)
+
+	manifestContent := []byte(`{"root_of_trust": {"product_line": "Milan"}}`)
+	manifestPath := path.Join(tempDir, "manifest.json")
+	err = os.WriteFile(manifestPath, manifestContent, 0o644)
+	assert.NoError(t, err)
+
+	cmd := cli.NewCABundleCmd(tempDir)
+	cmd.SetArgs([]string{manifestPath})
+	output := &bytes.Buffer{}
+	cmd.SetOutput(output)
+	err = cmd.Execute()
+
+	assert.NoError(t, err)
+
+	expectedFilePath := path.Join(tempDir, "Milan", caBundleName)
+	_, err = os.Stat(expectedFilePath)
+	assert.NoError(t, err)
+
+	content, err := os.ReadFile(expectedFilePath)
+	assert.NoError(t, err)
+	assert.NotNil(t, content)
+}
+
+func TestSaveToFile(t *testing.T) {
+	tempDir, err := os.MkdirTemp("", "save-to-file-test")
+	assert.NoError(t, err)
+	defer os.RemoveAll(tempDir)
+
+	filePath := path.Join(tempDir, "test-file.txt")
+	content := []byte("test content")
+
+	err = saveToFile(filePath, content)
+	assert.NoError(t, err)
+
+	savedContent, err := os.ReadFile(filePath)
+	assert.NoError(t, err)
+	assert.Equal(t, content, savedContent)
+
+	_, err = os.Stat(filePath)
+	assert.NoError(t, err)
+}

cli/checksum.go

@@ -0,0 +1,95 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"encoding/base64"
+	"encoding/hex"
+	"encoding/json"
+	"os"
+
+	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/internal"
+	"golang.org/x/crypto/sha3"
+)
+
+var (
+	ismanifest bool
+	toBase64   bool
+)
+
+func (cli *CLI) NewFileHashCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:     "checksum",
+		Short:   "Compute the sha3-256 hash of a file",
+		Example: "checksum <file>",
+		Args:    cobra.ExactArgs(1),
+		Run: func(cmd *cobra.Command, args []string) {
+			path := args[0]
+
+			if ismanifest {
+				hash, err := manifestChecksum(path)
+				if err != nil {
+					printError(cmd, "Error computing hash: %v ❌ ", err)
+					return
+				}
+
+				cmd.Println("Hash of manifest file:", hashOut(hash))
+				return
+			}
+
+			hash, err := internal.ChecksumHex(path)
+			if err != nil {
+				printError(cmd, "Error computing hash: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println("Hash of file:", hashOut(hash))
+		},
+	}
+
+	cmd.Flags().BoolVarP(&ismanifest, "manifest", "m", false, "Compute the hash of the manifest file")
+	cmd.Flags().BoolVarP(&toBase64, "base64", "b", false, "Output the hash in base64")
+
+	return cmd
+}
+
+func manifestChecksum(path string) (string, error) {
+	file, err := os.ReadFile(path)
+	if err != nil {
+		return "", err
+	}
+
+	var cmp agent.Computation
+
+	if err := json.Unmarshal(file, &cmp); err != nil {
+		return "", err
+	}
+
+	jsonBytes, err := json.Marshal(cmp)
+	if err != nil {
+		return "", err
+	}
+
+	sum := sha3.Sum256(jsonBytes)
+
+	return hex.EncodeToString(sum[:]), nil
+}
+
+func hashOut(hashHex string) string {
+	if toBase64 {
+		return hexToBase64(hashHex)
+	}
+
+	return hashHex
+}
+
+func hexToBase64(hexStr string) string {
+	decoded, err := hex.DecodeString(hexStr)
+	if err != nil {
+		return ""
+	}
+
+	return base64.StdEncoding.EncodeToString(decoded)
+}

cli/checksum_test.go

@@ -0,0 +1,230 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewFileHashCmd(t *testing.T) {
+	cli := &CLI{}
+	cmd := cli.NewFileHashCmd()
+
+	if cmd.Use != "checksum" {
+		t.Errorf("Expected Use to be 'checksum', got %s", cmd.Use)
+	}
+
+	if cmd.Short != "Compute the sha3-256 hash of a file" {
+		t.Errorf("Expected Short to be 'Compute the sha3-256 hash of a file', got %s", cmd.Short)
+	}
+
+	if cmd.Example != "checksum <file>" {
+		t.Errorf("Expected Example to be 'checksum <file>', got %s", cmd.Example)
+	}
+}
+
+func TestNewFileHashCmdRun(t *testing.T) {
+	testCases := []struct {
+		name        string
+		isManifest  bool
+		toBase64    bool
+		expectedOut string
+		expectedErr string
+	}{
+		{
+			name:        "Valid file",
+			isManifest:  false,
+			toBase64:    false,
+			expectedOut: "Hash of file:",
+			expectedErr: "",
+		},
+		{
+			name:        "Valid manifest file",
+			isManifest:  true,
+			toBase64:    false,
+			expectedOut: "Hash of manifest file:",
+			expectedErr: "",
+		},
+		{
+			name:        "Valid file with base64 output",
+			isManifest:  false,
+			toBase64:    true,
+			expectedOut: "Hash of file:",
+			expectedErr: "",
+		},
+		{
+			name:        "Non-existent file",
+			isManifest:  false,
+			toBase64:    false,
+			expectedOut: "Error computing hash:",
+			expectedErr: "",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			cli := &CLI{}
+			cmd := cli.NewFileHashCmd()
+
+			var output bytes.Buffer
+			cmd.SetOut(&output)
+			cmd.SetErr(&output)
+
+			err := cmd.Flags().Set("manifest", fmt.Sprint(tc.isManifest))
+			assert.Nil(t, err)
+			err = cmd.Flags().Set("base64", fmt.Sprint(tc.toBase64))
+			assert.Nil(t, err)
+
+			if tc.name == "Non-existent file" {
+				cmd.SetArgs([]string{"non_existent_file.txt"})
+			} else {
+				content := []byte("{}")
+				tmpfile, err := os.CreateTemp("", "example")
+				if err != nil {
+					t.Fatal(err)
+				}
+				defer os.Remove(tmpfile.Name())
+
+				if _, err := tmpfile.Write(content); err != nil {
+					t.Fatal(err)
+				}
+				if err := tmpfile.Close(); err != nil {
+					t.Fatal(err)
+				}
+
+				cmd.SetArgs([]string{tmpfile.Name()})
+			}
+
+			err = cmd.Execute()
+			if err != nil {
+				t.Fatalf("Error executing command: %v", err)
+			}
+
+			out := output.String()
+			if !strings.Contains(out, tc.expectedOut) {
+				t.Errorf("Expected output to contain '%s', got '%s'", tc.expectedOut, out)
+			}
+
+			if tc.expectedErr != "" && !strings.Contains(out, tc.expectedErr) {
+				t.Errorf("Expected output to contain '%s', got '%s'", tc.expectedErr, out)
+			}
+		})
+	}
+}
+
+func TestManifestChecksum(t *testing.T) {
+	testCases := []struct {
+		name        string
+		jsonContent string
+		expectedSum string
+	}{
+		{
+			name: "Valid manifest file",
+			jsonContent: `{
+				"id": "1234",
+				"name": "Example Computation",
+				"description": "This is an example computation"
+			}`,
+			expectedSum: "a99683e4d22ba54cefa51aa49fb2e97a92b828c088395992ddff16a6236f3299",
+		},
+		{
+			name:        "Invalid JSON",
+			jsonContent: `{`,
+			expectedSum: "",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			f, err := os.CreateTemp("", "test")
+			assert.Nil(t, err)
+
+			t.Cleanup(func() {
+				os.Remove(f.Name())
+			})
+
+			_, err = f.WriteString(tc.jsonContent)
+			assert.NoError(t, err)
+
+			err = f.Close()
+			assert.Nil(t, err)
+
+			hash, err := manifestChecksum(f.Name())
+			if tc.expectedSum == "" && err == nil {
+				t.Errorf("Expected error, got nil")
+			}
+			if tc.expectedSum != "" && err != nil {
+				t.Errorf("Unexpected error: %v", err)
+			}
+			if hash != tc.expectedSum {
+				t.Errorf("Expected hash %s, got %s", tc.expectedSum, hash)
+			}
+		})
+	}
+}
+
+func TestHexToBase64(t *testing.T) {
+	testCases := []struct {
+		name        string
+		hexInput    string
+		expectedOut string
+	}{
+		{
+			name:        "Valid hex input",
+			hexInput:    "48656c6c6f",
+			expectedOut: "SGVsbG8=",
+		},
+		{
+			name:        "Invalid hex input",
+			hexInput:    "invalid-hex",
+			expectedOut: "",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			out := hexToBase64(tc.hexInput)
+			if out != tc.expectedOut {
+				t.Errorf("Expected %s, got %s", tc.expectedOut, out)
+			}
+		})
+	}
+}
+
+func TestHashOut(t *testing.T) {
+	testCases := []struct {
+		name        string
+		hashHex     string
+		toBase64    bool
+		expectedOut string
+	}{
+		{
+			name:        "Hex output",
+			hashHex:     "48656c6c6f",
+			toBase64:    false,
+			expectedOut: "48656c6c6f",
+		},
+		{
+			name:        "Base64 output",
+			hashHex:     "48656c6c6f",
+			toBase64:    true,
+			expectedOut: "SGVsbG8=",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			toBase64 = tc.toBase64
+			out := hashOut(tc.hashHex)
+			if out != tc.expectedOut {
+				t.Errorf("Expected %s, got %s", tc.expectedOut, out)
+			}
+		})
+	}
+}

cli/datasets.go

@@ -3,72 +3,117 @@
 package cli
 
 import (
+	"context"
 	"crypto/x509"
 	"encoding/pem"
-	"log"
 	"os"
+	"path"
 
+	"github.com/absmach/supermq/pkg/errors"
+	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/internal"
+	"google.golang.org/grpc/metadata"
 )
 
+var decompressDataset bool
+
 func (cli *CLI) NewDatasetsCmd() *cobra.Command {
-	return &cobra.Command{
+	cmd := &cobra.Command{
 		Use:     "data",
-		Short:   "Upload a dataset CSV file",
-		Example: "data <dataset.csv> <private_key_file_path>",
+		Short:   "Upload a dataset",
+		Example: "data <dataset_path> <private_key_file_path>",
 		Args:    cobra.ExactArgs(2),
 		Run: func(cmd *cobra.Command, args []string) {
-			datasetFile := args[0]
-
-			log.Println("Uploading dataset CSV:", datasetFile)
-
-			dataset, err := os.ReadFile(datasetFile)
-			if err != nil {
-				log.Fatalf("Error reading dataset file: %v", err)
+			if cli.connectErr != nil {
+				printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
+				return
 			}
 
-			dataReq := agent.Dataset{
-				Dataset: dataset,
+			datasetPath := args[0]
+
+			cmd.Println("Uploading dataset:", datasetPath)
+
+			f, err := os.Stat(datasetPath)
+			if err != nil {
+				printError(cmd, "Error reading dataset file: %v ❌ ", err)
+				return
+			}
+
+			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)
+					return
+				}
+				defer dataset.Close()
+				defer os.Remove(dataset.Name())
+			} else {
+				dataset, err = os.Open(datasetPath)
+				if err != nil {
+					printError(cmd, "Error reading dataset file: %v ❌ ", err)
+					return
+				}
+				defer dataset.Close()
 			}
 
 			privKeyFile, err := os.ReadFile(args[1])
 			if err != nil {
-				log.Fatalf("Error reading private key file: %v", err)
+				printError(cmd, "Error reading private key file: %v ❌ ", err)
+				return
 			}
 
 			pemBlock, _ := pem.Decode(privKeyFile)
 
-			privKey := decodeKey(pemBlock)
-
-			if err := cli.agentSDK.Data(cmd.Context(), dataReq, privKey); err != nil {
-				log.Fatalf("Error uploading dataset: %v", err)
+			privKey, err := decodeKey(pemBlock)
+			if err != nil {
+				printError(cmd, "Error decoding private key: %v ❌ ", err)
+				return
 			}
 
-			log.Println("Successfully uploaded dataset")
+			ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
+			if err := cli.agentSDK.Data(addDatasetMetadata(ctx), dataset, path.Base(datasetPath), privKey); err != nil {
+				printError(cmd, "Failed to upload dataset due to error: %v ❌ ", err)
+				return
+			}
+
+			cmd.Println(color.New(color.FgGreen).Sprint("Successfully uploaded dataset! ✔ "))
 		},
 	}
+
+	cmd.Flags().BoolVarP(&decompressDataset, "decompress", "d", false, "Decompress the dataset on agent")
+	return cmd
 }
 
-func decodeKey(b *pem.Block) interface{} {
+func decodeKey(b *pem.Block) (any, error) {
+	if b == nil {
+		return nil, errors.New("error decoding key")
+	}
 	switch b.Type {
 	case rsaKeyType:
 		privKey, err := x509.ParsePKCS8PrivateKey(b.Bytes)
 		if err != nil {
 			privKey, err = x509.ParsePKCS1PrivateKey(b.Bytes)
 			if err != nil {
-				log.Fatalf("Error parsing private key: %v", err)
+				return nil, err
 			}
 		}
-		return privKey
+		return privKey, nil
 	case ecdsaKeyType:
 		privKey, err := x509.ParseECPrivateKey(b.Bytes)
 		if err != nil {
-			log.Fatalf("Error parsing private key: %v", err)
+			return nil, err
 		}
-		return privKey
+		return privKey, nil
 	default:
-		log.Fatalf("Error decoding key")
-		return nil
+		return nil, errors.New("error decoding key")
 	}
 }
+
+func addDatasetMetadata(ctx context.Context) context.Context {
+	return agent.DecompressToContext(ctx, decompressDataset)
+}

cli/datasets_test.go

@@ -0,0 +1,158 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"bytes"
+	"errors"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
+)
+
+func createTempDatasetFile(content string) (string, error) {
+	tmpFile, err := os.CreateTemp("", "dataset-*.txt")
+	if err != nil {
+		return "", err
+	}
+	defer tmpFile.Close()
+
+	_, err = tmpFile.WriteString(content)
+	if err != nil {
+		return "", err
+	}
+	return tmpFile.Name(), nil
+}
+
+func TestDatasetsCmd(t *testing.T) {
+	tests := []struct {
+		name           string
+		setupMock      func(*mocks.SDK)
+		setupFiles     func() (string, error)
+		connectErr     error
+		expectedOutput string
+		cleanup        func(string, string)
+	}{
+		{
+			name: "successful upload",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Data", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			},
+			setupFiles: func() (string, error) {
+				datasetFile, err := createTempDatasetFile("test dataset content")
+				if err != nil {
+					return "", err
+				}
+				err = generateRSAPrivateKeyFile(privateKeyFile)
+				return datasetFile, err
+			},
+			expectedOutput: "Successfully uploaded dataset",
+			cleanup: func(datasetFile, privateKeyFile string) {
+				os.Remove(datasetFile)
+				os.Remove(privateKeyFile)
+			},
+		},
+		{
+			name: "missing dataset file",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Data", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			},
+			setupFiles: func() (string, error) {
+				return "", nil
+			},
+			expectedOutput: "Error reading dataset file",
+		},
+		{
+			name: "missing private key file",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Data", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			},
+			setupFiles: func() (string, error) {
+				return createTempDatasetFile("test dataset content")
+			},
+			expectedOutput: "Error reading private key file",
+			cleanup: func(datasetFile, _ string) {
+				os.Remove(datasetFile)
+			},
+		},
+		{
+			name: "upload failure",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Data", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(errors.New("failed to upload algorithm due to error"))
+			},
+			setupFiles: func() (string, error) {
+				datasetFile, err := createTempDatasetFile("test dataset content")
+				if err != nil {
+					return "", err
+				}
+				err = generateRSAPrivateKeyFile(privateKeyFile)
+				return datasetFile, err
+			},
+			expectedOutput: "Failed to upload dataset due to error",
+			cleanup: func(datasetFile, privateKeyFile string) {
+				os.Remove(datasetFile)
+				os.Remove(privateKeyFile)
+			},
+		},
+		{
+			name: "invalid private key",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Data", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
+			},
+			setupFiles: func() (string, error) {
+				datasetFile, err := createTempDatasetFile("test dataset content")
+				if err != nil {
+					return "", err
+				}
+				err = os.WriteFile(privateKeyFile, []byte("invalid private key"), 0o644)
+				return datasetFile, err
+			},
+			expectedOutput: "Error decoding private key",
+			cleanup: func(datasetFile, privateKeyFile string) {
+				os.Remove(datasetFile)
+				os.Remove(privateKeyFile)
+			},
+		},
+		{
+			name: "connection error",
+			setupMock: func(m *mocks.SDK) {
+			},
+			setupFiles:     func() (string, error) { return "", nil },
+			connectErr:     errors.New("failed to connect to agent"),
+			expectedOutput: "Failed to connect to agent",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockSDK := new(mocks.SDK)
+			if tt.setupMock != nil {
+				tt.setupMock(mockSDK)
+			}
+
+			testCLI := CLI{
+				agentSDK:   mockSDK,
+				connectErr: tt.connectErr,
+			}
+
+			datasetFile, err := tt.setupFiles()
+			require.NoError(t, err)
+
+			cmd := testCLI.NewDatasetsCmd()
+			buf := new(bytes.Buffer)
+			cmd.SetOut(buf)
+			cmd.SetArgs([]string{datasetFile, privateKeyFile})
+			err = cmd.Execute()
+			require.NoError(t, err)
+
+			require.Contains(t, buf.String(), tt.expectedOutput)
+
+			if tt.cleanup != nil {
+				tt.cleanup(datasetFile, privateKeyFile)
+			}
+		})
+	}
+}