Core certificate authority, part 1: Proxy

Cargo.toml

@@ -8,6 +8,7 @@ repository = "https://github.com/DefGuard/proxy"
 
 [dependencies]
 defguard_version = { git = "https://github.com/DefGuard/defguard.git", rev = "8649a9ba225d7bd2066a09c9e1347705c34bd158" }
+defguard_certs = { path = "../defguard/crates/defguard_certs" }
 # base `axum` deps
 axum = { version = "0.8", features = ["ws"] }
 axum-client-ip = "0.7"

src/config.rs

@@ -59,6 +59,14 @@ pub struct Config {
 
     #[arg(long, env = "DEFGUARD_GRPC_BIND_ADDRESS")]
     pub grpc_bind_address: Option<IpAddr>,
+
+    // TODO: On different platforms this may be different
+    #[arg(
+        long,
+        env = "DEFGUARD_PROXY_CERT_DIR",
+        default_value = "/etc/defguard/certs"
+    )]
+    pub cert_dir: PathBuf,
 }
 
 #[derive(thiserror::Error, Debug)]

src/grpc.rs

@@ -8,26 +8,44 @@ use std::{
     },
 };
 
-use defguard_version::{get_tracing_variables, ComponentInfo, DefguardComponent, Version};
+use defguard_version::{
+    get_tracing_variables,
+    server::{grpc::DefguardVersionInterceptor, DefguardVersionLayer},
+    ComponentInfo, DefguardComponent, Version,
+};
 use tokio::sync::{mpsc, oneshot};
 use tokio_stream::wrappers::UnboundedReceiverStream;
-use tonic::{Request, Response, Status, Streaming};
+use tonic::{
+    transport::{Identity, Server, ServerTlsConfig},
+    Request, Response, Status, Streaming,
+};
+use tower::ServiceBuilder;
 use tracing::Instrument;
 
 use crate::{
     error::ApiError,
+    http::GRPC_SERVER_RESTART_CHANNEL,
     proto::{core_request, core_response, proxy_server, CoreRequest, CoreResponse, DeviceInfo},
+    MIN_CORE_VERSION, VERSION,
 };
 
 // connected clients
 type ClientMap = HashMap<SocketAddr, mpsc::UnboundedSender<Result<CoreRequest, Status>>>;
 
+#[derive(Debug, Clone, Default)]
+pub(crate) struct Configuration {
+    pub(crate) grpc_key_pem: String,
+    pub(crate) grpc_cert_pem: String,
+}
+
 pub(crate) struct ProxyServer {
     current_id: Arc<AtomicU64>,
     clients: Arc<Mutex<ClientMap>>,
     results: Arc<Mutex<HashMap<u64, oneshot::Sender<core_response::Payload>>>>,
     pub(crate) connected: Arc<AtomicBool>,
     pub(crate) core_version: Arc<Mutex<Option<Version>>>,
+    config: Arc<Mutex<Option<Configuration>>>,
+    setup_in_progress: Arc<AtomicBool>,
 }
 
 impl ProxyServer {
@@ -40,18 +58,94 @@ impl ProxyServer {
             results: Arc::new(Mutex::new(HashMap::new())),
             connected: Arc::new(AtomicBool::new(false)),
             core_version: Arc::new(Mutex::new(None)),
+            config: Arc::new(Mutex::new(None)),
+            setup_in_progress: Arc::new(AtomicBool::new(false)),
         }
     }
 
-    /// Sends message to the other side of RPC, with given `payload` and optional `device_info`.
-    /// Returns `tokio::sync::oneshot::Reveicer` to let the caller await reply.
-    #[instrument(name = "send_grpc_message", level = "debug", skip(self, payload))]
+    pub(crate) fn set_tls_config(&self, cert_pem: String, key_pem: String) -> Result<(), ApiError> {
+        let mut lock = self
+            .config
+            .lock()
+            .expect("Failed to acquire lock on config mutex when updating TLS configuration");
+        let config = lock.get_or_insert_with(Configuration::default);
+        config.grpc_cert_pem = cert_pem;
+        config.grpc_key_pem = key_pem;
+        Ok(())
+    }
+
+    pub(crate) fn configure(&self, config: Configuration) {
+        let mut lock = self
+            .config
+            .lock()
+            .expect("Failed to acquire lock on config mutex when applying proxy configuration");
+        *lock = Some(config);
+    }
+
+    pub(crate) fn get_configuration(&self) -> Option<Configuration> {
+        let lock = self
+            .config
+            .lock()
+            .expect("Failed to acquire lock on config mutex when retrieving proxy configuration");
+        lock.clone()
+    }
+
+    pub(crate) async fn run(self, addr: SocketAddr) -> Result<(), anyhow::Error> {
+        info!("Starting gRPC server on {addr}");
+        let config = self.get_configuration();
+        let (grpc_cert, grpc_key) = if let Some(cfg) = config {
+            (cfg.grpc_cert_pem, cfg.grpc_key_pem)
+        } else {
+            return Err(anyhow::anyhow!("gRPC server configuration is missing"));
+        };
+
+        let identity = Identity::from_pem(grpc_cert, grpc_key);
+        let mut builder =
+            Server::builder().tls_config(ServerTlsConfig::new().identity(identity))?;
+
+        let own_version = Version::parse(VERSION)?;
+        let versioned_service = ServiceBuilder::new()
+            .layer(tonic::service::InterceptorLayer::new(
+                DefguardVersionInterceptor::new(
+                    own_version.clone(),
+                    DefguardComponent::Core,
+                    MIN_CORE_VERSION,
+                    false,
+                ),
+            ))
+            .layer(DefguardVersionLayer::new(own_version))
+            .service(proxy_server::ProxyServer::new(self.clone()));
+
+        builder
+            .add_service(versioned_service)
+            .serve_with_shutdown(addr, async move {
+                let mut rx_lock = GRPC_SERVER_RESTART_CHANNEL.1.lock().await;
+                rx_lock.recv().await;
+                info!("Shutting down gRPC server for restart...");
+            })
+            .await
+            .map_err(|err| {
+                error!("gRPC server error: {err}");
+                err
+            })?;
+
+        Ok(())
+    }
+
+    /// Sends message to the other side of RPC, with given `payload` and `device_info`.
+    #[instrument(level = "debug", skip(self, payload))]
     pub(crate) fn send(
         &self,
         payload: core_request::Payload,
         device_info: DeviceInfo,
     ) -> Result<oneshot::Receiver<core_response::Payload>, ApiError> {
-        if let Some(client_tx) = self.clients.lock().unwrap().values().next() {
+        if let Some(client_tx) = self
+            .clients
+            .lock()
+            .expect("Failed to acquire lock on clients hashmap when sending message to core")
+            .values()
+            .next()
+        {
             let id = self.current_id.fetch_add(1, Ordering::Relaxed);
             let res = CoreRequest {
                 id,
@@ -63,8 +157,10 @@ impl ProxyServer {
                 return Err(ApiError::Unexpected("Failed to send CoreRequest".into()));
             }
             let (tx, rx) = oneshot::channel();
-            let mut results = self.results.lock().unwrap();
-            results.insert(id, tx);
+            self.results
+                .lock()
+                .expect("Failed to acquire lock on results hashmap when sending CoreRequest")
+                .insert(id, tx);
             self.connected.store(true, Ordering::Relaxed);
             Ok(rx)
         } else {
@@ -75,6 +171,14 @@ impl ProxyServer {
             ))
         }
     }
+
+    pub(crate) fn setup_completed(&self) -> bool {
+        let lock = self
+            .config
+            .lock()
+            .expect("Failed to acquire lock on config mutex when checking setup status");
+        lock.is_some()
+    }
 }
 
 impl Clone for ProxyServer {
@@ -85,6 +189,8 @@ impl Clone for ProxyServer {
             results: Arc::clone(&self.results),
             connected: Arc::clone(&self.connected),
             core_version: Arc::clone(&self.core_version),
+            config: Arc::clone(&self.config),
+            setup_in_progress: Arc::clone(&self.setup_in_progress),
         }
     }
 }
@@ -99,14 +205,22 @@ impl proxy_server::Proxy for ProxyServer {
         &self,
         request: Request<Streaming<CoreResponse>>,
     ) -> Result<Response<Self::BidiStream>, Status> {
+        if !self.setup_completed() {
+            error!("Received bidi connection before setup completion");
+            return Err(Status::failed_precondition(
+                "Setup must be completed before establishing bidi connection",
+            ));
+        }
+
         let Some(address) = request.remote_addr() else {
             error!("Failed to determine client address for request: {request:?}");
             return Err(Status::internal("Failed to determine client address"));
         };
         let maybe_info = ComponentInfo::from_metadata(request.metadata());
         let (version, info) = get_tracing_variables(&maybe_info);
-        let mut core_version = self.core_version.lock().unwrap();
-        *core_version = Some(version.clone());
+        *self.core_version.lock().expect(
+            "Failed to acquire lock on core_version mutex when storing version information",
+        ) = Some(version.clone());
 
         let span = tracing::info_span!("core_bidi_stream", component = %DefguardComponent::Core,
             version = version.to_string(), info);
@@ -115,7 +229,12 @@ impl proxy_server::Proxy for ProxyServer {
         info!("Defguard Core gRPC client connected from: {address}");
 
         let (tx, rx) = mpsc::unbounded_channel();
-        self.clients.lock().unwrap().insert(address, tx);
+        self.clients
+            .lock()
+            .expect(
+                "Failed to acquire lock on clients hashmap when registering new core connection",
+            )
+            .insert(address, tx);
         self.connected.store(true, Ordering::Relaxed);
 
         let clients = Arc::clone(&self.clients);
@@ -129,9 +248,9 @@ impl proxy_server::Proxy for ProxyServer {
                         Ok(Some(response)) => {
                             debug!("Received message from Defguard Core ID={}", response.id);
                             connected.store(true, Ordering::Relaxed);
-                            // Discard empty payloads.
                             if let Some(payload) = response.payload {
-                                if let Some(rx) = results.lock().unwrap().remove(&response.id) {
+                                let maybe_rx = results.lock().expect("Failed to acquire lock on results hashmap when processing response").remove(&response.id);
+                                if let Some(rx) = maybe_rx {
                                     if let Err(err) = rx.send(payload) {
                                         error!("Failed to send message to rx {:?}", err.type_id());
                                     }
@@ -152,7 +271,7 @@ impl proxy_server::Proxy for ProxyServer {
                 }
                 info!("Defguard core client disconnected: {address}");
                 connected.store(false, Ordering::Relaxed);
-                clients.lock().unwrap().remove(&address);
+                clients.lock().expect("Failed to acquire lock on clients hashmap when removing disconnected client").remove(&address);
             }
             .instrument(tracing::Span::current()),
         );

src/handlers/desktop_client_mfa.rs

@@ -1,3 +1,5 @@
+use std::collections::hash_map::Entry;
+
 use axum::{
     extract::{
         ws::{Message, WebSocket},
@@ -10,7 +12,6 @@ use axum::{
 use futures_util::{sink::SinkExt, stream::StreamExt};
 use serde::Deserialize;
 use serde_json::json;
-use std::collections::hash_map::Entry;
 use tokio::{sync::oneshot, task::JoinSet};
 
 use crate::{

src/handlers/enrollment.rs

@@ -3,7 +3,6 @@ use axum_extra::extract::{cookie::Cookie, PrivateCookieJar};
 use time::OffsetDateTime;
 
 use super::register_mfa::router as register_mfa_router;
-
 use crate::{
     error::ApiError,
     handlers::{get_core_response, mobile_client::register_mobile_auth},

src/handlers/register_mfa.rs

@@ -1,7 +1,6 @@
-use serde::Deserialize;
-
 use axum::{extract::State, response::IntoResponse, routing::post, Json, Router};
 use axum_extra::extract::PrivateCookieJar;
+use serde::Deserialize;
 
 use crate::{
     error::ApiError,