NOISSUE - Support non-chunked computation requests and add KBS decryption for uploaded algorithms and datasets. (#608)

* feat: support non-chunked computation requests and add KBS decryption for uploaded algorithms and datasets. Signed-off-by: Sammy Oina <sammyoina@gmail.com> * feat: update dataset resolution logic with context-based index verification and add extensive service error handling tests Signed-off-by: Sammy Oina <sammyoina@gmail.com> * feat: add AES encryption script and update package sources to connector-mods fork Signed-off-by: Sammy Oina <sammyoina@gmail.com> * fix: update file permission syntax in encrypt.go Signed-off-by: Sammy Oina <sammyoina@gmail.com> --------- Signed-off-by: Sammy Oina <sammyoina@gmail.com>
2026-08-07 07:14:50 +00:00 · 2026-07-01 16:52:43 +03:00
parent 6169766666
commit 5c3561f85d
4 changed files with 258 additions and 39 deletions
@@ -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:
@@ -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 {
@@ -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)
+}
@@ -0,0 +1,74 @@
+// Copyright (c) Ultraviolet
+// SPDX-License-Identifier: Apache-2.0
+
+package main
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"fmt"
+	"io"
+	"os"
+)
+
+func main() {
+	if len(os.Args) < 4 {
+		fmt.Println("Usage: go run encrypt.go <key_file> <input_file> <output_file>")
+		os.Exit(1)
+	}
+
+	keyFile := os.Args[1]
+	inputFile := os.Args[2]
+	outputFile := os.Args[3]
+
+	// Read key
+	key, err := os.ReadFile(keyFile)
+	if err != nil {
+		fmt.Printf("Failed to read key file: %v\n", err)
+		os.Exit(1)
+	}
+	if len(key) != 32 {
+		fmt.Printf("Key must be 32 bytes, got %d\n", len(key))
+		os.Exit(1)
+	}
+
+	// Read plaintext
+	plaintext, err := os.ReadFile(inputFile)
+	if err != nil {
+		fmt.Printf("Failed to read input file: %v\n", err)
+		os.Exit(1)
+	}
+
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		fmt.Printf("Failed to create cipher: %v\n", err)
+		os.Exit(1)
+	}
+
+	aesgcm, err := cipher.NewGCM(block)
+	if err != nil {
+		fmt.Printf("Failed to create GCM: %v\n", err)
+		os.Exit(1)
+	}
+
+	nonce := make([]byte, 12)
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		fmt.Printf("Failed to generate nonce: %v\n", err)
+		os.Exit(1)
+	}
+
+	ciphertext := aesgcm.Seal(nil, nonce, plaintext, nil)
+
+	// Combine nonce + ciphertext + tag
+	// Seal returns ciphertext || tag, so we just append it to nonce
+	output := append(nonce, ciphertext...)
+
+	err = os.WriteFile(outputFile, output, 0o644)
+	if err != nil {
+		fmt.Printf("Failed to write output file: %v\n", err)
+		os.Exit(1)
+	}
+
+	fmt.Printf("Successfully encrypted %s to %s\n", inputFile, outputFile)
+}