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

feox_server/network/
connection.rs
