diff --git a/core/embed/rust/build.rs b/core/embed/rust/build.rs index ce44f1075..47f90b548 100644 --- a/core/embed/rust/build.rs +++ b/core/embed/rust/build.rs @@ -219,6 +219,7 @@ fn generate_micropython_bindings() { .allowlist_function("mp_obj_new_exception_args") .allowlist_function("trezor_obj_call_protected") .allowlist_var("mp_type_AttributeError") + .allowlist_var("mp_type_EOFError") .allowlist_var("mp_type_IndexError") .allowlist_var("mp_type_KeyError") .allowlist_var("mp_type_MemoryError") diff --git a/core/embed/rust/src/error.rs b/core/embed/rust/src/error.rs index 417b4dfb3..b949f0539 100644 --- a/core/embed/rust/src/error.rs +++ b/core/embed/rust/src/error.rs @@ -14,6 +14,7 @@ pub enum Error { OutOfRange, MissingKwargs, AllocationFailed, + EOFError, IndexError, #[cfg(feature = "micropython")] CaughtException(Obj), @@ -71,6 +72,7 @@ impl Error { Error::AttributeError(attr) => { ffi::mp_obj_new_exception_args(&ffi::mp_type_AttributeError, 1, &attr.into()) } + Error::EOFError => ffi::mp_obj_new_exception(&ffi::mp_type_EOFError), } } } diff --git a/core/embed/rust/src/io.rs b/core/embed/rust/src/io.rs new file mode 100644 index 000000000..dbef3552e --- /dev/null +++ b/core/embed/rust/src/io.rs @@ -0,0 +1,71 @@ +use crate::error::Error; + +pub struct InputStream<'a> { + buf: &'a [u8], + pos: usize, +} + +impl<'a> InputStream<'a> { + pub fn new(buf: &'a [u8]) -> Self { + Self { buf, pos: 0 } + } + + pub fn remaining(&self) -> usize { + self.buf.len().saturating_sub(self.pos) + } + + pub fn tell(&self) -> usize { + self.pos + } + + pub fn read_stream(&mut self, len: usize) -> Result { + self.read(len).map(Self::new) + } + + pub fn read(&mut self, len: usize) -> Result<&'a [u8], Error> { + let buf = self + .buf + .get(self.pos..self.pos + len) + .ok_or(Error::EOFError)?; + self.pos += len; + Ok(buf) + } + + pub fn rest(self) -> &'a [u8] { + if self.pos > self.buf.len() { + &[] + } else { + &self.buf[self.pos..] + } + } + + pub fn read_byte(&mut self) -> Result { + let val = self.buf.get(self.pos).ok_or(Error::EOFError)?; + self.pos += 1; + Ok(*val) + } + + pub fn read_u16_le(&mut self) -> Result { + let buf = self.read(2)?; + Ok(u16::from_le_bytes(unwrap!(buf.try_into()))) + } + + pub fn read_u32_le(&mut self) -> Result { + let buf = self.read(4)?; + Ok(u32::from_le_bytes(unwrap!(buf.try_into()))) + } + + pub fn read_uvarint(&mut self) -> Result { + let mut uint = 0; + let mut shift = 0; + loop { + let byte = self.read_byte()?; + uint += (byte as u64 & 0x7F) << shift; + shift += 7; + if byte & 0x80 == 0 { + break; + } + } + Ok(uint) + } +} diff --git a/core/embed/rust/src/lib.rs b/core/embed/rust/src/lib.rs index e0dee4280..9e45c37ac 100644 --- a/core/embed/rust/src/lib.rs +++ b/core/embed/rust/src/lib.rs @@ -17,6 +17,8 @@ mod error; // use trezorhal for its macros early #[macro_use] mod trezorhal; + +mod io; mod maybe_trace; #[cfg(feature = "micropython")] #[macro_use] diff --git a/core/embed/rust/src/protobuf/decode.rs b/core/embed/rust/src/protobuf/decode.rs index 7d537f450..aadcffadd 100644 --- a/core/embed/rust/src/protobuf/decode.rs +++ b/core/embed/rust/src/protobuf/decode.rs @@ -5,6 +5,7 @@ use core::{ use crate::{ error::Error, + io::InputStream, micropython::{buffer, gc::Gc, list::List, map::Map, obj::Obj, qstr::Qstr, util}, }; @@ -250,56 +251,3 @@ impl Decoder { } } } - -pub struct InputStream<'a> { - buf: &'a [u8], - pos: usize, -} - -impl<'a> InputStream<'a> { - pub fn new(buf: &'a [u8]) -> Self { - Self { buf, pos: 0 } - } - - pub fn read_stream(&mut self, len: usize) -> Result { - let buf = self - .buf - .get(self.pos..self.pos + len) - .ok_or_else(error::end_of_buffer)?; - self.pos += len; - Ok(Self::new(buf)) - } - - pub fn read(&mut self, len: usize) -> Result<&[u8], Error> { - let buf = self - .buf - .get(self.pos..self.pos + len) - .ok_or_else(error::end_of_buffer)?; - self.pos += len; - Ok(buf) - } - - pub fn read_byte(&mut self) -> Result { - let val = self - .buf - .get(self.pos) - .copied() - .ok_or_else(error::end_of_buffer)?; - self.pos += 1; - Ok(val) - } - - pub fn read_uvarint(&mut self) -> Result { - let mut uint = 0; - let mut shift = 0; - loop { - let byte = self.read_byte()?; - uint += (byte as u64 & 0x7F) << shift; - shift += 7; - if byte & 0x80 == 0 { - break; - } - } - Ok(uint) - } -}