Version exchange and logging

Cargo.toml

@@ -7,6 +7,7 @@ homepage = "https://github.com/DefGuard/proxy"
 repository = "https://github.com/DefGuard/proxy"
 
 [dependencies]
+defguard_version = { path = "../defguard/crates/defguard_version" }
 # base `axum` deps
 axum = { version = "0.7", features = ["macros", "tracing"] }
 axum-client-ip = "0.6"
@@ -48,6 +49,7 @@ tower_governor = "0.4"
 rust-embed = { version = "8.5", features = ["include-exclude"] }
 mime_guess = "2.0"
 base64 = "0.22.1"
+tower = "0.5.2"
 
 [build-dependencies]
 tonic-prost-build = "0.14"

src/grpc.rs

@@ -1,3 +1,4 @@
+use defguard_version::{version_info_from_metadata, DefguardComponent};
 use std::{
     collections::HashMap,
     net::SocketAddr,
@@ -6,10 +7,10 @@ use std::{
         Arc, Mutex,
     },
 };
-
 use tokio::sync::{mpsc, oneshot};
 use tokio_stream::wrappers::UnboundedReceiverStream;
 use tonic::{Request, Response, Status, Streaming};
+use tracing::Instrument;
 
 use crate::{
     error::ApiError,
@@ -89,7 +90,7 @@ impl proxy_server::Proxy for ProxyServer {
     type BidiStream = UnboundedReceiverStream<Result<CoreRequest, Status>>;
 
     /// Handle bidirectional communication with Defguard core.
-    #[instrument(name = "bidirectional_communication", level = "debug", skip(self))]
+    #[instrument(name = "bidirectional_communication", level = "info", skip(self))]
     async fn bidi(
         &self,
         request: Request<Streaming<CoreResponse>>,
@@ -98,6 +99,9 @@ impl proxy_server::Proxy for ProxyServer {
             error!("Failed to determine client address for request: {request:?}");
             return Err(Status::internal("Failed to determine client address"));
         };
+        let (version, info) = version_info_from_metadata(request.metadata());
+        let span = tracing::info_span!("core_bidi_stream", component = %DefguardComponent::Core, version, info);
+        let _guard = span.enter();
         info!("Defguard Core gRPC client connected from: {address}");
 
         let (tx, rx) = mpsc::unbounded_channel();
@@ -108,37 +112,40 @@ impl proxy_server::Proxy for ProxyServer {
         let results = Arc::clone(&self.results);
         let connected = Arc::clone(&self.connected);
         let mut stream = request.into_inner();
-        tokio::spawn(async move {
-            loop {
-                match stream.message().await {
-                    Ok(Some(response)) => {
-                        debug!("Received message from Defguard core: {response:?}");
-                        connected.store(true, Ordering::Relaxed);
-                        // Discard empty payloads.
-                        if let Some(payload) = response.payload {
-                            if let Some(rx) = results.lock().unwrap().remove(&response.id) {
-                                if let Err(err) = rx.send(payload) {
-                                    error!("Failed to send message to rx: {err:?}");
+        tokio::spawn(
+            async move {
+                loop {
+                    match stream.message().await {
+                        Ok(Some(response)) => {
+                            debug!("Received message from Defguard core: {response:?}");
+                            connected.store(true, Ordering::Relaxed);
+                            // Discard empty payloads.
+                            if let Some(payload) = response.payload {
+                                if let Some(rx) = results.lock().unwrap().remove(&response.id) {
+                                    if let Err(err) = rx.send(payload) {
+                                        error!("Failed to send message to rx: {err:?}");
+                                    }
+                                } else {
+                                    error!("Missing receiver for response #{}", response.id);
                                 }
-                            } else {
-                                error!("Missing receiver for response #{}", response.id);
                             }
                         }
-                    }
-                    Ok(None) => {
-                        info!("gRPC stream has been closed");
-                        break;
-                    }
-                    Err(err) => {
-                        error!("gRPC client error: {err}");
-                        break;
+                        Ok(None) => {
+                            info!("gRPC stream has been closed");
+                            break;
+                        }
+                        Err(err) => {
+                            error!("gRPC client error: {err}");
+                            break;
+                        }
                     }
                 }
+                info!("Defguard core client disconnected: {address}");
+                connected.store(false, Ordering::Relaxed);
+                clients.lock().unwrap().remove(&address);
             }
-            info!("Defguard core client disconnected: {address}");
-            connected.store(false, Ordering::Relaxed);
-            clients.lock().unwrap().remove(&address);
-        });
+            .instrument(tracing::Span::current()),
+        );
 
         Ok(Response::new(UnboundedReceiverStream::new(rx)))
     }

src/http.rs

@@ -15,9 +15,11 @@ use axum::{
 };
 use axum_extra::extract::cookie::Key;
 use clap::crate_version;
+use defguard_version::{server::DefguardVersionLayer, Version};
 use serde::Serialize;
 use tokio::{net::TcpListener, task::JoinSet};
 use tonic::transport::{Identity, Server, ServerTlsConfig};
+use tower::ServiceBuilder;
 use tower_governor::{
     governor::GovernorConfigBuilder, key_extractor::SmartIpKeyExtractor, GovernorLayer,
 };
@@ -33,6 +35,7 @@ use crate::{
     grpc::ProxyServer,
     handlers::{desktop_client_mfa, enrollment, password_reset, polling},
     proto::proxy_server,
+    VERSION,
 };
 
 pub(crate) static ENROLLMENT_COOKIE_NAME: &str = "defguard_proxy";
@@ -162,8 +165,11 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> {
         } else {
             Server::builder()
         };
+        let versioned_service = ServiceBuilder::new()
+            .layer(DefguardVersionLayer::new(Version::parse(VERSION)?))
+            .service(proxy_server::ProxyServer::new(grpc_server));
         builder
-            .add_service(proxy_server::ProxyServer::new(grpc_server))
+            .add_service(versioned_service)
             .serve(addr)
             .await
             .context("Error running gRPC server")