From 878e47d8162cc9c637d5cee71412f51900ee4c9d Mon Sep 17 00:00:00 2001
From: Reinis Martinsons <reinis@umaproject.org>
Date: Tue, 23 Jul 2024 08:33:50 +0000
Subject: [PATCH] fix: flashbots signature header verification

Signed-off-by: Reinis Martinsons <reinis@umaproject.org>
---
 src/lib/helpers.ts | 42 ++++++++++++++++++++++++++++--------------
 1 file changed, 28 insertions(+), 14 deletions(-)

diff --git a/src/lib/helpers.ts b/src/lib/helpers.ts
index 96860b0..d86a449 100644
--- a/src/lib/helpers.ts
+++ b/src/lib/helpers.ts
@@ -176,14 +176,12 @@ function isOvalConfig(input: unknown): input is OvalConfig {
     typeof input === "object" &&
     input !== null &&
     !Array.isArray(input) &&
-    (
-      ("unlockerKey" in input && typeof input["unlockerKey"] === "string" &&
-        ((!input["unlockerKey"].startsWith("0x") && isHexString("0x" + input["unlockerKey"], 32)) ||
-          isHexString(input["unlockerKey"], 32)) &&
-        !("gckmsKeyId" in input)) ||
-      ("gckmsKeyId" in input && typeof input["gckmsKeyId"] === "string" &&
-        !("unlockerKey" in input))
-    ) &&
+    (("unlockerKey" in input &&
+      typeof input["unlockerKey"] === "string" &&
+      ((!input["unlockerKey"].startsWith("0x") && isHexString("0x" + input["unlockerKey"], 32)) ||
+        isHexString(input["unlockerKey"], 32)) &&
+      !("gckmsKeyId" in input)) ||
+      ("gckmsKeyId" in input && typeof input["gckmsKeyId"] === "string" && !("unlockerKey" in input))) &&
     "refundAddress" in input &&
     typeof input["refundAddress"] === "string" &&
     isAddress(input["refundAddress"]) &&
@@ -204,7 +202,8 @@ function isOvalConfigs(input: unknown): input is OvalConfigs {
     Object.keys(input).length === new Set(Object.keys(input)).size &&
     Object.keys(input).every((key) => isAddress(key)) &&
     Object.values(input).every((value) => isOvalConfig(value)) &&
-    Object.values(input).length === new Set(Object.values(input).map((value) => value.unlockerKey || value.gckmsKeyId)).size
+    Object.values(input).length ===
+      new Set(Object.values(input).map((value) => value.unlockerKey || value.gckmsKeyId)).size
   );
 }
 
@@ -283,8 +282,23 @@ export function verifyBundleSignature(
     return null;
   }
 
-  const bundleSignaturePublicKey = xFlashbotsSignatureHeader.split(":")[0];
-  const bundleSignedMessage = xFlashbotsSignatureHeader.split(":")[1];
+  const xFlashbotsSignatureHeaderParts = xFlashbotsSignatureHeader.split(":");
+  if (xFlashbotsSignatureHeaderParts.length !== 2) {
+    Logger.debug(
+      req.transactionId,
+      `Invalid signature header: ${xFlashbotsSignatureHeader}, expected address and signature separated by a colon`,
+    );
+    return null;
+  }
+
+  const bundleSignaturePublicKey = xFlashbotsSignatureHeaderParts[0];
+  if (!isAddress(bundleSignaturePublicKey)) {
+    Logger.debug(req.transactionId, `Invalid signature header: ${xFlashbotsSignatureHeader}, expected valid address`);
+    return null;
+  }
+  const bundleSignatureAddress = getAddress(bundleSignaturePublicKey);
+
+  const bundleSignedMessage = xFlashbotsSignatureHeaderParts[1];
 
   const serializedBody = JSON.stringify(body);
 
@@ -292,7 +306,7 @@ export function verifyBundleSignature(
 
   const recoveredAddress = ethers.verifyMessage(hash, bundleSignedMessage);
 
-  const verified = recoveredAddress === bundleSignaturePublicKey;
+  const verified = recoveredAddress === bundleSignatureAddress;
 
   return verified ? recoveredAddress : null;
 }
@@ -314,7 +328,7 @@ export class WalletManager {
   private static instance: WalletManager;
   private wallets: Record<string, Wallet> = {};
 
-  private constructor() { }
+  private constructor() {}
 
   public static getInstance(): WalletManager {
     if (!WalletManager.instance) {
@@ -346,4 +360,4 @@ export class WalletManager {
     }
     return this.wallets[checkSummedAddress].connect(provider);
   }
-}
\ No newline at end of file
+}