feat: add API key support to prover client (#1912)

sergeytimoshin · web-flow · commit 4174716f9103 · 2025-08-18T16:44:13.000+01:00
* feat: add API key support to prover client

* feat: implement construct_prover_url utility for dynamic prover URL generation

* feat: add prover URL configurations for address-append, append and update operations

* feat: add prover URL fields for append and update operations in tests
diff --git a/forester-utils/src/instructions/address_batch_update.rs b/forester-utils/src/instructions/address_batch_update.rs
@@ -29,6 +29,7 @@ pub struct AddressUpdateConfig<R: Rpc> {
     pub rpc_pool: Arc<SolanaRpcPool<R>>,
     pub merkle_tree_pubkey: Pubkey,
     pub prover_url: String,
+    pub prover_api_key: Option<String>,
     pub polling_interval: Duration,
     pub max_wait_time: Duration,
     pub ixs_per_tx: usize,
@@ -39,6 +40,7 @@ async fn stream_instruction_data<'a, R: Rpc>(
     rpc_pool: Arc<SolanaRpcPool<R>>,
     merkle_tree_pubkey: Pubkey,
     prover_url: String,
+    prover_api_key: Option<String>,
     polling_interval: Duration,
     max_wait_time: Duration,
     leaves_hash_chains: Vec<[u8; 32]>,
@@ -49,7 +51,7 @@ async fn stream_instruction_data<'a, R: Rpc>(
 ) -> impl Stream<Item = Result<Vec<InstructionDataAddressAppendInputs>, ForesterUtilsError>> + Send + 'a
 {
     stream! {
-        let proof_client = Arc::new(ProofClient::with_config(prover_url, polling_interval, max_wait_time));
+        let proof_client = Arc::new(ProofClient::with_config(prover_url, polling_interval, max_wait_time, prover_api_key));
         let max_zkp_batches_per_call = calculate_max_zkp_batches_per_call(zkp_batch_size);
         let total_chunks = leaves_hash_chains.len().div_ceil(max_zkp_batches_per_call);
 
@@ -324,6 +326,7 @@ pub async fn get_address_update_instruction_stream<'a, R: Rpc>(
         config.rpc_pool,
         config.merkle_tree_pubkey,
         config.prover_url,
+        config.prover_api_key,
         config.polling_interval,
         config.max_wait_time,
         leaves_hash_chains,
diff --git a/forester-utils/src/instructions/state_batch_append.rs b/forester-utils/src/instructions/state_batch_append.rs
@@ -48,6 +48,7 @@ pub async fn get_append_instruction_stream<'a, R: Rpc>(
     rpc_pool: Arc<SolanaRpcPool<R>>,
     merkle_tree_pubkey: Pubkey,
     prover_url: String,
+    prover_api_key: Option<String>,
     polling_interval: Duration,
     max_wait_time: Duration,
     merkle_tree_data: ParsedMerkleTreeData,
@@ -126,7 +127,7 @@ pub async fn get_append_instruction_stream<'a, R: Rpc>(
         }
 
         let mut all_changelogs: Vec<ChangelogEntry<{ DEFAULT_BATCH_STATE_TREE_HEIGHT as usize }>> = Vec::new();
-        let proof_client = Arc::new(ProofClient::with_config(prover_url.clone(), polling_interval, max_wait_time));
+        let proof_client = Arc::new(ProofClient::with_config(prover_url.clone(), polling_interval, max_wait_time, prover_api_key.clone()));
         let mut futures_ordered = FuturesOrdered::new();
         let mut pending_count = 0;
 
diff --git a/forester-utils/src/instructions/state_batch_nullify.rs b/forester-utils/src/instructions/state_batch_nullify.rs
@@ -47,6 +47,7 @@ pub async fn get_nullify_instruction_stream<'a, R: Rpc>(
     rpc_pool: Arc<SolanaRpcPool<R>>,
     merkle_tree_pubkey: Pubkey,
     prover_url: String,
+    prover_api_key: Option<String>,
     polling_interval: Duration,
     max_wait_time: Duration,
     merkle_tree_data: ParsedMerkleTreeData,
@@ -125,7 +126,7 @@ pub async fn get_nullify_instruction_stream<'a, R: Rpc>(
         }
 
         let mut all_changelogs = Vec::new();
-        let proof_client = Arc::new(ProofClient::with_config(prover_url.clone(), polling_interval, max_wait_time));
+        let proof_client = Arc::new(ProofClient::with_config(prover_url.clone(), polling_interval, max_wait_time, prover_api_key.clone()));
         let mut futures_ordered = FuturesOrdered::new();
         let mut pending_count = 0;
 
diff --git a/forester/src/cli.rs b/forester/src/cli.rs
@@ -35,6 +35,30 @@ pub struct StartArgs {
     #[arg(long, env = "FORESTER_PROVER_URL")]
     pub prover_url: Option<String>,
 
+    #[arg(
+        long,
+        env = "FORESTER_PROVER_APPEND_URL",
+        help = "Prover URL for append operations. If not specified, uses prover_url"
+    )]
+    pub prover_append_url: Option<String>,
+
+    #[arg(
+        long,
+        env = "FORESTER_PROVER_UPDATE_URL",
+        help = "Prover URL for update operations. If not specified, uses prover_url"
+    )]
+    pub prover_update_url: Option<String>,
+
+    #[arg(
+        long,
+        env = "FORESTER_PROVER_ADDRESS_APPEND_URL",
+        help = "Prover URL for address-append operations. If not specified, uses prover_url"
+    )]
+    pub prover_address_append_url: Option<String>,
+
+    #[arg(long, env = "FORESTER_PROVER_API_KEY")]
+    pub prover_api_key: Option<String>,
+
     #[arg(long, env = "FORESTER_PAYER")]
     pub payer: Option<String>,
 
diff --git a/forester/src/config.rs b/forester/src/config.rs
@@ -35,6 +35,10 @@ pub struct ExternalServicesConfig {
     pub ws_rpc_url: Option<String>,
     pub indexer_url: Option<String>,
     pub prover_url: Option<String>,
+    pub prover_append_url: Option<String>,
+    pub prover_update_url: Option<String>,
+    pub prover_address_append_url: Option<String>,
+    pub prover_api_key: Option<String>,
     pub photon_api_key: Option<String>,
     pub pushgateway_url: Option<String>,
     pub pagerduty_routing_key: Option<String>,
@@ -210,6 +214,19 @@ impl ForesterConfig {
                 ws_rpc_url: args.ws_rpc_url.clone(),
                 indexer_url: args.indexer_url.clone(),
                 prover_url: args.prover_url.clone(),
+                prover_append_url: args
+                    .prover_append_url
+                    .clone()
+                    .or_else(|| args.prover_url.clone()),
+                prover_update_url: args
+                    .prover_update_url
+                    .clone()
+                    .or_else(|| args.prover_url.clone()),
+                prover_address_append_url: args
+                    .prover_address_append_url
+                    .clone()
+                    .or_else(|| args.prover_url.clone()),
+                prover_api_key: args.prover_api_key.clone(),
                 photon_api_key: args.photon_api_key.clone(),
                 pushgateway_url: args.push_gateway_url.clone(),
                 pagerduty_routing_key: args.pagerduty_routing_key.clone(),
@@ -280,6 +297,10 @@ impl ForesterConfig {
                 ws_rpc_url: None,
                 indexer_url: None,
                 prover_url: None,
+                prover_append_url: None,
+                prover_update_url: None,
+                prover_address_append_url: None,
+                prover_api_key: None,
                 photon_api_key: None,
                 pushgateway_url: args.push_gateway_url.clone(),
                 pagerduty_routing_key: args.pagerduty_routing_key.clone(),
diff --git a/forester/src/epoch_manager.rs b/forester/src/epoch_manager.rs
@@ -1185,6 +1185,7 @@ impl<R: Rpc> EpochManager<R> {
     }
 
     async fn process_v2(&self, epoch_info: &Epoch, tree_accounts: &TreeAccounts) -> Result<usize> {
+        let default_prover_url = "http://127.0.0.1:3001".to_string();
         let batch_context = BatchContext {
             rpc_pool: self.rpc_pool.clone(),
             authority: self.config.payer_keypair.insecure_clone(),
@@ -1193,12 +1194,25 @@ impl<R: Rpc> EpochManager<R> {
             merkle_tree: tree_accounts.merkle_tree,
             output_queue: tree_accounts.queue,
             ixs_per_tx: self.config.transaction_config.batch_ixs_per_tx,
-            prover_url: self
+            prover_append_url: self
                 .config
                 .external_services
-                .prover_url
+                .prover_append_url
                 .clone()
-                .unwrap_or_else(|| "http://127.0.0.1:3001".to_string()),
+                .unwrap_or_else(|| default_prover_url.clone()),
+            prover_update_url: self
+                .config
+                .external_services
+                .prover_update_url
+                .clone()
+                .unwrap_or_else(|| default_prover_url.clone()),
+            prover_address_append_url: self
+                .config
+                .external_services
+                .prover_address_append_url
+                .clone()
+                .unwrap_or_else(|| default_prover_url.clone()),
+            prover_api_key: self.config.external_services.prover_api_key.clone(),
             prover_polling_interval: Duration::from_secs(1),
             prover_max_wait_time: Duration::from_secs(120),
             ops_cache: self.ops_cache.clone(),
diff --git a/forester/src/processor/v2/address.rs b/forester/src/processor/v2/address.rs
@@ -27,7 +27,8 @@ where
     let config = AddressUpdateConfig {
         rpc_pool: ctx.rpc_pool.clone(),
         merkle_tree_pubkey: ctx.merkle_tree,
-        prover_url: ctx.prover_url.clone(),
+        prover_url: ctx.prover_address_append_url.clone(),
+        prover_api_key: ctx.prover_api_key.clone(),
         polling_interval: ctx.prover_polling_interval,
         max_wait_time: ctx.prover_max_wait_time,
         ixs_per_tx: ctx.ixs_per_tx,
diff --git a/forester/src/processor/v2/common.rs b/forester/src/processor/v2/common.rs
@@ -40,7 +40,10 @@ pub struct BatchContext<R: Rpc> {
     pub merkle_tree: Pubkey,
     pub output_queue: Pubkey,
     pub ixs_per_tx: usize,
-    pub prover_url: String,
+    pub prover_append_url: String,
+    pub prover_update_url: String,
+    pub prover_address_append_url: String,
+    pub prover_api_key: Option<String>,
     pub prover_polling_interval: Duration,
     pub prover_max_wait_time: Duration,
     pub ops_cache: Arc<Mutex<ProcessedHashCache>>,
diff --git a/forester/src/processor/v2/state.rs b/forester/src/processor/v2/state.rs
@@ -32,7 +32,8 @@ where
     let (stream, size) = get_nullify_instruction_stream(
         ctx.rpc_pool.clone(),
         ctx.merkle_tree,
-        ctx.prover_url.clone(),
+        ctx.prover_update_url.clone(),
+        ctx.prover_api_key.clone(),
         ctx.prover_polling_interval,
         ctx.prover_max_wait_time,
         merkle_tree_data,
@@ -58,7 +59,8 @@ where
     let (stream, size) = get_append_instruction_stream(
         ctx.rpc_pool.clone(),
         ctx.merkle_tree,
-        ctx.prover_url.clone(),
+        ctx.prover_append_url.clone(),
+        ctx.prover_api_key.clone(),
         ctx.prover_polling_interval,
         ctx.prover_max_wait_time,
         merkle_tree_data,
diff --git a/forester/tests/e2e_test.rs b/forester/tests/e2e_test.rs
@@ -114,13 +114,20 @@ fn get_prover_url() -> String {
     }
 }
 
-fn get_api_key() -> Option<String> {
+fn get_photon_api_key() -> Option<String> {
     match TestMode::from_env() {
         TestMode::Local => None,
         TestMode::Devnet => Some(get_env_var("PHOTON_API_KEY")),
     }
 }
 
+fn get_prover_api_key() -> Option<String> {
+    match TestMode::from_env() {
+        TestMode::Local => None,
+        TestMode::Devnet => Some(get_env_var("PROVER_API_KEY")),
+    }
+}
+
 fn get_forester_keypair() -> Keypair {
     match TestMode::from_env() {
         TestMode::Local => Keypair::new(),
@@ -189,7 +196,11 @@ async fn e2e_test() {
             ws_rpc_url: Some(get_ws_rpc_url()),
             indexer_url: Some(get_indexer_url()),
             prover_url: Some(get_prover_url()),
-            photon_api_key: get_api_key(),
+            prover_append_url: None,
+            prover_update_url: None,
+            prover_address_append_url: None,
+            prover_api_key: get_prover_api_key(),
+            photon_api_key: get_photon_api_key(),
             pushgateway_url: None,
             pagerduty_routing_key: None,
             rpc_rate_limit: None,
@@ -453,7 +464,7 @@ async fn setup_rpc_connection(forester: &Keypair) -> LightClient {
     let mut rpc = LightClient::new(if TestMode::from_env() == TestMode::Local {
         LightClientConfig::local()
     } else {
-        LightClientConfig::new(get_rpc_url(), Some(get_indexer_url()), get_api_key())
+        LightClientConfig::new(get_rpc_url(), Some(get_indexer_url()), get_photon_api_key())
     })
     .await
     .unwrap();
diff --git a/forester/tests/priority_fee_test.rs b/forester/tests/priority_fee_test.rs
@@ -38,6 +38,10 @@ async fn test_priority_fee_request() {
             std::env::var("FORESTER_PROVER_URL")
                 .expect("FORESTER_PROVER_URL must be set in environment"),
         ),
+        prover_append_url: None,
+        prover_update_url: None,
+        prover_address_append_url: None,
+        prover_api_key: None,
         payer: Some(
             std::env::var("FORESTER_PAYER").expect("FORESTER_PAYER must be set in environment"),
         ),
diff --git a/forester/tests/test_utils.rs b/forester/tests/test_utils.rs
@@ -84,6 +84,10 @@ pub fn forester_config() -> ForesterConfig {
             ws_rpc_url: Some("ws://localhost:8900".to_string()),
             indexer_url: Some("http://localhost:8784".to_string()),
             prover_url: Some("http://localhost:3001".to_string()),
+            prover_append_url: None,
+            prover_update_url: None,
+            prover_address_append_url: None,
+            prover_api_key: None,
             photon_api_key: None,
             pushgateway_url: None,
             pagerduty_routing_key: None,
diff --git a/prover/client/src/proof_client.rs b/prover/client/src/proof_client.rs
@@ -51,6 +51,7 @@ pub struct ProofClient {
     server_address: String,
     polling_interval: Duration,
     max_wait_time: Duration,
+    api_key: Option<String>,
 }
 
 impl ProofClient {
@@ -60,6 +61,7 @@ impl ProofClient {
             server_address: DEFAULT_LOCAL_SERVER.to_string(),
             polling_interval: Duration::from_secs(DEFAULT_POLLING_INTERVAL_SECS),
             max_wait_time: Duration::from_secs(DEFAULT_MAX_WAIT_TIME_SECS),
+            api_key: None,
         }
     }
 
@@ -68,12 +70,14 @@ impl ProofClient {
         server_address: String,
         polling_interval: Duration,
         max_wait_time: Duration,
+        api_key: Option<String>,
     ) -> Self {
         Self {
             client: Client::new(),
             server_address,
             polling_interval,
             max_wait_time,
+            api_key,
         }
     }
 
@@ -157,9 +161,16 @@ impl ProofClient {
     ) -> Result<reqwest::Response, ProverClientError> {
         let url = format!("{}{}", self.server_address, PROVE_PATH);
 
-        self.client
+        let mut request = self
+            .client
             .post(&url)
-            .header("Content-Type", "application/json")
+            .header("Content-Type", "application/json");
+
+        if let Some(api_key) = &self.api_key {
+            request = request.header("X-API-Key", api_key);
+        }
+
+        request
             .body(inputs_json.to_string())
             .send()
             .await
@@ -369,7 +380,13 @@ impl ProofClient {
         job_id: &str,
         poll_count: u32,
     ) -> Result<JobStatusResponse, ProverClientError> {
-        let response = self.client.get(status_url).send().await.map_err(|e| {
+        let mut request = self.client.get(status_url);
+
+        if let Some(api_key) = &self.api_key {
+            request = request.header("X-API-Key", api_key);
+        }
+
+        let response = request.send().await.map_err(|e| {
             error!("Failed to send status request for job {}: {}", job_id, e);
             ProverClientError::ProverServerError(format!("Failed to check job status: {}", e))
         })?;
