newton_cli/commands/
policy_data.rs

1use alloy::primitives::{Address, FixedBytes};
2use clap::{Parser, Subcommand};
3use eyre::{Context, Result};
4use newton_prover_chainio::policy_data::PolicyDataController;
5use newton_prover_core::{
6    common::chain::get_block_time_ms, config::NewtonAvsConfig, newton_policy_data::INewtonPolicyData,
7};
8use serde_json::Value;
9use std::path::PathBuf;
10use tracing::{self, info, warn};
11
12use crate::{commands::utils, config::NewtonCliConfig, task_generator};
13
14/// Policy data commands
15#[derive(Debug, Parser)]
16#[command(name = "policy-data")]
17pub struct PolicyDataCommand {
18    #[command(subcommand)]
19    pub subcommand: PolicyDataSubcommand,
20}
21
22#[derive(Debug, Subcommand)]
23pub enum PolicyDataSubcommand {
24    /// Deploy policy data
25    Deploy(DeployCommand),
26    /// Simulate WASM execution
27    Simulate(SimulateCommand),
28}
29
30/// Simulate WASM execution command
31#[derive(Debug, Parser)]
32pub struct SimulateCommand {
33    /// Path to the WASM component file
34    #[arg(long)]
35    wasm_file: PathBuf,
36
37    /// Input JSON string
38    #[arg(long)]
39    input_json: String,
40}
41
42/// Deploy policy data command
43#[derive(Debug, Parser)]
44pub struct DeployCommand {
45    #[arg(long, env = "PRIVATE_KEY")]
46    private_key: Option<String>,
47
48    #[arg(long, env = "RPC_URL")]
49    rpc_url: Option<String>,
50
51    #[arg(long, default_value = "policy-files/policy_cids.json")]
52    policy_cids: PathBuf,
53
54    /// Expiration in blocks. If not set, defaults to 300 seconds converted to blocks using the chain's block time.
55    #[arg(long)]
56    expire_after_blocks: Option<u32>,
57}
58
59impl DeployCommand {
60    /// Convert seconds to blocks based on chain's block time
61    async fn seconds_to_blocks(rpc_url: &str, seconds: u32) -> Result<u32> {
62        let block_time_ms = get_block_time_ms(rpc_url).await?;
63        let block_time_seconds = block_time_ms / 1000;
64
65        if block_time_seconds == 0 {
66            return Err(eyre::eyre!("Block time is zero, cannot convert seconds to blocks"));
67        }
68
69        let blocks = (seconds as u64 / block_time_seconds) as u32;
70        tracing::info!(
71            "Converting {} seconds to blocks: block_time={}ms ({}s), result={} blocks",
72            seconds,
73            block_time_ms,
74            block_time_seconds,
75            blocks
76        );
77
78        Ok(blocks)
79    }
80
81    /// Deploy policy data contract
82    ///
83    /// This function deploys a policy data contract using the Rust chainio functions.
84    /// It returns the deployed policy data address.
85    /// Note: expire_after_blocks should be in blocks, not seconds.
86    async fn deploy_policy_data(
87        private_key: &str,
88        rpc_url: &str,
89        wasm_cid: &str,
90        secrets_schema_cid: &str,
91        policy_data_metadata_cid: &str,
92        expire_after_blocks: u32,
93    ) -> Result<Address> {
94        let controller = PolicyDataController::new(private_key.to_string(), rpc_url.to_string());
95
96        tracing::info!(
97            "Deploying policy data:\n  wasmCid: {}\n  secretsSchemaCid: {}\n  metadataCid: {}\n  expireAfter: {} blocks \n",
98            wasm_cid,
99            secrets_schema_cid,
100            policy_data_metadata_cid,
101            expire_after_blocks
102        );
103
104        let (_receipt, policy_data_address) = controller
105            .deploy_policy_data(
106                wasm_cid.to_string(),
107                secrets_schema_cid.to_string(),
108                expire_after_blocks,
109                policy_data_metadata_cid.to_string(),
110            )
111            .await
112            .map_err(|e| eyre::eyre!("Failed to deploy policy data: {}", e))?;
113
114        tracing::info!("Policy data deployed successfully at address: {}", policy_data_address);
115
116        Ok(policy_data_address)
117    }
118
119    /// Set attestation info on deployed policy data
120    ///
121    /// This function sets the attestation info on a deployed policy data contract.
122    /// It creates an AttestationInfo struct with:
123    /// - attestationType: ECDSA (0)
124    /// - verifier: address(0)
125    /// - verificationKey: bytes32(0)
126    async fn set_attestation_info(
127        private_key: &str,
128        rpc_url: &str,
129        policy_data_address: Address,
130        attester_address: Address,
131        chain_id: u64,
132        deployment_env: &str,
133    ) -> Result<()> {
134        let controller = PolicyDataController::new(private_key.to_string(), rpc_url.to_string());
135
136        // Determine which environment's taskGenerator[] we should use.
137        // We infer this by matching the attester (policy_cids.json) against `taskGenerator[0]`.
138        let prod0 = task_generator::task_generator_0("prod", chain_id)
139            .and_then(|s| s.parse::<Address>().ok());
140        let stagef0 = task_generator::task_generator_0("stagef", chain_id)
141            .and_then(|s| s.parse::<Address>().ok());
142
143        let inferred_env = if prod0 == Some(attester_address) {
144            "prod"
145        } else if stagef0 == Some(attester_address) {
146            "stagef"
147        } else {
148            deployment_env
149        };
150
151        let task_generators = task_generator::task_generators(inferred_env, chain_id).ok_or_else(
152            || eyre::eyre!("No hardcoded taskGenerator[] for env={} chain_id={}", inferred_env, chain_id)
153        )?;
154
155        // Parse task generator addresses.
156        let task_generator_addresses: Vec<Address> = task_generators
157            .iter()
158            .filter_map(|s| s.parse::<Address>().ok())
159            .collect();
160
161        // Create attesters array: [attester, ...task_generator_addresses]
162        let mut attesters = vec![attester_address];
163        attesters.extend(task_generator_addresses);
164
165        // Create AttestationInfo struct
166        let attestation_info = INewtonPolicyData::AttestationInfo {
167            attesters,
168            attestationType: 0u8, // ECDSA
169            verifier: Address::ZERO,
170            verificationKey: FixedBytes::ZERO,
171        };
172
173        tracing::info!(
174            "Setting attestation info on policy data {} with attesters: {:?}",
175            policy_data_address,
176            attestation_info.attesters
177        );
178
179        controller
180            .set_attestation_info(policy_data_address, attestation_info)
181            .await
182            .map_err(|e| eyre::eyre!("Failed to set attestation info: {}", e))?;
183
184        tracing::info!(
185            "Attestation info set successfully on policy data: {}",
186            policy_data_address
187        );
188
189        Ok(())
190    }
191
192    /// Execute the deploy command
193    pub async fn execute(
194        self: Box<Self>,
195        config: NewtonAvsConfig<NewtonCliConfig>,
196    ) -> eyre::Result<()> {
197        // Get values from args or env, with error if still missing
198        let private_key = self
199            .private_key
200            .ok_or_else(|| eyre::eyre!("private_key is required (use --private-key or PRIVATE_KEY env var)"))?;
201
202        let rpc_url = self
203            .rpc_url
204            .ok_or_else(|| eyre::eyre!("rpc_url is required (use --rpc-url or RPC_URL env var)"))?;
205
206        // Read from policy_cids.json file
207        let json_content = std::fs::read_to_string(&self.policy_cids)
208            .with_context(|| format!("Failed to read policy_cids.json: {:?}", self.policy_cids))?;
209        let json: Value = serde_json::from_str(&json_content)
210            .with_context(|| format!("Failed to parse policy_cids.json: {:?}", self.policy_cids))?;
211
212        // Extract values from policy_cids.json
213        let wasm_cid = json.get("wasmCid").and_then(|v| v.as_str()).unwrap_or("");
214        let secrets_schema_cid = json.get("secretsSchemaCid").and_then(|v| v.as_str()).unwrap_or("");
215        let policy_data_metadata_cid = json.get("policyDataMetadataCid").and_then(|v| v.as_str()).unwrap_or("");
216        let attester = json.get("attester").and_then(|v| v.as_str()).unwrap_or("");
217        let attester_address: Address = attester.parse().unwrap_or(Address::ZERO);
218
219        // Use --expire-after-blocks if provided, otherwise convert 300 seconds to blocks
220        let expire_after_blocks = match self.expire_after_blocks {
221            Some(blocks) => {
222                tracing::info!("Using expire_after_blocks from command: {} blocks", blocks);
223                blocks
224            }
225            None => Self::seconds_to_blocks(&rpc_url, 300).await?,
226        };
227
228        // Deploy policy data
229        let policy_data_address = Self::deploy_policy_data(
230            &private_key,
231            &rpc_url,
232            wasm_cid,
233            secrets_schema_cid,
234            policy_data_metadata_cid,
235            expire_after_blocks,
236        )
237        .await?;
238
239        // Allow time for the deploy tx to be reflected in RPC state (nonce update).
240        // This prevents "replacement transaction underpriced" when setting attestation info.
241        tracing::info!("Waiting for RPC to reflect updated nonce before setting attestation info...");
242        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
243
244        tracing::info!("Policy data deployment completed! Setting attestation info...");
245        // Set attestation info on policy data with retries for nonce/underpriced errors
246        const MAX_RETRIES: u32 = 3;
247        let mut attempt = 0;
248        loop {
249            match Self::set_attestation_info(
250                &private_key,
251                &rpc_url,
252                policy_data_address,
253                attester_address,
254                config.chain_id,
255                config.env.as_str(),
256            )
257            .await
258            {
259                Ok(()) => break,
260                Err(e) => {
261                    let error_str = e.to_string();
262                    let is_nonce_error = error_str.contains("underpriced") || error_str.contains("nonce");
263                    if is_nonce_error && attempt < MAX_RETRIES {
264                        attempt += 1;
265                        warn!(
266                            "Set attestation info nonce issue (attempt {}/{}), retrying in 2s: {}",
267                            attempt, MAX_RETRIES, e
268                        );
269                        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
270                    } else {
271                        return Err(e);
272                    }
273                }
274            }
275        }
276
277        tracing::info!("Successfully deployed policy data and set attestation");
278        tracing::info!("Policy Data Address: {}", policy_data_address);
279
280        Ok(())
281    }
282}
283
284impl SimulateCommand {
285    /// Execute the simulate command
286    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
287        info!("Executing WASM simulation...");
288        let output = utils::execute_wasm(self.wasm_file, self.input_json, config).await?;
289        info!("WASM simulation output: {}", output);
290        Ok(())
291    }
292}
293
294impl PolicyDataCommand {
295    /// Execute the policy-data command
296    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
297        match self.subcommand {
298            PolicyDataSubcommand::Deploy(command) => {
299                Box::new(command).execute(config).await?;
300            }
301            PolicyDataSubcommand::Simulate(command) => {
302                Box::new(command).execute(config).await?;
303            }
304        }
305        Ok(())
306    }
307}
newton_cli/commands/policy_data.rs

newton_cli/commands/
policy_data.rs
