From 84fb9e4d4211c36056fbd8d8fa7ea87d04a1a9a8 Mon Sep 17 00:00:00 2001
From: Victor Graf <victor@risczero.com>
Date: Fri, 27 Sep 2024 14:33:33 -0700
Subject: [PATCH] ensure no extra data is every written to stdout

---
 ffi/Cargo.toml  |  1 +
 ffi/src/main.rs | 55 ++++++++++++++++++++++++++++++++++++++-----------
 2 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/ffi/Cargo.toml b/ffi/Cargo.toml
index ee60c620..51c58a9e 100644
--- a/ffi/Cargo.toml
+++ b/ffi/Cargo.toml
@@ -12,6 +12,7 @@ alloy = { workspace = true }
 anyhow = { workspace = true }
 clap = { version = "4.5", features = ["derive", "env"] }
 hex = { version = "0.4" }
+libc = "0.2.159"
 risc0-ethereum-contracts = { workspace = true }
 risc0-zkvm = { workspace = true, features = ["client"] }
 
diff --git a/ffi/src/main.rs b/ffi/src/main.rs
index 4c491eb0..76276e38 100644
--- a/ffi/src/main.rs
+++ b/ffi/src/main.rs
@@ -12,10 +12,15 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::io::Write;
+use std::{
+    fs::{File, OpenOptions},
+    io,
+    io::Write,
+    os::unix::io::{AsRawFd, FromRawFd},
+};
 
 use alloy::{primitives::Bytes, sol_types::SolValue};
-use anyhow::{Context, Result};
+use anyhow::{ensure, Context, Result};
 use clap::Parser;
 use risc0_ethereum_contracts::encode_seal;
 use risc0_zkvm::{default_prover, ExecutorEnv, ProverOpts, VerifierContext};
@@ -35,7 +40,10 @@ enum Command {
 
 /// Run the CLI.
 pub fn main() -> Result<()> {
-    match Command::parse() {
+    // Take stdout is ensure no extra data is written to it.
+    let mut stdout = take_stdout()?;
+
+    let output = match Command::parse() {
         Command::Prove {
             guest_binary_path,
             input,
@@ -45,22 +53,19 @@ pub fn main() -> Result<()> {
         )?,
     };
 
+    // Forge test FFI calls expect hex encoded bytes sent to stdout
+    write!(&mut stdout, "{}", hex::encode(output)).context("failed to write to stdout")?;
+    stdout.flush().context("failed to flush stdout")?;
+
     Ok(())
 }
 
 /// Prints on stdio the Ethereum ABI and hex encoded proof.
-fn prove_ffi(elf_path: String, input: Vec<u8>) -> Result<()> {
+fn prove_ffi(elf_path: String, input: Vec<u8>) -> Result<Vec<u8>> {
     let elf = std::fs::read(elf_path).expect("failed to read guest ELF");
     let (journal, seal) = prove(&elf, &input)?;
     let calldata = (Bytes(journal.into()), Bytes(seal.into()));
-    let output = hex::encode(calldata.abi_encode());
-
-    // Forge test FFI calls expect hex encoded bytes sent to stdout
-    print!("{output}");
-    std::io::stdout()
-        .flush()
-        .context("failed to flush stdout buffer")?;
-    Ok(())
+    Ok(calldata.abi_encode())
 }
 
 /// Generates journal and snark seal as a pair (`Vec<u8>`, `Vec<u8>)
@@ -78,3 +83,29 @@ fn prove(elf: &[u8], input: &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
     let seal = encode_seal(&receipt)?;
     Ok((journal, seal))
 }
+
+/// "Takes" stdout such, returning a handle and ensuring no other code in this process can write to
+/// it. This is used to ensure that no additional data (e.g. log lines) is written to stdout, as
+/// any extra will cause a decoding failure.
+fn take_stdout() -> Result<File> {
+    let stdout = io::stdout();
+    let mut handle = stdout.lock();
+    // Ensure all buffered data is written before redirection
+    handle.flush()?;
+
+    let devnull = OpenOptions::new().write(true).open("/dev/null")?;
+
+    unsafe {
+        // Create a copy of stdout to use for our output.
+        let dup_fd = libc::dup(handle.as_raw_fd());
+        ensure!(dup_fd >= 0, "call to libc::dup failed: {}", dup_fd);
+        // Redirect stdout to the fd we opened for /dev/null
+        let dup2_result = libc::dup2(devnull.as_raw_fd(), libc::STDOUT_FILENO);
+        ensure!(
+            dup2_result >= 0,
+            "call to libc::dup2 failed: {}",
+            dup2_result
+        );
+        Ok(File::from_raw_fd(dup_fd))
+    }
+}