Skip to content

Commit c8b00a4

Browse files
committed
feat(core): add support for tron vote witness
1 parent 839118e commit c8b00a4

File tree

24 files changed

+366
-58
lines changed

24 files changed

+366
-58
lines changed

common/protob/check.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,15 @@
1515
for fn in sorted(glob(os.path.join(MYDIR, "messages-*.proto"))):
1616
with open(fn, "rt") as f:
1717
prefix = EXPECTED_PREFIX_RE.search(fn).group(1).capitalize()
18-
if prefix in ("Bitcoin", "Bootloader", "Common", "Crypto", "Management", "Ton"):
18+
if prefix in (
19+
"Bitcoin",
20+
"Bootloader",
21+
"Common",
22+
"Crypto",
23+
"Management",
24+
"Ton",
25+
"Tron",
26+
):
1927
continue
2028
if prefix == "Nem":
2129
prefix = "NEM"

common/protob/messages-tron.proto

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ message TronSignTx {
3535
required bytes ref_block_bytes = 2; // Reference block number
3636
required bytes ref_block_hash = 3; // Reference block hash
3737
required uint64 expiration = 4; // Transaction expiration
38-
optional string data = 5; // Extra transaction info
38+
optional bytes data = 5; // Extra transaction info
3939
required TronContract contract = 6; // Contract messages
4040
required uint64 timestamp = 7; // UTC timestamp
4141
optional uint64 fee_limit = 8; // Fee limit for smartcontracts
@@ -63,6 +63,7 @@ message TronSignTx {
6363
enum TronResourceCode {
6464
BANDWIDTH = 0x00;
6565
ENERGY = 0x01;
66+
TRON_POWER = 0x02;
6667
}
6768

6869
// Freeze TRX balance
@@ -105,7 +106,17 @@ message TronSignTx {
105106
optional string receiver_address = 4;
106107
}
107108

109+
// Vote Witness Contract
110+
message TronVoteWitnessContract {
111+
message Vote {
112+
required string vote_address = 1;
113+
required uint32 vote_count = 2;
114+
}
115+
repeated Vote votes = 2;
116+
optional bool support = 3;
117+
}
108118
optional TronTransferContract transfer_contract = 2;
119+
optional TronVoteWitnessContract vote_witness_contract = 4;
109120
optional TronFreezeBalanceContract freeze_balance_contract = 11;
110121
optional TronUnfreezeBalanceContract unfreeze_balance_contract = 12;
111122
optional TronWithdrawBalanceContract withdraw_balance_contract = 13;
@@ -115,6 +126,9 @@ message TronSignTx {
115126
optional TronWithdrawExpireUnfreezeContract withdraw_expire_unfreeze_contract = 56;
116127
optional TronDelegateResourceContract delegate_resource_contract = 57;
117128
optional TronUnDelegateResourceContract undelegate_resource_contract = 58;
129+
optional bytes provider = 3;
130+
optional bytes contract_name = 5;
131+
optional uint32 permission_id = 6;
118132
}
119133
}
120134

core/src/apps/tron/address.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,14 @@
22
from trezor.crypto.hashlib import sha3_256
33

44

5-
def get_address_from_public_key(pubkey):
5+
def get_address_from_public_key(pubkey: bytes) -> str:
66
address = b"\x41" + sha3_256(pubkey[1:65], keccak=True).digest()[12:32]
7-
return _address_base58(address)
7+
return address_base58(address)
88

99

10-
def _address_base58(address):
10+
def address_base58(address: bytes) -> str:
1111
return base58.encode_check(address)
1212

1313

14-
def address_to_bytes(address):
14+
def address_to_bytes(address: str) -> bytes:
1515
return base58.decode_check(address)

core/src/apps/tron/layout.py

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
from trezor.strings import format_amount
77
from trezor.ui.layouts import confirm_address, confirm_output, should_show_details
88
from trezor.ui.layouts.lvgl.altcoin import confirm_total_tron
9-
from trezor.utils import chunks
109

1110
from . import tokens
1211

@@ -236,6 +235,24 @@ def require_confirm_undelegate(
236235
)
237236

238237

238+
def require_confirm_vote_witness(
239+
ctx: Context,
240+
signer: str,
241+
votes: list[tuple[str, int]],
242+
support: bool | None,
243+
) -> Awaitable[None]:
244+
from trezor.ui.layouts.lvgl import confirm_tron_vote
245+
246+
return confirm_tron_vote(
247+
ctx,
248+
"Vote for Witness"
249+
if (support is None or support)
250+
else "Remove Vote for Witness",
251+
signer,
252+
votes,
253+
)
254+
255+
239256
def format_amount_trx(value: int, token: tokens.TokenInfo | None) -> str:
240257
if token:
241258
suffix = token.symbol
@@ -245,11 +262,3 @@ def format_amount_trx(value: int, token: tokens.TokenInfo | None) -> str:
245262
decimals = 6
246263

247264
return f"{format_amount(value, decimals)} {suffix}"
248-
249-
250-
def split_address(address):
251-
return chunks(address, 16)
252-
253-
254-
def split_text(text):
255-
return chunks(text, 18)

core/src/apps/tron/serialize.py

Lines changed: 64 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,7 @@
77
# PROTOBUF3 types
88
TYPE_VARINT = 0
99
TYPE_DOUBLE = 1
10-
TYPE_STRING = 2
11-
TYPE_GROUPS = 3
12-
TYPE_GROUPE = 4
10+
TYPE_LEN = 2
1311
TYPE_FLOAT = 5
1412

1513

@@ -56,22 +54,39 @@ def pack_contract(contract, owner_address):
5654
write_varint(retc, 1)
5755
api = "TransferContract"
5856

59-
add_field(cmessage, 1, TYPE_STRING)
57+
add_field(cmessage, 1, TYPE_LEN)
6058
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
61-
add_field(cmessage, 2, TYPE_STRING)
59+
add_field(cmessage, 2, TYPE_LEN)
6260
write_bytes_with_length(
6361
cmessage, base58.decode_check(contract.transfer_contract.to_address)
6462
)
6563
add_field(cmessage, 3, TYPE_VARINT)
6664
write_varint(cmessage, contract.transfer_contract.amount)
65+
elif contract.vote_witness_contract:
66+
write_varint(retc, 4)
67+
api = "VoteWitnessContract"
6768

68-
if contract.trigger_smart_contract:
69+
add_field(cmessage, 1, TYPE_LEN)
70+
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
71+
for vote in contract.vote_witness_contract.votes:
72+
v_message = bytearray()
73+
add_field(cmessage, 2, TYPE_LEN)
74+
add_field(v_message, 1, TYPE_LEN)
75+
write_bytes_with_length(v_message, base58.decode_check(vote.vote_address))
76+
add_field(v_message, 2, TYPE_VARINT)
77+
write_varint(v_message, vote.vote_count)
78+
write_bytes_with_length(cmessage, v_message)
79+
if contract.vote_witness_contract.support is not None:
80+
add_field(cmessage, 3, TYPE_VARINT)
81+
write_varint(cmessage, int(contract.vote_witness_contract.support))
82+
83+
elif contract.trigger_smart_contract:
6984
write_varint(retc, 31)
7085
api = "TriggerSmartContract"
7186

72-
add_field(cmessage, 1, TYPE_STRING)
87+
add_field(cmessage, 1, TYPE_LEN)
7388
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
74-
add_field(cmessage, 2, TYPE_STRING)
89+
add_field(cmessage, 2, TYPE_LEN)
7590
write_bytes_with_length(
7691
cmessage,
7792
base58.decode_check(contract.trigger_smart_contract.contract_address),
@@ -81,7 +96,7 @@ def pack_contract(contract, owner_address):
8196
write_varint(cmessage, contract.trigger_smart_contract.call_value)
8297

8398
# Contract data
84-
add_field(cmessage, 4, TYPE_STRING)
99+
add_field(cmessage, 4, TYPE_LEN)
85100
write_bytes_with_length(cmessage, contract.trigger_smart_contract.data)
86101

87102
if contract.trigger_smart_contract.call_token_value:
@@ -90,11 +105,11 @@ def pack_contract(contract, owner_address):
90105
add_field(cmessage, 6, TYPE_VARINT)
91106
write_varint(cmessage, contract.trigger_smart_contract.asset_id)
92107

93-
if contract.freeze_balance_contract:
108+
elif contract.freeze_balance_contract:
94109
write_varint(retc, 11)
95110
api = "FreezeBalanceContract"
96111

97-
add_field(cmessage, 1, TYPE_STRING)
112+
add_field(cmessage, 1, TYPE_LEN)
98113
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
99114
add_field(cmessage, 2, TYPE_VARINT)
100115
write_varint(cmessage, contract.freeze_balance_contract.frozen_balance)
@@ -104,41 +119,41 @@ def pack_contract(contract, owner_address):
104119
add_field(cmessage, 4, TYPE_VARINT)
105120
write_varint(cmessage, contract.freeze_balance_contract.resource)
106121
if contract.freeze_balance_contract.receiver_address is not None:
107-
add_field(cmessage, 5, TYPE_STRING)
122+
add_field(cmessage, 5, TYPE_LEN)
108123
write_bytes_with_length(
109124
cmessage,
110125
base58.decode_check(contract.freeze_balance_contract.receiver_address),
111126
)
112127

113-
if contract.unfreeze_balance_contract:
128+
elif contract.unfreeze_balance_contract:
114129
write_varint(retc, 12)
115130
api = "UnfreezeBalanceContract"
116131

117-
add_field(cmessage, 1, TYPE_STRING)
132+
add_field(cmessage, 1, TYPE_LEN)
118133
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
119134

120135
if contract.unfreeze_balance_contract.resource is not None:
121136
add_field(cmessage, 2, TYPE_VARINT)
122137
write_varint(cmessage, contract.unfreeze_balance_contract.resource)
123138
if contract.unfreeze_balance_contract.receiver_address is not None:
124-
add_field(cmessage, 3, TYPE_STRING)
139+
add_field(cmessage, 3, TYPE_LEN)
125140
write_bytes_with_length(
126141
cmessage,
127142
base58.decode_check(
128143
contract.unfreeze_balance_contract.receiver_address
129144
),
130145
)
131146

132-
if contract.withdraw_balance_contract:
147+
elif contract.withdraw_balance_contract:
133148
write_varint(retc, 13)
134149
api = "WithdrawBalanceContract"
135-
add_field(cmessage, 1, TYPE_STRING)
150+
add_field(cmessage, 1, TYPE_LEN)
136151
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
137152

138-
if contract.freeze_balance_v2_contract:
153+
elif contract.freeze_balance_v2_contract:
139154
write_varint(retc, 54)
140155
api = "FreezeBalanceV2Contract"
141-
add_field(cmessage, 1, TYPE_STRING)
156+
add_field(cmessage, 1, TYPE_LEN)
142157
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
143158

144159
add_field(cmessage, 2, TYPE_VARINT)
@@ -147,10 +162,10 @@ def pack_contract(contract, owner_address):
147162
add_field(cmessage, 3, TYPE_VARINT)
148163
write_varint(cmessage, contract.freeze_balance_v2_contract.resource)
149164

150-
if contract.unfreeze_balance_v2_contract:
165+
elif contract.unfreeze_balance_v2_contract:
151166
write_varint(retc, 55)
152167
api = "UnfreezeBalanceV2Contract"
153-
add_field(cmessage, 1, TYPE_STRING)
168+
add_field(cmessage, 1, TYPE_LEN)
154169
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
155170

156171
add_field(cmessage, 2, TYPE_VARINT)
@@ -159,23 +174,23 @@ def pack_contract(contract, owner_address):
159174
add_field(cmessage, 3, TYPE_VARINT)
160175
write_varint(cmessage, contract.unfreeze_balance_v2_contract.resource)
161176

162-
if contract.withdraw_expire_unfreeze_contract:
177+
elif contract.withdraw_expire_unfreeze_contract:
163178
write_varint(retc, 56)
164179
api = "WithdrawExpireUnfreezeContract"
165-
add_field(cmessage, 1, TYPE_STRING)
180+
add_field(cmessage, 1, TYPE_LEN)
166181
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
167182

168-
if contract.delegate_resource_contract:
183+
elif contract.delegate_resource_contract:
169184
write_varint(retc, 57)
170185
api = "DelegateResourceContract"
171-
add_field(cmessage, 1, TYPE_STRING)
186+
add_field(cmessage, 1, TYPE_LEN)
172187
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
173188

174189
add_field(cmessage, 2, TYPE_VARINT)
175190
write_varint(cmessage, contract.delegate_resource_contract.resource)
176191
add_field(cmessage, 3, TYPE_VARINT)
177192
write_varint(cmessage, contract.delegate_resource_contract.balance)
178-
add_field(cmessage, 4, TYPE_STRING)
193+
add_field(cmessage, 4, TYPE_LEN)
179194
write_bytes_with_length(
180195
cmessage,
181196
base58.decode_check(contract.delegate_resource_contract.receiver_address),
@@ -184,55 +199,67 @@ def pack_contract(contract, owner_address):
184199
add_field(cmessage, 5, TYPE_VARINT)
185200
write_varint(cmessage, contract.delegate_resource_contract.lock)
186201

187-
if contract.undelegate_resource_contract:
202+
elif contract.undelegate_resource_contract:
188203
write_varint(retc, 58)
189204
api = "UnDelegateResourceContract"
190-
add_field(cmessage, 1, TYPE_STRING)
205+
add_field(cmessage, 1, TYPE_LEN)
191206
write_bytes_with_length(cmessage, base58.decode_check(owner_address))
192207

193208
add_field(cmessage, 2, TYPE_VARINT)
194209
write_varint(cmessage, contract.undelegate_resource_contract.resource)
195210
add_field(cmessage, 3, TYPE_VARINT)
196211
write_varint(cmessage, contract.undelegate_resource_contract.balance)
197-
add_field(cmessage, 4, TYPE_STRING)
212+
add_field(cmessage, 4, TYPE_LEN)
198213
write_bytes_with_length(
199214
cmessage,
200215
base58.decode_check(contract.undelegate_resource_contract.receiver_address),
201216
)
217+
else:
218+
raise ValueError("Invalid contract type")
202219

203220
# write API
204221
capi = bytearray()
205-
add_field(capi, 1, TYPE_STRING)
222+
add_field(capi, 1, TYPE_LEN)
206223
# write_bytes_with_length(capi, "type.googleapis.com/protocol." + api)
207224
write_bytes_with_length(capi, bytes("type.googleapis.com/protocol." + api, "ascii"))
208225

209226
# extend to capi
210-
add_field(capi, 2, TYPE_STRING)
227+
add_field(capi, 2, TYPE_LEN)
211228
write_bytes_with_length(capi, cmessage)
212229

213230
# extend to contract
214-
add_field(retc, 2, TYPE_STRING)
231+
add_field(retc, 2, TYPE_LEN)
215232
write_bytes_with_length(retc, capi)
233+
234+
if contract.provider:
235+
add_field(retc, 3, TYPE_LEN)
236+
write_bytes_with_length(retc, contract.provider)
237+
if contract.contract_name:
238+
add_field(retc, 4, TYPE_LEN)
239+
write_bytes_with_length(retc, contract.contract_name)
240+
if contract.permission_id is not None:
241+
add_field(retc, 5, TYPE_VARINT)
242+
write_varint(retc, contract.permission_id)
216243
return retc
217244

218245

219246
def serialize(transaction: TronSignTx, owner_address: str):
220247
# transaction parameters
221248
ret = bytearray()
222-
add_field(ret, 1, TYPE_STRING)
249+
add_field(ret, 1, TYPE_LEN)
223250
write_bytes_with_length(ret, transaction.ref_block_bytes)
224-
add_field(ret, 4, TYPE_STRING)
251+
add_field(ret, 4, TYPE_LEN)
225252
write_bytes_with_length(ret, transaction.ref_block_hash)
226253
add_field(ret, 8, TYPE_VARINT)
227254
write_varint(ret, transaction.expiration)
228255
if transaction.data is not None:
229-
add_field(ret, 10, TYPE_STRING)
230-
write_bytes_with_length(ret, bytes(transaction.data, "ascii"))
256+
add_field(ret, 10, TYPE_LEN)
257+
write_bytes_with_length(ret, transaction.data)
231258

232259
# add Contract
233260
retc = pack_contract(transaction.contract, owner_address)
234261

235-
add_field(ret, 11, TYPE_STRING)
262+
add_field(ret, 11, TYPE_LEN)
236263
write_bytes_with_length(ret, retc)
237264
# add timestamp
238265
add_field(ret, 14, TYPE_VARINT)

0 commit comments

Comments
 (0)