From 5691c37f8d0ec7369b9d64c5b91f395a59b001cf Mon Sep 17 00:00:00 2001
From: Atticus <atticus@macbook36gb.local>
Date: Fri, 15 Mar 2024 14:09:49 -0600
Subject: [PATCH] fix(jwk): add validation for jwks

---
 src/common/common.ts  | 16 ++++++++++++++++
 src/common/wallets.ts | 16 +++++++++++++++-
 test/wallets.ts       | 11 +++++++++++
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/src/common/common.ts b/src/common/common.ts
index 0f0f6854..ee4edcfb 100644
--- a/src/common/common.ts
+++ b/src/common/common.ts
@@ -85,6 +85,13 @@ export default class Arweave {
     attributes: Partial<CreateTransactionInterface>,
     jwk?: JWKInterface | "use_wallet"
   ): Promise<Transaction> {
+
+    const isValidJWK = jwk && jwk !== "use_wallet" ? await this.validateJWK(jwk) : true;
+
+    if (!isValidJWK) {
+      throw new Error(`The JWK is not valid.`);
+    }
+
     const transaction: Partial<CreateTransactionInterface> = {};
 
     Object.assign(transaction, attributes);
@@ -216,4 +223,13 @@ export default class Arweave {
       .post("/arql", query)
       .then((response) => response.data || []);
   }
+
+  public async validateJWK(jwk: JWKInterface): Promise<boolean> {
+    const transaction = await this.createTransaction({
+      data: "test",
+    });
+    await this.transactions.sign(transaction, jwk);
+
+    return await this.transactions.verify(transaction);
+  }
 }
diff --git a/src/common/wallets.ts b/src/common/wallets.ts
index 595923b0..31433474 100644
--- a/src/common/wallets.ts
+++ b/src/common/wallets.ts
@@ -3,6 +3,10 @@ import CryptoInterface from "./lib/crypto/crypto-interface";
 import { JWKInterface } from "./lib/wallet";
 import * as ArweaveUtils from "./lib/utils";
 import "arconnect";
+import Transaction from "lib/transaction";
+import Transactions from "transactions";
+import Chunks from "chunks";
+import Arweave from "./common";
 
 export default class Wallets {
   private api: Api;
@@ -66,13 +70,23 @@ export default class Wallets {
       // @ts-ignore
       return arweaveWallet.getActiveAddress();
     } else {
+      // verify the JWK produces valid transactions
+      const isValid = await new Arweave(this.api.getConfig()).validateJWK(jwk);
+      if (!isValid) {
+        throw new Error("The JWK is not valid.");
+      }
       return this.ownerToAddress(jwk.n);
     }
   }
 
   public async ownerToAddress(owner: string): Promise<string> {
+    const buffer = await ArweaveUtils.b64UrlToBuffer(owner);
+    // RSA 4096 keys have an n length of 512 bytes. Validate JWK should handle this, so this is just a sanity check.
+    if (buffer.byteLength !== 512) {
+      throw new Error("Invalid JWK");
+    }
     return ArweaveUtils.bufferTob64Url(
-      await this.crypto.hash(ArweaveUtils.b64UrlToBuffer(owner))
+      await this.crypto.hash(buffer)
     );
   }
 }
diff --git a/test/wallets.ts b/test/wallets.ts
index 89470b87..0bc05ff2 100644
--- a/test/wallets.ts
+++ b/test/wallets.ts
@@ -99,4 +99,15 @@ describe("Wallets and keys", function () {
       .to.be.a("string")
       .and.equal("fOVzBRTBnyt4VrUUYadBH8yras_-jhgpmNgg-5b3vEw");
   });
+
+  it("should throw on bad JWK", async function () {
+    const jwk = await arweave.wallets.generate();
+    jwk.n = jwk.n.slice(0, -1);
+    try {
+      await arweave.wallets.jwkToAddress(jwk);
+    } catch (e: any) {
+      expect(e).to.be.an("error");
+      expect(e.message).to.equal("The JWK is not valid.");
+    }
+  });
 });