extmod/modbtree: Add checks for already-closed database.

Fixes use-after-free when accessing the database after it is closed with
`btree_close`.  `btree_close` always succeeds when called with an
already-closed database.

The new test checks that operations that access the underlying database
(get, set, flush, seq) fail with a `ValueError` when the btree is already
closed.  It also checks that closing and printing the btree succeed when
the btree is already closed.

Fixes issue micropython#12543.

Signed-off-by: Michael Vornovitsky <[email protected]>

Author: BottleRocketeer

extmod/modbtree.c

@@ -89,6 +89,12 @@ void __dbpanic(DB *db) {
     mp_printf(&mp_plat_print, "__dbpanic(%p)\n", db);
 }
 
+static void check_btree_is_open(mp_obj_btree_t *self) {
+    if (!self->db) {
+        mp_raise_ValueError(MP_ERROR_TEXT("database closed"));
+    }
+}
+
 static mp_obj_btree_t *btree_new(DB *db, mp_obj_t stream) {
     mp_obj_btree_t *o = mp_obj_malloc(mp_obj_btree_t, (mp_obj_type_t *)&btree_type);
     o->stream = stream;
@@ -114,19 +120,28 @@ static void btree_print(const mp_print_t *print, mp_obj_t self_in, mp_print_kind
 
 static mp_obj_t btree_flush(mp_obj_t self_in) {
     mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
+    check_btree_is_open(self);
     return MP_OBJ_NEW_SMALL_INT(__bt_sync(self->db, 0));
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(btree_flush_obj, btree_flush);
 
 static mp_obj_t btree_close(mp_obj_t self_in) {
     mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
-    return MP_OBJ_NEW_SMALL_INT(__bt_close(self->db));
+    int res;
+    if (self->db) {
+        res = __bt_close(self->db);
+        self->db = NULL;
+    } else {
+        res = RET_SUCCESS; // Closing an already-closed DB always succeeds.
+    }
+    return MP_OBJ_NEW_SMALL_INT(res);
 }
 static MP_DEFINE_CONST_FUN_OBJ_1(btree_close_obj, btree_close);
 
 static mp_obj_t btree_put(size_t n_args, const mp_obj_t *args) {
     (void)n_args;
     mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
+    check_btree_is_open(self);
     DBT key, val;
     buf_to_dbt(args[1], &key);
     buf_to_dbt(args[2], &val);
@@ -136,6 +151,7 @@ static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_put_obj, 3, 4, btree_put);
 
 static mp_obj_t btree_get(size_t n_args, const mp_obj_t *args) {
     mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
+    check_btree_is_open(self);
     DBT key, val;
     buf_to_dbt(args[1], &key);
     int res = __bt_get(self->db, &key, &val, 0);
@@ -153,6 +169,7 @@ static MP_DEFINE_CONST_FUN_OBJ_VAR_BETWEEN(btree_get_obj, 2, 3, btree_get);
 
 static mp_obj_t btree_seq(size_t n_args, const mp_obj_t *args) {
     mp_obj_btree_t *self = MP_OBJ_TO_PTR(args[0]);
+    check_btree_is_open(self);
     int flags = MP_OBJ_SMALL_INT_VALUE(args[1]);
     DBT key, val;
     if (n_args > 2) {
@@ -225,6 +242,7 @@ static mp_obj_t btree_getiter(mp_obj_t self_in, mp_obj_iter_buf_t *iter_buf) {
 
 static mp_obj_t btree_iternext(mp_obj_t self_in) {
     mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
+    check_btree_is_open(self);
     DBT key, val;
     int res;
     bool desc = self->flags & FLAG_DESC;
@@ -281,6 +299,7 @@ static mp_obj_t btree_iternext(mp_obj_t self_in) {
 
 static mp_obj_t btree_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
     mp_obj_btree_t *self = MP_OBJ_TO_PTR(self_in);
+    check_btree_is_open(self);
     if (value == MP_OBJ_NULL) {
         // delete
         DBT key;
@@ -314,6 +333,7 @@ static mp_obj_t btree_subscr(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
 
 static mp_obj_t btree_binary_op(mp_binary_op_t op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
     mp_obj_btree_t *self = MP_OBJ_TO_PTR(lhs_in);
+    check_btree_is_open(self);
     switch (op) {
         case MP_BINARY_OP_CONTAINS: {
             DBT key, val;

tests/extmod/btree_closed.py

@@ -0,0 +1,39 @@
+try:
+    import btree
+    import io
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+f = io.BytesIO()
+db = btree.open(f)
+
+db[b"foo"] = b"42"
+
+db.close()
+
+# Accessing an already-closed database should fail.
+try:
+    print(db[b"foo"])
+except ValueError:
+    print("ValueError")
+
+try:
+    db[b"bar"] = b"43"
+except ValueError:
+    print("ValueError")
+
+try:
+    db.flush()
+except ValueError:
+    print("ValueError")
+
+try:
+    for k, v in db.items():
+        pass
+except ValueError:
+    print("ValueError")
+
+# Closing and printing an already-closed database should not fail.
+db.close()
+print(db)

tests/extmod/btree_closed.py.exp

@@ -0,0 +1,5 @@
+ValueError
+ValueError
+ValueError
+ValueError
+<btree 0>