Skip to main content

airbender_host/
vk.rs

1use crate::error::{HostError, Result};
2use crate::prover::ProverLevel;
3use crate::security::SecurityLevel;
4use airbender_core::guest::Commit;
5use execution_utils::setups;
6use execution_utils::unified_circuit::verify_proof_in_unified_layer;
7use execution_utils::unrolled::{
8    compute_setup_for_machine_configuration, get_unrolled_circuits_artifacts_for_machine_type,
9    verify_unrolled_layer_proof, UnrolledProgramProof, UnrolledProgramSetup,
10};
11use riscv_transpiler::cycle::{
12    IMStandardIsaConfigWithUnsignedMulDiv, IWithoutByteAccessIsaConfigWithDelegation,
13};
14use sha3::Digest;
15use std::fs;
16use std::path::{Path, PathBuf};
17
18/// Unified verification key bundle for recursion.
19#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
20pub struct UnifiedVk {
21    pub security: SecurityLevel,
22    pub app_bin_hash: [u8; 32],
23    pub unified_setup: UnrolledProgramSetup,
24    pub unified_layouts: setups::CompiledCircuitsSet,
25}
26
27/// Unrolled verification key bundle for base or recursion-unrolled layers.
28#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
29pub struct UnrolledVk {
30    pub security: SecurityLevel,
31    pub app_bin_hash: [u8; 32],
32    pub setup: UnrolledProgramSetup,
33    pub compiled_layouts: setups::CompiledCircuitsSet,
34}
35
36pub fn compute_unified_vk(app_bin_path: &Path, security: SecurityLevel) -> Result<UnifiedVk> {
37    #[cfg(not(feature = "gpu-prover"))]
38    {
39        let _ = (app_bin_path, security);
40        return Err(HostError::Verification(
41            "recursion-unified verification key generation requires the `gpu-prover` feature"
42                .to_string(),
43        ));
44    }
45
46    #[cfg(feature = "gpu-prover")]
47    {
48        let app_bin_hash = hash_app_bin(app_bin_path)?;
49
50        // TODO: cache unified setup/layout artifacts on disk to avoid recomputing on every run.
51        let security_model = security.into();
52        let (binary, binary_u32) = setups::pad_binary(
53            execution_utils::verifier_binaries::recursion_artifact(
54                security_model,
55                execution_utils::RecursionLayer::Unified,
56                execution_utils::RecursionArtifact::Bin,
57            )
58            .to_vec(),
59        );
60        let (text, _) = setups::pad_binary(
61            execution_utils::verifier_binaries::recursion_artifact(
62                security_model,
63                execution_utils::RecursionLayer::Unified,
64                execution_utils::RecursionArtifact::Txt,
65            )
66            .to_vec(),
67        );
68
69        let unified_setup =
70            execution_utils::unified_circuit::compute_unified_setup_for_machine_configuration::<
71                IWithoutByteAccessIsaConfigWithDelegation,
72            >(&binary, &text);
73        let unified_layouts =
74            execution_utils::setups::get_unified_circuit_artifact_for_machine_type::<
75                IWithoutByteAccessIsaConfigWithDelegation,
76            >(&binary_u32);
77
78        Ok(UnifiedVk {
79            security,
80            app_bin_hash,
81            unified_setup,
82            unified_layouts,
83        })
84    }
85}
86
87pub fn compute_unrolled_vk(
88    app_bin_path: &Path,
89    level: ProverLevel,
90    security: SecurityLevel,
91) -> Result<UnrolledVk> {
92    if level == ProverLevel::RecursionUnified {
93        return Err(HostError::Verification(
94            "unified verification keys must be generated with compute_unified_vk".to_string(),
95        ));
96    }
97
98    let resolved_bin_path = resolve_bin_path(app_bin_path)?;
99    let app_bin_hash = hash_app_bin(&resolved_bin_path)?;
100
101    let (binary, binary_u32, text) = match level {
102        ProverLevel::Base => {
103            let app_text_path = resolve_text_path(&resolved_bin_path)?;
104            let (binary, binary_u32) = setups::read_and_pad_binary(&resolved_bin_path);
105            let (text, _) = setups::read_and_pad_binary(&app_text_path);
106            (binary, binary_u32, text)
107        }
108        ProverLevel::RecursionUnrolled => {
109            #[cfg(not(feature = "gpu-prover"))]
110            {
111                return Err(HostError::Verification(
112                    "recursion-unrolled verification key generation requires the `gpu-prover` feature"
113                        .to_string(),
114                ));
115            }
116
117            #[cfg(feature = "gpu-prover")]
118            {
119                let security_model = security.into();
120                let (binary, binary_u32) = setups::pad_binary(
121                    execution_utils::verifier_binaries::recursion_artifact(
122                        security_model,
123                        execution_utils::RecursionLayer::Unrolled,
124                        execution_utils::RecursionArtifact::Bin,
125                    )
126                    .to_vec(),
127                );
128                let (text, _) = setups::pad_binary(
129                    execution_utils::verifier_binaries::recursion_artifact(
130                        security_model,
131                        execution_utils::RecursionLayer::Unrolled,
132                        execution_utils::RecursionArtifact::Txt,
133                    )
134                    .to_vec(),
135                );
136                (binary, binary_u32, text)
137            }
138        }
139        ProverLevel::RecursionUnified => {
140            return Err(HostError::Verification(
141                "unified verification keys must be generated with compute_unified_vk".to_string(),
142            ));
143        }
144    };
145
146    let (setup, compiled_layouts) = match level {
147        ProverLevel::Base => {
148            let setup = compute_setup_for_machine_configuration::<
149                IMStandardIsaConfigWithUnsignedMulDiv,
150            >(&binary, &text);
151            let compiled_layouts = get_unrolled_circuits_artifacts_for_machine_type::<
152                IMStandardIsaConfigWithUnsignedMulDiv,
153            >(&binary_u32);
154            (setup, compiled_layouts)
155        }
156        ProverLevel::RecursionUnrolled => {
157            let setup = compute_setup_for_machine_configuration::<
158                IWithoutByteAccessIsaConfigWithDelegation,
159            >(&binary, &text);
160            let compiled_layouts = get_unrolled_circuits_artifacts_for_machine_type::<
161                IWithoutByteAccessIsaConfigWithDelegation,
162            >(&binary_u32);
163            (setup, compiled_layouts)
164        }
165        ProverLevel::RecursionUnified => {
166            return Err(HostError::Verification(
167                "unified verification keys must be generated with compute_unified_vk".to_string(),
168            ));
169        }
170    };
171
172    Ok(UnrolledVk {
173        security,
174        app_bin_hash,
175        setup,
176        compiled_layouts,
177    })
178}
179
180pub fn verify_proof(
181    proof: &UnrolledProgramProof,
182    vk: &UnifiedVk,
183    expected_app_bin_hash: Option<[u8; 32]>,
184    expected_output: Option<&dyn Commit>,
185) -> Result<()> {
186    verify_app_bin_hash(expected_app_bin_hash, vk.app_bin_hash)?;
187
188    let verifier_output = verify_proof_in_unified_layer(
189        proof,
190        &vk.unified_setup,
191        &vk.unified_layouts,
192        false,
193        vk.security.into(),
194    )
195    .map_err(|_| HostError::Verification("proof verification failed".to_string()))?;
196    verify_expected_output(expected_output, verifier_output)?;
197    Ok(())
198}
199
200pub fn verify_unrolled_proof(
201    proof: &UnrolledProgramProof,
202    vk: &UnrolledVk,
203    level: ProverLevel,
204    expected_app_bin_hash: Option<[u8; 32]>,
205    expected_output: Option<&dyn Commit>,
206) -> Result<()> {
207    verify_app_bin_hash(expected_app_bin_hash, vk.app_bin_hash)?;
208
209    let is_base_layer = match level {
210        ProverLevel::Base => true,
211        ProverLevel::RecursionUnrolled => false,
212        ProverLevel::RecursionUnified => {
213            return Err(HostError::Verification(
214                "recursion-unified proofs must be verified with unified verification keys"
215                    .to_string(),
216            ));
217        }
218    };
219
220    let verifier_output = verify_unrolled_layer_proof(
221        proof,
222        &vk.setup,
223        &vk.compiled_layouts,
224        is_base_layer,
225        vk.security.into(),
226    )
227    .map_err(|_| HostError::Verification("proof verification failed".to_string()))?;
228    verify_expected_output(expected_output, verifier_output)?;
229    Ok(())
230}
231
232fn verify_expected_output(
233    expected_output: Option<&dyn Commit>,
234    verifier_output: [u32; 16],
235) -> Result<()> {
236    let Some(expected_output) = expected_output else {
237        return Ok(());
238    };
239
240    let expected_words = expected_output.commit_words();
241    let mut actual_words = [0u32; 8];
242    actual_words.copy_from_slice(&verifier_output[..8]);
243
244    if expected_words != actual_words {
245        return Err(HostError::Verification(format!(
246            "public output mismatch: expected {expected_words:?}, got {actual_words:?}"
247        )));
248    }
249
250    Ok(())
251}
252
253fn verify_app_bin_hash(
254    expected_app_bin_hash: Option<[u8; 32]>,
255    actual_app_bin_hash: [u8; 32],
256) -> Result<()> {
257    if let Some(expected) = expected_app_bin_hash {
258        if expected != actual_app_bin_hash {
259            return Err(HostError::Verification(
260                "app.bin hash does not match verification key".to_string(),
261            ));
262        }
263    }
264    Ok(())
265}
266
267fn hash_app_bin(path: &Path) -> Result<[u8; 32]> {
268    let app_bin_bytes = fs::read(path)?;
269    Ok(sha3::Keccak256::digest(&app_bin_bytes).into())
270}
271
272fn resolve_bin_path(path: &Path) -> Result<PathBuf> {
273    let base_path = base_path(path)?;
274    let app_bin_path = PathBuf::from(format!("{base_path}.bin"));
275
276    if !app_bin_path.exists() {
277        return Err(HostError::Verification(format!(
278            "binary not found: {}",
279            app_bin_path.display()
280        )));
281    }
282
283    Ok(app_bin_path)
284}
285
286fn resolve_text_path(app_bin_path: &Path) -> Result<PathBuf> {
287    let mut app_text_path = app_bin_path.to_path_buf();
288    app_text_path.set_extension("text");
289
290    if !app_text_path.exists() {
291        return Err(HostError::Verification(format!(
292            "text file not found: {}",
293            app_text_path.display()
294        )));
295    }
296
297    Ok(app_text_path)
298}
299
300fn base_path(app_bin_path: &Path) -> Result<String> {
301    let path_str = app_bin_path
302        .to_str()
303        .ok_or_else(|| HostError::Verification("app path is not valid UTF-8".to_string()))?;
304    if let Some(stripped) = path_str.strip_suffix(".bin") {
305        Ok(stripped.to_string())
306    } else {
307        Ok(path_str.to_string())
308    }
309}
310
311#[cfg(test)]
312mod tests {
313    use super::verify_expected_output;
314
315    #[test]
316    fn verify_expected_output_accepts_matching_words() {
317        let mut verifier_output = [0u32; 16];
318        verifier_output[0] = 42;
319
320        verify_expected_output(Some(&42u32), verifier_output).expect("matching output must verify");
321    }
322
323    #[test]
324    fn verify_expected_output_rejects_mismatch() {
325        let verifier_output = [0u32; 16];
326
327        let err = verify_expected_output(Some(&1u32), verifier_output)
328            .expect_err("mismatching output must fail verification");
329        assert!(err.to_string().contains("public output mismatch"));
330    }
331}