use bytes::Bytes;
use http_body_util::Full;
use hyper::{Method, Request, Response, StatusCode, server::conn::http1};
use hyper_util::rt::TokioIo;
use log::{debug, error, info, warn};
use serde::{Deserialize, Serialize};
use std::convert::Infallible;
use std::sync::Arc;
use tokio::net::TcpListener;
use tokio::sync::Semaphore;
use crate::cache::SyncDnsCache;
use crate::stats::SharedStats;
#[derive(Serialize, Deserialize)]
struct ApiResponse {
success: bool,
message: String,
}
#[derive(Serialize, Deserialize)]
struct CacheStatusResponse {
success: bool,
size: usize,
capacity: usize,
hit_rate: f64,
hits: u64,
misses: u64,
}
fn is_domain_in_zone(domain: &str, zone: &str) -> bool {
let normalized_domain = domain.trim_end_matches('.').to_lowercase();
let normalized_zone = zone.trim_end_matches('.').to_lowercase();
if normalized_zone.is_empty() {
return true;
}
if normalized_domain == normalized_zone {
return true;
}
normalized_domain.ends_with(&format!(".{normalized_zone}"))
}
async fn handle_request(
req: Request<hyper::body::Incoming>,
control_path: String,
dns_cache: Option<Arc<SyncDnsCache>>,
stats: Option<Arc<SharedStats>>,
) -> Result<Response<Full<Bytes>>, Infallible> {
let path = req.uri().path();
let method = req.method();
let is_control_path = path == control_path
|| path
.strip_prefix(&control_path)
.map(|suffix| control_path.ends_with('/') || suffix.starts_with('/'))
.unwrap_or(false);
if !is_control_path {
let mut response = Response::new(Full::new(Bytes::from("Not Found")));
*response.status_mut() = StatusCode::NOT_FOUND;
return Ok(response);
}
let endpoint = path
.strip_prefix(&control_path)
.unwrap_or("")
.trim_start_matches('/');
match (method, endpoint) {
(&Method::GET, "status") => {
let response = ApiResponse {
success: true,
message: "Server is running".to_string(),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
Ok(response)
}
(&Method::GET, "cache") => {
if let Some(cache) = dns_cache {
let size = cache.len();
let capacity = cache.capacity();
let (hits, misses, hit_rate) = if let Some(stats_ref) = &stats {
let global_stats = stats_ref.get_stats().await;
let hits = global_stats.cache_hits;
let misses = global_stats.cache_misses;
let hit_rate = if hits + misses > 0 {
hits as f64 / (hits + misses) as f64
} else {
0.0
};
(hits, misses, hit_rate)
} else {
let hits = 0;
let misses = 0;
let hit_rate = 0.0;
(hits, misses, hit_rate)
};
let response = CacheStatusResponse {
success: true,
size,
capacity,
hit_rate,
hits,
misses,
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
Ok(response)
} else {
let response = ApiResponse {
success: false,
message: "DNS cache not available".to_string(),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
*response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
Ok(response)
}
}
(&Method::DELETE, "cache") => {
if let Some(cache) = dns_cache {
cache.clear();
let cache_size = cache.len();
info!("DNS cache cleared, current size: {cache_size}");
let response = ApiResponse {
success: true,
message: format!("Cache cleared successfully. Current size: {cache_size}"),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
Ok(response)
} else {
let response = ApiResponse {
success: false,
message: "DNS cache not available".to_string(),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
*response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
Ok(response)
}
}
(&Method::DELETE, zone_path) if zone_path.starts_with("cache/zone/") => {
if let Some(cache) = dns_cache {
let zone = zone_path.strip_prefix("cache/zone/").unwrap_or("");
if zone.is_empty() {
let response = ApiResponse {
success: false,
message: "Zone parameter is required".to_string(),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
*response.status_mut() = StatusCode::BAD_REQUEST;
return Ok(response);
}
let normalized_zone = match crate::dns_key::DNSKey::normalize_name(zone) {
Ok(zone) => zone,
Err(e) => {
let mut response = Response::new(Full::new(Bytes::from(format!(
"Invalid domain name: {e}"
))));
*response.status_mut() = StatusCode::BAD_REQUEST;
return Ok(response);
}
};
let (removed_entries, new_size) = cache.with_lock(|inner| {
let initial_size = inner.len();
inner.retain(|key, _| !is_domain_in_zone(&key.name, &normalized_zone));
let new_size = inner.len();
(initial_size.saturating_sub(new_size), new_size)
});
info!("Cleared {removed_entries} entries for zone {zone} from DNS cache");
let response = ApiResponse {
success: true,
message: format!(
"Cleared {removed_entries} entries for zone {zone} from DNS cache. New size: {new_size}"
),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
Ok(response)
} else {
let response = ApiResponse {
success: false,
message: "DNS cache not available".to_string(),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
*response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
Ok(response)
}
}
(&Method::DELETE, name_path) if name_path.starts_with("cache/name/") => {
if let Some(cache) = dns_cache {
let name = name_path.strip_prefix("cache/name/").unwrap_or("");
if name.is_empty() {
let response = ApiResponse {
success: false,
message: "Name parameter is required".to_string(),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
*response.status_mut() = StatusCode::BAD_REQUEST;
return Ok(response);
}
let normalized_name = match crate::dns_key::DNSKey::normalize_name(name) {
Ok(name) => name,
Err(e) => {
error!("Failed to normalize domain name: {}", e);
let response = ApiResponse {
success: false,
message: format!("Invalid domain name: {}", e),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {}", e);
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
*response.status_mut() = StatusCode::BAD_REQUEST;
return Ok(response);
}
};
let (removed_entries, new_size) = cache.with_lock(|inner| {
let initial_size = inner.len();
inner.retain(|key, _| key.name != normalized_name);
let new_size = inner.len();
(initial_size.saturating_sub(new_size), new_size)
});
info!("Cleared {removed_entries} entries for name {name} from DNS cache");
let response = ApiResponse {
success: true,
message: format!(
"Cleared {removed_entries} entries for name {name} from DNS cache. New size: {new_size}"
),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
Ok(response)
} else {
let response = ApiResponse {
success: false,
message: "DNS cache not available".to_string(),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
*response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
Ok(response)
}
}
_ => {
let response = ApiResponse {
success: false,
message: "Endpoint not found".to_string(),
};
let json = serde_json::to_string(&response).unwrap_or_else(|e| {
error!("Failed to serialize API response: {e}");
r#"{"success":false,"message":"Internal server error"}"#.to_string()
});
let mut response = Response::new(Full::new(Bytes::from(json)));
response.headers_mut().insert(
hyper::header::CONTENT_TYPE,
hyper::header::HeaderValue::from_static("application/json"),
);
*response.status_mut() = StatusCode::NOT_FOUND;
Ok(response)
}
}
}
pub async fn start_control_server(
listener: TcpListener,
control_path: String,
max_connections: usize,
dns_cache: Option<Arc<SyncDnsCache>>,
stats: Option<Arc<SharedStats>>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let addr = listener.local_addr()?;
info!("Control server listening on {addr}, base path: {control_path}");
if dns_cache.is_some() {
info!(
"DNS cache control API enabled at {control_path}/cache (GET/DELETE), {control_path}/cache/zone/<zone> (DELETE), and {control_path}/cache/name/<name> (DELETE)"
);
} else {
warn!("DNS cache not available for control API");
}
let semaphore = Arc::new(Semaphore::new(max_connections));
loop {
let (stream, client_addr) = crate::net::accept_with_retry(&listener, "control").await;
let io = TokioIo::new(stream);
let control_path = control_path.clone();
let semaphore = semaphore.clone();
let dns_cache = dns_cache.clone();
let stats = stats.clone();
tokio::spawn(async move {
let _permit = match semaphore.try_acquire() {
Ok(permit) => permit,
Err(_) => {
warn!("Too many control connections, rejecting connection from {client_addr}");
return;
}
};
debug!("Accepted control connection from {client_addr}");
let service = hyper::service::service_fn(move |req| {
let control_path = control_path.clone();
let dns_cache = dns_cache.clone();
let stats = stats.clone();
async move { handle_request(req, control_path, dns_cache, stats).await }
});
if let Ok(Err(err)) = tokio::time::timeout(
std::time::Duration::from_secs(30),
http1::Builder::new().serve_connection(io, service),
)
.await
{
error!("Error serving control connection from {client_addr}: {err}");
}
debug!("Closed control connection from {client_addr}");
});
}
}
#[cfg(test)]
mod tests {
use super::is_domain_in_zone;
#[test]
fn zone_matching_normalizes_case_and_trailing_dots() {
assert!(is_domain_in_zone("WWW.Example.COM.", "example.com."));
assert!(is_domain_in_zone("example.com", "EXAMPLE.COM."));
assert!(!is_domain_in_zone("notexample.com", "example.com."));
assert!(is_domain_in_zone("www.example.com", "."));
}
}