o1vm/
test_preimage_read.rs

1use crate::{
2    cannon::{PreimageKey, VmConfiguration},
3    cli,
4    preimage_oracle::{PreImageOracle, PreImageOracleT},
5};
6use log::{debug, error};
7use std::{
8    io::{self},
9    path::{Path, PathBuf},
10    process::ExitCode,
11    str::FromStr,
12};
13
14pub fn main(args: cli::cannon::RunArgs) -> ExitCode {
15    use rand::Rng;
16
17    let configuration: VmConfiguration = args.vm_cfg.into();
18    let preimage_db_dir = args.preimage_db_dir;
19
20    if let Some(preimage_key_dir) = preimage_db_dir {
21        let host_program = configuration.host.expect("No host program specified");
22        let mut po = PreImageOracle::create(host_program);
23        #[allow(unknown_lints, clippy::zombie_processes)]
24        let _child = po.start();
25        debug!("Let server start");
26        std::thread::sleep(std::time::Duration::from_secs(5));
27
28        debug!("Reading from {}", preimage_key_dir);
29        // Get all files under the preimage db dir
30        let paths = std::fs::read_dir(preimage_key_dir)
31            .unwrap()
32            .map(|res| res.map(|e| e.path()))
33            .collect::<Result<Vec<_>, io::Error>>()
34            .unwrap();
35
36        // Just take 10 elements at random to test
37        let how_many = 10_usize;
38        let max = paths.len();
39        let mut rng = rand::thread_rng();
40
41        let mut selected: Vec<PathBuf> = vec![PathBuf::new(); how_many];
42        for pb in selected.iter_mut() {
43            let idx = rng.gen_range(0..max);
44            *pb = paths[idx].to_path_buf();
45        }
46
47        for (idx, path) in selected.into_iter().enumerate() {
48            let preimage_key = Path::new(&path)
49                .file_name()
50                .unwrap()
51                .to_str()
52                .unwrap()
53                .split('.')
54                .collect::<Vec<_>>()[0];
55
56            let hash = PreimageKey::from_str(preimage_key).unwrap();
57            let mut key = hash.0;
58            key[0] = 2; // Keccak
59
60            debug!(
61                "Generating OP Keccak key for {} at index {}",
62                preimage_key, idx
63            );
64
65            let expected = std::fs::read_to_string(path).unwrap();
66
67            debug!("Asking for preimage");
68            let preimage = po.get_preimage(key);
69            let got = hex::encode(preimage.get()).to_string();
70
71            assert_eq!(expected, got);
72        }
73        ExitCode::SUCCESS
74    } else {
75        error!("Unset command-line argument --preimage-db-dir. Cannot run test. Please set parameter and rerun.");
76        ExitCode::FAILURE
77    }
78}