systemprompt_api/services/middleware/analytics/

mod.rs

mod detection;
mod events;

use axum::extract::Request;
use axum::http::StatusCode;
use axum::middleware::Next;
use axum::response::Response;
use std::sync::Arc;

use systemprompt_analytics::SessionRepository;
use systemprompt_identifiers::SessionId;
use systemprompt_logging::AnalyticsRepository;
use systemprompt_models::{RequestContext, RouteClassifier};
use systemprompt_runtime::AppContext;
use systemprompt_security::ScannerDetector;

pub use events::AnalyticsEventParams;

struct TrackingParams<'a> {
    req_ctx: &'a RequestContext,
    uri: &'a http::Uri,
    method: &'a http::Method,
    status_code: u16,
    response_time_ms: u64,
    user_agent: Option<String>,
    referer: Option<String>,
    is_scanner: bool,
}

#[derive(Debug, Clone)]
pub struct AnalyticsMiddleware {
    session_repo: Arc<SessionRepository>,
    analytics_repo: Arc<AnalyticsRepository>,
    route_classifier: Arc<RouteClassifier>,
}

impl AnalyticsMiddleware {
    pub fn new(app_context: &AppContext) -> anyhow::Result<Self> {
        let db_pool = app_context.db_pool();
        let session_repo = Arc::new(SessionRepository::new(db_pool)?);
        let analytics_repo = Arc::new(AnalyticsRepository::new(db_pool)?);
        let route_classifier = Arc::clone(app_context.route_classifier());

        Ok(Self {
            session_repo,
            analytics_repo,
            route_classifier,
        })
    }

    pub async fn track_request(
        &self,
        request: Request,
        next: Next,
    ) -> Result<Response, StatusCode> {
        let method = request.method().clone();
        let uri = request.uri().clone();

        let Some(req_ctx) = request.extensions().get::<RequestContext>().cloned() else {
            return Ok(next.run(request).await);
        };

        if !req_ctx.request.is_tracked {
            return Ok(next.run(request).await);
        }

        let user_agent = request
            .headers()
            .get("user-agent")
            .and_then(|v| v.to_str().ok())
            .map(ToString::to_string);

        let referer = request
            .headers()
            .get("referer")
            .and_then(|v| v.to_str().ok())
            .map(ToString::to_string);

        let start_time = std::time::Instant::now();
        let response = next.run(request).await;
        let response_time_ms = start_time.elapsed().as_millis() as u64;
        let status_code = response.status();

        let should_track = self
            .route_classifier
            .should_track_analytics(uri.path(), method.as_str());

        let is_scanner =
            ScannerDetector::is_scanner(Some(uri.path()), user_agent.as_deref(), None, None);

        if should_track {
            self.spawn_tracking_tasks(TrackingParams {
                req_ctx: &req_ctx,
                uri: &uri,
                method: &method,
                status_code: status_code.as_u16(),
                response_time_ms,
                user_agent,
                referer,
                is_scanner,
            });
        }

        Ok(response)
    }

    fn spawn_tracking_tasks(&self, params: TrackingParams<'_>) {
        let req_ctx = params.req_ctx;
        let uri = params.uri;
        let method = params.method;
        let status_code = params.status_code;
        let response_time_ms = params.response_time_ms;
        let user_agent = params.user_agent;
        let referer = params.referer;
        let is_scanner = params.is_scanner;
        let endpoint = format!("{} {}", method, uri.path());
        let path = uri.path().to_string();

        if is_scanner {
            self.spawn_mark_scanner_task(req_ctx.request.session_id.clone());
        }

        self.spawn_velocity_scanner_check(req_ctx.request.session_id.clone());

        self.spawn_session_tracking_task(req_ctx.request.session_id.clone());

        detection::spawn_behavioral_detection_task(
            Arc::clone(&self.session_repo),
            req_ctx.request.session_id.clone(),
            req_ctx.request.fingerprint_hash.clone(),
            user_agent.clone(),
            1,
        );

        events::spawn_analytics_event_task(
            Arc::clone(&self.analytics_repo),
            Arc::clone(&self.route_classifier),
            AnalyticsEventParams {
                req_ctx: req_ctx.clone(),
                endpoint,
                path,
                method: method.to_string(),
                uri: uri.clone(),
                status_code,
                response_time_ms,
                user_agent,
                referer,
            },
        );
    }

    fn spawn_session_tracking_task(&self, session_id: SessionId) {
        let session_repo = Arc::clone(&self.session_repo);

        tokio::spawn(async move {
            if let Err(e) = session_repo.update_activity(&session_id).await {
                tracing::error!(error = %e, "Failed to update session activity");
            }

            if let Err(e) = session_repo.increment_request_count(&session_id).await {
                tracing::error!(error = %e, "Failed to increment request count");
            }
        });
    }

    fn spawn_velocity_scanner_check(&self, session_id: SessionId) {
        let session_repo = Arc::clone(&self.session_repo);

        tokio::spawn(async move {
            let (request_count, duration_seconds) = session_repo
                .get_session_velocity(&session_id)
                .await
                .unwrap_or((None, None));

            if let (Some(count), Some(duration)) = (request_count, duration_seconds) {
                if ScannerDetector::is_high_velocity(count, duration) {
                    if let Err(e) = session_repo.mark_as_scanner(&session_id).await {
                        tracing::warn!(
                            error = %e,
                            session_id = %session_id,
                            "Failed to mark high-velocity session as scanner"
                        );
                    }
                }
            }
        });
    }

    fn spawn_mark_scanner_task(&self, session_id: SessionId) {
        let session_repo = Arc::clone(&self.session_repo);

        tokio::spawn(async move {
            if let Err(e) = session_repo.mark_as_scanner(&session_id).await {
                tracing::warn!(error = %e, session_id = %session_id, "Failed to mark session as scanner");
            }
        });
    }
}