COCOS-144 - Add Agent SDK Tests (#167)

* add tests and mocks

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

fix ci

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

update test

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

fix(agent/grpc): revert change

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

fix ci

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

* refactor attestation and report tests

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

refactor tests

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

remove commented code

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

remove comment

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

remove comments

* add test cases

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

export agent errors

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

remove comm

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

* fix tests

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

---------

Signed-off-by: WashingtonKK <washingtonkigan@gmail.com>

.github/workflows/checkproto.yaml

@@ -33,9 +33,9 @@ jobs:
 
       - name: Set up protoc
         run: |
-          PROTOC_VERSION=25.3
-          PROTOC_GEN_VERSION=v1.33.0
-          PROTOC_GRPC_VERSION=v1.3.0
+          PROTOC_VERSION=27.2
+          PROTOC_GEN_VERSION=v1.34.2
+          PROTOC_GRPC_VERSION=v1.4.0
 
           # Download and install protoc
           PROTOC_ZIP=protoc-$PROTOC_VERSION-linux-x86_64.zip

.github/workflows/main.yaml

@@ -22,9 +22,9 @@ jobs:
           go-version: 1.22.x
 
       - name: golangci-lint
-        uses: golangci/golangci-lint-action@v3
+        uses: golangci/golangci-lint-action@v6
         with:
-          version: latest
+          version: v1.59.1
 
       - name: Build
         run: |

.golangci.yaml

@@ -1,6 +1,6 @@
 run:
   timeout: 3m
-  
+
 issues:
   max-issues-per-linter: 10
   max-same-issues: 10
@@ -15,10 +15,7 @@ linters-settings:
 
   gocritic:
     enabled-checks:
-      - captLocal
       - dupImport
-      - singleCaseSwitch
-      - switchTrue
       - importShadow
       - httpNoBody
       - paramTypeCombine
@@ -34,9 +31,6 @@ linters-settings:
       - style
       - experimental
       - opinionated
-  misspell:
-    ignore-words:
-      - "mosquitto"
   stylecheck:
     checks: ["-ST1000", "-ST1003", "-ST1020", "-ST1021", "-ST1022"]
   goheader:
@@ -83,4 +77,4 @@ linters:
     - makezero
     - mirror
     - nakedret
-    - dupword
+    - dupword

Makefile

@@ -6,8 +6,6 @@ GOARCH ?= amd64
 VERSION ?= $(shell git describe --abbrev=0 --tags --always)
 COMMIT ?= $(shell git rev-parse HEAD)
 TIME ?= $(shell date +%F_%T)
-empty:=
-space:= $(empty) $(empty)
 
 define compile_service
 	CGO_ENABLED=$(CGO_ENABLED) GOOS=$(GOOS) GOARCH=$(GOARCH) GOARM=$(GOARM) \
@@ -31,3 +29,6 @@ $(PLATFORM_INFO):
 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
+
+mocks:
+	go generate ./...

agent/agent.pb.go

@@ -3,8 +3,8 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.33.0
-// 	protoc        v4.25.3
+// 	protoc-gen-go v1.34.2
+// 	protoc        v5.27.2
 // source: agent/agent.proto
 
 package agent
@@ -426,7 +426,7 @@ func file_agent_agent_proto_rawDescGZIP() []byte {
 }
 
 var file_agent_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 8)
-var file_agent_agent_proto_goTypes = []interface{}{
+var file_agent_agent_proto_goTypes = []any{
 	(*AlgoRequest)(nil),         // 0: agent.AlgoRequest
 	(*AlgoResponse)(nil),        // 1: agent.AlgoResponse
 	(*DataRequest)(nil),         // 2: agent.DataRequest
@@ -458,7 +458,7 @@ func file_agent_agent_proto_init() {
 		return
 	}
 	if !protoimpl.UnsafeEnabled {
-		file_agent_agent_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
+		file_agent_agent_proto_msgTypes[0].Exporter = func(v any, i int) any {
 			switch v := v.(*AlgoRequest); i {
 			case 0:
 				return &v.state
@@ -470,7 +470,7 @@ func file_agent_agent_proto_init() {
 				return nil
 			}
 		}
-		file_agent_agent_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} {
+		file_agent_agent_proto_msgTypes[1].Exporter = func(v any, i int) any {
 			switch v := v.(*AlgoResponse); i {
 			case 0:
 				return &v.state
@@ -482,7 +482,7 @@ func file_agent_agent_proto_init() {
 				return nil
 			}
 		}
-		file_agent_agent_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} {
+		file_agent_agent_proto_msgTypes[2].Exporter = func(v any, i int) any {
 			switch v := v.(*DataRequest); i {
 			case 0:
 				return &v.state
@@ -494,7 +494,7 @@ func file_agent_agent_proto_init() {
 				return nil
 			}
 		}
-		file_agent_agent_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} {
+		file_agent_agent_proto_msgTypes[3].Exporter = func(v any, i int) any {
 			switch v := v.(*DataResponse); i {
 			case 0:
 				return &v.state
@@ -506,7 +506,7 @@ func file_agent_agent_proto_init() {
 				return nil
 			}
 		}
-		file_agent_agent_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} {
+		file_agent_agent_proto_msgTypes[4].Exporter = func(v any, i int) any {
 			switch v := v.(*ResultRequest); i {
 			case 0:
 				return &v.state
@@ -518,7 +518,7 @@ func file_agent_agent_proto_init() {
 				return nil
 			}
 		}
-		file_agent_agent_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} {
+		file_agent_agent_proto_msgTypes[5].Exporter = func(v any, i int) any {
 			switch v := v.(*ResultResponse); i {
 			case 0:
 				return &v.state
@@ -530,7 +530,7 @@ func file_agent_agent_proto_init() {
 				return nil
 			}
 		}
-		file_agent_agent_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} {
+		file_agent_agent_proto_msgTypes[6].Exporter = func(v any, i int) any {
 			switch v := v.(*AttestationRequest); i {
 			case 0:
 				return &v.state
@@ -542,7 +542,7 @@ func file_agent_agent_proto_init() {
 				return nil
 			}
 		}
-		file_agent_agent_proto_msgTypes[7].Exporter = func(v interface{}, i int) interface{} {
+		file_agent_agent_proto_msgTypes[7].Exporter = func(v any, i int) any {
 			switch v := v.(*AttestationResponse); i {
 			case 0:
 				return &v.state

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.3.0
-// - protoc             v4.25.3
+// - protoc-gen-go-grpc v1.4.0
+// - protoc             v5.27.2
 // 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.32.0 or later.
-const _ = grpc.SupportPackageIsVersion7
+// Requires gRPC-Go v1.62.0 or later.
+const _ = grpc.SupportPackageIsVersion8
 
 const (
 	AgentService_Algo_FullMethodName        = "/agent.AgentService/Algo"
@@ -47,11 +47,12 @@ func NewAgentServiceClient(cc grpc.ClientConnInterface) AgentServiceClient {
 }
 
 func (c *agentServiceClient) Algo(ctx context.Context, opts ...grpc.CallOption) (AgentService_AlgoClient, error) {
-	stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[0], AgentService_Algo_FullMethodName, opts...)
+	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{stream}
+	x := &agentServiceAlgoClient{ClientStream: stream}
 	return x, nil
 }
 
@@ -81,11 +82,12 @@ func (x *agentServiceAlgoClient) CloseAndRecv() (*AlgoResponse, error) {
 }
 
 func (c *agentServiceClient) Data(ctx context.Context, opts ...grpc.CallOption) (AgentService_DataClient, error) {
-	stream, err := c.cc.NewStream(ctx, &AgentService_ServiceDesc.Streams[1], AgentService_Data_FullMethodName, opts...)
+	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{stream}
+	x := &agentServiceDataClient{ClientStream: stream}
 	return x, nil
 }
 
@@ -115,8 +117,9 @@ func (x *agentServiceDataClient) CloseAndRecv() (*DataResponse, error) {
 }
 
 func (c *agentServiceClient) Result(ctx context.Context, in *ResultRequest, opts ...grpc.CallOption) (*ResultResponse, error) {
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(ResultResponse)
-	err := c.cc.Invoke(ctx, AgentService_Result_FullMethodName, in, out, opts...)
+	err := c.cc.Invoke(ctx, AgentService_Result_FullMethodName, in, out, cOpts...)
 	if err != nil {
 		return nil, err
 	}
@@ -124,8 +127,9 @@ func (c *agentServiceClient) Result(ctx context.Context, in *ResultRequest, opts
 }
 
 func (c *agentServiceClient) Attestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (*AttestationResponse, error) {
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
 	out := new(AttestationResponse)
-	err := c.cc.Invoke(ctx, AgentService_Attestation_FullMethodName, in, out, opts...)
+	err := c.cc.Invoke(ctx, AgentService_Attestation_FullMethodName, in, out, cOpts...)
 	if err != nil {
 		return nil, err
 	}
@@ -173,7 +177,7 @@ func RegisterAgentServiceServer(s grpc.ServiceRegistrar, srv AgentServiceServer)
 }
 
 func _AgentService_Algo_Handler(srv interface{}, stream grpc.ServerStream) error {
-	return srv.(AgentServiceServer).Algo(&agentServiceAlgoServer{stream})
+	return srv.(AgentServiceServer).Algo(&agentServiceAlgoServer{ServerStream: stream})
 }
 
 type AgentService_AlgoServer interface {
@@ -199,7 +203,7 @@ func (x *agentServiceAlgoServer) Recv() (*AlgoRequest, error) {
 }
 
 func _AgentService_Data_Handler(srv interface{}, stream grpc.ServerStream) error {
-	return srv.(AgentServiceServer).Data(&agentServiceDataServer{stream})
+	return srv.(AgentServiceServer).Data(&agentServiceDataServer{ServerStream: stream})
 }
 
 type AgentService_DataServer interface {

agent/algorithm/algorithm.go

@@ -0,0 +1,9 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package algorithm
+
+// Algorithm is an interface that specifies the API for an algorithm.
+type Algorithm interface {
+	// Run executes the algorithm and returns the result.
+	Run() ([]byte, error)
+}

agent/algorithm/binary/binary.go

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

agent/algorithm/logging.go

@@ -0,0 +1,73 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package algorithm
+
+import (
+	"bytes"
+	"io"
+	"log/slog"
+
+	"github.com/ultravioletrs/cocos/agent/events"
+)
+
+var (
+	_ io.Writer = &Stdout{}
+	_ io.Writer = &Stderr{}
+)
+
+const bufSize = 1024
+
+type Stdout struct {
+	Logger *slog.Logger
+}
+
+// Write implements io.Writer.
+func (s *Stdout) Write(p []byte) (n int, err error) {
+	inBuf := bytes.NewBuffer(p)
+
+	buf := make([]byte, bufSize)
+
+	for {
+		n, err := inBuf.Read(buf)
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return len(p) - inBuf.Len(), err
+		}
+
+		s.Logger.Debug(string(buf[:n]))
+	}
+
+	return len(p), nil
+}
+
+type Stderr struct {
+	Logger   *slog.Logger
+	EventSvc events.Service
+}
+
+// Write implements io.Writer.
+func (s *Stderr) Write(p []byte) (n int, err error) {
+	inBuf := bytes.NewBuffer(p)
+
+	buf := make([]byte, bufSize)
+
+	for {
+		n, err := inBuf.Read(buf)
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return len(p) - inBuf.Len(), err
+		}
+
+		s.Logger.Error(string(buf[:n]))
+	}
+
+	if err := s.EventSvc.SendEvent("algorithm-run", "failed", nil); err != nil {
+		return len(p), err
+	}
+
+	return len(p), nil
+}

agent/algorithm/logging_test.go

@@ -0,0 +1,85 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package algorithm
+
+import (
+	"strings"
+	"testing"
+
+	mglog "github.com/absmach/magistrala/logger"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent/events/mocks"
+)
+
+func TestStdoutWrite(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected []string
+	}{
+		{
+			name:     "Single line",
+			input:    "Hello, World!",
+			expected: []string{"Hello, World!"},
+		},
+		{
+			name:     "Multiple lines",
+			input:    "Line 1\nLine 2\nLine 3",
+			expected: []string{"Line 1\nLine 2\nLine 3"},
+		},
+		{
+			name:     "Long input",
+			input:    strings.Repeat("a", bufSize+100),
+			expected: []string{strings.Repeat("a", bufSize), strings.Repeat("a", 100)},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			stdout := &Stdout{Logger: mglog.NewMock()}
+			n, err := stdout.Write([]byte(tt.input))
+
+			assert.NoError(t, err)
+			assert.Equal(t, len(tt.input), n)
+		})
+	}
+}
+
+func TestStderrWrite(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected []string
+	}{
+		{
+			name:     "Single line",
+			input:    "Error: Something went wrong",
+			expected: []string{"Error: Something went wrong"},
+		},
+		{
+			name:     "Multiple lines",
+			input:    "Error 1\nError 2\nError 3",
+			expected: []string{"Error 1\nError 2\nError 3"},
+		},
+		{
+			name:     "Long input",
+			input:    strings.Repeat("e", bufSize+100),
+			expected: []string{strings.Repeat("e", bufSize), strings.Repeat("e", 100)},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mockEventService := mocks.NewService(t)
+			mockEventService.On("SendEvent", "algorithm-run", "failed", mock.Anything).Return(nil)
+
+			stderr := &Stderr{Logger: mglog.NewMock(), EventSvc: mockEventService}
+			n, err := stderr.Write([]byte(tt.input))
+
+			assert.NoError(t, err)
+			assert.Equal(t, len(tt.input), n)
+			mockEventService.AssertExpectations(t)
+		})
+	}
+}

agent/api/grpc/interceptors.go

@@ -0,0 +1,67 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package grpc
+
+import (
+	"context"
+
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/agent/auth"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+type authInterceptor struct {
+	auth auth.Authenticator
+}
+
+type wrappedServerStream struct {
+	grpc.ServerStream
+	ctx context.Context
+}
+
+func (s *wrappedServerStream) Context() context.Context {
+	return s.ctx
+}
+
+func NewAuthInterceptor(authSvc auth.Authenticator) (grpc.UnaryServerInterceptor, grpc.StreamServerInterceptor) {
+	ai := &authInterceptor{auth: authSvc}
+	return ai.AuthUnaryInterceptor(), ai.AuthStreamInterceptor()
+}
+
+func (s *authInterceptor) AuthStreamInterceptor() grpc.StreamServerInterceptor {
+	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		switch info.FullMethod {
+		case agent.AgentService_Algo_FullMethodName:
+			if _, err := s.auth.AuthenticateUser(stream.Context(), auth.AlgorithmProviderRole); err != nil {
+				return status.Errorf(codes.Unauthenticated, err.Error())
+			}
+			return 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())
+			}
+			wrapped := &wrappedServerStream{ServerStream: stream, ctx: ctx}
+			return handler(srv, wrapped)
+		default:
+			return handler(srv, stream)
+		}
+	}
+}
+
+func (s *authInterceptor) AuthUnaryInterceptor() grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+		switch info.FullMethod {
+		case agent.AgentService_Result_FullMethodName:
+			ctx, err := s.auth.AuthenticateUser(ctx, auth.ConsumerRole)
+			if err != nil {
+				return nil, status.Errorf(codes.Unauthenticated, err.Error())
+			}
+			return handler(ctx, req)
+		default:
+			return handler(ctx, req)
+		}
+	}
+}

agent/api/grpc/interceptors_test.go

@@ -0,0 +1,152 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package grpc
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/agent/auth"
+	"github.com/ultravioletrs/cocos/agent/mocks"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestAuthUnaryInterceptor(t *testing.T) {
+	authmock := new(mocks.Authenticator)
+	tests := []struct {
+		name       string
+		authorized bool
+		method     string
+		role       auth.UserRole
+		wantErr    bool
+	}{
+		{
+			name:       "authorized result method",
+			authorized: true,
+			method:     agent.AgentService_Result_FullMethodName,
+			role:       auth.ConsumerRole,
+			wantErr:    false,
+		},
+		{
+			name:       "unauthorized result method",
+			authorized: false,
+			method:     agent.AgentService_Result_FullMethodName,
+			role:       auth.ConsumerRole,
+			wantErr:    true,
+		},
+		{
+			name:       "other method",
+			authorized: false,
+			method:     "OtherMethod",
+			role:       auth.ConsumerRole,
+			wantErr:    false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			switch tt.authorized {
+			case true:
+				mockCall := authmock.On("AuthenticateUser", context.Background(), tt.role).Return(context.Background(), nil)
+				mockCall.Once()
+			case false:
+				mockCall := authmock.On("AuthenticateUser", context.Background(), tt.role).Return(context.Background(), auth.ErrMissingMetadata)
+				mockCall.Once()
+			}
+			unaryInt, _ := NewAuthInterceptor(authmock)
+
+			_, err := unaryInt(context.Background(), nil, &grpc.UnaryServerInfo{FullMethod: tt.method}, func(ctx context.Context, req interface{}) (interface{}, error) {
+				return nil, nil
+			})
+
+			if tt.wantErr && err == nil {
+				t.Errorf("expected error, got none")
+			} else if !tt.wantErr && err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+		})
+	}
+}
+
+func TestAuthStreamInterceptor(t *testing.T) {
+	authmock := new(mocks.Authenticator)
+	tests := []struct {
+		name       string
+		authorized bool
+		method     string
+		role       auth.UserRole
+		wantErr    bool
+	}{
+		{
+			name:       "authorized algo method",
+			authorized: true,
+			method:     agent.AgentService_Algo_FullMethodName,
+			role:       auth.AlgorithmProviderRole,
+			wantErr:    false,
+		},
+		{
+			name:       "unauthorized algo method",
+			authorized: false,
+			method:     agent.AgentService_Algo_FullMethodName,
+			role:       auth.AlgorithmProviderRole,
+			wantErr:    true,
+		},
+		{
+			name:       "authorized data method",
+			authorized: true,
+			method:     agent.AgentService_Data_FullMethodName,
+			role:       auth.DataProviderRole,
+			wantErr:    false,
+		},
+		{
+			name:       "unauthorized data method",
+			authorized: false,
+			method:     agent.AgentService_Data_FullMethodName,
+			role:       auth.DataProviderRole,
+			wantErr:    true,
+		},
+		{
+			name:       "other method",
+			authorized: false,
+			method:     "OtherMethod",
+			role:       auth.DataProviderRole,
+			wantErr:    false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			switch tt.authorized {
+			case true:
+				mockCall := authmock.On("AuthenticateUser", mock.Anything, tt.role).Return(context.Background(), nil)
+				mockCall.Once()
+			case false:
+				mockCall := authmock.On("AuthenticateUser", mock.Anything, tt.role).Return(context.Background(), auth.ErrMissingMetadata)
+				mockCall.Once()
+			}
+			_, streamInt := NewAuthInterceptor(authmock)
+
+			err := streamInt(nil, &mockServerStream{ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs())}, &grpc.StreamServerInfo{FullMethod: tt.method}, func(srv interface{}, stream grpc.ServerStream) error {
+				return nil
+			})
+
+			if tt.wantErr && err == nil {
+				t.Errorf("expected error, got none")
+			} else if !tt.wantErr && err != nil {
+				t.Errorf("unexpected error: %v", err)
+			}
+		})
+	}
+}
+
+type mockServerStream struct {
+	grpc.ServerStream
+	ctx context.Context
+}
+
+func (m *mockServerStream) Context() context.Context {
+	return m.ctx
+}

agent/auth/auth.go

@@ -6,14 +6,15 @@ package auth
 import (
 	"context"
 	"crypto"
+	"crypto/ecdsa"
+	"crypto/ed25519"
 	"crypto/rsa"
 	"crypto/sha256"
 	"crypto/x509"
 	"encoding/base64"
-	"errors"
 
+	"github.com/absmach/magistrala/pkg/errors"
 	"github.com/ultravioletrs/cocos/agent"
-	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/status"
@@ -29,37 +30,37 @@ const (
 	AlgorithmProviderRole UserRole = "algorithm-provider"
 )
 
-var errNotRSAPublicKey = errors.New("not an RSA public key")
+var (
+	ErrMissingMetadata             = errors.New("missing metadata")
+	ErrInvalidMetadata             = errors.New("invalid metadata")
+	ErrSignatureVerificationFailed = errors.New("signature verification failed")
+)
 
-type wrappedServerStream struct {
-	grpc.ServerStream
-	ctx context.Context
+//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)
 }
 
-func (s *wrappedServerStream) Context() context.Context {
-	return s.ctx
+type service struct {
+	resultConsumers   []interface{}
+	datasetProviders  []interface{}
+	algorithmProvider interface{}
 }
 
-type Service struct {
-	resultConsumers   []*rsa.PublicKey
-	datasetProviders  []*rsa.PublicKey
-	algorithmProvider *rsa.PublicKey
-}
-
-func New(manifest agent.Computation) (*Service, error) {
-	s := &Service{}
+func New(manifest agent.Computation) (Authenticator, error) {
+	s := &service{}
 	for _, rc := range manifest.ResultConsumers {
 		pubKey, err := x509.ParsePKIXPublicKey(rc.UserKey)
 		if err != nil {
 			return nil, err
 		}
 
-		rsaPubKey, ok := pubKey.(*rsa.PublicKey)
-		if !ok {
-			return nil, errNotRSAPublicKey
+		pKey, err := decodePublicKey(pubKey)
+		if err != nil {
+			return nil, err
 		}
 
-		s.resultConsumers = append(s.resultConsumers, rsaPubKey)
+		s.resultConsumers = append(s.resultConsumers, pKey)
 	}
 
 	for _, dp := range manifest.Datasets {
@@ -68,12 +69,12 @@ func New(manifest agent.Computation) (*Service, error) {
 			return nil, err
 		}
 
-		rsaPubKey, ok := pubKey.(*rsa.PublicKey)
-		if !ok {
-			return nil, errNotRSAPublicKey
+		pKey, err := decodePublicKey(pubKey)
+		if err != nil {
+			return nil, err
 		}
 
-		s.datasetProviders = append(s.datasetProviders, rsaPubKey)
+		s.datasetProviders = append(s.datasetProviders, pKey)
 	}
 
 	pubKey, err := x509.ParsePKIXPublicKey(manifest.Algorithm.UserKey)
@@ -81,82 +82,15 @@ func New(manifest agent.Computation) (*Service, error) {
 		return nil, err
 	}
 
-	rsaPubKey, ok := pubKey.(*rsa.PublicKey)
-	if !ok {
-		return nil, errNotRSAPublicKey
+	pKey, err := decodePublicKey(pubKey)
+	if err != nil {
+		return nil, err
 	}
 
-	s.algorithmProvider = rsaPubKey
+	s.algorithmProvider = pKey
 	return s, nil
 }
 
-func (s *Service) AuthStreamInterceptor() grpc.StreamServerInterceptor {
-	return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
-		switch info.FullMethod {
-		case agent.AgentService_Algo_FullMethodName:
-			md, ok := metadata.FromIncomingContext(stream.Context())
-			if !ok {
-				return status.Errorf(codes.Unauthenticated, "missing metadata")
-			}
-			signature, err := extractSignature(md)
-			if err != nil {
-				return status.Errorf(codes.Unauthenticated, "invalid metadata")
-			}
-			isValid, err := verifySignature(AlgorithmProviderRole, signature, s.algorithmProvider)
-			if err != nil || !isValid {
-				return status.Errorf(codes.Unauthenticated, "signature verification failed")
-			}
-		case agent.AgentService_Data_FullMethodName:
-			md, ok := metadata.FromIncomingContext(stream.Context())
-			if !ok {
-				return status.Errorf(codes.Unauthenticated, "missing metadata")
-			}
-			signature, err := extractSignature(md)
-			if err != nil {
-				return status.Errorf(codes.Unauthenticated, "invalid metadata")
-			}
-			for index, dp := range s.datasetProviders {
-				isValid, err := verifySignature(DataProviderRole, signature, dp)
-				if err == nil || isValid {
-					ctx := agent.IndexToContext(stream.Context(), index)
-					wrapped := &wrappedServerStream{ServerStream: stream, ctx: ctx}
-					return handler(srv, wrapped)
-				}
-			}
-			return status.Errorf(codes.Unauthenticated, "signature verification failed")
-		default:
-			return handler(srv, stream)
-		}
-		return handler(srv, stream)
-	}
-}
-
-func (s *Service) AuthUnaryInterceptor() grpc.UnaryServerInterceptor {
-	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-		switch info.FullMethod {
-		case agent.AgentService_Result_FullMethodName:
-			md, ok := metadata.FromIncomingContext(ctx)
-			if !ok {
-				return nil, status.Errorf(codes.Unauthenticated, "missing metadata")
-			}
-			signature, err := extractSignature(md)
-			if err != nil {
-				return nil, status.Errorf(codes.Unauthenticated, "invalid metadata")
-			}
-			for index, rc := range s.resultConsumers {
-				isValid, err := verifySignature(ConsumerRole, signature, rc)
-				if err == nil || isValid {
-					ctx := agent.IndexToContext(ctx, index)
-					return handler(ctx, req)
-				}
-			}
-			return nil, status.Errorf(codes.Unauthenticated, "signature verification failed")
-		default:
-			return handler(ctx, req)
-		}
-	}
-}
-
 func extractSignature(md metadata.MD) (string, error) {
 	signature := md.Get(SignatureMetadataKey)
 	if len(signature) != 1 {
@@ -166,15 +100,75 @@ func extractSignature(md metadata.MD) (string, error) {
 	return signature[0], nil
 }
 
-func verifySignature(role UserRole, signature string, publicKey *rsa.PublicKey) (bool, error) {
+func verifySignature(role UserRole, signature string, publicKey any) error {
 	hash := sha256.Sum256([]byte(role))
 	sigByte, err := base64.StdEncoding.DecodeString(signature)
 	if err != nil {
-		return false, err
+		return err
 	}
 
-	if err := rsa.VerifyPKCS1v15(publicKey, crypto.SHA256, hash[:], sigByte); err != nil {
-		return false, err
+	var ok bool
+
+	switch publicKey := publicKey.(type) {
+	case *rsa.PublicKey:
+		if err = rsa.VerifyPKCS1v15(publicKey, crypto.SHA256, hash[:], sigByte); err != nil {
+			return err
+		}
+		return nil
+	case *ecdsa.PublicKey:
+		ok = ecdsa.VerifyASN1(publicKey, hash[:], sigByte)
+	case ed25519.PublicKey:
+		ok = ed25519.Verify(publicKey, []byte(role), sigByte)
+	}
+
+	if !ok {
+		return ErrSignatureVerificationFailed
+	}
+
+	return nil
+}
+
+func (s *service) AuthenticateUser(ctx context.Context, role UserRole) (context.Context, error) {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		return nil, ErrMissingMetadata
+	}
+	signature, err := extractSignature(md)
+	if err != nil {
+		return nil, errors.Wrap(err, ErrInvalidMetadata)
+	}
+
+	switch role {
+	case ConsumerRole:
+		for i, rc := range s.resultConsumers {
+			if err := verifySignature(role, signature, rc); err == nil {
+				return agent.IndexToContext(ctx, i), nil
+			}
+		}
+	case DataProviderRole:
+		for i, dp := range s.datasetProviders {
+			if err := verifySignature(role, signature, dp); err == nil {
+				return agent.IndexToContext(ctx, i), nil
+			}
+		}
+	case AlgorithmProviderRole:
+		if err := verifySignature(role, signature, s.algorithmProvider); err == nil {
+			return ctx, nil
+		}
+	}
+
+	return ctx, ErrSignatureVerificationFailed
+}
+
+func decodePublicKey(key any) (pubKey any, err error) {
+	switch key := key.(type) {
+	case *rsa.PublicKey:
+		return key, nil
+	case *ecdsa.PublicKey:
+		return key, nil
+	case ed25519.PublicKey:
+		return key, nil
+	default:
+		return nil, errors.New("unsupported public key type")
 	}
-	return true, nil
 }

agent/auth/auth_test.go

@@ -0,0 +1,158 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package auth
+
+import (
+	"context"
+	"crypto"
+	"crypto/ecdsa"
+	"crypto/ed25519"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/sha256"
+	"crypto/x509"
+	"encoding/base64"
+	"testing"
+
+	"github.com/absmach/magistrala/pkg/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/agent"
+	"google.golang.org/grpc/metadata"
+)
+
+func TestAuthenticateUser(t *testing.T) {
+	resultConsumerKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	require.NoError(t, err)
+
+	dataProviderKey, err := rsa.GenerateKey(rand.Reader, 2048)
+	require.NoError(t, err)
+
+	pubEd25519Key, algorithmProviderKey, err := ed25519.GenerateKey(rand.Reader)
+	require.NoError(t, err)
+
+	resultConsumerPubKey, err := x509.MarshalPKIXPublicKey(&resultConsumerKey.PublicKey)
+	require.NoError(t, err)
+
+	dataProviderPubKey, err := x509.MarshalPKIXPublicKey(&dataProviderKey.PublicKey)
+	require.NoError(t, err)
+
+	algorithmProviderPubKey, err := x509.MarshalPKIXPublicKey(pubEd25519Key)
+	require.NoError(t, err)
+
+	manifest := agent.Computation{
+		ResultConsumers: []agent.ResultConsumer{{UserKey: resultConsumerPubKey}},
+		Datasets:        []agent.Dataset{{UserKey: dataProviderPubKey}},
+		Algorithm:       agent.Algorithm{UserKey: algorithmProviderPubKey},
+	}
+
+	auth, err := New(manifest)
+	if err != nil {
+		t.Fatalf("failed to create authenticator: %v", err)
+	}
+
+	testCases := []struct {
+		name        string
+		role        UserRole
+		key         any
+		expectedErr error
+	}{
+		{
+			name:        "valid result consumer",
+			role:        ConsumerRole,
+			key:         resultConsumerKey,
+			expectedErr: nil,
+		},
+		{
+			name:        "valid data provider",
+			role:        DataProviderRole,
+			key:         dataProviderKey,
+			expectedErr: nil,
+		},
+		{
+			name:        "valid algorithm provider",
+			role:        AlgorithmProviderRole,
+			key:         algorithmProviderKey,
+			expectedErr: nil,
+		},
+		{
+			name:        "invalid role",
+			role:        "invalid-role",
+			key:         resultConsumerKey,
+			expectedErr: ErrSignatureVerificationFailed,
+		},
+		{
+			name:        "invalid key",
+			role:        ConsumerRole,
+			key:         dataProviderKey,
+			expectedErr: ErrSignatureVerificationFailed,
+		},
+		{
+			name:        "missing signature",
+			role:        ConsumerRole,
+			key:         resultConsumerKey,
+			expectedErr: ErrInvalidMetadata,
+		},
+		{
+			name:        "missing metadata",
+			role:        ConsumerRole,
+			key:         resultConsumerKey,
+			expectedErr: ErrMissingMetadata,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			signature, err := signRole(tc.role, tc.key)
+			if err != nil {
+				t.Fatalf("failed to sign role: %v", err)
+			}
+
+			ctx := context.Background()
+
+			switch tc.name {
+			case "missing signature":
+				ctx = metadata.NewIncomingContext(ctx, metadata.Pairs())
+			case "missing metadata":
+			default:
+				ctx = metadata.NewIncomingContext(ctx, metadata.Pairs(SignatureMetadataKey, signature))
+			}
+
+			ctx, err = auth.AuthenticateUser(ctx, tc.role)
+			assert.True(t, errors.Contains(err, tc.expectedErr), "expected error %v, got %v", tc.expectedErr, err)
+
+			if err == nil {
+				switch id, ok := agent.IndexFromContext(ctx); {
+				case tc.role == ConsumerRole, tc.role == DataProviderRole:
+					assert.True(t, ok, "expected index in context")
+					assert.Equal(t, 0, id, "expected index 0 in context")
+				default:
+					assert.False(t, ok, "expected no index in context")
+				}
+			}
+		})
+	}
+}
+
+func signRole(role UserRole, key crypto.PrivateKey) (string, error) {
+	var signature []byte
+	var err error
+
+	switch k := key.(type) {
+	case ed25519.PrivateKey:
+		signature, err = k.Sign(rand.Reader, []byte(role), crypto.Hash(0))
+	case *rsa.PrivateKey, *ecdsa.PrivateKey:
+		hash := sha256.Sum256([]byte(role))
+		signer := key.(crypto.Signer)
+		signature, err = signer.Sign(rand.Reader, hash[:], crypto.SHA256)
+	default:
+		return "", errors.New("unsupported key type")
+	}
+
+	if err != nil {
+		return "", err
+	}
+
+	return base64.StdEncoding.EncodeToString(signature), nil
+}

agent/events/events.go

@@ -27,6 +27,7 @@ type AgentEvent struct {
 	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

agent/events/mocks/events.go

@@ -0,0 +1,67 @@
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package mocks
+
+import (
+	json "encoding/json"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// Service is an autogenerated mock type for the Service type
+type Service struct {
+	mock.Mock
+}
+
+// Close provides a mock function with given fields:
+func (_m *Service) Close() error {
+	ret := _m.Called()
+
+	if len(ret) == 0 {
+		panic("no return value specified for Close")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func() error); ok {
+		r0 = rf()
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// SendEvent provides a mock function with given fields: event, status, details
+func (_m *Service) SendEvent(event string, status string, details json.RawMessage) error {
+	ret := _m.Called(event, status, details)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SendEvent")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(string, string, json.RawMessage) error); ok {
+		r0 = rf(event, status, details)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewService(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Service {
+	mock := &Service{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

agent/mocks/agent.go

@@ -0,0 +1,129 @@
+// Code generated by mockery v2.42.3. DO NOT EDIT.
+
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package mocks
+
+import (
+	context "context"
+
+	agent "github.com/ultravioletrs/cocos/agent"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// Service is an autogenerated mock type for the Service type
+type Service struct {
+	mock.Mock
+}
+
+// Algo provides a mock function with given fields: ctx, algorithm
+func (_m *Service) Algo(ctx context.Context, algorithm agent.Algorithm) error {
+	ret := _m.Called(ctx, algorithm)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Algo")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, agent.Algorithm) error); ok {
+		r0 = rf(ctx, algorithm)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Attestation provides a mock function with given fields: ctx, reportData
+func (_m *Service) Attestation(ctx context.Context, reportData [64]byte) ([]byte, error) {
+	ret := _m.Called(ctx, reportData)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Attestation")
+	}
+
+	var r0 []byte
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, [64]byte) ([]byte, error)); ok {
+		return rf(ctx, reportData)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, [64]byte) []byte); ok {
+		r0 = rf(ctx, reportData)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]byte)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, [64]byte) error); ok {
+		r1 = rf(ctx, reportData)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// Data provides a mock function with given fields: ctx, dataset
+func (_m *Service) Data(ctx context.Context, dataset agent.Dataset) error {
+	ret := _m.Called(ctx, dataset)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Data")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(context.Context, agent.Dataset) error); ok {
+		r0 = rf(ctx, dataset)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Result provides a mock function with given fields: ctx
+func (_m *Service) Result(ctx context.Context) ([]byte, error) {
+	ret := _m.Called(ctx)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Result")
+	}
+
+	var r0 []byte
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context) ([]byte, error)); ok {
+		return rf(ctx)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context) []byte); ok {
+		r0 = rf(ctx)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).([]byte)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context) error); ok {
+		r1 = rf(ctx)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// NewService creates a new instance of Service. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewService(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Service {
+	mock := &Service{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

agent/mocks/auth.go

@@ -0,0 +1,63 @@
+// Code generated by mockery v2.42.3. DO NOT EDIT.
+
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package mocks
+
+import (
+	context "context"
+
+	auth "github.com/ultravioletrs/cocos/agent/auth"
+
+	mock "github.com/stretchr/testify/mock"
+)
+
+// Authenticator is an autogenerated mock type for the Authenticator type
+type Authenticator struct {
+	mock.Mock
+}
+
+// AuthenticateUser provides a mock function with given fields: ctx, role
+func (_m *Authenticator) AuthenticateUser(ctx context.Context, role auth.UserRole) (context.Context, error) {
+	ret := _m.Called(ctx, role)
+
+	if len(ret) == 0 {
+		panic("no return value specified for AuthenticateUser")
+	}
+
+	var r0 context.Context
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, auth.UserRole) (context.Context, error)); ok {
+		return rf(ctx, role)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, auth.UserRole) context.Context); ok {
+		r0 = rf(ctx, role)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(context.Context)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, auth.UserRole) error); ok {
+		r1 = rf(ctx, role)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// NewAuthenticator creates a new instance of Authenticator. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewAuthenticator(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Authenticator {
+	mock := &Authenticator{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

agent/service.go

@@ -4,19 +4,17 @@
 package agent
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"log/slog"
 	"os"
-	"os/exec"
 	"slices"
 
 	"github.com/google/go-sev-guest/client"
+	"github.com/ultravioletrs/cocos/agent/algorithm/binary"
 	"github.com/ultravioletrs/cocos/agent/events"
-	"github.com/ultravioletrs/cocos/pkg/socket"
 	"golang.org/x/crypto/sha3"
 )
 
@@ -25,7 +23,6 @@ var _ Service = (*agentService)(nil)
 const (
 	// ReportDataSize is the size of the report data expected by the attestation service.
 	ReportDataSize     = 64
-	socketPath         = "unix_socket"
 	algoFilePermission = 0o700
 )
 
@@ -37,21 +34,23 @@ var (
 	// when accessing a protected resource.
 	ErrUnauthorizedAccess = errors.New("missing or invalid credentials provided")
 	// errUndeclaredAlgorithm indicates algorithm was not declared in computation manifest.
-	errUndeclaredDataset = errors.New("dataset not declared in computation manifest")
+	ErrUndeclaredDataset = errors.New("dataset not declared in computation manifest")
 	// errAllManifestItemsReceived indicates no new computation manifest items expected.
-	errAllManifestItemsReceived = errors.New("all expected manifest Items have been received")
+	ErrAllManifestItemsReceived = errors.New("all expected manifest Items have been received")
 	// errUndeclaredConsumer indicates the consumer requesting results in not declared in computation manifest.
-	errUndeclaredConsumer = errors.New("result consumer is undeclared in computation manifest")
+	ErrUndeclaredConsumer = errors.New("result consumer is undeclared in computation manifest")
 	// errResultsNotReady indicates the computation results are not ready.
-	errResultsNotReady = errors.New("computation results are not yet ready")
+	ErrResultsNotReady = errors.New("computation results are not yet ready")
 	// errStateNotReady agent received a request in the wrong state.
-	errStateNotReady = errors.New("agent not expecting this operation in the current state")
+	ErrStateNotReady = errors.New("agent not expecting this operation in the current state")
 	// errHashMismatch provided algorithm/dataset does not match hash in manifest.
-	errHashMismatch = errors.New("malformed data, hash does not match manifest")
+	ErrHashMismatch = errors.New("malformed data, hash does not match manifest")
 )
 
 // 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
@@ -95,16 +94,16 @@ func New(ctx context.Context, logger *slog.Logger, eventSvc events.Service, cmp
 
 func (as *agentService) Algo(ctx context.Context, algorithm Algorithm) error {
 	if as.sm.GetState() != receivingAlgorithm {
-		return errStateNotReady
+		return ErrStateNotReady
 	}
 	if as.algorithm != "" {
-		return errAllManifestItemsReceived
+		return ErrAllManifestItemsReceived
 	}
 
 	hash := sha3.Sum256(algorithm.Algorithm)
 
 	if hash != as.computation.Algorithm.Hash {
-		return errHashMismatch
+		return ErrHashMismatch
 	}
 
 	f, err := os.CreateTemp("", "algorithm")
@@ -135,21 +134,21 @@ func (as *agentService) Algo(ctx context.Context, algorithm Algorithm) error {
 
 func (as *agentService) Data(ctx context.Context, dataset Dataset) error {
 	if as.sm.GetState() != receivingData {
-		return errStateNotReady
+		return ErrStateNotReady
 	}
 	if len(as.computation.Datasets) == 0 {
-		return errAllManifestItemsReceived
+		return ErrAllManifestItemsReceived
 	}
 
 	hash := sha3.Sum256(dataset.Dataset)
 
 	index, ok := IndexFromContext(ctx)
 	if !ok {
-		return errUndeclaredDataset
+		return ErrUndeclaredDataset
 	}
 
 	if hash != as.computation.Datasets[index].Hash {
-		return errHashMismatch
+		return ErrHashMismatch
 	}
 	as.computation.Datasets = slices.Delete(as.computation.Datasets, index, index+1)
 
@@ -176,14 +175,14 @@ func (as *agentService) Data(ctx context.Context, dataset Dataset) error {
 
 func (as *agentService) Result(ctx context.Context) ([]byte, error) {
 	if as.sm.GetState() != resultsReady {
-		return []byte{}, errResultsNotReady
+		return []byte{}, ErrResultsNotReady
 	}
 	if len(as.computation.ResultConsumers) == 0 {
-		return []byte{}, errAllManifestItemsReceived
+		return []byte{}, ErrAllManifestItemsReceived
 	}
 	index, ok := IndexFromContext(ctx)
 	if !ok {
-		return []byte{}, errUndeclaredConsumer
+		return []byte{}, ErrUndeclaredConsumer
 	}
 	as.computation.ResultConsumers = slices.Delete(as.computation.ResultConsumers, index, index+1)
 
@@ -212,7 +211,8 @@ func (as *agentService) runComputation() {
 	as.sm.logger.Debug("computation run started")
 	defer as.sm.SendEvent(runComplete)
 	as.publishEvent("in-progress", json.RawMessage{})()
-	result, err := as.run(as.algorithm, as.datasets)
+	algorithm := binary.New(as.sm.logger, as.eventSvc, as.algorithm, as.datasets...)
+	result, err := algorithm.Run()
 	if err != nil {
 		as.runError = err
 		as.sm.logger.Warn(fmt.Sprintf("computation failed with error: %s", err.Error()))
@@ -230,49 +230,3 @@ func (as *agentService) publishEvent(status string, details json.RawMessage) fun
 		}
 	}
 }
-
-func (as *agentService) run(algoFile string, dataFiles []string) ([]byte, error) {
-	defer os.Remove(algoFile)
-	defer func() {
-		for _, file := range dataFiles {
-			os.Remove(file)
-		}
-	}()
-	listener, err := socket.StartUnixSocketServer(socketPath)
-	if err != nil {
-		return nil, fmt.Errorf("error creating stdout pipe: %v", err)
-	}
-	defer listener.Close()
-
-	// Create channels for received data and errors
-	dataChannel := make(chan []byte)
-	errorChannel := make(chan error)
-
-	var result []byte
-
-	var outStd, outErr bytes.Buffer
-
-	go socket.AcceptConnection(listener, dataChannel, errorChannel)
-
-	args := append([]string{socketPath}, dataFiles...)
-	cmd := exec.Command(algoFile, args...)
-	cmd.Stderr = &outErr
-	cmd.Stdout = &outStd
-
-	if err := cmd.Start(); err != nil {
-		return nil, fmt.Errorf("error starting algorithm: %v", err)
-	}
-
-	if err := cmd.Wait(); err != nil {
-		as.sm.logger.Debug(outErr.String())
-		return nil, fmt.Errorf("algorithm execution error: %v", err)
-	}
-
-	select {
-	case result = <-dataChannel:
-		as.sm.logger.Debug(outStd.String())
-		return result, nil
-	case err = <-errorChannel:
-		return nil, fmt.Errorf("error receiving data: %v", err)
-	}
-}

cli/algorithms.go

@@ -3,7 +3,6 @@
 package cli
 
 import (
-	"crypto/x509"
 	"encoding/pem"
 	"log"
 	"os"
@@ -39,10 +38,7 @@ func (cli *CLI) NewAlgorithmCmd() *cobra.Command {
 
 			pemBlock, _ := pem.Decode(privKeyFile)
 
-			privKey, err := x509.ParsePKCS1PrivateKey(pemBlock.Bytes)
-			if err != nil {
-				log.Fatalf("Error parsing private key: %v", err)
-			}
+			privKey := decodeKey(pemBlock)
 
 			if err := cli.agentSDK.Algo(cmd.Context(), algoReq, privKey); err != nil {
 				log.Fatalf("Error uploading algorithm with error: %v", err)

cli/backend_info.go

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

cli/cache.go

@@ -0,0 +1,69 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package cli
+
+import (
+	"log"
+	"os"
+	"path"
+
+	"github.com/google/go-sev-guest/abi"
+	"github.com/google/go-sev-guest/kds"
+	"github.com/google/go-sev-guest/verify/trust"
+	"github.com/spf13/cobra"
+	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
+)
+
+const (
+	caBundleName      = "ask_ark.pem"
+	filePermisionKeys = 0o766
+)
+
+func (cli *CLI) NewCABundleCmd(fileSavePath string) *cobra.Command {
+	return &cobra.Command{
+		Use:     "ca-bundle",
+		Short:   "Fetch AMD SEV-SNPs CA Bundle (ASK and ARK)",
+		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)
+			if err != nil {
+				log.Fatalf("Error while reading manifest: %v", err)
+			}
+
+			product := attestationConfiguration.RootOfTrust.Product
+
+			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)
+			}
+
+			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)
+			}
+
+			bundleFilePath := path.Join(fileSavePath, product, caBundleName)
+			if err = saveToFile(bundleFilePath, bundle); err != nil {
+				log.Fatalf("Error while saving ARK-ASK to file: %v", err)
+			}
+		},
+	}
+}
+
+func saveToFile(fileSavePath string, content []byte) error {
+	file, err := os.OpenFile(fileSavePath, os.O_CREATE|os.O_RDWR|os.O_TRUNC, filePermisionKeys)
+	if err != nil {
+		return err
+	}
+
+	if _, err := file.Write(content); err != nil {
+		return err
+	}
+
+	return nil
+}

cli/datasets.go

@@ -39,10 +39,7 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
 
 			pemBlock, _ := pem.Decode(privKeyFile)
 
-			privKey, err := x509.ParsePKCS1PrivateKey(pemBlock.Bytes)
-			if err != nil {
-				log.Fatalf("Error parsing private key: %v", err)
-			}
+			privKey := decodeKey(pemBlock)
 
 			if err := cli.agentSDK.Data(cmd.Context(), dataReq, privKey); err != nil {
 				log.Fatalf("Error uploading dataset: %v", err)
@@ -52,3 +49,26 @@ func (cli *CLI) NewDatasetsCmd() *cobra.Command {
 		},
 	}
 }
+
+func decodeKey(b *pem.Block) interface{} {
+	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 privKey
+	case ecdsaKeyType:
+		privKey, err := x509.ParseECPrivateKey(b.Bytes)
+		if err != nil {
+			log.Fatalf("Error parsing private key: %v", err)
+		}
+		return privKey
+	default:
+		log.Fatalf("Error decoding key")
+		return nil
+	}
+}

cli/keys.go

@@ -3,66 +3,126 @@
 package cli
 
 import (
+	"crypto/ecdsa"
+	"crypto/ed25519"
+	"crypto/elliptic"
 	"crypto/rand"
 	"crypto/rsa"
 	"crypto/x509"
 	"encoding/pem"
 	"log"
 	"os"
+	"reflect"
 
 	"github.com/spf13/cobra"
 )
 
 const (
 	keyBitSize     = 4096
-	privateKeyType = "RSA PRIVATE KEY"
+	rsaKeyType     = "PRIVATE KEY"
+	ecdsaKeyType   = "EC PRIVATE KEY"
+	ed25519KeyType = "PRIVATE KEY"
 	publicKeyType  = "PUBLIC KEY"
 	publicKeyFile  = "public.pem"
 	privateKeyFile = "private.pem"
 )
 
+var KeyType string
+
 func (cli *CLI) NewKeysCmd() *cobra.Command {
 	return &cobra.Command{
 		Use:   "keys",
 		Short: "Generate a new public/private key pair",
+		Long: "Generates a new public/private key pair using an algorithm of the users choice.\n" +
+			"Supported algorithms are RSA, ecdsa, and ed25519.",
+		Example: "./build/cocos-cli keys -k rsa",
+		Args:    cobra.ExactArgs(0),
 		Run: func(cmd *cobra.Command, args []string) {
-			privKey, err := rsa.GenerateKey(rand.Reader, keyBitSize)
-			if err != nil {
-				log.Fatalf("Error generating public key: %v", err)
-			}
+			switch KeyType {
+			case "ecdsa":
+				privEcdsaKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+				if err != nil {
+					log.Fatalf("Error generating keys: %v", err)
+				}
 
-			pubKey, err := x509.MarshalPKIXPublicKey(&privKey.PublicKey)
-			if err != nil {
-				log.Fatalf("Error marshalling public key: %v", err)
-			}
+				pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privEcdsaKey.PublicKey)
+				if err != nil {
+					log.Fatalf("Error marshalling public key: %v", err)
+				}
 
-			privFile, err := os.Create(privateKeyFile)
-			if err != nil {
-				log.Fatalf("Error creating private key file: %v", err)
-			}
-			defer privFile.Close()
+				generateAndWriteKeys(privEcdsaKey, pubKeyBytes, ecdsaKeyType)
 
-			if err := pem.Encode(privFile, &pem.Block{
-				Type:  privateKeyType,
-				Bytes: x509.MarshalPKCS1PrivateKey(privKey),
-			}); err != nil {
-				log.Fatalf("Error encoding private key: %v", err)
-			}
+			case "ed25519":
+				pubEd25519Key, privEd25519Key, err := ed25519.GenerateKey(rand.Reader)
+				if err != nil {
+					log.Fatalf("Error generating keys: %v", err)
+				}
+				pubKey, err := x509.MarshalPKIXPublicKey(pubEd25519Key)
+				if err != nil {
+					log.Fatalf("Error marshalling public key: %v", err)
+				}
+				generateAndWriteKeys(privEd25519Key, pubKey, ed25519KeyType)
 
-			pubFile, err := os.Create(publicKeyFile)
-			if err != nil {
-				log.Fatalf("Error creating public key file: %v", err)
-			}
-			defer pubFile.Close()
+			// Default to RSA
+			default:
+				privKey, err := rsa.GenerateKey(rand.Reader, keyBitSize)
+				if err != nil {
+					log.Fatalf("Error generating keys: %v", err)
+				}
 
-			if err := pem.Encode(pubFile, &pem.Block{
-				Type:  publicKeyType,
-				Bytes: pubKey,
-			}); err != nil {
-				log.Fatalf("Error encoding public key: %v", err)
+				pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privKey.PublicKey)
+				if err != nil {
+					log.Fatalf("Error marshalling public key: %v", err)
+				}
+				generateAndWriteKeys(privKey, pubKeyBytes, rsaKeyType)
 			}
-
-			log.Println("Successfully generated public/private key pair")
 		},
 	}
 }
+
+func generateAndWriteKeys(privKey interface{}, pubKeyBytes []byte, keyType string) {
+	privFile, err := os.Create(privateKeyFile)
+	if err != nil {
+		log.Fatalf("Error creating private key file: %v", err)
+	}
+	defer privFile.Close()
+
+	var b []byte
+	switch privKey := privKey.(type) {
+	case *rsa.PrivateKey:
+		b = x509.MarshalPKCS1PrivateKey(privKey)
+	case *ecdsa.PrivateKey:
+		b, err = x509.MarshalECPrivateKey(privKey)
+	case ed25519.PrivateKey:
+		b, err = x509.MarshalPKCS8PrivateKey(privKey)
+	}
+	if err != nil {
+		log.Printf("Error marshalling private key: %v", err)
+		return
+	}
+
+	if err := pem.Encode(privFile, &pem.Block{
+		Type:  keyType,
+		Bytes: b,
+	}); err != nil {
+		log.Printf("Error encoding private key: %v", err)
+		return
+	}
+
+	pubFile, err := os.Create(publicKeyFile)
+	if err != nil {
+		log.Printf("Error creating public key file: %v", err)
+		return
+	}
+	defer pubFile.Close()
+
+	if err := pem.Encode(pubFile, &pem.Block{
+		Type:  publicKeyType,
+		Bytes: pubKeyBytes,
+	}); err != nil {
+		log.Printf("Error encoding public key: %v", err)
+		return
+	}
+
+	log.Printf("Successfully generated public/private key pair of type: %s", reflect.TypeOf(privKey).String())
+}

cli/measurement.go

@@ -1,62 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package cli
-
-import (
-	"encoding/base64"
-	"encoding/json"
-	"log"
-	"os"
-
-	"github.com/google/go-sev-guest/proto/check"
-	"github.com/spf13/cobra"
-)
-
-const filePermision = 0o755
-
-type AttestationConfiguration struct {
-	SNPPolicy   *check.Policy      `json:"snp_policy,omitempty"`
-	RootOFTrust *check.RootOfTrust `json:"root_of_trust,omitempty"`
-}
-
-func (cli *CLI) NewAddMeasurementCmd() *cobra.Command {
-	return &cobra.Command{
-		Use:     "measurement",
-		Short:   "Add measurement to the platform info file. The value should be in base64. The second parameter is platform_info.json file",
-		Example: "measurement <measurement> <platform_info.json>",
-		Args:    cobra.ExactArgs(2),
-		Run: func(cmd *cobra.Command, args []string) {
-			measurement, err := base64.StdEncoding.DecodeString(args[0])
-			if err != nil {
-				log.Fatalf("Error could not decode base64: %v", err)
-			}
-
-			attestationConfiguration := AttestationConfiguration{}
-
-			manifest, err := os.OpenFile(args[1], os.O_RDWR, filePermision)
-			if err != nil {
-				log.Fatalf("Error opening the platform information file: %v", err)
-			}
-			defer manifest.Close()
-
-			decoder := json.NewDecoder(manifest)
-			err = decoder.Decode(&attestationConfiguration)
-			if err != nil {
-				log.Fatalf("Error decoding the platform information file: %v", err)
-			}
-
-			attestationConfiguration.SNPPolicy.Measurement = measurement
-			if err = manifest.Truncate(0); err != nil {
-				log.Fatalf("Error could not truncate platform information JSON file: %v", err)
-			}
-
-			fileJson, err := json.MarshalIndent(attestationConfiguration, "", " ")
-			if err != nil {
-				log.Fatalf("Error marshaling the platform information JSON: %v", err)
-			}
-			if err = os.WriteFile(manifest.Name(), fileJson, filePermision); err != nil {
-				log.Fatalf("Error writing into platform information JSON file: %v", err)
-			}
-		},
-	}
-}

cli/result.go

@@ -3,7 +3,6 @@
 package cli
 
 import (
-	"crypto/x509"
 	"encoding/pem"
 	"log"
 	"os"
@@ -29,12 +28,10 @@ func (cli *CLI) NewResultsCmd() *cobra.Command {
 
 			pemBlock, _ := pem.Decode(privKeyFile)
 
-			privKey, err := x509.ParsePKCS1PrivateKey(pemBlock.Bytes)
-			if err != nil {
-				log.Fatalf("Error parsing private key: %v", err)
-			}
+			var result []byte
 
-			result, err := cli.agentSDK.Result(cmd.Context(), privKey)
+			privKey := decodeKey(pemBlock)
+			result, err = cli.agentSDK.Result(cmd.Context(), privKey)
 			if err != nil {
 				log.Fatalf("Error retrieving computation result: %v", err)
 			}

cmd/agent/main.go

@@ -10,17 +10,18 @@ import (
 	"log"
 	"log/slog"
 
+	"github.com/absmach/magistrala/pkg/prometheus"
 	"github.com/mdlayher/vsock"
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/agent/api"
 	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
 	"github.com/ultravioletrs/cocos/agent/auth"
 	"github.com/ultravioletrs/cocos/agent/events"
-	"github.com/ultravioletrs/cocos/internal"
 	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"
+	"github.com/ultravioletrs/cocos/manager/qemu"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/reflection"
@@ -45,6 +46,7 @@ func main() {
 		log.Fatal(err)
 	}
 	defer conn.Close()
+
 	var level slog.Level
 	if err := level.UnmarshalText([]byte(cfg.AgentConfig.LogLevel)); err != nil {
 		log.Println(err)
@@ -59,6 +61,7 @@ func main() {
 		return
 	}
 	defer eventSvc.Close()
+
 	svc := newService(ctx, logger, eventSvc, cfg)
 
 	grpcServerConfig := server.Config{
@@ -81,6 +84,7 @@ func main() {
 		logger.Error(fmt.Sprintf("failed to create auth service %s", err.Error()))
 		return
 	}
+
 	gs := grpcserver.New(ctx, cancel, svcName, grpcServerConfig, registerAgentServiceServer, logger, svc, authSvc)
 
 	g.Go(func() error {
@@ -100,18 +104,19 @@ func newService(ctx context.Context, logger *slog.Logger, eventSvc events.Servic
 	svc := agent.New(ctx, logger, eventSvc, cmp)
 
 	svc = api.LoggingMiddleware(svc, logger)
-	counter, latency := internal.MakeMetrics(svcName, "api")
+	counter, latency := prometheus.MakeMetrics(svcName, "api")
 	svc = api.MetricsMiddleware(svc, counter, latency)
 
 	return svc
 }
 
 func readConfig() (agent.Computation, error) {
-	l, err := vsock.Listen(manager.VsockConfigPort, nil)
+	l, err := vsock.Listen(qemu.VsockConfigPort, nil)
 	if err != nil {
 		return agent.Computation{}, err
 	}
 	defer l.Close()
+
 	conn, err := l.Accept()
 	if err != nil {
 		return agent.Computation{}, err

cmd/cli/main.go

@@ -6,12 +6,13 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"path"
 
 	mglog "github.com/absmach/magistrala/logger"
+	"github.com/caarlos0/env/v11"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 	"github.com/ultravioletrs/cocos/cli"
-	"github.com/ultravioletrs/cocos/internal/env"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc"
 	"github.com/ultravioletrs/cocos/pkg/clients/grpc/agent"
 	"github.com/ultravioletrs/cocos/pkg/sdk"
@@ -21,6 +22,8 @@ const (
 	svcName            = "cli"
 	envPrefixAgentGRPC = "AGENT_GRPC_"
 	completion         = "completion"
+	filePermision      = 0o755
+	cocosDirectory     = ".cocos"
 )
 
 type config struct {
@@ -33,13 +36,24 @@ func main() {
 		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.Parse(&agentGRPCConfig, env.Options{Prefix: envPrefixAgentGRPC}); err != nil {
+	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
 	}
@@ -83,19 +97,36 @@ func main() {
 		},
 	}
 
+	keysCmd := cliSVC.NewKeysCmd()
+	attestationCmd := cliSVC.NewAttestationCmd()
+	backendCmd := cliSVC.NewBackendCmd()
+
 	// Agent Commands
 	rootCmd.AddCommand(cliSVC.NewAlgorithmCmd())
 	rootCmd.AddCommand(cliSVC.NewDatasetsCmd())
 	rootCmd.AddCommand(cliSVC.NewResultsCmd())
-	attestaionCmd := cliSVC.NewAttestationCmd()
-	rootCmd.AddCommand(attestaionCmd)
+	rootCmd.AddCommand(attestationCmd)
 	rootCmd.AddCommand(cliSVC.NewFileHashCmd())
-	rootCmd.AddCommand(cliSVC.NewAddMeasurementCmd())
-	rootCmd.AddCommand(cliSVC.NewKeysCmd())
+	rootCmd.AddCommand(backendCmd)
+	rootCmd.AddCommand(keysCmd)
+	rootCmd.AddCommand(cliSVC.NewCABundleCmd(directoryCachePath))
 
 	// Attestation commands
-	attestaionCmd.AddCommand(cliSVC.NewGetAttestationCmd())
-	attestaionCmd.AddCommand(cliSVC.NewValidateAttestationValidationCmd())
+	attestationCmd.AddCommand(cliSVC.NewGetAttestationCmd())
+	attestationCmd.AddCommand(cliSVC.NewValidateAttestationValidationCmd())
+
+	// Flags
+	keysCmd.PersistentFlags().StringVarP(
+		&cli.KeyType,
+		"key-type",
+		"k",
+		"rsa",
+		"User Key type",
+	)
+
+	// Backend information commands
+	backendCmd.AddCommand(cliSVC.NewAddMeasurementCmd())
+	backendCmd.AddCommand(cliSVC.NewAddHostDataCmd())
 
 	if err := rootCmd.Execute(); err != nil {
 		logger.Error(fmt.Sprintf("Command execution failed: %s", err))

cmd/manager/main.go

@@ -8,14 +8,16 @@ import (
 	"fmt"
 	"log"
 	"log/slog"
+	"net/url"
 	"os"
 	"strings"
 
 	mglog "github.com/absmach/magistrala/logger"
+	"github.com/absmach/magistrala/pkg/jaeger"
+	"github.com/absmach/magistrala/pkg/prometheus"
 	"github.com/absmach/magistrala/pkg/uuid"
+	"github.com/caarlos0/env/v11"
 	"github.com/ultravioletrs/cocos/internal"
-	"github.com/ultravioletrs/cocos/internal/env"
-	jaegerclient "github.com/ultravioletrs/cocos/internal/jaeger"
 	"github.com/ultravioletrs/cocos/manager"
 	"github.com/ultravioletrs/cocos/manager/api"
 	managerapi "github.com/ultravioletrs/cocos/manager/api/grpc"
@@ -35,9 +37,10 @@ const (
 )
 
 type config struct {
-	LogLevel   string `env:"MANAGER_LOG_LEVEL"          envDefault:"info"`
-	JaegerURL  string `env:"COCOS_JAEGER_URL"           envDefault:"http://localhost:14268/api/traces"`
-	InstanceID string `env:"MANAGER_INSTANCE_ID"        envDefault:""`
+	LogLevel   string  `env:"MANAGER_LOG_LEVEL"     envDefault:"info"`
+	JaegerURL  url.URL `env:"COCOS_JAEGER_URL"      envDefault:"http://localhost:14268/api/traces"`
+	TraceRatio float64 `env:"MG_JAEGER_TRACE_RATIO" envDefault:"1.0"`
+	InstanceID string  `env:"MANAGER_INSTANCE_ID"   envDefault:""`
 }
 
 func main() {
@@ -55,14 +58,13 @@ func main() {
 	}
 
 	if cfg.InstanceID == "" {
-		cfg.InstanceID, err = uuid.New().ID()
-		if err != nil {
+		if cfg.InstanceID, err = uuid.New().ID(); err != nil {
 			logger.Error(fmt.Sprintf("Failed to generate instance ID: %s", err))
 			return
 		}
 	}
 
-	tp, err := jaegerclient.NewProvider(ctx, svcName, cfg.JaegerURL, cfg.InstanceID)
+	tp, err := jaeger.NewProvider(ctx, svcName, cfg.JaegerURL, cfg.InstanceID, cfg.TraceRatio)
 	if err != nil {
 		logger.Error(fmt.Sprintf("Failed to init Jaeger: %s", err))
 	}
@@ -74,19 +76,15 @@ func main() {
 	tracer := tp.Tracer(svcName)
 
 	qemuCfg := qemu.Config{}
-	if err := env.Parse(&qemuCfg, env.Options{Prefix: envPrefixQemu}); err != nil {
+	if err := env.ParseWithOptions(&qemuCfg, env.Options{Prefix: envPrefixQemu}); err != nil {
 		logger.Error(fmt.Sprintf("failed to load QEMU configuration: %s", err))
 		return
 	}
-	exe, args, err := qemu.ExecutableAndArgs(qemuCfg)
-	if err != nil {
-		logger.Error(fmt.Sprintf("failed to parse QEMU configuration: %s", err))
-		return
-	}
-	logger.Info(fmt.Sprintf("%s %s", exe, strings.Join(args, " ")))
+	args := qemuCfg.ConstructQemuArgs()
+	logger.Info(strings.Join(args, " "))
 
 	managerGRPCConfig := grpc.Config{}
-	if err := env.Parse(&managerGRPCConfig, env.Options{Prefix: envPrefixGRPC}); err != nil {
+	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))
 		return
 	}
@@ -117,17 +115,16 @@ func main() {
 		logger.Error(fmt.Sprintf("%s service terminated: %s", svcName, err))
 	}
 
-	err = internal.DeleteFilesInDir(qemuCfg.TmpFileLoc)
-	if err != nil {
+	if err = internal.DeleteFilesInDir(qemuCfg.TmpFileLoc); err != nil {
 		logger.Error(err.Error())
 	}
 }
 
 func newService(logger *slog.Logger, tracer trace.Tracer, qemuCfg qemu.Config, eventsChan chan *pkgmanager.ClientStreamMessage) manager.Service {
-	svc := manager.New(qemuCfg, logger, eventsChan)
+	svc := manager.New(qemuCfg, logger, eventsChan, qemu.NewVM)
 	go svc.RetrieveAgentEventsLogs()
 	svc = api.LoggingMiddleware(svc, logger)
-	counter, latency := internal.MakeMetrics(svcName, "api")
+	counter, latency := prometheus.MakeMetrics(svcName, "api")
 	svc = api.MetricsMiddleware(svc, counter, latency)
 	svc = tracing.New(svc, tracer)
 

go.mod

@@ -1,58 +1,65 @@
 module github.com/ultravioletrs/cocos
 
-go 1.22.0
+go 1.22.4
 
 require (
-	github.com/absmach/magistrala v0.0.0-20240119191055-d95283d31472
-	github.com/caarlos0/env/v7 v7.1.0
-	github.com/cenkalti/backoff/v4 v4.2.1
-	github.com/digitalocean/go-libvirt v0.0.0-20221205150000-2939327a8519
+	github.com/absmach/magistrala v0.14.1-0.20240709113739-04c359462746
+	github.com/caarlos0/env/v11 v11.1.0
+	github.com/cenkalti/backoff/v4 v4.3.0
+	github.com/digitalocean/go-libvirt v0.0.0-20240709142323-d8406205c752
 	github.com/go-kit/kit v0.13.0
 	github.com/gofrs/uuid v4.4.0+incompatible
-	github.com/google/go-sev-guest v0.10.1
+	github.com/google/go-sev-guest v0.11.1
 	github.com/mdlayher/vsock v1.2.1
-	github.com/prometheus/client_golang v1.18.0
-	github.com/spf13/cobra v1.8.0
+	github.com/spf13/cobra v1.8.1
 	github.com/spf13/pflag v1.0.5
-	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.46.1
-	go.opentelemetry.io/otel v1.21.0
-	go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.21.0
-	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.21.0
-	go.opentelemetry.io/otel/sdk v1.21.0
-	go.opentelemetry.io/otel/trace v1.21.0
-	golang.org/x/crypto v0.19.0
-	golang.org/x/sync v0.6.0
-	google.golang.org/grpc v1.60.1
-	google.golang.org/protobuf v1.32.0
+	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
+)
+
+require (
+	go.opentelemetry.io/otel v1.28.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
 )
 
 require (
-	cloud.google.com/go/compute v1.23.3 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
-	github.com/cespare/xxhash/v2 v2.2.0 // indirect
+	github.com/cespare/xxhash/v2 v2.3.0 // indirect
+	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	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.1 // 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.3 // 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/uuid v1.5.0 // indirect
-	github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.0 // indirect
+	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
-	github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0 // indirect
 	github.com/mdlayher/socket v0.4.1 // indirect
 	github.com/pborman/uuid v1.2.1 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
-	github.com/prometheus/client_model v0.5.0 // indirect
-	github.com/prometheus/common v0.45.0 // indirect
-	github.com/prometheus/procfs v0.12.0 // indirect
-	go.opentelemetry.io/otel/metric v1.21.0 // indirect
-	go.opentelemetry.io/proto/otlp v1.0.0 // indirect
+	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
+	github.com/prometheus/client_golang v1.19.1 // indirect
+	github.com/prometheus/client_model v0.6.1 // indirect
+	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/proto/otlp v1.3.1 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
-	golang.org/x/net v0.20.0 // indirect
-	golang.org/x/sys v0.17.0 // indirect
-	golang.org/x/text v0.14.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20240108191215-35c7eff3a6b1 // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20240108191215-35c7eff3a6b1 // indirect
+	golang.org/x/net v0.26.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
+	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -1,31 +1,20 @@
-cloud.google.com/go v0.110.8 h1:tyNdfIxjzaWctIiLYOTalaLKZ17SI44SKFW26QbOhME=
-cloud.google.com/go/compute v1.23.3 h1:6sVlXXBmbd7jNX0Ipq0trII3e4n1/MsADLK6a+aiVlk=
-cloud.google.com/go/compute v1.23.3/go.mod h1:VCgBUoMnIVIR0CscqQiPJLAG25E3ZRZMzcFZeQ+h8CI=
-cloud.google.com/go/compute/metadata v0.2.3 h1:mg4jlk7mCAj6xXp9UJ4fjI9VUI5rubuGBW5aJ7UnBMY=
-cloud.google.com/go/compute/metadata v0.2.3/go.mod h1:VAV5nSsACxMJvgaAuX6Pk2AawlZn8kiOGuCv6gTkwuA=
 github.com/VividCortex/gohistogram v1.0.0 h1:6+hBz+qvs0JOrrNhhmR7lFxo5sINxBCGXrdtl/UvroE=
 github.com/VividCortex/gohistogram v1.0.0/go.mod h1:Pf5mBqqDxYaXu3hDrrU+w6nw50o/4+TcAqDqk/vUH7g=
-github.com/absmach/magistrala v0.0.0-20240119191055-d95283d31472 h1:MdS3yTTJP/OSFkMJ4+cJ0LnRS9OyPtPxsR0mSFYglXw=
-github.com/absmach/magistrala v0.0.0-20240119191055-d95283d31472/go.mod h1:02Ip3Ow7JlPT4nb2TdZcOuf7C3ac8tzyEVwB+H1G40g=
+github.com/absmach/magistrala v0.14.1-0.20240709113739-04c359462746 h1:Tj567KeGVygjTsSCxn4++skKiz9GkPugM1KMdIFxvfw=
+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/v7 v7.1.0 h1:9lzTF5amyQeWHZzuZeKlCb5FWSUxpG1js43mhbY8ozg=
-github.com/caarlos0/env/v7 v7.1.0/go.mod h1:LPPWniDUq4JaO6Q41vtlyikhMknqymCLBw0eX4dcH1E=
-github.com/cenkalti/backoff/v4 v4.2.1 h1:y4OZtCnogmCPw98Zjyt5a6+QwPLGkiQsYW5oUqylYbM=
-github.com/cenkalti/backoff/v4 v4.2.1/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
-github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
-github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
-github.com/cncf/xds/go v0.0.0-20230607035331-e9ce68804cb4 h1:/inchEIKaYC1Akx+H+gqO04wryn5h75LSazbRlnya1k=
-github.com/cncf/xds/go v0.0.0-20230607035331-e9ce68804cb4/go.mod h1:eXthEFrGJvWHgFFCl3hGmgk+/aYT6PnTQLykKQRLhEs=
-github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
-github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+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/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/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-20221205150000-2939327a8519 h1:OpkN/n40cmKenDQS+IOAeW9DLhYy4DADSeZnouCEV/E=
-github.com/digitalocean/go-libvirt v0.0.0-20221205150000-2939327a8519/go.mod h1:WyJJyfmJ0gWJvjV+ZH4DOgtOYZc1KOvYyBXWCLKxsUU=
-github.com/envoyproxy/protoc-gen-validate v1.0.2 h1:QkIBuU5k+x7/QXPvPPnWXWlCdaBFApVqftFV6k087DA=
-github.com/envoyproxy/protoc-gen-validate v1.0.2/go.mod h1:GpiZQP3dDbg4JouG/NNS7QWXpgx6x8QiMKdmN72jogE=
+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/go-kit/kit v0.13.0 h1:OoneCcHKHQ03LfBpoQCUfCluwd2Vt3ohz+kvbJneZAU=
 github.com/go-kit/kit v0.13.0/go.mod h1:phqEHMMUbyrCFCTgH48JueqrM3md2HcAZ8N3XE4FKDg=
 github.com/go-kit/log v0.2.1 h1:MRVx0/zhvdseW+Gza6N9rVzU/IVzaeE1SFI4raAhmBU=
@@ -33,33 +22,33 @@ github.com/go-kit/log v0.2.1/go.mod h1:NwTd00d/i8cPZ3xOwwiv2PO5MOcx78fFErGNcVmBj
 github.com/go-logfmt/logfmt v0.6.0 h1:wGYYu3uicYdqXVgoYbvnkrPVXkuLM1p1ifugDMEdRi4=
 github.com/go-logfmt/logfmt v0.6.0/go.mod h1:WYhtIu8zTZfxdn5+rREduYbwxfcBr/Vr6KEVveWlfTs=
 github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A=
-github.com/go-logr/logr v1.4.1 h1:pKouT5E8xu9zeFC39JXRDukb6JFQPXM5p5I91188VAQ=
-github.com/go-logr/logr v1.4.1/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
+github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY=
+github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
 github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
 github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
 github.com/gofrs/uuid v4.4.0+incompatible h1:3qXRTX8/NbyulANqlc0lchS1gqAVxRgsuW1YrTJupqA=
 github.com/gofrs/uuid v4.4.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM=
-github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
-github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg=
-github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
-github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
+github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-configfs-tsm v0.2.2 h1:YnJ9rXIOj5BYD7/0DNnzs8AOp7UcvjfTvt215EWcs98=
 github.com/google/go-configfs-tsm v0.2.2/go.mod h1:EL1GTDFMb5PZQWDviGfZV9n87WeGTR/JUg13RfwkgRo=
-github.com/google/go-sev-guest v0.10.1 h1:kZIcGD513eqInbezqxchTtOLJlVXALyWlIXu5E0tNJE=
-github.com/google/go-sev-guest v0.10.1/go.mod h1:/5hrgGWqG7+MPTXKhQz+v9ZE+Eh4MCBdItCNwgshlrk=
+github.com/google/go-sev-guest v0.11.1 h1:gnww4U8fHV5DCPz4gykr1s8SEX1fFNcxCBy+vvXN24k=
+github.com/google/go-sev-guest v0.11.1/go.mod h1:qBOfb+JmgsUI3aUyzQoGC13Kpp9zwLeWvuyXmA9q77w=
 github.com/google/logger v1.1.1 h1:+6Z2geNxc9G+4D4oDO9njjjn2d0wN5d7uOo0vOIW1NQ=
 github.com/google/logger v1.1.1/go.mod h1:BkeJZ+1FhQ+/d087r4dzojEg1u2ZX+ZqG1jTUrLM+zQ=
 github.com/google/uuid v1.0.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/google/uuid v1.5.0 h1:1p67kYwdtXjb0gL0BPiP1Av9wiZPo5A8z2cWkTZ+eyU=
-github.com/google/uuid v1.5.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.0 h1:Wqo399gCIufwto+VfwCSvsnfGpF/w5E9CNxSwbpD6No=
-github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.0/go.mod h1:qmOFXW2epJhM0qSnUUYpldc7gVz2KMQwJ/QYCDIa7XU=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0 h1:bkypFPDjIYGfCYD5mRBvpqxfYX1YCS1PXdKYWi8FsN0=
+github.com/grpc-ecosystem/grpc-gateway/v2 v2.20.0/go.mod h1:P+Lt/0by1T8bfcF3z737NnSbmxQAppXMRziHUxPOC8k=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
-github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0 h1:jWpvCLoY8Z/e3VKvlsiIGKtc+UG6U5vzxaoagmhXfyg=
-github.com/matttproud/golang_protobuf_extensions/v2 v2.0.0/go.mod h1:QUyp042oQthUoa9bqDv0ER0wrtXnBruoNd7aNjkbP+k=
+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/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=
@@ -68,100 +57,68 @@ github.com/pborman/uuid v1.2.1 h1:+ZZIw58t/ozdjRaXh/3awHfmWRbzYxJoAdNJxe/3pvw=
 github.com/pborman/uuid v1.2.1/go.mod h1:X/NO0urCmaxf9VXbdlT7C2Yzkj2IKimNn4k+gtPdI/k=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
-github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/prometheus/client_golang v1.18.0 h1:HzFfmkOzH5Q8L8G+kSJKUx5dtG87sewO+FoDDqP5Tbk=
-github.com/prometheus/client_golang v1.18.0/go.mod h1:T+GXkCk5wSJyOqMIzVgvvjFDlkOQntgjkJWKrN5txjA=
-github.com/prometheus/client_model v0.5.0 h1:VQw1hfvPvk3Uv6Qf29VrPF32JB6rtbgI6cYPYQjL0Qw=
-github.com/prometheus/client_model v0.5.0/go.mod h1:dTiFglRmd66nLR9Pv9f0mZi7B7fk5Pm3gvsjB5tr+kI=
-github.com/prometheus/common v0.45.0 h1:2BGz0eBc2hdMDLnO/8n0jeB3oPrt2D08CekT0lneoxM=
-github.com/prometheus/common v0.45.0/go.mod h1:YJmSTw9BoKxJplESWWxlbyttQR4uaEcGyv9MZjVOJsY=
-github.com/prometheus/procfs v0.12.0 h1:jluTpSng7V9hY0O2R9DzzJHYb2xULk9VTR1V1R/k6Bo=
-github.com/prometheus/procfs v0.12.0/go.mod h1:pcuDEFsWDnvcgNzo4EEweacyhjeA9Zk3cnaOZAZEfOo=
+github.com/prometheus/client_golang v1.19.1 h1:wZWJDwK+NameRJuPGDhlnFgx8e8HN3XHQeLaYJFJBOE=
+github.com/prometheus/client_golang v1.19.1/go.mod h1:mP78NwGzrVks5S2H6ab8+ZZGJLZUq1hoULYBAYBw1Ho=
+github.com/prometheus/client_model v0.6.1 h1:ZKSh/rekM+n3CeS952MLRAdFwIKqeY8b62p8ais2e9E=
+github.com/prometheus/client_model v0.6.1/go.mod h1:OrxVMOVHjw3lKMa8+x6HeMGkHMQyHDk9E3jmP2AmGiY=
+github.com/prometheus/common v0.52.2 h1:LW8Vk7BccEdONfrJBDffQGRtpSzi5CQaRZGtboOO2ck=
+github.com/prometheus/common v0.52.2/go.mod h1:lrWtQx+iDfn2mbH5GUzlH9TSHyfZpHkSiG1W7y3sF2Q=
+github.com/prometheus/procfs v0.13.0 h1:GqzLlQyfsPbaEHaQkO7tbDlriv/4o5Hudv6OXHGKX7o=
+github.com/prometheus/procfs v0.13.0/go.mod h1:cd4PFCR54QLnGKPaKGA6l+cfuNXtht43ZKY6tow0Y1g=
+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/spf13/cobra v1.8.0 h1:7aJaZx1B85qltLMc546zn58BxxfZdR/W22ej9CFoEf0=
-github.com/spf13/cobra v1.8.0/go.mod h1:WXLWApfZ71AjXPya3WOlMsY9yMs7YeiHhFVlvLyhcho=
+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=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
-github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
-github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
-github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
-github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
-github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
-github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
-github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.46.1 h1:SpGay3w+nEwMpfVnbqOLH5gY52/foP8RE8UzTZ1pdSE=
-go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.46.1/go.mod h1:4UoMYEZOC0yN/sPGH76KPkkU7zgiEWYWL9vwmbnTJPE=
-go.opentelemetry.io/otel v1.21.0 h1:hzLeKBZEL7Okw2mGzZ0cc4k/A7Fta0uoPgaJCr8fsFc=
-go.opentelemetry.io/otel v1.21.0/go.mod h1:QZzNPQPm1zLX4gZK4cMi+71eaorMSGT3A4znnUvNNEo=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.21.0 h1:cl5P5/GIfFh4t6xyruOgJP5QiA1pw4fYYdv6nc6CBWw=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.21.0/go.mod h1:zgBdWWAu7oEEMC06MMKc5NLbA/1YDXV1sMpSqEeLQLg=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.21.0 h1:digkEZCJWobwBqMwC0cwCq8/wkkRy/OowZg5OArWZrM=
-go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.21.0/go.mod h1:/OpE/y70qVkndM0TrxT4KBoN3RsFZP0QaofcfYrj76I=
-go.opentelemetry.io/otel/metric v1.21.0 h1:tlYWfeo+Bocx5kLEloTjbcDwBuELRrIFxwdQ36PlJu4=
-go.opentelemetry.io/otel/metric v1.21.0/go.mod h1:o1p3CA8nNHW8j5yuQLdc1eeqEaPfzug24uvsyIEJRWM=
-go.opentelemetry.io/otel/sdk v1.21.0 h1:FTt8qirL1EysG6sTQRZ5TokkU8d0ugCj8htOgThZXQ8=
-go.opentelemetry.io/otel/sdk v1.21.0/go.mod h1:Nna6Yv7PWTdgJHVRD9hIYywQBRx7pbox6nwBnZIxl/E=
-go.opentelemetry.io/otel/trace v1.21.0 h1:WD9i5gzvoUPuXIXH24ZNBudiarZDKuekPqi/E8fpfLc=
-go.opentelemetry.io/otel/trace v1.21.0/go.mod h1:LGbsEB0f9LGjN+OZaQQ26sohbOmiMR+BaslueVtS/qQ=
-go.opentelemetry.io/proto/otlp v1.0.0 h1:T0TX0tmXU8a3CbNXzEKGeU5mIVOdf0oykP+u2lIVU/I=
-go.opentelemetry.io/proto/otlp v1.0.0/go.mod h1:Sy6pihPLfYHkr3NkUbEhGHFhINUSI/v80hjKIs5JXpM=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+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/otel v1.28.0 h1:/SqNcYk+idO0CxKEUOtKQClMK/MimZihKYMruSMViUo=
+go.opentelemetry.io/otel v1.28.0/go.mod h1:q68ijF8Fc8CnMHKyzqL6akLO46ePnjkgfIMIjUIX9z4=
+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/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/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=
 go.uber.org/multierr v1.11.0/go.mod h1:20+QtiLqy0Nd6FdQB9TLXag12DsQkrbs3htMFfDN80Y=
-golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
-golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
-golang.org/x/crypto v0.19.0 h1:ENy+Az/9Y1vSrlrvBSyna3PITt4tiZLf7sgCjZBX7Wo=
-golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
-golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
-golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
-golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
-golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
-golang.org/x/net v0.20.0 h1:aCL9BSgETF1k+blQaYUBx9hJ9LOGP3gAVemcZlf1Kpo=
-golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY=
-golang.org/x/oauth2 v0.15.0 h1:s8pnnxNVzjWyrvYdFUQq5llS1PX2zhPXmccZv99h7uQ=
-golang.org/x/oauth2 v0.15.0/go.mod h1:q48ptWNTY5XWf+JNten23lcvHpLJ0ZSxF5ttTHKVCAM=
-golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
-golang.org/x/sync v0.6.0 h1:5BMeUDZ7vkXGfEr1x9B4bRcTH4lpkTkpdh0T/J+qjbQ=
-golang.org/x/sync v0.6.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-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+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/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ=
+golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE=
+golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M=
+golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
 golang.org/x/sys v0.0.0-20210426230700-d19ff857e887/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.17.0 h1:25cE3gD+tdBA7lp7QfhuV+rJiE9YXTcS3VG1SqssI/Y=
-golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
-golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
-golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
-golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
-golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
-golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
-golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
-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.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
-golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-google.golang.org/appengine v1.6.8 h1:IhEN5q69dyKagZPYMSdIjS2HqprW324FRQZJcGqPAsM=
-google.golang.org/appengine v1.6.8/go.mod h1:1jJ3jBArFh5pcgW8gCtRJnepW8FzD1V44FJffLiz/Ds=
-google.golang.org/genproto v0.0.0-20240102182953-50ed04b92917 h1:nz5NESFLZbJGPFxDT/HCn+V1mZ8JGNoY4nUpmW/Y2eg=
-google.golang.org/genproto v0.0.0-20240102182953-50ed04b92917/go.mod h1:pZqR+glSb11aJ+JQcczCvgf47+duRuzNSKqE8YAQnV0=
-google.golang.org/genproto/googleapis/api v0.0.0-20240108191215-35c7eff3a6b1 h1:OPXtXn7fNMaXwO3JvOmF1QyTc00jsSFFz1vXXBOdCDo=
-google.golang.org/genproto/googleapis/api v0.0.0-20240108191215-35c7eff3a6b1/go.mod h1:B5xPO//w8qmBDjGReYLpR6UJPnkldGkCSMoH/2vxJeg=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20240108191215-35c7eff3a6b1 h1:gphdwh0npgs8elJ4T6J+DQJHPVF7RsuJHCfwztUb4J4=
-google.golang.org/genproto/googleapis/rpc v0.0.0-20240108191215-35c7eff3a6b1/go.mod h1:daQN87bsDqDoe316QbbvX60nMoJQa4r6Ds0ZuoAe5yA=
-google.golang.org/grpc v1.60.1 h1:26+wFr+cNqSGFcOXcabYC0lUVJVRa2Sb2ortSK7VrEU=
-google.golang.org/grpc v1.60.1/go.mod h1:OlCHIeLYqSSsLi6i49B5QGdzaMZK9+M7LXN2FKz4eGM=
-google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
-google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
-google.golang.org/protobuf v1.32.0 h1:pPC6BG5ex8PDFnkbrGU3EixyhKcQ2aDuBS36lqK/C7I=
-google.golang.org/protobuf v1.32.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
+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/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4=
+golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI=
+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=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+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=

hal/linux/Config.in

@@ -1 +1,2 @@
-source "$BR2_EXTERNAL_COCOS_PATH/package/agent/Config.in"
+source "$BR2_EXTERNAL_COCOS_PATH/package/agent/Config.in"
+source "$BR2_EXTERNAL_COCOS_PATH/package/wasmtime/Config.in"

hal/linux/configs/cocos_defconfig

@@ -43,3 +43,13 @@ BR2_LINUX_KERNEL_NEEDS_HOST_LIBELF=y
 # host-qemu for gitlab testing
 BR2_PACKAGE_HOST_QEMU=y
 BR2_PACKAGE_HOST_QEMU_SYSTEM_MODE=y
+
+# Docker
+BR2_PACKAGE_LIBSECCOMP_ARCH_SUPPORTS=y 
+BR2_PACKAGE_LIBSECCOMP=y 
+BR2_PACKAGE_CA_CERTIFICATES=y  
+BR2_PACKAGE_DOCKER_CLI=y  
+BR2_PACKAGE_DOCKER_COMPOSE=y  
+BR2_PACKAGE_DOCKER_ENGINE=y  
+BR2_PACKAGE_CONTAINERD=y  
+BR2_PACKAGE_RUNC=y

hal/linux/package/wasmtime/Config.in

@@ -0,0 +1,5 @@
+config BR2_PACKAGE_WASMTIME
+    bool "wasmtime"
+    help
+      Wasmtime is a standalone runtime for WebAssembly.
+      https://github.com/bytecodealliance/wasmtime

hal/linux/package/wasmtime/wasmtime.mk

@@ -0,0 +1,11 @@
+WASMTIME_SITE = https://wasmtime.dev/install.sh
+
+define WASMTIME_BUILD_CMDS
+    curl $(WASMTIME_SITE) -sSf | bash
+endef
+
+define WASMTIME_INSTALL_TARGET_CMDS
+    $(INSTALL) -D -m 0755 ~/.wasmtime/bin/wasmtime $(TARGET_DIR)/usr/bin/wasmtime
+endef
+
+$(eval $(generic-package))

internal/env/doc.go

@@ -1,6 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-
-// Package env contains utility tools for handling
-// environment-based service configuration.
-package env

internal/env/parser.go

@@ -1,38 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package env
-
-import "github.com/caarlos0/env/v7"
-
-type Options struct {
-	// Environment keys and values that will be accessible for the service
-	Environment map[string]string
-
-	// TagName specifies another tagname to use rather than the default env
-	TagName string
-
-	// RequiredIfNoDef automatically sets all env as required if they do not declare 'envDefault'
-	RequiredIfNoDef bool
-
-	// OnSet allows to run a function when a value is set
-	OnSet env.OnSetFn
-
-	// Prefix define a prefix for each key
-	Prefix string
-}
-
-func Parse(v interface{}, opts ...Options) error {
-	var actOpts []env.Options
-
-	for _, opt := range opts {
-		actOpts = append(actOpts, env.Options{
-			Environment:     opt.Environment,
-			TagName:         opt.TagName,
-			RequiredIfNoDef: opt.RequiredIfNoDef,
-			OnSet:           opt.OnSet,
-			Prefix:          opt.Prefix,
-		})
-	}
-
-	return env.Parse(v, actOpts...)
-}

internal/jaeger/doc.go

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

internal/jaeger/provider.go

@@ -1,80 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package jaeger
-
-import (
-	"context"
-	"errors"
-	"net/url"
-
-	"go.opentelemetry.io/otel"
-	"go.opentelemetry.io/otel/attribute"
-	"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
-	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp"
-	"go.opentelemetry.io/otel/propagation"
-	"go.opentelemetry.io/otel/sdk/resource"
-	"go.opentelemetry.io/otel/sdk/trace"
-	semconv "go.opentelemetry.io/otel/semconv/v1.21.0"
-)
-
-var (
-	errNoURL                     = errors.New("URL is empty")
-	errNoSvcName                 = errors.New("Service Name is empty")
-	errUnsupportedTraceURLScheme = errors.New("unsupported tracing url scheme")
-)
-
-// NewProvider initializes Jaeger TraceProvider.
-func NewProvider(ctx context.Context, svcName, jaegerurl, instanceID string) (*trace.TracerProvider, error) {
-	if jaegerurl == "" {
-		return nil, errNoURL
-	}
-
-	if svcName == "" {
-		return nil, errNoSvcName
-	}
-
-	jurl, err := url.Parse(jaegerurl)
-	if err != nil {
-		return nil, err
-	}
-
-	var exporter *otlptrace.Exporter
-	switch jurl.Scheme {
-	case "http":
-		exporter, err = otlptracehttp.New(ctx, otlptracehttp.WithEndpoint(jurl.Host), otlptracehttp.WithURLPath(jurl.Path), otlptracehttp.WithInsecure())
-		if err != nil {
-			return nil, err
-		}
-	case "https":
-		exporter, err = otlptracehttp.New(ctx, otlptracehttp.WithEndpoint(jurl.Host), otlptracehttp.WithURLPath(jurl.Path))
-		if err != nil {
-			return nil, err
-		}
-	default:
-		return nil, errUnsupportedTraceURLScheme
-	}
-
-	attributes := []attribute.KeyValue{
-		semconv.ServiceNameKey.String(svcName),
-		attribute.String("host.id", instanceID),
-	}
-
-	hostAttr, err := resource.New(ctx, resource.WithHost(), resource.WithOSDescription(), resource.WithContainer())
-	if err != nil {
-		return nil, err
-	}
-	attributes = append(attributes, hostAttr.Attributes()...)
-
-	tp := trace.NewTracerProvider(
-		trace.WithSampler(trace.AlwaysSample()),
-		trace.WithBatcher(exporter),
-		trace.WithResource(resource.NewWithAttributes(
-			semconv.SchemaURL,
-			attributes...,
-		)),
-	)
-	otel.SetTracerProvider(tp)
-	otel.SetTextMapPropagator(propagation.TraceContext{})
-
-	return tp, nil
-}

internal/metrics.go

@@ -1,27 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package internal
-
-import (
-	kitprometheus "github.com/go-kit/kit/metrics/prometheus"
-	stdprometheus "github.com/prometheus/client_golang/prometheus"
-)
-
-// MakeMetrics returns an instance of metrics.
-func MakeMetrics(namespace, subsystem string) (*kitprometheus.Counter, *kitprometheus.Summary) {
-	counter := kitprometheus.NewCounterFrom(stdprometheus.CounterOpts{
-		Namespace: namespace,
-		Subsystem: subsystem,
-		Name:      "request_count",
-		Help:      "Number of requests received.",
-	}, []string{"method"})
-	latency := kitprometheus.NewSummaryFrom(stdprometheus.SummaryOpts{
-		Namespace:  namespace,
-		Subsystem:  subsystem,
-		Objectives: map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
-		Name:       "request_latency_microseconds",
-		Help:       "Total duration of requests in microseconds.",
-	}, []string{"method"})
-
-	return counter, latency
-}

internal/server/grpc/grpc.go

@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"github.com/ultravioletrs/cocos/agent"
+	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
 	"github.com/ultravioletrs/cocos/agent/auth"
 	"github.com/ultravioletrs/cocos/internal/server"
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
@@ -48,14 +49,14 @@ type Server struct {
 	server          *grpc.Server
 	registerService serviceRegister
 	agent           agent.Service
-	authSvc         *auth.Service
+	authSvc         auth.Authenticator
 }
 
 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, agentSvc agent.Service, authSvc *auth.Service) server.Server {
+func New(ctx context.Context, cancel context.CancelFunc, name string, config server.Config, registerService serviceRegister, logger *slog.Logger, agentSvc agent.Service, authSvc auth.Authenticator) server.Server {
 	listenFullAddress := fmt.Sprintf("%s:%s", config.Host, config.Port)
 	return &Server{
 		BaseServer: server.BaseServer{
@@ -79,8 +80,9 @@ func (s *Server) Start() error {
 	}
 
 	if s.authSvc != nil {
-		grpcServerOptions = append(grpcServerOptions, grpc.UnaryInterceptor(s.authSvc.AuthUnaryInterceptor()))
-		grpcServerOptions = append(grpcServerOptions, grpc.StreamInterceptor(s.authSvc.AuthStreamInterceptor()))
+		unary, stream := agentgrpc.NewAuthInterceptor(s.authSvc)
+		grpcServerOptions = append(grpcServerOptions, grpc.UnaryInterceptor(unary))
+		grpcServerOptions = append(grpcServerOptions, grpc.StreamInterceptor(stream))
 	}
 
 	listener, err := net.Listen("tcp", s.Address)

manager/api/grpc/client.go

@@ -50,6 +50,10 @@ func (client ManagerClient) Process(ctx context.Context, cancel context.CancelFu
 			case *pkgmanager.ServerStreamMessage_TerminateReq:
 				cancel()
 				return errors.Join(errTerminationFromServer, errors.New(mes.TerminateReq.Message))
+			case *pkgmanager.ServerStreamMessage_StopComputation:
+				if err := client.svc.Stop(ctx, mes.StopComputation.ComputationId); err != nil {
+					return err
+				}
 			}
 		}
 	})

manager/api/logging.go

@@ -41,6 +41,19 @@ func (lm *loggingMiddleware) Run(ctx context.Context, mc *pkgmanager.Computation
 	return lm.svc.Run(ctx, mc)
 }
 
+func (lm *loggingMiddleware) Stop(ctx context.Context, computationID string) (err error) {
+	defer func(begin time.Time) {
+		message := fmt.Sprintf("Method Stop for computation took %s to complete", time.Since(begin))
+		if err != nil {
+			lm.logger.Warn(fmt.Sprintf("%s with error: %s.", message, err))
+			return
+		}
+		lm.logger.Info(message)
+	}(time.Now())
+
+	return lm.svc.Stop(ctx, computationID)
+}
+
 func (lm *loggingMiddleware) RetrieveAgentEventsLogs() {
 	lm.svc.RetrieveAgentEventsLogs()
 }

manager/api/metrics.go

@@ -42,6 +42,15 @@ func (ms *metricsMiddleware) Run(ctx context.Context, mc *pkgmanager.Computation
 	return ms.svc.Run(ctx, mc)
 }
 
+func (ms *metricsMiddleware) Stop(ctx context.Context, computationID string) error {
+	defer func(begin time.Time) {
+		ms.counter.With("method", "Stop").Add(1)
+		ms.latency.With("method", "Stop").Observe(time.Since(begin).Seconds())
+	}(time.Now())
+
+	return ms.svc.Stop(ctx, computationID)
+}
+
 func (ms *metricsMiddleware) RetrieveAgentEventsLogs() {
 	ms.svc.RetrieveAgentEventsLogs()
 }

manager/manager.proto

@@ -17,6 +17,10 @@ message Terminate {
   string message = 1;
 }
 
+message StopComputation {
+  string computation_id = 1;
+}
+
 message RunResponse{
   string agent_port = 1;
   string computation_id = 2;
@@ -50,6 +54,7 @@ message ServerStreamMessage {
   oneof message {
     ComputationRunReq runReq = 1;
     Terminate terminateReq = 2;
+    StopComputation stopComputation = 3;
   }
 }
 

manager/manager_test.go

@@ -0,0 +1,232 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package manager
+
+import (
+	"context"
+	"encoding/json"
+	"log/slog"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/ultravioletrs/cocos/manager/qemu"
+	"github.com/ultravioletrs/cocos/manager/vm"
+	"github.com/ultravioletrs/cocos/manager/vm/mocks"
+	"github.com/ultravioletrs/cocos/pkg/manager"
+)
+
+func TestNew(t *testing.T) {
+	qemuCfg := qemu.Config{}
+	logger := slog.Default()
+	eventsChan := make(chan *manager.ClientStreamMessage)
+	vmf := new(mocks.Provider)
+
+	service := New(qemuCfg, logger, eventsChan, vmf.Execute)
+
+	assert.NotNil(t, service)
+	assert.IsType(t, &managerService{}, service)
+}
+
+func TestRun(t *testing.T) {
+	vmf := new(mocks.Provider)
+	vmMock := new(mocks.VM)
+	vmf.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return(vmMock)
+	tests := []struct {
+		name          string
+		req           *manager.ComputationRunReq
+		vmStartError  error
+		expectedError error
+	}{
+		{
+			name: "Successful run",
+			req: &manager.ComputationRunReq{
+				Id:   "test-computation",
+				Name: "Test Computation",
+				Algorithm: &manager.Algorithm{
+					Hash: make([]byte, hashLength),
+				},
+				AgentConfig: &manager.AgentConfig{},
+			},
+			vmStartError:  nil,
+			expectedError: nil,
+		},
+		{
+			name: "VM start failure",
+			req: &manager.ComputationRunReq{
+				Id:   "test-computation",
+				Name: "Test Computation",
+				Algorithm: &manager.Algorithm{
+					Hash: make([]byte, hashLength),
+				},
+				AgentConfig: &manager.AgentConfig{},
+			},
+			vmStartError:  assert.AnError,
+			expectedError: assert.AnError,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.vmStartError == nil {
+				vmMock.On("Start").Return(nil).Once()
+			} else {
+				vmMock.On("Start").Return(tt.vmStartError).Once()
+			}
+
+			vmMock.On("SendAgentConfig", mock.Anything).Return(nil)
+
+			qemuCfg := qemu.Config{
+				VSockConfig: qemu.VSockConfig{
+					GuestCID: 3,
+					Vnc:      5900,
+				},
+			}
+			logger := slog.Default()
+			eventsChan := make(chan *manager.ClientStreamMessage, 10)
+
+			ms := &managerService{
+				qemuCfg:    qemuCfg,
+				logger:     logger,
+				agents:     make(map[int]string),
+				vms:        make(map[string]vm.VM),
+				eventsChan: eventsChan,
+				vmFactory:  vmf.Execute,
+			}
+
+			ctx := context.Background()
+
+			port, err := ms.Run(ctx, tt.req)
+
+			if tt.expectedError != nil {
+				assert.Error(t, err)
+				assert.ErrorIs(t, err, tt.expectedError)
+				assert.Empty(t, port)
+			} else {
+				assert.NoError(t, err)
+				assert.NotEmpty(t, port)
+				assert.Len(t, ms.vms, 1)
+				assert.Len(t, ms.agents, 1)
+			}
+
+			vmf.AssertExpectations(t)
+
+			// Clear the events channel
+			for len(eventsChan) > 0 {
+				<-eventsChan
+			}
+		})
+	}
+}
+
+func TestStop(t *testing.T) {
+	tests := []struct {
+		name           string
+		computationID  string
+		vmStopError    error
+		expectedError  error
+		initialVMCount int
+	}{
+		{
+			name:           "Successful stop",
+			computationID:  "existing-computation",
+			vmStopError:    nil,
+			expectedError:  nil,
+			initialVMCount: 1,
+		},
+		{
+			name:           "Non-existent computation",
+			computationID:  "non-existent-computation",
+			vmStopError:    nil,
+			expectedError:  ErrNotFound,
+			initialVMCount: 0,
+		},
+		{
+			name:           "VM stop error",
+			computationID:  "error-computation",
+			vmStopError:    assert.AnError,
+			expectedError:  assert.AnError,
+			initialVMCount: 1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ms := &managerService{
+				vms: make(map[string]vm.VM),
+			}
+			vmMock := new(mocks.VM)
+
+			if tt.vmStopError == nil {
+				vmMock.On("Stop").Return(nil).Once()
+			} else {
+				vmMock.On("Stop").Return(assert.AnError).Once()
+			}
+
+			if tt.initialVMCount > 0 {
+				ms.vms[tt.computationID] = vmMock
+			}
+
+			err := ms.Stop(context.Background(), tt.computationID)
+
+			if tt.expectedError != nil {
+				assert.Error(t, err)
+				assert.ErrorIs(t, err, tt.expectedError)
+			} else {
+				assert.NoError(t, err)
+				assert.Len(t, ms.vms, 0)
+			}
+		})
+	}
+}
+
+func TestGetFreePort(t *testing.T) {
+	port, err := getFreePort()
+
+	assert.NoError(t, err)
+	assert.Greater(t, port, 0)
+}
+
+func TestPublishEvent(t *testing.T) {
+	tests := []struct {
+		name          string
+		event         string
+		computationID string
+		status        string
+		details       json.RawMessage
+	}{
+		{
+			name:          "Standard event",
+			event:         "test-event",
+			computationID: "test-computation",
+			status:        "test-status",
+			details:       nil,
+		},
+		{
+			name:          "Event with details",
+			event:         "detailed-event",
+			computationID: "detailed-computation",
+			status:        "detailed-status",
+			details:       json.RawMessage(`{"key": "value"}`),
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			eventsChan := make(chan *manager.ClientStreamMessage, 1)
+			ms := &managerService{
+				eventsChan: eventsChan,
+			}
+
+			ms.publishEvent(tt.event, tt.computationID, tt.status, tt.details)
+
+			assert.Len(t, eventsChan, 1)
+			event := <-eventsChan
+			assert.Equal(t, tt.event, event.GetAgentEvent().EventType)
+			assert.Equal(t, tt.computationID, event.GetAgentEvent().ComputationId)
+			assert.Equal(t, tt.status, event.GetAgentEvent().Status)
+			assert.Equal(t, "manager", event.GetAgentEvent().Originator)
+			assert.Equal(t, tt.details, json.RawMessage(event.GetAgentEvent().Details))
+		})
+	}
+}

manager/qemu/config.go

@@ -8,24 +8,24 @@ import (
 )
 
 type MemoryConfig struct {
-	Size  string `env:"MEMORY_SIZE" envDefault:"2048M"`
+	Size  string `env:"MEMORY_SIZE"  envDefault:"2048M"`
 	Slots int    `env:"MEMORY_SLOTS" envDefault:"5"`
-	Max   string `env:"MAX_MEMORY" envDefault:"30G"`
+	Max   string `env:"MAX_MEMORY"   envDefault:"30G"`
 }
 
 type OVMFCodeConfig struct {
-	If       string `env:"OVMF_CODE_IF" envDefault:"pflash"`
-	Format   string `env:"OVMF_CODE_FORMAT" envDefault:"raw"`
-	Unit     int    `env:"OVMF_CODE_UNIT" envDefault:"0"`
-	File     string `env:"OVMF_CODE_FILE" envDefault:"/usr/share/OVMF/OVMF_CODE.fd"`
+	If       string `env:"OVMF_CODE_IF"       envDefault:"pflash"`
+	Format   string `env:"OVMF_CODE_FORMAT"   envDefault:"raw"`
+	Unit     int    `env:"OVMF_CODE_UNIT"     envDefault:"0"`
+	File     string `env:"OVMF_CODE_FILE"     envDefault:"/usr/share/OVMF/OVMF_CODE.fd"`
 	ReadOnly string `env:"OVMF_CODE_READONLY" envDefault:"on"`
 }
 
 type OVMFVarsConfig struct {
-	If     string `env:"OVMF_VARS_IF" envDefault:"pflash"`
+	If     string `env:"OVMF_VARS_IF"     envDefault:"pflash"`
 	Format string `env:"OVMF_VARS_FORMAT" envDefault:"raw"`
-	Unit   int    `env:"OVMF_VARS_UNIT" envDefault:"1"`
-	File   string `env:"OVMF_VARS_FILE" envDefault:"/usr/share/OVMF/OVMF_VARS.fd"`
+	Unit   int    `env:"OVMF_VARS_UNIT"   envDefault:"1"`
+	File   string `env:"OVMF_VARS_FILE"   envDefault:"/usr/share/OVMF/OVMF_VARS.fd"`
 }
 
 type NetDevConfig struct {
@@ -37,7 +37,7 @@ type NetDevConfig struct {
 type VirtioNetPciConfig struct {
 	DisableLegacy string `env:"VIRTIO_NET_PCI_DISABLE_LEGACY" envDefault:"on"`
 	IOMMUPlatform bool   `env:"VIRTIO_NET_PCI_IOMMU_PLATFORM" envDefault:"true"`
-	Addr          string `env:"VIRTIO_NET_PCI_ADDR" envDefault:"0x2"`
+	Addr          string `env:"VIRTIO_NET_PCI_ADDR"           envDefault:"0x2"`
 	ROMFile       string `env:"VIRTIO_NET_PCI_ROMFILE"`
 }
 
@@ -47,9 +47,10 @@ type DiskImgConfig struct {
 }
 
 type SevConfig struct {
-	ID              string `env:"SEV_ID" envDefault:"sev0"`
-	CBitPos         int    `env:"SEV_CBITPOS" envDefault:"51"`
+	ID              string `env:"SEV_ID"                envDefault:"sev0"`
+	CBitPos         int    `env:"SEV_CBITPOS"           envDefault:"51"`
 	ReducedPhysBits int    `env:"SEV_REDUCED_PHYS_BITS" envDefault:"1"`
+	HostData        string `env:"HOST_DATA" envDefault:""`
 }
 
 type VSockConfig struct {
@@ -59,20 +60,20 @@ type VSockConfig struct {
 }
 
 type Config struct {
-	QemuBinPath  string `env:"BIN_PATH" envDefault:"qemu-system-x86_64"`
-	TmpFileLoc   string `env:"TMP_FILE_LOC" envDefault:"tmp"`
-	UseSudo      bool   `env:"USE_SUDO" envDefault:"false"`
-	EnableSEV    bool   `env:"ENABLE_SEV" envDefault:"false"`
+	QemuBinPath  string `env:"BIN_PATH"       envDefault:"qemu-system-x86_64"`
+	TmpFileLoc   string `env:"TMP_FILE_LOC"   envDefault:"tmp"`
+	UseSudo      bool   `env:"USE_SUDO"       envDefault:"false"`
+	EnableSEV    bool   `env:"ENABLE_SEV"     envDefault:"false"`
 	EnableSEVSNP bool   `env:"ENABLE_SEV_SNP" envDefault:"true"`
 
 	EnableKVM bool `env:"ENABLE_KVM" envDefault:"true"`
 
 	// machine, CPU, RAM
-	Machine  string `env:"MACHINE" envDefault:"q35"`
-	CPU      string `env:"CPU" envDefault:"EPYC"`
-	SMPCount int    `env:"SMP_COUNT" envDefault:"4"`
+	Machine  string `env:"MACHINE"     envDefault:"q35"`
+	CPU      string `env:"CPU"         envDefault:"EPYC"`
+	SMPCount int    `env:"SMP_COUNT"   envDefault:"4"`
 	MaxCPUs  int    `env:"SMP_MAXCPUS" envDefault:"64"`
-	MemID    string `env:"MEM_ID" envDefault:"ram1"`
+	MemID    string `env:"MEM_ID"      envDefault:"ram1"`
 	MemoryConfig
 
 	// Kernel hash
@@ -97,10 +98,10 @@ type Config struct {
 
 	// display
 	NoGraphic bool   `env:"NO_GRAPHIC" envDefault:"true"`
-	Monitor   string `env:"MONITOR" envDefault:"pty"`
+	Monitor   string `env:"MONITOR"    envDefault:"pty"`
 }
 
-func constructQemuArgs(config Config) []string {
+func (config Config) ConstructQemuArgs() []string {
 	args := []string{}
 
 	// virtualization
@@ -177,9 +178,14 @@ func constructQemuArgs(config Config) []string {
 	if config.EnableSEV || config.EnableSEVSNP {
 		sevType := "sev-guest"
 		kernelHash := ""
+		hostData := ""
 
 		if config.EnableSEVSNP {
 			sevType = "sev-snp-guest"
+
+			if config.SevConfig.HostData != "" {
+				hostData = fmt.Sprintf(",host-data=%s", config.SevConfig.HostData)
+			}
 		}
 
 		if config.KernelHash {
@@ -187,12 +193,13 @@ func constructQemuArgs(config Config) []string {
 		}
 
 		args = append(args, "-object",
-			fmt.Sprintf("%s,id=%s,cbitpos=%d,reduced-phys-bits=%d%s",
+			fmt.Sprintf("%s,id=%s,cbitpos=%d,reduced-phys-bits=%d%s%s",
 				sevType,
 				config.SevConfig.ID,
 				config.SevConfig.CBitPos,
 				config.SevConfig.ReducedPhysBits,
-				kernelHash))
+				kernelHash,
+				hostData))
 
 		args = append(args, "-machine",
 			fmt.Sprintf("memory-encryption=%s", config.SevConfig.ID))

manager/qemu/qemu.go

@@ -1,93 +0,0 @@
-// Copyright (c) Ultraviolet
-// SPDX-License-Identifier: Apache-2.0
-package qemu
-
-import (
-	"context"
-	"fmt"
-	"os/exec"
-
-	"github.com/gofrs/uuid"
-	"github.com/ultravioletrs/cocos/internal"
-)
-
-const (
-	firmwareVars = "OVMF_VARS"
-	KernelFile   = "bzImage"
-	rootfsFile   = "rootfs.cpio"
-)
-
-func CreateVM(ctx context.Context, cfg Config) (*exec.Cmd, error) {
-	// Create unique emu device identifiers
-	id, err := uuid.NewV4()
-	if err != nil {
-		return &exec.Cmd{}, err
-	}
-	qemuCfg := cfg
-	qemuCfg.NetDevConfig.ID = fmt.Sprintf("%s-%s", qemuCfg.NetDevConfig.ID, id)
-	qemuCfg.SevConfig.ID = fmt.Sprintf("%s-%s", qemuCfg.SevConfig.ID, id)
-
-	if !cfg.KernelHash {
-		// Copy firmware vars file
-		srcFile := qemuCfg.OVMFVarsConfig.File
-		dstFile := fmt.Sprintf("%s/%s-%s.fd", cfg.TmpFileLoc, firmwareVars, id)
-		err = internal.CopyFile(srcFile, dstFile)
-		if err != nil {
-			return &exec.Cmd{}, err
-		}
-		qemuCfg.OVMFVarsConfig.File = dstFile
-	}
-
-	// Copy img files
-	srcFile := qemuCfg.DiskImgConfig.KernelFile
-	dstFile := fmt.Sprintf("%s/%s-%s", cfg.TmpFileLoc, KernelFile, id)
-	err = internal.CopyFile(srcFile, dstFile)
-	if err != nil {
-		return &exec.Cmd{}, err
-	}
-	qemuCfg.DiskImgConfig.KernelFile = dstFile
-
-	srcFile = qemuCfg.DiskImgConfig.RootFsFile
-	dstFile = fmt.Sprintf("%s/%s-%s.gz", cfg.TmpFileLoc, rootfsFile, id)
-	err = internal.CopyFile(srcFile, dstFile)
-	if err != nil {
-		return &exec.Cmd{}, err
-	}
-	qemuCfg.DiskImgConfig.RootFsFile = dstFile
-
-	exe, args, err := ExecutableAndArgs(qemuCfg)
-	if err != nil {
-		return &exec.Cmd{}, err
-	}
-	cmd, err := runQemuVM(exe, args)
-	if err != nil {
-		return cmd, err
-	}
-
-	return cmd, nil
-}
-
-func ExecutableAndArgs(cfg Config) (string, []string, error) {
-	exe, err := exec.LookPath(cfg.QemuBinPath)
-	if err != nil {
-		return "", nil, err
-	}
-
-	args := constructQemuArgs(cfg)
-
-	if cfg.UseSudo {
-		args = append([]string{exe}, args...)
-		exe = "sudo"
-	}
-
-	return exe, args, nil
-}
-
-func runQemuVM(exe string, args []string) (*exec.Cmd, error) {
-	cmd, err := internal.RunCmdStart(exe, args...)
-	if err != nil {
-		return nil, err
-	}
-
-	return cmd, nil
-}

manager/qemu/vm.go

@@ -0,0 +1,104 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package qemu
+
+import (
+	"fmt"
+	"os/exec"
+
+	"github.com/gofrs/uuid"
+	"github.com/ultravioletrs/cocos/internal"
+	"github.com/ultravioletrs/cocos/manager/vm"
+	"github.com/ultravioletrs/cocos/pkg/manager"
+)
+
+const (
+	firmwareVars = "OVMF_VARS"
+	KernelFile   = "bzImage"
+	rootfsFile   = "rootfs.cpio"
+)
+
+type qemuVM struct {
+	config        Config
+	cmd           *exec.Cmd
+	logsChan      chan *manager.ClientStreamMessage
+	computationId string
+}
+
+func NewVM(config interface{}, logsChan chan *manager.ClientStreamMessage, computationId string) vm.VM {
+	return &qemuVM{
+		config:        config.(Config),
+		logsChan:      logsChan,
+		computationId: computationId,
+	}
+}
+
+func (v *qemuVM) Start() error {
+	// Create unique qemu device identifiers
+	id, err := uuid.NewV4()
+	if err != nil {
+		return err
+	}
+	qemuCfg := v.config
+	qemuCfg.NetDevConfig.ID = fmt.Sprintf("%s-%s", qemuCfg.NetDevConfig.ID, id)
+	qemuCfg.SevConfig.ID = fmt.Sprintf("%s-%s", qemuCfg.SevConfig.ID, id)
+
+	if !v.config.KernelHash {
+		// Copy firmware vars file
+		srcFile := qemuCfg.OVMFVarsConfig.File
+		dstFile := fmt.Sprintf("%s/%s-%s.fd", v.config.TmpFileLoc, firmwareVars, id)
+		err = internal.CopyFile(srcFile, dstFile)
+		if err != nil {
+			return err
+		}
+		qemuCfg.OVMFVarsConfig.File = dstFile
+	}
+
+	// Copy img files
+	srcFile := qemuCfg.DiskImgConfig.KernelFile
+	dstFile := fmt.Sprintf("%s/%s-%s", v.config.TmpFileLoc, KernelFile, id)
+	err = internal.CopyFile(srcFile, dstFile)
+	if err != nil {
+		return err
+	}
+	qemuCfg.DiskImgConfig.KernelFile = dstFile
+
+	srcFile = qemuCfg.DiskImgConfig.RootFsFile
+	dstFile = fmt.Sprintf("%s/%s-%s.gz", v.config.TmpFileLoc, rootfsFile, id)
+	err = internal.CopyFile(srcFile, dstFile)
+	if err != nil {
+		return err
+	}
+	qemuCfg.DiskImgConfig.RootFsFile = dstFile
+
+	exe, args, err := v.executableAndArgs()
+	if err != nil {
+		return err
+	}
+
+	v.cmd = exec.Command(exe, args...)
+	v.cmd.Stdout = &vm.Stdout{LogsChan: v.logsChan, ComputationId: v.computationId}
+	v.cmd.Stderr = &vm.Stderr{LogsChan: v.logsChan, ComputationId: v.computationId}
+
+	return v.cmd.Start()
+}
+
+func (v *qemuVM) Stop() error {
+	return v.cmd.Process.Kill()
+}
+
+func (v *qemuVM) executableAndArgs() (string, []string, error) {
+	exe, err := exec.LookPath(v.config.QemuBinPath)
+	if err != nil {
+		return "", nil, err
+	}
+
+	args := v.config.ConstructQemuArgs()
+
+	if v.config.UseSudo {
+		args = append([]string{exe}, args...)
+		exe = "sudo"
+	}
+
+	return exe, args, nil
+}

manager/qemu/vm_test.go

@@ -0,0 +1,39 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package qemu
+
+import (
+	"os/exec"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/ultravioletrs/cocos/pkg/manager"
+)
+
+func TestNewVM(t *testing.T) {
+	config := Config{}
+	logsChan := make(chan *manager.ClientStreamMessage)
+	computationId := "test-computation"
+
+	nvm := NewVM(config, logsChan, computationId)
+
+	assert.NotNil(t, nvm)
+	assert.IsType(t, &qemuVM{}, nvm)
+}
+
+func TestVM_Stop(t *testing.T) {
+	// Setup
+	v := &qemuVM{
+		cmd: exec.Command("sleep", "1"),
+	}
+
+	err := v.cmd.Start()
+	assert.NoError(t, err)
+
+	// Test
+	err = v.Stop()
+
+	// Assert
+	assert.NoError(t, err)
+	assert.Error(t, v.cmd.Wait()) // Process should have been killed
+}

manager/qemu/vsock.go

@@ -1,6 +1,6 @@
 // Copyright (c) Ultraviolet
 // SPDX-License-Identifier: Apache-2.0
-package manager
+package qemu
 
 import (
 	"encoding/json"
@@ -11,8 +11,8 @@ import (
 
 const VsockConfigPort uint32 = 9999
 
-func SendAgentConfig(cid uint32, ac agent.Computation) error {
-	conn, err := vsock.Dial(cid, VsockConfigPort, nil)
+func (v *qemuVM) SendAgentConfig(ac agent.Computation) error {
+	conn, err := vsock.Dial(uint32(v.config.GuestCID), VsockConfigPort, nil)
 	if err != nil {
 		return err
 	}

manager/service.go

@@ -4,6 +4,7 @@ package manager
 
 import (
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"log/slog"
@@ -14,7 +15,9 @@ import (
 	"github.com/cenkalti/backoff/v4"
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/manager/qemu"
+	"github.com/ultravioletrs/cocos/manager/vm"
 	"github.com/ultravioletrs/cocos/pkg/manager"
+	"golang.org/x/crypto/sha3"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
@@ -36,6 +39,9 @@ var (
 	ErrFailedToAllocatePort = errors.New("failed to allocate free port on host")
 
 	errInvalidHashLength = errors.New("hash must be of byte length 32")
+
+	// ErrFailedToCalculateHash indicates that agent computation returned an error while calculating the hash of the computation.
+	ErrFailedToCalculateHash = errors.New("error while calculating the hash of the computation")
 )
 
 // Service specifies an API that must be fulfilled by the domain service
@@ -43,6 +49,8 @@ var (
 type Service interface {
 	// Run create a computation.
 	Run(ctx context.Context, c *manager.ComputationRunReq) (string, error)
+	// Stop stops a computation.
+	Stop(ctx context.Context, computationID string) error
 	// RetrieveAgentEventsLogs Retrieve and forward agent logs and events via vsock.
 	RetrieveAgentEventsLogs()
 }
@@ -52,17 +60,21 @@ type managerService struct {
 	logger     *slog.Logger
 	agents     map[int]string // agent map of vsock cid to computationID.
 	eventsChan chan *manager.ClientStreamMessage
+	vms        map[string]vm.VM
+	vmFactory  vm.Provider
 }
 
 var _ Service = (*managerService)(nil)
 
 // New instantiates the manager service implementation.
-func New(qemuCfg qemu.Config, logger *slog.Logger, eventsChan chan *manager.ClientStreamMessage) Service {
+func New(qemuCfg qemu.Config, logger *slog.Logger, eventsChan chan *manager.ClientStreamMessage, vmFactory vm.Provider) Service {
 	ms := &managerService{
 		qemuCfg:    qemuCfg,
 		logger:     logger,
 		agents:     make(map[int]string),
+		vms:        make(map[string]vm.VM),
 		eventsChan: eventsChan,
+		vmFactory:  vmFactory,
 	}
 	return ms
 }
@@ -104,16 +116,27 @@ func (ms *managerService) Run(ctx context.Context, c *manager.ComputationRunReq)
 	}
 	ms.qemuCfg.HostFwdAgent = agentPort
 
+	ch, err := computationHash(ac)
+	if err != nil {
+		ms.publishEvent("vm-provision", c.Id, "failed", json.RawMessage{})
+		return "", errors.Wrap(ErrFailedToCalculateHash, err)
+	}
+
+	// Define host-data value of QEMU for SEV-SNP, with a base64 encoding of the computation hash.
+	ms.qemuCfg.SevConfig.HostData = base64.StdEncoding.EncodeToString(ch[:])
+
+	cvm := ms.vmFactory(ms.qemuCfg, ms.eventsChan, c.Id)
 	ms.publishEvent("vm-provision", c.Id, "in-progress", json.RawMessage{})
-	if _, err = qemu.CreateVM(ctx, ms.qemuCfg); err != nil {
+	if err = cvm.Start(); err != nil {
 		ms.publishEvent("vm-provision", c.Id, "failed", json.RawMessage{})
 		return "", err
 	}
+	ms.vms[c.Id] = cvm
 
 	ms.agents[ms.qemuCfg.VSockConfig.GuestCID] = c.Id
 
 	err = backoff.Retry(func() error {
-		return SendAgentConfig(uint32(ms.qemuCfg.VSockConfig.GuestCID), ac)
+		return cvm.SendAgentConfig(ac)
 	}, backoff.NewExponentialBackOff())
 	if err != nil {
 		return "", err
@@ -125,6 +148,18 @@ func (ms *managerService) Run(ctx context.Context, c *manager.ComputationRunReq)
 	return fmt.Sprint(ms.qemuCfg.HostFwdAgent), nil
 }
 
+func (ms *managerService) Stop(ctx context.Context, computationID string) error {
+	cvm, ok := ms.vms[computationID]
+	if !ok {
+		return ErrNotFound
+	}
+	if err := cvm.Stop(); err != nil {
+		return err
+	}
+	delete(ms.vms, computationID)
+	return nil
+}
+
 func getFreePort() (int, error) {
 	listener, err := net.Listen("tcp", "")
 	if err != nil {
@@ -156,3 +191,12 @@ func (ms *managerService) publishEvent(event, cmpID, status string, details json
 		},
 	}
 }
+
+func computationHash(ac agent.Computation) ([32]byte, error) {
+	jsonData, err := json.Marshal(ac)
+	if err != nil {
+		return [32]byte{}, err
+	}
+
+	return sha3.Sum256(jsonData), nil
+}

manager/tracing/tracing.go

@@ -29,6 +29,13 @@ func (tm *tracingMiddleware) Run(ctx context.Context, mc *pkgmanager.Computation
 	return tm.svc.Run(ctx, mc)
 }
 
+func (tm *tracingMiddleware) Stop(ctx context.Context, computationID string) error {
+	ctx, span := tm.tracer.Start(ctx, "stop")
+	defer span.End()
+
+	return tm.svc.Stop(ctx, computationID)
+}
+
 func (tm *tracingMiddleware) RetrieveAgentEventsLogs() {
 	tm.svc.RetrieveAgentEventsLogs()
 }

manager/vm/logging.go

@@ -0,0 +1,133 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package vm
+
+import (
+	"bytes"
+	"errors"
+	"io"
+
+	"github.com/ultravioletrs/cocos/pkg/manager"
+	"google.golang.org/protobuf/types/known/timestamppb"
+)
+
+var (
+	_                      io.Writer = &Stdout{}
+	_                      io.Writer = &Stderr{}
+	ErrFailedToSendMessage           = errors.New("failed to send message to channel")
+	ErrPanicRecovered                = errors.New("panic recovered: channel may be closed")
+)
+
+const bufSize = 1024
+
+type Stdout struct {
+	LogsChan      chan *manager.ClientStreamMessage
+	ComputationId string
+}
+
+// safeSend safely sends a message to the channel and returns an error on failure.
+func safeSend(ch chan *manager.ClientStreamMessage, msg *manager.ClientStreamMessage) (err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			// Recover from panic if the channel is closed
+			err = ErrPanicRecovered
+		}
+	}()
+	select {
+	case ch <- msg:
+		return nil
+	default:
+		// Channel is full or closed
+		return ErrFailedToSendMessage
+	}
+}
+
+// Write implements io.Writer.
+func (s *Stdout) Write(p []byte) (n int, err error) {
+	inBuf := bytes.NewBuffer(p)
+
+	buf := make([]byte, bufSize)
+
+	for {
+		n, err := inBuf.Read(buf)
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return len(p) - inBuf.Len(), err
+		}
+
+		msg := &manager.ClientStreamMessage{
+			Message: &manager.ClientStreamMessage_AgentLog{
+				AgentLog: &manager.AgentLog{
+					Message:       string(buf[:n]),
+					ComputationId: s.ComputationId,
+					Level:         "debug",
+					Timestamp:     timestamppb.Now(),
+				},
+			},
+		}
+
+		if err := safeSend(s.LogsChan, msg); err != nil {
+			return len(p) - inBuf.Len(), err
+		}
+	}
+
+	return len(p), nil
+}
+
+type Stderr struct {
+	LogsChan      chan *manager.ClientStreamMessage
+	ComputationId string
+}
+
+// Write implements io.Writer.
+func (s *Stderr) Write(p []byte) (n int, err error) {
+	inBuf := bytes.NewBuffer(p)
+
+	buf := make([]byte, bufSize)
+
+	for {
+		n, err := inBuf.Read(buf)
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+			return len(p) - inBuf.Len(), err
+		}
+
+		msg := &manager.ClientStreamMessage{
+			Message: &manager.ClientStreamMessage_AgentLog{
+				AgentLog: &manager.AgentLog{
+					Message:       string(buf[:n]),
+					ComputationId: s.ComputationId,
+					Level:         "error",
+					Timestamp:     timestamppb.Now(),
+				},
+			},
+		}
+
+		if err := safeSend(s.LogsChan, msg); err != nil {
+			return len(p) - inBuf.Len(), err
+		}
+	}
+
+	// Ensure vm-provision failure message is sent
+	eventMsg := &manager.ClientStreamMessage{
+		Message: &manager.ClientStreamMessage_AgentEvent{
+			AgentEvent: &manager.AgentEvent{
+				ComputationId: s.ComputationId,
+				EventType:     "vm-provision",
+				Timestamp:     timestamppb.Now(),
+				Originator:    "manager",
+				Status:        "failed",
+			},
+		},
+	}
+
+	if err := safeSend(s.LogsChan, eventMsg); err != nil {
+		return len(p), err
+	}
+
+	return len(p), nil
+}

manager/vm/logging_test.go

@@ -0,0 +1,171 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package vm
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/ultravioletrs/cocos/pkg/manager"
+)
+
+func TestStdoutWrite(t *testing.T) {
+	tests := []struct {
+		name           string
+		input          string
+		expectedWrites int
+	}{
+		{
+			name:           "Single write within buffer size",
+			input:          "Hello, World!",
+			expectedWrites: 1,
+		},
+		{
+			name:           "Multiple writes within buffer size",
+			input:          "This is a longer message that will be split into multiple writes.",
+			expectedWrites: 1,
+		},
+		{
+			name:           "Large write exceeding buffer size",
+			input:          string(make([]byte, bufSize*2+1)),
+			expectedWrites: 3,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			logsChan := make(chan *manager.ClientStreamMessage, 10)
+			s := &Stdout{
+				LogsChan:      logsChan,
+				ComputationId: "test-computation",
+			}
+
+			n, err := s.Write([]byte(tt.input))
+
+			assert.NoError(t, err)
+			assert.Equal(t, len(tt.input), n)
+
+			var receivedWrites int
+			for i := 0; i < tt.expectedWrites; i++ {
+				select {
+				case msg := <-logsChan:
+					receivedWrites++
+					agentLog := msg.GetAgentLog()
+					assert.NotNil(t, agentLog)
+					assert.Equal(t, "test-computation", agentLog.ComputationId)
+					assert.Equal(t, "debug", agentLog.Level)
+					assert.NotEmpty(t, agentLog.Message)
+					assert.NotNil(t, agentLog.Timestamp)
+				case <-time.After(time.Second):
+					t.Fatal("Timed out waiting for log message")
+				}
+			}
+
+			assert.Equal(t, tt.expectedWrites, receivedWrites)
+		})
+	}
+}
+
+func TestStderrWrite(t *testing.T) {
+	tests := []struct {
+		name           string
+		input          string
+		expectedWrites int
+	}{
+		{
+			name:           "Single write within buffer size",
+			input:          "Error: Something went wrong",
+			expectedWrites: 1,
+		},
+		{
+			name:           "Multiple writes within buffer size",
+			input:          "This is a longer error message that will be split into multiple writes.",
+			expectedWrites: 1,
+		},
+		{
+			name:           "Large write exceeding buffer size",
+			input:          string(make([]byte, bufSize*2)),
+			expectedWrites: 3,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			logsChan := make(chan *manager.ClientStreamMessage, 10)
+			s := &Stderr{
+				LogsChan:      logsChan,
+				ComputationId: "test-computation",
+			}
+
+			n, err := s.Write([]byte(tt.input))
+
+			assert.NoError(t, err)
+			assert.Equal(t, len(tt.input), n)
+
+			var receivedWrites int
+			for i := 0; i < tt.expectedWrites; i++ {
+				select {
+				case msg := <-logsChan:
+					receivedWrites++
+					switch msg.Message.(type) {
+					case *manager.ClientStreamMessage_AgentLog:
+						agentLog := msg.GetAgentLog()
+						assert.NotNil(t, agentLog)
+						assert.Equal(t, "test-computation", agentLog.ComputationId)
+						assert.Equal(t, "error", agentLog.Level)
+						assert.NotEmpty(t, agentLog.Message)
+						assert.NotNil(t, agentLog.Timestamp)
+					case *manager.ClientStreamMessage_AgentEvent:
+						agentEvent := msg.GetAgentEvent()
+						assert.NotNil(t, agentEvent)
+						assert.Equal(t, "test-computation", agentEvent.ComputationId)
+						assert.Equal(t, "vm-provision", agentEvent.EventType)
+						assert.Equal(t, "failed", agentEvent.Status)
+						assert.NotNil(t, agentEvent.Timestamp)
+					}
+				case <-time.After(time.Second):
+					t.Fatal("Timed out waiting for log message")
+				}
+			}
+
+			assert.Equal(t, tt.expectedWrites, receivedWrites)
+		})
+	}
+}
+
+func TestStdoutWriteErrorHandling(t *testing.T) {
+	logsChan := make(chan *manager.ClientStreamMessage, 1)
+	s := &Stdout{
+		LogsChan:      logsChan,
+		ComputationId: "test-computation",
+	}
+
+	// Test with a closed channel to simulate an error condition
+	close(logsChan)
+
+	message := []byte("This should fail")
+	n, err := s.Write(message)
+
+	assert.Error(t, err)
+	assert.Equal(t, len(message), n)
+	assert.Equal(t, ErrPanicRecovered, err)
+}
+
+func TestStderrWriteErrorHandling(t *testing.T) {
+	logsChan := make(chan *manager.ClientStreamMessage, 1)
+	s := &Stderr{
+		LogsChan:      logsChan,
+		ComputationId: "test-computation",
+	}
+
+	// Test with a closed channel to simulate an error condition
+	close(logsChan)
+
+	message := []byte("This should fail")
+	n, err := s.Write(message)
+
+	assert.Error(t, err)
+	assert.Equal(t, len(message), n)
+	assert.Equal(t, ErrPanicRecovered, err)
+}

manager/vm/mocks/provider.go

@@ -0,0 +1,52 @@
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package mocks
+
+import (
+	mock "github.com/stretchr/testify/mock"
+	manager "github.com/ultravioletrs/cocos/pkg/manager"
+
+	vm "github.com/ultravioletrs/cocos/manager/vm"
+)
+
+// Provider is an autogenerated mock type for the Provider type
+type Provider struct {
+	mock.Mock
+}
+
+// Execute provides a mock function with given fields: config, logsChan, computationId
+func (_m *Provider) Execute(config interface{}, logsChan chan *manager.ClientStreamMessage, computationId string) vm.VM {
+	ret := _m.Called(config, logsChan, computationId)
+
+	if len(ret) == 0 {
+		panic("no return value specified for Execute")
+	}
+
+	var r0 vm.VM
+	if rf, ok := ret.Get(0).(func(interface{}, chan *manager.ClientStreamMessage, string) vm.VM); ok {
+		r0 = rf(config, logsChan, computationId)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(vm.VM)
+		}
+	}
+
+	return r0
+}
+
+// NewProvider creates a new instance of Provider. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewProvider(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *Provider {
+	mock := &Provider{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

manager/vm/mocks/vm.go

@@ -0,0 +1,84 @@
+// Code generated by mockery v2.43.2. DO NOT EDIT.
+
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package mocks
+
+import (
+	mock "github.com/stretchr/testify/mock"
+	agent "github.com/ultravioletrs/cocos/agent"
+)
+
+// VM is an autogenerated mock type for the VM type
+type VM struct {
+	mock.Mock
+}
+
+// SendAgentConfig provides a mock function with given fields: ac
+func (_m *VM) SendAgentConfig(ac agent.Computation) error {
+	ret := _m.Called(ac)
+
+	if len(ret) == 0 {
+		panic("no return value specified for SendAgentConfig")
+	}
+
+	var r0 error
+	if rf, ok := ret.Get(0).(func(agent.Computation) error); ok {
+		r0 = rf(ac)
+	} else {
+		r0 = ret.Error(0)
+	}
+
+	return r0
+}
+
+// Start provides a mock function with given fields:
+func (_m *VM) 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
+}
+
+// Stop provides a mock function with given fields:
+func (_m *VM) 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
+}
+
+// NewVM creates a new instance of VM. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
+// The first argument is typically a *testing.T value.
+func NewVM(t interface {
+	mock.TestingT
+	Cleanup(func())
+}) *VM {
+	mock := &VM{}
+	mock.Mock.Test(t)
+
+	t.Cleanup(func() { mock.AssertExpectations(t) })
+
+	return mock
+}

manager/vm/vm.go

@@ -0,0 +1,20 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package vm
+
+import (
+	"github.com/ultravioletrs/cocos/agent"
+	"github.com/ultravioletrs/cocos/pkg/manager"
+)
+
+// VM represents a virtual machine.
+//
+//go:generate mockery --name VM --output=./mocks --filename vm.go --quiet --note "Copyright (c) Ultraviolet \n // SPDX-License-Identifier: Apache-2.0"
+type VM interface {
+	Start() error
+	Stop() error
+	SendAgentConfig(ac agent.Computation) error
+}
+
+//go:generate mockery --name Provider --output=./mocks --filename provider.go --quiet --note "Copyright (c) Ultraviolet \n // SPDX-License-Identifier: Apache-2.0"
+type Provider func(config interface{}, logsChan chan *manager.ClientStreamMessage, computationId string) VM

pkg/clients/grpc/connect.go

@@ -9,11 +9,13 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"path"
 	"time"
 
 	"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/google/go-sev-guest/validate"
 	"github.com/google/go-sev-guest/verify"
 	"github.com/google/go-sev-guest/verify/trust"
@@ -32,6 +34,13 @@ const (
 	withmTLS
 )
 
+const (
+	cocosDirectory   = ".cocos"
+	caBundleName     = "ask_ark.pem"
+	productNameMilan = "Milan"
+	productNameGenoa = "Genoa"
+)
+
 var (
 	errGrpcConnect      = errors.New("failed to connect to grpc server")
 	errGrpcClose        = errors.New("failed to close grpc connection")
@@ -63,7 +72,7 @@ type Config struct {
 
 type AttestationConfiguration struct {
 	SNPPolicy   *check.Policy      `json:"snp_policy,omitempty"`
-	RootOFTrust *check.RootOfTrust `json:"root_of_trust,omitempty"`
+	RootOfTrust *check.RootOfTrust `json:"root_of_trust,omitempty"`
 }
 
 type Client interface {
@@ -132,7 +141,7 @@ func connect(cfg Config) (*grpc.ClientConn, security, error) {
 	tc := insecure.NewCredentials()
 
 	if cfg.AttestedTLS {
-		err := readManifest(cfg)
+		err := ReadManifest(cfg.Manifest, &attestationConfiguration)
 		if err != nil {
 			return nil, secure, fmt.Errorf("failed to read Manifest %w", err)
 		}
@@ -183,16 +192,16 @@ func connect(cfg Config) (*grpc.ClientConn, security, error) {
 	return conn, secure, nil
 }
 
-func readManifest(cfg Config) error {
-	if cfg.Manifest != "" {
-		manifest, err := os.Open(cfg.Manifest)
+func ReadManifest(manifestPath string, attestationConfiguration *AttestationConfiguration) error {
+	if manifestPath != "" {
+		manifest, err := os.Open(manifestPath)
 		if err != nil {
 			return errors.Wrap(errManifestOpen, err)
 		}
 		defer manifest.Close()
 
 		decoder := json.NewDecoder(manifest)
-		err = decoder.Decode(&attestationConfiguration)
+		err = decoder.Decode(attestationConfiguration)
 		if err != nil {
 			return errors.Wrap(errManifestDecode, err)
 		}
@@ -226,7 +235,7 @@ func verifyAttestationReportTLS(rawCerts [][]byte, verifiedChains [][]*x509.Cert
 			attestationConfiguration.SNPPolicy.ReportData = expectedReportData[:]
 
 			// Attestation verification and validation
-			sopts, err := verify.RootOfTrustToOptions(attestationConfiguration.RootOFTrust)
+			sopts, err := verify.RootOfTrustToOptions(attestationConfiguration.RootOfTrust)
 			if err != nil {
 				return errors.Wrap(errAttVerification, err)
 			}
@@ -243,6 +252,10 @@ func verifyAttestationReportTLS(rawCerts [][]byte, verifiedChains [][]*x509.Cert
 				return errors.Wrap(errAttVerification, err)
 			}
 
+			if err := fillInAttestationLocal(attestationPB); err != nil {
+				return err
+			}
+
 			if err = verify.SnpAttestation(attestationPB, sopts); err != nil {
 				return errors.Wrap(errAttVerification, err)
 			}
@@ -278,3 +291,32 @@ func checkIfCertificateSelfSigned(cert *x509.Certificate) error {
 
 	return nil
 }
+
+func fillInAttestationLocal(attestation *sevsnp.Attestation) error {
+	product := attestationConfiguration.RootOfTrust.Product
+
+	chain := attestation.GetCertificateChain()
+	if chain == nil {
+		chain = &sevsnp.CertificateChain{}
+		attestation.CertificateChain = chain
+	}
+	if len(chain.GetAskCert()) == 0 || len(chain.GetArkCert()) == 0 {
+		homePath, err := os.UserHomeDir()
+		if err != nil {
+			return err
+		}
+
+		bundleFilePath := path.Join(homePath, cocosDirectory, product, caBundleName)
+		if _, err := os.Stat(bundleFilePath); err == nil {
+			amdRootCerts := trust.AMDRootCerts{}
+			if err := amdRootCerts.FromKDSCert(bundleFilePath); err != nil {
+				return err
+			}
+
+			chain.ArkCert = amdRootCerts.ProductCerts.Ark.Raw
+			chain.AskCert = amdRootCerts.ProductCerts.Ask.Raw
+		}
+	}
+
+	return nil
+}

pkg/manager/manager.pb.go

@@ -3,8 +3,8 @@
 
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.33.0
-// 	protoc        v4.25.3
+// 	protoc-gen-go v1.34.2
+// 	protoc        v5.27.2
 // source: manager/manager.proto
 
 package manager
@@ -71,6 +71,53 @@ func (x *Terminate) GetMessage() string {
 	return ""
 }
 
+type StopComputation struct {
+	state         protoimpl.MessageState
+	sizeCache     protoimpl.SizeCache
+	unknownFields protoimpl.UnknownFields
+
+	ComputationId string `protobuf:"bytes,1,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
+}
+
+func (x *StopComputation) Reset() {
+	*x = StopComputation{}
+	if protoimpl.UnsafeEnabled {
+		mi := &file_manager_manager_proto_msgTypes[1]
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		ms.StoreMessageInfo(mi)
+	}
+}
+
+func (x *StopComputation) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*StopComputation) ProtoMessage() {}
+
+func (x *StopComputation) ProtoReflect() protoreflect.Message {
+	mi := &file_manager_manager_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 StopComputation.ProtoReflect.Descriptor instead.
+func (*StopComputation) Descriptor() ([]byte, []int) {
+	return file_manager_manager_proto_rawDescGZIP(), []int{1}
+}
+
+func (x *StopComputation) GetComputationId() string {
+	if x != nil {
+		return x.ComputationId
+	}
+	return ""
+}
+
 type RunResponse struct {
 	state         protoimpl.MessageState
 	sizeCache     protoimpl.SizeCache
@@ -83,7 +130,7 @@ type RunResponse struct {
 func (x *RunResponse) Reset() {
 	*x = RunResponse{}
 	if protoimpl.UnsafeEnabled {
-		mi := &file_manager_manager_proto_msgTypes[1]
+		mi := &file_manager_manager_proto_msgTypes[2]
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		ms.StoreMessageInfo(mi)
 	}
@@ -96,7 +143,7 @@ func (x *RunResponse) String() string {
 func (*RunResponse) ProtoMessage() {}
 
 func (x *RunResponse) ProtoReflect() protoreflect.Message {
-	mi := &file_manager_manager_proto_msgTypes[1]
+	mi := &file_manager_manager_proto_msgTypes[2]
 	if protoimpl.UnsafeEnabled && x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -109,7 +156,7 @@ func (x *RunResponse) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use RunResponse.ProtoReflect.Descriptor instead.
 func (*RunResponse) Descriptor() ([]byte, []int) {
-	return file_manager_manager_proto_rawDescGZIP(), []int{1}
+	return file_manager_manager_proto_rawDescGZIP(), []int{2}
 }
 
 func (x *RunResponse) GetAgentPort() string {
@@ -142,7 +189,7 @@ type AgentEvent struct {
 func (x *AgentEvent) Reset() {
 	*x = AgentEvent{}
 	if protoimpl.UnsafeEnabled {
-		mi := &file_manager_manager_proto_msgTypes[2]
+		mi := &file_manager_manager_proto_msgTypes[3]
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		ms.StoreMessageInfo(mi)
 	}
@@ -155,7 +202,7 @@ func (x *AgentEvent) String() string {
 func (*AgentEvent) ProtoMessage() {}
 
 func (x *AgentEvent) ProtoReflect() protoreflect.Message {
-	mi := &file_manager_manager_proto_msgTypes[2]
+	mi := &file_manager_manager_proto_msgTypes[3]
 	if protoimpl.UnsafeEnabled && x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -168,7 +215,7 @@ func (x *AgentEvent) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use AgentEvent.ProtoReflect.Descriptor instead.
 func (*AgentEvent) Descriptor() ([]byte, []int) {
-	return file_manager_manager_proto_rawDescGZIP(), []int{2}
+	return file_manager_manager_proto_rawDescGZIP(), []int{3}
 }
 
 func (x *AgentEvent) GetEventType() string {
@@ -227,7 +274,7 @@ type AgentLog struct {
 func (x *AgentLog) Reset() {
 	*x = AgentLog{}
 	if protoimpl.UnsafeEnabled {
-		mi := &file_manager_manager_proto_msgTypes[3]
+		mi := &file_manager_manager_proto_msgTypes[4]
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		ms.StoreMessageInfo(mi)
 	}
@@ -240,7 +287,7 @@ func (x *AgentLog) String() string {
 func (*AgentLog) ProtoMessage() {}
 
 func (x *AgentLog) ProtoReflect() protoreflect.Message {
-	mi := &file_manager_manager_proto_msgTypes[3]
+	mi := &file_manager_manager_proto_msgTypes[4]
 	if protoimpl.UnsafeEnabled && x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -253,7 +300,7 @@ func (x *AgentLog) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use AgentLog.ProtoReflect.Descriptor instead.
 func (*AgentLog) Descriptor() ([]byte, []int) {
-	return file_manager_manager_proto_rawDescGZIP(), []int{3}
+	return file_manager_manager_proto_rawDescGZIP(), []int{4}
 }
 
 func (x *AgentLog) GetMessage() string {
@@ -300,7 +347,7 @@ type ClientStreamMessage struct {
 func (x *ClientStreamMessage) Reset() {
 	*x = ClientStreamMessage{}
 	if protoimpl.UnsafeEnabled {
-		mi := &file_manager_manager_proto_msgTypes[4]
+		mi := &file_manager_manager_proto_msgTypes[5]
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		ms.StoreMessageInfo(mi)
 	}
@@ -313,7 +360,7 @@ func (x *ClientStreamMessage) String() string {
 func (*ClientStreamMessage) ProtoMessage() {}
 
 func (x *ClientStreamMessage) ProtoReflect() protoreflect.Message {
-	mi := &file_manager_manager_proto_msgTypes[4]
+	mi := &file_manager_manager_proto_msgTypes[5]
 	if protoimpl.UnsafeEnabled && x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -326,7 +373,7 @@ func (x *ClientStreamMessage) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use ClientStreamMessage.ProtoReflect.Descriptor instead.
 func (*ClientStreamMessage) Descriptor() ([]byte, []int) {
-	return file_manager_manager_proto_rawDescGZIP(), []int{4}
+	return file_manager_manager_proto_rawDescGZIP(), []int{5}
 }
 
 func (m *ClientStreamMessage) GetMessage() isClientStreamMessage_Message {
@@ -388,13 +435,14 @@ type ServerStreamMessage struct {
 	//
 	//	*ServerStreamMessage_RunReq
 	//	*ServerStreamMessage_TerminateReq
+	//	*ServerStreamMessage_StopComputation
 	Message isServerStreamMessage_Message `protobuf_oneof:"message"`
 }
 
 func (x *ServerStreamMessage) Reset() {
 	*x = ServerStreamMessage{}
 	if protoimpl.UnsafeEnabled {
-		mi := &file_manager_manager_proto_msgTypes[5]
+		mi := &file_manager_manager_proto_msgTypes[6]
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		ms.StoreMessageInfo(mi)
 	}
@@ -407,7 +455,7 @@ func (x *ServerStreamMessage) String() string {
 func (*ServerStreamMessage) ProtoMessage() {}
 
 func (x *ServerStreamMessage) ProtoReflect() protoreflect.Message {
-	mi := &file_manager_manager_proto_msgTypes[5]
+	mi := &file_manager_manager_proto_msgTypes[6]
 	if protoimpl.UnsafeEnabled && x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -420,7 +468,7 @@ func (x *ServerStreamMessage) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use ServerStreamMessage.ProtoReflect.Descriptor instead.
 func (*ServerStreamMessage) Descriptor() ([]byte, []int) {
-	return file_manager_manager_proto_rawDescGZIP(), []int{5}
+	return file_manager_manager_proto_rawDescGZIP(), []int{6}
 }
 
 func (m *ServerStreamMessage) GetMessage() isServerStreamMessage_Message {
@@ -444,6 +492,13 @@ func (x *ServerStreamMessage) GetTerminateReq() *Terminate {
 	return nil
 }
 
+func (x *ServerStreamMessage) GetStopComputation() *StopComputation {
+	if x, ok := x.GetMessage().(*ServerStreamMessage_StopComputation); ok {
+		return x.StopComputation
+	}
+	return nil
+}
+
 type isServerStreamMessage_Message interface {
 	isServerStreamMessage_Message()
 }
@@ -456,10 +511,16 @@ type ServerStreamMessage_TerminateReq struct {
 	TerminateReq *Terminate `protobuf:"bytes,2,opt,name=terminateReq,proto3,oneof"`
 }
 
+type ServerStreamMessage_StopComputation struct {
+	StopComputation *StopComputation `protobuf:"bytes,3,opt,name=stopComputation,proto3,oneof"`
+}
+
 func (*ServerStreamMessage_RunReq) isServerStreamMessage_Message() {}
 
 func (*ServerStreamMessage_TerminateReq) isServerStreamMessage_Message() {}
 
+func (*ServerStreamMessage_StopComputation) isServerStreamMessage_Message() {}
+
 type ComputationRunReq struct {
 	state         protoimpl.MessageState
 	sizeCache     protoimpl.SizeCache
@@ -477,7 +538,7 @@ type ComputationRunReq struct {
 func (x *ComputationRunReq) Reset() {
 	*x = ComputationRunReq{}
 	if protoimpl.UnsafeEnabled {
-		mi := &file_manager_manager_proto_msgTypes[6]
+		mi := &file_manager_manager_proto_msgTypes[7]
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		ms.StoreMessageInfo(mi)
 	}
@@ -490,7 +551,7 @@ func (x *ComputationRunReq) String() string {
 func (*ComputationRunReq) ProtoMessage() {}
 
 func (x *ComputationRunReq) ProtoReflect() protoreflect.Message {
-	mi := &file_manager_manager_proto_msgTypes[6]
+	mi := &file_manager_manager_proto_msgTypes[7]
 	if protoimpl.UnsafeEnabled && x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -503,7 +564,7 @@ func (x *ComputationRunReq) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use ComputationRunReq.ProtoReflect.Descriptor instead.
 func (*ComputationRunReq) Descriptor() ([]byte, []int) {
-	return file_manager_manager_proto_rawDescGZIP(), []int{6}
+	return file_manager_manager_proto_rawDescGZIP(), []int{7}
 }
 
 func (x *ComputationRunReq) GetId() string {
@@ -566,7 +627,7 @@ type ResultConsumer struct {
 func (x *ResultConsumer) Reset() {
 	*x = ResultConsumer{}
 	if protoimpl.UnsafeEnabled {
-		mi := &file_manager_manager_proto_msgTypes[7]
+		mi := &file_manager_manager_proto_msgTypes[8]
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		ms.StoreMessageInfo(mi)
 	}
@@ -579,7 +640,7 @@ func (x *ResultConsumer) String() string {
 func (*ResultConsumer) ProtoMessage() {}
 
 func (x *ResultConsumer) ProtoReflect() protoreflect.Message {
-	mi := &file_manager_manager_proto_msgTypes[7]
+	mi := &file_manager_manager_proto_msgTypes[8]
 	if protoimpl.UnsafeEnabled && x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -592,7 +653,7 @@ func (x *ResultConsumer) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use ResultConsumer.ProtoReflect.Descriptor instead.
 func (*ResultConsumer) Descriptor() ([]byte, []int) {
-	return file_manager_manager_proto_rawDescGZIP(), []int{7}
+	return file_manager_manager_proto_rawDescGZIP(), []int{8}
 }
 
 func (x *ResultConsumer) GetUserKey() []byte {
@@ -614,7 +675,7 @@ type Dataset struct {
 func (x *Dataset) Reset() {
 	*x = Dataset{}
 	if protoimpl.UnsafeEnabled {
-		mi := &file_manager_manager_proto_msgTypes[8]
+		mi := &file_manager_manager_proto_msgTypes[9]
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		ms.StoreMessageInfo(mi)
 	}
@@ -627,7 +688,7 @@ func (x *Dataset) String() string {
 func (*Dataset) ProtoMessage() {}
 
 func (x *Dataset) ProtoReflect() protoreflect.Message {
-	mi := &file_manager_manager_proto_msgTypes[8]
+	mi := &file_manager_manager_proto_msgTypes[9]
 	if protoimpl.UnsafeEnabled && x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -640,7 +701,7 @@ func (x *Dataset) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use Dataset.ProtoReflect.Descriptor instead.
 func (*Dataset) Descriptor() ([]byte, []int) {
-	return file_manager_manager_proto_rawDescGZIP(), []int{8}
+	return file_manager_manager_proto_rawDescGZIP(), []int{9}
 }
 
 func (x *Dataset) GetHash() []byte {
@@ -669,7 +730,7 @@ type Algorithm struct {
 func (x *Algorithm) Reset() {
 	*x = Algorithm{}
 	if protoimpl.UnsafeEnabled {
-		mi := &file_manager_manager_proto_msgTypes[9]
+		mi := &file_manager_manager_proto_msgTypes[10]
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		ms.StoreMessageInfo(mi)
 	}
@@ -682,7 +743,7 @@ func (x *Algorithm) String() string {
 func (*Algorithm) ProtoMessage() {}
 
 func (x *Algorithm) ProtoReflect() protoreflect.Message {
-	mi := &file_manager_manager_proto_msgTypes[9]
+	mi := &file_manager_manager_proto_msgTypes[10]
 	if protoimpl.UnsafeEnabled && x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -695,7 +756,7 @@ func (x *Algorithm) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use Algorithm.ProtoReflect.Descriptor instead.
 func (*Algorithm) Descriptor() ([]byte, []int) {
-	return file_manager_manager_proto_rawDescGZIP(), []int{9}
+	return file_manager_manager_proto_rawDescGZIP(), []int{10}
 }
 
 func (x *Algorithm) GetHash() []byte {
@@ -730,7 +791,7 @@ type AgentConfig struct {
 func (x *AgentConfig) Reset() {
 	*x = AgentConfig{}
 	if protoimpl.UnsafeEnabled {
-		mi := &file_manager_manager_proto_msgTypes[10]
+		mi := &file_manager_manager_proto_msgTypes[11]
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		ms.StoreMessageInfo(mi)
 	}
@@ -743,7 +804,7 @@ func (x *AgentConfig) String() string {
 func (*AgentConfig) ProtoMessage() {}
 
 func (x *AgentConfig) ProtoReflect() protoreflect.Message {
-	mi := &file_manager_manager_proto_msgTypes[10]
+	mi := &file_manager_manager_proto_msgTypes[11]
 	if protoimpl.UnsafeEnabled && x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -756,7 +817,7 @@ func (x *AgentConfig) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use AgentConfig.ProtoReflect.Descriptor instead.
 func (*AgentConfig) Descriptor() ([]byte, []int) {
-	return file_manager_manager_proto_rawDescGZIP(), []int{10}
+	return file_manager_manager_proto_rawDescGZIP(), []int{11}
 }
 
 func (x *AgentConfig) GetPort() string {
@@ -824,57 +885,65 @@ var file_manager_manager_proto_rawDesc = []byte{
 	0x66, 0x2f, 0x74, 0x69, 0x6d, 0x65, 0x73, 0x74, 0x61, 0x6d, 0x70, 0x2e, 0x70, 0x72, 0x6f, 0x74,
 	0x6f, 0x22, 0x25, 0x0a, 0x09, 0x54, 0x65, 0x72, 0x6d, 0x69, 0x6e, 0x61, 0x74, 0x65, 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, 0x22, 0x53, 0x0a, 0x0b, 0x52, 0x75, 0x6e, 0x52,
-	0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x1d, 0x0a, 0x0a, 0x61, 0x67, 0x65, 0x6e, 0x74,
-	0x5f, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x61, 0x67, 0x65,
-	0x6e, 0x74, 0x50, 0x6f, 0x72, 0x74, 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, 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, 0xbb, 0x01, 0x0a,
-	0x13, 0x43, 0x6c, 0x69, 0x65, 0x6e, 0x74, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73,
-	0x73, 0x61, 0x67, 0x65, 0x12, 0x30, 0x0a, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x6f,
-	0x67, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x11, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65,
-	0x72, 0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x48, 0x00, 0x52, 0x08, 0x61, 0x67,
-	0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67, 0x12, 0x36, 0x0a, 0x0b, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f,
-	0x65, 0x76, 0x65, 0x6e, 0x74, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x13, 0x2e, 0x6d, 0x61,
-	0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x41, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74,
-	0x48, 0x00, 0x52, 0x0a, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x12, 0x2f,
-	0x0a, 0x07, 0x72, 0x75, 0x6e, 0x5f, 0x72, 0x65, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b, 0x32,
-	0x14, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x52, 0x75, 0x6e, 0x52, 0x65, 0x73,
-	0x70, 0x6f, 0x6e, 0x73, 0x65, 0x48, 0x00, 0x52, 0x06, 0x72, 0x75, 0x6e, 0x52, 0x65, 0x73, 0x42,
-	0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x90, 0x01, 0x0a, 0x13, 0x53,
-	0x65, 0x72, 0x76, 0x65, 0x72, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73, 0x73, 0x61,
-	0x67, 0x65, 0x12, 0x34, 0x0a, 0x06, 0x72, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x18, 0x01, 0x20, 0x01,
-	0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x43, 0x6f, 0x6d,
-	0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x48, 0x00,
-	0x52, 0x06, 0x72, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x12, 0x38, 0x0a, 0x0c, 0x74, 0x65, 0x72, 0x6d,
-	0x69, 0x6e, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x12,
-	0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x54, 0x65, 0x72, 0x6d, 0x69, 0x6e, 0x61,
-	0x74, 0x65, 0x48, 0x00, 0x52, 0x0c, 0x74, 0x65, 0x72, 0x6d, 0x69, 0x6e, 0x61, 0x74, 0x65, 0x52,
-	0x65, 0x71, 0x42, 0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0xb6, 0x02,
+	0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0x38, 0x0a, 0x0f, 0x53, 0x74, 0x6f, 0x70,
+	0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x12, 0x25, 0x0a, 0x0e, 0x63,
+	0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x5f, 0x69, 0x64, 0x18, 0x01, 0x20,
+	0x01, 0x28, 0x09, 0x52, 0x0d, 0x63, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e,
+	0x49, 0x64, 0x22, 0x53, 0x0a, 0x0b, 0x52, 0x75, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73,
+	0x65, 0x12, 0x1d, 0x0a, 0x0a, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x70, 0x6f, 0x72, 0x74, 0x18,
+	0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x50, 0x6f, 0x72, 0x74,
+	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, 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, 0xbb, 0x01, 0x0a, 0x13, 0x43, 0x6c, 0x69, 0x65, 0x6e,
+	0x74, 0x53, 0x74, 0x72, 0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x30,
+	0x0a, 0x09, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x6c, 0x6f, 0x67, 0x18, 0x01, 0x20, 0x01, 0x28,
+	0x0b, 0x32, 0x11, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x41, 0x67, 0x65, 0x6e,
+	0x74, 0x4c, 0x6f, 0x67, 0x48, 0x00, 0x52, 0x08, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x4c, 0x6f, 0x67,
+	0x12, 0x36, 0x0a, 0x0b, 0x61, 0x67, 0x65, 0x6e, 0x74, 0x5f, 0x65, 0x76, 0x65, 0x6e, 0x74, 0x18,
+	0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x13, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e,
+	0x41, 0x67, 0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x48, 0x00, 0x52, 0x0a, 0x61, 0x67,
+	0x65, 0x6e, 0x74, 0x45, 0x76, 0x65, 0x6e, 0x74, 0x12, 0x2f, 0x0a, 0x07, 0x72, 0x75, 0x6e, 0x5f,
+	0x72, 0x65, 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x14, 0x2e, 0x6d, 0x61, 0x6e, 0x61,
+	0x67, 0x65, 0x72, 0x2e, 0x52, 0x75, 0x6e, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x48,
+	0x00, 0x52, 0x06, 0x72, 0x75, 0x6e, 0x52, 0x65, 0x73, 0x42, 0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73,
+	0x73, 0x61, 0x67, 0x65, 0x22, 0xd6, 0x01, 0x0a, 0x13, 0x53, 0x65, 0x72, 0x76, 0x65, 0x72, 0x53,
+	0x74, 0x72, 0x65, 0x61, 0x6d, 0x4d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x12, 0x34, 0x0a, 0x06,
+	0x72, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x18, 0x01, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x1a, 0x2e, 0x6d,
+	0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69,
+	0x6f, 0x6e, 0x52, 0x75, 0x6e, 0x52, 0x65, 0x71, 0x48, 0x00, 0x52, 0x06, 0x72, 0x75, 0x6e, 0x52,
+	0x65, 0x71, 0x12, 0x38, 0x0a, 0x0c, 0x74, 0x65, 0x72, 0x6d, 0x69, 0x6e, 0x61, 0x74, 0x65, 0x52,
+	0x65, 0x71, 0x18, 0x02, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x12, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67,
+	0x65, 0x72, 0x2e, 0x54, 0x65, 0x72, 0x6d, 0x69, 0x6e, 0x61, 0x74, 0x65, 0x48, 0x00, 0x52, 0x0c,
+	0x74, 0x65, 0x72, 0x6d, 0x69, 0x6e, 0x61, 0x74, 0x65, 0x52, 0x65, 0x71, 0x12, 0x44, 0x0a, 0x0f,
+	0x73, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x18,
+	0x03, 0x20, 0x01, 0x28, 0x0b, 0x32, 0x18, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e,
+	0x53, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x48,
+	0x00, 0x52, 0x0f, 0x73, 0x74, 0x6f, 0x70, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69,
+	0x6f, 0x6e, 0x42, 0x09, 0x0a, 0x07, 0x6d, 0x65, 0x73, 0x73, 0x61, 0x67, 0x65, 0x22, 0xb6, 0x02,
 	0x0a, 0x11, 0x43, 0x6f, 0x6d, 0x70, 0x75, 0x74, 0x61, 0x74, 0x69, 0x6f, 0x6e, 0x52, 0x75, 0x6e,
 	0x52, 0x65, 0x71, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52,
 	0x02, 0x69, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28,
@@ -942,40 +1011,42 @@ func file_manager_manager_proto_rawDescGZIP() []byte {
 	return file_manager_manager_proto_rawDescData
 }
 
-var file_manager_manager_proto_msgTypes = make([]protoimpl.MessageInfo, 11)
-var file_manager_manager_proto_goTypes = []interface{}{
+var file_manager_manager_proto_msgTypes = make([]protoimpl.MessageInfo, 12)
+var file_manager_manager_proto_goTypes = []any{
 	(*Terminate)(nil),             // 0: manager.Terminate
-	(*RunResponse)(nil),           // 1: manager.RunResponse
-	(*AgentEvent)(nil),            // 2: manager.AgentEvent
-	(*AgentLog)(nil),              // 3: manager.AgentLog
-	(*ClientStreamMessage)(nil),   // 4: manager.ClientStreamMessage
-	(*ServerStreamMessage)(nil),   // 5: manager.ServerStreamMessage
-	(*ComputationRunReq)(nil),     // 6: manager.ComputationRunReq
-	(*ResultConsumer)(nil),        // 7: manager.ResultConsumer
-	(*Dataset)(nil),               // 8: manager.Dataset
-	(*Algorithm)(nil),             // 9: manager.Algorithm
-	(*AgentConfig)(nil),           // 10: manager.AgentConfig
-	(*timestamppb.Timestamp)(nil), // 11: google.protobuf.Timestamp
+	(*StopComputation)(nil),       // 1: manager.StopComputation
+	(*RunResponse)(nil),           // 2: manager.RunResponse
+	(*AgentEvent)(nil),            // 3: manager.AgentEvent
+	(*AgentLog)(nil),              // 4: manager.AgentLog
+	(*ClientStreamMessage)(nil),   // 5: manager.ClientStreamMessage
+	(*ServerStreamMessage)(nil),   // 6: manager.ServerStreamMessage
+	(*ComputationRunReq)(nil),     // 7: manager.ComputationRunReq
+	(*ResultConsumer)(nil),        // 8: manager.ResultConsumer
+	(*Dataset)(nil),               // 9: manager.Dataset
+	(*Algorithm)(nil),             // 10: manager.Algorithm
+	(*AgentConfig)(nil),           // 11: manager.AgentConfig
+	(*timestamppb.Timestamp)(nil), // 12: google.protobuf.Timestamp
 }
 var file_manager_manager_proto_depIdxs = []int32{
-	11, // 0: manager.AgentEvent.timestamp:type_name -> google.protobuf.Timestamp
-	11, // 1: manager.AgentLog.timestamp:type_name -> google.protobuf.Timestamp
-	3,  // 2: manager.ClientStreamMessage.agent_log:type_name -> manager.AgentLog
-	2,  // 3: manager.ClientStreamMessage.agent_event:type_name -> manager.AgentEvent
-	1,  // 4: manager.ClientStreamMessage.run_res:type_name -> manager.RunResponse
-	6,  // 5: manager.ServerStreamMessage.runReq:type_name -> manager.ComputationRunReq
+	12, // 0: manager.AgentEvent.timestamp:type_name -> google.protobuf.Timestamp
+	12, // 1: manager.AgentLog.timestamp:type_name -> google.protobuf.Timestamp
+	4,  // 2: manager.ClientStreamMessage.agent_log:type_name -> manager.AgentLog
+	3,  // 3: manager.ClientStreamMessage.agent_event:type_name -> manager.AgentEvent
+	2,  // 4: manager.ClientStreamMessage.run_res:type_name -> manager.RunResponse
+	7,  // 5: manager.ServerStreamMessage.runReq:type_name -> manager.ComputationRunReq
 	0,  // 6: manager.ServerStreamMessage.terminateReq:type_name -> manager.Terminate
-	8,  // 7: manager.ComputationRunReq.datasets:type_name -> manager.Dataset
-	9,  // 8: manager.ComputationRunReq.algorithm:type_name -> manager.Algorithm
-	7,  // 9: manager.ComputationRunReq.result_consumers:type_name -> manager.ResultConsumer
-	10, // 10: manager.ComputationRunReq.agent_config:type_name -> manager.AgentConfig
-	4,  // 11: manager.ManagerService.Process:input_type -> manager.ClientStreamMessage
-	5,  // 12: manager.ManagerService.Process:output_type -> manager.ServerStreamMessage
-	12, // [12:13] is the sub-list for method output_type
-	11, // [11:12] is the sub-list for method input_type
-	11, // [11:11] is the sub-list for extension type_name
-	11, // [11:11] is the sub-list for extension extendee
-	0,  // [0:11] is the sub-list for field type_name
+	1,  // 7: manager.ServerStreamMessage.stopComputation:type_name -> manager.StopComputation
+	9,  // 8: manager.ComputationRunReq.datasets:type_name -> manager.Dataset
+	10, // 9: manager.ComputationRunReq.algorithm:type_name -> manager.Algorithm
+	8,  // 10: manager.ComputationRunReq.result_consumers:type_name -> manager.ResultConsumer
+	11, // 11: manager.ComputationRunReq.agent_config:type_name -> manager.AgentConfig
+	5,  // 12: manager.ManagerService.Process:input_type -> manager.ClientStreamMessage
+	6,  // 13: manager.ManagerService.Process:output_type -> manager.ServerStreamMessage
+	13, // [13:14] is the sub-list for method output_type
+	12, // [12:13] is the sub-list for method input_type
+	12, // [12:12] is the sub-list for extension type_name
+	12, // [12:12] is the sub-list for extension extendee
+	0,  // [0:12] is the sub-list for field type_name
 }
 
 func init() { file_manager_manager_proto_init() }
@@ -984,7 +1055,7 @@ func file_manager_manager_proto_init() {
 		return
 	}
 	if !protoimpl.UnsafeEnabled {
-		file_manager_manager_proto_msgTypes[0].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[0].Exporter = func(v any, i int) any {
 			switch v := v.(*Terminate); i {
 			case 0:
 				return &v.state
@@ -996,7 +1067,19 @@ func file_manager_manager_proto_init() {
 				return nil
 			}
 		}
-		file_manager_manager_proto_msgTypes[1].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[1].Exporter = func(v any, i int) any {
+			switch v := v.(*StopComputation); i {
+			case 0:
+				return &v.state
+			case 1:
+				return &v.sizeCache
+			case 2:
+				return &v.unknownFields
+			default:
+				return nil
+			}
+		}
+		file_manager_manager_proto_msgTypes[2].Exporter = func(v any, i int) any {
 			switch v := v.(*RunResponse); i {
 			case 0:
 				return &v.state
@@ -1008,7 +1091,7 @@ func file_manager_manager_proto_init() {
 				return nil
 			}
 		}
-		file_manager_manager_proto_msgTypes[2].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[3].Exporter = func(v any, i int) any {
 			switch v := v.(*AgentEvent); i {
 			case 0:
 				return &v.state
@@ -1020,7 +1103,7 @@ func file_manager_manager_proto_init() {
 				return nil
 			}
 		}
-		file_manager_manager_proto_msgTypes[3].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[4].Exporter = func(v any, i int) any {
 			switch v := v.(*AgentLog); i {
 			case 0:
 				return &v.state
@@ -1032,7 +1115,7 @@ func file_manager_manager_proto_init() {
 				return nil
 			}
 		}
-		file_manager_manager_proto_msgTypes[4].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[5].Exporter = func(v any, i int) any {
 			switch v := v.(*ClientStreamMessage); i {
 			case 0:
 				return &v.state
@@ -1044,7 +1127,7 @@ func file_manager_manager_proto_init() {
 				return nil
 			}
 		}
-		file_manager_manager_proto_msgTypes[5].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[6].Exporter = func(v any, i int) any {
 			switch v := v.(*ServerStreamMessage); i {
 			case 0:
 				return &v.state
@@ -1056,7 +1139,7 @@ func file_manager_manager_proto_init() {
 				return nil
 			}
 		}
-		file_manager_manager_proto_msgTypes[6].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[7].Exporter = func(v any, i int) any {
 			switch v := v.(*ComputationRunReq); i {
 			case 0:
 				return &v.state
@@ -1068,7 +1151,7 @@ func file_manager_manager_proto_init() {
 				return nil
 			}
 		}
-		file_manager_manager_proto_msgTypes[7].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[8].Exporter = func(v any, i int) any {
 			switch v := v.(*ResultConsumer); i {
 			case 0:
 				return &v.state
@@ -1080,7 +1163,7 @@ func file_manager_manager_proto_init() {
 				return nil
 			}
 		}
-		file_manager_manager_proto_msgTypes[8].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[9].Exporter = func(v any, i int) any {
 			switch v := v.(*Dataset); i {
 			case 0:
 				return &v.state
@@ -1092,7 +1175,7 @@ func file_manager_manager_proto_init() {
 				return nil
 			}
 		}
-		file_manager_manager_proto_msgTypes[9].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[10].Exporter = func(v any, i int) any {
 			switch v := v.(*Algorithm); i {
 			case 0:
 				return &v.state
@@ -1104,7 +1187,7 @@ func file_manager_manager_proto_init() {
 				return nil
 			}
 		}
-		file_manager_manager_proto_msgTypes[10].Exporter = func(v interface{}, i int) interface{} {
+		file_manager_manager_proto_msgTypes[11].Exporter = func(v any, i int) any {
 			switch v := v.(*AgentConfig); i {
 			case 0:
 				return &v.state
@@ -1117,14 +1200,15 @@ func file_manager_manager_proto_init() {
 			}
 		}
 	}
-	file_manager_manager_proto_msgTypes[4].OneofWrappers = []interface{}{
+	file_manager_manager_proto_msgTypes[5].OneofWrappers = []any{
 		(*ClientStreamMessage_AgentLog)(nil),
 		(*ClientStreamMessage_AgentEvent)(nil),
 		(*ClientStreamMessage_RunRes)(nil),
 	}
-	file_manager_manager_proto_msgTypes[5].OneofWrappers = []interface{}{
+	file_manager_manager_proto_msgTypes[6].OneofWrappers = []any{
 		(*ServerStreamMessage_RunReq)(nil),
 		(*ServerStreamMessage_TerminateReq)(nil),
+		(*ServerStreamMessage_StopComputation)(nil),
 	}
 	type x struct{}
 	out := protoimpl.TypeBuilder{
@@ -1132,7 +1216,7 @@ func file_manager_manager_proto_init() {
 			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
 			RawDescriptor: file_manager_manager_proto_rawDesc,
 			NumEnums:      0,
-			NumMessages:   11,
+			NumMessages:   12,
 			NumExtensions: 0,
 			NumServices:   1,
 		},

pkg/manager/manager_grpc.pb.go

@@ -3,8 +3,8 @@
 
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
-// - protoc-gen-go-grpc v1.3.0
-// - protoc             v4.25.3
+// - protoc-gen-go-grpc v1.4.0
+// - protoc             v5.27.2
 // source: manager/manager.proto
 
 package manager
@@ -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.32.0 or later.
-const _ = grpc.SupportPackageIsVersion7
+// Requires gRPC-Go v1.62.0 or later.
+const _ = grpc.SupportPackageIsVersion8
 
 const (
 	ManagerService_Process_FullMethodName = "/manager.ManagerService/Process"
@@ -41,11 +41,12 @@ func NewManagerServiceClient(cc grpc.ClientConnInterface) ManagerServiceClient {
 }
 
 func (c *managerServiceClient) Process(ctx context.Context, opts ...grpc.CallOption) (ManagerService_ProcessClient, error) {
-	stream, err := c.cc.NewStream(ctx, &ManagerService_ServiceDesc.Streams[0], ManagerService_Process_FullMethodName, opts...)
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
+	stream, err := c.cc.NewStream(ctx, &ManagerService_ServiceDesc.Streams[0], ManagerService_Process_FullMethodName, cOpts...)
 	if err != nil {
 		return nil, err
 	}
-	x := &managerServiceProcessClient{stream}
+	x := &managerServiceProcessClient{ClientStream: stream}
 	return x, nil
 }
 
@@ -100,7 +101,7 @@ func RegisterManagerServiceServer(s grpc.ServiceRegistrar, srv ManagerServiceSer
 }
 
 func _ManagerService_Process_Handler(srv interface{}, stream grpc.ServerStream) error {
-	return srv.(ManagerServiceServer).Process(&managerServiceProcessServer{stream})
+	return srv.(ManagerServiceServer).Process(&managerServiceProcessServer{ServerStream: stream})
 }
 
 type ManagerService_ProcessServer interface {

pkg/progressbar/progressbar.go

@@ -0,0 +1,237 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+package progressbar
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/ultravioletrs/cocos/agent"
+	"golang.org/x/term"
+)
+
+const (
+	progressBarDots = "... "
+	leftBracket     = "["
+	rightBracket    = "]"
+	head            = ">"
+	body            = "="
+	bodyPadding     = "."
+	bufferSize      = 1024 * 1024
+)
+
+var (
+	_ streamSender = (*algoClientWrapper)(nil)
+	_ streamSender = (*dataClientWrapper)(nil)
+)
+
+type streamSender interface {
+	Send(interface{}) error
+	CloseAndRecv() (interface{}, error)
+}
+
+type algoClientWrapper struct {
+	client *agent.AgentService_AlgoClient
+}
+
+func (a *algoClientWrapper) Send(req interface{}) error {
+	algoReq, ok := req.(*agent.AlgoRequest)
+	if !ok {
+		return fmt.Errorf("expected *AlgoRequest, got %T", req)
+	}
+
+	return (*a.client).Send(algoReq)
+}
+
+func (a *algoClientWrapper) CloseAndRecv() (interface{}, error) {
+	return (*a.client).CloseAndRecv()
+}
+
+type dataClientWrapper struct {
+	client *agent.AgentService_DataClient
+}
+
+func (a *dataClientWrapper) Send(req interface{}) error {
+	dataReq, ok := req.(*agent.DataRequest)
+	if !ok {
+		return fmt.Errorf("expected *DataRequest, got %T", req)
+	}
+
+	return (*a.client).Send(dataReq)
+}
+
+func (a *dataClientWrapper) CloseAndRecv() (interface{}, error) {
+	return (*a.client).CloseAndRecv()
+}
+
+type ProgressBar struct {
+	numberOfBytes           int
+	currentUploadedBytes    int
+	currentUploadPercentage int
+	description             string
+	maxWidth                int
+}
+
+func New() *ProgressBar {
+	return &ProgressBar{}
+}
+
+func (p *ProgressBar) SendAlgorithm(description string, buffer *bytes.Buffer, stream *agent.AgentService_AlgoClient) error {
+	return p.sendData(description, buffer, &algoClientWrapper{client: stream}, func(data []byte) interface{} {
+		return &agent.AlgoRequest{Algorithm: data}
+	})
+}
+
+func (p *ProgressBar) SendData(description string, buffer *bytes.Buffer, stream *agent.AgentService_DataClient) error {
+	return p.sendData(description, buffer, &dataClientWrapper{client: stream}, func(data []byte) interface{} {
+		return &agent.DataRequest{Dataset: data}
+	})
+}
+
+func (p *ProgressBar) sendData(description string, buffer *bytes.Buffer, stream streamSender, createRequest func([]byte) interface{}) error {
+	p.reset(description, buffer.Len())
+
+	buf := make([]byte, bufferSize)
+
+	for {
+		n, err := buffer.Read(buf)
+		if err == io.EOF {
+			if _, err := io.WriteString(os.Stdout, "\n"); err != nil {
+				return err
+			}
+			break
+		}
+		if err != nil {
+			return err
+		}
+
+		p.updateProgress(n)
+
+		if err := stream.Send(createRequest(buf[:n])); err != nil {
+			return err
+		}
+
+		if err := p.renderProgressBar(); err != nil {
+			return err
+		}
+	}
+
+	_, err := stream.CloseAndRecv()
+	return err
+}
+
+func (p *ProgressBar) reset(description string, totalBytes int) {
+	p.currentUploadedBytes = 0
+	p.currentUploadPercentage = 0
+	p.numberOfBytes = totalBytes
+	p.description = description
+}
+
+func (p *ProgressBar) updateProgress(bytesRead int) {
+	if p.currentUploadedBytes < p.numberOfBytes {
+		p.currentUploadedBytes += bytesRead
+		p.currentUploadPercentage = p.currentUploadedBytes * 100 / p.numberOfBytes
+	}
+}
+
+// Progress bar example: Uploading algorithm... 25% [==>   ].
+func (p *ProgressBar) renderProgressBar() error {
+	var builder strings.Builder
+
+	// Get terminal width.
+	width, err := terminalWidth()
+	if err != nil {
+		return err
+	}
+
+	if p.maxWidth < width {
+		p.maxWidth = width
+	}
+
+	if err := p.clearProgressBar(); err != nil {
+		return err
+	}
+
+	// The progress bar starts with the description.
+	if _, err := builder.WriteString(p.description); err != nil {
+		return err
+	}
+
+	// Add dots to progress bar.
+	if _, err := builder.WriteString(progressBarDots); err != nil {
+		return err
+	}
+
+	// Add uploaded percentage.
+	strCurrentUploadPercentage := fmt.Sprintf("%4d%% ", p.currentUploadPercentage)
+	if _, err := builder.WriteString(strCurrentUploadPercentage); err != nil {
+		return err
+	}
+
+	// Add letf bracket and space to progress bar.
+	if _, err := builder.WriteString(leftBracket); err != nil {
+		return err
+	}
+
+	progressWidth := width - builder.Len() - len(rightBracket+" ")
+	numOfCharactersBody := progressWidth * p.currentUploadPercentage / 100
+	if numOfCharactersBody == 0 {
+		numOfCharactersBody = 1
+	}
+
+	numOfCharactersPadding := progressWidth - numOfCharactersBody
+
+	// Add body which represents the percentage.
+	progress := strings.Repeat(body, numOfCharactersBody-1)
+
+	// Add progress to the progress bar.
+	if _, err := builder.WriteString(progress); err != nil {
+		return err
+	}
+
+	// Add head to progress bar.
+	if _, err := builder.WriteString(head); err != nil {
+		return err
+	}
+
+	// Add padding to end of bar.
+	padding := strings.Repeat(bodyPadding, numOfCharactersPadding)
+
+	// Add padding to progress bar.
+	if _, err := builder.WriteString(padding); err != nil {
+		return err
+	}
+
+	// Add right bracket to progress bar.
+	if _, err := builder.WriteString(rightBracket); err != nil {
+		return err
+	}
+
+	// Write progress bar.
+	if _, err := io.WriteString(os.Stdout, builder.String()); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func terminalWidth() (int, error) {
+	width, _, err := term.GetSize(int(os.Stdout.Fd()))
+	if err == nil {
+		return width, nil
+	}
+
+	return 0, err
+}
+
+func (p *ProgressBar) clearProgressBar() error {
+	emptySpace := fmt.Sprintf("\r%s\r", strings.Repeat(" ", p.maxWidth))
+	if _, err := io.WriteString(os.Stdout, emptySpace); err != nil {
+		return err
+	}
+
+	return nil
+}

pkg/sdk/agent.go

@@ -6,28 +6,32 @@ import (
 	"bytes"
 	"context"
 	"crypto"
+	"crypto/ecdsa"
+	"crypto/ed25519"
 	"crypto/rand"
 	"crypto/rsa"
 	"crypto/sha256"
 	"encoding/base64"
-	"io"
+	"errors"
 	"log/slog"
 
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/agent/auth"
+	"github.com/ultravioletrs/cocos/pkg/progressbar"
 	"google.golang.org/grpc/metadata"
 )
 
 type SDK interface {
-	Algo(ctx context.Context, algorithm agent.Algorithm, privKey *rsa.PrivateKey) error
-	Data(ctx context.Context, dataset agent.Dataset, privKey *rsa.PrivateKey) error
-	Result(ctx context.Context, privKey *rsa.PrivateKey) ([]byte, error)
+	Algo(ctx context.Context, algorithm agent.Algorithm, privKey any) error
+	Data(ctx context.Context, dataset agent.Dataset, privKey any) error
+	Result(ctx context.Context, privKey any) ([]byte, error)
 	Attestation(ctx context.Context, reportData [size64]byte) ([]byte, error)
 }
 
 const (
-	size64     = 64
-	bufferSize = 1024 * 1024
+	size64                     = 64
+	algoProgressBarDescription = "Uploading algorithm"
+	dataProgressBarDescription = "Uploading data"
 )
 
 type agentSDK struct {
@@ -42,7 +46,7 @@ func NewAgentSDK(log *slog.Logger, agentClient agent.AgentServiceClient) SDK {
 	}
 }
 
-func (sdk *agentSDK) Algo(ctx context.Context, algorithm agent.Algorithm, privKey *rsa.PrivateKey) error {
+func (sdk *agentSDK) Algo(ctx context.Context, algorithm agent.Algorithm, privKey any) error {
 	md, err := generateMetadata(string(auth.AlgorithmProviderRole), privKey)
 	if err != nil {
 		sdk.logger.Error("Failed to generate metadata")
@@ -57,30 +61,16 @@ func (sdk *agentSDK) Algo(ctx context.Context, algorithm agent.Algorithm, privKe
 	}
 	algoBuffer := bytes.NewBuffer(algorithm.Algorithm)
 
-	buf := make([]byte, bufferSize)
-	for {
-		n, err := algoBuffer.Read(buf)
-		if err == io.EOF {
-			break
-		}
-		if err != nil {
-			return err
-		}
-
-		err = stream.Send(&agent.AlgoRequest{Algorithm: buf[:n]})
-		if err != nil {
-			return err
-		}
-	}
-
-	if _, err := stream.CloseAndRecv(); err != nil {
+	pb := progressbar.New()
+	if err := pb.SendAlgorithm(algoProgressBarDescription, algoBuffer, &stream); err != nil {
+		sdk.logger.Error("Failed to send Algorithm")
 		return err
 	}
 
 	return nil
 }
 
-func (sdk *agentSDK) Data(ctx context.Context, dataset agent.Dataset, privKey *rsa.PrivateKey) error {
+func (sdk *agentSDK) Data(ctx context.Context, dataset agent.Dataset, privKey any) error {
 	md, err := generateMetadata(string(auth.DataProviderRole), privKey)
 	if err != nil {
 		sdk.logger.Error("Failed to generate metadata")
@@ -90,35 +80,21 @@ func (sdk *agentSDK) Data(ctx context.Context, dataset agent.Dataset, privKey *r
 	ctx = metadata.NewOutgoingContext(ctx, md)
 	stream, err := sdk.client.Data(ctx)
 	if err != nil {
-		sdk.logger.Error("Failed to call Algo RPC")
+		sdk.logger.Error("Failed to call Data RPC")
 		return err
 	}
 	dataBuffer := bytes.NewBuffer(dataset.Dataset)
 
-	buf := make([]byte, bufferSize)
-	for {
-		n, err := dataBuffer.Read(buf)
-		if err == io.EOF {
-			break
-		}
-		if err != nil {
-			return err
-		}
-
-		err = stream.Send(&agent.DataRequest{Dataset: buf[:n]})
-		if err != nil {
-			return err
-		}
-	}
-
-	if _, err := stream.CloseAndRecv(); err != nil {
+	pb := progressbar.New()
+	if err := pb.SendData(dataProgressBarDescription, dataBuffer, &stream); err != nil {
+		sdk.logger.Error("Failed to send Data")
 		return err
 	}
 
 	return nil
 }
 
-func (sdk *agentSDK) Result(ctx context.Context, privKey *rsa.PrivateKey) ([]byte, error) {
+func (sdk *agentSDK) Result(ctx context.Context, privKey any) ([]byte, error) {
 	request := &agent.ResultRequest{}
 
 	md, err := generateMetadata(string(auth.ConsumerRole), privKey)
@@ -151,9 +127,20 @@ func (sdk *agentSDK) Attestation(ctx context.Context, reportData [size64]byte) (
 	return response.File, nil
 }
 
-func signData(userID string, privKey *rsa.PrivateKey) ([]byte, error) {
-	hash := sha256.Sum256([]byte(userID))
-	signature, err := rsa.SignPKCS1v15(rand.Reader, privKey, crypto.SHA256, hash[:])
+func signData(userID string, privKey crypto.Signer) ([]byte, error) {
+	var signature []byte
+	var err error
+
+	switch k := privKey.(type) {
+	case ed25519.PrivateKey:
+		signature, err = k.Sign(rand.Reader, []byte(userID), crypto.Hash(0))
+	case *rsa.PrivateKey, *ecdsa.PrivateKey:
+		hash := sha256.Sum256([]byte(userID))
+		signature, err = privKey.Sign(rand.Reader, hash[:], crypto.SHA256)
+	default:
+		return nil, errors.New("unsupported key type")
+	}
+
 	if err != nil {
 		return nil, err
 	}
@@ -161,8 +148,8 @@ func signData(userID string, privKey *rsa.PrivateKey) ([]byte, error) {
 	return signature, nil
 }
 
-func generateMetadata(userID string, privateKey *rsa.PrivateKey) (metadata.MD, error) {
-	signature, err := signData(userID, privateKey)
+func generateMetadata(userID string, privateKey crypto.PrivateKey) (metadata.MD, error) {
+	signature, err := signData(userID, privateKey.(crypto.Signer))
 	if err != nil {
 		return nil, err
 	}

pkg/sdk/agent_test.go

@@ -0,0 +1,435 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package sdk_test
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"crypto/ed25519"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"errors"
+	"os"
+	"testing"
+
+	mglog "github.com/absmach/magistrala/logger"
+	"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"
+	"golang.org/x/crypto/sha3"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/status"
+)
+
+var (
+	algoPath = "../../test/manual/algo/lin_reg.py"
+	dataPath = "../../test/manual/data/iris.csv"
+
+	errInappropriateIoctl = errors.New("inappropriate ioctl for device")
+)
+
+func TestAlgo(t *testing.T) {
+	logger, err := mglog.New(os.Stdout, "info")
+	require.NoError(t, err)
+
+	conn, err := grpc.DialContext(context.Background(), "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
+	if err != nil {
+		t.Fatalf("Failed to dial bufnet: %v", err)
+	}
+	defer conn.Close()
+
+	client := agent.NewAgentServiceClient(conn)
+
+	sdk := sdk.NewAgentSDK(logger, client)
+	algo, err := os.ReadFile(algoPath)
+	require.NoError(t, err)
+
+	algoHash := sha3.Sum256(algo)
+
+	algorithmProviderKey, algorithmProviderPubKey := generateKeys(t, "ed25519")
+
+	algoProvider1Key, algoProvider1PubKey := generateKeys(t, "ed25519")
+
+	algorithm := agent.Algorithm{
+		Algorithm: algo,
+		Hash:      algoHash,
+		UserKey:   algorithmProviderPubKey,
+	}
+
+	cases := []struct {
+		name    string
+		err     error
+		algo    agent.Algorithm
+		userKey any
+	}{
+		{
+			name:    "Test Algo successfully",
+			algo:    algorithm,
+			userKey: algorithmProviderKey,
+			err:     nil,
+		},
+		{
+			name: "hash mismatch",
+			algo: agent.Algorithm{
+				Algorithm: algo,
+				Hash:      sha3.Sum256([]byte("algo")),
+				UserKey:   algoProvider1PubKey,
+			},
+			userKey: algoProvider1Key,
+			err:     errInappropriateIoctl,
+		},
+		{
+			name: "no manifest expected",
+			algo: agent.Algorithm{
+				Algorithm: algo,
+				UserKey:   algorithmProviderPubKey,
+				Hash:      algoHash,
+			},
+			userKey: algorithmProviderKey,
+			err:     errInappropriateIoctl,
+		},
+		{
+			name: "state not ready",
+			algo: agent.Algorithm{
+				Algorithm: algo,
+				UserKey:   algorithmProviderPubKey,
+				Hash:      algoHash,
+			},
+			userKey: algorithmProviderKey,
+			err:     errInappropriateIoctl,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			svcCall := svc.On("Algo", mock.Anything, mock.Anything).Return(tc.err)
+			err = sdk.Algo(context.Background(), tc.algo, tc.userKey)
+
+			st, _ := status.FromError(err)
+
+			if tc.err != nil {
+				if st.Message() != tc.err.Error() {
+					t.Errorf("%s : Expected error message %q, but got %q", tc.name, tc.err.Error(), st.Message())
+				}
+			}
+
+			svcCall.Unset()
+		})
+	}
+}
+
+func TestData(t *testing.T) {
+	logger, err := mglog.New(os.Stdout, "info")
+	require.NoError(t, err)
+
+	conn, err := grpc.DialContext(context.Background(), "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
+	if err != nil {
+		t.Fatalf("Failed to dial bufnet: %v", err)
+	}
+	defer conn.Close()
+
+	client := agent.NewAgentServiceClient(conn)
+
+	sdk := sdk.NewAgentSDK(logger, client)
+
+	data, err := os.ReadFile(dataPath)
+	require.NoError(t, err)
+
+	dataHash := sha3.Sum256(data)
+
+	dataProviderKey, dataProviderPubKey := generateKeys(t, "ecdsa")
+
+	dataProvider1Key, dataProvider1PubKey := generateKeys(t, "ed25519")
+
+	dataset := agent.Dataset{
+		Hash:    dataHash,
+		Dataset: data,
+		UserKey: dataProviderPubKey,
+	}
+
+	cases := []struct {
+		name    string
+		data    agent.Dataset
+		userKey any
+		svcErr  error
+	}{
+		{
+			name:    "Test data successfully",
+			data:    dataset,
+			userKey: dataProviderKey,
+		},
+		{
+			name: "undeclared dataset",
+			data: agent.Dataset{
+				Dataset: data,
+				UserKey: dataProvider1PubKey,
+				Hash:    dataHash,
+			},
+			userKey: dataProvider1Key,
+			svcErr:  errInappropriateIoctl,
+		},
+		{
+			name: "hash mismatch",
+			data: agent.Dataset{
+				Dataset: data,
+				UserKey: dataProvider1PubKey,
+				Hash:    dataHash,
+			},
+			userKey: dataProvider1Key,
+			svcErr:  errInappropriateIoctl,
+		},
+		{
+			name: "all manifest items received",
+			data: agent.Dataset{
+				Dataset: data,
+				UserKey: dataProvider1PubKey,
+				Hash:    dataHash,
+			},
+			userKey: dataProvider1Key,
+			svcErr:  errInappropriateIoctl,
+		},
+		{
+			name: "missing dataset file",
+			data: agent.Dataset{
+				UserKey: dataProvider1PubKey,
+				Hash:    dataHash,
+			},
+			userKey: dataProvider1Key,
+			svcErr:  errors.New("dataset CSV file is required"),
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			dataCall := svc.On("Data", mock.Anything, mock.Anything).Return(tc.svcErr)
+
+			err = sdk.Data(context.Background(), tc.data, tc.userKey)
+
+			st, _ := status.FromError(err)
+
+			if tc.svcErr != nil {
+				if st.Message() != tc.svcErr.Error() {
+					t.Errorf("%s: Expected error message %q, but got %q", tc.name, tc.svcErr.Error(), st.Message())
+				}
+			}
+
+			dataCall.Unset()
+		})
+	}
+}
+
+func TestResult(t *testing.T) {
+	logger, err := mglog.New(os.Stdout, "info")
+	require.NoError(t, err)
+
+	conn, err := grpc.DialContext(context.Background(), "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
+	if err != nil {
+		t.Fatalf("Failed to dial bufnet: %v", err)
+	}
+	defer conn.Close()
+
+	client := agent.NewAgentServiceClient(conn)
+
+	sdk := sdk.NewAgentSDK(logger, client)
+
+	resultConsumerKey, _ := generateKeys(t, "ecdsa")
+	resultConsumer1Key, _ := generateKeys(t, "ed25519")
+
+	response := &agent.ResultResponse{
+		File: []byte{
+			0x01, 0x02, 0x03, 0x04,
+			0x05, 0x06, 0x07, 0x08,
+		},
+	}
+
+	cases := []struct {
+		name     string
+		userKey  any
+		response *agent.ResultResponse
+		svcRes   []byte
+		err      error
+	}{
+		{
+			name:     "Test result successfully",
+			userKey:  resultConsumerKey,
+			response: response,
+			svcRes:   response.File,
+			err:      nil,
+		},
+		{
+			name:     "Test result successfully with ed25519 key type",
+			userKey:  resultConsumer1Key,
+			response: response,
+			svcRes:   response.File,
+			err:      nil,
+		},
+		{
+			name:    "Results not ready",
+			userKey: resultConsumer1Key,
+			response: &agent.ResultResponse{
+				File: []byte(nil),
+			},
+			svcRes: nil,
+			err:    agent.ErrResultsNotReady,
+		},
+		{
+			name:    "All manifest items received",
+			userKey: resultConsumer1Key,
+			response: &agent.ResultResponse{
+				File: []byte(nil),
+			},
+			svcRes: nil,
+			err:    agent.ErrAllManifestItemsReceived,
+		},
+		{
+			name:    "Undeclared consumer",
+			userKey: resultConsumer1Key,
+			response: &agent.ResultResponse{
+				File: []byte(nil),
+			},
+			svcRes: nil,
+			err:    agent.ErrUndeclaredConsumer,
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			svcCall := svc.On("Result", mock.Anything, mock.Anything).Return(tc.svcRes, tc.err)
+			res, err := sdk.Result(context.Background(), tc.userKey)
+
+			st, ok := status.FromError(err)
+			if !ok {
+				t.Fatalf("Expected gRPC status error, but got: %v", err)
+			}
+
+			if tc.err != nil {
+				if st.Message() != tc.err.Error() {
+					t.Errorf("%s: Expected error message %q, but got %q", tc.name, tc.err.Error(), st.Message())
+				}
+			}
+			assert.Equal(t, tc.response.File, res, tc.name)
+
+			svcCall.Unset()
+		})
+	}
+}
+
+func TestAttestation(t *testing.T) {
+	logger, err := mglog.New(os.Stdout, "info")
+	require.NoError(t, err)
+
+	resultConsumerKey, _ := generateKeys(t, "rsa")
+	resultConsumer1Key, _ := generateKeys(t, "ed25519")
+
+	reportData := make([]byte, 64)
+	report := []byte{
+		0x01, 0x02, 0x03, 0x04,
+		0x05, 0x06, 0x07, 0x08,
+	}
+
+	conn, err := grpc.DialContext(context.Background(), "bufnet", grpc.WithContextDialer(bufDialer), grpc.WithInsecure())
+	if err != nil {
+		t.Fatalf("Failed to dial bufnet: %v", err)
+	}
+	defer conn.Close()
+
+	client := agent.NewAgentServiceClient(conn)
+
+	sdk := sdk.NewAgentSDK(logger, client)
+
+	_, err = rand.Read(reportData)
+	require.NoError(t, err)
+
+	cases := []struct {
+		name       string
+		userKey    any
+		reportData [agent.ReportDataSize]byte
+		response   *agent.AttestationResponse
+		svcRes     []byte
+		err        error
+	}{
+		{
+			name:       "fetch attestation report successfully",
+			userKey:    resultConsumerKey,
+			reportData: [agent.ReportDataSize]byte(reportData),
+			response: &agent.AttestationResponse{
+				File: report,
+			},
+			svcRes: report,
+			err:    nil,
+		},
+		{
+			name:       "fetch attestation report with different key type",
+			userKey:    resultConsumer1Key,
+			reportData: [agent.ReportDataSize]byte(reportData),
+			response: &agent.AttestationResponse{
+				File: report,
+			},
+			svcRes: report,
+			err:    nil,
+		},
+		{
+			name:       "failed to fetch attestation report",
+			userKey:    resultConsumerKey,
+			reportData: [agent.ReportDataSize]byte(reportData),
+			response: &agent.AttestationResponse{
+				File: nil,
+			},
+			err: nil,
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			svcCall := svc.On("Attestation", mock.Anything, mock.Anything).Return(tc.svcRes, tc.err)
+
+			res, err := sdk.Attestation(context.Background(), tc.reportData)
+
+			st, ok := status.FromError(err)
+			if !ok {
+				t.Fatalf("Expected gRPC status error, but got: %v", err)
+			}
+
+			if tc.err != nil {
+				if st.Message() != tc.err.Error() {
+					t.Errorf("%s: Expected error message %q, but got %q", tc.name, tc.err.Error(), st.Message())
+				}
+			}
+
+			assert.Equal(t, tc.response.File, res, tc.name)
+
+			svcCall.Unset()
+		})
+	}
+}
+
+func generateKeys(t *testing.T, keyType string) (priv any, pub []byte) {
+	switch keyType {
+	case "ecdsa":
+		privEcdsaKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+		require.NoError(t, err)
+
+		pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privEcdsaKey.PublicKey)
+		require.NoError(t, err)
+		return privEcdsaKey, pubKeyBytes
+
+	case "ed25519":
+		pubEd25519Key, privEd25519Key, err := ed25519.GenerateKey(rand.Reader)
+		require.NoError(t, err)
+
+		pubKey, err := x509.MarshalPKIXPublicKey(pubEd25519Key)
+		require.NoError(t, err)
+		return privEd25519Key, pubKey
+
+	default:
+		privKey, err := rsa.GenerateKey(rand.Reader, 2048)
+		require.NoError(t, err)
+
+		pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privKey.PublicKey)
+		require.NoError(t, err)
+		return privKey, pubKeyBytes
+	}
+}

pkg/sdk/setup_test.go

@@ -0,0 +1,48 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package sdk_test
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"os"
+	"testing"
+
+	"github.com/ultravioletrs/cocos/agent"
+	agentgrpc "github.com/ultravioletrs/cocos/agent/api/grpc"
+	"github.com/ultravioletrs/cocos/agent/mocks"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/test/bufconn"
+)
+
+const bufSize = 1024 * 1024
+
+var (
+	lis *bufconn.Listener
+	svc = &mocks.Service{}
+)
+
+func TestMain(m *testing.M) {
+	lis = bufconn.Listen(bufSize)
+	s := grpc.NewServer()
+
+	agent.RegisterAgentServiceServer(s, agentgrpc.NewServer(svc))
+
+	go func() {
+		if err := s.Serve(lis); err != nil {
+			fmt.Println("Server exited with error:", err)
+		}
+	}()
+
+	code := m.Run()
+
+	s.Stop()
+	lis.Close()
+	os.Exit(code)
+}
+
+func bufDialer(context.Context, string) (net.Conn, error) {
+	return lis.Dial()
+}

scripts/backend_info/Cargo.toml

@@ -1,5 +1,5 @@
 [package]
-name = "platform_info"
+name = "backend_info"
 version = "0.1.0"
 edition = "2021"
 

scripts/backend_info/Makefile

@@ -1,7 +1,7 @@
 CARGO = cargo
 TARGET = target
 BUILD_DIR = $(TARGET)/release
-BIN_NAME = platform_info
+BIN_NAME = backend_info
 
 all: build
 

scripts/backend_info/README.md

@@ -1,4 +1,4 @@
-# Rust project for fetching platform info
+# Rust project for fetching backend info
 This rust project fetches information from the host system needed for validation of the attestation report. It outputs a JSON file that contains the said information.
 The JSON file is in a format that can be used with the [go-sev-guest](https://github.com/google/go-sev-guest) library.
 
@@ -6,7 +6,7 @@ The JSON file is in a format that can be used with the [go-sev-guest](https://gi
 Clone `cocos` repository:
 ```bash
 git clone git@github.com:ultravioletrs/cocos.git
-cd ./cocos/scripts/platform_info 
+cd ./cocos/scripts/backend_info 
 make
 ```
 
@@ -15,5 +15,5 @@ Then run the binary. Keep in mind that you have to specify the policy of the Gue
 cd ./target/releas
 
 # Run with option --policy (policy is 64 bit number) 
-./platform_info --policy 196608
+./backend_info --policy 196608
 ```

scripts/backend_info/src/main.rs

@@ -6,7 +6,7 @@ use sysinfo::System;
 use regex::Regex;
 use sev::firmware::host::*;
 
-const PLATFORM_INFO_JSON : &str = "platform_info.json";
+const BACKEND_INFO_JSON : &str = "backend_info.json";
 
 #[derive(Serialize)]
 struct Vmpl {
@@ -70,7 +70,7 @@ fn get_uint64_from_tcb(tcb_version : &TcbVersion) -> u64 {
 }
 
 fn main() {
-    let matches = Command::new("Platform info")
+    let matches = Command::new("Backend info")
         .about("Processes command line options and outputs a JSON file for Attestation verification")
         .arg(Arg::new("policy")
             .long("policy")
@@ -130,8 +130,8 @@ fn main() {
     };
 
     let json = serde_json::to_string_pretty(&computation).expect("Failed to serialize to JSON");
-    let mut file = File::create(PLATFORM_INFO_JSON).expect("Failed to create file");
+    let mut file = File::create(BACKEND_INFO_JSON).expect("Failed to create file");
     file.write_all(json.as_bytes()).expect("Failed to write to file");
 
-    println!("Computation JSON has been written to {}", PLATFORM_INFO_JSON);
+    println!("Computation JSON has been written to {}", BACKEND_INFO_JSON);
 }

test/computations/main.go

@@ -12,7 +12,7 @@ import (
 	"strconv"
 
 	mglog "github.com/absmach/magistrala/logger"
-	"github.com/ultravioletrs/cocos/internal/env"
+	"github.com/caarlos0/env/v11"
 	"github.com/ultravioletrs/cocos/internal/server"
 	grpcserver "github.com/ultravioletrs/cocos/internal/server/grpc"
 	managergrpc "github.com/ultravioletrs/cocos/manager/api/grpc"
@@ -122,7 +122,7 @@ func main() {
 		manager.RegisterManagerServiceServer(srv, managergrpc.NewServer(incomingChan, &svc{logger: logger}))
 	}
 	grpcServerConfig := server.Config{Port: defaultPort}
-	if err := env.Parse(&grpcServerConfig, env.Options{}); err != nil {
+	if err := env.ParseWithOptions(&grpcServerConfig, env.Options{}); err != nil {
 		logger.Error(fmt.Sprintf("failed to load %s gRPC client configuration : %s", svcName, err))
 		return
 	}

test/manual/README.md

@@ -30,24 +30,38 @@ LINE="earlyprintk=serial console=ttyS0"
 sev-snp-measure --mode snp --vcpus 4 --vcpu-type EPYC-v4 --ovmf $OVMF_CODE --kernel $KERNEL --initrd $INITRD --append "$LINE" --output-format base64
 ```
 
-```sh
+To speed up the verification process of attested TLS, download the ARK and ASK certificates using the CLI tool. The CLI tool will download the certificates under your home directory in the `.cocos` directory.
+```bash
+go run cmd/cli/main.go ca-bundle <path/to/platfrom_info.json>
+```
+
+In the following text, we can see an example of how the CLI tool is used.
+```bash
 export AGENT_GRPC_URL=localhost:7002
 
-# For attested TLS, the CLI should also be aware of the VM measurement. To 
-# add the measurement to the .json file that contains the information about 
-# the platform, run CLI with the measurement in base64 format and the path 
-# of the platform_info.json file.:
-go run cmd/cli/main.go measurement '<measurement>' '<platform_info.json>'
-
-# The platform_info.json file can be generated using Rust by running:
-cd scripts/platform_info
+# For attested TLS, the CLI needs a file containing the necessary information 
+# about the SEV-SNP capable backend. This information will be used to verify 
+# the attestation report received from the agent.
+# The backend_info.json file can be generated using Rust by running:
+cd scripts/backend_info
 make
-sudo ./target/release/platform_info --policy 196608 # Default value of the policy should be 196608
-# The output file platform_info.json will be generated in the directory from which the executable has been called.
+sudo ./target/release/backend_info --policy 196608 # Default value of the policy should be 196608
+# The output file backend_info.json will be generated in the directory from which the executable has been called.
 cd ../..
 
-# For attested TLS, also define the path to the platform_info.json that contains reference values for the fields of the attestation report
-export AGENT_GRPC_MANIFEST=./scripts/platform_info/platform_info.json
+# The CLI should also be aware of the VM measurement. To add the measurement 
+# to the .json file that contains the information about the platform, run CLI 
+# with the measurement in base64 format and the path of the backend_info.json file.:
+go run cmd/cli/main.go backend measurement '<measurement>' '<backend_info.json>'
+
+# If the VM is booted with the QEMU host data option, the CLI should also know 
+# the host data information. To add the host data to the .json file that contains 
+# the information about the platform, run CLI with the host data in base64 format 
+# and the path of the backend_info.json file.:
+go run cmd/cli/main.go backend measurement '<host-data>' '<backend_info.json>'
+
+# For attested TLS, also define the path to the backend_info.json that contains reference values for the fields of the attestation report
+export AGENT_GRPC_MANIFEST=./scripts/backend_info/backend_info.json
 export AGENT_GRPC_ATTESTED_TLS=true
 
 # Retieve Attestation

test/manual/agent-config/main.go

@@ -16,6 +16,7 @@ import (
 	"github.com/mdlayher/vsock"
 	"github.com/ultravioletrs/cocos/agent"
 	"github.com/ultravioletrs/cocos/manager"
+	"github.com/ultravioletrs/cocos/manager/qemu"
 	pkgmanager "github.com/ultravioletrs/cocos/pkg/manager"
 	"golang.org/x/crypto/sha3"
 	"google.golang.org/protobuf/proto"
@@ -89,7 +90,7 @@ func main() {
 }
 
 func SendAgentConfig(cid uint32, ac agent.Computation) error {
-	conn, err := vsock.Dial(cid, manager.VsockConfigPort, nil)
+	conn, err := vsock.Dial(cid, qemu.VsockConfigPort, nil)
 	if err != nil {
 		return err
 	}

test/manual/algo/covid_19/README.md

@@ -1,19 +0,0 @@
-## Algorihm covid19.py
-
-To download dataset first you have to sign in Keggle. https://www.kaggle.com/
-Then follow this documentation. https://github.com/Kaggle/kaggle-api
-
-Go to settings and download API key. Copy it into keggle folder.
-Run divide_save_data.py
-You should have now 3 datasets, named Hospital_1, Hospital_2, Hospital_3
-
-`covid19.py` trains the the model and produces model.
-run :
-`python covid19.py Hospital_1 Hospital_2 Hospital_3 --model model.pth`
-in this example dataset is Hospital_1, Hospital_2, and Hospital_3 and it produces model.pth
-
-`covid.19.py` produced model, and now you should be able to use that model in `predict.py`
-run :
-`python predict.py --model model.pth --image Hospital_2/COVID/images/COVID-3.png`
-in this example you have trained model (model.pth) that runs tests on images (this example is image from Hospital_2 in COVID/images/COVID-3.png)
-

test/manual/algo/covid_19/covid_19.py

@@ -1,190 +0,0 @@
-import os
-import shutil
-import random
-import torch
-import torchvision
-import numpy as np
-import matplotlib.pyplot as plt
-from PIL import Image
-import argparse
-import zipfile
-import socket
-import sys
-import ssl
-
-def prepare_test_set(root_dir, class_names):
-    test_dir = os.path.join(root_dir, 'test')
-    if os.path.isdir(test_dir):
-        shutil.rmtree(test_dir)
-    
-    os.mkdir(test_dir)
-    for class_name in class_names:
-        os.mkdir(os.path.join(test_dir, class_name))
-    
-    for class_name in class_names:
-        images = [x for x in os.listdir(os.path.join(root_dir, class_name, "images")) if x.lower().endswith('png')]
-        selected_images = random.sample(images, 30)
-        for image in selected_images:
-            source_path = os.path.join(root_dir, class_name, "images", image)
-            target_path = os.path.join(test_dir, class_name, image)
-            shutil.copy(source_path, target_path)
-    
-    return test_dir
-
-class ChestXRayDataset(torch.utils.data.Dataset):
-    def __init__(self, image_dirs, transform):
-        def get_images(class_name):
-            images = []
-            for dir_path in image_dirs[class_name]:
-                images += [os.path.join(dir_path, x) for x in os.listdir(dir_path) if x.lower().endswith('png')]
-            print(f'Found {len(images)} {class_name} examples')
-            return images
-
-        self.images = {}
-        self.class_names = ['Normal', 'Viral Pneumonia', 'COVID']
-        for class_name in self.class_names:
-            self.images[class_name] = get_images(class_name)
-        self.transform = transform
-
-    def __len__(self):
-        return sum([len(self.images[class_name]) for class_name in self.class_names])
-
-    def __getitem__(self, index):
-        class_name = random.choice(self.class_names)
-        index = index % len(self.images[class_name])
-        image_path = self.images[class_name][index]
-        image = Image.open(image_path).convert('RGB')
-        return self.transform(image), self.class_names.index(class_name)
-
-def train_and_evaluate_model(dl_train, dl_test, class_names, model_file_name, epochs=1):
-    resnet18 = torchvision.models.resnet18(pretrained=True)
-    resnet18.fc = torch.nn.Linear(in_features=512, out_features=len(class_names))
-    loss_fn = torch.nn.CrossEntropyLoss()
-    optimizer = torch.optim.Adam(resnet18.parameters(), lr=3e-5)
-    print(resnet18)
-
-    print('Starting training..')
-    for e in range(0, epochs):
-        print('=' * 20)
-        print(f'Starting epoch {e + 1}/{epochs}')
-        print('=' * 20)
-        train_loss = 0.
-        val_loss = 0.
-        resnet18.train()  # set model to training phase
-        for train_step, (images, labels) in enumerate(dl_train):
-            optimizer.zero_grad()
-            outputs = resnet18(images)
-            loss = loss_fn(outputs, labels)
-            loss.backward()
-            optimizer.step()
-            train_loss += loss.item()
-            if train_step % 20 == 0:
-                print('Evaluating at step', train_step)
-                accuracy = 0
-                resnet18.eval()  # set model to eval phase
-                for val_step, (images, labels) in enumerate(dl_test):
-                    outputs = resnet18(images)
-                    loss = loss_fn(outputs, labels)
-                    val_loss += loss.item()
-                    _, preds = torch.max(outputs, 1)
-                    accuracy += sum((preds == labels).numpy())
-                val_loss /= (val_step + 1)
-                accuracy = accuracy / len(dl_test.dataset)
-                print(f'Validation Loss: {val_loss:.4f}, Accuracy: {accuracy:.4f}')
-                #show_preds()
-                resnet18.train()
-                if accuracy >= 0.95:
-                    print('Performance condition satisfied, stopping..')
-                    torch.save(resnet18.state_dict(), model_file_name)
-                    return
-        train_loss /= (train_step + 1)
-        print(f'Training Loss: {train_loss:.4f}')
-    print('Training complete..')
-    torch.save(resnet18.state_dict(), model_file_name)
-
-def main():
-    hospitals_zip = []
-    for i, arg in enumerate(sys.argv[2:]):
-        hospitals_zip.append(arg)
-
-    model_file_name = "model.pth"
-
-    class_names = ['Normal', 'Viral Pneumonia', 'COVID']
-
-    # Combine datasets from multiple hospitals
-    train_transform = torchvision.transforms.Compose([
-        torchvision.transforms.Resize(size=(224, 224)),
-        torchvision.transforms.RandomHorizontalFlip(),
-        torchvision.transforms.ToTensor(),
-        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])
-    test_transform = torchvision.transforms.Compose([
-        torchvision.transforms.Resize(size=(224, 224)),
-        torchvision.transforms.ToTensor(),
-        torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ])
-
-    hospital_path = '/tmp/work'
-    if not os.path.isdir(hospital_path):
-        os.mkdir(hospital_path, mode = 0o777)
-    
-    for hospital in hospitals_zip:
-        with zipfile.ZipFile(hospital, 'r') as zip_ref:
-            zip_ref.extractall(hospital_path)
-
-    hospitals = []
-    for hospital in os.listdir(hospital_path):
-        hospitals.append(os.path.join(hospital_path, hospital))    
-
-    # Prepare combined training dataset
-    train_dirs = {class_name: [] for class_name in class_names}
-    for hospital in hospitals:
-        for class_name in class_names:
-            train_dirs[class_name].append(os.path.join(hospital, class_name, 'images'))
-
-    train_image_dirs = {class_name: train_dirs[class_name] for class_name in class_names}
-    train_dataset = ChestXRayDataset(train_image_dirs, train_transform)
-
-    print(f'Total number of training images: {len(train_dataset)}')
-
-    # Prepare test dataset
-    test_dirs = {class_name: [] for class_name in class_names}
-    for hospital in hospitals:
-        test_dir = prepare_test_set(hospital, class_names)
-        for class_name in class_names:
-            test_dirs[class_name].append(os.path.join(test_dir, class_name))
-
-    test_image_dirs = {class_name: test_dirs[class_name] for class_name in class_names}
-    test_dataset = ChestXRayDataset(test_image_dirs, test_transform)
-    
-    batch_size = 6
-    dl_train = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-    dl_test = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
-    
-    print('Number of training batches:', len(dl_train))
-    print('Number of test batches:', len(dl_test))
-
-    train_and_evaluate_model(dl_train, dl_test, class_names, model_file_name, epochs=1)
-
-    # Define the path for the Unix domain socket
-    socket_path = sys.argv[1]
-
-    # Create a Unix domain socket client
-    client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-
-    try:
-        # Connect to the server
-        client.connect(socket_path)
-
-        # Send the serialized model over the socket
-        with open(model_file_name, 'rb') as f:
-            data = f.read()
-            client.sendall(data)
-
-    finally:
-        # Close the socket
-        client.close()
-
-if __name__ == '__main__':
-    ssl._create_default_https_context = ssl._create_unverified_context
-    main()

test/manual/algo/covid_19/divide_save_data.py

@@ -1,190 +0,0 @@
-import os
-import shutil
-import random
-import torch
-import torchvision
-import numpy as np
-import matplotlib.pyplot as plt
-from PIL import Image
-import argparse
-import zipfile
-import socket
-import sys
-
-def prepare_test_set(root_dir, class_names):
-    test_dir = os.path.join(root_dir, 'test')
-    if os.path.isdir(test_dir):
-        shutil.rmtree(test_dir)
-    
-    os.mkdir(test_dir)
-    for class_name in class_names:
-        os.mkdir(os.path.join(test_dir, class_name))
-    
-    for class_name in class_names:
-        images = [x for x in os.listdir(os.path.join(root_dir, class_name, "images")) if x.lower().endswith('png')]
-        selected_images = random.sample(images, 30)
-        for image in selected_images:
-            source_path = os.path.join(root_dir, class_name, "images", image)
-            target_path = os.path.join(test_dir, class_name, image)
-            shutil.copy(source_path, target_path)
-    
-    return test_dir
-
-class ChestXRayDataset(torch.utils.data.Dataset):
-    def __init__(self, image_dirs, transform):
-        def get_images(class_name):
-            images = []
-            for dir_path in image_dirs[class_name]:
-                images += [os.path.join(dir_path, x) for x in os.listdir(dir_path) if x.lower().endswith('png')]
-            print(f'Found {len(images)} {class_name} examples')
-            return images
-
-        self.images = {}
-        self.class_names = ['Normal', 'Viral Pneumonia', 'COVID']
-        for class_name in self.class_names:
-            self.images[class_name] = get_images(class_name)
-        self.transform = transform
-
-    def __len__(self):
-        return sum([len(self.images[class_name]) for class_name in self.class_names])
-
-    def __getitem__(self, index):
-        class_name = random.choice(self.class_names)
-        index = index % len(self.images[class_name])
-        image_path = self.images[class_name][index]
-        image = Image.open(image_path).convert('RGB')
-        return self.transform(image), self.class_names.index(class_name)
-
-def train_and_evaluate_model(dl_train, dl_test, class_names, model_file_name, epochs=1):
-    resnet18 = torchvision.models.resnet18(pretrained=True)
-    resnet18.fc = torch.nn.Linear(in_features=512, out_features=len(class_names))
-    loss_fn = torch.nn.CrossEntropyLoss()
-    optimizer = torch.optim.Adam(resnet18.parameters(), lr=3e-5)
-    print(resnet18)
-
-    print('Starting training..')
-    for e in range(0, epochs):
-        print('=' * 20)
-        print(f'Starting epoch {e + 1}/{epochs}')
-        print('=' * 20)
-        train_loss = 0.
-        val_loss = 0.
-        resnet18.train()  # set model to training phase
-        for train_step, (images, labels) in enumerate(dl_train):
-            optimizer.zero_grad()
-            outputs = resnet18(images)
-            loss = loss_fn(outputs, labels)
-            loss.backward()
-            optimizer.step()
-            train_loss += loss.item()
-            if train_step % 20 == 0:
-                print('Evaluating at step', train_step)
-                accuracy = 0
-                resnet18.eval()  # set model to eval phase
-                for val_step, (images, labels) in enumerate(dl_test):
-                    outputs = resnet18(images)
-                    loss = loss_fn(outputs, labels)
-                    val_loss += loss.item()
-                    _, preds = torch.max(outputs, 1)
-                    accuracy += sum((preds == labels).numpy())
-                val_loss /= (val_step + 1)
-                accuracy = accuracy / len(dl_test.dataset)
-                print(f'Validation Loss: {val_loss:.4f}, Accuracy: {accuracy:.4f}')
-                #show_preds()
-                resnet18.train()
-                if accuracy >= 0.95:
-                    print('Performance condition satisfied, stopping..')
-                    torch.save(resnet18.state_dict(), model_file_name)
-                    return
-        train_loss /= (train_step + 1)
-        print(f'Training Loss: {train_loss:.4f}')
-    print('Training complete..')
-    torch.save(resnet18.state_dict(), model_file_name)
-
-def main():
-    # parser = argparse.ArgumentParser(description='Process hospital datasets and save the model.')
-    # parser.add_argument('hospitals', metavar='-H', type=str, nargs='+',
-    #                     help='paths to hospital datasets')
-    # parser.add_argument('--model', type=str, required=True,
-    #                     help='name of the output model file')
-    
-    # args = parser.parse_args()
-    hospitals_zip = []
-    for i, arg in enumerate(sys.argv[2:]):
-        hospitals_zip.append(arg)
-
-    # hospitals_zip = args.hospitals
-    model_file_name = "model.pth"
-
-    class_names = ['Normal', 'Viral Pneumonia', 'COVID']
-
-    # Combine datasets from multiple hospitals
-    train_transform = torchvision.transforms.Compose([
-        torchvision.transforms.Resize(size=(224, 224)),
-        torchvision.transforms.RandomHorizontalFlip(),
-        torchvision.transforms.ToTensor(),
-        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])
-    test_transform = torchvision.transforms.Compose([
-        torchvision.transforms.Resize(size=(224, 224)),
-        torchvision.transforms.ToTensor(),
-        torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-    ])
-
-    hospitals = []
-    for hospital in hospitals_zip:
-        hospitals.append(os.path.basename(os.path.splitext(hospital)[0]))
-        with zipfile.ZipFile(hospital, 'r') as zip_ref:
-            zip_ref.extractall('./')
-
-    # Prepare combined training dataset
-    train_dirs = {class_name: [] for class_name in class_names}
-    for hospital in hospitals:
-        for class_name in class_names:
-            train_dirs[class_name].append(os.path.join(hospital, class_name, 'images'))
-
-    train_image_dirs = {class_name: train_dirs[class_name] for class_name in class_names}
-    train_dataset = ChestXRayDataset(train_image_dirs, train_transform)
-
-    print(f'Total number of training images: {len(train_dataset)}')
-
-    # Prepare test dataset
-    test_dirs = {class_name: [] for class_name in class_names}
-    for hospital in hospitals:
-        test_dir = prepare_test_set(hospital, class_names)
-        for class_name in class_names:
-            test_dirs[class_name].append(os.path.join(test_dir, class_name))
-
-    test_image_dirs = {class_name: test_dirs[class_name] for class_name in class_names}
-    test_dataset = ChestXRayDataset(test_image_dirs, test_transform)
-    
-    batch_size = 6
-    dl_train = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-    dl_test = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
-    
-    print('Number of training batches:', len(dl_train))
-    print('Number of test batches:', len(dl_test))
-
-    train_and_evaluate_model(dl_train, dl_test, class_names, model_file_name, epochs=1)
-
-    # Define the path for the Unix domain socket
-    socket_path = sys.argv[1]
-
-    # Create a Unix domain socket client
-    client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
-
-    try:
-        # Connect to the server
-        client.connect(socket_path)
-
-        # Send the serialized model over the socket
-        with open(model_file_name, 'rb') as f:
-            data = f.read()
-            client.sendall(data)
-
-    finally:
-        # Close the socket
-        client.close()
-
-if __name__ == '__main__':
-    main()

test/manual/algo/covid_19/predict.py

@@ -1,59 +0,0 @@
-import torch
-import torchvision
-import numpy as np
-import matplotlib.pyplot as plt
-from PIL import Image
-import argparse
-
-def load_model(model_path, class_names):
-    model = torchvision.models.resnet18(weights=None)
-    model.fc = torch.nn.Linear(in_features=512, out_features=len(class_names))
-    model.load_state_dict(torch.load(model_path))
-    model.eval()  # Set the model to evaluation mode
-    return model
-
-def preprocess_image(image_path):
-    transform = torchvision.transforms.Compose([
-        torchvision.transforms.Resize(size=(224, 224)),
-        torchvision.transforms.ToTensor(),
-        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])
-    image = Image.open(image_path).convert('RGB')
-    image = transform(image)
-    image = image.unsqueeze(0)  # Add batch dimension
-    return image
-
-def predict(model, image_path, class_names):
-    image = preprocess_image(image_path)
-    outputs = model(image)
-    _, preds = torch.max(outputs, 1)
-    predicted_class = class_names[preds.item()]
-    return predicted_class
-
-def show_image_with_prediction(image_path, predicted_class):
-    image = Image.open(image_path)
-    plt.imshow(image)
-    plt.title(f'Predicted: {predicted_class}')
-    plt.axis('off')
-    plt.show()
-
-def main():
-    parser = argparse.ArgumentParser(description='Load model and make a prediction.')
-    parser.add_argument('--model', type=str, required=True,
-                        help='path to the model file')
-    parser.add_argument('--image', type=str, required=True,
-                        help='path to the image for prediction')
-    
-    args = parser.parse_args()
-    model_path = args.model
-    image_path = args.image
-
-    class_names = ['Normal', 'Viral Pneumonia', 'COVID']
-
-    model = load_model(model_path, class_names)
-    predicted_class = predict(model, image_path, class_names)
-    print(f'The predicted class for the image is: {predicted_class}')
-    show_image_with_prediction(image_path, predicted_class)
-
-if __name__ == '__main__':
-    main()