Reject unsafe matchspecs on ets:select_replace/2

Preemptively fail operation with badarg if the replacement object
might have a different key.

4 files changed, 177 insertions, 33 deletions

diff --git a/erts/emulator/beam/erl_db_hash.c b/erts/emulator/beam/erl_db_hash.c
index 44d11dd5e7..ac31569181 100644
--- a/erts/emulator/beam/erl_db_hash.c
+++ b/erts/emulator/beam/erl_db_hash.c
@@ -355,6 +355,9 @@ static ERTS_INLINE void SET_SEGTAB(DbTableHash* tb,
 	erts_smp_atomic_set_nob(&tb->segtab, (erts_aint_t) segtab);
 }
 
+/* Used by select_replace on analyze_pattern */
+typedef int (*extra_match_validator_t)(int keypos, Eterm match, Eterm guard, Eterm body);
+
 /*
 ** Forward decl's (static functions)
 */
@@ -369,8 +372,9 @@ static void shrink(DbTableHash* tb, int nitems);
 static void grow(DbTableHash* tb, int nitems);
 static Eterm build_term_list(Process* p, HashDbTerm* ptr1, HashDbTerm* ptr2,
 			   Uint sz, DbTableHash*);
-static int analyze_pattern(DbTableHash *tb, Eterm pattern, 
-			   struct mp_info *mpi);
+static int analyze_pattern(DbTableHash *tb, Eterm pattern,
+                           extra_match_validator_t extra_validator, /* Optional callback */
+                           struct mp_info *mpi);
 
 /*
  *  Method interface functions
@@ -1202,7 +1206,9 @@ typedef int (*mtraversal_on_trap_t)(void* context_ptr, Sint slot_ix, Sint got, B
 /*
  * Begin hash table match traversal
  */
-static int match_traverse(Process* p, DbTableHash* tb, Eterm pattern,
+static int match_traverse(Process* p, DbTableHash* tb,
+                          Eterm pattern,
+                          extra_match_validator_t extra_match_validator, /* Optional */
                           Sint chunk_size,      /* If 0, no chunking */
                           Sint iterations_left, /* Nr. of iterations left */
                           Eterm** hpp,          /* Heap */
@@ -1240,7 +1246,9 @@ static int match_traverse(Process* p, DbTableHash* tb, Eterm pattern,
 
     *ret = NIL;
 
-    if ((ret_value = analyze_pattern(tb, pattern, &mpi)) != DB_ERROR_NONE) {
+    if ((ret_value = analyze_pattern(tb, pattern, extra_match_validator, &mpi))
+            != DB_ERROR_NONE)
+    {
         goto done;
     }
 
@@ -1713,7 +1721,9 @@ static int db_select_chunk_hash(Process *p, DbTable *tbl, Eterm tid, Eterm patte
     sc_context.prev_continuation_tptr = NULL;
 
     return match_traverse(
-            sc_context.p, sc_context.tb, pattern, sc_context.chunk_size,
+            sc_context.p, sc_context.tb,
+            pattern, NULL,
+            sc_context.chunk_size,
             MAX_SELECT_CHUNK_ITERATIONS,
             &sc_context.hp, 0,
             mtraversal_select_chunk_on_nothing_can_match,
@@ -1906,8 +1916,8 @@ static int db_select_count_hash(Process *p, DbTable *tbl, Eterm tid, Eterm patte
 
     return match_traverse(
             scnt_context.p, scnt_context.tb,
-            pattern, chunk_size,
-            iterations_left, NULL, 0,
+            pattern, NULL,
+            chunk_size, iterations_left, NULL, 0,
             mtraversal_select_count_on_nothing_can_match,
             mtraversal_select_count_on_match_res,
             mtraversal_select_count_on_loop_ended,
@@ -2046,7 +2056,9 @@ static int db_select_delete_hash(Process *p, DbTable *tbl, Eterm tid, Eterm patt
     sd_context.last_pseudo_delete = (Uint) -1;
 
     return match_traverse(
-            sd_context.p, sd_context.tb, pattern, chunk_size,
+            sd_context.p, sd_context.tb,
+            pattern, NULL,
+            chunk_size,
             MAX_SELECT_DELETE_ITERATIONS, NULL, 1,
             mtraversal_select_delete_on_nothing_can_match,
             mtraversal_select_delete_on_match_res,
@@ -2120,15 +2132,18 @@ static int mtraversal_select_replace_on_match_res(void* context_ptr, Sint slot_i
 {
     mtraversal_select_replace_context_t* sr_context_ptr = (mtraversal_select_replace_context_t*) context_ptr;
     DbTableHash* tb = sr_context_ptr->tb;
+#ifdef DEBUG
     Eterm key = NIL;
+#endif
     HashDbTerm* new = NULL;
     HashDbTerm* next = NULL;
     HashValue hval = INVALID_HASH;
 
-    if (is_value(match_res) &&
-            is_value(key = db_getkey(tb->common.keypos, match_res)) &&
-            eq(key, GETKEY(tb, (**current_ptr_ptr)->dbterm.tpl)))
-    {
+    if (is_value(match_res)) {
+#ifdef DEBUG
+        ASSERT(is_value(key = db_getkey(tb->common.keypos, match_res)));
+        ASSERT(eq(key, GETKEY(tb, (**current_ptr_ptr)->dbterm.tpl)));
+#endif
         next = (**current_ptr_ptr)->next;
         hval = (**current_ptr_ptr)->hvalue;
         new = replace_dbterm(tb, **current_ptr_ptr, match_res);
@@ -2184,7 +2199,9 @@ static int db_select_replace_hash(Process *p, DbTable *tbl, Eterm pattern, Eterm
     sr_context.prev_continuation_tptr = NULL;
 
     return match_traverse(
-            sr_context.p, sr_context.tb, pattern, chunk_size,
+            sr_context.p, sr_context.tb,
+            pattern, db_match_keeps_key,
+            chunk_size,
             MAX_SELECT_REPLACE_ITERATIONS, NULL, 1,
             mtraversal_select_replace_on_nothing_can_match,
             mtraversal_select_replace_on_match_res,
@@ -2428,7 +2445,8 @@ static SWord db_free_table_continue_hash(DbTable *tbl, SWord reds)
 ** slots should be searched. Also compiles the match program
 */
 static int analyze_pattern(DbTableHash *tb, Eterm pattern, 
-			   struct mp_info *mpi)
+                           extra_match_validator_t extra_validator, /* Optional callback */
+                           struct mp_info *mpi)
 {
     Eterm *ptpl;
     Eterm lst, tpl, ttpl;
@@ -2466,7 +2484,10 @@ static int analyze_pattern(DbTableHash *tb, Eterm pattern,
 
     i = 0;
     for(lst = pattern; is_list(lst); lst = CDR(list_val(lst))) {
-	Eterm body;
+        Eterm match = NIL;
+        Eterm guard = NIL;
+        Eterm body = NIL;
+
 	ttpl = CAR(list_val(lst));
 	if (!is_tuple(ttpl)) {
 	    if (buff != sbuff) { 
@@ -2481,9 +2502,17 @@ static int analyze_pattern(DbTableHash *tb, Eterm pattern,
 	    }
 	    return DB_ERROR_BADPARAM;
 	}
-	matches[i] = tpl = ptpl[1];
-	guards[i] = ptpl[2];
+	matches[i] = match = tpl = ptpl[1];
+	guards[i] = guard = ptpl[2];
 	bodies[i] = body = ptpl[3];
+
+        if(extra_validator != NULL && !extra_validator(tb->common.keypos, match, guard, body)) {
+	    if (buff != sbuff) {
+		erts_free(ERTS_ALC_T_DB_TMP, buff);
+	    }
+            return DB_ERROR_BADPARAM;
+        }
+
 	if (!is_list(body) || CDR(list_val(body)) != NIL ||
 	    CAR(list_val(body)) != am_DollarUnderscore) {
 	    mpi->all_objects = 0;
diff --git a/erts/emulator/beam/erl_db_tree.c b/erts/emulator/beam/erl_db_tree.c
index fea888b12d..949f8c46b6 100644
--- a/erts/emulator/beam/erl_db_tree.c
+++ b/erts/emulator/beam/erl_db_tree.c
@@ -300,6 +300,9 @@ struct select_replace_context {
     Sint replaced;
 };
 
+/* Used by select_replace on analyze_pattern */
+typedef int (*extra_match_validator_t)(int keypos, Eterm match, Eterm guard, Eterm body);
+
 /*
 ** Forward declarations 
 */
@@ -351,7 +354,8 @@ static Sint cmp_partly_bound(Eterm partly_bound_key, Eterm bound_key);
 static Sint do_cmp_partly_bound(Eterm a, Eterm b, int *done);
 
 static int analyze_pattern(DbTableTree *tb, Eterm pattern, 
-			   struct mp_info *mpi);
+                           extra_match_validator_t extra_validator, /* Optional callback */
+                           struct mp_info *mpi);
 static int doit_select(DbTableTree *tb,
 		       TreeDbTerm *this,
 		       void *ptr,
@@ -1132,7 +1136,7 @@ static int db_select_tree(Process *p, DbTable *tbl, Eterm tid,
     sc.got = 0;
     sc.chunk_size = 0;
 
-    if ((errcode = analyze_pattern(tb, pattern, &mpi)) != DB_ERROR_NONE) {
+    if ((errcode = analyze_pattern(tb, pattern, NULL, &mpi)) != DB_ERROR_NONE) {
 	RET_TO_BIF(NIL,errcode);
     }
 
@@ -1335,7 +1339,7 @@ static int db_select_count_tree(Process *p, DbTable *tbl, Eterm tid,
     sc.keypos = tb->common.keypos;
     sc.got = 0;
 
-    if ((errcode = analyze_pattern(tb, pattern, &mpi)) != DB_ERROR_NONE) {
+    if ((errcode = analyze_pattern(tb, pattern, NULL, &mpi)) != DB_ERROR_NONE) {
 	RET_TO_BIF(NIL,errcode);
     }
 
@@ -1438,7 +1442,7 @@ static int db_select_chunk_tree(Process *p, DbTable *tbl, Eterm tid,
     sc.got = 0;
     sc.chunk_size = chunk_size;
 
-    if ((errcode = analyze_pattern(tb, pattern, &mpi)) != DB_ERROR_NONE) {
+    if ((errcode = analyze_pattern(tb, pattern, NULL, &mpi)) != DB_ERROR_NONE) {
 	RET_TO_BIF(NIL,errcode);
     }
 
@@ -1680,7 +1684,7 @@ static int db_select_delete_tree(Process *p, DbTable *tbl, Eterm tid,
     sc.keypos = tb->common.keypos;
     sc.tb = tb;
     
-    if ((errcode = analyze_pattern(tb, pattern, &mpi)) != DB_ERROR_NONE) {
+    if ((errcode = analyze_pattern(tb, pattern, NULL, &mpi)) != DB_ERROR_NONE) {
 	RET_TO_BIF(0,errcode);
     }
 
@@ -1872,7 +1876,7 @@ static int db_select_replace_tree(Process *p, DbTable *tbl,
     sc.keypos = tb->common.keypos;
     sc.replaced = 0;
 
-    if ((errcode = analyze_pattern(tb, pattern, &mpi)) != DB_ERROR_NONE) {
+    if ((errcode = analyze_pattern(tb, pattern, db_match_keeps_key, &mpi)) != DB_ERROR_NONE) {
         RET_TO_BIF(NIL,errcode);
     }
 
@@ -2193,8 +2197,9 @@ static TreeDbTerm *linkout_object_tree(DbTableTree *tb,
 ** For the select functions, analyzes the pattern and determines which
 ** part of the tree should be searched. Also compiles the match program
 */
-static int analyze_pattern(DbTableTree *tb, Eterm pattern, 
-			   struct mp_info *mpi)
+static int analyze_pattern(DbTableTree *tb, Eterm pattern,
+                           extra_match_validator_t extra_validator, /* Optional callback */
+                           struct mp_info *mpi)
 {
     Eterm lst, tpl, ttpl;
     Eterm *matches,*guards, *bodies;
@@ -2232,7 +2237,10 @@ static int analyze_pattern(DbTableTree *tb, Eterm pattern,
 
     i = 0;
     for(lst = pattern; is_list(lst); lst = CDR(list_val(lst))) {
-	Eterm body;
+        Eterm match = NIL;
+        Eterm guard = NIL;
+        Eterm body = NIL;
+
 	ttpl = CAR(list_val(lst));
 	if (!is_tuple(ttpl)) {
 	    if (buff != sbuff) { 
@@ -2247,9 +2255,17 @@ static int analyze_pattern(DbTableTree *tb, Eterm pattern,
 	    }
 	    return DB_ERROR_BADPARAM;
 	}
-	matches[i] = tpl = ptpl[1];
-	guards[i] = ptpl[2];
+	matches[i] = match = tpl = ptpl[1];
+	guards[i] = guard = ptpl[2];
 	bodies[i] = body = ptpl[3];
+
+        if(extra_validator != NULL && !extra_validator(tb->common.keypos, match, guard, body)) {
+	    if (buff != sbuff) {
+		erts_free(ERTS_ALC_T_DB_TMP, buff);
+	    }
+            return DB_ERROR_BADPARAM;
+        }
+
 	if (!is_list(body) || CDR(list_val(body)) != NIL ||
 	    CAR(list_val(body)) != am_DollarUnderscore) {
 	    mpi->all_objects = 0;
@@ -3443,7 +3459,10 @@ static int doit_select_replace(DbTableTree *tb, TreeDbTerm **this, void *ptr,
                                int forward)
 {
     struct select_replace_context *sc = (struct select_replace_context *) ptr;
-    Eterm ret, key;
+    Eterm ret = NIL;
+#ifdef DEBUG
+    Eterm key = NIL;
+#endif
 
     sc->lastobj = (*this)->dbterm.tpl;
 
@@ -3456,10 +3475,11 @@ static int doit_select_replace(DbTableTree *tb, TreeDbTerm **this, void *ptr,
     ret = db_match_dbterm(&tb->common, sc->p, sc->mp, 0,
 			  &(*this)->dbterm, NULL, 0);
 
-    if (is_value(ret) &&
-        is_value(key = db_getkey(tb->common.keypos, ret)) &&
-        (cmp_key(tb, key, *this) == 0))
-    {
+    if (is_value(ret)) {
+#ifdef DEBUG
+        ASSERT(is_value(key = db_getkey(tb->common.keypos, ret)));
+        ASSERT(cmp_key(tb, key, *this) == 0);
+#endif
         *this = replace_dbterm(tb, *this, ret);
         sc->lastobj = (*this)->dbterm.tpl;
         ++(sc->replaced);
diff --git a/erts/emulator/beam/erl_db_util.c b/erts/emulator/beam/erl_db_util.c
index 6f30b1d3dd..3b2e7f4407 100644
--- a/erts/emulator/beam/erl_db_util.c
+++ b/erts/emulator/beam/erl_db_util.c
@@ -1119,6 +1119,100 @@ error:
     return NULL;
 }
 
+/* This is used by select_replace */
+int db_match_keeps_key(int keypos, Eterm match, Eterm guard, Eterm body) {
+    Eterm match_key = NIL;
+    int match_key_variable = -1;
+    Eterm* body_list = NULL;
+    Eterm single_body_term = NIL;
+    Eterm* single_body_term_tpl = NULL;
+    Eterm single_body_subterm = NIL;
+    Eterm single_body_subterm_key = NIL;
+    int single_body_subterm_key_variable = -1;
+    Eterm* single_body_subterm_key_tpl = NULL;
+
+    if (!is_list(body)) {
+        return 0;
+    }
+
+    body_list = list_val(body);
+    if (CDR(body_list) != NIL) {
+        return 0;
+    }
+
+    single_body_term = CAR(body_list);
+    if (single_body_term == am_DollarUnderscore) {
+        /* same tuple is returned */
+        return 1;
+    }
+
+    if (!is_tuple(single_body_term)) {
+        return 0;
+    }
+
+    single_body_term_tpl = tuple_val(single_body_term);
+    if (arityval(*single_body_term_tpl) != 1) {
+        // not the 1-element tuple we're expecting
+        return 0;
+    }
+
+    match_key = db_getkey(keypos, match);
+    if (!is_value(match_key)) {
+        // can't get key out of match
+        return 0;
+    }
+
+    single_body_subterm = single_body_term_tpl[1];
+    single_body_subterm_key = db_getkey(keypos, single_body_subterm);
+    if (!is_value(single_body_subterm_key)) {
+        // can't get key out of single body subterm
+        return 0;
+    }
+
+    match_key_variable = db_is_variable(match_key);
+    single_body_subterm_key_variable = db_is_variable(single_body_subterm_key);
+    if (match_key_variable != -1 && match_key_variable == single_body_subterm_key_variable) {
+        /* tuple with same key is returned */
+        return 1;
+    }
+
+    if (!is_tuple(single_body_subterm_key)) {
+        /* can't possibly be an element instruction */
+        return 0;
+    }
+
+    single_body_subterm_key_tpl = tuple_val(single_body_subterm_key);
+    if (arityval(*single_body_subterm_key_tpl) != 3) {
+        /* can't possibly be an element instruction */
+        return 0;
+    }
+
+    if (single_body_subterm_key_tpl[1] != am_element) {
+        /* tag is not of an element instruction */
+        return 0;
+    }
+    if (single_body_subterm_key_tpl[3] != am_DollarUnderscore) {
+        /* even if it's an element instruction, it's not fetching from the original tuple */
+        return 0;
+    }
+
+    if (is_big(single_body_subterm_key_tpl[2])
+        && (big_to_uint32(single_body_subterm_key_tpl[2]) != keypos))
+    {
+        /* the key comes from a different position */
+        return 0;
+    }
+
+    if (is_small(single_body_subterm_key_tpl[2])
+        && (unsigned_val(single_body_subterm_key_tpl[2]) != keypos))
+    {
+        /* the key comes from a different position */
+        return 0;
+    }
+
+    return 1;
+}
+
 /* This is used when tracing */
 Eterm erts_match_set_lint(Process *p, Eterm matchexpr) {
     return db_match_set_lint(p, matchexpr, DCOMP_TRACE);
diff --git a/erts/emulator/beam/erl_db_util.h b/erts/emulator/beam/erl_db_util.h
index 89ef79d2dc..1a84f20e5d 100644
--- a/erts/emulator/beam/erl_db_util.h
+++ b/erts/emulator/beam/erl_db_util.h
@@ -382,6 +382,7 @@ Eterm db_add_counter(Eterm** hpp, Wterm counter, Eterm incr);
 Eterm db_match_set_lint(Process *p, Eterm matchexpr, Uint flags);
 Binary *db_match_set_compile(Process *p, Eterm matchexpr, 
 			     Uint flags);
+int db_match_keeps_key(int keypos, Eterm match, Eterm guard, Eterm body);
 int erts_db_match_prog_destructor(Binary *);
 
 typedef struct match_prog {