mirror of
https://github.com/ultravioletrs/cocos.git
synced 2026-06-23 04:10:25 +00:00
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
* 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 toc28cefaeIncludes 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 tof6981ac5Includes 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>
654 lines
20 KiB
Go
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")
|
|
}
|