feat: vaults agent handlers are using raw streams now

tegefaulkes · tegefaulkes · commit b4f49a95bd2f · 2023-08-16T17:06:29.000+10:00
diff --git a/src/agent/handlers/clientManifest.ts b/src/agent/handlers/clientManifest.ts
@@ -2,12 +2,10 @@ import type { AgentRPCRequestParams, AgentRPCResponseResult } from '../types';
 import type {
   AgentClaimMessage,
   ClaimIdMessage,
-  GitPackMessage,
   HolePunchRelayMessage,
   NodeAddressMessage,
   NodeIdMessage,
   SignedNotificationEncoded,
-  VaultsGitPackGetMessage,
   VaultsScanMessage,
 } from './types';
 import {
@@ -44,10 +42,7 @@ const notificationsSend = new UnaryCaller<
 
 const vaultsGitInfoGet = new RawCaller();
 
-const vaultsGitPackGet = new ServerCaller<
-  AgentRPCRequestParams<VaultsGitPackGetMessage>,
-  AgentRPCResponseResult<GitPackMessage>
->();
+const vaultsGitPackGet = new RawCaller();
 
 const vaultsScan = new ServerCaller<
   AgentRPCRequestParams,
diff --git a/src/agent/handlers/types.ts b/src/agent/handlers/types.ts
@@ -40,21 +40,3 @@ export type VaultInfo = {
 export type VaultsScanMessage = VaultInfo & {
   vaultPermissions: Array<VaultAction>;
 };
-
-export type VaultsGitInfoGetMessage = {
-  vaultNameOrId: VaultIdEncoded | VaultName;
-  action: VaultAction;
-};
-
-export type GitPackMessage = {
-  /**
-   * Chunk of data in binary form;
-   */
-  chunk: string;
-};
-
-export type VaultsGitPackGetMessage = {
-  body: string;
-  nameOrId: VaultIdEncoded | VaultName;
-  vaultAction: VaultAction;
-};
diff --git a/src/agent/handlers/vaultsGitInfoGet.ts b/src/agent/handlers/vaultsGitInfoGet.ts
@@ -30,12 +30,7 @@ class VaultsGitInfoGetHandler extends RawHandler<{
   ): Promise<[JSONValue, ReadableStream<Uint8Array>]> {
     const { db, vaultManager, acl } = this.container;
     const [headerMessage, inputStream] = input;
-    const readableProm = (async () => {
-      for await (const _ of inputStream) {
-        // Input stream is not used here, wait for finish.
-        // It should be closed by the caller immediately
-      }
-    })();
+    await inputStream.cancel();
     const params = headerMessage.params;
     if (params == null || !utils.isObject(params)) never();
     if (
@@ -82,29 +77,30 @@ class VaultsGitInfoGetHandler extends RawHandler<{
           )}`,
         );
       }
-
       return {
         vaultId,
         vaultName,
       };
     });
 
-    // TODO: Needs to handle cancellation
+    let handleInfoRequestGen: AsyncGenerator<Buffer>;
     const stream = new ReadableStream({
-      start: async (controller) => {
-        for await (const buffer of vaultManager.handleInfoRequest(
-          data.vaultId,
-        )) {
-          if (buffer != null) {
-            controller.enqueue(buffer);
-          } else {
-            break;
-          }
+      start: async () => {
+        handleInfoRequestGen = vaultManager.handleInfoRequest(data.vaultId);
+      },
+      pull: async (controller) => {
+        const result = await handleInfoRequestGen.next();
+        if (result.done) {
+          controller.close();
+          return;
+        } else {
+          controller.enqueue(result.value);
         }
-        controller.close();
+      },
+      cancel: async (reason) => {
+        await handleInfoRequestGen.throw(reason).catch(() => {});
       },
     });
-    await readableProm;
     return [
       {
         vaultName: data.vaultName,
diff --git a/src/agent/handlers/vaultsGitPackGet.ts b/src/agent/handlers/vaultsGitPackGet.ts
@@ -1,106 +1,108 @@
 import type { DB } from '@matrixai/db';
-import type { GitPackMessage, VaultsGitPackGetMessage } from './types';
-import type { AgentRPCRequestParams, AgentRPCResponseResult } from '../types';
-import type { VaultAction, VaultName } from '../../vaults/types';
+import type { VaultName } from '../../vaults/types';
 import type VaultManager from '../../vaults/VaultManager';
 import type ACL from '../../acl/ACL';
+import type { JSONValue } from '../../types';
+import type { PassThrough } from 'readable-stream';
+import type { JSONRPCRequest } from '../../rpc/types';
+import { ReadableStream } from 'stream/web';
+import * as utils from '../../utils';
 import * as agentErrors from '../errors';
 import * as agentUtils from '../utils';
 import * as nodesUtils from '../../nodes/utils';
 import * as vaultsUtils from '../../vaults/utils';
 import * as vaultsErrors from '../../vaults/errors';
-import { validateSync } from '../../validation';
-import { matchSync } from '../../utils';
+import { never } from '../../utils';
 import * as validationUtils from '../../validation/utils';
-import { ServerHandler } from '../../rpc/handlers';
+import { RawHandler } from '../../rpc/handlers';
 
-// TODO: This needs to be a raw handler
-class VaultsGitPackGetHandler extends ServerHandler<
-  {
-    vaultManager: VaultManager;
-    acl: ACL;
-    db: DB;
-  },
-  AgentRPCRequestParams<VaultsGitPackGetMessage>,
-  AgentRPCResponseResult<GitPackMessage>
-> {
-  public async *handle(
-    input: AgentRPCRequestParams<VaultsGitPackGetMessage>,
+class VaultsGitPackGetHandler extends RawHandler<{
+  vaultManager: VaultManager;
+  acl: ACL;
+  db: DB;
+}> {
+  public async handle(
+    input: [JSONRPCRequest, ReadableStream<Uint8Array>],
     _cancel,
     meta,
-  ): AsyncGenerator<AgentRPCResponseResult<GitPackMessage>> {
+  ): Promise<[JSONValue, ReadableStream<Uint8Array>]> {
     const { vaultManager, acl, db } = this.container;
+    const [headerMessage, inputStream] = input;
     const requestingNodeId = agentUtils.nodeIdFromMeta(meta);
     if (requestingNodeId == null) {
       throw new agentErrors.ErrorAgentNodeIdMissing();
     }
     const nodeIdEncoded = nodesUtils.encodeNodeId(requestingNodeId);
-    const nameOrId = input.nameOrId;
-    yield* db.withTransactionG(async function* (
-      tran,
-    ): AsyncGenerator<AgentRPCResponseResult<GitPackMessage>> {
-      const vaultIdFromName = await vaultManager.getVaultId(
-        nameOrId as VaultName,
-        tran,
-      );
-      const vaultId = vaultIdFromName ?? vaultsUtils.decodeVaultId(nameOrId);
-      if (vaultId == null) {
-        throw new vaultsErrors.ErrorVaultsVaultUndefined();
-      }
-      const {
-        actionType,
-      }: {
-        actionType: VaultAction;
-      } = validateSync(
-        (keyPath, value) => {
-          return matchSync(keyPath)(
-            [['actionType'], () => validationUtils.parseVaultAction(value)],
-            () => value,
-          );
-        },
-        {
-          actionType: input.vaultAction,
-        },
-      );
-      // Checking permissions
-      const permissions = await acl.getNodePerm(requestingNodeId, tran);
-      const vaultPerms = permissions?.vaults[vaultId];
-      if (vaultPerms?.[actionType] !== null) {
-        throw new vaultsErrors.ErrorVaultsPermissionDenied(
-          `${nodeIdEncoded} does not have permission to ${actionType} from vault ${vaultsUtils.encodeVaultId(
-            vaultId,
-          )}`,
+    const params = headerMessage.params;
+    if (params == null || !utils.isObject(params)) never();
+    if (!('nameOrId' in params) || typeof params.nameOrId != 'string') {
+      never();
+    }
+    if (!('vaultAction' in params) || typeof params.vaultAction != 'string') {
+      never();
+    }
+    const nameOrId = params.nameOrId;
+    const actionType = validationUtils.parseVaultAction(params.vaultAction);
+    const [vaultIdFromName, permissions] = await db.withTransactionF(
+      async (tran) => {
+        const vaultIdFromName = await vaultManager.getVaultId(
+          nameOrId as VaultName,
+          tran,
         );
-      }
-      const [sideBand, progressStream] = await vaultManager.handlePackRequest(
-        vaultId,
-        Buffer.from(input.body, 'utf-8'),
-        tran,
+        const permissions = await acl.getNodePerm(requestingNodeId, tran);
+
+        return [vaultIdFromName, permissions];
+      },
+    );
+    const vaultId = vaultIdFromName ?? vaultsUtils.decodeVaultId(nameOrId);
+    if (vaultId == null) {
+      throw new vaultsErrors.ErrorVaultsVaultUndefined();
+    }
+    // Checking permissions
+    const vaultPerms = permissions?.vaults[vaultId];
+    if (vaultPerms?.[actionType] !== null) {
+      throw new vaultsErrors.ErrorVaultsPermissionDenied(
+        `${nodeIdEncoded} does not have permission to ${actionType} from vault ${vaultsUtils.encodeVaultId(
+          vaultId,
+        )}`,
       );
-      yield {
-        chunk: Buffer.from('0008NAK\n').toString('binary'),
-      };
-      const responseBuffers: Uint8Array[] = [];
-      // FIXME: this WHOLE thing needs to change, why are we streaming when we send monolithic messages?
-      const result = await new Promise<string>((resolve, reject) => {
+    }
+
+    // Getting data
+    let sideBand: PassThrough;
+    let progressStream: PassThrough;
+    const outputStream = new ReadableStream({
+      start: async (controller) => {
+        const body = new Array<Uint8Array>();
+        for await (const message of inputStream) {
+          body.push(message);
+        }
+        [sideBand, progressStream] = await vaultManager.handlePackRequest(
+          vaultId,
+          Buffer.concat(body),
+        );
+        controller.enqueue(Buffer.from('0008NAK\n'));
         sideBand.on('data', async (data: Uint8Array) => {
-          responseBuffers.push(data);
+          controller.enqueue(data);
+          sideBand.pause();
         });
         sideBand.on('end', async () => {
-          const result = Buffer.concat(responseBuffers).toString('binary');
-          resolve(result);
+          controller.close();
         });
-        sideBand.on('error', (err) => {
-          reject(err);
+        sideBand.on('error', (e) => {
+          controller.error(e);
         });
         progressStream.write(Buffer.from('0014progress is at 50%\n'));
         progressStream.end();
-      });
-      yield {
-        chunk: result,
-      };
+      },
+      pull: () => {
+        sideBand.resume();
+      },
+      cancel: (e) => {
+        sideBand.destroy(e);
+      },
     });
-    return;
+    return [null, outputStream];
   }
 }
 
diff --git a/src/rpc/RPCServer.ts b/src/rpc/RPCServer.ts
@@ -586,6 +586,7 @@ class RPCServer extends EventTarget {
           id: null,
         };
         await headerWriter.write(Buffer.from(JSON.stringify(rpcErrorMessage)));
+        await headerWriter.close();
         // Clean up and return
         timer.cancel(cleanupReason);
         abortController.signal.removeEventListener('abort', handleAbort);
diff --git a/src/vaults/VaultInternal.ts b/src/vaults/VaultInternal.ts
@@ -767,13 +767,11 @@ class VaultInternal {
       typeof vaultNameOrId === 'string'
         ? vaultNameOrId
         : vaultsUtils.encodeVaultId(vaultNameOrId);
-    const response = await client.methods.vaultsGitInfoGet({
+    const vaultsGitInfoGetStream = await client.methods.vaultsGitInfoGet({
       vaultNameOrId: vaultNameOrId_,
       action: vaultAction,
     });
-    await response.writable.close();
-
-    const result = response.meta?.result;
+    const result = vaultsGitInfoGetStream.meta?.result;
     if (result == null || !utils.isObject(result)) never();
     if (!('vaultName' in result) || typeof result.vaultName != 'string') {
       never();
@@ -789,7 +787,7 @@ class VaultInternal {
 
     // Collect the response buffers from the GET request
     const infoResponse: Uint8Array[] = [];
-    for await (const chunk of response.readable) {
+    for await (const chunk of vaultsGitInfoGetStream.readable) {
       infoResponse.push(chunk);
     }
     return [
@@ -816,13 +814,15 @@ class VaultInternal {
           };
         } else if (method === 'POST') {
           const responseBuffers: Array<Uint8Array> = [];
-          const stream = await client.methods.vaultsGitPackGet({
-            body: body[0].toString('binary'),
+          const vaultsGitPackGetStream = await client.methods.vaultsGitPackGet({
             nameOrId: result.vaultIdEncoded as string,
             vaultAction,
           });
-          for await (const value of stream) {
-            responseBuffers.push(Buffer.from(value.chunk, 'binary'));
+          const writer = vaultsGitPackGetStream.writable.getWriter();
+          await writer.write(body[0]);
+          await writer.close();
+          for await (const value of vaultsGitPackGetStream.readable) {
+            responseBuffers.push(value);
           }
           return {
             url: url,
diff --git a/tests/vaults/VaultManager.test.ts b/tests/vaults/VaultManager.test.ts
@@ -469,8 +469,7 @@ describe('VaultManager', () => {
       await vaultManager?.destroy();
     }
   });
-  // TODO: disabled until feature is addressed in agent migration stage 2
-  describe.skip('with remote agents', () => {
+  describe('with remote agents', () => {
     let allDataDir: string;
     let keyRing: KeyRing;
     let nodeGraph: NodeGraph;
@@ -653,7 +652,6 @@ describe('VaultManager', () => {
           localNodeId,
           'pull',
         );
-
         await vaultManager.cloneVault(remoteKeynode1Id, vaultName);
         const vaultId = await vaultManager.getVaultId(vaultName);
         if (vaultId === undefined) fail('VaultId is not found.');
