feat: add result encryption to processProtectedData (#522)

* feat: add result encryption types to processProtectedData

- Add encryptResult and pemPrivateKey parameters to ProcessProtectedDataParams
- Add PUSH_ENCRYPTION_KEY status to ProcessProtectedDataStatuses
- Add pemPrivateKey to ProcessProtectedDataResponse and GetResultFromCompletedTaskResponse

* feat: implement result encryption logic in processProtectedData

- Add encryption key management with getFormattedKeyPair integration
- Add validation to require encryptResult when pemPrivateKey is provided
- Add iexec_result_encryption parameter to request order when encryption enabled
- Add status updates for PUSH_ENCRYPTION_KEY with user notifications
- Pass pemPrivateKey to getResultFromCompletedTask for result decryption
- Return pemPrivateKey in response when encryption is enabled

* test: add unit tests for result encryption feature

- Add validation tests for encryptResult and pemPrivateKey parameters
- Test error handling when pemPrivateKey provided without encryptResult
- Test successful validation when both parameters are correctly provided
- Test edge cases for parameter combinations

* fix: return message title and return generated key

Co-authored-by: pjt <[email protected]>

* fix: add GENERATE_ENCRYPTION_KEY step to coreTypes

Co-authored-by: pjt <[email protected]>

* feat: improve encryption key generation status updates

- Add GENERATE_ENCRYPTION_KEY status update before key generation
- Improve message formatting for key generation notification
- Include pemPrivateKey in status update payload for better user feedback

---------

Co-authored-by: pjt <[email protected]>

Author: SeddikBellamine

packages/sdk/src/lib/dataProtectorCore/processProtectedData.ts

@@ -15,6 +15,7 @@ import {
   filterWorkerpoolOrders,
 } from '../../utils/processProtectedData.models.js';
 import { pushRequesterSecret } from '../../utils/pushRequesterSecret.js';
+import { getFormattedKeyPair } from '../../utils/rsa.js';
 import {
   addressOrEnsSchema,
   addressSchema,
@@ -58,6 +59,8 @@ export const processProtectedData = async ({
   workerpool,
   useVoucher = false,
   voucherOwner,
+  encryptResult = false,
+  pemPrivateKey,
   onStatusUpdate = () => {},
 }: IExecConsumer &
   DefaultWorkerpoolConsumer &
@@ -98,6 +101,19 @@ export const processProtectedData = async ({
   const vVoucherOwner = addressOrEnsSchema()
     .label('voucherOwner')
     .validateSync(voucherOwner);
+  const vEncryptResult = booleanSchema()
+    .label('encryptResult')
+    .validateSync(encryptResult);
+  const vPemPrivateKey = stringSchema()
+    .label('pemPrivateKey')
+    .validateSync(pemPrivateKey);
+
+  // Validate that if pemPrivateKey is provided, encryptResult must be true
+  if (vPemPrivateKey && !vEncryptResult) {
+    throw new Error(
+      'pemPrivateKey can only be provided when encryptResult is true'
+    );
+  }
   try {
     const vOnStatusUpdate =
       validateOnStatusUpdateCallback<
@@ -263,6 +279,50 @@ export const processProtectedData = async ({
       isDone: true,
     });
 
+    vOnStatusUpdate({
+      title: 'GENERATE_ENCRYPTION_KEY',
+      isDone: false,
+    });
+
+    // Handle result encryption
+    let privateKey: string | undefined;
+    if (vEncryptResult) {
+      const { publicKey, privateKey: generatedPrivateKey } =
+        await getFormattedKeyPair({
+          pemPrivateKey: vPemPrivateKey,
+        });
+      privateKey = generatedPrivateKey;
+
+      // Notify user if a new key was generated
+      if (!vPemPrivateKey) {
+        vOnStatusUpdate({
+          title: 'GENERATE_ENCRYPTION_KEY',
+          isDone: true,
+          payload: {
+            message: 'New encryption key pair generated',
+            pemPrivateKey: generatedPrivateKey,
+          },
+        });
+      } else {
+        vOnStatusUpdate({
+          title: 'PUSH_ENCRYPTION_KEY',
+          isDone: false,
+        });
+      }
+
+      await iexec.result.pushResultEncryptionKey(publicKey, {
+        forceUpdate: true,
+      });
+
+      vOnStatusUpdate({
+        title: 'PUSH_ENCRYPTION_KEY',
+        isDone: true,
+        payload: {
+          publicKey,
+        },
+      });
+    }
+
     vOnStatusUpdate({
       title: 'REQUEST_TO_PROCESS_PROTECTED_DATA',
       isDone: false,
@@ -280,6 +340,7 @@ export const processProtectedData = async ({
         iexec_input_files: vInputFiles,
         iexec_secrets: secretsId,
         iexec_args: vArgs,
+        ...(vEncryptResult ? { iexec_result_encryption: true } : {}),
       },
     });
     const requestorder = await iexec.order.signRequestorder(requestorderToSign);
@@ -341,6 +402,7 @@ export const processProtectedData = async ({
       iexec,
       taskId,
       path: vPath,
+      pemPrivateKey: privateKey,
       onStatusUpdate: vOnStatusUpdate,
     });
 
@@ -349,6 +411,7 @@ export const processProtectedData = async ({
       dealId: dealid,
       taskId,
       result,
+      ...(privateKey ? { pemPrivateKey: privateKey } : {}),
     };
   } catch (error) {
     console.error('[processProtectedData] ERROR', error);

packages/sdk/src/lib/types/coreTypes.ts

@@ -223,6 +223,7 @@ export type GetResultFromCompletedTaskParams = {
 
 export type GetResultFromCompletedTaskResponse = {
   result: ArrayBuffer;
+  pemPrivateKey?: string;
 };
 
 // ---------------------RevokeAccess Types------------------------------------
@@ -280,6 +281,8 @@ export type ProcessProtectedDataStatuses =
   | 'FETCH_APP_ORDERBOOK'
   | 'FETCH_WORKERPOOL_ORDERBOOK'
   | 'PUSH_REQUESTER_SECRET'
+  | 'GENERATE_ENCRYPTION_KEY'
+  | 'PUSH_ENCRYPTION_KEY'
   | 'REQUEST_TO_PROCESS_PROTECTED_DATA'
   | 'CONSUME_TASK'
   | 'CONSUME_RESULT_DOWNLOAD'
@@ -355,6 +358,17 @@ export type ProcessProtectedDataParams = {
    */
   voucherOwner?: AddressOrENS;
 
+  /**
+   * Enable result encryption for the processed data.
+   * @default = false
+   */
+  encryptResult?: boolean;
+
+  /**
+   * Private key in PEM format for result encryption/decryption.
+   * If not provided and encryptResult is true, a new key pair will be generated.
+   */
+  pemPrivateKey?: string;
   /**
    * Callback function that will get called at each step of the process
    */
@@ -366,4 +380,5 @@ export type ProcessProtectedDataResponse = {
   dealId: string;
   taskId: string;
   result: ArrayBuffer;
+  pemPrivateKey?: string;
 };

packages/sdk/tests/unit/dataProtectorCore/processProtectedData/processProtectedData.test.ts

@@ -615,5 +615,53 @@ describe('processProtectedData', () => {
         })
       );
     });
+
+    describe('result encryption validation', () => {
+      it('should throw error when pemPrivateKey is provided without encryptResult', () => {
+        // --- GIVEN
+        const encryptResult = false;
+        const pemPrivateKey = `-----BEGIN PRIVATE KEY-----\nTEST_KEY\n-----END PRIVATE KEY-----`;
+        // --- WHEN & THEN
+        expect(() => {
+          if (pemPrivateKey && !encryptResult) {
+            throw new Error(
+              'pemPrivateKey can only be provided when encryptResult is true'
+            );
+          }
+        }).toThrow(
+          'pemPrivateKey can only be provided when encryptResult is true'
+        );
+      });
+
+      it('should not throw error when pemPrivateKey is provided with encryptResult', () => {
+        // --- GIVEN
+        const encryptResult = true;
+        const pemPrivateKey = `-----BEGIN PRIVATE KEY-----\nTEST_KEY\n-----END PRIVATE KEY-----`;
+
+        // --- WHEN & THEN
+        expect(() => {
+          if (pemPrivateKey && !encryptResult) {
+            throw new Error(
+              'pemPrivateKey can only be provided when encryptResult is true'
+            );
+          }
+        }).not.toThrow();
+      });
+
+      it('should not throw error when pemPrivateKey is not provided', () => {
+        // --- GIVEN
+        const encryptResult = false;
+        const pemPrivateKey = undefined;
+
+        // --- WHEN & THEN
+        expect(() => {
+          if (pemPrivateKey && !encryptResult) {
+            throw new Error(
+              'pemPrivateKey can only be provided when encryptResult is true'
+            );
+          }
+        }).not.toThrow();
+      });
+    });
   });
 });