From d55e6a34836db3c94c2df2ee77b32aee55107269 Mon Sep 17 00:00:00 2001
From: Kyle Peacock <kylpeacock@gmail.com>
Date: Thu, 16 Nov 2023 13:39:48 -0800
Subject: [PATCH 1/6] feat: retry query signature verification in case cache is
 stale

---
 docs/generated/changelog.html          |  4 ++++
 packages/agent/src/agent/http/index.ts | 17 +++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/docs/generated/changelog.html b/docs/generated/changelog.html
index ae13f1a43..5b59aa642 100644
--- a/docs/generated/changelog.html
+++ b/docs/generated/changelog.html
@@ -11,6 +11,10 @@ <h1>Agent-JS Changelog</h1>
 
     <section>
       <h2>Version x.x.x</h2>
+      <ul>
+        <li>feat: retry query signature verification in case cache is stale</li>
+      </ul> 
+      <h2>Version 0.20.0</h2>
       <ul>
         <li>feat: uses expirable map for subnet keys in agent-js, with a timeout of 1 hour</li>
         <li>chore: cleanup for node 20 development in agent-js</li>
diff --git a/packages/agent/src/agent/http/index.ts b/packages/agent/src/agent/http/index.ts
index db0dc78b3..af7c8f8b1 100644
--- a/packages/agent/src/agent/http/index.ts
+++ b/packages/agent/src/agent/http/index.ts
@@ -532,7 +532,20 @@ export class HttpAgent implements Agent {
     if (!this.#verifyQuerySignatures) {
       return query;
     }
-    return this.#verifyQueryResponse(query, subnetStatus);
+
+    try {
+      return this.#verifyQueryResponse(query, subnetStatus);
+    } catch (error) {
+      // In case the node signatures have changed, refresh the subnet keys and try again
+      console.warn('Query response verification failed. Retrying with fresh subnet keys.');
+      const nodeId = Principal.from(query.signatures?.[0].identity);
+      this.#subnetKeys.delete(nodeId.toText());
+      await this.fetchSubnetKeys(nodeId);
+
+      const updatedSubnetStatus = this.#subnetKeys.get(nodeId.toText());
+      return this.#verifyQueryResponse(query, updatedSubnetStatus);
+    }
+
   }
 
   /**
@@ -557,7 +570,7 @@ export class HttpAgent implements Agent {
     const { status, signatures, requestId } = queryResponse;
 
     const domainSeparator = new TextEncoder().encode('\x0Bic-response');
-    signatures?.forEach(sig => {
+    signatures?.forEach(async sig => {
       const { timestamp, identity } = sig;
       const nodeId = Principal.fromUint8Array(identity).toText();
       let hash: ArrayBuffer;

From 10297dd23c393f3eaf98c6245d5556a587a4c4f6 Mon Sep 17 00:00:00 2001
From: Kyle Peacock <kylpeacock@gmail.com>
Date: Thu, 16 Nov 2023 13:44:35 -0800
Subject: [PATCH 2/6] todo test

---
 packages/agent/src/agent/http/http.test.ts | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/packages/agent/src/agent/http/http.test.ts b/packages/agent/src/agent/http/http.test.ts
index 0af2b0695..b83fad136 100644
--- a/packages/agent/src/agent/http/http.test.ts
+++ b/packages/agent/src/agent/http/http.test.ts
@@ -790,3 +790,5 @@ test('retry requests that fail due to a network failure', async () => {
     expect(mockFetch.mock.calls.length).toBe(4);
   }
 });
+
+test.todo('retry query signature validation after refreshing the subnet node keys');

From 08b3c481191c2d2e5094c7519fdf1ef176256b44 Mon Sep 17 00:00:00 2001
From: Kyle Peacock <kylpeacock@gmail.com>
Date: Thu, 16 Nov 2023 14:36:09 -0800
Subject: [PATCH 3/6] 5 minute expiration, canisterId as key

---
 packages/agent/src/agent/http/index.ts | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/packages/agent/src/agent/http/index.ts b/packages/agent/src/agent/http/index.ts
index af7c8f8b1..b3a7cf9f5 100644
--- a/packages/agent/src/agent/http/index.ts
+++ b/packages/agent/src/agent/http/index.ts
@@ -188,7 +188,7 @@ export class HttpAgent implements Agent {
   #updatePipeline: HttpAgentRequestTransformFn[] = [];
 
   #subnetKeys: ExpirableMap<string, SubnetStatus> = new ExpirableMap({
-    expirationTime: 60 * 60 * 1000, // 1 hour
+    expirationTime: 5 * 60 * 1000, // 5 minutes
   });
   #verifyQuerySignatures = true;
 
@@ -532,20 +532,17 @@ export class HttpAgent implements Agent {
     if (!this.#verifyQuerySignatures) {
       return query;
     }
-
     try {
       return this.#verifyQueryResponse(query, subnetStatus);
     } catch (error) {
       // In case the node signatures have changed, refresh the subnet keys and try again
       console.warn('Query response verification failed. Retrying with fresh subnet keys.');
       const nodeId = Principal.from(query.signatures?.[0].identity);
-      this.#subnetKeys.delete(nodeId.toText());
+      this.#subnetKeys.delete(canisterId.toString());
       await this.fetchSubnetKeys(nodeId);
-
-      const updatedSubnetStatus = this.#subnetKeys.get(nodeId.toText());
-      return this.#verifyQueryResponse(query, updatedSubnetStatus);
     }
-
+    const updatedSubnetStatus = this.#subnetKeys.get(canisterId.toString());
+    return this.#verifyQueryResponse(query, updatedSubnetStatus);
   }
 
   /**

From af224118c20a894bf403de9548ea16eddc569c71 Mon Sep 17 00:00:00 2001
From: Kyle Peacock <kylpeacock@gmail.com>
Date: Thu, 16 Nov 2023 14:52:05 -0800
Subject: [PATCH 4/6] fetchSubnetKeys

---
 packages/agent/src/agent/http/index.ts | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/packages/agent/src/agent/http/index.ts b/packages/agent/src/agent/http/index.ts
index b3a7cf9f5..c07cd59e7 100644
--- a/packages/agent/src/agent/http/index.ts
+++ b/packages/agent/src/agent/http/index.ts
@@ -537,12 +537,23 @@ export class HttpAgent implements Agent {
     } catch (error) {
       // In case the node signatures have changed, refresh the subnet keys and try again
       console.warn('Query response verification failed. Retrying with fresh subnet keys.');
-      const nodeId = Principal.from(query.signatures?.[0].identity);
       this.#subnetKeys.delete(canisterId.toString());
-      await this.fetchSubnetKeys(nodeId);
+      await this.fetchSubnetKeys(canisterId.toString());
     }
     const updatedSubnetStatus = this.#subnetKeys.get(canisterId.toString());
-    return this.#verifyQueryResponse(query, updatedSubnetStatus);
+    if (!updatedSubnetStatus) {
+      throw new CertificateVerificationError(
+        'Invalid signature from replica signed query: no matching node key found.',
+      );
+    }
+    const isValid = this.#verifyQueryResponse(query, updatedSubnetStatus);
+    if (isValid) {
+      return query;
+    } else {
+      throw new CertificateVerificationError(
+        'Invalid signature from replica signed query: no matching node key found.',
+      );
+    }
   }
 
   /**

From 860298ae4e00baab7ff49a82ff9207d5c391b144 Mon Sep 17 00:00:00 2001
From: Kyle Peacock <kylpeacock@gmail.com>
Date: Thu, 16 Nov 2023 16:46:48 -0800
Subject: [PATCH 5/6] foreach async wasn't correctly awaited

---
 e2e/node/basic/mitm.test.ts            | 26 +++++------
 packages/agent/src/agent/http/index.ts | 63 ++++++++++++++------------
 2 files changed, 48 insertions(+), 41 deletions(-)

diff --git a/e2e/node/basic/mitm.test.ts b/e2e/node/basic/mitm.test.ts
index a1171e00f..cfc100b4c 100644
--- a/e2e/node/basic/mitm.test.ts
+++ b/e2e/node/basic/mitm.test.ts
@@ -6,19 +6,19 @@ let mitmTest: TestAPI | typeof test.skip = test;
 if (!process.env['MITM']) {
   mitmTest = test.skip;
 }
-mitmTest(
-  'mitm greet',
-  async () => {
-    const counter = await createActor('tnnnb-2yaaa-aaaab-qaiiq-cai', {
-      agent: await makeAgent({
-        host: 'http://127.0.0.1:8888',
-      }),
-    });
-    await expect(counter.greet('counter')).rejects.toThrow(/Invalid certificate/);
-    expect(await counter.queryGreet('counter')).toEqual('Hullo, counter!');
-  },
-  { timeout: 30000 },
-);
+// mitmTest(
+//   'mitm greet',
+//   async () => {
+//     const counter = await createActor('tnnnb-2yaaa-aaaab-qaiiq-cai', {
+//       agent: await makeAgent({
+//         host: 'http://127.0.0.1:8888',
+//       }),
+//     });
+//     await expect(counter.greet('counter')).rejects.toThrow(/Invalid certificate/);
+//     expect(await counter.queryGreet('counter')).toEqual('Hullo, counter!');
+//   },
+//   { timeout: 30000 },
+// );
 
 mitmTest('mitm with query verification', async () => {
   const counter = await createActor('tnnnb-2yaaa-aaaab-qaiiq-cai', {
diff --git a/packages/agent/src/agent/http/index.ts b/packages/agent/src/agent/http/index.ts
index c07cd59e7..04811ef87 100644
--- a/packages/agent/src/agent/http/index.ts
+++ b/packages/agent/src/agent/http/index.ts
@@ -532,27 +532,30 @@ export class HttpAgent implements Agent {
     if (!this.#verifyQuerySignatures) {
       return query;
     }
-    try {
-      return this.#verifyQueryResponse(query, subnetStatus);
-    } catch (error) {
+    console.log("here");
+    const isValid = this.#verifyQueryResponse(query, subnetStatus);
+    console.log("isValid 1", isValid);
+    if (isValid === true) {
+      return query;
+    } else {
       // In case the node signatures have changed, refresh the subnet keys and try again
       console.warn('Query response verification failed. Retrying with fresh subnet keys.');
       this.#subnetKeys.delete(canisterId.toString());
       await this.fetchSubnetKeys(canisterId.toString());
-    }
-    const updatedSubnetStatus = this.#subnetKeys.get(canisterId.toString());
-    if (!updatedSubnetStatus) {
-      throw new CertificateVerificationError(
-        'Invalid signature from replica signed query: no matching node key found.',
-      );
-    }
-    const isValid = this.#verifyQueryResponse(query, updatedSubnetStatus);
-    if (isValid) {
-      return query;
-    } else {
-      throw new CertificateVerificationError(
-        'Invalid signature from replica signed query: no matching node key found.',
-      );
+
+      const updatedSubnetStatus = this.#subnetKeys.get(canisterId.toString());
+      if (!updatedSubnetStatus) {
+        throw new CertificateVerificationError(
+          'Invalid signature from replica signed query: no matching node key found.',
+        );
+      }
+      const isValid = this.#verifyQueryResponse(query, updatedSubnetStatus);
+      console.log('isValid 2', isValid);
+      if (isValid === true) {
+        return query;
+      } else {
+        throw isValid;
+      }
     }
   }
 
@@ -565,20 +568,24 @@ export class HttpAgent implements Agent {
   #verifyQueryResponse = (
     queryResponse: ApiQueryResponse,
     subnetStatus: SubnetStatus | void,
-  ): ApiQueryResponse => {
+  ): true | CertificateVerificationError => {
+    let result: true | CertificateVerificationError = true;
+
     if (this.#verifyQuerySignatures === false) {
       // This should not be called if the user has disabled verification
-      return queryResponse;
+      result = true;
     }
     if (!subnetStatus) {
-      throw new CertificateVerificationError(
+      return new CertificateVerificationError(
         'Invalid signature from replica signed query: no matching node key found.',
       );
     }
-    const { status, signatures, requestId } = queryResponse;
+    const { status, signatures = [], requestId } = queryResponse;
 
     const domainSeparator = new TextEncoder().encode('\x0Bic-response');
-    signatures?.forEach(async sig => {
+
+
+    for (const sig of signatures) {
       const { timestamp, identity } = sig;
       const nodeId = Principal.fromUint8Array(identity).toText();
       let hash: ArrayBuffer;
@@ -603,7 +610,7 @@ export class HttpAgent implements Agent {
           request_id: requestId,
         });
       } else {
-        throw new Error(`Unknown status: ${status}`);
+        return new CertificateVerificationError(`Unknown status: ${status}`);
       }
 
       const separatorWithHash = concat(domainSeparator, new Uint8Array(hash));
@@ -611,7 +618,7 @@ export class HttpAgent implements Agent {
       // FIX: check for match without verifying N times
       const pubKey = subnetStatus?.nodeKeys.get(nodeId);
       if (!pubKey) {
-        throw new CertificateVerificationError(
+        return new CertificateVerificationError(
           'Invalid signature from replica signed query: no matching node key found.',
         );
       }
@@ -621,13 +628,13 @@ export class HttpAgent implements Agent {
         new Uint8Array(separatorWithHash),
         new Uint8Array(rawKey),
       );
-      if (valid) return queryResponse;
+      if (valid) result = true;
 
-      throw new CertificateVerificationError(
+      return new CertificateVerificationError(
         `Invalid signature from replica ${nodeId} signed query.`,
       );
-    });
-    return queryResponse;
+    }
+    return result;
   };
 
   public async createReadStateRequest(

From ce7e135995c889f8c9e6c718daa4ba3a8ccb30af Mon Sep 17 00:00:00 2001
From: Kyle Peacock <kylpeacock@gmail.com>
Date: Thu, 16 Nov 2023 16:53:54 -0800
Subject: [PATCH 6/6] cleaning up

---
 e2e/node/basic/mitm.test.ts            | 26 +++++++++---------
 packages/agent/src/agent/http/index.ts | 37 +++++++++-----------------
 2 files changed, 25 insertions(+), 38 deletions(-)

diff --git a/e2e/node/basic/mitm.test.ts b/e2e/node/basic/mitm.test.ts
index cfc100b4c..a1171e00f 100644
--- a/e2e/node/basic/mitm.test.ts
+++ b/e2e/node/basic/mitm.test.ts
@@ -6,19 +6,19 @@ let mitmTest: TestAPI | typeof test.skip = test;
 if (!process.env['MITM']) {
   mitmTest = test.skip;
 }
-// mitmTest(
-//   'mitm greet',
-//   async () => {
-//     const counter = await createActor('tnnnb-2yaaa-aaaab-qaiiq-cai', {
-//       agent: await makeAgent({
-//         host: 'http://127.0.0.1:8888',
-//       }),
-//     });
-//     await expect(counter.greet('counter')).rejects.toThrow(/Invalid certificate/);
-//     expect(await counter.queryGreet('counter')).toEqual('Hullo, counter!');
-//   },
-//   { timeout: 30000 },
-// );
+mitmTest(
+  'mitm greet',
+  async () => {
+    const counter = await createActor('tnnnb-2yaaa-aaaab-qaiiq-cai', {
+      agent: await makeAgent({
+        host: 'http://127.0.0.1:8888',
+      }),
+    });
+    await expect(counter.greet('counter')).rejects.toThrow(/Invalid certificate/);
+    expect(await counter.queryGreet('counter')).toEqual('Hullo, counter!');
+  },
+  { timeout: 30000 },
+);
 
 mitmTest('mitm with query verification', async () => {
   const counter = await createActor('tnnnb-2yaaa-aaaab-qaiiq-cai', {
diff --git a/packages/agent/src/agent/http/index.ts b/packages/agent/src/agent/http/index.ts
index 04811ef87..8c942a403 100644
--- a/packages/agent/src/agent/http/index.ts
+++ b/packages/agent/src/agent/http/index.ts
@@ -532,12 +532,9 @@ export class HttpAgent implements Agent {
     if (!this.#verifyQuerySignatures) {
       return query;
     }
-    console.log("here");
-    const isValid = this.#verifyQueryResponse(query, subnetStatus);
-    console.log("isValid 1", isValid);
-    if (isValid === true) {
-      return query;
-    } else {
+    try {
+      return this.#verifyQueryResponse(query, subnetStatus);
+    } catch (_) {
       // In case the node signatures have changed, refresh the subnet keys and try again
       console.warn('Query response verification failed. Retrying with fresh subnet keys.');
       this.#subnetKeys.delete(canisterId.toString());
@@ -549,13 +546,7 @@ export class HttpAgent implements Agent {
           'Invalid signature from replica signed query: no matching node key found.',
         );
       }
-      const isValid = this.#verifyQueryResponse(query, updatedSubnetStatus);
-      console.log('isValid 2', isValid);
-      if (isValid === true) {
-        return query;
-      } else {
-        throw isValid;
-      }
+      return this.#verifyQueryResponse(query, updatedSubnetStatus);
     }
   }
 
@@ -568,23 +559,19 @@ export class HttpAgent implements Agent {
   #verifyQueryResponse = (
     queryResponse: ApiQueryResponse,
     subnetStatus: SubnetStatus | void,
-  ): true | CertificateVerificationError => {
-    let result: true | CertificateVerificationError = true;
-
+  ): ApiQueryResponse => {
     if (this.#verifyQuerySignatures === false) {
       // This should not be called if the user has disabled verification
-      result = true;
+      return queryResponse;
     }
     if (!subnetStatus) {
-      return new CertificateVerificationError(
+      throw new CertificateVerificationError(
         'Invalid signature from replica signed query: no matching node key found.',
       );
     }
     const { status, signatures = [], requestId } = queryResponse;
 
     const domainSeparator = new TextEncoder().encode('\x0Bic-response');
-
-
     for (const sig of signatures) {
       const { timestamp, identity } = sig;
       const nodeId = Principal.fromUint8Array(identity).toText();
@@ -610,7 +597,7 @@ export class HttpAgent implements Agent {
           request_id: requestId,
         });
       } else {
-        return new CertificateVerificationError(`Unknown status: ${status}`);
+        throw new Error(`Unknown status: ${status}`);
       }
 
       const separatorWithHash = concat(domainSeparator, new Uint8Array(hash));
@@ -618,7 +605,7 @@ export class HttpAgent implements Agent {
       // FIX: check for match without verifying N times
       const pubKey = subnetStatus?.nodeKeys.get(nodeId);
       if (!pubKey) {
-        return new CertificateVerificationError(
+        throw new CertificateVerificationError(
           'Invalid signature from replica signed query: no matching node key found.',
         );
       }
@@ -628,13 +615,13 @@ export class HttpAgent implements Agent {
         new Uint8Array(separatorWithHash),
         new Uint8Array(rawKey),
       );
-      if (valid) result = true;
+      if (valid) return queryResponse;
 
-      return new CertificateVerificationError(
+      throw new CertificateVerificationError(
         `Invalid signature from replica ${nodeId} signed query.`,
       );
     }
-    return result;
+    return queryResponse;
   };
 
   public async createReadStateRequest(