1
use crate::common::{proof_get_inputs, ZkProgramManifest};
2
use anyhow::{anyhow, Result};
3
use bonsol_prover::image::Image;
4
use bonsol_prover::prover::{get_risc0_prover, new_risc0_exec_env};
5
use bonsol_sdk::BonsolClient;
6
use bytes::Bytes;
7
use risc0_zkvm::VerifierContext;
8
use std::fs::{read, File};
9
use std::io::Write;
10
use std::path::Path;
11

            
12
pub async fn prove(
13
    sdk: &BonsolClient,
14
    execution_id: String,
15
    manifest_path: Option<String>,
16
    program_id: Option<String>,
17
    input_file: Option<String>,
18
    output_location: Option<String>,
19
    stdin: Option<String>,
20
) -> Result<()> {
21
    let pwd = std::env::current_dir()?;
22
    let image_bytes = match (&program_id, manifest_path) {
23
        (Some(i), None) => {
24
            let bytes: Bytes = sdk.download_program(i).await?;
25
            Ok(bytes)
26
        }
27
        (None, Some(m)) => {
28
            let manifest_path = Path::new(&m);
29
            let manifest_file = if manifest_path.is_relative() {
30
                File::open(pwd.join(manifest_path))?
31
            } else {
32
                File::open(manifest_path)?
33
            };
34
            let manifest: ZkProgramManifest = serde_json::from_reader(manifest_file)?;
35
            let binary_path = Path::new(&manifest.binary_path);
36
            let bytes =
37
                read(binary_path).map_err(|_| anyhow!("Failed to read binary in manifest file"))?;
38
            Ok(Bytes::from(bytes))
39
        }
40
        _ => Err(anyhow!("Please provide a program id or a manifest path")),
41
    }?;
42
    let ext = Path::new(&execution_id).with_extension("bin");
43
    let output_binary_path = output_location
44
        .map(|o| Path::new(&o).join(&ext))
45
        .unwrap_or(ext);
46
    let image = Image::from_bytes(image_bytes)?;
47
    let memory_image = image.get_memory_image()?;
48
    let program_inputs = proof_get_inputs(input_file, stdin)?;
49
    let mut exec = new_risc0_exec_env(memory_image, program_inputs)?;
50
    let session = exec.run()?;
51
    let prover = get_risc0_prover()?;
52
    let ctx = VerifierContext::default();
53
    println!("Generating proof");
54
    let info = prover.prove_session(&ctx, &session);
55
    match info {
56
        Ok(proveinfo) => {
57
            let proof = bincode::serialize(&proveinfo.receipt)?;
58
            let mut file = File::create(&output_binary_path)?;
59
            file.write_all(&proof)?;
60
            println!("Proof written to {}", output_binary_path.to_string_lossy());
61
        }
62
        Err(e) => {
63
            println!("Error generating proof: {:?}", e);
64
        }
65
    }
66
    Ok(())
67
}