From c492bcc85f4c876192e019bad447bdd7a9cae7e4 Mon Sep 17 00:00:00 2001 From: upstream Date: Fri, 11 Sep 2026 23:33:21 +0200 Subject: [PATCH] Upstream valeriansaliou/sonic state of the touched files before 534733e716 --- server/src/channel/handle.rs | 317 +++++++++++++++++++++++++++++++++++ 1 file changed, 317 insertions(+) create mode 100644 server/src/channel/handle.rs diff --git a/server/src/channel/handle.rs b/server/src/channel/handle.rs new file mode 100644 index 0000000..74a5f90 --- /dev/null +++ b/server/src/channel/handle.rs @@ -0,0 +1,317 @@ +// 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) => { + tracing::error!("closing channel thread with traceback: {}", err); + + panic!("closing channel"); + } + } + } + } + + 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), + } + } +}