From 049e61d471733d38cd37aacd4b9b8d33feeceae9 Mon Sep 17 00:00:00 2001
From: Oliver He <heliuchuan@gmail.com>
Date: Mon, 8 Jul 2024 15:50:03 -0400
Subject: [PATCH] Implement a base version of verifySignature for
 KeylessAccounts (#454)

* Implement a base version of verifySignature for KeylessAccounts

* lint

* fmt

* lint

* update changelog
---
 CHANGELOG.md                  |  1 +
 src/account/KeylessAccount.ts | 24 ++++++++++++++++++++----
 src/core/crypto/ephemeral.ts  |  2 +-
 tests/e2e/api/keyless.test.ts | 13 +++++++++++++
 4 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4fdb48693..7f042952f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@ All notable changes to the Aptos TypeScript SDK will be captured in this file. T
 
 # Unreleased
 
+- Adds a base implementation of verify signature for Keyless Accounts
 - [`Fix`] Support migrated coins in coin balance lookup indexer queries
 - Add support for BlockEpilogueTransaction
 - [`Fix`] Fixes a bug with ANS not returning subdomains with an expiration policy of 1 when the subdomain is expired but the parent domain is not.
diff --git a/src/account/KeylessAccount.ts b/src/account/KeylessAccount.ts
index e222816e9..b73950e43 100644
--- a/src/account/KeylessAccount.ts
+++ b/src/account/KeylessAccount.ts
@@ -11,7 +11,6 @@ import {
   KeylessPublicKey,
   KeylessSignature,
   EphemeralCertificate,
-  Signature,
   ZeroKnowledgeSig,
   ZkProof,
 } from "../core/crypto";
@@ -269,9 +268,26 @@ export class KeylessAccount extends Serializable implements Account {
     return this.sign(signMess);
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars, class-methods-use-this
-  verifySignature(args: { message: HexInput; signature: Signature }): boolean {
-    throw new Error("Not implemented");
+  /**
+   * Note - This function is currently incomplete and should only be used to verify ownership of the KeylessAccount
+   *
+   * Verifies a signature given the message.
+   *
+   * TODO: Groth16 proof verification
+   *
+   * @param args.message the message that was signed.
+   * @param args.signature the KeylessSignature to verify
+   * @returns boolean
+   */
+  verifySignature(args: { message: HexInput; signature: KeylessSignature }): boolean {
+    const { message, signature } = args;
+    if (this.isExpired()) {
+      return false;
+    }
+    if (!this.ephemeralKeyPair.getPublicKey().verifySignature({ message, signature: signature.ephemeralSignature })) {
+      return false;
+    }
+    return true;
   }
 
   static fromBytes(bytes: Uint8Array): KeylessAccount {
diff --git a/src/core/crypto/ephemeral.ts b/src/core/crypto/ephemeral.ts
index 07b1ddcf1..a7e9274c6 100644
--- a/src/core/crypto/ephemeral.ts
+++ b/src/core/crypto/ephemeral.ts
@@ -52,7 +52,7 @@ export class EphemeralPublicKey extends PublicKey {
    */
   verifySignature(args: { message: HexInput; signature: EphemeralSignature }): boolean {
     const { message, signature } = args;
-    return this.publicKey.verifySignature({ message, signature });
+    return this.publicKey.verifySignature({ message, signature: signature.signature });
   }
 
   serialize(serializer: Serializer): void {
diff --git a/tests/e2e/api/keyless.test.ts b/tests/e2e/api/keyless.test.ts
index ebb0189fe..188b32ceb 100644
--- a/tests/e2e/api/keyless.test.ts
+++ b/tests/e2e/api/keyless.test.ts
@@ -203,6 +203,19 @@ describe("keyless api", () => {
       KEYLESS_TEST_TIMEOUT,
     );
 
+    test(
+      "keyless account verifies signature for arbitrary message correctly",
+      async () => {
+        // Select a random test token.  Using the same one may encounter rate limits
+        const jwt = TEST_JWT_TOKENS[Math.floor(Math.random() * TEST_JWT_TOKENS.length)];
+        const sender = await aptos.deriveKeylessAccount({ jwt, ephemeralKeyPair });
+        const message = "hello world";
+        const signature = sender.sign(message);
+        expect(sender.verifySignature({ message, signature })).toBe(true);
+      },
+      KEYLESS_TEST_TIMEOUT,
+    );
+
     test(
       "serializes and deserializes",
       async () => {