use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use bytes::BytesMut;
use parking_lot::RwLock;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::time;
use crate::command::handle_command;
use crate::config::Config;
use crate::protocol::{encode_frame, parse_frame, RespFrame};
use crate::store::Store;
#[derive(Clone)]
struct AuthAttempt {
count: u32,
locked_until: Option<Instant>,
}
struct AuthRateLimiter {
attempts: Arc<RwLock<HashMap<String, AuthAttempt>>>,
}
impl AuthRateLimiter {
fn new() -> Self {
Self {
attempts: Arc::new(RwLock::new(HashMap::new())),
}
}
fn check_locked(&self, ip: &str, _lockout_duration: Duration) -> Option<Duration> {
let mut map = self.attempts.write();
if let Some(attempt) = map.get(ip) {
if let Some(locked_until) = attempt.locked_until {
if Instant::now() < locked_until {
return Some(locked_until.saturating_duration_since(Instant::now()));
} else {
map.remove(ip);
}
}
}
None
}
fn record_failure(&self, ip: &str, max_attempts: u32, lockout_duration: Duration) -> bool {
let mut map = self.attempts.write();
let attempt = map.entry(ip.to_string()).or_insert_with(|| AuthAttempt {
count: 0,
locked_until: None,
});
attempt.count += 1;
if attempt.count >= max_attempts {
attempt.locked_until = Some(Instant::now() + lockout_duration);
true } else {
false }
}
fn record_success(&self, ip: &str) {
let mut map = self.attempts.write();
map.remove(ip);
}
fn cleanup_old_entries(&self) {
let mut map = self.attempts.write();
let now = Instant::now();
map.retain(|_, attempt| {
if let Some(locked_until) = attempt.locked_until {
now < locked_until
} else {
true }
});
}
}
pub async fn run_server(
addr: SocketAddr,
store: Arc<Store>,
read_buffer_size: usize,
config: Arc<Config>,
) -> anyhow::Result<()> {
let rate_limiter = Arc::new(AuthRateLimiter::new());
let cleanup_limiter = rate_limiter.clone();
tokio::spawn(async move {
let mut interval = time::interval(Duration::from_secs(60));
loop {
interval.tick().await;
cleanup_limiter.cleanup_old_entries();
}
});
let listener = TcpListener::bind(addr).await?;
loop {
let (socket, peer_addr) = listener.accept().await?;
let store = store.clone();
let read_buffer_size = read_buffer_size;
let config = config.clone();
let rate_limiter = rate_limiter.clone();
tokio::spawn(async move {
if let Err(e) = handle_connection(socket, peer_addr, store, read_buffer_size, config, rate_limiter).await {
eprintln!("connection error: {e}");
}
});
}
}
async fn handle_connection(
mut stream: TcpStream,
peer_addr: SocketAddr,
store: Arc<Store>,
read_buffer_size: usize,
config: Arc<Config>,
rate_limiter: Arc<AuthRateLimiter>,
) -> anyhow::Result<()> {
let mut buf = BytesMut::with_capacity(read_buffer_size * 2);
let mut authenticated = !config.requires_auth();
let client_ip = peer_addr.ip().to_string();
loop {
let mut temp = vec![0u8; read_buffer_size];
let n = stream.read(&mut temp).await?;
if n == 0 {
break;
}
buf.extend_from_slice(&temp[..n]);
loop {
match parse_frame(&mut buf) {
Ok(Some(frame)) => {
if let RespFrame::Array(Some(items)) = &frame {
if !items.is_empty() {
let cmd_name = match &items[0] {
RespFrame::BulkString(Some(b)) => {
String::from_utf8_lossy(b).to_ascii_uppercase()
}
RespFrame::SimpleString(s) => s.to_ascii_uppercase(),
_ => String::new(),
};
if cmd_name == "AUTH" {
let resp = handle_auth(
&config,
&items,
&mut authenticated,
&client_ip,
&rate_limiter,
);
let encoded = encode_frame(&resp);
stream.write_all(&encoded).await?;
continue;
}
}
}
if !authenticated {
let err = RespFrame::Error("NOAUTH Authentication required.".into());
let encoded = encode_frame(&err);
stream.write_all(&encoded).await?;
continue;
}
let resp = handle_command(store.clone(), frame).await;
let encoded = encode_frame(&resp);
stream.write_all(&encoded).await?;
}
Ok(None) => break,
Err(e) => {
let err = RespFrame::Error(format!("ERR {}", e));
let encoded = encode_frame(&err);
stream.write_all(&encoded).await?;
return Ok(());
}
}
}
}
Ok(())
}
fn handle_auth(
config: &Config,
items: &[RespFrame],
authenticated: &mut bool,
client_ip: &str,
rate_limiter: &Arc<AuthRateLimiter>,
) -> RespFrame {
if !config.requires_auth() {
return RespFrame::Error("ERR Client sent AUTH, but no password is set".into());
}
let lockout_duration = Duration::from_secs(config.lockout_duration_seconds);
if let Some(remaining) = rate_limiter.check_locked(client_ip, lockout_duration) {
let minutes = remaining.as_secs() / 60;
let seconds = remaining.as_secs() % 60;
return RespFrame::Error(format!(
"ERR Too many failed login attempts. Please try again in {}m {}s.",
minutes, seconds
));
}
let auth_result = if items.len() == 2 {
let password = match &items[1] {
RespFrame::BulkString(Some(b)) => String::from_utf8_lossy(b).to_string(),
RespFrame::SimpleString(s) => s.clone(),
_ => return RespFrame::Error("ERR invalid password".into()),
};
config.check_auth(None, Some(&password))
} else if items.len() == 3 {
let username = match &items[1] {
RespFrame::BulkString(Some(b)) => String::from_utf8_lossy(b).to_string(),
RespFrame::SimpleString(s) => s.clone(),
_ => return RespFrame::Error("ERR invalid username".into()),
};
let password = match &items[2] {
RespFrame::BulkString(Some(b)) => String::from_utf8_lossy(b).to_string(),
RespFrame::SimpleString(s) => s.clone(),
_ => return RespFrame::Error("ERR invalid password".into()),
};
config.check_auth(Some(&username), Some(&password))
} else {
return RespFrame::Error("ERR wrong number of arguments for 'auth' command".into());
};
if auth_result {
rate_limiter.record_success(client_ip);
*authenticated = true;
RespFrame::SimpleString("OK".into())
} else {
let is_locked = rate_limiter.record_failure(client_ip, config.max_failed_attempts, lockout_duration);
if is_locked {
let minutes = lockout_duration.as_secs() / 60;
RespFrame::Error(format!(
"ERR Too many failed login attempts. Account locked for {} minutes.",
minutes
))
} else {
RespFrame::Error("ERR invalid password".into())
}
}
}