2 files changed, 28 insertions, 10 deletions

diff --git a/zkvms_guest_io/default.nix b/zkvms_guest_io/default.nix
index 20c75a7..bce5a9c 100644
--- a/zkvms_guest_io/default.nix
+++ b/zkvms_guest_io/default.nix
@@ -1,11 +1,11 @@
-{ craneLib-default, guest, zkvms, hostPackages, lib, }:
+{ craneLib-default, guest, zkvms, hostPackages, lib, benchexec, }:
 let
   commonArgs = {
     name = "${guest}";
 
-    buildInputs =
-      lib.foldr (zkvm: accum: accum ++ [ hostPackages."${zkvm}/${guest}" ]) [ ]
-      zkvms;
+    buildInputs = [ benchexec ] ++
+      (lib.foldr (zkvm: accum: accum ++ [ hostPackages."${zkvm}/${guest}" ]) [ ]
+      zkvms);
 
     src = lib.fileset.toSource {
       root = ./.;
@@ -15,6 +15,10 @@ let
     PROGRAMS = lib.foldr (zkvm: accum:
       hostPackages."${zkvm}/${guest}" + "/bin/${zkvm}_${guest}," + accum) ""
       zkvms;
+
+    postPatch = ''
+      sed -i 's|"runexec"|"${benchexec}/bin/runexec"|' ./src/main.rs
+    '';
   };
 
   cargoArtifacts = craneLib-default.buildDepsOnly commonArgs;
diff --git a/zkvms_guest_io/src/main.rs b/zkvms_guest_io/src/main.rs
index 04d33d1..ab20a32 100644
--- a/zkvms_guest_io/src/main.rs
+++ b/zkvms_guest_io/src/main.rs
@@ -30,14 +30,21 @@ struct Cli {
 }
 
 fn run_command(zkvm_guest_command: &str, operation: &str) -> Result<std::process::Output, Error> {
-    Command::new(zkvm_guest_command)
-        .arg(operation)
+    Command::new("runexec")
+        .args(["--no-container", "--"])
+        .args([zkvm_guest_command, operation])
         .arg("--benchmark")
         .args([ "--metrics-output", "/tmp/current_metrics" ])
         .stdout(Stdio::piped())
         .output()
 }
 
+fn get_runexec_value(output: &String, name: &str, end: char) -> String {
+    let start_bytes = output.find(name).unwrap();
+    let right_half = &output[start_bytes + name.len() + 1..];
+    right_half[..right_half.find(end).unwrap()].to_string()
+}
+
 fn main() {
     let cli = Cli::parse();
 
@@ -65,7 +72,7 @@ fn main() {
 
             let output = run_command(zkvm_guest_command, operation);
 
-            // Couldn't run the command
+            // Couldn't run runexec
             if let Err(msg) = output {
                 println!("Failed to run command!");
                 println!("{msg}");
@@ -75,10 +82,10 @@ fn main() {
                 continue;
             }
 
-            // The command ran and therefore produced some output
+            // runexec ran and therefore produced some output
             let output = output.unwrap();
 
-            // The command ran but exited with non-zero status code
+            // runexec exited with non-zero status code
             if !output.status.success() {
                 println!("Command failed!");
             }
@@ -93,15 +100,22 @@ fn main() {
                 );
             }
 
-            // The command ran but exited with non-zero status code
+            // runexec ran but exited with non-zero status code
             if !output.status.success() {
                 break 'guest_iter;
             }
 
+            // The guest program ran but exited with non-zero status code
+            if get_runexec_value(&stdout, "returnvalue", '\n') != "0" {
+                run[operation] = Null;
+                continue;
+            }
+
             let raw_data = &read_to_string("/tmp/current_metrics")
                 .ok()
                 .unwrap();
             run[operation] = json::parse(raw_data).unwrap();
+            run[operation]["memory"] = get_runexec_value(&stdout, "memory", 'B').parse::<u64>().unwrap().into();
         }
 
         runs["benchmarking"].push(run);