refactor as own middleware

Author: Devdutt Shenoi

src/audit.rs

@@ -2,12 +2,8 @@ use std::{collections::HashMap, fmt::Debug, sync::Arc};
 
 use crate::about::current;
 use crate::handlers::http::modal::utils::rbac_utils::get_metadata;
-use crate::rbac::map::SessionKey;
-use crate::rbac::Users;
 
 use super::option::CONFIG;
-use actix_web::dev::ServiceRequest;
-use actix_web_httpauth::extractors::basic::BasicAuth;
 use chrono::{DateTime, Utc};
 use reqwest::Client;
 use serde::Serialize;
@@ -156,8 +152,6 @@ impl Default for ResponseLog {
     }
 }
 
-const DROP_HEADERS: [&str; 4] = ["authorization", "cookie", "user-agent", "x-p-stream"];
-
 pub struct AuditLogBuilder {
     version: AuditLogVersion,
     deployment_id: Ulid,
@@ -196,59 +190,6 @@ impl AuditLogBuilder {
     pub fn set_stream_name(&mut self, stream: String) {
         self.stream = stream;
     }
-
-    pub fn update_from_http(&mut self, req: &mut ServiceRequest) {
-        let mut username = "Unknown".to_owned();
-        let mut authorization_method = "None".to_owned();
-
-        // Extract authorization details from request, either from basic auth
-        // header or cookie, else use default value.
-        if let Ok(creds) = req.extract::<BasicAuth>().into_inner() {
-            username = creds.user_id().trim().to_owned();
-            authorization_method = "Basic Auth".to_owned();
-        } else if let Some(cookie) = req.cookie("session") {
-            authorization_method = "Session Cookie".to_owned();
-            if let Some(user_id) = Ulid::from_string(cookie.value())
-                .ok()
-                .and_then(|ulid| Users.get_username_from_session(&SessionKey::SessionId(ulid)))
-            {
-                username = user_id;
-            }
-        }
-
-        let conn = req.connection_info();
-        self.request = RequestLog {
-            method: req.method().to_string(),
-            path: req.path().to_string(),
-            protocol: conn.scheme().to_owned(),
-            headers: req
-                .headers()
-                .iter()
-                .filter_map(|(name, value)| match name.as_str() {
-                    // NOTE: drop headers that are not required
-                    name if DROP_HEADERS.contains(&name.to_lowercase().as_str()) => None,
-                    name => {
-                        // NOTE: Drop headers that can't be parsed as string
-                        value
-                            .to_str()
-                            .map(|value| (name.to_owned(), value.to_string()))
-                            .ok()
-                    }
-                })
-                .collect(),
-        };
-        self.actor = ActorLog {
-            remote_host: conn.realip_remote_addr().unwrap_or_default().to_owned(),
-            user_agent: req
-                .headers()
-                .get("User-Agent")
-                .and_then(|a| a.to_str().ok())
-                .unwrap_or_default()
-                .to_owned(),
-            username,
-            authorization_method,
-        }
-    }
 }
 
 impl Drop for AuditLogBuilder {

src/handlers/http/audit.rs

@@ -0,0 +1,109 @@
+use super::middleware::Message;
+use actix_web::{
+    body::MessageBody,
+    dev::{ServiceRequest, ServiceResponse},
+    middleware::Next,
+};
+use actix_web_httpauth::extractors::basic::BasicAuth;
+use ulid::Ulid;
+
+use crate::{
+    audit::{ActorLog, AuditLogBuilder, RequestLog},
+    handlers::{KINESIS_COMMON_ATTRIBUTES_KEY, STREAM_NAME_HEADER_KEY},
+    rbac::{map::SessionKey, Users},
+};
+
+const DROP_HEADERS: [&str; 4] = ["authorization", "cookie", "user-agent", "x-p-stream"];
+
+pub async fn audit_log_middleware(
+    mut req: ServiceRequest,
+    next: Next<impl MessageBody>,
+) -> Result<ServiceResponse<impl MessageBody>, actix_web::Error> {
+    // Ensures that log will be pushed to subscriber on drop
+    let mut log_builder = AuditLogBuilder::default();
+
+    if let Some(kinesis_common_attributes) =
+        req.request().headers().get(KINESIS_COMMON_ATTRIBUTES_KEY)
+    {
+        let attribute_value: &str = kinesis_common_attributes.to_str().unwrap();
+        let message: Message = serde_json::from_str(attribute_value).unwrap();
+        log_builder.set_stream_name(message.common_attributes.x_p_stream);
+    } else if let Some(stream) = req.match_info().get("logstream") {
+        log_builder.set_stream_name(stream.to_owned());
+    } else if let Some(value) = req.headers().get(STREAM_NAME_HEADER_KEY) {
+        if let Ok(stream) = value.to_str() {
+            log_builder.set_stream_name(stream.to_owned());
+        }
+    }
+    let mut username = "Unknown".to_owned();
+    let mut authorization_method = "None".to_owned();
+
+    // Extract authorization details from request, either from basic auth
+    // header or cookie, else use default value.
+    if let Ok(creds) = req.extract::<BasicAuth>().into_inner() {
+        username = creds.user_id().trim().to_owned();
+        authorization_method = "Basic Auth".to_owned();
+    } else if let Some(cookie) = req.cookie("session") {
+        authorization_method = "Session Cookie".to_owned();
+        if let Some(user_id) = Ulid::from_string(cookie.value())
+            .ok()
+            .and_then(|ulid| Users.get_username_from_session(&SessionKey::SessionId(ulid)))
+        {
+            username = user_id;
+        }
+    }
+
+    log_builder.request = RequestLog {
+        method: req.method().to_string(),
+        path: req.path().to_string(),
+        protocol: req.connection_info().scheme().to_owned(),
+        headers: req
+            .headers()
+            .iter()
+            .filter_map(|(name, value)| match name.as_str() {
+                // NOTE: drop headers that are not required
+                name if DROP_HEADERS.contains(&name.to_lowercase().as_str()) => None,
+                name => {
+                    // NOTE: Drop headers that can't be parsed as string
+                    value
+                        .to_str()
+                        .map(|value| (name.to_owned(), value.to_string()))
+                        .ok()
+                }
+            })
+            .collect(),
+    };
+    log_builder.actor = ActorLog {
+        remote_host: req
+            .connection_info()
+            .realip_remote_addr()
+            .unwrap_or_default()
+            .to_owned(),
+        user_agent: req
+            .headers()
+            .get("User-Agent")
+            .and_then(|a| a.to_str().ok())
+            .unwrap_or_default()
+            .to_owned(),
+        username,
+        authorization_method,
+    };
+    log_builder.set_deployment_id().await;
+
+    let res = next.call(req).await;
+
+    // Capture status_code and error information from response
+    match &res {
+        Ok(res) => {
+            let status = res.status();
+            log_builder.response.status_code = status.as_u16();
+            // Use error information from reponse object if an error
+            if let Some(err) = res.response().error() {
+                log_builder.set_response_error(err.to_string());
+            }
+        }
+        Err(err) => log_builder.set_response_error(err.to_string()),
+    }
+
+    res
+}

src/handlers/http/middleware.rs

@@ -28,7 +28,6 @@ use actix_web::{
 use futures_util::future::LocalBoxFuture;
 
 use crate::{
-    audit::AuditLogBuilder,
     handlers::{
         AUTHORIZATION_KEY, KINESIS_COMMON_ATTRIBUTES_KEY, LOG_SOURCE_KEY, LOG_SOURCE_KINESIS,
         STREAM_NAME_HEADER_KEY,
@@ -46,16 +45,16 @@ use serde::{Deserialize, Serialize};
 
 #[derive(Serialize, Deserialize, Debug)]
 #[serde(rename_all = "camelCase")]
-struct Message {
-    common_attributes: CommonAttributes,
+pub struct Message {
+    pub common_attributes: CommonAttributes,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
-struct CommonAttributes {
+pub struct CommonAttributes {
     #[serde(rename = "Authorization")]
     authorization: String,
     #[serde(rename = "X-P-Stream")]
-    x_p_stream: String,
+    pub x_p_stream: String,
 }
 
 pub trait RouteExt {
@@ -138,7 +137,6 @@ where
         For requests made from other clients, no change.
 
          ## Section start */
-        let mut stream_name = None;
         if let Some(kinesis_common_attributes) =
             req.request().headers().get(KINESIS_COMMON_ATTRIBUTES_KEY)
         {
@@ -156,27 +154,14 @@ where
                 HeaderName::from_static(LOG_SOURCE_KEY),
                 header::HeaderValue::from_static(LOG_SOURCE_KINESIS),
             );
-            stream_name.replace(message.common_attributes.x_p_stream);
-        } else if let Some(stream) = req.match_info().get("logstream") {
-            stream_name.replace(stream.to_owned());
-        } else if let Some(value) = req.headers().get(STREAM_NAME_HEADER_KEY) {
-            if let Ok(stream) = value.to_str() {
-                stream_name.replace(stream.to_owned());
-            }
         }
 
         /* ## Section end */
 
         let auth_result: Result<_, Error> = (self.auth_method)(&mut req, self.action);
 
-        // Ensures that log will be pushed to subscriber on drop
-        let mut log_builder = AuditLogBuilder::default();
-        log_builder.set_stream_name(stream_name.unwrap_or_default());
-        log_builder.update_from_http(&mut req);
         let fut = self.service.call(req);
         Box::pin(async move {
-            log_builder.set_deployment_id().await;
-
             match auth_result? {
                 rbac::Response::UnAuthorized => return Err(
                     ErrorForbidden("You don't have permission to access this resource. Please contact your administrator for assistance.")
@@ -186,22 +171,8 @@ where
                 ),
                 _ => {}
             }
-            let res = fut.await;
-
-            // Capture status_code and error information from response
-            match &res {
-                Ok(res) => {
-                    let status = res.status();
-                    log_builder.response.status_code = status.as_u16();
-                    // Use error information from reponse object if an error
-                    if let Some(err) = res.response().error() {
-                        log_builder.set_response_error(err.to_string());
-                    }
-                }
-                Err(err) => log_builder.set_response_error(err.to_string()),
-            }
 
-            res
+            fut.await
         })
     }
 }

src/handlers/http/mod.rs

@@ -28,6 +28,7 @@ use crate::{option::CONFIG, storage::STREAM_ROOT_DIRECTORY};
 use self::{cluster::get_ingestor_info, query::Query};
 
 pub mod about;
+mod audit;
 pub mod cluster;
 pub mod correlation;
 pub mod health_check;

src/handlers/http/modal/mod.rs

@@ -41,6 +41,7 @@ use ssl_acceptor::get_ssl_acceptor;
 use tokio::sync::{oneshot, Mutex};
 use tracing::{error, info, warn};
 
+use super::audit;
 use super::cross_origin_config;
 use super::API_BASE_PATH;
 use super::API_VERSION;
@@ -101,6 +102,7 @@ pub trait ParseableServer {
                 .wrap(prometheus.clone())
                 .configure(|config| Self::configure_routes(config, oidc_client.clone()))
                 .wrap(from_fn(health_check::check_shutdown_middleware))
+                .wrap(from_fn(audit::audit_log_middleware))
                 .wrap(actix_web::middleware::Logger::default())
                 .wrap(actix_web::middleware::Compress::default())
                 .wrap(cross_origin_config())