Files
cocos/agent/cvms/api/grpc/client_test.go
T
Sammy Kerata Oina da31d76c94
CI / checkproto (push) Has been cancelled
CI / lint (push) Has been cancelled
Rust CI Pipeline / rust-check (push) Has been cancelled
CI / test (agent) (push) Has been cancelled
CI / test (cli) (push) Has been cancelled
CI / test (cmd) (push) Has been cancelled
CI / test (internal) (push) Has been cancelled
CI / test (manager, true) (push) Has been cancelled
CI / test (pkg) (push) Has been cancelled
CI / upload-coverage (push) Has been cancelled
NOISSUE - Agent Pull mode for remote resources (#575)
* feat(kbs): implement KBS client for attestation and resource retrieval

- Added KBS client implementation in pkg/kbs/client.go with methods for attestation and resource retrieval.
- Introduced necessary data structures for requests and responses.
- Implemented error handling for various scenarios.

test(kbs): add unit tests for KBS client

- Created comprehensive tests for the KBS client in pkg/kbs/client_test.go.
- Included tests for attestation success and failure cases, as well as resource retrieval.

feat(registry): introduce HTTP and S3 registry implementations

- Added HTTPRegistry for downloading resources over HTTP/HTTPS with retry logic in pkg/registry/http.go.
- Implemented S3Registry for downloading resources from AWS S3 and S3-compatible services in pkg/registry/s3.go.
- Included error handling and configuration options for both registries.

chore(registry): define registry interface and configuration

- Created registry interface and configuration struct in pkg/registry/registry.go.
- Added default configuration settings for registry clients.

docs(cvms): update README for CVMS server configuration and usage

- Enhanced documentation for CVMS server with detailed command-line flags and usage examples.
- Clarified direct upload and remote resource modes, including KBS integration.

fix(cvms): integrate KBS for remote resource handling in main.go

- Updated main.go to support remote datasets and algorithms using KBS.
- Added validation for command-line flags to ensure proper configuration.

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

* fix: Move ifeq conditional outside define block in attestation-service.mk

Make conditionals cannot be evaluated inside define...endef blocks
when used as recipe bodies. Restructured to define the
ATTESTATION_SERVICE_INSTALL_INIT_SYSTEMD block conditionally based
on BR2_PACKAGE_CC_ATTESTATION_AGENT configuration.

* feat: Implement remote resource downloading for algorithms and datasets using AWS S3/MinIO credentials.

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

* feat: Add comprehensive documentation and agent support for testing remote resource download with KBS attestation.

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

* feat: Improve agent logging for remote resource configuration and KBS status, and add a testing guide for remote resource downloads with KBS attestation.

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

* feat: Add a comprehensive guide for testing remote resource download with KBS attestation and update multiple package versions to a specific commit.

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

* feat: Add failure transitions for resource reception states and a comprehensive guide for testing remote resource downloads with KBS attestation.

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

* feat: Implement remote resource download with KBS attestation in the agent and add a comprehensive testing guide.

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

* test: Add comprehensive guide for testing remote resource download with KBS attestation and include a debug log in the attestation client.

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

* feat: Delegate KBS attestation and token retrieval to a new attestation-agent service and document remote resource testing.

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

* client fixes

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

* raw evidence

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

* fix: Build all Go files in cmd directories, not just main.go

This fixes the issue where fetch_raw_evidence.go wasn't being included
in the attestation-service build.

* fix: Wrap binary evidence in JSON for KBS compatibility

Fixes 'invalid character' error by wrapping raw binary evidence
in a JSON structure with base64 encoding, as expected by KBS.

* chore: Update buildroot packages to c28cefae

Includes fixes for:
1. attestation-service build (including fetch_raw_evidence.go)
2. Agent KBS evidence format (wrapping binary in JSON)

* fix: Implement KBS RCAR handshake with cookies

Fixes 'cookie not found' error (401) from KBS by:
1. Adding CookieJar support to KBS client
2. Implementing GetChallenge() to perform /auth handshake and capture session cookie
3. Updating Agent to get challenge, decode nonce, and use it for evidence generation
4. Regenerating mocks

* chore: Update buildroot packages to f6981ac5

Includes KBS RCAR handshake fix (cookie support + GetChallenge loop)

* fix: Update KBS client JSON tags to kebab-case

Fixes deserialization error (401) from KBS by:
1. Using kebab-case (e.g. extra-params) for JSON tags as per protocol.
2. Initializing ExtraParams as empty object {} instead of null/omitted.

* fix: Wrap attestation evidence in primary_evidence format

Updates Agent to construct 'tee-evidence' payload with:
- primary_evidence: containing the actual quote/data
- additional_evidence: empty JSON object

This matches the Confidential Containers KBS Attestation Protocol requirements.

* fix: Update KBS protocol version to 0.4.0

KBS rejected 0.1.0 with a version mismatch error. Bumping to 0.4.0 to match server expectation.

* fix: Generate ephemeral key for KBS RuntimeData

Updates RuntimeData to include a valid ephemeral EC P-256 public key in JWK format, as required by the KBS RCAR protocol.
Also fixes the KBS client struct to support TEEPubKey as an object.

* fix: Update sample attestation quote to valid JSON

The default attestation.bin was binary, but the KBS Sample Verifier expects a valid JSON quote containing 'svn' and 'report_data'.
Updated the embedded bin file to contain this JSON structure.

* fix: Generate dynamic JSON quote for Sample TEE in FetchRawEvidence

The KBS Sample Verifier expects a JSON object with 'svn' and 'report_data'.
Previously, we were returning raw binary data (reportData+nonce).
This commit updates FetchRawEvidence to return a marshaled JSON structure with:
- svn: "1"
- report_data: base64(req.ReportData)

* refactor: Delegate Sample Attestation to Provider

Refactored sample attestation logic:
- Moved JSON Quote generation into EmptyProvider (standalone mode).
- Updated FetchRawEvidence to call provider.TeeAttestation instead of manual generation.
This enables using the real CC Attestation Agent for UNSPECIFIED platform if configured.

* feat: Add comprehensive debug logging and enforce CC AA usage

Changes:
- Updated EmptyProvider to return error instead of generating mock data
  This forces proper use of CC Attestation Agent's sample attester
- Added detailed logging to attestation-service FetchRawEvidence:
  * Hex dump of evidence (first 200 bytes)
  * String preview of evidence
  * Total evidence length
- Added detailed logging to agent service:
  * Raw evidence hex and string previews
  * KBS evidence JSON preview (first 500 bytes)
  * Evidence lengths at each transformation step

This logging will help diagnose why KBS Sample Verifier is rejecting evidence.

* fix: Enable CC AA by default and add attestation-service log forwarding

Changes:
- Set USE_CC_ATTESTATION_AGENT=true by default in systemd service
- Added StandardOutput/StandardError to forward logs to /var/log/cocos/
- Updated HAL makefile to handle new default value
- This ensures attestation-service uses CC AA's sample attester
- Logs will now be visible in CVMS output for debugging

* feat: Add gRPC log forwarding to attestation-service

Implemented the same log forwarding mechanism used by the agent:
- Added ProtoHandler to write logs to both stdout and logQueue
- Connected to log client (/run/cocos/log.sock) for gRPC forwarding
- Added goroutine to forward logs to CVMS via log client
- Logs will now appear in CVMS output during computation runs

This enables visibility into attestation-service debug output including:
- CC AA connection status
- Evidence generation details (hex dumps, string previews)
- Any errors from providers

* fix: Parse sample evidence JSON instead of base64-encoding it

The attestation-service returns sample evidence as JSON:
{"svn":"1","report_data":"base64..."}

The agent was incorrectly base64-encoding this JSON string again.
KBS Sample Verifier expects the parsed JSON object directly.

Fixed by:
- Parsing the JSON evidence from attestation-service
- Passing the parsed object directly in primary_evidence.evidence
- This matches what KBS Sample Verifier expects

* debug: Increase KBS evidence logging preview to 1000 bytes

Show the complete JSON structure being sent to KBS to debug
the attestation failure.

* debug: Add comprehensive CC AA configuration logging

Added debug logs to show:
- Whether CC AA is enabled in config
- CC AA address being used
- Connection success/failure
- Which provider is ultimately selected
- Warning when falling back to EmptyProvider

This will help diagnose why EmptyProvider is being used
instead of CC Attestation Agent.

* debug: Add startup logging for log client connection

Added log message to show if log client connection succeeds
at attestation-service startup. This will help diagnose why
logs aren't appearing in CVMS output.

* feat: Add retry logic with exponential backoff to log client

Added simple retry mechanism to handle concurrent log requests:
- 3 retry attempts with exponential backoff (10ms, 20ms, 40ms)
- Applies to both SendLog and SendEvent methods
- Centralized in log client so all services benefit
- Should eliminate 'failed to send log' errors from concurrent requests

This fixes the issue where attestation-service logs weren't
appearing in CVMS output due to dropped messages.

* fix: Flatten sample evidence fields in primary_evidence for KBS

KBS Sample Verifier expects svn and report_data at the top level
of primary_evidence, not nested under an 'evidence' key.

Changed structure from:
{"primary_evidence": {"tee": "sample", "evidence": {"svn": "1", ...}}}

To:
{"primary_evidence": {"tee": "sample", "svn": "1", "report_data": "...", ...}}

This matches what KBS expects when deserializing the Quote structure.

* fix: Use sample quote directly as primary_evidence per KBS protocol

According to KBS attestation protocol spec, for sample TEE type,
primary_evidence should be the sample quote JSON directly:
{"svn": "1", "report_data": "..."}

Removed extra 'tee' and 'platform' fields that were causing KBS
to fail deserializing the Quote structure. The 'tee' field is
already sent in the Request payload during RCAR handshake.

Refs:
- https://github.com/confidential-containers/trustee/blob/main/kbs/docs/kbs_attestation_protocol.md
- https://github.com/confidential-containers/guest-components/blob/main/attestation-agent/attester/src/sample/mod.rs

* fix: Make CC AA required for sample attestation when configured

When USE_CC_ATTESTATION_AGENT=true, attestation-service now
requires AA to be available for NoCC/sample platform. This ensures
sample evidence always comes from AA with the correct KBS format.

Changes:
- Error out if AA connection fails for NoCC platform when AA is configured
- Only use EmptyProvider if AA is explicitly NOT configured
- Prevents incorrect sample evidence format from EmptyProvider

This ensures attestation-service delegates to AA for sample evidence
generation instead of creating it itself.

* fix: Implement proper RCAR protocol with tee-pubkey and runtime-data hash

Fixed KBS attestation error 'REPORT_DATA is different from that in Sample Quote'

Changes:
1. Generate ephemeral EC key pair BEFORE getting evidence from AA
2. Create runtime-data with nonce + tee-pubkey (JWK format)
3. Hash runtime-data (SHA-256) and use as report_data for AA
4. This binds the tee-pubkey to the TEE evidence per RCAR protocol

The report_data in the evidence now matches what KBS expects:
hash(runtime-data) instead of computation ID.

This completes the full RCAR protocol implementation:
- Request → Challenge → Attestation (with bound tee-pubkey) → Response

* fix(agent): use simple nonce for Sample attestation report_data

For Sample/NoCC attestation, use the raw nonce bytes directly as
report_data instead of hashing runtime-data. This avoids JSON
serialization mismatches with the KBS Sample verifier.

Real TEEs (TDX/SNP) still use runtime-data hash binding to
cryptographically bind the ephemeral tee-pubkey to the evidence.

* fix(agent): use RFC 8785 canonical JSON for runtime-data hashing

The KBS Sample attestation verifier (and likely others) expects the
report_data to be the SHA-256 hash of the *canonical* JSON serialization
(RFC 8785) of the runtime-data. Standard Go JSON marshaling does not
guarantee key ordering, leading to hash mismatches.

This change uses github.com/gowebpki/jcs to canonicalize the runtime-data
before hashing, ensuring compatibility with the KBS RCAR implementation.
Also reverted the temporary 'simple nonce' workaround.

* feat(hal): add CoCo Keyprovider and Skopeo packages

- Add coco-keyprovider buildroot package with systemd service
- Add skopeo buildroot package for OCI image handling
- Add ocicrypt_keyprovider.conf for encrypted image decryption
- Update Config.in to include new packages

This enables standard CoCo ecosystem integration for encrypted
OCI images instead of custom S3/HTTP registry clients.

* feat(oci): add OCI image handling package with Skopeo integration

- Add pkg/oci/types.go with ResourceSource and ImageManifest types
- Add pkg/oci/skopeo.go with Skopeo wrapper for pull/decrypt
- Add pkg/oci/extract.go for extracting algorithms and datasets from layers

This package provides OCI image handling using Skopeo and CoCo
Keyprovider for encrypted image decryption, replacing custom
S3/HTTP registry clients.

* chore: regenerate protobuf files for updated cvms.proto

* refactor(agent): replace S3/HTTP/KBS with OCI package

- Remove pkg/kbs and pkg/registry imports
- Add pkg/oci import for OCI image handling
- Replace downloadAndDecryptResource with OCI-based implementation
- Use Skopeo + CoCo Keyprovider for automatic decryption
- Reduce code from ~240 lines to ~70 lines

This eliminates custom KBS RCAR handshake, S3/HTTP registry clients,
and manual decryption logic. CoCo Keyprovider handles all decryption
automatically via ocicrypt protocol.

* chore: remove obsolete pkg/kbs and pkg/registry packages

- Delete pkg/kbs/ (custom KBS client, ~300 lines)
- Delete pkg/registry/ (S3/HTTP registry clients, ~400 lines)
- Remove unused imports from agent/service.go
- Run go mod tidy to clean up dependencies

These packages have been replaced by pkg/oci with Skopeo and
CoCo Keyprovider for standard CoCo ecosystem integration.

* fix(agent): update ResourceSource struct to include type and encryption fields

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

* fix(hal): update CoCo Keyprovider to v0.16.0 and fix build path

- Update version from v0.11.0 to v0.16.0 (matches attestation agent)
- Fix install path: target is at repo root, not in coco_keyprovider subdir
- This fixes the build error where coco_keyprovider binary wasn't found

The cargo workspace in guest-components builds to a shared target/
directory at the repository root, not within each crate's subdirectory.

* feat: Update remote resources testing guide to use kbs-client and coco-keyprovider for key management and encryption, enable insecure TLS for Skopeo, and enhance CVMS with

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

* feat: Update component versions, revise image encryption documentation, and sanitize OCI image paths for Skopeo compatibility.

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

* feat: Add `decompress` option to Dataset and `algo_type`/`algo_args` to Algorithm protobuf messages, updating client, test, and build configurations.

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

* Update multiple package versions and enhance OCI image extraction error reporting for missing algorithm files.

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

* chore: Bump package versions, improve OCI image extraction debugging by returning seen files, and remove unused dataset type parsing from test code.

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

* refactor: Migrate OCI extraction to use structured logging with `slog` and `context`, and update package versions.

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

* feat: Bump multiple component versions, add encrypted status for computation inputs and algorithms, and refine OCI layer extraction warnings.

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

* logging

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

* feat: Add `Encrypted` field to algorithm and dataset resource sources and update all component versions.

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

* feat: update component versions, integrate coco-keyprovider service, and configure ocicrypt key provider.

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

* feat: add support for KBS parameters and dataset/algorithm hash calculations in CVMS

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

* feat: update resource download and extraction logic to support requirements.txt and improve hash verification

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

* chore: Update dependencies, improve code style, and add GetRawEvidence to attestation client mocks.

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

* Refactor code structure for improved readability and maintainability

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

* fix: update golangci configuration to include errcheck for build path and remove unnecessary exclusions

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

* fix: streamline kernel command line handling in QEMU args construction

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

* feat: add attestation binary and update checksum tests and policy structure

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

* Add unit tests for attestation agent, attestation, log, crypto, OCI, and Skopeo clients

- Implement tests for the attestation agent client including Unix socket and TCP address handling, token retrieval, and error scenarios.
- Enhance attestation client tests to cover fetching raw evidence for various platforms (SNP, TDX, VTPM, SNPvTPM) and validate error handling.
- Introduce log client tests to verify retry behavior for sending logs and events.
- Create comprehensive tests for crypto package focusing on AES-GCM decryption, encrypted resource parsing, and key unwrapping.
- Add tests for OCI package to validate algorithm and dataset extraction, including JSON serialization of OCILayout.
- Implement Skopeo client tests to ensure proper functionality for image pulling, inspecting, and resource source handling.

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

* fix: handle JSON marshal errors in test cases for decrypt and extract functions

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

* test: add comprehensive tests for algorithm and dataset extraction with various scenarios

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

* refactor: replace hardcoded Python script content with constant variable

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

* fix: remove redundant mock expectation for SendAgentConfig in TestCreateVMWithAaKbsParams

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

* test: add tests for event sending failure, dataset extraction with path traversal, and Skopeo client behavior

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

* test: add tests for download and decryption of resources with various URL formats

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

* refactor: Introduce OCIClient interface for agent service to improve testability of OCI image operations and enhance related tests.

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

* refactor: Change `get_uint64_from_tcb` to accept `TcbVersion` by value and use `u64::from` for type conversions.

---------

Signed-off-by: Sammy Oina <sammyoina@gmail.com>
2026-03-16 14:48:55 +01:00

654 lines
20 KiB
Go

// Copyright (c) Ultraviolet
// SPDX-License-Identifier: Apache-2.0
package grpc
import (
"context"
"testing"
"time"
mglog "github.com/absmach/supermq/logger"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/ultravioletrs/cocos/agent"
"github.com/ultravioletrs/cocos/agent/cvms"
"github.com/ultravioletrs/cocos/agent/cvms/api/grpc/storage"
servermocks "github.com/ultravioletrs/cocos/agent/cvms/server/mocks"
"github.com/ultravioletrs/cocos/agent/mocks"
pkggrpc "github.com/ultravioletrs/cocos/pkg/clients/grpc"
clientmocks "github.com/ultravioletrs/cocos/pkg/clients/grpc/mocks"
"github.com/ultravioletrs/cocos/pkg/ingress"
"golang.org/x/crypto/sha3"
"google.golang.org/grpc"
"google.golang.org/protobuf/proto"
)
type mockStream struct {
mock.Mock
grpc.ClientStream
}
func (m *mockStream) Recv() (*cvms.ServerStreamMessage, error) {
args := m.Called()
return args.Get(0).(*cvms.ServerStreamMessage), args.Error(1)
}
func (m *mockStream) Send(msg *cvms.ClientStreamMessage) error {
args := m.Called(msg)
return args.Error(0)
}
// mockIngressProxy is a mock implementation of the ingress proxy.
type mockIngressProxy struct {
mock.Mock
}
func (m *mockIngressProxy) Start(config ingress.ProxyConfig, ctx ingress.ProxyContext) error {
args := m.Called(config, ctx)
return args.Error(0)
}
func (m *mockIngressProxy) Stop() error {
args := m.Called()
return args.Error(0)
}
func TestManagerClient_Process(t *testing.T) {
tests := []struct {
name string
setupMocks func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client)
expectError bool
errorMsg string
}{
{
name: "Stop computation",
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client) {
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{
Message: &cvms.ServerStreamMessage_StopComputation{
StopComputation: &cvms.StopComputation{},
},
}, nil)
mockStream.On("Send", mock.Anything).Return(nil)
mockSvc.On("StopComputation", mock.Anything).Return(nil)
mockServerSvc.On("Stop").Return(nil)
},
expectError: true,
errorMsg: "context deadline exceeded",
},
{
name: "Run request chunks",
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client) {
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{
Message: &cvms.ServerStreamMessage_RunReqChunks{
RunReqChunks: &cvms.RunReqChunks{},
},
}, nil)
mockStream.On("Send", mock.Anything).Return(nil).Once()
mockSvc.On("Run", mock.Anything, mock.Anything).Return("", assert.AnError).Once()
},
expectError: true,
},
{
name: "Agent state request",
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client) {
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{
Message: &cvms.ServerStreamMessage_AgentStateReq{
AgentStateReq: &cvms.AgentStateReq{
Id: "test-agent",
},
},
}, nil)
mockStream.On("Send", mock.Anything).Return(nil)
mockSvc.On("State").Return("test-state")
},
expectError: true,
errorMsg: "context deadline exceeded",
},
{
name: "Disconnect request",
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client) {
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{
Message: &cvms.ServerStreamMessage_DisconnectReq{},
}, nil)
mockStream.On("Send", mock.Anything).Return(nil)
grpcClient.On("Close").Return(nil)
},
expectError: true,
errorMsg: "context deadline exceeded",
},
{
name: "Receive error",
setupMocks: func(mockStream *mockStream, mockSvc *mocks.Service, mockServerSvc *servermocks.AgentServer, grpcClient *clientmocks.Client) {
mockStream.On("Recv").Return(&cvms.ServerStreamMessage{}, assert.AnError)
},
expectError: true,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage)
logger := mglog.NewMock()
go func() {
<-messageQueue
}()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
tc.setupMocks(mockStream, mockSvc, mockServerSvc, grpcClient)
err = client.Process(ctx, cancel)
if tc.expectError {
assert.Error(t, err)
if tc.errorMsg != "" {
assert.Contains(t, err.Error(), tc.errorMsg)
}
} else {
assert.NoError(t, err)
}
})
}
}
func TestManagerClient_handleRunReqChunks(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
runReq := &cvms.ComputationRunReq{
Id: "test-id",
Datasets: []*cvms.Dataset{
{
Hash: sha3.New256().Sum([]byte("test-dataset")),
},
},
Algorithm: &cvms.Algorithm{
Hash: sha3.New256().Sum([]byte("test-algorithm")),
},
ResultConsumers: []*cvms.ResultConsumer{
{
UserKey: []byte("test-consumer"),
},
},
}
runReqBytes, _ := proto.Marshal(runReq)
chunk1 := &cvms.ServerStreamMessage_RunReqChunks{
RunReqChunks: &cvms.RunReqChunks{
Id: "chunk-1",
Data: runReqBytes[:len(runReqBytes)/2],
IsLast: false,
},
}
chunk2 := &cvms.ServerStreamMessage_RunReqChunks{
RunReqChunks: &cvms.RunReqChunks{
Id: "chunk-1",
Data: runReqBytes[len(runReqBytes)/2:],
IsLast: true,
},
}
mockSvc.On("State").Return("ReceivingManifest")
mockSvc.On("InitComputation", mock.Anything, mock.Anything).Return(nil)
mockServerSvc.On("Start", mock.Anything, mock.Anything, mock.Anything).Return(nil)
err = client.handleRunReqChunks(context.Background(), chunk1)
assert.NoError(t, err)
err = client.handleRunReqChunks(context.Background(), chunk2)
assert.NoError(t, err)
// Wait for the goroutine to finish
time.Sleep(50 * time.Millisecond)
mockSvc.AssertExpectations(t)
assert.Len(t, messageQueue, 1)
msg := <-messageQueue
runRes, ok := msg.Message.(*cvms.ClientStreamMessage_RunRes)
assert.True(t, ok)
assert.Equal(t, "test-id", runRes.RunRes.ComputationId)
}
func TestManagerClient_handleStopComputation(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
stopReq := &cvms.ServerStreamMessage_StopComputation{
StopComputation: &cvms.StopComputation{
ComputationId: "test-comp-id",
},
}
mockSvc.On("StopComputation", mock.Anything).Return(nil)
mockServerSvc.On("Stop").Return(nil)
client.handleStopComputation(context.Background(), stopReq)
// Wait for the goroutine to finish
time.Sleep(50 * time.Millisecond)
mockSvc.AssertExpectations(t)
assert.Len(t, messageQueue, 1)
msg := <-messageQueue
stopRes, ok := msg.Message.(*cvms.ClientStreamMessage_StopComputationRes)
assert.True(t, ok)
assert.Equal(t, "test-comp-id", stopRes.StopComputationRes.ComputationId)
assert.Empty(t, stopRes.StopComputationRes.Message)
}
func TestManagerClient_timeoutRequest(t *testing.T) {
rm := newRunRequestManager()
rm.requests["test-id"] = &runRequest{
timer: time.NewTimer(100 * time.Millisecond),
buffer: []byte("test-data"),
lastChunk: time.Now(),
}
rm.timeoutRequest("test-id")
assert.Len(t, rm.requests, 0)
}
// TestManagerClient_sendPendingMessages tests sending pending messages on reconnection.
func TestManagerClient_sendPendingMessages(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
// Add a pending message to storage
testMsg := &cvms.ClientStreamMessage{
Message: &cvms.ClientStreamMessage_RunRes{
RunRes: &cvms.RunResponse{
ComputationId: "test-id",
},
},
}
err = client.storage.Add(testMsg)
assert.NoError(t, err)
// Mock successful send
mockStream.On("Send", mock.Anything).Return(nil).Once()
// Load and send pending messages
pending, err := client.storage.Load()
assert.NoError(t, err)
assert.Len(t, pending, 1)
client.sendPendingMessages(pending)
mockStream.AssertExpectations(t)
}
// TestManagerClient_sendPendingMessagesWithError tests pending message send failure.
func TestManagerClient_sendPendingMessagesWithError(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
testMsg := &cvms.ClientStreamMessage{
Message: &cvms.ClientStreamMessage_RunRes{
RunRes: &cvms.RunResponse{
ComputationId: "test-id",
},
},
}
// Mock failed send
mockStream.On("Send", mock.Anything).Return(assert.AnError)
pending := []storage.Message{
{
Message: testMsg,
Time: time.Now(),
},
}
client.sendPendingMessages(pending)
mockStream.AssertExpectations(t)
}
// TestManagerClient_addChunkTimeout tests chunk timeout in runRequestManager.
func TestManagerClient_addChunkTimeout(t *testing.T) {
rm := newRunRequestManager()
// Add first chunk
chunk1 := []byte("chunk1")
buffer, complete := rm.addChunk("test-id", chunk1, false)
assert.Nil(t, buffer)
assert.False(t, complete)
// Verify request exists
rm.mu.Lock()
assert.Contains(t, rm.requests, "test-id")
rm.mu.Unlock()
// Wait for timeout
time.Sleep(35 * time.Second) // runReqTimeout is 30 seconds
// Verify request was removed
rm.mu.Lock()
assert.NotContains(t, rm.requests, "test-id")
rm.mu.Unlock()
}
// TestManagerClient_addChunkMultiple tests adding multiple chunks.
func TestManagerClient_addChunkMultiple(t *testing.T) {
rm := newRunRequestManager()
chunk1 := []byte("chunk1")
chunk2 := []byte("chunk2")
chunk3 := []byte("chunk3")
// Add chunks
buffer, complete := rm.addChunk("test-id", chunk1, false)
assert.Nil(t, buffer)
assert.False(t, complete)
buffer, complete = rm.addChunk("test-id", chunk2, false)
assert.Nil(t, buffer)
assert.False(t, complete)
buffer, complete = rm.addChunk("test-id", chunk3, true)
assert.NotNil(t, buffer)
assert.True(t, complete)
expected := append(append(chunk1, chunk2...), chunk3...)
assert.Equal(t, expected, buffer)
}
// TestManagerClient_handleStopComputationWithIngressProxy tests stop with ingress proxy.
func TestManagerClient_handleStopComputationWithIngressProxy(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
mockIngressProxy := new(mockIngressProxy)
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, mockIngressProxy, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
stopReq := &cvms.ServerStreamMessage_StopComputation{
StopComputation: &cvms.StopComputation{
ComputationId: "test-comp-id",
},
}
mockSvc.On("StopComputation", mock.Anything).Return(nil)
mockServerSvc.On("Stop").Return(nil)
mockIngressProxy.On("Stop").Return(nil)
client.handleStopComputation(context.Background(), stopReq)
time.Sleep(50 * time.Millisecond)
mockSvc.AssertExpectations(t)
mockServerSvc.AssertExpectations(t)
mockIngressProxy.AssertExpectations(t)
assert.Len(t, messageQueue, 1)
}
// TestManagerClient_handleStopComputationWithIngressProxyError tests stop with ingress proxy error.
func TestManagerClient_handleStopComputationWithIngressProxyError(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
mockIngressProxy := new(mockIngressProxy)
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, mockIngressProxy, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
stopReq := &cvms.ServerStreamMessage_StopComputation{
StopComputation: &cvms.StopComputation{
ComputationId: "test-comp-id",
},
}
mockSvc.On("StopComputation", mock.Anything).Return(nil)
mockServerSvc.On("Stop").Return(nil)
mockIngressProxy.On("Stop").Return(assert.AnError)
client.handleStopComputation(context.Background(), stopReq)
time.Sleep(50 * time.Millisecond)
mockIngressProxy.AssertExpectations(t)
}
// TestManagerClient_sendMessage tests sendMessage with timeout.
func TestManagerClient_sendMessage(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage, 1)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
msg := &cvms.ClientStreamMessage{
Message: &cvms.ClientStreamMessage_RunRes{
RunRes: &cvms.RunResponse{
ComputationId: "test-id",
},
},
}
client.sendMessage(msg)
select {
case received := <-messageQueue:
assert.Equal(t, msg, received)
case <-time.After(1 * time.Second):
t.Fatal("Message not received")
}
}
// TestManagerClient_sendMessageTimeout tests sendMessage timeout when queue is full.
func TestManagerClient_sendMessageTimeout(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage) // No buffer
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
msg := &cvms.ClientStreamMessage{
Message: &cvms.ClientStreamMessage_RunRes{
RunRes: &cvms.RunResponse{
ComputationId: "test-id",
},
},
}
// Don't read from queue, so sendMessage will timeout
client.sendMessage(msg)
// Should complete without blocking
time.Sleep(100 * time.Millisecond)
}
// TestManagerClient_handleRunReqChunksWithRemoteSource tests handling run request with remote source.
func TestManagerClient_handleRunReqChunksWithRemoteSource(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
runReq := &cvms.ComputationRunReq{
Id: "test-id-remote",
Name: "test-computation",
Description: "test description",
Datasets: []*cvms.Dataset{
{
Hash: sha3.New256().Sum([]byte("test-dataset")),
Filename: "data.csv",
Source: &cvms.Source{
Type: "oci-image",
Url: "docker://registry.example.com/data:v1",
KbsResourcePath: "default/key/data-key",
Encrypted: true,
},
Decompress: true,
},
},
Algorithm: &cvms.Algorithm{
Hash: sha3.New256().Sum([]byte("test-algorithm")),
AlgoType: "python",
AlgoArgs: []string{"--verbose"},
Source: &cvms.Source{
Type: "oci-image",
Url: "docker://registry.example.com/algo:v1",
KbsResourcePath: "default/key/algo-key",
Encrypted: true,
},
},
Kbs: &cvms.KBSConfig{
Url: "https://kbs.example.com:8080",
Enabled: true,
},
ResultConsumers: []*cvms.ResultConsumer{
{
UserKey: []byte("test-consumer"),
},
},
}
runReqBytes, _ := proto.Marshal(runReq)
chunk := &cvms.ServerStreamMessage_RunReqChunks{
RunReqChunks: &cvms.RunReqChunks{
Id: "chunk-remote-1",
Data: runReqBytes,
IsLast: true,
},
}
mockSvc.On("State").Return("ReceivingManifest")
mockSvc.On("InitComputation", mock.Anything, mock.MatchedBy(func(c agent.Computation) bool {
// Verify KBS config is passed
if !c.KBS.Enabled || c.KBS.URL != "https://kbs.example.com:8080" {
return false
}
// Verify algorithm source is passed
if c.Algorithm.Source == nil ||
c.Algorithm.Source.URL != "docker://registry.example.com/algo:v1" ||
c.Algorithm.Source.KBSResourcePath != "default/key/algo-key" ||
!c.Algorithm.Source.Encrypted {
return false
}
// Verify algorithm type and args
if c.Algorithm.AlgoType != "python" || len(c.Algorithm.AlgoArgs) != 1 || c.Algorithm.AlgoArgs[0] != "--verbose" {
return false
}
// Verify dataset source is passed
if len(c.Datasets) != 1 ||
c.Datasets[0].Source == nil ||
c.Datasets[0].Source.URL != "docker://registry.example.com/data:v1" ||
c.Datasets[0].Filename != "data.csv" ||
!c.Datasets[0].Decompress {
return false
}
return true
})).Return(nil)
mockServerSvc.On("Start", mock.Anything, mock.Anything, mock.Anything).Return(nil)
err = client.handleRunReqChunks(context.Background(), chunk)
assert.NoError(t, err)
// Wait for the goroutine to finish
time.Sleep(100 * time.Millisecond)
mockSvc.AssertExpectations(t)
}
// TestManagerClient_handleRunReqChunksAlreadyProcessing tests skipping init when already processing.
func TestManagerClient_handleRunReqChunksAlreadyProcessing(t *testing.T) {
mockStream := new(mockStream)
mockSvc := new(mocks.Service)
mockServerSvc := new(servermocks.AgentServer)
messageQueue := make(chan *cvms.ClientStreamMessage, 10)
logger := mglog.NewMock()
grpcClient := new(clientmocks.Client)
client, err := NewClient(mockStream, mockSvc, messageQueue, logger, mockServerSvc, nil, t.TempDir(), func(ctx context.Context) (pkggrpc.Client, cvms.Service_ProcessClient, error) { return nil, nil, nil }, grpcClient)
assert.NoError(t, err)
runReq := &cvms.ComputationRunReq{
Id: "test-id-processing",
Name: "test-computation",
}
runReqBytes, _ := proto.Marshal(runReq)
chunk := &cvms.ServerStreamMessage_RunReqChunks{
RunReqChunks: &cvms.RunReqChunks{
Id: "chunk-processing-1",
Data: runReqBytes,
IsLast: true,
},
}
// Simulate agent already processing a computation
mockSvc.On("State").Return("Running")
err = client.handleRunReqChunks(context.Background(), chunk)
assert.NoError(t, err)
// Wait for the goroutine to finish
time.Sleep(50 * time.Millisecond)
// InitComputation should NOT be called since state is not ReceivingManifest
mockSvc.AssertNotCalled(t, "InitComputation")
}