diff --git a/core/embed/rust/src/micropython/iter.rs b/core/embed/rust/src/micropython/iter.rs index 696b2dd878..afbfcf9ed6 100644 --- a/core/embed/rust/src/micropython/iter.rs +++ b/core/embed/rust/src/micropython/iter.rs @@ -6,6 +6,8 @@ use super::{ffi, runtime::catch_exception}; pub struct IterBuf { iter_buf: ffi::mp_obj_iter_buf_t, + error_checking: bool, + caught_exception: Obj, } impl IterBuf { @@ -17,20 +19,41 @@ impl IterBuf { }, buf: [Obj::const_null(), Obj::const_null(), Obj::const_null()], }, + error_checking: false, + caught_exception: Obj::const_null(), + } + } + + pub fn new_fallible() -> Self { + let mut new = Self::new(); + new.error_checking = true; + new + } + + pub fn try_iterate(&mut self, o: Obj) -> Result { + Iter::try_from_obj_with_buf(o, self) + } + + pub fn error(&self) -> Option { + if !self.error_checking || self.caught_exception.is_null() { + None + } else { + Some(self.caught_exception) } } } -#[allow(dead_code)] // iter_buf is not really used but needs to be there for SAFETY reasons pub struct Iter<'a> { iter: Obj, + // SAFETY: + // In the typical case, `iter` is literally a pointer to `iter_buf`. We need to hold + // an exclusive reference to `iter_buf` to avoid problems. iter_buf: &'a mut IterBuf, finished: bool, - caught_exception: Obj, } impl<'a> Iter<'a> { - pub fn try_from_obj_with_buf(o: Obj, iter_buf: &'a mut IterBuf) -> Result { + fn try_from_obj_with_buf(o: Obj, iter_buf: &'a mut IterBuf) -> Result { // SAFETY: // - In the common case, `ffi::mp_getiter` does not heap-allocate, but instead // uses memory from the passed `iter_buf`. We maintain this invariant by @@ -43,7 +66,6 @@ impl<'a> Iter<'a> { iter, iter_buf, finished: false, - caught_exception: Obj::const_null(), }) } } @@ -61,8 +83,8 @@ impl<'a> Iterator for Iter<'a> { // EXCEPTION: Can raise from the underlying iterator. let item = catch_exception(|| unsafe { ffi::mp_iternext(self.iter) }); match item { - Err(Error::CaughtException(exc)) => { - self.caught_exception = exc; + Err(Error::CaughtException(exc)) if self.iter_buf.error_checking => { + self.iter_buf.caught_exception = exc; None } Err(_) => panic!("unexpected error"), diff --git a/core/embed/rust/src/micropython/list.rs b/core/embed/rust/src/micropython/list.rs index a865b02b43..42da263139 100644 --- a/core/embed/rust/src/micropython/list.rs +++ b/core/embed/rust/src/micropython/list.rs @@ -144,10 +144,7 @@ impl TryFrom for Gc { #[cfg(test)] mod tests { - use crate::micropython::{ - iter::{Iter, IterBuf}, - testutil::mpy_init, - }; + use crate::micropython::{iter::IterBuf, testutil::mpy_init}; use super::*; use heapless::Vec; @@ -160,10 +157,10 @@ mod tests { let vec: Vec = (0..5).collect(); let list: Obj = List::from_iter(vec.iter().copied()).unwrap().into(); - let mut buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(list, &mut buf).unwrap(); // collect the elements into a Vec of maximum length 10, through an iterator - let retrieved_vec: Vec = iter + let retrieved_vec: Vec = IterBuf::new() + .try_iterate(list) + .unwrap() .map(TryInto::try_into) .collect::, Error>>() .unwrap(); @@ -199,9 +196,9 @@ mod tests { ); } - let mut buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(gc_list.into(), &mut buf).unwrap(); - let retrieved_vec: Vec = iter + let retrieved_vec: Vec = IterBuf::new() + .try_iterate(gc_list.into()) + .unwrap() .map(TryInto::try_into) .collect::, Error>>() .unwrap(); @@ -240,9 +237,9 @@ mod tests { slice[i] = ((i + 10) as u16).into(); } - let mut buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(list.into(), &mut buf).unwrap(); - let retrieved_vec: Vec = iter + let retrieved_vec: Vec = IterBuf::new() + .try_iterate(list.into()) + .unwrap() .map(TryInto::try_into) .collect::, Error>>() .unwrap(); diff --git a/core/embed/rust/src/micropython/obj.rs b/core/embed/rust/src/micropython/obj.rs index d8749b4ff7..ef75de97dd 100644 --- a/core/embed/rust/src/micropython/obj.rs +++ b/core/embed/rust/src/micropython/obj.rs @@ -447,11 +447,11 @@ impl Obj { // SAFETY: // Safe for pointers, for as long as MicroPython behaves sanely. // We assume that: - // * The pointer is a valid MicroPython object, which has `ObjBase` - // as its first element. + // * The pointer is a valid MicroPython object, which has `ObjBase` as its first + // element. // * The type pointer points to a valid type object. - // * The pointee has a 'static lifetime, i.e., either is - // ROM-based, or GC allocated. + // * The pointee has a 'static lifetime, i.e., either is ROM-based, or GC + // allocated. let base = self.as_ptr() as *const ObjBase; unsafe { (*base).type_.as_ref() } } else { diff --git a/core/embed/rust/src/protobuf/encode.rs b/core/embed/rust/src/protobuf/encode.rs index e9e8defe16..929a6927f2 100644 --- a/core/embed/rust/src/protobuf/encode.rs +++ b/core/embed/rust/src/protobuf/encode.rs @@ -2,15 +2,7 @@ use core::convert::{TryFrom, TryInto}; use crate::{ error::Error, - micropython::{ - buffer, - gc::Gc, - iter::{Iter, IterBuf}, - list::List, - obj::Obj, - qstr::Qstr, - util, - }, + micropython::{buffer, gc::Gc, iter::IterBuf, list::List, obj::Obj, qstr::Qstr, util}, }; use super::{ @@ -80,9 +72,7 @@ impl Encoder { }; if field.is_repeated() { - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(field_value, &mut iter_buf)?; - for iter_value in iter { + for iter_value in IterBuf::new().try_iterate(field_value)? { stream.write_uvarint(field_key)?; self.encode_field(stream, field, iter_value)?; } @@ -126,8 +116,7 @@ impl Encoder { // Serialize the total length of the buffer. let mut len = 0; - let iter = Iter::try_from_obj_with_buf(value, &mut iter_buf)?; - for value in iter { + for value in iter_buf.try_iterate(value)? { // SAFETY: buffer is dropped immediately. let buffer = unsafe { buffer::get_buffer(value) }?; len += buffer.len(); @@ -135,8 +124,7 @@ impl Encoder { stream.write_uvarint(len as u64)?; // Serialize the buffers one-by-one. - let iter = Iter::try_from_obj_with_buf(value, &mut iter_buf)?; - for value in iter { + for value in iter_buf.try_iterate(value)? { // SAFETY: buffer is dropped immediately. let buffer = unsafe { buffer::get_buffer(value) }?; stream.write(buffer)?; diff --git a/core/embed/rust/src/ui/layout/util.rs b/core/embed/rust/src/ui/layout/util.rs index 800cf44bbf..57b3a2f53f 100644 --- a/core/embed/rust/src/ui/layout/util.rs +++ b/core/embed/rust/src/ui/layout/util.rs @@ -3,7 +3,7 @@ use crate::{ micropython::{ buffer::{hexlify_bytes, StrBuffer}, gc::Gc, - iter::{Iter, IterBuf}, + iter::IterBuf, list::List, obj::Obj, util::try_or_raise, @@ -42,8 +42,7 @@ where Error: From, { let mut vec = Vec::::new(); - let mut iter_buf = IterBuf::new(); - for item in Iter::try_from_obj_with_buf(iterable, &mut iter_buf)? { + for item in IterBuf::new().try_iterate(iterable)? { vec.push(item.try_into()?) .map_err(|_| value_error!("Invalid iterable length"))?; } diff --git a/core/embed/rust/src/ui/model_tr/layout.rs b/core/embed/rust/src/ui/model_tr/layout.rs index f4944f1077..0ef823224b 100644 --- a/core/embed/rust/src/ui/model_tr/layout.rs +++ b/core/embed/rust/src/ui/model_tr/layout.rs @@ -6,15 +6,8 @@ use crate::{ error::Error, maybe_trace::MaybeTrace, micropython::{ - buffer::StrBuffer, - gc::Gc, - iter::{Iter, IterBuf}, - list::List, - map::Map, - module::Module, - obj::Obj, - qstr::Qstr, - util, + buffer::StrBuffer, gc::Gc, iter::IterBuf, list::List, map::Map, module::Module, obj::Obj, + qstr::Qstr, util, }, strutil::StringType, ui::{ @@ -367,9 +360,7 @@ extern "C" fn new_confirm_properties(n_args: usize, args: *const Obj, kwargs: *m let mut paragraphs = ParagraphVecLong::new(); - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?; - for para in iter { + for para in IterBuf::new().try_iterate(items)? { let [key, value, is_data]: [Obj; 3] = iter_into_array(para)?; let key = key.try_into_option::()?; let value = value.try_into_option::()?; @@ -430,12 +421,10 @@ extern "C" fn new_show_address_details(n_args: usize, args: *const Obj, kwargs: let path: Option = kwargs.get(Qstr::MP_QSTR_path)?.try_into_option()?; let xpubs: Obj = kwargs.get(Qstr::MP_QSTR_xpubs)?; - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(xpubs, &mut iter_buf)?; let mut ad = AddressDetails::new(address, case_sensitive, account, path)?; - for i in iter { + for i in IterBuf::new().try_iterate(xpubs)? { let [xtitle, text]: [StrBuffer; 2] = iter_into_array(i)?; ad.add_xpub(xtitle, text)?; } @@ -897,9 +886,7 @@ extern "C" fn new_confirm_with_info(n_args: usize, args: *const Obj, kwargs: *mu let mut paragraphs = ParagraphVecShort::new(); - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?; - for para in iter { + for para in IterBuf::new().try_iterate(items)? { let [font, text]: [Obj; 2] = iter_into_array(para)?; let style: &TextStyle = theme::textstyle_number_bold_or_mono(font.try_into()?); let text: StrBuffer = text.try_into()?; @@ -1062,10 +1049,8 @@ extern "C" fn new_show_checklist(n_args: usize, args: *const Obj, kwargs: *mut M let active: usize = kwargs.get(Qstr::MP_QSTR_active)?.try_into()?; let items: Obj = kwargs.get(Qstr::MP_QSTR_items)?; - let mut iter_buf = IterBuf::new(); let mut paragraphs = ParagraphVecLong::new(); - let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?; - for (i, item) in iter.enumerate() { + for (i, item) in IterBuf::new().try_iterate(items)?.enumerate() { let style = match i.cmp(&active) { Ordering::Less => &theme::TEXT_NORMAL, Ordering::Equal => &theme::TEXT_BOLD, diff --git a/core/embed/rust/src/ui/model_tt/layout.rs b/core/embed/rust/src/ui/model_tt/layout.rs index 2f2069dab1..98ee4b2da8 100644 --- a/core/embed/rust/src/ui/model_tt/layout.rs +++ b/core/embed/rust/src/ui/model_tt/layout.rs @@ -6,7 +6,7 @@ use crate::{ micropython::{ buffer::{get_buffer, StrBuffer}, gc::Gc, - iter::{Iter, IterBuf}, + iter::IterBuf, list::List, map::Map, module::Module, @@ -466,10 +466,8 @@ extern "C" fn new_confirm_emphasized(n_args: usize, args: *const Obj, kwargs: *m .try_into_option()?; let items: Obj = kwargs.get(Qstr::MP_QSTR_items)?; - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?; let mut ops = OpTextLayout::new(theme::TEXT_NORMAL); - for item in iter { + for item in IterBuf::new().try_iterate(items)? { if item.is_str() { ops = ops.text_normal(item.try_into()?) } else { @@ -748,12 +746,10 @@ extern "C" fn new_show_address_details(n_args: usize, args: *const Obj, kwargs: let path: Option = kwargs.get(Qstr::MP_QSTR_path)?.try_into_option()?; let xpubs: Obj = kwargs.get(Qstr::MP_QSTR_xpubs)?; - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(xpubs, &mut iter_buf)?; let mut ad = AddressDetails::new(address, case_sensitive, account, path)?; - for i in iter { + for i in IterBuf::new().try_iterate(xpubs)? { let [xtitle, text]: [StrBuffer; 2] = iter_into_array(i)?; ad.add_xpub(xtitle, text)?; } @@ -833,9 +829,7 @@ extern "C" fn new_confirm_total(n_args: usize, args: *const Obj, kwargs: *mut Ma let mut paragraphs = ParagraphVecShort::new(); - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?; - for pair in iter { + for pair in IterBuf::new().try_iterate(items)? { let [label, value]: [StrBuffer; 2] = iter_into_array(pair)?; paragraphs.add(Paragraph::new(&theme::TEXT_NORMAL, label)); paragraphs.add(Paragraph::new(&theme::TEXT_MONO, value)); @@ -1159,9 +1153,7 @@ extern "C" fn new_confirm_with_info(n_args: usize, args: *const Obj, kwargs: *mu let mut paragraphs = ParagraphVecShort::new(); - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?; - for para in iter { + for para in IterBuf::new().try_iterate(items)? { let [font, text]: [Obj; 2] = iter_into_array(para)?; let style: &TextStyle = theme::textstyle_number(font.try_into()?); let text: StrBuffer = text.try_into()?; @@ -1191,9 +1183,7 @@ extern "C" fn new_confirm_more(n_args: usize, args: *const Obj, kwargs: *mut Map let mut paragraphs = ParagraphVecLong::new(); - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?; - for para in iter { + for para in IterBuf::new().try_iterate(items)? { let [font, text]: [Obj; 2] = iter_into_array(para)?; let style: &TextStyle = theme::textstyle_number(font.try_into()?); let text: StrBuffer = text.try_into()?; @@ -1308,9 +1298,7 @@ extern "C" fn new_show_share_words(n_args: usize, args: *const Obj, kwargs: *mut let pages: Obj = kwargs.get(Qstr::MP_QSTR_pages)?; let mut paragraphs = ParagraphVecLong::new(); - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(pages, &mut iter_buf)?; - for page in iter { + for page in IterBuf::new().try_iterate(pages)? { let text: StrBuffer = page.try_into()?; paragraphs.add(Paragraph::new(&theme::TEXT_MONO, text).break_after()); } @@ -1360,10 +1348,8 @@ extern "C" fn new_show_checklist(n_args: usize, args: *const Obj, kwargs: *mut M let active: usize = kwargs.get(Qstr::MP_QSTR_active)?.try_into()?; let items: Obj = kwargs.get(Qstr::MP_QSTR_items)?; - let mut iter_buf = IterBuf::new(); let mut paragraphs = ParagraphVecLong::new(); - let iter = Iter::try_from_obj_with_buf(items, &mut iter_buf)?; - for (i, item) in iter.enumerate() { + for (i, item) in IterBuf::new().try_iterate(items)?.enumerate() { let style = match i.cmp(&active) { Ordering::Less => &theme::TEXT_CHECKLIST_DONE, Ordering::Equal => &theme::TEXT_CHECKLIST_SELECTED, @@ -1489,9 +1475,7 @@ extern "C" fn new_show_remaining_shares(n_args: usize, args: *const Obj, kwargs: let pages_iterable: Obj = kwargs.get(Qstr::MP_QSTR_pages)?; let mut paragraphs = ParagraphVecLong::new(); - let mut iter_buf = IterBuf::new(); - let iter = Iter::try_from_obj_with_buf(pages_iterable, &mut iter_buf)?; - for page in iter { + for page in IterBuf::new().try_iterate(pages_iterable)? { let [title, description]: [StrBuffer; 2] = iter_into_array(page)?; paragraphs .add(Paragraph::new(&theme::TEXT_DEMIBOLD, title))