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}