extend script coverage profiling to track whether conditionals evaluate to true/false

src/ScriptCoverageManager.cc

@@ -20,6 +20,8 @@ using namespace std;
 
 namespace zeek::detail {
 
+ScriptCoverageManager::ScriptCoverageManager() { pf = getenv("ZEEK_PROFILER_FILE"); }
+
 void ScriptCoverageManager::AddStmt(Stmt* s) {
     if ( ignoring != 0 || analysis_options.gen_ZAM )
         return;
@@ -34,14 +36,16 @@ void ScriptCoverageManager::AddFunction(IDPtr func_id, StmtPtr body) {
     func_instances.emplace_back(func_id, body);
 }
 
-bool ScriptCoverageManager::ReadStats() {
-    char* bf = getenv("ZEEK_PROFILER_FILE");
+void ScriptCoverageManager::AddConditional(Location cond_loc, std::string_view text, bool was_true) {
+    cond_instances.push_back({cond_loc, std::string(text), was_true});
+}
 
-    if ( ! bf )
+bool ScriptCoverageManager::ReadStats() {
+    if ( ! IsActive() )
         return false;
 
     std::ifstream ifs;
-    ifs.open(bf, std::ifstream::in);
+    ifs.open(pf, std::ifstream::in);
 
     if ( ! ifs )
         return false;
@@ -82,38 +86,43 @@ bool ScriptCoverageManager::ReadStats() {
 }
 
 bool ScriptCoverageManager::WriteStats() {
-    char* bf = getenv("ZEEK_PROFILER_FILE");
-
-    if ( ! bf )
+    if ( ! IsActive() )
         return false;
 
-    util::SafeDirname dirname{bf};
+    util::SafeDirname dirname{pf};
 
     if ( ! util::detail::ensure_intermediate_dirs(dirname.result.data()) ) {
-        reporter->Error("Failed to open ZEEK_PROFILER_FILE destination '%s' for writing", bf);
+        reporter->Error("Failed to open ZEEK_PROFILER_FILE destination '%s' for writing", pf);
         return false;
     }
 
     FILE* f;
-    const char* p = strstr(bf, "XXXXXX");
+    const char* p = strstr(pf, "XXXXXX");
 
     if ( p && ! p[6] ) {
         mode_t old_umask = umask(S_IXUSR | S_IRWXO | S_IRWXG);
-        int fd = mkstemp(bf);
+        auto pf_copy = strdup(pf);
+        if ( ! pf_copy ) {
+            reporter->InternalError("Memory exhausted in ScriptCoverageManager::WriteStats");
+            return false;
+        }
+
+        int fd = mkstemp(pf_copy);
+        free(pf_copy);
         umask(old_umask);
 
         if ( fd == -1 ) {
-            reporter->Error("Failed to generate unique file name from ZEEK_PROFILER_FILE: %s", bf);
+            reporter->Error("Failed to generate unique file name from ZEEK_PROFILER_FILE: %s", pf);
             return false;
         }
         f = fdopen(fd, "w");
     }
     else {
-        f = fopen(bf, "w");
+        f = fopen(pf, "w");
     }
 
     if ( ! f ) {
-        reporter->Error("Failed to open ZEEK_PROFILER_FILE destination '%s' for writing", bf);
+        reporter->Error("Failed to open ZEEK_PROFILER_FILE destination '%s' for writing", pf);
         return false;
     }
 
@@ -130,6 +139,9 @@ bool ScriptCoverageManager::WriteStats() {
         TrackUsage(body, desc, body->GetAccessCount());
     }
 
+    for ( const auto& [cond_loc, text, was_true] : cond_instances )
+        TrackUsage(&cond_loc, text, was_true ? 1 : 0);
+
     for ( auto& [location_info, cnt] : usage_map )
         Report(f, cnt, location_info.first, location_info.second);
 
@@ -137,9 +149,9 @@ bool ScriptCoverageManager::WriteStats() {
     return true;
 }
 
-void ScriptCoverageManager::TrackUsage(const ObjPtr& obj, std::string desc, uint64_t cnt) {
+void ScriptCoverageManager::TrackUsage(const Location* loc, std::string desc, uint64_t cnt) {
     ODesc location_info;
-    obj->GetLocationInfo()->Describe(&location_info);
+    loc->Describe(&location_info);
 
     static canonicalize_desc cd{delim};
     for_each(desc.begin(), desc.end(), cd);

src/ScriptCoverageManager.h

@@ -15,10 +15,19 @@ namespace zeek::detail {
 using ObjPtr = IntrusivePtr<Obj>;
 
 /**
- * A simple class for managing stats of Zeek script coverage across Zeek runs.
+ * A class for managing stats of Zeek script coverage across Zeek runs.
  */
 class ScriptCoverageManager {
 public:
+    ScriptCoverageManager();
+
+    /**
+     * Returns true if the manager is active (will do work), false if not.
+     *
+     * @return: true if active, false if not.
+     */
+    bool IsActive() const { return pf != nullptr; }
+
     /**
      * Imports Zeek script Stmt usage information from file pointed to by
      * environment variable ZEEK_PROFILER_FILE.
@@ -45,8 +54,14 @@ public:
 
     void AddStmt(Stmt* s);
     void AddFunction(IDPtr func_id, StmtPtr body);
+    void AddConditional(Location cond_loc, std::string_view text, bool was_true);
 
 private:
+    /**
+     * The name of the profile file, or nil if we're not profiling.
+     */
+    const char* pf;
+
     /**
      * The current, global ScriptCoverageManager instance creates this list at parse-time.
      */
@@ -57,6 +72,20 @@ private:
      */
     std::list<std::pair<IDPtr, StmtPtr>> func_instances;
 
+    /**
+     * Helper struct for tracking the result of @-directives.
+     */
+    struct Conditional {
+        Location loc;
+        std::string text;
+        bool result;
+    };
+
+    /**
+     * A similar list for tracking conditionals and whether they were true.
+     */
+    std::list<Conditional> cond_instances;
+
     /**
      * Indicates whether new statements will not be considered as part of
      * coverage statistics because it was marked with the @no-test tag.
@@ -95,7 +124,10 @@ private:
      * Tracks the usage of a given object with a given description
      * and a given coverage count.
      */
-    void TrackUsage(const ObjPtr& obj, std::string desc, uint64_t cnt);
+    void TrackUsage(const ObjPtr& obj, std::string desc, uint64_t cnt) {
+        TrackUsage(obj->GetLocationInfo(), std::move(desc), cnt);
+    }
+    void TrackUsage(const Location* loc, std::string desc, uint64_t cnt);
 
     /**
      * Reports a single coverage instance.

src/scan.l

@@ -27,6 +27,7 @@
 #include "zeek/Scope.h"
 #include "zeek/ZeekString.h"
 #include "zeek/DNS_Mgr.h"
+#include "zeek/Desc.h"
 #include "zeek/Expr.h"
 #include "zeek/Func.h"
 #include "zeek/Stmt.h"
@@ -38,6 +39,7 @@
 #include "zeek/Reporter.h"
 #include "zeek/RE.h"
 #include "zeek/RunState.h"
+#include "zeek/ScriptCoverageManager.h"
 #include "zeek/Traverse.h"
 #include "zeek/module_util.h"
 #include "zeek/ScannedFile.h"
@@ -866,7 +868,7 @@ static int load_files(const char* orig_file) {
     yy_switch_to_buffer(buffer);
     yylloc.first_line = yylloc.last_line = line_number = 1;
 
-	return switch_to(file_path.c_str(), buffer);
+    return switch_to(file_path.c_str(), buffer);
 }
 
 void begin_RE() { BEGIN(RE); }
@@ -891,9 +893,23 @@ public:
     std::vector<const zeek::detail::NameExpr*> local_names;
 };
 
-static void begin_ignoring() {
-    if_stack.push_back(conditional_depth);
-    BEGIN(IGNORE);
+static void set_ignoring(bool start_ignoring, std::string_view directive, std::optional<std::string_view> value = std::nullopt) {
+    if ( script_coverage_mgr.IsActive() ) {
+        std::string text = std::string(directive);
+
+        if ( value ) {
+             text += " ( ";
+             text += *value;
+             text += " )";
+        }
+
+        script_coverage_mgr.AddConditional(GetCurrentLocation(), text, ! start_ignoring);
+    }
+
+    if ( start_ignoring ) {
+        if_stack.push_back(conditional_depth);
+        BEGIN(IGNORE);
+    }
 }
 
 static void resume_processing() {
@@ -920,8 +936,7 @@ void do_atif(zeek::detail::Expr* expr) {
         return;
     }
 
-    if ( ! val->AsBool() )
-        begin_ignoring();
+    set_ignoring(! val->AsBool(), "@if", zeek::obj_desc_short(expr));
 }
 
 void do_atifdef(const char* id) {
@@ -934,8 +949,7 @@ void do_atifdef(const char* id) {
     if ( ! found )
         found = (zeek::detail::module_names().count(id) != 0);
 
-    if ( ! found )
-        begin_ignoring();
+    set_ignoring(! found, "@ifdef", id);
 }
 
 void do_atifndef(const char* id) {
@@ -948,8 +962,7 @@ void do_atifndef(const char* id) {
     if ( ! found )
         found = (zeek::detail::module_names().count(id) != 0);
 
-    if ( found )
-        begin_ignoring();
+    set_ignoring(found, "@ifndef", id);
 }
 
 void do_atelse() {
@@ -960,9 +973,11 @@ void do_atelse() {
         return;
 
     if ( YY_START == INITIAL )
-        begin_ignoring();
-    else
+        set_ignoring(true, "@else");
+    else {
+        set_ignoring(false, "@else");
         resume_processing();
+    }
 }
 
 void do_atendif() {

testing/btest/Baseline/coverage.zeek-profiler-file/step1.out

@@ -1,3 +1,4 @@
 ### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
-1	./profiling-test1.zeek, line 2	print new conn;
-1	./profiling-test1.zeek, lines 1-2	event new_connection BODY
+1	./profiling-test1.zeek, line 2	@if ( T )
+1	./profiling-test1.zeek, line 4	print new conn;
+1	./profiling-test1.zeek, lines 3-4	event new_connection BODY

testing/btest/Baseline/coverage.zeek-profiler-file/step2.out

@@ -1,3 +1,4 @@
 ### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
-2	./profiling-test1.zeek, line 2	print new conn;
-2	./profiling-test1.zeek, lines 1-2	event new_connection BODY
+2	./profiling-test1.zeek, line 2	@if ( T )
+2	./profiling-test1.zeek, line 4	print new conn;
+2	./profiling-test1.zeek, lines 3-4	event new_connection BODY

testing/btest/Baseline/coverage.zeek-profiler-file/step3.out

@@ -1,5 +1,16 @@
 ### BTest baseline data generated by btest-diff. Do not edit. Use "btest -U/-u" to update. Requires BTest >= 0.63.
-2	./profiling-test1.zeek, line 2	print new conn;
-2	./profiling-test1.zeek, lines 1-2	event new_connection BODY
-1	./profiling-test2.zeek, line 2	print new conn;
-1	./profiling-test2.zeek, lines 1-2	event new_connection BODY
+2	./profiling-test1.zeek, line 2	@if ( T )
+2	./profiling-test1.zeek, line 4	print new conn;
+2	./profiling-test1.zeek, lines 3-4	event new_connection BODY
+0	./profiling-test2.zeek, line 1	@if ( F )
+0	./profiling-test2.zeek, line 11	@ifndef ( Conn::Info )
+1	./profiling-test2.zeek, line 13	@else
+1	./profiling-test2.zeek, line 14	event zeek_init BODY
+1	./profiling-test2.zeek, line 14	print Conn::Info;
+1	./profiling-test2.zeek, line 2	@else
+1	./profiling-test2.zeek, line 4	print new conn;
+1	./profiling-test2.zeek, line 6	@ifdef ( Conn::Info )
+1	./profiling-test2.zeek, line 7	event zeek_init BODY
+1	./profiling-test2.zeek, line 7	print Conn::Info;
+0	./profiling-test2.zeek, line 8	@else
+1	./profiling-test2.zeek, lines 3-4	event new_connection BODY

testing/btest/coverage/zeek-profiler-file.zeek

@@ -13,11 +13,27 @@
 # @TEST-EXEC: btest-diff step3.out
 
 # @TEST-START-FILE profiling-test1.zeek
+
+@if ( T )
 event new_connection(c: connection)
 	{ print "new conn"; }
+@endif
 # @TEST-END-FILE
 
 # @TEST-START-FILE profiling-test2.zeek
+@if ( F )
+@else
 event new_connection(c: connection)
 	{ print "new conn"; }
+@endif
+@ifdef ( Conn::Info )
+event zeek_init() { print Conn::Info; }
+@else
+event zeek_init() { print "No Conn::Info"; }
+@endif
+@ifndef ( Conn::Info )
+event zeek_init() { print "No Conn::Info"; }
+@else
+event zeek_init() { print Conn::Info; }
+@endif
 # @TEST-END-FILE