Skip to main content

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};
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    /// Get task generator addresses
82    ///
83    /// Returns all 12 task generator addresses for all networks.
84    async fn get_task_generator_addresses(_rpc_url: &str) -> Result<Vec<Address>> {
85        let addresses = vec![
86            "0x4883282094755C01cd0d15dFE74753c9E189d194",
87            "0x51eBfB4c0441b0D8898d968975073c24B9190227",
88            "0x19F9dEC9928417B6Ea6aD52291acCa9deef59E3F",
89            "0x9Ccd55A0Fb8aBD14919bb5B2BCDD39C1Dd40E883",
90            "0xb3785B3CA4d6b175518c9a2A51AE836f2dE3016B",
91            "0xC6aA2638873e90eCaF0306cbCC4370C6AC893E68",
92            "0x64E33a6b2874129f2C6FB61832b763F65C8342fb",
93            "0xA21d5CCa5771DAB5a049b1FAA524c41ae911e185",
94            "0x22748005349aD3e261Fc6B1C83956781d6d661B4",
95            "0x3F15Da29D0de8c45bf625299AA20e43970153622",
96            "0x494221AaA9B9B273EB171512Ea000bEaE868d332",
97            "0xD45062003a4626a532F30A4596aB253c45AE0647",
98        ];
99
100        addresses
101            .iter()
102            .map(|addr| {
103                addr.parse()
104                    .map_err(|e| eyre::eyre!("Failed to parse task generator address {}: {}", addr, e))
105            })
106            .collect()
107    }
108
109    /// Deploy policy data contract
110    ///
111    /// This function deploys a policy data contract using the Rust chainio functions.
112    /// It returns the deployed policy data address.
113    /// Note: expire_after_blocks should be in blocks, not seconds.
114    async fn deploy_policy_data(
115        private_key: &str,
116        rpc_url: &str,
117        wasm_cid: &str,
118        secrets_schema_cid: &str,
119        policy_data_metadata_cid: &str,
120        expire_after_blocks: u32,
121    ) -> Result<Address> {
122        let controller = PolicyDataController::new(private_key.to_string(), rpc_url.to_string());
123
124        tracing::info!(
125            "Deploying policy data:\n  wasmCid: {}\n  secretsSchemaCid: {}\n  metadataCid: {}\n  expireAfter: {} blocks \n",
126            wasm_cid,
127            secrets_schema_cid,
128            policy_data_metadata_cid,
129            expire_after_blocks
130        );
131
132        let (_receipt, policy_data_address) = controller
133            .deploy_policy_data(
134                wasm_cid.to_string(),
135                secrets_schema_cid.to_string(),
136                expire_after_blocks,
137                policy_data_metadata_cid.to_string(),
138            )
139            .await
140            .map_err(|e| eyre::eyre!("Failed to deploy policy data: {}", e))?;
141
142        tracing::info!("Policy data deployed successfully at address: {}", policy_data_address);
143
144        Ok(policy_data_address)
145    }
146
147    /// Set attestation info on deployed policy data
148    ///
149    /// This function sets the attestation info on a deployed policy data contract.
150    /// It creates an AttestationInfo struct with:
151    /// - attestationType: ECDSA (0)
152    /// - verifier: address(0)
153    /// - verificationKey: bytes32(0)
154    async fn set_attestation_info(
155        private_key: &str,
156        rpc_url: &str,
157        policy_data_address: Address,
158        attester_address: Address,
159    ) -> Result<()> {
160        let controller = PolicyDataController::new(private_key.to_string(), rpc_url.to_string());
161
162        // Get task generator addresses based on chain ID
163        let task_generator_addresses = Self::get_task_generator_addresses(rpc_url).await?;
164
165        // Create attesters array: [attester, ...task_generator_addresses]
166        let mut attesters = vec![attester_address];
167        attesters.extend(task_generator_addresses);
168
169        // Create AttestationInfo struct
170        let attestation_info = INewtonPolicyData::AttestationInfo {
171            attesters,
172            attestationType: 0u8, // ECDSA
173            verifier: Address::ZERO,
174            verificationKey: FixedBytes::ZERO,
175        };
176
177        tracing::info!(
178            "Setting attestation info on policy data {} with attesters: {:?}",
179            policy_data_address,
180            attestation_info.attesters
181        );
182
183        controller
184            .set_attestation_info(policy_data_address, attestation_info)
185            .await
186            .map_err(|e| eyre::eyre!("Failed to set attestation info: {}", e))?;
187
188        tracing::info!(
189            "Attestation info set successfully on policy data: {}",
190            policy_data_address
191        );
192
193        Ok(())
194    }
195
196    /// Execute the deploy command
197    pub async fn execute(self: Box<Self>, _config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
198        // Get values from args or env, with error if still missing
199        let private_key = self
200            .private_key
201            .ok_or_else(|| eyre::eyre!("private_key is required (use --private-key or PRIVATE_KEY env var)"))?;
202
203        let rpc_url = self
204            .rpc_url
205            .ok_or_else(|| eyre::eyre!("rpc_url is required (use --rpc-url or RPC_URL env var)"))?;
206
207        // Read from policy_cids.json file
208        let json_content = std::fs::read_to_string(&self.policy_cids)
209            .with_context(|| format!("Failed to read policy_cids.json: {:?}", self.policy_cids))?;
210        let json: Value = serde_json::from_str(&json_content)
211            .with_context(|| format!("Failed to parse policy_cids.json: {:?}", self.policy_cids))?;
212
213        // Extract values from policy_cids.json
214        let wasm_cid = json.get("wasmCid").and_then(|v| v.as_str()).unwrap_or("");
215        let secrets_schema_cid = json.get("secretsSchemaCid").and_then(|v| v.as_str()).unwrap_or("");
216        let policy_data_metadata_cid = json.get("policyDataMetadataCid").and_then(|v| v.as_str()).unwrap_or("");
217        let attester = json.get("attester").and_then(|v| v.as_str()).unwrap_or("");
218        let attester_address: Address = attester.parse().unwrap_or(Address::ZERO);
219
220        // Use --expire-after-blocks if provided, otherwise convert 300 seconds to blocks
221        let expire_after_blocks = match self.expire_after_blocks {
222            Some(blocks) => {
223                tracing::info!("Using expire_after_blocks from command: {} blocks", blocks);
224                blocks
225            }
226            None => Self::seconds_to_blocks(&rpc_url, 300).await?,
227        };
228
229        // Deploy policy data
230        let policy_data_address = Self::deploy_policy_data(
231            &private_key,
232            &rpc_url,
233            wasm_cid,
234            secrets_schema_cid,
235            policy_data_metadata_cid,
236            expire_after_blocks,
237        )
238        .await?;
239
240        // Allow time for the deploy tx to be reflected in RPC state (nonce update).
241        // This prevents "replacement transaction underpriced" when setting attestation info.
242        tracing::info!("Waiting for RPC to reflect updated nonce before setting attestation info...");
243        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
244
245        tracing::info!("Policy data deployment completed! Setting attestation info...");
246        // Set attestation info on policy data with retries for nonce/underpriced errors
247        const MAX_RETRIES: u32 = 3;
248        let mut attempt = 0;
249        loop {
250            match Self::set_attestation_info(&private_key, &rpc_url, policy_data_address, attester_address).await {
251                Ok(()) => break,
252                Err(e) => {
253                    let error_str = e.to_string();
254                    let is_nonce_error = error_str.contains("underpriced") || error_str.contains("nonce");
255                    if is_nonce_error && attempt < MAX_RETRIES {
256                        attempt += 1;
257                        warn!(
258                            "Set attestation info nonce issue (attempt {}/{}), retrying in 2s: {}",
259                            attempt, MAX_RETRIES, e
260                        );
261                        tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
262                    } else {
263                        return Err(e);
264                    }
265                }
266            }
267        }
268
269        tracing::info!("Successfully deployed policy data and set attestation");
270        tracing::info!("Policy Data Address: {}", policy_data_address);
271
272        Ok(())
273    }
274}
275
276impl SimulateCommand {
277    /// Execute the simulate command
278    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
279        info!("Executing WASM simulation...");
280        let output = utils::execute_wasm(self.wasm_file, self.input_json, config).await?;
281        info!("WASM simulation output: {}", output);
282        Ok(())
283    }
284}
285
286impl PolicyDataCommand {
287    /// Execute the policy-data command
288    pub async fn execute(self: Box<Self>, config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
289        match self.subcommand {
290            PolicyDataSubcommand::Deploy(command) => {
291                Box::new(command).execute(config).await?;
292            }
293            PolicyDataSubcommand::Simulate(command) => {
294                Box::new(command).execute(config).await?;
295            }
296        }
297        Ok(())
298    }
299}