ultrafast_gateway/plugins/
input_validation.rs

1use crate::gateway_error::GatewayError;
2use crate::plugins::{PluginHooks, PluginLifecycle, PluginMetadata, PluginState};
3use axum::body::Body;
4use axum::http::Request;
5use axum::response::Response;
6use serde::{Deserialize, Serialize};
7use std::str::FromStr;
8
9#[derive(Debug, Clone, Serialize, Deserialize, Default)]
10pub struct ValidationConfig {
11    pub enabled: bool,
12    pub max_request_size: usize,
13    pub max_model_name_length: usize,
14}
15
16#[derive(Debug, Clone)]
17pub struct InputValidationPlugin {
18    meta: PluginMetadata,
19    config: ValidationConfig,
20}
21
22impl InputValidationPlugin {
23    pub fn new(config: ValidationConfig) -> Self {
24        Self {
25            meta: PluginMetadata {
26                id: uuid::Uuid::new_v4().to_string(),
27                name: "input_validation".to_string(),
28                version: "1.0.0".to_string(),
29                enabled: config.enabled,
30                state: PluginState::Inactive,
31                dependencies: vec![],
32                priority: 4,
33                last_error: None,
34            },
35            config,
36        }
37    }
38
39    pub fn enabled(&self) -> bool {
40        self.meta.enabled
41    }
42}
43
44#[async_trait::async_trait]
45impl PluginLifecycle for InputValidationPlugin {
46    async fn initialize(&mut self) -> Result<(), GatewayError> {
47        Ok(())
48    }
49    async fn start(&mut self) -> Result<(), GatewayError> {
50        Ok(())
51    }
52    async fn stop(&mut self) -> Result<(), GatewayError> {
53        Ok(())
54    }
55    async fn cleanup(&mut self) -> Result<(), GatewayError> {
56        Ok(())
57    }
58    async fn health_check(&self) -> Result<(), GatewayError> {
59        Ok(())
60    }
61    fn metadata(&self) -> &PluginMetadata {
62        &self.meta
63    }
64    fn metadata_mut(&mut self) -> &mut PluginMetadata {
65        &mut self.meta
66    }
67}
68
69#[async_trait::async_trait]
70impl PluginHooks for InputValidationPlugin {
71    async fn before_request(&self, request: &mut Request<Body>) -> Result<(), GatewayError> {
72        if !self.config.enabled {
73            return Ok(());
74        }
75
76        // Lightweight check: respect Content-Length if present
77        if let Some(len) = request
78            .headers()
79            .get(axum::http::header::CONTENT_LENGTH)
80            .and_then(|v| v.to_str().ok())
81            .and_then(|s| usize::from_str(s).ok())
82        {
83            if len > self.config.max_request_size {
84                return Err(GatewayError::InvalidRequest {
85                    message: "Request too large".into(),
86                });
87            }
88        }
89
90        // Skip deep JSON validation to avoid consuming the body; plugin is intentionally permissive
91        Ok(())
92    }
93
94    async fn after_response(&self, _response: &mut Response<Body>) -> Result<(), GatewayError> {
95        Ok(())
96    }
97    async fn on_error(&self, _error: &GatewayError) -> Result<(), GatewayError> {
98        Ok(())
99    }
100}
101
102pub fn build_input_validation_plugin(enabled: bool) -> InputValidationPlugin {
103    InputValidationPlugin::new(ValidationConfig {
104        enabled,
105        max_request_size: 50 * 1024 * 1024,
106        max_model_name_length: 200,
107    })
108}
ultrafast_gateway/plugins/input_validation.rs

ultrafast_gateway/plugins/
input_validation.rs
