Add --save option and limit vinpgen length

Author: Merricx

src/zkregex_fuzzer/cli.py

@@ -4,10 +4,13 @@
 
 import argparse
 import os
+import random
+import uuid
 from pathlib import Path
 from zkregex_fuzzer.fuzzer import fuzz_with_database, fuzz_with_grammar
 from zkregex_fuzzer.grammar import REGEX_GRAMMAR
 from zkregex_fuzzer.configs import TARGETS, VALID_INPUT_GENERATORS, GENERATORS
+from zkregex_fuzzer.harness import HarnessStatus
 from zkregex_fuzzer.logger import logger
 from zkregex_fuzzer.runner.circom import CircomSubprocess, SnarkjsSubprocess, ZkRegexSubprocess
 
@@ -43,7 +46,13 @@ def main():
         help=f"The valid input generator to use for the fuzzer (options: {list(VALID_INPUT_GENERATORS.keys())})."
     )
     parser.add_argument(
-        "--output",
+        "--save",
+        choices=[status.name for status in HarnessStatus],
+        nargs="*",
+        help="Save reproducible files according to the specified Harness status",
+    )
+    parser.add_argument(
+        "--save-output",
         type=str,
         default=os.getcwd(),
         help=f"The output path where the reproducible files will be stored (default: .)"
@@ -148,6 +157,13 @@ def main():
     print(f"Max depth: {args.grammar_max_depth}")
     print("-" * 80)
 
+    kwargs = vars(args)
+
+    # set global seed
+    seed = str(uuid.uuid4())
+    kwargs['seed'] = seed
+    random.seed(seed)
+
     if args.fuzzer == "grammar":
         fuzz_with_grammar(
             target_grammar="basic",
@@ -156,15 +172,15 @@ def main():
             regex_num=args.regex_num,
             inputs_num=args.inputs_num,
             max_depth=args.grammar_max_depth,
-            kwargs=vars(args)
+            kwargs=kwargs
         )
     elif args.fuzzer == "database":
         fuzz_with_database(
             target_implementation=args.target,
             oracle_params=(args.oracle == "valid", args.valid_input_generator),
             regex_num=args.regex_num,
             inputs_num=args.inputs_num,
-            kwargs=vars(args)
+            kwargs=kwargs
         )
 
 

src/zkregex_fuzzer/fuzzer.py

@@ -2,7 +2,6 @@
 Implements the logic for generating regexes using The Fuzzing Book's GrammarFuzzer.
 """
 
-import re
 from fuzzingbook.Grammars import simple_grammar_fuzzer, Grammar
 from zkregex_fuzzer.runner.base_runner import Runner
 from zkregex_fuzzer.transformers import regex_to_grammar
@@ -31,15 +30,8 @@ def fuzz_with_grammar(
     regex_generator = GrammarRegexGenerator(grammar, "<start>")
     regexes = regex_generator.generate_many(regex_num)
     logger.info(f"Generated {len(regexes)} regexes.")
-    oracle, oracle_generator = oracle_params
-    if oracle:
-        generator = VALID_INPUT_GENERATORS[oracle_generator]
-        logger.info(f"Generating {inputs_num} inputs for each regex.")
-        regexes_inputs = [generator(regex).generate_many(inputs_num) for regex in regexes]
-    else:
-        raise NotImplementedError("Oracle not implemented")
-
-    fuzz_with_regexes(regexes, regexes_inputs, target_runner, oracle, kwargs)
+    
+    fuzz_with_regexes(regexes, inputs_num, target_runner, oracle_params, kwargs)
 
 def fuzz_with_database(
         target_implementation: str,
@@ -56,26 +48,35 @@ def fuzz_with_database(
     regex_generator = DatabaseRegexGenerator()
     regexes = regex_generator.generate_many(regex_num)
     logger.info(f"Generated {len(regexes)} regexes.")
-    oracle, oracle_generator = oracle_params
-    if oracle:
-        generator = VALID_INPUT_GENERATORS[oracle_generator]
-        logger.info(f"Generating {inputs_num} inputs for each regex.")
-        regexes_inputs = [generator(regex).generate_many(inputs_num) for regex in regexes]
-    else:
-        raise NotImplementedError("Oracle not implemented")
 
-    fuzz_with_regexes(regexes, regexes_inputs, target_runner, oracle, kwargs)
+    fuzz_with_regexes(regexes, inputs_num, target_runner, oracle_params, kwargs)
 
 def fuzz_with_regexes(
         regexes: list[str],
-        regexes_inputs: list[str],
-        target_runner: Runner,
-        oracle: bool,
+        inputs_num: int,
+        target_runner: type[Runner],
+        oracle_params: tuple[bool, str],
         kwargs: dict,
     ):
     """
     Fuzz test with pre-seeded regexes.
     """
+    max_input_size = kwargs.get("circom_max_input_size", None)
+    oracle, oracle_generator = oracle_params
+    if oracle:
+        generator = VALID_INPUT_GENERATORS[oracle_generator]
+        logger.info(f"Generating {inputs_num} inputs for each regex.")
+        regexes_inputs = []
+        for regex in regexes:
+            try:
+                regex_inputs = generator(regex).generate_many(inputs_num, max_input_size)
+            except ValueError as e:
+                logger.warning(e)
+                regex_inputs = []
+            regexes_inputs.append(regex_inputs)
+    else:
+        raise NotImplementedError("Oracle not implemented")
+
     # We should use the PythonReRunner to check the validity of the regexes and the inputs.
     # If there is a bug in the PythonReRunner, we might not find it as we will think that 
     # either the regex or the input is invalid.

src/zkregex_fuzzer/harness.py

@@ -6,7 +6,9 @@
     In the future, we could pass invalid regexes and check if the secondary runner fails to compile.
 """
 
-from typing import Type, List
+import json
+from pathlib import Path
+from typing import Type, List, Union
 from enum import Enum
 from dataclasses import dataclass
 from zkregex_fuzzer.logger import logger
@@ -31,6 +33,37 @@ class HarnessResult:
     # Error message (if any)
     error_message: str = ""
 
+def _return_harness_result(
+        result: HarnessResult,
+        status_to_save: list,
+        output_path: str,
+        runner: Union[Runner, None],
+        kwargs: dict,
+    ):
+    if runner:
+        if result.status.name in status_to_save:
+
+            metadata = {
+                "seed": kwargs.get("seed"),
+                "target": kwargs.get("target"),
+                "oracle": kwargs.get("oracle"),
+                "input_generator": kwargs.get("valid_input_generator"),
+                "fuzzer": kwargs.get("fuzzer"),
+                "regex_num": kwargs.get("regex_num"),
+                "inputs_num": kwargs.get("inputs_num"),
+            }
+
+            dir_path = runner.save(output_path)
+
+            metadata_json = json.dumps(metadata)
+            metadata_path = Path(dir_path) / "metadata.json"
+            with open(metadata_path.absolute(), "w") as f:
+                f.write(metadata_json)
+
+        runner.clean()
+    
+    return result
+
 def harness(
         regex: str,
         primary_runner_cls: Type[Runner],
@@ -54,26 +87,53 @@ def harness(
     """
     regex = regex
     inp_num = len(inputs)
-    output_path = kwargs.get("output")
+    output_path = kwargs.get("save_output", "")
+    status_to_save = kwargs.get("save", None) or []
+    
     try:
         primary_runner = primary_runner_cls(regex, {})
     except RegexCompileError as e:
-        return HarnessResult(regex, inp_num, oracle, [], HarnessStatus.INVALID_SEED, str(e))
+        return _return_harness_result(
+            HarnessResult(regex, inp_num, oracle, [], HarnessStatus.INVALID_SEED, str(e)),
+            status_to_save,
+            output_path,
+            None,
+            kwargs,
+        )
 
     try:
         secondary_runner = secondary_runner_cls(regex, kwargs)
     except RegexCompileError as e:
-        return HarnessResult(regex, inp_num, oracle, [], HarnessStatus.COMPILE_ERROR, str(e))
+        return _return_harness_result(
+            HarnessResult(regex, inp_num, oracle, [], HarnessStatus.COMPILE_ERROR, str(e)),
+            status_to_save,
+            output_path,
+            None,
+            kwargs
+        )
 
     failed_inputs = []
+    
     for input in inputs:
         primary_runner_str = None
         try:
             primary_runner_status, primary_runner_str = primary_runner.match(input)
             if primary_runner_status != oracle:
-                return HarnessResult(regex, inp_num, oracle, [], HarnessStatus.INVALID_SEED)
+                return _return_harness_result(
+                    HarnessResult(regex, inp_num, oracle, [], HarnessStatus.INVALID_SEED),
+                    status_to_save,
+                    output_path,
+                    None,
+                    kwargs
+                )
         except RegexRunError as e:
-            return HarnessResult(regex, inp_num, oracle, [], HarnessStatus.INVALID_SEED, str(e))
+            return _return_harness_result(
+                HarnessResult(regex, inp_num, oracle, [], HarnessStatus.INVALID_SEED, str(e)),
+                status_to_save,
+                output_path,
+                None,
+                kwargs
+            )
         try:
             secondary_runner_status, secondary_runner_str = secondary_runner.match(input)
             if secondary_runner_status != oracle:
@@ -84,11 +144,27 @@ def harness(
                 # failed_inputs.append(input)
         except RegexRunError as e:
             secondary_runner.save(output_path)
-            return HarnessResult(regex, inp_num, oracle, [input], HarnessStatus.RUN_ERROR, str(e))
+            return _return_harness_result(
+                HarnessResult(regex, inp_num, oracle, [input], HarnessStatus.RUN_ERROR, str(e)),
+                status_to_save,
+                output_path,
+                secondary_runner,
+                kwargs
+            )
 
     if len(failed_inputs) > 0:
-        secondary_runner.save(output_path)
-        return HarnessResult(regex, inp_num, oracle, failed_inputs, HarnessStatus.FAILED)
+        return _return_harness_result(
+            HarnessResult(regex, inp_num, oracle, failed_inputs, HarnessStatus.FAILED),
+            status_to_save,
+            output_path,
+            secondary_runner,
+            kwargs
+        )
 
-    secondary_runner.clean()
-    return HarnessResult(regex, inp_num, oracle, [], HarnessStatus.SUCCESS)
+    return _return_harness_result(
+        HarnessResult(regex, inp_num, oracle, [], HarnessStatus.SUCCESS),
+        status_to_save,
+        output_path,
+        secondary_runner,
+        kwargs
+    )

src/zkregex_fuzzer/regexgen.py

@@ -78,7 +78,7 @@ class DatabaseRegexGenerator(RegexGenerator):
     Generate regexes using a database of regexes.
     """
 
-    def __init__(self, dir_path: str = None):
+    def __init__(self, dir_path: str = ""):
         dir_path = dir_path or self._get_default_path()
         self.database = self._get_database_from_path(dir_path)
 

src/zkregex_fuzzer/runner/base_runner.py

@@ -56,9 +56,9 @@ def clean(self) -> None:
         """
         pass
 
-    def save(self, path: str) -> None:
+    def save(self, path: str) -> str:
         """
         Save any produced temporary files.
         """
-        pass
+        return ""
 

src/zkregex_fuzzer/runner/circom.py

@@ -152,6 +152,7 @@ def witness_gen(cls, wasm_file_path: str, input_path: str) -> str:
         ]
         result = subprocess.run(cmd, capture_output=True, text=True)
 
+        logger.debug(" ".join(cmd))
         if result.returncode != 0:
             raise RegexRunError(f"Error running with SnarkJS: {result.stdout}")
 
@@ -233,6 +234,8 @@ class CircomRunner(Runner):
     """
 
     def __init__(self, regex: str, kwargs: dict):
+        self._circom_path = ""
+        self._input_path = ""
         self._wasm_path = ""
         self._r1cs_path = ""
         self._zkey_path = ""
@@ -279,6 +282,8 @@ def compile(self, regex: str) -> None:
             f.write("\n\n")
             f.write("component main {public [msg]} = " + f"{self._template_name}({self._circom_max_input_size});")
 
+        self._circom_path = circom_file_path
+
         # Compile the circom code to wasm
         logger.debug(f"Compiling circom code starts")
         self._wasm_path, self._r1cs_path = CircomSubprocess.compile(circom_file_path, self._link_path)
@@ -302,6 +307,8 @@ def match(self, input: str) -> tuple[bool, str]:
         with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as tmp_file:
             tmp_file.write(json.dumps({"msg": numeric_input}).encode())
             input_path = tmp_file.name
+        
+        self._input_path = input_path
 
         # Skip if input is larger than circuit max input size
         if len(numeric_input) > self._circom_max_input_size:
@@ -311,8 +318,6 @@ def match(self, input: str) -> tuple[bool, str]:
         logger.debug(f"Generating witness starts")
         witness_path = SnarkjsSubprocess.witness_gen(self._wasm_path, input_path)
         logger.debug(f"Generating witness ends")
-        # Remove input file
-        Path(input_path).unlink()
 
         # Also run the proving backend if the flag is set
         if self._run_the_prover:
@@ -336,16 +341,22 @@ def clean(self):
         # Remove all temporary files
         if self._wasm_path:
             shutil.rmtree(Path(self._wasm_path).parent)
-        Path.unlink(self._r1cs_path, True)
-        Path.unlink(self._zkey_path, True)
-        Path.unlink(self._vkey_path, True)
+        if self._circom_path: Path(self._circom_path).unlink(True)
+        if self._input_path: Path(self._input_path).unlink(True)
+        if self._r1cs_path: Path(self._r1cs_path).unlink(True)
+        if self._zkey_path: Path(self._zkey_path).unlink(True)
+        if self._vkey_path: Path(self._vkey_path).unlink(True)
 
-    def save(self, path):
+    def save(self, path) -> str:
+        circom_path = Path(self._circom_path)
+        input_path = Path(self._input_path)
         r1cs_path = Path(self._r1cs_path)
         wasm_path = Path(self._wasm_path)
         target_path = Path(path).resolve() / f"output_{r1cs_path.stem}"
         target_path.mkdir()
 
+        circom_path.replace(target_path / circom_path.name)
+        input_path.replace(target_path / input_path.name)
         r1cs_path.replace(target_path / r1cs_path.name)
         wasm_path.replace(target_path / wasm_path.name)
 
@@ -354,4 +365,6 @@ def save(self, path):
             vkey_path = Path(self._vkey_path)
 
             zkey_path.replace(target_path / zkey_path.name)
-            vkey_path.replace(target_path / vkey_path.name)
+            vkey_path.replace(target_path / vkey_path.name)
+
+        return str(target_path)