NOISSUE - Fix TLS connection (#329)

* fix tls

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

* refactor

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

---------

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

.env

@@ -1,28 +0,0 @@
-## Jaeger
-COCOS_JAEGER_PORT=6831
-COCOS_JAEGER_FRONTEND=16686
-COCOS_JAEGER_COLLECTOR=14268
-COCOS_JAEGER_CONFIGS=5778
-COCOS_JAEGER_URL=http://jaeger:4318
-COCOS_JAEGER_TRACE_RATIO=1.0
-COCOS_JAEGER_COLLECTOR_OTLP_ENABLED=true
-COCOS_JAEGER_OLTP_HTTP_PORT=4318
-
-## Core Services
-
-### Manager
-MANAGER_GRPC_HOST=""
-MANAGER_GRPC_PORT=7003
-MANAGER_GRPC_SERVER_CERT=""
-MANAGER_GRPC_SERVER_KEY=""
-AGENT_GRPC_URL="localhost: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

.github/dependabot.yaml

@@ -0,0 +1,22 @@
+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"
+    groups:
+      go-dependency:
+        patterns:
+          - "*"

.github/workflows/checkproto.yaml

@@ -29,13 +29,13 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: 1.22.x
+          go-version: 1.23.x
 
       - name: Set up protoc
         run: |
-          PROTOC_VERSION=27.2
+          PROTOC_VERSION=28.1
           PROTOC_GEN_VERSION=v1.34.2
-          PROTOC_GRPC_VERSION=v1.4.0
+          PROTOC_GRPC_VERSION=v1.5.1
 
           # Download and install protoc
           PROTOC_ZIP=protoc-$PROTOC_VERSION-linux-x86_64.zip

.github/workflows/hal.yml

@@ -8,66 +8,66 @@ on:
 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.23.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: 2024.11-rc2
+
+      - name: Build
+        run: |
+          cd buildroot
+          make BR2_EXTERNAL=../cocos/hal/linux cocos_defconfig
+          make
+
+      - name: Release
+        uses: softprops/action-gh-release@v2
+        with:
+          files: |
+            buildroot/output/images/bzImage
+            buildroot/output/images/rootfs.cpio.gz
 
-    - 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

.github/workflows/main.yaml

@@ -19,16 +19,43 @@ jobs:
       - name: Install Go
         uses: actions/setup-go@v5
         with:
-          go-version: 1.22.x
+          go-version: 1.23.x
 
       - name: golangci-lint
         uses: golangci/golangci-lint-action@v6
         with:
-          version: v1.59.1
+          version: v1.60
 
       - name: Build
         run: |
           make
 
-      - name: Run tests
-        run: go test -v --race -covermode=atomic -coverprofile cover.out ./...
+      - name: Create coverage directory
+        run: |
+          mkdir coverage
+
+      - name: Run Agent tests
+        run: go test -v --race -covermode=atomic -coverprofile coverage/agent.out ./agent/...
+
+      - name: Run cli tests
+        run: go test -v --race -covermode=atomic -coverprofile coverage/cli.out ./cli/...
+
+      - name: Run cmd tests
+        run: go test -v --race -covermode=atomic -coverprofile coverage/cmd.out ./cmd/...
+
+      - name: Run internal tests
+        run: go test -v --race -covermode=atomic -coverprofile coverage/internal.out ./internal/...
+
+      - name: Run pkg tests
+        run: go test -v --race -covermode=atomic -coverprofile coverage/pkg.out ./pkg/...
+
+      - name: Run manager tests
+        run: sudo go test -v --race -covermode=atomic -coverprofile coverage/manager.out ./manager/...
+
+      - name: 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,6 +1,5 @@
 build
 
-build
 cmd/manager/img
 
 .cov
@@ -8,5 +7,21 @@ cmd/manager/img
 *.pem
 
 dist/
-results.zip
+*.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

@@ -68,7 +68,7 @@ linters:
     - errchkjson
     - errname
     - execinquery
-    - exportloopref
+    - copyloopvar
     - ginkgolinter
     - gocheckcompilerdirectives
     - gofumpt

Makefile

@@ -1,12 +1,17 @@
 BUILD_DIR = build
 SERVICES = manager agent cli
-BACKEND_INFO = backend_info
-CGO_ENABLED ?= 0
+ATTESTATION_POLICY = attestation_policy
+CGO_ENABLED ?= 1
 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
 
 define compile_service
 	CGO_ENABLED=$(CGO_ENABLED) GOOS=$(GOOS) GOARCH=$(GOARCH) GOARM=$(GOARM) \
@@ -18,19 +23,40 @@ define compile_service
 	-o ${BUILD_DIR}/cocos-$(1) cmd/$(1)/main.go
 endef
 
-.PHONY: all $(SERVICES) $(BACKEND_INFO)
+.PHONY: all $(SERVICES) $(ATTESTATION_POLICY) install clean
 
 all: $(SERVICES)
 
 $(SERVICES):
 	$(call compile_service,$@)
 
-$(BACKEND_INFO):
-	$(MAKE) -C ./scripts/backend_info
+$(ATTESTATION_POLICY):
+	$(MAKE) -C ./scripts/attestation_policy
 
 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
 
 mocks:
-	go generate ./...
+	mockery --config ./mockery.yml
+
+install: $(SERVICES)
+	install -d $(INSTALL_DIR)
+	install $(BUILD_DIR)/cocos-cli $(INSTALL_DIR)/cocos-cli
+	install $(BUILD_DIR)/cocos-manager $(INSTALL_DIR)/cocos-manager
+	install -d $(CONFIG_DIR)
+	install cocos-manager.env $(CONFIG_DIR)/cocos-manager.env
+
+clean:
+	rm -rf $(BUILD_DIR)
+
+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

README.md

@@ -1,4 +1,8 @@
 # Cocos AI
+
+[![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)
+
 [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].
 

agent/agent.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.34.2
-// 	protoc        v5.27.2
+// 	protoc        v5.28.1
 // source: agent/agent.proto
 
 package agent
@@ -413,7 +413,7 @@ var file_agent_agent_proto_rawDesc = []byte{
 	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, 0xfb, 0x01, 0x0a, 0x0c, 0x41, 0x67,
+	0x01, 0x28, 0x0c, 0x52, 0x04, 0x66, 0x69, 0x6c, 0x65, 0x32, 0xfd, 0x01, 0x0a, 0x0c, 0x41, 0x67,
 	0x65, 0x6e, 0x74, 0x53, 0x65, 0x72, 0x76, 0x69, 0x63, 0x65, 0x12, 0x33, 0x0a, 0x04, 0x41, 0x6c,
 	0x67, 0x6f, 0x12, 0x12, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x6c, 0x67, 0x6f, 0x52,
 	0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x13, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41,
@@ -425,12 +425,12 @@ var file_agent_agent_proto_rawDesc = []byte{
 	0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73, 0x75, 0x6c, 0x74, 0x52, 0x65, 0x71,
 	0x75, 0x65, 0x73, 0x74, 0x1a, 0x15, 0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x52, 0x65, 0x73,
 	0x75, 0x6c, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x12,
-	0x46, 0x0a, 0x0b, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x19,
+	0x48, 0x0a, 0x0b, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x19,
 	0x2e, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69,
 	0x6f, 0x6e, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1a, 0x2e, 0x61, 0x67, 0x65, 0x6e,
 	0x74, 0x2e, 0x41, 0x74, 0x74, 0x65, 0x73, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x65, 0x73,
-	0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x09, 0x5a, 0x07, 0x2e, 0x2f, 0x61, 0x67, 0x65,
-	0x6e, 0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
+	0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x30, 0x01, 0x42, 0x09, 0x5a, 0x07, 0x2e, 0x2f, 0x61,
+	0x67, 0x65, 0x6e, 0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
 }
 
 var (

agent/agent.proto

@@ -11,7 +11,7 @@ service AgentService {
   rpc Algo(stream AlgoRequest) returns (AlgoResponse) {}
   rpc Data(stream DataRequest) returns (DataResponse) {}
   rpc Result(ResultRequest) returns (stream ResultResponse) {}
-  rpc Attestation(AttestationRequest) returns (AttestationResponse) {}
+  rpc Attestation(AttestationRequest) returns (stream AttestationResponse) {}
 }
 
 message AlgoRequest {

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.28.1
 // source: agent/agent.proto
 
 package agent
@@ -18,8 +18,8 @@ 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"
@@ -32,10 +32,10 @@ const (
 //
 // 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) (AgentService_ResultClient, 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)
 }
 
 type agentServiceClient struct {
@@ -46,83 +46,39 @@ 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) (AgentService_ResultClient, error) {
+func (c *agentServiceClient) Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (grpc.ServerStreamingClient[ResultResponse], error) {
 	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[2], AgentService_Result_FullMethodName, cOpts...)
 	if err != nil {
 		return nil, err
 	}
-	x := &agentServiceResultClient{ClientStream: stream}
+	x := &grpc.GenericClientStream[ResultRequest, ResultResponse]{ClientStream: stream}
 	if err := x.ClientStream.SendMsg(in); err != nil {
 		return nil, err
 	}
@@ -132,61 +88,60 @@ func (c *agentServiceClient) Result(ctx context.Context, in *ResultRequest, opts
 	return x, nil
 }
 
-type AgentService_ResultClient interface {
-	Recv() (*ResultResponse, 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_ResultClient = grpc.ServerStreamingClient[ResultResponse]
 
-type agentServiceResultClient struct {
-	grpc.ClientStream
-}
-
-func (x *agentServiceResultClient) Recv() (*ResultResponse, error) {
-	m := new(ResultResponse)
-	if err := x.ClientStream.RecvMsg(m); err != nil {
-		return nil, err
-	}
-	return m, nil
-}
-
-func (c *agentServiceClient) Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (*AttestationResponse, error) {
+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
 	}
-	return out, nil
+	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]
+
 // 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(*ResultRequest, AgentService_ResultServer) 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
 	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(*ResultRequest, AgentService_ResultServer) error {
+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) 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
@@ -196,99 +151,51 @@ 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
-}
-
-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
-	}
-	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_DataServer = grpc.ClientStreamingServer[DataRequest, DataResponse]
 
 func _AgentService_Result_Handler(srv interface{}, stream grpc.ServerStream) error {
 	m := new(ResultRequest)
 	if err := stream.RecvMsg(m); err != nil {
 		return err
 	}
-	return srv.(AgentServiceServer).Result(m, &agentServiceResultServer{ServerStream: stream})
+	return srv.(AgentServiceServer).Result(m, &grpc.GenericServerStream[ResultRequest, ResultResponse]{ServerStream: stream})
 }
 
-type AgentService_ResultServer interface {
-	Send(*ResultResponse) 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_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})
 }
 
-type agentServiceResultServer struct {
-	grpc.ServerStream
-}
-
-func (x *agentServiceResultServer) Send(m *ResultResponse) error {
-	return x.ServerStream.SendMsg(m)
-}
-
-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 interceptor(ctx, in, info, handler)
-}
+// This type alias is provided for backwards compatibility with existing code that references the prior non-generic stream type by name.
+type AgentService_AttestationServer = grpc.ServerStreamingServer[AttestationResponse]
 
 // AgentService_ServiceDesc is the grpc.ServiceDesc for AgentService service.
 // It's only intended for direct use with grpc.RegisterService,
@@ -296,12 +203,7 @@ func _AgentService_Attestation_Handler(srv interface{}, ctx context.Context, dec
 var AgentService_ServiceDesc = grpc.ServiceDesc{
 	ServiceName: "agent.AgentService",
 	HandlerType: (*AgentServiceServer)(nil),
-	Methods: []grpc.MethodDesc{
-		{
-			MethodName: "Attestation",
-			Handler:    _AgentService_Attestation_Handler,
-		},
-	},
+	Methods:     []grpc.MethodDesc{},
 	Streams: []grpc.StreamDesc{
 		{
 			StreamName:    "Algo",
@@ -318,6 +220,11 @@ var AgentService_ServiceDesc = grpc.ServiceDesc{
 			Handler:       _AgentService_Result_Handler,
 			ServerStreams: true,
 		},
+		{
+			StreamName:    "Attestation",
+			Handler:       _AgentService_Attestation_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/binary/binary.go

@@ -9,6 +9,7 @@ import (
 	"os/exec"
 
 	"github.com/ultravioletrs/cocos/agent/algorithm"
+	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
 	"github.com/ultravioletrs/cocos/agent/events"
 )
 
@@ -24,8 +25,8 @@ type binary struct {
 func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string, args []string) algorithm.Algorithm {
 	return &binary{
 		algoFile: algoFile,
-		stderr:   &algorithm.Stderr{Logger: logger, EventSvc: eventsSvc},
-		stdout:   &algorithm.Stdout{Logger: logger},
+		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
+		stdout:   &logging.Stdout{Logger: logger},
 		args:     args,
 	}
 }

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

@@ -3,6 +3,7 @@
 package docker
 
 import (
+	"bufio"
 	"context"
 	"fmt"
 	"io"
@@ -15,6 +16,7 @@ import (
 	"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"
 )
 
@@ -37,8 +39,8 @@ func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string
 	d := &docker{
 		algoFile: algoFile,
 		logger:   logger,
-		stderr:   &algorithm.Stderr{Logger: logger, EventSvc: eventsSvc},
-		stdout:   &algorithm.Stdout{Logger: logger},
+		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
+		stdout:   &logging.Stdout{Logger: logger},
 	}
 
 	return d
@@ -89,7 +91,10 @@ func (d *docker) Run() error {
 
 	// Create and start the container.
 	respContainer, err := cli.ContainerCreate(ctx, &container.Config{
-		Image: dockerImageName,
+		Image:        dockerImageName,
+		Tty:          true,
+		AttachStdout: true,
+		AttachStderr: true,
 	}, &container.HostConfig{
 		Mounts: []mount.Mount{
 			{
@@ -112,6 +117,30 @@ func (d *docker) Run() error {
 		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:
@@ -121,28 +150,6 @@ func (d *docker) Run() error {
 	case <-statusCh:
 	}
 
-	stdout, err := cli.ContainerLogs(ctx, respContainer.ID, container.LogsOptions{ShowStdout: true})
-	if err != nil {
-		return fmt.Errorf("could not read stdout from the container: %v", err)
-	}
-	defer stdout.Close()
-
-	err = writeToOut(stdout, d.stdout)
-	if 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})
-	if err != nil {
-		d.logger.Warn(fmt.Sprintf("could not read stderr from the container: %v", err))
-	}
-	defer stderr.Close()
-
-	err = writeToOut(stderr, d.stderr)
-	if err != nil {
-		d.logger.Warn(fmt.Sprintf("could not write to stderr: %v", err))
-	}
-
 	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))
@@ -157,13 +164,14 @@ func (d *docker) Run() error {
 }
 
 func writeToOut(readCloser io.ReadCloser, ioWriter io.Writer) error {
-	content, err := io.ReadAll(readCloser)
-	if err != nil {
-		return fmt.Errorf("could not convert content from the container: %v", err)
+	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 := ioWriter.Write(content); err != nil {
-		return fmt.Errorf("could not write to output: %v", err)
+	if err := scanner.Err(); err != nil {
+		return fmt.Errorf("error reading container logs error: %v", err)
 	}
 
 	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
@@ -65,7 +70,7 @@ func (s *Stderr) Write(p []byte) (n int, err error) {
 		s.Logger.Error(string(buf[:n]))
 	}
 
-	if err := s.EventSvc.SendEvent("algorithm-run", "error", nil); err != nil {
+	if err := s.EventSvc.SendEvent(algorithmRun, warningStatus, json.RawMessage{}); err != nil {
 		return len(p), err
 	}
 

agent/algorithm/logging/logging_test.go

@@ -1,6 +1,6 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
-package algorithm
+package logging
 
 import (
 	"strings"
@@ -10,6 +10,7 @@ import (
 	"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", "error", mock.Anything).Return(nil)
+			mockEventService.On("SendEvent", "AlgorithmRun", manager.Warning.String(), mock.Anything).Return(nil)
 
 			stderr := &Stderr{Logger: mglog.NewMock(), EventSvc: mockEventService}
 			n, err := stderr.Write([]byte(tt.input))

agent/algorithm/python/python.go

@@ -12,21 +12,22 @@ import (
 	"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"
+	PyRuntimeKey = "python_runtime"
 )
 
 func PythonRunTimeToContext(ctx context.Context, runtime string) context.Context {
-	return metadata.AppendToOutgoingContext(ctx, pyRuntimeKey, runtime)
+	return metadata.AppendToOutgoingContext(ctx, PyRuntimeKey, runtime)
 }
 
 func PythonRunTimeFromContext(ctx context.Context) string {
-	return metadata.ValueFromIncomingContext(ctx, pyRuntimeKey)[0]
+	return metadata.ValueFromIncomingContext(ctx, PyRuntimeKey)[0]
 }
 
 var _ algorithm.Algorithm = (*python)(nil)
@@ -43,8 +44,8 @@ type python struct {
 func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, runtime, requirementsFile, algoFile string, args []string) algorithm.Algorithm {
 	p := &python{
 		algoFile:         algoFile,
-		stderr:           &algorithm.Stderr{Logger: logger, EventSvc: eventsSvc},
-		stdout:           &algorithm.Stdout{Logger: logger},
+		stderr:           &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
+		stdout:           &logging.Stdout{Logger: logger},
 		requirementsFile: requirementsFile,
 		args:             args,
 	}
@@ -67,6 +68,13 @@ func (p *python) Run() error {
 
 	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

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/wasm/wasm.go

@@ -9,6 +9,7 @@ import (
 	"os/exec"
 
 	"github.com/ultravioletrs/cocos/agent/algorithm"
+	"github.com/ultravioletrs/cocos/agent/algorithm/logging"
 	"github.com/ultravioletrs/cocos/agent/events"
 )
 
@@ -28,8 +29,8 @@ type wasm struct {
 func NewAlgorithm(logger *slog.Logger, eventsSvc events.Service, algoFile string, args []string) algorithm.Algorithm {
 	return &wasm{
 		algoFile: algoFile,
-		stderr:   &algorithm.Stderr{Logger: logger, EventSvc: eventsSvc},
-		stdout:   &algorithm.Stdout{Logger: logger},
+		stderr:   &logging.Stderr{Logger: logger, EventSvc: eventsSvc},
+		stdout:   &logging.Stdout{Logger: logger},
 		args:     args,
 	}
 }

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, algoFile, args)
+
+	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, algoFile, args).(*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_test.go

@@ -0,0 +1,173 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package grpc
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/agent/mocks"
+	"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{ReportData: sha3.Sum512([]byte("report data"))},
+		},
+		{
+			name:        "Service Error",
+			req:         attestationReq{ReportData: sha3.Sum512([]byte("report data"))},
+			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.ReportData).Return([]byte{}, errors.New("")).Once()
+			} else {
+				svc.On("Attestation", context.Background(), tt.req.ReportData).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)
+				}
+			}
+		})
+	}
+}

agent/api/grpc/interceptors.go

@@ -35,20 +35,20 @@ func (s *authInterceptor) AuthStreamInterceptor() grpc.StreamServerInterceptor {
 		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, err.Error())
+				return status.Errorf(codes.Unauthenticated, "%v", err.Error())
 			}
 			wrapped := &wrappedServerStream{ServerStream: stream, ctx: ctx}
 			return handler(srv, wrapped)
@@ -64,7 +64,7 @@ func (s *authInterceptor) AuthUnaryInterceptor() grpc.UnaryServerInterceptor {
 		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"
 )

agent/api/grpc/server.go

@@ -6,15 +6,20 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"io"
 
 	"github.com/go-kit/kit/transport/grpc"
 	"github.com/ultravioletrs/cocos/agent"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
 )
 
-const bufferSize = 1024 * 1024
+const (
+	bufferSize  = 1024 * 1024
+	FileSizeKey = "file-size"
+)
 
 var _ agent.AgentServiceServer = (*grpcServer)(nil)
 
@@ -156,12 +161,16 @@ func (s *grpcServer) Result(req *agent.ResultRequest, stream agent.AgentService_
 	}
 	rr := res.(*agent.ResultResponse)
 
-	reusltBuffer := bytes.NewBuffer(rr.File)
+	if err := stream.SetHeader(metadata.New(map[string]string{FileSizeKey: fmt.Sprint(len(rr.File))})); err != nil {
+		return status.Error(codes.Internal, err.Error())
+	}
+
+	resultBuffer := bytes.NewBuffer(rr.File)
 
 	buf := make([]byte, bufferSize)
 
 	for {
-		n, err := reusltBuffer.Read(buf)
+		n, err := resultBuffer.Read(buf)
 		if err == io.EOF {
 			break
 		}
@@ -177,11 +186,34 @@ func (s *grpcServer) Result(req *agent.ResultRequest, stream agent.AgentService_
 	return nil
 }
 
-func (s *grpcServer) Attestation(ctx context.Context, req *agent.AttestationRequest) (*agent.AttestationResponse, error) {
-	_, res, err := s.attestation.ServeGRPC(ctx, req)
+func (s *grpcServer) Attestation(req *agent.AttestationRequest, stream agent.AgentService_AttestationServer) error {
+	_, res, err := s.attestation.ServeGRPC(stream.Context(), req)
 	if err != nil {
-		return nil, err
+		return err
 	}
 	rr := res.(*agent.AttestationResponse)
-	return rr, nil
+
+	if err := stream.SetHeader(metadata.New(map[string]string{FileSizeKey: fmt.Sprint(len(rr.File))})); err != nil {
+		return status.Error(codes.Internal, err.Error())
+	}
+
+	attestationBuffer := bytes.NewBuffer(rr.File)
+
+	buf := make([]byte, bufferSize)
+
+	for {
+		n, err := attestationBuffer.Read(buf)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			return status.Error(codes.Internal, err.Error())
+		}
+
+		if err := stream.Send(&agent.AttestationResponse{File: buf[:n]}); err != nil {
+			return status.Error(codes.Internal, err.Error())
+		}
+	}
+
+	return nil
 }

agent/api/grpc/server_test.go

@@ -0,0 +1,213 @@
+// 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"
+	"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(metadata.MD) error {
+	return nil
+}
+
+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(metadata.MD) error {
+	return nil
+}
+
+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)
+	mockStream.On("SendAndClose", &agent.AlgoResponse{}).Return(nil)
+
+	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 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)
+	mockStream.On("SendAndClose", &agent.DataResponse{}).Return(nil)
+
+	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)
+
+	mockStream := &MockAgentService_ResultServer{ctx: context.Background()}
+	mockService.On("Result", mock.Anything).Return([]byte("result data"), nil)
+	mockStream.On("Send", mock.AnythingOfType("*agent.ResultResponse")).Return(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)
+
+	mockStream := &MockAgentService_AttestationServer{ctx: context.Background()}
+	mockStream.On("Send", mock.AnythingOfType("*agent.AttestationResponse")).Return(nil)
+
+	reportData := [agent.ReportDataSize]byte{}
+	mockService.On("Attestation", mock.Anything, reportData).Return([]byte("attestation data"), nil)
+
+	err := server.Attestation(&agent.AttestationRequest{ReportData: reportData[:]}, mockStream)
+	assert.NoError(t, err)
+
+	mockService.AssertExpectations(t)
+}
+
+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) {
+	reportData := [agent.ReportDataSize]byte{}
+	req := &agent.AttestationRequest{ReportData: reportData[:]}
+	decoded, err := decodeAttestationRequest(context.Background(), req)
+	assert.NoError(t, err)
+	assert.Equal(t, attestationReq{ReportData: reportData}, decoded)
+}
+
+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)
+}

agent/auth/auth.go

@@ -36,7 +36,6 @@ 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)
 }

agent/auth/mocks/auth.go

@@ -1,8 +1,8 @@
-// Code generated by mockery v2.43.2. DO NOT EDIT.
-
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
 package mocks
 
 import (
@@ -18,6 +18,14 @@ 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 with given fields: ctx, role
 func (_m *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRole) (context.Context, error) {
 	ret := _m.Called(ctx, role)
@@ -48,6 +56,35 @@ func (_m *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRol
 	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) {
+		run(args[0].(context.Context), args[1].(auth.UserRole))
+	})
+	return _c
+}
+
+func (_c *Authenticator_AuthenticateUser_Call) Return(_a0 context.Context, _a1 error) *Authenticator_AuthenticateUser_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *Authenticator_AuthenticateUser_Call) RunAndReturn(run func(context.Context, auth.UserRole) (context.Context, error)) *Authenticator_AuthenticateUser_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // NewAuthenticator creates a new instance of Authenticator. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
 // The first argument is typically a *testing.T value.
 func NewAuthenticator(t interface {

agent/computations.go

@@ -13,14 +13,14 @@ import (
 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"`
+	LogLevel     string `json:"log_level,omitempty"`
+	Host         string `json:"host,omitempty"`
+	Port         string `json:"port,omitempty"`
+	CertFile     string `json:"cert_file,omitempty"`
+	KeyFile      string `json:"server_key,omitempty"`
+	ServerCAFile string `json:"server_ca_file,omitempty"`
+	ClientCAFile string `json:"client_ca_file,omitempty"`
+	AttestedTls  bool   `json:"attested_tls,omitempty"`
 }
 
 type Computation struct {

agent/computations_test.go

@@ -0,0 +1,133 @@
+// 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{
+		LogLevel:     "info",
+		Host:         "localhost",
+		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/events/events.go

@@ -4,10 +4,8 @@ package events
 
 import (
 	"encoding/json"
-	"time"
+	"io"
 
-	"github.com/mdlayher/vsock"
-	"github.com/ultravioletrs/cocos/pkg/manager"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -15,29 +13,14 @@ import (
 type service struct {
 	service       string
 	computationID string
-	conn          *vsock.Conn
+	conn          io.Writer
 }
 
-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
 }
 
-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, computationID string, conn io.Writer) (Service, error) {
 	return &service{
 		service:       svc,
 		computationID: computationID,
@@ -46,7 +29,7 @@ func New(svc, computationID string, sockPort uint32) (Service, error) {
 }
 
 func (s *service) SendEvent(event, status string, details json.RawMessage) error {
-	body := manager.ClientStreamMessage{Message: &manager.ClientStreamMessage_AgentEvent{AgentEvent: &manager.AgentEvent{
+	body := EventsLogs{Message: &EventsLogs_AgentEvent{AgentEvent: &AgentEvent{
 		EventType:     event,
 		Timestamp:     timestamppb.Now(),
 		ComputationId: s.computationID,
@@ -58,12 +41,6 @@ func (s *service) SendEvent(event, status string, details json.RawMessage) error
 	if err != nil {
 		return err
 	}
-	if _, err := s.conn.Write(protoBody); err != nil {
-		return err
-	}
-	return nil
-}
-
-func (s *service) Close() error {
-	return s.conn.Close()
+	_, err = s.conn.Write(protoBody)
+	return err
 }

agent/events/events.pb.go

@@ -0,0 +1,405 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by protoc-gen-go. DO NOT EDIT.
+// versions:
+// 	protoc-gen-go v1.34.2
+// 	protoc        v5.28.1
+// 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"
+)
+
+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
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	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"`
+}
+
+func (x *AgentEvent) Reset() {
+	*x = AgentEvent{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_agent_events_events_proto_msgTypes[0]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+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 protoimpl.UnsafeEnabled && 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
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	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"`
+}
+
+func (x *AgentLog) Reset() {
+	*x = AgentLog{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_agent_events_events_proto_msgTypes[1]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+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 protoimpl.UnsafeEnabled && 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
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	// Types that are assignable to Message:
+	//
+	//	*EventsLogs_AgentLog
+	//	*EventsLogs_AgentEvent
+	Message isEventsLogs_Message `protobuf_oneof:"message"`
+}
+
+func (x *EventsLogs) Reset() {
+	*x = EventsLogs{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_agent_events_events_proto_msgTypes[2]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+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 protoimpl.UnsafeEnabled && 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 (m *EventsLogs) GetMessage() isEventsLogs_Message {
+	if m != nil {
+		return m.Message
+	}
+	return nil
+}
+
+func (x *EventsLogs) GetAgentLog() *AgentLog {
+	if x, ok := x.GetMessage().(*EventsLogs_AgentLog); ok {
+		return x.AgentLog
+	}
+	return nil
+}
+
+func (x *EventsLogs) GetAgentEvent() *AgentEvent {
+	if x, ok := x.GetMessage().(*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
+
+var file_agent_events_events_proto_rawDesc = []byte{
+	0x0a, 0x19, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x2f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x2f, 0x65,
+	0x76, 0x65, 0x6e, 0x74, 0x73, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, 0x06, 0x65, 0x76, 0x65,
+	0x6e, 0x74, 0x73, 0x1a, 0x1f, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2f, 0x70, 0x72, 0x6f, 0x74,
+	0x6f, 0x62, 0x75, 0x66, 0x2f, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x2e, 0x70,
+	0x72, 0x6f, 0x74, 0x6f, 0x22, 0xde, 0x01, 0x0a, 0x0a, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76,
+	0x65, 0x6e, 0x74, 0x12, 0x1d, 0x0a, 0x0a, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x5f, 0x74, 0x79, 0x70,
+	0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x54, 0x79,
+	0x70, 0x65, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18,
+	0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70,
+	0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d,
+	0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x12, 0x25, 0x0a, 0x0e,
+	0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x03,
+	0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f,
+	0x6e, 0x49, 0x64, 0x12, 0x18, 0x0a, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, 0x73, 0x18, 0x04,
+	0x20, 0x01, 0x28, 0x0c, 0x52, 0x07, 0x64, 0x65, 0x74, 0x61, 0x69, 0x6c, 0x73, 0x12, 0x1e, 0x0a,
+	0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x18, 0x05, 0x20, 0x01, 0x28,
+	0x09, 0x52, 0x0a, 0x6f, 0x72, 0x69, 0x67, 0x69, 0x6e, 0x61, 0x74, 0x6f, 0x72, 0x12, 0x16, 0x0a,
+	0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x06, 0x73,
+	0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0x9b, 0x01, 0x0a, 0x08, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x4c,
+	0x6f, 0x67, 0x12, 0x18, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x18, 0x01, 0x20,
+	0x01, 0x28, 0x09, 0x52, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x25, 0x0a, 0x0e,
+	0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x02,
+	0x20, 0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f,
+	0x6e, 0x49, 0x64, 0x12, 0x14, 0x0a, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x18, 0x03, 0x20, 0x01,
+	0x28, 0x09, 0x52, 0x05, 0x6c, 0x65, 0x76, 0x65, 0x6c, 0x12, 0x38, 0x0a, 0x09, 0x74, 0x69, 0x6d,
+	0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x18, 0x04, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x67,
+	0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x54,
+	0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x52, 0x09, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74,
+	0x61, 0x6d, 0x70, 0x22, 0x7f, 0x0a, 0x0a, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x4c, 0x6f, 0x67,
+	0x73, 0x12, 0x2f, 0x0a, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x6f, 0x67, 0x18, 0x01,
+	0x20, 0x01, 0x28, 0x0b, 0x32, 0x10, 0x2e, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73, 0x2e, 0x41, 0x67,
+	0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x48, 0x00, 0x52, 0x08, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x4c,
+	0x6f, 0x67, 0x12, 0x35, 0x0a, 0x0b, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x65, 0x76, 0x65, 0x6e,
+	0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x12, 0x2e, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73,
+	0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x48, 0x00, 0x52, 0x0a, 0x61,
+	0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x42, 0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73,
+	0x73, 0x61, 0x67, 0x65, 0x42, 0x0a, 0x5a, 0x08, 0x2e, 0x2f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x73,
+	0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
+}
+
+var (
+	file_agent_events_events_proto_rawDescOnce sync.Once
+	file_agent_events_events_proto_rawDescData = file_agent_events_events_proto_rawDesc
+)
+
+func file_agent_events_events_proto_rawDescGZIP() []byte {
+	file_agent_events_events_proto_rawDescOnce.Do(func() {
+		file_agent_events_events_proto_rawDescData = protoimpl.X.CompressGZIP(file_agent_events_events_proto_rawDescData)
+	})
+	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
+	}
+	if !protoimpl.UnsafeEnabled {
+		file_agent_events_events_proto_msgTypes[0].Exporter = func(v any, i int) any {
+			switch v := v.(*AgentEvent); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+		file_agent_events_events_proto_msgTypes[1].Exporter = func(v any, i int) any {
+			switch v := v.(*AgentLog); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+		file_agent_events_events_proto_msgTypes[2].Exporter = func(v any, i int) any {
+			switch v := v.(*EventsLogs); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+	}
+	file_agent_events_events_proto_msgTypes[2].OneofWrappers = []any{
+		(*EventsLogs_AgentLog)(nil),
+		(*EventsLogs_AgentEvent)(nil),
+	}
+	type x struct{}
+	out := protoimpl.TypeBuilder{
+		File: protoimpl.DescBuilder{
+			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
+			RawDescriptor: 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_rawDesc = nil
+	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,64 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package events
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/protobuf/proto"
+)
+
+type mockConn struct {
+	writeErr error
+	buf      bytes.Buffer
+}
+
+func (m *mockConn) Write(p []byte) (n int, err error) {
+	if m.writeErr != nil {
+		return 0, m.writeErr
+	}
+	return m.buf.Write(p)
+}
+
+func TestSendEventSuccess(t *testing.T) {
+	mockConnection := &mockConn{}
+
+	svc, err := New("test_service", "12345", mockConnection)
+	assert.NoError(t, err)
+
+	details := json.RawMessage(`{"key": "value"}`)
+
+	err = svc.SendEvent("test_event", "success", details)
+	assert.NoError(t, err)
+
+	var writtenMessage EventsLogs
+	err = proto.Unmarshal(mockConnection.buf.Bytes(), &writtenMessage)
+	assert.NoError(t, err)
+
+	assert.Equal(t, "test_event", writtenMessage.GetAgentEvent().EventType)
+	assert.Equal(t, "12345", writtenMessage.GetAgentEvent().ComputationId)
+	assert.Equal(t, "test_service", writtenMessage.GetAgentEvent().Originator)
+	assert.Equal(t, "success", writtenMessage.GetAgentEvent().Status)
+
+	now := time.Now()
+	eventTimestamp := writtenMessage.GetAgentEvent().GetTimestamp().AsTime()
+	assert.WithinDuration(t, now, eventTimestamp, 1*time.Second)
+}
+
+func TestSendEventFailure(t *testing.T) {
+	mockConnection := &mockConn{writeErr: errors.New("write error")}
+
+	svc, err := New("test_service", "12345", mockConnection)
+	assert.NoError(t, err)
+
+	details := json.RawMessage(`{"key": "value"}`)
+
+	err = svc.SendEvent("test_event", "failure", details)
+	assert.Error(t, err)
+	assert.Equal(t, "write error", err.Error())
+}

agent/events/mocks/events.go

@@ -1,8 +1,8 @@
-// Code generated by mockery v2.43.2. DO NOT EDIT.
-
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
 package mocks
 
 import (
@@ -16,22 +16,12 @@ type Service struct {
 	mock.Mock
 }
 
-// Close provides a mock function with given fields:
-func (_m *Service) Close() error {
-	ret := _m.Called()
+type Service_Expecter struct {
+	mock *mock.Mock
+}
 
-	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
+func (_m *Service) EXPECT() *Service_Expecter {
+	return &Service_Expecter{mock: &_m.Mock}
 }
 
 // SendEvent provides a mock function with given fields: event, status, details
@@ -52,6 +42,36 @@ func (_m *Service) SendEvent(event string, status string, details json.RawMessag
 	return r0
 }
 
+// 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
+//   - event string
+//   - status string
+//   - details json.RawMessage
+func (_e *Service_Expecter) SendEvent(event interface{}, status interface{}, details interface{}) *Service_SendEvent_Call {
+	return &Service_SendEvent_Call{Call: _e.mock.On("SendEvent", event, status, details)}
+}
+
+func (_c *Service_SendEvent_Call) Run(run func(event string, status string, details json.RawMessage)) *Service_SendEvent_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(string), args[1].(string), args[2].(json.RawMessage))
+	})
+	return _c
+}
+
+func (_c *Service_SendEvent_Call) Return(_a0 error) *Service_SendEvent_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Service_SendEvent_Call) RunAndReturn(run func(string, string, json.RawMessage) error) *Service_SendEvent_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
 // The first argument is typically a *testing.T value.
 func NewService(t interface {

agent/mocks/agent.go

@@ -1,8 +1,8 @@
-// Code generated by mockery v2.43.2. DO NOT EDIT.
-
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
 package mocks
 
 import (
@@ -18,6 +18,14 @@ 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 with given fields: ctx, algorithm
 func (_m *Service) Algo(ctx context.Context, algorithm agent.Algorithm) error {
 	ret := _m.Called(ctx, algorithm)
@@ -36,6 +44,35 @@ func (_m *Service) Algo(ctx context.Context, algorithm agent.Algorithm) error {
 	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) {
+		run(args[0].(context.Context), args[1].(agent.Algorithm))
+	})
+	return _c
+}
+
+func (_c *Service_Algo_Call) Return(_a0 error) *Service_Algo_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Service_Algo_Call) RunAndReturn(run func(context.Context, agent.Algorithm) error) *Service_Algo_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Attestation provides a mock function with given fields: ctx, reportData
 func (_m *Service) Attestation(ctx context.Context, reportData [64]byte) ([]byte, error) {
 	ret := _m.Called(ctx, reportData)
@@ -66,6 +103,35 @@ func (_m *Service) Attestation(ctx context.Context, reportData [64]byte) ([]byte
 	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
+func (_e *Service_Expecter) Attestation(ctx interface{}, reportData interface{}) *Service_Attestation_Call {
+	return &Service_Attestation_Call{Call: _e.mock.On("Attestation", ctx, reportData)}
+}
+
+func (_c *Service_Attestation_Call) Run(run func(ctx context.Context, reportData [64]byte)) *Service_Attestation_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].([64]byte))
+	})
+	return _c
+}
+
+func (_c *Service_Attestation_Call) Return(_a0 []byte, _a1 error) *Service_Attestation_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *Service_Attestation_Call) RunAndReturn(run func(context.Context, [64]byte) ([]byte, error)) *Service_Attestation_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Data provides a mock function with given fields: ctx, dataset
 func (_m *Service) Data(ctx context.Context, dataset agent.Dataset) error {
 	ret := _m.Called(ctx, dataset)
@@ -84,6 +150,35 @@ func (_m *Service) Data(ctx context.Context, dataset agent.Dataset) error {
 	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) {
+		run(args[0].(context.Context), args[1].(agent.Dataset))
+	})
+	return _c
+}
+
+func (_c *Service_Data_Call) Return(_a0 error) *Service_Data_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Service_Data_Call) RunAndReturn(run func(context.Context, agent.Dataset) error) *Service_Data_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Result provides a mock function with given fields: ctx
 func (_m *Service) Result(ctx context.Context) ([]byte, error) {
 	ret := _m.Called(ctx)
@@ -114,6 +209,34 @@ func (_m *Service) Result(ctx context.Context) ([]byte, error) {
 	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) {
+		run(args[0].(context.Context))
+	})
+	return _c
+}
+
+func (_c *Service_Result_Call) Return(_a0 []byte, _a1 error) *Service_Result_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *Service_Result_Call) RunAndReturn(run func(context.Context) ([]byte, error)) *Service_Result_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
 // The first argument is typically a *testing.T value.
 func NewService(t interface {

agent/mocks/agent_grpc_algo.go

@@ -0,0 +1,207 @@
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package mocks
+
+import (
+	context "context"
+
+	agent "github.com/ultravioletrs/cocos/agent"
+
+	metadata "google.golang.org/grpc/metadata"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// AgentService_AlgoClient is an autogenerated mock type for the AgentService_AlgoClient type
+type AgentService_AlgoClient struct {
+	mock.Mock
+}
+
+// CloseAndRecv provides a mock function with given fields:
+func (_m *AgentService_AlgoClient) CloseAndRecv() (*agent.AlgoResponse, error) {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseAndRecv")
+	}
+
+	var r0 *agent.AlgoResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func() (*agent.AlgoResponse, error)); ok {
+		return rf()
+	}
+	if rf, ok := ret.Get(0).(func() *agent.AlgoResponse); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*agent.AlgoResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// CloseSend provides a mock function with given fields:
+func (_m *AgentService_AlgoClient) CloseSend() error {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseSend")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Context provides a mock function with given fields:
+func (_m *AgentService_AlgoClient) Context() context.Context {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Context")
+	}
+
+	var r0 context.Context
+	if rf, ok := ret.Get(0).(func() context.Context); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+
+	return r0
+}
+
+// Header provides a mock function with given fields:
+func (_m *AgentService_AlgoClient) Header() (metadata.MD, error) {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Header")
+	}
+
+	var r0 metadata.MD
+	var r1 error
+	if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
+		return rf()
+	}
+	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// RecvMsg provides a mock function with given fields: m
+func (_m *AgentService_AlgoClient) RecvMsg(m interface{}) error {
+	ret := _m.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RecvMsg")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
+		r0 = rf(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Send provides a mock function with given fields: _a0
+func (_m *AgentService_AlgoClient) Send(_a0 *agent.AlgoRequest) error {
+	ret := _m.Called(_a0)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Send")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(*agent.AlgoRequest) error); ok {
+		r0 = rf(_a0)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// SendMsg provides a mock function with given fields: m
+func (_m *AgentService_AlgoClient) SendMsg(m interface{}) error {
+	ret := _m.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SendMsg")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
+		r0 = rf(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Trailer provides a mock function with given fields:
+func (_m *AgentService_AlgoClient) Trailer() metadata.MD {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Trailer")
+	}
+
+	var r0 metadata.MD
+	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+
+	return r0
+}
+
+// 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
+}

agent/mocks/agent_grpc_data.go

@@ -0,0 +1,207 @@
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package mocks
+
+import (
+	context "context"
+
+	agent "github.com/ultravioletrs/cocos/agent"
+
+	metadata "google.golang.org/grpc/metadata"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// AgentService_DataClient is an autogenerated mock type for the AgentService_DataClient type
+type AgentService_DataClient struct {
+	mock.Mock
+}
+
+// CloseAndRecv provides a mock function with given fields:
+func (_m *AgentService_DataClient) CloseAndRecv() (*agent.DataResponse, error) {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseAndRecv")
+	}
+
+	var r0 *agent.DataResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func() (*agent.DataResponse, error)); ok {
+		return rf()
+	}
+	if rf, ok := ret.Get(0).(func() *agent.DataResponse); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*agent.DataResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// CloseSend provides a mock function with given fields:
+func (_m *AgentService_DataClient) CloseSend() error {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for CloseSend")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Context provides a mock function with given fields:
+func (_m *AgentService_DataClient) Context() context.Context {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Context")
+	}
+
+	var r0 context.Context
+	if rf, ok := ret.Get(0).(func() context.Context); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+
+	return r0
+}
+
+// Header provides a mock function with given fields:
+func (_m *AgentService_DataClient) Header() (metadata.MD, error) {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Header")
+	}
+
+	var r0 metadata.MD
+	var r1 error
+	if rf, ok := ret.Get(0).(func() (metadata.MD, error)); ok {
+		return rf()
+	}
+	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func() error); ok {
+		r1 = rf()
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// RecvMsg provides a mock function with given fields: m
+func (_m *AgentService_DataClient) RecvMsg(m interface{}) error {
+	ret := _m.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for RecvMsg")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
+		r0 = rf(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Send provides a mock function with given fields: _a0
+func (_m *AgentService_DataClient) Send(_a0 *agent.DataRequest) error {
+	ret := _m.Called(_a0)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Send")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(*agent.DataRequest) error); ok {
+		r0 = rf(_a0)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// SendMsg provides a mock function with given fields: m
+func (_m *AgentService_DataClient) SendMsg(m interface{}) error {
+	ret := _m.Called(m)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SendMsg")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(interface{}) error); ok {
+		r0 = rf(m)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Trailer provides a mock function with given fields:
+func (_m *AgentService_DataClient) Trailer() metadata.MD {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Trailer")
+	}
+
+	var r0 metadata.MD
+	if rf, ok := ret.Get(0).(func() metadata.MD); ok {
+		r0 = rf()
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(metadata.MD)
+		}
+	}
+
+	return r0
+}
+
+// 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
+}

agent/quoteprovider/sev.go

@@ -1,13 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-//go:build !embed
-// +build !embed
-
-package quoteprovider
-
-import "github.com/google/go-sev-guest/client"
-
-func GetQuoteProvider() (client.QuoteProvider, error) {
-	return client.GetQuoteProvider()
-}

agent/service.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"path/filepath"
 	"slices"
+	sync "sync"
 
 	"github.com/google/go-sev-guest/client"
 	"github.com/ultravioletrs/cocos/agent/algorithm"
@@ -20,12 +21,53 @@ import (
 	"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"
 	"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
@@ -55,12 +97,12 @@ var (
 	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")
 )
 
 // 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 {
 	Algo(ctx context.Context, algorithm Algorithm) error
 	Data(ctx context.Context, dataset Dataset) error
@@ -69,46 +111,79 @@ type Service interface {
 }
 
 type agentService struct {
-	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        // Manages the state transitions of the agent service.
-	runError      error                // Stores any error encountered during the computation run.
-	eventSvc      events.Service       // Service for publishing events related to computation.
-	quoteProvider client.QuoteProvider // Provider for generating attestation quotes.
+	mu              sync.Mutex
+	computation     Computation               // Holds the current computation request details.
+	algorithm       algorithm.Algorithm       // Filepath to the algorithm received for the computation.
+	result          []byte                    // Stores the result of the computation.
+	sm              statemachine.StateMachine // Manages the state transitions of the agent service.
+	runError        error                     // Stores any error encountered during the computation run.
+	eventSvc        events.Service            // Service for publishing events related to computation.
+	quoteProvider   client.QuoteProvider      // Provider for generating attestation quotes.
+	logger          *slog.Logger              // Logger for the agent service.
+	resultsConsumed bool                      // Indicates if the results have been consumed.
 }
 
 var _ Service = (*agentService)(nil)
 
 // New instantiates the agent service implementation.
 func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp Computation, quoteProvider client.QuoteProvider) Service {
+	sm := statemachine.NewStateMachine(Idle)
 	svc := &agentService{
-		sm:            NewStateMachine(logger, cmp),
+		sm:            sm,
 		eventSvc:      eventSvc,
 		quoteProvider: quoteProvider,
+		logger:        logger,
+		computation:   cmp,
 	}
 
-	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
-	svc.sm.StateFunctions[failed] = svc.publishEvent("failed", json.RawMessage{})
+	transitions := []statemachine.Transition{
+		{From: Idle, Event: Start, To: ReceivingManifest},
+		{From: ReceivingManifest, Event: ManifestReceived, To: ReceivingAlgorithm},
+	}
 
-	svc.computation = cmp
+	if len(cmp.Datasets) == 0 {
+		transitions = append(transitions, statemachine.Transition{From: ReceivingAlgorithm, Event: AlgorithmReceived, To: Running})
+	} else {
+		transitions = append(transitions, statemachine.Transition{From: ReceivingAlgorithm, Event: AlgorithmReceived, To: ReceivingData})
+		transitions = append(transitions, statemachine.Transition{From: ReceivingData, Event: DataReceived, To: Running})
+	}
+
+	transitions = append(transitions, []statemachine.Transition{
+		{From: Running, Event: RunComplete, To: ConsumingResults},
+		{From: Running, Event: RunFailed, To: Failed},
+		{From: ConsumingResults, Event: ResultsConsumed, To: Complete},
+	}...)
+
+	for _, t := range transitions {
+		sm.AddTransition(t)
+	}
+
+	sm.SetAction(Idle, svc.publishEvent(IdleState.String()))
+	sm.SetAction(ReceivingManifest, svc.publishEvent(InProgress.String()))
+	sm.SetAction(ReceivingAlgorithm, svc.publishEvent(InProgress.String()))
+	sm.SetAction(ReceivingData, svc.publishEvent(InProgress.String()))
+	sm.SetAction(Running, svc.runComputation)
+	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())
+		}
+	}()
+	sm.SendEvent(Start)
+	defer sm.SendEvent(ManifestReceived)
 
-	svc.sm.SendEvent(manifestReceived)
 	return svc
 }
 
 func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
-	if as.sm.GetState() != receivingAlgorithm {
+	if as.sm.GetState() != ReceivingAlgorithm {
 		return ErrStateNotReady
 	}
+	as.mu.Lock()
+	defer as.mu.Unlock()
 	if as.algorithm != nil {
 		return ErrAllManifestItemsReceived
 	}
@@ -124,7 +199,7 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 		return fmt.Errorf("error getting current directory: %v", err)
 	}
 
-	f, err := os.Create(filepath.Join(currentDir, "algorithm"))
+	f, err := os.Create(filepath.Join(currentDir, "algo"))
 	if err != nil {
 		return fmt.Errorf("error creating algorithm file: %v", err)
 	}
@@ -150,7 +225,7 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 
 	switch algoType {
 	case string(algorithm.AlgoTypeBin):
-		as.algorithm = binary.NewAlgorithm(as.sm.logger, as.eventSvc, f.Name(), args)
+		as.algorithm = binary.NewAlgorithm(as.logger, as.eventSvc, f.Name(), args)
 	case string(algorithm.AlgoTypePython):
 		var requirementsFile string
 		if len(algo.Requirements) > 0 {
@@ -168,11 +243,11 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 			requirementsFile = fr.Name()
 		}
 		runtime := python.PythonRunTimeFromContext(ctx)
-		as.algorithm = python.NewAlgorithm(as.sm.logger, as.eventSvc, runtime, requirementsFile, f.Name(), args)
+		as.algorithm = python.NewAlgorithm(as.logger, as.eventSvc, runtime, requirementsFile, f.Name(), args)
 	case string(algorithm.AlgoTypeWasm):
-		as.algorithm = wasm.NewAlgorithm(as.sm.logger, as.eventSvc, f.Name(), args)
+		as.algorithm = wasm.NewAlgorithm(as.logger, as.eventSvc, f.Name(), args)
 	case string(algorithm.AlgoTypeDocker):
-		as.algorithm = docker.NewAlgorithm(as.sm.logger, as.eventSvc, f.Name())
+		as.algorithm = docker.NewAlgorithm(as.logger, as.eventSvc, f.Name())
 	}
 
 	if err := os.Mkdir(algorithm.DatasetsDir, 0o755); err != nil {
@@ -180,16 +255,18 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 	}
 
 	if as.algorithm != nil {
-		as.sm.SendEvent(algorithmReceived)
+		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
 	}
@@ -233,27 +310,32 @@ func (as *agentService) Data(ctx context.Context, dataset Dataset) error {
 	}
 
 	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 && as.sm.GetState() != failed {
+	currentState := as.sm.GetState()
+	if currentState != ConsumingResults && currentState != Complete && currentState != Failed {
 		return []byte{}, ErrResultsNotReady
 	}
-	if len(as.computation.ResultConsumers) == 0 {
-		return []byte{}, ErrAllResultsConsumed
-	}
+
 	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.GetState() == resultsReady {
-		as.sm.SendEvent(resultsConsumed)
+	as.mu.Lock()
+	defer as.mu.Unlock()
+	if index < 0 || index >= len(as.computation.ResultConsumers) {
+		return []byte{}, ErrUndeclaredConsumer
+	}
+
+	if !as.resultsConsumed && currentState == ConsumingResults {
+		as.resultsConsumed = true
+		defer as.sm.SendEvent(ResultsConsumed)
 	}
 
 	return as.result, as.runError
@@ -268,58 +350,58 @@ func (as *agentService) Attestation(ctx context.Context, reportData [ReportDataS
 	return rawQuote, nil
 }
 
-func (as *agentService) runComputation() {
-	as.publishEvent("starting", json.RawMessage{})()
-	as.sm.logger.Debug("computation run started")
+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)
+			as.sm.SendEvent(RunFailed)
 		} else {
-			as.sm.SendEvent(runComplete)
+			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.sm.logger.Warn(as.runError.Error())
-		as.publishEvent("failed", json.RawMessage{})()
+		as.logger.Warn(as.runError.Error())
+		as.publishEvent(Failed.String())(state)
 		return
 	}
 
 	defer func() {
 		if err := os.RemoveAll(algorithm.ResultsDir); err != nil {
-			as.sm.logger.Warn(fmt.Sprintf("error removing results directory and its contents: %s", err.Error()))
+			as.logger.Warn(fmt.Sprintf("error removing results directory and its contents: %s", err.Error()))
 		}
 		if err := os.RemoveAll(algorithm.DatasetsDir); err != nil {
-			as.sm.logger.Warn(fmt.Sprintf("error removing datasets directory and its contents: %s", err.Error()))
+			as.logger.Warn(fmt.Sprintf("error removing datasets directory and its contents: %s", err.Error()))
 		}
 	}()
 
-	as.publishEvent("in-progress", json.RawMessage{})()
+	as.publishEvent(InProgress.String())(state)
 	if err := as.algorithm.Run(); err != nil {
 		as.runError = err
-		as.sm.logger.Warn(fmt.Sprintf("failed to run computation: %s", err.Error()))
-		as.publishEvent("failed", json.RawMessage{})()
+		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.sm.logger.Warn(fmt.Sprintf("failed to zip results: %s", err.Error()))
-		as.publishEvent("failed", json.RawMessage{})()
+		as.logger.Warn(fmt.Sprintf("failed to zip results: %s", err.Error()))
+		as.publishEvent(Failed.String())(state)
 		return
 	}
 
-	as.publishEvent("complete", json.RawMessage{})()
+	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) {
+		if err := as.eventSvc.SendEvent(state.String(), status, json.RawMessage{}); err != nil {
+			as.logger.Warn(err.Error())
 		}
 	}
 }

agent/service_test.go

@@ -0,0 +1,406 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package agent
+
+import (
+	"context"
+	"crypto/rand"
+	"log"
+	"os"
+	"testing"
+	"time"
+
+	mglog "github.com/absmach/magistrala/logger"
+	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/agent/algorithm"
+	"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/quoteprovider"
+	mocks2 "github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider/mocks"
+	"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) {
+	events := new(mocks.Service)
+
+	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	defer evCall.Unset()
+
+	qp, err := quoteprovider.GetQuoteProvider()
+	require.NoError(t, err)
+
+	algo, err := os.ReadFile(algoPath)
+	require.NoError(t, err)
+
+	algoHash := sha3.Sum256(algo)
+
+	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),
+			)
+
+			ctx, cancel := context.WithCancel(ctx)
+			defer cancel()
+			svc := New(ctx, mglog.NewMock(), events, testComputation(t), qp)
+
+			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) {
+	events := new(mocks.Service)
+
+	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	defer evCall.Unset()
+
+	qp, err := quoteprovider.GetQuoteProvider()
+	require.NoError(t, err)
+
+	algo, err := os.ReadFile(algoPath)
+	require.NoError(t, err)
+
+	algoHash := sha3.Sum256(algo)
+
+	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),
+			)
+
+			if tc.err != ErrUndeclaredDataset {
+				ctx = IndexToContext(ctx, 0)
+			}
+
+			ctx, cancel := context.WithCancel(ctx)
+			defer cancel()
+
+			comp := testComputation(t)
+
+			svc := New(ctx, mglog.NewMock(), events, comp, qp)
+			time.Sleep(300 * time.Millisecond)
+
+			if tc.err != ErrStateNotReady {
+				_ = svc.Algo(ctx, alg)
+				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) {
+	events := new(mocks.Service)
+
+	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	defer evCall.Unset()
+
+	qp, err := quoteprovider.GetQuoteProvider()
+	require.NoError(t, err)
+
+	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 {
+		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,
+				quoteProvider: qp,
+				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) {
+	events := new(mocks.Service)
+	qp := new(mocks2.QuoteProvider)
+
+	evCall := events.On("SendEvent", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	defer evCall.Unset()
+
+	cases := []struct {
+		name       string
+		reportData [ReportDataSize]byte
+		rawQuote   []uint8
+		err        error
+	}{
+		{
+			name:       "Test attestation successful",
+			reportData: generateReportData(),
+			rawQuote:   make([]uint8, 0),
+			err:        nil,
+		},
+		{
+			name:       "Test attestation failed",
+			reportData: generateReportData(),
+			rawQuote:   nil,
+			err:        ErrAttestationFailed,
+		},
+	}
+	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),
+			)
+			ctx, cancel := context.WithCancel(ctx)
+			defer cancel()
+
+			getQuote := qp.On("GetRawQuote", mock.Anything).Return(tc.rawQuote, tc.err)
+			if tc.err != ErrAttestationFailed {
+				getQuote = qp.On("GetRawQuote", mock.Anything).Return(tc.reportData, nil)
+			}
+			defer getQuote.Unset()
+
+			svc := New(ctx, mglog.NewMock(), events, testComputation(t), qp)
+			time.Sleep(300 * time.Millisecond)
+			_, err := svc.Attestation(ctx, tc.reportData)
+			assert.True(t, errors.Contains(err, tc.err), "expected %v, got %v", tc.err, err)
+		})
+	}
+}
+
+func generateReportData() [ReportDataSize]byte {
+	bytes := make([]byte, ReportDataSize)
+	_, 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")}},
+		AgentConfig: AgentConfig{
+			Port:        "7002",
+			LogLevel:    "debug",
+			AttestedTls: false,
+		},
+	}
+}

agent/state.go

@@ -1,131 +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 uint8
-
-const (
-	idle state = iota
-	receivingManifest
-	receivingAlgorithm
-	receivingData
-	running
-	resultsReady
-	complete
-	failed
-)
-
-type event uint8
-
-const (
-	start event = iota
-	manifestReceived
-	algorithmReceived
-	dataReceived
-	runComplete
-	resultsConsumed
-	runFailed
-)
-
-// 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, cmp Computation) *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)
-	switch len(cmp.Datasets) {
-	case 0:
-		sm.Transitions[receivingAlgorithm][algorithmReceived] = running
-	default:
-		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[running][runFailed] = failed
-
-	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:
-			sm.mu.Lock()
-			nextState, valid := sm.Transitions[sm.State][event]
-			if valid {
-				sm.State = nextState
-				sm.logger.Debug(fmt.Sprintf("Transition: %v -> %v\n", sm.State, nextState))
-			} else {
-				sm.logger.Error(fmt.Sprintf("Invalid transition: %v -> ???\n", sm.State))
-			}
-			sm.mu.Unlock()
-
-			sm.mu.Lock()
-			stateFunc, exists := sm.StateFunctions[sm.State]
-			sm.mu.Unlock()
-			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()
-	defer sm.mu.Unlock()
-	return sm.State
-}
-
-func (sm *StateMachine) SetState(state state) {
-	sm.mu.Lock()
-	defer sm.mu.Unlock()
-	sm.State = state
-}

agent/state_string.go

@@ -1,30 +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]
-	_ = x[failed-7]
-}
-
-const _state_name = "idlereceivingManifestreceivingAlgorithmreceivingDatarunningresultsReadycompletefailed"
-
-var _state_index = [...]uint8{0, 4, 21, 39, 52, 59, 71, 79, 85}
-
-func (i state) String() string {
-	if 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,69 +4,249 @@ package agent
 
 import (
 	"context"
-	"fmt"
+	sync "sync"
 	"testing"
+	"time"
 
-	mglog "github.com/absmach/magistrala/logger"
+	"github.com/ultravioletrs/cocos/agent/statemachine"
 )
 
-var cmp = Computation{
-	Datasets: []Dataset{
-		{
-			Dataset: []byte("test"),
-			UserKey: []byte("test"),
-		},
-	},
+type MockState int
+
+type MockEvent int
+
+func (s MockState) String() string {
+	return []string{"State1", "State2", "State3"}[s]
 }
 
-func TestStateMachineTransitions(t *testing.T) {
-	cases := []struct {
-		fromState state
-		event     event
-		expected  state
-		cmp       Computation
+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 {
+			t.Errorf("Start returned error: %v", err)
+		}
+	}()
+
+	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(), 100*time.Millisecond)
+	defer cancel()
+
+	go func() {
+		if err := sm.Start(ctx); err != context.Canceled {
+			t.Errorf("Start returned error: %v", err)
+		}
+	}()
+
+	sm.SendEvent(Event1)
+
+	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 {
+			t.Errorf("Start returned error: %v", err)
+		}
+	}()
+
+	transitions := []struct {
+		event MockEvent
+		want  MockState
 	}{
-		{idle, start, receivingManifest, cmp},
-		{receivingManifest, manifestReceived, receivingAlgorithm, cmp},
-		{receivingAlgorithm, algorithmReceived, receivingData, cmp},
-		{receivingAlgorithm, algorithmReceived, running, Computation{}},
-		{receivingData, dataReceived, running, cmp},
-		{running, runComplete, resultsReady, cmp},
-		{resultsReady, resultsConsumed, complete, cmp},
+		{Event1, State2},
+		{Event2, State3},
+		{Event3, State1},
 	}
 
-	for _, tc := range cases {
-		t.Run(fmt.Sprintf("Transition from %v to %v", tc.fromState, tc.expected), func(t *testing.T) {
-			sm := NewStateMachine(mglog.NewMock(), tc.cmp)
-			ctx, cancel := context.WithCancel(context.Background())
-			go func() {
-				sm.Start(ctx)
-			}()
-			sm.wg.Wait()
-			sm.SetState(tc.fromState)
+	for _, tt := range transitions {
+		sm.SendEvent(tt.event)
+		time.Sleep(50 * time.Millisecond)
 
-			sm.SendEvent(tc.event)
-
-			if sm.GetState() != tc.expected {
-				t.Errorf("Expected state %v after the event, but got %v", tc.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(), cmp)
-	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/state.go

@@ -0,0 +1,233 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+package mocks
+
+import (
+	context "context"
+
+	mock "github.com/stretchr/testify/mock"
+	statemachine "github.com/ultravioletrs/cocos/agent/statemachine"
+)
+
+// 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 with given fields: t
+func (_m *StateMachine) AddTransition(t statemachine.Transition) {
+	_m.Called(t)
+}
+
+// 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) {
+		run(args[0].(statemachine.Transition))
+	})
+	return _c
+}
+
+func (_c *StateMachine_AddTransition_Call) Return() *StateMachine_AddTransition_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *StateMachine_AddTransition_Call) RunAndReturn(run func(statemachine.Transition)) *StateMachine_AddTransition_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// GetState provides a mock function with given fields:
+func (_m *StateMachine) GetState() statemachine.State {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for GetState")
+	}
+
+	var r0 statemachine.State
+	if rf, ok := ret.Get(0).(func() statemachine.State); ok {
+		r0 = rf()
+	} 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(_a0 statemachine.State) *StateMachine_GetState_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *StateMachine_GetState_Call) RunAndReturn(run func() statemachine.State) *StateMachine_GetState_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SendEvent provides a mock function with given fields: event
+func (_m *StateMachine) SendEvent(event statemachine.Event) {
+	_m.Called(event)
+}
+
+// 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) {
+		run(args[0].(statemachine.Event))
+	})
+	return _c
+}
+
+func (_c *StateMachine_SendEvent_Call) Return() *StateMachine_SendEvent_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *StateMachine_SendEvent_Call) RunAndReturn(run func(statemachine.Event)) *StateMachine_SendEvent_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// SetAction provides a mock function with given fields: state, action
+func (_m *StateMachine) SetAction(state statemachine.State, action statemachine.Action) {
+	_m.Called(state, action)
+}
+
+// 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) {
+		run(args[0].(statemachine.State), args[1].(statemachine.Action))
+	})
+	return _c
+}
+
+func (_c *StateMachine_SetAction_Call) Return() *StateMachine_SetAction_Call {
+	_c.Call.Return()
+	return _c
+}
+
+func (_c *StateMachine_SetAction_Call) RunAndReturn(run func(statemachine.State, statemachine.Action)) *StateMachine_SetAction_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Start provides a mock function with given fields: ctx
+func (_m *StateMachine) Start(ctx context.Context) error {
+	ret := _m.Called(ctx)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Start")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context) error); ok {
+		r0 = rf(ctx)
+	} 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) {
+		run(args[0].(context.Context))
+	})
+	return _c
+}
+
+func (_c *StateMachine_Start_Call) Return(_a0 error) *StateMachine_Start_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *StateMachine_Start_Call) RunAndReturn(run func(context.Context) error) *StateMachine_Start_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewStateMachine creates a new instance of StateMachine. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewStateMachine(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *StateMachine {
+	mock := &StateMachine{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

agent/statemachine/state.go

@@ -0,0 +1,112 @@
+// 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
+}
+
+type stateMachine struct {
+	mu           sync.Mutex
+	currentState State
+	transitions  map[State]map[Event]State
+	actions      map[State]Action
+	eventChan    chan Event
+}
+
+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),
+	}
+}
+
+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.eventChan <- event
+}
+
+func (sm *stateMachine) Start(ctx context.Context) error {
+	for {
+		select {
+		case event := <-sm.eventChan:
+			if err := sm.handleEvent(event); err != nil {
+				return err
+			}
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+}
+
+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/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]]
+}

boilerplate.txt

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

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

@@ -5,11 +5,10 @@ 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"
@@ -29,44 +28,55 @@ func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
 		Example: "algo <algo_file> <private_key_file_path>",
 		Args:    cobra.ExactArgs(2),
 		Run: func(cmd *cobra.Command, args []string) {
+			if cli.connectErr != nil {
+				printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
+				return
+			}
+
 			algorithmFile := args[0]
 
-			log.Println("Uploading algorithm file:", algorithmFile)
+			cmd.Println("Uploading algorithm file:", algorithmFile)
 
-			algorithm, err := os.ReadFile(algorithmFile)
+			algorithm, err := os.Open(algorithmFile)
 			if err != nil {
-				log.Fatalf("Error reading algorithm file: %v", err)
+				printError(cmd, "Error reading algorithm file: %v ❌ ", err)
+				return
 			}
 
-			var req []byte
+			defer algorithm.Close()
+
+			var req *os.File
 			if requirementsFile != "" {
-				req, err = os.ReadFile(requirementsFile)
+				req, err = os.Open(requirementsFile)
 				if err != nil {
-					log.Fatalf("Error reading requirments file: %v", err)
+					printError(cmd, "Error reading requirments file: %v ❌ ", err)
+					return
 				}
-			}
-
-			algoReq := agent.Algorithm{
-				Algorithm:    algorithm,
-				Requirements: req,
+				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)
+			privKey, err := decodeKey(pemBlock)
+			if err != nil {
+				printError(cmd, "Error decoding private key: %v ❌ ", err)
+				return
+			}
 
 			ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
 
-			if err := cli.agentSDK.Algo(addAlgoMetadata(ctx), algoReq, privKey); err != nil {
-				log.Fatalf("Error uploading algorithm with error: %v", err)
+			if err := cli.agentSDK.Algo(addAlgoMetadata(ctx), algorithm, req, privKey); err != nil {
+				printError(cmd, "Failed to upload algorithm due to error: %v ❌ ", err)
+				return
 			}
 
-			log.Println("Successfully uploaded algorithm")
+			cmd.Println(color.New(color.FgGreen).Sprint("Successfully uploaded algorithm! ✔ "))
 		},
 	}
 

cli/attestation.go

@@ -4,22 +4,23 @@ package cli
 
 import (
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
-	"log"
 	"os"
 	"strconv"
 	"strings"
 	"time"
 
+	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/fatih/color"
 	"github.com/google/go-sev-guest/abi"
 	"github.com/google/go-sev-guest/proto/check"
 	"github.com/google/go-sev-guest/proto/sevsnp"
-	"github.com/google/go-sev-guest/validate"
-	"github.com/google/go-sev-guest/verify"
-	"github.com/google/go-sev-guest/verify/trust"
+	"github.com/google/go-sev-guest/tools/lib/report"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/types/known/wrapperspb"
 )
@@ -41,6 +42,7 @@ const (
 	size48                  = 48
 	size64                  = 64
 	attestationFilePath     = "attestation.bin"
+	attestationJson         = "attestation.json"
 	sevProductNameMilan     = "Milan"
 	sevProductNameGenoa     = "Genoa"
 	exampleJSONConfig       = `
@@ -115,6 +117,8 @@ var (
 	empty32             = [size32]byte{}
 	empty64             = [size64]byte{}
 	defaultReportIdMa   = []byte{255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255}
+	getJsonAttestation  bool
+	errReportSize       = errors.New("attestation contents too small")
 )
 
 func (cli *CLI) NewAttestationCmd() *cobra.Command {
@@ -122,9 +126,9 @@ func (cli *CLI) NewAttestationCmd() *cobra.Command {
 		Use:   "attestation [command]",
 		Short: "Get and validate attestations",
 		Run: func(cmd *cobra.Command, args []string) {
-			fmt.Printf("Get and validate attestations\n\n")
-			fmt.Printf("Usage:\n  %s [command]\n\n", cmd.CommandPath())
-			fmt.Printf("Available Commands:\n")
+			cmd.Printf("Get and validate attestations\n\n")
+			cmd.Printf("Usage:\n  %s [command]\n\n", cmd.CommandPath())
+			cmd.Printf("Available Commands:\n")
 
 			// Filter out "completion" command
 			availableCommands := make([]*cobra.Command, 0)
@@ -135,47 +139,115 @@ func (cli *CLI) NewAttestationCmd() *cobra.Command {
 			}
 
 			for _, subCmd := range availableCommands {
-				fmt.Printf("  %-15s%s\n", subCmd.Name(), subCmd.Short)
+				cmd.Printf("  %-15s%s\n", subCmd.Name(), subCmd.Short)
 			}
 
-			fmt.Printf("\nFlags:\n")
+			cmd.Printf("\nFlags:\n")
 			cmd.Flags().VisitAll(func(flag *pflag.Flag) {
-				fmt.Printf("  -%s, --%s %s\n", flag.Shorthand, flag.Name, flag.Usage)
+				cmd.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())
+			cmd.Printf("\nUse \"%s [command] --help\" for more information about a command.\n", cmd.CommandPath())
 		},
 	}
 }
 
 func (cli *CLI) NewGetAttestationCmd() *cobra.Command {
-	return &cobra.Command{
+	cmd := &cobra.Command{
 		Use:     "get",
 		Short:   "Retrieve attestation information from agent. Report data expected in hex enoded string of length 64 bytes.",
 		Example: "get <report_data>",
 		Args:    cobra.ExactArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
-			log.Println("Getting attestation")
+			if cli.connectErr != nil {
+				printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
+				return
+			}
+
+			cmd.Println("Getting attestation")
 
 			reportData, err := hex.DecodeString(args[0])
 			if err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error decoding report data: %v ❌ ", err)
+				return
 			}
 			if len(reportData) != agent.ReportDataSize {
-				log.Fatalf("report data must be a hex encoded string of length %d bytes", agent.ReportDataSize)
+				msg := color.New(color.FgRed).Sprintf("report data must be a hex encoded string of length %d bytes ❌ ", agent.ReportDataSize)
+				cmd.Println(msg)
+				return
 			}
 
-			result, err := cli.agentSDK.Attestation(cmd.Context(), [agent.ReportDataSize]byte(reportData))
+			filename := attestationFilePath
+			if getJsonAttestation {
+				filename = attestationJson
+			}
+
+			attestationFile, err := os.Create(filename)
 			if err != nil {
-				log.Fatalf("Error retrieving attestation: %v", err)
+				printError(cmd, "Error creating attestation file: %v ❌ ", err)
+				return
 			}
 
-			if err = os.WriteFile(attestationFilePath, result, 0o644); err != nil {
-				log.Fatalf("Error saving attestation result: %v", err)
+			if err := cli.agentSDK.Attestation(cmd.Context(), [agent.ReportDataSize]byte(reportData), attestationFile); err != nil {
+				printError(cmd, "Failed to get attestation due to error: %v ❌ ", err)
+				return
 			}
 
-			log.Println("Attestation result retrieved and saved successfully!")
+			if err := attestationFile.Close(); err != nil {
+				printError(cmd, "Error closing attestation file: %v ❌ ", err)
+				return
+			}
+
+			if getJsonAttestation {
+				result, err := os.ReadFile(filename)
+				if err != nil {
+					printError(cmd, "Error reading attestation file: %v ❌ ", err)
+					return
+				}
+
+				result, err = attesationToJSON(result)
+				if err != nil {
+					printError(cmd, "Error converting attestation to json: %v ❌ ", err)
+					return
+				}
+
+				if err := os.WriteFile(filename, result, 0o644); err != nil {
+					printError(cmd, "Error writing attestation file: %v ❌ ", err)
+					return
+				}
+			}
+
+			cmd.Println("Attestation result retrieved and saved successfully!")
 		},
 	}
+
+	cmd.Flags().BoolVarP(&getJsonAttestation, "json", "j", false, "Get attestation in json format")
+
+	return cmd
+}
+
+func attesationToJSON(report []byte) ([]byte, error) {
+	if len(report) < abi.ReportSize {
+		return nil, errors.Wrap(errReportSize, fmt.Errorf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes", len(report), abi.ReportSize))
+	}
+	attestationPB, err := abi.ReportCertsToProto(report[:abi.ReportSize])
+	if err != nil {
+		return nil, err
+	}
+
+	return json.MarshalIndent(attestationPB, "", "	")
+}
+
+func attesationFromJSON(reportFile []byte) ([]byte, error) {
+	var attestationPB sevsnp.Attestation
+	if err := json.Unmarshal(reportFile, &attestationPB); err != nil {
+		return nil, err
+	}
+
+	return report.Transform(&attestationPB, "bin")
+}
+
+func isFileJSON(filename string) bool {
+	return strings.HasSuffix(filename, ".json")
 }
 
 func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
@@ -185,37 +257,45 @@ func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
 		Example: "validate <attestation_report_file_path>",
 		Args:    cobra.ExactArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
-			log.Println("Checking attestation")
+			cmd.Println("Checking attestation")
 
 			attestationFile = string(args[0])
 
 			if err := parseConfig(); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error parsing config: %v ❌ ", err)
+				return
 			}
 			if err := parseHashes(); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error parsing hashes: %v ❌ ", err)
+				return
 			}
 			if err := parseFiles(); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error parsing files: %v ❌ ", err)
+				return
 			}
 			// This format is the attestation report in AMD's specified ABI format, immediately
 			// followed by the certificate table bytes.
 			if len(attestation) < abi.ReportSize {
-				log.Fatalf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes", len(attestation), abi.ReportSize)
+				msg := color.New(color.FgRed).Sprintf("attestation contents too small (0x%x bytes). Want at least 0x%x bytes ❌ ", len(attestation), abi.ReportSize)
+				cmd.Println(msg)
+				return
 			}
 			if err := parseUints(); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error parsing uints: %v ❌ ", err)
+				return
 			}
 			cfg.Policy.Vmpl = wrapperspb.UInt32(0)
 
 			if err := validateInput(); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+				printError(cmd, "Error validating input: %v ❌ ", err)
+				return
 			}
 
-			if err := verifyAndValidateAttestation(attestation); err != nil {
-				log.Fatalf("attestation validation and verification failed with error: %s", err)
+			if err := quoteprovider.VerifyAndValidate(attestation, &cfg); err != nil {
+				printError(cmd, "Attestation validation and verification failed with error: %v ❌ ", err)
+				return
 			}
-			log.Println("Attestation validation and verification is successful!")
+			cmd.Println("Attestation validation and verification is successful!")
 		},
 	}
 	cmd.Flags().StringVar(
@@ -394,67 +474,18 @@ func (cli *CLI) NewValidateAttestationValidationCmd() *cobra.Command {
 	)
 
 	if err := cmd.MarkFlagRequired("report_data"); err != nil {
-		log.Fatalf("Failed to mark flag as required: %s", err)
+		printError(cmd, "Failed to mark flag as required: %v ❌ ", err)
+		return nil
 	}
 
 	if err := cmd.MarkFlagRequired("product"); err != nil {
-		log.Fatalf("Failed to mark flag as required: %s", err)
+		printError(cmd, "Failed to mark flag as required: %v ❌ ", err)
+		return nil
 	}
 
 	return cmd
 }
 
-func verifyAndValidateAttestation(attestation []byte) error {
-	sopts, err := verify.RootOfTrustToOptions(cfg.RootOfTrust)
-	if err != nil {
-		return err
-	}
-
-	if cfg.Policy.Product == nil {
-		productName := sevsnp.SevProduct_SEV_PRODUCT_UNKNOWN
-		switch cfg.RootOfTrust.ProductLine {
-		case sevProductNameMilan:
-			productName = sevsnp.SevProduct_SEV_PRODUCT_MILAN
-		case sevProductNameGenoa:
-			productName = sevsnp.SevProduct_SEV_PRODUCT_GENOA
-		default:
-		}
-
-		if productName == sevsnp.SevProduct_SEV_PRODUCT_UNKNOWN {
-			return fmt.Errorf("product name must be %s or %s", sevProductNameMilan, sevProductNameGenoa)
-		}
-
-		sopts.Product = &sevsnp.SevProduct{
-			Name: productName,
-		}
-	} else {
-		sopts.Product = cfg.Policy.Product
-	}
-
-	sopts.Getter = &trust.RetryHTTPSGetter{
-		Timeout:       timeout,
-		MaxRetryDelay: maxRetryDelay,
-		Getter:        &trust.SimpleHTTPSGetter{},
-	}
-
-	// Only take the attestation report and ignore everything else.
-	attestationPB, err := abi.ReportCertsToProto(attestation[:abi.ReportSize])
-	if err != nil {
-		return err
-	}
-	if err = verify.SnpAttestation(attestationPB, sopts); err != nil {
-		return err
-	}
-	opts, err := validate.PolicyToOptions(cfg.Policy)
-	if err != nil {
-		return err
-	}
-	if err = validate.SnpAttestation(attestationPB, opts); err != nil {
-		return err
-	}
-	return nil
-}
-
 // parseConfig decodes config passed as json for check.Config struct.
 // example
 /* {
@@ -549,6 +580,13 @@ func parseFiles() error {
 		return err
 	}
 	attestation = file
+	if isFileJSON(attestationFile) {
+		attestation, err = attesationFromJSON(attestation)
+		if err != nil {
+			return err
+		}
+	}
+
 	for _, path := range trustedAuthorKeys {
 		file, err := os.ReadFile(path)
 		if err != nil {
@@ -614,7 +652,7 @@ func getBase(val string) int {
 }
 
 func validateInput() error {
-	if len(cfg.RootOfTrust.CabundlePaths) != 0 || len(cfg.RootOfTrust.Cabundles) != 0 && cfg.RootOfTrust.Product == "" {
+	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")
 	}
 

cli/attestation_policy.go

@@ -4,15 +4,14 @@ 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"
+	"google.golang.org/protobuf/encoding/protojson"
 )
 
 type fieldType int
@@ -31,26 +30,21 @@ const (
 )
 
 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")
+	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")
 )
 
-type AttestationConfiguration struct {
-	SNPPolicy   *check.Policy      `json:"snp_policy,omitempty"`
-	RootOfTrust *check.RootOfTrust `json:"root_of_trust,omitempty"`
-}
-
-func (cli *CLI) NewBackendCmd() *cobra.Command {
+func (cli *CLI) NewAttestationPolicyCmd() *cobra.Command {
 	return &cobra.Command{
-		Use:   "backend [command]",
-		Short: "Change backend information",
+		Use:   "policy [command]",
+		Short: "Change attestation policy",
 		Run: func(cmd *cobra.Command, args []string) {
-			fmt.Printf("Change backend information\n\n")
+			fmt.Printf("Change attestation policy\n\n")
 			fmt.Printf("Usage:\n  %s [command]\n\n", cmd.CommandPath())
 			fmt.Printf("Available Commands:\n")
 
@@ -78,12 +72,13 @@ func (cli *CLI) NewBackendCmd() *cobra.Command {
 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>",
+		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 {
-				log.Fatalf("Error could not change measurement data %v", err)
+				printError(cmd, "Error could not change measurement data: %v ❌ ", err)
+				return
 			}
 		},
 	}
@@ -92,12 +87,13 @@ func (cli *CLI) NewAddMeasurementCmd() *cobra.Command {
 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>",
+		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 {
-				log.Fatalf("Error could not change host data %v", err)
+				printError(cmd, "Error could not change host data: %v ❌ ", err)
+				return
 			}
 		},
 	}
@@ -113,27 +109,27 @@ func changeAttestationConfiguration(fileName, base64Data string, expectedLength
 		return errDataLength
 	}
 
-	ac := AttestationConfiguration{}
+	ac := check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
 
-	backendInfo, err := os.ReadFile(fileName)
+	attestationPolicy, err := os.ReadFile(fileName)
 	if err != nil {
-		return errors.Wrap(errReadingBackendInfoFile, err)
+		return errors.Wrap(errReadingAttestationPolicyFile, err)
 	}
 
-	if err = json.Unmarshal(backendInfo, &ac); err != nil {
+	if err = protojson.Unmarshal(attestationPolicy, &ac); err != nil {
 		return errors.Wrap(errUnmarshalJSON, err)
 	}
 
 	switch field {
 	case measurementField:
-		ac.SNPPolicy.Measurement = data
+		ac.Policy.Measurement = data
 	case hostDataField:
-		ac.SNPPolicy.HostData = data
+		ac.Policy.HostData = data
 	default:
-		return errBackendField
+		return errAttestationPolicyField
 	}
 
-	fileJson, err := json.MarshalIndent(ac, "", " ")
+	fileJson, err := protojson.Marshal(&ac)
 	if err != nil {
 		return errors.Wrap(errMarshalJSON, err)
 	}

cli/attestation_policy_test.go

@@ -0,0 +1,131 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"encoding/base64"
+	"os"
+	"testing"
+
+	"github.com/google/go-sev-guest/proto/check"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/encoding/protojson"
+)
+
+func TestChangeAttestationConfiguration(t *testing.T) {
+	tmpfile, err := os.CreateTemp("", "attestation_policy.json")
+	require.NoError(t, err)
+	defer os.Remove(tmpfile.Name())
+
+	initialConfig := check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+
+	initialJSON, err := protojson.Marshal(&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)
+
+				config := check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+				err = protojson.Unmarshal(content, &config)
+				require.NoError(t, err)
+
+				decodedData, _ := base64.StdEncoding.DecodeString(tt.base64Data)
+				if tt.field == measurementField {
+					assert.Equal(t, decodedData, config.Policy.Measurement)
+				} else if tt.field == hostDataField {
+					assert.Equal(t, decodedData, 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)
+}

cli/attestation_test.go

@@ -0,0 +1,450 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"bytes"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"os"
+	"testing"
+
+	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/google/go-sev-guest/abi"
+	"github.com/google/go-sev-guest/proto/check"
+	"github.com/google/go-sev-guest/proto/sevsnp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
+)
+
+func TestNewAttestationCmd(t *testing.T) {
+	mockSDK := new(mocks.SDK)
+	cli := &CLI{agentSDK: mockSDK}
+	cmd := cli.NewAttestationCmd()
+
+	assert.Equal(t, "attestation [command]", cmd.Use)
+	assert.Equal(t, "Get and validate attestations", cmd.Short)
+
+	var buf bytes.Buffer
+	cmd.SetOut(&buf)
+
+	cmd.SetOutput(&buf)
+
+	reportData := bytes.Repeat([]byte{0x01}, agent.ReportDataSize)
+	mockSDK.On("Attestation", mock.Anything, [agent.ReportDataSize]byte(reportData), mock.Anything).Return(nil)
+
+	cmd.SetArgs([]string{hex.EncodeToString(reportData)})
+	err := cmd.Execute()
+	assert.NoError(t, err)
+	assert.Contains(t, buf.String(), "Get and validate attestations")
+}
+
+func TestNewGetAttestationCmd(t *testing.T) {
+	validattestation, err := os.ReadFile("../attestation.bin")
+	require.NoError(t, err)
+	testCases := []struct {
+		name         string
+		args         []string
+		mockResponse []byte
+		mockError    error
+		expectedErr  string
+		expectedOut  string
+	}{
+		{
+			name:         "successful attestation retrieval",
+			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
+			mockResponse: []byte("mock attestation"),
+			mockError:    nil,
+			expectedOut:  "Attestation result retrieved and saved successfully!",
+		},
+		{
+			name:         "invalid report data (decoding error)",
+			args:         []string{"invalid"},
+			mockResponse: nil,
+			mockError:    errors.New("error"),
+			expectedErr:  "Error decoding report data",
+		},
+		{
+			name:         "invalid report data size",
+			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, 32))},
+			mockResponse: nil,
+			mockError:    errors.New("error"),
+			expectedErr:  "report data must be a hex encoded string of length 64 bytes",
+		},
+		{
+			name:         "invalid report data hex",
+			args:         []string{"invalid"},
+			mockResponse: nil,
+			mockError:    errors.New("error"),
+			expectedErr:  "Error decoding report data",
+		},
+		{
+			name:         "failed to get attestation",
+			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
+			mockResponse: nil,
+			mockError:    errors.New("error"),
+			expectedErr:  "Failed to get attestation due to error",
+		},
+		{
+			name:         "JSON report error",
+			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), "--json"},
+			mockResponse: []byte("mock attestation"),
+			mockError:    nil,
+			expectedErr:  "Error converting attestation to json",
+		},
+		{
+			name:         "successful JSON report",
+			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), "--json"},
+			mockResponse: validattestation,
+			mockError:    nil,
+			expectedOut:  "Attestation result retrieved and saved successfully!",
+		},
+		{
+			name:         "connection error",
+			args:         []string{hex.EncodeToString(bytes.Repeat([]byte{0x01}, agent.ReportDataSize))},
+			mockResponse: nil,
+			mockError:    errors.New("failed to connect to agent"),
+			expectedErr:  "Failed to connect to agent",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Cleanup(func() {
+				os.Remove(attestationFilePath)
+				os.Remove(attestationJson)
+			})
+			mockSDK := new(mocks.SDK)
+			cli := &CLI{agentSDK: mockSDK}
+			if tc.name == "connection error" {
+				cli.connectErr = errors.New("failed to connect to agent")
+			}
+			cmd := cli.NewGetAttestationCmd()
+			var buf bytes.Buffer
+			cmd.SetOutput(&buf)
+
+			mockSDK.On("Attestation", mock.Anything, [agent.ReportDataSize]byte(bytes.Repeat([]byte{0x01}, agent.ReportDataSize)), mock.Anything).Return(tc.mockError).Run(func(args mock.Arguments) {
+				_, err := args.Get(2).(*os.File).Write(tc.mockResponse)
+				require.NoError(t, err)
+			})
+
+			cmd.SetArgs(tc.args)
+			err := cmd.Execute()
+
+			if tc.expectedErr != "" {
+				assert.Contains(t, buf.String(), tc.expectedErr)
+			} else {
+				assert.NoError(t, err)
+				assert.Contains(t, buf.String(), tc.expectedOut)
+			}
+		})
+	}
+}
+
+func TestNewValidateAttestationValidationCmd(t *testing.T) {
+	cli := &CLI{}
+	cmd := cli.NewValidateAttestationValidationCmd()
+
+	assert.Equal(t, "validate", cmd.Use)
+	assert.Equal(t, "Validate and verify attestation information. The report is provided as a file path.", cmd.Short)
+
+	assert.Equal(t, fmt.Sprint(defaultMinimumTcb), cmd.Flag("minimum_tcb").Value.String())
+	assert.Equal(t, fmt.Sprint(defaultMinimumLaunchTcb), cmd.Flag("minimum_lauch_tcb").Value.String())
+	assert.Equal(t, fmt.Sprint(defaultGuestPolicy), cmd.Flag("guest_policy").Value.String())
+	assert.Equal(t, fmt.Sprint(defaultMinimumGuestSvn), cmd.Flag("minimum_guest_svn").Value.String())
+	assert.Equal(t, fmt.Sprint(defaultMinimumBuild), cmd.Flag("minimum_build").Value.String())
+	assert.Equal(t, defaultCheckCrl, cmd.Flag("check_crl").Value.String() == "true")
+	assert.Equal(t, fmt.Sprint(defaultTimeout), cmd.Flag("timeout").Value.String())
+	assert.Equal(t, fmt.Sprint(defaultMaxRetryDelay), cmd.Flag("max_retry_delay").Value.String())
+}
+
+func TestParseConfig(t *testing.T) {
+	cfgString = ""
+	err := parseConfig()
+	assert.NoError(t, err)
+	assert.NotNil(t, cfg.RootOfTrust)
+	assert.NotNil(t, cfg.Policy)
+
+	cfgString = `{"rootOfTrust":{"product":"test_product"},"policy":{"minimumGuestSvn":1}}`
+	err = parseConfig()
+	assert.NoError(t, err)
+	assert.Equal(t, "test_product", cfg.RootOfTrust.Product)
+	assert.Equal(t, uint32(1), cfg.Policy.MinimumGuestSvn)
+
+	cfgString = `{"invalid_json"`
+	err = parseConfig()
+	assert.Error(t, err)
+}
+
+func TestParseHashes(t *testing.T) {
+	trustedAuthorHashes = []string{"0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"}
+	trustedIdKeyHashes = []string{"fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210fedcba9876543210"}
+
+	cfg = check.Config{}
+	if cfg.Policy == nil {
+		cfg.Policy = &check.Policy{}
+	}
+
+	err := parseHashes()
+	assert.NoError(t, err)
+	assert.Len(t, cfg.Policy.TrustedAuthorKeyHashes, 1)
+	assert.Len(t, cfg.Policy.TrustedIdKeyHashes, 1)
+
+	trustedAuthorHashes = []string{"invalid_hash"}
+	err = parseHashes()
+	assert.Error(t, err)
+}
+
+func TestParseFiles(t *testing.T) {
+	attestationFile = "test_attestation.bin"
+	authorKeyFile := "test_author_key.pem"
+	idKeyFile := "test_id_key.pem"
+
+	err := os.WriteFile(attestationFile, []byte("test attestation"), 0o644)
+	assert.NoError(t, err)
+	err = os.WriteFile(authorKeyFile, []byte("test author key"), 0o644)
+	assert.NoError(t, err)
+	err = os.WriteFile(idKeyFile, []byte("test id key"), 0o644)
+	assert.NoError(t, err)
+
+	trustedAuthorKeys = []string{authorKeyFile}
+	trustedIdKeys = []string{idKeyFile}
+
+	err = parseFiles()
+	assert.NoError(t, err)
+	assert.Equal(t, []byte("test attestation"), attestation)
+	assert.Len(t, cfg.Policy.TrustedAuthorKeys, 1)
+	assert.Len(t, cfg.Policy.TrustedIdKeys, 1)
+
+	os.Remove(attestationFile)
+	os.Remove(authorKeyFile)
+	os.Remove(idKeyFile)
+
+	attestationFile = "non_existent_file.bin"
+	err = parseFiles()
+	assert.Error(t, err)
+}
+
+func TestParseUints(t *testing.T) {
+	stepping = "10"
+	platformInfo = "0xFF"
+
+	cfg = check.Config{}
+	if cfg.Policy == nil {
+		cfg.Policy = &check.Policy{
+			Product: &sevsnp.SevProduct{},
+		}
+	}
+	err := parseUints()
+	assert.NoError(t, err)
+	assert.Equal(t, uint32(10), cfg.Policy.Product.MachineStepping.Value)
+	assert.Equal(t, uint64(255), cfg.Policy.PlatformInfo.Value)
+
+	stepping = "invalid"
+	err = parseUints()
+	assert.Error(t, err)
+
+	stepping = "10"
+	platformInfo = "invalid"
+	err = parseUints()
+	assert.Error(t, err)
+}
+
+func TestValidateInput(t *testing.T) {
+	cfg = check.Config{}
+	if cfg.Policy == nil {
+		cfg.Policy = &check.Policy{}
+	}
+	if cfg.RootOfTrust == nil {
+		cfg.RootOfTrust = &check.RootOfTrust{}
+	}
+	cfg.Policy.ReportData = make([]byte, 64)
+	cfg.Policy.HostData = make([]byte, 32)
+	cfg.Policy.FamilyId = make([]byte, 16)
+	cfg.Policy.ImageId = make([]byte, 16)
+	cfg.Policy.ReportId = make([]byte, 32)
+	cfg.Policy.ReportIdMa = make([]byte, 32)
+	cfg.Policy.Measurement = make([]byte, 48)
+	cfg.Policy.ChipId = make([]byte, 64)
+
+	err := validateInput()
+	assert.NoError(t, err)
+
+	cfg.Policy.ReportData = make([]byte, 32)
+	err = validateInput()
+	assert.Error(t, err)
+}
+
+func TestGetBase(t *testing.T) {
+	assert.Equal(t, 16, getBase("0xFF"))
+	assert.Equal(t, 8, getBase("0o77"))
+	assert.Equal(t, 2, getBase("0b1010"))
+	assert.Equal(t, 10, getBase("123"))
+}
+
+func TestAttestationToJSON(t *testing.T) {
+	validReport, err := os.ReadFile("../attestation.bin")
+	require.NoError(t, err)
+	tests := []struct {
+		name  string
+		input []byte
+		err   error
+	}{
+		{
+			name:  "Valid report",
+			input: validReport,
+			err:   nil,
+		},
+		{
+			name:  "Invalid report size",
+			input: make([]byte, abi.ReportSize-1),
+			err:   errReportSize,
+		},
+		{
+			name:  "Nil input",
+			input: nil,
+			err:   errReportSize,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := attesationToJSON(tt.input)
+			assert.True(t, errors.Contains(err, tt.err))
+			if tt.err != nil {
+				assert.Nil(t, got)
+				return
+			}
+
+			require.NotNil(t, got)
+
+			var js map[string]interface{}
+			err = json.Unmarshal(got, &js)
+			assert.NoError(t, err)
+		})
+	}
+}
+
+func TestAttestationFromJSON(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    []byte
+		err      error
+		validate func(t *testing.T, output []byte)
+	}{
+		{
+			name: "Valid JSON",
+			input: func() []byte {
+				att := &sevsnp.Attestation{
+					Report: &sevsnp.Report{
+						CurrentTcb:      1,
+						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),
+					},
+				}
+				data, err := json.Marshal(att)
+				require.NoError(t, err)
+				return data
+			}(),
+			err: nil,
+			validate: func(t *testing.T, output []byte) {
+				assert.NotEmpty(t, output)
+			},
+		},
+		{
+			name:  "Invalid JSON",
+			input: []byte(`{"invalid": json`),
+			err:   errors.New("invalid character 'j' looking for beginning of value"),
+			validate: func(t *testing.T, output []byte) {
+				assert.Nil(t, output)
+			},
+		},
+		{
+			name:  "Empty input",
+			input: []byte{},
+			err:   errors.New("unexpected end of JSON input"),
+			validate: func(t *testing.T, output []byte) {
+				assert.Nil(t, output)
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := attesationFromJSON(tt.input)
+			assert.True(t, errors.Contains(err, tt.err))
+			tt.validate(t, got)
+		})
+	}
+}
+
+func TestIsFileJSON(t *testing.T) {
+	tests := []struct {
+		name     string
+		filename string
+		want     bool
+	}{
+		{
+			name:     "Valid JSON extension",
+			filename: "test.json",
+			want:     true,
+		},
+		{
+			name:     "Valid JSON extension with path",
+			filename: "/path/to/test.json",
+			want:     true,
+		},
+		{
+			name:     "Invalid extension",
+			filename: "test.txt",
+			want:     false,
+		},
+		{
+			name:     "No extension",
+			filename: "test",
+			want:     false,
+		},
+		{
+			name:     "JSON in filename",
+			filename: "json.txt",
+			want:     false,
+		},
+		{
+			name:     "Empty string",
+			filename: "",
+			want:     false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := isFileJSON(tt.filename)
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}
+
+func TestRoundTrip(t *testing.T) {
+	originalReport, err := os.ReadFile("../attestation.bin")
+	require.NoError(t, err)
+	jsonData, err := attesationToJSON(originalReport)
+	require.NoError(t, err)
+	require.NotNil(t, jsonData)
+
+	roundTripReport, err := attesationFromJSON(jsonData)
+	require.NoError(t, err)
+	require.NotNil(t, roundTripReport)
+}

cli/cache.go

@@ -3,12 +3,13 @@
 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"
@@ -26,30 +27,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 := check.Config{Policy: &check.Policy{}, RootOfTrust: &check.RootOfTrust{}}
+			err := grpc.ReadAttestationPolicy(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.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

@@ -3,14 +3,24 @@
 package cli
 
 import (
-	"log"
+	"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 {
-	return &cobra.Command{
+	cmd := &cobra.Command{
 		Use:     "checksum",
 		Short:   "Compute the sha3-256 hash of a file",
 		Example: "checksum <file>",
@@ -18,12 +28,68 @@ func (cli *CLI) NewFileHashCmd() *cobra.Command {
 		Run: func(cmd *cobra.Command, args []string) {
 			path := args[0]
 
-			hash, err := internal.ChecksumHex(path)
-			if err != nil {
-				log.Fatalf("Error computing hash: %v", err)
+			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
 			}
 
-			log.Println("Hash of file:", hash)
+			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: "868825367c32c4b6d621d5d95e2890f233d8554df2348ab743aac2663a936f08",
+		},
+		{
+			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

@@ -6,10 +6,11 @@ import (
 	"context"
 	"crypto/x509"
 	"encoding/pem"
-	"log"
 	"os"
 	"path"
 
+	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/internal"
@@ -25,49 +26,61 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
 		Example: "data <dataset_path> <private_key_file_path>",
 		Args:    cobra.ExactArgs(2),
 		Run: func(cmd *cobra.Command, args []string) {
+			if cli.connectErr != nil {
+				printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
+				return
+			}
+
 			datasetPath := args[0]
 
-			log.Println("Uploading dataset:", datasetPath)
+			cmd.Println("Uploading dataset:", datasetPath)
 
 			f, err := os.Stat(datasetPath)
 			if err != nil {
-				log.Fatalf("Error reading dataset file: %v", err)
+				printError(cmd, "Error reading dataset file: %v ❌ ", err)
+				return
 			}
 
-			var dataset []byte
+			var dataset *os.File
 
 			if f.IsDir() {
-				dataset, err = internal.ZipDirectoryToMemory(datasetPath)
+				dataset, err = internal.ZipDirectoryToTempFile(datasetPath)
 				if err != nil {
-					log.Fatalf("Error zipping dataset directory: %v", err)
+					printError(cmd, "Error zipping dataset directory: %v ❌ ", err)
+					return
 				}
+				defer dataset.Close()
+				defer os.Remove(dataset.Name())
 			} else {
-				dataset, err = os.ReadFile(datasetPath)
+				dataset, err = os.Open(datasetPath)
 				if err != nil {
-					log.Fatalf("Error reading dataset file: %v", err)
+					printError(cmd, "Error reading dataset file: %v ❌ ", err)
+					return
 				}
-			}
-
-			dataReq := agent.Dataset{
-				Dataset:  dataset,
-				Filename: path.Base(datasetPath),
+				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)
-
-			ctx := metadata.NewOutgoingContext(cmd.Context(), metadata.New(make(map[string]string)))
-			if err := cli.agentSDK.Data(addDatasetMetadata(ctx), 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! ✔ "))
 		},
 	}
 
@@ -75,26 +88,28 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
 	return cmd
 }
 
-func decodeKey(b *pem.Block) interface{} {
+func decodeKey(b *pem.Block) (interface{}, 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")
 	}
 }
 

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)
+			}
+		})
+	}
+}

cli/errors.go

@@ -0,0 +1,49 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/fatih/color"
+	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/agent/auth"
+	"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+var (
+	errAgentUnavailable                   = errors.New("agent is unavailable on the current address")
+	errDigitalSignatureVerificationFailed = errors.New("digital signature verification failed, check the provided public key")
+)
+
+func decodeErros(err error) error {
+	statusErr, ok := status.FromError(err)
+	if ok {
+		switch statusErr.Code() {
+		case codes.PermissionDenied:
+			return errDigitalSignatureVerificationFailed
+		case codes.Unavailable:
+			return errAgentUnavailable
+		case codes.Unknown:
+			return err
+		}
+	}
+	switch {
+	case errors.Contains(err, auth.ErrSignatureVerificationFailed):
+		return auth.ErrSignatureVerificationFailed
+
+	case errors.Contains(err, agent.ErrAgentServiceUnavailable):
+		return agent.ErrAgentServiceUnavailable
+	default:
+		return err
+	}
+}
+
+func printError(cmd *cobra.Command, message string, err error) {
+	if !Verbose {
+		err = decodeErros(err)
+	}
+	msg := color.New(color.FgRed).Sprintf(message, err)
+	cmd.Println(msg)
+}

cli/errors_test.go

@@ -0,0 +1,110 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"bytes"
+	"errors"
+	"testing"
+
+	mgerrors "github.com/absmach/magistrala/pkg/errors"
+	"github.com/fatih/color"
+	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/agent/auth"
+	"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func TestDecodeErros(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    error
+		expected error
+	}{
+		{
+			name:     "Permission Denied",
+			input:    status.Error(codes.PermissionDenied, "permission denied"),
+			expected: errDigitalSignatureVerificationFailed,
+		},
+		{
+			name:     "Unavailable",
+			input:    status.Error(codes.Unavailable, "service unavailable"),
+			expected: errAgentUnavailable,
+		},
+		{
+			name:     "Unknown",
+			input:    status.Error(codes.Unknown, "unknown error"),
+			expected: status.Error(codes.Unknown, "unknown error"),
+		},
+		{
+			name:     "Signature Verification Failed",
+			input:    mgerrors.Wrap(auth.ErrSignatureVerificationFailed, errors.New("wrapped error")),
+			expected: auth.ErrSignatureVerificationFailed,
+		},
+		{
+			name:     "Other Error",
+			input:    errors.New("other error"),
+			expected: errors.New("other error"),
+		},
+		{
+			name:     "Agent Service Unavailable",
+			input:    mgerrors.Wrap(agent.ErrAgentServiceUnavailable, errors.New("wrapped error")),
+			expected: agent.ErrAgentServiceUnavailable,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := decodeErros(tt.input)
+			if result.Error() != tt.expected.Error() {
+				t.Errorf("decodeErros(%v) = %v, want %v", tt.input, result, tt.expected)
+			}
+		})
+	}
+}
+
+func TestPrintError(t *testing.T) {
+	// Save the original color.NoColor value and restore it after the test
+	origNoColor := color.NoColor
+	color.NoColor = true
+	defer func() { color.NoColor = origNoColor }()
+
+	tests := []struct {
+		name     string
+		message  string
+		err      error
+		verbose  bool
+		expected string
+	}{
+		{
+			name:     "Non-verbose mode",
+			message:  "Error: %s",
+			err:      status.Error(codes.PermissionDenied, "permission denied"),
+			verbose:  false,
+			expected: "Error: digital signature verification failed, check the provided public key\n",
+		},
+		{
+			name:     "Verbose mode",
+			message:  "Error: %s",
+			err:      status.Error(codes.PermissionDenied, "permission denied"),
+			verbose:  true,
+			expected: "Error: rpc error: code = PermissionDenied desc = permission denied\n",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			Verbose = tt.verbose
+			cmd := &cobra.Command{}
+			buf := new(bytes.Buffer)
+			cmd.SetOut(buf)
+
+			printError(cmd, tt.message, tt.err)
+
+			if got := buf.String(); got != tt.expected {
+				t.Errorf("printError() output = %q, want %q", got, tt.expected)
+			}
+		})
+	}
+}

cli/keys.go

@@ -10,9 +10,7 @@ import (
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/pem"
-	"log"
 	"os"
-	"reflect"
 
 	"github.com/spf13/cobra"
 )
@@ -25,6 +23,8 @@ const (
 	publicKeyType  = "PUBLIC KEY"
 	publicKeyFile  = "public.pem"
 	privateKeyFile = "private.pem"
+	ECDSA          = "ecdsa"
+	ED25519        = "ed25519"
 )
 
 var KeyType string
@@ -39,51 +39,67 @@ func (cli *CLI) NewKeysCmd() *cobra.Command {
 		Args:    cobra.ExactArgs(0),
 		Run: func(cmd *cobra.Command, args []string) {
 			switch KeyType {
-			case "ecdsa":
+			case ECDSA:
 				privEcdsaKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 				if err != nil {
-					log.Fatalf("Error generating keys: %v", err)
+					printError(cmd, "Error generating keys: %v ❌ ", err)
+					return
 				}
 
 				pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privEcdsaKey.PublicKey)
 				if err != nil {
-					log.Fatalf("Error marshalling public key: %v", err)
+					printError(cmd, "Error marshalling public key: %v ❌ ", err)
+					return
 				}
 
-				generateAndWriteKeys(privEcdsaKey, pubKeyBytes, ecdsaKeyType)
+				if err := generateAndWriteKeys(privEcdsaKey, pubKeyBytes, ecdsaKeyType); err != nil {
+					printError(cmd, "Error generating and writing keys: %v ❌ ", err)
+					return
+				}
 
-			case "ed25519":
+			case ED25519:
 				pubEd25519Key, privEd25519Key, err := ed25519.GenerateKey(rand.Reader)
 				if err != nil {
-					log.Fatalf("Error generating keys: %v", err)
+					printError(cmd, "Error generating keys: %v ❌ ", err)
+					return
 				}
 				pubKey, err := x509.MarshalPKIXPublicKey(pubEd25519Key)
 				if err != nil {
-					log.Fatalf("Error marshalling public key: %v", err)
+					printError(cmd, "Error marshalling public key: %v ❌ ", err)
+					return
+				}
+				if err := generateAndWriteKeys(privEd25519Key, pubKey, ed25519KeyType); err != nil {
+					printError(cmd, "Error generating and writing keys: %v ❌ ", err)
+					return
 				}
-				generateAndWriteKeys(privEd25519Key, pubKey, ed25519KeyType)
 
-			// Default to RSA
 			default:
 				privKey, err := rsa.GenerateKey(rand.Reader, keyBitSize)
 				if err != nil {
-					log.Fatalf("Error generating keys: %v", err)
+					printError(cmd, "Error generating keys: %v ❌ ", err)
+					return
 				}
 
 				pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privKey.PublicKey)
 				if err != nil {
-					log.Fatalf("Error marshalling public key: %v", err)
+					printError(cmd, "Error marshalling public key: %v ❌ ", err)
+					return
+				}
+				if err := generateAndWriteKeys(privKey, pubKeyBytes, rsaKeyType); err != nil {
+					printError(cmd, "Error generating and writing keys: %v ❌ ", err)
+					return
 				}
-				generateAndWriteKeys(privKey, pubKeyBytes, rsaKeyType)
 			}
+
+			cmd.Printf("Successfully generated public/private key pair of type: %s", KeyType)
 		},
 	}
 }
 
-func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType string) {
+func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType string) error {
 	privFile, err := os.Create(privateKeyFile)
 	if err != nil {
-		log.Fatalf("Error creating private key file: %v", err)
+		return err
 	}
 	defer privFile.Close()
 
@@ -97,22 +113,19 @@ func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType strin
 		b, err = x509.MarshalPKCS8PrivateKey(privKey)
 	}
 	if err != nil {
-		log.Printf("Error marshalling private key: %v", err)
-		return
+		return err
 	}
 
 	if err := pem.Encode(privFile, &pem.Block{
 		Type:  keyType,
 		Bytes: b,
 	}); err != nil {
-		log.Printf("Error encoding private key: %v", err)
-		return
+		return err
 	}
 
 	pubFile, err := os.Create(publicKeyFile)
 	if err != nil {
-		log.Printf("Error creating public key file: %v", err)
-		return
+		return err
 	}
 	defer pubFile.Close()
 
@@ -120,9 +133,8 @@ func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType strin
 		Type:  publicKeyType,
 		Bytes: pubKeyBytes,
 	}); err != nil {
-		log.Printf("Error encoding public key: %v", err)
-		return
+		return err
 	}
 
-	log.Printf("Successfully generated public/private key pair of type: %s", reflect.TypeOf(privKey).String())
+	return nil
 }

cli/keys_test.go

@@ -0,0 +1,92 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"crypto/ecdsa"
+	"crypto/ed25519"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/pem"
+	"os"
+	"testing"
+)
+
+func TestNewKeysCmd(t *testing.T) {
+	cli := &CLI{}
+	cmd := cli.NewKeysCmd()
+
+	if cmd.Use != "keys" {
+		t.Errorf("Expected Use to be 'keys', got %s", cmd.Use)
+	}
+
+	if cmd.Short != "Generate a new public/private key pair" {
+		t.Errorf("Unexpected Short description: %s", cmd.Short)
+	}
+}
+
+func TestGenerateAndWriteKeys(t *testing.T) {
+	tests := []struct {
+		name    string
+		keyType string
+	}{
+		{"RSA", "rsa"},
+		{"ECDSA", "ecdsa"},
+		{"ED25519", "ed25519"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			KeyType = tt.keyType
+			cmd := (&CLI{}).NewKeysCmd()
+			cmd.Run(cmd, []string{})
+
+			if _, err := os.Stat(privateKeyFile); os.IsNotExist(err) {
+				t.Errorf("Private key file was not created")
+			}
+			if _, err := os.Stat(publicKeyFile); os.IsNotExist(err) {
+				t.Errorf("Public key file was not created")
+			}
+
+			privKeyData, err := os.ReadFile(privateKeyFile)
+			if err != nil {
+				t.Fatalf("Failed to read private key file: %v", err)
+			}
+			privPem, _ := pem.Decode(privKeyData)
+			if privPem == nil {
+				t.Fatalf("Failed to decode private key PEM")
+			}
+
+			var privKey interface{}
+			switch tt.keyType {
+			case "rsa":
+				privKey, err = x509.ParsePKCS1PrivateKey(privPem.Bytes)
+			case "ecdsa":
+				privKey, err = x509.ParseECPrivateKey(privPem.Bytes)
+			case "ed25519":
+				privKey, err = x509.ParsePKCS8PrivateKey(privPem.Bytes)
+			}
+			if err != nil {
+				t.Fatalf("Failed to parse private key: %v", err)
+			}
+
+			switch tt.keyType {
+			case "rsa":
+				if _, ok := privKey.(*rsa.PrivateKey); !ok {
+					t.Errorf("Expected RSA private key, got %T", privKey)
+				}
+			case "ecdsa":
+				if _, ok := privKey.(*ecdsa.PrivateKey); !ok {
+					t.Errorf("Expected ECDSA private key, got %T", privKey)
+				}
+			case "ed25519":
+				if _, ok := privKey.(ed25519.PrivateKey); !ok {
+					t.Errorf("Expected ED25519 private key, got %T", privKey)
+				}
+			}
+
+			os.Remove(privateKeyFile)
+			os.Remove(publicKeyFile)
+		})
+	}
+}

cli/result.go

@@ -4,43 +4,64 @@ package cli
 
 import (
 	"encoding/pem"
-	"log"
 	"os"
 
+	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 )
 
-const resultFilePath = "results.zip"
+const (
+	resultFilePrefix = "results"
+	resultFileExt    = ".zip"
+	resultfilename   = "results.zip"
+)
 
 func (cli *CLI) NewResultsCmd() *cobra.Command {
 	return &cobra.Command{
 		Use:     "result",
 		Short:   "Retrieve computation result file",
-		Example: "result <private_key_file_path>",
-		Args:    cobra.ExactArgs(1),
+		Example: "result <private_key_file_path> <optional_file_name.zip>",
+		Args:    cobra.MinimumNArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
-			log.Println("Retrieving computation result file")
+			if cli.connectErr != nil {
+				printError(cmd, "Failed to connect to agent: %v ❌ ", cli.connectErr)
+				return
+			}
+
+			cmd.Println("⏳ Retrieving computation result file")
 
 			privKeyFile, err := os.ReadFile(args[0])
 			if err != nil {
-				log.Fatalf("Error reading private key file: %v", err)
+				printError(cmd, "Error reading private key file: %v ❌ ", err)
+				return
+			}
+
+			filename := resultfilename
+			if len(args) > 1 {
+				filename = args[1]
 			}
 
 			pemBlock, _ := pem.Decode(privKeyFile)
 
-			var result []byte
-
-			privKey := decodeKey(pemBlock)
-			result, err = cli.agentSDK.Result(cmd.Context(), privKey)
+			privKey, err := decodeKey(pemBlock)
 			if err != nil {
-				log.Fatalf("Error retrieving computation result: %v", err)
+				printError(cmd, "Error decoding private key: %v ❌ ", err)
+				return
 			}
 
-			if err := os.WriteFile(resultFilePath, result, 0o644); err != nil {
-				log.Fatalf("Error saving computation result to %s: %v", resultFilePath, err)
+			resultFile, err := os.Create(filename)
+			if err != nil {
+				printError(cmd, "Error creating result file: %v ❌ ", err)
+				return
+			}
+			defer resultFile.Close()
+
+			if err = cli.agentSDK.Result(cmd.Context(), privKey, resultFile); err != nil {
+				printError(cmd, "Error retrieving computation result: %v ❌ ", err)
+				return
 			}
 
-			log.Println("Computation result retrieved and saved successfully!")
+			cmd.Println(color.New(color.FgGreen).Sprintf("Computation result retrieved and saved successfully as %s! ✔ ", filename))
 		},
 	}
 }

cli/result_test.go

@@ -0,0 +1,177 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package cli
+
+import (
+	"bytes"
+	"errors"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/pkg/sdk/mocks"
+)
+
+const compResult = "Test computation result"
+
+func TestResultsCmd_MultipleExecutions(t *testing.T) {
+	mockSDK := new(mocks.SDK)
+	mockSDK.On("Result", mock.Anything, mock.Anything, mock.Anything).Return(nil).Run(func(args mock.Arguments) {
+		_, err := args.Get(2).(*os.File).WriteString(compResult)
+		require.NoError(t, err)
+	})
+	testCLI := CLI{agentSDK: mockSDK}
+
+	err := generateRSAPrivateKeyFile(privateKeyFile)
+	require.NoError(t, err)
+
+	cmd := testCLI.NewResultsCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetArgs([]string{privateKeyFile})
+
+	for i := 0; i < 3; i++ {
+		err = cmd.Execute()
+		require.NoError(t, err)
+		require.Contains(t, buf.String(), "Computation result retrieved and saved successfully")
+		buf.Reset()
+	}
+
+	files, err := filepath.Glob("results*.zip")
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		for _, file := range files {
+			os.Remove(file)
+		}
+		os.Remove(privateKeyFile)
+	})
+}
+
+func TestResultsCmd_InvalidPrivateKey(t *testing.T) {
+	mockSDK := new(mocks.SDK)
+	mockSDK.On("Result", mock.Anything, mock.Anything, mock.Anything).Return(nil).Run(func(args mock.Arguments) {
+		_, err := args.Get(2).(*os.File).WriteString(compResult)
+		require.NoError(t, err)
+	})
+	testCLI := CLI{agentSDK: mockSDK}
+
+	invalidPrivateKey, err := os.CreateTemp("", "invalid_private_key.pem")
+	require.NoError(t, err)
+	err = invalidPrivateKey.Close()
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		err := os.Remove(invalidPrivateKey.Name())
+		require.NoError(t, err)
+	})
+
+	cmd := testCLI.NewResultsCmd()
+	buf := new(bytes.Buffer)
+	cmd.SetOut(buf)
+	cmd.SetArgs([]string{invalidPrivateKey.Name()})
+	err = cmd.Execute()
+	require.NoError(t, err)
+
+	require.Contains(t, buf.String(), "Error decoding private key")
+	mockSDK.AssertNotCalled(t, "Result", mock.Anything, mock.Anything, mock.Anything)
+}
+
+func TestResultsCmd(t *testing.T) {
+	tests := []struct {
+		name           string
+		setupMock      func(*mocks.SDK)
+		setupFiles     func() (string, error)
+		connectErr     error
+		expectedOutput string
+		cleanup        func()
+	}{
+		{
+			name: "successful result retrieval",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Result", mock.Anything, mock.Anything, mock.Anything).Return(nil).Run(func(args mock.Arguments) {
+					_, err := args.Get(2).(*os.File).WriteString(compResult)
+					require.NoError(t, err)
+				})
+			},
+			setupFiles: func() (string, error) {
+				return privateKeyFile, generateRSAPrivateKeyFile(privateKeyFile)
+			},
+			expectedOutput: "Computation result retrieved and saved successfully",
+			cleanup: func() {
+				files, _ := filepath.Glob("results*.zip")
+				for _, file := range files {
+					os.Remove(file)
+				}
+				os.Remove(privateKeyFile)
+			},
+		},
+		{
+			name: "missing private key file",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Result", mock.Anything, mock.Anything, mock.Anything).Return(nil).Run(func(args mock.Arguments) {
+					_, err := args.Get(2).(*os.File).WriteString(compResult)
+					require.NoError(t, err)
+				})
+			},
+			setupFiles: func() (string, error) {
+				return "non_existent_private_key.pem", nil
+			},
+			expectedOutput: "Error reading private key file",
+		},
+		{
+			name: "result retrieval failure",
+			setupMock: func(m *mocks.SDK) {
+				m.On("Result", mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error retrieving computation result"))
+			},
+			setupFiles: func() (string, error) {
+				return privateKeyFile, generateRSAPrivateKeyFile(privateKeyFile)
+			},
+			expectedOutput: "error retrieving computation result",
+			cleanup: func() {
+				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,
+			}
+
+			file, err := tt.setupFiles()
+			require.NoError(t, err)
+
+			cmd := testCLI.NewResultsCmd()
+			buf := new(bytes.Buffer)
+			cmd.SetOut(buf)
+			cmd.SetArgs([]string{file})
+			err = cmd.Execute()
+			require.NoError(t, err)
+
+			require.Contains(t, buf.String(), tt.expectedOutput)
+
+			if tt.cleanup != nil {
+				tt.cleanup()
+			}
+		})
+	}
+}

cli/sdk.go

@@ -2,14 +2,43 @@
 // SPDX-License-Identifier: Apache-2.0
 package cli
 
-import "github.com/ultravioletrs/cocos/pkg/sdk"
+import (
+	"context"
+
+	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
+	"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
+	"github.com/ultravioletrs/cocos/pkg/sdk"
+)
+
+var Verbose bool
 
 type CLI struct {
-	agentSDK sdk.SDK
+	agentSDK   sdk.SDK
+	config     grpc.AgentClientConfig
+	client     grpc.Client
+	connectErr error
 }
 
-func New(agentSDK sdk.SDK) *CLI {
+func New(config grpc.AgentClientConfig) *CLI {
 	return &CLI{
-		agentSDK: agentSDK,
+		config: config,
 	}
 }
+
+func (c *CLI) InitializeSDK(cmd *cobra.Command) error {
+	agentGRPCClient, agentClient, err := agent.NewAgentClient(context.Background(), c.config)
+	if err != nil {
+		c.connectErr = err
+		return err
+	}
+	cmd.Println("🔗 Connected to agent using ", agentGRPCClient.Secure())
+	c.client = agentGRPCClient
+
+	c.agentSDK = sdk.NewAgentSDK(agentClient)
+	return nil
+}
+
+func (c *CLI) Close() {
+	c.client.Close()
+}

cmd/agent/main.go

@@ -3,14 +3,20 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
 	"io"
 	"log"
 	"log/slog"
+	"os"
+	"time"
 
+	mglog "github.com/absmach/magistrala/logger"
 	"github.com/absmach/magistrala/pkg/prometheus"
+	"github.com/cenkalti/backoff/v4"
+	"github.com/google/go-sev-guest/abi"
 	"github.com/google/go-sev-guest/client"
 	"github.com/mdlayher/vsock"
 	"github.com/ultravioletrs/cocos/agent"
@@ -18,12 +24,14 @@ import (
 	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
 	"github.com/ultravioletrs/cocos/agent/auth"
 	"github.com/ultravioletrs/cocos/agent/events"
-	"github.com/ultravioletrs/cocos/agent/quoteprovider"
 	agentlogger "github.com/ultravioletrs/cocos/internal/logger"
 	"github.com/ultravioletrs/cocos/internal/server"
 	grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
-	"github.com/ultravioletrs/cocos/manager"
+	ackvsock "github.com/ultravioletrs/cocos/internal/vsock"
+	managerevents "github.com/ultravioletrs/cocos/manager/events"
 	"github.com/ultravioletrs/cocos/manager/qemu"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider"
+	"golang.org/x/crypto/sha3"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
@@ -32,6 +40,7 @@ import (
 const (
 	svcName        = "agent"
 	defSvcGRPCPort = "7002"
+	retryInterval  = 5 * time.Second
 )
 
 func main() {
@@ -43,43 +52,63 @@ func main() {
 		log.Fatalf("failed to read agent configuration from vsock %s", err.Error())
 	}
 
-	conn, err := vsock.Dial(vsock.Host, manager.ManagerVsockPort, nil)
+	conn, err := dialVsock()
 	if err != nil {
 		log.Fatal(err)
 	}
 	defer conn.Close()
 
+	ackConn := ackvsock.NewAckWriter(conn)
+
+	var exitCode int
+	defer mglog.ExitWithError(&exitCode)
+
 	var level slog.Level
 	if err := level.UnmarshalText([]byte(cfg.AgentConfig.LogLevel)); err != nil {
 		log.Println(err)
+		exitCode = 1
 		return
 	}
-	handler := agentlogger.NewProtoHandler(conn, &slog.HandlerOptions{Level: level})
+
+	handler := agentlogger.NewProtoHandler(ackConn, &slog.HandlerOptions{Level: level}, cfg.ID)
 	logger := slog.New(handler)
 
-	eventSvc, err := events.New(svcName, cfg.ID, manager.ManagerVsockPort)
+	eventSvc, err := events.New(svcName, cfg.ID, ackConn)
 	if err != nil {
 		logger.Error(fmt.Sprintf("failed to create events service %s", err.Error()))
+		exitCode = 1
 		return
 	}
-	defer eventSvc.Close()
 
 	qp, err := quoteprovider.GetQuoteProvider()
 	if err != nil {
 		logger.Error(fmt.Sprintf("failed to create quote provider %s", err.Error()))
+		exitCode = 1
 		return
 	}
 
+	if err := verifyManifest(cfg, qp); err != nil {
+		logger.Error(err.Error())
+		exitCode = 1
+		return
+	}
+
+	setDefaultValues(&cfg)
+
 	svc := newService(ctx, logger, eventSvc, cfg, qp)
 
-	grpcServerConfig := server.Config{
-		Port:         cfg.AgentConfig.Port,
-		Host:         cfg.AgentConfig.Host,
-		CertFile:     cfg.AgentConfig.CertFile,
-		KeyFile:      cfg.AgentConfig.KeyFile,
-		ServerCAFile: cfg.AgentConfig.ServerCAFile,
-		ClientCAFile: cfg.AgentConfig.ClientCAFile,
-		AttestedTLS:  cfg.AgentConfig.AttestedTls,
+	agentGrpcServerConfig := server.AgentConfig{
+		ServerConfig: server.ServerConfig{
+			BaseConfig: server.BaseConfig{
+				Host:         cfg.AgentConfig.Host,
+				Port:         cfg.AgentConfig.Port,
+				CertFile:     cfg.AgentConfig.CertFile,
+				KeyFile:      cfg.AgentConfig.KeyFile,
+				ServerCAFile: cfg.AgentConfig.ServerCAFile,
+				ClientCAFile: cfg.AgentConfig.ClientCAFile,
+			},
+		},
+		AttestedTLS: cfg.AgentConfig.AttestedTls,
 	}
 
 	registerAgentServiceServer := func(srv *grpc.Server) {
@@ -90,10 +119,25 @@ func main() {
 	authSvc, err := auth.New(cfg)
 	if err != nil {
 		logger.Error(fmt.Sprintf("failed to create auth service %s", err.Error()))
+		exitCode = 1
 		return
 	}
 
-	gs := grpcserver.New(ctx, cancel, svcName, grpcServerConfig, registerAgentServiceServer, logger, qp, authSvc)
+	gs := grpcserver.New(ctx, cancel, svcName, agentGrpcServerConfig, registerAgentServiceServer, logger, qp, authSvc)
+
+	g.Go(func() error {
+		for {
+			if _, err := io.Copy(io.Discard, conn); err != nil {
+				log.Printf("vsock connection lost: %v, reconnecting...", err)
+				conn.Close()
+				conn, err = dialVsock()
+				if err != nil {
+					log.Fatal("failed to reconnect: ", err)
+				}
+			}
+			time.Sleep(retryInterval)
+		}
+	})
 
 	g.Go(func() error {
 		return gs.Start()
@@ -150,11 +194,71 @@ func readConfig() (agent.Computation, error) {
 	if err := json.Unmarshal(buffer, &ac); err != nil {
 		return agent.Computation{}, err
 	}
-	if ac.AgentConfig.LogLevel == "" {
-		ac.AgentConfig.LogLevel = "info"
-	}
-	if ac.AgentConfig.Port == "" {
-		ac.AgentConfig.Port = defSvcGRPCPort
-	}
 	return ac, nil
 }
+
+func setDefaultValues(cfg *agent.Computation) {
+	if cfg.AgentConfig.LogLevel == "" {
+		cfg.AgentConfig.LogLevel = "info"
+	}
+	if cfg.AgentConfig.Port == "" {
+		cfg.AgentConfig.Port = defSvcGRPCPort
+	}
+}
+
+func isTEE() bool {
+	_, err := os.Stat("/dev/sev-guest")
+	return !os.IsNotExist(err)
+}
+
+func dialVsock() (*vsock.Conn, error) {
+	var conn *vsock.Conn
+	var err error
+
+	err = backoff.Retry(func() error {
+		conn, err = vsock.Dial(vsock.Host, managerevents.ManagerVsockPort, nil)
+		if err == nil {
+			log.Println("vsock connection established")
+			return nil
+		}
+		log.Printf("vsock connection failed, retrying in %s... Error: %v", retryInterval, err)
+		return err
+	}, backoff.NewExponentialBackOff())
+	if err != nil {
+		return nil, err
+	}
+
+	return conn, nil
+}
+
+func verifyManifest(cfg agent.Computation, qp client.QuoteProvider) error {
+	if !isTEE() {
+		return nil
+	}
+
+	ar, err := qp.GetRawQuote(sha3.Sum512([]byte(cfg.ID)))
+	if err != nil {
+		return err
+	}
+
+	arProto, err := abi.ReportCertsToProto(ar[:abi.ReportSize])
+	if err != nil {
+		return err
+	}
+
+	cfgBytes, err := json.Marshal(cfg)
+	if err != nil {
+		return err
+	}
+
+	mcHash := sha3.Sum256(cfgBytes)
+
+	if arProto.Report.HostData == nil {
+		return fmt.Errorf("manifest verification failed: HostData is nil")
+	}
+	if !bytes.Equal(arProto.Report.HostData, mcHash[:]) {
+		return fmt.Errorf("manifest verification failed")
+	}
+
+	return nil
+}

cmd/agent/main_test.go

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

cmd/cli/main.go

@@ -4,18 +4,18 @@ package main
 
 import (
 	"fmt"
-	"log"
 	"os"
+	"os/signal"
 	"path"
+	"syscall"
 
-	mglog "github.com/absmach/magistrala/logger"
 	"github.com/caarlos0/env/v11"
+	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"github.com/ultravioletrs/cocos/cli"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
-	"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
-	"github.com/ultravioletrs/cocos/pkg/sdk"
+	cmd "github.com/virtee/sev-snp-measure-go/sevsnpmeasure/cmd"
 )
 
 const (
@@ -31,44 +31,6 @@ type config struct {
 }
 
 func main() {
-	var cfg config
-	if err := env.Parse(&cfg); err != nil {
-		log.Fatalf("failed to load %s configuration : %s", svcName, err)
-	}
-
-	homePath, err := os.UserHomeDir()
-	if err != nil {
-		log.Fatalf("Error fetching user home directory: %v", err)
-	}
-
-	directoryCachePath := path.Join(homePath, cocosDirectory)
-
-	if err := os.MkdirAll(directoryCachePath, filePermision); err != nil {
-		log.Fatalf("Error while creating directory %s, error: %v", directoryCachePath, err)
-	}
-
-	logger, err := mglog.New(os.Stdout, cfg.LogLevel)
-	if err != nil {
-		log.Fatalf("Error creating logger: %s", err)
-	}
-
-	agentGRPCConfig := grpc.Config{}
-	if err := env.ParseWithOptions(&agentGRPCConfig, env.Options{Prefix: envPrefixAgentGRPC}); err != nil {
-		logger.Error(fmt.Sprintf("failed to load %s gRPC client configuration : %s", svcName, err))
-		return
-	}
-
-	agentGRPCClient, agentClient, err := agent.NewAgentClient(agentGRPCConfig)
-	if err != nil {
-		logger.Error(err.Error())
-		return
-	}
-	defer agentGRPCClient.Close()
-
-	agentSDK := sdk.NewAgentSDK(logger, agentClient)
-
-	cliSVC := cli.New(agentSDK)
-
 	rootCmd := &cobra.Command{
 		Use:   "cocos-cli [command]",
 		Short: "CLI application for CoCos Service API",
@@ -97,9 +59,56 @@ func main() {
 		},
 	}
 
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
+
+	go func() {
+		<-signalChan
+		fmt.Println()
+		rootCmd.Println(color.New(color.FgRed).Sprint("Operation aborted by user!"))
+		os.Exit(2)
+	}()
+
+	var cfg config
+	if err := env.Parse(&cfg); err != nil {
+		message := color.New(color.FgRed).Sprintf("failed to load %s configuration : %s", svcName, err)
+		rootCmd.Println(message)
+		return
+	}
+
+	homePath, err := os.UserHomeDir()
+	if err != nil {
+		message := color.New(color.FgRed).Sprintf("failed to fetch user home directory: %s", err)
+		rootCmd.Println(message)
+		return
+	}
+
+	directoryCachePath := path.Join(homePath, cocosDirectory)
+
+	if err := os.MkdirAll(directoryCachePath, filePermision); err != nil {
+		message := color.New(color.FgRed).Sprintf("failed to create directory %s : %s", directoryCachePath, err)
+		rootCmd.Println(message)
+		return
+	}
+
+	agentGRPCConfig := grpc.AgentClientConfig{}
+	if err := env.ParseWithOptions(&agentGRPCConfig, env.Options{Prefix: envPrefixAgentGRPC}); err != nil {
+		message := color.New(color.FgRed).Sprintf("failed to load %s gRPC client configuration : %s", svcName, err)
+		rootCmd.Println(message)
+		return
+	}
+
+	cliSVC := cli.New(agentGRPCConfig)
+
+	if err := cliSVC.InitializeSDK(rootCmd); err == nil {
+		defer cliSVC.Close()
+	}
+
+	rootCmd.PersistentFlags().BoolVarP(&cli.Verbose, "verbose", "v", false, "Enable verbose output")
+
 	keysCmd := cliSVC.NewKeysCmd()
 	attestationCmd := cliSVC.NewAttestationCmd()
-	backendCmd := cliSVC.NewBackendCmd()
+	attestationPolicyCmd := cliSVC.NewAttestationPolicyCmd()
 
 	// Agent Commands
 	rootCmd.AddCommand(cliSVC.NewAlgorithmCmd())
@@ -107,7 +116,7 @@ func main() {
 	rootCmd.AddCommand(cliSVC.NewResultsCmd())
 	rootCmd.AddCommand(attestationCmd)
 	rootCmd.AddCommand(cliSVC.NewFileHashCmd())
-	rootCmd.AddCommand(backendCmd)
+	rootCmd.AddCommand(attestationPolicyCmd)
 	rootCmd.AddCommand(keysCmd)
 	rootCmd.AddCommand(cliSVC.NewCABundleCmd(directoryCachePath))
 
@@ -115,6 +124,9 @@ func main() {
 	attestationCmd.AddCommand(cliSVC.NewGetAttestationCmd())
 	attestationCmd.AddCommand(cliSVC.NewValidateAttestationValidationCmd())
 
+	// measure.
+	rootCmd.AddCommand(cmd.NewRootCmd())
+
 	// Flags
 	keysCmd.PersistentFlags().StringVarP(
 		&cli.KeyType,
@@ -124,12 +136,19 @@ func main() {
 		"User Key type",
 	)
 
-	// Backend information commands
-	backendCmd.AddCommand(cliSVC.NewAddMeasurementCmd())
-	backendCmd.AddCommand(cliSVC.NewAddHostDataCmd())
+	// Attestation Policy commands
+	attestationPolicyCmd.AddCommand(cliSVC.NewAddMeasurementCmd())
+	attestationPolicyCmd.AddCommand(cliSVC.NewAddHostDataCmd())
 
 	if err := rootCmd.Execute(); err != nil {
-		logger.Error(fmt.Sprintf("Command execution failed: %s", err))
+		logErrorCmd(*rootCmd, err)
 		return
 	}
 }
+
+func logErrorCmd(cmd cobra.Command, err error) {
+	boldRed := color.New(color.FgRed, color.Bold)
+	boldRed.Fprintf(cmd.ErrOrStderr(), "\nerror: ")
+
+	fmt.Fprintf(cmd.ErrOrStderr(), "%s\n\n", color.RedString(err.Error()))
+}

cmd/manager/main.go

@@ -10,7 +10,9 @@ import (
 	"log/slog"
 	"net/url"
 	"os"
+	"os/signal"
 	"strings"
+	"syscall"
 
 	mglog "github.com/absmach/magistrala/logger"
 	"github.com/absmach/magistrala/pkg/jaeger"
@@ -20,27 +22,29 @@ import (
 	"github.com/ultravioletrs/cocos/manager"
 	"github.com/ultravioletrs/cocos/manager/api"
 	managerapi "github.com/ultravioletrs/cocos/manager/api/grpc"
+	"github.com/ultravioletrs/cocos/manager/events"
 	"github.com/ultravioletrs/cocos/manager/qemu"
 	"github.com/ultravioletrs/cocos/manager/tracing"
-	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
+	pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	managergrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/manager"
-	pkgmanager "github.com/ultravioletrs/cocos/pkg/manager"
 	"go.opentelemetry.io/otel/trace"
 	"golang.org/x/sync/errgroup"
 )
 
 const (
-	svcName       = "manager"
-	envPrefixGRPC = "MANAGER_GRPC_"
-	envPrefixQemu = "MANAGER_QEMU_"
+	svcName          = "manager"
+	envPrefixGRPC    = "MANAGER_GRPC_"
+	envPrefixQemu    = "MANAGER_QEMU_"
+	clientBufferSize = 100
 )
 
 type config struct {
-	LogLevel                 string  `env:"MANAGER_LOG_LEVEL"                  envDefault:"info"`
-	JaegerURL                url.URL `env:"COCOS_JAEGER_URL"                   envDefault:"http://localhost:4318"`
-	TraceRatio               float64 `env:"COCOS_JAEGER_TRACE_RATIO"           envDefault:"1.0"`
-	InstanceID               string  `env:"MANAGER_INSTANCE_ID"                envDefault:""`
-	BackendMeasurementBinary string  `env:"MANAGER_BACKEND_MEASUREMENT_BINARY" envDefault:"../../build"`
+	LogLevel                string  `env:"MANAGER_LOG_LEVEL"                  envDefault:"info"`
+	JaegerURL               url.URL `env:"COCOS_JAEGER_URL"                   envDefault:"http://localhost:4318"`
+	TraceRatio              float64 `env:"COCOS_JAEGER_TRACE_RATIO"           envDefault:"1.0"`
+	InstanceID              string  `env:"MANAGER_INSTANCE_ID"                envDefault:""`
+	AttestationPolicyBinary string  `env:"MANAGER_ATTESTATION_POLICY_BINARY"  envDefault:"../../build"`
+	EosVersion              string  `env:"MANAGER_EOS_VERSION"                envDefault:""`
 }
 
 func main() {
@@ -54,12 +58,16 @@ func main() {
 
 	logger, err := mglog.New(os.Stdout, cfg.LogLevel)
 	if err != nil {
-		log.Fatalf(err.Error())
+		log.Fatal(err.Error())
 	}
 
+	var exitCode int
+	defer mglog.ExitWithError(&exitCode)
+
 	if cfg.InstanceID == "" {
 		if cfg.InstanceID, err = uuid.New().ID(); err != nil {
 			logger.Error(fmt.Sprintf("Failed to generate instance ID: %s", err))
+			exitCode = 1
 			return
 		}
 	}
@@ -78,20 +86,23 @@ func main() {
 	qemuCfg := qemu.Config{}
 	if err := env.ParseWithOptions(&qemuCfg, env.Options{Prefix: envPrefixQemu}); err != nil {
 		logger.Error(fmt.Sprintf("failed to load QEMU configuration: %s", err))
+		exitCode = 1
 		return
 	}
 	args := qemuCfg.ConstructQemuArgs()
 	logger.Info(strings.Join(args, " "))
 
-	managerGRPCConfig := grpc.Config{}
+	managerGRPCConfig := pkggrpc.ManagerClientConfig{}
 	if err := env.ParseWithOptions(&managerGRPCConfig, env.Options{Prefix: envPrefixGRPC}); err != nil {
 		logger.Error(fmt.Sprintf("failed to load %s gRPC client configuration : %s", svcName, err))
+		exitCode = 1
 		return
 	}
 
 	managerGRPCClient, managerClient, err := managergrpc.NewManagerClient(managerGRPCConfig)
 	if err != nil {
 		logger.Error(err.Error())
+		exitCode = 1
 		return
 	}
 	defer managerGRPCClient.Close()
@@ -99,17 +110,43 @@ func main() {
 	pc, err := managerClient.Process(ctx)
 	if err != nil {
 		logger.Error(err.Error())
+		exitCode = 1
 		return
 	}
 
-	eventsChan := make(chan *pkgmanager.ClientStreamMessage)
-	svc, err := newService(logger, tracer, qemuCfg, eventsChan, cfg.BackendMeasurementBinary)
+	eventsChan := make(chan *manager.ClientStreamMessage, clientBufferSize)
+	svc, err := newService(logger, tracer, qemuCfg, eventsChan, cfg.AttestationPolicyBinary, cfg.EosVersion)
 	if err != nil {
 		logger.Error(err.Error())
+		exitCode = 1
 		return
 	}
 
-	mc := managerapi.NewClient(pc, svc, eventsChan)
+	eventsSvc, err := events.New(logger, svc.ReportBrokenConnection, eventsChan)
+	if err != nil {
+		logger.Error(err.Error())
+		exitCode = 1
+		return
+	}
+
+	go eventsSvc.Listen(ctx)
+
+	mc := managerapi.NewClient(pc, svc, eventsChan, logger)
+
+	g.Go(func() error {
+		ch := make(chan os.Signal, 1)
+		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
+		defer signal.Stop(ch)
+
+		select {
+		case <-ch:
+			logger.Info("Received signal, shutting down...")
+			cancel()
+			return nil
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	})
 
 	g.Go(func() error {
 		return mc.Process(ctx, cancel)
@@ -120,12 +157,11 @@ func main() {
 	}
 }
 
-func newService(logger *slog.Logger, tracer trace.Tracer, qemuCfg qemu.Config, eventsChan chan *pkgmanager.ClientStreamMessage, backendMeasurementPath string) (manager.Service, error) {
-	svc, err := manager.New(qemuCfg, backendMeasurementPath, logger, eventsChan, qemu.NewVM)
+func newService(logger *slog.Logger, tracer trace.Tracer, qemuCfg qemu.Config, eventsChan chan *manager.ClientStreamMessage, attestationPolicyPath string, eosVersion string) (manager.Service, error) {
+	svc, err := manager.New(qemuCfg, attestationPolicyPath, logger, eventsChan, qemu.NewVM, eosVersion)
 	if err != nil {
 		return nil, err
 	}
-	go svc.RetrieveAgentEventsLogs()
 	svc = api.LoggingMiddleware(svc, logger)
 	counter, latency := prometheus.MakeMetrics(svcName, "api")
 	svc = api.MetricsMiddleware(svc, counter, latency)

cmd/manager/start_VM.sh

@@ -1,62 +0,0 @@
-#!/bin/bash
-
-# Set your default values for sudo and sev
-sudo_option=false
-sev_option=false
-
-# Parse command line arguments
-while [[ $# -gt 0 ]]; do
-    key="$1"
-
-    case $key in
-        --sudo)
-            sudo_option=true
-            shift
-            ;;
-        --sev)
-            sev_option=true
-            shift
-            ;;
-        *)
-            echo "Unknown option: $key"
-            exit 1
-            ;;
-    esac
-done
-
-build_qemu_command() {
-    local qemu_command="/usr/bin/qemu-system-x86_64 -enable-kvm -machine q35 -cpu EPYC -smp 4,maxcpus=64 -m 2048M,slots=5,maxmem=30G -drive if=pflash,format=raw,unit=0,file=$MANAGER_QEMU_OVMF_CODE_FILE,readonly=on -drive if=pflash,format=raw,unit=1,file=img/OVMF_VARS.fd -device virtio-scsi-pci,id=scsi,disable-legacy=on,iommu_platform=true -drive file=img/focal-server-cloudimg-amd64.img,if=none,id=disk0,format=qcow2 -device scsi-hd,drive=disk0 -netdev user,id=vmnic,hostfwd=tcp::2222-:22,hostfwd=tcp::9301-:9031,hostfwd=tcp::7020-:7002 -device virtio-net-pci,disable-legacy=on,iommu_platform=true,netdev=vmnic,romfile= -nographic -monitor pty"
-
-    if [ "$sev_option" = true ]; then
-        qemu_command="$qemu_command -object sev-guest,id=sev0,cbitpos=51,reduced-phys-bits=1 -machine memory-encryption=sev0"
-    fi
-
-    echo "$qemu_command"
-}
-
-if [ ! -f "img/OVMF_VARS.fd" ]; then
-    cp "$MANAGER_QEMU_OVMF_VARS_FILE" "img/OVMF_VARS.fd"
-    echo "Copied $MANAGER_QEMU_OVMF_VARS_FILE to img/OVMF_VARS.fd"
-else
-    echo "img/OVMF_VARS.fd already exists. No need to copy."
-fi
-
-echo "Launching VM ..."
-
-qemu_command=$(build_qemu_command)
-echo "$qemu_command"
-
-echo "Mapping CTRL-C to CTRL-]"
-stty intr ^]
-
-if [ "$sudo_option" = true ]; then
-    # Split the command and arguments into an array; << operator is known as a "here string"
-    IFS=" " read -r -a qemu_command_array <<< "$qemu_command"
-    # Treat each element in the array as a separate word, preserving spaces within each element
-    sudo "${qemu_command_array[@]}"
-else
-    $qemu_command
-fi
-
-# Restore the mapping
-stty intr ^c

cmd/manager/xml/dom.xml

@@ -1,62 +0,0 @@
-<domain type="kvm">
-  <name>QEmu-alpine-standard-x86_64</name>
-  <uuid>c7a5fdbd-cdaf-9455-926a-d65c16db1809</uuid>
-  <metadata>
-    <libosinfo:libosinfo xmlns:libosinfo="http://libosinfo.org/xmlns/libvirt/domain/1.0">
-      <libosinfo:os id="http://alpinelinux.org/alpinelinux/3.15"/>
-    </libosinfo:libosinfo>
-  </metadata>
-  <memory unit="KiB">4194304</memory>
-  <currentMemory unit="KiB">4194304</currentMemory>
-  <vcpu placement="static">1</vcpu>
-  <os>
-    <type arch="x86_64" machine="q35">hvm</type>
-    <bootmenu enable="yes"/>
-    <loader readonly="yes" type="pflash">/usr/share/OVMF/OVMF_CODE.fd</loader>
-    <nvram template='/usr/share/OVMF/OVMF_VARS.fd'>./img/OVMF_VARS.fd</nvram>
-    <!-- <boot dev='hd'/> -->
-  </os>
-  <features>
-    <acpi/>
-    <apic/>
-    <vmport state="off"/>
-  </features>
-  <cpu mode="host-passthrough" check="none" migratable="on"/>
-  <clock offset="utc">
-    <timer name="rtc" tickpolicy="catchup"/>
-    <timer name="pit" tickpolicy="delay"/>
-    <timer name="hpet" present="no"/>
-  </clock>
-  <on_poweroff>destroy</on_poweroff>
-  <on_reboot>restart</on_reboot>
-  <on_crash>destroy</on_crash>
-  <pm>
-    <suspend-to-mem enabled="no"/>
-    <suspend-to-disk enabled="no"/>
-  </pm>
-  <devices>
-    <emulator>/usr/bin/qemu-system-x86_64</emulator>
-    <disk type="file" device="disk">
-      <driver name="qemu" type="qcow2" discard="unmap"/>
-      <source file="./img/focal-server-cloudimg-amd64.qcow2"/>
-      <target dev="vda" bus="virtio"/>
-      <address type="pci" domain="0x0000" bus="0x04" slot="0x00" function="0x0"/>
-      <boot order="1"/>
-    </disk>
-    <graphics type="spice" autoport="yes">
-      <listen type="address"/>
-      <image compression="off"/>
-    </graphics>
-    <video>
-      <model type="qxl" ram="65536" vram="65536" vgamem="16384" heads="1" primary="yes"/>
-      <address type="pci" domain="0x0000" bus="0x00" slot="0x01" function="0x0"/>
-    </video>
-    <interface type="network">
-      <mac address="52:54:00:03:7b:5f"/>
-      <source network="default"/>
-      <model type="virtio"/>
-      <address type="pci" domain="0x0000" bus="0x01" slot="0x00" function="0x0"/>
-    </interface>
-  </devices>
-
-</domain>

cmd/manager/xml/pool.xml

@@ -1,6 +0,0 @@
-<pool type="dir">
-<name>virtimages</name>
-<target>
-  <path>./img</path>
-</target>
-</pool>

cmd/manager/xml/vol.xml

@@ -1,15 +0,0 @@
-<volume>
-<name>boot.img</name>
-<allocation>0</allocation>
-<capacity unit="G">1</capacity>
-<target>
-  <format type="qcow2"/>
-  <path>./img/boot.img</path>
-  <permissions>
-	<owner>107</owner>
-	<group>107</group>
-	<mode>0744</mode>
-	<label>virt_image_t</label>
-  </permissions>
-</target>
-</volume>

cocos-manager.env

@@ -0,0 +1,60 @@
+# Environment Configuration for Cocos
+
+# Jaeger Tracing
+COCOS_JAEGER_URL=http://localhost:4318
+COCOS_JAEGER_TRACE_RATIO=1.0
+
+# Manager Service Configuration
+MANAGER_INSTANCE_ID=
+MANAGER_ATTESTATION_POLICY_BINARY=../../build
+MANAGER_GRPC_CLIENT_CERT=
+MANAGER_GRPC_CLIENT_KEY=
+MANAGER_GRPC_SERVER_CA_CERTS=
+MANAGER_GRPC_URL=localhost:7001
+MANAGER_GRPC_TIMEOUT=60s
+MANAGER_EOS_VERSION=""
+
+# QEMU Configuration
+MANAGER_QEMU_MEMORY_SIZE=25G
+MANAGER_QEMU_MEMORY_SLOTS=5
+MANAGER_QEMU_MAX_MEMORY=30G
+MANAGER_QEMU_OVMF_CODE_IF=pflash
+MANAGER_QEMU_OVMF_CODE_FORMAT=raw
+MANAGER_QEMU_OVMF_CODE_UNIT=0
+MANAGER_QEMU_OVMF_CODE_FILE=/usr/share/OVMF/x64/OVMF_CODE.fd
+MANAGER_QEMU_OVMF_VERSION=edk2-stable202408
+MANAGER_QEMU_OVMF_CODE_READONLY=on
+MANAGER_QEMU_OVMF_VARS_IF=pflash
+MANAGER_QEMU_OVMF_VARS_FORMAT=raw
+MANAGER_QEMU_OVMF_VARS_UNIT=1
+MANAGER_QEMU_OVMF_VARS_FILE=/usr/share/OVMF/x64/OVMF_VARS.fd
+MANAGER_QEMU_NETDEV_ID=vmnic
+MANAGER_QEMU_HOST_FWD_AGENT=7020
+MANAGER_QEMU_GUEST_FWD_AGENT=7002
+MANAGER_QEMU_VIRTIO_NET_PCI_DISABLE_LEGACY=on
+MANAGER_QEMU_VIRTIO_NET_PCI_IOMMU_PLATFORM=true
+MANAGER_QEMU_VIRTIO_NET_PCI_ADDR=0x2
+MANAGER_QEMU_VIRTIO_NET_PCI_ROMFILE=
+MANAGER_QEMU_DISK_IMG_KERNEL_FILE=/home/sammyk/Documents/cocos-ai/cmd/manager/img/bzImage
+MANAGER_QEMU_DISK_IMG_ROOTFS_FILE=/home/sammyk/Documents/cocos-ai/cmd/manager/img/rootfs.cpio.gz
+MANAGER_QEMU_SEV_ID=sev0
+MANAGER_QEMU_SEV_CBITPOS=51
+MANAGER_QEMU_SEV_REDUCED_PHYS_BITS=1
+MANAGER_QEMU_HOST_DATA=
+MANAGER_QEMU_VSOCK_ID=vhost-vsock-pci0
+MANAGER_QEMU_VSOCK_GUEST_CID=3
+MANAGER_QEMU_VSOCK_VNC=0
+MANAGER_QEMU_BIN_PATH=qemu-system-x86_64
+MANAGER_QEMU_USE_SUDO=true
+MANAGER_QEMU_ENABLE_SEV=false
+MANAGER_QEMU_ENABLE_SEV_SNP=false
+MANAGER_QEMU_ENABLE_KVM=true
+MANAGER_QEMU_MACHINE=q35
+MANAGER_QEMU_CPU=EPYC
+MANAGER_QEMU_SMP_COUNT=4
+MANAGER_QEMU_SMP_MAXCPUS=16
+MANAGER_QEMU_MEM_ID=ram1
+MANAGER_QEMU_KERNEL_HASH=false
+MANAGER_QEMU_NO_GRAPHIC=true
+MANAGER_QEMU_MONITOR=pty
+MANAGER_QEMU_HOST_FWD_RANGE=6100-6200

codecov.yml

@@ -0,0 +1,13 @@
+# Copyright (c) Ultraviolet
+# SPDX-License-Identifier: Apache-2.0
+
+coverage:
+  ignore:
+    - "test/*"
+    - "cmd/**"
+    - "**/mocks/**"
+    - "mocks/**"
+    - "**/*.pb.go"
+    - "**/logging.go"
+    - "**/metrics.go"
+    - "**/tracing.go"

go.mod

@@ -1,12 +1,12 @@
 module github.com/ultravioletrs/cocos
 
-go 1.22.4
+go 1.23.0
 
 require (
 	github.com/absmach/magistrala v0.14.1-0.20240709113739-04c359462746
-	github.com/caarlos0/env/v11 v11.1.0
+	github.com/caarlos0/env/v11 v11.2.2
 	github.com/cenkalti/backoff/v4 v4.3.0
-	github.com/digitalocean/go-libvirt v0.0.0-20240709142323-d8406205c752
+	github.com/fatih/color v1.18.0
 	github.com/go-kit/kit v0.13.0
 	github.com/gofrs/uuid v4.4.0+incompatible
 	github.com/google/go-sev-guest v0.11.1
@@ -14,45 +14,51 @@ require (
 	github.com/spf13/cobra v1.8.1
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.9.0
-	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.53.0
-	go.opentelemetry.io/otel/trace v1.28.0
-	golang.org/x/crypto v0.25.0
-	golang.org/x/sync v0.7.0
-	google.golang.org/grpc v1.65.0
-	google.golang.org/protobuf v1.34.2
+	github.com/virtee/sev-snp-measure-go v0.0.0-20240530153610-e6e8dc9b6877
+	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.57.0
+	go.opentelemetry.io/otel/trace v1.32.0
+	golang.org/x/crypto v0.29.0
+	golang.org/x/sync v0.9.0
+	google.golang.org/grpc v1.68.0
+	google.golang.org/protobuf v1.35.2
 )
 
 require (
 	github.com/Microsoft/go-winio v0.6.1 // indirect
+	github.com/containerd/log v0.1.0 // indirect
 	github.com/distribution/reference v0.6.0 // indirect
 	github.com/docker/go-connections v0.5.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/felixge/httpsnoop v1.0.4 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/moby/docker-image-spec v1.3.1 // indirect
+	github.com/morikuni/aec v1.0.0 // indirect
 	github.com/opencontainers/go-digest v1.0.0 // indirect
 	github.com/opencontainers/image-spec v1.1.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.53.0 // indirect
-	go.opentelemetry.io/otel v1.28.0 // indirect
+	go.opentelemetry.io/otel v1.32.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.28.0 // indirect
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.28.0 // indirect
 	go.opentelemetry.io/otel/sdk v1.28.0 // indirect
 	golang.org/x/mod v0.19.0 // indirect
 	golang.org/x/tools v0.23.0 // indirect
+	gotest.tools/v3 v3.5.1 // indirect
 )
 
 require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
-	github.com/docker/docker v27.1.0+incompatible
+	github.com/docker/docker v27.3.1+incompatible
 	github.com/go-kit/log v0.2.1 // indirect
 	github.com/go-logfmt/logfmt v0.6.0 // indirect
 	github.com/go-logr/logr v1.4.2 // indirect
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/go-configfs-tsm v0.2.2 // indirect
-	github.com/google/logger v1.1.1 // indirect
+	github.com/google/logger v1.1.1
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
@@ -65,14 +71,16 @@ require (
 	github.com/prometheus/common v0.52.2 // indirect
 	github.com/prometheus/procfs v0.13.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
-	go.opentelemetry.io/otel/metric v1.28.0 // indirect
+	go.opentelemetry.io/otel/metric v1.32.0 // indirect
 	go.opentelemetry.io/proto/otlp v1.3.1 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	golang.org/x/net v0.27.0 // indirect
-	golang.org/x/sys v0.22.0 // indirect
-	golang.org/x/term v0.22.0
-	golang.org/x/text v0.16.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20240701130421-f6361c86f094 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20240701130421-f6361c86f094 // indirect
+	golang.org/x/net v0.30.0 // indirect
+	golang.org/x/sys v0.27.0 // indirect
+	golang.org/x/term v0.26.0
+	golang.org/x/text v0.20.0 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20240903143218-8af14fe29dc1 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20241104194629-dd2ea8efbc28 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
+
+replace github.com/virtee/sev-snp-measure-go => github.com/sammyoina/sev-snp-measure-go v0.0.0-20241107163739-38915ab517c7

go.sum

@@ -1,3 +1,5 @@
+github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 h1:L/gRVlceqvL25UVaW/CKtUDjefjrs0SPonmDGUVOYP0=
+github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
 github.com/Microsoft/go-winio v0.6.1 h1:9/kr64B9VUZrLm5YYwbGtUJnMgqWVOdUAXu6Migciow=
 github.com/Microsoft/go-winio v0.6.1/go.mod h1:LRdKpFKfdobln8UmuiYcKPot9D2v6svN5+sAH+4kjUM=
 github.com/VividCortex/gohistogram v1.0.0 h1:6+hBz+qvs0JOrrNhhmR7lFxo5sINxBCGXrdtl/UvroE=
@@ -6,25 +8,27 @@ github.com/absmach/magistrala v0.14.1-0.20240709113739-04c359462746 h1:Tj567KeGV
 github.com/absmach/magistrala v0.14.1-0.20240709113739-04c359462746/go.mod h1:CIx3OsPFc4doJZmBWSA6LNWefcznKv9c3cLOxNxL4q4=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
-github.com/caarlos0/env/v11 v11.1.0 h1:a5qZqieE9ZfzdvbbdhTalRrHT5vu/4V1/ad1Ka6frhI=
-github.com/caarlos0/env/v11 v11.1.0/go.mod h1:LwgkYk1kDvfGpHthrWWLof3Ny7PezzFwS4QrsJdHTMo=
+github.com/caarlos0/env/v11 v11.2.2 h1:95fApNrUyueipoZN/EhA8mMxiNxrBwDa+oAZrMWl3Kg=
+github.com/caarlos0/env/v11 v11.2.2/go.mod h1:JBfcdeQiBoI3Zh1QRAWfe+tpiNTmDtcCj/hHHHMx0vc=
 github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
 github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
 github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
 github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
+github.com/containerd/log v0.1.0 h1:TCJt7ioM2cr/tfR8GPbGf9/VRAX8D2B4PjzCpfX540I=
+github.com/containerd/log v0.1.0/go.mod h1:VRRf09a7mHDIRezVKTRCrOq78v577GXq3bSa3EhrzVo=
 github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/digitalocean/go-libvirt v0.0.0-20240709142323-d8406205c752 h1:NI7XEcHzWVvBfVjSVK6Qk4wmrUfoyQxCNpBjrHelZFk=
-github.com/digitalocean/go-libvirt v0.0.0-20240709142323-d8406205c752/go.mod h1:/Ok8PA2qi/ve0Py38+oL+VxoYmlowigYRyLEODRYdgc=
 github.com/distribution/reference v0.6.0 h1:0IXCQ5g4/QMHHkarYzh5l+u8T3t73zM5QvfrDyIgxBk=
 github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E=
-github.com/docker/docker v27.1.0+incompatible h1:rEHVQc4GZ0MIQKifQPHSFGV/dVgaZafgRf8fCPtDYBs=
-github.com/docker/docker v27.1.0+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
+github.com/docker/docker v27.3.1+incompatible h1:KttF0XoteNTicmUtBO0L2tP+J7FGRFTjaEF4k6WdhfI=
+github.com/docker/docker v27.3.1+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk=
 github.com/docker/go-connections v0.5.0 h1:USnMq7hx7gwdVZq1L49hLXaFtUdTADjXGp+uj1Br63c=
 github.com/docker/go-connections v0.5.0/go.mod h1:ov60Kzw0kKElRwhNs9UlUHAE/F9Fe6GLaXnqyDdmEXc=
 github.com/docker/go-units v0.5.0 h1:69rxXcBk27SvSaaxTtLh/8llcHD8vYHT7WSdRZ/jvr4=
 github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
+github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
+github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
 github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg=
 github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/go-kit/kit v0.13.0 h1:OoneCcHKHQ03LfBpoQCUfCluwd2Vt3ohz+kvbJneZAU=
@@ -65,12 +69,21 @@ github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
+github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
+github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
+github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
+github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
+github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/mdlayher/socket v0.4.1 h1:eM9y2/jlbs1M615oshPQOHZzj6R6wMT7bX5NPiQvn2U=
 github.com/mdlayher/socket v0.4.1/go.mod h1:cAqeGjoufqdxWkD7DkpyS+wcefOtmu5OQ8KuoJGIReA=
 github.com/mdlayher/vsock v1.2.1 h1:pC1mTJTvjo1r9n9fbm7S1j04rCgCzhCOS5DY0zqHlnQ=
 github.com/mdlayher/vsock v1.2.1/go.mod h1:NRfCibel++DgeMD8z/hP+PPTjlNJsdPOmxcnENvE+SE=
 github.com/moby/docker-image-spec v1.3.1 h1:jMKff3w6PgbfSa69GfNg+zN/XLhfXJGnEx3Nl2EsFP0=
 github.com/moby/docker-image-spec v1.3.1/go.mod h1:eKmb5VW8vQEh/BAr2yvVNvuiJuY6UIocYsFu/DxxRpo=
+github.com/moby/term v0.5.0 h1:xt8Q1nalod/v7BqbG21f8mQPqH+xAaC9C3N3wfWbVP0=
+github.com/moby/term v0.5.0/go.mod h1:8FzsFHVUBGZdbDsJw/ot+X+d5HLUbvklYLJ9uGfcI3Y=
+github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A=
+github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc=
 github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/image-spec v1.1.0 h1:8SG7/vwALn54lVB/0yZ/MMwhFrPYtpEHQb2IpWsCzug=
@@ -92,6 +105,10 @@ github.com/prometheus/procfs v0.13.0/go.mod h1:cd4PFCR54QLnGKPaKGA6l+cfuNXtht43Z
 github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
 github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/sammyoina/sev-snp-measure-go v0.0.0-20241107163739-38915ab517c7 h1:g+a3hLU4pl41mhP+CQu3+bDhL4HSfrnPI1BvLjKfD8Y=
+github.com/sammyoina/sev-snp-measure-go v0.0.0-20241107163739-38915ab517c7/go.mod h1:dEkBe8JnxU5itNjZDEQINFd7f7l4DtjfqRuzPQcit4w=
+github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
+github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
 github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM=
 github.com/spf13/cobra v1.8.1/go.mod h1:wHxEcudfqmLYa8iTfL+OuZPbBZkmvliBWKIezN3kD9Y=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
@@ -102,22 +119,22 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.53.0 h1:9G6E0TXzGFVfTnawRzrPl83iHOAV7L8NJiR8RSGYV1g=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.53.0/go.mod h1:azvtTADFQJA8mX80jIH/akaE7h+dbm/sVuaHqN13w74=
+go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.57.0 h1:qtFISDHKolvIxzSs0gIaiPUPR0Cucb0F2coHC7ZLdps=
+go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.57.0/go.mod h1:Y+Pop1Q6hCOnETWTW4NROK/q1hv50hM7yDaUTjG8lp8=
 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.53.0 h1:4K4tsIXefpVJtvA/8srF4V4y0akAoPHkIslgAkjixJA=
 go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.53.0/go.mod h1:jjdQuTGVsXV4vSs+CJ2qYDeDPf9yIJV23qlIzBm73Vg=
-go.opentelemetry.io/otel v1.28.0 h1:/SqNcYk+idO0CxKEUOtKQClMK/MimZihKYMruSMViUo=
-go.opentelemetry.io/otel v1.28.0/go.mod h1:q68ijF8Fc8CnMHKyzqL6akLO46ePnjkgfIMIjUIX9z4=
+go.opentelemetry.io/otel v1.32.0 h1:WnBN+Xjcteh0zdk01SVqV55d/m62NJLJdIyb4y/WO5U=
+go.opentelemetry.io/otel v1.32.0/go.mod h1:00DCVSB0RQcnzlwyTfqtxSm+DRr9hpYrHjNGiBHVQIg=
 go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.28.0 h1:3Q/xZUyC1BBkualc9ROb4G8qkH90LXEIICcs5zv1OYY=
 go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.28.0/go.mod h1:s75jGIWA9OfCMzF0xr+ZgfrB5FEbbV7UuYo32ahUiFI=
 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.28.0 h1:j9+03ymgYhPKmeXGk5Zu+cIZOlVzd9Zv7QIiyItjFBU=
 go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.28.0/go.mod h1:Y5+XiUG4Emn1hTfciPzGPJaSI+RpDts6BnCIir0SLqk=
-go.opentelemetry.io/otel/metric v1.28.0 h1:f0HGvSl1KRAU1DLgLGFjrwVyismPlnuU6JD6bOeuA5Q=
-go.opentelemetry.io/otel/metric v1.28.0/go.mod h1:Fb1eVBFZmLVTMb6PPohq3TO9IIhUisDsbJoL/+uQW4s=
+go.opentelemetry.io/otel/metric v1.32.0 h1:xV2umtmNcThh2/a/aCP+h64Xx5wsj8qqnkYZktzNa0M=
+go.opentelemetry.io/otel/metric v1.32.0/go.mod h1:jH7CIbbK6SH2V2wE16W05BHCtIDzauciCRLoc/SyMv8=
 go.opentelemetry.io/otel/sdk v1.28.0 h1:b9d7hIry8yZsgtbmM0DKyPWMMUMlK9NEKuIG4aBqWyE=
 go.opentelemetry.io/otel/sdk v1.28.0/go.mod h1:oYj7ClPUA7Iw3m+r7GeEjz0qckQRJK2B8zjcZEfu7Pg=
-go.opentelemetry.io/otel/trace v1.28.0 h1:GhQ9cUuQGmNDd5BTCP2dAvv75RdMxEfTmYejp+lkx9g=
-go.opentelemetry.io/otel/trace v1.28.0/go.mod h1:jPyXzNPg6da9+38HEwElrQiHlVMTnVfM3/yv2OlIHaI=
+go.opentelemetry.io/otel/trace v1.32.0 h1:WIC9mYrXf8TmY/EXuULKc8hR17vE+Hjv2cssQDe03fM=
+go.opentelemetry.io/otel/trace v1.32.0/go.mod h1:+i4rkvCraA+tG6AzwloGaCtkx53Fa+L+V8e9a7YvhT8=
 go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0=
 go.opentelemetry.io/proto/otlp v1.3.1/go.mod h1:0X1WI4de4ZsLrrJNLAQbFeLCm3T7yBkR0XqQ7niQU+8=
 go.uber.org/multierr v1.11.0 h1:blXXJkSxSSfBVBlC76pxqeO+LN3aDfLQo+309xJstO0=
@@ -125,8 +142,8 @@ go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN8
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
-golang.org/x/crypto v0.25.0 h1:ypSNr+bnYL2YhwoMt2zPxHFmbAN1KZs/njMG3hxUp30=
-golang.org/x/crypto v0.25.0/go.mod h1:T+wALwcMOSE0kXgUAnPAHqTLW+XHgcELELW8VaDgm/M=
+golang.org/x/crypto v0.29.0 h1:L5SG1JTTXupVV3n6sUqMTeWbjAyfPwoda2DLX8J8FrQ=
+golang.org/x/crypto v0.29.0/go.mod h1:+F4F4N5hv6v38hfeYwTdx20oUvLLc+QfrE9Ax9HtgRg=
 golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.19.0 h1:fEdghXQSo20giMthA7cd28ZC+jts4amQ3YMXiP5oMQ8=
@@ -135,27 +152,29 @@ golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
-golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ=
-golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE=
-golang.org/x/net v0.27.0 h1:5K3Njcw06/l2y9vpGCSdcxWOYHOUk3dVNGDXN+FvAys=
-golang.org/x/net v0.27.0/go.mod h1:dDi0PyhWNoiUOrAS8uXv/vnScO4wnHQO4mj9fn/RytE=
+golang.org/x/net v0.30.0 h1:AcW1SDZMkb8IpzCdQUaIq2sP4sZ4zw+55h6ynffypl4=
+golang.org/x/net v0.30.0/go.mod h1:2wGyMJ5iFasEhkwi13ChkO/t1ECNC4X4eBKkVFyYFlU=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M=
-golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
+golang.org/x/sync v0.9.0 h1:fEo0HyrW1GIgZdpbhCRO0PkJajUS5H9IFUztCgEo2jQ=
+golang.org/x/sync v0.9.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20210426230700-d19ff857e887/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI=
-golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/term v0.22.0 h1:BbsgPEJULsl2fV/AT3v15Mjva5yXKQDyKf+TbDz7QJk=
-golang.org/x/term v0.22.0/go.mod h1:F3qCibpT5AMpCRfhfT53vVJwhLtIVHhB9XDjfFvnMI4=
+golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.27.0 h1:wBqf8DvsY9Y/2P8gAfPDEYNuS30J4lPHJxXSb/nJZ+s=
+golang.org/x/sys v0.27.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/term v0.26.0 h1:WEQa6V3Gja/BhNxg540hBip/kkaYtRg3cxg4oXSw4AU=
+golang.org/x/term v0.26.0/go.mod h1:Si5m1o57C5nBNQo5z1iq+XDijt21BDBDp2bK0QI8e3E=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
-golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
-golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=
+golang.org/x/text v0.20.0 h1:gK/Kv2otX8gz+wn7Rmb3vT96ZwuoxnQlY+HlJVj7Qug=
+golang.org/x/text v0.20.0/go.mod h1:D4IsuqiFMhST5bX19pQ9ikHC2GsaKyk/oF+pn3ducp4=
+golang.org/x/time v0.5.0 h1:o7cqy6amK/52YcAKIPlM3a+Fpj35zvRj2TP+e1xFSfk=
+golang.org/x/time v0.5.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=
@@ -166,16 +185,18 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-google.golang.org/genproto/googleapis/api v0.0.0-20240701130421-f6361c86f094 h1:0+ozOGcrp+Y8Aq8TLNN2Aliibms5LEzsq99ZZmAGYm0=
-google.golang.org/genproto/googleapis/api v0.0.0-20240701130421-f6361c86f094/go.mod h1:fJ/e3If/Q67Mj99hin0hMhiNyCRmt6BQ2aWIJshUSJw=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20240701130421-f6361c86f094 h1:BwIjyKYGsK9dMCBOorzRri8MQwmi7mT9rGHsCEinZkA=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20240701130421-f6361c86f094/go.mod h1:Ue6ibwXGpU+dqIcODieyLOcgj7z8+IcskoNIgZxtrFY=
-google.golang.org/grpc v1.65.0 h1:bs/cUb4lp1G5iImFFd3u5ixQzweKizoZJAwBNLR42lc=
-google.golang.org/grpc v1.65.0/go.mod h1:WgYC2ypjlB0EiQi6wdKixMqukr6lBc0Vo+oOgjrM5ZQ=
-google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg=
-google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw=
+google.golang.org/genproto/googleapis/api v0.0.0-20240903143218-8af14fe29dc1 h1:hjSy6tcFQZ171igDaN5QHOw2n6vx40juYbC/x67CEhc=
+google.golang.org/genproto/googleapis/api v0.0.0-20240903143218-8af14fe29dc1/go.mod h1:qpvKtACPCQhAdu3PyQgV4l3LMXZEtft7y8QcarRsp9I=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20241104194629-dd2ea8efbc28 h1:XVhgTWWV3kGQlwJHR3upFWZeTsei6Oks1apkZSeonIE=
+google.golang.org/genproto/googleapis/rpc v0.0.0-20241104194629-dd2ea8efbc28/go.mod h1:GX3210XPVPUjJbTUbvwI8f2IpZDMZuPJWDzDuebbviI=
+google.golang.org/grpc v1.68.0 h1:aHQeeJbo8zAkAa3pRzrVjZlbz6uSfeOXlJNQM0RAbz0=
+google.golang.org/grpc v1.68.0/go.mod h1:fmSPC5AsjSBCK54MyHRx48kpOti1/jRfOlwEWywNjWA=
+google.golang.org/protobuf v1.35.2 h1:8Ar7bF+apOIoThw1EdZl0p1oWvMqTHmpA2fRTyZO8io=
+google.golang.org/protobuf v1.35.2/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
 gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gotest.tools/v3 v3.5.1 h1:EENdUnS3pdur5nybKYIh2Vfgc8IUNBjxDPSjtiJcOzU=
+gotest.tools/v3 v3.5.1/go.mod h1:isy3WKz7GK6uNw/sbHzfKBLvlvXwUyV06n6brMxxopU=

hal/linux/README.md

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

hal/linux/configs/cocos_defconfig

@@ -23,18 +23,16 @@ BR2_ROOTFS_POST_IMAGE_SCRIPT="$(BR2_EXTERNAL_COCOS_PATH)/board/cocos/post-image.
 BR2_ROOTFS_POST_SCRIPT_ARGS="$(BR2_DEFCONFIG)"
 
 # Linux headers same as kernel
-BR2_PACKAGE_HOST_LINUX_HEADERS_CUSTOM_6_6=y
-BR2_TOOLCHAIN_HEADERS_AT_LEAST_6_5=y
-BR2_TOOLCHAIN_HEADERS_AT_LEAST_6_6=y
+BR2_PACKAGE_HOST_LINUX_HEADERS_CUSTOM_6_11=y
 BR2_TOOLCHAIN_HEADERS_LATEST=y
-BR2_TOOLCHAIN_HEADERS_AT_LEAST="6.6"
+BR2_TOOLCHAIN_HEADERS_AT_LEAST="6.12-rc6"
 
 # Kernel
 BR2_LINUX_KERNEL=y
 BR2_LINUX_KERNEL_CUSTOM_GIT=y
 BR2_LINUX_KERNEL_CUSTOM_REPO_URL="https://github.com/torvalds/linux.git"
-BR2_LINUX_KERNEL_CUSTOM_REPO_VERSION="v6.6"
-BR2_LINUX_KERNEL_VERSION="v6.6"
+BR2_LINUX_KERNEL_CUSTOM_REPO_VERSION="v6.12-rc6"
+BR2_LINUX_KERNEL_VERSION="v6.12-rc6"
 BR2_LINUX_KERNEL_PATCH=""
 BR2_LINUX_KERNEL_USE_CUSTOM_CONFIG=y
 BR2_LINUX_KERNEL_CUSTOM_CONFIG_FILE="$(BR2_EXTERNAL_COCOS_PATH)/board/cocos/linux.config"

init/systemd/agent_setup.sh

@@ -1,6 +1,6 @@
 #!/bin/sh
 
-MOUNT_POINT="/cocos"
+WORK_DIR="/cocos"
 
 # IFACES are all network interfaces excluding lo (LOOPBACK) and sit interfaces 
 IFACES=$(ip link show | grep -vE 'LOOPBACK|sit*' | awk -F': ' '{print $2}')
@@ -18,8 +18,9 @@ for IFACE in $IFACES; do
     fi
 done
 
-if [ ! -d "$MOUNT_POINT" ]; then
-    mkdir -p $MOUNT_POINT
+if [ ! -d "$WORK_DIR" ]; then
+    mkdir -p $WORK_DIR
 fi
 
-mount -t ramfs ramfs $MOUNT_POINT
+# Resize the root file system to 100%
+mount -o remount,size=100% /

init/systemd/cocos-manager.service

@@ -0,0 +1,15 @@
+[Unit]
+Description=Cocos Manager
+After=network.target
+
+[Service]
+ExecStart=cocos-manager
+Restart=always
+RestartSec=5s
+EnvironmentFile=/etc/cocos/cocos-manager.env
+StandardOutput=syslog
+StandardError=syslog
+SyslogIdentifier=cocos-manager
+
+[Install]
+WantedBy=multi-user.target

internal/cmd.go

@@ -1,75 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package internal
-
-import (
-	"bytes"
-	"fmt"
-	"io"
-	"os"
-	"os/exec"
-	"strings"
-)
-
-// ExeShCmdStdout executes a shell command capturing the standard output.
-func ExeShCmdStdout(command string, args ...string) (string, error) {
-	var stdoutBuf, stderrBuf bytes.Buffer
-
-	cmd := exec.Command(command, args...)
-
-	// Capture stdout and stderr using buffers
-	cmd.Stdout = io.MultiWriter(&stdoutBuf, os.Stdout)
-	cmd.Stderr = io.MultiWriter(&stderrBuf, os.Stderr)
-
-	err := cmd.Run()
-	if err != nil {
-		return "", fmt.Errorf("error executing command '%s': %s", cmd.String(), err)
-	}
-
-	return stdoutBuf.String(), nil
-}
-
-// ExtractCmdAndArgs extracts the command and its arguments from the output string.
-func ExtractCmdAndArgs(cmdLine string, sudo bool) (string, []string) {
-	lines := strings.Split(cmdLine, "\n")
-	if len(lines) == 0 {
-		return "", nil
-	}
-
-	parts := strings.Fields(lines[0])
-	if len(parts) == 0 {
-		return "", nil
-	}
-
-	if sudo {
-		parts = append([]string{"sudo"}, parts...)
-	}
-
-	cmd := parts[0]
-	args := parts[1:]
-
-	return cmd, args
-}
-
-// RunCmdOutput runs the specified command and returns its standard output as a string.
-func RunCmdOutput(command string, args ...string) (string, error) {
-	cmd := exec.Command(command, args...)
-
-	output, err := cmd.Output()
-	if err != nil {
-		return "", fmt.Errorf("error executing command '%s': %s", cmd.String(), err)
-	}
-
-	return string(output), nil
-}
-
-// RunCmdStart starts the specified command and returns the *exec.Cmd for the running process.
-func RunCmdStart(command string, args ...string) (*exec.Cmd, error) {
-	cmd := exec.Command(command, args...)
-
-	if err := cmd.Start(); err != nil {
-		return nil, fmt.Errorf("error starting command '%s': %s", cmd.String(), err)
-	}
-
-	return cmd, nil
-}

internal/file_test.go

@@ -0,0 +1,157 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package internal
+
+import (
+	"bytes"
+	"encoding/hex"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+func TestCopyFile(t *testing.T) {
+	tempDir, err := os.MkdirTemp("", "copyfile_test")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	srcPath := filepath.Join(tempDir, "source.txt")
+	content := []byte("Hello, World!")
+	if err := os.WriteFile(srcPath, content, 0o644); err != nil {
+		t.Fatalf("Failed to create source file: %v", err)
+	}
+
+	dstPath := filepath.Join(tempDir, "destination.txt")
+	if err := CopyFile(srcPath, dstPath); err != nil {
+		t.Fatalf("CopyFile failed: %v", err)
+	}
+
+	copiedContent, err := os.ReadFile(dstPath)
+	if err != nil {
+		t.Fatalf("Failed to read destination file: %v", err)
+	}
+
+	if !bytes.Equal(content, copiedContent) {
+		t.Errorf("Copied content does not match original. Got %s, want %s", copiedContent, content)
+	}
+}
+
+func TestCopyFile_NonExistentSource(t *testing.T) {
+	err := CopyFile("nonexistent.txt", "destination.txt")
+	if err == nil {
+		t.Error("CopyFile did not return an error for a nonexistent source file")
+	}
+}
+
+func TestDeleteFilesInDir(t *testing.T) {
+	tempDir, err := os.MkdirTemp("", "deletefiles_test")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	filenames := []string{"file1.txt", "file2.txt", "file3.txt"}
+	for _, filename := range filenames {
+		filepath := filepath.Join(tempDir, filename)
+		if err := os.WriteFile(filepath, []byte("test"), 0o644); err != nil {
+			t.Fatalf("Failed to create test file: %v", err)
+		}
+	}
+
+	if err := DeleteFilesInDir(tempDir); err != nil {
+		t.Fatalf("DeleteFilesInDir failed: %v", err)
+	}
+
+	remainingFiles, err := os.ReadDir(tempDir)
+	if err != nil {
+		t.Fatalf("Failed to read directory: %v", err)
+	}
+
+	if len(remainingFiles) != 0 {
+		t.Errorf("Directory not empty after deletion. %d files remain", len(remainingFiles))
+	}
+}
+
+func TestChecksum(t *testing.T) {
+	tempDir, err := os.MkdirTemp("", "checksum_test")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	filePath := filepath.Join(tempDir, "test.txt")
+	content := []byte("Hello, World!")
+	if err := os.WriteFile(filePath, content, 0o644); err != nil {
+		t.Fatalf("Failed to create test file: %v", err)
+	}
+
+	checksum, err := Checksum(filePath)
+	if err != nil {
+		t.Fatalf("Checksum failed: %v", err)
+	}
+
+	expectedChecksum, _ := hex.DecodeString("1af17a664e3fa8e419b8ba05c2a173169df76162a5a286e0c405b460d478f7ef")
+	if !bytes.Equal(checksum, expectedChecksum) {
+		t.Errorf("File checksum mismatch. Got %x, want %x", checksum, expectedChecksum)
+	}
+
+	dirPath := filepath.Join(tempDir, "testdir")
+	if err := os.Mkdir(dirPath, 0o755); err != nil {
+		t.Fatalf("Failed to create test directory: %v", err)
+	}
+	if err := os.WriteFile(filepath.Join(dirPath, "file1.txt"), []byte("File 1"), 0o644); err != nil {
+		t.Fatalf("Failed to create test file: %v", err)
+	}
+	if err := os.WriteFile(filepath.Join(dirPath, "file2.txt"), []byte("File 2"), 0o644); err != nil {
+		t.Fatalf("Failed to create test file: %v", err)
+	}
+
+	dirChecksum, err := Checksum(dirPath)
+	if err != nil {
+		t.Fatalf("Directory Checksum failed: %v", err)
+	}
+
+	if len(dirChecksum) != 32 { // SHA3-256 produces a 32-byte hash
+		t.Errorf("Unexpected directory checksum length. Got %d bytes, want 32 bytes", len(dirChecksum))
+	}
+}
+
+func TestChecksum_NonExistentFile(t *testing.T) {
+	_, err := Checksum("nonexistent.txt")
+	if err == nil {
+		t.Error("Checksum did not return an error for a nonexistent file")
+	}
+}
+
+func TestChecksumHex(t *testing.T) {
+	tempFile, err := os.CreateTemp("", "checksumhex_test")
+	if err != nil {
+		t.Fatalf("Failed to create temp file: %v", err)
+	}
+	defer os.Remove(tempFile.Name())
+
+	content := []byte("Hello, World!")
+	if _, err := tempFile.Write(content); err != nil {
+		t.Fatalf("Failed to write to test file: %v", err)
+	}
+	tempFile.Close()
+
+	checksumHex, err := ChecksumHex(tempFile.Name())
+	if err != nil {
+		t.Fatalf("ChecksumHex failed: %v", err)
+	}
+
+	expectedChecksumHex := "1af17a664e3fa8e419b8ba05c2a173169df76162a5a286e0c405b460d478f7ef"
+	if checksumHex != expectedChecksumHex {
+		t.Errorf("ChecksumHex mismatch. Got %s, want %s", checksumHex, expectedChecksumHex)
+	}
+}
+
+func TestChecksumHex_NonExistentFile(t *testing.T) {
+	_, err := ChecksumHex("nonexistent.txt")
+	if err == nil {
+		t.Error("ChecksumHex did not return an error for a nonexistent file")
+	}
+}

internal/libvirt/connect.go

@@ -1,45 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package libvirt
-
-import (
-	"fmt"
-	"log"
-	"log/slog"
-	"net"
-	"time"
-
-	"github.com/digitalocean/go-libvirt"
-)
-
-func Connect(logger *slog.Logger) *libvirt.Libvirt {
-	// This dials libvirt on the local machine, but you can substitute the first
-	// two parameters with "tcp", "<ip address>:<port>" to connect to libvirt on
-	// a remote machine.
-	c, err := net.DialTimeout("unix", "/var/run/libvirt/libvirt-sock", 2*time.Second)
-	if err != nil {
-		log.Fatalf("failed to dial libvirt: %v", err)
-	}
-
-	l := libvirt.New(c)
-	if err := l.Connect(); err != nil {
-		log.Fatalf("failed to connect: %v", err)
-	}
-
-	v, err := l.Version()
-	if err != nil {
-		logger.Error(fmt.Sprintf("failed to retrieve libvirt version: %v", err))
-	}
-	logger.Info(fmt.Sprintf("Retrieved libvirt version: %s", v))
-
-	domains, err := l.Domains()
-	if err != nil {
-		logger.Error(fmt.Sprintf("failed to retrieve domains: %v", err))
-	}
-
-	for _, d := range domains {
-		logger.Info(fmt.Sprintf("%d\t%s\t%x\n", d.ID, d.Name, d.UUID))
-	}
-
-	return l
-}

internal/libvirt/libvirt.go

@@ -1,167 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package libvirt
-
-import (
-	"context"
-	"errors"
-	"fmt"
-	"os"
-	"regexp"
-	"strings"
-
-	golibvirt "github.com/digitalocean/go-libvirt"
-)
-
-var re = regexp.MustCompile(`'([^']*)'`)
-
-func CreateDomain(ctx context.Context, libvirt *golibvirt.Libvirt, poolXML, volXML, domXML string) (string, error) {
-	wd, err := os.Getwd()
-	if err != nil {
-		return "", err
-	}
-
-	poolStr, err := readXMLFile(poolXML, "pool.xml")
-	if err != nil {
-		return "", err
-	}
-	poolStr = replaceSubstring(poolStr, "./", wd+"/")
-
-	volStr, err := readXMLFile(volXML, "vol.xml")
-	if err != nil {
-		return "", err
-	}
-	volStr = replaceSubstring(volStr, "./", wd+"/")
-
-	domStr, err := readXMLFile(domXML, "dom.xml")
-	if err != nil {
-		return "", err
-	}
-	domStr = replaceSubstring(domStr, "./", wd+"/")
-
-	dom, err := createDomain(libvirt, poolStr, volStr, domStr)
-	if err != nil {
-		return "", fmt.Errorf("failed to create domain: %s", err)
-	}
-
-	return dom.Name, nil
-}
-
-func createDomain(libvirtConn *golibvirt.Libvirt, poolXML, volXML, domXML string) (golibvirt.Domain, error) {
-	pool, err := libvirtConn.StoragePoolCreateXML(poolXML, 0)
-	_ = pool
-	if err != nil {
-		lvErr := err.(golibvirt.Error)
-		if lvErr.Code == 9 {
-			name, err := entityName(lvErr.Message)
-			if err != nil {
-				return golibvirt.Domain{}, err
-			}
-			pool, err = libvirtConn.StoragePoolLookupByName(name)
-			if err != nil {
-				return golibvirt.Domain{}, err
-			}
-
-			goto pool_exists
-		}
-
-		return golibvirt.Domain{}, err
-	}
-pool_exists:
-
-	_, err = libvirtConn.StorageVolCreateXML(pool, volXML, 0)
-	if err != nil {
-		lvErr := err.(golibvirt.Error)
-		if lvErr.Code == 90 {
-			name, err := entityName(lvErr.Message)
-			if err != nil {
-				return golibvirt.Domain{}, err
-			}
-			_, err = libvirtConn.StorageVolLookupByName(pool, name)
-			if err != nil {
-				return golibvirt.Domain{}, err
-			}
-
-			goto vol_exists
-		}
-
-		return golibvirt.Domain{}, err
-	}
-
-vol_exists:
-
-	dom, err := libvirtConn.DomainDefineXMLFlags(domXML, 0)
-	if err != nil {
-		return golibvirt.Domain{}, err
-	}
-
-	err = libvirtConn.DomainCreate(dom)
-	if err != nil {
-		lvErr := err.(golibvirt.Error)
-		if lvErr.Code == 55 {
-			return dom, nil
-		}
-
-		return golibvirt.Domain{}, err
-	}
-
-	// extra flags; not used yet, so callers should always pass 0
-	current, err := libvirtConn.DomainSnapshotCurrent(dom, 0)
-	if err != nil {
-		lvErr := err.(golibvirt.Error)
-		if lvErr.Code == 72 {
-			return dom, nil
-		}
-
-		return golibvirt.Domain{}, err
-	}
-
-	err = libvirtConn.DomainRevertToSnapshot(current, uint32(golibvirt.DomainSnapshotRevertRunning))
-	if err != nil {
-		return golibvirt.Domain{}, err
-	}
-
-	return dom, nil
-}
-
-func entityName(msg string) (string, error) {
-	match := re.FindStringSubmatch(msg)
-	if len(match) < 1 {
-		return "", errors.New("entity not found")
-	}
-
-	return match[1], nil
-}
-
-func readXMLFile(filename, defaultFilename string) (string, error) {
-	if filename == "" {
-		filename = "./xml/" + defaultFilename
-	}
-
-	xmlBytes, err := os.ReadFile(filename)
-	if err != nil {
-		return "", fmt.Errorf("failed to read XML file: %s", err)
-	}
-
-	return string(xmlBytes), nil
-}
-
-func replaceSubstring(xml, substring, replacement string) string {
-	// Split the file text into lines
-	lines := strings.Split(xml, "\n")
-
-	// Create a variable to hold the resulting string
-	var result strings.Builder
-
-	// Iterate over each line
-	for _, line := range lines {
-		// Replace the substring with the replacement
-		newLine := strings.ReplaceAll(line, substring, replacement)
-
-		// Append the modified line to the resulting string
-		result.WriteString(newLine)
-		result.WriteString("\n")
-	}
-
-	return result.String()
-}

internal/logger/protohandler.go

@@ -7,7 +7,7 @@ import (
 	"io"
 	"log/slog"
 
-	"github.com/ultravioletrs/cocos/pkg/manager"
+	"github.com/ultravioletrs/cocos/agent/events"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -15,21 +15,24 @@ import (
 var _ slog.Handler = (*handler)(nil)
 
 type handler struct {
-	opts slog.HandlerOptions
-	w    io.Writer
+	opts  slog.HandlerOptions
+	w     io.Writer
+	cmpID string
 }
 
-func NewProtoHandler(w io.Writer, opts *slog.HandlerOptions) slog.Handler {
+func NewProtoHandler(conn io.Writer, opts *slog.HandlerOptions, cmpID string) slog.Handler {
 	if opts == nil {
 		opts = &slog.HandlerOptions{}
 	}
-	return &handler{
-		opts: *opts,
-		w:    w,
+	h := &handler{
+		opts:  *opts,
+		w:     conn,
+		cmpID: cmpID,
 	}
+
+	return h
 }
 
-// Enabled implements slog.Handler.
 func (h *handler) Enabled(_ context.Context, l slog.Level) bool {
 	minLevel := slog.LevelInfo
 	if h.opts.Level != nil {
@@ -38,13 +41,11 @@ func (h *handler) Enabled(_ context.Context, l slog.Level) bool {
 	return l >= minLevel
 }
 
-// Handle implements slog.Handler.
 func (h *handler) Handle(_ context.Context, r slog.Record) error {
 	message := r.Message
 	timestamp := timestamppb.New(r.Time)
 	level := r.Level.String()
 
-	// Calculate the number of chunks
 	chunkSize := 500
 	numChunks := (len(message) + chunkSize - 1) / chunkSize
 
@@ -55,28 +56,26 @@ func (h *handler) Handle(_ context.Context, r slog.Record) error {
 			end = len(message)
 		}
 
-		// Create a chunk of the message
 		chunk := message[start:end]
 
-		// Create the agent log with the chunk
-		agentLog := manager.ClientStreamMessage{
-			Message: &manager.ClientStreamMessage_AgentLog{
-				AgentLog: &manager.AgentLog{
-					Timestamp: timestamp,
-					Message:   chunk,
-					Level:     level,
+		agentLog := events.EventsLogs{
+			Message: &events.EventsLogs_AgentLog{
+				AgentLog: &events.AgentLog{
+					Timestamp:     timestamp,
+					Message:       chunk,
+					Level:         level,
+					ComputationId: h.cmpID,
 				},
 			},
 		}
 
-		// Marshal the chunk to protobuf
 		b, err := proto.Marshal(&agentLog)
 		if err != nil {
 			return err
 		}
 
-		// Write the chunk to the writer
-		if _, err := h.w.Write(b); err != nil {
+		_, err = h.w.Write(b)
+		if err != nil {
 			return err
 		}
 	}
@@ -84,12 +83,14 @@ func (h *handler) Handle(_ context.Context, r slog.Record) error {
 	return nil
 }
 
-// WithAttrs implements slog.Handler.
-func (*handler) WithAttrs(attrs []slog.Attr) slog.Handler {
+func (h *handler) WithAttrs(attrs []slog.Attr) slog.Handler {
 	panic("unimplemented")
 }
 
-// WithGroup implements slog.Handler.
-func (*handler) WithGroup(name string) slog.Handler {
+func (h *handler) WithGroup(name string) slog.Handler {
 	panic("unimplemented")
 }
+
+func (h *handler) Close() error {
+	return nil
+}

internal/logger/protohandler_test.go

@@ -0,0 +1,76 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package logger
+
+import (
+	"context"
+	"io"
+	"log/slog"
+	"testing"
+	"time"
+
+	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/stretchr/testify/assert"
+)
+
+type failedWriter struct{}
+
+func (f *failedWriter) Write(p []byte) (n int, err error) {
+	return 0, io.ErrUnexpectedEOF
+}
+
+// TestNewProtoHandler tests the initialization of the ProtoHandler.
+func TestNewProtoHandler(t *testing.T) {
+	handler := NewProtoHandler(io.Discard, nil, "testCmpID")
+
+	assert.NotNil(t, handler, "Handler should not be nil")
+}
+
+// TestHandleMessageSuccess tests the handling of a message when the write succeeds.
+func TestHandleMessageSuccess(t *testing.T) {
+	handler := NewProtoHandler(io.Discard, nil, "testCmpID")
+	record := slog.Record{
+		Time:    time.Now(),
+		Message: "Test message",
+		Level:   slog.LevelInfo,
+	}
+
+	err := handler.Handle(context.Background(), record)
+
+	assert.NoError(t, err, "Handle should not return an error")
+}
+
+// TestHandleMessageFailure tests the caching mechanism when the write fails.
+func TestHandleMessageFailure(t *testing.T) {
+	protohandler := NewProtoHandler(&failedWriter{}, nil, "testCmpID")
+	record := slog.Record{
+		Time:    time.Now(),
+		Message: "Test message",
+		Level:   slog.LevelInfo,
+	}
+
+	err := protohandler.Handle(context.Background(), record)
+
+	assert.True(t, errors.Contains(err, io.ErrUnexpectedEOF), "Handle should return an error")
+}
+
+// TestEnabled tests that the handler enables logging based on level.
+func TestEnabled(t *testing.T) {
+	handler := NewProtoHandler(io.Discard, nil, "testCmpID")
+
+	assert.True(t, handler.Enabled(context.Background(), slog.LevelInfo), "Logging should be enabled for LevelInfo")
+	assert.False(t, handler.Enabled(context.Background(), slog.LevelDebug), "Logging should be disabled for LevelDebug by default")
+}
+
+// TestPeriodicRetry stops retry after close.
+func TestCloseStopsRetry(t *testing.T) {
+	mockWriter := io.Discard
+
+	handler := NewProtoHandler(mockWriter, nil, "testCmpID").(*handler)
+
+	time.Sleep(2 * time.Second)
+	err := handler.Close()
+
+	assert.NoError(t, err, "Close should not return an error")
+	time.Sleep(1 * time.Second) // Ensure no retry after close
+}

internal/server/doc.go

@@ -1,5 +1,5 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
 
-// Package server contains the HTTP, gRPC and CoAP server implementation.
+// Package server contains the gRPC server implementation.
 package server

internal/server/grpc/grpc.go

@@ -11,24 +11,26 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"crypto/x509/pkix"
-	"encoding/asn1"
 	"encoding/pem"
 	"fmt"
 	"log/slog"
 	"math/big"
 	"net"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/google/go-sev-guest/client"
 	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
 	"github.com/ultravioletrs/cocos/agent/auth"
 	"github.com/ultravioletrs/cocos/internal/server"
+	"github.com/ultravioletrs/cocos/pkg/atls"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
-	"golang.org/x/crypto/sha3"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/health"
+	grpchealth "google.golang.org/grpc/health/grpc_health_v1"
 )
 
 const (
@@ -42,6 +44,7 @@ const (
 	notAfterYear  = 1
 	notAfterMonth = 0
 	notAfterDay   = 0
+	nonceSize     = 32
 )
 
 type Server struct {
@@ -50,14 +53,16 @@ type Server struct {
 	registerService serviceRegister
 	quoteProvider   client.QuoteProvider
 	authSvc         auth.Authenticator
+	health          *health.Server
 }
 
 type serviceRegister func(srv *grpc.Server)
 
 var _ server.Server = (*Server)(nil)
 
-func New(ctx context.Context, cancel context.CancelFunc, name string, config server.Config, registerService serviceRegister, logger *slog.Logger, qp client.QuoteProvider, authSvc auth.Authenticator) server.Server {
-	listenFullAddress := fmt.Sprintf("%s:%s", config.Host, config.Port)
+func New(ctx context.Context, cancel context.CancelFunc, name string, config server.ServerConfiguration, registerService serviceRegister, logger *slog.Logger, qp client.QuoteProvider, authSvc auth.Authenticator) server.Server {
+	base := config.GetBaseConfig()
+	listenFullAddress := fmt.Sprintf("%s:%s", base.Host, base.Port)
 	return &Server{
 		BaseServer: server.BaseServer{
 			Ctx:     ctx,
@@ -85,15 +90,11 @@ func (s *Server) Start() error {
 		grpcServerOptions = append(grpcServerOptions, grpc.StreamInterceptor(stream))
 	}
 
-	listener, err := net.Listen("tcp", s.Address)
-	if err != nil {
-		return fmt.Errorf("failed to listen on port %s: %w", s.Address, err)
-	}
 	creds := grpc.Creds(insecure.NewCredentials())
+	var listener net.Listener
 
-	switch {
-	case s.Config.AttestedTLS:
-		certificateBytes, privateKeyBytes, err := generateCertificatesForATLS(s.quoteProvider)
+	if agCfg, ok := s.Config.(server.AgentConfig); ok && agCfg.AttestedTLS {
+		certificateBytes, privateKeyBytes, err := generateCertificatesForATLS()
 		if err != nil {
 			return fmt.Errorf("failed to create certificate: %w", err)
 		}
@@ -109,62 +110,91 @@ func (s *Server) Start() error {
 		}
 
 		creds = grpc.Creds(credentials.NewTLS(tlsConfig))
+
+		listener, err = atls.Listen(
+			s.Address,
+			certificateBytes,
+			privateKeyBytes,
+		)
+		if err != nil {
+			return fmt.Errorf("failed to create Listener for aTLS: %w", err)
+		}
 		s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s with Attested TLS", s.Name, s.Address))
-	case s.Config.CertFile != "" || s.Config.KeyFile != "":
-		certificate, err := loadX509KeyPair(s.Config.CertFile, s.Config.KeyFile)
-		if err != nil {
-			return fmt.Errorf("failed to load auth certificates: %w", err)
-		}
-		tlsConfig := &tls.Config{
-			ClientAuth:   tls.RequireAndVerifyClientCert,
-			Certificates: []tls.Certificate{certificate},
-		}
-
-		var mtlsCA string
-		// Loading Server CA file
-		rootCA, err := loadCertFile(s.Config.ServerCAFile)
-		if err != nil {
-			return fmt.Errorf("failed to load root ca file: %w", err)
-		}
-		if len(rootCA) > 0 {
-			if tlsConfig.RootCAs == nil {
-				tlsConfig.RootCAs = x509.NewCertPool()
-			}
-			if !tlsConfig.RootCAs.AppendCertsFromPEM(rootCA) {
-				return fmt.Errorf("failed to append root ca to tls.Config")
-			}
-			mtlsCA = fmt.Sprintf("root ca %s", s.Config.ServerCAFile)
-		}
-
-		// Loading Client CA File
-		clientCA, err := loadCertFile(s.Config.ClientCAFile)
-		if err != nil {
-			return fmt.Errorf("failed to load client ca file: %w", err)
-		}
-		if len(clientCA) > 0 {
-			if tlsConfig.ClientCAs == nil {
-				tlsConfig.ClientCAs = x509.NewCertPool()
-			}
-			if !tlsConfig.ClientCAs.AppendCertsFromPEM(clientCA) {
-				return fmt.Errorf("failed to append client ca to tls.Config")
-			}
-			mtlsCA = fmt.Sprintf("%s client ca %s", mtlsCA, s.Config.ClientCAFile)
-		}
-		creds = grpc.Creds(credentials.NewTLS(tlsConfig))
+	} else {
+		c := s.Config.GetBaseConfig()
 		switch {
-		case mtlsCA != "":
-			s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s with TLS/mTLS cert %s , key %s and %s", s.Name, s.Address, s.Config.CertFile, s.Config.KeyFile, mtlsCA))
+		case c.CertFile != "" || c.KeyFile != "":
+			certificate, err := loadX509KeyPair(c.CertFile, c.KeyFile)
+			if err != nil {
+				return fmt.Errorf("failed to load auth certificates: %w", err)
+			}
+			tlsConfig := &tls.Config{
+				ClientAuth:   tls.NoClientCert,
+				Certificates: []tls.Certificate{certificate},
+			}
+
+			var mtlsCA string
+			// Loading Server CA file
+			rootCA, err := loadCertFile(c.ServerCAFile)
+			if err != nil {
+				return fmt.Errorf("failed to load root ca file: %w", err)
+			}
+			if len(rootCA) > 0 {
+				if tlsConfig.RootCAs == nil {
+					tlsConfig.RootCAs = x509.NewCertPool()
+				}
+				if !tlsConfig.RootCAs.AppendCertsFromPEM(rootCA) {
+					return fmt.Errorf("failed to append root ca to tls.Config")
+				}
+				mtlsCA = fmt.Sprintf("root ca %s", c.ServerCAFile)
+			}
+
+			// Loading Client CA File
+			clientCA, err := loadCertFile(c.ClientCAFile)
+			if err != nil {
+				return fmt.Errorf("failed to load client ca file: %w", err)
+			}
+			if len(clientCA) > 0 {
+				if tlsConfig.ClientCAs == nil {
+					tlsConfig.ClientCAs = x509.NewCertPool()
+				}
+				if !tlsConfig.ClientCAs.AppendCertsFromPEM(clientCA) {
+					return fmt.Errorf("failed to append client ca to tls.Config")
+				}
+				mtlsCA = fmt.Sprintf("%s client ca %s", mtlsCA, c.ClientCAFile)
+			}
+			creds = grpc.Creds(credentials.NewTLS(tlsConfig))
+			switch {
+			case mtlsCA != "":
+				tlsConfig.ClientAuth = tls.RequireAndVerifyClientCert
+				creds = grpc.Creds(credentials.NewTLS(tlsConfig))
+				s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s with TLS/mTLS cert %s , key %s and %s", s.Name, s.Address, c.CertFile, c.KeyFile, mtlsCA))
+			default:
+				s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s with TLS cert %s and key %s", s.Name, s.Address, c.CertFile, c.KeyFile))
+			}
+
+			listener, err = net.Listen("tcp", s.Address)
+			if err != nil {
+				return fmt.Errorf("failed to listen on port %s: %w", s.Address, err)
+			}
 		default:
-			s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s with TLS cert %s and key %s", s.Name, s.Address, s.Config.CertFile, s.Config.KeyFile))
+			var err error
+
+			listener, err = net.Listen("tcp", s.Address)
+			if err != nil {
+				return fmt.Errorf("failed to listen on port %s: %w", s.Address, err)
+			}
+			s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s without TLS", s.Name, s.Address))
 		}
-	default:
-		s.Logger.Info(fmt.Sprintf("%s service gRPC server listening at %s without TLS", s.Name, s.Address))
 	}
 
 	grpcServerOptions = append(grpcServerOptions, creds)
 
 	s.server = grpc.NewServer(grpcServerOptions...)
+	s.health = health.NewServer()
+	grpchealth.RegisterHealthServer(s.server, s.health)
 	s.registerService(s.server)
+	s.health.SetServingStatus(s.Name, grpchealth.HealthCheckResponse_SERVING)
 
 	go func() {
 		errCh <- s.server.Serve(listener)
@@ -184,6 +214,7 @@ func (s *Server) Stop() error {
 	c := make(chan bool)
 	go func() {
 		defer close(c)
+		s.health.Shutdown()
 		s.server.GracefulStop()
 	}()
 	select {
@@ -197,55 +228,42 @@ func (s *Server) Stop() error {
 
 func loadCertFile(certFile string) ([]byte, error) {
 	if certFile != "" {
-		return os.ReadFile(certFile)
+		return readFileOrData(certFile)
 	}
 	return []byte{}, nil
 }
 
+func readFileOrData(input string) ([]byte, error) {
+	if len(input) < 1000 && !strings.Contains(input, "\n") {
+		data, err := os.ReadFile(input)
+		if err == nil {
+			return data, nil
+		}
+	}
+	return []byte(input), nil
+}
+
 func loadX509KeyPair(certfile, keyfile string) (tls.Certificate, error) {
-	var cert, key []byte
-	var err error
-	if _, err = os.Stat(certfile); err == nil {
-		cert, err = os.ReadFile(certfile)
-		if err != nil {
-			return tls.Certificate{}, err
-		}
-	} else if os.IsNotExist(err) {
-		cert = []byte(certfile)
-	} else {
-		return tls.Certificate{}, err
+	cert, err := readFileOrData(certfile)
+	if err != nil {
+		return tls.Certificate{}, fmt.Errorf("failed to read cert: %v", err)
 	}
-	if _, err := os.Stat(keyfile); err == nil {
-		key, err = os.ReadFile(keyfile)
-		if err != nil {
-			return tls.Certificate{}, err
-		}
-	} else if os.IsNotExist(err) {
-		key = []byte(keyfile)
-	} else {
-		return tls.Certificate{}, err
+
+	key, err := readFileOrData(keyfile)
+	if err != nil {
+		return tls.Certificate{}, fmt.Errorf("failed to read key: %v", err)
 	}
+
 	return tls.X509KeyPair(cert, key)
 }
 
-func generateCertificatesForATLS(qp client.QuoteProvider) ([]byte, []byte, error) {
+func generateCertificatesForATLS() ([]byte, []byte, error) {
 	curve := elliptic.P256()
 	privateKey, err := ecdsa.GenerateKey(curve, rand.Reader)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to generate private/public key: %w", err)
 	}
 
-	publicKeyBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
-	if err != nil {
-		return nil, nil, fmt.Errorf("failed to marshal the public key: %w", err)
-	}
-
-	// The Attestation Report will be added as an X.509 certificate extension
-	attestationReport, err := qp.GetRawQuote(sha3.Sum512(publicKeyBytes))
-	if err != nil {
-		return nil, nil, fmt.Errorf("failed to fetch the attestation report: %w", err)
-	}
-
 	certTemplate := &x509.Certificate{
 		SerialNumber: big.NewInt(202403311),
 		Subject: pkix.Name{
@@ -261,13 +279,6 @@ func generateCertificatesForATLS(qp client.QuoteProvider) ([]byte, []byte, error
 		KeyUsage:              x509.KeyUsageDigitalSignature,
 		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
 		BasicConstraintsValid: true,
-		ExtraExtensions: []pkix.Extension{
-			{
-				Id:       asn1.ObjectIdentifier{1, 2, 3, 4, 5, 6},
-				Critical: false,
-				Value:    attestationReport,
-			},
-		},
 	}
 
 	certDERBytes, err := x509.CreateCertificate(rand.Reader, certTemplate, certTemplate, &privateKey.PublicKey, privateKey)

internal/server/grpc/grpc_test.go

@@ -0,0 +1,534 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package grpc
+
+import (
+	"context"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/pem"
+	"fmt"
+	"log/slog"
+	"math/big"
+	"os"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	authmocks "github.com/ultravioletrs/cocos/agent/auth/mocks"
+	"github.com/ultravioletrs/cocos/internal/server"
+	"github.com/ultravioletrs/cocos/pkg/attestation/quoteprovider/mocks"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+const bufSize = 1024 * 1024
+
+var lis *bufconn.Listener
+
+func init() {
+	lis = bufconn.Listen(bufSize)
+}
+
+func TestNew(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	config := server.AgentConfig{
+		ServerConfig: server.ServerConfig{
+			BaseConfig: server.BaseConfig{
+				Host: "localhost",
+				Port: "50051",
+			},
+		},
+	}
+	logger := slog.Default()
+	qp := new(mocks.QuoteProvider)
+	authSvc := new(authmocks.Authenticator)
+
+	srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
+
+	assert.NotNil(t, srv)
+	assert.IsType(t, &Server{}, srv)
+}
+
+func TestServerStartWithTLSFile(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	cert, key, err := generateSelfSignedCert()
+	assert.NoError(t, err)
+
+	certFile, err := os.CreateTemp("", "cert*.pem")
+	assert.NoError(t, err)
+
+	keyFile, err := os.CreateTemp("", "key*.pem")
+	assert.NoError(t, err)
+
+	t.Cleanup(func() {
+		os.Remove(certFile.Name())
+		os.Remove(keyFile.Name())
+	})
+
+	_, err = certFile.Write(cert)
+	assert.NoError(t, err)
+
+	_, err = keyFile.Write(key)
+	assert.NoError(t, err)
+
+	err = certFile.Close()
+	assert.NoError(t, err)
+	err = keyFile.Close()
+	assert.NoError(t, err)
+
+	config := server.AgentConfig{
+		ServerConfig: server.ServerConfig{
+			BaseConfig: server.BaseConfig{
+				Host:     "localhost",
+				Port:     "0",
+				CertFile: certFile.Name(),
+				KeyFile:  keyFile.Name(),
+			},
+		},
+	}
+
+	logBuffer := &ThreadSafeBuffer{}
+	logger := slog.New(slog.NewTextHandler(logBuffer, &slog.HandlerOptions{Level: slog.LevelDebug}))
+	qp := new(mocks.QuoteProvider)
+	authSvc := new(authmocks.Authenticator)
+
+	srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	go func() {
+		wg.Done()
+		err := srv.Start()
+		assert.NoError(t, err)
+	}()
+
+	wg.Wait()
+
+	time.Sleep(200 * time.Millisecond)
+
+	cancel()
+
+	time.Sleep(200 * time.Millisecond)
+
+	logContent := logBuffer.String()
+	fmt.Println(logContent)
+	assert.Contains(t, logContent, "TestServer service gRPC server listening at localhost:0 with TLS")
+}
+
+func TestServerStartWithmTLSFile(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	caCertFile, clientCertFile, clientKeyFile, err := createCertificatesFiles()
+	assert.NoError(t, err)
+
+	config := server.AgentConfig{
+		ServerConfig: server.ServerConfig{
+			BaseConfig: server.BaseConfig{
+				Host:         "localhost",
+				Port:         "0",
+				CertFile:     string(clientCertFile),
+				KeyFile:      string(clientKeyFile),
+				ServerCAFile: caCertFile,
+			},
+		},
+	}
+
+	logBuffer := &ThreadSafeBuffer{}
+	logger := slog.New(slog.NewTextHandler(logBuffer, &slog.HandlerOptions{Level: slog.LevelDebug}))
+	qp := new(mocks.QuoteProvider)
+	authSvc := new(authmocks.Authenticator)
+
+	srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	go func() {
+		wg.Done()
+		err := srv.Start()
+		assert.NoError(t, err)
+	}()
+
+	wg.Wait()
+
+	time.Sleep(200 * time.Millisecond)
+
+	cancel()
+
+	time.Sleep(200 * time.Millisecond)
+
+	logContent := logBuffer.String()
+	fmt.Println(logContent)
+	assert.Contains(t, logContent, "TestServer service gRPC server listening at localhost:0 with TLS")
+}
+
+func TestServerStop(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	config := server.AgentConfig{
+		ServerConfig: server.ServerConfig{
+			BaseConfig: server.BaseConfig{
+				Host: "localhost",
+				Port: "0",
+			},
+		},
+	}
+	buf := &ThreadSafeBuffer{}
+	logger := slog.New(slog.NewTextHandler(buf, &slog.HandlerOptions{Level: slog.LevelDebug}))
+	qp := new(mocks.QuoteProvider)
+	authSvc := new(authmocks.Authenticator)
+
+	srv := New(ctx, cancel, "TestServer", config, func(srv *grpc.Server) {}, logger, qp, authSvc)
+
+	go func() {
+		err := srv.Start()
+		assert.NoError(t, err)
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+
+	cancel()
+
+	time.Sleep(100 * time.Millisecond)
+
+	err := srv.Stop()
+	assert.NoError(t, err)
+
+	assert.Contains(t, buf.String(), "TestServer gRPC service shutdown at localhost:0")
+}
+
+func generateSelfSignedCert() ([]byte, []byte, error) {
+	key, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	cert, err := generateSelfSignedCertFromKey(key)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return cert, pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)}), nil
+}
+
+func generateSelfSignedCertFromKey(key *rsa.PrivateKey) ([]byte, error) {
+	template := x509.Certificate{
+		SerialNumber: big.NewInt(1),
+		Subject: pkix.Name{
+			Organization: []string{"Test"},
+		},
+		NotBefore:             time.Now(),
+		NotAfter:              time.Now().AddDate(1, 0, 0),
+		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
+		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
+		BasicConstraintsValid: true,
+	}
+
+	certBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key)
+	if err != nil {
+		return nil, err
+	}
+
+	return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certBytes}), nil
+}
+
+type ThreadSafeBuffer struct {
+	buffer strings.Builder
+	mu     sync.Mutex
+}
+
+func (b *ThreadSafeBuffer) Write(p []byte) (n int, err error) {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buffer.Write(p)
+}
+
+func (b *ThreadSafeBuffer) String() string {
+	b.mu.Lock()
+	defer b.mu.Unlock()
+	return b.buffer.String()
+}
+
+func TestServerInitializationAndStartup(t *testing.T) {
+	testCases := []struct {
+		name          string
+		config        server.AgentConfig
+		expectedLog   string
+		expectError   bool
+		setupCallback func(*testing.T, *server.AgentConfig, *ThreadSafeBuffer)
+	}{
+		{
+			name: "Non-TLS Server Startup",
+			config: server.AgentConfig{
+				ServerConfig: server.ServerConfig{
+					BaseConfig: server.BaseConfig{
+						Host: "localhost",
+						Port: "0",
+					},
+				},
+			},
+			expectedLog: "TestServer service gRPC server listening at localhost:0 without TLS",
+		},
+		{
+			name: "TLS Server Startup with Self-Signed Certificate",
+			config: server.AgentConfig{
+				ServerConfig: server.ServerConfig{
+					BaseConfig: server.BaseConfig{
+						Host: "localhost",
+						Port: "0",
+					},
+				},
+			},
+			setupCallback: setupTLSConfig,
+			expectedLog:   "TestServer service gRPC server listening at localhost:0 with TLS",
+		},
+		{
+			name: "TLS Server Startup with Invalid Certificates",
+			config: server.AgentConfig{
+				ServerConfig: server.ServerConfig{
+					BaseConfig: server.BaseConfig{
+						Host:     "localhost",
+						Port:     "0",
+						CertFile: "invalid",
+						KeyFile:  "invalid",
+					},
+				},
+			},
+			expectError: true,
+			expectedLog: "failed to load auth certificates",
+		},
+		{
+			name: "mTLS Server Startup",
+			config: server.AgentConfig{
+				ServerConfig: server.ServerConfig{
+					BaseConfig: server.BaseConfig{
+						Host: "localhost",
+						Port: "0",
+					},
+				},
+			},
+			setupCallback: setupMTLSConfig,
+			expectedLog:   "TestServer service gRPC server listening at localhost:0 with TLS",
+		},
+		{
+			name: "mTLS Server Startup with Invalid Root CA",
+			config: server.AgentConfig{
+				ServerConfig: server.ServerConfig{
+					BaseConfig: server.BaseConfig{
+						Host:         "localhost",
+						Port:         "0",
+						ServerCAFile: "invalid",
+					},
+				},
+			},
+			setupCallback: setupInvalidRootCAConfig,
+			expectError:   true,
+			expectedLog:   "failed to append root ca to tls.Config",
+		},
+		{
+			name: "mTLS Server Startup with Invalid Client CA",
+			config: server.AgentConfig{
+				ServerConfig: server.ServerConfig{
+					BaseConfig: server.BaseConfig{
+						Host:         "localhost",
+						Port:         "0",
+						ServerCAFile: "invalid",
+					},
+				},
+			},
+			setupCallback: setupInvalidClientCAConfig,
+			expectError:   true,
+			expectedLog:   "failed to append client ca to tls.Config",
+		},
+		{
+			name: "Attested TLS Server Startup",
+			config: server.AgentConfig{
+				ServerConfig: server.ServerConfig{
+					BaseConfig: server.BaseConfig{
+						Host: "localhost",
+						Port: "0",
+					},
+				},
+				AttestedTLS: true,
+			},
+			expectedLog: "TestServer service gRPC server listening at localhost:0 with Attested TLS",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			if tc.setupCallback != nil {
+				tc.setupCallback(t, &tc.config, nil)
+			}
+
+			logBuffer := &ThreadSafeBuffer{}
+			logger := slog.New(slog.NewTextHandler(logBuffer, &slog.HandlerOptions{Level: slog.LevelDebug}))
+			qp := new(mocks.QuoteProvider)
+			authSvc := new(authmocks.Authenticator)
+
+			srv := New(ctx, cancel, "TestServer", tc.config, func(srv *grpc.Server) {}, logger, qp, authSvc)
+			var wg sync.WaitGroup
+			wg.Add(1)
+
+			go func() {
+				wg.Done()
+				err := srv.Start()
+				if tc.expectError {
+					assert.Error(t, err)
+					assert.Contains(t, err.Error(), tc.expectedLog)
+				} else {
+					assert.NoError(t, err)
+				}
+			}()
+
+			wg.Wait()
+
+			time.Sleep(200 * time.Millisecond)
+
+			cancel()
+
+			time.Sleep(200 * time.Millisecond)
+
+			if !tc.expectError {
+				logContent := logBuffer.String()
+				fmt.Println(logContent)
+				assert.Contains(t, logContent, tc.expectedLog)
+			}
+		})
+	}
+}
+
+func setupTLSConfig(t *testing.T, config *server.AgentConfig, _ *ThreadSafeBuffer) {
+	cert, key, err := generateSelfSignedCert()
+	assert.NoError(t, err)
+
+	config.CertFile = string(cert)
+	config.KeyFile = string(key)
+}
+
+func setupMTLSConfig(t *testing.T, config *server.AgentConfig, _ *ThreadSafeBuffer) {
+	cert, key, err := generateSelfSignedCert()
+	assert.NoError(t, err)
+
+	config.CertFile = string(cert)
+	config.KeyFile = string(key)
+	config.ServerCAFile = string(cert)
+	config.ClientCAFile = string(cert)
+}
+
+func setupInvalidRootCAConfig(t *testing.T, config *server.AgentConfig, _ *ThreadSafeBuffer) {
+	cert, key, err := generateSelfSignedCert()
+	assert.NoError(t, err)
+
+	config.CertFile = string(cert)
+	config.KeyFile = string(key)
+	config.ServerCAFile = "invalid"
+	config.ClientCAFile = string(cert)
+}
+
+func setupInvalidClientCAConfig(t *testing.T, config *server.AgentConfig, _ *ThreadSafeBuffer) {
+	cert, key, err := generateSelfSignedCert()
+	assert.NoError(t, err)
+
+	config.CertFile = string(cert)
+	config.KeyFile = string(key)
+	config.ClientCAFile = "invalid"
+	config.ServerCAFile = string(cert)
+}
+
+func createCertificatesFiles() (string, string, string, error) {
+	caKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		return "", "", "", err
+	}
+
+	caTemplate := x509.Certificate{
+		SerialNumber: big.NewInt(1),
+		Subject: pkix.Name{
+			Organization: []string{"Test Org"},
+		},
+		NotBefore:             time.Now(),
+		NotAfter:              time.Now().Add(time.Hour * 24),
+		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
+		BasicConstraintsValid: true,
+		IsCA:                  true,
+	}
+
+	caCertDER, err := x509.CreateCertificate(rand.Reader, &caTemplate, &caTemplate, &caKey.PublicKey, caKey)
+	if err != nil {
+		return "", "", "", err
+	}
+
+	caCertFile, err := createTempFile(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: caCertDER}))
+	if err != nil {
+		return "", "", "", err
+	}
+
+	clientKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	if err != nil {
+		return "", "", "", err
+	}
+
+	clientTemplate := x509.Certificate{
+		SerialNumber: big.NewInt(2),
+		Subject: pkix.Name{
+			Organization: []string{"Test Org"},
+		},
+		NotBefore:             time.Now(),
+		NotAfter:              time.Now().Add(time.Hour * 24),
+		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
+		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
+		BasicConstraintsValid: true,
+	}
+
+	clientCertDER, err := x509.CreateCertificate(rand.Reader, &clientTemplate, &caTemplate, &clientKey.PublicKey, caKey)
+	if err != nil {
+		return "", "", "", err
+	}
+
+	clientCertFile, err := createTempFile(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: clientCertDER}))
+	if err != nil {
+		return "", "", "", err
+	}
+
+	clientKeyFile, err := createTempFile(pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(clientKey)}))
+	if err != nil {
+		return "", "", "", err
+	}
+
+	return caCertFile, clientCertFile, clientKeyFile, nil
+}
+
+func createTempFile(data []byte) (string, error) {
+	file, err := createTempFileHandle()
+	if err != nil {
+		return "", err
+	}
+
+	_, err = file.Write(data)
+	if err != nil {
+		return "", err
+	}
+
+	err = file.Close()
+	if err != nil {
+		return "", err
+	}
+
+	return file.Name(), nil
+}
+
+func createTempFileHandle() (*os.File, error) {
+	return os.CreateTemp("", "test")
+}

internal/server/mocks/server.go

@@ -0,0 +1,125 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+package mocks
+
+import mock "github.com/stretchr/testify/mock"
+
+// Server is an autogenerated mock type for the Server type
+type Server struct {
+	mock.Mock
+}
+
+type Server_Expecter struct {
+	mock *mock.Mock
+}
+
+func (_m *Server) EXPECT() *Server_Expecter {
+	return &Server_Expecter{mock: &_m.Mock}
+}
+
+// Start provides a mock function with given fields:
+func (_m *Server) Start() error {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Start")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Server_Start_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Start'
+type Server_Start_Call struct {
+	*mock.Call
+}
+
+// Start is a helper method to define mock.On call
+func (_e *Server_Expecter) Start() *Server_Start_Call {
+	return &Server_Start_Call{Call: _e.mock.On("Start")}
+}
+
+func (_c *Server_Start_Call) Run(run func()) *Server_Start_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *Server_Start_Call) Return(_a0 error) *Server_Start_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Server_Start_Call) RunAndReturn(run func() error) *Server_Start_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// Stop provides a mock function with given fields:
+func (_m *Server) Stop() error {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Stop")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Server_Stop_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Stop'
+type Server_Stop_Call struct {
+	*mock.Call
+}
+
+// Stop is a helper method to define mock.On call
+func (_e *Server_Expecter) Stop() *Server_Stop_Call {
+	return &Server_Stop_Call{Call: _e.mock.On("Stop")}
+}
+
+func (_c *Server_Stop_Call) Run(run func()) *Server_Stop_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run()
+	})
+	return _c
+}
+
+func (_c *Server_Stop_Call) Return(_a0 error) *Server_Stop_Call {
+	_c.Call.Return(_a0)
+	return _c
+}
+
+func (_c *Server_Stop_Call) RunAndReturn(run func() error) *Server_Stop_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
+// NewServer creates a new instance of Server. 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 NewServer(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Server {
+	mock := &Server{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

internal/server/server.go

@@ -16,14 +16,25 @@ type Server interface {
 	Stop() error
 }
 
-type Config struct {
-	Host         string `env:"HOST"               envDefault:""`
-	Port         string `env:"PORT"               envDefault:""`
+type ServerConfiguration interface {
+	GetBaseConfig() ServerConfig
+}
+
+type BaseConfig struct {
+	Host         string `env:"HOST"               envDefault:"localhost"`
+	Port         string `env:"PORT"               envDefault:"7001"`
+	ServerCAFile string `env:"SERVER_CA_CERTS" envDefault:""`
 	CertFile     string `env:"SERVER_CERT"        envDefault:""`
 	KeyFile      string `env:"SERVER_KEY"         envDefault:""`
-	ServerCAFile string `env:"SERVER_CA_CERTS"    envDefault:""`
 	ClientCAFile string `env:"CLIENT_CA_CERTS"    envDefault:""`
-	AttestedTLS  bool   `env:"ATTESTED_TLS"       envDefault:"false"`
+}
+
+type ServerConfig struct {
+	BaseConfig
+}
+type AgentConfig struct {
+	ServerConfig
+	AttestedTLS bool `env:"ATTESTED_TLS"       envDefault:"false"`
 }
 
 type BaseServer struct {
@@ -31,11 +42,19 @@ type BaseServer struct {
 	Cancel   context.CancelFunc
 	Name     string
 	Address  string
-	Config   Config
+	Config   ServerConfiguration
 	Logger   *slog.Logger
 	Protocol string
 }
 
+func (s ServerConfig) GetBaseConfig() ServerConfig {
+	return s
+}
+
+func (a AgentConfig) GetBaseConfig() ServerConfig {
+	return a.ServerConfig
+}
+
 func stopAllServer(servers ...Server) error {
 	var errs []error
 	for _, server := range servers {

internal/server/server_test.go

@@ -0,0 +1,138 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package server
+
+import (
+	"context"
+	"errors"
+	"log/slog"
+	"os"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/ultravioletrs/cocos/internal/server/mocks"
+)
+
+func TestStopAllServer(t *testing.T) {
+	server1 := new(mocks.Server)
+	server2 := new(mocks.Server)
+	server1.On("Stop").Return(nil)
+	server2.On("Stop").Return(errors.New("failed to stop"))
+	tests := []struct {
+		name          string
+		servers       []Server
+		expectedError bool
+	}{
+		{
+			name: "All servers stop successfully",
+			servers: []Server{
+				server1,
+				server1,
+			},
+			expectedError: false,
+		},
+		{
+			name: "One server fails to stop",
+			servers: []Server{
+				server1,
+				server2,
+			},
+			expectedError: true,
+		},
+		{
+			name:          "No servers",
+			servers:       []Server{},
+			expectedError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := stopAllServer(tt.servers...)
+			if (err != nil) != tt.expectedError {
+				t.Errorf("stopAllServer() error = %v, expectedError %v", err, tt.expectedError)
+			}
+		})
+	}
+}
+
+func TestStopHandler(t *testing.T) {
+	mockServer := new(mocks.Server)
+	mockServer.On("Stop").Return(nil)
+	tests := []struct {
+		name           string
+		setupFunc      func() (context.Context, context.CancelFunc, *slog.Logger, string, []Server)
+		triggerSignal  bool
+		expectedError  bool
+		expectCanceled bool
+	}{
+		{
+			name: "Graceful shutdown on signal",
+			setupFunc: func() (context.Context, context.CancelFunc, *slog.Logger, string, []Server) {
+				ctx, cancel := context.WithCancel(context.Background())
+				logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+				return ctx, cancel, logger, "test", []Server{mockServer}
+			},
+			triggerSignal:  true,
+			expectedError:  false,
+			expectCanceled: true,
+		},
+		{
+			name: "Context canceled",
+			setupFunc: func() (context.Context, context.CancelFunc, *slog.Logger, string, []Server) {
+				ctx, cancel := context.WithCancel(context.Background())
+				logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
+				go func() {
+					time.Sleep(100 * time.Millisecond)
+					cancel()
+				}()
+				return ctx, cancel, logger, "test", []Server{mockServer}
+			},
+			triggerSignal:  false,
+			expectedError:  false,
+			expectCanceled: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx, cancel, logger, svcName, servers := tt.setupFunc()
+			defer cancel()
+
+			errChan := make(chan error)
+			go func() {
+				errChan <- StopHandler(ctx, cancel, logger, svcName, servers...)
+			}()
+
+			if tt.triggerSignal {
+				// Simulate SIGINT
+				go func() {
+					time.Sleep(100 * time.Millisecond)
+					err := syscall.Kill(syscall.Getpid(), syscall.SIGINT)
+					if err != nil {
+						t.Errorf("failed to send signal: %v", err)
+					}
+				}()
+			}
+
+			select {
+			case err := <-errChan:
+				if (err != nil) != tt.expectedError {
+					t.Errorf("StopHandler() error = %v, expectedError %v", err, tt.expectedError)
+				}
+			case <-time.After(2 * time.Second):
+				t.Error("StopHandler() timed out")
+			}
+
+			if tt.expectCanceled {
+				select {
+				case <-ctx.Done():
+					// Context was canceled as expected
+				default:
+					t.Error("Context was not canceled")
+				}
+			}
+		})
+	}
+}