Send IP address and user agent over gRPC (#27)

* Send IP address and user agent over gRPC --------- Co-authored-by: bgasztka <[email protected]> Co-authored-by: bgasztka <[email protected]>
DefGuard · Nov 23, 2023 · 8d42bda · 8d42bda
1 parent 5ec42d4
commit 8d42bda
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -5,7 +5,7 @@ edition = "2021"
 
 [dependencies]
 # base `axum` deps
-axum = { version = "0.6", features = ["json", "tracing", "macros"] }
+axum = { version = "0.6", features = ["json", "tracing", "macros", "headers"] }
 hyper = { version = "0.14", features = ["full"] }
 tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
 tower = "0.4"

diff --git a/src/handlers/enrollment.rs b/src/handlers/enrollment.rs
@@ -1,3 +1,5 @@
+use std::net::SocketAddr;
+
 use crate::error::ApiError;
 use crate::server::{COOKIE_NAME, SECRET_KEY};
 use crate::{
@@ -8,7 +10,12 @@ use crate::{
     handlers::ApiResult,
     server::AppState,
 };
-use axum::{extract::State, routing::post, Json, Router};
+use axum::{
+    extract::{ConnectInfo, State},
+    headers::UserAgent,
+    routing::post,
+    Json, Router, TypedHeader,
+};
 use tonic::metadata::MetadataValue;
 use tower_cookies::{cookie::time::OffsetDateTime, Cookie, Cookies};
 use tracing::{debug, error, info};
@@ -41,6 +48,23 @@ fn add_auth_header<T>(cookies: Cookies, request: &mut tonic::Request<T>) -> Resu
     Ok(())
 }
 
+fn add_device_info_header<T>(
+    request: &mut tonic::Request<T>,
+    ip_address: String,
+    user_agent: Option<TypedHeader<UserAgent>>,
+) -> Result<(), ApiError> {
+    let user_agent_string: String = user_agent.map(|v| v.to_string()).unwrap_or_default();
+
+    request
+        .metadata_mut()
+        .insert("ip_address", MetadataValue::try_from(ip_address)?);
+    request
+        .metadata_mut()
+        .insert("user_agent", MetadataValue::try_from(user_agent_string)?);
+
+    Ok(())
+}
+
 pub async fn start_enrollment_process(
     State(state): State<AppState>,
     cookies: Cookies,
@@ -72,29 +96,37 @@ pub async fn start_enrollment_process(
 
 pub async fn activate_user(
     State(state): State<AppState>,
+    ConnectInfo(connect_info): ConnectInfo<SocketAddr>,
+    user_agent: Option<TypedHeader<UserAgent>>,
     cookies: Cookies,
     Json(req): Json<ActivateUserRequest>,
 ) -> ApiResult<()> {
     info!("Activating user");
 
+    let ip_address = connect_info.ip().to_string();
     let mut client = state.client.lock().await;
     let mut request = tonic::Request::new(req);
     add_auth_header(cookies, &mut request)?;
+    add_device_info_header(&mut request, ip_address, user_agent)?;
     client.activate_user(request).await?;
 
     Ok(())
 }
 
 pub async fn create_device(
     State(state): State<AppState>,
+    ConnectInfo(connect_info): ConnectInfo<SocketAddr>,
+    user_agent: Option<TypedHeader<UserAgent>>,
     cookies: Cookies,
     Json(req): Json<NewDevice>,
 ) -> ApiResult<Json<DeviceConfigResponse>> {
     info!("Adding new device");
 
+    let ip_address = connect_info.ip().to_string();
     let mut client = state.client.lock().await;
     let mut request = tonic::Request::new(req);
     add_auth_header(cookies, &mut request)?;
+    add_device_info_header(&mut request, ip_address, user_agent)?;
     let response = client.create_device(request).await?;
 
     Ok(Json(response.into_inner()))

diff --git a/src/server.rs b/src/server.rs
@@ -104,7 +104,7 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> {
     let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), http_port);
     info!("Listening on {}", addr);
     axum::Server::bind(&addr)
-        .serve(app.into_make_service())
+        .serve(app.into_make_service_with_connect_info::<SocketAddr>())
         .await
         .context("Error running HTTP server")
 }