cargo_openvm/commands/
keygen.rs

1use std::{
2    fs::{copy, create_dir_all},
3    path::{Path, PathBuf},
4};
5
6use clap::Parser;
7use eyre::{Context, Result};
8use openvm_sdk::{config::AggregationSystemParams, fs::write_object_to_file, Sdk};
9
10use crate::{
11    args::ManifestArgs,
12    default::{
13        default_app_config, DEFAULT_AGG_PREFIX_PK_NAME, DEFAULT_APP_PK_NAME, DEFAULT_APP_VK_NAME,
14        OPENVM_CONFIG_FILENAME,
15    },
16    util::{
17        get_agg_prefix_pk_path, get_app_pk_path, get_app_vk_path, get_manifest_path_and_dir,
18        get_target_dir, read_config_toml_or_default,
19    },
20};
21
22#[derive(Parser)]
23#[command(name = "keygen", about = "Generate application and aggregation keys")]
24pub struct KeygenCmd {
25    #[arg(
26        long,
27        help = "Path to the OpenVM config .toml file that specifies the VM extensions, by default will search for the file at ${manifest_dir}/openvm.toml",
28        help_heading = "OpenVM Options"
29    )]
30    config: Option<PathBuf>,
31
32    #[arg(
33        long,
34        help = "Output directory that OpenVM proving artifacts will be copied to",
35        help_heading = "OpenVM Options"
36    )]
37    output_dir: Option<PathBuf>,
38
39    #[arg(
40        long,
41        help = "Only generate app keys (app.pk/app.vk), skip the aggregation prefix proving key (agg_prefix.pk)",
42        help_heading = "OpenVM Options"
43    )]
44    app_only: bool,
45
46    #[command(flatten)]
47    cargo_args: KeygenCargoArgs,
48}
49
50#[derive(Parser)]
51pub struct KeygenCargoArgs {
52    #[clap(flatten)]
53    pub(crate) manifest: ManifestArgs,
54}
55
56impl KeygenCmd {
57    pub fn run(&self) -> Result<()> {
58        let (manifest_path, manifest_dir) =
59            get_manifest_path_and_dir(&self.cargo_args.manifest.manifest_path)?;
60        let target_dir = get_target_dir(&self.cargo_args.manifest.target_dir, &manifest_path);
61        let app_pk_path = get_app_pk_path(&target_dir);
62        let app_vk_path = get_app_vk_path(&target_dir);
63        let agg_prefix_pk_path = get_agg_prefix_pk_path(&target_dir);
64
65        keygen(
66            self.config
67                .to_owned()
68                .unwrap_or_else(|| manifest_dir.join(OPENVM_CONFIG_FILENAME)),
69            &app_pk_path,
70            &app_vk_path,
71            &agg_prefix_pk_path,
72            self.output_dir.as_ref(),
73            !self.app_only,
74        )?;
75        println!(
76            "Successfully generated {} in {}",
77            if self.app_only {
78                "app pk and vk"
79            } else {
80                "app pk/vk and agg_prefix pk"
81            },
82            if let Some(output_dir) = self.output_dir.as_ref() {
83                output_dir.display()
84            } else {
85                app_pk_path.parent().unwrap().display()
86            }
87        );
88        Ok(())
89    }
90}
91
92pub(crate) fn keygen(
93    config: impl AsRef<Path>,
94    app_pk_path: impl AsRef<Path>,
95    app_vk_path: impl AsRef<Path>,
96    agg_prefix_pk_path: impl AsRef<Path>,
97    output_dir: Option<impl AsRef<Path>>,
98    generate_agg: bool,
99) -> Result<()> {
100    let app_config = read_config_toml_or_default(config)?;
101    assert_default_root_shape(&app_config);
102    let sdk = Sdk::new(app_config, AggregationSystemParams::default())?;
103    let (app_pk, app_vk) = sdk.app_keygen();
104    write_object_to_file(&app_vk_path, app_vk)?;
105    write_object_to_file(&app_pk_path, app_pk)?;
106    if generate_agg {
107        write_object_to_file(&agg_prefix_pk_path, sdk.agg_prefix_pk())?;
108    }
109
110    if let Some(output_dir) = output_dir {
111        let output_dir = output_dir.as_ref();
112        create_dir_all(output_dir)
113            .with_context(|| format!("failed to create directory {}", output_dir.display()))?;
114        copy(&app_pk_path, output_dir.join(DEFAULT_APP_PK_NAME))
115            .with_context(|| format!("failed to copy app pk to {}", output_dir.display()))?;
116        copy(&app_vk_path, output_dir.join(DEFAULT_APP_VK_NAME))
117            .with_context(|| format!("failed to copy app vk to {}", output_dir.display()))?;
118        if generate_agg {
119            copy(
120                &agg_prefix_pk_path,
121                output_dir.join(DEFAULT_AGG_PREFIX_PK_NAME),
122            )
123            .with_context(|| format!("failed to copy agg prefix pk to {}", output_dir.display()))?;
124        }
125    }
126
127    Ok(())
128}
129
130fn assert_default_root_shape(
131    app_config: &openvm_sdk::config::AppConfig<openvm_sdk_config::SdkVmConfig>,
132) {
133    let default_system_config = default_app_config().app_vm_config;
134    let actual_system_config = app_config.app_vm_config.as_ref();
135    let default_system_config = default_system_config.as_ref();
136
137    assert_eq!(
138        actual_system_config.num_public_values, default_system_config.num_public_values,
139        "cargo openvm keygen only supports the default num_public_values"
140    );
141    let actual_memory_dimensions = actual_system_config.memory_config.memory_dimensions();
142    let default_memory_dimensions = default_system_config.memory_config.memory_dimensions();
143    assert!(
144        actual_memory_dimensions.addr_space_height == default_memory_dimensions.addr_space_height
145            && actual_memory_dimensions.address_height == default_memory_dimensions.address_height,
146        "cargo openvm keygen only supports the default memory_dimensions"
147    );
148}