Merge pull request #19 from Shopify/chore/listener-simplification

[chore] ♻️ Refactor onConnect listener creation to a util function

packages/connect-wallet/src/hooks/useConnectWallet/useConnectWalletCallbacks.ts

@@ -1,58 +1,24 @@
-import {isAnyOf} from '@reduxjs/toolkit';
-import {useContext, useEffect} from 'react';
+import {useEffect} from 'react';
 
-import {ConnectWalletContext} from '../../providers/ConnectWalletProvider';
-import {
-  addWallet,
-  removeWallet,
-  setActiveWallet,
-  validatePendingWallet,
-} from '../../slices/walletSlice';
+import {buildOnConnectMiddleware} from '../../middleware/onConnectMiddleware';
+import {removeWallet} from '../../slices/walletSlice';
 import {addListener} from '../../store/listenerMiddleware';
-import {Wallet} from '../../types/wallet';
 import {useAppDispatch} from '../useAppState';
 
 import {useConnectWalletProps} from './types';
 
 export const useConnectWalletCallbacks = (props?: useConnectWalletProps) => {
   const {onConnect, onDisconnect} = props || {};
-  const {requireSignature} = useContext(ConnectWalletContext);
   const dispatch = useAppDispatch();
 
   // Add the onConnect callback listeners.
   useEffect(() => {
-    const unsubscribeToOnConnectListener = dispatch(
-      addListener({
-        matcher: isAnyOf(
-          addWallet,
-          validatePendingWallet.fulfilled,
-          setActiveWallet,
-        ),
-        effect: (action, state) => {
-          let walletToDispatch: Wallet | undefined;
-
-          if (action.type === 'wallet/validatePendingWallet/fulfilled') {
-            const signatureResponse = action.meta.arg;
-            const {address, signature} = signatureResponse;
-
-            const {connectedWallets} = state.getState().wallet;
-            walletToDispatch = connectedWallets.find(
-              (wallet) =>
-                wallet.address === address && wallet.signature === signature,
-            );
-          } else {
-            walletToDispatch = action.payload;
-          }
-
-          if (walletToDispatch) {
-            onConnect?.(walletToDispatch);
-          }
-        },
-      }),
-    );
+    const listener = buildOnConnectMiddleware(({wallet}) => {
+      onConnect?.(wallet);
+    });
 
-    return unsubscribeToOnConnectListener;
-  }, [dispatch, onConnect, requireSignature]);
+    return dispatch(listener);
+  }, [dispatch, onConnect]);
 
   // Add the onDisconnect callback listeners.
   useEffect(() => {

packages/connect-wallet/src/hooks/useMiddleware.ts

@@ -0,0 +1,135 @@
+import {useEffect} from 'react';
+import {useAccount} from 'wagmi';
+
+import {buildOnConnectMiddleware} from '../middleware/onConnectMiddleware';
+import {
+  addWallet,
+  attributeOrder,
+  setActiveWallet,
+  setPendingWallet,
+} from '../slices/walletSlice';
+import {addListener} from '../store/listenerMiddleware';
+import {OrderAttributionMode} from '../types/orderAttribution';
+import {SignatureResponse, Wallet} from '../types/wallet';
+
+import {useAppDispatch, useAppSelector} from './useAppState';
+
+interface UseMiddlewareProps {
+  orderAttributionMode: OrderAttributionMode;
+  requestSignature: (wallet: Wallet) => Promise<SignatureResponse | undefined>;
+  requireSignature?: boolean;
+}
+
+export const useMiddleware = ({
+  orderAttributionMode,
+  requestSignature,
+  requireSignature,
+}: UseMiddlewareProps) => {
+  const dispatch = useAppDispatch();
+  const {connectedWallets, pendingConnector} = useAppSelector(
+    (state) => state.wallet,
+  );
+
+  useAccount({
+    onConnect: ({address, connector, isReconnected}) => {
+      if (!address) {
+        return;
+      }
+
+      const reconnectedWallet: Wallet | undefined = connectedWallets.find(
+        (wallet) => wallet.address === address,
+      );
+
+      if (requireSignature) {
+        /**
+         * Check if the wallet has already signed. If so, we can set
+         * the active wallet and not require a new signature.
+         */
+        if (isReconnected && reconnectedWallet?.signature) {
+          return dispatch(setActiveWallet(reconnectedWallet));
+        }
+
+        /**
+         * Check to ensure we have connector data before proceeding. We
+         * need connector for injected connectors such as Coinbase Wallet
+         * and MetaMask. Otherwise, utilize pendingConnector value.
+         */
+        if (!pendingConnector && !connector) {
+          return;
+        }
+
+        const wallet: Wallet = {
+          address,
+          connectorId: pendingConnector?.id || connector?.id,
+          connectorName: pendingConnector?.name || connector?.name,
+        };
+
+        return dispatch(setPendingWallet(wallet));
+      }
+
+      /**
+       * If we don't require a signature and we have a reconnected
+       * wallet then we can set the active wallet.
+       */
+      if (reconnectedWallet) {
+        return dispatch(setActiveWallet(reconnectedWallet));
+      }
+
+      // Exit if we don't have pendingConnector information.
+      if (!pendingConnector) {
+        return;
+      }
+
+      // This means that the user just connected their wallet.
+      dispatch(
+        addWallet({
+          address,
+          connectorId: pendingConnector.id,
+          connectorName: pendingConnector.name,
+        }),
+      );
+    },
+  });
+
+  /**
+   * Pending wallet listener
+   *
+   * Utilized to request signature verification.
+   */
+  useEffect(() => {
+    if (requireSignature) {
+      return dispatch(
+        addListener({
+          actionCreator: setPendingWallet,
+          effect: (action, _) => {
+            const wallet = action.payload;
+
+            if (!wallet) {
+              return;
+            }
+
+            requestSignature(wallet);
+          },
+        }),
+      );
+    }
+  }, [dispatch, requestSignature, requireSignature]);
+
+  /**
+   * onConnect listener (internal)
+   *
+   * This listener will run order attribution functionality.
+   */
+  useEffect(() => {
+    const listener = buildOnConnectMiddleware(({state, wallet}) => {
+      state.dispatch(
+        attributeOrder({
+          orderAttributionMode,
+          wallet,
+        }),
+      );
+    });
+
+    return dispatch(listener);
+  }, [dispatch, orderAttributionMode]);
+};

packages/connect-wallet/src/middleware/onConnectMiddleware.test.ts

@@ -0,0 +1,73 @@
+import {
+  addWallet,
+  setActiveWallet,
+  setPendingWallet,
+  validatePendingWallet,
+} from '../slices/walletSlice';
+import {store} from '../test/configureStore';
+import {
+  INVALID_SIGNATURE_RESPONSE,
+  VALID_SIGNATURE_RESPONSE,
+} from '../test/fixtures/signature';
+import {DEFAULT_WALLET} from '../test/fixtures/wallet';
+import {ConnectWalletError} from '../utils/error';
+
+import {buildOnConnectMiddleware} from './onConnectMiddleware';
+
+describe('onConnectMiddleware', () => {
+  const effectFn = jest.fn();
+
+  beforeEach(() => {
+    const listener = buildOnConnectMiddleware(({wallet}) => effectFn(wallet));
+
+    store.dispatch(listener);
+  });
+
+  afterEach(() => jest.clearAllMocks());
+
+  describe('addWallet', () => {
+    it('runs the effect when addWallet is dispatched', () => {
+      store.dispatch(addWallet(DEFAULT_WALLET));
+
+      expect(effectFn).toHaveBeenCalledWith(DEFAULT_WALLET);
+    });
+  });
+
+  describe('setActiveWallet', () => {
+    it('runs the effect when setActiveWallet payload is a wallet', () => {
+      store.dispatch(setActiveWallet(DEFAULT_WALLET));
+
+      expect(effectFn).toHaveBeenCalledWith(DEFAULT_WALLET);
+    });
+
+    it('does not run the effect when setActiveWallet payload is undefined', () => {
+      store.dispatch(setActiveWallet(undefined));
+
+      expect(effectFn).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('validatePendingWallet', () => {
+    it('runs the effect when validatePendingWallet payload is a valid signed message', async () => {
+      store.dispatch(setPendingWallet(DEFAULT_WALLET));
+
+      await store.dispatch(validatePendingWallet(VALID_SIGNATURE_RESPONSE));
+
+      expect(effectFn).toHaveBeenCalledWith(
+        expect.objectContaining({...DEFAULT_WALLET}),
+      );
+    });
+
+    it('does not run the effect when validatePendingWallet payload is not a valid signed message', async () => {
+      store.dispatch(setPendingWallet(DEFAULT_WALLET));
+
+      await expect(
+        store.dispatch(validatePendingWallet(INVALID_SIGNATURE_RESPONSE)),
+      ).rejects.toThrow(
+        new ConnectWalletError(
+          'Address that signed message does not match the connected address',
+        ),
+      );
+    });
+  });
+});

packages/connect-wallet/src/middleware/onConnectMiddleware.ts

@@ -0,0 +1,79 @@
+import {isAnyOf, ListenerEffect} from '@reduxjs/toolkit';
+
+import {
+  addWallet,
+  setActiveWallet,
+  validatePendingWallet,
+} from '../slices/walletSlice';
+import {AppDispatch, RootState} from '../store/configureStore';
+import {addListener} from '../store/listenerMiddleware';
+import {GuardedType} from '../types/generics';
+import {Wallet} from '../types/wallet';
+
+const matcher = isAnyOf(
+  addWallet,
+  setActiveWallet,
+  validatePendingWallet.fulfilled,
+);
+
+type ListenerEffectType = ListenerEffect<
+  GuardedType<typeof matcher>,
+  RootState,
+  AppDispatch
+>;
+
+type EffectAction = Parameters<ListenerEffectType>['0'];
+type EffectState = Parameters<ListenerEffectType>['1'];
+interface EffectCallbackProps {
+  state: EffectState;
+  wallet: Wallet;
+}
+
+const isWalletPayload = (
+  payload: EffectAction['payload'],
+): payload is Wallet => {
+  return payload !== undefined && 'address' in payload;
+};
+
+/**
+ * Creates a listener that runs an effect when:
+ * - a wallet is added via `addWallet`
+ * - a wallet is set as the activeWallet via `setActiveWallet`
+ * - a wallet is validated after signing via `validatePendingWallet`
+ *
+ * **NOTE**: The listener created is not automatically dispatched, this must be
+ * done manually.
+ *
+ * @param effect Callback function which receives a Wallet type.
+ */
+export const buildOnConnectMiddleware = (
+  effect?: ({state, wallet}: EffectCallbackProps) => void,
+) => {
+  return addListener({
+    matcher,
+    effect: (action, state) => {
+      let walletToDispatch: Wallet | undefined;
+
+      if (action.type === 'wallet/validatePendingWallet/fulfilled') {
+        // Since the `validatePendingWallet` action ran we know that we
+        // need to find the wallet from the list of connected wallets.
+        const signatureResponse = action.meta.arg;
+        const {address, signature} = signatureResponse;
+
+        const {connectedWallets} = state.getState().wallet;
+        walletToDispatch = connectedWallets.find(
+          (wallet) =>
+            wallet.address === address && wallet.signature === signature,
+        );
+      }
+
+      if (isWalletPayload(action.payload)) {
+        walletToDispatch = action.payload;
+      }
+
+      if (walletToDispatch) {
+        return effect?.({state, wallet: walletToDispatch});
+      }
+    },
+  });
+};