From f004871b4d18558dbe76ea60a526fde32ff6b14e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lloren=C3=A7=20Muntaner?= <llorenc.muntaner@gmail.com>
Date: Mon, 11 Jul 2022 18:31:17 +0200
Subject: [PATCH] feat: reuse signed request when reading state (#584)

* Feat: Reuse signed request when reading state
---
 docs/generated/changelog.html              | 11 ++++
 e2e/node/basic/basic.test.ts               | 36 +++++++++++++
 packages/agent/src/agent/api.ts            | 16 ++++++
 packages/agent/src/agent/http/http.test.ts | 61 +++++++++++++++++++++-
 packages/agent/src/agent/http/index.ts     | 22 +++++---
 packages/agent/src/polling/index.ts        |  8 ++-
 6 files changed, 145 insertions(+), 9 deletions(-)

diff --git a/docs/generated/changelog.html b/docs/generated/changelog.html
index 2982ac9ab..94a10ef28 100644
--- a/docs/generated/changelog.html
+++ b/docs/generated/changelog.html
@@ -13,6 +13,17 @@ <h1>Agent-JS Changelog</h1>
       <h2>Version 0.12.1</h2>
       <ul>
         <li>Adds UTF-8 as an encoding option for CanisterStatus custom paths</li>
+        <li>
+          Adds a public method "createReadStateRequest" that creates the request for "readState".
+        </li>
+        <li>
+          Add an extra parameter to "readState" to pass a created request. If this parameter is
+          passed, the method does the request directly without creating a new one.
+        </li>
+        <li>
+          Use the "createReadStateRequest" and the extra parameter when polling for the response to
+          avoid signing requests during polling.
+        </li>
         <li>
           Adds derivationOrigin to auth-client login to support the ability to login using the
           identity derived from a different origin. See
diff --git a/e2e/node/basic/basic.test.ts b/e2e/node/basic/basic.test.ts
index a085873f4..e42dae552 100644
--- a/e2e/node/basic/basic.test.ts
+++ b/e2e/node/basic/basic.test.ts
@@ -36,6 +36,42 @@ test('read_state', async () => {
   expect(Math.abs(time - now) < 5).toBe(true);
 });
 
+test('read_state with passed request', async () => {
+  const resolvedAgent = await agent;
+  const now = Date.now() / 1000;
+  const path = [new TextEncoder().encode('time')];
+  const canisterId = Principal.fromHex('00000000000000000001');
+  const request = await resolvedAgent.createReadStateRequest({ paths: [path] });
+  const response = await resolvedAgent.readState(
+    canisterId,
+    {
+      paths: [path],
+    },
+    undefined,
+    request,
+  );
+  if (resolvedAgent.rootKey == null) throw new Error(`The agent doesn't have a root key yet`);
+  const cert = await Certificate.create({
+    certificate: response.certificate,
+    rootKey: resolvedAgent.rootKey,
+    canisterId: canisterId,
+  });
+  expect(cert.lookup([new TextEncoder().encode('Time')])).toBe(undefined);
+  // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+  const rawTime = cert.lookup(path)!;
+  const decoded = IDL.decode(
+    [IDL.Nat],
+    new Uint8Array([
+      ...new TextEncoder().encode('DIDL\x00\x01\x7d'),
+      ...(new Uint8Array(rawTime) || []),
+    ]),
+  )[0];
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const time = Number(decoded as any) / 1e9;
+  // The diff between decoded time and local time is within 5s
+  expect(Math.abs(time - now) < 5).toBe(true);
+});
+
 test('createCanister', async () => {
   // Make sure this doesn't fail.
   await getManagementCanister({
diff --git a/packages/agent/src/agent/api.ts b/packages/agent/src/agent/api.ts
index 4f8d2a6b3..6323bd914 100644
--- a/packages/agent/src/agent/api.ts
+++ b/packages/agent/src/agent/api.ts
@@ -111,16 +111,32 @@ export interface Agent {
    */
   getPrincipal(): Promise<Principal>;
 
+  /**
+   * Create the request for the read state call.
+   * `readState` uses this internally.
+   * Useful to avoid signing the same request multiple times.
+   */
+  createReadStateRequest?(
+    options: ReadStateOptions,
+    identity?: Identity,
+    // eslint-disable-next-line  @typescript-eslint/no-explicit-any
+  ): Promise<any>;
+
   /**
    * Send a read state query to the replica. This includes a list of paths to return,
    * and will return a Certificate. This will only reject on communication errors,
    * but the certificate might contain less information than requested.
    * @param effectiveCanisterId A Canister ID related to this call.
    * @param options The options for this call.
+   * @param identity Identity for the call. If not specified, uses the instance identity.
+   * @param request The request to send in case it has already been created.
    */
   readState(
     effectiveCanisterId: Principal | string,
     options: ReadStateOptions,
+    identity?: Identity,
+    // eslint-disable-next-line  @typescript-eslint/no-explicit-any
+    request?: any,
   ): Promise<ReadStateResponse>;
 
   call(canisterId: Principal | string, fields: CallOptions): Promise<SubmitResponse>;
diff --git a/packages/agent/src/agent/http/http.test.ts b/packages/agent/src/agent/http/http.test.ts
index 21cf8d002..012662699 100644
--- a/packages/agent/src/agent/http/http.test.ts
+++ b/packages/agent/src/agent/http/http.test.ts
@@ -1,7 +1,13 @@
 import { HttpAgent, Nonce } from '../index';
 import * as cbor from '../../cbor';
 import { Expiry, makeNonceTransform } from './transforms';
-import { CallRequest, Envelope, makeNonce, SubmitRequestType } from './types';
+import {
+  CallRequest,
+  Envelope,
+  HttpAgentRequestTransformFn,
+  makeNonce,
+  SubmitRequestType,
+} from './types';
 import { Principal } from '@dfinity/principal';
 import { requestIdOf } from '../../request_id';
 
@@ -154,6 +160,59 @@ test('queries with the same content should have the same signature', async () =>
   expect(response3).toEqual(response4);
 });
 
+test('readState should not call transformers if request is passed', async () => {
+  const mockResponse = {
+    status: 'replied',
+    reply: { arg: new Uint8Array([]) },
+  };
+
+  const mockFetch: jest.Mock = jest.fn((resource, init) => {
+    const body = cbor.encode(mockResponse);
+    return Promise.resolve(
+      new Response(body, {
+        status: 200,
+      }),
+    );
+  });
+
+  const canisterIdent = '2chl6-4hpzw-vqaaa-aaaaa-c';
+  const nonce = new Uint8Array([0, 1, 2, 3, 4, 5, 6, 7]) as Nonce;
+
+  const principal = await Principal.anonymous();
+
+  const httpAgent = new HttpAgent({
+    fetch: mockFetch,
+    host: 'http://localhost',
+    disableNonce: true,
+  });
+  httpAgent.addTransform(makeNonceTransform(() => nonce));
+  const transformMock: HttpAgentRequestTransformFn = jest
+    .fn()
+    .mockImplementation(d => Promise.resolve(d));
+  httpAgent.addTransform(transformMock);
+
+  const methodName = 'greet';
+  const arg = new Uint8Array([]);
+
+  const requestId = await requestIdOf({
+    request_type: SubmitRequestType.Call,
+    nonce,
+    canister_id: Principal.fromText(canisterIdent).toString(),
+    method_name: methodName,
+    arg,
+    sender: principal,
+  });
+
+  const paths = [
+    [new TextEncoder().encode('request_status'), requestId, new TextEncoder().encode('reply')],
+  ];
+
+  const request = await httpAgent.createReadStateRequest({ paths });
+  expect(transformMock).toBeCalledTimes(1);
+  await httpAgent.readState(canisterIdent, { paths }, undefined, request);
+  expect(transformMock).toBeCalledTimes(1);
+});
+
 test('redirect avoid', async () => {
   function checkUrl(base: string, result: string) {
     const httpAgent = new HttpAgent({ host: base });
diff --git a/packages/agent/src/agent/http/index.ts b/packages/agent/src/agent/http/index.ts
index f59e96b53..6e140ce0d 100644
--- a/packages/agent/src/agent/http/index.ts
+++ b/packages/agent/src/agent/http/index.ts
@@ -352,12 +352,11 @@ export class HttpAgent implements Agent {
     return cbor.decode(await response.arrayBuffer());
   }
 
-  public async readState(
-    canisterId: Principal | string,
+  public async createReadStateRequest(
     fields: ReadStateOptions,
     identity?: Identity | Promise<Identity>,
-  ): Promise<ReadStateResponse> {
-    const canister = typeof canisterId === 'string' ? Principal.fromText(canisterId) : canisterId;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  ): Promise<any> {
     const id = await (identity !== undefined ? await identity : await this._identity);
     if (!id) {
       throw new IdentityInvalidError(
@@ -368,7 +367,7 @@ export class HttpAgent implements Agent {
 
     // TODO: remove this any. This can be a Signed or UnSigned request.
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    let transformedRequest: any = await this._transform({
+    const transformedRequest: any = await this._transform({
       request: {
         method: 'POST',
         headers: {
@@ -386,8 +385,19 @@ export class HttpAgent implements Agent {
     });
 
     // Apply transform for identity.
-    transformedRequest = await id?.transformRequest(transformedRequest);
+    return id?.transformRequest(transformedRequest);
+  }
+
+  public async readState(
+    canisterId: Principal | string,
+    fields: ReadStateOptions,
+    identity?: Identity | Promise<Identity>,
+    // eslint-disable-next-line
+    request?: any,
+  ): Promise<ReadStateResponse> {
+    const canister = typeof canisterId === 'string' ? Principal.fromText(canisterId) : canisterId;
 
+    const transformedRequest = request ?? (await this.createReadStateRequest(fields, identity));
     const body = cbor.encode(transformedRequest.body);
 
     const response = await this._fetch(
diff --git a/packages/agent/src/polling/index.ts b/packages/agent/src/polling/index.ts
index f71123a90..788550583 100644
--- a/packages/agent/src/polling/index.ts
+++ b/packages/agent/src/polling/index.ts
@@ -20,15 +20,19 @@ export type PollStrategyFactory = () => PollStrategy;
  * @param canisterId The effective canister ID.
  * @param requestId The Request ID to poll status for.
  * @param strategy A polling strategy.
+ * @param request Request for the readState call.
  */
 export async function pollForResponse(
   agent: Agent,
   canisterId: Principal,
   requestId: RequestId,
   strategy: PollStrategy,
+  // eslint-disable-next-line
+  request?: any,
 ): Promise<ArrayBuffer> {
   const path = [new TextEncoder().encode('request_status'), requestId];
-  const state = await agent.readState(canisterId, { paths: [path] });
+  const currentRequest = request ?? (await agent.createReadStateRequest?.({ paths: [path] }));
+  const state = await agent.readState(canisterId, { paths: [path] }, undefined, currentRequest);
   if (agent.rootKey == null) throw new Error('Agent root key not initialized before polling');
   const cert = await Certificate.create({
     certificate: state.certificate,
@@ -54,7 +58,7 @@ export async function pollForResponse(
     case RequestStatusResponseStatus.Processing:
       // Execute the polling strategy, then retry.
       await strategy(canisterId, requestId, status);
-      return pollForResponse(agent, canisterId, requestId, strategy);
+      return pollForResponse(agent, canisterId, requestId, strategy, currentRequest);
 
     case RequestStatusResponseStatus.Rejected: {
       const rejectCode = new Uint8Array(cert.lookup([...path, 'reject_code'])!)[0];