From 1a1d3edd503d1f529dc053c57fa40db3a93797a5 Mon Sep 17 00:00:00 2001
From: Tobias Wienand <tobiaswienand11235@gmail.com>
Date: Fri, 23 Aug 2024 12:56:57 +0200
Subject: [PATCH] Implements robust break context identification and tests

---
 Sources/Fuzzilli/FuzzIL/Analyzer.swift |  8 +++++
 Tests/FuzzilliTests/AnalyzerTest.swift | 43 ++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/Sources/Fuzzilli/FuzzIL/Analyzer.swift b/Sources/Fuzzilli/FuzzIL/Analyzer.swift
index c43f9904..3297a691 100644
--- a/Sources/Fuzzilli/FuzzIL/Analyzer.swift
+++ b/Sources/Fuzzilli/FuzzIL/Analyzer.swift
@@ -168,6 +168,14 @@ struct ContextAnalyzer: Analyzer {
 
                 newContext.formUnion(contextStack.secondToTop)
             }
+            
+            // If we are in a loop, we don't want to propagate the switch context and vice versa.
+            if (instr.op.contextOpened.contains(.switchBlock) || instr.op.contextOpened.contains(.switchCase)) {
+                newContext.remove(.loop)
+            } else if (instr.op.contextOpened.contains(.loop)) {
+                newContext.remove(.switchBlock) 
+                newContext.remove(.switchCase)
+            }
             contextStack.push(newContext)
         }
     }
diff --git a/Tests/FuzzilliTests/AnalyzerTest.swift b/Tests/FuzzilliTests/AnalyzerTest.swift
index fd7afffd..1e1c4eb7 100644
--- a/Tests/FuzzilliTests/AnalyzerTest.swift
+++ b/Tests/FuzzilliTests/AnalyzerTest.swift
@@ -304,4 +304,47 @@ class AnalyzerTests: XCTestCase {
 
         let _  = b.finalize()
     }
+
+    func testBreakContext() {
+        let fuzzer = makeMockFuzzer()
+        let b = fuzzer.makeBuilder()
+
+        let case1 = b.loadInt(1337)
+        let case2 = b.loadInt(9001)
+        
+        // Test case 1: switch -> loop -> switch
+        b.buildSwitch(on: case1) { outer_switch in
+            XCTAssertEqual(b.context, .switchBlock)
+            outer_switch.addCase(case1) {
+                XCTAssertEqual(b.context, [.javascript, .switchCase])
+                b.buildWhileLoop({ b.loadBool(true) }) {        
+                    XCTAssertEqual(b.context, [.javascript, .loop])
+                    b.buildSwitch(on: case2) { inner_switch in
+                        XCTAssertEqual(b.context, .switchBlock)
+                        inner_switch.addCase(case2) {
+                            XCTAssertEqual(b.context, [.javascript, .switchCase])
+                        }
+                    }
+                } 
+            }
+        }
+        XCTAssertEqual(b.context, .javascript)
+
+        // Test case 2: loop -> switch -> loop
+        b.buildWhileLoop({ b.loadBool(true) }) {
+            XCTAssertEqual(b.context, [.javascript, .loop])
+            b.buildSwitch(on: case1) { swtch in
+                XCTAssertEqual(b.context, .switchBlock)
+                swtch.addCase(case1) {
+                    XCTAssertEqual(b.context, [.javascript, .switchCase])
+                    b.buildWhileLoop({ b.loadBool(true) }) {
+                        XCTAssertEqual(b.context, [.javascript, .loop])
+                    }
+                }
+            }
+        }
+        XCTAssertEqual(b.context, .javascript)
+        
+        let _ = b.finalize()
+    }
 }