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

feox_server/network/
connection.rs
