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#[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(DeployCommand),
26 Simulate(SimulateCommand),
28}
29
30#[derive(Debug, Parser)]
32pub struct SimulateCommand {
33 #[arg(long)]
35 wasm_file: PathBuf,
36
37 #[arg(long)]
39 input_json: String,
40}
41
42#[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 #[arg(long)]
56 expire_after_blocks: Option<u32>,
57}
58
59impl DeployCommand {
60 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 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 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 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 let task_generator_addresses = Self::get_task_generator_addresses(rpc_url).await?;
164
165 let mut attesters = vec![attester_address];
167 attesters.extend(task_generator_addresses);
168
169 let attestation_info = INewtonPolicyData::AttestationInfo {
171 attesters,
172 attestationType: 0u8, 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 pub async fn execute(self: Box<Self>, _config: NewtonAvsConfig<NewtonCliConfig>) -> eyre::Result<()> {
198 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 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 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 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 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 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 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 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 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}