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::VecDeque;
8use std::net::SocketAddr;
9use std::os::fd::RawFd;
10use std::sync::Arc;
11use std::time::{SystemTime, UNIX_EPOCH};
12
13#[derive(Debug)]
14pub enum PubSubOp {
15    Subscribe(Vec<Vec<u8>>),
16    Unsubscribe(Option<Vec<Vec<u8>>>),
17    PSubscribe(Vec<Vec<u8>>),
18    PUnsubscribe(Option<Vec<Vec<u8>>>),
19    Publish { channel: Vec<u8>, message: Vec<u8> },
20    PubSubChannels { pattern: Option<Vec<u8>> },
21    PubSubNumSub { channels: Vec<Vec<u8>> },
22    PubSubNumPat,
23}
24
25/// Manages a client connection with RESP protocol handling
26///
27/// Handles command parsing, execution, and response buffering.
28pub struct Connection {
29    fd: RawFd,
30
31    // Protocol parser
32    parser: RespParser,
33    executor: CommandExecutor,
34
35    // Authentication state
36    authenticated: bool,
37    auth_required: bool,
38
39    // Single consolidated write buffer for better performance
40    pub write_buffer: Vec<u8>,
41    write_position: usize,
42
43    // Pipeline tracking
44    pipeline_depth: usize,
45
46    // Connection state
47    closed: bool,
48
49    // Pub/Sub state
50    pub connection_id: usize,
51    pub subscription_count: usize,
52    pending_pubsub_messages: VecDeque<PubSubMessage>,
53
54    // Client metadata
55    pub client_name: Option<String>,
56    pub client_addr: Option<SocketAddr>,
57    pub connected_at: u64,    // Unix timestamp in seconds
58    pub last_command_at: u64, // Unix timestamp in seconds
59    pub commands_processed: u64,
60    pub flags: Vec<String>, // Client flags (e.g., "pubsub", "master", "replica")
61}
62
63impl Connection {
64    /// Create a new connection handler
65    pub fn new(fd: RawFd, buffer_size: usize, store: Arc<FeoxStore>, config: &Config) -> Self {
66        Self::new_with_addr(fd, buffer_size, store, config, None)
67    }
68
69    /// Set the executor with client registry info
70    pub fn set_client_registry(&mut self, registry: Arc<crate::client_registry::ClientRegistry>) {
71        self.executor = self
72            .executor
73            .clone()
74            .with_client_info(registry, self.connection_id);
75    }
76
77    /// Create a new connection handler with address
78    pub fn new_with_addr(
79        fd: RawFd,
80        buffer_size: usize,
81        store: Arc<FeoxStore>,
82        config: &Config,
83        addr: Option<SocketAddr>,
84    ) -> Self {
85        let executor = CommandExecutor::new(store, config);
86        let auth_required = config.auth_required();
87
88        static CONNECTION_ID: std::sync::atomic::AtomicUsize =
89            std::sync::atomic::AtomicUsize::new(0);
90        let connection_id = CONNECTION_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
91
92        let now = SystemTime::now()
93            .duration_since(UNIX_EPOCH)
94            .unwrap_or_default()
95            .as_secs();
96
97        Self {
98            fd,
99            parser: RespParser::new(),
100            executor,
101            authenticated: !auth_required, // If no auth required, consider authenticated
102            auth_required,
103            write_buffer: Vec::with_capacity(buffer_size),
104            write_position: 0,
105            pipeline_depth: 0,
106            closed: false,
107            connection_id,
108            subscription_count: 0,
109            pending_pubsub_messages: VecDeque::new(),
110            client_name: None,
111            client_addr: addr,
112            connected_at: now,
113            last_command_at: now,
114            commands_processed: 0,
115            flags: Vec::new(),
116        }
117    }
118
119    pub fn fd(&self) -> RawFd {
120        self.fd
121    }
122
123    pub fn is_closed(&self) -> bool {
124        self.closed
125    }
126
127    pub fn close(&mut self) {
128        if !self.closed {
129            self.closed = true;
130            unsafe {
131                libc::close(self.fd);
132            }
133        }
134    }
135
136    /// Set authentication status
137    pub fn set_authenticated(&mut self, authenticated: bool) {
138        self.authenticated = authenticated;
139    }
140
141    /// Check if connection is authenticated
142    pub fn is_authenticated(&self) -> bool {
143        !self.auth_required || self.authenticated
144    }
145
146    /// Process incoming data with inline execution
147    /// Returns pub/sub operations that need to be executed
148    pub fn process_read(&mut self, data: &[u8]) -> crate::error::Result<Vec<PubSubOp>> {
149        let mut pubsub_ops = Vec::new();
150
151        // Feed data to parser
152        self.parser.feed(data);
153
154        // Clear buffer for new batch of responses
155        self.write_buffer.clear();
156        self.write_position = 0;
157
158        // Parse and execute commands inline
159        while let Some(resp_value) = self
160            .parser
161            .parse_next()
162            .map_err(crate::error::Error::Protocol)?
163        {
164            // Update last command timestamp
165            self.last_command_at = SystemTime::now()
166                .duration_since(UNIX_EPOCH)
167                .unwrap_or_default()
168                .as_secs();
169            self.commands_processed += 1;
170            // Parse command
171            let command = Command::from_resp(resp_value).map_err(crate::error::Error::Protocol)?;
172
173            // Check for quit
174            if matches!(command, Command::Quit) {
175                self.closed = true;
176                self.write_buffer.extend_from_slice(b"+OK\r\n");
177                return Ok(pubsub_ops);
178            }
179
180            // Check if in pub/sub mode and restrict commands
181            if self.is_in_pubsub_mode() && !command.is_allowed_in_pubsub_mode() {
182                write_resp_value(
183                    &mut self.write_buffer,
184                    &RespValue::Error(
185                        "-ERR only (P)SUBSCRIBE / (P)UNSUBSCRIBE / PING / QUIT allowed in this context".to_string(),
186                    ),
187                );
188                continue;
189            }
190
191            // Special handling for CLIENT SETNAME - update connection metadata
192            if let Command::Client {
193                ref subcommand,
194                ref args,
195            } = command
196            {
197                if subcommand.to_uppercase() == "SETNAME" && !args.is_empty() {
198                    self.client_name = Some(String::from_utf8_lossy(&args[0]).to_string());
199                }
200            }
201
202            // Check authentication for non-AUTH commands
203            let response = if !self.authenticated && !matches!(command, Command::Auth(_)) {
204                // Allow PING without auth (Redis-compatible)
205                if matches!(command, Command::Ping(_)) {
206                    self.executor.execute(command)
207                } else {
208                    RespValue::Error("-NOAUTH Authentication required.".to_string())
209                }
210            } else {
211                // Special handling for AUTH command
212                if let Command::Auth(password) = &command {
213                    // Check if password is configured
214                    if !self.auth_required {
215                        RespValue::Error(
216                            "-ERR Client sent AUTH, but no password is set".to_string(),
217                        )
218                    } else {
219                        let password_str = String::from_utf8_lossy(password);
220                        if self.executor.check_auth(&password_str) {
221                            self.set_authenticated(true);
222                            RespValue::SimpleString(Bytes::from_static(b"OK"))
223                        } else {
224                            RespValue::Error("-ERR invalid password".to_string())
225                        }
226                    }
227                } else if command.is_pubsub_command() {
228                    // Capture subcommand for error message if needed
229                    let subcommand_str = if let Command::PubSub { ref subcommand, .. } = command {
230                        Some(subcommand.clone())
231                    } else {
232                        None
233                    };
234
235                    if let Some(pubsub_op) = command.to_pubsub_op() {
236                        pubsub_ops.push(pubsub_op);
237                        // Response will be sent after processing by pub/sub manager
238                        continue;
239                    } else if let Some(subcommand) = subcommand_str {
240                        // Unknown PUBSUB subcommand
241                        RespValue::Error(format!("ERR Unknown PUBSUB subcommand '{}'", subcommand))
242                    } else {
243                        RespValue::Error("ERR Failed to process pub/sub command".to_string())
244                    }
245                } else {
246                    self.executor.execute(command)
247                }
248            };
249
250            write_resp_value(&mut self.write_buffer, &response);
251
252            self.pipeline_depth += 1;
253        }
254
255        Ok(pubsub_ops)
256    }
257
258    /// Get pending write data as a single buffer slice
259    pub fn pending_writes(&mut self) -> Option<&[u8]> {
260        if self.write_position < self.write_buffer.len() {
261            Some(&self.write_buffer[self.write_position..])
262        } else {
263            None
264        }
265    }
266
267    /// Mark bytes as written
268    pub fn consume_writes(&mut self, n: usize) {
269        self.write_position += n;
270    }
271
272    /// Add a pub/sub message to the pending queue
273    pub fn queue_pubsub_message(&mut self, message: PubSubMessage) {
274        self.pending_pubsub_messages.push_back(message);
275    }
276
277    /// Check if connection is in pub/sub mode
278    pub fn is_in_pubsub_mode(&self) -> bool {
279        self.subscription_count > 0
280    }
281
282    /// Update subscription count
283    pub fn set_subscription_count(&mut self, count: usize) {
284        self.subscription_count = count;
285        // Update flags based on subscription status
286        if count > 0 && !self.flags.contains(&"pubsub".to_string()) {
287            self.flags.push("pubsub".to_string());
288        } else if count == 0 {
289            self.flags.retain(|f| f != "pubsub");
290        }
291    }
292
293    /// Process pending pub/sub messages
294    pub fn process_pubsub_messages(&mut self) {
295        while let Some(message) = self.pending_pubsub_messages.pop_front() {
296            let resp = message.to_resp();
297            write_resp_value(&mut self.write_buffer, &resp);
298        }
299    }
300}
301
302impl Drop for Connection {
303    fn drop(&mut self) {
304        self.close();
305    }
306}
feox_server/network/connection.rs

feox_server/network/
connection.rs
