NOISSUE - Simplify atls provider initialization by removing external certs SDK dependency (#614)

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

.github/workflows/checkproto.yaml

@@ -34,7 +34,7 @@ jobs:
 
       - name: Set up protoc
         run: |
-          PROTOC_VERSION=33.1
+          PROTOC_VERSION=35.1
           PROTOC_GEN_VERSION=v1.36.11
           PROTOC_GRPC_VERSION=v1.6.0
 

.github/workflows/hal.yml

@@ -17,6 +17,9 @@ jobs:
           sudo rm -rf /opt/ghc
           sudo rm -rf "/usr/local/share/boost"
           sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+          mkdir -p "$RUNNER_TEMP/hostedtoolcache"
+          echo "RUNNER_TOOL_CACHE=$RUNNER_TEMP/hostedtoolcache" >> "$GITHUB_ENV"
+          echo "AGENT_TOOLSDIRECTORY=$RUNNER_TEMP/hostedtoolcache" >> "$GITHUB_ENV"
           df -h
 
       - name: Update Ubuntu
@@ -39,17 +42,16 @@ jobs:
           echo "Free space:"
           df -h
 
+      - name: Checkout cocos
+        uses: actions/checkout@v4
+        with:
+          path: cocos
+
       - name: Install Go
         uses: actions/setup-go@v5
         with:
           go-version: 1.26.x
-          cache-dependency-path: "go.sum"
-
-      - name: Checkout cocos
-        uses: actions/checkout@v4
-        with:
-          repository: "ultravioletrs/cocos"
-          path: cocos
+          cache-dependency-path: "cocos/go.sum"
 
       - name: Checkout buildroot
         uses: actions/checkout@v4
@@ -78,3 +80,9 @@ jobs:
             cocos/build/cocos-agent
             cocos/build/cocos-cli
             cocos/build/cocos-manager
+            cocos/build/cocos-attestation-service
+            cocos/build/cocos-log-forwarder
+            cocos/build/cocos-computation-runner
+            cocos/build/cocos-egress-proxy
+            cocos/build/cocos-ingress-proxy
+            cocos/build/igvmmeasure

agent/agent.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.36.11
-// 	protoc        v6.33.1
+// 	protoc        v7.35.1
 // source: agent/agent.proto
 
 package agent

agent/agent_grpc.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
 // - protoc-gen-go-grpc v1.6.0
-// - protoc             v6.33.1
+// - protoc             v7.35.1
 // source: agent/agent.proto
 
 package agent

agent/algorithm/algorithm.go

@@ -18,11 +18,14 @@ const (
 	AlgoTypeKey                 = "algo_type"
 	AlgoArgsKey                 = "algo_args"
 
-	ResultsDir     = "results"
-	DatasetsDir    = "datasets"
-	AlgoWorkingDir = "/cocos"
+	ResultsDir  = "results"
+	DatasetsDir = "datasets"
 )
 
+// AlgoWorkingDir is the base directory used by algorithm runners (e.g. docker)
+// to create datasets/results mounts. It is a variable so tests can override it.
+var AlgoWorkingDir = "/cocos"
+
 func AlgorithmTypeToContext(ctx context.Context, algoType string) context.Context {
 	return metadata.AppendToOutgoingContext(ctx, AlgoTypeKey, algoType)
 }

agent/algorithm/docker/docker.go

@@ -90,6 +90,16 @@ func (d *docker) Run() error {
 		return fmt.Errorf("could not find image ID")
 	}
 
+	datasetsDir := path.Join(algorithm.AlgoWorkingDir, algorithm.DatasetsDir)
+	resultsDir := path.Join(algorithm.AlgoWorkingDir, algorithm.ResultsDir)
+
+	if err := os.MkdirAll(datasetsDir, 0o755); err != nil {
+		return fmt.Errorf("could not create datasets directory %s: %v", datasetsDir, err)
+	}
+	if err := os.MkdirAll(resultsDir, 0o755); err != nil {
+		return fmt.Errorf("could not create results directory %s: %v", resultsDir, err)
+	}
+
 	// Create and start the container.
 	respContainer, err := cli.ContainerCreate(ctx, &container.Config{
 		Image:        dockerImageName,

agent/cvms/cvms.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.36.11
-// 	protoc        v6.33.1
+// 	protoc        v7.35.1
 // source: agent/cvms/cvms.proto
 
 package cvms
@@ -1123,8 +1123,8 @@ func (x *Algorithm) GetKbs() *KBSConfig {
 
 type Source struct {
 	state           protoimpl.MessageState `protogen:"open.v1"`
-	Type            string                 `protobuf:"bytes,1,opt,name=type,proto3" json:"type,omitempty"`                                                // Type of source: "oci-image" (only OCI images supported for CoCo)
-	Url             string                 `protobuf:"bytes,2,opt,name=url,proto3" json:"url,omitempty"`                                                  // URL of the OCI image (e.g., docker://registry/repo:tag)
+	Type            string                 `protobuf:"bytes,1,opt,name=type,proto3" json:"type,omitempty"`                                                // Type of source: "oci-image", "s3", "gcs", "https", "http"
+	Url             string                 `protobuf:"bytes,2,opt,name=url,proto3" json:"url,omitempty"`                                                  // URL of the resource (e.g., docker://registry/repo:tag, s3://bucket/key, https://host/path)
 	KbsResourcePath string                 `protobuf:"bytes,3,opt,name=kbs_resource_path,json=kbsResourcePath,proto3" json:"kbs_resource_path,omitempty"` // Path to decryption key in KBS (e.g., "default/key/my-key")
 	Encrypted       bool                   `protobuf:"varint,4,opt,name=encrypted,proto3" json:"encrypted,omitempty"`                                     // Whether the resource is encrypted (requires KBS)
 	unknownFields   protoimpl.UnknownFields

agent/cvms/cvms_grpc.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
 // - protoc-gen-go-grpc v1.6.0
-// - protoc             v6.33.1
+// - protoc             v7.35.1
 // source: agent/cvms/cvms.proto
 
 package cvms

agent/events/events.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.36.11
-// 	protoc        v6.33.1
+// 	protoc        v7.35.1
 // source: agent/events/events.proto
 
 package events

agent/log/log.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.36.11
-// 	protoc        v6.33.1
+// 	protoc        v7.35.1
 // source: agent/log/log.proto
 
 package log

agent/log/log_grpc.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
 // - protoc-gen-go-grpc v1.6.0
-// - protoc             v6.33.1
+// - protoc             v7.35.1
 // source: agent/log/log.proto
 
 package log

agent/mock_attestation_client_test.go

@@ -28,6 +28,11 @@ func (m *MockAttestationClient) GetAzureToken(ctx context.Context, nonce [32]byt
 	return args.Get(0).([]byte), args.Error(1)
 }
 
+func (m *MockAttestationClient) GetKbsToken(ctx context.Context) ([]byte, error) {
+	args := m.Called(ctx)
+	return args.Get(0).([]byte), args.Error(1)
+}
+
 func (m *MockAttestationClient) Close() error {
 	args := m.Called()
 	return args.Error(0)

agent/resource_test.go

@@ -51,9 +51,13 @@ func TestDownloadAndDecryptGenericResource(t *testing.T) {
 	mockDownloader.On("Type").Return(resource.SourceTypeHTTP)
 	registry.Register(mockDownloader)
 
+	attestationClient := new(MockAttestationClient)
+	attestationClient.On("GetKbsToken", mock.Anything).Return([]byte("mockToken"), nil).Maybe()
+
 	svc := &agentService{
-		logger:           slog.Default(),
-		resourceRegistry: registry,
+		logger:            slog.Default(),
+		resourceRegistry:  registry,
+		attestationClient: attestationClient,
 		computation: Computation{
 			Algorithm: &Algorithm{
 				KBS: &KBSConfig{
@@ -122,8 +126,12 @@ func TestDownloadAndDecryptGenericResource(t *testing.T) {
 }
 
 func TestGetKeyFromKBS(t *testing.T) {
+	attestationClient := new(MockAttestationClient)
+	attestationClient.On("GetKbsToken", mock.Anything).Return([]byte("mockToken"), nil).Maybe()
+
 	svc := &agentService{
-		logger: slog.Default(),
+		logger:            slog.Default(),
+		attestationClient: attestationClient,
 		computation: Computation{
 			Algorithm: &Algorithm{
 				KBS: &KBSConfig{

agent/runner/runner.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.36.11
-// 	protoc        v6.33.1
+// 	protoc        v7.35.1
 // source: agent/runner/runner.proto
 
 package runner
@@ -26,15 +26,15 @@ const (
 )
 
 type RunRequest struct {
-	state         protoimpl.MessageState `protogen:"open.v1"`
-	ComputationId string                 `protobuf:"bytes,1,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
-	AlgoType      string                 `protobuf:"bytes,2,opt,name=algo_type,json=algoType,proto3" json:"algo_type,omitempty"` // "binary", "python", "wasm", "docker"
-	Algorithm     []byte                 `protobuf:"bytes,3,opt,name=algorithm,proto3" json:"algorithm,omitempty"`               // The algorithm binary/script content
-	Requirements  []byte                 `protobuf:"bytes,4,opt,name=requirements,proto3" json:"requirements,omitempty"`         // Python requirements.txt content
-	Args          []string               `protobuf:"bytes,5,rep,name=args,proto3" json:"args,omitempty"`
-	Datasets      []*Dataset             `protobuf:"bytes,6,rep,name=datasets,proto3" json:"datasets,omitempty"`
-	unknownFields protoimpl.UnknownFields
-	sizeCache     protoimpl.SizeCache
+	state            protoimpl.MessageState `protogen:"open.v1"`
+	ComputationId    string                 `protobuf:"bytes,1,opt,name=computation_id,json=computationId,proto3" json:"computation_id,omitempty"`
+	AlgoType         string                 `protobuf:"bytes,2,opt,name=algo_type,json=algoType,proto3" json:"algo_type,omitempty"`                         // "binary", "python", "wasm", "docker"
+	AlgorithmPath    string                 `protobuf:"bytes,3,opt,name=algorithm_path,json=algorithmPath,proto3" json:"algorithm_path,omitempty"`          // Path to the staged algorithm artifact
+	RequirementsPath string                 `protobuf:"bytes,4,opt,name=requirements_path,json=requirementsPath,proto3" json:"requirements_path,omitempty"` // Optional path to staged Python requirements.txt
+	Args             []string               `protobuf:"bytes,5,rep,name=args,proto3" json:"args,omitempty"`
+	Datasets         []*Dataset             `protobuf:"bytes,6,rep,name=datasets,proto3" json:"datasets,omitempty"`
+	unknownFields    protoimpl.UnknownFields
+	sizeCache        protoimpl.SizeCache
 }
 
 func (x *RunRequest) Reset() {
@@ -81,18 +81,18 @@ func (x *RunRequest) GetAlgoType() string {
 	return ""
 }
 
-func (x *RunRequest) GetAlgorithm() []byte {
+func (x *RunRequest) GetAlgorithmPath() string {
 	if x != nil {
-		return x.Algorithm
+		return x.AlgorithmPath
 	}
-	return nil
+	return ""
 }
 
-func (x *RunRequest) GetRequirements() []byte {
+func (x *RunRequest) GetRequirementsPath() string {
 	if x != nil {
-		return x.Requirements
+		return x.RequirementsPath
 	}
-	return nil
+	return ""
 }
 
 func (x *RunRequest) GetArgs() []string {
@@ -261,13 +261,13 @@ var File_agent_runner_runner_proto protoreflect.FileDescriptor
 
 const file_agent_runner_runner_proto_rawDesc = "" +
 	"\n" +
-	"\x19agent/runner/runner.proto\x12\x06runner\x1a\x1bgoogle/protobuf/empty.proto\"\xd3\x01\n" +
+	"\x19agent/runner/runner.proto\x12\x06runner\x1a\x1bgoogle/protobuf/empty.proto\"\xe5\x01\n" +
 	"\n" +
 	"RunRequest\x12%\n" +
 	"\x0ecomputation_id\x18\x01 \x01(\tR\rcomputationId\x12\x1b\n" +
-	"\talgo_type\x18\x02 \x01(\tR\balgoType\x12\x1c\n" +
-	"\talgorithm\x18\x03 \x01(\fR\talgorithm\x12\"\n" +
-	"\frequirements\x18\x04 \x01(\fR\frequirements\x12\x12\n" +
+	"\talgo_type\x18\x02 \x01(\tR\balgoType\x12%\n" +
+	"\x0ealgorithm_path\x18\x03 \x01(\tR\ralgorithmPath\x12+\n" +
+	"\x11requirements_path\x18\x04 \x01(\tR\x10requirementsPath\x12\x12\n" +
 	"\x04args\x18\x05 \x03(\tR\x04args\x12+\n" +
 	"\bdatasets\x18\x06 \x03(\v2\x0f.runner.DatasetR\bdatasets\"9\n" +
 	"\aDataset\x12\x1a\n" +

agent/runner/runner.proto

@@ -17,8 +17,8 @@ service ComputationRunner {
 message RunRequest {
   string computation_id = 1;
   string algo_type = 2; // "binary", "python", "wasm", "docker"
-  bytes algorithm = 3; // The algorithm binary/script content
-  bytes requirements = 4; // Python requirements.txt content
+  string algorithm_path = 3; // Path to the staged algorithm artifact
+  string requirements_path = 4; // Optional path to staged Python requirements.txt
   repeated string args = 5;
   repeated Dataset datasets = 6;
 }

agent/runner/runner_grpc.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
 // - protoc-gen-go-grpc v1.6.0
-// - protoc             v6.33.1
+// - protoc             v7.35.1
 // source: agent/runner/runner.proto
 
 package runner

agent/runner/service/service.go

@@ -6,8 +6,6 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
-	"os"
-	"path/filepath"
 	"sync"
 
 	"github.com/ultravioletrs/cocos/agent/algorithm"
@@ -20,10 +18,6 @@ import (
 	"google.golang.org/protobuf/types/known/emptypb"
 )
 
-const (
-	algoFilePermission = 0o700
-)
-
 var _ pb.ComputationRunnerServer = (*RunnerService)(nil)
 
 type RunnerService struct {
@@ -58,65 +52,24 @@ func (s *RunnerService) Run(ctx context.Context, req *pb.RunRequest) (*pb.RunRes
 		s.mu.Unlock()
 	}()
 
-	currentDir, err := os.Getwd()
-	if err != nil {
-		return nil, fmt.Errorf("error getting current directory: %v", err)
+	if req.AlgorithmPath == "" {
+		return nil, fmt.Errorf("algorithm path is required")
 	}
 
-	// Write Algo File
-	algoPath := filepath.Join(currentDir, "algo")
-	f, err := os.Create(algoPath)
-	if err != nil {
-		return nil, fmt.Errorf("error creating algorithm file: %v", err)
-	}
-	if _, err := f.Write(req.Algorithm); err != nil {
-		return nil, fmt.Errorf("error writing algorithm to file: %v", err)
-	}
-	if err := os.Chmod(algoPath, algoFilePermission); err != nil {
-		return nil, fmt.Errorf("error changing file permissions: %v", err)
-	}
-	if err := f.Close(); err != nil {
-		return nil, fmt.Errorf("error closing file: %v", err)
-	}
-	defer func() {
-		if err := os.Remove(algoPath); err != nil {
-			s.logger.Warn("error removing algorithm file", "error", err)
-		}
-	}()
-
 	var algo algorithm.Algorithm
 
 	switch req.AlgoType {
 	case string(algorithm.AlgoTypeBin):
-		algo = binary.NewAlgorithm(s.logger, s.eventSvc, algoPath, req.Args, req.ComputationId)
+		algo = binary.NewAlgorithm(s.logger, s.eventSvc, req.AlgorithmPath, req.Args, req.ComputationId)
 	case string(algorithm.AlgoTypePython):
-		var requirementsFile string
-		if len(req.Requirements) > 0 {
-			fr, err := os.CreateTemp("", "requirements.txt")
-			if err != nil {
-				return nil, fmt.Errorf("error creating requirments file: %v", err)
-			}
-			defer func() {
-				if err := os.Remove(fr.Name()); err != nil {
-					s.logger.Warn("error removing requirements file", "error", err)
-				}
-			}()
-			if _, err := fr.Write(req.Requirements); err != nil {
-				return nil, fmt.Errorf("error writing requirements to file: %v", err)
-			}
-			if err := fr.Close(); err != nil {
-				return nil, fmt.Errorf("error closing file: %v", err)
-			}
-			requirementsFile = fr.Name()
-		}
 		// Assuming default python runtime if not specified in request (proto doesn't have runtime field yet)
 		// We can add it or assume.
 		runtime := python.PyRuntime
-		algo = python.NewAlgorithm(s.logger, s.eventSvc, runtime, requirementsFile, algoPath, req.Args, req.ComputationId)
+		algo = python.NewAlgorithm(s.logger, s.eventSvc, runtime, req.RequirementsPath, req.AlgorithmPath, req.Args, req.ComputationId)
 	case string(algorithm.AlgoTypeWasm):
-		algo = wasm.NewAlgorithm(s.logger, s.eventSvc, req.Args, algoPath, req.ComputationId)
+		algo = wasm.NewAlgorithm(s.logger, s.eventSvc, req.Args, req.AlgorithmPath, req.ComputationId)
 	case string(algorithm.AlgoTypeDocker):
-		algo = docker.NewAlgorithm(s.logger, s.eventSvc, algoPath, req.ComputationId)
+		algo = docker.NewAlgorithm(s.logger, s.eventSvc, req.AlgorithmPath, req.ComputationId)
 	default:
 		return nil, fmt.Errorf("unsupported algorithm type: %s", req.AlgoType)
 	}

agent/runner/service/service_test.go

@@ -8,11 +8,13 @@ import (
 	"fmt"
 	"log/slog"
 	"os"
+	"path/filepath"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/ultravioletrs/cocos/agent/algorithm"
 	pb "github.com/ultravioletrs/cocos/agent/runner"
 )
 
@@ -30,6 +32,15 @@ func (m *MockEventService) SendEvent(cmpID, event, status string, details json.R
 	})
 }
 
+func writeRunnerTestFile(t *testing.T, dir, name string, data []byte, mode os.FileMode) string {
+	t.Helper()
+
+	path := filepath.Join(dir, name)
+	require.NoError(t, os.WriteFile(path, data, mode))
+
+	return path
+}
+
 // TestNewRunnerService tests the creation of a new runner service.
 func TestNewRunnerService(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
@@ -48,6 +59,7 @@ func TestRunWithBinaryAlgorithm(t *testing.T) {
 	tmpDir := t.TempDir()
 	require.NoError(t, os.Chdir(tmpDir))
 	defer func() { require.NoError(t, os.Chdir(origDir)) }()
+	algoPath := writeRunnerTestFile(t, tmpDir, "algo", []byte("#!/bin/bash\necho 'test'"), 0o700)
 
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	eventSvc := &MockEventService{}
@@ -56,7 +68,7 @@ func TestRunWithBinaryAlgorithm(t *testing.T) {
 	req := &pb.RunRequest{
 		ComputationId: "test-1",
 		AlgoType:      "bin",
-		Algorithm:     []byte("#!/bin/bash\necho 'test'"),
+		AlgorithmPath: algoPath,
 		Args:          []string{"arg1", "arg2"},
 	}
 
@@ -72,13 +84,16 @@ func TestRunWithPythonAlgorithm(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	eventSvc := &MockEventService{}
 	rs := New(logger, eventSvc)
+	tmpDir := t.TempDir()
+	algoPath := writeRunnerTestFile(t, tmpDir, "algo.py", []byte("print('hello')"), 0o600)
+	requirementsPath := writeRunnerTestFile(t, tmpDir, "requirements.txt", []byte("numpy==2.2.0"), 0o600)
 
 	req := &pb.RunRequest{
-		ComputationId: "test-python",
-		AlgoType:      "python",
-		Algorithm:     []byte("print('hello')"),
-		Args:          []string{},
-		Requirements:  []byte("numpy==2.2.0"),
+		ComputationId:    "test-python",
+		AlgoType:         "python",
+		AlgorithmPath:    algoPath,
+		Args:             []string{},
+		RequirementsPath: requirementsPath,
 	}
 
 	resp, err := rs.Run(context.Background(), req)
@@ -86,9 +101,6 @@ func TestRunWithPythonAlgorithm(t *testing.T) {
 	require.NotNil(t, resp)
 	assert.Empty(t, resp.Error)
 	assert.Equal(t, "test-python", resp.ComputationId)
-	t.Cleanup(func() {
-		_ = os.Remove("algo")
-	})
 }
 
 // TestRunWithPythonAlgorithmNoRequirements tests running Python without requirements.
@@ -96,11 +108,13 @@ func TestRunWithPythonAlgorithmNoRequirements(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	eventSvc := &MockEventService{}
 	rs := New(logger, eventSvc)
+	tmpDir := t.TempDir()
+	algoPath := writeRunnerTestFile(t, tmpDir, "algo.py", []byte("print('hello')"), 0o600)
 
 	req := &pb.RunRequest{
 		ComputationId: "test-python-noreq",
 		AlgoType:      "python",
-		Algorithm:     []byte("print('hello')"),
+		AlgorithmPath: algoPath,
 		Args:          []string{},
 	}
 
@@ -109,9 +123,6 @@ func TestRunWithPythonAlgorithmNoRequirements(t *testing.T) {
 	require.NotNil(t, resp)
 	assert.Empty(t, resp.Error)
 	assert.Equal(t, "test-python-noreq", resp.ComputationId)
-	t.Cleanup(func() {
-		_ = os.Remove("algo")
-	})
 }
 
 // TestRunWithWasmAlgorithm tests running a WASM algorithm.
@@ -119,11 +130,13 @@ func TestRunWithWasmAlgorithm(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	eventSvc := &MockEventService{}
 	rs := New(logger, eventSvc)
+	tmpDir := t.TempDir()
+	algoPath := writeRunnerTestFile(t, tmpDir, "algo.wasm", []byte{0x00, 0x61, 0x73, 0x6d}, 0o600)
 
 	req := &pb.RunRequest{
 		ComputationId: "test-wasm",
 		AlgoType:      "wasm",
-		Algorithm:     []byte{0x00, 0x61, 0x73, 0x6d},
+		AlgorithmPath: algoPath,
 		Args:          []string{},
 	}
 
@@ -135,9 +148,6 @@ func TestRunWithWasmAlgorithm(t *testing.T) {
 		t.Skip("wasmedge not found, skipping test")
 	}
 	assert.Equal(t, "test-wasm", resp.ComputationId)
-	t.Cleanup(func() {
-		_ = os.Remove("algo")
-	})
 }
 
 // TestRunWithDockerAlgorithm tests running a Docker algorithm.
@@ -145,11 +155,17 @@ func TestRunWithDockerAlgorithm(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	eventSvc := &MockEventService{}
 	rs := New(logger, eventSvc)
+	tmpDir := t.TempDir()
+	algoPath := writeRunnerTestFile(t, tmpDir, "Dockerfile", []byte("FROM ubuntu:latest\nRUN echo 'test'"), 0o600)
+
+	origWorkingDir := algorithm.AlgoWorkingDir
+	algorithm.AlgoWorkingDir = tmpDir
+	t.Cleanup(func() { algorithm.AlgoWorkingDir = origWorkingDir })
 
 	req := &pb.RunRequest{
 		ComputationId: "test-docker",
 		AlgoType:      "docker",
-		Algorithm:     []byte("FROM ubuntu:latest\nRUN echo 'test'"),
+		AlgorithmPath: algoPath,
 		Args:          []string{},
 	}
 
@@ -161,9 +177,6 @@ func TestRunWithDockerAlgorithm(t *testing.T) {
 		t.Skip("Docker issue, skipping test")
 	}
 	assert.Equal(t, "test-docker", resp.ComputationId)
-	t.Cleanup(func() {
-		_ = os.Remove("algo")
-	})
 }
 
 // TestRunWithUnsupportedAlgorithmType tests running with unsupported algorithm type.
@@ -175,7 +188,7 @@ func TestRunWithUnsupportedAlgorithmType(t *testing.T) {
 	req := &pb.RunRequest{
 		ComputationId: "test-unsupported",
 		AlgoType:      "unsupported",
-		Algorithm:     []byte("test"),
+		AlgorithmPath: "/tmp/test",
 		Args:          []string{},
 	}
 
@@ -189,12 +202,14 @@ func TestRunAlreadyRunning(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	eventSvc := &MockEventService{}
 	rs := New(logger, eventSvc)
+	tmpDir := t.TempDir()
+	algoPath := writeRunnerTestFile(t, tmpDir, "algo", []byte("#!/bin/bash\nsleep 30"), 0o700)
 
 	// Use a long-running bash script
 	req := &pb.RunRequest{
 		ComputationId: "test-running",
 		AlgoType:      "bin",
-		Algorithm:     []byte("#!/bin/bash\nsleep 30"),
+		AlgorithmPath: algoPath,
 		Args:          []string{},
 	}
 
@@ -211,9 +226,6 @@ func TestRunAlreadyRunning(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, resp)
 	assert.Equal(t, "computation already running", resp.Error)
-	t.Cleanup(func() {
-		_ = os.Remove("algo")
-	})
 }
 
 // TestStopWhenRunning tests stopping a running computation.
@@ -221,11 +233,13 @@ func TestStopWhenRunning(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	eventSvc := &MockEventService{}
 	rs := New(logger, eventSvc)
+	tmpDir := t.TempDir()
+	algoPath := writeRunnerTestFile(t, tmpDir, "algo", []byte("#!/bin/bash\nsleep 10"), 0o700)
 
 	req := &pb.RunRequest{
 		ComputationId: "test-stop",
 		AlgoType:      "bin",
-		Algorithm:     []byte("#!/bin/bash\nsleep 10"),
+		AlgorithmPath: algoPath,
 		Args:          []string{},
 	}
 
@@ -243,9 +257,6 @@ func TestStopWhenRunning(t *testing.T) {
 	stopResp, err := rs.Stop(context.Background(), stopReq)
 	require.NoError(t, err)
 	require.NotNil(t, stopResp)
-	t.Cleanup(func() {
-		_ = os.Remove("algo")
-	})
 }
 
 // TestRunErrors tests error paths in Run.
@@ -255,59 +266,27 @@ func TestRunErrors(t *testing.T) {
 	rs := New(logger, eventSvc)
 
 	t.Run("create algo file failure", func(t *testing.T) {
-		// Create a directory named "algo" to make os.Create("algo") fail
-		err := os.Mkdir("algo", 0o755)
-		require.NoError(t, err)
-		defer os.RemoveAll("algo")
-
+		var err error
 		req := &pb.RunRequest{
 			ComputationId: "test-err",
 			AlgoType:      "bin",
-			Algorithm:     []byte("test"),
+			AlgorithmPath: "",
 		}
 		_, err = rs.Run(context.Background(), req)
 		assert.Error(t, err)
-		assert.Contains(t, err.Error(), "error creating algorithm file")
-	})
-
-	t.Run("getwd failure", func(t *testing.T) {
-		origDir, _ := os.Getwd()
-		tmpDir := t.TempDir()
-		err := os.Chdir(tmpDir)
-		require.NoError(t, err)
-
-		// Remove the current working directory to trigger Getwd failure
-		err = os.RemoveAll(tmpDir)
-		require.NoError(t, err)
-
-		req := &pb.RunRequest{
-			ComputationId: "test-err-getwd",
-			AlgoType:      "bin",
-			Algorithm:     []byte("test"),
-		}
-		_, err = rs.Run(context.Background(), req)
-		assert.Error(t, err)
-		assert.Contains(t, err.Error(), "error getting current directory")
-
-		// Restore working directory
-		_ = os.Chdir(origDir)
+		assert.Contains(t, err.Error(), "algorithm path is required")
 	})
 
 	t.Run("requirements file creation failure", func(t *testing.T) {
-		// This one is harder because it uses os.CreateTemp("", "requirements.txt")
-		// We can't easily make this fail without reaching into the system's temp dir.
-		// Skipping for now as it's a very unlikely edge case.
+		// Requirements are now staged by the agent, so the runner no longer creates temp files.
 	})
 
 	t.Run("chmod failure", func(t *testing.T) {
-		// We can't easily mock os.Chmod, but we can try to make the file unmodifiable
-		// On Linux, we can set the immutable attribute, but that requires root.
-		// Alternatively, we can try to use a directory with permissions that prevent chmod?
-		// No, chmod usually works if you own the file.
+		// Permission management is now the agent's responsibility during staging.
 	})
 
 	t.Run("write algorithm failure", func(t *testing.T) {
-		// This is also hard without mocking os.File.Write or reaching internal limits.
+		// Write failures are now handled by the agent before invoking the runner.
 	})
 }
 
@@ -316,11 +295,13 @@ func TestConcurrentRun(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	eventSvc := &MockEventService{}
 	rs := New(logger, eventSvc)
+	tmpDir := t.TempDir()
+	algoPath := writeRunnerTestFile(t, tmpDir, "algo", []byte("#!/bin/bash\nsleep 15"), 0o700)
 
 	req := &pb.RunRequest{
 		ComputationId: "test-concurrent",
 		AlgoType:      "bin",
-		Algorithm:     []byte("#!/bin/bash\nsleep 15"),
+		AlgorithmPath: algoPath,
 		Args:          []string{},
 	}
 
@@ -336,9 +317,6 @@ func TestConcurrentRun(t *testing.T) {
 	resp2, err := rs.Run(context.Background(), req)
 	require.NoError(t, err)
 	assert.Equal(t, "computation already running", resp2.Error)
-	t.Cleanup(func() {
-		_ = os.Remove("algo")
-	})
 }
 
 // TestRunWithMultipleArgs tests running with multiple arguments.
@@ -346,11 +324,13 @@ func TestRunWithMultipleArgs(t *testing.T) {
 	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	eventSvc := &MockEventService{}
 	rs := New(logger, eventSvc)
+	tmpDir := t.TempDir()
+	algoPath := writeRunnerTestFile(t, tmpDir, "algo", []byte("#!/bin/bash\necho $@"), 0o700)
 
 	req := &pb.RunRequest{
 		ComputationId: "test-multi-args",
 		AlgoType:      "bin",
-		Algorithm:     []byte("#!/bin/bash\necho $@"),
+		AlgorithmPath: algoPath,
 		Args:          []string{"arg1", "arg2", "arg3", "arg4"},
 	}
 
@@ -359,9 +339,6 @@ func TestRunWithMultipleArgs(t *testing.T) {
 	require.NotNil(t, resp)
 	assert.Empty(t, resp.Error)
 	assert.Equal(t, "test-multi-args", resp.ComputationId)
-	t.Cleanup(func() {
-		_ = os.Remove("algo")
-	})
 }
 
 func TestStopFailure(t *testing.T) {

agent/service.go

@@ -5,7 +5,11 @@ package agent
 
 import (
 	"context"
+	"crypto/ecdh"
+	"crypto/ecdsa"
+	"crypto/x509"
 	"encoding/json"
+	"encoding/pem"
 	"fmt"
 	"io"
 	"log/slog"
@@ -29,6 +33,7 @@ import (
 	"github.com/ultravioletrs/cocos/pkg/attestation/vtpm"
 	attestation_client "github.com/ultravioletrs/cocos/pkg/clients/grpc/attestation"
 	runner_client "github.com/ultravioletrs/cocos/pkg/clients/grpc/runner"
+	"github.com/ultravioletrs/cocos/pkg/crypto"
 	"github.com/ultravioletrs/cocos/pkg/oci"
 	"github.com/ultravioletrs/cocos/pkg/resource"
 	"golang.org/x/crypto/sha3"
@@ -77,7 +82,9 @@ const (
 )
 
 const (
-	algoFilePermission = 0o700
+	algoFilePermission   = 0o700
+	algoFileName         = "algo"
+	requirementsFileName = "requirements.txt"
 )
 
 var (
@@ -108,6 +115,30 @@ func ensureDir(path string, mode os.FileMode) error {
 	return nil
 }
 
+func writeFile(path string, data []byte, mode os.FileMode) error {
+	if err := os.WriteFile(path, data, mode); err != nil {
+		return fmt.Errorf("writing file %q: %w", path, err)
+	}
+
+	return nil
+}
+
+func stageRequirementsFile(workDir string, requirements []byte) (string, error) {
+	requirementsPath := filepath.Join(workDir, requirementsFileName)
+	if len(requirements) == 0 {
+		if err := os.Remove(requirementsPath); err != nil && !os.IsNotExist(err) {
+			return "", fmt.Errorf("removing stale requirements file: %w", err)
+		}
+		return "", nil
+	}
+
+	if err := writeFile(requirementsPath, requirements, 0o600); err != nil {
+		return "", err
+	}
+
+	return requirementsPath, nil
+}
+
 var (
 	// ErrMalformedEntity indicates malformed entity specification (e.g.
 	// invalid username or password).
@@ -359,10 +390,14 @@ func (as *agentService) StopComputation(ctx context.Context) error {
 		return fmt.Errorf("error removing results directory: %v", err)
 	}
 
-	if err := os.Remove("algo"); err != nil && !os.IsNotExist(err) {
+	if err := os.Remove(algoFileName); err != nil && !os.IsNotExist(err) {
 		as.logger.Warn("error removing algorithm file", "error", err)
 	}
 
+	if err := os.Remove(requirementsFileName); err != nil && !os.IsNotExist(err) {
+		as.logger.Warn("error removing requirements file", "error", err)
+	}
+
 	as.sm.Reset(Idle)
 
 	as.computation = Computation{}
@@ -443,9 +478,9 @@ func (as *agentService) downloadAlgorithmIfRemote(state statemachine.State) {
 		}
 
 		// Write algorithm to file
-		currentDir, err := os.Getwd()
-		if err != nil {
-			as.runError = fmt.Errorf("error getting current directory: %w", err)
+		currentDir, getwdErr := os.Getwd()
+		if getwdErr != nil {
+			as.runError = fmt.Errorf("error getting current directory: %w", getwdErr)
 			as.logger.Error(as.runError.Error())
 			as.sm.SendEvent(RunFailed)
 			return
@@ -499,7 +534,14 @@ func (as *agentService) downloadAlgorithmIfRemote(state statemachine.State) {
 		}
 
 		as.algoReceived = true
-		as.algoRequirements = res.Requirements // Store requirements for installation
+		as.algoRequirements = res.Requirements
+
+		if _, err := stageRequirementsFile(currentDir, as.algoRequirements); err != nil {
+			as.runError = fmt.Errorf("error staging requirements file: %w", err)
+			as.logger.Error(as.runError.Error())
+			as.sm.SendEvent(RunFailed)
+			return
+		}
 
 		// The initramfs may have already provisioned /cocos/datasets.
 		if err := ensureDir(algorithm.DatasetsDir, 0o755); err != nil {
@@ -762,26 +804,91 @@ func (as *agentService) getKeyFromKBS(ctx context.Context, kbsURL, resourcePath
 
 	as.logger.Info("fetching key from KBS", "url", kbsResourceURL)
 
-	// Use a simple HTTP GET to KBS for now.
-	// In a full CoCo deployment, this would go through the Attestation Agent
-	// which performs attestation before KBS releases the key.
-	// For non-OCI resources, the AA/KBS handshake may need to be handled
-	// differently than via ocicrypt.
-	resp, err := kbsHTTPGet(ctx, kbsResourceURL)
+	// Fetch token from attestation service via the initialized client
+	tokenBytes, err := as.attestationClient.GetKbsToken(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to retrieve token from attestation service: %w", err)
+	}
+
+	// Unmarshal the JSON from attestation service containing token and private key
+	var msg struct {
+		Token      string `json:"token"`
+		TeeKeyPair string `json:"tee_keypair"`
+	}
+
+	var token string
+	var teeKeyPairPem string
+	if err := json.Unmarshal(tokenBytes, &msg); err == nil {
+		token = strings.TrimSpace(msg.Token)
+		teeKeyPairPem = msg.TeeKeyPair
+		as.logger.Info("retrieved token and keypair from attestation service successfully")
+	} else {
+		// Fallback for non-JSON token bytes
+		token = strings.TrimSpace(string(tokenBytes))
+		as.logger.Warn("attestation service token is not in JSON format; using raw token without local JWE decryption fallback", "error", err)
+	}
+
+	resp, err := kbsHTTPGet(ctx, kbsResourceURL, token)
 	if err != nil {
 		return nil, fmt.Errorf("failed to fetch key from KBS at %s: %w", kbsResourceURL, err)
 	}
 
-	return resp, nil
+	// If we don't have a private key, return the response as is (could be plaintext)
+	if teeKeyPairPem == "" {
+		as.logger.Info("no session keypair found, returning KBS response directly")
+		return resp, nil
+	}
+
+	// Parse JWE response
+	encryptedRes, err := crypto.ParseEncryptedResource(resp)
+	if err != nil {
+		as.logger.Warn("failed to parse JWE resource response, returning raw response", "error", err)
+		return resp, nil
+	}
+
+	// Parse PKCS#8 private key PEM
+	block, _ := pem.Decode([]byte(teeKeyPairPem))
+	if block == nil {
+		return nil, fmt.Errorf("failed to decode private key PEM")
+	}
+
+	privKey, err := x509.ParsePKCS8PrivateKey(block.Bytes)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse private key PKCS8: %w", err)
+	}
+
+	var ecdhKey *ecdh.PrivateKey
+	switch k := privKey.(type) {
+	case *ecdsa.PrivateKey:
+		ecdhKey, err = k.ECDH()
+		if err != nil {
+			return nil, fmt.Errorf("failed to get ECDH private key: %w", err)
+		}
+	default:
+		return nil, fmt.Errorf("private key is not ECDSA (unsupported curve or type)")
+	}
+
+	// Decrypt JWE resource using the private key
+	decryptedKey, err := crypto.DecryptWithWrappedKey(*encryptedRes, ecdhKey)
+	if err != nil {
+		return nil, fmt.Errorf("failed to decrypt KBS resource: %w", err)
+	}
+
+	as.logger.Info("successfully decrypted KBS key", "key_len", len(decryptedKey))
+	return decryptedKey, nil
 }
 
 // kbsHTTPGet performs an HTTP GET to the KBS endpoint.
-func kbsHTTPGet(ctx context.Context, url string) ([]byte, error) {
+func kbsHTTPGet(ctx context.Context, url string, token string) ([]byte, error) {
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
 	if err != nil {
 		return nil, err
 	}
 
+	if token != "" {
+		req.Header.Set("Authorization", "Bearer "+token)
+	}
+
 	client := &http.Client{}
 	resp, err := client.Do(req)
 	if err != nil {
@@ -983,21 +1090,9 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 		return fmt.Errorf("error getting current directory: %v", err)
 	}
 
-	f, err := os.Create(filepath.Join(currentDir, "algo"))
-	if err != nil {
-		return fmt.Errorf("error creating algorithm file: %v", err)
-	}
-
-	if _, err := f.Write(algoData); err != nil {
-		return fmt.Errorf("error writing algorithm to file: %v", err)
-	}
-
-	if err := os.Chmod(f.Name(), algoFilePermission); err != nil {
-		return fmt.Errorf("error changing file permissions: %v", err)
-	}
-
-	if err := f.Close(); err != nil {
-		return fmt.Errorf("error closing file: %v", err)
+	algoPath := filepath.Join(currentDir, algoFileName)
+	if err := writeFile(algoPath, algoData, algoFilePermission); err != nil {
+		return fmt.Errorf("error staging algorithm file: %w", err)
 	}
 
 	algoType := algorithm.AlgorithmTypeFromContext(ctx)
@@ -1009,7 +1104,13 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 
 	as.algoType = algoType
 	as.algoArgs = args
-	as.algoRequirements = algo.Requirements
+	if len(as.algoRequirements) == 0 {
+		as.algoRequirements = algo.Requirements
+	}
+
+	if _, err := stageRequirementsFile(currentDir, as.algoRequirements); err != nil {
+		return fmt.Errorf("error staging requirements file: %w", err)
+	}
 	as.algoReceived = true
 
 	if err := ensureDir(algorithm.DatasetsDir, 0o755); err != nil {
@@ -1206,9 +1307,9 @@ func (as *agentService) runComputation(state statemachine.State) {
 		}
 	}()
 
-	// Read algo file
 	currentDir, _ := os.Getwd()
-	algoFile := filepath.Join(currentDir, "algo")
+	algoFile := filepath.Join(currentDir, algoFileName)
+	requirementsFile := filepath.Join(currentDir, requirementsFileName)
 
 	defer func() {
 		if err := os.RemoveAll(algorithm.ResultsDir); err != nil {
@@ -1220,6 +1321,9 @@ func (as *agentService) runComputation(state statemachine.State) {
 		if err := os.Remove(algoFile); err != nil && !os.IsNotExist(err) {
 			as.logger.Warn(fmt.Sprintf("error removing algorithm file: %s", err.Error()))
 		}
+		if err := os.Remove(requirementsFile); err != nil && !os.IsNotExist(err) {
+			as.logger.Warn(fmt.Sprintf("error removing requirements file: %s", err.Error()))
+		}
 	}()
 
 	if err := ensureDir(algorithm.ResultsDir, 0o755); err != nil {
@@ -1231,10 +1335,21 @@ func (as *agentService) runComputation(state statemachine.State) {
 		return
 	}
 
-	algoBytes, err := os.ReadFile(algoFile)
-	if err != nil {
+	if _, err := os.Stat(algoFile); err != nil {
 		as.mu.Lock()
-		as.runError = fmt.Errorf("failed to read algo file: %w", err)
+		as.runError = fmt.Errorf("failed to stat algo file: %w", err)
+		as.mu.Unlock()
+		as.logger.Warn(as.runError.Error())
+		as.publishEvent(Failed.String())(state)
+		return
+	}
+
+	requirementsPath := ""
+	if _, err := os.Stat(requirementsFile); err == nil {
+		requirementsPath = requirementsFile
+	} else if !os.IsNotExist(err) {
+		as.mu.Lock()
+		as.runError = fmt.Errorf("failed to stat requirements file: %w", err)
 		as.mu.Unlock()
 		as.logger.Warn(as.runError.Error())
 		as.publishEvent(Failed.String())(state)
@@ -1245,11 +1360,11 @@ func (as *agentService) runComputation(state statemachine.State) {
 
 	// Call Runner
 	resp, err := as.runnerClient.Run(context.Background(), &runnerpb.RunRequest{
-		ComputationId: as.computation.ID,
-		AlgoType:      as.algoType,
-		Algorithm:     algoBytes,
-		Requirements:  as.algoRequirements,
-		Args:          as.algoArgs,
+		ComputationId:    as.computation.ID,
+		AlgoType:         as.algoType,
+		AlgorithmPath:    algoFile,
+		RequirementsPath: requirementsPath,
+		Args:             as.algoArgs,
 		// Datasets implicit on shared FS
 	})
 	if err != nil {

agent/service_test.go

@@ -1014,7 +1014,7 @@ func TestRunComputation(t *testing.T) {
 		svc.runComputation(Running)
 
 		assert.Error(t, svc.runError)
-		assert.Contains(t, svc.runError.Error(), "failed to read algo file")
+		assert.Contains(t, svc.runError.Error(), "failed to stat algo file")
 		sm.AssertExpectations(t)
 	})
 
@@ -1026,7 +1026,9 @@ func TestRunComputation(t *testing.T) {
 		require.NoError(t, os.WriteFile("algo", []byte("#!/bin/sh\necho ok\n"), 0o755))
 
 		runnerCli := new(runnermocks.Client)
-		runnerCli.On("Run", mock.Anything, mock.Anything).Return((*runnerpb.RunResponse)(nil), fmt.Errorf("runner unavailable"))
+		runnerCli.On("Run", mock.Anything, mock.MatchedBy(func(req *runnerpb.RunRequest) bool {
+			return req != nil && filepath.Base(req.AlgorithmPath) == algoFileName && req.RequirementsPath == "" && len(req.Args) == 0
+		})).Return((*runnerpb.RunResponse)(nil), fmt.Errorf("runner unavailable"))
 
 		eventsSvc := new(mocks.Service)
 		eventsSvc.EXPECT().SendEvent(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return().Maybe()
@@ -1893,6 +1895,6 @@ func TestEnsureDir_Error(t *testing.T) {
 }
 
 func TestKbsHTTPGet_Error(t *testing.T) {
-	_, err := kbsHTTPGet(context.Background(), "%%")
+	_, err := kbsHTTPGet(context.Background(), "%%", "")
 	assert.Error(t, err)
 }

cmd/agent/main.go

@@ -15,8 +15,8 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
-	"github.com/absmach/certs/sdk"
 	mglog "github.com/absmach/magistrala/logger"
 	"github.com/absmach/magistrala/pkg/prometheus"
 	"github.com/caarlos0/env/v11"
@@ -39,6 +39,7 @@ import (
 	runnerclient "github.com/ultravioletrs/cocos/pkg/clients/grpc/runner"
 	"github.com/ultravioletrs/cocos/pkg/ingress"
 	"golang.org/x/sync/errgroup"
+	"google.golang.org/grpc/metadata"
 )
 
 const (
@@ -51,9 +52,7 @@ type config struct {
 	LogLevel                 string `env:"AGENT_LOG_LEVEL"              envDefault:"debug"`
 	Vmpl                     int    `env:"AGENT_VMPL"                   envDefault:"2"`
 	AgentGrpcHost            string `env:"AGENT_GRPC_HOST"              envDefault:"0.0.0.0"`
-	CAUrl                    string `env:"AGENT_CVM_CA_URL"             envDefault:""`
 	CVMId                    string `env:"AGENT_CVM_ID"                 envDefault:""`
-	CertsToken               string `env:"AGENT_CERTS_TOKEN"            envDefault:""`
 	AgentMaaURL              string `env:"AGENT_MAA_URL"                envDefault:"https://sharedeus2.eus2.attest.azure.net"`
 	AgentOSBuild             string `env:"AGENT_OS_BUILD"               envDefault:"UVC"`
 	AgentOSDistro            string `env:"AGENT_OS_DISTRO"              envDefault:"UVC"`
@@ -170,6 +169,9 @@ func main() {
 		}
 		// Don't defer close here as we want to keep the connection open
 
+		if cfg.CVMId != "" {
+			ctx = metadata.AppendToOutgoingContext(ctx, "job-id", cfg.CVMId)
+		}
 		pc, err := newClient.Process(ctx)
 		if err != nil {
 			grpcClient.Close()
@@ -211,13 +213,7 @@ func main() {
 	var certProvider atls.CertificateProvider
 	if ccPlatform != attestation.NoCC {
 		logger.Info(fmt.Sprintf("Initializing aTLS for platform %v with attestation service at %s", ccPlatform, cfg.AttestationServiceSocket))
-		var certsSDK sdk.SDK
-		if cfg.CAUrl != "" {
-			certsSDK = sdk.NewSDK(sdk.Config{
-				CertsURL: cfg.CAUrl,
-			})
-		}
-		certProvider, err = atls.NewProvider(attClient, ccPlatform, cfg.CertsToken, cfg.CVMId, certsSDK)
+		certProvider, err = atls.NewProvider(attClient, ccPlatform)
 		if err != nil {
 			logger.Error(fmt.Sprintf("failed to create certificate provider for aTLS: %s. Continuing without attested TLS.", err))
 		} else {
@@ -236,7 +232,11 @@ func main() {
 	}
 	ingressProxy := ingress.NewProxyServer(logger, backendURL, certProvider)
 
-	pc, err := cvmsClient.Process(ctx)
+	agentCtx := ctx
+	if cfg.CVMId != "" {
+		agentCtx = metadata.AppendToOutgoingContext(ctx, "job-id", cfg.CVMId)
+	}
+	pc, err := cvmsClient.Process(agentCtx)
 	if err != nil {
 		logger.Error(fmt.Sprintf("failed to connect to cvm server: %s", err))
 		exitCode = 1
@@ -266,7 +266,15 @@ func main() {
 	})
 
 	g.Go(func() error {
-		return mc.Process(ctx, cancel)
+		for {
+			if err := mc.Process(ctx, cancel); err != nil {
+				if ctx.Err() != nil {
+					return ctx.Err()
+				}
+				logger.Warn("CVMS client process connection loop exited, retrying in 2s...", "error", err)
+				time.Sleep(2 * time.Second)
+			}
+		}
 	})
 
 	attest, certSerialNumber, err := attestationFromCert(ctx, cvmGrpcConfig.ClientCert, svc, ccPlatform)

cmd/attestation-service/main.go

@@ -409,7 +409,7 @@ func convertPlatformType(pt attestationpb.PlatformType) attestation.PlatformType
 	}
 }
 
-func (s *service) GetAzureToken(ctx context.Context, req *attestationpb.AzureTokenRequest) (*attestationpb.AzureTokenResponse, error) {
+func (s *service) FetchAzureToken(ctx context.Context, req *attestationpb.AzureTokenRequest) (*attestationpb.AzureTokenResponse, error) {
 	var nonce [32]byte
 	copy(nonce[:], req.Nonce)
 	token, err := s.provider.AzureAttestationToken(nonce[:])
@@ -418,3 +418,14 @@ func (s *service) GetAzureToken(ctx context.Context, req *attestationpb.AzureTok
 	}
 	return &attestationpb.AzureTokenResponse{Token: token}, nil
 }
+
+func (s *service) FetchKbsToken(ctx context.Context, req *attestationpb.KbsTokenRequest) (*attestationpb.KbsTokenResponse, error) {
+	s.logger.Info("[ATTESTATION-SERVICE] FetchKbsToken RPC called")
+	token, err := s.provider.KbsToken()
+	if err != nil {
+		s.logger.Error(fmt.Sprintf("[ATTESTATION-SERVICE] FetchKbsToken failed: %v", err))
+		return nil, err
+	}
+	s.logger.Info(fmt.Sprintf("[ATTESTATION-SERVICE] FetchKbsToken succeeded, token len: %d", len(token)))
+	return &attestationpb.KbsTokenResponse{Token: token}, nil
+}

cmd/ingress-proxy/main.go

@@ -11,7 +11,6 @@ import (
 	"os/signal"
 	"syscall"
 
-	"github.com/absmach/certs/sdk"
 	"github.com/caarlos0/env/v11"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -36,9 +35,7 @@ type config struct {
 	Backend  string `env:"COCOS_INGRESS_BACKEND"           envDefault:"http://localhost:7001"`
 
 	// ATLS Config
-	CAUrl         string `env:"AGENT_CVM_CA_URL"             envDefault:""`
 	CVMId         string `env:"AGENT_CVM_ID"                 envDefault:""`
-	CertsToken    string `env:"AGENT_CERTS_TOKEN"            envDefault:""`
 	AgentMaaURL   string `env:"AGENT_MAA_URL"                envDefault:"https://sharedeus2.eus2.attest.azure.net"`
 	AgentOSBuild  string `env:"AGENT_OS_BUILD"               envDefault:"UVC"`
 	AgentOSDistro string `env:"AGENT_OS_DISTRO"              envDefault:"UVC"`
@@ -143,13 +140,7 @@ func run(cfg config) error {
 		}
 		defer attClient.Close()
 
-		var certsSDK sdk.SDK
-		if cfg.CAUrl != "" {
-			certsSDK = sdk.NewSDK(sdk.Config{
-				CertsURL: cfg.CAUrl,
-			})
-		}
-		certProvider, err = atls.NewProvider(attClient, ccPlatform, cfg.CertsToken, cfg.CVMId, certsSDK)
+		certProvider, err = atls.NewProvider(attClient, ccPlatform)
 		if err != nil {
 			return fmt.Errorf("failed to create certificate provider: %w", err)
 		}

cmd/log-forwarder/main.go

@@ -20,6 +20,7 @@ import (
 	cvmsgrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc/cvm"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
 )
 
 const (
@@ -30,6 +31,7 @@ const (
 
 type config struct {
 	LogLevel string `env:"LOG_FORWARDER_LOG_LEVEL" envAlternate:"AGENT_LOG_LEVEL" envDefault:"debug"`
+	CVMId    string `env:"AGENT_CVM_ID"                 envDefault:""`
 }
 
 func main() {
@@ -100,6 +102,9 @@ func main() {
 	defer cvmClient.Close()
 
 	// Create stream to Manager
+	if cfg.CVMId != "" {
+		ctx = metadata.AppendToOutgoingContext(ctx, "job-id", cfg.CVMId, "connection-type", "log-forwarder")
+	}
 	stream, err := cvmsClient.Process(ctx)
 	if err != nil {
 		logger.Error(fmt.Sprintf("failed to create stream to manager: %s", err))
@@ -122,12 +127,23 @@ func main() {
 			case msg := <-logQueue:
 				if err := stream.Send(msg); err != nil {
 					logger.Error(fmt.Sprintf("failed to send log to manager: %s", err))
-					// Reconnect logic would go here
+					return err
 				}
 			}
 		}
 	})
 
+	// Stream Receiver Goroutine
+	g.Go(func() error {
+		for {
+			_, err := stream.Recv()
+			if err != nil {
+				logger.Error(fmt.Sprintf("stream connection lost: %s", err))
+				return err
+			}
+		}
+	})
+
 	g.Go(func() error {
 		ch := make(chan os.Signal, 1)
 		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)

go.mod

@@ -48,7 +48,6 @@ require (
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.55.0 // indirect
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.55.0 // indirect
 	github.com/Microsoft/go-winio v0.6.2 // indirect
-	github.com/absmach/supermq v0.19.2-0.20260317185610-fade98b84ee4 // indirect
 	github.com/cenkalti/backoff/v5 v5.0.3 // indirect
 	github.com/cncf/xds/go v0.0.0-20251210132809-ee656c7534f5 // indirect
 	github.com/containerd/errdefs v1.0.0 // indirect
@@ -103,11 +102,9 @@ require (
 	golang.org/x/time v0.15.0 // indirect
 	google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 // indirect
 	gotest.tools/v3 v3.5.1 // indirect
-	moul.io/http2curl v1.0.0 // indirect
 )
 
 require (
-	github.com/absmach/certs v0.18.5
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect

go.sum

@@ -34,12 +34,8 @@ github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERo
 github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
 github.com/VividCortex/gohistogram v1.0.0 h1:6+hBz+qvs0JOrrNhhmR7lFxo5sINxBCGXrdtl/UvroE=
 github.com/VividCortex/gohistogram v1.0.0/go.mod h1:Pf5mBqqDxYaXu3hDrrU+w6nw50o/4+TcAqDqk/vUH7g=
-github.com/absmach/certs v0.18.5 h1:eYlvitou+LoDtt7ETVLTp6d/1xCejGL3EmVOg+rHGTU=
-github.com/absmach/certs v0.18.5/go.mod h1:31dtVe1VYF16W+IvjAE/uPAIz4f3uLHgh+moBezjqIc=
 github.com/absmach/magistrala v0.20.0 h1:3AQ0C2AMoOCc1UuJLhPNJLMrNRLZoN0ibSOERqEkM98=
 github.com/absmach/magistrala v0.20.0/go.mod h1:lnuO4fSngMiRYyNYL4yz5UP8DX3bbXRm87b2KHFGwJU=
-github.com/absmach/supermq v0.19.2-0.20260317185610-fade98b84ee4 h1:533pRc6R7perWDqJuZq+ofBQfYfmyj7n49V4LFY4zpo=
-github.com/absmach/supermq v0.19.2-0.20260317185610-fade98b84ee4/go.mod h1:xDAX/O3VcOsHWCx2fk85VD7FI17hAUOvoOhho7DA7g0=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/caarlos0/env/v10 v10.0.0 h1:yIHUBZGsyqCnpTkbjk8asUlx6RFhhEs+h7TOBdgdzXA=
@@ -94,8 +90,6 @@ github.com/fxamacker/cbor/v2 v2.9.0 h1:NpKPmjDBgUfBms6tr6JZkTHtfFGcMKsw3eGcmD/sa
 github.com/fxamacker/cbor/v2 v2.9.0/go.mod h1:vM4b+DJCtHn+zz7h3FFp/hDAI9WNWCsZj23V5ytsSxQ=
 github.com/go-chi/chi/v5 v5.2.5 h1:Eg4myHZBjyvJmAFjFvWgrqDTXFyOzjj7YIm3L3mu6Ug=
 github.com/go-chi/chi/v5 v5.2.5/go.mod h1:X7Gx4mteadT3eDOMTsXzmI4/rwUpOwBHLpAfupzFJP0=
-github.com/go-gorp/gorp/v3 v3.1.0 h1:ItKF/Vbuj31dmV4jxA1qblpSwkl9g1typ24xoe70IGs=
-github.com/go-gorp/gorp/v3 v3.1.0/go.mod h1:dLEjIyyRNiXvNZ8PSmzpt1GsWAUK8kjVhEpjH8TixEw=
 github.com/go-jose/go-jose/v3 v3.0.4 h1:Wp5HA7bLQcKnf6YYao/4kpRpVMp/yf6+pJKV8WFSaNY=
 github.com/go-jose/go-jose/v3 v3.0.4/go.mod h1:5b+7YgP7ZICgJDBdfjZaIt+H/9L9T/YQrVfLAMboGkQ=
 github.com/go-jose/go-jose/v4 v4.1.4 h1:moDMcTHmvE6Groj34emNPLs/qtYXRVcd6S7NHbHz3kA=
@@ -156,24 +150,10 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.14 h1:yh8ncqsbUY4shRD5dA
 github.com/googleapis/enterprise-certificate-proxy v0.3.14/go.mod h1:vqVt9yG9480NtzREnTlmGSBmFrA+bzb0yl0TxoBQXOg=
 github.com/googleapis/gax-go/v2 v2.21.0 h1:h45NjjzEO3faG9Lg/cFrBh2PgegVVgzqKzuZl/wMbiI=
 github.com/googleapis/gax-go/v2 v2.21.0/go.mod h1:But/NJU6TnZsrLai/xBAQLLz+Hc7fHZJt/hsCz3Fih4=
-github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g=
-github.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0 h1:HWRh5R2+9EifMyIHV7ZV+MIZqgz+PMpZ14Jynv3O2Zs=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.28.0/go.mod h1:JfhWUomR1baixubs02l85lZYYOm7LV6om4ceouMv45c=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
-github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
-github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
-github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
-github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
-github.com/jackc/pgx/v5 v5.9.1 h1:uwrxJXBnx76nyISkhr33kQLlUqjv7et7b9FjCen/tdc=
-github.com/jackc/pgx/v5 v5.9.1/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4=
-github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
-github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
-github.com/jmoiron/sqlx v1.4.0 h1:1PLqN7S1UYp5t4SrVVnt4nUVNemrDAtxlulVe+Qgm3o=
-github.com/jmoiron/sqlx v1.4.0/go.mod h1:ZrZ7UsYB/weZdl2Bxg6jCRO9c3YHl8r3ahlKmRT4JLY=
-github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo=
-github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
 github.com/klauspost/compress v1.18.5 h1:/h1gH5Ce+VWNLSWqPzOVn6XBO+vJbCNGvjoaGBFW2IE=
 github.com/klauspost/compress v1.18.5/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ=
 github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
@@ -231,8 +211,6 @@ github.com/prometheus/procfs v0.19.2 h1:zUMhqEW66Ex7OXIiDkll3tl9a1ZdilUOd/F6ZXw4
 github.com/prometheus/procfs v0.19.2/go.mod h1:M0aotyiemPhBCM0z5w87kL22CxfcH05ZpYlu+b4J7mw=
 github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ=
 github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc=
-github.com/rubenv/sql-migrate v1.8.1 h1:EPNwCvjAowHI3TnZ+4fQu3a915OpnQoPAjTXCGOy2U0=
-github.com/rubenv/sql-migrate v1.8.1/go.mod h1:BTIKBORjzyxZDS6dzoiw6eAFYJ1iNlGAtjn4LGeVjS8=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/sammyoina/sev-snp-measure-go v0.0.0-20241202151803-ef189f0ff825 h1:SqNaL9udBIc026SGNEuEuiVL0/hw9fXxM5qrFhWGkdE=
 github.com/sammyoina/sev-snp-measure-go v0.0.0-20241202151803-ef189f0ff825/go.mod h1:dEkBe8JnxU5itNjZDEQINFd7f7l4DtjfqRuzPQcit4w=
@@ -240,10 +218,6 @@ github.com/segmentio/asm v1.2.1 h1:DTNbBqs57ioxAD4PrArqftgypG4/qNpXoJx8TVXxPR0=
 github.com/segmentio/asm v1.2.1/go.mod h1:BqMnlJP91P8d+4ibuonYZw9mfnzI9HfxselHZr5aAcs=
 github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ=
 github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
-github.com/smarty/assertions v1.16.0 h1:EvHNkdRA4QHMrn75NZSoUQ/mAUXAYWfatfB01yTCzfY=
-github.com/smarty/assertions v1.16.0/go.mod h1:duaaFdCS0K9dnoM50iyek/eYINOZ64gbh1Xlf6LG7AI=
-github.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY=
-github.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60=
 github.com/spf13/cast v1.10.0 h1:h2x0u2shc1QuLHfxi+cTJvs30+ZAHOGRic8uyGTDWxY=
 github.com/spf13/cast v1.10.0/go.mod h1:jNfB8QC9IA6ZuY2ZjDp0KtFO2LZZlg4S/7bzP6qqeHo=
 github.com/spf13/cobra v1.10.2 h1:DMTTonx5m65Ic0GOoRY2c16WCbHxOOw6xxezuLaBpcU=
@@ -382,5 +356,3 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gotest.tools/v3 v3.5.1 h1:EENdUnS3pdur5nybKYIh2Vfgc8IUNBjxDPSjtiJcOzU=
 gotest.tools/v3 v3.5.1/go.mod h1:isy3WKz7GK6uNw/sbHzfKBLvlvXwUyV06n6brMxxopU=
-moul.io/http2curl v1.0.0 h1:6XwpyZOYsgZJrU8exnG87ncVkU1FVCcTRpwzOkTDUi8=
-moul.io/http2curl v1.0.0/go.mod h1:f6cULg+e4Md/oW1cYmwW4IWQOVl2lGbmCNGOHvzX2kE=

hal/disk/configs/cocos_defconfig

@@ -38,9 +38,9 @@ BR2_ROOTFS_POST_IMAGE_SCRIPT="$(BR2_EXTERNAL_COCOS_PATH)/board/cocos/post-image.
 BR2_ROOTFS_POST_SCRIPT_ARGS="$(BR2_DEFCONFIG)"
 
 # Linux headers same as kernel
-BR2_PACKAGE_HOST_LINUX_HEADERS_CUSTOM_6_11=y
+BR2_PACKAGE_HOST_LINUX_HEADERS_CUSTOM_7_0=y
 BR2_TOOLCHAIN_HEADERS_LATEST=y
-BR2_TOOLCHAIN_HEADERS_AT_LEAST="6.11-rc7"
+BR2_TOOLCHAIN_HEADERS_AT_LEAST="7.0"
 
 # Kernel
 BR2_LINUX_KERNEL=y

hal/linux/board/cocos/linux.config

@@ -1,8 +1,18 @@
 CONFIG_SYSVIPC=y
 CONFIG_CGROUPS=y
+CONFIG_CGROUP_DEVICE=y
+CONFIG_CGROUP_PIDS=y
+CONFIG_MEMCG=y
+CONFIG_BLK_CGROUP=y
 CONFIG_MODULES=y
 CONFIG_MODULE_UNLOAD=y
 CONFIG_SMP=y
+CONFIG_NAMESPACES=y
+CONFIG_UTS_NS=y
+CONFIG_IPC_NS=y
+CONFIG_USER_NS=y
+CONFIG_PID_NS=y
+CONFIG_NET_NS=y
 CONFIG_HYPERVISOR_GUEST=y
 CONFIG_PARAVIRT=y
 CONFIG_NET=y
@@ -14,12 +24,17 @@ CONFIG_PCI=y
 CONFIG_DEVTMPFS=y
 CONFIG_DEVTMPFS_MOUNT=y
 CONFIG_VIRTIO_BLK=y
+CONFIG_BLK_DEV_LOOP=y
+CONFIG_BLK_DEV_LOOP_MIN_COUNT=8
 CONFIG_BLK_DEV_SD=y
 CONFIG_SCSI_VIRTIO=y
 CONFIG_ATA=y
 CONFIG_ATA_PIIX=y
 CONFIG_NETDEVICES=y
 CONFIG_VIRTIO_NET=y
+CONFIG_VETH=y
+CONFIG_BRIDGE=y
+CONFIG_BRIDGE_NETFILTER=y
 CONFIG_NE2K_PCI=y
 CONFIG_8139CP=y
 # CONFIG_WLAN is not set
@@ -34,15 +49,43 @@ CONFIG_VIRTIO_INPUT=y
 CONFIG_VIRTIO_MMIO=y
 CONFIG_VIRTIO_MMIO_CMDLINE_DEVICES=y
 CONFIG_EXT4_FS=y
+CONFIG_OVERLAY_FS=y
 CONFIG_AUTOFS4_FS=y
 CONFIG_TMPFS=y
 CONFIG_TMPFS_POSIX_ACL=y
+CONFIG_FHANDLE=y
+CONFIG_INOTIFY_USER=y
+CONFIG_SIGNALFD=y
+CONFIG_TIMERFD=y
+CONFIG_EPOLL=y
+CONFIG_POSIX_MQUEUE=y
+CONFIG_POSIX_MQUEUE_SYSCTL=y
 CONFIG_UNWINDER_FRAME_POINTER=y
 CONFIG_VSOCKETS=y
 CONFIG_VIRTIO_VSOCKETS=y
+CONFIG_NETFILTER=y
+CONFIG_NETFILTER_ADVANCED=y
+CONFIG_NETFILTER_XTABLES_LEGACY=y
+CONFIG_NF_CONNTRACK=y
+CONFIG_NF_CONNTRACK_MARK=y
+CONFIG_NF_NAT=y
+CONFIG_NF_NAT_MASQUERADE=y
+CONFIG_NETFILTER_XTABLES=y
+CONFIG_NETFILTER_XT_NAT=y
+CONFIG_NETFILTER_XT_TARGET_MASQUERADE=y
+CONFIG_NF_DEFRAG_IPV4=y
+CONFIG_IP_NF_IPTABLES_LEGACY=y
+CONFIG_IP_NF_IPTABLES=y
+CONFIG_IP_NF_FILTER=y
+CONFIG_IP_NF_NAT=y
+CONFIG_IP_NF_TARGET_MASQUERADE=y
+CONFIG_NETFILTER_XT_MATCH_ADDRTYPE=y
+CONFIG_NETFILTER_XT_MATCH_CONNTRACK=y
 CONFIG_NF_TABLES=y
 CONFIG_BPF_SYSCALL=y
 CONFIG_CGROUP_BPF=y
+CONFIG_SECCOMP=y
+CONFIG_SECCOMP_FILTER=y
 
 ###
 # AMD SEV-SNP

hal/linux/configs/cocos_defconfig

@@ -20,14 +20,17 @@ BR2_PACKAGE_9PFS=y
 # Image
 BR2_ROOTFS_POST_BUILD_SCRIPT="$(BR2_EXTERNAL_COCOS_PATH)/board/cocos/post-build.sh"
  
+# Patches for existing Buildroot packages
+BR2_GLOBAL_PATCH_DIR="$(BR2_EXTERNAL_COCOS_PATH)/patches"
+
 # Image
 BR2_ROOTFS_POST_IMAGE_SCRIPT="$(BR2_EXTERNAL_COCOS_PATH)/board/cocos/post-image.sh"
 BR2_ROOTFS_POST_SCRIPT_ARGS="$(BR2_DEFCONFIG)"
 
 # Linux headers same as kernel
-BR2_PACKAGE_HOST_LINUX_HEADERS_CUSTOM_6_11=y
+BR2_PACKAGE_HOST_LINUX_HEADERS_CUSTOM_7_0=y
 BR2_TOOLCHAIN_HEADERS_LATEST=y
-BR2_TOOLCHAIN_HEADERS_AT_LEAST="6.11-rc7"
+BR2_TOOLCHAIN_HEADERS_AT_LEAST="7.0"
 
 # Kernel
 BR2_LINUX_KERNEL=y

hal/linux/package/agent/agent.mk

@@ -22,6 +22,12 @@ define AGENT_INSTALL_INIT_SYSTEMD
 	$(INSTALL) -D -m 0640 $(@D)/init/systemd/cocos-agent.service $(TARGET_DIR)/usr/lib/systemd/system/cocos-agent.service
 	$(INSTALL) -D -m 0750 $(@D)/init/systemd/agent_setup.sh $(TARGET_DIR)/cocos_init/agent_setup.sh
 	$(INSTALL) -D -m 0750 $(@D)/init/systemd/agent_start_script.sh $(TARGET_DIR)/cocos_init/agent_start_script.sh
+	mkdir -p $(TARGET_DIR)/usr/lib/systemd/system/multi-user.target.wants
+	ln -sf ../docker.service $(TARGET_DIR)/usr/lib/systemd/system/multi-user.target.wants/docker.service
+	ln -sf ../log-forwarder.service $(TARGET_DIR)/usr/lib/systemd/system/multi-user.target.wants/log-forwarder.service
+	ln -sf ../computation-runner.service $(TARGET_DIR)/usr/lib/systemd/system/multi-user.target.wants/computation-runner.service
+	ln -sf ../egress-proxy.service $(TARGET_DIR)/usr/lib/systemd/system/multi-user.target.wants/egress-proxy.service
+	ln -sf ../cocos-agent.service $(TARGET_DIR)/usr/lib/systemd/system/multi-user.target.wants/cocos-agent.service
 endef
 
 $(eval $(generic-package))

hal/linux/patches/gcc/14.3.0/0007-libsanitizer-guard-linux-scc.h-for-linux-7.1.patch

@@ -0,0 +1,33 @@
+--- a/libsanitizer/sanitizer_common/sanitizer_platform_limits_posix.cpp
++++ b/libsanitizer/sanitizer_common/sanitizer_platform_limits_posix.cpp
+@@ -151,7 +151,15 @@
+ #include <linux/lp.h>
+ #include <linux/mroute.h>
+ #include <linux/mroute6.h>
+-#include <linux/scc.h>
++#if defined(__has_include)
++# if __has_include(<linux/scc.h>)
++#  include <linux/scc.h>
++#  define SANITIZER_HAS_LINUX_SCC_H 1
++# endif
++#else
++# include <linux/scc.h>
++# define SANITIZER_HAS_LINUX_SCC_H 1
++#endif
+ #include <linux/serial.h>
+ #include <sys/msg.h>
+ #include <sys/ipc.h>
+@@ -517,8 +525,13 @@
+   unsigned struct_kbsentry_sz = sizeof(struct kbsentry);
+   unsigned struct_mtconfiginfo_sz = sizeof(struct mtconfiginfo);
+   unsigned struct_nr_parms_struct_sz = sizeof(struct nr_parms_struct);
++#if SANITIZER_HAS_LINUX_SCC_H
+   unsigned struct_scc_modem_sz = sizeof(struct scc_modem);
+   unsigned struct_scc_stat_sz = sizeof(struct scc_stat);
++#else
++  unsigned struct_scc_modem_sz = 0;
++  unsigned struct_scc_stat_sz = 0;
++#endif
+   unsigned struct_serial_multiport_struct_sz
+       = sizeof(struct serial_multiport_struct);
+   unsigned struct_serial_struct_sz = sizeof(struct serial_struct);

init/systemd/agent_setup.sh

@@ -1,6 +1,24 @@
 #!/bin/sh
 
 WORK_DIR="/cocos"
+DOCKER_OVERRIDE_DIR="/etc/systemd/system/docker.service.d"
+DOCKER_OVERRIDE_FILE="$DOCKER_OVERRIDE_DIR/override.conf"
+
+# Docker runs on a writable bind mount while the rootfs itself stays ephemeral,
+# so force RAM-disk mode to avoid pivot_root failures inside containers.
+mkdir -p "$DOCKER_OVERRIDE_DIR"
+if ! grep -qs '^Environment=DOCKER_RAMDISK=true$' "$DOCKER_OVERRIDE_FILE"; then
+    tee "$DOCKER_OVERRIDE_FILE" > /dev/null <<EOF
+[Service]
+Environment=DOCKER_RAMDISK=true
+EOF
+    systemctl daemon-reload
+fi
+
+# Docker is required by the computation runner. Start the runtime stack here so
+# later units don't depend on passive boot enablement or service ordering races.
+systemctl start containerd.service
+systemctl start docker.service
 
 # IFACES are all network interfaces excluding lo (LOOPBACK) and sit interfaces 
 IFACES=$(ip link show | grep -vE 'LOOPBACK|sit*' | awk -F': ' '{print $2}')

init/systemd/cocos-agent.service

@@ -2,7 +2,6 @@
 Description=Cocos AI agent
 After=network.target attestation-service.service log-forwarder.service computation-runner.service egress-proxy.service
 Requires=log-forwarder.service computation-runner.service egress-proxy.service
-Before=docker.service
 
 [Service]
 WorkingDirectory=/cocos
@@ -17,4 +16,4 @@ Restart=always
 RestartSec=5s
 
 [Install]
-WantedBy=default.target
+WantedBy=multi-user.target

init/systemd/computation-runner.service

@@ -1,8 +1,9 @@
 [Unit]
 Description=Cocos AI Computation Runner
-After=network.target log-forwarder.service
+After=network.target log-forwarder.service docker.service
 Before=cocos-agent.service
 Requires=log-forwarder.service
+Wants=docker.service
 
 [Service]
 WorkingDirectory=/cocos
@@ -17,4 +18,4 @@ Restart=always
 RestartSec=5s
 
 [Install]
-WantedBy=default.target
+WantedBy=multi-user.target

init/systemd/log-forwarder.service

@@ -14,4 +14,4 @@ Restart=always
 RestartSec=5s
 
 [Install]
-WantedBy=default.target
+WantedBy=multi-user.target

internal/proto/attestation-agent/attestation-agent.pb.go

@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.36.11
-// 	protoc        v6.33.1
+// 	protoc        v7.35.1
 // source: internal/proto/attestation-agent/attestation-agent.proto
 
 package attestation_agent

internal/proto/attestation-agent/attestation-agent_grpc.pb.go

@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
 // - protoc-gen-go-grpc v1.6.0
-// - protoc             v6.33.1
+// - protoc             v7.35.1
 // source: internal/proto/attestation-agent/attestation-agent.proto
 
 package attestation_agent

internal/proto/attestation/v1/attestation.pb.go

@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.36.11
-// 	protoc        v6.33.1
+// 	protoc        v7.35.1
 // source: internal/proto/attestation/v1/attestation.proto
 
 package attestation
@@ -82,6 +82,86 @@ func (PlatformType) EnumDescriptor() ([]byte, []int) {
 	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{0}
 }
 
+type KbsTokenRequest struct {
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *KbsTokenRequest) Reset() {
+	*x = KbsTokenRequest{}
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[0]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *KbsTokenRequest) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*KbsTokenRequest) ProtoMessage() {}
+
+func (x *KbsTokenRequest) ProtoReflect() protoreflect.Message {
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[0]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use KbsTokenRequest.ProtoReflect.Descriptor instead.
+func (*KbsTokenRequest) Descriptor() ([]byte, []int) {
+	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{0}
+}
+
+type KbsTokenResponse struct {
+	state         protoimpl.MessageState `protogen:"open.v1"`
+	Token         []byte                 `protobuf:"bytes,1,opt,name=token,proto3" json:"token,omitempty"`
+	unknownFields protoimpl.UnknownFields
+	sizeCache     protoimpl.SizeCache
+}
+
+func (x *KbsTokenResponse) Reset() {
+	*x = KbsTokenResponse{}
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[1]
+	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+	ms.StoreMessageInfo(mi)
+}
+
+func (x *KbsTokenResponse) String() string {
+	return protoimpl.X.MessageStringOf(x)
+}
+
+func (*KbsTokenResponse) ProtoMessage() {}
+
+func (x *KbsTokenResponse) ProtoReflect() protoreflect.Message {
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[1]
+	if x != nil {
+		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
+		if ms.LoadMessageInfo() == nil {
+			ms.StoreMessageInfo(mi)
+		}
+		return ms
+	}
+	return mi.MessageOf(x)
+}
+
+// Deprecated: Use KbsTokenResponse.ProtoReflect.Descriptor instead.
+func (*KbsTokenResponse) Descriptor() ([]byte, []int) {
+	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{1}
+}
+
+func (x *KbsTokenResponse) GetToken() []byte {
+	if x != nil {
+		return x.Token
+	}
+	return nil
+}
+
 type AttestationRequest struct {
 	state         protoimpl.MessageState `protogen:"open.v1"`
 	ReportData    []byte                 `protobuf:"bytes,1,opt,name=report_data,json=reportData,proto3" json:"report_data,omitempty"` // 64 bytes for SNP/TDX
@@ -93,7 +173,7 @@ type AttestationRequest struct {
 
 func (x *AttestationRequest) Reset() {
 	*x = AttestationRequest{}
-	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[0]
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[2]
 	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 	ms.StoreMessageInfo(mi)
 }
@@ -105,7 +185,7 @@ func (x *AttestationRequest) String() string {
 func (*AttestationRequest) ProtoMessage() {}
 
 func (x *AttestationRequest) ProtoReflect() protoreflect.Message {
-	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[0]
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[2]
 	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -118,7 +198,7 @@ func (x *AttestationRequest) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use AttestationRequest.ProtoReflect.Descriptor instead.
 func (*AttestationRequest) Descriptor() ([]byte, []int) {
-	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{0}
+	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{2}
 }
 
 func (x *AttestationRequest) GetReportData() []byte {
@@ -151,7 +231,7 @@ type AttestationResponse struct {
 
 func (x *AttestationResponse) Reset() {
 	*x = AttestationResponse{}
-	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[1]
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[3]
 	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 	ms.StoreMessageInfo(mi)
 }
@@ -163,7 +243,7 @@ func (x *AttestationResponse) String() string {
 func (*AttestationResponse) ProtoMessage() {}
 
 func (x *AttestationResponse) ProtoReflect() protoreflect.Message {
-	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[1]
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[3]
 	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -176,7 +256,7 @@ func (x *AttestationResponse) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use AttestationResponse.ProtoReflect.Descriptor instead.
 func (*AttestationResponse) Descriptor() ([]byte, []int) {
-	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{1}
+	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{3}
 }
 
 func (x *AttestationResponse) GetEatToken() []byte {
@@ -195,7 +275,7 @@ type RawEvidenceResponse struct {
 
 func (x *RawEvidenceResponse) Reset() {
 	*x = RawEvidenceResponse{}
-	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[2]
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[4]
 	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 	ms.StoreMessageInfo(mi)
 }
@@ -207,7 +287,7 @@ func (x *RawEvidenceResponse) String() string {
 func (*RawEvidenceResponse) ProtoMessage() {}
 
 func (x *RawEvidenceResponse) ProtoReflect() protoreflect.Message {
-	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[2]
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[4]
 	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -220,7 +300,7 @@ func (x *RawEvidenceResponse) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use RawEvidenceResponse.ProtoReflect.Descriptor instead.
 func (*RawEvidenceResponse) Descriptor() ([]byte, []int) {
-	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{2}
+	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{4}
 }
 
 func (x *RawEvidenceResponse) GetEvidence() []byte {
@@ -239,7 +319,7 @@ type AzureTokenRequest struct {
 
 func (x *AzureTokenRequest) Reset() {
 	*x = AzureTokenRequest{}
-	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[3]
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[5]
 	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 	ms.StoreMessageInfo(mi)
 }
@@ -251,7 +331,7 @@ func (x *AzureTokenRequest) String() string {
 func (*AzureTokenRequest) ProtoMessage() {}
 
 func (x *AzureTokenRequest) ProtoReflect() protoreflect.Message {
-	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[3]
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[5]
 	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -264,7 +344,7 @@ func (x *AzureTokenRequest) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use AzureTokenRequest.ProtoReflect.Descriptor instead.
 func (*AzureTokenRequest) Descriptor() ([]byte, []int) {
-	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{3}
+	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{5}
 }
 
 func (x *AzureTokenRequest) GetNonce() []byte {
@@ -283,7 +363,7 @@ type AzureTokenResponse struct {
 
 func (x *AzureTokenResponse) Reset() {
 	*x = AzureTokenResponse{}
-	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[4]
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[6]
 	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 	ms.StoreMessageInfo(mi)
 }
@@ -295,7 +375,7 @@ func (x *AzureTokenResponse) String() string {
 func (*AzureTokenResponse) ProtoMessage() {}
 
 func (x *AzureTokenResponse) ProtoReflect() protoreflect.Message {
-	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[4]
+	mi := &file_internal_proto_attestation_v1_attestation_proto_msgTypes[6]
 	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -308,7 +388,7 @@ func (x *AzureTokenResponse) ProtoReflect() protoreflect.Message {
 
 // Deprecated: Use AzureTokenResponse.ProtoReflect.Descriptor instead.
 func (*AzureTokenResponse) Descriptor() ([]byte, []int) {
-	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{4}
+	return file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP(), []int{6}
 }
 
 func (x *AzureTokenResponse) GetToken() []byte {
@@ -322,7 +402,10 @@ var File_internal_proto_attestation_v1_attestation_proto protoreflect.FileDescri
 
 const file_internal_proto_attestation_v1_attestation_proto_rawDesc = "" +
 	"\n" +
-	"/internal/proto/attestation/v1/attestation.proto\x12\x0eattestation.v1\"\x8e\x01\n" +
+	"/internal/proto/attestation/v1/attestation.proto\x12\x0eattestation.v1\"\x11\n" +
+	"\x0fKbsTokenRequest\"(\n" +
+	"\x10KbsTokenResponse\x12\x14\n" +
+	"\x05token\x18\x01 \x01(\fR\x05token\"\x8e\x01\n" +
 	"\x12AttestationRequest\x12\x1f\n" +
 	"\vreport_data\x18\x01 \x01(\fR\n" +
 	"reportData\x12\x14\n" +
@@ -343,11 +426,12 @@ const file_internal_proto_attestation_v1_attestation_proto_rawDesc = "" +
 	"\x12PLATFORM_TYPE_VTPM\x10\x03\x12\x1a\n" +
 	"\x16PLATFORM_TYPE_SNP_VTPM\x10\x04\x12\x17\n" +
 	"\x13PLATFORM_TYPE_AZURE\x10\x05\x12\x17\n" +
-	"\x13PLATFORM_TYPE_NO_CC\x10\x062\xa8\x02\n" +
+	"\x13PLATFORM_TYPE_NO_CC\x10\x062\xfc\x02\n" +
 	"\x12AttestationService\x12[\n" +
 	"\x10FetchAttestation\x12\".attestation.v1.AttestationRequest\x1a#.attestation.v1.AttestationResponse\x12[\n" +
 	"\x10FetchRawEvidence\x12\".attestation.v1.AttestationRequest\x1a#.attestation.v1.RawEvidenceResponse\x12X\n" +
-	"\x0fFetchAzureToken\x12!.attestation.v1.AzureTokenRequest\x1a\".attestation.v1.AzureTokenResponseBJZHgithub.com/ultravioletrs/cocos/internal/proto/attestation/v1;attestationb\x06proto3"
+	"\x0fFetchAzureToken\x12!.attestation.v1.AzureTokenRequest\x1a\".attestation.v1.AzureTokenResponse\x12R\n" +
+	"\rFetchKbsToken\x12\x1f.attestation.v1.KbsTokenRequest\x1a .attestation.v1.KbsTokenResponseBJZHgithub.com/ultravioletrs/cocos/internal/proto/attestation/v1;attestationb\x06proto3"
 
 var (
 	file_internal_proto_attestation_v1_attestation_proto_rawDescOnce sync.Once
@@ -362,25 +446,29 @@ func file_internal_proto_attestation_v1_attestation_proto_rawDescGZIP() []byte {
 }
 
 var file_internal_proto_attestation_v1_attestation_proto_enumTypes = make([]protoimpl.EnumInfo, 1)
-var file_internal_proto_attestation_v1_attestation_proto_msgTypes = make([]protoimpl.MessageInfo, 5)
+var file_internal_proto_attestation_v1_attestation_proto_msgTypes = make([]protoimpl.MessageInfo, 7)
 var file_internal_proto_attestation_v1_attestation_proto_goTypes = []any{
 	(PlatformType)(0),           // 0: attestation.v1.PlatformType
-	(*AttestationRequest)(nil),  // 1: attestation.v1.AttestationRequest
-	(*AttestationResponse)(nil), // 2: attestation.v1.AttestationResponse
-	(*RawEvidenceResponse)(nil), // 3: attestation.v1.RawEvidenceResponse
-	(*AzureTokenRequest)(nil),   // 4: attestation.v1.AzureTokenRequest
-	(*AzureTokenResponse)(nil),  // 5: attestation.v1.AzureTokenResponse
+	(*KbsTokenRequest)(nil),     // 1: attestation.v1.KbsTokenRequest
+	(*KbsTokenResponse)(nil),    // 2: attestation.v1.KbsTokenResponse
+	(*AttestationRequest)(nil),  // 3: attestation.v1.AttestationRequest
+	(*AttestationResponse)(nil), // 4: attestation.v1.AttestationResponse
+	(*RawEvidenceResponse)(nil), // 5: attestation.v1.RawEvidenceResponse
+	(*AzureTokenRequest)(nil),   // 6: attestation.v1.AzureTokenRequest
+	(*AzureTokenResponse)(nil),  // 7: attestation.v1.AzureTokenResponse
 }
 var file_internal_proto_attestation_v1_attestation_proto_depIdxs = []int32{
 	0, // 0: attestation.v1.AttestationRequest.platform_type:type_name -> attestation.v1.PlatformType
-	1, // 1: attestation.v1.AttestationService.FetchAttestation:input_type -> attestation.v1.AttestationRequest
-	1, // 2: attestation.v1.AttestationService.FetchRawEvidence:input_type -> attestation.v1.AttestationRequest
-	4, // 3: attestation.v1.AttestationService.FetchAzureToken:input_type -> attestation.v1.AzureTokenRequest
-	2, // 4: attestation.v1.AttestationService.FetchAttestation:output_type -> attestation.v1.AttestationResponse
-	3, // 5: attestation.v1.AttestationService.FetchRawEvidence:output_type -> attestation.v1.RawEvidenceResponse
-	5, // 6: attestation.v1.AttestationService.FetchAzureToken:output_type -> attestation.v1.AzureTokenResponse
-	4, // [4:7] is the sub-list for method output_type
-	1, // [1:4] is the sub-list for method input_type
+	3, // 1: attestation.v1.AttestationService.FetchAttestation:input_type -> attestation.v1.AttestationRequest
+	3, // 2: attestation.v1.AttestationService.FetchRawEvidence:input_type -> attestation.v1.AttestationRequest
+	6, // 3: attestation.v1.AttestationService.FetchAzureToken:input_type -> attestation.v1.AzureTokenRequest
+	1, // 4: attestation.v1.AttestationService.FetchKbsToken:input_type -> attestation.v1.KbsTokenRequest
+	4, // 5: attestation.v1.AttestationService.FetchAttestation:output_type -> attestation.v1.AttestationResponse
+	5, // 6: attestation.v1.AttestationService.FetchRawEvidence:output_type -> attestation.v1.RawEvidenceResponse
+	7, // 7: attestation.v1.AttestationService.FetchAzureToken:output_type -> attestation.v1.AzureTokenResponse
+	2, // 8: attestation.v1.AttestationService.FetchKbsToken:output_type -> attestation.v1.KbsTokenResponse
+	5, // [5:9] is the sub-list for method output_type
+	1, // [1:5] is the sub-list for method input_type
 	1, // [1:1] is the sub-list for extension type_name
 	1, // [1:1] is the sub-list for extension extendee
 	0, // [0:1] is the sub-list for field type_name
@@ -397,7 +485,7 @@ func file_internal_proto_attestation_v1_attestation_proto_init() {
 			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
 			RawDescriptor: unsafe.Slice(unsafe.StringData(file_internal_proto_attestation_v1_attestation_proto_rawDesc), len(file_internal_proto_attestation_v1_attestation_proto_rawDesc)),
 			NumEnums:      1,
-			NumMessages:   5,
+			NumMessages:   7,
 			NumExtensions: 0,
 			NumServices:   1,
 		},

internal/proto/attestation/v1/attestation.proto

@@ -8,6 +8,13 @@ service AttestationService {
   rpc FetchAttestation (AttestationRequest) returns (AttestationResponse);
   rpc FetchRawEvidence (AttestationRequest) returns (RawEvidenceResponse);
   rpc FetchAzureToken (AzureTokenRequest) returns (AzureTokenResponse);
+  rpc FetchKbsToken (KbsTokenRequest) returns (KbsTokenResponse);
+}
+
+message KbsTokenRequest {}
+
+message KbsTokenResponse {
+  bytes token = 1;
 }
 
 message AttestationRequest {

internal/proto/attestation/v1/attestation_grpc.pb.go

@@ -1,14 +1,13 @@
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
 // - protoc-gen-go-grpc v1.6.0
-// - protoc             v6.33.1
+// - protoc             v7.35.1
 // source: internal/proto/attestation/v1/attestation.proto
 
 package attestation
 
 import (
 	context "context"
-
 	grpc "google.golang.org/grpc"
 	codes "google.golang.org/grpc/codes"
 	status "google.golang.org/grpc/status"
@@ -23,6 +22,7 @@ const (
 	AttestationService_FetchAttestation_FullMethodName = "/attestation.v1.AttestationService/FetchAttestation"
 	AttestationService_FetchRawEvidence_FullMethodName = "/attestation.v1.AttestationService/FetchRawEvidence"
 	AttestationService_FetchAzureToken_FullMethodName  = "/attestation.v1.AttestationService/FetchAzureToken"
+	AttestationService_FetchKbsToken_FullMethodName    = "/attestation.v1.AttestationService/FetchKbsToken"
 )
 
 // AttestationServiceClient is the client API for AttestationService service.
@@ -32,6 +32,7 @@ type AttestationServiceClient interface {
 	FetchAttestation(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (*AttestationResponse, error)
 	FetchRawEvidence(ctx context.Context, in *AttestationRequest, opts ...grpc.CallOption) (*RawEvidenceResponse, error)
 	FetchAzureToken(ctx context.Context, in *AzureTokenRequest, opts ...grpc.CallOption) (*AzureTokenResponse, error)
+	FetchKbsToken(ctx context.Context, in *KbsTokenRequest, opts ...grpc.CallOption) (*KbsTokenResponse, error)
 }
 
 type attestationServiceClient struct {
@@ -72,6 +73,16 @@ func (c *attestationServiceClient) FetchAzureToken(ctx context.Context, in *Azur
 	return out, nil
 }
 
+func (c *attestationServiceClient) FetchKbsToken(ctx context.Context, in *KbsTokenRequest, opts ...grpc.CallOption) (*KbsTokenResponse, error) {
+	cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
+	out := new(KbsTokenResponse)
+	err := c.cc.Invoke(ctx, AttestationService_FetchKbsToken_FullMethodName, in, out, cOpts...)
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
 // AttestationServiceServer is the server API for AttestationService service.
 // All implementations must embed UnimplementedAttestationServiceServer
 // for forward compatibility.
@@ -79,6 +90,7 @@ type AttestationServiceServer interface {
 	FetchAttestation(context.Context, *AttestationRequest) (*AttestationResponse, error)
 	FetchRawEvidence(context.Context, *AttestationRequest) (*RawEvidenceResponse, error)
 	FetchAzureToken(context.Context, *AzureTokenRequest) (*AzureTokenResponse, error)
+	FetchKbsToken(context.Context, *KbsTokenRequest) (*KbsTokenResponse, error)
 	mustEmbedUnimplementedAttestationServiceServer()
 }
 
@@ -98,6 +110,9 @@ func (UnimplementedAttestationServiceServer) FetchRawEvidence(context.Context, *
 func (UnimplementedAttestationServiceServer) FetchAzureToken(context.Context, *AzureTokenRequest) (*AzureTokenResponse, error) {
 	return nil, status.Error(codes.Unimplemented, "method FetchAzureToken not implemented")
 }
+func (UnimplementedAttestationServiceServer) FetchKbsToken(context.Context, *KbsTokenRequest) (*KbsTokenResponse, error) {
+	return nil, status.Error(codes.Unimplemented, "method FetchKbsToken not implemented")
+}
 func (UnimplementedAttestationServiceServer) mustEmbedUnimplementedAttestationServiceServer() {}
 func (UnimplementedAttestationServiceServer) testEmbeddedByValue()                            {}
 
@@ -173,6 +188,24 @@ func _AttestationService_FetchAzureToken_Handler(srv interface{}, ctx context.Co
 	return interceptor(ctx, in, info, handler)
 }
 
+func _AttestationService_FetchKbsToken_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
+	in := new(KbsTokenRequest)
+	if err := dec(in); err != nil {
+		return nil, err
+	}
+	if interceptor == nil {
+		return srv.(AttestationServiceServer).FetchKbsToken(ctx, in)
+	}
+	info := &grpc.UnaryServerInfo{
+		Server:     srv,
+		FullMethod: AttestationService_FetchKbsToken_FullMethodName,
+	}
+	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
+		return srv.(AttestationServiceServer).FetchKbsToken(ctx, req.(*KbsTokenRequest))
+	}
+	return interceptor(ctx, in, info, handler)
+}
+
 // AttestationService_ServiceDesc is the grpc.ServiceDesc for AttestationService service.
 // It's only intended for direct use with grpc.RegisterService,
 // and not to be introspected or modified (even as a copy)
@@ -192,6 +225,10 @@ var AttestationService_ServiceDesc = grpc.ServiceDesc{
 			MethodName: "FetchAzureToken",
 			Handler:    _AttestationService_FetchAzureToken_Handler,
 		},
+		{
+			MethodName: "FetchKbsToken",
+			Handler:    _AttestationService_FetchKbsToken_Handler,
+		},
 	},
 	Streams:  []grpc.StreamDesc{},
 	Metadata: "internal/proto/attestation/v1/attestation.proto",

manager/manager.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.36.11
-// 	protoc        v6.33.1
+// 	protoc        v7.35.1
 // source: manager/manager.proto
 
 package manager

manager/manager_grpc.pb.go

@@ -4,7 +4,7 @@
 // Code generated by protoc-gen-go-grpc. DO NOT EDIT.
 // versions:
 // - protoc-gen-go-grpc v1.6.0
-// - protoc             v6.33.1
+// - protoc             v7.35.1
 // source: manager/manager.proto
 
 package manager

pkg/atls/provider.go

@@ -11,7 +11,6 @@ import (
 	"crypto/x509"
 	"fmt"
 
-	"github.com/absmach/certs/sdk"
 	"github.com/ultravioletrs/cocos/pkg/atls/ea"
 	eaattestation "github.com/ultravioletrs/cocos/pkg/atls/eaattestation"
 	cocosattestation "github.com/ultravioletrs/cocos/pkg/attestation"
@@ -30,7 +29,7 @@ type provider struct {
 	platformType cocosattestation.PlatformType
 }
 
-func NewProvider(attClient attestationclient.Client, platformType cocosattestation.PlatformType, _ string, _ string, _ sdk.SDK) (CertificateProvider, error) {
+func NewProvider(attClient attestationclient.Client, platformType cocosattestation.PlatformType) (CertificateProvider, error) {
 	if attClient == nil {
 		return nil, fmt.Errorf("atls: missing attestation client")
 	}

pkg/attestation/attestation.go

@@ -42,6 +42,7 @@ type Provider interface {
 	TeeAttestation(teeNonce []byte) ([]byte, error)
 	VTpmAttestation(vTpmNonce []byte) ([]byte, error)
 	AzureAttestationToken(tokenNonce []byte) ([]byte, error)
+	KbsToken() ([]byte, error)
 }
 
 type Verifier interface {

pkg/attestation/azure/snp.go

@@ -134,6 +134,10 @@ func (a provider) AzureAttestationToken(tokenNonce []byte) ([]byte, error) {
 	return []byte(token), nil
 }
 
+func (a provider) KbsToken() ([]byte, error) {
+	return nil, fmt.Errorf("KbsToken not supported on Azure provider")
+}
+
 type verifier struct {
 	writer io.Writer
 }

pkg/attestation/ccaa/provider.go

@@ -129,3 +129,21 @@ func (p *Provider) AzureAttestationToken(nonce []byte) ([]byte, error) {
 
 	return resp.Token, nil
 }
+
+// KbsToken retrieves the official KBS-signed token from the attestation agent.
+func (p *Provider) KbsToken() ([]byte, error) {
+	fmt.Println("[CCAA-PROVIDER] KbsToken called, querying attestation-agent...")
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+
+	resp, err := p.client.GetToken(ctx, &attestation_agent.GetTokenRequest{
+		TokenType: "kbs",
+	})
+	if err != nil {
+		fmt.Printf("[CCAA-PROVIDER] GetToken failed: %v\n", err)
+		return nil, fmt.Errorf("CC attestation-agent GetToken failed for kbs: %w", err)
+	}
+
+	fmt.Printf("[CCAA-PROVIDER] GetToken succeeded! Return raw JSON message, len: %d\n", len(resp.Token))
+	return resp.Token, nil
+}

pkg/attestation/emptyprovider.go

@@ -31,3 +31,7 @@ func (e *EmptyProvider) VTpmAttestation(vTpmNonce []byte) ([]byte, error) {
 func (e *EmptyProvider) AzureAttestationToken(nonce []byte) ([]byte, error) {
 	return nil, nil
 }
+
+func (e *EmptyProvider) KbsToken() ([]byte, error) {
+	return nil, fmt.Errorf("KbsToken not supported on EmptyProvider")
+}

pkg/attestation/tdx/tdx.go

@@ -72,6 +72,10 @@ func (v provider) AzureAttestationToken(tokenNonce []byte) ([]byte, error) {
 	return nil, errors.New("Azure attestation token is not supported")
 }
 
+func (v provider) KbsToken() ([]byte, error) {
+	return nil, errors.New("KbsToken is not supported on TDX provider")
+}
+
 type verifier struct {
 	Policy *checkconfig.Config
 }

pkg/attestation/vtpm/vtpm.go

@@ -120,6 +120,10 @@ func (v provider) AzureAttestationToken(tokenNonce []byte) ([]byte, error) {
 	return nil, errors.New("Azure attestation token is not supported")
 }
 
+func (v provider) KbsToken() ([]byte, error) {
+	return nil, errors.New("KbsToken is not supported on vTPM provider")
+}
+
 type verifier struct {
 	writer io.Writer
 }

pkg/clients/grpc/attestation/client.go

@@ -17,6 +17,7 @@ type Client interface {
 	GetAttestation(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error)
 	GetRawEvidence(ctx context.Context, reportData [64]byte, nonce [32]byte, attType attestation.PlatformType) ([]byte, error)
 	GetAzureToken(ctx context.Context, nonce [32]byte) ([]byte, error)
+	GetKbsToken(ctx context.Context) ([]byte, error)
 	Close() error
 }
 
@@ -132,3 +133,15 @@ func (c *client) GetAzureToken(ctx context.Context, nonce [32]byte) ([]byte, err
 
 	return resp.Token, nil
 }
+
+func (c *client) GetKbsToken(ctx context.Context) ([]byte, error) {
+	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+
+	resp, err := c.client.FetchKbsToken(ctx, &attestation_v1.KbsTokenRequest{})
+	if err != nil {
+		return nil, err
+	}
+
+	return resp.Token, nil
+}

pkg/crypto/decrypt.go

@@ -10,9 +10,9 @@ import (
 	"crypto/sha256"
 	"encoding/base64"
 	"encoding/json"
+	"fmt"
 
 	"github.com/absmach/magistrala/pkg/errors"
-	"golang.org/x/crypto/hkdf"
 )
 
 var (
@@ -117,12 +117,26 @@ func DecryptWithWrappedKey(encryptedResource EncryptedResource, privateKey *ecdh
 		return nil, errors.Wrap(ErrDecryptionFailed, err)
 	}
 
-	// Derive KEK (Key Encryption Key) using HKDF
-	kek := make([]byte, 32)
-	kdf := hkdf.New(sha256.New, sharedSecret, nil, nil)
-	if _, err := kdf.Read(kek); err != nil {
-		return nil, errors.Wrap(ErrDecryptionFailed, err)
-	}
+	// Derive KEK (Key Encryption Key) using Concat KDF (NIST SP 800-56A)
+	algStr := "ECDH-ES+A256KW"
+	otherInfo := make([]byte, 0, 4+len(algStr)+4+4+4)
+	algLen := uint32(len(algStr))
+	otherInfo = append(otherInfo, byte(algLen>>24), byte(algLen>>16), byte(algLen>>8), byte(algLen))
+	otherInfo = append(otherInfo, algStr...)
+	otherInfo = append(otherInfo, 0, 0, 0, 0) // PartyUInfo
+	otherInfo = append(otherInfo, 0, 0, 0, 0) // PartyVInfo
+	otherInfo = append(otherInfo, 0, 0, 1, 0) // SuppPubInfo (256 bits BE)
+
+	// Since we need a 32-byte KEK, and SHA-256 produces 32 bytes, we run exactly 1 iteration (counter = 1)
+	counter := uint32(1)
+	hashInput := make([]byte, 0, 4+len(sharedSecret)+len(otherInfo))
+	hashInput = append(hashInput, byte(counter>>24), byte(counter>>16), byte(counter>>8), byte(counter))
+	hashInput = append(hashInput, sharedSecret...)
+	hashInput = append(hashInput, otherInfo...)
+
+	h := sha256.New()
+	h.Write(hashInput)
+	kek := h.Sum(nil)
 
 	// Unwrap the content encryption key (CEK)
 	cek, err := unwrapKey(encryptedResource.EncryptedKey, kek)
@@ -201,13 +215,134 @@ func unwrapKey(wrappedKey, kek []byte) ([]byte, error) {
 	return unwrapped, nil
 }
 
+func decodeBase64(s string) ([]byte, error) {
+	if d, err := base64.StdEncoding.DecodeString(s); err == nil {
+		return d, nil
+	}
+	if d, err := base64.URLEncoding.DecodeString(s); err == nil {
+		return d, nil
+	}
+	if d, err := base64.RawURLEncoding.DecodeString(s); err == nil {
+		return d, nil
+	}
+	if d, err := base64.RawStdEncoding.DecodeString(s); err == nil {
+		return d, nil
+	}
+	return nil, errors.New("invalid base64 encoding")
+}
+
 // ParseEncryptedResource parses a JSON-encoded encrypted resource.
 func ParseEncryptedResource(data []byte) (*EncryptedResource, error) {
-	var resource EncryptedResource
-	if err := json.Unmarshal(data, &resource); err != nil {
+	var jwe struct {
+		Protected    string `json:"protected"`
+		EncryptedKey string `json:"encrypted_key"`
+		IV           string `json:"iv"`
+		Ciphertext   string `json:"ciphertext"`
+		Tag          string `json:"tag"`
+	}
+	if err := json.Unmarshal(data, &jwe); err != nil {
 		return nil, errors.Wrap(ErrInvalidFormat, err)
 	}
-	return &resource, nil
+
+	// JWE structure check: if it lacks protected header, try legacy standard struct unmarshal
+	if jwe.Protected == "" {
+		var legacy struct {
+			Ciphertext   string              `json:"ciphertext"`
+			EncryptedKey string              `json:"encrypted_key"`
+			IV           string              `json:"iv"`
+			Tag          string              `json:"tag"`
+			AAD          string              `json:"aad,omitempty"`
+			EPK          *EphemeralPublicKey `json:"epk,omitempty"`
+		}
+		if err := json.Unmarshal(data, &legacy); err != nil {
+			return nil, errors.Wrap(ErrInvalidFormat, err)
+		}
+
+		ciphertext, err := decodeBase64(legacy.Ciphertext)
+		if err != nil {
+			return nil, errors.Wrap(ErrInvalidFormat, err)
+		}
+
+		encryptedKey, err := decodeBase64(legacy.EncryptedKey)
+		if err != nil {
+			return nil, errors.Wrap(ErrInvalidFormat, err)
+		}
+
+		iv, err := decodeBase64(legacy.IV)
+		if err != nil {
+			return nil, errors.Wrap(ErrInvalidFormat, err)
+		}
+
+		tag, err := decodeBase64(legacy.Tag)
+		if err != nil {
+			return nil, errors.Wrap(ErrInvalidFormat, err)
+		}
+
+		var aad []byte
+		if legacy.AAD != "" {
+			aad, err = decodeBase64(legacy.AAD)
+			if err != nil {
+				return nil, errors.Wrap(ErrInvalidFormat, err)
+			}
+		}
+
+		return &EncryptedResource{
+			Ciphertext:   ciphertext,
+			EncryptedKey: encryptedKey,
+			IV:           iv,
+			Tag:          tag,
+			AAD:          aad,
+			EPK:          legacy.EPK,
+		}, nil
+	}
+
+	// 1. Decode Protected Header JSON
+	protectedJSON, err := decodeBase64(jwe.Protected)
+	if err != nil {
+		return nil, errors.Wrap(ErrInvalidFormat, fmt.Errorf("failed to decode JWE protected header: %w", err))
+	}
+
+	// 2. Parse Ephemeral Public Key (EPK) from Protected Header
+	var header struct {
+		Alg string              `json:"alg"`
+		EPK *EphemeralPublicKey `json:"epk"`
+	}
+	if err := json.Unmarshal(protectedJSON, &header); err != nil {
+		return nil, errors.Wrap(ErrInvalidFormat, fmt.Errorf("failed to parse JWE header JSON: %w", err))
+	}
+
+	// 3. Decode main crypto fields
+	ciphertext, err := decodeBase64(jwe.Ciphertext)
+	if err != nil {
+		return nil, errors.Wrap(ErrInvalidFormat, err)
+	}
+
+	encryptedKey, err := decodeBase64(jwe.EncryptedKey)
+	if err != nil {
+		return nil, errors.Wrap(ErrInvalidFormat, err)
+	}
+
+	iv, err := decodeBase64(jwe.IV)
+	if err != nil {
+		return nil, errors.Wrap(ErrInvalidFormat, err)
+	}
+
+	tag, err := decodeBase64(jwe.Tag)
+	if err != nil {
+		return nil, errors.Wrap(ErrInvalidFormat, err)
+	}
+
+	// In JWE, AAD is the ASCII bytes of the protected header string
+	aad := []byte(jwe.Protected)
+
+	return &EncryptedResource{
+		Ciphertext:   ciphertext,
+		EncryptedKey: encryptedKey,
+		IV:           iv,
+		Tag:          tag,
+		AAD:          aad,
+		EPK:          header.EPK,
+	}, nil
 }
 
 // zeroBytes securely zeros out a byte slice.

pkg/crypto/decrypt_test.go

@@ -15,9 +15,31 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"golang.org/x/crypto/hkdf"
 )
 
+// testConcatKDF derives a KEK using Concat KDF (NIST SP 800-56A), matching
+// the exact implementation in DecryptWithWrappedKey.
+func testConcatKDF(sharedSecret []byte) []byte {
+	algStr := "ECDH-ES+A256KW"
+	otherInfo := make([]byte, 0, 4+len(algStr)+4+4+4)
+	algLen := uint32(len(algStr))
+	otherInfo = append(otherInfo, byte(algLen>>24), byte(algLen>>16), byte(algLen>>8), byte(algLen))
+	otherInfo = append(otherInfo, algStr...)
+	otherInfo = append(otherInfo, 0, 0, 0, 0) // PartyUInfo
+	otherInfo = append(otherInfo, 0, 0, 0, 0) // PartyVInfo
+	otherInfo = append(otherInfo, 0, 0, 1, 0) // SuppPubInfo (256 bits BE)
+
+	counter := uint32(1)
+	hashInput := make([]byte, 0, 4+len(sharedSecret)+len(otherInfo))
+	hashInput = append(hashInput, byte(counter>>24), byte(counter>>16), byte(counter>>8), byte(counter))
+	hashInput = append(hashInput, sharedSecret...)
+	hashInput = append(hashInput, otherInfo...)
+
+	h := sha256.New()
+	h.Write(hashInput)
+	return h.Sum(nil)
+}
+
 // testAESKeyWrap implements RFC 3394 AES Key Wrap for use in test setup.
 func testAESKeyWrap(kek, key []byte) ([]byte, error) {
 	block, err := aes.NewCipher(kek)
@@ -527,11 +549,8 @@ func TestDecryptWithWrappedKeyFullRoundTrip(t *testing.T) {
 		sharedSecret, err := ephemeralKey.ECDH(recipientKey.PublicKey())
 		require.NoError(t, err)
 
-		// Derive KEK using HKDF (same as in DecryptWithWrappedKey)
-		kek := make([]byte, 32)
-		kdf := hkdf.New(sha256.New, sharedSecret, nil, nil)
-		_, err = kdf.Read(kek)
-		require.NoError(t, err)
+		// Derive KEK using Concat KDF (same as in DecryptWithWrappedKey)
+		kek := testConcatKDF(sharedSecret)
 
 		// Generate random CEK (32 bytes)
 		cek := make([]byte, 32)
@@ -589,10 +608,7 @@ func TestDecryptWithWrappedKeyFullRoundTrip(t *testing.T) {
 		sharedSecret, err := ephemeralKey.ECDH(recipientKey.PublicKey())
 		require.NoError(t, err)
 
-		kek := make([]byte, 32)
-		kdf := hkdf.New(sha256.New, sharedSecret, nil, nil)
-		_, err = kdf.Read(kek)
-		require.NoError(t, err)
+		kek := testConcatKDF(sharedSecret)
 
 		cek := make([]byte, 16) // 16-byte CEK (AES-128)
 		_, err = rand.Read(cek)
@@ -650,10 +666,7 @@ func TestDecryptWithWrappedKeyFullRoundTrip(t *testing.T) {
 		sharedSecret, err := ephemeralKey.ECDH(recipientKey.PublicKey())
 		require.NoError(t, err)
 
-		kek := make([]byte, 32)
-		kdf := hkdf.New(sha256.New, sharedSecret, nil, nil)
-		_, err = kdf.Read(kek)
-		require.NoError(t, err)
+		kek := testConcatKDF(sharedSecret)
 
 		cek := make([]byte, 32)
 		_, err = rand.Read(cek)

test/manual/algo/lin_reg.py

@@ -1,5 +1,25 @@
-import os
+import subprocess
 import sys
+
+def install(package):
+    subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+
+try:
+    import pandas
+except ImportError:
+    install("pandas")
+
+try:
+    import sklearn
+except ImportError:
+    install("scikit-learn")
+
+try:
+    import joblib
+except ImportError:
+    install("joblib")
+
+import os
 import joblib
 import pandas as pd
 from sklearn.model_selection import train_test_split