core/crypto: introduce SecureContext

core/embed/extmod/modtrezorcrypto/modtrezorcrypto-secbytes.h

@@ -17,27 +17,18 @@
  * along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-// #include "py/objstr.h"
+#include "py/objlist.h"
 
 #include "memzero.h"
+#include "secbool.h"
 
 /// package: trezorcrypto.__init__
 
 #define SECBYTES_DEBUG 1
 
-#if SECBYTES_DEBUG
-STATIC uint32_t id_counter = 0;
-#endif
-
-/// class secbytes:
-///     """
-///     secbytes
-///     """
 typedef struct _mp_obj_secbytes_t {
   mp_obj_base_t base;
-#if SECBYTES_DEBUG
+  secbool valid;
-  uint32_t id;
-#endif
   uint32_t len;
   byte *ptr;
 } mp_obj_secbytes_t;
@@ -47,84 +38,38 @@ STATIC const mp_obj_type_t secbytes_type; // forward declaration
 STATIC mp_obj_secbytes_t *m_new_secbytes(const byte *data, uint32_t len) {
   mp_obj_secbytes_t *o = m_new_obj_with_finaliser(mp_obj_secbytes_t);
   o->base.type = &secbytes_type;
-#if SECBYTES_DEBUG
-  o->id = ++id_counter;
-#endif
   o->len = len;
   o->ptr = m_new(byte, o->len);
   if (data) {
     memcpy(o->ptr, data, len);
   }
+  o->valid = sectrue;
   return o;
 }
 
-/// def __init__(self) -> None:
+#if SECBYTES_DEBUG
-///     """
-///     Creates a secbytes object.
-///     """
-STATIC mp_obj_t secbytes_make_new(const mp_obj_type_t *type, size_t n_args,
-                                  size_t n_kw, const mp_obj_t *args) {
-  (void)type;
-  if (n_args != 1) {
-    mp_raise_ValueError("missing argument");
-  }
-  mp_buffer_info_t input;
-  mp_get_buffer_raise(args[0], &input, MP_BUFFER_READ);
-  mp_obj_secbytes_t *o = m_new_secbytes(input.buf, input.len);
-  return MP_OBJ_FROM_PTR(o);
-}
 
 STATIC void secbytes_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind_t kind) {
     (void)kind;
     mp_obj_secbytes_t *self = MP_OBJ_TO_PTR(self_in);
-#if SECBYTES_DEBUG
+    mp_printf(print, "secbytes{valid=" UINT_FMT ", len=" UINT_FMT ", data=\"", self->valid, self->len);
-    mp_printf(print, "secbytes{id=" UINT_FMT ", len=" UINT_FMT ", data=\"", self->id, self->len);
     for (uint32_t i = 0; i < self->len; i++) {
       mp_printf(print, "%02x", self->ptr[i]);
     }
     mp_printf(print, "\"}");
-#else
+}
-    mp_printf(print, "secbytes{}");
+
 #endif
-}
-
-STATIC mp_obj_t secbytes_binary_op(mp_binary_op_t op, mp_obj_t lhs_in,
-                                   mp_obj_t rhs_in) {
-  mp_obj_type_t *lhs_type = mp_obj_get_type(lhs_in);
-  mp_obj_type_t *rhs_type = mp_obj_get_type(rhs_in);
-  if (lhs_type != &secbytes_type) {
-    mp_raise_TypeError("lhs");
-  }
-  if (rhs_type != &secbytes_type) {
-    mp_raise_TypeError("rhs");
-  }
-
-  switch (op) {
-    case MP_BINARY_OP_ADD:
-    case MP_BINARY_OP_INPLACE_ADD: {
-      mp_obj_secbytes_t *lhs = MP_OBJ_TO_PTR(lhs_in);
-      mp_obj_secbytes_t *rhs = MP_OBJ_TO_PTR(rhs_in);
-      mp_obj_secbytes_t *o = m_new_secbytes(NULL, lhs->len + rhs->len);
-      memcpy(o->ptr, lhs->ptr, lhs->len);
-      memcpy(o->ptr + lhs->len, rhs->ptr, rhs->len);
-      return MP_OBJ_FROM_PTR(o);
-    }
-    default:
-      return MP_OBJ_NULL; // op not supported
-  }
-}
 
 STATIC mp_obj_t secbytes___del__(mp_obj_t self) {
   mp_obj_secbytes_t *o = MP_OBJ_TO_PTR(self);
 #if SECBYTES_DEBUG
-  mp_printf(&mp_stderr_print, "del secbytes{id=%u}\n", o->id);
+  mp_printf(&mp_stderr_print, "del secbytes\n");
-#endif
-  memzero(&(o->ptr), o->len);
-#if SECBYTES_DEBUG
-  o->id = 0;
 #endif
+  memzero(o->ptr, o->len);
   o->len = 0;
   // o->ptr should be freed by MicroPython
+  o->valid = secfalse;
   return mp_const_none;
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(secbytes___del___obj, secbytes___del__);
@@ -137,8 +82,71 @@ STATIC MP_DEFINE_CONST_DICT(secbytes_locals_dict, secbytes_locals_dict_table);
 STATIC const mp_obj_type_t secbytes_type = {
     {&mp_type_type},
     .name = MP_QSTR_secbytes,
+#if SECBYTES_DEBUG
     .print = secbytes_print,
-    .make_new = secbytes_make_new,
+#endif
-    .binary_op = secbytes_binary_op,
     .locals_dict = (void *)&secbytes_locals_dict,
 };
+
+typedef struct _mp_obj_SecureContext_t {
+  mp_obj_base_t base;
+  mp_obj_list_t list;
+} mp_obj_SecureContext_t;
+
+STATIC mp_obj_t SecureContext_make_new(const mp_obj_type_t *type, size_t n_args,
+                                  size_t n_kw, const mp_obj_t *args) {
+  mp_obj_SecureContext_t *o = m_new_obj(mp_obj_SecureContext_t);
+  o->base.type = type;
+  mp_obj_list_init(&(o->list), 0);
+  return MP_OBJ_FROM_PTR(o);
+}
+
+STATIC mp_obj_t SecureContext___exit__(size_t n_args, const mp_obj_t *args) {
+  mp_obj_SecureContext_t *o = MP_OBJ_TO_PTR(args[0]);
+  // iterate the list and memzero its members
+  for (size_t i = 0; i < o->list.len; i++) {
+    mp_obj_secbytes_t *b = MP_OBJ_TO_PTR(o->list.items[i]);
+    memzero(b->ptr, b->len);
+    b->len = 0;
+    b->valid = secfalse;
+  }
+  return mp_const_none;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(SecureContext___exit___obj, 4, 4, SecureContext___exit__);
+
+STATIC mp_obj_t SecureContext_new(mp_obj_t self, mp_obj_t data) {
+  mp_buffer_info_t input;
+  mp_get_buffer_raise(data, &input, MP_BUFFER_READ);
+  mp_obj_secbytes_t *b = m_new_secbytes(input.buf, input.len);
+  mp_obj_SecureContext_t *o = MP_OBJ_TO_PTR(self);
+  mp_obj_list_append(MP_OBJ_FROM_PTR(&(o->list)), MP_OBJ_FROM_PTR(b));
+  return MP_OBJ_FROM_PTR(b);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_2(SecureContext_new_obj, SecureContext_new);
+
+STATIC mp_obj_t SecureContext_concat(mp_obj_t self, mp_obj_t bytes1, mp_obj_t bytes2) {
+  mp_obj_secbytes_t *b1 = MP_OBJ_TO_PTR(bytes1);
+  mp_obj_secbytes_t *b2 = MP_OBJ_TO_PTR(bytes2);
+  mp_obj_secbytes_t *b = m_new_secbytes(NULL, b1->len + b2->len);
+  memcpy(b->ptr, b1->ptr, b1->len);
+  memcpy(b->ptr + b1->len, b2->ptr, b2->len);
+  mp_obj_SecureContext_t *o = MP_OBJ_TO_PTR(self);
+  mp_obj_list_append(MP_OBJ_FROM_PTR(&(o->list)), MP_OBJ_FROM_PTR(b));
+  return MP_OBJ_FROM_PTR(b);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_3(SecureContext_concat_obj, SecureContext_concat);
+
+STATIC const mp_rom_map_elem_t SecureContext_locals_dict_table[] = {
+    {MP_ROM_QSTR(MP_QSTR___enter__), MP_ROM_PTR(&mp_identity_obj)},
+    {MP_ROM_QSTR(MP_QSTR___exit__), MP_ROM_PTR(&SecureContext___exit___obj)},
+    {MP_ROM_QSTR(MP_QSTR_new), MP_ROM_PTR(&SecureContext_new_obj)},
+    {MP_ROM_QSTR(MP_QSTR_concat), MP_ROM_PTR(&SecureContext_concat_obj)},
+};
+STATIC MP_DEFINE_CONST_DICT(SecureContext_locals_dict, SecureContext_locals_dict_table);
+
+STATIC const mp_obj_type_t SecureContext_type = {
+    {&mp_type_type},
+    .name = MP_QSTR_SecureContext,
+    .make_new = SecureContext_make_new,
+    .locals_dict = (void *)&SecureContext_locals_dict,
+};

core/embed/extmod/modtrezorcrypto/modtrezorcrypto.c

@@ -87,7 +87,7 @@ STATIC const mp_rom_map_elem_t mp_module_trezorcrypto_globals_table[] = {
     {MP_ROM_QSTR(MP_QSTR_rfc6979), MP_ROM_PTR(&mod_trezorcrypto_Rfc6979_type)},
     {MP_ROM_QSTR(MP_QSTR_ripemd160),
      MP_ROM_PTR(&mod_trezorcrypto_Ripemd160_type)},
-    {MP_ROM_QSTR(MP_QSTR_secbytes), MP_ROM_PTR(&secbytes_type)},
+    {MP_ROM_QSTR(MP_QSTR_SecureContext), MP_ROM_PTR(&SecureContext_type)},
     {MP_ROM_QSTR(MP_QSTR_secp256k1),
      MP_ROM_PTR(&mod_trezorcrypto_secp256k1_module)},
 #if !BITCOIN_ONLY

core/src/trezor/crypto/__init__.py

@@ -8,6 +8,7 @@ from trezorcrypto import (  # noqa: F401
     pbkdf2,
     random,
     rfc6979,
+    SecureContext,
 )
 
 if not utils.BITCOIN_ONLY: