From 257897c60be7d7ac35ed993794d95b6a9f4c7573 Mon Sep 17 00:00:00 2001
From: Dave Thaler <dthaler1968@gmail.com>
Date: Sat, 5 Oct 2024 12:58:01 -0700
Subject: [PATCH] Fix bpf2bpf bug when multiple branches in a subprogram go to
 exit (#715)

* Fix bpf2bpf bug when multiple branches in a subprogram go to exit

The bug was that the successor label could be processed multiple
times if there were multiple predecessors, resulting in multiple
exit instructions appearing in the block, which in turn would
result in restore_callee_saved_registers() being called multiple
times, giving incorrect results.

* Added a YAML test case for this.
* clang format and tidy asm_cfg.cpp

---------

Signed-off-by: Dave Thaler <dthaler1968@gmail.com>
Co-authored-by: Elazar Gershuni <elazarg@gmail.com>
---
 src/asm_cfg.cpp          | 16 +++++++++-------
 test-data/calllocal.yaml | 24 ++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/asm_cfg.cpp b/src/asm_cfg.cpp
index 84fd81790..5d6439bff 100644
--- a/src/asm_cfg.cpp
+++ b/src/asm_cfg.cpp
@@ -5,7 +5,6 @@
 #include <algorithm>
 #include <map>
 #include <optional>
-#include <queue>
 #include <string>
 #include <vector>
 
@@ -59,11 +58,11 @@ static void add_cfg_nodes(cfg_t& cfg, const label_t& caller_label, const label_t
 
     // Walk the transitive closure of CFG nodes starting at entry_label and ending at
     // any exit instruction.
-    std::queue<label_t> macro_labels{{entry_label}};
+    std::set macro_labels{entry_label};
     std::set seen_labels{entry_label};
     while (!macro_labels.empty()) {
-        label_t macro_label = macro_labels.front();
-        macro_labels.pop();
+        label_t macro_label = *macro_labels.begin();
+        macro_labels.erase(macro_label);
 
         if (stack_frame_prefix == macro_label.stack_frame_prefix) {
             throw std::runtime_error{stack_frame_prefix + ": illegal recursion"};
@@ -105,7 +104,9 @@ static void add_cfg_nodes(cfg_t& cfg, const label_t& caller_label, const label_t
                 bb >> exit_to_node;
             } else if (!seen_labels.contains(next_macro_label)) {
                 // Push any other unprocessed successor label onto the list to be processed.
-                macro_labels.push(next_macro_label);
+                if (!macro_labels.contains(next_macro_label)) {
+                    macro_labels.insert(next_macro_label);
+                }
                 seen_labels.insert(macro_label);
             }
         }
@@ -249,7 +250,7 @@ static cfg_t to_nondet(const cfg_t& cfg) {
         auto nextlist = bb.next_blocks_set();
         if (nextlist.size() == 2) {
             label_t mid_label = this_label;
-            Jmp jmp = std::get<Jmp>(*bb.rbegin());
+            auto jmp = std::get<Jmp>(*bb.rbegin());
 
             nextlist.erase(jmp.target);
             label_t fallthrough = *nextlist.begin();
@@ -358,7 +359,8 @@ std::map<std::string, int> collect_stats(const cfg_t& cfg) {
     return res;
 }
 
-cfg_t prepare_cfg(const InstructionSeq& prog, const program_info& info, bool simplify, bool must_have_exit) {
+cfg_t prepare_cfg(const InstructionSeq& prog, const program_info& info, const bool simplify,
+                  const bool must_have_exit) {
     // Convert the instruction sequence to a deterministic control-flow graph.
     cfg_t det_cfg = instruction_seq_to_cfg(prog, must_have_exit);
 
diff --git a/test-data/calllocal.yaml b/test-data/calllocal.yaml
index 4bb735ad0..fa2c811c1 100644
--- a/test-data/calllocal.yaml
+++ b/test-data/calllocal.yaml
@@ -203,6 +203,30 @@ code:
     r6 = 1
     exit
 
+post:
+  - r0.type=number
+  - r0.svalue=0
+  - r0.uvalue=0
+  - r6.type=number
+  - r6.svalue=0
+  - r6.uvalue=0
+---
+test-case: call local with multiple branches
+
+pre: ["r0.type=number"]
+
+code:
+  <start>: |
+    r6 = 0
+    call <sub>
+    r0 = 0
+    exit
+  <sub>: |
+    if r0 == 0 goto <done>
+    r6 = 1
+  <done>: |
+    exit
+
 post:
   - r0.type=number
   - r0.svalue=0