Fix non-iteratrable input of yieldEach

googleprojectzero · Oct 24, 2024 · 28659d0 · 28659d0
1 parent 0e20cd5
commit 28659d0
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/Sources/Fuzzilli/Base/ProgramBuilder.swift b/Sources/Fuzzilli/Base/ProgramBuilder.swift
@@ -2711,4 +2711,13 @@ public class ProgramBuilder {
             break
         }
     }
+
+    public func buildIteratorVariable(_ b: ProgramBuilder, _ it: Variable) -> Variable{
+        if(b.type(of: it).Is(.iterable)){
+            return it
+        }
+        let initialValues = (0..<Int.random(in: 1...5)).map({ _ in b.randomVariable() })
+        return b.createArray(with: initialValues)
+    }
+
 }
diff --git a/Sources/Fuzzilli/CodeGen/CodeGenerators.swift b/Sources/Fuzzilli/CodeGen/CodeGenerators.swift
@@ -769,13 +769,13 @@ public let CodeGenerators: [CodeGenerator] = [
         // These are "typically" used as arguments, so we don't directly generate a call operation here.
     },
 
-    RecursiveCodeGenerator("GeneratorFunctionGenerator") { b in
+    RecursiveCodeGenerator("GeneratorFunctionGenerator", inputs: .preferred(.iterable)) { b, it in
         let f = b.buildGeneratorFunction(with: b.randomParameters(), isStrict: probability(0.1)) { _ in
             b.buildRecursive()
             if probability(0.5) {
                 b.yield(b.randomVariable())
             } else {
-                b.yieldEach(b.randomVariable())
+                b.yieldEach(b.buildIteratorVariable(b, it))
             }
             b.doReturn(b.randomVariable())
         }
@@ -800,14 +800,14 @@ public let CodeGenerators: [CodeGenerator] = [
         // These are "typically" used as arguments, so we don't directly generate a call operation here.
     },
 
-    RecursiveCodeGenerator("AsyncGeneratorFunctionGenerator") { b in
+    RecursiveCodeGenerator("AsyncGeneratorFunctionGenerator", inputs: .preferred(.iterable)) { b, it in
         let f = b.buildAsyncGeneratorFunction(with: b.randomParameters(), isStrict: probability(0.1)) { _ in
             b.buildRecursive()
             b.await(b.randomVariable())
             if probability(0.5) {
                 b.yield(b.randomVariable())
             } else {
-                b.yieldEach(b.randomVariable())
+                b.yieldEach(b.buildIteratorVariable(b, it))
             }
             b.doReturn(b.randomVariable())
         }
@@ -1085,7 +1085,7 @@ public let CodeGenerators: [CodeGenerator] = [
             }
         } else {
             // TODO only do this when the value is iterable?
-            b.yieldEach(val)
+            b.yieldEach(b.buildIteratorVariable(b, val))
         }
     },