communicate with python over a pipe to increase fuzzing performance

2 files changed, 70 insertions, 36 deletions

diff --git a/fuzz/fuzz.rs b/fuzz/fuzz.rs
index 72f751b..fd01914 100644
--- a/fuzz/fuzz.rs
+++ b/fuzz/fuzz.rs
@@ -2,49 +2,83 @@ use core::slice;
 use gentoo_utils::{Parseable, atom::Atom};
 use mon::{Parser, ParserFinishedError, input::InputIter};
 use std::{
-    io::Write,
-    process::{Command, Stdio},
+    io::{BufRead, BufReader, Write},
+    process::{ChildStdin, ChildStdout, Command, Stdio},
+    sync::{LazyLock, Mutex},
 };
 
+struct PyProcess {
+    stdin: Mutex<ChildStdin>,
+    stdout: Mutex<BufReader<ChildStdout>>,
+    buffer: Mutex<String>,
+}
+
 #[allow(clippy::missing_safety_doc, clippy::needless_return)]
 #[unsafe(no_mangle)]
 pub unsafe extern "C" fn LLVMFuzzerTestOneInput(input: *const u8, len: usize) -> i32 {
+    static PY_PROCESS: LazyLock<PyProcess> = LazyLock::new(|| {
+        #[allow(clippy::zombie_processes)]
+        let mut proc = Command::new("atom.py")
+            .stdin(Stdio::piped())
+            .stdout(Stdio::piped())
+            .spawn()
+            .expect("failed to spawn atom.py");
+
+        let stdin = Mutex::new(proc.stdin.take().unwrap());
+        let stdout = Mutex::new(BufReader::new(proc.stdout.take().unwrap()));
+
+        PyProcess {
+            stdin,
+            stdout,
+            buffer: Mutex::new(String::new()),
+        }
+    });
+
     let slice = unsafe { slice::from_raw_parts(input, len) };
 
-    let atom = match str::from_utf8(slice) {
-        Ok(str) => str.trim(),
-        _ => return -1,
+    let str = match str::from_utf8(slice) {
+        Ok(str) => str,
+        Err(_) => return -1,
     };
 
-    let mut proc = Command::new("atom.py")
-        .stdin(Stdio::piped())
-        .spawn()
-        .expect("failed to start atom.py");
+    let atom = str.trim();
+
+    let mut stdin = PY_PROCESS.stdin.lock().expect("failed to get stdin lock");
+
+    writeln!(&mut stdin, "{atom}").expect("failed to write to python stdin");
 
-    proc.stdin
-        .as_mut()
-        .unwrap()
-        .write_all(atom.as_bytes())
-        .unwrap();
+    let mut stdout = PY_PROCESS.stdout.lock().expect("failed to get stdout lock");
 
-    let status = proc.wait().unwrap();
+    let mut buffer = PY_PROCESS.buffer.lock().expect("failed to get buffer lock");
 
-    let result = Atom::parser().check_finished(InputIter::new(atom));
+    buffer.clear();
 
-    match (status.success(), result) {
-        (true, Ok(_)) => {
+    stdout
+        .read_line(&mut buffer)
+        .expect("failed to readline from python");
+
+    let portage_result = match buffer.as_str().trim() {
+        "0" => true,
+        "1" => false,
+        result => panic!("got unexpected result from python: {result}"),
+    };
+
+    let gentoo_utils_result = Atom::parser().check_finished(InputIter::new(atom)).is_ok();
+
+    match (portage_result, gentoo_utils_result) {
+        (true, true) => {
             eprintln!("agreement that {atom} is valid");
-            return 0;
         }
-        (true, Err(ParserFinishedError::Err(it) | ParserFinishedError::Unfinished(it))) => {
-            panic!("gentoo-utils rejected valid atom: {atom}: {}", it.rest());
-        }
-        (false, Err(_)) => {
+        (false, false) => {
             eprintln!("agreement that {atom} is invalid");
-            return -1;
         }
-        (false, Ok(_)) => {
-            panic!("gentoo-utils accepted invalid atom: {atom}");
+        (true, false) => {
+            panic!("rejected valid atom: {atom}");
+        }
+        (false, true) => {
+            panic!("accpeted invalid atom: {atom}")
         }
     }
+
+    return 0;
 }
diff --git a/scripts/atom.py b/scripts/atom.py
index a1a59bd..2e8ca7f 100755
--- a/scripts/atom.py
+++ b/scripts/atom.py
@@ -1,13 +1,13 @@
 #!/usr/bin/env python3
 
 import sys
-import portage.dep as dep
-
-input = sys.stdin.read().strip()
-
-try:
-    dep.Atom(input)
-except Exception as e:
-    sys.exit(1)
-
-sys.exit(0)
+from portage.dep import Atom
+
+for line in sys.stdin.buffer:
+    try:
+        Atom(line.decode().strip())
+        sys.stdout.buffer.write(b"0\n")
+    except:
+        sys.stdout.buffer.write(b"1\n")
+    finally:
+        sys.stdout.buffer.flush()