Add Authenticator flow for Permissioned Keys

I only fixed and cleaned previous PR #328 by konichuvak

Author: pnowosie <[email protected]>, konichuvak <[email protected]>

v4-client-py-v2/dydx_v4_client/config.py

@@ -1 +1 @@
-GAS_MULTIPLIER = 1.4
+GAS_MULTIPLIER = 1.7

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",

v4-client-py-v2/dydx_v4_client/node/authenticators.py

@@ -0,0 +1,75 @@
+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.encode() for sa in sub_authenticators],
+            separators=(",", ":"),
+        )
+
+        return Authenticator(
+            auth_type,
+            composed_config.encode(),
+        )
+
+    def encode(self):
+        """Prepare object for composition."""
+        dicls = asdict(self)
+        dicls["config"] = list(dicls["config"])
+        return dicls

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,
+    ):
+        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,
+    ):
+        return self.build_transaction(wallet, [as_any(message)], fee, tx_options)

v4-client-py-v2/dydx_v4_client/node/client.py

@@ -25,6 +25,9 @@
     SimulateRequest,
 )
 from v4_proto.cosmos.tx.v1beta1.tx_pb2 import Tx
+from v4_proto.dydxprotocol.accountplus import query_pb2 as accountplus_query
+from v4_proto.dydxprotocol.accountplus import query_pb2_grpc as accountplus_query_grpc
+from v4_proto.dydxprotocol.accountplus import tx_pb2 as accountplus_tx
 from v4_proto.dydxprotocol.bridge import query_pb2 as bridge_query
 from v4_proto.dydxprotocol.bridge import query_pb2_grpc as bridge_query_grpc
 from v4_proto.dydxprotocol.clob import clob_pair_pb2 as clob_pair_type
@@ -69,7 +72,8 @@
 from v4_proto.dydxprotocol.clob.tx_pb2 import OrderBatch
 
 from dydx_v4_client.network import NodeConfig
-from dydx_v4_client.node.builder import Builder
+from dydx_v4_client.node.authenticators import *
+from dydx_v4_client.node.builder import Builder, TxOptions
 from dydx_v4_client.node.fee import Coin, Fee, calculate_fee, Denom
 from dydx_v4_client.node.message import (
     cancel_order,
@@ -79,6 +83,8 @@
     transfer,
     withdraw,
     batch_cancel,
+    add_authenticator,
+    remove_authenticator,
 )
 from dydx_v4_client.wallet import Wallet
 
@@ -429,6 +435,14 @@ async def get_rewards_params(self) -> rewards_query.QueryParamsResponse:
         stub = rewards_query_grpc.QueryStub(self.channel)
         return stub.Params(rewards_query.QueryParamsRequest())
 
+    async def get_authenticators(
+        self, address: str
+    ) -> accountplus_query.GetAuthenticatorsResponse:
+        stub = accountplus_query_grpc.QueryStub(self.channel)
+        return stub.GetAuthenticators(
+            accountplus_query.GetAuthenticatorsRequest(account=address)
+        )
+
 
 class SequenceManager:
     def __init__(self, query_node_client: QueryNodeClient):
@@ -528,7 +542,11 @@ async def send_message(
         return response
 
     async def broadcast_message(
-        self, wallet: Wallet, message: Message, mode=BroadcastMode.BROADCAST_MODE_SYNC
+        self,
+        wallet: Wallet,
+        message: Message,
+        mode=BroadcastMode.BROADCAST_MODE_SYNC,
+        tx_options: Optional[TxOptions] = None,
     ):
         """
         Broadcasts a message.
@@ -537,16 +555,20 @@ async def broadcast_message(
             wallet (Wallet): The wallet to use for signing the transaction.
             message (Message): The message to broadcast.
             mode (BroadcastMode, optional): The broadcast mode. Defaults to BroadcastMode.BROADCAST_MODE_SYNC.
+            tx_options (TxOptions, optional): Options for transaction to support authenticators.
 
         Returns:
             The response from the broadcast.
         """
-        if self.sequence_manager:
+        if not tx_options and self.sequence_manager:
             await self.sequence_manager.before_send(wallet)
 
-        response = await self.broadcast(self.builder.build(wallet, message), mode)
+        response = await self.broadcast(
+            self.builder.build(wallet, message, tx_options=tx_options),
+            mode,
+        )
 
-        if self.sequence_manager:
+        if not tx_options and self.sequence_manager:
             await self.sequence_manager.after_send(wallet)
 
         return response
@@ -710,25 +732,36 @@ async def transfer(
             ),
         )
 
-    async def place_order(self, wallet: Wallet, order: Order):
+    async def place_order(
+        self,
+        wallet: Wallet,
+        order: Order,
+        tx_options: Optional[TxOptions] = None,
+    ):
         """
         Places an order.
 
         Args:
             wallet (Wallet): The wallet to use for signing the transaction.
             order (Order): The order to place.
+            tx_options (TxOptions, optional): Options for transaction to support authenticators.
 
         Returns:
             The response from the transaction broadcast.
         """
-        return await self.broadcast_message(wallet, place_order(order))
+        return await self.broadcast_message(
+            wallet,
+            place_order(order),
+            tx_options=tx_options,
+        )
 
     async def cancel_order(
         self,
         wallet: Wallet,
         order_id: OrderId,
         good_til_block: int = None,
         good_til_block_time: int = None,
+        tx_options: Optional[TxOptions] = None,
     ):
         """
         Cancels an order.
@@ -738,12 +771,15 @@ async def cancel_order(
             order_id (OrderId): The ID of the order to cancel.
             good_til_block (int, optional): The block number until which the order is valid. Defaults to None.
             good_til_block_time (int, optional): The block time until which the order is valid. Defaults to None.
+            tx_options (TxOptions, optional): Options for transaction to support authenticators.
 
         Returns:
             The response from the transaction broadcast.
         """
         return await self.broadcast_message(
-            wallet, cancel_order(order_id, good_til_block, good_til_block_time)
+            wallet,
+            cancel_order(order_id, good_til_block, good_til_block_time),
+            tx_options=tx_options,
         )
 
     async def batch_cancel_orders(
@@ -752,6 +788,7 @@ async def batch_cancel_orders(
         subaccount_id: SubaccountId,
         short_term_cancels: List[OrderBatch],
         good_til_block: int,
+        tx_options: Optional[TxOptions] = None,
     ):
         """
         Batch cancels orders for a subaccount.
@@ -761,6 +798,7 @@ async def batch_cancel_orders(
             subaccount_id (SubaccountId): The subaccount ID for which to cancel orders.
             short_term_cancels (List[OrderBatch]): List of OrderBatch objects containing the orders to cancel.
             good_til_block (int): The last block the short term order cancellations can be executed at.
+            tx_options (TxOptions, optional): Options for transaction to support authenticators.
 
         Returns:
             The response from the transaction broadcast.
@@ -770,4 +808,56 @@ async def batch_cancel_orders(
             short_term_cancels=short_term_cancels,
             good_til_block=good_til_block,
         )
-        return await self.broadcast_message(wallet, batch_cancel_msg)
+        return await self.broadcast_message(
+            wallet,
+            batch_cancel_msg,
+            tx_options=tx_options,
+        )
+
+    async def add_authenticator(
+        self,
+        wallet: Wallet,
+        authenticator: Authenticator,
+    ):
+        """
+        Adds authenticator to a subaccount.
+
+        Args:
+            wallet (Wallet): The wallet to use for signing the transaction or authenticating the request.
+            authenticator Authenticator: The authenticator to be added.
+
+        Returns:
+            The response from the transaction broadcast.
+        """
+        add_authenticator_msg = add_authenticator(
+            wallet.address,
+            authenticator.type,
+            authenticator.config,
+        )
+
+        return await self.send_message(
+            wallet,
+            add_authenticator_msg,
+        )
+
+    async def remove_authenticator(self, wallet: Wallet, authenticator_id: int):
+        """
+        Removes authenticator from a subaccount.
+
+        Args:
+            wallet (Wallet): The wallet to use for signing the transaction or authenticating the request.
+            authenticator_id (int): The authenticator identifier.
+
+        Returns:
+            The response from the transaction broadcast.
+        """
+
+        remove_authenticator_msg = remove_authenticator(
+            wallet.address,
+            authenticator_id,
+        )
+
+        return await self.send_message(
+            wallet,
+            remove_authenticator_msg,
+        )