From 6d14dad02a4746c69c9e0e5a5c9f6a794027ab44 Mon Sep 17 00:00:00 2001
From: sancar <sancar@upstash.com>
Date: Fri, 25 Oct 2024 14:42:13 +0300
Subject: [PATCH] DX-1388 Verify Next Signing Key

---
 src/receiver.test.ts | 76 ++++++++++++++++++++++++++++++++++++++++++++
 src/receiver.ts      | 21 +++++++-----
 2 files changed, 89 insertions(+), 8 deletions(-)
 create mode 100644 src/receiver.test.ts

diff --git a/src/receiver.test.ts b/src/receiver.test.ts
new file mode 100644
index 00000000..cc5f8a41
--- /dev/null
+++ b/src/receiver.test.ts
@@ -0,0 +1,76 @@
+/* eslint-disable @typescript-eslint/no-magic-numbers */
+/**
+ * Tests the Receiver functionality.
+ */
+
+import { nanoid } from "ai";
+import { describe, test } from "bun:test";
+import { SignJWT } from "jose";
+import { createHash } from "node:crypto";
+import { Receiver } from ".";
+
+async function createUpstashSingature({
+  url,
+  body,
+  key,
+}: {
+  url: string;
+  body: string;
+  key: string;
+}) {
+  const payload = {
+    iss: "Upstash",
+    sub: url,
+    exp: Math.floor(Date.now() / 1000) + 300, // expires in 5 minutes
+    nbf: Math.floor(Date.now() / 1000),
+    iat: Math.floor(Date.now() / 1000),
+    jti: `jwt_${Math.random().toString(36).slice(2, 15)}`,
+    body: createHash("sha256").update(body).digest("base64url"),
+  };
+
+  const jwt = await new SignJWT(payload)
+    .setProtectedHeader({ alg: "HS256", typ: "JWT" })
+    .sign(Buffer.from(key, "utf8"));
+
+  return jwt;
+}
+
+const currentSigningKey = nanoid();
+const nextSigningKey = nanoid();
+
+const randomBody = btoa(nanoid());
+const url = "example.com";
+
+describe("receiver", () => {
+  test("verify signed with currentSigningKey", async () => {
+    const receiver = new Receiver({ currentSigningKey, nextSigningKey });
+
+    const upstashSignature = await createUpstashSingature({
+      url: url,
+      body: randomBody,
+      key: currentSigningKey,
+    });
+
+    await receiver.verify({
+      signature: upstashSignature,
+      body: randomBody,
+      url: url,
+    });
+  });
+
+  test("verify signed with nextSigninKey", async () => {
+    const receiver = new Receiver({ currentSigningKey, nextSigningKey });
+
+    const upstashSignature = await createUpstashSingature({
+      url: url,
+      body: randomBody,
+      key: nextSigningKey,
+    });
+
+    await receiver.verify({
+      signature: upstashSignature,
+      body: randomBody,
+      url: url,
+    });
+  });
+});
diff --git a/src/receiver.ts b/src/receiver.ts
index 213d3a37..9f54a80a 100644
--- a/src/receiver.ts
+++ b/src/receiver.ts
@@ -69,17 +69,20 @@ export class Receiver {
    * If that fails, the signature is invalid and a `SignatureError` is thrown.
    */
   public async verify(request: VerifyRequest): Promise<boolean> {
-    const isValid = await this.verifyWithKey(this.currentSigningKey, request);
-    if (isValid) {
-      return true;
+    let payload: jose.JWTPayload;
+    try {
+      payload = await this.verifyWithKey(this.currentSigningKey, request);
+    } catch {
+      payload = await this.verifyWithKey(this.nextSigningKey, request);
     }
-    return this.verifyWithKey(this.nextSigningKey, request);
+    this.verifyBodyAndUrl(payload, request);
+    return true;
   }
 
   /**
    * Verify signature with a specific signing key
    */
-  private async verifyWithKey(key: string, request: VerifyRequest): Promise<boolean> {
+  private async verifyWithKey(key: string, request: VerifyRequest): Promise<jose.JWTPayload> {
     const jwt = await jose
       .jwtVerify(request.signature, new TextEncoder().encode(key), {
         issuer: "Upstash",
@@ -89,7 +92,11 @@ export class Receiver {
         throw new SignatureError((error as Error).message);
       });
 
-    const p = jwt.payload as {
+    return jwt.payload;
+  }
+
+  private verifyBodyAndUrl(payload: jose.JWTPayload, request: VerifyRequest) {
+    const p = payload as {
       iss: string;
       sub: string;
       exp: number;
@@ -110,7 +117,5 @@ export class Receiver {
     if (p.body.replace(padding, "") !== bodyHash.replace(padding, "")) {
       throw new SignatureError(`body hash does not match, want: ${p.body}, got: ${bodyHash}`);
     }
-
-    return true;
   }
 }