From 1870d266848bc8f146151b02be35e706fd736823 Mon Sep 17 00:00:00 2001
From: Tim Wojtulewicz <tim@corelight.com>
Date: Tue, 2 Aug 2022 12:19:36 -0700
Subject: [PATCH] GH-2034: Store module names and use them in lookups for ifdef

---
 src/parse.y                                   |  3 ++
 src/scan.l                                    | 16 +++++++--
 testing/btest/Baseline/language.at-ifdef/out  |  2 ++
 testing/btest/Baseline/language.at-ifndef/out |  2 ++
 testing/btest/language/at-ifdef.zeek          | 34 +++++++++++++++++-
 testing/btest/language/at-ifndef.zeek         | 35 ++++++++++++++++++-
 6 files changed, 88 insertions(+), 4 deletions(-)

diff --git a/src/parse.y b/src/parse.y
index 1336d177cf..6a43ba9cf9 100644
--- a/src/parse.y
+++ b/src/parse.y
@@ -156,6 +156,8 @@ static int func_hdr_cond_epoch = 0;
 EnumType* cur_enum_type = nullptr;
 static ID* cur_decl_type_id = nullptr;
 
+std::set<std::string> module_names;
+
 static void parse_new_enum(void)
 	{
 	// Starting a new enum definition.
@@ -1269,6 +1271,7 @@ decl:
 		TOK_MODULE TOK_ID ';'
 			{
 			current_module = $2;
+			module_names.insert($2);
 			zeekygen_mgr->ModuleUsage(::filename, current_module);
 			}
 
diff --git a/src/scan.l b/src/scan.l
index 169af2932b..c9397e7431 100644
--- a/src/scan.l
+++ b/src/scan.l
@@ -80,6 +80,8 @@ const char* last_tok_filename = 0;
 const char* last_last_tok_filename = 0;
 char last_tok[128];
 
+extern std::set<std::string> module_names;
+
 #define YY_USER_ACTION	strncpy(last_tok, yytext, sizeof(last_tok) - 1); \
                         last_last_tok_filename = last_tok_filename; \
                         last_tok_filename = ::filename;
@@ -796,7 +798,12 @@ void do_atifdef(const char* id)
 
 	const auto& i = zeek::detail::lookup_ID(id, zeek::detail::current_module.c_str());
 
-	if ( ! i )
+	// If the ID didn't exist, look to see if this is a module name instead.
+	bool found = (i != nullptr);
+	if ( ! found )
+		found = (module_names.count(id) != 0);
+
+	if ( ! found )
 		begin_ignoring();
 	}
 
@@ -806,7 +813,12 @@ void do_atifndef(const char *id)
 
 	const auto& i = zeek::detail::lookup_ID(id, zeek::detail::current_module.c_str());
 
-	if ( i )
+	// If the ID didn't exist, look to see if this is a module name instead.
+	bool found = (i != nullptr);
+	if ( ! found )
+		found = (module_names.count(id) != 0);
+
+	if ( found )
 		begin_ignoring();
 	}
 
diff --git a/testing/btest/Baseline/language.at-ifdef/out b/testing/btest/Baseline/language.at-ifdef/out
index e66a13b4f4..baeef79e22 100644
--- a/testing/btest/Baseline/language.at-ifdef/out
+++ b/testing/btest/Baseline/language.at-ifdef/out
@@ -2,3 +2,5 @@
 @ifdef (PASS)
 @ifdef...@else (PASS)
 @ifdef...@else (PASS)
+@ifdef module name (PASS)
+@ifdef child variable (PASS)
diff --git a/testing/btest/Baseline/language.at-ifndef/out b/testing/btest/Baseline/language.at-ifndef/out
index 7dd86b5463..e9de53fc54 100644
--- a/testing/btest/Baseline/language.at-ifndef/out
+++ b/testing/btest/Baseline/language.at-ifndef/out
@@ -2,3 +2,5 @@
 @ifndef (PASS)
 @ifndef...@else (PASS)
 @ifndef...@else (PASS)
+@ifndef module name (PASS)
+@ifndef child variable (PASS)
diff --git a/testing/btest/language/at-ifdef.zeek b/testing/btest/language/at-ifdef.zeek
index df20c972de..78dc6a48f5 100644
--- a/testing/btest/language/at-ifdef.zeek
+++ b/testing/btest/language/at-ifdef.zeek
@@ -1,7 +1,9 @@
-# @TEST-EXEC: zeek -b %INPUT >out
+# @TEST-EXEC: zeek -b bar.zeek main.zeek >out
 # @TEST-EXEC: btest-diff out
 # @TEST-EXEC: btest-diff .stderr
 
+@TEST-START-FILE main.zeek
+
 function test_case(msg: string, expect: bool)
         {
         print fmt("%s (%s)", msg, expect ? "PASS" : "FAIL");
@@ -46,5 +48,35 @@ event zeek_init()
 
 	test_case( "@ifdef...@else", xyz == 1 );
 
+	xyz = 0;
+
+	@ifdef ( Bar )
+		xyz += 1;
+	@else
+		xyz += 2;
+	@endif
+
+	test_case( "@ifdef module name", xyz == 1 );
+
+	xyz = 0;
+
+	@ifdef ( Bar::exists )
+		xyz += 1;
+	@else
+		xyz += 2;
+	@endif
+
+	test_case( "@ifdef child variable", xyz == 1 );
 }
 
+@TEST-END-FILE
+
+@TEST-START-FILE bar.zeek
+
+module Bar;
+
+export {
+  option exists = T;
+}
+
+@TEST-END-FILE
diff --git a/testing/btest/language/at-ifndef.zeek b/testing/btest/language/at-ifndef.zeek
index 6d3c3325dc..5debbe765a 100644
--- a/testing/btest/language/at-ifndef.zeek
+++ b/testing/btest/language/at-ifndef.zeek
@@ -1,7 +1,9 @@
-# @TEST-EXEC: zeek -b %INPUT >out
+# @TEST-EXEC: zeek -b bar.zeek main.zeek >out
 # @TEST-EXEC: btest-diff out
 # @TEST-EXEC: btest-diff .stderr
 
+@TEST-START-FILE main.zeek
+
 function test_case(msg: string, expect: bool)
         {
         print fmt("%s (%s)", msg, expect ? "PASS" : "FAIL");
@@ -46,5 +48,36 @@ event zeek_init()
 
 	test_case( "@ifndef...@else", xyz == 2 );
 
+	xyz = 0;
+
+	@ifndef ( Bar )
+		xyz += 1;
+	@else
+		xyz += 2;
+	@endif
+
+	test_case( "@ifndef module name", xyz == 2 );
+
+	xyz = 0;
+
+	@ifndef ( Bar::exists )
+		xyz += 1;
+	@else
+		xyz += 2;
+	@endif
+
+	test_case( "@ifndef child variable", xyz == 2 );
+
 }
 
+@TEST-END-FILE
+
+@TEST-START-FILE bar.zeek
+
+module Bar;
+
+export {
+  option exists = T;
+}
+
+@TEST-END-FILE