versioned core communication

j-chmielewski · j-chmielewski · commit c77d6cf724dc · 2025-08-08T11:08:51.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/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"
@@ -44,6 +45,7 @@ tower_governor = "0.4"
 rust-embed = { version = "8.5", features = ["include-exclude"] }
 mime_guess = "2.0"
 base64 = "0.22.1"
+tonic-middleware = "0.2"
 
 [build-dependencies]
 tonic-build = { version = "0.12" }
diff --git a/src/http.rs b/src/http.rs
@@ -1,7 +1,7 @@
 use std::{
     fs::read_to_string,
     net::{IpAddr, Ipv4Addr, SocketAddr},
-    sync::atomic::Ordering,
+    sync::{atomic::Ordering, Arc},
     time::Duration,
 };
 
@@ -15,9 +15,11 @@ use axum::{
 };
 use axum_extra::extract::cookie::Key;
 use clap::crate_version;
+use defguard_version::{server::DefguardVersionServerMiddleware, DefguardVersionSet};
 use serde::Serialize;
 use tokio::{net::TcpListener, task::JoinSet};
 use tonic::transport::{Identity, Server, ServerTlsConfig};
+use tonic_middleware::MiddlewareFor;
 use tower_governor::{
     governor::GovernorConfigBuilder, key_extractor::SmartIpKeyExtractor, GovernorLayer,
 };
@@ -116,7 +118,10 @@ fn get_client_addr(request: &Request<Body>) -> String {
         )
 }
 
-pub async fn run_server(config: Config) -> anyhow::Result<()> {
+pub async fn run_server(
+    config: Config,
+    version_set: Arc<DefguardVersionSet>,
+) -> anyhow::Result<()> {
     info!("Starting Defguard Proxy server");
     debug!("Using config: {config:?}");
 
@@ -162,8 +167,15 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> {
         } else {
             Server::builder()
         };
+        let versioned_server = MiddlewareFor::new(
+            proxy_server::ProxyServer::new(grpc_server),
+            DefguardVersionServerMiddleware::new(
+                version_set.own.clone(),
+                Arc::clone(&version_set.proxy),
+            ),
+        );
         builder
-            .add_service(proxy_server::ProxyServer::new(grpc_server))
+            .add_service(versioned_server)
             .serve(addr)
             .await
             .context("Error running gRPC server")
diff --git a/src/main.rs b/src/main.rs
@@ -1,4 +1,4 @@
-use defguard_proxy::{config::get_config, http::run_server, logging::init_tracing, VERSION};
+use defguard_proxy::{config::get_config, http::run_server, VERSION};
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
@@ -9,11 +9,15 @@ async fn main() -> anyhow::Result<()> {
 
     // read config from env
     let config = get_config()?;
-    init_tracing(&config.log_level);
+    let version_set = defguard_version::tracing::init(
+        VERSION,
+        &config.log_level.to_string(),
+        &["send_grpc_message", "bidirectional_communication"],
+    );
     tracing::info!("Starting ... version v{}", VERSION);
 
     // run API web server
-    run_server(config).await?;
+    run_server(config, version_set).await?;
 
     Ok(())
 }
