COCOS-584 - Support multiple kbs

agent/TESTING_REMOTE_RESOURCES.md

@@ -280,17 +280,22 @@ go build -o build/cvms-test ./test/cvms/main.go
 HOST=$HOST_IP PORT=7001 ./build/cvms-test \
   -public-key-path ./public.pem \
   -attested-tls-bool false \
-  -kbs-url http://$HOST_IP:8080 \
   -algo-type python \
   -algo-source-url docker://$HOST_IP:5000/encrypted-lin-reg:v1.0 \
   -algo-kbs-path default/key/algo-key \
+  -algo-kbs-url http://$HOST_IP:8080 \
   -algo-hash $ALGO_HASH \
   -algo-args datasets/dataset_0.csv \
   -dataset-source-urls docker://$HOST_IP:5000/encrypted-iris:v1.0 \
   -dataset-kbs-paths default/key/dataset-key \
+  -dataset-kbs-urls http://$HOST_IP:8080 \
   -dataset-hash $DATASET_HASH
 ```
 
+> [!NOTE]
+> You must specify the KBS URL for each encrypted resource using `-algo-kbs-url` and `-dataset-kbs-urls`. A global KBS is no longer supported.
+
+
 ### 3. Create VM via CLI (Host)
 
 ```bash
@@ -356,17 +361,31 @@ The CVMS server sends this manifest to the agent:
     "type": "oci-image",
     "uri": "docker://localhost:5000/encrypted-lin-reg:v1.0",
     "encrypted": true,
-    "kbs_resource_path": "default/key/algo-key"
+    "kbs_resource_path": "default/key/algo-key",
+    "kbs": {
+      "url": "http://192.168.100.15:8080",
+      "enabled": true
+    }
   },
   "datasets": [
     {
-      "type": "oci-image",
-      "uri": "docker://localhost:5000/encrypted-iris:v1.0",
-      "encrypted": true,
-      "kbs_resource_path": "default/key/dataset-key"
+      "filename": "iris.csv",
+      "source": {
+        "type": "oci-image",
+        "url": "docker://localhost:5000/encrypted-iris:v1.0",
+        "encrypted": true,
+        "kbs_resource_path": "default/key/dataset-key"
+      },
+      "kbs": {
+        "url": "http://192.168.100.20:8080",
+        "enabled": true
+      }
     }
   ],
-  "kbs_url": "http://192.168.100.15:8080"
+  "kbs": {
+    "url": "http://192.168.100.15:8080",
+    "enabled": true
+  }
 }
 ```
 

agent/auth/auth_test.go

@@ -44,7 +44,7 @@ func TestAuthenticateUser(t *testing.T) {
 	manifest := agent.Computation{
 		ResultConsumers: []agent.ResultConsumer{{UserKey: resultConsumerPubKey}},
 		Datasets:        []agent.Dataset{{UserKey: dataProviderPubKey}},
-		Algorithm:       agent.Algorithm{UserKey: algorithmProviderPubKey},
+		Algorithm:       &agent.Algorithm{UserKey: algorithmProviderPubKey},
 	}
 
 	auth, err := New(manifest)

agent/computations.go

@@ -45,9 +45,8 @@ type Computation struct {
 	Name            string           `json:"name,omitempty"`
 	Description     string           `json:"description,omitempty"`
 	Datasets        Datasets         `json:"datasets,omitempty"`
-	Algorithm       Algorithm        `json:"algorithm,omitempty"`
+	Algorithm       *Algorithm       `json:"algorithm,omitempty"`
 	ResultConsumers []ResultConsumer `json:"result_consumers,omitempty"`
-	KBS             KBSConfig        `json:"kbs,omitempty"`
 }
 
 type ResultConsumer struct {
@@ -69,6 +68,7 @@ type Dataset struct {
 	Filename   string          `json:"filename,omitempty"`
 	Source     *ResourceSource `json:"source,omitempty"` // Optional remote source
 	Decompress bool            `json:"decompress,omitempty"`
+	KBS        *KBSConfig      `json:"kbs,omitempty"`
 }
 
 type Datasets []Dataset
@@ -81,6 +81,7 @@ type Algorithm struct {
 	Source       *ResourceSource `json:"source,omitempty"` // Optional remote source
 	AlgoType     string          `json:"algo_type,omitempty"`
 	AlgoArgs     []string        `json:"algo_args,omitempty"`
+	KBS          *KBSConfig      `json:"kbs,omitempty"`
 }
 
 type ManifestIndexKey struct{}

agent/cvms/api/grpc/client.go

@@ -234,9 +234,10 @@ func (client *CVMSClient) executeRun(ctx context.Context, runReq *cvms.Computati
 	}
 
 	if runReq.Algorithm != nil {
-		ac.Algorithm = agent.Algorithm{
-			Hash:    [32]byte(runReq.Algorithm.Hash),
-			UserKey: runReq.Algorithm.UserKey,
+		ac.Algorithm = &agent.Algorithm{
+			Hash:     [32]byte(runReq.Algorithm.Hash),
+			UserKey:  runReq.Algorithm.UserKey,
+			AlgoType: runReq.Algorithm.AlgoType,
 		}
 		// Copy remote source if configured
 		if runReq.Algorithm.Source != nil {
@@ -246,8 +247,13 @@ func (client *CVMSClient) executeRun(ctx context.Context, runReq *cvms.Computati
 				Encrypted:       runReq.Algorithm.Source.Encrypted,
 			}
 		}
-		ac.Algorithm.AlgoType = runReq.Algorithm.AlgoType
 		ac.Algorithm.AlgoArgs = runReq.Algorithm.AlgoArgs
+		if runReq.Algorithm.Kbs != nil {
+			ac.Algorithm.KBS = &agent.KBSConfig{
+				URL:     runReq.Algorithm.Kbs.Url,
+				Enabled: runReq.Algorithm.Kbs.Enabled,
+			}
+		}
 	}
 
 	for _, ds := range runReq.Datasets {
@@ -265,6 +271,12 @@ func (client *CVMSClient) executeRun(ctx context.Context, runReq *cvms.Computati
 			}
 		}
 		dataset.Decompress = ds.Decompress
+		if ds.Kbs != nil {
+			dataset.KBS = &agent.KBSConfig{
+				URL:     ds.Kbs.Url,
+				Enabled: ds.Kbs.Enabled,
+			}
+		}
 		ac.Datasets = append(ac.Datasets, dataset)
 	}
 
@@ -274,14 +286,6 @@ func (client *CVMSClient) executeRun(ctx context.Context, runReq *cvms.Computati
 		})
 	}
 
-	// Copy KBS configuration
-	if runReq.Kbs != nil {
-		ac.KBS = agent.KBSConfig{
-			URL:     runReq.Kbs.Url,
-			Enabled: runReq.Kbs.Enabled,
-		}
-	}
-
 	// Check if the agent is in the correct state to initialize a new computation.
 	// If the agent is already processing this computation (e.g., after a reconnection),
 	// skip initialization to avoid state errors.

agent/cvms/api/grpc/client_test.go

@@ -554,11 +554,12 @@ func TestManagerClient_handleRunReqChunksWithRemoteSource(t *testing.T) {
 				KbsResourcePath: "default/key/algo-key",
 				Encrypted:       true,
 			},
+			Kbs: &cvms.KBSConfig{
+				Url:     "https://kbs.example.com:8080",
+				Enabled: true,
+			},
 		},
-		Kbs: &cvms.KBSConfig{
-			Url:     "https://kbs.example.com:8080",
-			Enabled: true,
-		},
+
 		ResultConsumers: []*cvms.ResultConsumer{
 			{
 				UserKey: []byte("test-consumer"),
@@ -577,8 +578,8 @@ func TestManagerClient_handleRunReqChunksWithRemoteSource(t *testing.T) {
 
 	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" {
+		// Verify Algorithm KBS config is passed
+		if c.Algorithm.KBS == nil || !c.Algorithm.KBS.Enabled || c.Algorithm.KBS.URL != "https://kbs.example.com:8080" {
 			return false
 		}
 		// Verify algorithm source is passed

agent/cvms/cvms.proto

@@ -92,7 +92,6 @@ message ComputationRunReq {
   Algorithm algorithm = 5;
   repeated ResultConsumer result_consumers = 6;
   AgentConfig agent_config = 7;
-  KBSConfig kbs = 8; // Optional KBS configuration for remote resources
 }
 
 message ResultConsumer {
@@ -105,6 +104,7 @@ message Dataset {
   string filename = 3;
   Source source = 4; // Optional remote source for encrypted dataset
   bool decompress = 5;
+  KBSConfig kbs = 6; // Optional KBS configuration override
 }
 
 message Algorithm {
@@ -113,6 +113,7 @@ message Algorithm {
   Source source = 3; // Optional remote source for encrypted algorithm
   string algo_type = 4;
   repeated string algo_args = 5;
+  KBSConfig kbs = 6; // Optional KBS configuration override
 }
 
 message Source {