fix review comment

rhinestonewtf · Sep 10, 2024 · 16799eb · 16799eb
1 parent b33465f
commit 16799eb
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 33 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -107,9 +107,10 @@ export {
   getSpendingLimitsPolicy,
   getSudoPolicy,
   getUniversalActionPolicy,
-  encodeUseSmartSessionSignature,
-  encodeEnableSmartSessionSignature,
+  encodeSmartSessionSignature,
+  encodeUseOrEnableSmartSessionSignature,
   isSessionEnabled,
+  hashChainSessions,
   SmartSessionMode,
   Session,
   SessionEIP712,

diff --git a/src/module/index.ts b/src/module/index.ts
@@ -130,9 +130,10 @@ export {
   getSpendingLimitsPolicy,
   getSudoPolicy,
   getUniversalActionPolicy,
-  encodeUseSmartSessionSignature,
-  encodeEnableSmartSessionSignature,
+  encodeSmartSessionSignature,
+  encodeUseOrEnableSmartSessionSignature,
   isSessionEnabled,
+  hashChainSessions,
   SmartSessionMode,
   Session,
   SessionEIP712,

diff --git a/src/module/smart-sessions/usage.ts b/src/module/smart-sessions/usage.ts
@@ -104,51 +104,104 @@ export const getSessionDigest = async ({
   })) as Hex
 }
 
-export const encodeUseSmartSessionSignature = ({
+export const encodeSmartSessionSignature = ({
+  mode,
   permissionId,
   signature,
+  enableSessionData,
 }: {
+  mode: SmartSessionModeType
   permissionId: Hex
   signature: Hex
+  enableSessionData?: EnableSessionData
 }) => {
-  return encodePacked(
-    ['bytes1', 'bytes32', 'bytes'],
-    [
-      SmartSessionMode.USE,
-      permissionId,
-      LibZip.flzCompress(
-        encodeAbiParameters(
-          [
-            {
-              type: 'bytes',
-            },
-          ],
-          [signature],
-        ),
-      ) as Hex,
-    ],
-  )
+  switch (mode) {
+    case SmartSessionMode.USE:
+      return encodePacked(
+        ['bytes1', 'bytes32', 'bytes'],
+        [
+          mode,
+          permissionId,
+          LibZip.flzCompress(
+            encodeAbiParameters(
+              [
+                {
+                  type: 'bytes',
+                },
+              ],
+              [signature],
+            ),
+          ) as Hex,
+        ],
+      )
+    case SmartSessionMode.ENABLE:
+    case SmartSessionMode.UNSAFE_ENABLE:
+      if (!enableSessionData) {
+        throw new Error('enableSession is required for ENABLE mode')
+      }
+
+      return encodePacked(
+        ['bytes1', 'bytes32', 'bytes'],
+        [
+          mode,
+          permissionId,
+          LibZip.flzCompress(
+            encodeEnableSessionSignature({ enableSessionData, signature }),
+          ) as Hex,
+        ],
+      )
+    default:
+      throw new Error(`Unknown mode ${mode}`)
+  }
 }
 
-export const encodeEnableSmartSessionSignature = ({
+export const encodeUseOrEnableSmartSessionSignature = async ({
+  account,
+  client,
   permissionId,
   signature,
   enableSessionData,
 }: {
+  account: Account
+  client: PublicClient
   permissionId: Hex
   signature: Hex
   enableSessionData: EnableSessionData
 }) => {
-  return encodePacked(
-    ['bytes1', 'bytes32', 'bytes'],
-    [
-      SmartSessionMode.ENABLE,
-      permissionId,
-      LibZip.flzCompress(
-        encodeEnableSessionSignature({ enableSessionData, signature }),
-      ) as Hex,
-    ],
-  )
+  const sessionEnabled = await isSessionEnabled({
+    account,
+    client,
+    permissionId,
+  })
+
+  return sessionEnabled
+    ? encodePacked(
+        ['bytes1', 'bytes32', 'bytes'],
+        [
+          SmartSessionMode.USE,
+          permissionId,
+          LibZip.flzCompress(
+            encodeAbiParameters(
+              [
+                {
+                  type: 'bytes',
+                },
+              ],
+              [signature],
+            ),
+          ) as Hex,
+        ],
+      )
+    : encodePacked(
+        ['bytes1', 'bytes32', 'bytes'],
+        [
+          SmartSessionMode.ENABLE,
+          permissionId,
+          LibZip.flzCompress(
+            encodeEnableSessionSignature({ enableSessionData, signature }),
+          ) as Hex,
+        ],
+      )
 }
 
 export const hashChainSessions = (chainSessions: ChainSession[]): Hex => {