agent: fix decrypted HMAC keys in proxy_connectors and controlplane.rs

Author: mdibaiee

crates/agent-sql/src/data_plane.rs

@@ -56,7 +56,7 @@ pub async fn fetch_data_planes<'a, 'b>(
             ops_logs_name as "ops_logs_name: models::Collection",
             ops_stats_name as "ops_stats_name: models::Collection"
         from data_planes
-        where (array_length($1::flowid[], 1) = 0 OR id in (select id from unnest($1::flowid[]) id))
+        where (array_length($1::flowid[], 1) is null or id in (select id from unnest($1::flowid[]) id))
            or data_plane_name = $2
         "#,
         &data_plane_ids as &[Id],

crates/agent/src/api/snapshot.rs

@@ -540,6 +540,16 @@ async fn try_fetch(
         "fetched authorization snapshot",
     );
 
+    data_planes
+        .iter_mut()
+        .filter(|dp| decrypted_hmac_keys.contains_key(&dp.data_plane_name))
+        .for_each(|dp| {
+            dp.hmac_keys = decrypted_hmac_keys
+                .get(&dp.data_plane_name)
+                .unwrap()
+                .clone()
+        });
+
     futures::future::try_join_all(
         data_planes
             .iter_mut()

crates/agent/src/main.rs

@@ -213,30 +213,7 @@ async fn async_main(args: Args) -> Result<(), anyhow::Error> {
         connectors,
     );
 
-    let mut data_planes =
-        agent_sql::data_plane::fetch_data_planes(&pg_pool, Vec::new(), "", uuid::Uuid::nil())
-            .await?;
-
-    futures::future::try_join_all(
-        data_planes
-            .iter_mut()
-            .filter(|dp| {
-                !dp.encrypted_hmac_keys
-                    .to_value()
-                    .as_object()
-                    .unwrap()
-                    .is_empty()
-            })
-            .map(|dp| agent::decrypt_hmac_keys(dp)),
-    )
-    .await?;
-
-    let decrypted_hmac_keys = Arc::new(RwLock::new(
-        data_planes
-            .iter()
-            .map(|dp| (dp.data_plane_name.clone(), dp.hmac_keys.clone()))
-            .collect(),
-    ));
+    let decrypted_hmac_keys = Arc::new(RwLock::new(HashMap::new()));
 
     tokio::spawn(refresh_decrypted_hmac_keys(
         pg_pool.clone(),
@@ -307,13 +284,10 @@ async fn refresh_decrypted_hmac_keys(
     const REFRESH_INTERVAL: chrono::TimeDelta = chrono::TimeDelta::seconds(60);
 
     loop {
-        let mut data_planes =
+        let mut data_planes: Vec<_> =
             agent_sql::data_plane::fetch_data_planes(&pg_pool, Vec::new(), "", uuid::Uuid::nil())
-                .await?;
-
-        futures::future::try_join_all(
-            data_planes
-                .iter_mut()
+                .await?
+                .into_iter()
                 .filter(|dp| {
                     !decrypted_hmac_keys
                         .read()
@@ -327,6 +301,11 @@ async fn refresh_decrypted_hmac_keys(
                         .unwrap()
                         .is_empty()
                 })
+                .collect();
+
+        futures::future::try_join_all(
+            data_planes
+                .iter_mut()
                 .map(|dp| agent::decrypt_hmac_keys(dp)),
         )
         .await?;

crates/agent/src/proxy_connectors.rs

@@ -215,19 +215,22 @@ impl<L: runtime::LogHandler> ProxyConnectors<L> {
             impl Future<Output = anyhow::Result<()>> + 'a,
         ),
     )> {
+        let mut dp = data_plane.clone();
+        crate::decrypt_hmac_keys(&mut dp).await?;
+
         let tables::DataPlane {
             reactor_address,
             hmac_keys,
             data_plane_fqdn,
             ..
-        } = data_plane;
+        } = dp;
 
         let mut metadata = gazette::Metadata::default();
 
         metadata
             .signed_claims(
                 proto_flow::capability::PROXY_CONNECTOR,
-                data_plane_fqdn,
+                &data_plane_fqdn,
                 *CONNECTOR_TIMEOUT * 2,
                 &hmac_keys,
                 Default::default(),
@@ -241,7 +244,7 @@ impl<L: runtime::LogHandler> ProxyConnectors<L> {
 
         let mut proxy_client =
             proto_grpc::runtime::connector_proxy_client::ConnectorProxyClient::with_interceptor(
-                gazette::dial_channel(reactor_address)?,
+                gazette::dial_channel(&reactor_address)?,
                 metadata.clone(),
             );
         let mut proxy_responses = proxy_client