// Sonic // // Fast, lightweight and schema-less search backend // Copyright: 2019, Valerian Saliou // Copyright: 2026, Rémi Bardon // License: Mozilla Public License v2.0 (MPL v2.0) use std::collections::VecDeque; use std::io::{ErrorKind, Read, Write}; use std::net::TcpStream; use std::result::Result; use std::str; use std::sync::Arc; use std::time::Duration; use sonic::Executor; use super::message::{ ChannelMessageModeControl, ChannelMessageModeIngest, ChannelMessageModeSearch, ChannelMessageResult, }; use super::mode::ChannelMode; use super::statistics::CLIENTS_CONNECTED; use crate::LINE_FEED; pub struct ChannelHandle { pub app_conf: Arc, pub executor: Executor, } enum ChannelHandleError { Closed, InvalidMode, AuthenticationRequired, AuthenticationFailed, NotRecognized, TimedOut, ConnectionAborted, Interrupted, Unknown, } const LINE_END_GAP: usize = 1; const BUFFER_SIZE: usize = 20000; const MAX_LINE_SIZE: usize = BUFFER_SIZE + LINE_END_GAP + 1; const TCP_TIMEOUT_NON_ESTABLISHED: u64 = 10; const PROTOCOL_REVISION: u8 = 1; const BUFFER_LINE_SEPARATOR: u8 = b'\n'; const CONNECTED_BANNER: &str = concat!( "CONNECTED <", env!("CARGO_PKG_NAME"), " v", env!("CARGO_PKG_VERSION"), ">" ); impl ChannelHandleError { pub fn to_str(&self) -> &'static str { match *self { ChannelHandleError::Closed => "closed", ChannelHandleError::InvalidMode => "invalid_mode", ChannelHandleError::AuthenticationRequired => "authentication_required", ChannelHandleError::AuthenticationFailed => "authentication_failed", ChannelHandleError::NotRecognized => "not_recognized", ChannelHandleError::TimedOut => "timed_out", ChannelHandleError::ConnectionAborted => "connection_aborted", ChannelHandleError::Interrupted => "interrupted", ChannelHandleError::Unknown => "unknown", } } } impl ChannelHandle { pub fn client(&self, mut stream: TcpStream) { // Configure stream (non-established) self.configure_stream(&stream, false); // Send connected banner write!(stream, "{}{}", CONNECTED_BANNER, LINE_FEED).expect("write failed"); // Increment connected clients count *CLIENTS_CONNECTED.write().unwrap() += 1; // Ensure channel mode is set match self.ensure_start(&stream) { Ok(mode) => { // Configure stream (established) self.configure_stream(&stream, true); // Send started acknowledgement (with environment variables) write!( stream, "STARTED {} protocol({}) buffer({}){}", mode.to_str(), PROTOCOL_REVISION, BUFFER_SIZE, LINE_FEED ) .expect("write failed"); self.handle_stream(mode, stream); } Err(err) => match write!(stream, "ENDED {}{}", err.to_str(), LINE_FEED) { Ok(()) => {} Err(err) if err.kind() == ErrorKind::BrokenPipe => {} Err(err) => panic!("write failed: {err:?}"), }, } // Decrement connected clients count *CLIENTS_CONNECTED.write().unwrap() -= 1; } fn configure_stream(&self, stream: &TcpStream, is_established: bool) { let tcp_timeout = if is_established { self.app_conf.channel.tcp_timeout } else { TCP_TIMEOUT_NON_ESTABLISHED }; assert!(stream.set_nodelay(true).is_ok()); assert!( stream .set_read_timeout(Some(Duration::new(tcp_timeout, 0))) .is_ok() ); assert!( stream .set_write_timeout(Some(Duration::new(tcp_timeout, 0))) .is_ok() ); } fn handle_stream(&self, mode: ChannelMode, mut stream: TcpStream) { // Initialize packet buffer let mut buffer: VecDeque = VecDeque::with_capacity(MAX_LINE_SIZE); // Wait for incoming messages 'handler: loop { let mut read = [0; MAX_LINE_SIZE]; match stream.read(&mut read) { Ok(n) => { // Should close? if n == 0 { break; } let (mut chunk, mut read) = read[..n].split_at(n.min(MAX_LINE_SIZE - buffer.len())); // Add chunk to buffer buffer.extend(chunk); assert!(buffer.len() <= MAX_LINE_SIZE); // Handle full lines from buffer (keep the last incomplete line in buffer) { let mut processed_line = Vec::with_capacity(MAX_LINE_SIZE); while let Some(byte) = buffer.pop_front() { // Commit line and start a new one? if byte == BUFFER_LINE_SEPARATOR { if self.on_message(&mode, &stream, &processed_line) == ChannelMessageResult::Close { // Should close? break 'handler; } // Important: clear the contents of the line, as it has just been \ // processed. processed_line.clear(); (chunk, read) = read.split_at(read.len().min(MAX_LINE_SIZE - buffer.len())); // Add chunk to buffer buffer.extend(chunk); assert!(buffer.len() <= MAX_LINE_SIZE); } else { // Append current byte to processed line processed_line.push(byte); } } // Incomplete line remaining? Put it back in buffer. if !processed_line.is_empty() { buffer.extend(processed_line); assert!(buffer.len() <= MAX_LINE_SIZE); } } // Check for buffer overflow. // NOTE: To avoid a needless read of `MAX_LINE_SIZE` bytes, // we also ensure there is enough space for the line // separator. If there isn’t, next loop cycle will abort // because the line is too long anyway. let separator_len = char::from(BUFFER_LINE_SEPARATOR).len_utf8(); if (buffer.len() + read.len()) < (MAX_LINE_SIZE - separator_len) { buffer.extend(read); } else { // Do not continue, as there is too much pending data // in the buffer. Most likely the client does not // implement a proper back-pressure management system, // thus we terminate it. tracing::error!("closing channel thread because of buffer overflow"); panic!( "buffer overflow ({}/{} bytes)", buffer.len() + read.len(), MAX_LINE_SIZE ); } } Err(err) => { // NOTE: Panicking here would unwind past the connected clients // counter decrement in `client()`, leaking a client slot. tracing::debug!("closing channel thread: {}", err); break 'handler; } } } } fn ensure_start(&self, mut stream: &TcpStream) -> Result { #[allow(clippy::never_loop)] loop { let mut read = [0; MAX_LINE_SIZE]; match stream.read(&mut read) { Ok(n) => { if n == 0 { return Err(ChannelHandleError::Closed); } let mut parts = str::from_utf8(&read[0..n]).unwrap_or("").split_whitespace(); if parts.next().unwrap_or("").to_uppercase().as_str() == "START" { if let Some(res_mode) = parts.next() { tracing::debug!("got mode response: {}", res_mode); // Extract mode if let Ok(mode) = ChannelMode::from_str(res_mode) { // Check if authenticated? if let Some(ref auth_password) = self.app_conf.channel.auth_password { if let Some(provided_auth) = parts.next() { // Compare provided password with configured password if provided_auth != auth_password { tracing::info!("password provided, but does not match"); return Err(ChannelHandleError::AuthenticationFailed); } } else { tracing::info!("no password provided, but one required"); // No password was provided, but we require one return Err(ChannelHandleError::AuthenticationRequired); } } return Ok(mode); } } return Err(ChannelHandleError::InvalidMode); } return Err(ChannelHandleError::NotRecognized); } Err(err) => { let err_reason = match err.kind() { ErrorKind::TimedOut => ChannelHandleError::TimedOut, ErrorKind::ConnectionAborted => ChannelHandleError::ConnectionAborted, ErrorKind::Interrupted => ChannelHandleError::Interrupted, _ => ChannelHandleError::Unknown, }; return Err(err_reason); } } } } fn on_message( &self, mode: &ChannelMode, stream: &TcpStream, message_slice: &[u8], ) -> ChannelMessageResult { use crate::channel::message::ChannelMessageMode as _; match mode { ChannelMode::Search => ChannelMessageModeSearch { executor: &self.executor, search_config: &self.app_conf.sonic.search, normalization_config: &self.app_conf.sonic.normalization, tokenization_config: &self.app_conf.sonic.tokenization, stopwords_config: &self.app_conf.sonic.stopwords, } .on(stream, message_slice), ChannelMode::Ingest => ChannelMessageModeIngest { executor: &self.executor, normalization_config: &self.app_conf.sonic.normalization, tokenization_config: &self.app_conf.sonic.tokenization, stopwords_config: &self.app_conf.sonic.stopwords, } .on(stream, message_slice), ChannelMode::Control => ChannelMessageModeControl { executor: &self.executor, } .on(stream, message_slice), } } }