feat(nano): Add address balance in global state

Author: msbrogli

hathor/dag_builder/builder.py

@@ -45,6 +45,8 @@
 
 NC_DEPOSIT_KEY = 'nc_deposit'
 NC_WITHDRAWAL_KEY = 'nc_withdrawal'
+NC_TRANSFER_INPUT_KEY = 'nc_transfer_input'
+NC_TRANSFER_OUTPUT_KEY = 'nc_transfer_output'
 
 
 class DAGBuilder:
@@ -240,6 +242,22 @@ def _add_nc_attribute(self, name: str, key: str, value: str) -> None:
             actions.append((token, amount))
             node.attrs[key] = actions
 
+        elif key == NC_TRANSFER_INPUT_KEY:
+            transfer_inputs = node.get_attr_list(key, default=[])
+            token, amount, (wallet,) = parse_amount_token(value)
+            if amount < 0:
+                raise SyntaxError(f'unexpected negative amount in `{key}`')
+            transfer_inputs.append((wallet, token, amount))
+            node.attrs[key] = transfer_inputs
+
+        elif key == NC_TRANSFER_OUTPUT_KEY:
+            transfer_outputs = node.get_attr_list(key, default=[])
+            token, amount, (wallet,) = parse_amount_token(value)
+            if amount < 0:
+                raise SyntaxError(f'unexpected negative amount in `{key}`')
+            transfer_outputs.append((wallet, token, amount))
+            node.attrs[key] = transfer_outputs
+
         else:
             node.attrs[key] = value
 

hathor/dag_builder/vertex_exporter.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import ast
+import hashlib
 import re
 from collections import defaultdict
 from types import ModuleType
@@ -23,7 +24,14 @@
 from hathor.conf.settings import HathorSettings
 from hathor.crypto.util import decode_address, get_address_from_public_key_bytes
 from hathor.daa import DifficultyAdjustmentAlgorithm
-from hathor.dag_builder.builder import NC_DEPOSIT_KEY, NC_WITHDRAWAL_KEY, DAGBuilder, DAGNode
+from hathor.dag_builder.builder import (
+    NC_DEPOSIT_KEY,
+    NC_TRANSFER_INPUT_KEY,
+    NC_TRANSFER_OUTPUT_KEY,
+    NC_WITHDRAWAL_KEY,
+    DAGBuilder,
+    DAGNode,
+)
 from hathor.dag_builder.types import DAGNodeType, VertexResolverType, WalletFactoryType
 from hathor.dag_builder.utils import get_literal, is_literal
 from hathor.nanocontracts import Blueprint, OnChainBlueprint
@@ -34,6 +42,7 @@
 from hathor.transaction import BaseTransaction, Block, Transaction
 from hathor.transaction.base_transaction import TxInput, TxOutput
 from hathor.transaction.headers.nano_header import ADDRESS_LEN_BYTES
+from hathor.transaction.headers.transfer_header import TxTransferInput, TxTransferOutput
 from hathor.transaction.scripts.p2pkh import P2PKH
 from hathor.transaction.token_creation_tx import TokenCreationTransaction
 from hathor.wallet import BaseWallet, HDWallet, KeyPair
@@ -229,6 +238,7 @@ def create_vertex_token(self, node: DAGNode) -> TokenCreationTransaction:
         vertex.token_symbol = node.name
         vertex.timestamp = self.get_min_timestamp(node)
         self.add_nano_header_if_needed(node, vertex)
+        self.add_transfer_header_if_needed(node, vertex)
         self.sign_all_inputs(vertex, node=node)
         if 'weight' in node.attrs:
             vertex.weight = float(node.attrs['weight'])
@@ -253,6 +263,7 @@ def create_vertex_block(self, node: DAGNode) -> Block:
 
         blk = Block(parents=parents, outputs=outputs)
         self.add_nano_header_if_needed(node, blk)
+        self.add_transfer_header_if_needed(node, blk)
         blk.timestamp = self.get_min_timestamp(node) + self._settings.AVG_TIME_BETWEEN_BLOCKS
         blk.get_height = lambda: height  # type: ignore[method-assign]
         blk.update_hash()  # the next call fails is blk.hash is None
@@ -302,6 +313,65 @@ def _get_next_nc_seqnum(self, nc_pubkey: bytes) -> int:
         self._next_nc_seqnum[address] = cur + 1
         return cur
 
+    def _get_token_index(self, token_name: str, node: DAGNode, vertex: BaseTransaction) -> None:
+        token_index = 0
+        if token_name != 'HTR':
+            assert isinstance(vertex, Transaction)
+            token_creation_tx = self._vertices[token_name]
+            if token_creation_tx.hash not in vertex.tokens:
+                vertex.tokens.append(token_creation_tx.hash)
+            token_index = 1 + vertex.tokens.index(token_creation_tx.hash)
+        return token_index
+
+    def add_transfer_header_if_needed(self, node: DAGNode, vertex: BaseTransaction) -> None:
+        inputs = node.get_attr_list(NC_TRANSFER_INPUT_KEY)
+        outputs = node.get_attr_list(NC_TRANSFER_OUTPUT_KEY)
+
+        if not inputs and not outputs:
+            return
+
+        transfer_inputs: list[TxTransferInput] = []
+        for wallet_name, token_name, amount in inputs:
+            wallet = self.get_wallet(wallet_name)
+            assert isinstance(wallet, HDWallet)
+            privkey = wallet.get_key_at_index(0)
+            pubkey_bytes = privkey.sec()
+            address = get_address_from_public_key_bytes(pubkey_bytes)
+            token_index = self._get_token_index(token_name, node, vertex)
+
+            sighash_data = vertex.get_sighash_all_data()
+            sighash_data_hash = hashlib.sha256(sighash_data).digest()
+            signature = privkey.sign(sighash_data_hash)
+            script = P2PKH.create_input_data(public_key_bytes=pubkey_bytes, signature=signature)
+
+            transfer_inputs.append(TxTransferInput(
+                address=address,
+                amount=amount,
+                token_index=token_index,
+                script=script,
+            ))
+
+        transfer_outputs: list[TxTransferOutput] = []
+        for wallet_name, token_name, amount in outputs:
+            wallet = self.get_wallet(wallet_name)
+            assert isinstance(wallet, HDWallet)
+            privkey = wallet.get_key_at_index(0)
+            pubkey_bytes = privkey.sec()
+            address = get_address_from_public_key_bytes(pubkey_bytes)
+            token_index = self._get_token_index(token_name, node, vertex)
+            transfer_outputs.append(TxTransferOutput(
+                address=address,
+                amount=amount,
+                token_index=token_index,
+            ))
+
+        from hathor.transaction.headers import TransferHeader
+        transfer_header = TransferHeader(
+            inputs=transfer_inputs,
+            outputs=transfer_outputs,
+        )
+        vertex.headers.append(transfer_header)
+
     def add_nano_header_if_needed(self, node: DAGNode, vertex: BaseTransaction) -> None:
         if 'nc_id' not in node.attrs:
             return
@@ -405,6 +475,7 @@ def create_vertex_on_chain_blueprint(self, node: DAGNode) -> OnChainBlueprint:
         assert len(block_parents) == 0
         ocb = OnChainBlueprint(parents=txs_parents, inputs=inputs, outputs=outputs, tokens=tokens)
         self.add_nano_header_if_needed(node, ocb)
+        self.add_transfer_header_if_needed(node, ocb)
         code_attr = node.get_attr_str('ocb_code')
 
         if is_literal(code_attr):
@@ -453,6 +524,7 @@ def create_vertex_transaction(self, node: DAGNode, *, cls: type[Transaction] = T
         tx = cls(parents=txs_parents, inputs=inputs, outputs=outputs, tokens=tokens)
         tx.timestamp = self.get_min_timestamp(node)
         self.add_nano_header_if_needed(node, tx)
+        self.add_transfer_header_if_needed(node, tx)
         self.sign_all_inputs(tx, node=node)
         if 'weight' in node.attrs:
             tx.weight = float(node.attrs['weight'])

hathor/nanocontracts/blueprint_env.py

@@ -24,7 +24,7 @@
     from hathor.nanocontracts.nc_exec_logs import NCLogger
     from hathor.nanocontracts.rng import NanoRNG
     from hathor.nanocontracts.runner import Runner
-    from hathor.nanocontracts.types import NCArgs
+    from hathor.nanocontracts.types import Address, NCArgs
 
 
 NCAttrCache: TypeAlias = dict[bytes, Any] | None
@@ -309,3 +309,7 @@ def get_contract(
         """
         from hathor.nanocontracts.contract_accessor import ContractAccessor
         return ContractAccessor(runner=self.__runner, contract_id=contract_id, blueprint_id=blueprint_id)
+
+    def transfer_to_address(self, address: Address, amount: Amount) -> None:
+        """Transfer a given amount to an address balance."""
+        self.__runner.syscall_transfer_to_address(address, amount)

hathor/nanocontracts/runner/runner.py

@@ -61,6 +61,7 @@
     NC_FALLBACK_METHOD,
     NC_INITIALIZE_METHOD,
     Address,
+    Amount,
     BaseTokenAction,
     BlueprintId,
     ContractId,
@@ -512,6 +513,15 @@ def _validate_balances(self, ctx: Context) -> None:
     def _commit_all_changes_to_storage(self) -> None:
         """Commit all change trackers."""
         assert self._call_info is not None
+
+        transfer_header = tx.get_transfer_header()
+        for input_ in transfer_header.inputs:
+            token_uid = tx.get_token_uid(input_.token_index)
+            self.block_storage.add_address_balance(-input_.amount, token_uid)
+        for output_ in transfer_header.outputs:
+            token_uid = tx.get_token_uid(output_.token_index)
+            self.block_storage.add_address_balance(output_.amount, token_uid)
+
         for nc_id, change_trackers in self._call_info.change_trackers.items():
             assert len(change_trackers) == 1
             change_tracker = change_trackers[0]
@@ -1105,6 +1115,23 @@ def syscall_change_blueprint(self, blueprint_id: BlueprintId) -> None:
         nc_storage = self.get_current_changes_tracker(last_call_record.contract_id)
         nc_storage.set_blueprint_id(blueprint_id)
 
+    @_forbid_syscall_from_view('transfer_to_address')
+    def syscall_transfer_to_address(self, address: Address, amount: Amount) -> None:
+        if amount < 0:
+            raise NCInvalidSyscall('amount cannot be negative')
+
+        if amount == 0:
+            # XXX Should we fail?
+            return
+
+        # XXX Should we check for the size to prevent miscalling with a contract id?
+        if not isinstance(address, Address):
+            raise NCInvalidSyscall('only addresses are allowed')
+
+        call_record = self.get_current_call_record()
+        changes_tracker = self.get_current_changes_tracker(call_record.contract_id)
+        changes_tracker.add_address_balance(address, amount)
+
 
 class RunnerFactory:
     __slots__ = ('reactor', 'settings', 'tx_storage', 'nc_storage_factory')

hathor/nanocontracts/storage/block_storage.py

@@ -22,15 +22,16 @@
 from hathor.nanocontracts.storage.contract_storage import NCContractStorage
 from hathor.nanocontracts.storage.patricia_trie import NodeId, PatriciaTrie
 from hathor.nanocontracts.storage.token_proxy import TokenProxy
-from hathor.nanocontracts.types import Address, ContractId, TokenUid
+from hathor.nanocontracts.types import Address, Amount, ContractId, TokenUid
 from hathor.transaction.headers.nano_header import ADDRESS_SEQNUM_SIZE
 from hathor.utils import leb128
 
 
 class _Tag(Enum):
     CONTRACT = b'\0'
     TOKEN = b'\1'
-    ADDRESS = b'\2'
+    ADDRESS_SEQNUM = b'\2'
+    ADDRESS_BALANCE = b'\3'
 
 
 class ContractKey(NamedTuple):
@@ -47,11 +48,18 @@ def __bytes__(self):
         return _Tag.TOKEN.value + self.token_id
 
 
-class AddressKey(NamedTuple):
+class AddressSeqnumKey(NamedTuple):
     address: Address
 
     def __bytes__(self):
-        return _Tag.ADDRESS.value + self.address
+        return _Tag.ADDRESS_SEQNUM.value + self.address
+
+
+class AddressBalanceKey(NamedTuple):
+    address: Address
+
+    def __bytes__(self):
+        return _Tag.ADDRESS_BALANCE.value + self.address
 
 
 class NCBlockStorage:
@@ -142,11 +150,30 @@ def create_token(self, token_id: TokenUid, token_name: str, token_symbol: str) -
         token_description_bytes = self._TOKEN_DESCRIPTION_NC_TYPE.to_bytes(token_description)
         self._block_trie.update(bytes(key), token_description_bytes)
 
+    def get_address_balance(self, address: Address, token_id) -> Amount:
+        key = AddressBalanceKey(address, token_id)
+        try:
+            balance_bytes = self._block_trie.get(bytes(key))
+        except KeyError:
+            return Amount(0)
+        else:
+            balance, buf = leb128.decode_unsigned(balance_bytes)
+            assert len(buf) == 0
+            return balance
+
+    def add_address_balance(self, address: Address, amount: Amount, token_id: TokenUid) -> None:
+        key = AddressBalanceKey(address, token_id)
+        balance = self.get_address_balance(address, token_id)
+        balance += amount
+        assert balance >= 0
+        balance_bytes = leb128.encode_unsigned(amount)
+        self._block_trie.update(bytes(key), balance_bytes)
+
     def get_address_seqnum(self, address: Address) -> int:
         """Get the latest seqnum for an address.
 
         For clarity, new transactions must have a GREATER seqnum to be able to be executed."""
-        key = AddressKey(address)
+        key = AddressSeqnumKey(address)
         try:
             seqnum_bytes = self._block_trie.get(bytes(key))
         except KeyError:
@@ -161,6 +188,6 @@ def set_address_seqnum(self, address: Address, seqnum: int) -> None:
         assert seqnum >= 0
         old_seqnum = self.get_address_seqnum(address)
         assert seqnum > old_seqnum
-        key = AddressKey(address)
+        key = AddressSeqnumKey(address)
         seqnum_bytes = leb128.encode_unsigned(seqnum, max_bytes=ADDRESS_SEQNUM_SIZE)
         self._block_trie.update(bytes(key), seqnum_bytes)

hathor/nanocontracts/storage/changes_tracker.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import itertools
+from collections import defaultdict
 from dataclasses import dataclass
 from enum import Enum
 from types import MappingProxyType
@@ -31,7 +32,7 @@
     NCContractStorage,
 )
 from hathor.nanocontracts.storage.types import _NOT_PROVIDED, DeletedKey, DeletedKeyType
-from hathor.nanocontracts.types import BlueprintId, ContractId, TokenUid
+from hathor.nanocontracts.types import Address, Amount, BlueprintId, ContractId, TokenUid
 from hathor.transaction.token_info import TokenDescription
 
 T = TypeVar('T')
@@ -81,6 +82,7 @@ def __init__(self, nc_id: ContractId, storage: NCContractStorage):
         self._balance_diff: dict[BalanceKey, int] = {}
         self._authorities_diff: dict[BalanceKey, _NCAuthorityDiff] = {}
         self._created_tokens: dict[TokenUid, TokenDescription] = {}
+        self._transfers: defaultdict[Address, int] = defaultdict(int)
         self._blueprint_id: BlueprintId | None = None
 
         self.has_been_commited = False
@@ -102,6 +104,10 @@ def has_token(self, token_id: TokenUid) -> bool:
             return True
         return self.storage.has_token(token_id)
 
+    def add_address_balance(self, address: Address, amount: Amount) -> None:
+        assert amount >= 0
+        self._transfers[address] += amount
+
     def get_balance_diff(self) -> MappingProxyType[BalanceKey, int]:
         """Return the balance diff of this change tracker."""
         return MappingProxyType(self._balance_diff)
@@ -188,6 +194,9 @@ def commit(self) -> None:
         for td in self._created_tokens.values():
             self.storage.create_token(TokenUid(td.token_id), td.token_name, td.token_symbol)
 
+        for address, amount in self._transfers.items():
+            self.stoarge.add_address_balance(address, amount)
+
         if self._blueprint_id is not None:
             self.storage.set_blueprint_id(self._blueprint_id)
 

hathor/nanocontracts/storage/contract_storage.py

@@ -27,7 +27,7 @@
 from hathor.nanocontracts.storage.patricia_trie import PatriciaTrie
 from hathor.nanocontracts.storage.token_proxy import TokenProxy
 from hathor.nanocontracts.storage.types import _NOT_PROVIDED, DeletedKey, DeletedKeyType
-from hathor.nanocontracts.types import BlueprintId, TokenUid, VertexId
+from hathor.nanocontracts.types import Address, Amount, BlueprintId, TokenUid, VertexId
 from hathor.serialization import Deserializer, Serializer
 
 T = TypeVar('T')
@@ -138,7 +138,7 @@ class NCContractStorage:
 
     This implementation works for both memory and rocksdb backends."""
 
-    def __init__(self, *, trie: PatriciaTrie, nc_id: VertexId, token_proxy: TokenProxy) -> None:
+    def __init__(self, *, trie: PatriciaTrie, nc_id: VertexId, block_proxy: TokenProxy) -> None:
         # State (balances, metadata and attributes)
         self._trie: PatriciaTrie = trie
 
@@ -148,15 +148,18 @@ def __init__(self, *, trie: PatriciaTrie, nc_id: VertexId, token_proxy: TokenPro
         # Flag to check whether any change or commit can be executed.
         self.is_locked = False
 
-        self._token_proxy = token_proxy
+        self._block_proxy = block_proxy
 
     def has_token(self, token_id: TokenUid) -> bool:
         """Return True if token_id exists in the current block."""
-        return self._token_proxy.has_token(token_id)
+        return self._block_proxy.has_token(token_id)
 
     def create_token(self, token_id: TokenUid, token_name: str, token_symbol: str) -> None:
         """Create a new token in the current block."""
-        self._token_proxy.create_token(token_id, token_name, token_symbol)
+        self._block_proxy.create_token(token_id, token_name, token_symbol)
+
+    def add_address_balance(self, address: Address, amount: Amount) -> None:
+        self._block_proxy.add_address_balance(address, amount)
 
     def lock(self) -> None:
         """Lock the storage for changes or commits."""