From a441472b7dd89ef9ecf34d2478c546c1ea9fe5f1 Mon Sep 17 00:00:00 2001
From: Kyle Peacock <kylpeacock@gmail.com>
Date: Tue, 20 Aug 2024 08:58:24 -0700
Subject: [PATCH] feat: new option for setting rootKey during agent creation

---
 docs/CHANGELOG.md                          |  1 +
 e2e/node/basic/mainnet.test.ts             | 18 ++++++++++++++++++
 packages/agent/src/agent/http/http.test.ts |  2 ++
 packages/agent/src/agent/http/index.ts     |  8 +++++++-
 4 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
index d4bb990e..42cff552 100644
--- a/docs/CHANGELOG.md
+++ b/docs/CHANGELOG.md
@@ -7,6 +7,7 @@
 - feat: management canister interface updates for schnorr signatures
 - feat: ensure that identity-secp256k1 seed phrase must produce a 64 byte seed
 - docs: documentation and metadata for use-auth-client
+- feat: adds optional `rootKey` to `HttpAgentOptions` to allow for a custom root key to be used for verifying signatures from other networks
 
 ### Changed
 - feat: replaces hdkey and bip32 implementations with `@scure/bip39` and `@scure/bip32` due to vulnerability and lack of maintenance for `elliptic`
diff --git a/e2e/node/basic/mainnet.test.ts b/e2e/node/basic/mainnet.test.ts
index 49787b2d..b3a530ad 100644
--- a/e2e/node/basic/mainnet.test.ts
+++ b/e2e/node/basic/mainnet.test.ts
@@ -141,3 +141,21 @@ describe('call forwarding', () => {
     expect(reply).toBeTruthy();
   }, 15_000);
 });
+
+
+test('it should allow you to set an incorrect root key', async () => {
+  const agent = HttpAgent.createSync({
+    rootKey: new Uint8Array(31),
+  });
+  const idlFactory = ({ IDL }) =>
+    IDL.Service({
+      whoami: IDL.Func([], [IDL.Principal], ['query']),
+    });
+
+  const actor = Actor.createActor(idlFactory, {
+    agent,
+    canisterId: Principal.fromText('rrkah-fqaaa-aaaaa-aaaaq-cai'),
+  });
+
+  expect(actor.whoami).rejects.toThrowError(`Invalid certificate:`);
+});
diff --git a/packages/agent/src/agent/http/http.test.ts b/packages/agent/src/agent/http/http.test.ts
index c040dae7..dbe2f730 100644
--- a/packages/agent/src/agent/http/http.test.ts
+++ b/packages/agent/src/agent/http/http.test.ts
@@ -793,6 +793,8 @@ test('retry requests that fail due to a network failure', async () => {
     fetch: mockFetch,
   });
 
+  agent.rootKey = new Uint8Array(32);
+
   try {
     await agent.call(Principal.managementCanister(), {
       methodName: 'test',
diff --git a/packages/agent/src/agent/http/index.ts b/packages/agent/src/agent/http/index.ts
index b32a9152..387962e0 100644
--- a/packages/agent/src/agent/http/index.ts
+++ b/packages/agent/src/agent/http/index.ts
@@ -138,6 +138,11 @@ export interface HttpAgentOptions {
    * Whether to log to the console. Defaults to false.
    */
   logToConsole?: boolean;
+
+  /**
+   * Alternate root key to use for verifying certificates. If not provided, the default IC root key will be used.
+   */
+  rootKey?: ArrayBuffer;
 }
 
 function getDefaultFetch(): typeof fetch {
@@ -233,7 +238,7 @@ other computations so that this class can stay as simple as possible while
 allowing extensions.
  */
 export class HttpAgent implements Agent {
-  public rootKey = fromHex(IC_ROOT_KEY);
+  public rootKey: ArrayBuffer;
   #identity: Promise<Identity> | null;
   readonly #fetch: typeof fetch;
   readonly #fetchOptions?: Record<string, unknown>;
@@ -275,6 +280,7 @@ export class HttpAgent implements Agent {
     this.#fetch = options.fetch || getDefaultFetch() || fetch.bind(global);
     this.#fetchOptions = options.fetchOptions;
     this.#callOptions = options.callOptions;
+    this.rootKey = options.rootKey ? options.rootKey : fromHex(IC_ROOT_KEY);
 
     const host = determineHost(options.host);
     this.host = new URL(host);