NOISSUE - Add FetchKbsToken RPC support, update protobuf generation, and include additional binaries in CI workflow. (#610)

* feat: add FetchKbsToken RPC support, update protobuf generation, and include additional binaries in CI workflow.

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

* chore: update protoc version and add GetKbsToken mock method with updated kbsHTTPGet signature

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

* test: inject mock attestation client into agentService for resource and KBS tests

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

* test: update key derivation in tests to use Concat KDF instead of HKDF

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

---------

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.0
           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
@@ -75,6 +77,12 @@ jobs:
           files: |
             buildroot/output/images/bzImage
             buildroot/output/images/rootfs.cpio.gz
-            cocos/build/cocos-agent  
-            cocos/build/cocos-cli  
+            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.0
 // 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.0
 // source: agent/agent.proto
 
 package agent

agent/cvms/api/grpc/client.go

@@ -178,6 +178,9 @@ func (client *CVMSClient) processIncomingMessage(ctx context.Context, req *cvms.
 	switch mes := req.Message.(type) {
 	case *cvms.ServerStreamMessage_RunReqChunks:
 		return client.handleRunReqChunks(ctx, mes)
+	case *cvms.ServerStreamMessage_RunReq:
+		client.logger.Info("Starting computation execution from non-chunked run request", "computationId", mes.RunReq.Id, "name", mes.RunReq.Name)
+		go client.executeRun(ctx, mes.RunReq)
 	case *cvms.ServerStreamMessage_StopComputation:
 		go client.handleStopComputation(ctx, mes)
 	case *cvms.ServerStreamMessage_AgentStateReq:

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.0
 // 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.0
 // 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.0
 // 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.0
 // 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.0
 // 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.0
 // source: agent/runner/runner.proto
 
 package runner

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.0
 // source: agent/runner/runner.proto
 
 package runner

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"
@@ -416,7 +421,7 @@ func (as *agentService) downloadAlgorithmIfRemote(state statemachine.State) {
 		"kbs_enabled", kbsEnabled)
 
 	// Check if algorithm should be downloaded from remote source
-	if as.computation.Algorithm.Source != nil && kbsEnabled {
+	if as.computation.Algorithm.Source != nil && kbsEnabled && as.computation.Algorithm.Source.URL != "" {
 		as.logger.Info("downloading algorithm from remote source",
 			"url", as.computation.Algorithm.Source.URL,
 			"kbs_resource_path", as.computation.Algorithm.Source.KBSResourcePath,
@@ -536,7 +541,7 @@ func (as *agentService) downloadDatasetsIfRemote(state statemachine.State) {
 	hasRemoteDatasets := false
 	for _, d := range as.computation.Datasets {
 		kbsEnabled := d.KBS != nil && d.KBS.Enabled
-		if d.Source != nil && kbsEnabled {
+		if d.Source != nil && kbsEnabled && d.Source.URL != "" {
 			hasRemoteDatasets = true
 			break
 		}
@@ -557,7 +562,7 @@ func (as *agentService) downloadDatasetsIfRemote(state statemachine.State) {
 			kbsURL = d.KBS.URL
 		}
 
-		if d.Source != nil && kbsEnabled {
+		if d.Source != nil && kbsEnabled && d.Source.URL != "" {
 			as.logger.Info("downloading dataset from remote source", "filename", d.Filename, "kbs_url", kbsURL)
 
 			res, err := as.downloadAndDecryptResource(ctx, d.Source, kbsURL, "dataset")
@@ -762,26 +767,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 {
@@ -944,7 +1014,7 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 		kbsURL = as.computation.Algorithm.KBS.URL
 	}
 
-	if as.computation.Algorithm.Source != nil && kbsEnabled {
+	if as.computation.Algorithm.Source != nil && kbsEnabled && as.computation.Algorithm.Source.URL != "" {
 		as.logger.Info("downloading algorithm from remote source", "kbs_url", kbsURL)
 
 		res, err := as.downloadAndDecryptResource(ctx, as.computation.Algorithm.Source, kbsURL, "algorithm")
@@ -957,6 +1027,19 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error {
 	} else {
 		// Use directly uploaded algorithm
 		algoData = algo.Algorithm
+
+		if as.computation.Algorithm.Source != nil && as.computation.Algorithm.Source.Encrypted && kbsEnabled {
+			as.logger.Info("directly uploaded algorithm is encrypted, retrieving key from KBS")
+			key, err := as.getKeyFromKBS(ctx, kbsURL, as.computation.Algorithm.Source.KBSResourcePath)
+			if err != nil {
+				return fmt.Errorf("failed to retrieve key from KBS for uploaded algorithm: %w", err)
+			}
+			decrypted, err := resource.DecryptData(algoData, key)
+			if err != nil {
+				return fmt.Errorf("failed to decrypt uploaded algorithm: %w", err)
+			}
+			algoData = decrypted
+		}
 	}
 
 	hash := sha3.Sum256(algoData)
@@ -1032,7 +1115,7 @@ func (as *agentService) Data(ctx context.Context, dataset Dataset) error {
 			kbsURL = d.KBS.URL
 		}
 
-		if d.Source != nil && kbsEnabled {
+		if d.Source != nil && kbsEnabled && d.Source.URL != "" {
 			as.logger.Info("downloading dataset from remote source", "filename", d.Filename, "kbs_url", kbsURL)
 
 			downloadedData, err := as.downloadAndDecryptResource(ctx, d.Source, kbsURL, "dataset")
@@ -1051,44 +1134,85 @@ func (as *agentService) Data(ctx context.Context, dataset Dataset) error {
 	if matchedIndex == -1 {
 		datasetData = dataset.Dataset
 		datasetFilename = dataset.Filename
+
+		index, ok := IndexFromContext(ctx)
+		if ok {
+			if index < 0 || index >= len(as.computation.Datasets) {
+				return ErrUndeclaredDataset
+			}
+			if as.computation.Datasets[index].Filename != datasetFilename {
+				return ErrFileNameMismatch
+			}
+			matchedIndex = index
+		} else {
+			matchedIndex = -1
+			for i, d := range as.computation.Datasets {
+				if d.Filename == datasetFilename {
+					matchedIndex = i
+					break
+				}
+			}
+			if matchedIndex == -1 {
+				return ErrUndeclaredDataset
+			}
+		}
+	} else {
+		remoteIndex := -1
+		for i, d := range as.computation.Datasets {
+			if d.Filename == datasetFilename {
+				remoteIndex = i
+				break
+			}
+		}
+		if remoteIndex == -1 {
+			return ErrUndeclaredDataset
+		}
+		matchedIndex = remoteIndex
+	}
+
+	d := as.computation.Datasets[matchedIndex]
+
+	kbsEnabled := d.KBS != nil && d.KBS.Enabled
+	kbsURL := ""
+	if d.KBS != nil {
+		kbsURL = d.KBS.URL
+	}
+	if d.Source != nil && d.Source.Encrypted && kbsEnabled {
+		as.logger.Info("directly uploaded dataset is encrypted, retrieving key from KBS", "filename", d.Filename)
+		key, err := as.getKeyFromKBS(ctx, kbsURL, d.Source.KBSResourcePath)
+		if err != nil {
+			return fmt.Errorf("failed to retrieve key from KBS for dataset %s: %w", d.Filename, err)
+		}
+		decrypted, err := resource.DecryptData(datasetData, key)
+		if err != nil {
+			return fmt.Errorf("failed to decrypt dataset %s: %w", d.Filename, err)
+		}
+		datasetData = decrypted
 	}
 
 	hash := sha3.Sum256(datasetData)
-
-	matched := false
-	for i, d := range as.computation.Datasets {
-		if hash == d.Hash {
-			if d.Filename != "" && d.Filename != datasetFilename {
-				return ErrFileNameMismatch
-			}
-
-			as.computation.Datasets = slices.Delete(as.computation.Datasets, i, i+1)
-
-			if DecompressFromContext(ctx) {
-				if err := internal.UnzipFromMemory(datasetData, algorithm.DatasetsDir); err != nil {
-					return fmt.Errorf("error decompressing dataset: %v", err)
-				}
-			} else {
-				f, err := os.Create(fmt.Sprintf("%s/%s", algorithm.DatasetsDir, datasetFilename))
-				if err != nil {
-					return fmt.Errorf("error creating dataset file: %v", err)
-				}
-
-				if _, err := f.Write(datasetData); err != nil {
-					return fmt.Errorf("error writing dataset to file: %v", err)
-				}
-				if err := f.Close(); err != nil {
-					return fmt.Errorf("error closing file: %v", err)
-				}
-			}
-
-			matched = true
-			break
-		}
+	if hash != d.Hash {
+		return ErrHashMismatch
 	}
 
-	if !matched {
-		return ErrUndeclaredDataset
+	as.computation.Datasets = slices.Delete(as.computation.Datasets, matchedIndex, matchedIndex+1)
+
+	if DecompressFromContext(ctx) {
+		if err := internal.UnzipFromMemory(datasetData, algorithm.DatasetsDir); err != nil {
+			return fmt.Errorf("error decompressing dataset: %v", err)
+		}
+	} else {
+		f, err := os.Create(fmt.Sprintf("%s/%s", algorithm.DatasetsDir, datasetFilename))
+		if err != nil {
+			return fmt.Errorf("error creating dataset file: %v", err)
+		}
+
+		if _, err := f.Write(datasetData); err != nil {
+			return fmt.Errorf("error writing dataset to file: %v", err)
+		}
+		if err := f.Close(); err != nil {
+			return fmt.Errorf("error closing file: %v", err)
+		}
 	}
 
 	if len(as.computation.Datasets) == 0 {

agent/service_test.go

@@ -218,7 +218,7 @@ func TestData(t *testing.T) {
 		{
 			name: "Test dataset not declared in manifest",
 			data: Dataset{
-				Filename: datasetFile,
+				Filename: "undeclared.csv",
 			},
 			err: ErrUndeclaredDataset,
 		},
@@ -1808,3 +1808,91 @@ func TestInferSourceType(t *testing.T) {
 		})
 	}
 }
+
+func TestInitComputation_Errors(t *testing.T) {
+	sm := &smmocks.StateMachine{}
+	sm.On("GetState").Return(ReceivingAlgorithm) // Not ReceivingManifest
+
+	svc := &agentService{
+		sm: sm,
+	}
+
+	err := svc.InitComputation(context.Background(), Computation{})
+	assert.ErrorIs(t, err, ErrStateNotReady)
+}
+
+func TestAlgo_Errors(t *testing.T) {
+	t.Run("state not ready", func(t *testing.T) {
+		sm := &smmocks.StateMachine{}
+		sm.On("GetState").Return(ReceivingManifest) // Not ReceivingAlgorithm
+
+		svc := &agentService{
+			sm: sm,
+		}
+
+		err := svc.Algo(context.Background(), Algorithm{})
+		assert.ErrorIs(t, err, ErrStateNotReady)
+	})
+
+	t.Run("all manifest items received", func(t *testing.T) {
+		sm := &smmocks.StateMachine{}
+		sm.On("GetState").Return(ReceivingAlgorithm)
+
+		svc := &agentService{
+			sm:           sm,
+			algoReceived: true,
+		}
+
+		err := svc.Algo(context.Background(), Algorithm{})
+		assert.ErrorIs(t, err, ErrAllManifestItemsReceived)
+	})
+
+	t.Run("undeclared algorithm", func(t *testing.T) {
+		sm := &smmocks.StateMachine{}
+		sm.On("GetState").Return(ReceivingAlgorithm)
+
+		svc := &agentService{
+			sm:           sm,
+			algoReceived: false,
+			computation: Computation{
+				Algorithm: nil, // Not declared
+			},
+		}
+
+		err := svc.Algo(context.Background(), Algorithm{})
+		assert.ErrorIs(t, err, ErrUndeclaredAlgorithm)
+	})
+}
+
+func TestData_ErrorsExtra(t *testing.T) {
+	t.Run("all manifest items received", func(t *testing.T) {
+		sm := &smmocks.StateMachine{}
+		sm.On("GetState").Return(ReceivingData)
+
+		svc := &agentService{
+			sm: sm,
+			computation: Computation{
+				Datasets: nil, // length 0
+			},
+		}
+
+		err := svc.Data(context.Background(), Dataset{})
+		assert.ErrorIs(t, err, ErrAllManifestItemsReceived)
+	})
+}
+
+func TestEnsureDir_Error(t *testing.T) {
+	tmpFile, err := os.CreateTemp("", "ensureDirTest")
+	require.NoError(t, err)
+	defer os.Remove(tmpFile.Name())
+	tmpFile.Close()
+
+	// ensureDir should fail because the parent path is a file, not a directory
+	err = ensureDir(filepath.Join(tmpFile.Name(), "subdir"), 0o755)
+	assert.Error(t, err)
+}
+
+func TestKbsHTTPGet_Error(t *testing.T) {
+	_, err := kbsHTTPGet(context.Background(), "%%", "")
+	assert.Error(t, err)
+}

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

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.0
 // 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.0
 // 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.0
 // 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.0
 // 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.0
 // 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.0
 // source: manager/manager.proto
 
 package manager

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)

scripts/encrypt.go

@@ -0,0 +1,74 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package main
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"fmt"
+	"io"
+	"os"
+)
+
+func main() {
+	if len(os.Args) < 4 {
+		fmt.Println("Usage: go run encrypt.go <key_file> <input_file> <output_file>")
+		os.Exit(1)
+	}
+
+	keyFile := os.Args[1]
+	inputFile := os.Args[2]
+	outputFile := os.Args[3]
+
+	// Read key
+	key, err := os.ReadFile(keyFile)
+	if err != nil {
+		fmt.Printf("Failed to read key file: %v\n", err)
+		os.Exit(1)
+	}
+	if len(key) != 32 {
+		fmt.Printf("Key must be 32 bytes, got %d\n", len(key))
+		os.Exit(1)
+	}
+
+	// Read plaintext
+	plaintext, err := os.ReadFile(inputFile)
+	if err != nil {
+		fmt.Printf("Failed to read input file: %v\n", err)
+		os.Exit(1)
+	}
+
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		fmt.Printf("Failed to create cipher: %v\n", err)
+		os.Exit(1)
+	}
+
+	aesgcm, err := cipher.NewGCM(block)
+	if err != nil {
+		fmt.Printf("Failed to create GCM: %v\n", err)
+		os.Exit(1)
+	}
+
+	nonce := make([]byte, 12)
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		fmt.Printf("Failed to generate nonce: %v\n", err)
+		os.Exit(1)
+	}
+
+	ciphertext := aesgcm.Seal(nil, nonce, plaintext, nil)
+
+	// Combine nonce + ciphertext + tag
+	// Seal returns ciphertext || tag, so we just append it to nonce
+	output := append(nonce, ciphertext...)
+
+	err = os.WriteFile(outputFile, output, 0o644)
+	if err != nil {
+		fmt.Printf("Failed to write output file: %v\n", err)
+		os.Exit(1)
+	}
+
+	fmt.Printf("Successfully encrypted %s to %s\n", inputFile, outputFile)
+}

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