feat(tron): Add support for VoteWitnessContract - core

- Main logic. Scaffolding in previous commit.
- Supports up to 9 candidates.
- Each layout needs some variation
- Figma design available.

[no changelog]

Author: PrisionMike

core/src/apps/tron/README.md

@@ -32,7 +32,7 @@ Therefore, we do not show any fees to the user on the device. The host applicati
 - [X] `FreezeBalanceV2`: Stake `TRX` to get more `Bandwidth` or `Energy`
 - [X] `UnfreezeBalanceV2`: Unstake frozen `TRX`
 - [X] `WithdrawExpireUnfreeze`: Claim Unstaked `TRX` past their lock-up period
-- [ ] `VoteWitnessAccount`: Vote using earned TRON Power. **Coming soon**
+- [X] `VoteWitnessAccount`: Vote using earned TRON Power.
 
 ### Notes
 

core/src/apps/tron/consts.py

@@ -6,15 +6,11 @@
     from buffer_types import AnyBytes
     from typing import Iterator, Tuple
 
-    from trezor.messages import TronTransferContract
-
-    TronMessageType = TronTransferContract
-
 TYPE_URL_TEMPLATE = "type.googleapis.com/protocol."
 
-# TODO: Use TypeVar like ethereum/keychain.py:MsgInSignTx
 CONTRACT_TYPES = (
     MessageType.TronTransferContract,
+    MessageType.TronVoteWitnessContract,
     MessageType.TronTriggerSmartContract,
     MessageType.TronFreezeBalanceV2Contract,
     MessageType.TronUnfreezeBalanceV2Contract,

core/src/apps/tron/layout.py

@@ -9,7 +9,11 @@
 if TYPE_CHECKING:
     from buffer_types import AnyBytes
 
-    from trezor.messages import TronTransferContract, TronTriggerSmartContract
+    from trezor.messages import (
+        TronTransferContract,
+        TronTriggerSmartContract,
+        TronVoteWitnessContract,
+    )
 
 
 def format_trx_amount(amount: int) -> str:
@@ -154,3 +158,13 @@ async def confirm_withdraw_unfreeze(owner_address: AnyBytes) -> None:
         br_name="tron/claim",
         cancel=True,
     )
+
+
+async def confirm_votes(contract: TronVoteWitnessContract) -> None:
+    from trezor.ui.layouts import confirm_tron_voting
+
+    voting_list: list[tuple[int, str]] = [
+        (vote.count, get_encoded_address(vote.address))
+        for vote in contract.votes
+    ]
+    await confirm_tron_voting(voting_list)

core/src/apps/tron/sign_tx.py

@@ -56,6 +56,7 @@ async def sign_tx(msg: TronSignTx, keychain: Keychain) -> TronSignature:
     if fee_limit > _MAX_FEE_LIMIT:
         raise DataError("Tron: fees too high")
 
+    # TODO: CONTRACT_TYPES are all the Tron message types. Can we do better than manually writing them again in consts.py?
     contract = await call_any(messages.TronContractRequest(), *consts.CONTRACT_TYPES)
     raw_contract = await process_contract(contract, fee_limit)
 
@@ -89,11 +90,12 @@ async def process_contract(
     # But it causes type error in messages.TronRawContract.type.
     from trezor import TR
     from trezor.enums import TronRawContractType
-    from trezor.ui.layouts import confirm_tron_send
 
     _INT64_MAX = const(9_223_372_036_854_775_807)
 
     if messages.TronTransferContract.is_type_of(contract):
+        from trezor.ui.layouts import confirm_tron_send
+
         contract_type = TronRawContractType.TransferContract
         await layout.confirm_transfer_contract(contract)
         if contract.amount > _INT64_MAX:
@@ -148,6 +150,9 @@ async def process_contract(
         contract_type = TronRawContractType.WithdrawExpireUnfreezeContract
         await layout.confirm_withdraw_unfreeze(contract.owner_address)
 
+    elif messages.TronVoteWitnessContract.is_type_of(contract):
+        contract_type = TronRawContractType.VoteWitnessContract
+        await layout.confirm_votes(contract)
     else:
         raise DataError("Tron: contract type unknown")
 

core/src/trezor/ui/layouts/bolt/__init__.py

@@ -1609,6 +1609,17 @@ async def confirm_tron_transfer(
             None,
         )
 
+    async def confirm_tron_voting(voting_list: list[tuple[int, str]]) -> None:
+        await raise_if_not_confirmed(
+            trezorui_api.confirm_properties(
+                title=TR.words__review,
+                items=[(f"{TR.words__votes}: {vote[0]}", f"{vote[1]}\n", True) for vote in voting_list],
+                hold=True,
+            ),
+            br_name="tron/vote",
+            br_code=ButtonRequestType.SignTx,
+        )
+
 
 def confirm_joint_total(spending_amount: str, total_amount: str) -> Awaitable[None]:
     return raise_if_not_confirmed(

core/src/trezor/ui/layouts/caesar/__init__.py

@@ -1672,6 +1672,18 @@ async def confirm_tron_transfer(
             br_name=br_name,
         )
 
+    async def confirm_tron_voting(voting_list: list[tuple[int, str]]) -> None:
+        await raise_if_not_confirmed(
+        trezorui_api.confirm_properties(
+            title=TR.words__review,
+            subtitle=TR.words__voting,
+            items=[(f"{TR.words__votes}: {vote[0]}", vote[1], True) for vote in voting_list],
+            hold=True,
+        ),
+        br_name="tron/vote",
+        br_code=ButtonRequestType.SignTx,
+    )
+
 
 def confirm_joint_total(spending_amount: str, total_amount: str) -> Awaitable[None]:
     return confirm_properties(

core/src/trezor/ui/layouts/delizia/__init__.py

@@ -1565,6 +1565,24 @@ async def confirm_tron_approve(
             None,
         )
 
+    async def confirm_tron_voting(voting_list: list[tuple[int, str]]) -> None:
+
+        item_list : list[StrPropertyType] = []
+        for vote_count, address in voting_list:
+            item_list.append((TR.words__address, address, True))
+            item_list.append((f"\n{TR.words__votes}", f"{vote_count}", False))
+
+        await raise_if_not_confirmed(
+        trezorui_api.confirm_properties(
+            title=TR.words__review,
+            subtitle=TR.words__voting,
+            items=item_list,
+            hold=True,
+        ),
+        br_name="tron/vote",
+        br_code=ButtonRequestType.SignTx,
+    )
+
 
 def confirm_joint_total(spending_amount: str, total_amount: str) -> Awaitable[None]:
     return _confirm_summary(

core/src/trezor/ui/layouts/eckhart/__init__.py

@@ -1662,6 +1662,24 @@ async def confirm_tron_approve(
             None,
         )
 
+    async def confirm_tron_voting(voting_list: list[tuple[int, str]]) -> None:
+
+        item_list : list[StrPropertyType] = []
+        for vote_count, address in voting_list:
+            item_list.append((TR.words__address, address, True))
+            item_list.append((f"\n{TR.words__votes}", f"{vote_count}", False))
+
+        await raise_if_not_confirmed(
+        trezorui_api.confirm_properties(
+            title=TR.words__review,
+            subtitle=TR.words__voting,
+            items=item_list,
+            hold=True,
+        ),
+        br_name="tron/vote",
+        br_code=ButtonRequestType.SignTx,
+    )
+
 
 def confirm_joint_total(spending_amount: str, total_amount: str) -> Awaitable[None]:
     return _confirm_summary(

python/src/trezorlib/tron.py

@@ -14,6 +14,7 @@
         messages.TronFreezeBalanceV2Contract,
         messages.TronUnfreezeBalanceV2Contract,
         messages.TronWithdrawUnfreeze,
+        messages.TronVoteWitnessContract,
     ]
 
 DEFAULT_BIP32_PATH = "m/44h/195h/0h/0/0"
@@ -86,6 +87,15 @@ def from_raw_data(
         contract = messages.TronWithdrawUnfreeze(
             owner_address=raw_contract.owner_address,
         )
+    elif contract_type == messages.TronRawContractType.VoteWitnessContract:
+        raw_contract = load_message(
+            io.BytesIO(parameter_value),
+            messages.TronVoteWitnessContract,
+        )
+        contract = messages.TronVoteWitnessContract(
+            owner_address=raw_contract.owner_address,
+            votes=raw_contract.votes
+        )
     else:
         raise ValueError(f"Unsupported contract type: {contract_type}")