lsp_bridge/
hardening.rs

1//! Production hardening and resource limits
2//!
3//! This module provides production-grade hardening features including
4//! resource limits, rate limiting, input validation, and enhanced error recovery.
5
6use crate::error::{LspBridgeError, Result};
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::sync::atomic::{AtomicUsize, Ordering};
10use std::sync::Arc;
11use std::time::{Duration, Instant};
12use tokio::sync::{RwLock, Semaphore};
13use tokio::time::{sleep, timeout};
14
15/// Production hardening configuration
16#[derive(Debug, Clone, Serialize, Deserialize, Default)]
17pub struct HardeningConfig {
18    /// Resource limits configuration
19    pub resource_limits: ResourceLimits,
20    /// Rate limiting configuration
21    pub rate_limits: RateLimits,
22    /// Security configuration
23    pub security: SecurityConfig,
24    /// Error recovery configuration
25    pub error_recovery: ErrorRecoveryConfig,
26    /// Monitoring configuration
27    pub monitoring: MonitoringConfig,
28}
29
30/// Resource limits configuration
31#[derive(Debug, Clone, Serialize, Deserialize)]
32pub struct ResourceLimits {
33    /// Maximum memory usage in MB
34    pub max_memory_mb: usize,
35    /// Maximum number of concurrent connections
36    pub max_connections: usize,
37    /// Maximum number of concurrent requests per connection
38    pub max_requests_per_connection: usize,
39    /// Maximum request size in bytes
40    pub max_request_size_bytes: usize,
41    /// Maximum response size in bytes
42    pub max_response_size_bytes: usize,
43    /// Request timeout duration
44    pub request_timeout: Duration,
45    /// Connection timeout duration
46    pub connection_timeout: Duration,
47    /// Maximum number of retries for failed operations
48    pub max_retries: usize,
49}
50
51/// Rate limiting configuration
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct RateLimits {
54    /// Requests per second per connection
55    pub requests_per_second: f64,
56    /// Burst capacity for rate limiter
57    pub burst_capacity: usize,
58    /// Rate limit window duration
59    pub window_duration: Duration,
60    /// Enable adaptive rate limiting
61    pub adaptive: bool,
62}
63
64/// Security configuration
65#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct SecurityConfig {
67    /// Enable input validation
68    pub enable_input_validation: bool,
69    /// Enable output sanitization
70    pub enable_output_sanitization: bool,
71    /// Maximum depth for nested JSON structures
72    pub max_json_depth: usize,
73    /// Maximum string length in requests
74    pub max_string_length: usize,
75    /// Enable request logging for security audit
76    pub enable_security_logging: bool,
77}
78
79/// Error recovery configuration
80#[derive(Debug, Clone, Serialize, Deserialize)]
81pub struct ErrorRecoveryConfig {
82    /// Enable automatic retry for transient errors
83    pub enable_auto_retry: bool,
84    /// Base delay for exponential backoff
85    pub retry_base_delay: Duration,
86    /// Maximum delay for exponential backoff
87    pub retry_max_delay: Duration,
88    /// Enable circuit breaker pattern
89    pub enable_circuit_breaker: bool,
90    /// Circuit breaker failure threshold
91    pub circuit_breaker_threshold: usize,
92    /// Circuit breaker timeout duration
93    pub circuit_breaker_timeout: Duration,
94}
95
96/// Monitoring configuration
97#[derive(Debug, Clone, Serialize, Deserialize)]
98pub struct MonitoringConfig {
99    /// Enable detailed metrics collection
100    pub enable_detailed_metrics: bool,
101    /// Metrics collection interval
102    pub metrics_interval: Duration,
103    /// Enable health checks
104    pub enable_health_checks: bool,
105    /// Health check interval
106    pub health_check_interval: Duration,
107    /// Enable alerting
108    pub enable_alerting: bool,
109}
110
111/// Production hardening manager
112#[derive(Debug)]
113pub struct HardeningManager {
114    config: HardeningConfig,
115    /// Connection semaphore for limiting concurrent connections
116    connection_semaphore: Arc<Semaphore>,
117    /// Request semaphores per connection
118    request_semaphores: Arc<RwLock<HashMap<String, Arc<Semaphore>>>>,
119    /// Rate limiters per connection
120    rate_limiters: Arc<RwLock<HashMap<String, Arc<RateLimiter>>>>,
121    /// Circuit breakers per service
122    circuit_breakers: Arc<RwLock<HashMap<String, Arc<CircuitBreaker>>>>,
123    /// Resource usage tracker
124    resource_tracker: Arc<ResourceTracker>,
125    /// Input validator
126    input_validator: Arc<InputValidator>,
127}
128
129/// Rate limiter implementation using token bucket algorithm
130#[derive(Debug)]
131pub struct RateLimiter {
132    capacity: usize,
133    tokens: Arc<AtomicUsize>,
134    refill_rate: f64,
135    last_refill: Arc<RwLock<Instant>>,
136}
137
138/// Circuit breaker implementation
139#[derive(Debug)]
140pub struct CircuitBreaker {
141    state: Arc<RwLock<CircuitBreakerState>>,
142    failure_count: Arc<AtomicUsize>,
143    failure_threshold: usize,
144    timeout: Duration,
145    last_failure: Arc<RwLock<Option<Instant>>>,
146}
147
148/// Circuit breaker states
149#[derive(Debug, Clone)]
150pub enum CircuitBreakerState {
151    /// Circuit is closed, requests flow normally
152    Closed,
153    /// Circuit is open, requests are rejected
154    Open,
155    /// Circuit is half-open, testing if service has recovered
156    HalfOpen,
157}
158
159/// Resource usage tracker
160#[derive(Debug)]
161pub struct ResourceTracker {
162    memory_usage: Arc<AtomicUsize>,
163    connection_count: Arc<AtomicUsize>,
164    request_count: Arc<AtomicUsize>,
165    start_time: Instant,
166}
167
168/// Input validator for security hardening
169#[derive(Debug)]
170pub struct InputValidator {
171    config: SecurityConfig,
172}
173
174/// Validation result
175#[derive(Debug)]
176pub struct ValidationResult {
177    pub is_valid: bool,
178    pub errors: Vec<String>,
179    pub sanitized_input: Option<String>,
180}
181
182impl HardeningManager {
183    /// Create a new hardening manager with the given configuration
184    pub fn new(config: HardeningConfig) -> Self {
185        let connection_semaphore = Arc::new(Semaphore::new(config.resource_limits.max_connections));
186
187        Self {
188            connection_semaphore,
189            request_semaphores: Arc::new(RwLock::new(HashMap::new())),
190            rate_limiters: Arc::new(RwLock::new(HashMap::new())),
191            circuit_breakers: Arc::new(RwLock::new(HashMap::new())),
192            resource_tracker: Arc::new(ResourceTracker::new()),
193            input_validator: Arc::new(InputValidator::new(config.security.clone())),
194            config,
195        }
196    }
197
198    /// Acquire a connection permit (blocks if limit exceeded)
199    pub async fn acquire_connection_permit(&self) -> Result<ConnectionPermit> {
200        let permit = timeout(
201            self.config.resource_limits.connection_timeout,
202            self.connection_semaphore.clone().acquire_owned(),
203        )
204        .await
205        .map_err(|_| LspBridgeError::ResourceLimitExceeded("Connection timeout".to_string()))?
206        .map_err(|_| {
207            LspBridgeError::ResourceLimitExceeded("Connection semaphore closed".to_string())
208        })?;
209
210        self.resource_tracker.increment_connections();
211
212        Ok(ConnectionPermit {
213            _permit: permit,
214            tracker: Arc::clone(&self.resource_tracker),
215        })
216    }
217
218    /// Acquire a request permit for a specific connection
219    pub async fn acquire_request_permit(&self, connection_id: &str) -> Result<RequestPermit> {
220        // Get or create request semaphore for this connection
221        let semaphore = {
222            let mut semaphores = self.request_semaphores.write().await;
223            let entry = semaphores.entry(connection_id.to_string());
224            entry
225                .or_insert_with(|| {
226                    Arc::new(Semaphore::new(
227                        self.config.resource_limits.max_requests_per_connection,
228                    ))
229                })
230                .clone()
231        };
232
233        let permit = timeout(
234            self.config.resource_limits.request_timeout,
235            semaphore.acquire_owned(),
236        )
237        .await
238        .map_err(|_| LspBridgeError::ResourceLimitExceeded("Request timeout".to_string()))?
239        .map_err(|_| {
240            LspBridgeError::ResourceLimitExceeded("Request semaphore closed".to_string())
241        })?;
242
243        self.resource_tracker.increment_requests();
244
245        Ok(RequestPermit {
246            _permit: permit,
247            tracker: Arc::clone(&self.resource_tracker),
248        })
249    }
250
251    /// Check rate limit for a connection
252    pub async fn check_rate_limit(&self, connection_id: &str) -> Result<()> {
253        let rate_limiter = {
254            let mut limiters = self.rate_limiters.write().await;
255            let entry = limiters.entry(connection_id.to_string());
256            entry
257                .or_insert_with(|| {
258                    Arc::new(RateLimiter::new(
259                        self.config.rate_limits.requests_per_second,
260                        self.config.rate_limits.burst_capacity,
261                    ))
262                })
263                .clone()
264        };
265
266        if !rate_limiter.allow_request().await {
267            return Err(LspBridgeError::RateLimitExceeded(
268                "Too many requests".to_string(),
269            ));
270        }
271
272        Ok(())
273    }
274
275    /// Check circuit breaker for a service
276    pub async fn check_circuit_breaker(&self, service_name: &str) -> Result<()> {
277        if !self.config.error_recovery.enable_circuit_breaker {
278            return Ok(());
279        }
280
281        let circuit_breaker = {
282            let mut breakers = self.circuit_breakers.write().await;
283            let entry = breakers.entry(service_name.to_string());
284            entry
285                .or_insert_with(|| {
286                    Arc::new(CircuitBreaker::new(
287                        self.config.error_recovery.circuit_breaker_threshold,
288                        self.config.error_recovery.circuit_breaker_timeout,
289                    ))
290                })
291                .clone()
292        };
293
294        circuit_breaker.check_request().await
295    }
296
297    /// Record a successful operation for circuit breaker
298    pub async fn record_success(&self, service_name: &str) {
299        if let Some(circuit_breaker) = self.circuit_breakers.read().await.get(service_name) {
300            circuit_breaker.record_success().await;
301        }
302    }
303
304    /// Record a failed operation for circuit breaker
305    pub async fn record_failure(&self, service_name: &str) {
306        if let Some(circuit_breaker) = self.circuit_breakers.read().await.get(service_name) {
307            circuit_breaker.record_failure().await;
308        }
309    }
310
311    /// Validate input data
312    pub async fn validate_input(&self, input: &str) -> ValidationResult {
313        self.input_validator.validate(input).await
314    }
315
316    /// Execute operation with automatic retry and circuit breaker
317    pub async fn execute_with_retry<F, Fut, T>(&self, service_name: &str, operation: F) -> Result<T>
318    where
319        F: Fn() -> Fut + Send + Sync,
320        Fut: std::future::Future<Output = Result<T>> + Send,
321        T: Send,
322    {
323        self.check_circuit_breaker(service_name).await?;
324
325        let mut last_error = None;
326        let mut delay = self.config.error_recovery.retry_base_delay;
327
328        for attempt in 0..=self.config.resource_limits.max_retries {
329            match operation().await {
330                Ok(result) => {
331                    self.record_success(service_name).await;
332                    return Ok(result);
333                }
334                Err(error) => {
335                    last_error = Some(error);
336
337                    if attempt < self.config.resource_limits.max_retries {
338                        // Exponential backoff with jitter
339                        let jitter = Duration::from_millis((rand::random::<f64>() * 100.0) as u64);
340                        sleep(delay + jitter).await;
341                        delay =
342                            std::cmp::min(delay * 2, self.config.error_recovery.retry_max_delay);
343                    }
344                }
345            }
346        }
347
348        self.record_failure(service_name).await;
349        Err(last_error.unwrap())
350    }
351
352    /// Get current resource usage
353    pub fn get_resource_usage(&self) -> ResourceUsage {
354        self.resource_tracker.get_usage()
355    }
356
357    /// Check if memory usage is within limits
358    pub fn check_memory_limit(&self, additional_mb: usize) -> bool {
359        let current_usage = self.resource_tracker.memory_usage.load(Ordering::Relaxed);
360        current_usage + additional_mb <= self.config.resource_limits.max_memory_mb
361    }
362
363    /// Update memory usage
364    pub fn update_memory_usage(&self, usage_mb: usize) {
365        self.resource_tracker
366            .memory_usage
367            .store(usage_mb, Ordering::Relaxed);
368    }
369}
370
371/// Connection permit that automatically releases when dropped
372pub struct ConnectionPermit {
373    _permit: tokio::sync::OwnedSemaphorePermit,
374    tracker: Arc<ResourceTracker>,
375}
376
377impl Drop for ConnectionPermit {
378    fn drop(&mut self) {
379        self.tracker.decrement_connections();
380    }
381}
382
383/// Request permit that automatically releases when dropped
384pub struct RequestPermit {
385    _permit: tokio::sync::OwnedSemaphorePermit,
386    tracker: Arc<ResourceTracker>,
387}
388
389impl Drop for RequestPermit {
390    fn drop(&mut self) {
391        self.tracker.decrement_requests();
392    }
393}
394
395/// Current resource usage snapshot
396#[derive(Debug, Clone)]
397pub struct ResourceUsage {
398    pub memory_mb: usize,
399    pub connection_count: usize,
400    pub request_count: usize,
401    pub uptime: Duration,
402}
403
404impl RateLimiter {
405    fn new(requests_per_second: f64, capacity: usize) -> Self {
406        Self {
407            capacity,
408            tokens: Arc::new(AtomicUsize::new(capacity)),
409            refill_rate: requests_per_second,
410            last_refill: Arc::new(RwLock::new(Instant::now())),
411        }
412    }
413
414    async fn allow_request(&self) -> bool {
415        self.refill_tokens().await;
416
417        let current_tokens = self.tokens.load(Ordering::Relaxed);
418        if current_tokens > 0 {
419            self.tokens.fetch_sub(1, Ordering::Relaxed);
420            true
421        } else {
422            false
423        }
424    }
425
426    async fn refill_tokens(&self) {
427        let now = Instant::now();
428        let mut last_refill = self.last_refill.write().await;
429
430        let elapsed = now.duration_since(*last_refill);
431        let tokens_to_add = (elapsed.as_secs_f64() * self.refill_rate) as usize;
432
433        if tokens_to_add > 0 {
434            let current_tokens = self.tokens.load(Ordering::Relaxed);
435            let new_tokens = std::cmp::min(current_tokens + tokens_to_add, self.capacity);
436            self.tokens.store(new_tokens, Ordering::Relaxed);
437            *last_refill = now;
438        }
439    }
440}
441
442impl CircuitBreaker {
443    fn new(failure_threshold: usize, timeout: Duration) -> Self {
444        Self {
445            state: Arc::new(RwLock::new(CircuitBreakerState::Closed)),
446            failure_count: Arc::new(AtomicUsize::new(0)),
447            failure_threshold,
448            timeout,
449            last_failure: Arc::new(RwLock::new(None)),
450        }
451    }
452
453    async fn check_request(&self) -> Result<()> {
454        let state = self.state.read().await.clone();
455
456        match state {
457            CircuitBreakerState::Closed => Ok(()),
458            CircuitBreakerState::Open => {
459                // Check if timeout has passed
460                let last_failure = self.last_failure.read().await;
461                if let Some(last_fail_time) = *last_failure {
462                    if last_fail_time.elapsed() >= self.timeout {
463                        // Transition to half-open
464                        drop(last_failure);
465                        *self.state.write().await = CircuitBreakerState::HalfOpen;
466                        Ok(())
467                    } else {
468                        Err(LspBridgeError::CircuitBreakerOpen(
469                            "Service is temporarily unavailable".to_string(),
470                        ))
471                    }
472                } else {
473                    Ok(())
474                }
475            }
476            CircuitBreakerState::HalfOpen => Ok(()),
477        }
478    }
479
480    async fn record_success(&self) {
481        self.failure_count.store(0, Ordering::Relaxed);
482        *self.state.write().await = CircuitBreakerState::Closed;
483    }
484
485    async fn record_failure(&self) {
486        let failures = self.failure_count.fetch_add(1, Ordering::Relaxed) + 1;
487        *self.last_failure.write().await = Some(Instant::now());
488
489        if failures >= self.failure_threshold {
490            *self.state.write().await = CircuitBreakerState::Open;
491        }
492    }
493}
494
495impl ResourceTracker {
496    fn new() -> Self {
497        Self {
498            memory_usage: Arc::new(AtomicUsize::new(0)),
499            connection_count: Arc::new(AtomicUsize::new(0)),
500            request_count: Arc::new(AtomicUsize::new(0)),
501            start_time: Instant::now(),
502        }
503    }
504
505    fn increment_connections(&self) {
506        self.connection_count.fetch_add(1, Ordering::Relaxed);
507    }
508
509    fn decrement_connections(&self) {
510        self.connection_count.fetch_sub(1, Ordering::Relaxed);
511    }
512
513    fn increment_requests(&self) {
514        self.request_count.fetch_add(1, Ordering::Relaxed);
515    }
516
517    fn decrement_requests(&self) {
518        self.request_count.fetch_sub(1, Ordering::Relaxed);
519    }
520
521    fn get_usage(&self) -> ResourceUsage {
522        ResourceUsage {
523            memory_mb: self.memory_usage.load(Ordering::Relaxed),
524            connection_count: self.connection_count.load(Ordering::Relaxed),
525            request_count: self.request_count.load(Ordering::Relaxed),
526            uptime: self.start_time.elapsed(),
527        }
528    }
529}
530
531impl InputValidator {
532    fn new(config: SecurityConfig) -> Self {
533        Self { config }
534    }
535
536    async fn validate(&self, input: &str) -> ValidationResult {
537        let mut errors = Vec::new();
538        let mut is_valid = true;
539
540        if !self.config.enable_input_validation {
541            return ValidationResult {
542                is_valid: true,
543                errors: Vec::new(),
544                sanitized_input: None,
545            };
546        }
547
548        // Check string length
549        if input.len() > self.config.max_string_length {
550            errors.push(format!(
551                "Input exceeds maximum length of {} characters",
552                self.config.max_string_length
553            ));
554            is_valid = false;
555        }
556
557        // Validate JSON structure if applicable
558        if let Ok(json_value) = serde_json::from_str::<serde_json::Value>(input) {
559            if Self::get_json_depth(&json_value) > self.config.max_json_depth {
560                errors.push(format!(
561                    "JSON structure exceeds maximum depth of {}",
562                    self.config.max_json_depth
563                ));
564                is_valid = false;
565            }
566        }
567
568        // Basic sanitization
569        let sanitized = if self.config.enable_output_sanitization {
570            Some(self.sanitize_string(input))
571        } else {
572            None
573        };
574
575        ValidationResult {
576            is_valid,
577            errors,
578            sanitized_input: sanitized,
579        }
580    }
581
582    fn get_json_depth(value: &serde_json::Value) -> usize {
583        match value {
584            serde_json::Value::Object(map) => {
585                1 + map.values().map(Self::get_json_depth).max().unwrap_or(0)
586            }
587            serde_json::Value::Array(arr) => {
588                1 + arr.iter().map(Self::get_json_depth).max().unwrap_or(0)
589            }
590            _ => 0,
591        }
592    }
593
594    fn sanitize_string(&self, input: &str) -> String {
595        // Basic sanitization - remove control characters
596        input
597            .chars()
598            .filter(|c| !c.is_control() || *c == '\n' || *c == '\r' || *c == '\t')
599            .collect()
600    }
601}
602
603impl Default for ResourceLimits {
604    fn default() -> Self {
605        Self {
606            max_memory_mb: 1024, // 1GB
607            max_connections: 100,
608            max_requests_per_connection: 50,
609            max_request_size_bytes: 10 * 1024 * 1024, // 10MB
610            max_response_size_bytes: 10 * 1024 * 1024, // 10MB
611            request_timeout: Duration::from_secs(30),
612            connection_timeout: Duration::from_secs(10),
613            max_retries: 3,
614        }
615    }
616}
617
618impl Default for RateLimits {
619    fn default() -> Self {
620        Self {
621            requests_per_second: 100.0,
622            burst_capacity: 200,
623            window_duration: Duration::from_secs(1),
624            adaptive: true,
625        }
626    }
627}
628
629impl Default for SecurityConfig {
630    fn default() -> Self {
631        Self {
632            enable_input_validation: true,
633            enable_output_sanitization: true,
634            max_json_depth: 10,
635            max_string_length: 1024 * 1024, // 1MB
636            enable_security_logging: true,
637        }
638    }
639}
640
641impl Default for ErrorRecoveryConfig {
642    fn default() -> Self {
643        Self {
644            enable_auto_retry: true,
645            retry_base_delay: Duration::from_millis(100),
646            retry_max_delay: Duration::from_secs(30),
647            enable_circuit_breaker: true,
648            circuit_breaker_threshold: 5,
649            circuit_breaker_timeout: Duration::from_secs(60),
650        }
651    }
652}
653
654impl Default for MonitoringConfig {
655    fn default() -> Self {
656        Self {
657            enable_detailed_metrics: true,
658            metrics_interval: Duration::from_secs(30),
659            enable_health_checks: true,
660            health_check_interval: Duration::from_secs(60),
661            enable_alerting: true,
662        }
663    }
664}
665
666#[cfg(test)]
667mod tests {
668    use super::*;
669
670    #[tokio::test]
671    async fn test_connection_limits() {
672        let config = HardeningConfig {
673            resource_limits: ResourceLimits {
674                max_connections: 2,
675                ..Default::default()
676            },
677            ..Default::default()
678        };
679
680        let manager = HardeningManager::new(config);
681
682        // Acquire first connection
683        let _permit1 = manager.acquire_connection_permit().await.unwrap();
684
685        // Acquire second connection
686        let _permit2 = manager.acquire_connection_permit().await.unwrap();
687
688        // Third connection should timeout quickly
689        let config_quick_timeout = HardeningConfig {
690            resource_limits: ResourceLimits {
691                max_connections: 2,
692                connection_timeout: Duration::from_millis(100),
693                ..Default::default()
694            },
695            ..Default::default()
696        };
697
698        let manager_quick = HardeningManager::new(config_quick_timeout);
699        let _permit1_quick = manager_quick.acquire_connection_permit().await.unwrap();
700        let _permit2_quick = manager_quick.acquire_connection_permit().await.unwrap();
701
702        // This should fail due to timeout
703        let result = manager_quick.acquire_connection_permit().await;
704        assert!(result.is_err());
705    }
706
707    #[tokio::test]
708    async fn test_rate_limiting() {
709        let rate_limiter = RateLimiter::new(2.0, 2); // 2 requests per second, burst of 2
710
711        // First two requests should succeed (burst)
712        assert!(rate_limiter.allow_request().await);
713        assert!(rate_limiter.allow_request().await);
714
715        // Third request should fail (rate limited)
716        assert!(!rate_limiter.allow_request().await);
717    }
718
719    #[tokio::test]
720    async fn test_circuit_breaker() {
721        let circuit_breaker = CircuitBreaker::new(2, Duration::from_millis(100));
722
723        // Normal operation
724        assert!(circuit_breaker.check_request().await.is_ok());
725
726        // Record failures
727        circuit_breaker.record_failure().await;
728        circuit_breaker.record_failure().await;
729
730        // Circuit should be open now
731        assert!(circuit_breaker.check_request().await.is_err());
732
733        // Wait for timeout
734        tokio::time::sleep(Duration::from_millis(150)).await;
735
736        // Should be half-open now
737        assert!(circuit_breaker.check_request().await.is_ok());
738
739        // Record success to close circuit
740        circuit_breaker.record_success().await;
741        assert!(circuit_breaker.check_request().await.is_ok());
742    }
743
744    #[tokio::test]
745    async fn test_input_validation() {
746        let config = SecurityConfig {
747            enable_input_validation: true,
748            max_string_length: 10,
749            max_json_depth: 2,
750            ..Default::default()
751        };
752
753        let validator = InputValidator::new(config);
754
755        // Valid input
756        let result = validator.validate("hello").await;
757        assert!(result.is_valid);
758
759        // Too long input
760        let result = validator.validate("this is too long").await;
761        assert!(!result.is_valid);
762        assert!(!result.errors.is_empty());
763    }
764}
lsp_bridge/hardening.rs

lsp_bridge/
hardening.rs
