fix: align generator contracts (#17)

* fix: align generator contracts

* fix: satisfy clippy in generator core

Author: berisher

README.md

@@ -23,7 +23,7 @@ A structure-aware test case generator for Python pickle parsers and validators.
 - **Comprehensive Opcode Coverage**: Supports all standard pickle opcodes including FRAME, EXT, GLOBAL, etc.
 - **Parallel Generation**: Generate multiple pickle files concurrently
 - **Configurable Output**: Single file or batch generation modes
-- **Deterministic Fuzzing**: Optional seed-based generation for reproducibility
+- **Deterministic Fuzzing**: Optional seed-based generation for reproducibility, including seeded batch mode with deterministic per-sample fan-out
 
 ## Installation
 
@@ -134,22 +134,28 @@ Options:
       --seed <SEED>                    Seed for reproducible generation
       --min-opcodes <MIN_OPCODES>      Minimum opcodes to generate [default: 60]
       --max-opcodes <MAX_OPCODES>      Maximum opcodes to generate [default: 300]
-      --mutators <MUTATOR>...          Enable mutators (all, bitflip, boundary, offbyone,
+      --mutators <MUTATOR>             Enable mutators (all, bitflip, boundary, offbyone,
                                        stringlen, character, memoindex, typeconfusion)
       --mutation-rate <MUTATION_RATE>  Mutation probability 0.0-1.0 [default: 0.1]
       --unsafe-mutations               Allow mutations that may produce invalid pickles
       --allow-ext                      Allow EXT* opcodes (requires extension registry)
       --allow-buffer                   Allow buffer opcodes (requires buffer support)
+      --allow-persistent-ids           Allow PERSID/BINPERSID opcodes (requires persistent_load support)
   -h, --help                           Print help
   -V, --version                        Print version
 ```
 
 **Special Opcodes:**
 - `--allow-ext`: Enables EXT1/EXT2/EXT4 opcodes. Only use if your unpickler has a configured extension registry, otherwise unpickling will fail.
 - `--allow-buffer`: Enables NEXT_BUFFER/READONLY_BUFFER opcodes. Only use if your unpickler has out-of-band buffer callbacks configured.
+- `--allow-persistent-ids`: Enables PERSID/BINPERSID opcodes. Only use if your unpickler provides a `persistent_load` callback.
 
 By default, these opcodes are disabled to ensure generated pickles work with standard Python's `pickle` module without additional configuration.
 
+Seeded batch mode derives a deterministic per-sample seed from the base `--seed`,
+so repeated runs reproduce the same corpus without collapsing every file to the
+same bytes.
+
 ## Python Bindings
 
 `pickle-fuzzer` provides Python bindings for integration with Python-based fuzzing tools like Atheris.
@@ -173,16 +179,22 @@ gen = Generator(protocol=3)
 
 # Generate a random pickle
 pickle_bytes = gen.generate()
+pickle_bytes = gen.generate(max_size=4096)
 
 # Generate from fuzzer input (deterministic)
 fuzzer_data = b"some_fuzzer_input"
 pickle_bytes = gen.generate_from_bytes(fuzzer_data)
+pickle_bytes = gen.generate_from_bytes(fuzzer_data, max_size=4096)
 
 # Configure generation
 gen.set_opcode_range(10, 50)  # Control pickle complexity
-gen.reset()  # Reset internal state
+gen = Generator(protocol=4, allow_persistent_ids=True)  # Opt in to persistent IDs
 ```
 
+`generate()` and `generate_from_bytes()` reset internal generator state before
+each run, so repeated calls on the same `Generator` remain deterministic for the
+same seed or fuzzer input.
+
 ### Integration with Atheris
 
 Use the `PickleMutator` class for structure-aware fuzzing:
@@ -196,7 +208,7 @@ mutator = PickleMutator(protocol=3)
 
 @atheris.instrument_func
 def test_one_input(data: bytes):
-    # Generate valid pickle from fuzzer input
+    # Generate a valid pickle from fuzzer input within the requested budget
     pickle_bytes = mutator.mutate(data, max_size=10000)
 
     try:

python/pickle_fuzzer/_native.pyi

@@ -16,8 +16,13 @@
 from typing import Optional
 
 class Generator:
-    def __init__(self, protocol: int = 3, seed: Optional[int] = None) -> None: ...
-    def generate(self) -> bytes: ...
-    def generate_from_bytes(self, data: bytes) -> bytes: ...
+    def __init__(
+        self,
+        protocol: int = 3,
+        seed: Optional[int] = None,
+        allow_persistent_ids: bool = False,
+    ) -> None: ...
+    def generate(self, max_size: Optional[int] = None) -> bytes: ...
+    def generate_from_bytes(self, data: bytes, max_size: Optional[int] = None) -> bytes: ...
     def set_opcode_range(self, min: int, max: int) -> None: ...
     def reset(self) -> None: ...

python/pickle_fuzzer/fuzzer.py

@@ -15,6 +15,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 """Atheris integration utilities for pickle-fuzzer fuzzing."""
+import pickle
 import sys
 from typing import Optional, Callable
 
@@ -56,14 +57,17 @@ def __init__(self, protocol: int = 3, seed: Optional[int] = None):
             seed: optional seed for deterministic generation
         """
         self.protocol = protocol
+        self.seed = seed
         self.generator = Generator(protocol=protocol, seed=seed)
+        self._fallback_pickle = pickle.dumps(None, protocol=protocol)
+        self._minimum_fallback_pickle = pickle.dumps(None, protocol=0)
 
     def mutate(self, data: bytes, max_size: int) -> bytes:
         """mutate pickle data using structure-aware mutations.
         
         uses the fuzzer-provided data as a seed for generating new pickle
-        bytecode. the generated pickle will be valid according to the
-        specified protocol version.
+        bytecode. when max_size can accommodate a pickle for the configured
+        protocol, the returned pickle remains structurally valid.
         
         args:
             data: fuzzer-provided input data
@@ -73,14 +77,19 @@ def mutate(self, data: bytes, max_size: int) -> bytes:
             generated pickle bytecode
         """
         try:
-            # use fuzzer bytes to generate new pickle
-            result = self.generator.generate_from_bytes(data)
-            if len(result) <= max_size:
-                return result
-            # if too large, truncate to max_size (may be invalid)
-            return result[:max_size]
-        except Exception:
-            # if generation fails, return original data
+            return self.generator.generate_from_bytes(data, max_size=max_size)
+        except Exception as exc:
+            print(
+                f"PickleMutator generation failed for protocol {self.protocol}: {exc}",
+                file=sys.stderr,
+            )
+
+            if len(self._fallback_pickle) <= max_size:
+                return self._fallback_pickle
+            if len(self._minimum_fallback_pickle) <= max_size:
+                return self._minimum_fallback_pickle
+
+            # No valid pickle fits in the requested budget.
             return data[:max_size] if len(data) > max_size else data
 
     def reset(self):

python/tests/test_fuzzer.py

@@ -0,0 +1,51 @@
+# Copyright 2025 Cisco Systems, Inc. and its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import pickletools
+
+from pickle_fuzzer.fuzzer import PickleMutator
+
+
+def assert_whole_pickle_consumed(data: bytes) -> None:
+    stop_position = None
+    for _, _, position in pickletools.genops(data):
+        stop_position = position + 1
+    assert stop_position is not None
+    assert stop_position == len(data)
+
+
+def test_pickle_mutator_respects_max_size():
+    mutator = PickleMutator(protocol=4, seed=123)
+
+    data = mutator.mutate(b"fuzzer_input_bytes", max_size=48)
+
+    assert len(data) <= 48
+    assert_whole_pickle_consumed(data)
+
+
+def test_pickle_mutator_fallback_stays_valid(monkeypatch):
+    mutator = PickleMutator(protocol=4, seed=123)
+
+    class FailingGenerator:
+        def generate_from_bytes(self, *args, **kwargs):
+            raise RuntimeError("boom")
+
+    monkeypatch.setattr(mutator, "generator", FailingGenerator())
+
+    data = mutator.mutate(b"fuzzer_input_bytes", max_size=32)
+
+    assert len(data) <= 32
+    assert_whole_pickle_consumed(data)

python/tests/test_generation.py

@@ -23,6 +23,7 @@ def test_basic_generation():
     assert len(data) > 0
     assert data[-1] == ord(".")
 
+
 def test_deterministic_generation():
     gen1 = pickle_fuzzer.Generator(protocol=3, seed=42)
     gen2 = pickle_fuzzer.Generator(protocol=3, seed=42)
@@ -32,13 +33,30 @@ def test_deterministic_generation():
 
     assert data1 == data2
 
+
+def test_set_opcode_range_preserves_seeded_determinism():
+    gen1 = pickle_fuzzer.Generator(protocol=3, seed=42)
+    gen2 = pickle_fuzzer.Generator(protocol=3, seed=42)
+
+    gen1.set_opcode_range(50, 10)
+    gen2.set_opcode_range(50, 10)
+
+    assert gen1.generate() == gen2.generate()
+
+
 def test_generate_from_bytes():
     gen = pickle_fuzzer.Generator(protocol=3)
     fuzzer_input = b"test_fuzzer_input_bytes"
 
     data1 = gen.generate_from_bytes(fuzzer_input)
-    gen.reset()
     data2 = gen.generate_from_bytes(fuzzer_input)
 
     assert data1 == data2  # same input bytes means same output
 
+
+def test_generate_from_bytes_respects_max_size():
+    gen = pickle_fuzzer.Generator(protocol=4, seed=123)
+    data = gen.generate_from_bytes(b"test_fuzzer_input_bytes", max_size=32)
+
+    assert len(data) <= 32
+    assert data[-1] == ord(".")

src/cli.rs

@@ -14,9 +14,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::path::PathBuf;
+use std::{ffi::OsString, path::PathBuf};
 
-use clap::Parser;
+use clap::{Parser, ValueEnum};
 
 /// Parse and validate a pickle protocol version string.
 ///
@@ -32,6 +32,50 @@ fn parse_version(s: &str) -> Result<usize, String> {
     }
 }
 
+fn normalize_mutator_args<I>(args: I) -> Vec<OsString>
+where
+    I: IntoIterator<Item = OsString>,
+{
+    let args: Vec<OsString> = args.into_iter().collect();
+    let mut normalized = Vec::with_capacity(args.len());
+    let mut idx = 0;
+
+    while idx < args.len() {
+        let arg = &args[idx];
+        if arg == "--mutators" {
+            normalized.push(arg.clone());
+            idx += 1;
+
+            let mut saw_mutator = false;
+            while idx < args.len() {
+                let Some(value) = args[idx].to_str() else {
+                    break;
+                };
+                if value == "--" || value.starts_with('-') {
+                    break;
+                }
+                if crate::mutators::MutatorKind::from_str(value, true).is_err() {
+                    break;
+                }
+
+                if saw_mutator {
+                    normalized.push(OsString::from("--mutators"));
+                }
+                normalized.push(args[idx].clone());
+                saw_mutator = true;
+                idx += 1;
+            }
+
+            continue;
+        }
+
+        normalized.push(arg.clone());
+        idx += 1;
+    }
+
+    normalized
+}
+
 /// Command-line interface for pickle-fuzzer.
 ///
 /// Supports two modes:
@@ -68,7 +112,7 @@ pub struct Cli {
     #[arg(short, long, default_value_t = 10_000, requires = "dir")]
     pub samples: usize,
 
-    /// seed for random number generator (for reproducible, byte-identical generation)
+    /// seed for reproducible generation
     #[arg(long)]
     pub seed: Option<u64>,
 
@@ -80,8 +124,12 @@ pub struct Cli {
     #[arg(long, default_value_t = 300)]
     pub max_opcodes: usize,
 
-    /// enable specific mutators (can be specified multiple times)
-    #[arg(long = "mutators", value_name = "MUTATOR", num_args = 1..)]
+    /// enable specific mutators (repeat the flag or list mutators after one occurrence)
+    #[arg(
+        long = "mutators",
+        value_name = "MUTATOR",
+        action = clap::ArgAction::Append
+    )]
     pub mutators: Vec<crate::mutators::MutatorKind>,
 
     /// mutation rate (0.0-1.0, probability of applying mutation)
@@ -99,9 +147,17 @@ pub struct Cli {
     /// allow NEXT_BUFFER/READONLY_BUFFER opcodes (requires out-of-band buffer support in unpickler)
     #[arg(long)]
     pub allow_buffer: bool,
+
+    /// allow PERSID/BINPERSID opcodes (requires persistent_load support in unpickler)
+    #[arg(long)]
+    pub allow_persistent_ids: bool,
 }
 
 impl Cli {
+    pub fn parse_args() -> Self {
+        Self::parse_from(normalize_mutator_args(std::env::args_os()))
+    }
+
     /// Check if running in batch mode (generating multiple files).
     pub fn is_batch_mode(&self) -> bool {
         self.dir.is_some()
@@ -154,6 +210,7 @@ mod tests {
             unsafe_mutations: false,
             allow_ext: false,
             allow_buffer: false,
+            allow_persistent_ids: false,
         };
 
         assert!(cli_single.is_single_file_mode());
@@ -172,9 +229,33 @@ mod tests {
             unsafe_mutations: false,
             allow_ext: false,
             allow_buffer: false,
+            allow_persistent_ids: false,
         };
 
         assert!(!cli_batch.is_single_file_mode());
         assert!(cli_batch.is_batch_mode());
     }
+
+    #[test]
+    fn test_normalize_mutator_args_keeps_output_path_positional() {
+        let normalized = normalize_mutator_args([
+            OsString::from("pickle-fuzzer"),
+            OsString::from("--mutators"),
+            OsString::from("bitflip"),
+            OsString::from("boundary"),
+            OsString::from("output.pkl"),
+        ]);
+
+        assert_eq!(
+            normalized,
+            vec![
+                OsString::from("pickle-fuzzer"),
+                OsString::from("--mutators"),
+                OsString::from("bitflip"),
+                OsString::from("--mutators"),
+                OsString::from("boundary"),
+                OsString::from("output.pkl"),
+            ]
+        );
+    }
 }