M1nd3r
2 months ago
parent
4b68431f7c
commit
5394ce8df7
@ -0,0 +1,216 @@
|
||||
syntax = "proto2";
|
||||
package hw.trezor.messages.thp;
|
||||
|
||||
// Sugar for easier handling in Java
|
||||
option java_package = "com.satoshilabs.trezor.lib.protobuf";
|
||||
option java_outer_classname = "TrezorMessageThp";
|
||||
|
||||
|
||||
/**
|
||||
* Numeric identifiers of pairing methods.
|
||||
* @embed
|
||||
*/
|
||||
enum ThpPairingMethod {
|
||||
NoMethod = 1; // Trust without MITM protection.
|
||||
CodeEntry = 2; // User types code diplayed on Trezor into the host application.
|
||||
QrCode = 3; // User scans code displayed on Trezor into host application.
|
||||
NFC_Unidirectional = 4; // Trezor transmits an authentication key to the host device via NFC.
|
||||
}
|
||||
|
||||
/**
|
||||
* @embed
|
||||
*/
|
||||
message ThpDeviceProperties {
|
||||
optional string internal_model = 1; // Internal model name e.g. "T2B1".
|
||||
optional uint32 model_variant = 2; // Encodes the device properties such as color.
|
||||
optional bool bootloader_mode = 3; // Indicates whether the device is in bootloader or firmware mode.
|
||||
optional uint32 protocol_version = 4; // The communication protocol version supported by the firmware.
|
||||
repeated ThpPairingMethod pairing_methods = 5; // The pairing methods supported by the Trezor.
|
||||
}
|
||||
|
||||
/**
|
||||
* @embed
|
||||
*/
|
||||
message ThpHandshakeCompletionReqNoisePayload {
|
||||
optional bytes host_pairing_credential = 1; // Host's pairing credential
|
||||
repeated ThpPairingMethod pairing_methods = 2; // The pairing methods chosen by the host
|
||||
}
|
||||
|
||||
/**
|
||||
* Request: Ask device for a new session with given passphrase.
|
||||
* @start
|
||||
* @next ThpNewSession
|
||||
*/
|
||||
message ThpCreateNewSession{
|
||||
optional string passphrase = 1;
|
||||
optional bool on_device = 2; // User wants to enter passphrase on the device
|
||||
optional bool derive_cardano = 3; // If True, Cardano keys will be derived. Ignored with BTC-only
|
||||
}
|
||||
|
||||
/**
|
||||
* Response: Contains session_id of the newly created session.
|
||||
* @end
|
||||
*/
|
||||
message ThpNewSession{
|
||||
optional uint32 new_session_id = 1;
|
||||
}
|
||||
|
||||
/**
|
||||
* Request: Start pairing process.
|
||||
* @start
|
||||
* @next ThpCodeEntryCommitment
|
||||
* @next ThpPairingPreparationsFinished
|
||||
*/
|
||||
message ThpStartPairingRequest{
|
||||
optional string host_name = 1; // Human-readable host name
|
||||
}
|
||||
|
||||
/**
|
||||
* Response: Pairing is ready for user input / OOB communication.
|
||||
* @next ThpCodeEntryCpace
|
||||
* @next ThpQrCodeTag
|
||||
* @next ThpNfcUnidirectionalTag
|
||||
*/
|
||||
message ThpPairingPreparationsFinished{
|
||||
}
|
||||
|
||||
/**
|
||||
* Response: If Code Entry is an allowed pairing option, Trezor responds with a commitment.
|
||||
* @next ThpCodeEntryChallenge
|
||||
*/
|
||||
message ThpCodeEntryCommitment {
|
||||
optional bytes commitment = 1; // SHA-256 of Trezor's random 32-byte secret
|
||||
}
|
||||
|
||||
/**
|
||||
* Response: Host responds to Trezor's Code Entry commitment with a challenge.
|
||||
* @next ThpPairingPreparationsFinished
|
||||
*/
|
||||
message ThpCodeEntryChallenge {
|
||||
optional bytes challenge = 1; // host's random 32-byte challenge
|
||||
}
|
||||
|
||||
/**
|
||||
* Request: User selected Code Entry option in Host. Host starts CPACE protocol with Trezor.
|
||||
* @next ThpCodeEntryCpaceTrezor
|
||||
*/
|
||||
message ThpCodeEntryCpaceHost {
|
||||
optional bytes cpace_host_public_key = 1; // Host's ephemeral CPace public key
|
||||
}
|
||||
|
||||
/**
|
||||
* Response: Trezor continues with the CPACE protocol.
|
||||
* @next ThpCodeEntryTag
|
||||
*/
|
||||
message ThpCodeEntryCpaceTrezor {
|
||||
optional bytes cpace_trezor_public_key = 1; // Trezor's ephemeral CPace public key
|
||||
}
|
||||
|
||||
/**
|
||||
* Response: Host continues with the CPACE protocol.
|
||||
* @next ThpCodeEntrySecret
|
||||
*/
|
||||
message ThpCodeEntryTag {
|
||||
optional bytes tag = 2; // SHA-256 of shared secret
|
||||
}
|
||||
|
||||
/**
|
||||
* Response: Trezor finishes the CPACE protocol.
|
||||
* @next ThpCredentialRequest
|
||||
* @next ThpEndRequest
|
||||
*/
|
||||
message ThpCodeEntrySecret {
|
||||
optional bytes secret = 1; // Trezor's secret
|
||||
}
|
||||
|
||||
/**
|
||||
* Request: User selected QR Code pairing option. Host sends a QR Tag.
|
||||
* @next ThpQrCodeSecret
|
||||
*/
|
||||
message ThpQrCodeTag {
|
||||
optional bytes tag = 1; // SHA-256 of shared secret
|
||||
}
|
||||
|
||||
/**
|
||||
* Response: Trezor sends the QR secret.
|
||||
* @next ThpCredentialRequest
|
||||
* @next ThpEndRequest
|
||||
*/
|
||||
message ThpQrCodeSecret {
|
||||
optional bytes secret = 1; // Trezor's secret
|
||||
}
|
||||
|
||||
/**
|
||||
* Request: User selected Unidirectional NFC pairing option. Host sends an Unidirectional NFC Tag.
|
||||
* @next ThpNfcUnidirectionalSecret
|
||||
*/
|
||||
message ThpNfcUnidirectionalTag {
|
||||
optional bytes tag = 1; // SHA-256 of shared secret
|
||||
}
|
||||
|
||||
/**
|
||||
* Response: Trezor sends the Unidirectioal NFC secret.
|
||||
* @next ThpCredentialRequest
|
||||
* @next ThpEndRequest
|
||||
*/
|
||||
message ThpNfcUnidirectionalSecret {
|
||||
optional bytes secret = 1; // Trezor's secret
|
||||
}
|
||||
|
||||
/**
|
||||
* Request: Host requests issuance of a new pairing credential.
|
||||
* @start
|
||||
* @next ThpCredentialResponse
|
||||
*/
|
||||
message ThpCredentialRequest {
|
||||
optional bytes host_static_pubkey = 1; // Host's static public key used in the handshake.
|
||||
}
|
||||
|
||||
/**
|
||||
* Response: Trezor issues a new pairing credential.
|
||||
* @next ThpCredentialRequest
|
||||
* @next ThpEndRequest
|
||||
*/
|
||||
message ThpCredentialResponse {
|
||||
optional bytes trezor_static_pubkey = 1; // Trezor's static public key used in the handshake.
|
||||
optional bytes credential = 2; // The pairing credential issued by the Trezor to the host.
|
||||
}
|
||||
|
||||
/**
|
||||
* Request: Host requests transition to the encrypted traffic phase.
|
||||
* @start
|
||||
* @next ThpEndResponse
|
||||
*/
|
||||
message ThpEndRequest {}
|
||||
|
||||
/**
|
||||
* Response: Trezor approves transition to the encrypted traffic phase
|
||||
* @end
|
||||
*/
|
||||
message ThpEndResponse {}
|
||||
|
||||
/**
|
||||
* Only for internal use.
|
||||
* @embed
|
||||
*/
|
||||
message ThpCredentialMetadata {
|
||||
optional string host_name = 1; // Human-readable host name
|
||||
}
|
||||
|
||||
/**
|
||||
* Only for internal use.
|
||||
* @embed
|
||||
*/
|
||||
message ThpPairingCredential {
|
||||
optional ThpCredentialMetadata cred_metadata = 1; // Credential metadata
|
||||
optional bytes mac = 2; // Message authentication code generated by the Trezor
|
||||
}
|
||||
|
||||
/**
|
||||
* Only for internal use.
|
||||
* @embed
|
||||
*/
|
||||
message ThpAuthenticatedCredentialData {
|
||||
optional bytes host_static_pubkey = 1; // Host's static public key used in the handshake
|
||||
optional ThpCredentialMetadata cred_metadata = 2; // Credential metadata
|
||||
}
|
||||
@ -0,0 +1,23 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from trezor.wire import context
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Callable, ParamSpec
|
||||
|
||||
P = ParamSpec("P")
|
||||
ByteFunc = Callable[P, bytes]
|
||||
|
||||
|
||||
def stored(key: int) -> Callable[[ByteFunc[P]], ByteFunc[P]]:
|
||||
def decorator(func: ByteFunc[P]) -> ByteFunc[P]:
|
||||
def wrapper(*args: P.args, **kwargs: P.kwargs):
|
||||
value = context.cache_get(key)
|
||||
if value is None:
|
||||
value = func(*args, **kwargs)
|
||||
context.cache_set(key, value)
|
||||
return value
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
@ -0,0 +1,51 @@
|
||||
from trezor import log, loop
|
||||
from trezor.enums import FailureType
|
||||
from trezor.messages import Failure, ThpCreateNewSession, ThpNewSession
|
||||
from trezor.wire.context import get_context
|
||||
from trezor.wire.errors import ActionCancelled, DataError
|
||||
from trezor.wire.thp import SessionState
|
||||
|
||||
|
||||
async def create_new_session(message: ThpCreateNewSession) -> ThpNewSession | Failure:
|
||||
from trezor.wire.thp.session_manager import create_new_session
|
||||
|
||||
from apps.common.seed import derive_and_store_roots
|
||||
|
||||
ctx = get_context()
|
||||
|
||||
# Assert that context `ctx` is ManagementSessionContext
|
||||
from trezor.wire.thp.session_context import ManagementSessionContext
|
||||
|
||||
assert isinstance(ctx, ManagementSessionContext)
|
||||
|
||||
channel = ctx.channel
|
||||
|
||||
# Do not use `ctx` beyond this point, as it is techically
|
||||
# allowed to change inbetween await statements
|
||||
|
||||
new_session = create_new_session(channel)
|
||||
try:
|
||||
await derive_and_store_roots(new_session, message)
|
||||
except DataError as e:
|
||||
return Failure(code=FailureType.DataError, message=e.message)
|
||||
except ActionCancelled as e:
|
||||
return Failure(code=FailureType.ActionCancelled, message=e.message)
|
||||
# TODO handle other errors
|
||||
# TODO handle BITCOIN_ONLY
|
||||
|
||||
new_session.set_session_state(SessionState.ALLOCATED)
|
||||
channel.sessions[new_session.session_id] = new_session
|
||||
loop.schedule(new_session.handle())
|
||||
new_session_id: int = new_session.session_id
|
||||
# await get_seed() TODO
|
||||
|
||||
if __debug__:
|
||||
log.debug(
|
||||
__name__,
|
||||
"create_new_session - new session created. Passphrase: %s, Session id: %d\n%s",
|
||||
message.passphrase if message.passphrase is not None else "",
|
||||
new_session.session_id,
|
||||
str(channel.sessions),
|
||||
)
|
||||
|
||||
return ThpNewSession(new_session_id=new_session_id)
|
||||
@ -0,0 +1,110 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from trezor import protobuf
|
||||
from trezor.crypto import hmac
|
||||
from trezor.messages import (
|
||||
ThpAuthenticatedCredentialData,
|
||||
ThpCredentialMetadata,
|
||||
ThpPairingCredential,
|
||||
)
|
||||
from trezor.wire import message_handler
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from apps.common.paths import Slip21Path
|
||||
|
||||
|
||||
def derive_cred_auth_key() -> bytes:
|
||||
"""
|
||||
Derive current credential authentication mac-ing key from device secret.
|
||||
"""
|
||||
from storage.device import get_cred_auth_key_counter, get_device_secret
|
||||
|
||||
from apps.common.seed import Slip21Node
|
||||
|
||||
# Derive the key using SLIP-21 https://github.com/satoshilabs/slips/blob/master/slip-0021.md,
|
||||
# the derivation path is m/"Credential authentication key"/(counter 4-byte BE)
|
||||
|
||||
thp_secret = get_device_secret()
|
||||
label = b"Credential authentication key"
|
||||
counter = get_cred_auth_key_counter()
|
||||
path: Slip21Path = [label, counter]
|
||||
|
||||
symmetric_key_node: Slip21Node = Slip21Node(thp_secret)
|
||||
symmetric_key_node.derive_path(path)
|
||||
cred_auth_key = symmetric_key_node.key()
|
||||
|
||||
return cred_auth_key
|
||||
|
||||
|
||||
def invalidate_cred_auth_key() -> None:
|
||||
from storage.device import increment_cred_auth_key_counter
|
||||
|
||||
increment_cred_auth_key_counter()
|
||||
|
||||
|
||||
def issue_credential(
|
||||
host_static_pubkey: bytes,
|
||||
credential_metadata: ThpCredentialMetadata,
|
||||
) -> bytes:
|
||||
"""
|
||||
Issue a pairing credential binded to the provided host static public key
|
||||
and credential metadata.
|
||||
"""
|
||||
return _issue_credential(
|
||||
derive_cred_auth_key(), host_static_pubkey, credential_metadata
|
||||
)
|
||||
|
||||
|
||||
def validate_credential(
|
||||
encoded_pairing_credential_message: bytes,
|
||||
host_static_pubkey: bytes,
|
||||
) -> bool:
|
||||
"""
|
||||
Validate a pairing credential binded to the provided host static public key.
|
||||
"""
|
||||
return _validate_credential(
|
||||
derive_cred_auth_key(), encoded_pairing_credential_message, host_static_pubkey
|
||||
)
|
||||
|
||||
|
||||
def _issue_credential(
|
||||
cred_auth_key: bytes,
|
||||
host_static_pubkey: bytes,
|
||||
credential_metadata: ThpCredentialMetadata,
|
||||
) -> bytes:
|
||||
proto_msg = ThpAuthenticatedCredentialData(
|
||||
host_static_pubkey=host_static_pubkey,
|
||||
cred_metadata=credential_metadata,
|
||||
)
|
||||
authenticated_credential_data = _encode_message_into_new_buffer(proto_msg)
|
||||
mac = hmac(hmac.SHA256, cred_auth_key, authenticated_credential_data).digest()
|
||||
|
||||
proto_msg = ThpPairingCredential(cred_metadata=credential_metadata, mac=mac)
|
||||
credential_raw = _encode_message_into_new_buffer(proto_msg)
|
||||
return credential_raw
|
||||
|
||||
|
||||
def _validate_credential(
|
||||
cred_auth_key: bytes,
|
||||
encoded_pairing_credential_message: bytes,
|
||||
host_static_pubkey: bytes,
|
||||
) -> bool:
|
||||
expected_type = protobuf.type_for_name("ThpPairingCredential")
|
||||
credential = message_handler.wrap_protobuf_load(
|
||||
encoded_pairing_credential_message, expected_type
|
||||
)
|
||||
assert ThpPairingCredential.is_type_of(credential)
|
||||
proto_msg = ThpAuthenticatedCredentialData(
|
||||
host_static_pubkey=host_static_pubkey,
|
||||
cred_metadata=credential.cred_metadata,
|
||||
)
|
||||
authenticated_credential_data = _encode_message_into_new_buffer(proto_msg)
|
||||
mac = hmac(hmac.SHA256, cred_auth_key, authenticated_credential_data).digest()
|
||||
return mac == credential.mac
|
||||
|
||||
|
||||
def _encode_message_into_new_buffer(msg: protobuf.MessageType) -> bytes:
|
||||
msg_len = protobuf.encoded_length(msg)
|
||||
new_buffer = bytearray(msg_len)
|
||||
protobuf.encode(new_buffer, msg)
|
||||
return new_buffer
|
||||
@ -0,0 +1,379 @@
|
||||
from typing import TYPE_CHECKING
|
||||
from ubinascii import hexlify
|
||||
|
||||
from trezor import loop, protobuf
|
||||
from trezor.crypto.hashlib import sha256
|
||||
from trezor.enums import MessageType, ThpPairingMethod
|
||||
from trezor.messages import (
|
||||
ThpCodeEntryChallenge,
|
||||
ThpCodeEntryCommitment,
|
||||
ThpCodeEntryCpaceHost,
|
||||
ThpCodeEntryCpaceTrezor,
|
||||
ThpCodeEntrySecret,
|
||||
ThpCodeEntryTag,
|
||||
ThpCredentialMetadata,
|
||||
ThpCredentialRequest,
|
||||
ThpCredentialResponse,
|
||||
ThpEndRequest,
|
||||
ThpEndResponse,
|
||||
ThpNfcUnidirectionalSecret,
|
||||
ThpNfcUnidirectionalTag,
|
||||
ThpPairingPreparationsFinished,
|
||||
ThpQrCodeSecret,
|
||||
ThpQrCodeTag,
|
||||
ThpStartPairingRequest,
|
||||
)
|
||||
from trezor.wire.errors import ActionCancelled, UnexpectedMessage
|
||||
from trezor.wire.thp import ChannelState, ThpError, crypto
|
||||
from trezor.wire.thp.pairing_context import PairingContext
|
||||
|
||||
from .credential_manager import issue_credential
|
||||
|
||||
if __debug__:
|
||||
from trezor import log
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any, Callable, Concatenate, Container, ParamSpec, Tuple
|
||||
|
||||
P = ParamSpec("P")
|
||||
FuncWithContext = Callable[Concatenate[PairingContext, P], Any]
|
||||
|
||||
#
|
||||
# Helpers - decorators
|
||||
|
||||
|
||||
def check_state_and_log(
|
||||
*allowed_states: ChannelState,
|
||||
) -> Callable[[FuncWithContext], FuncWithContext]:
|
||||
def decorator(f: FuncWithContext) -> FuncWithContext:
|
||||
def inner(context: PairingContext, *args: P.args, **kwargs: P.kwargs) -> object:
|
||||
_check_state(context, *allowed_states)
|
||||
if __debug__:
|
||||
try:
|
||||
log.debug(__name__, "started %s", f.__name__)
|
||||
except AttributeError:
|
||||
log.debug(
|
||||
__name__,
|
||||
"started a function that cannot be named, because it raises AttributeError, eg. closure",
|
||||
)
|
||||
return f(context, *args, **kwargs)
|
||||
|
||||
return inner
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def check_method_is_allowed(
|
||||
pairing_method: ThpPairingMethod,
|
||||
) -> Callable[[FuncWithContext], FuncWithContext]:
|
||||
def decorator(f: FuncWithContext) -> FuncWithContext:
|
||||
def inner(context: PairingContext, *args: P.args, **kwargs: P.kwargs) -> object:
|
||||
_check_method_is_allowed(context, pairing_method)
|
||||
return f(context, *args, **kwargs)
|
||||
|
||||
return inner
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
#
|
||||
# Pairing handlers
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TP1)
|
||||
async def handle_pairing_request(
|
||||
ctx: PairingContext, message: protobuf.MessageType
|
||||
) -> ThpEndResponse:
|
||||
|
||||
if not ThpStartPairingRequest.is_type_of(message):
|
||||
raise UnexpectedMessage("Unexpected message")
|
||||
|
||||
ctx.host_name = message.host_name or ""
|
||||
|
||||
skip_pairing = _is_method_included(ctx, ThpPairingMethod.NoMethod)
|
||||
if skip_pairing:
|
||||
return await _end_pairing(ctx)
|
||||
|
||||
await _prepare_pairing(ctx)
|
||||
await ctx.write(ThpPairingPreparationsFinished())
|
||||
ctx.channel_ctx.set_channel_state(ChannelState.TP3)
|
||||
response = await show_display_data(ctx, _get_possible_pairing_methods(ctx))
|
||||
|
||||
# TODO disable NFC (if enabled)
|
||||
response = await _handle_different_pairing_methods(ctx, response)
|
||||
|
||||
while ThpCredentialRequest.is_type_of(response):
|
||||
response = await _handle_credential_request(ctx, response)
|
||||
|
||||
return await _handle_end_request(ctx, response)
|
||||
|
||||
|
||||
async def _prepare_pairing(ctx: PairingContext) -> None:
|
||||
|
||||
if _is_method_included(ctx, ThpPairingMethod.CodeEntry):
|
||||
await _handle_code_entry_is_included(ctx)
|
||||
|
||||
if _is_method_included(ctx, ThpPairingMethod.QrCode):
|
||||
_handle_qr_code_is_included(ctx)
|
||||
|
||||
if _is_method_included(ctx, ThpPairingMethod.NFC_Unidirectional):
|
||||
_handle_nfc_unidirectional_is_included(ctx)
|
||||
|
||||
|
||||
async def show_display_data(ctx: PairingContext, expected_types: Container[int] = ()):
|
||||
|
||||
read_task = ctx.read(expected_types)
|
||||
cancel_task = ctx.display_data.get_display_layout()
|
||||
race = loop.race(read_task, cancel_task)
|
||||
result = await race
|
||||
|
||||
if read_task in race.finished:
|
||||
return result
|
||||
if cancel_task in race.finished:
|
||||
raise ActionCancelled
|
||||
else:
|
||||
return Exception("Should not happen") # TODO
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TP1)
|
||||
async def _handle_code_entry_is_included(ctx: PairingContext) -> None:
|
||||
commitment = sha256(ctx.secret).digest()
|
||||
|
||||
challenge_message = await ctx.call( # noqa: F841
|
||||
ThpCodeEntryCommitment(commitment=commitment), ThpCodeEntryChallenge
|
||||
)
|
||||
ctx.channel_ctx.set_channel_state(ChannelState.TP2)
|
||||
|
||||
if not ThpCodeEntryChallenge.is_type_of(challenge_message):
|
||||
raise UnexpectedMessage("Unexpected message")
|
||||
|
||||
if challenge_message.challenge is None:
|
||||
raise Exception("Invalid message")
|
||||
|
||||
code_code_entry_hash = sha256(
|
||||
challenge_message.challenge
|
||||
+ ctx.secret
|
||||
+ bytes("PairingMethod_CodeEntry", "utf-8")
|
||||
).digest() # TODO add handshake hash
|
||||
ctx.display_data.code_code_entry = (
|
||||
int.from_bytes(code_code_entry_hash, "big") % 1000000
|
||||
)
|
||||
ctx.display_data.display_code_entry = True
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TP1, ChannelState.TP2)
|
||||
def _handle_qr_code_is_included(ctx: PairingContext) -> None:
|
||||
ctx.display_data.code_qr_code = sha256(
|
||||
ctx.secret + bytes("PairingMethod_QrCode", "utf-8")
|
||||
).digest()[
|
||||
:16
|
||||
] # TODO add handshake hash
|
||||
ctx.display_data.display_qr_code = True
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TP1, ChannelState.TP2)
|
||||
def _handle_nfc_unidirectional_is_included(ctx: PairingContext) -> None:
|
||||
ctx.display_data.code_nfc_unidirectional = sha256(
|
||||
ctx.secret + bytes("PairingMethod_NfcUnidirectional", "utf-8")
|
||||
).digest()[
|
||||
:16
|
||||
] # TODO add handshake hash
|
||||
ctx.display_data.display_nfc_unidirectional = True
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TP3)
|
||||
async def _handle_different_pairing_methods(
|
||||
ctx: PairingContext, response: protobuf.MessageType
|
||||
) -> protobuf.MessageType:
|
||||
if ThpCodeEntryCpaceHost.is_type_of(response):
|
||||
return await _handle_code_entry_cpace(ctx, response)
|
||||
if ThpQrCodeTag.is_type_of(response):
|
||||
return await _handle_qr_code_tag(ctx, response)
|
||||
if ThpNfcUnidirectionalTag.is_type_of(response):
|
||||
return await _handle_nfc_unidirectional_tag(ctx, response)
|
||||
raise UnexpectedMessage("Unexpected message")
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TP3)
|
||||
@check_method_is_allowed(ThpPairingMethod.CodeEntry)
|
||||
async def _handle_code_entry_cpace(
|
||||
ctx: PairingContext, message: protobuf.MessageType
|
||||
) -> protobuf.MessageType:
|
||||
from trezor.wire.thp.cpace import Cpace
|
||||
|
||||
# TODO check that ThpCodeEntryCpaceHost message is valid
|
||||
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(message, ThpCodeEntryCpaceHost)
|
||||
if message.cpace_host_public_key is None:
|
||||
raise ThpError("Message ThpCodeEntryCpaceHost has no public key")
|
||||
|
||||
ctx.cpace = Cpace(message.cpace_host_public_key)
|
||||
assert ctx.display_data.code_code_entry is not None
|
||||
ctx.cpace.generate_keys_and_secret(
|
||||
ctx.display_data.code_code_entry.to_bytes(6, "big")
|
||||
)
|
||||
|
||||
ctx.channel_ctx.set_channel_state(ChannelState.TP4)
|
||||
response = await ctx.call(
|
||||
ThpCodeEntryCpaceTrezor(cpace_trezor_public_key=ctx.cpace.trezor_public_key),
|
||||
ThpCodeEntryTag,
|
||||
)
|
||||
return await _handle_code_entry_tag(ctx, response)
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TP4)
|
||||
@check_method_is_allowed(ThpPairingMethod.CodeEntry)
|
||||
async def _handle_code_entry_tag(
|
||||
ctx: PairingContext, message: protobuf.MessageType
|
||||
) -> protobuf.MessageType:
|
||||
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(message, ThpCodeEntryTag)
|
||||
|
||||
expected_tag = sha256(ctx.cpace.shared_secret).digest()
|
||||
if expected_tag != message.tag:
|
||||
print(
|
||||
"expected code entry tag:", hexlify(expected_tag).decode()
|
||||
) # TODO remove after testing
|
||||
print(
|
||||
"expected code entry shared secret:",
|
||||
hexlify(ctx.cpace.shared_secret).decode(),
|
||||
) # TODO remove after testing
|
||||
raise ThpError("Unexpected Entry Code Tag")
|
||||
|
||||
return await _handle_secret_reveal(
|
||||
ctx,
|
||||
msg=ThpCodeEntrySecret(secret=ctx.secret),
|
||||
)
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TP3)
|
||||
@check_method_is_allowed(ThpPairingMethod.QrCode)
|
||||
async def _handle_qr_code_tag(
|
||||
ctx: PairingContext, message: protobuf.MessageType
|
||||
) -> protobuf.MessageType:
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(message, ThpQrCodeTag)
|
||||
expected_tag = sha256(ctx.display_data.code_qr_code).digest()
|
||||
if expected_tag != message.tag:
|
||||
print(
|
||||
"expected qr code tag:", hexlify(expected_tag).decode()
|
||||
) # TODO remove after testing
|
||||
raise ThpError("Unexpected QR Code Tag")
|
||||
|
||||
return await _handle_secret_reveal(
|
||||
ctx,
|
||||
msg=ThpQrCodeSecret(secret=ctx.secret),
|
||||
)
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TP3)
|
||||
@check_method_is_allowed(ThpPairingMethod.NFC_Unidirectional)
|
||||
async def _handle_nfc_unidirectional_tag(
|
||||
ctx: PairingContext, message: protobuf.MessageType
|
||||
) -> protobuf.MessageType:
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(message, ThpNfcUnidirectionalTag)
|
||||
|
||||
expected_tag = sha256(ctx.display_data.code_nfc_unidirectional).digest()
|
||||
if expected_tag != message.tag:
|
||||
print(
|
||||
"expected nfc tag:", hexlify(expected_tag).decode()
|
||||
) # TODO remove after testing
|
||||
raise ThpError("Unexpected NFC Unidirectional Tag")
|
||||
|
||||
return await _handle_secret_reveal(
|
||||
ctx,
|
||||
msg=ThpNfcUnidirectionalSecret(secret=ctx.secret),
|
||||
)
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TP3, ChannelState.TP4)
|
||||
async def _handle_secret_reveal(
|
||||
ctx: PairingContext,
|
||||
msg: protobuf.MessageType,
|
||||
) -> protobuf.MessageType:
|
||||
ctx.channel_ctx.set_channel_state(ChannelState.TC1)
|
||||
return await ctx.call_any(
|
||||
msg,
|
||||
MessageType.ThpCredentialRequest,
|
||||
MessageType.ThpEndRequest,
|
||||
)
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TC1)
|
||||
async def _handle_credential_request(
|
||||
ctx: PairingContext, message: protobuf.MessageType
|
||||
) -> protobuf.MessageType:
|
||||
ctx.secret
|
||||
|
||||
if not ThpCredentialRequest.is_type_of(message):
|
||||
raise UnexpectedMessage("Unexpected message")
|
||||
if message.host_static_pubkey is None:
|
||||
raise Exception("Invalid message") # TODO change failure type
|
||||
|
||||
trezor_static_pubkey = crypto.get_trezor_static_pubkey()
|
||||
credential_metadata = ThpCredentialMetadata(host_name=ctx.host_name)
|
||||
credential = issue_credential(message.host_static_pubkey, credential_metadata)
|
||||
|
||||
return await ctx.call_any(
|
||||
ThpCredentialResponse(
|
||||
trezor_static_pubkey=trezor_static_pubkey, credential=credential
|
||||
),
|
||||
MessageType.ThpCredentialRequest,
|
||||
MessageType.ThpEndRequest,
|
||||
)
|
||||
|
||||
|
||||
@check_state_and_log(ChannelState.TC1)
|
||||
async def _handle_end_request(
|
||||
ctx: PairingContext, message: protobuf.MessageType
|
||||
) -> ThpEndResponse:
|
||||
if not ThpEndRequest.is_type_of(message):
|
||||
raise UnexpectedMessage("Unexpected message")
|
||||
return await _end_pairing(ctx)
|
||||
|
||||
|
||||
async def _end_pairing(ctx: PairingContext) -> ThpEndResponse:
|
||||
ctx.channel_ctx.set_channel_state(ChannelState.ENCRYPTED_TRANSPORT)
|
||||
return ThpEndResponse()
|
||||
|
||||
|
||||
#
|
||||
# Helpers - checkers
|
||||
|
||||
|
||||
def _check_state(ctx: PairingContext, *allowed_states: ChannelState) -> None:
|
||||
if ctx.channel_ctx.get_channel_state() not in allowed_states:
|
||||
raise UnexpectedMessage("Unexpected message")
|
||||
|
||||
|
||||
def _check_method_is_allowed(ctx: PairingContext, method: ThpPairingMethod) -> None:
|
||||
if not _is_method_included(ctx, method):
|
||||
raise ThpError("Unexpected pairing method")
|
||||
|
||||
|
||||
def _is_method_included(ctx: PairingContext, method: ThpPairingMethod) -> bool:
|
||||
return method in ctx.channel_ctx.selected_pairing_methods
|
||||
|
||||
|
||||
#
|
||||
# Helpers - getters
|
||||
|
||||
|
||||
def _get_possible_pairing_methods(ctx: PairingContext) -> Tuple[int, ...]:
|
||||
return tuple(
|
||||
_get_message_type_for_method(method)
|
||||
for method in ctx.channel_ctx.selected_pairing_methods
|
||||
)
|
||||
|
||||
|
||||
def _get_message_type_for_method(method: int) -> int:
|
||||
if method is ThpPairingMethod.CodeEntry:
|
||||
return MessageType.ThpCodeEntryCpaceHost
|
||||
if method is ThpPairingMethod.NFC_Unidirectional:
|
||||
return MessageType.ThpNfcUnidirectionalTag
|
||||
if method is ThpPairingMethod.QrCode:
|
||||
return MessageType.ThpQrCodeTag
|
||||
raise ValueError("Unexpected pairing method - no message type available")
|
||||
@ -0,0 +1,145 @@
|
||||
import builtins
|
||||
from micropython import const
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from storage.cache_common import DataCache
|
||||
from trezor import utils
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
_MAX_SESSIONS_COUNT = const(10)
|
||||
SESSION_ID_LENGTH = const(32)
|
||||
|
||||
|
||||
class SessionCache(DataCache):
|
||||
def __init__(self) -> None:
|
||||
self.session_id = bytearray(SESSION_ID_LENGTH)
|
||||
if utils.BITCOIN_ONLY:
|
||||
self.fields = (
|
||||
64, # APP_COMMON_SEED
|
||||
2, # APP_COMMON_AUTHORIZATION_TYPE
|
||||
128, # APP_COMMON_AUTHORIZATION_DATA
|
||||
32, # APP_COMMON_NONCE
|
||||
)
|
||||
else:
|
||||
self.fields = (
|
||||
64, # APP_COMMON_SEED
|
||||
2, # APP_COMMON_AUTHORIZATION_TYPE
|
||||
128, # APP_COMMON_AUTHORIZATION_DATA
|
||||
32, # APP_COMMON_NONCE
|
||||
0, # APP_COMMON_DERIVE_CARDANO
|
||||
96, # APP_CARDANO_ICARUS_SECRET
|
||||
96, # APP_CARDANO_ICARUS_TREZOR_SECRET
|
||||
0, # APP_MONERO_LIVE_REFRESH
|
||||
)
|
||||
self.last_usage = 0
|
||||
super().__init__()
|
||||
|
||||
def export_session_id(self) -> bytes:
|
||||
from trezorcrypto import random # avoid pulling in trezor.crypto
|
||||
|
||||
# generate a new session id if we don't have it yet
|
||||
if not self.session_id:
|
||||
self.session_id[:] = random.bytes(SESSION_ID_LENGTH)
|
||||
# export it as immutable bytes
|
||||
return bytes(self.session_id)
|
||||
|
||||
def clear(self) -> None:
|
||||
super().clear()
|
||||
self.last_usage = 0
|
||||
self.session_id[:] = b""
|
||||
|
||||
|
||||
_SESSIONS: list[SessionCache] = []
|
||||
|
||||
|
||||
def initialize() -> None:
|
||||
global _SESSIONS
|
||||
for _ in range(_MAX_SESSIONS_COUNT):
|
||||
_SESSIONS.append(SessionCache())
|
||||
|
||||
for session in _SESSIONS:
|
||||
session.clear()
|
||||
|
||||
|
||||
initialize()
|
||||
|
||||
|
||||
_active_session_idx: int | None = None
|
||||
_session_usage_counter = 0
|
||||
|
||||
|
||||
def get_active_session() -> SessionCache | None:
|
||||
if _active_session_idx is None:
|
||||
return None
|
||||
return _SESSIONS[_active_session_idx]
|
||||
|
||||
|
||||
def start_session(received_session_id: bytes | None = None) -> bytes:
|
||||
global _active_session_idx
|
||||
global _session_usage_counter
|
||||
|
||||
if (
|
||||
received_session_id is not None
|
||||
and len(received_session_id) != SESSION_ID_LENGTH
|
||||
):
|
||||
# Prevent the caller from setting received_session_id=b"" and finding a cleared
|
||||
# session. More generally, short-circuit the session id search, because we know
|
||||
# that wrong-length session ids should not be in cache.
|
||||
# Reduce to "session id not provided" case because that's what we do when
|
||||
# caller supplies an id that is not found.
|
||||
received_session_id = None
|
||||
|
||||
_session_usage_counter += 1
|
||||
|
||||
# attempt to find specified session id
|
||||
if received_session_id:
|
||||
for i in range(_MAX_SESSIONS_COUNT):
|
||||
if _SESSIONS[i].session_id == received_session_id:
|
||||
_active_session_idx = i
|
||||
_SESSIONS[i].last_usage = _session_usage_counter
|
||||
return received_session_id
|
||||
|
||||
# allocate least recently used session
|
||||
lru_counter = _session_usage_counter
|
||||
lru_session_idx = 0
|
||||
for i in range(_MAX_SESSIONS_COUNT):
|
||||
if _SESSIONS[i].last_usage < lru_counter:
|
||||
lru_counter = _SESSIONS[i].last_usage
|
||||
lru_session_idx = i
|
||||
|
||||
_active_session_idx = lru_session_idx
|
||||
selected_session = _SESSIONS[lru_session_idx]
|
||||
selected_session.clear()
|
||||
selected_session.last_usage = _session_usage_counter
|
||||
return selected_session.export_session_id()
|
||||
|
||||
|
||||
def end_current_session() -> None:
|
||||
global _active_session_idx
|
||||
|
||||
if _active_session_idx is None:
|
||||
return
|
||||
|
||||
_SESSIONS[_active_session_idx].clear()
|
||||
_active_session_idx = None
|
||||
|
||||
|
||||
def get_int_all_sessions(key: int) -> builtins.set[int]:
|
||||
values = builtins.set()
|
||||
for session in _SESSIONS:
|
||||
encoded = session.get(key)
|
||||
if encoded is not None:
|
||||
values.add(int.from_bytes(encoded, "big"))
|
||||
return values
|
||||
|
||||
|
||||
def clear_all() -> None:
|
||||
global _active_session_idx
|
||||
_active_session_idx = None
|
||||
for session in _SESSIONS:
|
||||
session.clear()
|
||||
@ -0,0 +1,179 @@
|
||||
from micropython import const
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from trezor import utils
|
||||
|
||||
# Traditional cache keys
|
||||
APP_COMMON_SEED = const(0)
|
||||
APP_COMMON_AUTHORIZATION_TYPE = const(1)
|
||||
APP_COMMON_AUTHORIZATION_DATA = const(2)
|
||||
APP_COMMON_NONCE = const(3)
|
||||
if not utils.BITCOIN_ONLY:
|
||||
APP_COMMON_DERIVE_CARDANO = const(4)
|
||||
APP_CARDANO_ICARUS_SECRET = const(5)
|
||||
APP_CARDANO_ICARUS_TREZOR_SECRET = const(6)
|
||||
APP_MONERO_LIVE_REFRESH = const(7)
|
||||
|
||||
# Cache keys for THP channel
|
||||
if utils.USE_THP:
|
||||
CHANNEL_HANDSHAKE_HASH = const(0)
|
||||
CHANNEL_KEY_RECEIVE = const(1)
|
||||
CHANNEL_KEY_SEND = const(2)
|
||||
CHANNEL_NONCE_RECEIVE = const(3)
|
||||
CHANNEL_NONCE_SEND = const(4)
|
||||
|
||||
# Keys that are valid across sessions
|
||||
SESSIONLESS_FLAG = const(128)
|
||||
APP_COMMON_SEED_WITHOUT_PASSPHRASE = const(0 | SESSIONLESS_FLAG)
|
||||
APP_COMMON_SAFETY_CHECKS_TEMPORARY = const(1 | SESSIONLESS_FLAG)
|
||||
APP_COMMON_REQUEST_PIN_LAST_UNLOCK = const(2 | SESSIONLESS_FLAG)
|
||||
APP_COMMON_BUSY_DEADLINE_MS = const(3 | SESSIONLESS_FLAG)
|
||||
APP_MISC_COSI_NONCE = const(4 | SESSIONLESS_FLAG)
|
||||
APP_MISC_COSI_COMMITMENT = const(5 | SESSIONLESS_FLAG)
|
||||
APP_RECOVERY_REPEATED_BACKUP_UNLOCKED = const(6 | SESSIONLESS_FLAG)
|
||||
|
||||
|
||||
# === Homescreen storage ===
|
||||
# This does not logically belong to the "cache" functionality, but the cache module is
|
||||
# a convenient place to put this.
|
||||
# When a Homescreen layout is instantiated, it checks the value of `homescreen_shown`
|
||||
# to know whether it should render itself or whether the result of a previous instance
|
||||
# is still on. This way we can avoid unnecessary fadeins/fadeouts when a workflow ends.
|
||||
HOMESCREEN_ON = object()
|
||||
LOCKSCREEN_ON = object()
|
||||
BUSYSCREEN_ON = object()
|
||||
homescreen_shown: object | None = None
|
||||
|
||||
# Timestamp of last autolock activity.
|
||||
# Here to persist across main loop restart between workflows.
|
||||
autolock_last_touch: int | None = None
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Sequence, TypeVar, overload
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class InvalidSessionError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class DataCache:
|
||||
fields: Sequence[int]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.data = [bytearray(f + 1) for f in self.fields]
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@overload
|
||||
def get(self, key: int) -> bytes | None: # noqa: F811
|
||||
...
|
||||
|
||||
@overload
|
||||
def get(self, key: int, default: T) -> bytes | T: # noqa: F811
|
||||
...
|
||||
|
||||
def get(self, key: int, default: T | None = None) -> bytes | T | None: # noqa: F811
|
||||
utils.ensure(key < len(self.fields))
|
||||
if self.data[key][0] != 1:
|
||||
return default
|
||||
return bytes(self.data[key][1:])
|
||||
|
||||
def get_bool(self, key: int) -> bool: # noqa: F811
|
||||
return self.get(key) is not None
|
||||
|
||||
def get_int(
|
||||
self, key: int, default: T | None = None
|
||||
) -> int | T | None: # noqa: F811
|
||||
encoded = self.get(key)
|
||||
if encoded is None:
|
||||
return default
|
||||
else:
|
||||
return int.from_bytes(encoded, "big")
|
||||
|
||||
def is_set(self, key: int) -> bool:
|
||||
utils.ensure(key < len(self.fields))
|
||||
return self.data[key][0] == 1
|
||||
|
||||
def set(self, key: int, value: bytes) -> None:
|
||||
utils.ensure(key < len(self.fields))
|
||||
utils.ensure(len(value) <= self.fields[key])
|
||||
self.data[key][0] = 1
|
||||
self.data[key][1:] = value
|
||||
|
||||
def set_bool(self, key: int, value: bool) -> None:
|
||||
utils.ensure(
|
||||
self._get_length(key) == 0
|
||||
) # skipping get_length in production build
|
||||
if value:
|
||||
self.set(key, b"")
|
||||
else:
|
||||
self.delete(key)
|
||||
|
||||
def set_int(self, key: int, value: int) -> None:
|
||||
length = self.fields[key]
|
||||
encoded = value.to_bytes(length, "big")
|
||||
|
||||
# Ensure that the value fits within the length. Micropython's int.to_bytes()
|
||||
# doesn't raise OverflowError.
|
||||
assert int.from_bytes(encoded, "big") == value
|
||||
|
||||
self.set(key, encoded)
|
||||
|
||||
def delete(self, key: int) -> None:
|
||||
utils.ensure(key < len(self.fields))
|
||||
self.data[key][:] = b"\x00"
|
||||
|
||||
def clear(self) -> None:
|
||||
for i in range(len(self.fields)):
|
||||
self.delete(i)
|
||||
|
||||
def _get_length(self, key: int) -> int:
|
||||
utils.ensure(key < len(self.fields))
|
||||
return self.fields[key]
|
||||
|
||||
|
||||
class SessionlessCache(DataCache):
|
||||
def __init__(self) -> None:
|
||||
self.fields = (
|
||||
64, # APP_COMMON_SEED_WITHOUT_PASSPHRASE
|
||||
1, # APP_COMMON_SAFETY_CHECKS_TEMPORARY
|
||||
8, # APP_COMMON_REQUEST_PIN_LAST_UNLOCK
|
||||
8, # APP_COMMON_BUSY_DEADLINE_MS
|
||||
32, # APP_MISC_COSI_NONCE
|
||||
32, # APP_MISC_COSI_COMMITMENT
|
||||
0, # APP_RECOVERY_REPEATED_BACKUP_UNLOCKED
|
||||
)
|
||||
super().__init__()
|
||||
|
||||
def get(self, key: int, default: T | None = None) -> bytes | T | None: # noqa: F811
|
||||
return super().get(key & ~SESSIONLESS_FLAG, default)
|
||||
|
||||
def get_bool(self, key: int) -> bool: # noqa: F811
|
||||
return super().get_bool(key & ~SESSIONLESS_FLAG)
|
||||
|
||||
def get_int(
|
||||
self, key: int, default: T | None = None
|
||||
) -> int | T | None: # noqa: F811
|
||||
return super().get_int(key & ~SESSIONLESS_FLAG, default)
|
||||
|
||||
def is_set(self, key: int) -> bool:
|
||||
return super().is_set(key & ~SESSIONLESS_FLAG)
|
||||
|
||||
def set(self, key: int, value: bytes) -> None:
|
||||
super().set(key & ~SESSIONLESS_FLAG, value)
|
||||
|
||||
def set_bool(self, key: int, value: bool) -> None:
|
||||
super().set_bool(key & ~SESSIONLESS_FLAG, value)
|
||||
|
||||
def set_int(self, key: int, value: int) -> None:
|
||||
super().set_int(key & ~SESSIONLESS_FLAG, value)
|
||||
|
||||
def delete(self, key: int) -> None:
|
||||
super().delete(key & ~SESSIONLESS_FLAG)
|
||||
|
||||
def clear(self) -> None:
|
||||
for i in range(len(self.fields)):
|
||||
self.delete(i)
|
||||
@ -0,0 +1,325 @@
|
||||
import builtins
|
||||
from micropython import const # pyright: ignore[reportMissingModuleSource]
|
||||
from typing import TYPE_CHECKING # pyright: ignore[reportShadowedImports]
|
||||
|
||||
from storage.cache_common import DataCache
|
||||
from trezor import utils
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import TypeVar # pyright: ignore[reportShadowedImports]
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
if __debug__:
|
||||
from trezor import log
|
||||
|
||||
# THP specific constants
|
||||
_MAX_CHANNELS_COUNT = const(10)
|
||||
_MAX_SESSIONS_COUNT = const(20)
|
||||
|
||||
|
||||
_CHANNEL_STATE_LENGTH = const(1)
|
||||
_WIRE_INTERFACE_LENGTH = const(1)
|
||||
_SESSION_STATE_LENGTH = const(1)
|
||||
_CHANNEL_ID_LENGTH = const(2)
|
||||
SESSION_ID_LENGTH = const(1)
|
||||
BROADCAST_CHANNEL_ID = const(65535)
|
||||
KEY_LENGTH = const(32)
|
||||
TAG_LENGTH = const(16)
|
||||
_UNALLOCATED_STATE = const(0)
|
||||
MANAGEMENT_SESSION_ID = const(0)
|
||||
|
||||
|
||||
class ConnectionCache(DataCache):
|
||||
def __init__(self) -> None:
|
||||
self.channel_id = bytearray(_CHANNEL_ID_LENGTH)
|
||||
self.last_usage = 0
|
||||
super().__init__()
|
||||
|
||||
def clear(self) -> None:
|
||||
self.channel_id[:] = b""
|
||||
self.last_usage = 0
|
||||
super().clear()
|
||||
|
||||
|
||||
class ChannelCache(ConnectionCache):
|
||||
def __init__(self) -> None:
|
||||
self.host_ephemeral_pubkey = bytearray(KEY_LENGTH)
|
||||
self.state = bytearray(_CHANNEL_STATE_LENGTH)
|
||||
self.iface = bytearray(1) # TODO add decoding
|
||||
self.sync = 0x80 # can_send_bit | sync_receive_bit | sync_send_bit | rfu(5)
|
||||
self.session_id_counter = 0x00
|
||||
self.fields = (
|
||||
32, # CHANNEL_HANDSHAKE_HASH
|
||||
32, # CHANNEL_KEY_RECEIVE
|
||||
32, # CHANNEL_KEY_SEND
|
||||
8, # CHANNEL_NONCE_RECEIVE
|
||||
8, # CHANNEL_NONCE_SEND
|
||||
)
|
||||
super().__init__()
|
||||
|
||||
def clear(self) -> None:
|
||||
self.state[:] = bytearray(
|
||||
int.to_bytes(0, _CHANNEL_STATE_LENGTH, "big")
|
||||
) # Set state to UNALLOCATED
|
||||
# TODO clear all keys
|
||||
super().clear()
|
||||
|
||||
|
||||
class SessionThpCache(ConnectionCache):
|
||||
def __init__(self) -> None:
|
||||
self.session_id = bytearray(SESSION_ID_LENGTH)
|
||||
self.state = bytearray(_SESSION_STATE_LENGTH)
|
||||
if utils.BITCOIN_ONLY:
|
||||
self.fields = (
|
||||
64, # APP_COMMON_SEED
|
||||
2, # APP_COMMON_AUTHORIZATION_TYPE
|
||||
128, # APP_COMMON_AUTHORIZATION_DATA
|
||||
32, # APP_COMMON_NONCE
|
||||
)
|
||||
else:
|
||||
self.fields = (
|
||||
64, # APP_COMMON_SEED
|
||||
2, # APP_COMMON_AUTHORIZATION_TYPE
|
||||
128, # APP_COMMON_AUTHORIZATION_DATA
|
||||
32, # APP_COMMON_NONCE
|
||||
0, # APP_COMMON_DERIVE_CARDANO
|
||||
96, # APP_CARDANO_ICARUS_SECRET
|
||||
96, # APP_CARDANO_ICARUS_TREZOR_SECRET
|
||||
0, # APP_MONERO_LIVE_REFRESH
|
||||
)
|
||||
super().__init__()
|
||||
|
||||
def clear(self) -> None:
|
||||
self.state[:] = bytearray(int.to_bytes(0, 1, "big")) # Set state to UNALLOCATED
|
||||
self.session_id[:] = b""
|
||||
super().clear()
|
||||
|
||||
|
||||
_CHANNELS: list[ChannelCache] = []
|
||||
_SESSIONS: list[SessionThpCache] = []
|
||||
|
||||
|
||||
def initialize() -> None:
|
||||
global _CHANNELS
|
||||
global _SESSIONS
|
||||
|
||||
for _ in range(_MAX_CHANNELS_COUNT):
|
||||
_CHANNELS.append(ChannelCache())
|
||||
for _ in range(_MAX_SESSIONS_COUNT):
|
||||
_SESSIONS.append(SessionThpCache())
|
||||
|
||||
for channel in _CHANNELS:
|
||||
channel.clear()
|
||||
for session in _SESSIONS:
|
||||
session.clear()
|
||||
|
||||
|
||||
initialize()
|
||||
|
||||
|
||||
# THP vars
|
||||
_next_unauthenicated_session_index: int = 0 # TODO remove
|
||||
|
||||
# First unauthenticated channel will have index 0
|
||||
_usage_counter = 0
|
||||
|
||||
# with this (arbitrary) value=4659, the first allocated channel will have cid=1234 (hex)
|
||||
cid_counter: int = 4659 # TODO change to random value on start
|
||||
|
||||
|
||||
def get_new_unauthenticated_channel(iface: bytes) -> ChannelCache:
|
||||
if len(iface) != _WIRE_INTERFACE_LENGTH:
|
||||
raise Exception("Invalid WireInterface (encoded) length")
|
||||
|
||||
new_cid = get_next_channel_id()
|
||||
index = _get_next_unauthenticated_channel_index()
|
||||
|
||||
# clear sessions from replaced channel
|
||||
if _get_channel_state(_CHANNELS[index]) != _UNALLOCATED_STATE:
|
||||
old_cid = _CHANNELS[index].channel_id
|
||||
clear_sessions_with_channel_id(old_cid)
|
||||
|
||||
_CHANNELS[index] = ChannelCache()
|
||||
_CHANNELS[index].channel_id[:] = new_cid
|
||||
_CHANNELS[index].last_usage = _get_usage_counter_and_increment()
|
||||
_CHANNELS[index].state[:] = bytearray(
|
||||
_UNALLOCATED_STATE.to_bytes(_CHANNEL_STATE_LENGTH, "big")
|
||||
)
|
||||
_CHANNELS[index].iface[:] = bytearray(iface)
|
||||
return _CHANNELS[index]
|
||||
|
||||
|
||||
def get_all_allocated_channels() -> list[ChannelCache]:
|
||||
_list: list[ChannelCache] = []
|
||||
for channel in _CHANNELS:
|
||||
if _get_channel_state(channel) != _UNALLOCATED_STATE:
|
||||
_list.append(channel)
|
||||
return _list
|
||||
|
||||
|
||||
def get_all_allocated_sessions() -> list[SessionThpCache]:
|
||||
if __debug__:
|
||||
from trezor.utils import get_bytes_as_str
|
||||
_list: list[SessionThpCache] = []
|
||||
for session in _SESSIONS:
|
||||
if _get_session_state(session) != _UNALLOCATED_STATE:
|
||||
_list.append(session)
|
||||
if __debug__:
|
||||
log.debug(
|
||||
__name__,
|
||||
"session with channel_id: %s and session_id: %s is in ALLOCATED state",
|
||||
get_bytes_as_str(session.channel_id),
|
||||
get_bytes_as_str(session.session_id),
|
||||
)
|
||||
elif __debug__:
|
||||
log.debug(__name__, "session %s is in UNALLOCATED state", str(session))
|
||||
return _list
|
||||
|
||||
|
||||
def set_channel_host_ephemeral_key(channel: ChannelCache, key: bytearray) -> None:
|
||||
if len(key) != KEY_LENGTH:
|
||||
raise Exception("Invalid key length")
|
||||
channel.host_ephemeral_pubkey = key
|
||||
|
||||
|
||||
def get_new_session(channel: ChannelCache):
|
||||
new_sid = get_next_session_id(channel)
|
||||
index = _get_next_session_index()
|
||||
|
||||
_SESSIONS[index] = SessionThpCache()
|
||||
_SESSIONS[index].channel_id[:] = channel.channel_id
|
||||
_SESSIONS[index].session_id[:] = new_sid
|
||||
_SESSIONS[index].last_usage = _get_usage_counter_and_increment()
|
||||
channel.last_usage = (
|
||||
_get_usage_counter_and_increment()
|
||||
) # increment also use of the channel so it does not get replaced
|
||||
_SESSIONS[index].state[:] = bytearray(
|
||||
_UNALLOCATED_STATE.to_bytes(_SESSION_STATE_LENGTH, "big")
|
||||
)
|
||||
return _SESSIONS[index]
|
||||
|
||||
|
||||
def _get_usage_counter() -> int:
|
||||
global _usage_counter
|
||||
return _usage_counter
|
||||
|
||||
|
||||
def _get_usage_counter_and_increment() -> int:
|
||||
global _usage_counter
|
||||
_usage_counter += 1
|
||||
return _usage_counter
|
||||
|
||||
|
||||
def _get_next_unauthenticated_channel_index() -> int:
|
||||
idx = _get_unallocated_channel_index()
|
||||
if idx is not None:
|
||||
return idx
|
||||
return get_least_recently_used_item(_CHANNELS, max_count=_MAX_CHANNELS_COUNT)
|
||||
|
||||
|
||||
def _get_next_session_index() -> int:
|
||||
idx = _get_unallocated_session_index()
|
||||
if idx is not None:
|
||||
return idx
|
||||
return get_least_recently_used_item(_SESSIONS, _MAX_SESSIONS_COUNT)
|
||||
|
||||
|
||||
def _get_unallocated_channel_index() -> int | None:
|
||||
for i in range(_MAX_CHANNELS_COUNT):
|
||||
if _get_channel_state(_CHANNELS[i]) is _UNALLOCATED_STATE:
|
||||
return i
|
||||
return None
|
||||
|
||||
|
||||
def _get_unallocated_session_index() -> int | None:
|
||||
for i in range(_MAX_SESSIONS_COUNT):
|
||||
if (_SESSIONS[i]) is _UNALLOCATED_STATE:
|
||||
return i
|
||||
return None
|
||||
|
||||
|
||||
def _get_channel_state(channel: ChannelCache) -> int:
|
||||
if channel is None:
|
||||
return _UNALLOCATED_STATE
|
||||
return int.from_bytes(channel.state, "big")
|
||||
|
||||
|
||||
def _get_session_state(session: SessionThpCache) -> int:
|
||||
if session is None:
|
||||
return _UNALLOCATED_STATE
|
||||
return int.from_bytes(session.state, "big")
|
||||
|
||||
|
||||
def get_next_channel_id() -> bytes:
|
||||
global cid_counter
|
||||
while True:
|
||||
cid_counter += 1
|
||||
if cid_counter >= BROADCAST_CHANNEL_ID:
|
||||
cid_counter = 1
|
||||
if _is_cid_unique():
|
||||
break
|
||||
return cid_counter.to_bytes(_CHANNEL_ID_LENGTH, "big")
|
||||
|
||||
|
||||
def get_next_session_id(channel: ChannelCache) -> bytes:
|
||||
while True:
|
||||
if channel.session_id_counter >= 255:
|
||||
channel.session_id_counter = 1
|
||||
else:
|
||||
channel.session_id_counter += 1
|
||||
if _is_session_id_unique(channel):
|
||||
break
|
||||
new_sid = channel.session_id_counter
|
||||
return new_sid.to_bytes(SESSION_ID_LENGTH, "big")
|
||||
|
||||
|
||||
def _is_session_id_unique(channel: ChannelCache) -> bool:
|
||||
for session in _SESSIONS:
|
||||
if session.channel_id == channel.channel_id:
|
||||
if session.session_id == channel.session_id_counter:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _is_cid_unique() -> bool:
|
||||
for session in _SESSIONS:
|
||||
if cid_counter == _get_cid(session):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _get_cid(session: SessionThpCache) -> int:
|
||||
return int.from_bytes(session.session_id[2:], "big")
|
||||
|
||||
|
||||
def get_least_recently_used_item(
|
||||
list: list[ChannelCache] | list[SessionThpCache], max_count: int
|
||||
):
|
||||
lru_counter = _get_usage_counter()
|
||||
lru_item_index = 0
|
||||
for i in range(max_count):
|
||||
if list[i].last_usage < lru_counter:
|
||||
lru_counter = list[i].last_usage
|
||||
lru_item_index = i
|
||||
return lru_item_index
|
||||
|
||||
|
||||
def get_int_all_sessions(key: int) -> builtins.set[int]:
|
||||
values = builtins.set()
|
||||
for session in _SESSIONS:
|
||||
encoded = session.get(key)
|
||||
if encoded is not None:
|
||||
values.add(int.from_bytes(encoded, "big"))
|
||||
return values
|
||||
|
||||
|
||||
def clear_sessions_with_channel_id(channel_id: bytes):
|
||||
for session in _SESSIONS:
|
||||
if session.channel_id == channel_id:
|
||||
session.clear()
|
||||
|
||||
|
||||
def clear_all() -> None:
|
||||
for session in _SESSIONS:
|
||||
session.clear()
|
||||
@ -0,0 +1,8 @@
|
||||
# Automatically generated by pb2py
|
||||
# fmt: off
|
||||
# isort:skip_file
|
||||
|
||||
NoMethod = 1
|
||||
CodeEntry = 2
|
||||
QrCode = 3
|
||||
NFC_Unidirectional = 4
|
||||
@ -0,0 +1,255 @@
|
||||
from micropython import const
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from storage.cache_common import InvalidSessionError
|
||||
from trezor import log, loop, protobuf, utils, workflow
|
||||
from trezor.enums import FailureType
|
||||
from trezor.messages import Failure
|
||||
from trezor.wire.context import Context, UnexpectedMessageException, with_context
|
||||
from trezor.wire.errors import ActionCancelled, DataError, Error, UnexpectedMessage
|
||||
from trezor.wire.protocol_common import Message
|
||||
|
||||
# Import all errors into namespace, so that `wire.Error` is available from
|
||||
# other packages.
|
||||
from trezor.wire.errors import * # isort:skip # noqa: F401,F403
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any, Callable, Container
|
||||
|
||||
from trezor.wire import Handler, LoadedMessageType
|
||||
|
||||
HandlerFinder = Callable[[Any], Handler | None]
|
||||
|
||||
# If set to False protobuf messages marked with "experimental_message" option are rejected.
|
||||
EXPERIMENTAL_ENABLED = False
|
||||
|
||||
|
||||
def wrap_protobuf_load(
|
||||
buffer: bytes,
|
||||
expected_type: type[LoadedMessageType],
|
||||
) -> LoadedMessageType:
|
||||
try:
|
||||
msg = protobuf.decode(buffer, expected_type, EXPERIMENTAL_ENABLED)
|
||||
if __debug__ and utils.EMULATOR:
|
||||
log.debug(
|
||||
__name__, "received message contents:\n%s", utils.dump_protobuf(msg)
|
||||
)
|
||||
return msg
|
||||
except Exception as e:
|
||||
if __debug__:
|
||||
log.exception(__name__, e)
|
||||
if e.args:
|
||||
raise DataError("Failed to decode message: " + " ".join(e.args))
|
||||
else:
|
||||
raise DataError("Failed to decode message")
|
||||
|
||||
|
||||
_PROTOBUF_BUFFER_SIZE = const(8192)
|
||||
|
||||
WIRE_BUFFER = bytearray(_PROTOBUF_BUFFER_SIZE)
|
||||
|
||||
if __debug__:
|
||||
PROTOBUF_BUFFER_SIZE_DEBUG = 1024
|
||||
WIRE_BUFFER_DEBUG = bytearray(PROTOBUF_BUFFER_SIZE_DEBUG)
|
||||
|
||||
|
||||
async def handle_single_message(
|
||||
ctx: Context,
|
||||
msg: Message,
|
||||
handler_finder: HandlerFinder,
|
||||
) -> bool:
|
||||
"""Handle a message that was loaded from USB by the caller.
|
||||
|
||||
Find the appropriate handler, run it and write its result on the wire. In case
|
||||
a problem is encountered at any point, write the appropriate error on the wire.
|
||||
|
||||
The return value indicates whether to override the default restarting behavior. If
|
||||
`False` is returned, the caller is allowed to clear the loop and restart the
|
||||
MicroPython machine (see `session.py`). This would lose all state and incurs a cost
|
||||
in terms of repeated startup time. When handling the message didn't cause any
|
||||
significant fragmentation (e.g., if decoding the message was skipped), or if
|
||||
the type of message is supposed to be optimized and not disrupt the running state,
|
||||
this function will return `True`.
|
||||
"""
|
||||
if __debug__:
|
||||
try:
|
||||
msg_type = protobuf.type_for_wire(msg.type).MESSAGE_NAME
|
||||
except Exception:
|
||||
msg_type = f"{msg.type} - unknown message type"
|
||||
if ctx.channel_id is not None:
|
||||
sid = int.from_bytes(ctx.channel_id, "big")
|
||||
log.debug(
|
||||
__name__,
|
||||
"%s:%x receive: <%s>",
|
||||
ctx.iface.iface_num(),
|
||||
sid,
|
||||
msg_type,
|
||||
)
|
||||
else:
|
||||
log.debug(
|
||||
__name__,
|
||||
"%s:unknown_sid receive: <%s>",
|
||||
ctx.iface.iface_num(),
|
||||
msg_type,
|
||||
)
|
||||
|
||||
res_msg: protobuf.MessageType | None = None
|
||||
|
||||
# # We need to find a handler for this message type.
|
||||
# try:
|
||||
# handler = find_handler(ctx.iface, msg.type)
|
||||
# except Error as exc:
|
||||
# # Handlers are allowed to exception out. In that case, we can skip decoding
|
||||
# # and return the error.
|
||||
# await ctx.write(failure(exc))
|
||||
# return True
|
||||
|
||||
# We need to find a handler for this message type. Should not raise.
|
||||
|
||||
handler: Handler | None = handler_finder(msg.type)
|
||||
|
||||
if handler is None:
|
||||
# If no handler is found, we can skip decoding and directly
|
||||
# respond with failure.
|
||||
print("handler is none")
|
||||
await ctx.write(unexpected_message())
|
||||
return True
|
||||
|
||||
if msg.type in workflow.ALLOW_WHILE_LOCKED:
|
||||
workflow.autolock_interrupts_workflow = False
|
||||
|
||||
# Here we make sure we always respond with a Failure response
|
||||
# in case of any errors.
|
||||
try:
|
||||
# Find a protobuf.MessageType subclass that describes this
|
||||
# message. Raises if the type is not found.
|
||||
req_type = protobuf.type_for_wire(msg.type)
|
||||
|
||||
# Try to decode the message according to schema from
|
||||
# `req_type`. Raises if the message is malformed.
|
||||
req_msg = wrap_protobuf_load(msg.data, req_type)
|
||||
|
||||
# Create the handler task.
|
||||
task = handler(req_msg)
|
||||
|
||||
# Run the workflow task. Workflow can do more on-the-wire
|
||||
# communication inside, but it should eventually return a
|
||||
# response message, or raise an exception (a rather common
|
||||
# thing to do). Exceptions are handled in the code below.
|
||||
|
||||
# Spawn a workflow around the task. This ensures that concurrent
|
||||
# workflows are shut down.
|
||||
res_msg = await workflow.spawn(with_context(ctx, task))
|
||||
|
||||
except UnexpectedMessageException:
|
||||
# Workflow was trying to read a message from the wire, and
|
||||
# something unexpected came in. See Context.read() for
|
||||
# example, which expects some particular message and raises
|
||||
# UnexpectedMessage if another one comes in.
|
||||
# In order not to lose the message, we return it to the caller.
|
||||
|
||||
# We process the unexpected message by aborting the current workflow and
|
||||
# possibly starting a new one, initiated by that message. (The main usecase
|
||||
# being, the host does not finish the workflow, we want other callers to
|
||||
# be able to do their own thing.)
|
||||
#
|
||||
# The message is stored in the exception, which we re-raise for the caller
|
||||
# to process. It is not a standard exception that should be logged and a result
|
||||
# sent to the wire.
|
||||
raise
|
||||
except BaseException as exc:
|
||||
# Either:
|
||||
# - the message had a type that has a registered handler, but does not have
|
||||
# a protobuf class
|
||||
# - the message was not valid protobuf
|
||||
# - workflow raised some kind of an exception while running
|
||||
# - something canceled the workflow from the outside
|
||||
if __debug__:
|
||||
if isinstance(exc, ActionCancelled):
|
||||
log.debug(__name__, "cancelled: %s", exc.message)
|
||||
elif isinstance(exc, loop.TaskClosed):
|
||||
log.debug(__name__, "cancelled: loop task was closed")
|
||||
else:
|
||||
log.exception(__name__, exc)
|
||||
res_msg = failure(exc)
|
||||
|
||||
if res_msg is not None:
|
||||
# perform the write outside the big try-except block, so that usb write
|
||||
# problem bubbles up
|
||||
await ctx.write(res_msg)
|
||||
|
||||
# Look into `AVOID_RESTARTING_FOR` to see if this message should avoid restarting.
|
||||
return msg.type in AVOID_RESTARTING_FOR
|
||||
|
||||
|
||||
AVOID_RESTARTING_FOR: Container[int] = ()
|
||||
|
||||
|
||||
def failure(exc: BaseException) -> Failure:
|
||||
if isinstance(exc, Error):
|
||||
return Failure(code=exc.code, message=exc.message)
|
||||
elif isinstance(exc, loop.TaskClosed):
|
||||
return Failure(code=FailureType.ActionCancelled, message="Cancelled")
|
||||
elif isinstance(exc, InvalidSessionError):
|
||||
return Failure(code=FailureType.InvalidSession, message="Invalid session")
|
||||
else:
|
||||
# NOTE: when receiving generic `FirmwareError` on non-debug build,
|
||||
# change the `if __debug__` to `if True` to get the full error message.
|
||||
if __debug__:
|
||||
message = str(exc)
|
||||
else:
|
||||
message = "Firmware error"
|
||||
return Failure(code=FailureType.FirmwareError, message=message)
|
||||
|
||||
|
||||
def unexpected_message() -> Failure:
|
||||
return Failure(code=FailureType.UnexpectedMessage, message="Unexpected message")
|
||||
|
||||
|
||||
def find_handler(msg_type: int) -> Handler:
|
||||
from apps import workflow_handlers
|
||||
|
||||
handler = workflow_handlers.find_registered_handler(msg_type)
|
||||
if handler is None:
|
||||
raise UnexpectedMessage("Unexpected message")
|
||||
|
||||
for filter in filters:
|
||||
handler = filter(msg_type, handler)
|
||||
|
||||
return handler
|
||||
|
||||
|
||||
filters: list[Callable[[int, Handler], Handler]] = []
|
||||
"""Filters for the wire handler.
|
||||
|
||||
Filters are applied in order. Each filter gets a message id and a preceding handler. It
|
||||
must either return a handler (the same one or a modified one), or raise an exception
|
||||
that gets sent to wire directly.
|
||||
|
||||
Filters are not applied to debug sessions.
|
||||
|
||||
The filters are designed for:
|
||||
* rejecting messages -- while in Recovery mode, most messages are not allowed
|
||||
* adding additional behavior -- while device is soft-locked, a PIN screen will be shown
|
||||
before allowing a message to trigger its original behavior.
|
||||
|
||||
For this, the filters are effectively deny-first. If an earlier filter rejects the
|
||||
message, the later filters are not called. But if a filter adds behavior, the latest
|
||||
filter "wins" and the latest behavior triggers first.
|
||||
Please note that this behavior is really unsuited to anything other than what we are
|
||||
using it for now. It might be necessary to modify the semantics if we need more complex
|
||||
usecases.
|
||||
|
||||
NB: `filters` is currently public so callers can have control over where they insert
|
||||
new filters, but removal should be done using `remove_filter`!
|
||||
We should, however, change it such that filters must be added using an `add_filter`
|
||||
and `filters` becomes private!
|
||||
"""
|
||||
|
||||
|
||||
def remove_filter(filter):
|
||||
try:
|
||||
filters.remove(filter)
|
||||
except ValueError:
|
||||
pass
|
||||
@ -0,0 +1,70 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from trezor import protobuf
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from trezorio import WireInterface
|
||||
from typing import Container, TypeVar, overload
|
||||
|
||||
from storage.cache_common import DataCache
|
||||
|
||||
LoadedMessageType = TypeVar("LoadedMessageType", bound=protobuf.MessageType)
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class Message:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message_type: int,
|
||||
message_data: bytes,
|
||||
) -> None:
|
||||
self.data = message_data
|
||||
self.type = message_type
|
||||
|
||||
def to_bytes(self):
|
||||
return self.type.to_bytes(2, "big") + self.data
|
||||
|
||||
|
||||
class Context:
|
||||
def __init__(self, iface: WireInterface, channel_id: bytes) -> None:
|
||||
self.iface: WireInterface = iface
|
||||
self.channel_id: bytes = channel_id
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@overload
|
||||
async def read(
|
||||
self, expected_types: Container[int]
|
||||
) -> protobuf.MessageType: ...
|
||||
|
||||
@overload
|
||||
async def read(
|
||||
self, expected_types: Container[int], expected_type: type[LoadedMessageType]
|
||||
) -> LoadedMessageType: ...
|
||||
|
||||
async def read(
|
||||
self,
|
||||
expected_types: Container[int],
|
||||
expected_type: type[protobuf.MessageType] | None = None,
|
||||
) -> protobuf.MessageType: ...
|
||||
|
||||
async def write(self, msg: protobuf.MessageType) -> None: ...
|
||||
|
||||
async def call(
|
||||
self,
|
||||
msg: protobuf.MessageType,
|
||||
expected_type: type[LoadedMessageType],
|
||||
) -> LoadedMessageType:
|
||||
assert expected_type.MESSAGE_WIRE_TYPE is not None
|
||||
|
||||
await self.write(msg)
|
||||
del msg
|
||||
return await self.read((expected_type.MESSAGE_WIRE_TYPE,), expected_type)
|
||||
|
||||
@property
|
||||
def cache(self) -> DataCache: ...
|
||||
|
||||
|
||||
class WireError(Exception):
|
||||
pass
|
||||
@ -0,0 +1,75 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from trezor.wire.protocol_common import WireError
|
||||
|
||||
|
||||
class ThpError(WireError):
|
||||
pass
|
||||
|
||||
|
||||
class ThpDecryptionError(ThpError):
|
||||
pass
|
||||
|
||||
|
||||
class ThpUnallocatedSessionError(ThpError):
|
||||
def __init__(self, session_id: int):
|
||||
self.session_id = session_id
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from enum import IntEnum
|
||||
else:
|
||||
IntEnum = object
|
||||
|
||||
|
||||
class ThpErrorType(IntEnum):
|
||||
TRANSPORT_BUSY = 1
|
||||
UNALLOCATED_CHANNEL = 2
|
||||
DECRYPTION_FAILED = 3
|
||||
|
||||
|
||||
class ChannelState(IntEnum):
|
||||
UNALLOCATED = 0
|
||||
TH1 = 1
|
||||
TH2 = 2
|
||||
TP1 = 3
|
||||
TP2 = 4
|
||||
TP3 = 5
|
||||
TP4 = 6
|
||||
TC1 = 7
|
||||
ENCRYPTED_TRANSPORT = 8
|
||||
|
||||
|
||||
class SessionState(IntEnum):
|
||||
UNALLOCATED = 0
|
||||
ALLOCATED = 1
|
||||
MANAGEMENT = 2
|
||||
|
||||
|
||||
class WireInterfaceType(IntEnum):
|
||||
MOCK = 0
|
||||
USB = 1
|
||||
BLE = 2
|
||||
|
||||
|
||||
def is_channel_state_pairing(state: int) -> bool:
|
||||
if state in (
|
||||
ChannelState.TP1,
|
||||
ChannelState.TP2,
|
||||
ChannelState.TP3,
|
||||
ChannelState.TP4,
|
||||
ChannelState.TC1,
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
if __debug__:
|
||||
|
||||
def state_to_str(state: int) -> str:
|
||||
name = {
|
||||
v: k for k, v in ChannelState.__dict__.items() if not k.startswith("__")
|
||||
}.get(state)
|
||||
if name is not None:
|
||||
return name
|
||||
return "UNKNOWN_STATE"
|
||||
@ -0,0 +1,102 @@
|
||||
from storage.cache_thp import ChannelCache
|
||||
from trezor import log
|
||||
from trezor.wire.thp import ThpError
|
||||
|
||||
|
||||
def is_ack_valid(cache: ChannelCache, ack_bit: int) -> bool:
|
||||
"""
|
||||
Checks if:
|
||||
- an ACK message is expected
|
||||
- the received ACK message acknowledges correct sequence number (bit)
|
||||
"""
|
||||
if not _is_ack_expected(cache):
|
||||
return False
|
||||
|
||||
if not _has_ack_correct_sync_bit(cache, ack_bit):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def _is_ack_expected(cache: ChannelCache) -> bool:
|
||||
is_expected: bool = not is_sending_allowed(cache)
|
||||
if __debug__ and not is_expected:
|
||||
log.debug(__name__, "Received unexpected ACK message")
|
||||
return is_expected
|
||||
|
||||
|
||||
def _has_ack_correct_sync_bit(cache: ChannelCache, sync_bit: int) -> bool:
|
||||
is_correct: bool = get_send_seq_bit(cache) == sync_bit
|
||||
if __debug__ and not is_correct:
|
||||
log.debug(__name__, "Received ACK message with wrong ack bit")
|
||||
return is_correct
|
||||
|
||||
|
||||
def is_sending_allowed(cache: ChannelCache) -> bool:
|
||||
"""
|
||||
Checks whether sending a message in the provided channel is allowed.
|
||||
|
||||
Note: Sending a message in a channel before receipt of ACK message for the previously
|
||||
sent message (in the channel) is prohibited, as it can lead to desynchronization.
|
||||
"""
|
||||
return bool(cache.sync >> 7)
|
||||
|
||||
|
||||
def get_send_seq_bit(cache: ChannelCache) -> int:
|
||||
"""
|
||||
Returns the sequential number (bit) of the next message to be sent
|
||||
in the provided channel.
|
||||
"""
|
||||
return (cache.sync & 0x20) >> 5
|
||||
|
||||
|
||||
def get_expected_receive_seq_bit(cache: ChannelCache) -> int:
|
||||
"""
|
||||
Returns the (expected) sequential number (bit) of the next message
|
||||
to be received in the provided channel.
|
||||
"""
|
||||
return (cache.sync & 0x40) >> 6
|
||||
|
||||
|
||||
def set_sending_allowed(cache: ChannelCache, sending_allowed: bool) -> None:
|
||||
"""
|
||||
Set the flag whether sending a message in this channel is allowed or not.
|
||||
"""
|
||||
cache.sync &= 0x7F
|
||||
if sending_allowed:
|
||||
cache.sync |= 0x80
|
||||
|
||||
|
||||
def set_expected_receive_seq_bit(cache: ChannelCache, seq_bit: int) -> None:
|
||||
"""
|
||||
Set the expected sequential number (bit) of the next message to be received
|
||||
in the provided channel
|
||||
"""
|
||||
if __debug__:
|
||||
log.debug(__name__, "Set sync receive expected seq bit to %d", seq_bit)
|
||||
if seq_bit not in (0, 1):
|
||||
raise ThpError("Unexpected receive sync bit")
|
||||
|
||||
# set second bit to "seq_bit" value
|
||||
cache.sync &= 0xBF
|
||||
if seq_bit:
|
||||
cache.sync |= 0x40
|
||||
|
||||
|
||||
def _set_send_seq_bit(cache: ChannelCache, seq_bit: int) -> None:
|
||||
if seq_bit not in (0, 1):
|
||||
raise ThpError("Unexpected send seq bit")
|
||||
if __debug__:
|
||||
log.debug(__name__, "setting sync send seq bit to %d", seq_bit)
|
||||
# set third bit to "seq_bit" value
|
||||
cache.sync &= 0xDF
|
||||
if seq_bit:
|
||||
cache.sync |= 0x20
|
||||
|
||||
|
||||
def set_send_seq_bit_to_opposite(cache: ChannelCache) -> None:
|
||||
"""
|
||||
Set the sequential bit of the "next message to be send" to the opposite value,
|
||||
i.e. 1 -> 0 and 0 -> 1
|
||||
"""
|
||||
_set_send_seq_bit(cache=cache, seq_bit=1 - get_send_seq_bit(cache))
|
||||
@ -0,0 +1,308 @@
|
||||
import ustruct
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from storage.cache_common import (
|
||||
CHANNEL_HANDSHAKE_HASH,
|
||||
CHANNEL_KEY_RECEIVE,
|
||||
CHANNEL_KEY_SEND,
|
||||
CHANNEL_NONCE_RECEIVE,
|
||||
CHANNEL_NONCE_SEND,
|
||||
)
|
||||
from storage.cache_thp import TAG_LENGTH, ChannelCache, clear_sessions_with_channel_id
|
||||
from trezor import log, loop, protobuf, utils, workflow
|
||||
from trezor.wire.thp.transmission_loop import TransmissionLoop
|
||||
|
||||
from . import ChannelState, ThpDecryptionError, ThpError
|
||||
from . import alternating_bit_protocol as ABP
|
||||
from . import (
|
||||
control_byte,
|
||||
crypto,
|
||||
interface_manager,
|
||||
memory_manager,
|
||||
received_message_handler,
|
||||
session_manager,
|
||||
)
|
||||
from .checksum import CHECKSUM_LENGTH
|
||||
from .thp_messages import ENCRYPTED_TRANSPORT, PacketHeader
|
||||
from .writer import (
|
||||
CONT_HEADER_LENGTH,
|
||||
INIT_HEADER_LENGTH,
|
||||
write_payload_to_wire_and_add_checksum,
|
||||
)
|
||||
|
||||
if __debug__:
|
||||
from ubinascii import hexlify
|
||||
|
||||
from . import state_to_str
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from trezorio import WireInterface
|
||||
|
||||
from .pairing_context import PairingContext
|
||||
from .session_context import GenericSessionContext
|
||||
|
||||
|
||||
class Channel:
|
||||
def __init__(self, channel_cache: ChannelCache) -> None:
|
||||
if __debug__:
|
||||
log.debug(__name__, "channel initialization")
|
||||
self.iface: WireInterface = interface_manager.decode_iface(channel_cache.iface)
|
||||
self.channel_cache: ChannelCache = channel_cache
|
||||
self.is_cont_packet_expected: bool = False
|
||||
self.expected_payload_length: int = 0
|
||||
self.bytes_read: int = 0
|
||||
self.buffer: utils.BufferType
|
||||
self.channel_id: bytes = channel_cache.channel_id
|
||||
self.selected_pairing_methods = []
|
||||
self.sessions: dict[int, GenericSessionContext] = {}
|
||||
self.write_task_spawn: loop.spawn | None = None
|
||||
self.connection_context: PairingContext | None = None
|
||||
self.transmission_loop: TransmissionLoop | None = None
|
||||
self.handshake: crypto.Handshake | None = None
|
||||
|
||||
self._create_management_session()
|
||||
|
||||
def clear(self):
|
||||
clear_sessions_with_channel_id(self.channel_id)
|
||||
self.channel_cache.clear()
|
||||
|
||||
# ACCESS TO CHANNEL_DATA
|
||||
def get_channel_id_int(self) -> int:
|
||||
return int.from_bytes(self.channel_id, "big")
|
||||
|
||||
def get_channel_state(self) -> int:
|
||||
state = int.from_bytes(self.channel_cache.state, "big")
|
||||
if __debug__:
|
||||
log.debug(__name__, "get_channel_state: %s", state_to_str(state))
|
||||
return state
|
||||
|
||||
def set_channel_state(self, state: ChannelState) -> None:
|
||||
self.channel_cache.state = bytearray(state.to_bytes(1, "big"))
|
||||
if __debug__:
|
||||
log.debug(__name__, "set_channel_state: %s", state_to_str(state))
|
||||
|
||||
def set_buffer(self, buffer: utils.BufferType) -> None:
|
||||
self.buffer = buffer
|
||||
if __debug__:
|
||||
log.debug(__name__, "set_buffer: %s", type(self.buffer))
|
||||
|
||||
def _create_management_session(self) -> None:
|
||||
session = session_manager.create_new_management_session(self)
|
||||
self.sessions[session.session_id] = session
|
||||
loop.schedule(session.handle())
|
||||
|
||||
# CALLED BY THP_MAIN_LOOP
|
||||
|
||||
async def receive_packet(self, packet: utils.BufferType):
|
||||
if __debug__:
|
||||
log.debug(__name__, "receive_packet")
|
||||
|
||||
await self._handle_received_packet(packet)
|
||||
|
||||
if __debug__:
|
||||
log.debug(__name__, "self.buffer: %s", utils.get_bytes_as_str(self.buffer))
|
||||
|
||||
if self.expected_payload_length + INIT_HEADER_LENGTH == self.bytes_read:
|
||||
self._finish_message()
|
||||
await received_message_handler.handle_received_message(self, self.buffer)
|
||||
elif self.expected_payload_length + INIT_HEADER_LENGTH > self.bytes_read:
|
||||
self.is_cont_packet_expected = True
|
||||
else:
|
||||
raise ThpError(
|
||||
"Read more bytes than is the expected length of the message!"
|
||||
)
|
||||
|
||||
async def _handle_received_packet(self, packet: utils.BufferType) -> None:
|
||||
ctrl_byte = packet[0]
|
||||
if control_byte.is_continuation(ctrl_byte):
|
||||
await self._handle_cont_packet(packet)
|
||||
else:
|
||||
await self._handle_init_packet(packet)
|
||||
|
||||
async def _handle_init_packet(self, packet: utils.BufferType) -> None:
|
||||
if __debug__:
|
||||
log.debug(__name__, "handle_init_packet")
|
||||
# ctrl_byte, _, payload_length = ustruct.unpack(">BHH", packet) # TODO use this with single packet decryption
|
||||
_, _, payload_length = ustruct.unpack(">BHH", packet)
|
||||
self.expected_payload_length = payload_length
|
||||
packet_payload = memoryview(packet)[INIT_HEADER_LENGTH:]
|
||||
|
||||
# If the channel does not "own" the buffer lock, decrypt first packet
|
||||
# TODO do it only when needed!
|
||||
# TODO FIX: If "_decrypt_single_packet_payload" is implemented, it will (possibly) break "decrypt_buffer" and nonces incrementation.
|
||||
# On the other hand, without the single packet decryption, the "advanced" buffer selection cannot be implemented
|
||||
# in "memory_manager.select_buffer", because the session id is unknown (encrypted).
|
||||
|
||||
# if control_byte.is_encrypted_transport(ctrl_byte):
|
||||
# packet_payload = self._decrypt_single_packet_payload(packet_payload)
|
||||
|
||||
self.buffer = memory_manager.select_buffer(
|
||||
self.get_channel_state(),
|
||||
self.buffer,
|
||||
packet_payload,
|
||||
payload_length,
|
||||
)
|
||||
await self._buffer_packet_data(self.buffer, packet, 0)
|
||||
|
||||
if __debug__:
|
||||
log.debug(__name__, "handle_init_packet - payload len: %d", payload_length)
|
||||
log.debug(__name__, "handle_init_packet - buffer len: %d", len(self.buffer))
|
||||
|
||||
async def _handle_cont_packet(self, packet: utils.BufferType) -> None:
|
||||
if __debug__:
|
||||
log.debug(__name__, "handle_cont_packet")
|
||||
if not self.is_cont_packet_expected:
|
||||
raise ThpError("Continuation packet is not expected, ignoring")
|
||||
await self._buffer_packet_data(self.buffer, packet, CONT_HEADER_LENGTH)
|
||||
|
||||
def _decrypt_single_packet_payload(
|
||||
self, payload: utils.BufferType
|
||||
) -> utils.BufferType:
|
||||
# crypto.decrypt(b"\x00", b"\x00", payload_buffer, INIT_DATA_OFFSET, len(payload))
|
||||
return payload
|
||||
|
||||
def decrypt_buffer(
|
||||
self, message_length: int, offset: int = INIT_HEADER_LENGTH
|
||||
) -> None:
|
||||
noise_buffer = memoryview(self.buffer)[
|
||||
offset : message_length - CHECKSUM_LENGTH - TAG_LENGTH
|
||||
]
|
||||
tag = self.buffer[
|
||||
message_length
|
||||
- CHECKSUM_LENGTH
|
||||
- TAG_LENGTH : message_length
|
||||
- CHECKSUM_LENGTH
|
||||
]
|
||||
if utils.DISABLE_ENCRYPTION:
|
||||
is_tag_valid = tag == crypto.DUMMY_TAG
|
||||
else:
|
||||
key_receive = self.channel_cache.get(CHANNEL_KEY_RECEIVE)
|
||||
nonce_receive = self.channel_cache.get_int(CHANNEL_NONCE_RECEIVE)
|
||||
auth_data = self.channel_cache.get(CHANNEL_HANDSHAKE_HASH)
|
||||
|
||||
assert key_receive is not None
|
||||
assert nonce_receive is not None
|
||||
assert auth_data is not None
|
||||
|
||||
print("Buffer before decryption:", hexlify(noise_buffer))
|
||||
is_tag_valid = crypto.dec(
|
||||
noise_buffer, tag, key_receive, nonce_receive, auth_data
|
||||
)
|
||||
print("Buffer after decryption:", hexlify(noise_buffer))
|
||||
|
||||
self.channel_cache.set_int(CHANNEL_NONCE_RECEIVE, nonce_receive + 1)
|
||||
|
||||
if __debug__:
|
||||
log.debug(__name__, "Is decrypted tag valid? %s", str(is_tag_valid))
|
||||
log.debug(__name__, "Received tag: %s", (hexlify(tag).decode()))
|
||||
log.debug(__name__, "New nonce_receive: %i", nonce_receive + 1)
|
||||
|
||||
if not is_tag_valid:
|
||||
raise ThpDecryptionError()
|
||||
|
||||
def _encrypt(self, buffer: utils.BufferType, noise_payload_len: int) -> None:
|
||||
if __debug__:
|
||||
log.debug(__name__, "encrypt")
|
||||
assert len(buffer) >= noise_payload_len + TAG_LENGTH + CHECKSUM_LENGTH
|
||||
|
||||
noise_buffer = memoryview(buffer)[0:noise_payload_len]
|
||||
|
||||
if utils.DISABLE_ENCRYPTION:
|
||||
tag = crypto.DUMMY_TAG
|
||||
else:
|
||||
key_send = self.channel_cache.get(CHANNEL_KEY_SEND)
|
||||
nonce_send = self.channel_cache.get_int(CHANNEL_NONCE_SEND)
|
||||
auth_data = self.channel_cache.get(CHANNEL_HANDSHAKE_HASH)
|
||||
|
||||
assert key_send is not None
|
||||
assert nonce_send is not None
|
||||
assert auth_data is not None
|
||||
|
||||
tag = crypto.enc(noise_buffer, key_send, nonce_send, auth_data)
|
||||
|
||||
self.channel_cache.set_int(CHANNEL_NONCE_SEND, nonce_send + 1)
|
||||
if __debug__:
|
||||
log.debug(__name__, "New nonce_send: %i", nonce_send + 1)
|
||||
|
||||
buffer[noise_payload_len : noise_payload_len + TAG_LENGTH] = tag
|
||||
|
||||
async def _buffer_packet_data(
|
||||
self, payload_buffer: utils.BufferType, packet: utils.BufferType, offset: int
|
||||
):
|
||||
self.bytes_read += utils.memcpy(payload_buffer, self.bytes_read, packet, offset)
|
||||
|
||||
def _finish_message(self):
|
||||
self.bytes_read = 0
|
||||
self.expected_payload_length = 0
|
||||
self.is_cont_packet_expected = False
|
||||
|
||||
# CALLED BY WORKFLOW / SESSION CONTEXT
|
||||
|
||||
async def write(self, msg: protobuf.MessageType, session_id: int = 0) -> None:
|
||||
if __debug__:
|
||||
log.debug(__name__, "write message: %s", msg.MESSAGE_NAME)
|
||||
self.buffer = memory_manager.get_write_buffer(self.buffer, msg)
|
||||
noise_payload_len = memory_manager.encode_into_buffer(
|
||||
self.buffer, msg, session_id
|
||||
)
|
||||
await self.write_and_encrypt(self.buffer[:noise_payload_len])
|
||||
|
||||
async def write_error(self, err_type: int):
|
||||
msg_data = err_type.to_bytes(1, "big")
|
||||
length = len(msg_data) + CHECKSUM_LENGTH
|
||||
header = PacketHeader.get_error_header(self.get_channel_id_int(), length)
|
||||
await write_payload_to_wire_and_add_checksum(self.iface, header, msg_data)
|
||||
|
||||
async def write_and_encrypt(self, payload: bytes) -> None:
|
||||
payload_length = len(payload)
|
||||
|
||||
if not isinstance(self.buffer, bytearray):
|
||||
self.buffer = bytearray(self.buffer)
|
||||
self._encrypt(self.buffer, payload_length)
|
||||
payload_length = payload_length + TAG_LENGTH
|
||||
|
||||
if self.write_task_spawn is not None:
|
||||
self.write_task_spawn.close() # UPS TODO might break something
|
||||
print("\nCLOSED\n")
|
||||
self._prepare_write()
|
||||
self.write_task_spawn = loop.spawn(
|
||||
self._write_encrypted_payload_loop(
|
||||
ENCRYPTED_TRANSPORT, memoryview(self.buffer[:payload_length])
|
||||
)
|
||||
)
|
||||
|
||||
async def write_handshake_message(self, ctrl_byte: int, payload: bytes) -> None:
|
||||
self._prepare_write()
|
||||
self.write_task_spawn = loop.spawn(
|
||||
self._write_encrypted_payload_loop(ctrl_byte, payload)
|
||||
)
|
||||
|
||||
def _prepare_write(self) -> None:
|
||||
# TODO add condition that disallows to write when can_send_message is false
|
||||
ABP.set_sending_allowed(self.channel_cache, False)
|
||||
|
||||
async def _write_encrypted_payload_loop(
|
||||
self, ctrl_byte: int, payload: bytes
|
||||
) -> None:
|
||||
if __debug__:
|
||||
log.debug(__name__, "write_encrypted_payload_loop")
|
||||
payload_len = len(payload) + CHECKSUM_LENGTH
|
||||
sync_bit = ABP.get_send_seq_bit(self.channel_cache)
|
||||
ctrl_byte = control_byte.add_seq_bit_to_ctrl_byte(ctrl_byte, sync_bit)
|
||||
header = PacketHeader(ctrl_byte, self.get_channel_id_int(), payload_len)
|
||||
self.transmission_loop = TransmissionLoop(self, header, payload)
|
||||
await self.transmission_loop.start()
|
||||
|
||||
ABP.set_send_seq_bit_to_opposite(self.channel_cache)
|
||||
|
||||
# Let the main loop be restarted and clear loop, if there is no other
|
||||
# workflow and the state is ENCRYPTED_TRANSPORT
|
||||
if self._can_clear_loop():
|
||||
if __debug__:
|
||||
log.debug(__name__, "clearing loop from channel")
|
||||
loop.clear()
|
||||
|
||||
def _can_clear_loop(self) -> bool:
|
||||
return (
|
||||
not workflow.tasks
|
||||
) and self.get_channel_state() is ChannelState.ENCRYPTED_TRANSPORT
|
||||
@ -0,0 +1,36 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from storage import cache_thp
|
||||
from trezor import utils
|
||||
|
||||
from . import ChannelState, interface_manager
|
||||
from .channel import Channel
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from trezorio import WireInterface
|
||||
|
||||
|
||||
def create_new_channel(iface: WireInterface, buffer: utils.BufferType) -> Channel:
|
||||
"""
|
||||
Creates a new channel for the interface `iface` with the buffer `buffer`.
|
||||
"""
|
||||
channel_cache = cache_thp.get_new_unauthenticated_channel(
|
||||
interface_manager.encode_iface(iface)
|
||||
)
|
||||
r = Channel(channel_cache)
|
||||
r.set_buffer(buffer)
|
||||
r.set_channel_state(ChannelState.TH1)
|
||||
return r
|
||||
|
||||
|
||||
def load_cached_channels(buffer: utils.BufferType) -> dict[int, Channel]:
|
||||
"""
|
||||
Returns all allocated channels from cache.
|
||||
"""
|
||||
channels: dict[int, Channel] = {}
|
||||
cached_channels = cache_thp.get_all_allocated_channels()
|
||||
for c in cached_channels:
|
||||
channels[int.from_bytes(c.channel_id, "big")] = Channel(c)
|
||||
for c in channels.values():
|
||||
c.set_buffer(buffer)
|
||||
return channels
|
||||
@ -0,0 +1,22 @@
|
||||
from micropython import const
|
||||
|
||||
from trezor import utils
|
||||
from trezor.crypto import crc
|
||||
|
||||
CHECKSUM_LENGTH = const(4)
|
||||
|
||||
|
||||
def compute(data: bytes | utils.BufferType) -> bytes:
|
||||
"""
|
||||
Returns a CRC-32 checksum of the provided `data`.
|
||||
"""
|
||||
return crc.crc32(data).to_bytes(CHECKSUM_LENGTH, "big")
|
||||
|
||||
|
||||
def is_valid(checksum: bytes | utils.BufferType, data: bytes) -> bool:
|
||||
"""
|
||||
Checks whether the CRC-32 checksum of the `data` is the same
|
||||
as the checksum provided in `checksum`.
|
||||
"""
|
||||
data_checksum = compute(data)
|
||||
return checksum == data_checksum
|
||||
@ -0,0 +1,47 @@
|
||||
from trezor.wire.thp import ThpError
|
||||
from trezor.wire.thp.thp_messages import (
|
||||
ACK_MASK,
|
||||
ACK_MESSAGE,
|
||||
CONTINUATION_PACKET,
|
||||
CONTINUATION_PACKET_MASK,
|
||||
DATA_MASK,
|
||||
ENCRYPTED_TRANSPORT,
|
||||
HANDSHAKE_COMP_REQ,
|
||||
HANDSHAKE_INIT_REQ,
|
||||
)
|
||||
|
||||
|
||||
def add_seq_bit_to_ctrl_byte(ctrl_byte: int, seq_bit: int) -> int:
|
||||
if seq_bit == 0:
|
||||
return ctrl_byte & 0xEF
|
||||
if seq_bit == 1:
|
||||
return ctrl_byte | 0x10
|
||||
raise ThpError("Unexpected sequence bit")
|
||||
|
||||
|
||||
def add_ack_bit_to_ctrl_byte(ctrl_byte: int, ack_bit: int) -> int:
|
||||
if ack_bit == 0:
|
||||
return ctrl_byte & 0xF7
|
||||
if ack_bit == 1:
|
||||
return ctrl_byte | 0x08
|
||||
raise ThpError("Unexpected acknowledgement bit")
|
||||
|
||||
|
||||
def is_ack(ctrl_byte: int) -> bool:
|
||||
return ctrl_byte & ACK_MASK == ACK_MESSAGE
|
||||
|
||||
|
||||
def is_continuation(ctrl_byte: int) -> bool:
|
||||
return ctrl_byte & CONTINUATION_PACKET_MASK == CONTINUATION_PACKET
|
||||
|
||||
|
||||
def is_encrypted_transport(ctrl_byte: int) -> bool:
|
||||
return ctrl_byte & DATA_MASK == ENCRYPTED_TRANSPORT
|
||||
|
||||
|
||||
def is_handshake_init_req(ctrl_byte: int) -> bool:
|
||||
return ctrl_byte & DATA_MASK == HANDSHAKE_INIT_REQ
|
||||
|
||||
|
||||
def is_handshake_comp_req(ctrl_byte: int) -> bool:
|
||||
return ctrl_byte & DATA_MASK == HANDSHAKE_COMP_REQ
|
||||
@ -0,0 +1,34 @@
|
||||
from trezor.crypto import elligator2, random
|
||||
from trezor.crypto.curve import curve25519
|
||||
from trezor.crypto.hashlib import sha512
|
||||
|
||||
_PREFIX = b"\x08\x43\x50\x61\x63\x65\x32\x35\x35\x06"
|
||||
_PADDING = b"\x50\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x20"
|
||||
|
||||
|
||||
class Cpace:
|
||||
"""
|
||||
CPace, a balanced composable PAKE: https://datatracker.ietf.org/doc/draft-irtf-cfrg-cpace/
|
||||
"""
|
||||
|
||||
def __init__(self, cpace_host_public_key: bytes) -> None:
|
||||
self.host_public_key: bytes = cpace_host_public_key
|
||||
self.trezor_private_key: bytes
|
||||
self.trezor_public_key: bytes
|
||||
self.shared_secret: bytes
|
||||
|
||||
def generate_keys_and_secret(self, code_code_entry: bytes) -> None:
|
||||
"""
|
||||
Generate ephemeral key pair and a shared secret using Elligator2 with X25519.
|
||||
"""
|
||||
pregenerator = sha512(_PREFIX + code_code_entry + _PADDING).digest()[
|
||||
:32
|
||||
] # TODO add handshake hash
|
||||
generator = elligator2.map_to_curve25519(pregenerator)
|
||||
self.trezor_private_key = random.bytes(32)
|
||||
if __debug__:
|
||||
self.trezor_private_key = b"\xf8\xb9\xa1\x3a\xe1\x30\xb3\xe1\x5b\x8e\xd5\x80\x85\x4f\xfc\xaf\x63\x4d\x6b\x0a\x10\xfd\xe7\xba\xde\xfd\xc3\xd1\x8d\x1a\x83\xf5"
|
||||
self.trezor_public_key = curve25519.multiply(self.trezor_private_key, generator)
|
||||
self.shared_secret = curve25519.multiply(
|
||||
self.trezor_private_key, self.host_public_key
|
||||
)
|
||||
@ -0,0 +1,206 @@
|
||||
from micropython import const
|
||||
from trezorcrypto import aesgcm, bip32, curve25519, hmac
|
||||
|
||||
from storage import device
|
||||
from trezor import utils
|
||||
from trezor.crypto.hashlib import sha256
|
||||
from trezor.wire.thp import ThpDecryptionError
|
||||
|
||||
# The HARDENED flag is taken from apps.common.paths
|
||||
# It is not imported to save on resources
|
||||
HARDENED = const(0x8000_0000)
|
||||
PUBKEY_LENGTH = const(32)
|
||||
if utils.DISABLE_ENCRYPTION:
|
||||
DUMMY_TAG = b"\xA0\xA1\xA2\xA3\xA4\xA5\xA6\xA7\xA8\xA9\xB0\xB1\xB2\xB3\xB4\xB5"
|
||||
|
||||
if __debug__:
|
||||
from ubinascii import hexlify
|
||||
|
||||
|
||||
def enc(buffer: utils.BufferType, key: bytes, nonce: int, auth_data: bytes) -> bytes:
|
||||
"""
|
||||
Encrypts the provided `buffer` with AES-GCM (in place).
|
||||
Returns a 16-byte long encryption tag.
|
||||
"""
|
||||
print("ENCRYPT-------------> used key: ", hexlify(key))
|
||||
print("ENCRYPT-------------> used nonce:", nonce)
|
||||
iv = _get_iv_from_nonce(nonce)
|
||||
aes_ctx = aesgcm(key, iv)
|
||||
aes_ctx.auth(auth_data)
|
||||
aes_ctx.encrypt_in_place(buffer)
|
||||
return aes_ctx.finish()
|
||||
|
||||
|
||||
def dec(
|
||||
buffer: utils.BufferType, tag: bytes, key: bytes, nonce: int, auth_data: bytes
|
||||
) -> bool:
|
||||
"""
|
||||
Decrypts the provided buffer (in place). Returns `True` if the provided authentication `tag` is the same as
|
||||
the tag computed in decryption, otherwise it returns `False`.
|
||||
"""
|
||||
iv = _get_iv_from_nonce(nonce)
|
||||
print("DECRYPT-------------> used key: ", hexlify(key))
|
||||
print("DECRYPT-------------> used nonce:", nonce)
|
||||
aes_ctx = aesgcm(key, iv)
|
||||
aes_ctx.auth(auth_data)
|
||||
aes_ctx.decrypt_in_place(buffer)
|
||||
computed_tag = aes_ctx.finish()
|
||||
return computed_tag == tag
|
||||
|
||||
|
||||
class BusyDecoder:
|
||||
def __init__(self, key: bytes, nonce: int, auth_data: bytes) -> None:
|
||||
iv = _get_iv_from_nonce(nonce)
|
||||
self.aes_ctx = aesgcm(key, iv)
|
||||
self.aes_ctx.auth(auth_data)
|
||||
|
||||
def decrypt_part(self, part: utils.BufferType) -> None:
|
||||
self.aes_ctx.decrypt_in_place(part)
|
||||
|
||||
def finish_and_check_tag(self, tag: bytes) -> bool:
|
||||
computed_tag = self.aes_ctx.finish()
|
||||
return computed_tag == tag
|
||||
|
||||
|
||||
PROTOCOL_NAME = bytes("Noise_XX_25519_AESGCM_SHA256", "ascii")
|
||||
IV_1 = b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
|
||||
IV_2 = b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01"
|
||||
|
||||
|
||||
class Handshake:
|
||||
"""
|
||||
`Handshake` holds (temporary) values and keys that are used during the creation of an encrypted channel.
|
||||
The following values should be saved for future use before disposing of this object:
|
||||
- `h` (handshake hash, can be used to bind other values to the channel)
|
||||
- `key_receive` (key for decrypting incoming communication)
|
||||
- `key_send` (key for encrypting outgoing communication)
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.trezor_ephemeral_privkey: bytes
|
||||
self.ck: bytes
|
||||
self.k: bytes
|
||||
self.h: bytes
|
||||
self.key_receive: bytes
|
||||
self.key_send: bytes
|
||||
|
||||
def handle_th1_crypto(
|
||||
self,
|
||||
device_properties: bytes,
|
||||
host_ephemeral_pubkey: bytes,
|
||||
) -> tuple[bytes, bytes, bytes]:
|
||||
|
||||
trezor_static_privkey, trezor_static_pubkey = _derive_static_key_pair()
|
||||
self.trezor_ephemeral_privkey = curve25519.generate_secret()
|
||||
trezor_ephemeral_pubkey = curve25519.publickey(self.trezor_ephemeral_privkey)
|
||||
self.h = _hash_of_two(PROTOCOL_NAME, device_properties)
|
||||
self.h = _hash_of_two(self.h, host_ephemeral_pubkey)
|
||||
self.h = _hash_of_two(self.h, trezor_ephemeral_pubkey)
|
||||
point = curve25519.multiply(
|
||||
self.trezor_ephemeral_privkey, host_ephemeral_pubkey
|
||||
)
|
||||
self.ck, self.k = _hkdf(PROTOCOL_NAME, point)
|
||||
mask = _hash_of_two(trezor_static_pubkey, trezor_ephemeral_pubkey)
|
||||
trezor_masked_static_pubkey = curve25519.multiply(mask, trezor_static_pubkey)
|
||||
aes_ctx = aesgcm(self.k, IV_1)
|
||||
encrypted_trezor_static_pubkey = aes_ctx.encrypt(trezor_masked_static_pubkey)
|
||||
print("TH1-ENCRYPT---------> used key: ", hexlify(self.k))
|
||||
print("TH1-ENCRYPT---------> used nonce:", 0)
|
||||
aes_ctx.auth(self.h)
|
||||
tag_to_encrypted_key = aes_ctx.finish()
|
||||
encrypted_trezor_static_pubkey = (
|
||||
encrypted_trezor_static_pubkey + tag_to_encrypted_key
|
||||
)
|
||||
self.h = _hash_of_two(self.h, encrypted_trezor_static_pubkey)
|
||||
point = curve25519.multiply(trezor_static_privkey, host_ephemeral_pubkey)
|
||||
self.ck, self.k = _hkdf(self.ck, curve25519.multiply(mask, point))
|
||||
aes_ctx = aesgcm(self.k, IV_1)
|
||||
aes_ctx.auth(self.h)
|
||||
tag = aes_ctx.finish()
|
||||
self.h = _hash_of_two(self.h, tag)
|
||||
return (trezor_ephemeral_pubkey, encrypted_trezor_static_pubkey, tag)
|
||||
|
||||
def handle_th2_crypto(
|
||||
self,
|
||||
encrypted_host_static_pubkey: utils.BufferType,
|
||||
encrypted_payload: utils.BufferType,
|
||||
):
|
||||
|
||||
aes_ctx = aesgcm(self.k, IV_2)
|
||||
|
||||
# The new value of hash `h` MUST be computed before the `encrypted_host_static_pubkey` is decrypted.
|
||||
# However, decryption of `encrypted_host_static_pubkey` MUST use the previous value of `h` for
|
||||
# authentication of the gcm tag.
|
||||
aes_ctx.auth(self.h) # Authenticate with the previous value of `h`
|
||||
self.h = _hash_of_two(self.h, encrypted_host_static_pubkey) # Compute new value
|
||||
aes_ctx.decrypt_in_place(
|
||||
memoryview(encrypted_host_static_pubkey)[:PUBKEY_LENGTH]
|
||||
)
|
||||
print("TH2-DECRYPT---------> used key: ", hexlify(self.k))
|
||||
print("TH2-DECRYPT---------> used nonce:", 1)
|
||||
host_static_pubkey = memoryview(encrypted_host_static_pubkey)[:PUBKEY_LENGTH]
|
||||
tag = aes_ctx.finish()
|
||||
if tag != encrypted_host_static_pubkey[-16:]:
|
||||
raise ThpDecryptionError()
|
||||
|
||||
self.ck, self.k = _hkdf(
|
||||
self.ck,
|
||||
curve25519.multiply(self.trezor_ephemeral_privkey, host_static_pubkey),
|
||||
)
|
||||
aes_ctx = aesgcm(self.k, IV_1)
|
||||
aes_ctx.auth(self.h)
|
||||
aes_ctx.decrypt_in_place(memoryview(encrypted_payload)[:-16])
|
||||
print("TH2-DECRYPT---------> used key: ", hexlify(self.k))
|
||||
print("TH2-DECRYPT---------> used nonce:", 0)
|
||||
tag = aes_ctx.finish()
|
||||
if tag != encrypted_payload[-16:]:
|
||||
raise ThpDecryptionError()
|
||||
|
||||
self.h = _hash_of_two(self.h, memoryview(encrypted_payload)[:-16])
|
||||
self.key_receive, self.key_send = _hkdf(self.ck, b"")
|
||||
print("TREZOR_KEY_RECEIVE:", hexlify(self.key_receive))
|
||||
print("TREZOR_KEY_SEND: ", hexlify(self.key_send))
|
||||
|
||||
def get_handshake_completion_response(self, trezor_state: bytes) -> bytes:
|
||||
aes_ctx = aesgcm(self.key_send, IV_1)
|
||||
encrypted_trezor_state = aes_ctx.encrypt(trezor_state)
|
||||
tag = aes_ctx.finish()
|
||||
return encrypted_trezor_state + tag
|
||||
|
||||
|
||||
def _derive_static_key_pair() -> tuple[bytes, bytes]:
|
||||
node_int = HARDENED | int.from_bytes(b"\x00THP", "big")
|
||||
node = bip32.from_seed(device.get_device_secret(), "curve25519")
|
||||
node.derive(node_int)
|
||||
|
||||
trezor_static_privkey = node.private_key()
|
||||
trezor_static_pubkey = node.public_key()[1:33]
|
||||
# Note: the first byte (\x01) of the public key is removed, as it
|
||||
# only indicates the type of the elliptic curve used
|
||||
|
||||
return trezor_static_privkey, trezor_static_pubkey
|
||||
|
||||
|
||||
def get_trezor_static_pubkey() -> bytes:
|
||||
_, pubkey = _derive_static_key_pair()
|
||||
return pubkey
|
||||
|
||||
|
||||
def _hkdf(chaining_key, input: bytes):
|
||||
temp_key = hmac(hmac.SHA256, chaining_key, input).digest()
|
||||
output_1 = hmac(hmac.SHA256, temp_key, b"\x01").digest()
|
||||
ctx_output_2 = hmac(hmac.SHA256, temp_key, output_1)
|
||||
ctx_output_2.update(b"\x02")
|
||||
output_2 = ctx_output_2.digest()
|
||||
return (output_1, output_2)
|
||||
|
||||
|
||||
def _hash_of_two(part_1: bytes, part_2: bytes) -> bytes:
|
||||
ctx = sha256(part_1)
|
||||
ctx.update(part_2)
|
||||
return ctx.digest()
|
||||
|
||||
|
||||
def _get_iv_from_nonce(nonce: int) -> bytes:
|
||||
utils.ensure(nonce <= 0xFFFFFFFFFFFFFFFF, "Nonce overflow, terminate the channel")
|
||||
return bytes(4) + nonce.to_bytes(8, "big")
|
||||
@ -0,0 +1,33 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from trezor import protobuf
|
||||
from trezor.enums import MessageType
|
||||
from trezor.wire.errors import UnexpectedMessage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any, Callable, Coroutine
|
||||
|
||||
from trezor.messages import Features, GetFeatures
|
||||
|
||||
|
||||
def find_management_session_message_handler(
|
||||
msg_type: int,
|
||||
) -> Callable[[Any], Coroutine[Any, Any, protobuf.MessageType]]:
|
||||
if msg_type is MessageType.ThpCreateNewSession:
|
||||
from apps.thp.create_session import create_new_session
|
||||
|
||||
return create_new_session
|
||||
if msg_type is MessageType.GetFeatures:
|
||||
return handle_GetFeatures
|
||||
if __debug__:
|
||||
if msg_type is MessageType.LoadDevice:
|
||||
from apps.debug.load_device import load_device
|
||||
|
||||
return load_device
|
||||
raise UnexpectedMessage("There is no handler available for this message")
|
||||
|
||||
|
||||
async def handle_GetFeatures(msg: GetFeatures) -> Features:
|
||||
from apps.base import get_features
|
||||
|
||||
return get_features()
|
||||
@ -0,0 +1,32 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import usb
|
||||
|
||||
_MOCK_INTERFACE_HID = b"\x00"
|
||||
_WIRE_INTERFACE_USB = b"\x01"
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from trezorio import WireInterface
|
||||
|
||||
|
||||
def decode_iface(cached_iface: bytes) -> WireInterface:
|
||||
"""Decode the cached wire interface."""
|
||||
if cached_iface == _WIRE_INTERFACE_USB:
|
||||
iface = usb.iface_wire
|
||||
if iface is None:
|
||||
raise RuntimeError("There is no valid USB WireInterface")
|
||||
return iface
|
||||
if __debug__ and cached_iface == _MOCK_INTERFACE_HID:
|
||||
raise NotImplementedError("Should return MockHID WireInterface")
|
||||
# TODO implement bluetooth interface
|
||||
raise Exception("Unknown WireInterface")
|
||||
|
||||
|
||||
def encode_iface(iface: WireInterface) -> bytes:
|
||||
"""Encode wire interface into bytes."""
|
||||
if iface is usb.iface_wire:
|
||||
return _WIRE_INTERFACE_USB
|
||||
# TODO implement bluetooth interface
|
||||
if __debug__:
|
||||
return _MOCK_INTERFACE_HID
|
||||
raise Exception("Unknown WireInterface")
|
||||
@ -0,0 +1,124 @@
|
||||
from storage.cache_thp import SESSION_ID_LENGTH, TAG_LENGTH
|
||||
from trezor import log, protobuf, utils
|
||||
|
||||
from . import ChannelState, ThpError
|
||||
from .checksum import CHECKSUM_LENGTH
|
||||
from .writer import (
|
||||
INIT_HEADER_LENGTH,
|
||||
MAX_PAYLOAD_LEN,
|
||||
MESSAGE_TYPE_LENGTH,
|
||||
PACKET_LENGTH,
|
||||
)
|
||||
|
||||
|
||||
def select_buffer(
|
||||
channel_state: int,
|
||||
channel_buffer: utils.BufferType,
|
||||
packet_payload: utils.BufferType,
|
||||
payload_length: int,
|
||||
) -> utils.BufferType:
|
||||
|
||||
if channel_state is ChannelState.ENCRYPTED_TRANSPORT:
|
||||
session_id = packet_payload[0]
|
||||
if session_id == 0:
|
||||
pass
|
||||
# TODO use small buffer
|
||||
else:
|
||||
pass
|
||||
# TODO use big buffer but only if the channel owns the buffer lock.
|
||||
# Otherwise send BUSY message and return
|
||||
else:
|
||||
pass
|
||||
# TODO use small buffer
|
||||
try:
|
||||
# TODO for now, we create a new big buffer every time. It should be changed
|
||||
buffer: utils.BufferType = _get_buffer_for_message(
|
||||
payload_length, channel_buffer
|
||||
)
|
||||
return buffer
|
||||
except Exception as e:
|
||||
if __debug__:
|
||||
log.exception(__name__, e)
|
||||
raise Exception("Failed to create a buffer for channel") # TODO handle better
|
||||
|
||||
|
||||
def get_write_buffer(
|
||||
buffer: utils.BufferType, msg: protobuf.MessageType
|
||||
) -> utils.BufferType:
|
||||
msg_size = protobuf.encoded_length(msg)
|
||||
payload_size = SESSION_ID_LENGTH + MESSAGE_TYPE_LENGTH + msg_size
|
||||
required_min_size = payload_size + CHECKSUM_LENGTH + TAG_LENGTH
|
||||
|
||||
if required_min_size > len(buffer):
|
||||
# message is too big, we need to allocate a new buffer
|
||||
return bytearray(required_min_size)
|
||||
return buffer
|
||||
|
||||
|
||||
def encode_into_buffer(
|
||||
buffer: utils.BufferType, msg: protobuf.MessageType, session_id: int
|
||||
) -> int:
|
||||
|
||||
# cannot write message without wire type
|
||||
assert msg.MESSAGE_WIRE_TYPE is not None
|
||||
|
||||
msg_size = protobuf.encoded_length(msg)
|
||||
payload_size = SESSION_ID_LENGTH + MESSAGE_TYPE_LENGTH + msg_size
|
||||
|
||||
_encode_session_into_buffer(memoryview(buffer), session_id)
|
||||
_encode_message_type_into_buffer(
|
||||
memoryview(buffer), msg.MESSAGE_WIRE_TYPE, SESSION_ID_LENGTH
|
||||
)
|
||||
_encode_message_into_buffer(
|
||||
memoryview(buffer), msg, SESSION_ID_LENGTH + MESSAGE_TYPE_LENGTH
|
||||
)
|
||||
|
||||
return payload_size
|
||||
|
||||
|
||||
def _encode_session_into_buffer(
|
||||
buffer: memoryview, session_id: int, buffer_offset: int = 0
|
||||
) -> None:
|
||||
session_id_bytes = int.to_bytes(session_id, SESSION_ID_LENGTH, "big")
|
||||
utils.memcpy(buffer, buffer_offset, session_id_bytes, 0)
|
||||
|
||||
|
||||
def _encode_message_type_into_buffer(
|
||||
buffer: memoryview, message_type: int, offset: int = 0
|
||||
) -> None:
|
||||
msg_type_bytes = int.to_bytes(message_type, MESSAGE_TYPE_LENGTH, "big")
|
||||
utils.memcpy(buffer, offset, msg_type_bytes, 0)
|
||||
|
||||
|
||||
def _encode_message_into_buffer(
|
||||
buffer: memoryview, message: protobuf.MessageType, buffer_offset: int = 0
|
||||
) -> None:
|
||||
protobuf.encode(memoryview(buffer[buffer_offset:]), message)
|
||||
|
||||
|
||||
def _get_buffer_for_message(
|
||||
payload_length: int, existing_buffer: utils.BufferType, max_length=MAX_PAYLOAD_LEN
|
||||
) -> utils.BufferType:
|
||||
length = payload_length + INIT_HEADER_LENGTH
|
||||
if __debug__:
|
||||
log.debug(
|
||||
__name__,
|
||||
"get_buffer_for_message - length: %d, %s %s",
|
||||
length,
|
||||
"existing buffer type:",
|
||||
type(existing_buffer),
|
||||
)
|
||||
if length > max_length:
|
||||
raise ThpError("Message too large")
|
||||
|
||||
if length > len(existing_buffer):
|
||||
# allocate a new buffer to fit the message
|
||||
try:
|
||||
payload: utils.BufferType = bytearray(length)
|
||||
except MemoryError:
|
||||
payload = bytearray(PACKET_LENGTH)
|
||||
raise ThpError("Message too large")
|
||||
return payload
|
||||
|
||||
# reuse a part of the supplied buffer
|
||||
return memoryview(existing_buffer)[:length]
|
||||
@ -0,0 +1,250 @@
|
||||
from typing import TYPE_CHECKING
|
||||
from ubinascii import hexlify
|
||||
|
||||
import trezorui2
|
||||
from trezor import loop, protobuf, workflow
|
||||
from trezor.ui.layouts.tt import RustLayout
|
||||
from trezor.wire import context, message_handler, protocol_common
|
||||
from trezor.wire.context import UnexpectedMessageException
|
||||
from trezor.wire.errors import ActionCancelled
|
||||
from trezor.wire.protocol_common import Context, Message
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Container
|
||||
|
||||
from .channel import Channel
|
||||
from .cpace import Cpace
|
||||
|
||||
pass
|
||||
|
||||
if __debug__:
|
||||
from trezor import log
|
||||
|
||||
|
||||
class PairingDisplayData:
|
||||
def __init__(self) -> None:
|
||||
self.display_code_entry: bool = False
|
||||
self.display_qr_code: bool = False
|
||||
self.display_nfc_unidirectional: bool = False
|
||||
self.code_code_entry: int | None = None
|
||||
self.code_qr_code: bytes | None = None
|
||||
self.code_nfc_unidirectional: bytes | None = None
|
||||
|
||||
def get_display_layout(self) -> RustLayout:
|
||||
return RustLayout(
|
||||
trezorui2.show_address_details( # noqa
|
||||
qr_title="Scan QR code to pair",
|
||||
address=self._get_code_qr_code_str(),
|
||||
case_sensitive=True,
|
||||
details_title="",
|
||||
account="Code to rewrite:\n" + self._get_code_code_entry_str(),
|
||||
path="",
|
||||
xpubs=[],
|
||||
)
|
||||
)
|
||||
|
||||
def _get_code_code_entry_str(self) -> str:
|
||||
if self.display_code_entry and self.code_code_entry is not None:
|
||||
code_str = str(self.code_code_entry)
|
||||
print("code_code_entry:", code_str)
|
||||
|
||||
return code_str[:3] + " " + code_str[3:]
|
||||
return "NOT ALLOWED"
|
||||
|
||||
def _get_code_qr_code_str(self) -> str:
|
||||
if self.display_qr_code and self.code_qr_code is not None:
|
||||
code_str = (hexlify(self.code_qr_code)).decode("utf-8")
|
||||
print("code_qr_code_hexlified:", code_str)
|
||||
return code_str
|
||||
return "QR CODE IS NOT SUPPOSED TO BE DISPLAYED!!!!"
|
||||
|
||||
|
||||
class PairingContext(Context):
|
||||
|
||||
def __init__(self, channel_ctx: Channel) -> None:
|
||||
super().__init__(channel_ctx.iface, channel_ctx.channel_id)
|
||||
self.channel_ctx: Channel = channel_ctx
|
||||
self.incoming_message = loop.chan()
|
||||
self.secret: bytes = (
|
||||
b"\xBA\xDA\x55\xBA\xDA\x55\xBA\xDA\x55\xBA\xDA\x55\xDE\xAD\xBE\xEF" # TODO generate randomly
|
||||
)
|
||||
|
||||
self.display_data: PairingDisplayData = PairingDisplayData()
|
||||
self.cpace: Cpace
|
||||
self.host_name: str
|
||||
|
||||
async def handle(self, is_debug_session: bool = False) -> None:
|
||||
if __debug__:
|
||||
log.debug(__name__, "handle - start")
|
||||
if is_debug_session:
|
||||
import apps.debug
|
||||
|
||||
apps.debug.DEBUG_CONTEXT = self
|
||||
|
||||
take = self.incoming_message.take()
|
||||
next_message: Message | None = None
|
||||
|
||||
while True:
|
||||
try:
|
||||
if next_message is None:
|
||||
# If the previous run did not keep an unprocessed message for us,
|
||||
# wait for a new one.
|
||||
try:
|
||||
message: Message = await take
|
||||
except protocol_common.WireError as e:
|
||||
if __debug__:
|
||||
log.exception(__name__, e)
|
||||
await self.write(message_handler.failure(e))
|
||||
continue
|
||||
else:
|
||||
# Process the message from previous run.
|
||||
message = next_message
|
||||
next_message = None
|
||||
|
||||
try:
|
||||
next_message = await handle_pairing_request_message(
|
||||
self, message, use_workflow=not is_debug_session
|
||||
)
|
||||
except Exception as exc:
|
||||
# Log and ignore. The session handler can only exit explicitly in the
|
||||
# following finally block.
|
||||
if __debug__:
|
||||
log.exception(__name__, exc)
|
||||
finally:
|
||||
if not __debug__ or not is_debug_session:
|
||||
# Unload modules imported by the workflow. Should not raise.
|
||||
# This is not done for the debug session because the snapshot taken
|
||||
# in a debug session would clear modules which are in use by the
|
||||
# workflow running on wire.
|
||||
# TODO utils.unimport_end(modules)
|
||||
|
||||
if next_message is None:
|
||||
|
||||
# Shut down the loop if there is no next message waiting.
|
||||
return # pylint: disable=lost-exception
|
||||
|
||||
except Exception as exc:
|
||||
# Log and try again. The session handler can only exit explicitly via
|
||||
# loop.clear() above.
|
||||
if __debug__:
|
||||
log.exception(__name__, exc)
|
||||
|
||||
async def read(
|
||||
self,
|
||||
expected_types: Container[int],
|
||||
expected_type: type[protobuf.MessageType] | None = None,
|
||||
) -> protobuf.MessageType:
|
||||
if __debug__:
|
||||
exp_type: str = str(expected_type)
|
||||
if expected_type is not None:
|
||||
exp_type = expected_type.MESSAGE_NAME
|
||||
log.debug(
|
||||
__name__,
|
||||
"Read - with expected types %s and expected type %s",
|
||||
str(expected_types),
|
||||
exp_type,
|
||||
)
|
||||
|
||||
message: Message = await self.incoming_message.take()
|
||||
|
||||
if message.type not in expected_types:
|
||||
raise UnexpectedMessageException(message)
|
||||
|
||||
if expected_type is None:
|
||||
expected_type = protobuf.type_for_wire(message.type)
|
||||
|
||||
return message_handler.wrap_protobuf_load(message.data, expected_type)
|
||||
|
||||
async def write(self, msg: protobuf.MessageType) -> None:
|
||||
return await self.channel_ctx.write(msg)
|
||||
|
||||
async def call(
|
||||
self, msg: protobuf.MessageType, expected_type: type[protobuf.MessageType]
|
||||
) -> protobuf.MessageType:
|
||||
assert expected_type.MESSAGE_WIRE_TYPE is not None
|
||||
|
||||
await self.write(msg)
|
||||
del msg
|
||||
|
||||
return await self.read((expected_type.MESSAGE_WIRE_TYPE,), expected_type)
|
||||
|
||||
async def call_any(
|
||||
self, msg: protobuf.MessageType, *expected_types: int
|
||||
) -> protobuf.MessageType:
|
||||
await self.write(msg)
|
||||
del msg
|
||||
return await self.read(expected_types)
|
||||
|
||||
|
||||
async def handle_pairing_request_message(
|
||||
pairing_ctx: PairingContext,
|
||||
msg: protocol_common.Message,
|
||||
use_workflow: bool,
|
||||
) -> protocol_common.Message | None:
|
||||
|
||||
res_msg: protobuf.MessageType | None = None
|
||||
|
||||
from apps.thp.pairing import handle_pairing_request
|
||||
|
||||
if msg.type in workflow.ALLOW_WHILE_LOCKED:
|
||||
workflow.autolock_interrupts_workflow = False
|
||||
|
||||
# Here we make sure we always respond with a Failure response
|
||||
# in case of any errors.
|
||||
try:
|
||||
# Find a protobuf.MessageType subclass that describes this
|
||||
# message. Raises if the type is not found.
|
||||
req_type = protobuf.type_for_wire(msg.type)
|
||||
|
||||
# Try to decode the message according to schema from
|
||||
# `req_type`. Raises if the message is malformed.
|
||||
req_msg = message_handler.wrap_protobuf_load(msg.data, req_type)
|
||||
|
||||
# Create the handler task.
|
||||
task = handle_pairing_request(pairing_ctx, req_msg)
|
||||
|
||||
# Run the workflow task. Workflow can do more on-the-wire
|
||||
# communication inside, but it should eventually return a
|
||||
# response message, or raise an exception (a rather common
|
||||
# thing to do). Exceptions are handled in the code below.
|
||||
if use_workflow:
|
||||
# Spawn a workflow around the task. This ensures that concurrent
|
||||
# workflows are shut down.
|
||||
res_msg = await workflow.spawn(context.with_context(pairing_ctx, task))
|
||||
else:
|
||||
# For debug messages, ignore workflow processing and just await
|
||||
# results of the handler.
|
||||
res_msg = await task
|
||||
|
||||
except UnexpectedMessageException as exc:
|
||||
# Workflow was trying to read a message from the wire, and
|
||||
# something unexpected came in. See Context.read() for
|
||||
# example, which expects some particular message and raises
|
||||
# UnexpectedMessage if another one comes in.
|
||||
# In order not to lose the message, we return it to the caller.
|
||||
# TODO:
|
||||
# We might handle only the few common cases here, like
|
||||
# Initialize and Cancel.
|
||||
return exc.msg
|
||||
|
||||
except BaseException as exc:
|
||||
# Either:
|
||||
# - the message had a type that has a registered handler, but does not have
|
||||
# a protobuf class
|
||||
# - the message was not valid protobuf
|
||||
# - workflow raised some kind of an exception while running
|
||||
# - something canceled the workflow from the outside
|
||||
if __debug__:
|
||||
if isinstance(exc, ActionCancelled):
|
||||
log.debug(__name__, "cancelled: %s", exc.message)
|
||||
elif isinstance(exc, loop.TaskClosed):
|
||||
log.debug(__name__, "cancelled: loop task was closed")
|
||||
else:
|
||||
log.exception(__name__, exc)
|
||||
res_msg = message_handler.failure(exc)
|
||||
|
||||
if res_msg is not None:
|
||||
# perform the write outside the big try-except block, so that usb write
|
||||
# problem bubbles up
|
||||
await pairing_ctx.write(res_msg)
|
||||
return None
|
||||
@ -0,0 +1,382 @@
|
||||
import ustruct
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from storage.cache_common import (
|
||||
CHANNEL_HANDSHAKE_HASH,
|
||||
CHANNEL_KEY_RECEIVE,
|
||||
CHANNEL_KEY_SEND,
|
||||
CHANNEL_NONCE_RECEIVE,
|
||||
CHANNEL_NONCE_SEND,
|
||||
)
|
||||
from storage.cache_thp import KEY_LENGTH, SESSION_ID_LENGTH, TAG_LENGTH
|
||||
from trezor import log, loop, utils
|
||||
from trezor.enums import FailureType
|
||||
from trezor.messages import Failure
|
||||
|
||||
from ..errors import DataError
|
||||
from ..protocol_common import Message
|
||||
from . import (
|
||||
ChannelState,
|
||||
SessionState,
|
||||
ThpDecryptionError,
|
||||
ThpError,
|
||||
ThpErrorType,
|
||||
ThpUnallocatedSessionError,
|
||||
)
|
||||
from . import alternating_bit_protocol as ABP
|
||||
from . import checksum, control_byte, is_channel_state_pairing, thp_messages
|
||||
from .checksum import CHECKSUM_LENGTH
|
||||
from .crypto import PUBKEY_LENGTH, Handshake
|
||||
from .thp_messages import (
|
||||
ACK_MESSAGE,
|
||||
HANDSHAKE_COMP_RES,
|
||||
HANDSHAKE_INIT_RES,
|
||||
PacketHeader,
|
||||
)
|
||||
from .writer import (
|
||||
INIT_HEADER_LENGTH,
|
||||
MESSAGE_TYPE_LENGTH,
|
||||
write_payload_to_wire_and_add_checksum,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from trezor.messages import ThpHandshakeCompletionReqNoisePayload
|
||||
|
||||
from .channel import Channel
|
||||
|
||||
if __debug__:
|
||||
from ubinascii import hexlify
|
||||
|
||||
from . import state_to_str
|
||||
|
||||
|
||||
async def handle_received_message(
|
||||
ctx: Channel, message_buffer: utils.BufferType
|
||||
) -> None:
|
||||
"""Handle a message received from the channel."""
|
||||
|
||||
if __debug__:
|
||||
log.debug(__name__, "handle_received_message")
|
||||
ctrl_byte, _, payload_length = ustruct.unpack(">BHH", message_buffer)
|
||||
message_length = payload_length + INIT_HEADER_LENGTH
|
||||
|
||||
_check_checksum(message_length, message_buffer)
|
||||
|
||||
# Synchronization process
|
||||
seq_bit = (ctrl_byte & 0x10) >> 4
|
||||
ack_bit = (ctrl_byte & 0x08) >> 3
|
||||
if __debug__:
|
||||
log.debug(
|
||||
__name__,
|
||||
"handle_completed_message - seq bit of message: %d, ack bit of message: %d",
|
||||
seq_bit,
|
||||
ack_bit,
|
||||
)
|
||||
|
||||
# 1: Handle ACKs
|
||||
if control_byte.is_ack(ctrl_byte):
|
||||
await _handle_ack(ctx, ack_bit)
|
||||
return
|
||||
|
||||
if _should_have_ctrl_byte_encrypted_transport(
|
||||
ctx
|
||||
) and not control_byte.is_encrypted_transport(ctrl_byte):
|
||||
raise ThpError("Message is not encrypted. Ignoring")
|
||||
|
||||
# 2: Handle message with unexpected sequential bit
|
||||
if seq_bit != ABP.get_expected_receive_seq_bit(ctx.channel_cache):
|
||||
if __debug__:
|
||||
log.debug(__name__, "Received message with an unexpected sequential bit")
|
||||
await _send_ack(ctx, ack_bit=seq_bit)
|
||||
raise ThpError("Received message with an unexpected sequential bit")
|
||||
|
||||
# 3: Send ACK in response
|
||||
await _send_ack(ctx, ack_bit=seq_bit)
|
||||
|
||||
ABP.set_expected_receive_seq_bit(ctx.channel_cache, 1 - seq_bit)
|
||||
|
||||
try:
|
||||
await _handle_message_to_app_or_channel(
|
||||
ctx, payload_length, message_length, ctrl_byte
|
||||
)
|
||||
except ThpUnallocatedSessionError as e:
|
||||
error_message = Failure(code=FailureType.ThpUnallocatedSession)
|
||||
await ctx.write(error_message, e.session_id)
|
||||
print(e)
|
||||
except ThpDecryptionError as e:
|
||||
await ctx.write_error(ThpErrorType.DECRYPTION_FAILED)
|
||||
# TODO ctx.write_error, i.e., ctx.send_channelClosed_error
|
||||
# ctx.clear
|
||||
print(e)
|
||||
if __debug__:
|
||||
log.debug(__name__, "handle_received_message - end")
|
||||
|
||||
|
||||
async def _send_ack(ctx: Channel, ack_bit: int) -> None:
|
||||
ctrl_byte = control_byte.add_ack_bit_to_ctrl_byte(ACK_MESSAGE, ack_bit)
|
||||
header = PacketHeader(ctrl_byte, ctx.get_channel_id_int(), CHECKSUM_LENGTH)
|
||||
if __debug__:
|
||||
log.debug(
|
||||
__name__,
|
||||
"Writing ACK message to a channel with id: %d, ack_bit: %d",
|
||||
ctx.get_channel_id_int(),
|
||||
ack_bit,
|
||||
)
|
||||
await write_payload_to_wire_and_add_checksum(ctx.iface, header, b"")
|
||||
|
||||
|
||||
def _check_checksum(message_length: int, message_buffer: utils.BufferType):
|
||||
if __debug__:
|
||||
log.debug(__name__, "check_checksum")
|
||||
if not checksum.is_valid(
|
||||
checksum=message_buffer[message_length - CHECKSUM_LENGTH : message_length],
|
||||
data=memoryview(message_buffer)[: message_length - CHECKSUM_LENGTH],
|
||||
):
|
||||
if __debug__:
|
||||
log.debug(__name__, "Invalid checksum, ignoring message.")
|
||||
raise ThpError("Invalid checksum, ignoring message.")
|
||||
|
||||
|
||||
async def _handle_ack(ctx: Channel, ack_bit: int):
|
||||
if not ABP.is_ack_valid(ctx.channel_cache, ack_bit):
|
||||
return
|
||||
# ACK is expected and it has correct sync bit
|
||||
if __debug__:
|
||||
log.debug(__name__, "Received ACK message with correct ack bit")
|
||||
if ctx.transmission_loop is not None:
|
||||
ctx.transmission_loop.stop_immediately()
|
||||
if __debug__:
|
||||
log.debug(__name__, "Stopped transmission loop")
|
||||
|
||||
ABP.set_sending_allowed(ctx.channel_cache, True)
|
||||
|
||||
if ctx.write_task_spawn is not None:
|
||||
if __debug__:
|
||||
log.debug(__name__, 'Control to "write_encrypted_payload_loop" task')
|
||||
await ctx.write_task_spawn
|
||||
# Note that no the write_task_spawn could result in loop.clear(),
|
||||
# which will result in termination of this function - any code after
|
||||
# this await might not be executed
|
||||
|
||||
|
||||
async def _handle_message_to_app_or_channel(
|
||||
ctx: Channel,
|
||||
payload_length: int,
|
||||
message_length: int,
|
||||
ctrl_byte: int,
|
||||
) -> None:
|
||||
state = ctx.get_channel_state()
|
||||
if __debug__:
|
||||
log.debug(__name__, "state: %s", state_to_str(state))
|
||||
|
||||
if state is ChannelState.ENCRYPTED_TRANSPORT:
|
||||
await _handle_state_ENCRYPTED_TRANSPORT(ctx, message_length)
|
||||
return
|
||||
|
||||
if state is ChannelState.TH1:
|
||||
await _handle_state_TH1(ctx, payload_length, message_length, ctrl_byte)
|
||||
return
|
||||
|
||||
if state is ChannelState.TH2:
|
||||
await _handle_state_TH2(ctx, message_length, ctrl_byte)
|
||||
return
|
||||
|
||||
if is_channel_state_pairing(state):
|
||||
await _handle_pairing(ctx, message_length)
|
||||
return
|
||||
|
||||
raise ThpError("Unimplemented channel state")
|
||||
|
||||
|
||||
async def _handle_state_TH1(
|
||||
ctx: Channel,
|
||||
payload_length: int,
|
||||
message_length: int,
|
||||
ctrl_byte: int,
|
||||
) -> None:
|
||||
if __debug__:
|
||||
log.debug(__name__, "handle_state_TH1")
|
||||
if not control_byte.is_handshake_init_req(ctrl_byte):
|
||||
raise ThpError("Message received is not a handshake init request!")
|
||||
if not payload_length == PUBKEY_LENGTH + CHECKSUM_LENGTH:
|
||||
raise ThpError("Message received is not a valid handshake init request!")
|
||||
|
||||
ctx.handshake = Handshake()
|
||||
|
||||
host_ephemeral_pubkey = bytearray(
|
||||
ctx.buffer[INIT_HEADER_LENGTH : message_length - CHECKSUM_LENGTH]
|
||||
)
|
||||
trezor_ephemeral_pubkey, encrypted_trezor_static_pubkey, tag = (
|
||||
ctx.handshake.handle_th1_crypto(
|
||||
thp_messages.get_encoded_device_properties(), host_ephemeral_pubkey
|
||||
)
|
||||
)
|
||||
|
||||
if __debug__:
|
||||
log.debug(
|
||||
__name__,
|
||||
"trezor ephemeral pubkey: %s",
|
||||
hexlify(trezor_ephemeral_pubkey).decode(),
|
||||
)
|
||||
log.debug(
|
||||
__name__,
|
||||
"trezor masked static pubkey: %s",
|
||||
hexlify(encrypted_trezor_static_pubkey).decode(),
|
||||
)
|
||||
log.debug(__name__, "tag: %s", hexlify(tag))
|
||||
|
||||
payload = trezor_ephemeral_pubkey + encrypted_trezor_static_pubkey + tag
|
||||
|
||||
# send handshake init response message
|
||||
await ctx.write_handshake_message(HANDSHAKE_INIT_RES, payload)
|
||||
ctx.set_channel_state(ChannelState.TH2)
|
||||
return
|
||||
|
||||
|
||||
async def _handle_state_TH2(ctx: Channel, message_length: int, ctrl_byte: int) -> None:
|
||||
from apps.thp.credential_manager import validate_credential
|
||||
|
||||
if __debug__:
|
||||
log.debug(__name__, "handle_state_TH2")
|
||||
if not control_byte.is_handshake_comp_req(ctrl_byte):
|
||||
raise ThpError("Message received is not a handshake completion request!")
|
||||
if ctx.handshake is None:
|
||||
raise Exception("Handshake object is not prepared. Retry handshake.")
|
||||
|
||||
host_encrypted_static_pubkey = memoryview(ctx.buffer)[
|
||||
INIT_HEADER_LENGTH : INIT_HEADER_LENGTH + KEY_LENGTH + TAG_LENGTH
|
||||
]
|
||||
handshake_completion_request_noise_payload = memoryview(ctx.buffer)[
|
||||
INIT_HEADER_LENGTH + KEY_LENGTH + TAG_LENGTH : message_length - CHECKSUM_LENGTH
|
||||
]
|
||||
|
||||
ctx.handshake.handle_th2_crypto(
|
||||
host_encrypted_static_pubkey, handshake_completion_request_noise_payload
|
||||
)
|
||||
|
||||
ctx.channel_cache.set(CHANNEL_KEY_RECEIVE, ctx.handshake.key_receive)
|
||||
ctx.channel_cache.set(CHANNEL_KEY_SEND, ctx.handshake.key_send)
|
||||
ctx.channel_cache.set(CHANNEL_HANDSHAKE_HASH, ctx.handshake.h)
|
||||
ctx.channel_cache.set_int(CHANNEL_NONCE_RECEIVE, 0)
|
||||
ctx.channel_cache.set_int(CHANNEL_NONCE_SEND, 1)
|
||||
|
||||
noise_payload = thp_messages.decode_message(
|
||||
ctx.buffer[
|
||||
INIT_HEADER_LENGTH
|
||||
+ KEY_LENGTH
|
||||
+ TAG_LENGTH : message_length
|
||||
- CHECKSUM_LENGTH
|
||||
- TAG_LENGTH
|
||||
],
|
||||
0,
|
||||
"ThpHandshakeCompletionReqNoisePayload",
|
||||
)
|
||||
if TYPE_CHECKING:
|
||||
assert ThpHandshakeCompletionReqNoisePayload.is_type_of(noise_payload)
|
||||
for i in noise_payload.pairing_methods:
|
||||
if i not in ctx.selected_pairing_methods:
|
||||
ctx.selected_pairing_methods.append(i)
|
||||
if __debug__:
|
||||
log.debug(
|
||||
__name__,
|
||||
"host static pubkey: %s, noise payload: %s",
|
||||
utils.get_bytes_as_str(host_encrypted_static_pubkey),
|
||||
utils.get_bytes_as_str(handshake_completion_request_noise_payload),
|
||||
)
|
||||
|
||||
# key is decoded in handshake._handle_th2_crypto
|
||||
host_static_pubkey = host_encrypted_static_pubkey[:PUBKEY_LENGTH]
|
||||
|
||||
paired: bool = False
|
||||
|
||||
if noise_payload.host_pairing_credential is not None:
|
||||
try: # TODO change try-except for something better
|
||||
paired = validate_credential(
|
||||
noise_payload.host_pairing_credential,
|
||||
host_static_pubkey,
|
||||
)
|
||||
except DataError as e:
|
||||
if __debug__:
|
||||
log.exception(__name__, e)
|
||||
pass
|
||||
|
||||
trezor_state = thp_messages.TREZOR_STATE_UNPAIRED
|
||||
if paired:
|
||||
trezor_state = thp_messages.TREZOR_STATE_PAIRED
|
||||
# send hanshake completion response
|
||||
await ctx.write_handshake_message(
|
||||
HANDSHAKE_COMP_RES,
|
||||
ctx.handshake.get_handshake_completion_response(trezor_state),
|
||||
)
|
||||
|
||||
ctx.handshake = None
|
||||
|
||||
if paired:
|
||||
ctx.set_channel_state(ChannelState.ENCRYPTED_TRANSPORT)
|
||||
else:
|
||||
ctx.set_channel_state(ChannelState.TP1)
|
||||
|
||||
|
||||
async def _handle_state_ENCRYPTED_TRANSPORT(ctx: Channel, message_length: int) -> None:
|
||||
if __debug__:
|
||||
log.debug(__name__, "handle_state_ENCRYPTED_TRANSPORT")
|
||||
|
||||
ctx.decrypt_buffer(message_length)
|
||||
session_id, message_type = ustruct.unpack(
|
||||
">BH", memoryview(ctx.buffer)[INIT_HEADER_LENGTH:]
|
||||
)
|
||||
if session_id not in ctx.sessions:
|
||||
raise ThpUnallocatedSessionError(session_id)
|
||||
|
||||
session_state = ctx.sessions[session_id].get_session_state()
|
||||
if session_state is SessionState.UNALLOCATED:
|
||||
raise ThpUnallocatedSessionError(session_id)
|
||||
|
||||
ctx.sessions[session_id].incoming_message.publish(
|
||||
Message(
|
||||
message_type,
|
||||
ctx.buffer[
|
||||
INIT_HEADER_LENGTH
|
||||
+ MESSAGE_TYPE_LENGTH
|
||||
+ SESSION_ID_LENGTH : message_length
|
||||
- CHECKSUM_LENGTH
|
||||
- TAG_LENGTH
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _handle_pairing(ctx: Channel, message_length: int) -> None:
|
||||
from .pairing_context import PairingContext
|
||||
|
||||
if ctx.connection_context is None:
|
||||
ctx.connection_context = PairingContext(ctx)
|
||||
loop.schedule(ctx.connection_context.handle())
|
||||
|
||||
ctx.decrypt_buffer(message_length)
|
||||
message_type = ustruct.unpack(
|
||||
">H", ctx.buffer[INIT_HEADER_LENGTH + SESSION_ID_LENGTH :]
|
||||
)[0]
|
||||
|
||||
ctx.connection_context.incoming_message.publish(
|
||||
Message(
|
||||
message_type,
|
||||
ctx.buffer[
|
||||
INIT_HEADER_LENGTH
|
||||
+ MESSAGE_TYPE_LENGTH
|
||||
+ SESSION_ID_LENGTH : message_length
|
||||
- CHECKSUM_LENGTH
|
||||
- TAG_LENGTH
|
||||
],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _should_have_ctrl_byte_encrypted_transport(ctx: Channel) -> bool:
|
||||
if ctx.get_channel_state() in [
|
||||
ChannelState.UNALLOCATED,
|
||||
ChannelState.TH1,
|
||||
ChannelState.TH2,
|
||||
]:
|
||||
return False
|
||||
return True
|
||||
@ -0,0 +1,223 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from storage.cache_thp import MANAGEMENT_SESSION_ID, SessionThpCache
|
||||
from trezor import log, loop, protobuf, utils
|
||||
from trezor.wire import message_handler, protocol_common
|
||||
from trezor.wire.context import UnexpectedMessageException
|
||||
from trezor.wire.message_handler import AVOID_RESTARTING_FOR, failure, find_handler
|
||||
|
||||
from ..protocol_common import Context, Message
|
||||
from . import SessionState
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any, Awaitable, Container, TypeVar, overload
|
||||
|
||||
from storage.cache_common import DataCache
|
||||
|
||||
from ..message_handler import HandlerFinder
|
||||
from .channel import Channel
|
||||
|
||||
pass
|
||||
|
||||
_EXIT_LOOP = True
|
||||
_REPEAT_LOOP = False
|
||||
|
||||
if __debug__:
|
||||
from trezor.utils import get_bytes_as_str
|
||||
|
||||
|
||||
class GenericSessionContext(Context):
|
||||
|
||||
def __init__(self, channel: Channel, session_id: int) -> None:
|
||||
super().__init__(channel.iface, channel.channel_id)
|
||||
self.channel: Channel = channel
|
||||
self.session_id: int = session_id
|
||||
self.incoming_message = loop.chan()
|
||||
self.handler_finder: HandlerFinder = find_handler
|
||||
|
||||
async def handle(self, is_debug_session: bool = False) -> None:
|
||||
if __debug__:
|
||||
self._handle_debug(is_debug_session)
|
||||
|
||||
take = self.incoming_message.take()
|
||||
next_message: Message | None = None
|
||||
|
||||
# Take a mark of modules that are imported at this point, so we can
|
||||
# roll back and un-import any others.
|
||||
# TODO modules = utils.unimport_begin()
|
||||
while True:
|
||||
try:
|
||||
if await self._handle_message(take, next_message, is_debug_session):
|
||||
return
|
||||
except UnexpectedMessageException as unexpected:
|
||||
# The workflow was interrupted by an unexpected message. We need to
|
||||
# process it as if it was a new message...
|
||||
next_message = unexpected.msg
|
||||
continue
|
||||
except Exception as exc:
|
||||
# Log and try again.
|
||||
if __debug__:
|
||||
log.exception(__name__, exc)
|
||||
|
||||
def _handle_debug(self, is_debug_session: bool) -> None:
|
||||
log.debug(
|
||||
__name__,
|
||||
"handle - start (channel_id (bytes): %s, session_id: %d)",
|
||||
get_bytes_as_str(self.channel_id),
|
||||
self.session_id,
|
||||
)
|
||||
if is_debug_session:
|
||||
import apps.debug
|
||||
|
||||
apps.debug.DEBUG_CONTEXT = self
|
||||
|
||||
async def _handle_message(
|
||||
self,
|
||||
take: Awaitable[Any],
|
||||
next_message: Message | None,
|
||||
is_debug_session: bool,
|
||||
) -> bool:
|
||||
|
||||
try:
|
||||
message = await self._get_message(take, next_message)
|
||||
except protocol_common.WireError as e:
|
||||
if __debug__:
|
||||
log.exception(__name__, e)
|
||||
await self.write(failure(e))
|
||||
return _REPEAT_LOOP
|
||||
|
||||
try:
|
||||
await message_handler.handle_single_message(
|
||||
self,
|
||||
message,
|
||||
self.handler_finder,
|
||||
)
|
||||
except UnexpectedMessageException:
|
||||
raise
|
||||
except Exception as exc:
|
||||
# Log and ignore. The session handler can only exit explicitly in the
|
||||
# following finally block.
|
||||
if __debug__:
|
||||
log.exception(__name__, exc)
|
||||
finally:
|
||||
if not __debug__ or not is_debug_session:
|
||||
# Unload modules imported by the workflow. Should not raise.
|
||||
# This is not done for the debug session because the snapshot taken
|
||||
# in a debug session would clear modules which are in use by the
|
||||
# workflow running on wire.
|
||||
# TODO utils.unimport_end(modules)
|
||||
|
||||
if next_message is None and message.type not in AVOID_RESTARTING_FOR:
|
||||
# Shut down the loop if there is no next message waiting.
|
||||
return _EXIT_LOOP # pylint: disable=lost-exception
|
||||
return _REPEAT_LOOP # pylint: disable=lost-exception
|
||||
|
||||
async def _get_message(
|
||||
self, take: Awaitable[Any], next_message: Message | None
|
||||
) -> Message:
|
||||
if next_message is None:
|
||||
# If the previous run did not keep an unprocessed message for us,
|
||||
# wait for a new one.
|
||||
message: Message = await take
|
||||
else:
|
||||
# Process the message from previous run.
|
||||
message = next_message
|
||||
next_message = None
|
||||
return message
|
||||
|
||||
async def read(
|
||||
self,
|
||||
expected_types: Container[int],
|
||||
expected_type: type[protobuf.MessageType] | None = None,
|
||||
) -> protobuf.MessageType:
|
||||
if __debug__:
|
||||
exp_type: str = str(expected_type)
|
||||
if expected_type is not None:
|
||||
exp_type = expected_type.MESSAGE_NAME
|
||||
log.debug(
|
||||
__name__,
|
||||
"Read - with expected types %s and expected type %s",
|
||||
str(expected_types),
|
||||
exp_type,
|
||||
)
|
||||
message: Message = await self.incoming_message.take()
|
||||
if message.type not in expected_types:
|
||||
raise UnexpectedMessageException(message)
|
||||
|
||||
if expected_type is None:
|
||||
expected_type = protobuf.type_for_wire(message.type)
|
||||
|
||||
return message_handler.wrap_protobuf_load(message.data, expected_type)
|
||||
|
||||
async def write(self, msg: protobuf.MessageType) -> None:
|
||||
return await self.channel.write(msg, self.session_id)
|
||||
|
||||
def get_session_state(self) -> SessionState: ...
|
||||
|
||||
|
||||
class ManagementSessionContext(GenericSessionContext):
|
||||
|
||||
def __init__(self, channel_ctx: Channel) -> None:
|
||||
super().__init__(channel_ctx, MANAGEMENT_SESSION_ID)
|
||||
from trezor.wire.thp.handler_provider import (
|
||||
find_management_session_message_handler,
|
||||
)
|
||||
|
||||
self.handler_finder = find_management_session_message_handler
|
||||
|
||||
def get_session_state(self) -> SessionState:
|
||||
return SessionState.MANAGEMENT
|
||||
|
||||
|
||||
class SessionContext(GenericSessionContext):
|
||||
|
||||
def __init__(self, channel_ctx: Channel, session_cache: SessionThpCache) -> None:
|
||||
if channel_ctx.channel_id != session_cache.channel_id:
|
||||
raise Exception(
|
||||
"The session has different channel id than the provided channel context!"
|
||||
)
|
||||
session_id = int.from_bytes(session_cache.session_id, "big")
|
||||
super().__init__(channel_ctx, session_id)
|
||||
self.session_cache = session_cache
|
||||
|
||||
# ACCESS TO SESSION DATA
|
||||
|
||||
def get_session_state(self) -> SessionState:
|
||||
state = int.from_bytes(self.session_cache.state, "big")
|
||||
return SessionState(state)
|
||||
|
||||
def set_session_state(self, state: SessionState) -> None:
|
||||
self.session_cache.state = bytearray(state.to_bytes(1, "big"))
|
||||
|
||||
# ACCESS TO CACHE
|
||||
@property
|
||||
def cache(self) -> DataCache:
|
||||
return self.session_cache
|
||||
|
||||
if TYPE_CHECKING:
|
||||
T = TypeVar("T")
|
||||
|
||||
@overload
|
||||
def cache_get(self, key: int) -> bytes | None: # noqa: F811
|
||||
...
|
||||
|
||||
@overload
|
||||
def cache_get(self, key: int, default: T) -> bytes | T: # noqa: F811
|
||||
...
|
||||
|
||||
def cache_get(
|
||||
self, key: int, default: T | None = None
|
||||
) -> bytes | T | None: # noqa: F811
|
||||
utils.ensure(key < len(self.session_cache.fields))
|
||||
if self.session_cache.data[key][0] != 1:
|
||||
return default
|
||||
return bytes(self.session_cache.data[key][1:])
|
||||
|
||||
def cache_is_set(self, key: int) -> bool:
|
||||
return self.session_cache.is_set(key)
|
||||
|
||||
def cache_set(self, key: int, value: bytes) -> None:
|
||||
utils.ensure(key < len(self.session_cache.fields))
|
||||
utils.ensure(len(value) <= self.session_cache.fields[key])
|
||||
self.session_cache.data[key][0] = 1
|
||||
self.session_cache.data[key][1:] = value
|
||||
@ -0,0 +1,48 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from storage import cache_thp
|
||||
from trezor import loop
|
||||
|
||||
from .session_context import (
|
||||
GenericSessionContext,
|
||||
ManagementSessionContext,
|
||||
SessionContext,
|
||||
)
|
||||
|
||||
if __debug__:
|
||||
from trezor import log
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .channel import Channel
|
||||
|
||||
|
||||
def create_new_session(channel_ctx: Channel) -> SessionContext:
|
||||
session_cache = cache_thp.get_new_session(channel_ctx.channel_cache)
|
||||
return SessionContext(channel_ctx, session_cache)
|
||||
|
||||
|
||||
def create_new_management_session(
|
||||
channel_ctx: Channel,
|
||||
) -> ManagementSessionContext:
|
||||
return ManagementSessionContext(channel_ctx)
|
||||
|
||||
|
||||
def load_cached_sessions(
|
||||
channel_ctx: Channel,
|
||||
) -> dict[int, GenericSessionContext]:
|
||||
if __debug__:
|
||||
log.debug(__name__, "load_cached_sessions")
|
||||
sessions: dict[int, GenericSessionContext] = {}
|
||||
cached_sessions = cache_thp.get_all_allocated_sessions()
|
||||
if __debug__:
|
||||
log.debug(
|
||||
__name__,
|
||||
"load_cached_sessions - loaded a total of %d sessions from cache",
|
||||
len(cached_sessions),
|
||||
)
|
||||
for session in cached_sessions:
|
||||
if session.channel_id == channel_ctx.channel_id:
|
||||
sid = int.from_bytes(session.session_id, "big")
|
||||
sessions[sid] = SessionContext(channel_ctx, session)
|
||||
loop.schedule(sessions[sid].handle())
|
||||
return sessions
|
||||
@ -0,0 +1,120 @@
|
||||
import ustruct
|
||||
from micropython import const
|
||||
|
||||
from storage.cache_thp import BROADCAST_CHANNEL_ID
|
||||
from trezor import protobuf, utils
|
||||
from trezor.enums import ThpPairingMethod
|
||||
from trezor.messages import ThpDeviceProperties
|
||||
|
||||
from .. import message_handler
|
||||
|
||||
CODEC_V1 = const(0x3F)
|
||||
CONTINUATION_PACKET = const(0x80)
|
||||
HANDSHAKE_INIT_REQ = const(0x00)
|
||||
HANDSHAKE_INIT_RES = const(0x01)
|
||||
HANDSHAKE_COMP_REQ = const(0x02)
|
||||
HANDSHAKE_COMP_RES = const(0x03)
|
||||
ENCRYPTED_TRANSPORT = const(0x04)
|
||||
|
||||
CONTINUATION_PACKET_MASK = const(0x80)
|
||||
ACK_MASK = const(0xF7)
|
||||
DATA_MASK = const(0xE7)
|
||||
|
||||
ACK_MESSAGE = const(0x20)
|
||||
_ERROR = const(0x42)
|
||||
CHANNEL_ALLOCATION_REQ = const(0x40)
|
||||
_CHANNEL_ALLOCATION_RES = const(0x41)
|
||||
|
||||
TREZOR_STATE_UNPAIRED = b"\x00"
|
||||
TREZOR_STATE_PAIRED = b"\x01"
|
||||
|
||||
if __debug__:
|
||||
from trezor import log
|
||||
|
||||
|
||||
class PacketHeader:
|
||||
format_str_init = ">BHH"
|
||||
format_str_cont = ">BH"
|
||||
|
||||
def __init__(self, ctrl_byte: int, cid: int, length: int) -> None:
|
||||
self.ctrl_byte = ctrl_byte
|
||||
self.cid = cid
|
||||
self.length = length
|
||||
|
||||
def to_bytes(self) -> bytes:
|
||||
return ustruct.pack(self.format_str_init, self.ctrl_byte, self.cid, self.length)
|
||||
|
||||
def pack_to_init_buffer(self, buffer, buffer_offset=0) -> None:
|
||||
ustruct.pack_into(
|
||||
self.format_str_init,
|
||||
buffer,
|
||||
buffer_offset,
|
||||
self.ctrl_byte,
|
||||
self.cid,
|
||||
self.length,
|
||||
)
|
||||
|
||||
def pack_to_cont_buffer(self, buffer, buffer_offset=0) -> None:
|
||||
ustruct.pack_into(
|
||||
self.format_str_cont, buffer, buffer_offset, CONTINUATION_PACKET, self.cid
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_error_header(cls, cid, length):
|
||||
return cls(_ERROR, cid, length)
|
||||
|
||||
@classmethod
|
||||
def get_channel_allocation_response_header(cls, length):
|
||||
return cls(_CHANNEL_ALLOCATION_RES, BROADCAST_CHANNEL_ID, length)
|
||||
|
||||
|
||||
_ENCODED_DEVICE_PROPERTIES: bytes | None = None
|
||||
|
||||
_ENABLED_PAIRING_METHODS = [
|
||||
ThpPairingMethod.CodeEntry,
|
||||
ThpPairingMethod.QrCode,
|
||||
ThpPairingMethod.NFC_Unidirectional,
|
||||
]
|
||||
|
||||
|
||||
def _get_device_properties() -> ThpDeviceProperties:
|
||||
# TODO define model variants
|
||||
return ThpDeviceProperties(
|
||||
pairing_methods=_ENABLED_PAIRING_METHODS,
|
||||
internal_model=utils.INTERNAL_MODEL,
|
||||
model_variant=0,
|
||||
bootloader_mode=False,
|
||||
protocol_version=3,
|
||||
)
|
||||
|
||||
|
||||
def get_encoded_device_properties() -> bytes:
|
||||
global _ENCODED_DEVICE_PROPERTIES
|
||||
if _ENCODED_DEVICE_PROPERTIES is None:
|
||||
props = _get_device_properties()
|
||||
length = protobuf.encoded_length(props)
|
||||
_ENCODED_DEVICE_PROPERTIES = bytearray(length)
|
||||
protobuf.encode(_ENCODED_DEVICE_PROPERTIES, props)
|
||||
return _ENCODED_DEVICE_PROPERTIES
|
||||
|
||||
|
||||
def get_channel_allocation_response(nonce: bytes, new_cid: bytes) -> bytes:
|
||||
props_msg = get_encoded_device_properties()
|
||||
return nonce + new_cid + props_msg
|
||||
|
||||
|
||||
def get_codec_v1_error_message() -> bytes:
|
||||
return b"?##" # TODO add unsupported codec_v1 error message
|
||||
|
||||
|
||||
def decode_message(
|
||||
buffer: bytes, msg_type: int, message_name: str | None = None
|
||||
) -> protobuf.MessageType:
|
||||
if __debug__:
|
||||
log.debug(__name__, "decode message")
|
||||
if message_name is not None:
|
||||
expected_type = protobuf.type_for_name(message_name)
|
||||
else:
|
||||
expected_type = protobuf.type_for_wire(msg_type)
|
||||
x = message_handler.wrap_protobuf_load(buffer, expected_type)
|
||||
return x
|
||||
@ -0,0 +1,59 @@
|
||||
from micropython import const
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from trezor import loop
|
||||
from trezor.wire.thp.thp_messages import PacketHeader
|
||||
from trezor.wire.thp.writer import write_payload_to_wire_and_add_checksum
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from trezor.wire.thp.channel import Channel
|
||||
|
||||
MAX_RETRANSMISSION_COUNT = const(50)
|
||||
MIN_RETRANSMISSION_COUNT = const(2)
|
||||
|
||||
|
||||
class TransmissionLoop:
|
||||
|
||||
def __init__(
|
||||
self, channel: Channel, header: PacketHeader, transport_payload: bytes
|
||||
) -> None:
|
||||
self.channel: Channel = channel
|
||||
self.header: PacketHeader = header
|
||||
self.transport_payload: bytes = transport_payload
|
||||
self.wait_task: loop.spawn | None = None
|
||||
self.min_retransmisson_count_achieved: bool = False
|
||||
|
||||
async def start(self):
|
||||
self.min_retransmisson_count_achieved = False
|
||||
for i in range(MAX_RETRANSMISSION_COUNT):
|
||||
if i >= MIN_RETRANSMISSION_COUNT:
|
||||
self.min_retransmisson_count_achieved = True
|
||||
await write_payload_to_wire_and_add_checksum(
|
||||
self.channel.iface, self.header, self.transport_payload
|
||||
)
|
||||
self.wait_task = loop.spawn(self._wait(i))
|
||||
try:
|
||||
await self.wait_task
|
||||
except loop.TaskClosed:
|
||||
self.wait_task = None
|
||||
break
|
||||
|
||||
def stop_immediately(self):
|
||||
if self.wait_task is not None:
|
||||
self.wait_task.close()
|
||||
self.wait_task = None
|
||||
|
||||
async def stop_after_min_retransmission(self):
|
||||
while not self.min_retransmisson_count_achieved and self.wait_task is not None:
|
||||
await self._short_wait()
|
||||
self.stop_immediately()
|
||||
|
||||
async def _wait(self, counter: int = 0) -> None:
|
||||
timeout_ms = round(10200 - 1010000 / (counter + 100))
|
||||
await loop.sleep(timeout_ms)
|
||||
|
||||
async def _short_wait(self):
|
||||
loop.wait(50)
|
||||
|
||||
def __del__(self):
|
||||
self.stop_immediately()
|
||||
@ -0,0 +1,82 @@
|
||||
from micropython import const
|
||||
from trezorcrypto import crc
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from trezor import io, log, loop, utils
|
||||
from trezor.wire.thp.thp_messages import PacketHeader
|
||||
|
||||
INIT_HEADER_LENGTH = const(5)
|
||||
CONT_HEADER_LENGTH = const(3)
|
||||
PACKET_LENGTH = const(64)
|
||||
CHECKSUM_LENGTH = const(4)
|
||||
MAX_PAYLOAD_LEN = const(60000)
|
||||
MESSAGE_TYPE_LENGTH = const(2)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from trezorio import WireInterface
|
||||
from typing import Sequence
|
||||
|
||||
|
||||
async def write_payload_to_wire_and_add_checksum(
|
||||
iface: WireInterface, header: PacketHeader, transport_payload: bytes
|
||||
):
|
||||
header_checksum: int = crc.crc32(header.to_bytes())
|
||||
checksum: bytes = crc.crc32(transport_payload, header_checksum).to_bytes(
|
||||
CHECKSUM_LENGTH, "big"
|
||||
)
|
||||
data = (transport_payload, checksum)
|
||||
await write_payloads_to_wire(iface, header, data)
|
||||
|
||||
|
||||
async def write_payloads_to_wire(
|
||||
iface: WireInterface, header: PacketHeader, data: Sequence[bytes]
|
||||
):
|
||||
n_of_data = len(data)
|
||||
total_length = sum(len(item) for item in data)
|
||||
|
||||
current_data_idx = 0
|
||||
current_data_offset = 0
|
||||
|
||||
packet = bytearray(PACKET_LENGTH)
|
||||
header.pack_to_init_buffer(packet)
|
||||
packet_offset: int = INIT_HEADER_LENGTH
|
||||
packet_number = 0
|
||||
nwritten = 0
|
||||
while nwritten < total_length:
|
||||
if packet_number == 1:
|
||||
header.pack_to_cont_buffer(packet)
|
||||
if packet_number >= 1 and nwritten >= total_length - PACKET_LENGTH:
|
||||
packet[:] = bytearray(PACKET_LENGTH)
|
||||
header.pack_to_cont_buffer(packet)
|
||||
while True:
|
||||
n = utils.memcpy(
|
||||
packet, packet_offset, data[current_data_idx], current_data_offset
|
||||
)
|
||||
packet_offset += n
|
||||
current_data_offset += n
|
||||
nwritten += n
|
||||
|
||||
if packet_offset < PACKET_LENGTH:
|
||||
current_data_idx += 1
|
||||
current_data_offset = 0
|
||||
if current_data_idx >= n_of_data:
|
||||
break
|
||||
elif packet_offset == PACKET_LENGTH:
|
||||
break
|
||||
else:
|
||||
raise Exception("Should not happen!!!")
|
||||
packet_number += 1
|
||||
packet_offset = CONT_HEADER_LENGTH
|
||||
await write_packet_to_wire(iface, packet)
|
||||
|
||||
|
||||
async def write_packet_to_wire(iface: WireInterface, packet: bytes) -> None:
|
||||
while True:
|
||||
await loop.wait(iface.iface_num() | io.POLL_WRITE)
|
||||
if __debug__:
|
||||
log.debug(
|
||||
__name__, "write_packet_to_wire: %s", utils.get_bytes_as_str(packet)
|
||||
)
|
||||
n = iface.write(packet)
|
||||
if n == len(packet):
|
||||
return
|
||||
@ -0,0 +1,159 @@
|
||||
from ubinascii import hexlify
|
||||
import ustruct
|
||||
from micropython import const
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from storage.cache_thp import BROADCAST_CHANNEL_ID
|
||||
from trezor import io, log, loop, utils
|
||||
from trezor.wire.thp import writer
|
||||
|
||||
from .thp import (
|
||||
ChannelState,
|
||||
ThpError,
|
||||
ThpErrorType,
|
||||
channel_manager,
|
||||
checksum,
|
||||
session_manager,
|
||||
thp_messages,
|
||||
)
|
||||
from .thp.channel import Channel
|
||||
from .thp.checksum import CHECKSUM_LENGTH
|
||||
from .thp.thp_messages import CHANNEL_ALLOCATION_REQ, CODEC_V1, PacketHeader
|
||||
from .thp.writer import (
|
||||
MAX_PAYLOAD_LEN,
|
||||
PACKET_LENGTH,
|
||||
INIT_HEADER_LENGTH,
|
||||
write_payload_to_wire_and_add_checksum,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from trezorio import WireInterface
|
||||
|
||||
_CID_REQ_PAYLOAD_LENGTH = const(12)
|
||||
_BUFFER: bytearray
|
||||
_CHANNELS: dict[int, Channel] = {}
|
||||
|
||||
|
||||
def set_buffer(buffer):
|
||||
global _BUFFER
|
||||
_BUFFER = buffer
|
||||
|
||||
|
||||
async def thp_main_loop(iface: WireInterface, is_debug_session=False):
|
||||
global _CHANNELS
|
||||
global _BUFFER
|
||||
_CHANNELS = channel_manager.load_cached_channels(_BUFFER)
|
||||
for ch in _CHANNELS.values():
|
||||
ch.sessions.update(session_manager.load_cached_sessions(ch))
|
||||
|
||||
read = loop.wait(iface.iface_num() | io.POLL_READ)
|
||||
|
||||
while True:
|
||||
try:
|
||||
if __debug__:
|
||||
log.debug(__name__, "thp_main_loop")
|
||||
packet = await read
|
||||
ctrl_byte, cid = ustruct.unpack(">BH", packet)
|
||||
|
||||
if ctrl_byte == CODEC_V1:
|
||||
await _handle_codec_v1(iface, packet)
|
||||
|
||||
if cid == BROADCAST_CHANNEL_ID:
|
||||
await _handle_broadcast(iface, ctrl_byte, packet)
|
||||
continue
|
||||
|
||||
if cid in _CHANNELS:
|
||||
await _handle_allocated(iface, cid, packet)
|
||||
else:
|
||||
await _handle_unallocated(iface, cid)
|
||||
|
||||
except ThpError as e:
|
||||
if __debug__:
|
||||
log.exception(__name__, e)
|
||||
|
||||
|
||||
async def _handle_codec_v1(iface: WireInterface, packet):
|
||||
# If the received packet is not initial codec_v1 packet, do not send error message
|
||||
if not packet[1:3] == b"##":
|
||||
return
|
||||
error_message = thp_messages.get_codec_v1_error_message()
|
||||
await writer.write_packet_to_wire(iface, error_message)
|
||||
|
||||
|
||||
async def _handle_broadcast(
|
||||
iface: WireInterface, ctrl_byte: int, packet: utils.BufferType
|
||||
) -> None:
|
||||
global _BUFFER
|
||||
if ctrl_byte != CHANNEL_ALLOCATION_REQ:
|
||||
raise ThpError("Unexpected ctrl_byte in a broadcast channel packet")
|
||||
if __debug__:
|
||||
log.debug(__name__, "Received valid message on the broadcast channel")
|
||||
|
||||
length, nonce = ustruct.unpack(">H8s", packet[3:])
|
||||
payload = _get_buffer_for_payload(length, packet[5:], _CID_REQ_PAYLOAD_LENGTH)
|
||||
if not checksum.is_valid(
|
||||
payload[-4:],
|
||||
packet[: _CID_REQ_PAYLOAD_LENGTH + INIT_HEADER_LENGTH - CHECKSUM_LENGTH],
|
||||
):
|
||||
raise ThpError("Checksum is not valid")
|
||||
|
||||
new_channel: Channel = channel_manager.create_new_channel(iface, _BUFFER)
|
||||
cid = int.from_bytes(new_channel.channel_id, "big")
|
||||
_CHANNELS[cid] = new_channel
|
||||
|
||||
response_data = thp_messages.get_channel_allocation_response(
|
||||
nonce, new_channel.channel_id
|
||||
)
|
||||
response_header = PacketHeader.get_channel_allocation_response_header(
|
||||
len(response_data) + CHECKSUM_LENGTH,
|
||||
)
|
||||
if __debug__:
|
||||
log.debug(__name__, "New channel allocated with id %d", cid)
|
||||
|
||||
await write_payload_to_wire_and_add_checksum(iface, response_header, response_data)
|
||||
|
||||
|
||||
async def _handle_allocated(
|
||||
iface: WireInterface, cid: int, packet: utils.BufferType
|
||||
) -> None:
|
||||
channel = _CHANNELS[cid]
|
||||
if channel is None:
|
||||
_handle_unallocated(iface, cid)
|
||||
raise ThpError("Invalid state of a channel")
|
||||
if channel.iface is not iface:
|
||||
# TODO send error message to wire
|
||||
raise ThpError("Channel has different WireInterface")
|
||||
|
||||
if channel.get_channel_state() != ChannelState.UNALLOCATED:
|
||||
await channel.receive_packet(packet)
|
||||
|
||||
|
||||
async def _handle_unallocated(iface, cid) -> None:
|
||||
data = (ThpErrorType.UNALLOCATED_CHANNEL).to_bytes(1, "big")
|
||||
header = PacketHeader.get_error_header(cid, len(data) + CHECKSUM_LENGTH)
|
||||
await write_payload_to_wire_and_add_checksum(iface, header, data)
|
||||
|
||||
|
||||
def _get_buffer_for_payload(
|
||||
payload_length: int, existing_buffer: utils.BufferType, max_length=MAX_PAYLOAD_LEN
|
||||
) -> utils.BufferType:
|
||||
if payload_length > max_length:
|
||||
raise ThpError("Message too large")
|
||||
if payload_length > len(existing_buffer):
|
||||
return _try_allocate_new_buffer(payload_length)
|
||||
return _reuse_existing_buffer(payload_length, existing_buffer)
|
||||
|
||||
|
||||
def _try_allocate_new_buffer(payload_length: int) -> utils.BufferType:
|
||||
try:
|
||||
payload: utils.BufferType = bytearray(payload_length)
|
||||
except MemoryError:
|
||||
payload = bytearray(PACKET_LENGTH)
|
||||
raise ThpError("Message too large")
|
||||
return payload
|
||||
|
||||
|
||||
def _reuse_existing_buffer(
|
||||
payload_length: int, existing_buffer: utils.BufferType
|
||||
) -> utils.BufferType:
|
||||
return memoryview(existing_buffer)[:payload_length]
|
||||
@ -0,0 +1,42 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
declare -a results
|
||||
declare -i passed=0 failed=0 exit_code=0
|
||||
declare COLOR_GREEN='\e[32m' COLOR_RED='\e[91m' COLOR_RESET='\e[39m'
|
||||
MICROPYTHON="${MICROPYTHON:-../build/unix/trezor-emu-core -X heapsize=2M}"
|
||||
print_summary() {
|
||||
echo
|
||||
echo 'Summary:'
|
||||
echo '-------------------'
|
||||
printf '%b\n' "${results[@]}"
|
||||
if [ $exit_code == 0 ]; then
|
||||
echo -e "${COLOR_GREEN}PASSED:${COLOR_RESET} $passed/$num_of_tests tests OK!"
|
||||
else
|
||||
echo -e "${COLOR_RED}FAILED:${COLOR_RESET} $failed/$num_of_tests tests failed!"
|
||||
fi
|
||||
}
|
||||
|
||||
trap 'print_summary; echo -e "${COLOR_RED}Interrupted by user!${COLOR_RESET}"; exit 1' SIGINT
|
||||
|
||||
cd $(dirname $0)
|
||||
|
||||
[ -z "$*" ] && tests=(test_trezor.wire.t*.py ) || tests=($*)
|
||||
|
||||
declare -i num_of_tests=${#tests[@]}
|
||||
|
||||
for test_case in ${tests[@]}; do
|
||||
echo ${MICROPYTHON}
|
||||
echo ${test_case}
|
||||
echo
|
||||
if $MICROPYTHON $test_case; then
|
||||
results+=("${COLOR_GREEN}OK:${COLOR_RESET} $test_case")
|
||||
((passed++))
|
||||
else
|
||||
results+=("${COLOR_RED}FAIL:${COLOR_RESET} $test_case")
|
||||
((failed++))
|
||||
exit_code=1
|
||||
fi
|
||||
done
|
||||
|
||||
print_summary
|
||||
exit $exit_code
|
||||
@ -0,0 +1,95 @@
|
||||
from common import *
|
||||
from trezor import utils
|
||||
|
||||
if utils.USE_THP:
|
||||
from trezor.wire.thp import checksum
|
||||
|
||||
|
||||
@unittest.skipUnless(utils.USE_THP, "only needed for THP")
|
||||
class TestTrezorHostProtocolChecksum(unittest.TestCase):
|
||||
vectors_correct = [
|
||||
(
|
||||
b"",
|
||||
b"\x00\x00\x00\x00",
|
||||
),
|
||||
(
|
||||
b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00",
|
||||
b"\x19\x0A\x55\xAD",
|
||||
),
|
||||
(
|
||||
bytes("a", "ascii"),
|
||||
b"\xE8\xB7\xBE\x43",
|
||||
),
|
||||
(
|
||||
bytes("abc", "ascii"),
|
||||
b"\x35\x24\x41\xC2",
|
||||
),
|
||||
(
|
||||
bytes("123456789", "ascii"),
|
||||
b"\xCB\xF4\x39\x26",
|
||||
),
|
||||
(
|
||||
bytes(
|
||||
"12345678901234567890123456789012345678901234567890123456789012345678901234567890",
|
||||
"ascii",
|
||||
),
|
||||
b"\x7C\xA9\x4A\x72",
|
||||
),
|
||||
(
|
||||
b"\x76\x61\x72\x69\x6F\x75\x73\x20\x43\x52\x43\x20\x61\x6C\x67\x6F\x72\x69\x74\x68\x6D\x73\x20\x69\x6E\x70\x75\x74\x20\x64\x61\x74\x61",
|
||||
b"\x9B\xD3\x66\xAE",
|
||||
),
|
||||
(
|
||||
b"\x67\x3a\x5f\x0e\x39\xc0\x3c\x79\x58\x22\x74\x76\x64\x9e\x36\xe9\x0b\x04\x8c\xd2\xc0\x4d\x76\x63\x1a\xa2\x17\x85\xe8\x50\xa7\x14\x18\xfb\x86\xed\xa3\x59\x2d\x62\x62\x49\x64\x62\x26\x12\xdb\x95\x3d\xd6\xb5\xca\x4b\x22\x0d\xc5\x78\xb2\x12\x97\x8e\x54\x4e\x06\xb7\x9c\x90\xf5\xa0\x21\xa6\xc7\xd8\x39\xfd\xea\x3a\xf1\x7b\xa2\xe8\x71\x41\xd6\xcb\x1e\x5b\x0e\x29\xf7\x0c\xc7\x57\x8b\x53\x20\x1d\x2b\x41\x1c\x25\xf9\x07\xbb\xb4\x37\x79\x6a\x13\x1f\x6c\x43\x71\xc1\x1e\x70\xe6\x74\xd3\x9c\xbf\x32\x15\xee\xf2\xa7\x86\xbe\x59\x99\xc4\x10\x09\x8a\x6a\xaa\xd4\xd1\xd0\x71\xd2\x06\x1a\xdd\x2a\xa0\x08\xeb\x08\x6c\xfb\xd2\x2d\xfb\xaa\x72\x56\xeb\xd1\x92\x92\xe5\x0e\x95\x67\xf8\x38\xc3\xab\x59\x37\xe6\xfd\x42\xb0\xd0\x31\xd0\xcb\x8a\x66\xce\x2d\x53\x72\x1e\x72\xd3\x84\x25\xb0\xb8\x93\xd2\x61\x5b\x32\xd5\xe7\xe4\x0e\x31\x11\xaf\xdc\xb4\xb8\xee\xa4\x55\x16\x5f\x78\x86\x8b\x50\x4d\xc5\x6d\x6e\xfc\xe1\x6b\x06\x5b\x37\x84\x2a\x67\x95\x28\x00\xa4\xd1\x32\x9f\xbf\xe1\x64\xf8\x17\x47\xe1\xad\x8b\x72\xd2\xd9\x45\x5b\x73\x43\x3c\xe6\x21\xf7\x53\xa3\x73\xf9\x2a\xb0\xe9\x75\x5e\xa6\xbe\x9a\xad\xfc\xed\xb5\x46\x5b\x9f\xa9\x5a\x4f\xcb\xb6\x60\x96\x31\x91\x42\xca\xaf\xee\xa5\x0c\xe0\xab\x3e\x83\xb8\xac\x88\x10\x2c\x63\xd3\xc9\xd2\xf2\x44\xef\xea\x3d\x19\x24\x3c\x5b\xe7\x0c\x52\xfd\xfe\x47\x41\x14\xd5\x4c\x67\x8d\xdb\xe5\xd9\xfa\x67\x9c\x06\x31\x01\x92\xba\x96\xc4\x0d\xef\xf7\xc1\xe9\x23\x28\x0f\xae\x27\x9b\xff\x28\x0b\x3e\x85\x0c\xae\x02\xda\x27\xb6\x04\x51\x04\x43\x04\x99\x8c\xa3\x97\x1d\x84\xec\x55\x59\xfb\xf3\x84\xe5\xf8\x40\xf8\x5f\x81\x65\x92\x4c\x92\x7a\x07\x51\x8d\x6f\xff\x8d\x15\x36\x5c\x57\x7a\x5b\x3a\x63\x1c\x87\x65\xee\x54\xd5\x96\x50\x73\x1a\x9c\xff\x59\xe5\xea\x6f\x89\xd2\xbb\xa9\x6a\x12\x21\xf5\x08\x8e\x8a\xc0\xd8\xf5\x14\xe9\x9d\x7e\x99\x13\x88\x29\xa8\xb4\x22\x2a\x41\x7c\xc5\x10\xdf\x11\x5e\xf8\x8d\x0e\xd9\x98\xd5\xaf\xa8\xf9\x55\x1e\xe3\x29\xcd\x2c\x51\x7b\x8a\x8d\x52\xaa\x8b\x87\xae\x8e\xb2\xfa\x31\x27\x60\x90\xcb\x01\x6f\x7a\x79\x38\x04\x05\x7c\x11\x79\x10\x40\x33\x70\x75\xfd\x0b\x88\xa5\xcd\x35\xd8\xa6\x3b\xb0\x45\x82\x64\xd1\xb5\xdc\x06\xc9\x89\xf4\x16\x3e\xc7\xb3\xf1\x9d\xd3\xc5\xe3\xaf\xe8\x25\x86\x7a\x4a\xfd\x10\x5d\x20\xe5\x76\x5a\x22\x5f\x8f\xbc\xaa\x97\xee\xf2\xc2\x4c\x0e\xdc\x7b\xc4\xee\x53\xa3\xe0\xfa\xcd\x1e\x4e\x54\x1d\x5e\xe1\x51\x17\x1f\x1a\x75\x7f\xed\x12\xd7\xf7\xe3\x18\x56\x24\xcf\xc6\x96\x30\x77\x0d\x73\x98\x9c\x09\x69\xa3\xbc\x96\x5e\xaf\xde\x76\xa4\x66\x04\x6b\x36\x2a\xac\x6d\x37\xf8\x1e\xe1\x2a\x3e\x42\x2d\x1d\xe6\x46\xdd\x28\xb9\x08\x44\xa1\x9e\xb2\x22\x7a\x45\x8a\x37\x39\x74\xb4\xae\xc8\x3b\x40\xf7\xec\xbf\xfd\xe5\xde\xb2\x83\x5e\xa4\x46\x19\xa6\x9d\xb0\xe8\x76\x80\xbd\xc1\x80\x7a\xd9\xeb\xe7\x90\x5b\x81\x25\x21\xd9\x5b\x4a\x80\x48\x92\x71\x77\x04\xb2\xac\x05\xc9\xdf\x5e\x44\x5a\xae\x6e\xb3\xd8\x30\x5e\xdc\x77\x2f\x79\xc2\x8e\x8b\x28\x24\x06\x1b\x6f\x8d\x88\x53\x80\x55\x0c\x3a\x7b\x85\xb8\x96\x85\xe9\xf0\x57\x63\xfe\x32\x80\xff\x57\xc9\x3c\xdb\xf6\xcd\x67\x14\x47\x6c\x43\x3d\x6d\x48\x3f\x9c\x00\x60\x0e\xf5\x94\xe4\x52\x97\x86\xcd\xac\xbc\xe4\xe3\xe7\xee\xa2\x91\x6e\x92\xbb\xd1\x55\x0c\x5c\x0d\x63\xdb\x6b\xb8\x6e\x45\x48\x0f\xdf\x44\x48\xd2\xf5\xf7\x4d\x7b\xd4\x4d\xd3\xcd\xcd\x5b\x40\x60\xb1\xb2\x8e\xc9\x9a\x65\xc5\x06\x24\xcf\xe9\xcc\x5e\x2c\x49\x47\x38\x45\x5d\xc5\xc0\x0d\x8a\x07\x1c\xb3\xbb\xb1\x69\xf5\x6d\x0e\x9c\x96\x14\x93\x58\x0c\xc9\x48\x74\xfc\x35\xda\x7d\x4e\x32\x73\xa3\x77\x4a\x9e\xc5\xd1\x08\xfe\xa6\xa0\xf1\x66\x72\xea\xc7\xae\x21\x81\x0e\x8a\xba\x99\x06\x97\xfc\xc6\x2b\x69\x53\xc6\x67\xec\x5d\xa1\xfc\xa1\x3b\xdd\x2a\xd6\x8f\x31\xa7\x8d\xec\xfe\x0a\x3b\x6b\x39\x70\x70\x09\x72\x12\xbc\x84\x67\xca\xd2\x4a\x17\x33\x94\x45\x25\xc7\xfd\x1e\xa2\x4a\x9e\x27\x9d\xfb\x87\xea\xe4\xfd\xb0\x11\x06\x9d\x72\xb9\x1d\xea\x9b\x81\x2e\x6a\x36\x76\x62\xfa\xbe\x96\x67\x7d\x35\xdd\x5e\x5c\x4f\x41\x0d\xce\xdb\x13\xb0\x46\x89\x92\x45\x02\x39\x0f\xe6\xd1\x20\x96\x1c\x34\x00\x8c\xc9\xdf\xe3\xf0\xb6\x92\x3a\xda\x5c\x96\xd9\x0b\x7d\x57\xf5\x78\x11\xc0\xcf\xbf\xb0\x92\x3d\xe5\x6a\x67\x34\xce\xd9\x16\x08\xa0\x09\x42\x0b\x07\x13\x7c\x73\x0c\xc6\x50\x17\x42\xcf\xd9\x85\xd9\x23\x3c\xb1\x40\x40\x0f\x94\x20\xed\x2d\xbf\x10\x44\x6e\x64\x65\xe5\x1d\x5f\xec\x24\xd8\x4b\xe8\xc2\xfb\x06\x11\x24\x3f\xdf\x54\x2d\xe8\x4d\xc2\x1c\x27\x11\xb8\xb3\xd4",
|
||||
b"\x6B\xA4\xEC\x92",
|
||||
),
|
||||
]
|
||||
vectors_incorrect = [
|
||||
(
|
||||
b"",
|
||||
b"\x00\x00\x00\x00\x00",
|
||||
),
|
||||
(
|
||||
b"",
|
||||
b"",
|
||||
),
|
||||
(
|
||||
b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00",
|
||||
b"\x19\x0A\x55\xAE",
|
||||
),
|
||||
(
|
||||
bytes("A", "ascii"),
|
||||
b"\xE8\xB7\xBE\x43",
|
||||
),
|
||||
(
|
||||
bytes("abc ", "ascii"),
|
||||
b"\x35\x24\x41\xC2",
|
||||
),
|
||||
(
|
||||
bytes("1234567890", "ascii"),
|
||||
b"\xCB\xF4\x39\x26",
|
||||
),
|
||||
(
|
||||
bytes(
|
||||
"1234567890123456789012345678901234567890123456789012345678901234567890123456789",
|
||||
"ascii",
|
||||
),
|
||||
b"\x7C\xA9\x4A\x72",
|
||||
),
|
||||
]
|
||||
|
||||
def test_computation(self):
|
||||
for data, chksum in self.vectors_correct:
|
||||
self.assertEqual(checksum.compute(data), chksum)
|
||||
|
||||
def test_validation_correct(self):
|
||||
for data, chksum in self.vectors_correct:
|
||||
self.assertTrue(checksum.is_valid(chksum, data))
|
||||
|
||||
def test_validation_incorrect(self):
|
||||
for data, chksum in self.vectors_incorrect:
|
||||
self.assertFalse(checksum.is_valid(chksum, data))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@ -0,0 +1,63 @@
|
||||
from common import *
|
||||
from trezor import config, utils
|
||||
from trezor import log
|
||||
|
||||
if utils.USE_THP:
|
||||
from apps.thp import credential_manager
|
||||
from trezor.messages import ThpCredentialMetadata
|
||||
|
||||
def _issue_credential(host_name: str, host_static_pubkey: bytes) -> bytes:
|
||||
metadata = ThpCredentialMetadata(host_name=host_name)
|
||||
return credential_manager.issue_credential(host_static_pubkey, metadata)
|
||||
|
||||
def _dummy_log(name: str, msg: str, *args):
|
||||
pass
|
||||
|
||||
log.debug = _dummy_log
|
||||
|
||||
|
||||
@unittest.skipUnless(utils.USE_THP, "only needed for THP")
|
||||
class TestTrezorHostProtocolCredentialManager(unittest.TestCase):
|
||||
def setUp(self):
|
||||
config.init()
|
||||
config.wipe()
|
||||
|
||||
def test_derive_cred_auth_key(self):
|
||||
key1 = credential_manager.derive_cred_auth_key()
|
||||
key2 = credential_manager.derive_cred_auth_key()
|
||||
self.assertEqual(len(key1), 32)
|
||||
self.assertEqual(key1, key2)
|
||||
|
||||
def test_invalidate_cred_auth_key(self):
|
||||
key1 = credential_manager.derive_cred_auth_key()
|
||||
credential_manager.invalidate_cred_auth_key()
|
||||
key2 = credential_manager.derive_cred_auth_key()
|
||||
self.assertNotEqual(key1, key2)
|
||||
|
||||
def test_credentials(self):
|
||||
DUMMY_KEY_1 = b"\x00\x00"
|
||||
DUMMY_KEY_2 = b"\xff\xff"
|
||||
HOST_NAME_1 = "host_name"
|
||||
HOST_NAME_2 = "different host_name"
|
||||
|
||||
cred_1 = _issue_credential(HOST_NAME_1, DUMMY_KEY_1)
|
||||
cred_2 = _issue_credential(HOST_NAME_1, DUMMY_KEY_1)
|
||||
self.assertEqual(cred_1, cred_2)
|
||||
|
||||
cred_3 = _issue_credential(HOST_NAME_2, DUMMY_KEY_1)
|
||||
self.assertNotEqual(cred_1, cred_3)
|
||||
|
||||
self.assertTrue(credential_manager.validate_credential(cred_1, DUMMY_KEY_1))
|
||||
self.assertTrue(credential_manager.validate_credential(cred_3, DUMMY_KEY_1))
|
||||
self.assertFalse(credential_manager.validate_credential(cred_1, DUMMY_KEY_2))
|
||||
|
||||
credential_manager.invalidate_cred_auth_key()
|
||||
cred_4 = _issue_credential(HOST_NAME_1, DUMMY_KEY_1)
|
||||
self.assertNotEqual(cred_1, cred_4)
|
||||
self.assertFalse(credential_manager.validate_credential(cred_1, DUMMY_KEY_1))
|
||||
self.assertFalse(credential_manager.validate_credential(cred_3, DUMMY_KEY_1))
|
||||
self.assertTrue(credential_manager.validate_credential(cred_4, DUMMY_KEY_1))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@ -0,0 +1,150 @@
|
||||
from common import *
|
||||
import storage
|
||||
from trezor import utils
|
||||
from trezor.wire.thp.crypto import IV_1, IV_2, Handshake
|
||||
from trezorcrypto import aesgcm, curve25519
|
||||
|
||||
if utils.USE_THP:
|
||||
from trezor.wire.thp import crypto
|
||||
|
||||
|
||||
def get_dummy_device_secret():
|
||||
return b"\x01\x02\x03\x04\x05\x06\x07\x08\x01\x02\x03\x04\x05\x06\x07\x08"
|
||||
|
||||
|
||||
@unittest.skipUnless(utils.USE_THP, "only needed for THP")
|
||||
class TestTrezorHostProtocolCrypto(unittest.TestCase):
|
||||
key_1 = b"\x00\x01\x02\x03\x04\x05\x06\x07\x00\x01\x02\x03\x04\x05\x06\x07\x00\x01\x02\x03\x04\x05\x06\x07\x00\x01\x02\x03\x04\x05\x06\x07"
|
||||
handshake = Handshake()
|
||||
# 0:key, 1:nonce, 2:auth_data, 3:plaintext, 4:expected_ciphertext, 5:expected_tag
|
||||
vectors_enc = [
|
||||
(
|
||||
key_1,
|
||||
0,
|
||||
b"\x55\x64",
|
||||
b"\x00\x01\x02\x03\x04\05\x06\x07\x08\x09",
|
||||
b"e2c9dd152fbee5821ea7",
|
||||
b"10625812de81b14a46b9f1e5100a6d0c",
|
||||
),
|
||||
(
|
||||
key_1,
|
||||
1,
|
||||
b"\x55\x64",
|
||||
b"\x00\x01\x02\x03\x04\05\x06\x07\x08\x09",
|
||||
b"79811619ddb07c2b99f8",
|
||||
b"71c6b872cdc499a7e9a3c7441f053214",
|
||||
),
|
||||
(
|
||||
key_1,
|
||||
369,
|
||||
b"\x55\x64",
|
||||
b"\x00\x01\x02\x03\x04\05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f",
|
||||
b"03bd030390f2dfe815a61c2b157a064f",
|
||||
b"c1200f8a7ae9a6d32cef0fff878d55c2",
|
||||
),
|
||||
(
|
||||
key_1,
|
||||
369,
|
||||
b"\x55\x64\x73\x82\x91",
|
||||
b"\x00\x01\x02\x03\x04\05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f",
|
||||
b"03bd030390f2dfe815a61c2b157a064f",
|
||||
b"693ac160cd93a20f7fc255f049d808d0",
|
||||
),
|
||||
]
|
||||
# 0:chaining key, 1:input, 2:output_1, 3:output:2
|
||||
vectors_hkdf = [
|
||||
(
|
||||
crypto.PROTOCOL_NAME,
|
||||
b"\x01\x02",
|
||||
b"c784373a217d6be057cddc6068e6748f255fc8beb6f99b7b90cbc64aad947514",
|
||||
b"12695451e29bf08ffe5e4e6ab734b0c3d7cdd99b16cd409f57bd4eaa874944ba",
|
||||
),
|
||||
(
|
||||
b"\xc7\x84\x37\x3a\x21\x7d\x6b\xe0\x57\xcd\xdc\x60\x68\xe6\x74\x8f\x25\x5f\xc8\xbe\xb6\xf9\x9b\x7b\x90\xcb\xc6\x4a\xad\x94\x75\x14",
|
||||
b"\x31\x41\x59\x26\x52\x12\x34\x56\x78\x89\x04\xaa",
|
||||
b"f88c1e08d5c3bae8f6e4a3d3324c8cbc60a805603e399e69c4bf4eacb27c2f48",
|
||||
b"5f0216bdb7110ee05372286974da8c9c8b96e2efa15b4af430755f462bd79a76",
|
||||
),
|
||||
]
|
||||
vectors_iv = [
|
||||
(0, b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"),
|
||||
(1, b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01"),
|
||||
(7, b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x07"),
|
||||
(1025, b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x04\x01"),
|
||||
(4294967295, b"\x00\x00\x00\x00\x00\x00\x00\x00\xff\xff\xff\xff"),
|
||||
(0xFFFFFFFFFFFFFFFF, b"\x00\x00\x00\x00\xff\xff\xff\xff\xff\xff\xff\xff"),
|
||||
]
|
||||
|
||||
def setUp(self):
|
||||
utils.DISABLE_ENCRYPTION = False
|
||||
|
||||
def test_encryption(self):
|
||||
for v in self.vectors_enc:
|
||||
buffer = bytearray(v[3])
|
||||
tag = crypto.enc(buffer, v[0], v[1], v[2])
|
||||
self.assertEqual(hexlify(buffer), v[4])
|
||||
self.assertEqual(hexlify(tag), v[5])
|
||||
self.assertTrue(crypto.dec(buffer, tag, v[0], v[1], v[2]))
|
||||
self.assertEqual(buffer, v[3])
|
||||
|
||||
def test_hkdf(self):
|
||||
for v in self.vectors_hkdf:
|
||||
ck, k = crypto._hkdf(v[0], v[1])
|
||||
self.assertEqual(hexlify(ck), v[2])
|
||||
self.assertEqual(hexlify(k), v[3])
|
||||
|
||||
def test_iv_from_nonce(self):
|
||||
for v in self.vectors_iv:
|
||||
x = v[0]
|
||||
y = x.to_bytes(8, "big")
|
||||
iv = crypto._get_iv_from_nonce(v[0])
|
||||
self.assertEqual(iv, v[1])
|
||||
with self.assertRaises(AssertionError) as e:
|
||||
iv = crypto._get_iv_from_nonce(0xFFFFFFFFFFFFFFFF + 1)
|
||||
self.assertEqual(e.value.value, "Nonce overflow, terminate the channel")
|
||||
|
||||
def test_incorrect_vectors(self):
|
||||
pass
|
||||
|
||||
def test_th1_crypto(self):
|
||||
storage.device.get_device_secret = get_dummy_device_secret
|
||||
handshake = self.handshake
|
||||
|
||||
host_ephemeral_privkey = curve25519.generate_secret()
|
||||
host_ephemeral_pubkey = curve25519.publickey(host_ephemeral_privkey)
|
||||
handshake.handle_th1_crypto(b"", host_ephemeral_pubkey)
|
||||
|
||||
def test_th2_crypto(self):
|
||||
handshake = self.handshake
|
||||
|
||||
host_static_privkey = curve25519.generate_secret()
|
||||
host_static_pubkey = curve25519.publickey(host_static_privkey)
|
||||
aes_ctx = aesgcm(handshake.k, IV_2)
|
||||
aes_ctx.auth(handshake.h)
|
||||
encrypted_host_static_pubkey = bytearray(
|
||||
aes_ctx.encrypt(host_static_pubkey) + aes_ctx.finish()
|
||||
)
|
||||
|
||||
# Code to encrypt Host's noise encrypted payload correctly:
|
||||
protomsg = bytearray(b"\x10\x02\x10\x03")
|
||||
temp_k = handshake.k
|
||||
temp_h = handshake.h
|
||||
|
||||
temp_h = crypto._hash_of_two(temp_h, encrypted_host_static_pubkey)
|
||||
_, temp_k = crypto._hkdf(
|
||||
handshake.ck,
|
||||
curve25519.multiply(handshake.trezor_ephemeral_privkey, host_static_pubkey),
|
||||
)
|
||||
aes_ctx = aesgcm(temp_k, IV_1)
|
||||
aes_ctx.encrypt_in_place(protomsg)
|
||||
aes_ctx.auth(temp_h)
|
||||
tag = aes_ctx.finish()
|
||||
encrypted_payload = bytearray(protomsg + tag)
|
||||
# end of encrypted payload generation
|
||||
|
||||
handshake.handle_th2_crypto(encrypted_host_static_pubkey, encrypted_payload)
|
||||
self.assertEqual(encrypted_payload[:4], b"\x10\x02\x10\x03")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@ -0,0 +1,260 @@
|
||||
from common import *
|
||||
|
||||
from apps.thp import pairing
|
||||
from storage.cache_common import (
|
||||
CHANNEL_HANDSHAKE_HASH,
|
||||
CHANNEL_KEY_RECEIVE,
|
||||
CHANNEL_KEY_SEND,
|
||||
CHANNEL_NONCE_RECEIVE,
|
||||
CHANNEL_NONCE_SEND,
|
||||
)
|
||||
from trezor.enums import ThpPairingMethod, MessageType
|
||||
from trezor.wire.errors import UnexpectedMessage
|
||||
from trezor.wire.protocol_common import Message
|
||||
from trezor.wire.thp.crypto import Handshake
|
||||
from trezor.wire.thp.pairing_context import PairingContext
|
||||
from trezor.messages import (
|
||||
ThpCodeEntryChallenge,
|
||||
ThpCodeEntryCpaceHost,
|
||||
ThpCodeEntryTag,
|
||||
ThpCredentialRequest,
|
||||
ThpEndRequest,
|
||||
ThpStartPairingRequest,
|
||||
)
|
||||
from trezor import io, config, log, protobuf
|
||||
from trezor.loop import wait
|
||||
from trezor.wire import thp_v3
|
||||
from trezor.wire.thp import interface_manager
|
||||
from storage import cache_thp
|
||||
from trezor.wire.thp import ChannelState
|
||||
from trezor.crypto import elligator2
|
||||
from trezor.crypto.curve import curve25519
|
||||
|
||||
|
||||
# Disable log.debug for the test
|
||||
log.debug = lambda name, msg, *args: None
|
||||
class MockHID:
|
||||
def __init__(self, num):
|
||||
self.num = num
|
||||
self.data = []
|
||||
|
||||
def iface_num(self):
|
||||
return self.num
|
||||
|
||||
def write(self, msg):
|
||||
self.data.append(bytearray(msg))
|
||||
return len(msg)
|
||||
|
||||
def wait_object(self, mode):
|
||||
return wait(mode | self.num)
|
||||
|
||||
|
||||
def dummy_decode_iface(cached_iface: bytes):
|
||||
return MockHID(0xDEADBEEF)
|
||||
|
||||
|
||||
def getBytes(a):
|
||||
return hexlify(a).decode("utf-8")
|
||||
|
||||
|
||||
def get_dummy_key() -> bytes:
|
||||
return b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x10\x01\x02\x03\x04\x05\x06\x07\x08\x09\x20\x01\x02\x03\x04\x05\x06\x07\x08\x09\x30\x31"
|
||||
|
||||
|
||||
class TestTrezorHostProtocol(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.interface = MockHID(0xDEADBEEF)
|
||||
buffer = bytearray(64)
|
||||
thp_v3.set_buffer(buffer)
|
||||
interface_manager.decode_iface = dummy_decode_iface
|
||||
|
||||
def test_simple(self):
|
||||
self.assertTrue(True)
|
||||
|
||||
def test_channel_allocation(self):
|
||||
cid_req = (
|
||||
b"\x40\xff\xff\x00\x0c\x00\x11\x22\x33\x44\x55\x66\x77\x96\x64\x3c\x6c"
|
||||
)
|
||||
expected_response = "41ffff0020001122334455667712340a04543254311000180020032802280328048ed892b3000000000000000000000000000000000000000000000000000000"
|
||||
test_counter = cache_thp.cid_counter + 1
|
||||
self.assertEqual(len(thp_v3._CHANNELS), 0)
|
||||
self.assertFalse(test_counter in thp_v3._CHANNELS)
|
||||
gen = thp_v3.thp_main_loop(self.interface, is_debug_session=True)
|
||||
query = gen.send(None)
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_READ))
|
||||
gen.send(cid_req)
|
||||
gen.send(None)
|
||||
self.assertEqual(
|
||||
getBytes(self.interface.data[-1]),
|
||||
expected_response,
|
||||
)
|
||||
self.assertTrue(test_counter in thp_v3._CHANNELS)
|
||||
self.assertEqual(len(thp_v3._CHANNELS), 1)
|
||||
gen.send(cid_req)
|
||||
gen.send(None)
|
||||
gen.send(cid_req)
|
||||
gen.send(None)
|
||||
|
||||
def test_channel_default_state_is_TH1(self):
|
||||
self.assertEqual(thp_v3._CHANNELS[4660].get_channel_state(), ChannelState.TH1)
|
||||
|
||||
def test_channel_errors(self):
|
||||
gen = thp_v3.thp_main_loop(self.interface, is_debug_session=True)
|
||||
query = gen.send(None)
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_READ))
|
||||
message_to_channel_789a = (
|
||||
b"\x04\x78\x9a\x00\x0c\x00\x11\x22\x33\x44\x55\x66\x77\x96\x64\x3c\x6c"
|
||||
)
|
||||
gen.send(message_to_channel_789a)
|
||||
gen.send(None)
|
||||
unallocated_chanel_error_on_channel_789a = "42789a0005027b743563000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
|
||||
self.assertEqual(
|
||||
getBytes(self.interface.data[-1]),
|
||||
unallocated_chanel_error_on_channel_789a,
|
||||
)
|
||||
config.init()
|
||||
config.wipe()
|
||||
channel = thp_v3._CHANNELS[4661]
|
||||
channel.iface = self.interface
|
||||
channel.set_channel_state(ChannelState.ENCRYPTED_TRANSPORT)
|
||||
message_with_invalid_tag = b"\x04\x12\x35\x00\x14\x00\x11\x22\x33\x44\x55\x66\x77\x00\x11\x22\x33\x44\x55\x66\x77\xe1\xfc\xc6\xe0"
|
||||
|
||||
channel.channel_cache.set(CHANNEL_KEY_RECEIVE, get_dummy_key())
|
||||
channel.channel_cache.set_int(CHANNEL_NONCE_RECEIVE, 0)
|
||||
channel.channel_cache.set(CHANNEL_HANDSHAKE_HASH, b"")
|
||||
|
||||
gen.send(message_with_invalid_tag)
|
||||
gen.send(None)
|
||||
ack_on_received_message = "2012350004d83ea46f00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
|
||||
self.assertEqual(
|
||||
getBytes(self.interface.data[-1]),
|
||||
ack_on_received_message,
|
||||
)
|
||||
gen.send(None)
|
||||
decryption_failed_error_on_channel_1235 = "421235000503caf9634a000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
|
||||
self.assertEqual(
|
||||
getBytes(self.interface.data[-1]),
|
||||
decryption_failed_error_on_channel_1235,
|
||||
)
|
||||
|
||||
channel = thp_v3._CHANNELS[4662]
|
||||
channel.iface = self.interface
|
||||
|
||||
channel.set_channel_state(ChannelState.TH2)
|
||||
|
||||
message_with_invalid_tag = b"\x0a\x12\x36\x00\x14\x00\x11\x22\x33\x44\x55\x66\x77\x00\x11\x22\x33\x44\x55\x66\x77\x91\x65\x4c\xf9"
|
||||
|
||||
channel.channel_cache.set(CHANNEL_KEY_RECEIVE, get_dummy_key())
|
||||
channel.channel_cache.set_int(CHANNEL_NONCE_RECEIVE, 0)
|
||||
channel.channel_cache.set(CHANNEL_HANDSHAKE_HASH, b"")
|
||||
|
||||
# gen.send(message_with_invalid_tag)
|
||||
# gen.send(None)
|
||||
# gen.send(None)
|
||||
for i in self.interface.data:
|
||||
print(hexlify(i))
|
||||
|
||||
def test_skip_pairing(self):
|
||||
config.init()
|
||||
config.wipe()
|
||||
channel = thp_v3._CHANNELS[4660]
|
||||
channel.selected_pairing_methods = [
|
||||
ThpPairingMethod.NoMethod,
|
||||
ThpPairingMethod.CodeEntry,
|
||||
ThpPairingMethod.NFC_Unidirectional,
|
||||
ThpPairingMethod.QrCode,
|
||||
]
|
||||
pairing_ctx = PairingContext(channel)
|
||||
request_message = ThpStartPairingRequest()
|
||||
channel.set_channel_state(ChannelState.TP1)
|
||||
gen = pairing.handle_pairing_request(pairing_ctx, request_message)
|
||||
|
||||
with self.assertRaises(StopIteration):
|
||||
gen.send(None)
|
||||
self.assertEqual(channel.get_channel_state(), ChannelState.ENCRYPTED_TRANSPORT)
|
||||
|
||||
# Teardown: set back initial channel state value
|
||||
channel.set_channel_state(ChannelState.TH1)
|
||||
|
||||
def test_pairing(self):
|
||||
config.init()
|
||||
config.wipe()
|
||||
channel = thp_v3._CHANNELS[4660]
|
||||
channel.selected_pairing_methods = [
|
||||
ThpPairingMethod.CodeEntry,
|
||||
ThpPairingMethod.NFC_Unidirectional,
|
||||
ThpPairingMethod.QrCode,
|
||||
]
|
||||
pairing_ctx = PairingContext(channel)
|
||||
request_message = ThpStartPairingRequest()
|
||||
with self.assertRaises(UnexpectedMessage) as e:
|
||||
pairing.handle_pairing_request(pairing_ctx, request_message)
|
||||
print(e.value.message)
|
||||
channel.set_channel_state(ChannelState.TP1)
|
||||
gen = pairing.handle_pairing_request(pairing_ctx, request_message)
|
||||
|
||||
channel.channel_cache.set(CHANNEL_KEY_SEND, get_dummy_key())
|
||||
channel.channel_cache.set_int(CHANNEL_NONCE_SEND, 0)
|
||||
channel.channel_cache.set(CHANNEL_HANDSHAKE_HASH, b"")
|
||||
|
||||
gen.send(None)
|
||||
|
||||
async def _dummy(ctx: PairingContext, expected_types):
|
||||
return await ctx.read([1018, 1024])
|
||||
|
||||
pairing.show_display_data = _dummy
|
||||
|
||||
msg_code_entry = ThpCodeEntryChallenge(challenge=b"\x12\x34")
|
||||
buffer: bytearray = bytearray(protobuf.encoded_length(msg_code_entry))
|
||||
protobuf.encode(buffer, msg_code_entry)
|
||||
code_entry_challenge = Message(MessageType.ThpCodeEntryChallenge, buffer)
|
||||
gen.send(code_entry_challenge)
|
||||
|
||||
# tag_qrc = b"\x55\xdf\x6c\xba\x0b\xe9\x5e\xd1\x4b\x78\x61\xec\xfa\x07\x9b\x5d\x37\x60\xd8\x79\x9c\xd7\x89\xb4\x22\xc1\x6f\x39\xde\x8f\x3b\xc3"
|
||||
# tag_nfc = b"\x8f\xf0\xfa\x37\x0a\x5b\xdb\x29\x32\x21\xd8\x2f\x95\xdd\xb6\xb8\xee\xfd\x28\x6f\x56\x9f\xa9\x0b\x64\x8c\xfc\x62\x46\x5a\xdd\xd0"
|
||||
|
||||
pregenerator_host = b"\xf6\x94\xc3\x6f\xb3\xbd\xfb\xba\x2f\xfd\x0c\xd0\x71\xed\x54\x76\x73\x64\x37\xfa\x25\x85\x12\x8d\xcf\xb5\x6c\x02\xaf\x9d\xe8\xbe"
|
||||
generator_host = elligator2.map_to_curve25519(pregenerator_host)
|
||||
cpace_host_private_key = b"\x02\x80\x70\x3c\x06\x45\x19\x75\x87\x0c\x82\xe1\x64\x11\xc0\x18\x13\xb2\x29\x04\xb3\xf0\xe4\x1e\x6b\xfd\x77\x63\x11\x73\x07\xa9"
|
||||
cpace_host_public_key: bytes = curve25519.multiply(
|
||||
cpace_host_private_key, generator_host
|
||||
)
|
||||
msg = ThpCodeEntryCpaceHost(cpace_host_public_key=cpace_host_public_key)
|
||||
|
||||
# msg = ThpQrCodeTag(tag=tag_qrc)
|
||||
# msg = ThpNfcUnidirectionalTag(tag=tag_nfc)
|
||||
buffer: bytearray = bytearray(protobuf.encoded_length(msg))
|
||||
|
||||
protobuf.encode(buffer, msg)
|
||||
user_message = Message(MessageType.ThpCodeEntryCpaceHost, buffer)
|
||||
gen.send(user_message)
|
||||
|
||||
tag_ent = b"\xf5\x20\xee\xae\xb8\xa9\x65\x3e\x77\x89\x8f\x81\x8d\x03\x4d\xaa\x93\x79\xc3\xe4\x89\x3c\xb8\x31\x42\xdc\x01\x57\x2d\x5d\x11\xb5"
|
||||
|
||||
msg = ThpCodeEntryTag(tag=tag_ent)
|
||||
|
||||
buffer: bytearray = bytearray(protobuf.encoded_length(msg))
|
||||
|
||||
protobuf.encode(buffer, msg)
|
||||
user_message = Message(MessageType.ThpCodeEntryTag, buffer)
|
||||
gen.send(user_message)
|
||||
|
||||
host_static_pubkey = b"\x00\x11\x22\x33\x44\x55\x66\x77\x00\x11\x22\x33\x44\x55\x66\x77\x00\x11\x22\x33\x44\x55\x66\x77\x00\x11\x22\x33\x44\x55\x66\x77\x00\x11\x22\x33\x44\x55\x66\x77\x00\x11\x22\x33\x44\x55\x66\x77"
|
||||
msg = ThpCredentialRequest(host_static_pubkey=host_static_pubkey)
|
||||
buffer: bytearray = bytearray(protobuf.encoded_length(msg))
|
||||
protobuf.encode(buffer, msg)
|
||||
credential_request = Message(MessageType.ThpCredentialRequest, buffer)
|
||||
gen.send(credential_request)
|
||||
|
||||
msg = ThpEndRequest()
|
||||
|
||||
buffer: bytearray = bytearray(protobuf.encoded_length(msg))
|
||||
protobuf.encode(buffer, msg)
|
||||
end_request = Message(1012, buffer)
|
||||
with self.assertRaises(StopIteration) as e:
|
||||
gen.send(end_request)
|
||||
print("response message:", e.value.value.MESSAGE_NAME)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@ -0,0 +1,167 @@
|
||||
from common import *
|
||||
from trezor import utils
|
||||
|
||||
if utils.USE_THP:
|
||||
from trezor.wire.thp import writer
|
||||
from trezor.wire.thp.thp_messages import PacketHeader, ENCRYPTED_TRANSPORT
|
||||
|
||||
if __debug__:
|
||||
# Disable log.debug for the test
|
||||
from trezor import log
|
||||
|
||||
log.debug = lambda name, msg, *args: None
|
||||
|
||||
|
||||
class MockHID:
|
||||
def __init__(self, num):
|
||||
self.num = num
|
||||
self.data = []
|
||||
|
||||
def iface_num(self):
|
||||
return self.num
|
||||
|
||||
def write(self, msg):
|
||||
self.data.append(bytearray(msg))
|
||||
return len(msg)
|
||||
|
||||
def wait_object(self, mode):
|
||||
return wait(mode | self.num)
|
||||
|
||||
|
||||
@unittest.skipUnless(utils.USE_THP, "only needed for THP")
|
||||
class TestTrezorHostProtocolWriter(unittest.TestCase):
|
||||
short_payload_expected = b"04123400050700000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
|
||||
longer_payload_expected = [
|
||||
b"0412340100000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132333435363738393a",
|
||||
b"8012343b3c3d3e3f404142434445464748494a4b4c4d4e4f505152535455565758595a5b5c5d5e5f606162636465666768696a6b6c6d6e6f7071727374757677",
|
||||
b"80123478797a7b7c7d7e7f808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9fa0a1a2a3a4a5a6a7a8a9aaabacadaeafb0b1b2b3b4",
|
||||
b"801234b5b6b7b8b9babbbcbdbebfc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedfe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1",
|
||||
b"801234f2f3f4f5f6f7f8f9fafbfcfdfeff0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
|
||||
]
|
||||
eight_longer_payloads_expected = [
|
||||
b"0412340800000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132333435363738393a",
|
||||
b"8012343b3c3d3e3f404142434445464748494a4b4c4d4e4f505152535455565758595a5b5c5d5e5f606162636465666768696a6b6c6d6e6f7071727374757677",
|
||||
b"80123478797a7b7c7d7e7f808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9fa0a1a2a3a4a5a6a7a8a9aaabacadaeafb0b1b2b3b4",
|
||||
b"801234b5b6b7b8b9babbbcbdbebfc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedfe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1",
|
||||
b"801234f2f3f4f5f6f7f8f9fafbfcfdfeff000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e",
|
||||
b"8012342f303132333435363738393a3b3c3d3e3f404142434445464748494a4b4c4d4e4f505152535455565758595a5b5c5d5e5f606162636465666768696a6b",
|
||||
b"8012346c6d6e6f707172737475767778797a7b7c7d7e7f808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9fa0a1a2a3a4a5a6a7a8",
|
||||
b"801234a9aaabacadaeafb0b1b2b3b4b5b6b7b8b9babbbcbdbebfc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedfe0e1e2e3e4e5",
|
||||
b"801234e6e7e8e9eaebecedeeeff0f1f2f3f4f5f6f7f8f9fafbfcfdfeff000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122",
|
||||
b"801234232425262728292a2b2c2d2e2f303132333435363738393a3b3c3d3e3f404142434445464748494a4b4c4d4e4f505152535455565758595a5b5c5d5e5f",
|
||||
b"801234606162636465666768696a6b6c6d6e6f707172737475767778797a7b7c7d7e7f808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c",
|
||||
b"8012349d9e9fa0a1a2a3a4a5a6a7a8a9aaabacadaeafb0b1b2b3b4b5b6b7b8b9babbbcbdbebfc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9",
|
||||
b"801234dadbdcdddedfe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1f2f3f4f5f6f7f8f9fafbfcfdfeff000102030405060708090a0b0c0d0e0f10111213141516",
|
||||
b"8012341718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132333435363738393a3b3c3d3e3f404142434445464748494a4b4c4d4e4f50515253",
|
||||
b"8012345455565758595a5b5c5d5e5f606162636465666768696a6b6c6d6e6f707172737475767778797a7b7c7d7e7f808182838485868788898a8b8c8d8e8f90",
|
||||
b"8012349192939495969798999a9b9c9d9e9fa0a1a2a3a4a5a6a7a8a9aaabacadaeafb0b1b2b3b4b5b6b7b8b9babbbcbdbebfc0c1c2c3c4c5c6c7c8c9cacbcccd",
|
||||
b"801234cecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedfe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1f2f3f4f5f6f7f8f9fafbfcfdfeff000102030405060708090a",
|
||||
b"8012340b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132333435363738393a3b3c3d3e3f4041424344454647",
|
||||
b"80123448494a4b4c4d4e4f505152535455565758595a5b5c5d5e5f606162636465666768696a6b6c6d6e6f707172737475767778797a7b7c7d7e7f8081828384",
|
||||
b"80123485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9fa0a1a2a3a4a5a6a7a8a9aaabacadaeafb0b1b2b3b4b5b6b7b8b9babbbcbdbebfc0c1",
|
||||
b"801234c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedfe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1f2f3f4f5f6f7f8f9fafbfcfdfe",
|
||||
b"801234ff000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132333435363738393a3b",
|
||||
b"8012343c3d3e3f404142434445464748494a4b4c4d4e4f505152535455565758595a5b5c5d5e5f606162636465666768696a6b6c6d6e6f707172737475767778",
|
||||
b"801234797a7b7c7d7e7f808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9fa0a1a2a3a4a5a6a7a8a9aaabacadaeafb0b1b2b3b4b5",
|
||||
b"801234b6b7b8b9babbbcbdbebfc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedfe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1f2",
|
||||
b"801234f3f4f5f6f7f8f9fafbfcfdfeff000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f",
|
||||
b"801234303132333435363738393a3b3c3d3e3f404142434445464748494a4b4c4d4e4f505152535455565758595a5b5c5d5e5f606162636465666768696a6b6c",
|
||||
b"8012346d6e6f707172737475767778797a7b7c7d7e7f808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9fa0a1a2a3a4a5a6a7a8a9",
|
||||
b"801234aaabacadaeafb0b1b2b3b4b5b6b7b8b9babbbcbdbebfc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedfe0e1e2e3e4e5e6",
|
||||
b"801234e7e8e9eaebecedeeeff0f1f2f3f4f5f6f7f8f9fafbfcfdfeff000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20212223",
|
||||
b"8012342425262728292a2b2c2d2e2f303132333435363738393a3b3c3d3e3f404142434445464748494a4b4c4d4e4f505152535455565758595a5b5c5d5e5f60",
|
||||
b"8012346162636465666768696a6b6c6d6e6f707172737475767778797a7b7c7d7e7f808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d",
|
||||
b"8012349e9fa0a1a2a3a4a5a6a7a8a9aaabacadaeafb0b1b2b3b4b5b6b7b8b9babbbcbdbebfc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9da",
|
||||
b"801234dbdcdddedfe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1f2f3f4f5f6f7f8f9fafbfcfdfeff000000000000000000000000000000000000000000000000",
|
||||
]
|
||||
empty_payload_with_checksum_expected = b"0412340004edbd479c00000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"
|
||||
longer_payload_with_checksum_expected = [
|
||||
b"0412340100000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f303132333435363738393a",
|
||||
b"8012343b3c3d3e3f404142434445464748494a4b4c4d4e4f505152535455565758595a5b5c5d5e5f606162636465666768696a6b6c6d6e6f7071727374757677",
|
||||
b"80123478797a7b7c7d7e7f808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9fa0a1a2a3a4a5a6a7a8a9aaabacadaeafb0b1b2b3b4",
|
||||
b"801234b5b6b7b8b9babbbcbdbebfc0c1c2c3c4c5c6c7c8c9cacbcccdcecfd0d1d2d3d4d5d6d7d8d9dadbdcdddedfe0e1e2e3e4e5e6e7e8e9eaebecedeeeff0f1",
|
||||
b"801234f2f3f4f5f6f7f8f9fafbfcfdfefff40c65ee00000000000000000000000000000000000000000000000000000000000000000000000000000000000000",
|
||||
]
|
||||
|
||||
def setUp(self):
|
||||
self.interface = MockHID(0xDEADBEEF)
|
||||
|
||||
def test_write_empty_payload(self):
|
||||
header = PacketHeader(ENCRYPTED_TRANSPORT, 4660, 4)
|
||||
gen = writer.write_payloads_to_wire(self.interface, header, (b"",))
|
||||
|
||||
with self.assertRaises(StopIteration):
|
||||
gen.send(None)
|
||||
self.assertEqual(len(self.interface.data), 0)
|
||||
|
||||
def test_write_short_payload(self):
|
||||
header = PacketHeader(ENCRYPTED_TRANSPORT, 4660, 5)
|
||||
data = b"\x07"
|
||||
gen = writer.write_payloads_to_wire(self.interface, header, (data,))
|
||||
|
||||
gen.send(None)
|
||||
with self.assertRaises(StopIteration):
|
||||
gen.send(None)
|
||||
self.assertEqual(hexlify(self.interface.data[0]), self.short_payload_expected)
|
||||
|
||||
def test_write_longer_payload(self):
|
||||
data = bytearray(range(256))
|
||||
header = PacketHeader(ENCRYPTED_TRANSPORT, 4660, 256)
|
||||
gen = writer.write_payloads_to_wire(self.interface, header, (data,))
|
||||
|
||||
for i in range(5):
|
||||
gen.send(None)
|
||||
with self.assertRaises(StopIteration):
|
||||
gen.send(None)
|
||||
for i in range(len(self.longer_payload_expected)):
|
||||
self.assertEqual(
|
||||
hexlify(self.interface.data[i]), self.longer_payload_expected[i]
|
||||
)
|
||||
|
||||
def test_write_eight_longer_payloads(self):
|
||||
data = bytearray(range(256))
|
||||
header = PacketHeader(ENCRYPTED_TRANSPORT, 4660, 2048)
|
||||
gen = writer.write_payloads_to_wire(
|
||||
self.interface, header, (data, data, data, data, data, data, data, data)
|
||||
)
|
||||
|
||||
for i in range(34):
|
||||
gen.send(None)
|
||||
with self.assertRaises(StopIteration):
|
||||
gen.send(None)
|
||||
for i in range(len(self.eight_longer_payloads_expected)):
|
||||
self.assertEqual(
|
||||
hexlify(self.interface.data[i]), self.eight_longer_payloads_expected[i]
|
||||
)
|
||||
|
||||
def test_write_empty_payload_with_checksum(self):
|
||||
header = PacketHeader(ENCRYPTED_TRANSPORT, 4660, 4)
|
||||
gen = writer.write_payload_to_wire_and_add_checksum(self.interface, header, b"")
|
||||
|
||||
gen.send(None)
|
||||
with self.assertRaises(StopIteration):
|
||||
gen.send(None)
|
||||
self.assertEqual(
|
||||
hexlify(self.interface.data[0]), self.empty_payload_with_checksum_expected
|
||||
)
|
||||
|
||||
def test_write_longer_payload_with_checksum(self):
|
||||
data = bytearray(range(256))
|
||||
header = PacketHeader(ENCRYPTED_TRANSPORT, 4660, 256)
|
||||
gen = writer.write_payload_to_wire_and_add_checksum(
|
||||
self.interface, header, data
|
||||
)
|
||||
|
||||
for i in range(5):
|
||||
gen.send(None)
|
||||
with self.assertRaises(StopIteration):
|
||||
gen.send(None)
|
||||
for i in range(len(self.longer_payload_with_checksum_expected)):
|
||||
self.assertEqual(
|
||||
hexlify(self.interface.data[i]),
|
||||
self.longer_payload_with_checksum_expected[i],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@ -0,0 +1,352 @@
|
||||
from common import *
|
||||
from typing import TYPE_CHECKING
|
||||
from storage.cache_thp import BROADCAST_CHANNEL_ID
|
||||
import trezor.wire.thp
|
||||
from trezor.wire.thp import alternating_bit_protocol as ABP
|
||||
from trezor.wire.thp.writer import PACKET_LENGTH
|
||||
from ubinascii import hexlify
|
||||
import ustruct
|
||||
|
||||
from trezor import io, log, utils
|
||||
from trezor.loop import wait
|
||||
from trezor.utils import chunks
|
||||
from trezor.wire import thp_v3
|
||||
from trezor.wire.protocol_common import Message
|
||||
from trezor.wire.thp import checksum
|
||||
from trezor.wire.thp.checksum import CHECKSUM_LENGTH
|
||||
|
||||
# Disable log.debug for the test
|
||||
log.debug = lambda name, msg, *args: None
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from trezorio import WireInterface
|
||||
|
||||
class MockHID:
|
||||
def __init__(self, num):
|
||||
self.num = num
|
||||
self.data = []
|
||||
|
||||
def iface_num(self):
|
||||
return self.num
|
||||
|
||||
def write(self, msg):
|
||||
self.data.append(bytearray(msg))
|
||||
return len(msg)
|
||||
|
||||
def wait_object(self, mode):
|
||||
return wait(mode | self.num)
|
||||
|
||||
|
||||
MESSAGE_TYPE = 0x4242
|
||||
MESSAGE_TYPE_BYTES = b"\x42\x42"
|
||||
_MESSAGE_TYPE_LEN = 2
|
||||
PLAINTEXT_0 = 0x01
|
||||
PLAINTEXT_1 = 0x11
|
||||
COMMON_CID = 4660
|
||||
CONT = 0x80
|
||||
|
||||
HEADER_INIT_LENGTH = 5
|
||||
HEADER_CONT_LENGTH = 3
|
||||
INIT_MESSAGE_DATA_LENGTH = PACKET_LENGTH - HEADER_INIT_LENGTH - _MESSAGE_TYPE_LEN
|
||||
|
||||
|
||||
def make_header(ctrl_byte, cid, length):
|
||||
return ustruct.pack(">BHH", ctrl_byte, cid, length)
|
||||
|
||||
|
||||
def make_cont_header():
|
||||
return ustruct.pack(">BH", CONT, COMMON_CID)
|
||||
|
||||
|
||||
def makeSimpleMessage(header, message_type, message_data):
|
||||
return header + ustruct.pack(">H", message_type) + message_data
|
||||
|
||||
|
||||
def makeCidRequest(header, message_data):
|
||||
return header + message_data
|
||||
|
||||
|
||||
def printBytes(a):
|
||||
print(hexlify(a).decode("utf-8"))
|
||||
|
||||
|
||||
def getPlaintext() -> bytes:
|
||||
if ABP.get_expected_receive_seq_bit(THP.get_active_session()) == 1:
|
||||
return PLAINTEXT_1
|
||||
return PLAINTEXT_0
|
||||
|
||||
|
||||
async def deprecated_read_message(
|
||||
iface: WireInterface, buffer: utils.BufferType
|
||||
) -> Message:
|
||||
return Message(-1, b"\x00")
|
||||
|
||||
|
||||
async def deprecated_write_message(
|
||||
iface: WireInterface, message: Message, is_retransmission: bool = False
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
|
||||
# This test suite is an adaptation of test_trezor.wire.codec_v1
|
||||
class TestWireTrezorHostProtocolV1(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.interface = MockHID(0xDEADBEEF)
|
||||
if not utils.USE_THP:
|
||||
import storage.cache_thp # noqa: F401
|
||||
|
||||
def _simple(self):
|
||||
cid_req_header = make_header(
|
||||
ctrl_byte=0x40, cid=BROADCAST_CHANNEL_ID, length=12
|
||||
)
|
||||
cid_request_dummy_data = b"\x00\x11\x22\x33\x44\x55\x66\x77\x96\x64\x3c\x6c"
|
||||
cid_req_message = makeCidRequest(cid_req_header, cid_request_dummy_data)
|
||||
|
||||
message_header = make_header(ctrl_byte=0x01, cid=COMMON_CID, length=18)
|
||||
cid_request_dummy_data_checksum = b"\x67\x8e\xac\xe0"
|
||||
message = makeSimpleMessage(
|
||||
message_header,
|
||||
MESSAGE_TYPE,
|
||||
cid_request_dummy_data + cid_request_dummy_data_checksum,
|
||||
)
|
||||
|
||||
buffer = bytearray(64)
|
||||
gen = deprecated_read_message(self.interface, buffer)
|
||||
query = gen.send(None)
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_READ))
|
||||
gen.send(cid_req_message)
|
||||
gen.send(None)
|
||||
gen.send(message)
|
||||
with self.assertRaises(StopIteration) as e:
|
||||
gen.send(None)
|
||||
|
||||
# e.value is StopIteration. e.value.value is the return value of the call
|
||||
result = e.value.value
|
||||
self.assertEqual(result.type, MESSAGE_TYPE)
|
||||
self.assertEqual(result.data, cid_request_dummy_data)
|
||||
|
||||
buffer_without_zeroes = buffer[: len(message) - 5]
|
||||
message_without_header = message[5:]
|
||||
# message should have been read into the buffer
|
||||
self.assertEqual(buffer_without_zeroes, message_without_header)
|
||||
|
||||
def _read_one_packet(self):
|
||||
# zero length message - just a header
|
||||
PLAINTEXT = getPlaintext()
|
||||
header = make_header(
|
||||
PLAINTEXT, cid=COMMON_CID, length=_MESSAGE_TYPE_LEN + CHECKSUM_LENGTH
|
||||
)
|
||||
chksum = checksum.compute(header + MESSAGE_TYPE_BYTES)
|
||||
message = header + MESSAGE_TYPE_BYTES + chksum
|
||||
|
||||
buffer = bytearray(64)
|
||||
gen = deprecated_read_message(self.interface, buffer)
|
||||
|
||||
query = gen.send(None)
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_READ))
|
||||
gen.send(message)
|
||||
with self.assertRaises(StopIteration) as e:
|
||||
gen.send(None)
|
||||
|
||||
# e.value is StopIteration. e.value.value is the return value of the call
|
||||
result = e.value.value
|
||||
self.assertEqual(result.type, MESSAGE_TYPE)
|
||||
self.assertEqual(result.data, b"")
|
||||
|
||||
# message should have been read into the buffer
|
||||
self.assertEqual(buffer, MESSAGE_TYPE_BYTES + chksum + b"\x00" * 58)
|
||||
|
||||
def _read_many_packets(self):
|
||||
message = bytes(range(256))
|
||||
header = make_header(
|
||||
getPlaintext(),
|
||||
COMMON_CID,
|
||||
len(message) + _MESSAGE_TYPE_LEN + CHECKSUM_LENGTH,
|
||||
)
|
||||
chksum = checksum.compute(header + MESSAGE_TYPE_BYTES + message)
|
||||
# message = MESSAGE_TYPE_BYTES + message + checksum
|
||||
|
||||
# first packet is init header + 59 bytes of data
|
||||
# other packets are cont header + 61 bytes of data
|
||||
cont_header = make_cont_header()
|
||||
packets = [header + MESSAGE_TYPE_BYTES + message[:INIT_MESSAGE_DATA_LENGTH]] + [
|
||||
cont_header + chunk
|
||||
for chunk in chunks(
|
||||
message[INIT_MESSAGE_DATA_LENGTH:] + chksum,
|
||||
64 - HEADER_CONT_LENGTH,
|
||||
)
|
||||
]
|
||||
buffer = bytearray(262)
|
||||
gen = deprecated_read_message(self.interface, buffer)
|
||||
query = gen.send(None)
|
||||
for packet in packets:
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_READ))
|
||||
query = gen.send(packet)
|
||||
|
||||
# last packet will stop
|
||||
with self.assertRaises(StopIteration) as e:
|
||||
gen.send(None)
|
||||
|
||||
# e.value is StopIteration. e.value.value is the return value of the call
|
||||
result = e.value.value
|
||||
|
||||
self.assertEqual(result.type, MESSAGE_TYPE)
|
||||
self.assertEqual(result.data, message)
|
||||
|
||||
# message should have been read into the buffer )
|
||||
self.assertEqual(buffer, MESSAGE_TYPE_BYTES + message + chksum)
|
||||
|
||||
def _read_large_message(self):
|
||||
message = b"hello world"
|
||||
header = make_header(
|
||||
getPlaintext(),
|
||||
COMMON_CID,
|
||||
_MESSAGE_TYPE_LEN + len(message) + CHECKSUM_LENGTH,
|
||||
)
|
||||
|
||||
packet = (
|
||||
header
|
||||
+ MESSAGE_TYPE_BYTES
|
||||
+ message
|
||||
+ checksum.compute(header + MESSAGE_TYPE_BYTES + message)
|
||||
)
|
||||
|
||||
# make sure we fit into one packet, to make this easier
|
||||
self.assertTrue(len(packet) <= thp_v3.PACKET_LENGTH)
|
||||
|
||||
buffer = bytearray(1)
|
||||
self.assertTrue(len(buffer) <= len(packet))
|
||||
|
||||
gen = deprecated_read_message(self.interface, buffer)
|
||||
query = gen.send(None)
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_READ))
|
||||
gen.send(packet)
|
||||
with self.assertRaises(StopIteration) as e:
|
||||
gen.send(None)
|
||||
|
||||
# e.value is StopIteration. e.value.value is the return value of the call
|
||||
result = e.value.value
|
||||
self.assertEqual(result.type, MESSAGE_TYPE)
|
||||
self.assertEqual(result.data, message)
|
||||
|
||||
# read should have allocated its own buffer and not touch ours
|
||||
self.assertEqual(buffer, b"\x00")
|
||||
|
||||
def _roundtrip(self):
|
||||
message_payload = bytes(range(256))
|
||||
message = Message(
|
||||
MESSAGE_TYPE, message_payload, 1
|
||||
) # TODO use different session id
|
||||
gen = deprecated_write_message(self.interface, message)
|
||||
# exhaust the iterator:
|
||||
# (XXX we can only do this because the iterator is only accepting None and returns None)
|
||||
for query in gen:
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_WRITE))
|
||||
|
||||
buffer = bytearray(1024)
|
||||
gen = deprecated_read_message(self.interface, buffer)
|
||||
query = gen.send(None)
|
||||
for packet in self.interface.data:
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_READ))
|
||||
printBytes(packet)
|
||||
query = gen.send(packet)
|
||||
|
||||
with self.assertRaises(StopIteration) as e:
|
||||
gen.send(None)
|
||||
|
||||
result = e.value.value
|
||||
self.assertEqual(result.type, MESSAGE_TYPE)
|
||||
self.assertEqual(result.data, message.data)
|
||||
|
||||
def _write_one_packet(self):
|
||||
message = Message(MESSAGE_TYPE, b"")
|
||||
gen = deprecated_write_message(self.interface, message)
|
||||
|
||||
query = gen.send(None)
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_WRITE))
|
||||
with self.assertRaises(StopIteration):
|
||||
gen.send(None)
|
||||
|
||||
header = make_header(
|
||||
getPlaintext(), COMMON_CID, _MESSAGE_TYPE_LEN + CHECKSUM_LENGTH
|
||||
)
|
||||
expected_message = (
|
||||
header
|
||||
+ MESSAGE_TYPE_BYTES
|
||||
+ checksum.compute(header + MESSAGE_TYPE_BYTES)
|
||||
+ b"\x00" * (INIT_MESSAGE_DATA_LENGTH - CHECKSUM_LENGTH)
|
||||
)
|
||||
self.assertTrue(self.interface.data == [expected_message])
|
||||
|
||||
def _write_multiple_packets(self):
|
||||
message_payload = bytes(range(256))
|
||||
message = Message(MESSAGE_TYPE, message_payload)
|
||||
gen = deprecated_write_message(self.interface, message)
|
||||
|
||||
header = make_header(
|
||||
PLAINTEXT_1,
|
||||
COMMON_CID,
|
||||
len(message.data) + _MESSAGE_TYPE_LEN + CHECKSUM_LENGTH,
|
||||
)
|
||||
cont_header = make_cont_header()
|
||||
chksum = checksum.compute(
|
||||
header + message.type.to_bytes(2, "big") + message.data
|
||||
)
|
||||
packets = [
|
||||
header + MESSAGE_TYPE_BYTES + message.data[:INIT_MESSAGE_DATA_LENGTH]
|
||||
] + [
|
||||
cont_header + chunk
|
||||
for chunk in chunks(
|
||||
message.data[INIT_MESSAGE_DATA_LENGTH:] + chksum,
|
||||
thp_v3.PACKET_LENGTH - HEADER_CONT_LENGTH,
|
||||
)
|
||||
]
|
||||
|
||||
for _ in packets:
|
||||
# we receive as many queries as there are packets
|
||||
query = gen.send(None)
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_WRITE))
|
||||
|
||||
# the first sent None only started the generator. the len(packets)-th None
|
||||
# will finish writing and raise StopIteration
|
||||
with self.assertRaises(StopIteration):
|
||||
gen.send(None)
|
||||
|
||||
# packets must be identical up to the last one
|
||||
self.assertListEqual(packets[:-1], self.interface.data[:-1])
|
||||
# last packet must be identical up to message length. remaining bytes in
|
||||
# the 64-byte packets are garbage -- in particular, it's the bytes of the
|
||||
# previous packet
|
||||
last_packet = packets[-1] + packets[-2][len(packets[-1]) :]
|
||||
self.assertEqual(last_packet, self.interface.data[-1])
|
||||
|
||||
def _read_huge_packet(self):
|
||||
PACKET_COUNT = 1180
|
||||
# message that takes up 1 180 USB packets
|
||||
message_size = (PACKET_COUNT - 1) * (
|
||||
PACKET_LENGTH - HEADER_CONT_LENGTH - CHECKSUM_LENGTH - _MESSAGE_TYPE_LEN
|
||||
) + INIT_MESSAGE_DATA_LENGTH
|
||||
|
||||
# ensure that a message this big won't fit into memory
|
||||
# Note: this control is changed, because THP has only 2 byte length field
|
||||
self.assertTrue(message_size > thp_v3.MAX_PAYLOAD_LEN)
|
||||
# self.assertRaises(MemoryError, bytearray, message_size)
|
||||
header = make_header(PLAINTEXT_1, COMMON_CID, message_size)
|
||||
packet = header + MESSAGE_TYPE_BYTES + (b"\x00" * INIT_MESSAGE_DATA_LENGTH)
|
||||
buffer = bytearray(65536)
|
||||
gen = deprecated_read_message(self.interface, buffer)
|
||||
|
||||
query = gen.send(None)
|
||||
|
||||
# THP returns "Message too large" error after reading the message size,
|
||||
# it is different from codec_v1 as it does not allow big enough messages
|
||||
# to raise MemoryError in this test
|
||||
self.assertObjectEqual(query, self.interface.wait_object(io.POLL_READ))
|
||||
with self.assertRaises(trezor.wire.thp.ThpError) as e:
|
||||
query = gen.send(packet)
|
||||
|
||||
self.assertEqual(e.value.args[0], "Message too large")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@ -0,0 +1 @@
|
||||
../../vendor/trezor-common/protob/messages-thp.proto
|
||||
File diff suppressed because it is too large
Load Diff
Loading…
Reference in new issue