Add more complete key-safety check

2 files changed, 192 insertions, 35 deletions

diff --git a/erts/emulator/beam/erl_db_util.c b/erts/emulator/beam/erl_db_util.c
index 3b2e7f4407..ab45c3fe0f 100644
--- a/erts/emulator/beam/erl_db_util.c
+++ b/erts/emulator/beam/erl_db_util.c
@@ -1119,17 +1119,111 @@ error:
     return NULL;
 }
 
+/*
+ * Compare a matching term 'a' with a constructing term 'b' for equality.
+ *
+ * Returns true if 'b' is guaranteed to always construct
+ * the same term as 'a' has matched.
+ */
+static int db_match_eq_body(Eterm a, Eterm b)
+{
+    DECLARE_ESTACK(s);
+    Uint arity;
+    Eterm *ap, *bp;
+    int const_mode = 0;
+    const Eterm CONST_MODE_OFF = THE_NON_VALUE;
+
+    while (1) {
+        switch(b & _TAG_PRIMARY_MASK) {
+        case TAG_PRIMARY_LIST:
+            if (!is_list(a))
+                return 0;
+            ESTACK_PUSH2(s, CDR(list_val(a)), CDR(list_val(b)));
+            a = CAR(list_val(a));
+            b = CAR(list_val(b));
+            continue; /* loop without pop */
+
+        case TAG_PRIMARY_BOXED:
+            if (is_tuple(b)) {
+                bp = tuple_val(b);
+                if (!const_mode) {
+                    if (bp[0] == make_arityval(1) && is_tuple(bp[1])) {
+                        b = bp[1]; /* double-tuple syntax */
+                    }
+                    else if (bp[0] == make_arityval(2) && bp[1] == am_const) {
+                        ESTACK_PUSH(s, CONST_MODE_OFF);
+                        const_mode = 1;   /* {const, term()} syntax */
+                        b = bp[2];
+                        continue; /* loop without pop */
+                    }
+                    else
+                        return 0; /* function call or invalid tuple syntax */
+                }
+                if (!is_tuple(a))
+                    return 0;
+
+                ap = tuple_val(a);
+                bp = tuple_val(b);
+                if (ap[0] != bp[0])
+                    return 0;
+                arity = arityval(ap[0]);
+                if (arity > 0) {
+                    a = *(++ap);
+                    b = *(++bp);
+                    while(--arity) {
+                        ESTACK_PUSH2(s, *(++ap), *(++bp));
+                    }
+                    continue; /* loop without pop */
+                }
+            }
+            else if (is_map(b)) {
+                /* We don't know what other pairs the matched map may contain */
+                return 0;
+            }
+            else if (!eq(a,b)) /* other boxed */
+                return 0;
+            break;
+
+        case TAG_PRIMARY_IMMED1:
+            if (a != b || a == am_Underscore || a == am_DollarDollar
+                || a == am_DollarUnderscore
+                || (const_mode && db_is_variable(a) >= 0)) {
+
+                return 0;
+            }
+            break;
+        default:
+            erts_exit(ERTS_ABORT_EXIT, "db_compare: "
+                      "Bad object on ESTACK: 0x%bex\n", b);
+        }
+
+pop_next:
+        if (ESTACK_ISEMPTY(s))
+            break; /* done */
+
+        b = ESTACK_POP(s);
+        if (b == CONST_MODE_OFF) {
+            ASSERT(const_mode);
+            const_mode = 0;
+            goto pop_next;
+        }
+        a = ESTACK_POP(s);
+    }
+
+    DESTROY_ESTACK(s);
+    return 1;
+}
+
 /* 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;
+int db_match_keeps_key(int keypos, Eterm match, Eterm guard, Eterm body)
+{
+    Eterm match_key;
+    Eterm* body_list;
+    Eterm single_body_term;
+    Eterm* single_body_term_tpl;
+    Eterm single_body_subterm;
+    Eterm single_body_subterm_key;
+    Eterm* single_body_subterm_key_tpl;
 
     if (!is_list(body)) {
         return 0;
@@ -1169,9 +1263,7 @@ int db_match_keeps_key(int keypos, Eterm match, Eterm guard, Eterm body) {
         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) {
+    if (db_match_eq_body(match_key, single_body_subterm_key)) {
         /* tuple with same key is returned */
         return 1;
     }
@@ -1187,30 +1279,15 @@ int db_match_keeps_key(int keypos, Eterm match, Eterm guard, Eterm body) {
         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))
+    if (single_body_subterm_key_tpl[1] == am_element &&
+        single_body_subterm_key_tpl[3] == am_DollarUnderscore &&
+        single_body_subterm_key_tpl[2] == make_small(keypos))
     {
-        /* the key comes from a different position */
-        return 0;
+        /* {element, KeyPos, '$_'} */
+        return 1;
     }
 
-    return 1;
+    return 0;
 }
 
 /* This is used when tracing */
diff --git a/lib/stdlib/test/ets_SUITE.erl b/lib/stdlib/test/ets_SUITE.erl
index 68d73e78b0..84d90b3260 100644
--- a/lib/stdlib/test/ets_SUITE.erl
+++ b/lib/stdlib/test/ets_SUITE.erl
@@ -87,6 +87,7 @@
 -include_lib("common_test/include/ct.hrl").
 
 -define(m(A,B), assert_eq(A,B)).
+-define(heap_binary_size, 64).
 
 init_per_testcase(Case, Config) ->
     rand:seed(exsplus),
@@ -1192,6 +1193,18 @@ t_select_replace(Config) when is_list(Config) ->
                                      [{{'$1', '$2'}}]}],
                       2000 = ets:select_replace(Table, MatchSpec6),
 
+                      % Replacement uses {element,KeyPos,T} for key
+                      2000 = ets:select_replace(Table,
+                                                [{{'$1', '$2'},
+                                                  [{'>', {'rem', '$1', 5}, 3}],
+                                                  [{{{element, 1, '$_'}, '$2'}}]}]),
+
+                      % Replacement uses wrong {element,KeyPos,T} for key
+                      {'EXIT',{badarg,_}} = (catch ets:select_replace(Table,
+                                                                     [{{'$1', '$2'},
+                                                                       [],
+                                                                       [{{{element, 2, '$_'}, '$2'}}]}])),
+
                       check(Table,
                             fun ({N, [$x, C | _]}) when ((N rem 5) < 2) -> (C >= $0) andalso (C =< $9);
                                 ({N, [C | _]}) when is_float(N) -> (C >= $0) andalso (C =< $9);
@@ -1259,6 +1272,74 @@ t_select_replace(Config) when is_list(Config) ->
               ets:delete(Table)
       end,
       Tables),
+
+    %% Test key-safe match-specs are accepted
+    BigNum = (123 bsl 123),
+    RefcBin = list_to_binary(lists:seq(1,?heap_binary_size+1)),
+    Terms = [a, "hej", 123, 1.23, BigNum , <<"123">>, RefcBin, TestFun, self()],
+    EqPairs = fun(X,Y) ->
+                      [{ '$1', '$1'},
+                       { {X, Y}, {{X, Y}}},
+                       { {'$1', Y}, {{'$1', Y}}},
+                       { {{X, Y}}, {{{{X, Y}}}}},
+                       { {X}, {{X}}},
+                       { X, {const, X}},
+                       { {X,Y}, {const, {X,Y}}},
+                       { {X}, {const, {X}}},
+                       { {X, Y}, {{X, {const, Y}}}},
+                       { {X, {Y,'$1'}}, {{{const, X}, {{Y,'$1'}}}}},
+                       { [X, Y | '$1'], [X, Y | '$1']},
+                       { [{X, '$1'}, Y], [{{X, '$1'}}, Y]},
+                       { [{X, Y} | '$1'], [{const, {X, Y}} | '$1']},
+                       { [$p,$r,$e,$f,$i,$x | '$1'], [$p,$r,$e,$f,$i,$x | '$1']},
+                       { {[{X,Y}]}, {{[{{X,Y}}]}}},
+                       { {[{X,Y}]}, {{{const, [{X,Y}]}}}},
+                       { {[{X,Y}]}, {{[{const,{X,Y}}]}}}
+                      ]
+              end,
+
+    T2 = ets:new(x, []),
+    [lists:foreach(fun({A, B}) ->
+                           %% just check that matchspec is accepted
+                           0 = ets:select_replace(T2, [{{A, '$2', '$3'}, [], [{{B, '$3', '$2'}}]}])
+                   end,
+                   EqPairs(X,Y)) || X <- Terms, Y <- Terms],
+
+    %% Test key-unsafe matchspecs are rejected
+    NeqPairs = fun(X, Y) ->
+                      [{'$1', '$2'},
+                       {{X, Y}, {X, Y}},
+                       {{{X, Y}}, {{{X, Y}}}},
+                       {{X}, {{{X}}}},
+                       {{const, X}, {const, X}},
+                       {{const, {X,Y}}, {const, {X,Y}}},
+                       {'$1', {const, '$1'}},
+                       {{X}, {const, {{X}}}},
+                       {{X, {Y,'$1'}}, {{{const, X}, {Y,'$1'}}}},
+                       {[X, Y | '$1'], [X, Y]},
+                       {[X, Y], [X, Y | '$1']},
+                       {[{X, '$1'}, Y], [{X, '$1'}, Y]},
+                       {[$p,$r,$e,$f,$i,$x | '$1'], [$p,$r,$e,$f,$I,$x | '$1']},
+                       { {[{X,Y}]}, {{[{X,Y}]}}},
+                       { {[{X,Y}]}, {{{const, [{{X,Y}}]}}}},
+                       { {[{X,Y}]}, {{[{const,{{X,Y}}}]}}},
+                       {'_', '_'},
+                       {'$_', '$_'},
+                       {'$$', '$$'},
+                       {#{}, #{}},
+                       {#{X => '$1'}, #{X => '$1'}}
+                      ]
+              end,
+
+    [lists:foreach(fun({A, B}) ->
+                           %% just check that matchspec is rejected
+                           {'EXIT',{badarg,_}} = (catch ets:select_replace(T2, [{{A, '$2', '$3'}, [], [{{B, '$3', '$2'}}]}]))
+                   end,
+                   NeqPairs(X,Y)) || X <- Terms, Y <- Terms],
+
+
+    ets:delete(T2),
+
     verify_etsmem(EtsMem).
 
 %% Test that partly bound keys gives faster matches.
@@ -6120,7 +6201,6 @@ only_if_smp(Schedulers, Func) ->
     end.
 
 %% Copy-paste from emulator/test/binary_SUITE.erl
--define(heap_binary_size, 64).
 test_terms(Test_Func, Mode) ->
     garbage_collect(),
     Pib0 = process_info(self(),binary),