From c72d4a4427653a757438e83bdc8a07e7f57c0721 Mon Sep 17 00:00:00 2001
From: Arne Welzel <arne.welzel@corelight.com>
Date: Wed, 1 Nov 2023 17:27:47 +0100
Subject: [PATCH] Expr: Implement string in table[pattern] of X

Not sure how useful this is (and the implementation isn't optimized in
any way), but seems reasonable for consistency.

Vern suggested that set[pattern] can already be achieved via
set_to_regex(), so left out any set[pattern] variants.
---
 src/Expr.cc                                   | 22 +++++++++++++++++--
 .../Baseline/language.pattern-tables-when/out |  2 ++
 .../Baseline/language.pattern-tables/out      |  4 ++++
 .../btest/language/pattern-tables-when.zeek   | 13 ++++++++++-
 testing/btest/language/pattern-tables.zeek    |  4 ++++
 5 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/Expr.cc b/src/Expr.cc
index cf4e52f3cf..95e9f0d93d 100644
--- a/src/Expr.cc
+++ b/src/Expr.cc
@@ -3815,6 +3815,18 @@ InExpr::InExpr(ExprPtr arg_op1, ExprPtr arg_op2) : BinaryExpr(EXPR_IN, std::move
         }
     }
 
+    // Support <string> in table[pattern] of X
+    if ( op1->GetType()->Tag() == TYPE_STRING ) {
+        if ( op2->GetType()->Tag() == TYPE_TABLE ) {
+            const auto& table_type = op2->GetType()->AsTableType();
+
+            if ( table_type->IsPatternIndex() && table_type->Yield() ) {
+                SetType(base_type(TYPE_BOOL));
+                return;
+            }
+        }
+    }
+
     if ( op1->Tag() != EXPR_LIST )
         op1 = make_intrusive<ListExpr>(std::move(op1));
 
@@ -3853,8 +3865,14 @@ ValPtr InExpr::Fold(Val* v1, Val* v2) const {
         auto ind = v1->AsListVal()->Idx(0)->CoerceToUnsigned();
         res = ind < vv2->Size() && vv2->ValAt(ind);
     }
-    else
-        res = (bool)v2->AsTableVal()->Find({NewRef{}, v1});
+    else {
+        const auto& table_val = v2->AsTableVal();
+        const auto& table_type = table_val->GetType<zeek::TableType>();
+        if ( table_type->IsPatternIndex() && table_type->Yield() && v1->GetType()->Tag() == TYPE_STRING )
+            res = table_val->LookupPattern({NewRef{}, v1->AsStringVal()})->Size() > 0;
+        else
+            res = (bool)v2->AsTableVal()->Find({NewRef{}, v1});
+    }
 
     return val_mgr->Bool(res);
 }
diff --git a/testing/btest/Baseline/language.pattern-tables-when/out b/testing/btest/Baseline/language.pattern-tables-when/out
index d9e62df0e8..7367e825e6 100644
--- a/testing/btest/Baseline/language.pattern-tables-when/out
+++ b/testing/btest/Baseline/language.pattern-tables-when/out
@@ -4,3 +4,5 @@ populate_a()
 gotcha a, [42]
 populate_b()
 gotcha b, [4242]
+populate_c()
+gotcha c, [4711]
diff --git a/testing/btest/Baseline/language.pattern-tables/out b/testing/btest/Baseline/language.pattern-tables/out
index 209e9c48b2..a74c2d32eb 100644
--- a/testing/btest/Baseline/language.pattern-tables/out
+++ b/testing/btest/Baseline/language.pattern-tables/out
@@ -2,9 +2,13 @@
 indexing empty, 0
 single insert, match, [1]
 single insert, non-match, []
+single insert, in, T
+single insert, not-in, F
 multiple inserts, non-match, []
 multiple inserts, single match, [3]
 multiple inserts, double match, [1, 3]
+multiple insert, in, T
+multiple insert, not-in, F
 triple match, [1, 3, 4]
 embedded newline, /s operator, [6]
 no embedded newline, /s vs. no /s operator, [5, 6, 7]
diff --git a/testing/btest/language/pattern-tables-when.zeek b/testing/btest/language/pattern-tables-when.zeek
index df303418c2..d371cccf2c 100644
--- a/testing/btest/language/pattern-tables-when.zeek
+++ b/testing/btest/language/pattern-tables-when.zeek
@@ -6,11 +6,18 @@ global pt: table[pattern] of count;
 
 redef exit_only_after_terminate = T;
 
+event populate_c()
+	{
+	print "populate_c()";
+	pt[/c/] = 4711;
+	terminate();
+	}
+
 event populate_b()
 	{
 	print "populate_b()";
 	pt[/b/] = 4242;
-	terminate();
+	schedule 1msec { populate_c() };
 	}
 
 event populate_a()
@@ -38,6 +45,10 @@ event zeek_init()
 		print "gotcha b", pt["b"];
 	}
 
+	when ( "c" in pt ) {
+		print "gotcha c", pt["c"];
+	}
+
 	print "schedule populate";
 	schedule 1msec { populate_a() };
 	}
diff --git a/testing/btest/language/pattern-tables.zeek b/testing/btest/language/pattern-tables.zeek
index 7ca5c42e46..5331bd6315 100644
--- a/testing/btest/language/pattern-tables.zeek
+++ b/testing/btest/language/pattern-tables.zeek
@@ -13,6 +13,8 @@ event zeek_init()
 
 	print "single insert, match", pt["foo"];
 	print "single insert, non-match", pt["foox"];
+	print "single insert, in", "foo" in pt;
+	print "single insert, not-in", "foox" in pt;
 
 	pt[/bar/] = 2;
 	pt[/(foo|bletch)/] = 3;
@@ -20,6 +22,8 @@ event zeek_init()
 	print "multiple inserts, non-match", pt["x"];
 	print "multiple inserts, single match", pt["bletch"];
 	print "multiple inserts, double match", sort(pt["foo"]);
+	print "multiple insert, in", "foo" in pt;
+	print "multiple insert, not-in", "x" in pt;
 
 	pt[/(foo|bletch|xyz)/] = 4;
 	print "triple match", sort(pt["foo"]);