Merge pull request #72 from adrgs/fix/security-audit-aisafe

fix: address aisafe.io security audit findings

Author: adrgs

Dockerfile

@@ -2,7 +2,7 @@
 # Uses cargo-chef for optimized dependency caching
 
 # Stage 1: Chef - install cargo-chef
-FROM rust:1.85-slim-bookworm AS chef
+FROM rust:1.88-slim-bookworm AS chef
 RUN apt-get update && apt-get install -y --no-install-recommends \
     libssl-dev \
     pkg-config \

src/src/cache/mod.rs

@@ -418,13 +418,13 @@ impl Cache {
         let mut empty_lists: Vec<String> = Vec::new();
 
         if let Ok(mut store) = self.request_store.write() {
+            // Collect keys ONCE outside the while loop to avoid O(N*M) re-cloning
+            let keys: Vec<String> = store.keys().cloned().collect();
+
             // Keep evicting until we've freed enough memory
             while freed < bytes_to_free {
                 let mut made_progress = false;
 
-                // Get all keys (we need to collect to avoid borrow issues)
-                let keys: Vec<String> = store.keys().cloned().collect();
-
                 // Round-robin: pop one oldest request from each list
                 for key in &keys {
                     if freed >= bytes_to_free {
@@ -490,10 +490,14 @@ impl Clone for Cache {
 /// Extract subdomain from cache key for size tracking
 fn extract_subdomain_from_key(key: &str) -> Option<String> {
     // files:{subdomain} -> subdomain
-    // dns:{subdomain} -> subdomain
+    // file:{subdomain}:{path} -> subdomain
     if key.starts_with("files:") || key.starts_with("file:") {
         return Some(key.split(':').nth(1)?.to_string());
     }
+    // request:{subdomain}:{id} -> subdomain
+    if key.starts_with("request:") {
+        return Some(key.split(':').nth(1)?.to_string());
+    }
     if key.starts_with("dns:") {
         // dns:{subdomain} or dns:{type}:{domain}
         let parts: Vec<&str> = key.split(':').collect();
@@ -502,13 +506,17 @@ fn extract_subdomain_from_key(key: &str) -> Option<String> {
         }
         // dns:{type}:{domain} - extract subdomain from domain
         // e.g., dns:A:test.abc123.example.com. -> abc123
+        // Strip trailing dot and base domain, then take the last remaining label
         if parts.len() >= 3 {
-            let domain = parts[2];
-            // Try to extract subdomain from domain name
-            let domain_parts: Vec<&str> = domain.split('.').collect();
-            if domain_parts.len() >= 2 {
-                // Second-to-last part before the base domain might be the subdomain
-                return Some(domain_parts[1].to_string());
+            let domain = parts[2].trim_end_matches('.');
+            if let Some(stripped) = domain.strip_suffix(&CONFIG.server_domain) {
+                let stripped = stripped.trim_end_matches('.');
+                // The last label in the remaining string is the subdomain
+                if let Some(sub) = stripped.rsplit('.').next() {
+                    if sub.len() == CONFIG.subdomain_length {
+                        return Some(sub.to_string());
+                    }
+                }
             }
         }
     }

src/src/dns/mod.rs

@@ -1,12 +1,14 @@
 use anyhow::{anyhow, Result};
 use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
 use rand::seq::SliceRandom;
-use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
+use std::collections::HashMap;
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
 use std::str::FromStr;
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
+use std::time::Instant;
 use tokio::net::UdpSocket;
 use tokio::sync::broadcast;
-use tracing::{error, info};
+use tracing::{error, info, warn};
 use trust_dns_proto::op::{Message, MessageType, OpCode, Query, ResponseCode};
 use trust_dns_proto::rr::rdata::{CNAME, MX, TXT};
 use trust_dns_proto::rr::{Name, RData, Record, RecordType};
@@ -18,6 +20,54 @@ use crate::models::{CacheMessage, DnsRequestLog};
 use crate::utils::config::CONFIG;
 use crate::utils::{generate_request_id, get_current_timestamp, get_subdomain_from_hostname};
 
+/// Maximum UDP DNS response size (RFC 1035)
+const MAX_UDP_RESPONSE_SIZE: usize = 512;
+
+/// Maximum DNS queries per second per IP
+const DNS_RATE_LIMIT_PER_SECOND: u32 = 100;
+
+/// Rate limiter cleanup interval (number of queries between cleanups)
+const RATE_LIMITER_CLEANUP_INTERVAL: u64 = 1000;
+
+struct DnsRateLimiter {
+    limits: Mutex<HashMap<IpAddr, (Instant, u32)>>,
+    max_per_second: u32,
+}
+
+impl DnsRateLimiter {
+    fn new(max_per_second: u32) -> Self {
+        Self {
+            limits: Mutex::new(HashMap::new()),
+            max_per_second,
+        }
+    }
+
+    /// Returns true if the request should be allowed, false if rate-limited.
+    fn check(&self, ip: IpAddr) -> bool {
+        let mut limits = self.limits.lock().unwrap_or_else(|e| e.into_inner());
+        let now = Instant::now();
+
+        let entry = limits.entry(ip).or_insert((now, 0));
+
+        // If more than 1 second has passed, reset the counter
+        if now.duration_since(entry.0).as_secs() >= 1 {
+            entry.0 = now;
+            entry.1 = 1;
+            true
+        } else {
+            entry.1 += 1;
+            entry.1 <= self.max_per_second
+        }
+    }
+
+    /// Remove stale entries older than 10 seconds
+    fn cleanup(&self) {
+        let mut limits = self.limits.lock().unwrap_or_else(|e| e.into_inner());
+        let now = Instant::now();
+        limits.retain(|_, (instant, _)| now.duration_since(*instant).as_secs() < 10);
+    }
+}
+
 pub struct Server {
     cache: Arc<Cache>,
     tx: Arc<broadcast::Sender<CacheMessage>>,
@@ -33,10 +83,24 @@ impl Server {
 
         let socket = Arc::new(UdpSocket::bind(format!("0.0.0.0:{}", CONFIG.dns_port)).await?);
         let mut buf = vec![0u8; 512];
+        let rate_limiter = Arc::new(DnsRateLimiter::new(DNS_RATE_LIMIT_PER_SECOND));
+        let mut query_count: u64 = 0;
 
         loop {
             match socket.recv_from(&mut buf).await {
                 Ok((len, addr)) => {
+                    // Rate limit check
+                    if !rate_limiter.check(addr.ip()) {
+                        warn!("DNS rate limit exceeded for {}", addr.ip());
+                        continue;
+                    }
+
+                    // Periodic cleanup of stale rate limiter entries
+                    query_count += 1;
+                    if query_count.is_multiple_of(RATE_LIMITER_CLEANUP_INTERVAL) {
+                        rate_limiter.cleanup();
+                    }
+
                     let data = buf[..len].to_vec();
                     let cache = self.cache.clone();
                     let tx = self.tx.clone();
@@ -97,6 +161,22 @@ async fn handle_dns_request(
     let response_bytes = response
         .to_bytes()
         .map_err(|e| anyhow!("Failed to serialize DNS response: {}", e))?;
+
+    // If UDP response exceeds 512 bytes, set TC (truncation) bit and strip answers
+    // This forces the client to retry via TCP (which is not spoofable)
+    let response_bytes = if response_bytes.len() > MAX_UDP_RESPONSE_SIZE {
+        let mut truncated = response.clone();
+        truncated.set_truncated(true);
+        // Remove all answers to fit within 512 bytes - keep only header + question
+        let empty_answers: Vec<Record> = Vec::new();
+        truncated.insert_answers(empty_answers);
+        truncated
+            .to_bytes()
+            .map_err(|e| anyhow!("Failed to serialize truncated DNS response: {}", e))?
+    } else {
+        response_bytes
+    };
+
     socket.send_to(&response_bytes, addr).await?;
 
     Ok(())
@@ -201,20 +281,9 @@ async fn log_dns_request(
 
     let request_json = serde_json::to_string(&request_log)?;
 
-    // Push request to list and get the new length to calculate the correct index
+    // Push request to list
     let list_key = format!("requests:{subdomain}");
-    let index = cache
-        .rpush(&list_key, &request_json)
-        .await?
-        .saturating_sub(1);
-
-    // Store the index for this request ID (used by delete endpoint)
-    cache
-        .set(
-            &format!("request:{subdomain}:{request_id}"),
-            &index.to_string(),
-        )
-        .await?;
+    cache.rpush(&list_key, &request_json).await?;
 
     let message = CacheMessage {
         cmd: "new_request".to_string(),

src/src/http/routes.rs

@@ -182,21 +182,9 @@ pub async fn catch_all(
 
         let request_json = serde_json::to_string(&request_log).unwrap_or_default();
 
-        // Push request to list and get the new length to calculate the correct index
+        // Push request to list
         let list_key = format!("requests:{subdomain}");
-        let index = match state.cache.rpush(&list_key, &request_json).await {
-            Ok(len) => len.saturating_sub(1), // Index is length - 1 (0-based)
-            Err(_) => 0,
-        };
-
-        // Store the index for this request ID (used by delete endpoint)
-        let _ = state
-            .cache
-            .set(
-                &format!("request:{subdomain}:{request_id}"),
-                &index.to_string(),
-            )
-            .await;
+        let _ = state.cache.rpush(&list_key, &request_json).await;
 
         let message = crate::models::CacheMessage {
             cmd: "new_request".to_string(),
@@ -368,6 +356,11 @@ async fn serve_file(
             }
         }
 
+        // Security headers: prevent XSS via served content
+        // Applied after user headers so they cannot be overridden
+        response = response.header("X-Content-Type-Options", "nosniff");
+        response = response.header("Content-Security-Policy", "sandbox allow-scripts");
+
         return response
             .body(Body::from(content))
             .unwrap_or_else(|_| {