Skip to content

Commit 5fa0ab4

Browse files
berisherberisher
authored
Merge pull request #7 from cisco-ai-defense/add-more-unit-tests
moved get-modules/stdlib_complete into folders; added tests
2 parents 3792d0e + 594543c commit 5fa0ab4

15 files changed

Lines changed: 880 additions & 20 deletions

stdlib_complete.txt renamed to data/stdlib_complete.txt

File renamed without changes.

‎fuzz/Cargo.lock‎

Lines changed: 15 additions & 15 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎fuzz/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ cargo-fuzz = true
2626
[dependencies]
2727
libfuzzer-sys = "0.4"
2828

29-
[dependencies.pickle-fuzzer]
29+
[dependencies.cisco-ai-defense-pickle-fuzzer]
3030
path = ".."
3131

3232
[workspace]

get-modules.py renamed to scripts/get-modules.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151
"lib2to3", # Large parsing library - slow
5252
}
5353

54-
OUTPUT_FILE = "stdlib_complete.txt"
54+
OUTPUT_FILE = "../data/stdlib_complete.txt"
5555

5656

5757
def get_all_module_names() -> Set[str]:

‎src/cli.rs‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,3 +112,69 @@ impl Cli {
112112
self.file.is_some()
113113
}
114114
}
115+
116+
#[cfg(test)]
117+
mod tests {
118+
use super::*;
119+
120+
#[test]
121+
fn test_parse_version_valid() {
122+
assert_eq!(parse_version("0").unwrap(), 0);
123+
assert_eq!(parse_version("3").unwrap(), 3);
124+
assert_eq!(parse_version("5").unwrap(), 5);
125+
}
126+
127+
#[test]
128+
fn test_parse_version_invalid_too_high() {
129+
assert!(parse_version("6").is_err());
130+
assert!(parse_version("10").is_err());
131+
}
132+
133+
#[test]
134+
fn test_parse_version_invalid_format() {
135+
assert!(parse_version("abc").is_err());
136+
assert!(parse_version("3.5").is_err());
137+
assert!(parse_version("-1").is_err());
138+
}
139+
140+
#[test]
141+
fn test_cli_mode_detection() {
142+
use std::path::PathBuf;
143+
144+
let cli_single = Cli {
145+
file: Some(PathBuf::from("test.pkl")),
146+
dir: None,
147+
protocol: None,
148+
samples: 10_000,
149+
seed: None,
150+
min_opcodes: 60,
151+
max_opcodes: 300,
152+
mutators: vec![],
153+
mutation_rate: 0.1,
154+
unsafe_mutations: false,
155+
allow_ext: false,
156+
allow_buffer: false,
157+
};
158+
159+
assert!(cli_single.is_single_file_mode());
160+
assert!(!cli_single.is_batch_mode());
161+
162+
let cli_batch = Cli {
163+
file: None,
164+
dir: Some(PathBuf::from("output")),
165+
protocol: None,
166+
samples: 10_000,
167+
seed: None,
168+
min_opcodes: 60,
169+
max_opcodes: 300,
170+
mutators: vec![],
171+
mutation_rate: 0.1,
172+
unsafe_mutations: false,
173+
allow_ext: false,
174+
allow_buffer: false,
175+
};
176+
177+
assert!(!cli_batch.is_single_file_mode());
178+
assert!(cli_batch.is_batch_mode());
179+
}
180+
}

‎src/generator/emission.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ static STDLIB_MODULES: OnceLock<Vec<String>> = OnceLock::new();
5858

5959
fn load_stdlib_complete() -> &'static Vec<String> {
6060
STDLIB_MODULES.get_or_init(|| {
61-
let content = include_str!("../../stdlib_complete.txt");
61+
let content = include_str!("../../data/stdlib_complete.txt");
6262
content.lines().map(|s| s.to_string()).collect()
6363
})
6464
}

‎src/mutators/bitflip.rs‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,3 +42,108 @@ impl Mutator for BitFlipMutator {
4242
Some(value ^ (1 << bit_pos))
4343
}
4444
}
45+
46+
#[cfg(test)]
47+
mod tests {
48+
use super::*;
49+
use crate::generator::GenerationSource;
50+
use rand::SeedableRng;
51+
use rand_chacha::ChaCha8Rng;
52+
53+
#[test]
54+
fn test_bitflip_name() {
55+
let mutator = BitFlipMutator;
56+
assert_eq!(mutator.name(), "bitflip");
57+
}
58+
59+
#[test]
60+
fn test_bitflip_int_always_mutates_at_rate_1() {
61+
let mutator = BitFlipMutator;
62+
let mut rng = ChaCha8Rng::seed_from_u64(42);
63+
let mut source = GenerationSource::Rand(&mut rng);
64+
65+
let original = 0b1010_1010;
66+
let result = mutator.mutate_int(original, &mut source, 1.0);
67+
68+
assert!(result.is_some());
69+
let mutated = result.unwrap();
70+
assert_ne!(mutated, original);
71+
72+
// verify exactly one bit flipped
73+
let diff = original ^ mutated;
74+
assert_eq!(diff.count_ones(), 1);
75+
}
76+
77+
#[test]
78+
fn test_bitflip_int_never_mutates_at_rate_0() {
79+
let mutator = BitFlipMutator;
80+
let mut rng = ChaCha8Rng::seed_from_u64(42);
81+
let mut source = GenerationSource::Rand(&mut rng);
82+
83+
let result = mutator.mutate_int(100, &mut source, 0.0);
84+
assert!(result.is_none());
85+
}
86+
87+
#[test]
88+
fn test_bitflip_long_always_mutates_at_rate_1() {
89+
let mutator = BitFlipMutator;
90+
let mut rng = ChaCha8Rng::seed_from_u64(42);
91+
let mut source = GenerationSource::Rand(&mut rng);
92+
93+
let original = 0x1234_5678_9ABC_DEF0_i64;
94+
let result = mutator.mutate_long(original, &mut source, 1.0);
95+
96+
assert!(result.is_some());
97+
let mutated = result.unwrap();
98+
assert_ne!(mutated, original);
99+
100+
// verify exactly one bit flipped
101+
let diff = original ^ mutated;
102+
assert_eq!(diff.count_ones(), 1);
103+
}
104+
105+
#[test]
106+
fn test_bitflip_long_never_mutates_at_rate_0() {
107+
let mutator = BitFlipMutator;
108+
let mut rng = ChaCha8Rng::seed_from_u64(42);
109+
let mut source = GenerationSource::Rand(&mut rng);
110+
111+
let result = mutator.mutate_long(1000, &mut source, 0.0);
112+
assert!(result.is_none());
113+
}
114+
115+
#[test]
116+
fn test_bitflip_int_zero() {
117+
let mutator = BitFlipMutator;
118+
let mut rng = ChaCha8Rng::seed_from_u64(42);
119+
let mut source = GenerationSource::Rand(&mut rng);
120+
121+
let result = mutator.mutate_int(0, &mut source, 1.0);
122+
assert!(result.is_some());
123+
124+
// flipping any bit in 0 gives a power of 2
125+
let mutated = result.unwrap();
126+
assert_ne!(mutated, 0);
127+
assert_eq!(mutated.count_ones(), 1);
128+
}
129+
130+
#[test]
131+
fn test_bitflip_produces_different_results() {
132+
let mutator = BitFlipMutator;
133+
let mut rng = ChaCha8Rng::seed_from_u64(42);
134+
let mut source = GenerationSource::Rand(&mut rng);
135+
136+
let value = 0xFF00;
137+
let mut results = std::collections::HashSet::new();
138+
139+
// collect multiple mutations
140+
for _ in 0..20 {
141+
if let Some(mutated) = mutator.mutate_int(value, &mut source, 1.0) {
142+
results.insert(mutated);
143+
}
144+
}
145+
146+
// should produce multiple different results
147+
assert!(results.len() > 1);
148+
}
149+
}

‎src/mutators/boundary.rs‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,3 +59,75 @@ impl Mutator for BoundaryMutator {
5959
Some(boundaries[source.gen_range(0, boundaries.len())])
6060
}
6161
}
62+
63+
#[cfg(test)]
64+
mod tests {
65+
use super::*;
66+
use crate::generator::GenerationSource;
67+
use rand::SeedableRng;
68+
use rand_chacha::ChaCha8Rng;
69+
70+
#[test]
71+
fn test_boundary_name() {
72+
let mutator = BoundaryMutator;
73+
assert_eq!(mutator.name(), "boundary");
74+
}
75+
76+
#[test]
77+
fn test_boundary_int_returns_boundary_values() {
78+
let mutator = BoundaryMutator;
79+
let mut rng = ChaCha8Rng::seed_from_u64(42);
80+
let mut source = GenerationSource::Rand(&mut rng);
81+
82+
let boundaries = [0, -1, 1, i32::MAX, i32::MIN];
83+
84+
for _ in 0..20 {
85+
if let Some(result) = mutator.mutate_int(100, &mut source, 1.0) {
86+
assert!(boundaries.contains(&result));
87+
}
88+
}
89+
}
90+
91+
#[test]
92+
fn test_boundary_int_never_mutates_at_rate_0() {
93+
let mutator = BoundaryMutator;
94+
let mut rng = ChaCha8Rng::seed_from_u64(42);
95+
let mut source = GenerationSource::Rand(&mut rng);
96+
97+
let result = mutator.mutate_int(100, &mut source, 0.0);
98+
assert!(result.is_none());
99+
}
100+
101+
#[test]
102+
fn test_boundary_long_returns_boundary_values() {
103+
let mutator = BoundaryMutator;
104+
let mut rng = ChaCha8Rng::seed_from_u64(42);
105+
let mut source = GenerationSource::Rand(&mut rng);
106+
107+
let boundaries = [0, -1, 1, i64::MAX, i64::MIN];
108+
109+
for _ in 0..20 {
110+
if let Some(result) = mutator.mutate_long(1000, &mut source, 1.0) {
111+
assert!(boundaries.contains(&result));
112+
}
113+
}
114+
}
115+
116+
#[test]
117+
fn test_boundary_float_returns_boundary_values() {
118+
let mutator = BoundaryMutator;
119+
let mut rng = ChaCha8Rng::seed_from_u64(42);
120+
let mut source = GenerationSource::Rand(&mut rng);
121+
122+
for _ in 0..20 {
123+
if let Some(result) = mutator.mutate_float(1.5, &mut source, 1.0) {
124+
// check it's one of the expected boundary values
125+
assert!(
126+
result == 0.0 || result == -1.0 || result == 1.0 ||
127+
result == f64::MAX || result == f64::MIN ||
128+
result.is_infinite() || result.is_nan()
129+
);
130+
}
131+
}
132+
}
133+
}

0 commit comments

Comments
 (0)