rust_cache_server/
server.rs

1use std::collections::HashMap;
2use std::net::SocketAddr;
3use std::sync::Arc;
4use std::time::{Duration, Instant};
5
6use bytes::BytesMut;
7use parking_lot::RwLock;
8use tokio::io::{AsyncReadExt, AsyncWriteExt};
9use tokio::net::{TcpListener, TcpStream};
10use tokio::time;
11
12use crate::command::handle_command;
13use crate::config::Config;
14use crate::protocol::{encode_frame, parse_frame, RespFrame};
15use crate::store::Store;
16
17#[derive(Clone)]
18struct AuthAttempt {
19    count: u32,
20    locked_until: Option<Instant>,
21}
22
23struct AuthRateLimiter {
24    attempts: Arc<RwLock<HashMap<String, AuthAttempt>>>,
25}
26
27impl AuthRateLimiter {
28    fn new() -> Self {
29        Self {
30            attempts: Arc::new(RwLock::new(HashMap::new())),
31        }
32    }
33
34    fn check_locked(&self, ip: &str, _lockout_duration: Duration) -> Option<Duration> {
35        let mut map = self.attempts.write();
36        if let Some(attempt) = map.get(ip) {
37            if let Some(locked_until) = attempt.locked_until {
38                if Instant::now() < locked_until {
39                    return Some(locked_until.saturating_duration_since(Instant::now()));
40                } else {
41                    // Lock expired, remove entry
42                    map.remove(ip);
43                }
44            }
45        }
46        None
47    }
48
49    fn record_failure(&self, ip: &str, max_attempts: u32, lockout_duration: Duration) -> bool {
50        let mut map = self.attempts.write();
51        let attempt = map.entry(ip.to_string()).or_insert_with(|| AuthAttempt {
52            count: 0,
53            locked_until: None,
54        });
55
56        attempt.count += 1;
57
58        if attempt.count >= max_attempts {
59            attempt.locked_until = Some(Instant::now() + lockout_duration);
60            true // Now locked
61        } else {
62            false // Not locked yet
63        }
64    }
65
66    fn record_success(&self, ip: &str) {
67        let mut map = self.attempts.write();
68        map.remove(ip);
69    }
70
71    fn cleanup_old_entries(&self) {
72        let mut map = self.attempts.write();
73        let now = Instant::now();
74        map.retain(|_, attempt| {
75            if let Some(locked_until) = attempt.locked_until {
76                now < locked_until
77            } else {
78                true // Keep entries that aren't locked (they'll expire naturally)
79            }
80        });
81    }
82}
83
84pub async fn run_server(
85    addr: SocketAddr,
86    store: Arc<Store>,
87    read_buffer_size: usize,
88    config: Arc<Config>,
89) -> anyhow::Result<()> {
90    let rate_limiter = Arc::new(AuthRateLimiter::new());
91
92    // Spawn cleanup task for old rate limit entries
93    let cleanup_limiter = rate_limiter.clone();
94    tokio::spawn(async move {
95        let mut interval = time::interval(Duration::from_secs(60));
96        loop {
97            interval.tick().await;
98            cleanup_limiter.cleanup_old_entries();
99        }
100    });
101
102    let listener = TcpListener::bind(addr).await?;
103    loop {
104        let (socket, peer_addr) = listener.accept().await?;
105        let store = store.clone();
106        let read_buffer_size = read_buffer_size;
107        let config = config.clone();
108        let rate_limiter = rate_limiter.clone();
109        tokio::spawn(async move {
110            if let Err(e) = handle_connection(socket, peer_addr, store, read_buffer_size, config, rate_limiter).await {
111                eprintln!("connection error: {e}");
112            }
113        });
114    }
115}
116
117async fn handle_connection(
118    mut stream: TcpStream,
119    peer_addr: SocketAddr,
120    store: Arc<Store>,
121    read_buffer_size: usize,
122    config: Arc<Config>,
123    rate_limiter: Arc<AuthRateLimiter>,
124) -> anyhow::Result<()> {
125    let mut buf = BytesMut::with_capacity(read_buffer_size * 2);
126    let mut authenticated = !config.requires_auth();
127    let client_ip = peer_addr.ip().to_string();
128
129    loop {
130        let mut temp = vec![0u8; read_buffer_size];
131        let n = stream.read(&mut temp).await?;
132        if n == 0 {
133            break;
134        }
135        buf.extend_from_slice(&temp[..n]);
136
137        loop {
138            match parse_frame(&mut buf) {
139                Ok(Some(frame)) => {
140                    // Check if this is an AUTH command
141                    if let RespFrame::Array(Some(items)) = &frame {
142                        if !items.is_empty() {
143                            let cmd_name = match &items[0] {
144                                RespFrame::BulkString(Some(b)) => {
145                                    String::from_utf8_lossy(b).to_ascii_uppercase()
146                                }
147                                RespFrame::SimpleString(s) => s.to_ascii_uppercase(),
148                                _ => String::new(),
149                            };
150
151                            if cmd_name == "AUTH" {
152                                let resp = handle_auth(
153                                    &config,
154                                    &items,
155                                    &mut authenticated,
156                                    &client_ip,
157                                    &rate_limiter,
158                                );
159                                let encoded = encode_frame(&resp);
160                                stream.write_all(&encoded).await?;
161                                continue;
162                            }
163                        }
164                    }
165
166                    // Require authentication if enabled
167                    if !authenticated {
168                        let err = RespFrame::Error("NOAUTH Authentication required.".into());
169                        let encoded = encode_frame(&err);
170                        stream.write_all(&encoded).await?;
171                        continue;
172                    }
173
174                    let resp = handle_command(store.clone(), frame).await;
175                    let encoded = encode_frame(&resp);
176                    stream.write_all(&encoded).await?;
177                }
178                Ok(None) => break,
179                Err(e) => {
180                    let err = RespFrame::Error(format!("ERR {}", e));
181                    let encoded = encode_frame(&err);
182                    stream.write_all(&encoded).await?;
183                    return Ok(());
184                }
185            }
186        }
187    }
188    Ok(())
189}
190
191fn handle_auth(
192    config: &Config,
193    items: &[RespFrame],
194    authenticated: &mut bool,
195    client_ip: &str,
196    rate_limiter: &Arc<AuthRateLimiter>,
197) -> RespFrame {
198    if !config.requires_auth() {
199        return RespFrame::Error("ERR Client sent AUTH, but no password is set".into());
200    }
201
202    // Check if IP is locked
203    let lockout_duration = Duration::from_secs(config.lockout_duration_seconds);
204    if let Some(remaining) = rate_limiter.check_locked(client_ip, lockout_duration) {
205        let minutes = remaining.as_secs() / 60;
206        let seconds = remaining.as_secs() % 60;
207        return RespFrame::Error(format!(
208            "ERR Too many failed login attempts. Please try again in {}m {}s.",
209            minutes, seconds
210        ));
211    }
212
213    let auth_result = if items.len() == 2 {
214        // AUTH password (Redis style, no username)
215        let password = match &items[1] {
216            RespFrame::BulkString(Some(b)) => String::from_utf8_lossy(b).to_string(),
217            RespFrame::SimpleString(s) => s.clone(),
218            _ => return RespFrame::Error("ERR invalid password".into()),
219        };
220
221        config.check_auth(None, Some(&password))
222    } else if items.len() == 3 {
223        // AUTH username password
224        let username = match &items[1] {
225            RespFrame::BulkString(Some(b)) => String::from_utf8_lossy(b).to_string(),
226            RespFrame::SimpleString(s) => s.clone(),
227            _ => return RespFrame::Error("ERR invalid username".into()),
228        };
229        let password = match &items[2] {
230            RespFrame::BulkString(Some(b)) => String::from_utf8_lossy(b).to_string(),
231            RespFrame::SimpleString(s) => s.clone(),
232            _ => return RespFrame::Error("ERR invalid password".into()),
233        };
234
235        config.check_auth(Some(&username), Some(&password))
236    } else {
237        return RespFrame::Error("ERR wrong number of arguments for 'auth' command".into());
238    };
239
240    if auth_result {
241        // Success: reset rate limiter
242        rate_limiter.record_success(client_ip);
243        *authenticated = true;
244        RespFrame::SimpleString("OK".into())
245    } else {
246        // Failure: record attempt
247        let is_locked = rate_limiter.record_failure(client_ip, config.max_failed_attempts, lockout_duration);
248        if is_locked {
249            let minutes = lockout_duration.as_secs() / 60;
250            RespFrame::Error(format!(
251                "ERR Too many failed login attempts. Account locked for {} minutes.",
252                minutes
253            ))
254        } else {
255            RespFrame::Error("ERR invalid password".into())
256        }
257    }
258}
259
260
rust_cache_server/server.rs

rust_cache_server/
server.rs
