feox_server/network/
connection.rs

1use crate::config::Config;
2use crate::protocol::resp::{write_resp_value, RespValue};
3use crate::protocol::{Command, CommandExecutor, RespParser};
4use crate::pubsub::PubSubMessage;
5use crate::watch_registry::WatchRegistry;
6use bytes::Bytes;
7use feoxdb::FeoxStore;
8use std::collections::{HashMap, VecDeque};
9use std::net::SocketAddr;
10use std::os::fd::RawFd;
11use std::sync::Arc;
12use std::time::{SystemTime, UNIX_EPOCH};
13
14#[derive(Debug, PartialEq)]
15enum TransactionState {
16    None,
17    Queuing,
18}
19
20#[derive(Debug)]
21pub enum PubSubOp {
22    Subscribe(Vec<Vec<u8>>),
23    Unsubscribe(Option<Vec<Vec<u8>>>),
24    PSubscribe(Vec<Vec<u8>>),
25    PUnsubscribe(Option<Vec<Vec<u8>>>),
26    Publish { channel: Vec<u8>, message: Vec<u8> },
27    PubSubChannels { pattern: Option<Vec<u8>> },
28    PubSubNumSub { channels: Vec<Vec<u8>> },
29    PubSubNumPat,
30}
31
32/// Manages a client connection with RESP protocol handling
33///
34/// Handles command parsing, execution, and response buffering.
35pub struct Connection {
36    fd: RawFd,
37
38    // Protocol parser
39    parser: RespParser,
40    executor: CommandExecutor,
41
42    // Authentication state
43    authenticated: bool,
44    auth_required: bool,
45
46    // Single consolidated write buffer for better performance
47    pub write_buffer: Vec<u8>,
48    write_position: usize,
49
50    // Pipeline tracking
51    pipeline_depth: usize,
52
53    // Connection state
54    closed: bool,
55
56    // Pub/Sub state
57    pub connection_id: usize,
58    pub subscription_count: usize,
59    pending_pubsub_messages: VecDeque<PubSubMessage>,
60
61    // Client metadata
62    pub client_name: Option<String>,
63    pub client_addr: Option<SocketAddr>,
64    pub connected_at: u64, // Unix timestamp in seconds
65    pub commands_processed: u64,
66    pub flags: Vec<String>, // Client flags (e.g., "pubsub", "master", "replica")
67
68    // Transaction state
69    transaction_state: TransactionState,
70    queued_commands: Vec<Command>,
71    watched_keys: HashMap<Vec<u8>, u64>,
72}
73
74impl Connection {
75    /// Create a new connection handler
76    pub fn new(fd: RawFd, buffer_size: usize, store: Arc<FeoxStore>, config: &Config) -> Self {
77        Self::new_with_addr(fd, buffer_size, store, config, None)
78    }
79
80    /// Set the executor with client registry info
81    pub fn set_client_registry(&mut self, registry: Arc<crate::client_registry::ClientRegistry>) {
82        self.executor = self
83            .executor
84            .clone()
85            .with_client_info(registry, self.connection_id);
86    }
87
88    /// Create a new connection handler with address
89    pub fn new_with_addr(
90        fd: RawFd,
91        buffer_size: usize,
92        store: Arc<FeoxStore>,
93        config: &Config,
94        addr: Option<SocketAddr>,
95    ) -> Self {
96        let executor = CommandExecutor::new(store, config);
97        let auth_required = config.auth_required();
98
99        static CONNECTION_ID: std::sync::atomic::AtomicUsize =
100            std::sync::atomic::AtomicUsize::new(0);
101        let connection_id = CONNECTION_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
102
103        let now = SystemTime::now()
104            .duration_since(UNIX_EPOCH)
105            .unwrap_or_default()
106            .as_secs();
107
108        Self {
109            fd,
110            parser: RespParser::new(),
111            executor,
112            authenticated: !auth_required, // If no auth required, consider authenticated
113            auth_required,
114            write_buffer: Vec::with_capacity(buffer_size),
115            write_position: 0,
116            pipeline_depth: 0,
117            closed: false,
118            connection_id,
119            subscription_count: 0,
120            pending_pubsub_messages: VecDeque::new(),
121            client_name: None,
122            client_addr: addr,
123            connected_at: now,
124            commands_processed: 0,
125            flags: Vec::new(),
126            transaction_state: TransactionState::None,
127            queued_commands: Vec::new(),
128            watched_keys: HashMap::new(),
129        }
130    }
131
132    pub fn fd(&self) -> RawFd {
133        self.fd
134    }
135
136    pub fn is_closed(&self) -> bool {
137        self.closed
138    }
139
140    pub fn close(&mut self) {
141        if !self.closed {
142            self.closed = true;
143        }
144    }
145
146    /// Set authentication status
147    pub fn set_authenticated(&mut self, authenticated: bool) {
148        self.authenticated = authenticated;
149    }
150
151    /// Check if connection is authenticated
152    pub fn is_authenticated(&self) -> bool {
153        !self.auth_required || self.authenticated
154    }
155
156    /// Process incoming data with inline execution
157    /// Returns pub/sub operations that need to be executed
158    pub fn process_read(&mut self, data: &[u8]) -> crate::error::Result<Vec<PubSubOp>> {
159        let mut pubsub_ops = Vec::new();
160
161        // Feed data to parser
162        self.parser.feed(data);
163
164        // Only clear buffer if all previous writes have been consumed
165        if self.write_position >= self.write_buffer.len() {
166            self.write_buffer.clear();
167            self.write_position = 0;
168        }
169
170        // Parse and execute commands inline
171        while let Some(resp_value) = self
172            .parser
173            .parse_next()
174            .map_err(crate::error::Error::Protocol)?
175        {
176            // Update command counter
177            self.commands_processed += 1;
178
179            // Fast-path for common commands (SET/GET) if not in transaction
180            if self.transaction_state == TransactionState::None && self.try_fast_path(&resp_value) {
181                self.pipeline_depth += 1;
182                continue;
183            }
184
185            // Parse command (slow path)
186            let command = Command::from_resp(resp_value).map_err(crate::error::Error::Protocol)?;
187
188            // Check for quit
189            if matches!(command, Command::Quit) {
190                self.closed = true;
191                self.write_buffer.extend_from_slice(b"+OK\r\n");
192                return Ok(pubsub_ops);
193            }
194
195            // Check if in pub/sub mode and restrict commands
196            if self.is_in_pubsub_mode() && !command.is_allowed_in_pubsub_mode() {
197                write_resp_value(
198                    &mut self.write_buffer,
199                    &RespValue::Error(
200                        "-ERR only (P)SUBSCRIBE / (P)UNSUBSCRIBE / PING / QUIT allowed in this context".to_string(),
201                    ),
202                );
203                continue;
204            }
205
206            // Special handling for CLIENT SETNAME - update connection metadata
207            if let Command::Client {
208                ref subcommand,
209                ref args,
210            } = command
211            {
212                if subcommand.to_uppercase() == "SETNAME" && !args.is_empty() {
213                    self.client_name = Some(String::from_utf8_lossy(&args[0]).to_string());
214                }
215            }
216
217            // Handle transaction commands
218            match command {
219                Command::Multi => {
220                    if self.transaction_state == TransactionState::Queuing {
221                        write_resp_value(
222                            &mut self.write_buffer,
223                            &RespValue::Error("-ERR MULTI calls can not be nested".to_string()),
224                        );
225                        continue;
226                    }
227                    self.transaction_state = TransactionState::Queuing;
228                    self.queued_commands.clear();
229                    write_resp_value(
230                        &mut self.write_buffer,
231                        &RespValue::SimpleString(Bytes::from_static(b"OK")),
232                    );
233                    continue;
234                }
235                Command::Exec => {
236                    if self.transaction_state != TransactionState::Queuing {
237                        write_resp_value(
238                            &mut self.write_buffer,
239                            &RespValue::Error("-ERR EXEC without MULTI".to_string()),
240                        );
241                        continue;
242                    }
243
244                    // Check if any watched keys have been modified
245                    let watched_versions: Vec<(Vec<u8>, u64)> = self
246                        .watched_keys
247                        .iter()
248                        .map(|(k, v)| (k.clone(), *v))
249                        .collect();
250
251                    if !WatchRegistry::check_keys_unchanged(&watched_versions) {
252                        // Transaction aborted - watched key was modified
253                        self.transaction_state = TransactionState::None;
254                        self.queued_commands.clear();
255                        self.watched_keys.clear();
256                        write_resp_value(&mut self.write_buffer, &RespValue::BulkString(None));
257                        continue;
258                    }
259
260                    // Execute all queued commands
261                    let mut results = Vec::new();
262                    for queued_cmd in self.queued_commands.drain(..) {
263                        results.push(self.executor.execute(queued_cmd));
264                    }
265
266                    self.transaction_state = TransactionState::None;
267                    self.watched_keys.clear();
268
269                    write_resp_value(&mut self.write_buffer, &RespValue::Array(Some(results)));
270                    continue;
271                }
272                Command::Discard => {
273                    if self.transaction_state != TransactionState::Queuing {
274                        write_resp_value(
275                            &mut self.write_buffer,
276                            &RespValue::Error("-ERR DISCARD without MULTI".to_string()),
277                        );
278                        continue;
279                    }
280
281                    self.transaction_state = TransactionState::None;
282                    self.queued_commands.clear();
283                    self.watched_keys.clear();
284
285                    write_resp_value(
286                        &mut self.write_buffer,
287                        &RespValue::SimpleString(Bytes::from_static(b"OK")),
288                    );
289                    continue;
290                }
291                Command::Watch(ref keys) => {
292                    if self.transaction_state == TransactionState::Queuing {
293                        write_resp_value(
294                            &mut self.write_buffer,
295                            &RespValue::Error("-ERR WATCH inside MULTI is not allowed".to_string()),
296                        );
297                        continue;
298                    }
299                    for key in keys {
300                        let version = WatchRegistry::get_key_version(key);
301                        self.watched_keys.insert(key.clone(), version);
302                    }
303                    write_resp_value(
304                        &mut self.write_buffer,
305                        &RespValue::SimpleString(Bytes::from_static(b"OK")),
306                    );
307                    continue;
308                }
309                Command::Unwatch => {
310                    self.watched_keys.clear();
311                    write_resp_value(
312                        &mut self.write_buffer,
313                        &RespValue::SimpleString(Bytes::from_static(b"OK")),
314                    );
315                    continue;
316                }
317                _ => {}
318            }
319
320            // If in transaction, queue the command
321            if self.transaction_state == TransactionState::Queuing {
322                self.queued_commands.push(command);
323                write_resp_value(
324                    &mut self.write_buffer,
325                    &RespValue::SimpleString(Bytes::from_static(b"QUEUED")),
326                );
327                continue;
328            }
329
330            // Check authentication for non-AUTH commands
331            let response = if !self.authenticated && !matches!(command, Command::Auth(_)) {
332                // Allow PING without auth (Redis-compatible)
333                if matches!(command, Command::Ping(_)) {
334                    self.executor.execute(command)
335                } else {
336                    RespValue::Error("-NOAUTH Authentication required.".to_string())
337                }
338            } else {
339                // Special handling for AUTH command
340                if let Command::Auth(password) = &command {
341                    // Check if password is configured
342                    if !self.auth_required {
343                        RespValue::Error(
344                            "-ERR Client sent AUTH, but no password is set".to_string(),
345                        )
346                    } else {
347                        let password_str = String::from_utf8_lossy(password);
348                        if self.executor.check_auth(&password_str) {
349                            self.set_authenticated(true);
350                            RespValue::SimpleString(Bytes::from_static(b"OK"))
351                        } else {
352                            RespValue::Error("-ERR invalid password".to_string())
353                        }
354                    }
355                } else if command.is_pubsub_command() {
356                    // Capture subcommand for error message if needed
357                    let subcommand_str = if let Command::PubSub { ref subcommand, .. } = command {
358                        Some(subcommand.clone())
359                    } else {
360                        None
361                    };
362
363                    if let Some(pubsub_op) = command.to_pubsub_op() {
364                        pubsub_ops.push(pubsub_op);
365                        // Response will be sent after processing by pub/sub manager
366                        continue;
367                    } else if let Some(subcommand) = subcommand_str {
368                        // Unknown PUBSUB subcommand
369                        RespValue::Error(format!("ERR Unknown PUBSUB subcommand '{}'", subcommand))
370                    } else {
371                        RespValue::Error("ERR Failed to process pub/sub command".to_string())
372                    }
373                } else {
374                    self.executor.execute(command)
375                }
376            };
377
378            write_resp_value(&mut self.write_buffer, &response);
379
380            self.pipeline_depth += 1;
381        }
382
383        Ok(pubsub_ops)
384    }
385
386    /// Get pending write data as a single buffer slice
387    pub fn pending_writes(&mut self) -> Option<&[u8]> {
388        if self.write_position < self.write_buffer.len() {
389            Some(&self.write_buffer[self.write_position..])
390        } else {
391            None
392        }
393    }
394
395    /// Mark bytes as written
396    pub fn consume_writes(&mut self, n: usize) {
397        self.write_position += n;
398    }
399
400    /// Add a pub/sub message to the pending queue
401    pub fn queue_pubsub_message(&mut self, message: PubSubMessage) {
402        self.pending_pubsub_messages.push_back(message);
403    }
404
405    /// Check if connection is in pub/sub mode
406    pub fn is_in_pubsub_mode(&self) -> bool {
407        self.subscription_count > 0
408    }
409
410    /// Update subscription count
411    pub fn set_subscription_count(&mut self, count: usize) {
412        self.subscription_count = count;
413        // Update flags based on subscription status
414        if count > 0 && !self.flags.contains(&"pubsub".to_string()) {
415            self.flags.push("pubsub".to_string());
416        } else if count == 0 {
417            self.flags.retain(|f| f != "pubsub");
418        }
419    }
420
421    /// Process pending pub/sub messages
422    pub fn process_pubsub_messages(&mut self) {
423        while let Some(message) = self.pending_pubsub_messages.pop_front() {
424            let resp = message.to_resp();
425            write_resp_value(&mut self.write_buffer, &resp);
426        }
427    }
428
429    /// Try to handle common commands (SET/GET) via fast path
430    /// Returns true if handled, false otherwise
431    #[inline(always)]
432    fn try_fast_path(&mut self, resp_value: &RespValue) -> bool {
433        // Static responses
434        const OK_RESPONSE: &[u8] = b"+OK\r\n";
435        const NIL_RESPONSE: &[u8] = b"$-1\r\n";
436
437        // Must be an array with at least 2 elements
438        let args = match resp_value {
439            RespValue::Array(Some(args)) if args.len() >= 2 => args,
440            _ => return false,
441        };
442
443        // Get command name
444        let cmd = match &args[0] {
445            RespValue::BulkString(Some(cmd)) => cmd,
446            _ => return false,
447        };
448
449        // Check for SET command (3 args minimum: SET key value)
450        if cmd.len() == 3 && cmd.eq_ignore_ascii_case(b"SET") && args.len() >= 3 {
451            // Extract key and value
452            let (key, value_bytes) = match (&args[1], &args[2]) {
453                (RespValue::BulkString(Some(k)), RespValue::BulkString(Some(v))) => {
454                    (k.as_ref(), v.clone())
455                }
456                _ => return false,
457            };
458
459            // Simple SET without options
460            if args.len() == 3 {
461                match self.executor.fast_set_bytes(key, value_bytes) {
462                    Ok(_) => {
463                        self.write_buffer.extend_from_slice(OK_RESPONSE);
464                        return true;
465                    }
466                    Err(e) => {
467                        self.write_buffer.extend_from_slice(b"-ERR ");
468                        self.write_buffer
469                            .extend_from_slice(e.to_string().as_bytes());
470                        self.write_buffer.extend_from_slice(b"\r\n");
471                        return true;
472                    }
473                }
474            }
475        }
476
477        // Check for GET command (2 args: GET key)
478        if cmd.len() == 3 && cmd.eq_ignore_ascii_case(b"GET") && args.len() == 2 {
479            let key = match &args[1] {
480                RespValue::BulkString(Some(k)) => k,
481                _ => return false,
482            };
483
484            match self.executor.fast_get(key) {
485                Ok(value) => {
486                    let mut num_buf = itoa::Buffer::new();
487                    let len_str = num_buf.format(value.len());
488
489                    self.write_buffer.push(b'$');
490                    self.write_buffer.extend_from_slice(len_str.as_bytes());
491                    self.write_buffer.extend_from_slice(b"\r\n");
492                    self.write_buffer.extend_from_slice(&value);
493                    self.write_buffer.extend_from_slice(b"\r\n");
494                    return true;
495                }
496                Err(feoxdb::FeoxError::KeyNotFound) => {
497                    self.write_buffer.extend_from_slice(NIL_RESPONSE);
498                    return true;
499                }
500                Err(e) => {
501                    self.write_buffer.extend_from_slice(b"-ERR ");
502                    self.write_buffer
503                        .extend_from_slice(e.to_string().as_bytes());
504                    self.write_buffer.extend_from_slice(b"\r\n");
505                    return true;
506                }
507            }
508        }
509
510        false // Not a fast-path command
511    }
512}
feox_server/network/connection.rs

feox_server/network/
connection.rs
