refactor: code review polish — dedup, simplify, consistency

- Move current_time_millis() to agent-tunnel-proto (R1: eliminate duplication)
- Delete DomainInfo, use DomainAdvertisement directly in AgentInfo (R2)
- Merge enroll_agent/bootstrap_and_persist into single function (I1)
- Agent task_handles: Vec<JoinHandle> → JoinSet with reaping (I4)
- Same-epoch route refresh: mutate updated_at in place, no clone (I5)
- Add #[must_use] on enrollment_store::redeem() (I6)
- connect_via_agent: cleaner error extraction with if-let (I3)
- Add TODO for active_stream_count tracking (I2)
- SECS_PER_DAY constant replaces magic 86400 (P4)
- Consistent .context() for ProtoError instead of map_err (P7)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: irvingoujAtDevolution

crates/agent-tunnel-proto/src/lib.rs

@@ -24,3 +24,14 @@ pub use error::ProtoError;
 pub use session::{ConnectRequest, ConnectResponse, MAX_SESSION_MESSAGE_SIZE};
 pub use stream::{ControlRecvStream, ControlSendStream, ControlStream, SessionStream};
 pub use version::{CURRENT_PROTOCOL_VERSION, MIN_SUPPORTED_VERSION, validate_protocol_version};
+
+/// Current wall-clock time in milliseconds since UNIX epoch.
+pub fn current_time_millis() -> u64 {
+    u64::try_from(
+        std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .expect("system time should be after unix epoch")
+            .as_millis(),
+    )
+    .expect("millisecond timestamp should fit in u64")
+}

devolutions-agent/src/enrollment.rs

@@ -56,16 +56,6 @@ pub async fn enroll_agent(
     enrollment_token: &str,
     agent_name: &str,
     advertise_subnets: Vec<String>,
-) -> anyhow::Result<()> {
-    bootstrap_and_persist(gateway_url, enrollment_token, agent_name, advertise_subnets).await?;
-    Ok(())
-}
-
-pub async fn bootstrap_and_persist(
-    gateway_url: &str,
-    enrollment_token: &str,
-    agent_name: &str,
-    advertise_subnets: Vec<String>,
 ) -> anyhow::Result<PersistedEnrollment> {
     // Generate key pair and CSR locally — the private key never leaves this machine.
     let (key_pem, csr_pem) = generate_key_and_csr(agent_name)?;

devolutions-agent/src/main.rs

@@ -301,7 +301,7 @@ fn main() {
 
                 let rt = tokio::runtime::Runtime::new().expect("failed to create tokio runtime");
                 let result = rt.block_on(async {
-                    devolutions_agent::enrollment::bootstrap_and_persist(
+                    devolutions_agent::enrollment::enroll_agent(
                         &command.gateway_url,
                         &command.enrollment_token,
                         &command.agent_name,

devolutions-agent/src/tunnel.rs

@@ -6,15 +6,17 @@
 use std::sync::Arc;
 use std::time::Duration;
 
-use agent_tunnel_proto::{ConnectResponse, ControlMessage, ControlRecvStream, ControlStream, SessionStream};
+use agent_tunnel_proto::{
+    ConnectResponse, ControlMessage, ControlRecvStream, ControlStream, SessionStream, current_time_millis,
+};
 use anyhow::{Context as _, bail};
 use async_trait::async_trait;
 use devolutions_gateway_task::{ShutdownSignal, Task};
 use ipnetwork::Ipv4Network;
 use sha2::Digest as _;
 
 use crate::config::ConfHandle;
-use crate::tunnel_helpers::{Target, connect_to_target, current_time_millis, resolve_target};
+use crate::tunnel_helpers::{Target, connect_to_target, resolve_target};
 
 // ---------------------------------------------------------------------------
 // Custom TLS verifier: chain + hostname validation + SPKI pinning
@@ -349,8 +351,8 @@ async fn run_single_connection(conf_handle: &ConfHandle, shutdown_signal: &mut S
 
     // Split: recv half goes to a reader task, send half stays for periodic messages.
     let (mut ctrl_send, ctrl_recv) = ctrl.into_split();
-    let mut task_handles: Vec<tokio::task::JoinHandle<()>> = Vec::new();
-    task_handles.push(tokio::spawn(run_control_reader(ctrl_recv)));
+    let mut task_handles = tokio::task::JoinSet::new();
+    task_handles.spawn(run_control_reader(ctrl_recv));
 
     // -- Main loop: accept incoming session streams + periodic tasks --
 
@@ -380,6 +382,7 @@ async fn run_single_connection(conf_handle: &ConfHandle, shutdown_signal: &mut S
             }
 
             _ = heartbeat_tick.tick() => {
+                // TODO: track actual active_stream_count instead of hardcoded 0.
                 let msg = ControlMessage::heartbeat(current_time_millis(), 0);
                 let _ = ctrl_send.send(&msg).await
                     .inspect(|_| trace!("Sent Heartbeat"))
@@ -389,18 +392,15 @@ async fn run_single_connection(conf_handle: &ConfHandle, shutdown_signal: &mut S
             result = connection.accept_bi() => {
                 let (send, recv) = result.context("accept incoming bidi stream")?;
                 let subnets = advertise_subnets.clone();
-                task_handles.push(tokio::spawn(run_session_proxy(subnets, send, recv)));
+                task_handles.spawn(run_session_proxy(subnets, send, recv));
             }
+
+            // Reap completed session tasks.
+            Some(_) = task_handles.join_next() => {}
         }
     }
 
-    // Abort all spawned tasks on shutdown.
-    for handle in &task_handles {
-        handle.abort();
-    }
-    for handle in task_handles {
-        let _ = handle.await;
-    }
+    task_handles.shutdown().await;
 
     Ok(())
 }

devolutions-agent/src/tunnel_helpers.rs

@@ -83,12 +83,3 @@ pub(crate) async fn connect_to_target(candidates: &[SocketAddr]) -> anyhow::Resu
 
     Err(error).with_context(|| format!("TCP connect failed for {candidate}"))
 }
-
-/// Current wall-clock time in milliseconds since UNIX epoch.
-pub(crate) fn current_time_millis() -> u64 {
-    let elapsed = std::time::SystemTime::now()
-        .duration_since(std::time::UNIX_EPOCH)
-        .expect("system time should be after unix epoch");
-
-    u64::try_from(elapsed.as_millis()).expect("millisecond timestamp should fit in u64")
-}

devolutions-gateway/src/agent_tunnel/cert.rs

@@ -20,6 +20,7 @@ const CA_VALIDITY_DAYS: u32 = 3650; // ~10 years
 const SERVER_CERT_VALIDITY_DAYS: u32 = 365; // 1 year
 const AGENT_CERT_VALIDITY_DAYS: u32 = 365; // 1 year
 
+const SECS_PER_DAY: u64 = 86_400;
 const CA_COMMON_NAME: &str = "Devolutions Gateway Agent Tunnel CA";
 const CA_ORG_NAME: &str = "Devolutions Inc.";
 
@@ -33,7 +34,8 @@ fn make_ca_params() -> CertificateParams {
     params.key_usages.push(KeyUsagePurpose::KeyCertSign);
     params.key_usages.push(KeyUsagePurpose::CrlSign);
     params.not_before = time::OffsetDateTime::now_utc();
-    params.not_after = time::OffsetDateTime::now_utc() + Duration::from_secs(u64::from(CA_VALIDITY_DAYS) * 86400);
+    params.not_after =
+        time::OffsetDateTime::now_utc() + Duration::from_secs(u64::from(CA_VALIDITY_DAYS) * SECS_PER_DAY);
     params
 }
 
@@ -154,7 +156,7 @@ impl CaManager {
             .push(ExtendedKeyUsagePurpose::ClientAuth);
         agent_params.not_before = time::OffsetDateTime::now_utc();
         agent_params.not_after =
-            time::OffsetDateTime::now_utc() + Duration::from_secs(u64::from(AGENT_CERT_VALIDITY_DAYS) * 86400);
+            time::OffsetDateTime::now_utc() + Duration::from_secs(u64::from(AGENT_CERT_VALIDITY_DAYS) * SECS_PER_DAY);
 
         // Sign with the CA, embedding the public key from the CSR.
         let ca_cert = self.reconstruct_ca_cert()?;
@@ -201,7 +203,7 @@ impl CaManager {
             .push(ExtendedKeyUsagePurpose::ServerAuth);
         server_params.not_before = time::OffsetDateTime::now_utc();
         server_params.not_after =
-            time::OffsetDateTime::now_utc() + Duration::from_secs(u64::from(SERVER_CERT_VALIDITY_DAYS) * 86400);
+            time::OffsetDateTime::now_utc() + Duration::from_secs(u64::from(SERVER_CERT_VALIDITY_DAYS) * SECS_PER_DAY);
 
         let ca_cert = self.reconstruct_ca_cert()?;
 

devolutions-gateway/src/agent_tunnel/enrollment_store.rs

@@ -52,6 +52,7 @@ impl EnrollmentTokenStore {
     ///
     /// Returns `true` if the token was valid and has been redeemed (removed).
     /// Returns `false` if the token doesn't exist or is expired.
+    #[must_use = "check whether the token was valid"]
     pub fn redeem(&self, token: &str) -> bool {
         let now = current_time_secs();
 

devolutions-gateway/src/agent_tunnel/listener.rs

@@ -73,19 +73,15 @@ impl AgentTunnelHandle {
         session
             .send_request(&connect_msg)
             .await
-            .map_err(|e| anyhow::anyhow!("send ConnectRequest: {e}"))?;
+            .context("send ConnectRequest")?;
 
         // Read ConnectResponse (with timeout to prevent stalled peers).
         let response = tokio::time::timeout(Duration::from_secs(30), session.recv_response())
             .await
             .map_err(|_| anyhow::anyhow!("session handshake timeout (30s)"))?
-            .map_err(|e| anyhow::anyhow!("recv ConnectResponse: {e}"))?;
+            .context("recv ConnectResponse")?;
 
-        if !response.is_success() {
-            let reason = match &response {
-                ConnectResponse::Error { reason, .. } => reason.clone(),
-                _ => "unknown".to_owned(),
-            };
+        if let ConnectResponse::Error { reason, .. } = &response {
             anyhow::bail!("agent refused connection: {reason}");
         }
 

devolutions-gateway/src/agent_tunnel/registry.rs

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
-use std::time::{Duration, SystemTime, UNIX_EPOCH};
+use std::time::{Duration, SystemTime};
 
 use agent_tunnel_proto::DomainAdvertisement;
 use dashmap::DashMap;
@@ -124,16 +124,10 @@ impl AgentPeer {
                 agent_id = %self.agent_id,
                 epoch,
                 subnet_count = subnets.len(),
-                domain_count = current.domains.len(),
+                domain_count = state.domains.len(),
                 "Refreshing route advertisement (same epoch)"
             );
-            *state = RouteAdvertisementState {
-                epoch,
-                subnets: current.subnets.clone(),
-                domains: current.domains.clone(),
-                received_at: current.received_at,
-                updated_at: now,
-            };
+            state.updated_at = now;
         } else {
             // New epoch (or first advertisement): replace everything.
             info!(
@@ -234,13 +228,6 @@ impl Default for AgentRegistry {
     }
 }
 
-/// Domain info with source tracking for API responses.
-#[derive(Debug, Clone, Serialize)]
-pub struct DomainInfo {
-    pub domain: String,
-    pub auto_detected: bool,
-}
-
 /// Serializable snapshot of an agent's state, suitable for API responses.
 #[derive(Debug, Clone, Serialize)]
 pub struct AgentInfo {
@@ -250,7 +237,7 @@ pub struct AgentInfo {
     pub is_online: bool,
     pub last_seen_ms: u64,
     pub subnets: Vec<String>,
-    pub domains: Vec<DomainInfo>,
+    pub domains: Vec<DomainAdvertisement>,
     pub route_epoch: u64,
 }
 
@@ -264,29 +251,13 @@ impl From<&Arc<AgentPeer>> for AgentInfo {
             is_online: agent.is_online(AGENT_OFFLINE_TIMEOUT),
             last_seen_ms: agent.last_seen_ms(),
             subnets: route_state.subnets.iter().map(ToString::to_string).collect(),
-            domains: route_state
-                .domains
-                .iter()
-                .map(|d| DomainInfo {
-                    domain: d.domain.clone(),
-                    auto_detected: d.auto_detected,
-                })
-                .collect(),
+            domains: route_state.domains.clone(),
             route_epoch: route_state.epoch,
         }
     }
 }
 
-/// Returns the current time as milliseconds since UNIX epoch.
-fn current_time_millis() -> u64 {
-    u64::try_from(
-        SystemTime::now()
-            .duration_since(UNIX_EPOCH)
-            .unwrap_or(Duration::ZERO)
-            .as_millis(),
-    )
-    .expect("millisecond timestamp should fit in u64")
-}
+use agent_tunnel_proto::current_time_millis;
 
 #[cfg(test)]
 mod tests {