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>

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/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"
@@ -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 {

agent/service_test.go

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