Error handling for network services (#294)

Author: moubctez

Cargo.lock

@@ -1,4 +1,9 @@
-use std::{fs, net::IpAddr, path::PathBuf, time::Duration};
+use std::{
+    fs,
+    net::{IpAddr, Ipv4Addr, SocketAddr},
+    path::PathBuf,
+    time::Duration,
+};
 
 use clap::Parser;
 use serde::Deserialize;
@@ -123,6 +128,12 @@ impl Config {
     pub fn stats_period(&self) -> Duration {
         Duration::from_secs(self.stats_period)
     }
+
+    /// Return [`SocketAddr`] for gRPC to listen on.
+    #[must_use]
+    pub(crate) fn grpc_socket(&self) -> SocketAddr {
+        SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), self.grpc_port)
+    }
 }
 
 impl Default for Config {

src/config.rs

@@ -9,18 +9,12 @@ pub enum GatewayError {
     #[error("Command {command} execution failed. Error: {error}")]
     CommandExecutionFailed { command: String, error: String },
 
-    #[error("WireGuard key error")]
-    KeyDecode(#[from] base64::DecodeError),
-
     #[error("Logger error")]
     Logger(#[from] log::SetLoggerError),
 
     #[error("Syslog error")]
     Syslog(#[from] syslog::Error),
 
-    #[error("Token parsing error")]
-    Token(#[from] tonic::metadata::errors::InvalidMetadataValue),
-
     #[error("Tonic error")]
     Tonic(#[from] tonic::transport::Error),
 
@@ -33,9 +27,6 @@ pub enum GatewayError {
     #[error("WireGuard error {0}")]
     WireguardError(#[from] WireguardInterfaceError),
 
-    #[error("HTTP error")]
-    HttpServer(String),
-
     #[error("Invalid CA file. Error")]
     InvalidCaFile,
 

src/error.rs

@@ -1,6 +1,5 @@
 use std::{
     collections::HashMap,
-    net::{IpAddr, Ipv4Addr, SocketAddr},
     path::PathBuf,
     str::FromStr,
     sync::{
@@ -48,7 +47,6 @@ use crate::{
 /// restarted in setup mode when a purge request arrives.
 pub async fn run_gateway_loop(
     config: Config,
-    cert_dir: std::path::PathBuf,
     gateway: Arc<Mutex<Gateway>>,
     logs_rx: Arc<tokio::sync::Mutex<mpsc::Receiver<LogEntry>>>,
     mut tls_config: TlsConfig,
@@ -59,7 +57,7 @@ pub async fn run_gateway_loop(
         let (reset_tx, mut reset_rx) = oneshot::channel();
         // Build and start a gRPC server instance wired with the reset signal.
         let mut gateway_server =
-            GatewayServer::new(Arc::clone(&gateway), cert_dir.clone(), reset_tx);
+            GatewayServer::new(Arc::clone(&gateway), config.cert_dir.clone(), reset_tx);
         gateway_server.set_tls_config(tls_config.clone());
         let mut server_handle = tokio::spawn(gateway_server.start(config.clone()));
 
@@ -87,7 +85,7 @@ pub async fn run_gateway_loop(
 
                 // Run setup server to obtain new TLS certs, then loop to restart gRPC.
                 log::info!("Restarting setup server after purge request");
-                tls_config = run_setup(&config, &cert_dir, Arc::clone(&logs_rx)).await?;
+                tls_config = run_setup(&config, Arc::clone(&logs_rx)).await?;
             }
             // Server exited on its own (error or normal shutdown).
             result = &mut server_handle => {
@@ -609,7 +607,7 @@ impl GatewayServer {
             .map(|c| c.grpc_key_pem.clone());
 
         // Build gRPC server.
-        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), config.grpc_port);
+        let addr = config.grpc_socket();
         info!("gRPC server is listening on {addr}");
         let mut builder = if let (Some(cert), Some(key)) = (grpc_cert, grpc_key) {
             let identity = Identity::from_pem(cert, key);
@@ -623,10 +621,6 @@ impl GatewayServer {
         builder
             .add_service(
                 ServiceBuilder::new()
-                    // .layer(InterceptorLayer::new(CoreVersionInterceptor::new(
-                    //     MIN_CORE_VERSION,
-                    //     incompatible_components,
-                    // )))
                     .layer(DefguardVersionLayer::new(Version::parse(VERSION)?))
                     .service(gateway_server::GatewayServer::new(self)),
             )
@@ -636,7 +630,7 @@ impl GatewayServer {
         Ok(())
     }
 
-    pub fn set_tls_config(&mut self, tls_config: TlsConfig) {
+    pub(crate) fn set_tls_config(&mut self, tls_config: TlsConfig) {
         if let Ok(mut gateway) = self.gateway.lock() {
             gateway.tls_config = Some(tls_config);
         }

src/gateway.rs

@@ -15,7 +15,7 @@ use defguard_gateway::{
     gateway::{Gateway, TlsConfig, run_gateway_loop, run_stats},
     init_syslog,
     logging::init_tracing,
-    server::run_server,
+    server::run_http_server,
     setup::run_setup,
 };
 use defguard_version::Version;
@@ -93,7 +93,7 @@ async fn main() -> Result<(), GatewayError> {
 
     // Optionally, launch HTTP server to report gateway's health.
     if let Some(health_port) = config.health_port {
-        tasks.spawn(run_server(
+        tasks.spawn(run_http_server(
             health_port,
             config.http_bind_address,
             Arc::clone(&gateway.connected),
@@ -104,41 +104,40 @@ async fn main() -> Result<(), GatewayError> {
     let gateway = Arc::new(Mutex::new(gateway));
     tasks.spawn(run_stats(Arc::clone(&gateway), config.stats_period()));
 
-    let tls_config = if needs_setup {
-        log::info!(
-            "gRPC TLS certificates not found in {}. They will be generated during setup.",
-            cert_dir.display()
-        );
-        run_setup(&config, cert_dir, Arc::clone(&logs_rx)).await?
-    } else if let (Some(cert), Some(key)) = (grpc_cert, grpc_key) {
-        log::info!(
-            "Using existing gRPC TLS certificates from {}",
-            cert_dir.display()
-        );
-        TlsConfig {
-            grpc_cert_pem: cert,
-            grpc_key_pem: key,
-        }
-    } else {
-        return Err(GatewayError::SetupError(
-            "gRPC TLS certificates are missing after setup".to_string(),
-        ));
-    };
-
-    // Launch gRPC server (with purge-triggered setup loop).
-    tasks.spawn(run_gateway_loop(
-        config.clone(),
-        cert_dir.clone(),
-        gateway,
-        Arc::clone(&logs_rx),
-        tls_config,
-    ));
+    // Clone for later.
+    let post_down_clone = config.post_down.clone();
+
+    tasks.spawn(async move {
+        let tls_config = if needs_setup {
+            log::info!(
+                "gRPC TLS certificates not found in {}. They will be generated during setup.",
+                config.cert_dir.display()
+            );
+            run_setup(&config, Arc::clone(&logs_rx)).await?
+        } else if let (Some(cert), Some(key)) = (grpc_cert, grpc_key) {
+            log::info!(
+                "Using existing gRPC TLS certificates from {}",
+                config.cert_dir.display()
+            );
+            TlsConfig {
+                grpc_cert_pem: cert,
+                grpc_key_pem: key,
+            }
+        } else {
+            return Err(GatewayError::SetupError(
+                "gRPC TLS certificates are missing after setup".to_string(),
+            ));
+        };
+
+        // Launch gRPC server (with purge-triggered setup loop).
+        run_gateway_loop(config, gateway, Arc::clone(&logs_rx), tls_config).await
+    });
 
     while let Some(Ok(result)) = tasks.join_next().await {
         result?;
     }
 
-    if let Some(post_down) = &config.post_down {
+    if let Some(post_down) = &post_down_clone {
         log::info!("Executing specified POST_DOWN command: {post_down}");
         execute_command(post_down)?;
     }

src/main.rs

@@ -21,7 +21,8 @@ async fn healthcheck<'a>(
     }
 }
 
-pub async fn run_server(
+/// Launch HTTP server with the health check.
+pub async fn run_http_server(
     http_port: u16,
     http_bind_address: Option<IpAddr>,
     connected: Arc<AtomicBool>,
@@ -35,9 +36,13 @@ pub async fn run_server(
         http_bind_address.unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED)),
         http_port,
     );
-    let listener = TcpListener::bind(&addr).await?;
+    let listener = TcpListener::bind(&addr).await.inspect_err(|err| {
+        error!("Failed to bind to {addr}: {err}");
+    })?;
     info!("Health check listening on {addr}");
-    serve(listener, app.into_make_service())
-        .await
-        .map_err(|err| GatewayError::HttpServer(err.to_string()))
+
+    // From axum docs: this future will never actually complete or return an error.
+    let _ = serve(listener, app.into_make_service()).await;
+
+    Ok(())
 }

src/server.rs

@@ -1,8 +1,4 @@
-use std::{
-    net::{IpAddr, Ipv4Addr, SocketAddr},
-    path::Path,
-    sync::{Arc, Mutex},
-};
+use std::sync::{Arc, Mutex};
 
 use defguard_version::{Version, server::DefguardVersionLayer};
 use tokio::{
@@ -28,9 +24,9 @@ type LogsReceiver = Arc<tokio::sync::Mutex<mpsc::Receiver<LogEntry>>>;
 
 pub async fn run_setup(
     config: &Config,
-    cert_dir: &Path,
     logs_rx: Arc<tokio::sync::Mutex<mpsc::Receiver<LogEntry>>>,
 ) -> Result<TlsConfig, GatewayError> {
+    let cert_dir = &config.cert_dir;
     let setup_server = GatewaySetupServer::new(logs_rx);
     let tls_config = setup_server.await_setup(config).await?;
 
@@ -100,8 +96,7 @@ impl GatewaySetupServer {
         let mut server_config = None;
         let setup_rx = Arc::clone(&self.setup_rx);
 
-        let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), config.grpc_port);
-
+        let addr = config.grpc_socket();
         info!("Starting Gateway setup server on {addr} and awaiting configuration from Core");
 
         server_builder
@@ -338,11 +333,10 @@ impl gateway_setup_server::GatewaySetup for GatewaySetupServer {
             match defguard_certs::der_to_pem(&cert_der, defguard_certs::PemLabel::Certificate) {
                 Ok(pem) => pem,
                 Err(err) => {
-                    error!("Failed to convert certificate DER to PEM: {err}");
+                    let msg = format!("Failed to convert certificate DER to PEM: {err}");
+                    error!("{msg}");
                     self.clear_setup_session();
-                    return Err(Status::internal(format!(
-                        "Failed to convert certificate DER to PEM: {err}"
-                    )));
+                    return Err(Status::internal(msg));
                 }
             };
         debug!("Certificate processed successfully");
@@ -359,15 +353,11 @@ impl gateway_setup_server::GatewaySetup for GatewaySetupServer {
             if let Some(kp) = key_pair {
                 kp
             } else {
-                error!(
-                    "Key pair not found during Gateway setup. Key pair generation step might have \
-                    failed."
-                );
+                let msg = "Key pair not found during Gateway setup. Key pair generation step might \
+                    have failed.";
+                error!("{msg}");
                 self.clear_setup_session();
-                return Err(Status::internal(
-                    "Key pair not found during Gateway setup. Key pair generation step might have \
-                    failed.",
-                ));
+                return Err(Status::internal(msg));
             }
         };
 
@@ -382,8 +372,9 @@ impl gateway_setup_server::GatewaySetup for GatewaySetupServer {
         };
 
         sender.send(configuration).map_err(|_| {
-            error!("Failed to send setup configuration through channel");
-            Status::internal("Failed to send setup configuration through channel")
+            let msg = "Failed to send setup configuration through channel";
+            error!("{msg}");
+            Status::internal(msg)
         })?;
 
         debug!("Setup process completed successfully, cleaning up temporary session");

src/setup.rs

@@ -1,6 +1,6 @@
 use defguard_version::{Version, is_version_lower};
 
-const MIN_CORE_VERSION: Version = Version::new(1, 6, 0);
+const MIN_CORE_VERSION: Version = Version::new(2, 0, 0);
 
 /// Checks if Defguard Core's version meets minimum version requirements.
 pub(crate) fn is_core_version_supported(core_version: Option<&Version>) -> bool {