Fix non-iteratrable input of yieldEach (#463)

googleprojectzero · Nov 15, 2024 · 0f52182 · 0f52182
1 parent d6f77d7
commit 0f52182
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 11 deletions.
diff --git a/Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift b/Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift
@@ -90,7 +90,8 @@ public let codeGeneratorWeights = [
     "AsyncGeneratorFunctionGenerator":          1,
     "ConstructorGenerator":                     5,
     "SubroutineReturnGenerator":                3,
-    "YieldGenerator":                           2,
+    "YieldGenerator":                           5,
+    "YieldEachGenerator":                       5,
     "AwaitGenerator":                           2,
     "PropertyRetrievalGenerator":               20,
     "PropertyAssignmentGenerator":              20,

diff --git a/Sources/Fuzzilli/CodeGen/CodeGenerators.swift b/Sources/Fuzzilli/CodeGen/CodeGenerators.swift
@@ -775,7 +775,9 @@ public let CodeGenerators: [CodeGenerator] = [
             if probability(0.5) {
                 b.yield(b.randomVariable())
             } else {
-                b.yieldEach(b.randomVariable())
+                let randomVariables = b.randomVariables(n: Int.random(in: 1...5))
+                let array = b.createArray(with: randomVariables)
+                b.yieldEach(array)
             }
             b.doReturn(b.randomVariable())
         }
@@ -807,7 +809,9 @@ public let CodeGenerators: [CodeGenerator] = [
             if probability(0.5) {
                 b.yield(b.randomVariable())
             } else {
-                b.yieldEach(b.randomVariable())
+                let randomVariables = b.randomVariables(n: Int.random(in: 1...5))
+                let array = b.createArray(with: randomVariables)
+                b.yieldEach(array)
             }
             b.doReturn(b.randomVariable())
         }
@@ -1077,18 +1081,18 @@ public let CodeGenerators: [CodeGenerator] = [
 
     CodeGenerator("YieldGenerator", inContext: .generatorFunction, inputs: .one) { b, val in
         assert(b.context.contains(.generatorFunction))
-        if probability(0.5) {
-            if probability(0.9) {
-                b.yield(val)
-            } else {
-                b.yield()
-            }
+        if probability(0.9) {
+            b.yield(val)
         } else {
-            // TODO only do this when the value is iterable?
-            b.yieldEach(val)
+            b.yield()
         }
     },
 
+    CodeGenerator("YieldEachGenerator", inContext: .generatorFunction, inputs: .required(.iterable)) { b, val in
+        assert(b.context.contains(.generatorFunction))
+        b.yieldEach(val)
+    },
+
     CodeGenerator("AwaitGenerator", inContext: .asyncFunction, inputs: .one) { b, val in
         assert(b.context.contains(.asyncFunction))
         b.await(val)