diff --git a/agent/cvms/api/grpc/client.go b/agent/cvms/api/grpc/client.go index f6b41c0d..c787844a 100644 --- a/agent/cvms/api/grpc/client.go +++ b/agent/cvms/api/grpc/client.go @@ -178,6 +178,9 @@ func (client *CVMSClient) processIncomingMessage(ctx context.Context, req *cvms. switch mes := req.Message.(type) { case *cvms.ServerStreamMessage_RunReqChunks: return client.handleRunReqChunks(ctx, mes) + case *cvms.ServerStreamMessage_RunReq: + client.logger.Info("Starting computation execution from non-chunked run request", "computationId", mes.RunReq.Id, "name", mes.RunReq.Name) + go client.executeRun(ctx, mes.RunReq) case *cvms.ServerStreamMessage_StopComputation: go client.handleStopComputation(ctx, mes) case *cvms.ServerStreamMessage_AgentStateReq: diff --git a/agent/service.go b/agent/service.go index 0588459f..a9313088 100644 --- a/agent/service.go +++ b/agent/service.go @@ -416,7 +416,7 @@ func (as *agentService) downloadAlgorithmIfRemote(state statemachine.State) { "kbs_enabled", kbsEnabled) // Check if algorithm should be downloaded from remote source - if as.computation.Algorithm.Source != nil && kbsEnabled { + if as.computation.Algorithm.Source != nil && kbsEnabled && as.computation.Algorithm.Source.URL != "" { as.logger.Info("downloading algorithm from remote source", "url", as.computation.Algorithm.Source.URL, "kbs_resource_path", as.computation.Algorithm.Source.KBSResourcePath, @@ -536,7 +536,7 @@ func (as *agentService) downloadDatasetsIfRemote(state statemachine.State) { hasRemoteDatasets := false for _, d := range as.computation.Datasets { kbsEnabled := d.KBS != nil && d.KBS.Enabled - if d.Source != nil && kbsEnabled { + if d.Source != nil && kbsEnabled && d.Source.URL != "" { hasRemoteDatasets = true break } @@ -557,7 +557,7 @@ func (as *agentService) downloadDatasetsIfRemote(state statemachine.State) { kbsURL = d.KBS.URL } - if d.Source != nil && kbsEnabled { + if d.Source != nil && kbsEnabled && d.Source.URL != "" { as.logger.Info("downloading dataset from remote source", "filename", d.Filename, "kbs_url", kbsURL) res, err := as.downloadAndDecryptResource(ctx, d.Source, kbsURL, "dataset") @@ -944,7 +944,7 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error { kbsURL = as.computation.Algorithm.KBS.URL } - if as.computation.Algorithm.Source != nil && kbsEnabled { + if as.computation.Algorithm.Source != nil && kbsEnabled && as.computation.Algorithm.Source.URL != "" { as.logger.Info("downloading algorithm from remote source", "kbs_url", kbsURL) res, err := as.downloadAndDecryptResource(ctx, as.computation.Algorithm.Source, kbsURL, "algorithm") @@ -957,6 +957,19 @@ func (as *agentService) Algo(ctx context.Context, algo Algorithm) error { } else { // Use directly uploaded algorithm algoData = algo.Algorithm + + if as.computation.Algorithm.Source != nil && as.computation.Algorithm.Source.Encrypted && kbsEnabled { + as.logger.Info("directly uploaded algorithm is encrypted, retrieving key from KBS") + key, err := as.getKeyFromKBS(ctx, kbsURL, as.computation.Algorithm.Source.KBSResourcePath) + if err != nil { + return fmt.Errorf("failed to retrieve key from KBS for uploaded algorithm: %w", err) + } + decrypted, err := resource.DecryptData(algoData, key) + if err != nil { + return fmt.Errorf("failed to decrypt uploaded algorithm: %w", err) + } + algoData = decrypted + } } hash := sha3.Sum256(algoData) @@ -1032,7 +1045,7 @@ func (as *agentService) Data(ctx context.Context, dataset Dataset) error { kbsURL = d.KBS.URL } - if d.Source != nil && kbsEnabled { + if d.Source != nil && kbsEnabled && d.Source.URL != "" { as.logger.Info("downloading dataset from remote source", "filename", d.Filename, "kbs_url", kbsURL) downloadedData, err := as.downloadAndDecryptResource(ctx, d.Source, kbsURL, "dataset") @@ -1051,44 +1064,85 @@ func (as *agentService) Data(ctx context.Context, dataset Dataset) error { if matchedIndex == -1 { datasetData = dataset.Dataset datasetFilename = dataset.Filename + + index, ok := IndexFromContext(ctx) + if ok { + if index < 0 || index >= len(as.computation.Datasets) { + return ErrUndeclaredDataset + } + if as.computation.Datasets[index].Filename != datasetFilename { + return ErrFileNameMismatch + } + matchedIndex = index + } else { + matchedIndex = -1 + for i, d := range as.computation.Datasets { + if d.Filename == datasetFilename { + matchedIndex = i + break + } + } + if matchedIndex == -1 { + return ErrUndeclaredDataset + } + } + } else { + remoteIndex := -1 + for i, d := range as.computation.Datasets { + if d.Filename == datasetFilename { + remoteIndex = i + break + } + } + if remoteIndex == -1 { + return ErrUndeclaredDataset + } + matchedIndex = remoteIndex + } + + d := as.computation.Datasets[matchedIndex] + + kbsEnabled := d.KBS != nil && d.KBS.Enabled + kbsURL := "" + if d.KBS != nil { + kbsURL = d.KBS.URL + } + if d.Source != nil && d.Source.Encrypted && kbsEnabled { + as.logger.Info("directly uploaded dataset is encrypted, retrieving key from KBS", "filename", d.Filename) + key, err := as.getKeyFromKBS(ctx, kbsURL, d.Source.KBSResourcePath) + if err != nil { + return fmt.Errorf("failed to retrieve key from KBS for dataset %s: %w", d.Filename, err) + } + decrypted, err := resource.DecryptData(datasetData, key) + if err != nil { + return fmt.Errorf("failed to decrypt dataset %s: %w", d.Filename, err) + } + datasetData = decrypted } hash := sha3.Sum256(datasetData) - - matched := false - for i, d := range as.computation.Datasets { - if hash == d.Hash { - if d.Filename != "" && d.Filename != datasetFilename { - return ErrFileNameMismatch - } - - as.computation.Datasets = slices.Delete(as.computation.Datasets, i, i+1) - - if DecompressFromContext(ctx) { - if err := internal.UnzipFromMemory(datasetData, algorithm.DatasetsDir); err != nil { - return fmt.Errorf("error decompressing dataset: %v", err) - } - } else { - f, err := os.Create(fmt.Sprintf("%s/%s", algorithm.DatasetsDir, datasetFilename)) - if err != nil { - return fmt.Errorf("error creating dataset file: %v", err) - } - - if _, err := f.Write(datasetData); err != nil { - return fmt.Errorf("error writing dataset to file: %v", err) - } - if err := f.Close(); err != nil { - return fmt.Errorf("error closing file: %v", err) - } - } - - matched = true - break - } + if hash != d.Hash { + return ErrHashMismatch } - if !matched { - return ErrUndeclaredDataset + as.computation.Datasets = slices.Delete(as.computation.Datasets, matchedIndex, matchedIndex+1) + + if DecompressFromContext(ctx) { + if err := internal.UnzipFromMemory(datasetData, algorithm.DatasetsDir); err != nil { + return fmt.Errorf("error decompressing dataset: %v", err) + } + } else { + f, err := os.Create(fmt.Sprintf("%s/%s", algorithm.DatasetsDir, datasetFilename)) + if err != nil { + return fmt.Errorf("error creating dataset file: %v", err) + } + + if _, err := f.Write(datasetData); err != nil { + return fmt.Errorf("error writing dataset to file: %v", err) + } + if err := f.Close(); err != nil { + return fmt.Errorf("error closing file: %v", err) + } } if len(as.computation.Datasets) == 0 { diff --git a/agent/service_test.go b/agent/service_test.go index bf23e34e..c21fa85a 100644 --- a/agent/service_test.go +++ b/agent/service_test.go @@ -218,7 +218,7 @@ func TestData(t *testing.T) { { name: "Test dataset not declared in manifest", data: Dataset{ - Filename: datasetFile, + Filename: "undeclared.csv", }, err: ErrUndeclaredDataset, }, @@ -1808,3 +1808,91 @@ func TestInferSourceType(t *testing.T) { }) } } + +func TestInitComputation_Errors(t *testing.T) { + sm := &smmocks.StateMachine{} + sm.On("GetState").Return(ReceivingAlgorithm) // Not ReceivingManifest + + svc := &agentService{ + sm: sm, + } + + err := svc.InitComputation(context.Background(), Computation{}) + assert.ErrorIs(t, err, ErrStateNotReady) +} + +func TestAlgo_Errors(t *testing.T) { + t.Run("state not ready", func(t *testing.T) { + sm := &smmocks.StateMachine{} + sm.On("GetState").Return(ReceivingManifest) // Not ReceivingAlgorithm + + svc := &agentService{ + sm: sm, + } + + err := svc.Algo(context.Background(), Algorithm{}) + assert.ErrorIs(t, err, ErrStateNotReady) + }) + + t.Run("all manifest items received", func(t *testing.T) { + sm := &smmocks.StateMachine{} + sm.On("GetState").Return(ReceivingAlgorithm) + + svc := &agentService{ + sm: sm, + algoReceived: true, + } + + err := svc.Algo(context.Background(), Algorithm{}) + assert.ErrorIs(t, err, ErrAllManifestItemsReceived) + }) + + t.Run("undeclared algorithm", func(t *testing.T) { + sm := &smmocks.StateMachine{} + sm.On("GetState").Return(ReceivingAlgorithm) + + svc := &agentService{ + sm: sm, + algoReceived: false, + computation: Computation{ + Algorithm: nil, // Not declared + }, + } + + err := svc.Algo(context.Background(), Algorithm{}) + assert.ErrorIs(t, err, ErrUndeclaredAlgorithm) + }) +} + +func TestData_ErrorsExtra(t *testing.T) { + t.Run("all manifest items received", func(t *testing.T) { + sm := &smmocks.StateMachine{} + sm.On("GetState").Return(ReceivingData) + + svc := &agentService{ + sm: sm, + computation: Computation{ + Datasets: nil, // length 0 + }, + } + + err := svc.Data(context.Background(), Dataset{}) + assert.ErrorIs(t, err, ErrAllManifestItemsReceived) + }) +} + +func TestEnsureDir_Error(t *testing.T) { + tmpFile, err := os.CreateTemp("", "ensureDirTest") + require.NoError(t, err) + defer os.Remove(tmpFile.Name()) + tmpFile.Close() + + // ensureDir should fail because the parent path is a file, not a directory + err = ensureDir(filepath.Join(tmpFile.Name(), "subdir"), 0o755) + assert.Error(t, err) +} + +func TestKbsHTTPGet_Error(t *testing.T) { + _, err := kbsHTTPGet(context.Background(), "%%") + assert.Error(t, err) +} diff --git a/scripts/encrypt.go b/scripts/encrypt.go new file mode 100644 index 00000000..64296946 --- /dev/null +++ b/scripts/encrypt.go @@ -0,0 +1,74 @@ +// Copyright (c) Ultraviolet +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "crypto/aes" + "crypto/cipher" + "crypto/rand" + "fmt" + "io" + "os" +) + +func main() { + if len(os.Args) < 4 { + fmt.Println("Usage: go run encrypt.go ") + os.Exit(1) + } + + keyFile := os.Args[1] + inputFile := os.Args[2] + outputFile := os.Args[3] + + // Read key + key, err := os.ReadFile(keyFile) + if err != nil { + fmt.Printf("Failed to read key file: %v\n", err) + os.Exit(1) + } + if len(key) != 32 { + fmt.Printf("Key must be 32 bytes, got %d\n", len(key)) + os.Exit(1) + } + + // Read plaintext + plaintext, err := os.ReadFile(inputFile) + if err != nil { + fmt.Printf("Failed to read input file: %v\n", err) + os.Exit(1) + } + + block, err := aes.NewCipher(key) + if err != nil { + fmt.Printf("Failed to create cipher: %v\n", err) + os.Exit(1) + } + + aesgcm, err := cipher.NewGCM(block) + if err != nil { + fmt.Printf("Failed to create GCM: %v\n", err) + os.Exit(1) + } + + nonce := make([]byte, 12) + if _, err := io.ReadFull(rand.Reader, nonce); err != nil { + fmt.Printf("Failed to generate nonce: %v\n", err) + os.Exit(1) + } + + ciphertext := aesgcm.Seal(nil, nonce, plaintext, nil) + + // Combine nonce + ciphertext + tag + // Seal returns ciphertext || tag, so we just append it to nonce + output := append(nonce, ciphertext...) + + err = os.WriteFile(outputFile, output, 0o644) + if err != nil { + fmt.Printf("Failed to write output file: %v\n", err) + os.Exit(1) + } + + fmt.Printf("Successfully encrypted %s to %s\n", inputFile, outputFile) +}