risc0_zkvm/host/client/
env.rs

1// Copyright 2025 RISC Zero, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! This module defines the [ExecutorEnv] and [ExecutorEnvBuilder].
16
17use std::{
18    cell::RefCell,
19    collections::HashMap,
20    io::{BufRead, BufReader, Cursor, Read, Write},
21    mem,
22    path::{Path, PathBuf},
23    rc::Rc,
24    sync::Arc,
25};
26
27use anyhow::{bail, Result};
28use bytemuck::Pod;
29use bytes::Bytes;
30use risc0_circuit_keccak::{KeccakState, KECCAK_PO2_RANGE};
31use risc0_zkp::core::digest::Digest;
32use risc0_zkvm_platform::{self, fileno};
33use serde::Serialize;
34use tempfile::TempDir;
35
36use crate::{
37    host::client::{
38        posix_io::PosixIo,
39        slice_io::{slice_io_from_fn, SliceIo, SliceIoTable},
40    },
41    serde::to_vec,
42    AssumptionReceipt, TraceCallback,
43};
44
45/// A builder pattern used to construct an [ExecutorEnv].
46#[derive(Default)]
47pub struct ExecutorEnvBuilder<'a> {
48    inner: ExecutorEnv<'a>,
49}
50
51#[allow(dead_code)]
52#[derive(Clone)]
53pub enum SegmentPath {
54    TempDir(Arc<TempDir>),
55    Path(PathBuf),
56}
57
58impl SegmentPath {
59    pub(crate) fn path(&self) -> &Path {
60        match self {
61            Self::TempDir(dir) => dir.path(),
62            Self::Path(path) => path.as_path(),
63        }
64    }
65}
66
67/// A ZKR proof request.
68#[stability::unstable]
69pub struct ProveZkrRequest {
70    /// The digest of the claim that this ZKR program is expected to produce.
71    pub claim_digest: Digest,
72
73    /// The control ID uniquely identifies the ZKR program to be proven.
74    pub control_id: Digest,
75
76    /// The input that the ZKR program should operate on.
77    pub input: Vec<u8>,
78}
79
80/// A Keccak proof request.
81#[stability::unstable]
82#[derive(Clone)]
83pub struct ProveKeccakRequest {
84    /// The digest of the claim that this keccak input is expected to produce.
85    pub claim_digest: Digest,
86
87    /// The requested size of the keccak proof, in powers of 2.
88    pub po2: usize,
89
90    /// The control root which identifies a particular keccak circuit revision.
91    pub control_root: Digest,
92
93    /// Input transcript to provide to the keccak circuit.
94    pub input: Vec<KeccakState>,
95}
96
97/// A trait that supports the ability to be notified of proof requests
98/// on-demand.
99#[stability::unstable]
100pub trait CoprocessorCallback {
101    /// Request that a ZKR proof is produced.
102    fn prove_zkr(&mut self, request: ProveZkrRequest) -> Result<()>;
103
104    /// Request that a keccak proof is produced.
105    fn prove_keccak(&mut self, request: ProveKeccakRequest) -> Result<()>;
106}
107
108pub type CoprocessorCallbackRef<'a> = Rc<RefCell<dyn CoprocessorCallback + 'a>>;
109
110/// Container for assumptions in the executor environment.
111#[derive(Default)]
112pub(crate) struct AssumptionReceipts(pub(crate) Vec<AssumptionReceipt>);
113
114/// The [Executor][crate::Executor] is configured from this object.
115///
116/// The executor environment holds configuration details that inform how the
117/// guest environment is set up prior to guest program execution.
118#[derive(Default)]
119pub struct ExecutorEnv<'a> {
120    pub(crate) env_vars: HashMap<String, String>,
121    pub(crate) args: Vec<String>,
122    pub(crate) keccak_max_po2: Option<u32>,
123    pub(crate) segment_limit_po2: Option<u32>,
124    pub(crate) session_limit: Option<u64>,
125    pub(crate) posix_io: Rc<RefCell<PosixIo<'a>>>,
126    pub(crate) slice_io: Rc<RefCell<SliceIoTable<'a>>>,
127    pub(crate) input: Vec<u8>,
128    pub(crate) trace: Vec<Rc<RefCell<dyn TraceCallback + 'a>>>,
129    pub(crate) assumptions: Rc<RefCell<AssumptionReceipts>>,
130    pub(crate) segment_path: Option<SegmentPath>,
131    pub(crate) pprof_out: Option<PathBuf>,
132    pub(crate) input_digest: Option<Digest>,
133    pub(crate) coprocessor: Option<CoprocessorCallbackRef<'a>>,
134}
135
136impl<'a> ExecutorEnv<'a> {
137    /// Construct a [ExecutorEnvBuilder].
138    ///
139    /// # Example
140    ///
141    /// ```
142    /// use risc0_zkvm::ExecutorEnv;
143    ///
144    /// let env = ExecutorEnv::builder().build();
145    /// ```
146    pub fn builder() -> ExecutorEnvBuilder<'a> {
147        ExecutorEnvBuilder::default()
148    }
149}
150
151impl<'a> ExecutorEnvBuilder<'a> {
152    /// Finalize this builder to construct an [ExecutorEnv].
153    ///
154    /// # Example
155    ///
156    /// ```
157    /// use risc0_zkvm::ExecutorEnv;
158    ///
159    /// let env = ExecutorEnv::builder().build().unwrap();
160    /// ```
161    ///
162    /// After calling `build`, the [ExecutorEnvBuilder] will be reset to
163    /// default.
164    pub fn build(&mut self) -> Result<ExecutorEnv<'a>> {
165        let mut inner = mem::take(&mut self.inner);
166
167        if !inner.input.is_empty() {
168            let reader = Cursor::new(inner.input.clone());
169            inner
170                .posix_io
171                .borrow_mut()
172                .with_read_fd(fileno::STDIN, reader);
173        }
174
175        if inner.pprof_out.is_none() {
176            if let Ok(env_var) = std::env::var("RISC0_PPROF_OUT") {
177                inner.pprof_out = Some(env_var.into());
178            }
179        }
180
181        if let Ok(po2) = std::env::var("RISC0_KECCAK_PO2") {
182            let po2 = po2.parse::<u32>()?;
183            if !KECCAK_PO2_RANGE.contains(&(po2 as usize)) {
184                bail!(
185                    "invalid keccak po2 {po2}. Expected range: {:?}",
186                    KECCAK_PO2_RANGE
187                );
188            }
189            inner.keccak_max_po2 = Some(po2);
190        }
191
192        Ok(inner)
193    }
194
195    /// Set a segment limit, specified in powers of 2 cycles.
196    ///
197    /// Lowering this value will reduce the memory consumption of the prover. Memory consumption is
198    /// roughly linear with the segment size, so lowering this value by 1 will cut memory
199    /// consumpton by about half.
200    ///
201    /// The default value is chosen to be performant on commonly used hardware. Tuning this value,
202    /// either up or down, may result in better proving performance.
203    ///
204    /// Given value must be between [risc0_zkp::MIN_CYCLES_PO2] and
205    /// [risc0_zkp::MAX_CYCLES_PO2] (inclusive).
206    pub fn segment_limit_po2(&mut self, limit: u32) -> &mut Self {
207        self.inner.segment_limit_po2 = Some(limit);
208        self
209    }
210
211    /// Set a segment limit for keccak proofs, specified in powers of 2 cycles.
212    ///
213    /// Lowering this value will reduce the memory consumption of the prover. Memory consumption is
214    /// roughly linear with the segment size, so lowering this value by 1 will cut memory
215    /// consumpton by about half.
216    ///
217    /// The default value is chosen to be performant on commonly used hardware. Tuning this value,
218    /// either up or down, may result in better proving performance.
219    ///
220    /// Given value must be within [risc0_circuit_keccak::KECCAK_PO2_RANGE]
221    pub fn keccak_max_po2(&mut self, limit: u32) -> Result<&mut Self> {
222        if !KECCAK_PO2_RANGE.contains(&(limit as usize)) {
223            bail!(
224                "invalid keccak po2 {limit}. Expected range: {:?}",
225                KECCAK_PO2_RANGE
226            );
227        }
228        self.inner.keccak_max_po2 = Some(limit);
229        Ok(self)
230    }
231
232    /// Set a session limit, specified in number of cycles.
233    ///
234    /// # Example
235    ///
236    /// ```
237    /// use risc0_zkvm::ExecutorEnv;
238    ///
239    /// let env = ExecutorEnv::builder()
240    ///     .session_limit(Some(32 * 1024 * 1024)) // 32M cycles
241    ///     .build()
242    ///     .unwrap();
243    /// ```
244    pub fn session_limit(&mut self, limit: Option<u64>) -> &mut Self {
245        self.inner.session_limit = limit;
246        self
247    }
248
249    /// Add environment variables to the guest environment.
250    ///
251    /// # Example
252    ///
253    /// ```
254    /// use std::collections::HashMap;
255    /// use risc0_zkvm::ExecutorEnv;
256    ///
257    /// let mut vars = HashMap::new();
258    /// vars.insert("VAR1".to_string(), "SOME_VALUE".to_string());
259    /// vars.insert("VAR2".to_string(), "SOME_VALUE".to_string());
260    ///
261    /// let env = ExecutorEnv::builder()
262    ///     .env_vars(vars)
263    ///     .build()
264    ///     .unwrap();
265    /// ```
266    pub fn env_vars(&mut self, vars: HashMap<String, String>) -> &mut Self {
267        self.inner.env_vars = vars;
268        self
269    }
270
271    /// Add an argument array to the guest environment.
272    ///
273    /// # Example
274    /// ```
275    /// # use risc0_zkvm::ExecutorEnv;
276    ///
277    /// let env = ExecutorEnv::builder()
278    ///     .args(&["grep".to_string(), "-c".to_string(), "foo".to_string(), "-".to_string()])
279    ///     .build()
280    ///     .unwrap();
281    /// ```
282    pub fn args(&mut self, args: &[String]) -> &mut Self {
283        self.inner.args.extend_from_slice(args);
284        self
285    }
286
287    /// Add an environment variable to the guest environment.
288    ///
289    /// # Example
290    ///
291    /// ```
292    /// use risc0_zkvm::ExecutorEnv;
293    ///
294    /// let env = ExecutorEnv::builder()
295    ///     .env_var("VAR1", "SOME_VALUE")
296    ///     .build()
297    ///     .unwrap();
298    /// ```
299    pub fn env_var(&mut self, name: &str, val: &str) -> &mut Self {
300        self.inner
301            .env_vars
302            .insert(name.to_string(), val.to_string());
303        self
304    }
305
306    /// Write input data to the zkVM guest stdin.
307    ///
308    /// This function will serialize `data` using a zkVM-optimized codec that
309    /// can be deserialized in the guest with a corresponding `env::read` with
310    /// the same data type.
311    ///
312    /// # Example
313    ///
314    /// ```
315    /// use risc0_zkvm::ExecutorEnv;
316    /// use serde::Serialize;
317    ///
318    /// #[derive(Serialize)]
319    /// struct Input {
320    ///     a: u32,
321    ///     b: u32,
322    /// }
323    ///
324    /// let input1 = Input{ a: 1, b: 2 };
325    /// let input2 = Input{ a: 3, b: 4 };
326    /// let env = ExecutorEnv::builder()
327    ///     .write(&input1).unwrap()
328    ///     .write(&input2).unwrap()
329    ///     .build()
330    ///     .unwrap();
331    /// ```
332    pub fn write<T: Serialize>(&mut self, data: &T) -> Result<&mut Self> {
333        Ok(self.write_slice(&to_vec(data)?))
334    }
335
336    /// Write input data to the zkVM guest stdin.
337    ///
338    /// This function writes a slice directly to the underlying buffer. A
339    /// corresponding `env::read_slice` can be used within the guest to read the
340    /// data.
341    ///
342    /// # Example
343    ///
344    /// ```
345    /// use risc0_zkvm::ExecutorEnv;
346    ///
347    /// let slice1 = [0, 1, 2, 3];
348    /// let slice2 = [3, 2, 1, 0];
349    /// let env = ExecutorEnv::builder()
350    ///     .write_slice(&slice1)
351    ///     .write_slice(&slice2)
352    ///     .build()
353    ///     .unwrap();
354    /// ```
355    pub fn write_slice<T: Pod>(&mut self, slice: &[T]) -> &mut Self {
356        self.inner
357            .input
358            .extend_from_slice(bytemuck::cast_slice(slice));
359        self
360    }
361
362    /// Write a frame to the zkVM guest via stdin.
363    ///
364    /// A frame contains a length header along with the payload. Reading a frame
365    /// can be more efficient than deserializing a message on-demand. On-demand
366    /// deserialization can cause many syscalls, whereas a frame will only have
367    /// two.
368    #[stability::unstable]
369    pub fn write_frame(&mut self, payload: &[u8]) -> &mut Self {
370        let len = payload.len() as u32;
371        self.inner.input.extend_from_slice(&len.to_le_bytes());
372        self.inner.input.extend_from_slice(payload);
373        self
374    }
375
376    /// Add a posix-style standard input.
377    pub fn stdin(&mut self, reader: impl Read + 'a) -> &mut Self {
378        self.read_fd(fileno::STDIN, BufReader::new(reader))
379    }
380
381    /// Add a posix-style standard output.
382    pub fn stdout(&mut self, writer: impl Write + 'a) -> &mut Self {
383        self.write_fd(fileno::STDOUT, writer)
384    }
385
386    /// Add a posix-style standard error.
387    pub fn stderr(&mut self, writer: impl Write + 'a) -> &mut Self {
388        self.write_fd(fileno::STDERR, writer)
389    }
390
391    /// Add a posix-style file descriptor for reading.
392    pub fn read_fd(&mut self, fd: u32, reader: impl BufRead + 'a) -> &mut Self {
393        self.inner.posix_io.borrow_mut().with_read_fd(fd, reader);
394        self
395    }
396
397    /// Add a posix-style file descriptor for writing.
398    pub fn write_fd(&mut self, fd: u32, writer: impl Write + 'a) -> &mut Self {
399        self.inner.posix_io.borrow_mut().with_write_fd(fd, writer);
400        self
401    }
402
403    /// Add a handler for simple I/O handling.
404    pub fn slice_io(&mut self, channel: &str, handler: impl SliceIo + 'a) -> &mut Self {
405        self.inner
406            .slice_io
407            .borrow_mut()
408            .with_handler(channel, handler);
409        self
410    }
411
412    /// Add a handler for simple I/O handling.
413    pub fn io_callback<C: AsRef<str>>(
414        &mut self,
415        channel: C,
416        callback: impl Fn(Bytes) -> Result<Bytes> + 'a,
417    ) -> &mut Self {
418        self.inner
419            .slice_io
420            .borrow_mut()
421            .with_handler(channel.as_ref(), slice_io_from_fn(callback));
422        self
423    }
424
425    /// Add an [AssumptionReceipt] to the [ExecutorEnv], for use in [composition].
426    ///
427    /// During execution, when the guest calls `env::verify` or `env::verify_integrity`, this
428    /// collection will be searched for an [AssumptionReceipt] that corresponds the verification
429    /// call.
430    ///
431    /// Either a [crate::Receipt] or a [crate::ReceiptClaim] can be provided. If a [crate::Receipt]
432    /// is provided, then an [AssumptionReceipt::Proven] will be added to the [ExecutorEnv]
433    /// and the [crate::Receipt] generated by proving will be unconditional.
434    ///
435    /// [composition]: https://dev.risczero.com/terminology#composition
436    pub fn add_assumption(&mut self, assumption: impl Into<AssumptionReceipt>) -> &mut Self {
437        self.inner
438            .assumptions
439            .borrow_mut()
440            .0
441            .push(assumption.into());
442        self
443    }
444
445    /// Add a callback handler for raw trace messages.
446    pub fn trace_callback(&mut self, callback: impl TraceCallback + 'a) -> &mut Self {
447        self.inner.trace.push(Rc::new(RefCell::new(callback)));
448        self
449    }
450
451    /// Set the path where segments will be stored.
452    pub fn segment_path<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
453        self.inner.segment_path = Some(SegmentPath::Path(path.as_ref().to_path_buf()));
454        self
455    }
456
457    /// Enable the profiler and output results to the specified path.
458    pub fn enable_profiler<P: AsRef<Path>>(&mut self, path: P) -> &mut Self {
459        self.inner.pprof_out = Some(path.as_ref().to_path_buf());
460        self
461    }
462
463    /// Set the input digest.
464    pub fn input_digest(&mut self, digest: Digest) -> &mut Self {
465        self.inner.input_digest = Some(digest);
466        self
467    }
468
469    /// Add a callback for coprocessor requests.
470    #[stability::unstable]
471    pub fn coprocessor_callback(&mut self, callback: impl CoprocessorCallback + 'a) -> &mut Self {
472        self.inner.coprocessor = Some(Rc::new(RefCell::new(callback)));
473        self
474    }
475
476    /// Add a callback for coprocessor requests.
477    #[stability::unstable]
478    pub fn coprocessor_callback_ref(&mut self, callback: CoprocessorCallbackRef<'a>) -> &mut Self {
479        self.inner.coprocessor = Some(callback);
480        self
481    }
482}