Merge pull request #1383 from trilitech/fix-beacon-signing

Fix beacon signing

apps/desktop/src/components/DynamicModal.tsx

@@ -1,10 +1,13 @@
-import { Modal, ModalOverlay, type ModalProps, useDisclosure } from "@chakra-ui/react";
-import { type ReactElement, createContext, useState } from "react";
+import { Modal, ModalOverlay, type ThemingProps, useDisclosure } from "@chakra-ui/react";
+import { type ReactElement, createContext, useCallback, useState } from "react";
 import { RemoveScroll } from "react-remove-scroll";
 
 // this should be used in components as useContext(DynamicModalContext);
 export const DynamicModalContext = createContext<{
-  openWith: (content: ReactElement, size?: ModalProps["size"]) => Promise<void>;
+  openWith: (
+    content: ReactElement,
+    props?: ThemingProps & { onClose?: () => void | Promise<void> }
+  ) => Promise<void>;
   onClose: () => void;
   isOpen: boolean;
 }>({
@@ -23,16 +26,35 @@ export const DynamicModalContext = createContext<{
 // This hook should be used only once in the app. You place the modal wrapper and then
 // use the `openWith` provided by the `DynamicModalContext` in components.
 export const useDynamicModal = () => {
-  const { isOpen, onClose, onOpen } = useDisclosure();
+  const { isOpen, onClose: closeModal, onOpen } = useDisclosure();
   const [modalContent, setModalContent] = useState<ReactElement | null>(null);
-  const [size, setSize] = useState<ModalProps["size"]>("md");
+  const [modalProps, setModalProps] = useState<
+    ThemingProps & { onClose: () => void | Promise<void> }
+  >({
+    size: "md",
+    onClose: closeModal,
+  });
 
-  const openWith = async (content: ReactElement, size: ModalProps["size"] = "md") => {
-    setSize(size);
-    setModalContent(content);
-    onOpen();
-    return Promise.resolve();
-  };
+  const onClose = useCallback(() => {
+    closeModal();
+  }, [closeModal]);
+
+  const openWith = useCallback(
+    async (
+      content: ReactElement,
+      props: ThemingProps & { onClose?: () => void | Promise<void> } = {}
+    ) => {
+      const onClose = () => {
+        closeModal();
+        return props.onClose?.();
+      };
+      setModalProps({ size: "md", ...props, onClose });
+      setModalContent(content);
+      onOpen();
+      return Promise.resolve();
+    },
+    [onOpen, closeModal]
+  );
 
   return {
     isOpen,
@@ -47,8 +69,7 @@ export const useDynamicModal = () => {
         isOpen={isOpen}
         // this is used in e2e tests to decrease flakiness due to animations
         motionPreset={(localStorage.getItem("chakra-modal-motion-preset") as any) || undefined}
-        onClose={onClose}
-        size={size}
+        {...modalProps}
       >
         <ModalOverlay />
         <RemoveScroll enabled={isOpen}>{modalContent}</RemoveScroll>

apps/desktop/src/components/SendFlow/Beacon/BatchSignPage.tsx

@@ -36,7 +36,7 @@ export const BatchSignPage: React.FC<BeaconSignPageProps> = ({ operation, messag
           <Header message={message} mode="batch" operation={operation} />
 
           <ModalBody>
-            <Accordion allowToggle={true}>
+            <Accordion allowToggle>
               <AccordionItem background={colors.gray[800]} border="none" borderRadius="8px">
                 <AccordionButton>
                   <Heading flex="1" textAlign="left" paddingY="6px" size="sm">

apps/desktop/src/components/SendFlow/Beacon/FinalizeUnstakeSignPage.tsx

@@ -1,4 +1,5 @@
 import { Flex, FormLabel, ModalBody, ModalContent, ModalFooter } from "@chakra-ui/react";
+import { FormProvider } from "react-hook-form";
 
 import { type BeaconSignPageProps } from "./BeaconSignPageProps";
 import { Header } from "./Header";
@@ -11,35 +12,37 @@ import { SignPageFee } from "../SignPageFee";
 import { headerText } from "../SignPageHeader";
 
 export const FinalizeUnstakeSignPage: React.FC<BeaconSignPageProps> = ({ operation, message }) => {
-  const { isSigning, onSign, network, fee } = useSignWithBeacon(operation, message);
+  const { isSigning, onSign, network, fee, form } = useSignWithBeacon(operation, message);
   const totalFinalizableAmount = useAccountTotalFinalizableUnstakeAmount(
     operation.signer.address.pkh
   );
   return (
-    <ModalContent>
-      <form>
-        <Header message={message} mode="single" operation={operation} />
-        <ModalBody>
-          <Flex alignItems="center" justifyContent="end" marginTop="12px">
-            <SignPageFee fee={fee} />
-          </Flex>
+    <FormProvider {...form}>
+      <ModalContent>
+        <form>
+          <Header message={message} mode="single" operation={operation} />
+          <ModalBody>
+            <Flex alignItems="center" justifyContent="end" marginTop="12px">
+              <SignPageFee fee={fee} />
+            </Flex>
 
-          <FormLabel marginTop="24px">From</FormLabel>
-          <AddressTile address={operation.sender.address} />
+            <FormLabel marginTop="24px">From</FormLabel>
+            <AddressTile address={operation.sender.address} />
 
-          <FormLabel marginTop="24px">Finalize Unstake</FormLabel>
-          <TezTile mutezAmount={totalFinalizableAmount} />
-        </ModalBody>
-        <ModalFooter>
-          <SignButton
-            isLoading={isSigning}
-            network={network}
-            onSubmit={onSign}
-            signer={operation.signer}
-            text={headerText(operation.type, "single")}
-          />
-        </ModalFooter>
-      </form>
-    </ModalContent>
+            <FormLabel marginTop="24px">Finalize Unstake</FormLabel>
+            <TezTile mutezAmount={totalFinalizableAmount} />
+          </ModalBody>
+          <ModalFooter>
+            <SignButton
+              isLoading={isSigning}
+              network={network}
+              onSubmit={onSign}
+              signer={operation.signer}
+              text={headerText(operation.type, "single")}
+            />
+          </ModalFooter>
+        </form>
+      </ModalContent>
+    </FormProvider>
   );
 };

apps/desktop/src/components/SendFlow/Beacon/StakeSignPage.tsx

@@ -1,4 +1,5 @@
 import { Flex, FormLabel, ModalBody, ModalContent, ModalFooter } from "@chakra-ui/react";
+import { FormProvider } from "react-hook-form";
 
 import { type BeaconSignPageProps } from "./BeaconSignPageProps";
 import { Header } from "./Header";
@@ -13,33 +14,35 @@ import { headerText } from "../SignPageHeader";
 export const StakeSignPage: React.FC<BeaconSignPageProps> = ({ operation, message }) => {
   const { amount: mutezAmount } = operation.operations[0] as Stake;
 
-  const { isSigning, onSign, network, fee } = useSignWithBeacon(operation, message);
+  const { isSigning, onSign, network, fee, form } = useSignWithBeacon(operation, message);
 
   return (
-    <ModalContent>
-      <form>
-        <Header message={message} mode="single" operation={operation} />
-        <ModalBody>
-          <Flex alignItems="center" justifyContent="end" marginTop="12px">
-            <SignPageFee fee={fee} />
-          </Flex>
+    <FormProvider {...form}>
+      <ModalContent>
+        <form>
+          <Header message={message} mode="single" operation={operation} />
+          <ModalBody>
+            <Flex alignItems="center" justifyContent="end" marginTop="12px">
+              <SignPageFee fee={fee} />
+            </Flex>
 
-          <FormLabel marginTop="24px">From</FormLabel>
-          <AddressTile address={operation.sender.address} />
+            <FormLabel marginTop="24px">From</FormLabel>
+            <AddressTile address={operation.sender.address} />
 
-          <FormLabel marginTop="24px">Stake</FormLabel>
-          <TezTile mutezAmount={mutezAmount} />
-        </ModalBody>
-        <ModalFooter>
-          <SignButton
-            isLoading={isSigning}
-            network={network}
-            onSubmit={onSign}
-            signer={operation.signer}
-            text={headerText(operation.type, "single")}
-          />
-        </ModalFooter>
-      </form>
-    </ModalContent>
+            <FormLabel marginTop="24px">Stake</FormLabel>
+            <TezTile mutezAmount={mutezAmount} />
+          </ModalBody>
+          <ModalFooter>
+            <SignButton
+              isLoading={isSigning}
+              network={network}
+              onSubmit={onSign}
+              signer={operation.signer}
+              text={headerText(operation.type, "single")}
+            />
+          </ModalFooter>
+        </form>
+      </ModalContent>
+    </FormProvider>
   );
 };

apps/desktop/src/components/SendFlow/Beacon/UnstakeSignPage.tsx

@@ -1,4 +1,5 @@
 import { Flex, FormLabel, ModalBody, ModalContent, ModalFooter } from "@chakra-ui/react";
+import { FormProvider } from "react-hook-form";
 
 import { type BeaconSignPageProps } from "./BeaconSignPageProps";
 import { Header } from "./Header";
@@ -13,33 +14,35 @@ import { headerText } from "../SignPageHeader";
 export const UnstakeSignPage: React.FC<BeaconSignPageProps> = ({ operation, message }) => {
   const { amount: mutezAmount } = operation.operations[0] as Unstake;
 
-  const { isSigning, onSign, network, fee } = useSignWithBeacon(operation, message);
+  const { isSigning, onSign, network, fee, form } = useSignWithBeacon(operation, message);
 
   return (
-    <ModalContent>
-      <form>
-        <Header message={message} mode="single" operation={operation} />
-        <ModalBody>
-          <Flex alignItems="center" justifyContent="end" marginTop="12px">
-            <SignPageFee fee={fee} />
-          </Flex>
+    <FormProvider {...form}>
+      <ModalContent>
+        <form>
+          <Header message={message} mode="single" operation={operation} />
+          <ModalBody>
+            <Flex alignItems="center" justifyContent="end" marginTop="12px">
+              <SignPageFee fee={fee} />
+            </Flex>
 
-          <FormLabel marginTop="24px">From</FormLabel>
-          <AddressTile address={operation.sender.address} />
+            <FormLabel marginTop="24px">From</FormLabel>
+            <AddressTile address={operation.sender.address} />
 
-          <FormLabel marginTop="24px">Unstake</FormLabel>
-          <TezTile mutezAmount={mutezAmount} />
-        </ModalBody>
-        <ModalFooter>
-          <SignButton
-            isLoading={isSigning}
-            network={network}
-            onSubmit={onSign}
-            signer={operation.signer}
-            text={headerText(operation.type, "single")}
-          />
-        </ModalFooter>
-      </form>
-    </ModalContent>
+            <FormLabel marginTop="24px">Unstake</FormLabel>
+            <TezTile mutezAmount={mutezAmount} />
+          </ModalBody>
+          <ModalFooter>
+            <SignButton
+              isLoading={isSigning}
+              network={network}
+              onSubmit={onSign}
+              signer={operation.signer}
+              text={headerText(operation.type, "single")}
+            />
+          </ModalFooter>
+        </form>
+      </ModalContent>
+    </FormProvider>
   );
 };

apps/desktop/src/mocks/modal.tsx

@@ -0,0 +1,37 @@
+import { type PropsWithChildren, createContext, useContext } from "react";
+
+const ModalContext = createContext<{ onClose: () => void }>({ onClose: () => {} });
+
+export const MockModal = ({
+  children,
+  isOpen,
+  onClose,
+}: PropsWithChildren<{ isOpen: boolean; onClose: () => void }>) => (
+  <ModalContext.Provider value={{ onClose }}>
+    <div data-testid="mock-modal">{isOpen ? children : null}</div>
+  </ModalContext.Provider>
+);
+
+export const MockModalHeader = ({ children }: PropsWithChildren<object>) => (
+  <div id="modal-header">{children}</div>
+);
+
+export const MockModalContent = ({ children }: PropsWithChildren<object>) => (
+  <section aria-labelledby="modal-header" aria-modal role="dialog">
+    {children}
+  </section>
+);
+
+export const MockModalInnerComponent = ({ children }: PropsWithChildren<object>) => (
+  <div>{children}</div>
+);
+
+export const MockModalCloseButton = ({ children }: PropsWithChildren<object>) => {
+  const { onClose } = useContext(ModalContext);
+
+  return (
+    <button aria-label="Close" onClick={onClose} type="button">
+      {children}
+    </button>
+  );
+};

apps/desktop/src/setupTests.ts

@@ -10,8 +10,14 @@ import { TextDecoder, TextEncoder } from "util";
 import { setupJestCanvasMock } from "jest-canvas-mock";
 import failOnConsole from "jest-fail-on-console";
 import MockDate from "mockdate";
-import React from "react";
 
+import {
+  MockModal,
+  MockModalCloseButton,
+  MockModalContent,
+  MockModalHeader,
+  MockModalInnerComponent,
+} from "./mocks/modal";
 import { mockUseToast } from "./mocks/toast";
 import { accountsSlice } from "./utils/redux/slices/accountsSlice/accountsSlice";
 import { announcementSlice } from "./utils/redux/slices/announcementSlice";
@@ -71,24 +77,6 @@ beforeEach(() => {
   setupJestCanvasMock();
 });
 
-const MockModal = ({ children, isOpen }: any) =>
-  React.createElement("div", { "data-testid": "mock-modal" }, isOpen ? children : null);
-
-const MockModalHeader = ({ children }: any) =>
-  React.createElement("header", { id: "modal-header" }, children);
-
-const MockModalContent = ({ children }: any) =>
-  React.createElement(
-    "section",
-    { role: "dialog", "aria-labelledby": "modal-header", "aria-modal": true },
-    children
-  );
-
-const MockModalInnerComponent = ({ children }: any) => React.createElement("div", {}, children);
-
-const MockModalCloseButton = ({ children }: any) =>
-  React.createElement("button", { "aria-label": "Close" }, children);
-
 jest.mock("@chakra-ui/react", () => ({
   ...jest.requireActual("@chakra-ui/react"),
   // Mock toast since it has an erratic behavior in RTL

apps/desktop/src/utils/beacon/PermissionRequestModal.tsx

@@ -1,5 +1,4 @@
 import {
-  BeaconErrorType,
   BeaconMessageType,
   type BeaconResponseInputMessage,
   type PermissionRequestOutput,
@@ -29,7 +28,6 @@ import type React from "react";
 import { useContext } from "react";
 import { FormProvider, useForm } from "react-hook-form";
 
-import { useRemovePeerBySenderId } from "./beacon";
 import { WalletClient } from "./WalletClient";
 import { JsValueWrap } from "../../components/AccountDrawer/JsValueWrap";
 import { OwnedImplicitAccountsAutocomplete } from "../../components/AddressAutocomplete";
@@ -53,16 +51,6 @@ export const PermissionRequestModal: React.FC<{
     getValues,
     formState: { errors, isValid },
   } = form;
-  const removePeer = useRemovePeerBySenderId();
-
-  const onModalClose = () => {
-    void removePeer(request.senderId);
-    void WalletClient.respond({
-      id: request.id,
-      type: BeaconMessageType.Error,
-      errorType: BeaconErrorType.NOT_GRANTED_ERROR,
-    });
-  };
 
   const grant = () =>
     handleAsyncAction(async () => {
@@ -101,7 +89,7 @@ export const PermissionRequestModal: React.FC<{
           </Text>
         </Flex>
       </ModalHeader>
-      <ModalCloseButton onClick={onModalClose} />
+      <ModalCloseButton />
       <ModalBody data-testid="beacon-request-body">
         <Flex
           alignItems="center"