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        let _child = po.start();
24        debug!("Let server start");
25        std::thread::sleep(std::time::Duration::from_secs(5));
26
27        debug!("Reading from {}", preimage_key_dir);
28        // Get all files under the preimage db dir
29        let paths = std::fs::read_dir(preimage_key_dir)
30            .unwrap()
31            .map(|res| res.map(|e| e.path()))
32            .collect::<Result<Vec<_>, io::Error>>()
33            .unwrap();
34
35        // Just take 10 elements at random to test
36        let how_many = 10_usize;
37        let max = paths.len();
38        let mut rng = rand::thread_rng();
39
40        let mut selected: Vec<PathBuf> = vec![PathBuf::new(); how_many];
41        for pb in selected.iter_mut() {
42            let idx = rng.gen_range(0..max);
43            *pb = paths[idx].to_path_buf();
44        }
45
46        for (idx, path) in selected.into_iter().enumerate() {
47            let preimage_key = Path::new(&path)
48                .file_name()
49                .unwrap()
50                .to_str()
51                .unwrap()
52                .split('.')
53                .collect::<Vec<_>>()[0];
54
55            let hash = PreimageKey::from_str(preimage_key).unwrap();
56            let mut key = hash.0;
57            key[0] = 2; // Keccak
58
59            debug!(
60                "Generating OP Keccak key for {} at index {}",
61                preimage_key, idx
62            );
63
64            let expected = std::fs::read_to_string(path).unwrap();
65
66            debug!("Asking for preimage");
67            let preimage = po.get_preimage(key);
68            let got = hex::encode(preimage.get()).to_string();
69
70            assert_eq!(expected, got);
71        }
72        ExitCode::SUCCESS
73    } else {
74        error!("Unset command-line argument --preimage-db-dir. Cannot run test. Please set parameter and rerun.");
75        ExitCode::FAILURE
76    }
77}
o1vm/test_preimage_read.rs

o1vm/
test_preimage_read.rs
