Skip to content

Commit 2244617

Browse files
adrgsclaude
andcommitted
fix: address security audit findings from aisafe.io
Fix 15 security issues identified in the aisafe.io audit: Critical/High: - Safe UTF-8 byte slicing in subdomain extraction (RQS-015) - JWT type confusion: deny_unknown_fields on Claims (RQS-004) - Cache eviction O(N*M) loop: collect keys once outside loop (RQS-008) - Cross-tenant quota bypass: fix subdomain extraction for DNS keys (RQS-005) - Cache leak: recognize request: keys in extract_subdomain_from_key (RQS-009) - Request deletion index desync: scan list by ID instead (RQS-007) - DNS amplification: TXT record 512-char limit + UDP truncation bit (RQS-006) - TCP exhaustion: semaphore (100 conns) + 30s read timeout (RQS-013) - SMTP memory exhaustion: 64KB log cap, 500 cmd limit, 100 rcpt limit (RQS-003) - WebSocket memory burst: limit initial history to last 100 requests (RQS-014) - XSS mitigation: sandbox CSP + nosniff on served files (RQS-002) Low/Info: - Constant-time admin token comparison (RQS-001) - DNS per-IP rate limiting at 100 qps (RQS-012) - WebSocket message rate limiting at 30 msg/s (RQS-011) - Share token endpoint rate limiting (RQS-010) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent fb39b70 commit 2244617

10 files changed

Lines changed: 322 additions & 70 deletions

File tree

src/src/cache/mod.rs

Lines changed: 18 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -418,13 +418,13 @@ impl Cache {
418418
let mut empty_lists: Vec<String> = Vec::new();
419419

420420
if let Ok(mut store) = self.request_store.write() {
421+
// Collect keys ONCE outside the while loop to avoid O(N*M) re-cloning
422+
let keys: Vec<String> = store.keys().cloned().collect();
423+
421424
// Keep evicting until we've freed enough memory
422425
while freed < bytes_to_free {
423426
let mut made_progress = false;
424427

425-
// Get all keys (we need to collect to avoid borrow issues)
426-
let keys: Vec<String> = store.keys().cloned().collect();
427-
428428
// Round-robin: pop one oldest request from each list
429429
for key in &keys {
430430
if freed >= bytes_to_free {
@@ -490,10 +490,14 @@ impl Clone for Cache {
490490
/// Extract subdomain from cache key for size tracking
491491
fn extract_subdomain_from_key(key: &str) -> Option<String> {
492492
// files:{subdomain} -> subdomain
493-
// dns:{subdomain} -> subdomain
493+
// file:{subdomain}:{path} -> subdomain
494494
if key.starts_with("files:") || key.starts_with("file:") {
495495
return Some(key.split(':').nth(1)?.to_string());
496496
}
497+
// request:{subdomain}:{id} -> subdomain
498+
if key.starts_with("request:") {
499+
return Some(key.split(':').nth(1)?.to_string());
500+
}
497501
if key.starts_with("dns:") {
498502
// dns:{subdomain} or dns:{type}:{domain}
499503
let parts: Vec<&str> = key.split(':').collect();
@@ -502,13 +506,17 @@ fn extract_subdomain_from_key(key: &str) -> Option<String> {
502506
}
503507
// dns:{type}:{domain} - extract subdomain from domain
504508
// e.g., dns:A:test.abc123.example.com. -> abc123
509+
// Strip trailing dot and base domain, then take the last remaining label
505510
if parts.len() >= 3 {
506-
let domain = parts[2];
507-
// Try to extract subdomain from domain name
508-
let domain_parts: Vec<&str> = domain.split('.').collect();
509-
if domain_parts.len() >= 2 {
510-
// Second-to-last part before the base domain might be the subdomain
511-
return Some(domain_parts[1].to_string());
511+
let domain = parts[2].trim_end_matches('.');
512+
if let Some(stripped) = domain.strip_suffix(&CONFIG.server_domain) {
513+
let stripped = stripped.trim_end_matches('.');
514+
// The last label in the remaining string is the subdomain
515+
if let Some(sub) = stripped.rsplit('.').next() {
516+
if sub.len() == CONFIG.subdomain_length {
517+
return Some(sub.to_string());
518+
}
519+
}
512520
}
513521
}
514522
}

src/src/dns/mod.rs

Lines changed: 86 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
use anyhow::{anyhow, Result};
22
use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
33
use rand::seq::SliceRandom;
4-
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
4+
use std::collections::HashMap;
5+
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
56
use std::str::FromStr;
6-
use std::sync::Arc;
7+
use std::sync::{Arc, Mutex};
8+
use std::time::Instant;
79
use tokio::net::UdpSocket;
810
use tokio::sync::broadcast;
9-
use tracing::{error, info};
11+
use tracing::{error, info, warn};
1012
use trust_dns_proto::op::{Message, MessageType, OpCode, Query, ResponseCode};
1113
use trust_dns_proto::rr::rdata::{CNAME, MX, TXT};
1214
use trust_dns_proto::rr::{Name, RData, Record, RecordType};
@@ -18,6 +20,54 @@ use crate::models::{CacheMessage, DnsRequestLog};
1820
use crate::utils::config::CONFIG;
1921
use crate::utils::{generate_request_id, get_current_timestamp, get_subdomain_from_hostname};
2022

23+
/// Maximum UDP DNS response size (RFC 1035)
24+
const MAX_UDP_RESPONSE_SIZE: usize = 512;
25+
26+
/// Maximum DNS queries per second per IP
27+
const DNS_RATE_LIMIT_PER_SECOND: u32 = 100;
28+
29+
/// Rate limiter cleanup interval (number of queries between cleanups)
30+
const RATE_LIMITER_CLEANUP_INTERVAL: u64 = 1000;
31+
32+
struct DnsRateLimiter {
33+
limits: Mutex<HashMap<IpAddr, (Instant, u32)>>,
34+
max_per_second: u32,
35+
}
36+
37+
impl DnsRateLimiter {
38+
fn new(max_per_second: u32) -> Self {
39+
Self {
40+
limits: Mutex::new(HashMap::new()),
41+
max_per_second,
42+
}
43+
}
44+
45+
/// Returns true if the request should be allowed, false if rate-limited.
46+
fn check(&self, ip: IpAddr) -> bool {
47+
let mut limits = self.limits.lock().unwrap_or_else(|e| e.into_inner());
48+
let now = Instant::now();
49+
50+
let entry = limits.entry(ip).or_insert((now, 0));
51+
52+
// If more than 1 second has passed, reset the counter
53+
if now.duration_since(entry.0).as_secs() >= 1 {
54+
entry.0 = now;
55+
entry.1 = 1;
56+
true
57+
} else {
58+
entry.1 += 1;
59+
entry.1 <= self.max_per_second
60+
}
61+
}
62+
63+
/// Remove stale entries older than 10 seconds
64+
fn cleanup(&self) {
65+
let mut limits = self.limits.lock().unwrap_or_else(|e| e.into_inner());
66+
let now = Instant::now();
67+
limits.retain(|_, (instant, _)| now.duration_since(*instant).as_secs() < 10);
68+
}
69+
}
70+
2171
pub struct Server {
2272
cache: Arc<Cache>,
2373
tx: Arc<broadcast::Sender<CacheMessage>>,
@@ -33,10 +83,27 @@ impl Server {
3383

3484
let socket = Arc::new(UdpSocket::bind(format!("0.0.0.0:{}", CONFIG.dns_port)).await?);
3585
let mut buf = vec![0u8; 512];
86+
let rate_limiter = Arc::new(DnsRateLimiter::new(DNS_RATE_LIMIT_PER_SECOND));
87+
let mut query_count: u64 = 0;
3688

3789
loop {
3890
match socket.recv_from(&mut buf).await {
3991
Ok((len, addr)) => {
92+
// Rate limit check
93+
if !rate_limiter.check(addr.ip()) {
94+
warn!("DNS rate limit exceeded for {}", addr.ip());
95+
continue;
96+
}
97+
98+
// Periodic cleanup of stale rate limiter entries
99+
query_count += 1;
100+
if query_count % RATE_LIMITER_CLEANUP_INTERVAL == 0 {
101+
let rl = rate_limiter.clone();
102+
tokio::spawn(async move {
103+
rl.cleanup();
104+
});
105+
}
106+
40107
let data = buf[..len].to_vec();
41108
let cache = self.cache.clone();
42109
let tx = self.tx.clone();
@@ -97,6 +164,22 @@ async fn handle_dns_request(
97164
let response_bytes = response
98165
.to_bytes()
99166
.map_err(|e| anyhow!("Failed to serialize DNS response: {}", e))?;
167+
168+
// If UDP response exceeds 512 bytes, set TC (truncation) bit and strip answers
169+
// This forces the client to retry via TCP (which is not spoofable)
170+
let response_bytes = if response_bytes.len() > MAX_UDP_RESPONSE_SIZE {
171+
let mut truncated = response.clone();
172+
truncated.set_truncated(true);
173+
// Remove all answers to fit within 512 bytes - keep only header + question
174+
let empty_answers: Vec<Record> = Vec::new();
175+
truncated.insert_answers(empty_answers);
176+
truncated
177+
.to_bytes()
178+
.map_err(|e| anyhow!("Failed to serialize truncated DNS response: {}", e))?
179+
} else {
180+
response_bytes
181+
};
182+
100183
socket.send_to(&response_bytes, addr).await?;
101184

102185
Ok(())

src/src/http/routes.rs

Lines changed: 7 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -182,21 +182,9 @@ pub async fn catch_all(
182182

183183
let request_json = serde_json::to_string(&request_log).unwrap_or_default();
184184

185-
// Push request to list and get the new length to calculate the correct index
185+
// Push request to list
186186
let list_key = format!("requests:{subdomain}");
187-
let index = match state.cache.rpush(&list_key, &request_json).await {
188-
Ok(len) => len.saturating_sub(1), // Index is length - 1 (0-based)
189-
Err(_) => 0,
190-
};
191-
192-
// Store the index for this request ID (used by delete endpoint)
193-
let _ = state
194-
.cache
195-
.set(
196-
&format!("request:{subdomain}:{request_id}"),
197-
&index.to_string(),
198-
)
199-
.await;
187+
let _ = state.cache.rpush(&list_key, &request_json).await;
200188

201189
let message = crate::models::CacheMessage {
202190
cmd: "new_request".to_string(),
@@ -368,6 +356,11 @@ async fn serve_file(
368356
}
369357
}
370358

359+
// Security headers: prevent XSS via served content
360+
// Applied after user headers so they cannot be overridden
361+
response = response.header("X-Content-Type-Options", "nosniff");
362+
response = response.header("Content-Security-Policy", "sandbox allow-scripts");
363+
371364
return response
372365
.body(Body::from(content))
373366
.unwrap_or_else(|_| {

0 commit comments

Comments
 (0)