From a732edd295af60d495cfe7f8b74711f7b581e422 Mon Sep 17 00:00:00 2001
From: Zhidong Peng <zpeng@micrsoft.com>
Date: Tue, 8 Sep 2026 20:38:23 +0000
Subject: [PATCH] skip redirector policy udpate in update_redirector_policy if
 proxy sever has not started yet.

---
 proxy_agent/src/key_keeper.rs |  9 ++++++++-
 proxy_agent/src/provision.rs  | 17 +++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/proxy_agent/src/key_keeper.rs b/proxy_agent/src/key_keeper.rs
index a42f5bce..3c94353a 100644
--- a/proxy_agent/src/key_keeper.rs
+++ b/proxy_agent/src/key_keeper.rs
@@ -291,7 +291,7 @@ impl KeyKeeper {
             }
 
             // check and update the redirect policy if not updated successfully before, try again here
-            // this could happen when the eBPF/redirector module was not started yet before
+            // this could happen when the eBPF/redirector or proxy_server module was not started yet before
             if !redirect_policy_updated {
                 logger::write_warning(
                     "redirect policy was not update successfully before, retrying now".to_string(),
@@ -850,6 +850,13 @@ impl KeyKeeper {
     /// update the redirector/eBPF policy based on the secure channel status
     /// it should be called when the secure channel state is changed
     async fn update_redirector_policy(&self, status: &KeyStatus) -> bool {
+        if !provision::is_proxy_server_provisioned(&self.provision_shared_state).await {
+            logger::write_warning(
+                "Proxy server is not provisioned, skipping redirector policy update.".to_string(),
+            );
+            return false;
+        }
+
         // update the redirector policy map
         if !redirector::update_wire_server_redirect_policy(
             status.get_wire_server_mode() != DISABLE_STATE,
diff --git a/proxy_agent/src/provision.rs b/proxy_agent/src/provision.rs
index f8f3ce28..016a4ac4 100644
--- a/proxy_agent/src/provision.rs
+++ b/proxy_agent/src/provision.rs
@@ -114,6 +114,17 @@ pub async fn listener_started(event_threads_shared_state: EventThreadsSharedStat
     .await;
 }
 
+/// Check if the proxy server is provisioned
+/// It checks if the listener module is ready, which indicates that the proxy server is provisioned.
+/// Returns true if the proxy server is provisioned, false otherwise.
+pub async fn is_proxy_server_provisioned(provision_shared_state: &ProvisionSharedState) -> bool {
+    if let Ok(provisioned) = provision_shared_state.get_state().await {
+        provisioned.contains(ProvisionFlags::LISTENER_READY)
+    } else {
+        false
+    }
+}
+
 /// Update provision state for each module to shared_state
 async fn update_provision_state(
     state: ProvisionFlags,
@@ -945,6 +956,12 @@ mod tests {
             event_threads_shared_state.clone(),
         )
         .await;
+        assert_eq!(
+            true,
+            super::is_proxy_server_provisioned(&provision_shared_state).await,
+            "Proxy server must be provisioned after listener is ready"
+        );
+
         super::update_provision_state(
             ProvisionFlags::KEY_LATCH_READY,
             Some(temp_test_path.clone()),