Skip to content

Implement authentication rate limiting (fail2ban-like) #140

@inureyes

Description

@inureyes

Summary

Implement authentication rate limiting to prevent brute-force attacks, similar to fail2ban functionality. This reuses the existing rate limiter from the client code.

Parent Epic

Implementation Details

1. Authentication Rate Limiter

// src/server/security/rate_limit.rs
use crate::shared::rate_limit::{RateLimiter, RateLimitConfig};
use std::collections::HashMap;
use std::net::IpAddr;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;

/// Authentication rate limiter with ban support
pub struct AuthRateLimiter {
    /// Failed attempt counter per IP
    failures: RwLock<HashMap<IpAddr, FailureRecord>>,
    /// Banned IPs with expiration
    bans: RwLock<HashMap<IpAddr, Instant>>,
    /// Configuration
    config: AuthRateLimitConfig,
}

#[derive(Debug, Clone)]
pub struct AuthRateLimitConfig {
    /// Maximum failed attempts before ban
    pub max_attempts: u32,
    /// Time window for counting attempts (seconds)
    pub window: Duration,
    /// Ban duration
    pub ban_duration: Duration,
    /// Whitelist IPs (never banned)
    pub whitelist: Vec<IpAddr>,
}

impl Default for AuthRateLimitConfig {
    fn default() -> Self {
        Self {
            max_attempts: 5,
            window: Duration::from_secs(300),
            ban_duration: Duration::from_secs(300),
            whitelist: vec![],
        }
    }
}

#[derive(Debug)]
struct FailureRecord {
    count: u32,
    first_failure: Instant,
    last_failure: Instant,
}

impl AuthRateLimiter {
    pub fn new(config: AuthRateLimitConfig) -> Self {
        Self {
            failures: RwLock::new(HashMap::new()),
            bans: RwLock::new(HashMap::new()),
            config,
        }
    }

    /// Check if IP is banned
    pub async fn is_banned(&self, ip: &IpAddr) -> bool {
        // Whitelisted IPs are never banned
        if self.config.whitelist.contains(ip) {
            return false;
        }

        let bans = self.bans.read().await;
        if let Some(expiry) = bans.get(ip) {
            if Instant::now() < *expiry {
                return true;
            }
        }
        false
    }

    /// Record a failed authentication attempt
    pub async fn record_failure(&self, ip: IpAddr) -> bool {
        // Skip whitelisted IPs
        if self.config.whitelist.contains(&ip) {
            return false;
        }

        let mut failures = self.failures.write().await;
        let now = Instant::now();

        let record = failures.entry(ip).or_insert(FailureRecord {
            count: 0,
            first_failure: now,
            last_failure: now,
        });

        // Reset if window expired
        if now.duration_since(record.first_failure) > self.config.window {
            record.count = 1;
            record.first_failure = now;
        } else {
            record.count += 1;
        }
        record.last_failure = now;

        // Check if should ban
        if record.count >= self.config.max_attempts {
            drop(failures); // Release lock before acquiring ban lock
            self.ban(ip).await;
            return true;
        }

        false
    }

    /// Record successful authentication (reset failures)
    pub async fn record_success(&self, ip: &IpAddr) {
        let mut failures = self.failures.write().await;
        failures.remove(ip);
    }

    /// Ban an IP address
    pub async fn ban(&self, ip: IpAddr) {
        tracing::warn!("Banning IP {} for {:?}", ip, self.config.ban_duration);
        
        let mut bans = self.bans.write().await;
        let expiry = Instant::now() + self.config.ban_duration;
        bans.insert(ip, expiry);

        // Clean up failure record
        let mut failures = self.failures.write().await;
        failures.remove(&ip);
    }

    /// Manually unban an IP
    pub async fn unban(&self, ip: &IpAddr) {
        let mut bans = self.bans.write().await;
        bans.remove(ip);
        tracing::info!("Unbanned IP {}", ip);
    }

    /// Get remaining attempts before ban
    pub async fn remaining_attempts(&self, ip: &IpAddr) -> u32 {
        let failures = self.failures.read().await;
        if let Some(record) = failures.get(ip) {
            self.config.max_attempts.saturating_sub(record.count)
        } else {
            self.config.max_attempts
        }
    }

    /// Clean up expired records (call periodically)
    pub async fn cleanup(&self) {
        let now = Instant::now();

        // Clean expired bans
        {
            let mut bans = self.bans.write().await;
            bans.retain(|_, expiry| now < *expiry);
        }

        // Clean old failure records
        {
            let mut failures = self.failures.write().await;
            failures.retain(|_, record| {
                now.duration_since(record.last_failure) < self.config.window
            });
        }
    }

    /// Get current ban list
    pub async fn get_bans(&self) -> Vec<(IpAddr, Duration)> {
        let now = Instant::now();
        let bans = self.bans.read().await;
        bans.iter()
            .filter_map(|(ip, expiry)| {
                if now < *expiry {
                    Some((*ip, *expiry - now))
                } else {
                    None
                }
            })
            .collect()
    }
}

2. Integrate with SSH Handler

// Update src/server/handler.rs
impl Handler for SshHandler {
    async fn auth_password(&mut self, user: &str, password: &str) -> Result<Auth, Self::Error> {
        if let Some(ip) = self.peer_addr.map(|a| a.ip()) {
            // Check if banned
            if self.config.security.rate_limiter.is_banned(&ip).await {
                tracing::warn!("Rejected auth from banned IP: {}", ip);
                return Ok(Auth::Reject { proceed_with_methods: None });
            }
        }

        // ... authentication logic ...

        // On failure
        if let Some(ip) = self.peer_addr.map(|a| a.ip()) {
            let banned = self.config.security.rate_limiter.record_failure(ip).await;
            if banned {
                // Log audit event
                self.audit_manager.log(AuditEvent::new(
                    EventType::IpBlocked,
                    user.to_string(),
                    self.session_id.clone(),
                ).with_client_ip(ip)).await;
            }
        }

        Ok(Auth::Reject { proceed_with_methods: None })
    }
}

3. Background Cleanup Task

// Start periodic cleanup
tokio::spawn(async move {
    let mut interval = tokio::time::interval(Duration::from_secs(60));
    loop {
        interval.tick().await;
        rate_limiter.cleanup().await;
    }
});

Configuration

security:
  max_auth_attempts: 5
  auth_window: 300        # seconds
  ban_time: 300           # seconds
  whitelist_ips:
    - 127.0.0.1
    - 10.0.0.0/8          # CIDR support (future)

Files to Create/Modify

File Action
src/server/security/mod.rs Create - Security module
src/server/security/rate_limit.rs Create - Auth rate limiter
src/server/handler.rs Modify - Integrate rate limiting
src/server/config/types.rs Modify - Add security config

Testing Requirements

  1. Unit test: Failure counting
  2. Unit test: Ban after max attempts
  3. Unit test: Ban expiration
  4. Unit test: Whitelist IPs
  5. Unit test: Success resets failures
  6. Integration test: Multiple failed auths trigger ban

Acceptance Criteria

  • AuthRateLimiter implementation
  • Configurable max attempts and window
  • Configurable ban duration
  • IP whitelist support
  • Automatic cleanup of expired records
  • Integration with auth handlers
  • Audit logging for bans
  • Tests passing

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions