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        // Only clear buffer if all previous writes have been consumed
155        if self.write_position >= self.write_buffer.len() {
156            self.write_buffer.clear();
157            self.write_position = 0;
158        }
159
160        // Parse and execute commands inline
161        while let Some(resp_value) = self
162            .parser
163            .parse_next()
164            .map_err(crate::error::Error::Protocol)?
165        {
166            // Update last command timestamp
167            self.last_command_at = SystemTime::now()
168                .duration_since(UNIX_EPOCH)
169                .unwrap_or_default()
170                .as_secs();
171            self.commands_processed += 1;
172            // Parse command
173            let command = Command::from_resp(resp_value).map_err(crate::error::Error::Protocol)?;
174
175            // Check for quit
176            if matches!(command, Command::Quit) {
177                self.closed = true;
178                self.write_buffer.extend_from_slice(b"+OK\r\n");
179                return Ok(pubsub_ops);
180            }
181
182            // Check if in pub/sub mode and restrict commands
183            if self.is_in_pubsub_mode() && !command.is_allowed_in_pubsub_mode() {
184                write_resp_value(
185                    &mut self.write_buffer,
186                    &RespValue::Error(
187                        "-ERR only (P)SUBSCRIBE / (P)UNSUBSCRIBE / PING / QUIT allowed in this context".to_string(),
188                    ),
189                );
190                continue;
191            }
192
193            // Special handling for CLIENT SETNAME - update connection metadata
194            if let Command::Client {
195                ref subcommand,
196                ref args,
197            } = command
198            {
199                if subcommand.to_uppercase() == "SETNAME" && !args.is_empty() {
200                    self.client_name = Some(String::from_utf8_lossy(&args[0]).to_string());
201                }
202            }
203
204            // Check authentication for non-AUTH commands
205            let response = if !self.authenticated && !matches!(command, Command::Auth(_)) {
206                // Allow PING without auth (Redis-compatible)
207                if matches!(command, Command::Ping(_)) {
208                    self.executor.execute(command)
209                } else {
210                    RespValue::Error("-NOAUTH Authentication required.".to_string())
211                }
212            } else {
213                // Special handling for AUTH command
214                if let Command::Auth(password) = &command {
215                    // Check if password is configured
216                    if !self.auth_required {
217                        RespValue::Error(
218                            "-ERR Client sent AUTH, but no password is set".to_string(),
219                        )
220                    } else {
221                        let password_str = String::from_utf8_lossy(password);
222                        if self.executor.check_auth(&password_str) {
223                            self.set_authenticated(true);
224                            RespValue::SimpleString(Bytes::from_static(b"OK"))
225                        } else {
226                            RespValue::Error("-ERR invalid password".to_string())
227                        }
228                    }
229                } else if command.is_pubsub_command() {
230                    // Capture subcommand for error message if needed
231                    let subcommand_str = if let Command::PubSub { ref subcommand, .. } = command {
232                        Some(subcommand.clone())
233                    } else {
234                        None
235                    };
236
237                    if let Some(pubsub_op) = command.to_pubsub_op() {
238                        pubsub_ops.push(pubsub_op);
239                        // Response will be sent after processing by pub/sub manager
240                        continue;
241                    } else if let Some(subcommand) = subcommand_str {
242                        // Unknown PUBSUB subcommand
243                        RespValue::Error(format!("ERR Unknown PUBSUB subcommand '{}'", subcommand))
244                    } else {
245                        RespValue::Error("ERR Failed to process pub/sub command".to_string())
246                    }
247                } else {
248                    self.executor.execute(command)
249                }
250            };
251
252            write_resp_value(&mut self.write_buffer, &response);
253
254            self.pipeline_depth += 1;
255        }
256
257        Ok(pubsub_ops)
258    }
259
260    /// Get pending write data as a single buffer slice
261    pub fn pending_writes(&mut self) -> Option<&[u8]> {
262        if self.write_position < self.write_buffer.len() {
263            Some(&self.write_buffer[self.write_position..])
264        } else {
265            None
266        }
267    }
268
269    /// Mark bytes as written
270    pub fn consume_writes(&mut self, n: usize) {
271        self.write_position += n;
272    }
273
274    /// Add a pub/sub message to the pending queue
275    pub fn queue_pubsub_message(&mut self, message: PubSubMessage) {
276        self.pending_pubsub_messages.push_back(message);
277    }
278
279    /// Check if connection is in pub/sub mode
280    pub fn is_in_pubsub_mode(&self) -> bool {
281        self.subscription_count > 0
282    }
283
284    /// Update subscription count
285    pub fn set_subscription_count(&mut self, count: usize) {
286        self.subscription_count = count;
287        // Update flags based on subscription status
288        if count > 0 && !self.flags.contains(&"pubsub".to_string()) {
289            self.flags.push("pubsub".to_string());
290        } else if count == 0 {
291            self.flags.retain(|f| f != "pubsub");
292        }
293    }
294
295    /// Process pending pub/sub messages
296    pub fn process_pubsub_messages(&mut self) {
297        while let Some(message) = self.pending_pubsub_messages.pop_front() {
298            let resp = message.to_resp();
299            write_resp_value(&mut self.write_buffer, &resp);
300        }
301    }
302}
303
304impl Drop for Connection {
305    fn drop(&mut self) {
306        self.close();
307    }
308}
feox_server/network/connection.rs

feox_server/network/
connection.rs
