From b19ddde74c99bd7da1087e719280dcd14666d756 Mon Sep 17 00:00:00 2001
From: Winston Yeo <winstonyeo99@yahoo.com>
Date: Thu, 25 Jan 2024 00:27:16 +0900
Subject: [PATCH] refactor(thirdweb): dry storage and authentication auth token
 retrieval

---
 .../embedded-wallet/core/authentication.ts    | 17 +++++++-----
 .../wallets/embedded-wallet/core/storage.ts   | 27 ++++++++++++++++---
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/packages/thirdweb/src/wallets/embedded-wallet/core/authentication.ts b/packages/thirdweb/src/wallets/embedded-wallet/core/authentication.ts
index c25e8e3bec6..d73a77d4725 100644
--- a/packages/thirdweb/src/wallets/embedded-wallet/core/authentication.ts
+++ b/packages/thirdweb/src/wallets/embedded-wallet/core/authentication.ts
@@ -166,9 +166,18 @@ export const pre2FA = async (
   arg: MultiStepAuthProviderType & { storage: AuthTokenStorageType },
 ): Promise<InitiateAuthResultType> => {
   const { ROUTE_INITIATE_2FA_AUTH } = await import("./routes.js");
+  const { AuthenticationError } = await import("./authentication.error.js");
   const { THIRDWEB_AUTH_TOKEN_KEY } = await import(
     "./authentication.constant.js"
   );
+
+  const token = await arg.storage.fetchToken({
+    key: THIRDWEB_AUTH_TOKEN_KEY,
+  });
+  if (!token) {
+    throw new AuthenticationError("No authenticated user found!");
+  }
+
   switch (arg.provider) {
     case "email": {
       const { email } = arg;
@@ -178,9 +187,7 @@ export const pre2FA = async (
           email,
         }),
         headers: {
-          Authorization: `Bearer ${await arg.storage.fetchToken({
-            key: THIRDWEB_AUTH_TOKEN_KEY,
-          })}`,
+          Authorization: `Bearer ${token}`,
         },
       });
       const result = await resp.json();
@@ -194,9 +201,7 @@ export const pre2FA = async (
           phone,
         }),
         headers: {
-          Authorization: `Bearer ${await arg.storage.fetchToken({
-            key: THIRDWEB_AUTH_TOKEN_KEY,
-          })}`,
+          Authorization: `Bearer ${token}`,
         },
       });
       const result = await resp.json();
diff --git a/packages/thirdweb/src/wallets/embedded-wallet/core/storage.ts b/packages/thirdweb/src/wallets/embedded-wallet/core/storage.ts
index 3d6d4a8a42d..b912ded87c3 100644
--- a/packages/thirdweb/src/wallets/embedded-wallet/core/storage.ts
+++ b/packages/thirdweb/src/wallets/embedded-wallet/core/storage.ts
@@ -10,6 +10,18 @@ import type {
 } from "./storage.type.js";
 import type { WalletDetailType } from "./wallet.type.js";
 
+export const getUserAuthToken = async (authUser?: AuthUserType) => {
+  const { StorageError } = await import("./storage.error.js");
+
+  const token = authUser?.authToken;
+  if (!token) {
+    throw new StorageError(
+      "An authenticated user is required to save the key material with thirdweb",
+    );
+  }
+  return token;
+};
+
 /** Sends an encrypted share / key to thirdweb for storage
  * @throws if developer is not on thirdweb managed storage
  */
@@ -33,12 +45,13 @@ export const saveEncryptedInThirdweb = (arg: {
         "Invalid encryption. The encrypted value must not be the same as the original value",
       );
     }
+    const token = await getUserAuthToken(authUser);
 
     const saveResp = await fetch(ROUTE_STORAGE_ENCRYPTED(), {
       method: "POST",
       headers: {
         "x-secret-key": secretKey ?? "",
-        "x-auth-user-token": authUser?.authToken ?? "",
+        Authorization: `Bearer ${token}`,
       },
       body: JSON.stringify({
         // TODO: figure out what goes here
@@ -78,11 +91,13 @@ export const saveInThirdweb = (): SaveKeyType => {
       );
     }
 
+    const token = await getUserAuthToken(authUser);
+
     const saveResp = await fetch(ROUTE_STORAGE_BASIC(), {
       method: "POST",
       headers: {
         "x-secret-key": secretKey ?? "",
-        "x-auth-user-token": authUser?.authToken ?? "",
+        Authorization: `Bearer ${token}`,
       },
       body: JSON.stringify({
         // walletId: walletDetail.,
@@ -137,11 +152,13 @@ export const loadEncryptedFromThirdweb = (arg: {
       );
     }
 
+    const token = await getUserAuthToken(authUser);
+
     const encryptedKeyMaterialResp = await fetch(ROUTE_STORAGE_ENCRYPTED(), {
       method: "GET",
       headers: {
         "x-secret-key": secretKey ?? "",
-        "x-auth-user-token": authUser?.authToken ?? "",
+        Authorization: `Bearer ${token}`,
       },
     });
     if (!encryptedKeyMaterialResp.ok) {
@@ -178,11 +195,13 @@ export const loadFromThirdweb = (): LoadKeyType => {
       );
     }
 
+    const token = await getUserAuthToken(authUser);
+
     const keyMaterialResp = await fetch(ROUTE_STORAGE_BASIC(), {
       method: "GET",
       headers: {
         "x-secret-key": secretKey ?? "",
-        "x-auth-user-token": authUser?.authToken ?? "",
+        Authorization: `Bearer ${token}`,
       },
     });