Refactor and document

boa/integrations/jupyter/constants.py

@@ -0,0 +1,7 @@
+from datetime import timedelta
+
+CALLBACK_TOKEN_TIMEOUT = timedelta(minutes=3)
+ADDRESS_LENGTH = 42
+TRANSACTION_JSON_LENGTH = 2048  # max length of a transaction JSON
+CALLBACK_TOKEN_BYTES = 32
+NUL = b"\0"

boa/integrations/jupyter/handlers.py

@@ -1,14 +1,22 @@
 from http import HTTPStatus
 from multiprocessing.shared_memory import SharedMemory
+from jupyter_server.serverapp import ServerApp
 
 import tornado
 from jupyter_server.base.handlers import APIHandler
 from jupyter_server.utils import url_path_join
 
-BaseAPIHandler = APIHandler
+from boa.integrations.jupyter.constants import CALLBACK_TOKEN_BYTES
 
 
-class CallbackHandler(BaseAPIHandler):
+class CallbackHandler(APIHandler):
+    """
+    This handler receives a callback from jupyter.js when the user interacts with their wallet in the browser.
+    The token is used to identify the SharedMemory object to write the callback data to.
+    Besides, the token needs to fulfill the expected regex that ensures it is a valid format.
+    It is expected that the SharedMemory object has already been created via a BrowserSigner instance.
+    """
+
     @tornado.web.authenticated  # ensure only authorized user can request the Jupyter server
     def post(self, token: str):
         body = self.request.body
@@ -35,15 +43,17 @@ def post(self, token: str):
         return self.finish()
 
 
-def setup_handlers(server_app, name) -> None:
+def setup_handlers(server_app: ServerApp, name: str) -> None:
     """
     Register the handlers in the Jupyter server.
     :param server_app: The Jupyter server application.
+    :param name: The name of the extension.
     """
     web_app = server_app.web_app
     base_url = url_path_join(web_app.settings["base_url"], name)
+    token_regex = rf"{name}_[0-9a-fA-F]{{{CALLBACK_TOKEN_BYTES * 2}}})"
     web_app.add_handlers(
         host_pattern=".*$",
-        host_handlers=[(rf"{base_url}/callback/(\w+)", CallbackHandler)]
+        host_handlers=[(rf"{base_url}/callback/{token_regex}", CallbackHandler)]
     )
     server_app.log.info(f"Handlers registered in {base_url}")

boa/integrations/jupyter/signer.py

@@ -1,6 +1,5 @@
 import json
 from asyncio import sleep, get_running_loop
-from datetime import timedelta
 from multiprocessing.shared_memory import SharedMemory
 from os import urandom
 from os.path import realpath, join, dirname
@@ -10,12 +9,10 @@
 import requests
 from IPython.display import display, Javascript
 
+from boa.integrations.jupyter.constants import CALLBACK_TOKEN_TIMEOUT, ADDRESS_LENGTH, TRANSACTION_JSON_LENGTH, \
+    CALLBACK_TOKEN_BYTES, NUL
+
 nest_asyncio.apply()
-_TIMEOUT = timedelta(minutes=3)
-_ADDRESS_LENGTH = 45  # 42 + quotes + \0
-_TOKEN_LENGTH = 32
-_TX_LENGTH = 2048
-_NUL = b"\0"
 
 
 class BrowserSigner:
@@ -28,7 +25,8 @@ def __init__(self, address=None):
             self.address = address
         else:
             # wait for the address to be set via the API, otherwise boa crashes when trying to create a transaction
-            self.address = _create_and_wait(_load_signer_snippet, size=_ADDRESS_LENGTH)
+            memory_size = ADDRESS_LENGTH + 3  # address + quotes from json encode + \0
+            self.address = _create_and_wait(_load_signer_snippet, size=memory_size)
 
     def send_transaction(self, tx_data: dict) -> dict:
         """
@@ -40,7 +38,7 @@ def send_transaction(self, tx_data: dict) -> dict:
         """
         sign_data = _create_and_wait(
             _sign_transaction_snippet,
-            size=_TX_LENGTH,
+            size=TRANSACTION_JSON_LENGTH,
             tx_data=tx_data,
         )
         return {k: int(v) if isinstance(v, str) and v.isnumeric() else v for k, v in sign_data.items() if v}
@@ -57,7 +55,7 @@ def _create_and_wait(snippet: callable, size: int, **kwargs) -> dict:
     token = _generate_token()
     memory = SharedMemory(name=token, create=True, size=size)
     try:
-        memory.buf[:1] = _NUL
+        memory.buf[:1] = NUL
         javascript = snippet(token, **kwargs)
         display(javascript)
         return _wait_buffer_set(memory.buf)
@@ -66,7 +64,7 @@ def _create_and_wait(snippet: callable, size: int, **kwargs) -> dict:
 
 
 def _generate_token():
-    return f"titanoboa_jupyterlab_{urandom(_TOKEN_LENGTH).hex()}"
+    return f"titanoboa_jupyterlab_{urandom(CALLBACK_TOKEN_BYTES).hex()}"
 
 
 def _wait_buffer_set(buffer: memoryview):
@@ -82,14 +80,14 @@ async def _wait_value(deadline: float) -> Any:
         :return: The result of the Javascript snippet sent to the API.
         """
         inner_loop = get_running_loop()
-        while buffer.tobytes().startswith(_NUL):
+        while buffer.tobytes().startswith(NUL):
             if inner_loop.time() > deadline:
                 raise TimeoutError("Timeout while waiting for user to confirm transaction in the browser.")
             await sleep(0.01)
         return json.loads(buffer.tobytes().decode().split("\0")[0])
 
     loop = get_running_loop()
-    future = _wait_value(deadline=loop.time() + _TIMEOUT.total_seconds())
+    future = _wait_value(deadline=loop.time() + CALLBACK_TOKEN_TIMEOUT.total_seconds())
     task = loop.create_task(future)
     loop.run_until_complete(task)
     return task.result()