feat: Add Authenticator flow for Permissioned Keys (#330)

This is cleaned & fixed version of #328
Big thanks to the original contributor: @konichuvak

---------

Co-authored-by: konichuvak <[email protected]>

Author: pnowosie

Diff for: v4-client-py-v2/dydx_v4_client/config.py

@@ -1 +1 @@
-GAS_MULTIPLIER = 1.4
+GAS_MULTIPLIER = 1.8

Diff for: v4-client-py-v2/dydx_v4_client/network.py

@@ -46,7 +46,6 @@ def make_config(
 make_secure = partial(make_config, secure_channel)
 make_insecure = partial(make_config, insecure_channel)
 
-
 mainnet_node = partial(
     NodeConfig,
     "dydx-mainnet-1",
@@ -73,7 +72,6 @@ def make_config(
 TESTNET_FAUCET = "https://faucet.v4testnet.dydx.exchange"
 TESTNET_NOBLE = "https://rpc.testnet.noble.strange.love"
 
-
 local_node = partial(
     NodeConfig,
     "localdydxprotocol",

Diff for: v4-client-py-v2/dydx_v4_client/node/authenticators.py

@@ -0,0 +1,116 @@
+from enum import Enum
+import json
+from dataclasses import asdict, dataclass
+from typing import List
+
+
+class AuthenticatorType(str, Enum):
+    AllOf = "AllOf"
+    AnyOf = "AnyOf"
+    SignatureVerification = "SignatureVerification"
+    MessageFilter = "MessageFilter"
+    SubaccountFilter = "SubaccountFilter"
+    ClobPairIdFilter = "ClobPairIdFilter"
+
+
+@dataclass
+class Authenticator:
+    type: AuthenticatorType
+    config: bytes
+
+    # helpers to create Authenticator instances
+    @classmethod
+    def signature_verification(cls, pub_key: bytes) -> "Authenticator":
+        """Enables authentication via a specific key."""
+        return Authenticator(
+            AuthenticatorType.SignatureVerification,
+            pub_key,
+        )
+
+    @classmethod
+    def message_filter(cls, msg_type: str) -> "Authenticator":
+        """Restricts authentication to certain message types."""
+        return Authenticator(
+            AuthenticatorType.MessageFilter,
+            msg_type.encode(),
+        )
+
+    @classmethod
+    def subaccount_filter(cls, subaccounts: List[int]) -> "Authenticator":
+        """Restricts authentication to a specific subaccount."""
+        config = ",".join(map(str, subaccounts))
+        return Authenticator(
+            AuthenticatorType.SubaccountFilter,
+            config.encode(),
+        )
+
+    @classmethod
+    def clob_pair_id_filter(cls, clob_pair_ids: List[int]) -> "Authenticator":
+        """Restricts authentication to a specific clob pair id."""
+        config = ",".join(map(str, clob_pair_ids))
+        return Authenticator(
+            AuthenticatorType.ClobPairIdFilter,
+            config.encode(),
+        )
+
+    @classmethod
+    def compose(
+        cls, auth_type: AuthenticatorType, sub_authenticators: list["Authenticator"]
+    ) -> "Authenticator":
+        """Combines multiple sub-authenticators into a single one."""
+        composed_config = json.dumps(
+            [sa.todict() for sa in sub_authenticators],
+            separators=(",", ":"),
+        )
+        return Authenticator(
+            auth_type,
+            composed_config.encode(),
+        )
+
+    def todict(self) -> dict:
+        """Prepare object for composition."""
+        dicls = asdict(self)
+        dicls["config"] = list(dicls["config"])
+        return dicls
+
+
+def validate_authenticator(authenticator: Authenticator) -> bool:
+    """Validate the authenticator."""
+    if authenticator.config.startswith(b"["):
+        decoded_config = json.loads(authenticator.config.decode())
+    else:
+        decoded_config = authenticator.config
+
+    return check_authenticator(dict(type=authenticator.type, config=decoded_config))
+
+
+def check_authenticator(auth: dict) -> bool:
+    """
+    Check if the authenticator is safe to use.
+    Parameters:
+    - auth is a decoded authenticator object.
+    """
+    if not is_authenticator_alike(auth):
+        return False
+
+    if auth["type"] == AuthenticatorType.SignatureVerification:
+        # SignatureVerification authenticator is considered safe
+        return True
+
+    if not isinstance(auth["config"], list):
+        return False
+
+    if auth["type"] == AuthenticatorType.AnyOf:
+        # ANY_OF is safe only if ALL sub-authenticators return true
+        return all(check_authenticator(sub_auth) for sub_auth in auth["config"])
+
+    if auth["type"] == AuthenticatorType.AllOf:
+        # ALL_OF is safe if at least one sub-authenticator returns true
+        return any(check_authenticator(sub_auth) for sub_auth in auth["config"])
+
+    # If it's a base-case authenticator but not SignatureVerification, it's unsafe
+    return False
+
+
+def is_authenticator_alike(auth: dict) -> bool:
+    return isinstance(auth, dict) and auth.get("type") and auth.get("config")

Diff for: v4-client-py-v2/dydx_v4_client/node/builder.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import List
+from typing import List, Optional
 
 import google
 from google.protobuf.message import Message
@@ -17,6 +17,7 @@
 
 from dydx_v4_client.node.fee import calculate_fee, Denom
 from dydx_v4_client.wallet import Wallet
+from v4_proto.dydxprotocol.accountplus.tx_pb2 import TxExtension
 
 
 def as_any(message: Message):
@@ -50,6 +51,13 @@ def get_signature(key_pair, body, auth_info, account_number, chain_id):
 )
 
 
+@dataclass
+class TxOptions:
+    authenticators: List[int]
+    sequence: int
+    account_number: int
+
+
 @dataclass
 class Builder:
     chain_id: str
@@ -69,17 +77,48 @@ def fee(self, gas_limit: int, *amount: List[Coin]) -> Fee:
             gas_limit=gas_limit,
         )
 
-    def build_transaction(self, wallet: Wallet, messages: List[Message], fee: Fee):
-        body = TxBody(messages=messages, memo=self.memo)
+    def build_transaction(
+        self,
+        wallet: Wallet,
+        messages: List[Message],
+        fee: Fee,
+        tx_options: Optional[TxOptions] = None,
+    ) -> Tx:
+        non_critical_extension_options = []
+        if tx_options is not None:
+            tx_extension = TxExtension(
+                selected_authenticators=tx_options.authenticators,
+            )
+            non_critical_extension_options.append(as_any(tx_extension))
+        body = TxBody(
+            messages=messages,
+            memo=self.memo,
+            non_critical_extension_options=non_critical_extension_options,
+        )
         auth_info = AuthInfo(
-            signer_infos=[get_signer_info(wallet.public_key, wallet.sequence)],
+            signer_infos=[
+                get_signer_info(
+                    wallet.public_key,
+                    tx_options.sequence if tx_options else wallet.sequence,
+                )
+            ],
             fee=fee,
         )
         signature = get_signature(
-            wallet.key, body, auth_info, wallet.account_number, self.chain_id
+            wallet.key,
+            body,
+            auth_info,
+            tx_options.account_number if tx_options else wallet.account_number,
+            self.chain_id,
         )
 
         return Tx(body=body, auth_info=auth_info, signatures=[signature])
 
-    def build(self, wallet: Wallet, message: Message, fee: Fee = DEFAULT_FEE):
-        return self.build_transaction(wallet, [as_any(message)], fee)
+    def build(
+        self,
+        wallet: Wallet,
+        message: Message,
+        fee: Fee = DEFAULT_FEE,
+        tx_options: Optional[dict] = None,
+    ) -> Tx:
+        return self.build_transaction(wallet, [as_any(message)], fee, tx_options)