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 && self.try_fast_path(&resp_value) {
180                self.pipeline_depth += 1;
181                continue;
182            }
183
184            // Parse command (slow path)
185            let command = Command::from_resp(resp_value).map_err(crate::error::Error::Protocol)?;
186
187            // Check for quit
188            if matches!(command, Command::Quit) {
189                self.closed = true;
190                self.write_buffer.extend_from_slice(b"+OK\r\n");
191                return Ok(pubsub_ops);
192            }
193
194            // Check if in pub/sub mode and restrict commands
195            if self.is_in_pubsub_mode() && !command.is_allowed_in_pubsub_mode() {
196                write_resp_value(
197                    &mut self.write_buffer,
198                    &RespValue::Error(
199                        "-ERR only (P)SUBSCRIBE / (P)UNSUBSCRIBE / PING / QUIT allowed in this context".to_string(),
200                    ),
201                );
202                continue;
203            }
204
205            // Special handling for CLIENT SETNAME - update connection metadata
206            if let Command::Client {
207                ref subcommand,
208                ref args,
209            } = command
210            {
211                if subcommand.to_uppercase() == "SETNAME" && !args.is_empty() {
212                    self.client_name = Some(String::from_utf8_lossy(&args[0]).to_string());
213                }
214            }
215
216            // Handle transaction commands
217            match command {
218                Command::Multi => {
219                    if self.transaction_state == TransactionState::Queuing {
220                        write_resp_value(
221                            &mut self.write_buffer,
222                            &RespValue::Error("-ERR MULTI calls can not be nested".to_string()),
223                        );
224                        continue;
225                    }
226                    self.transaction_state = TransactionState::Queuing;
227                    self.queued_commands.clear();
228                    write_resp_value(
229                        &mut self.write_buffer,
230                        &RespValue::SimpleString(Bytes::from_static(b"OK")),
231                    );
232                    continue;
233                }
234                Command::Exec => {
235                    if self.transaction_state != TransactionState::Queuing {
236                        write_resp_value(
237                            &mut self.write_buffer,
238                            &RespValue::Error("-ERR EXEC without MULTI".to_string()),
239                        );
240                        continue;
241                    }
242
243                    // Execute all queued commands
244                    let mut results = Vec::new();
245                    for queued_cmd in self.queued_commands.drain(..) {
246                        results.push(self.executor.execute(queued_cmd));
247                    }
248
249                    self.transaction_state = TransactionState::None;
250                    self.watched_keys.clear();
251
252                    write_resp_value(&mut self.write_buffer, &RespValue::Array(Some(results)));
253                    continue;
254                }
255                Command::Discard => {
256                    if self.transaction_state != TransactionState::Queuing {
257                        write_resp_value(
258                            &mut self.write_buffer,
259                            &RespValue::Error("-ERR DISCARD without MULTI".to_string()),
260                        );
261                        continue;
262                    }
263
264                    self.transaction_state = TransactionState::None;
265                    self.queued_commands.clear();
266                    self.watched_keys.clear();
267
268                    write_resp_value(
269                        &mut self.write_buffer,
270                        &RespValue::SimpleString(Bytes::from_static(b"OK")),
271                    );
272                    continue;
273                }
274                Command::Watch(ref keys) => {
275                    if self.transaction_state == TransactionState::Queuing {
276                        write_resp_value(
277                            &mut self.write_buffer,
278                            &RespValue::Error("-ERR WATCH inside MULTI is not allowed".to_string()),
279                        );
280                        continue;
281                    }
282                    for key in keys {
283                        self.watched_keys.insert(key.clone());
284                    }
285                    write_resp_value(
286                        &mut self.write_buffer,
287                        &RespValue::SimpleString(Bytes::from_static(b"OK")),
288                    );
289                    continue;
290                }
291                Command::Unwatch => {
292                    self.watched_keys.clear();
293                    write_resp_value(
294                        &mut self.write_buffer,
295                        &RespValue::SimpleString(Bytes::from_static(b"OK")),
296                    );
297                    continue;
298                }
299                _ => {}
300            }
301
302            // If in transaction, queue the command
303            if self.transaction_state == TransactionState::Queuing {
304                self.queued_commands.push(command);
305                write_resp_value(
306                    &mut self.write_buffer,
307                    &RespValue::SimpleString(Bytes::from_static(b"QUEUED")),
308                );
309                continue;
310            }
311
312            // Check authentication for non-AUTH commands
313            let response = if !self.authenticated && !matches!(command, Command::Auth(_)) {
314                // Allow PING without auth (Redis-compatible)
315                if matches!(command, Command::Ping(_)) {
316                    self.executor.execute(command)
317                } else {
318                    RespValue::Error("-NOAUTH Authentication required.".to_string())
319                }
320            } else {
321                // Special handling for AUTH command
322                if let Command::Auth(password) = &command {
323                    // Check if password is configured
324                    if !self.auth_required {
325                        RespValue::Error(
326                            "-ERR Client sent AUTH, but no password is set".to_string(),
327                        )
328                    } else {
329                        let password_str = String::from_utf8_lossy(password);
330                        if self.executor.check_auth(&password_str) {
331                            self.set_authenticated(true);
332                            RespValue::SimpleString(Bytes::from_static(b"OK"))
333                        } else {
334                            RespValue::Error("-ERR invalid password".to_string())
335                        }
336                    }
337                } else if command.is_pubsub_command() {
338                    // Capture subcommand for error message if needed
339                    let subcommand_str = if let Command::PubSub { ref subcommand, .. } = command {
340                        Some(subcommand.clone())
341                    } else {
342                        None
343                    };
344
345                    if let Some(pubsub_op) = command.to_pubsub_op() {
346                        pubsub_ops.push(pubsub_op);
347                        // Response will be sent after processing by pub/sub manager
348                        continue;
349                    } else if let Some(subcommand) = subcommand_str {
350                        // Unknown PUBSUB subcommand
351                        RespValue::Error(format!("ERR Unknown PUBSUB subcommand '{}'", subcommand))
352                    } else {
353                        RespValue::Error("ERR Failed to process pub/sub command".to_string())
354                    }
355                } else {
356                    self.executor.execute(command)
357                }
358            };
359
360            write_resp_value(&mut self.write_buffer, &response);
361
362            self.pipeline_depth += 1;
363        }
364
365        Ok(pubsub_ops)
366    }
367
368    /// Get pending write data as a single buffer slice
369    pub fn pending_writes(&mut self) -> Option<&[u8]> {
370        if self.write_position < self.write_buffer.len() {
371            Some(&self.write_buffer[self.write_position..])
372        } else {
373            None
374        }
375    }
376
377    /// Mark bytes as written
378    pub fn consume_writes(&mut self, n: usize) {
379        self.write_position += n;
380    }
381
382    /// Add a pub/sub message to the pending queue
383    pub fn queue_pubsub_message(&mut self, message: PubSubMessage) {
384        self.pending_pubsub_messages.push_back(message);
385    }
386
387    /// Check if connection is in pub/sub mode
388    pub fn is_in_pubsub_mode(&self) -> bool {
389        self.subscription_count > 0
390    }
391
392    /// Update subscription count
393    pub fn set_subscription_count(&mut self, count: usize) {
394        self.subscription_count = count;
395        // Update flags based on subscription status
396        if count > 0 && !self.flags.contains(&"pubsub".to_string()) {
397            self.flags.push("pubsub".to_string());
398        } else if count == 0 {
399            self.flags.retain(|f| f != "pubsub");
400        }
401    }
402
403    /// Process pending pub/sub messages
404    pub fn process_pubsub_messages(&mut self) {
405        while let Some(message) = self.pending_pubsub_messages.pop_front() {
406            let resp = message.to_resp();
407            write_resp_value(&mut self.write_buffer, &resp);
408        }
409    }
410
411    /// Try to handle common commands (SET/GET) via fast path
412    /// Returns true if handled, false otherwise
413    #[inline(always)]
414    fn try_fast_path(&mut self, resp_value: &RespValue) -> bool {
415        // Static responses
416        const OK_RESPONSE: &[u8] = b"+OK\r\n";
417        const NIL_RESPONSE: &[u8] = b"$-1\r\n";
418
419        // Must be an array with at least 2 elements
420        let args = match resp_value {
421            RespValue::Array(Some(args)) if args.len() >= 2 => args,
422            _ => return false,
423        };
424
425        // Get command name
426        let cmd = match &args[0] {
427            RespValue::BulkString(Some(cmd)) => cmd,
428            _ => return false,
429        };
430
431        // Check for SET command (3 args minimum: SET key value)
432        if cmd.len() == 3 && cmd.eq_ignore_ascii_case(b"SET") && args.len() >= 3 {
433            // Extract key and value
434            let (key, value_bytes) = match (&args[1], &args[2]) {
435                (RespValue::BulkString(Some(k)), RespValue::BulkString(Some(v))) => {
436                    (k.as_ref(), v.clone())
437                }
438                _ => return false,
439            };
440
441            // Simple SET without options
442            if args.len() == 3 {
443                match self.executor.fast_set_bytes(key, value_bytes) {
444                    Ok(_) => {
445                        self.write_buffer.extend_from_slice(OK_RESPONSE);
446                        return true;
447                    }
448                    Err(e) => {
449                        self.write_buffer.extend_from_slice(b"-ERR ");
450                        self.write_buffer
451                            .extend_from_slice(e.to_string().as_bytes());
452                        self.write_buffer.extend_from_slice(b"\r\n");
453                        return true;
454                    }
455                }
456            }
457        }
458
459        // Check for GET command (2 args: GET key)
460        if cmd.len() == 3 && cmd.eq_ignore_ascii_case(b"GET") && args.len() == 2 {
461            let key = match &args[1] {
462                RespValue::BulkString(Some(k)) => k,
463                _ => return false,
464            };
465
466            match self.executor.fast_get(key) {
467                Ok(value) => {
468                    let mut num_buf = itoa::Buffer::new();
469                    let len_str = num_buf.format(value.len());
470
471                    self.write_buffer.push(b'$');
472                    self.write_buffer.extend_from_slice(len_str.as_bytes());
473                    self.write_buffer.extend_from_slice(b"\r\n");
474                    self.write_buffer.extend_from_slice(&value);
475                    self.write_buffer.extend_from_slice(b"\r\n");
476                    return true;
477                }
478                Err(feoxdb::FeoxError::KeyNotFound) => {
479                    self.write_buffer.extend_from_slice(NIL_RESPONSE);
480                    return true;
481                }
482                Err(e) => {
483                    self.write_buffer.extend_from_slice(b"-ERR ");
484                    self.write_buffer
485                        .extend_from_slice(e.to_string().as_bytes());
486                    self.write_buffer.extend_from_slice(b"\r\n");
487                    return true;
488                }
489            }
490        }
491
492        false // Not a fast-path command
493    }
494}
feox_server/network/connection.rs

feox_server/network/
connection.rs
