use std::sync::Arc;
use std::sync::atomic::{
AtomicBool,
Ordering,
};
use std::time::Duration;
use tokio::io::{
AsyncReadExt,
AsyncWriteExt,
};
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;
use tracing::{
debug,
error,
};
use crate::Logger;
use crate::kube::http_log_watcher::HttpLogStateWatcher;
const BUFFER_SIZE: usize = 65536;
#[derive(Default)]
struct UpgradeDetector {
headers: Vec<u8>,
body_remaining: usize,
opaque: bool,
}
impl UpgradeDetector {
fn observe(&mut self, mut bytes: &[u8]) -> bool {
while !bytes.is_empty() && !self.opaque {
if self.body_remaining > 0 {
let consumed = self.body_remaining.min(bytes.len());
self.body_remaining -= consumed;
bytes = &bytes[consumed..];
continue;
}
self.headers.push(bytes[0]);
bytes = &bytes[1..];
let prefix_len = self.headers.len().min(5);
if self.headers[..prefix_len] != b"HTTP/"[..prefix_len] {
self.opaque = true;
self.headers.clear();
return false;
}
if self.headers.ends_with(b"\r\n\r\n") {
if kftray_http_logs::http_response_analyzer::HttpResponseAnalyzer::is_websocket_upgrade(&self.headers) {
self.headers.clear();
return true;
}
let mut headers = [httparse::EMPTY_HEADER; 64];
let mut response = httparse::Response::new(&mut headers);
if matches!(
response.parse(&self.headers),
Ok(httparse::Status::Complete(_))
) {
let no_body = matches!(response.code, Some(100..=199) | Some(204) | Some(304));
let length = response
.headers
.iter()
.find(|h| h.name.eq_ignore_ascii_case("content-length"))
.and_then(|h| std::str::from_utf8(h.value).ok())
.and_then(|v| v.trim().parse::<usize>().ok());
let transfer_encoded = response
.headers
.iter()
.any(|h| h.name.eq_ignore_ascii_case("transfer-encoding"));
self.opaque = response.code == Some(101)
|| transfer_encoded
|| (!no_body && length.is_none());
self.body_remaining = if no_body { 0 } else { length.unwrap_or(0) };
} else {
self.opaque = true;
}
self.headers.clear();
} else if self.headers.len() == BUFFER_SIZE {
self.opaque = true;
self.headers.clear();
}
}
false
}
}
#[derive(Clone)]
pub struct TcpForwarder {
config_id: i64,
workload_type: String,
logger: Option<Logger>,
}
impl TcpForwarder {
pub fn new(config_id: i64, workload_type: String) -> Self {
Self {
config_id,
workload_type,
logger: None,
}
}
pub fn is_logging_enabled(&self) -> bool {
self.logger.is_some()
}
pub fn apply_socket_optimizations(stream: &tokio::net::TcpStream) {
use socket2::SockRef;
if let Err(e) = stream.set_nodelay(true) {
tracing::debug!("Failed to set TCP_NODELAY: {}", e);
}
let sock_ref = SockRef::from(stream);
if let Err(e) = sock_ref.set_recv_buffer_size(BUFFER_SIZE) {
tracing::debug!("Failed to set receive buffer size: {}", e);
}
if let Err(e) = sock_ref.set_send_buffer_size(BUFFER_SIZE) {
tracing::debug!("Failed to set send buffer size: {}", e);
}
let keepalive = socket2::TcpKeepalive::new()
.with_time(Duration::from_secs(60))
.with_interval(Duration::from_secs(10));
#[cfg(any(target_os = "linux", target_os = "macos"))]
let keepalive = keepalive.with_retries(3);
if let Err(e) = sock_ref.set_tcp_keepalive(&keepalive) {
tracing::debug!("Failed to set TCP keepalive: {}", e);
}
}
#[allow(clippy::too_many_arguments)]
pub async fn forward_streams(
&mut self, client_stream: tokio::net::TcpStream,
upstream_stream: impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
client_address: std::net::SocketAddr, cancellation_token: CancellationToken,
http_log_watcher: Arc<HttpLogStateWatcher>, local_port: u16,
) -> anyhow::Result<()> {
Self::apply_socket_optimizations(&client_stream);
let _log_subscriber = http_log_watcher.create_filtered_subscriber(self.config_id);
let current_logging_enabled = http_log_watcher.get_http_logs(self.config_id).await;
if current_logging_enabled
&& self.logger.is_none()
&& let Err(e) = self.initialize_logger(local_port).await
{
error!("Failed to initialize logger for {}: {}", client_address, e);
}
self.forward_connection(
Arc::new(Mutex::new(client_stream)),
upstream_stream,
cancellation_token,
http_log_watcher,
local_port,
)
.await
}
pub async fn initialize_logger(&mut self, local_port: u16) -> anyhow::Result<()> {
if self.workload_type != "service" && self.workload_type != "pod" {
return Ok(());
}
let http_logs_enabled =
match kftray_commons::utils::http_logs_config::get_http_logs_config(self.config_id)
.await
{
Ok(config) => config.enabled,
Err(_) => false,
};
if http_logs_enabled {
if self.logger.is_none() {
debug!(
"Initializing HTTP logger for config_id {} on port {}",
self.config_id, local_port
);
let logger =
kftray_http_logs::HttpLogger::for_config(self.config_id, local_port).await?;
self.logger = Some(logger);
}
} else if self.logger.is_some() {
debug!(
"HTTP logging disabled for config_id {}, clearing logger",
self.config_id
);
self.logger = None;
}
Ok(())
}
pub async fn forward_connection(
&mut self, client_conn: Arc<Mutex<TcpStream>>,
mut upstream_conn: impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
cancellation_token: CancellationToken, http_log_watcher: Arc<HttpLogStateWatcher>,
local_port: u16,
) -> anyhow::Result<()> {
let log_subscriber = http_log_watcher.create_filtered_subscriber(self.config_id);
let current_logging_enabled = http_log_watcher.get_http_logs(self.config_id).await;
if current_logging_enabled && self.logger.is_none() {
let _ = self.initialize_logger(local_port).await;
}
if current_logging_enabled || self.logger.is_some() {
let config_id = self.config_id;
let shared_logger: Arc<Mutex<Option<Arc<crate::Logger>>>> =
Arc::new(Mutex::new(self.logger.take().map(Arc::new)));
let request_id = Arc::new(Mutex::new(None));
let websocket_tunnel_mode = Arc::new(AtomicBool::new(false));
let mut client_conn_guard = client_conn.lock().await;
let (mut client_reader, mut client_writer) = tokio::io::split(&mut *client_conn_guard);
let (mut upstream_reader, mut upstream_writer) = tokio::io::split(upstream_conn);
let client_to_upstream = Self::forward_client_to_upstream(
Arc::clone(&shared_logger),
config_id,
&mut client_reader,
&mut upstream_writer,
Arc::clone(&request_id),
cancellation_token.clone(),
log_subscriber.resubscribe(),
local_port,
Arc::clone(&websocket_tunnel_mode),
);
let upstream_to_client = Self::forward_upstream_to_client(
Arc::clone(&shared_logger),
config_id,
&mut upstream_reader,
&mut client_writer,
Arc::clone(&request_id),
cancellation_token.clone(),
log_subscriber,
local_port,
Arc::clone(&websocket_tunnel_mode),
);
let result = tokio::try_join!(client_to_upstream, upstream_to_client);
self.logger = Arc::try_unwrap(shared_logger)
.ok()
.and_then(|mutex| mutex.into_inner())
.and_then(|arc_logger| Arc::try_unwrap(arc_logger).ok());
match result {
Ok(_) => {
debug!("HTTP-aware connection closed normally");
}
Err(e) => {
error!("HTTP-aware connection closed with error: {}", e);
return Err(e);
}
}
} else {
let mut client_conn_guard = client_conn.lock().await;
let copy_future =
tokio::io::copy_bidirectional(&mut *client_conn_guard, &mut upstream_conn);
tokio::select! {
result = copy_future => {
match result {
Ok(_) => debug!("Simple connection closed normally"),
Err(e) => {
error!("Simple connection closed with error: {}", e);
return Err(e.into());
}
}
}
_ = cancellation_token.cancelled() => {
debug!("Connection cancelled");
}
}
}
Ok(())
}
pub async fn forward_tls_streams(
&self, mut client: tokio_rustls::server::TlsStream<TcpStream>,
mut upstream: impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
cancellation_token: tokio_util::sync::CancellationToken,
) -> anyhow::Result<()> {
Self::apply_socket_optimizations(client.get_ref().0);
let copy_fut = tokio::io::copy_bidirectional(&mut client, &mut upstream);
tokio::select! {
res = copy_fut => match res {
Ok((to_upstream, to_client)) => {
debug!(
"TLS connection closed: sent {} bytes, received {} bytes",
to_upstream, to_client
);
Ok(())
}
Err(e) => {
error!("TLS connection closed with error: {}", e);
Err(e.into())
}
},
_ = cancellation_token.cancelled() => {
debug!("TLS connection canceled");
Err(anyhow::anyhow!("TLS connection canceled"))
}
}
}
#[allow(clippy::too_many_arguments)]
async fn forward_client_to_upstream<'a>(
logger: Arc<Mutex<Option<Arc<crate::Logger>>>>, config_id: i64,
client_reader: &'a mut (impl AsyncReadExt + Unpin),
upstream_writer: &'a mut (impl AsyncWriteExt + Unpin),
request_id: Arc<Mutex<Option<String>>>, cancellation_token: CancellationToken,
mut log_subscriber: tokio::sync::broadcast::Receiver<
crate::kube::http_log_watcher::HttpLogStateEvent,
>,
local_port: u16, websocket_tunnel_mode: Arc<AtomicBool>,
) -> anyhow::Result<()> {
let mut buffer = [0; BUFFER_SIZE];
let mut should_log = {
let logger_guard = logger.lock().await;
logger_guard.is_some()
};
let mut request_buffer = if should_log { Some(Vec::new()) } else { None };
loop {
tokio::select! {
result = client_reader.read(&mut buffer) => {
let n = match result {
Ok(n) => n,
Err(e) => return Err(e.into()),
};
if n == 0 {
break;
}
if websocket_tunnel_mode.load(Ordering::SeqCst) {
tokio::select! {
result = upstream_writer.write_all(&buffer[..n]) => result?,
_ = cancellation_token.cancelled() => break,
}
} else if should_log {
if let Some(ref mut req_buf) = request_buffer.as_mut() {
req_buf.extend_from_slice(&buffer[..n]);
let maybe_logger = {
let guard = logger.lock().await;
guard.clone()
};
if let Some(log) = maybe_logger {
let new_request_id = log.log_request(req_buf.clone().into()).await;
let mut req_id_guard = request_id.lock().await;
*req_id_guard = Some(new_request_id);
if let Err(e) = upstream_writer.write_all(req_buf).await {
return Err(e.into());
}
req_buf.clear();
}
}
} else if let Err(e) = upstream_writer.write_all(&buffer[..n]).await {
return Err(e.into());
}
},
log_event = log_subscriber.recv() => {
if websocket_tunnel_mode.load(Ordering::SeqCst) {
continue;
}
if let Ok(event) = log_event
&& event.config_id == config_id {
let needs_logger = event.enabled && {
let guard = logger.lock().await;
guard.is_none()
};
let new_should_log = if needs_logger {
let http_logs_enabled = match kftray_commons::utils::http_logs_config::get_http_logs_config(config_id).await {
Ok(config) => config.enabled,
Err(_) => false,
};
if http_logs_enabled {
match kftray_http_logs::HttpLogger::for_config(config_id, local_port).await {
Ok(new_logger) => {
let mut guard = logger.lock().await;
if guard.is_none() {
*guard = Some(Arc::new(new_logger));
}
true
}
_ => false,
}
} else {
false
}
} else if event.enabled {
let guard = logger.lock().await;
guard.is_some()
} else {
false
};
if new_should_log != should_log {
should_log = new_should_log;
if should_log {
request_buffer = Some(Vec::new());
debug!("Enabled HTTP logging for client-to-upstream");
} else {
request_buffer = None;
debug!("Disabled HTTP logging for client-to-upstream");
}
}
}
},
_ = cancellation_token.cancelled() => break,
}
}
let _ = upstream_writer.shutdown().await;
Ok(())
}
#[allow(clippy::too_many_arguments)]
async fn forward_upstream_to_client<'a>(
logger: Arc<Mutex<Option<Arc<crate::Logger>>>>, config_id: i64,
upstream_reader: &'a mut (impl AsyncReadExt + Unpin),
client_writer: &'a mut (impl AsyncWriteExt + Unpin),
request_id: Arc<Mutex<Option<String>>>, cancellation_token: CancellationToken,
mut log_subscriber: tokio::sync::broadcast::Receiver<
crate::kube::http_log_watcher::HttpLogStateEvent,
>,
local_port: u16, websocket_tunnel_mode: Arc<AtomicBool>,
) -> anyhow::Result<()> {
let mut buffer = [0; BUFFER_SIZE];
let mut should_log = {
let logger_guard = logger.lock().await;
logger_guard.is_some()
};
let mut response_state = if should_log {
Some(ResponseState {
buffer: Vec::new(),
is_chunked: false,
found_end_marker: false,
total_chunks_received: 0,
current_response_id: None,
current_response_logged: false,
first_chunk_time: None,
force_log_time: None,
})
} else {
None
};
let mut upgrade_detector = UpgradeDetector::default();
loop {
tokio::select! {
result = upstream_reader.read(&mut buffer) => {
let n = match result {
Ok(n) => n,
Err(e) => return Err(e.into()),
};
if n == 0 {
if should_log
&& !websocket_tunnel_mode.load(Ordering::SeqCst)
&& let Some(ref mut state) = response_state.as_mut() {
let maybe_logger = {
let guard = logger.lock().await;
guard.clone()
};
if let Some(log) = maybe_logger
&& !state.buffer.is_empty() && !state.current_response_logged {
let req_id_guard = request_id.lock().await;
if let Some(req_id) = &*req_id_guard {
log.log_response(state.buffer.clone().into(), req_id.clone()).await;
}
}
}
break;
}
if websocket_tunnel_mode.load(Ordering::SeqCst) {
tokio::select! {
result = client_writer.write_all(&buffer[..n]) => result?,
_ = cancellation_token.cancelled() => break,
}
continue;
}
let detected_upgrade = upgrade_detector.observe(&buffer[..n]);
let logged_upgrade = if should_log
&& let Some(ref mut state) = response_state.as_mut() {
Self::handle_response_logging_static(&buffer[..n], state, &logger, &request_id).await
} else {
false
};
if logged_upgrade || detected_upgrade {
websocket_tunnel_mode.store(true, Ordering::SeqCst);
response_state = None;
should_log = false;
}
if let Err(e) = client_writer.write_all(&buffer[..n]).await {
return Err(e.into());
}
},
log_event = log_subscriber.recv() => {
if websocket_tunnel_mode.load(Ordering::SeqCst) {
continue;
}
if let Ok(event) = log_event
&& event.config_id == config_id {
let needs_logger = event.enabled && {
let guard = logger.lock().await;
guard.is_none()
};
let new_should_log = if needs_logger {
let http_logs_enabled = match kftray_commons::utils::http_logs_config::get_http_logs_config(config_id).await {
Ok(config) => config.enabled,
Err(_) => false,
};
if http_logs_enabled {
match kftray_http_logs::HttpLogger::for_config(config_id, local_port).await {
Ok(new_logger) => {
let mut guard = logger.lock().await;
if guard.is_none() {
*guard = Some(Arc::new(new_logger));
}
true
}
_ => false,
}
} else {
false
}
} else if event.enabled {
let guard = logger.lock().await;
guard.is_some()
} else {
false
};
if new_should_log != should_log {
should_log = new_should_log;
if should_log {
response_state = Some(ResponseState {
buffer: Vec::new(),
is_chunked: false,
found_end_marker: false,
total_chunks_received: 0,
current_response_id: None,
current_response_logged: false,
first_chunk_time: None,
force_log_time: None,
});
debug!("Enabled HTTP logging for upstream-to-client");
} else {
response_state = None;
debug!("Disabled HTTP logging for upstream-to-client");
}
}
}
},
_ = cancellation_token.cancelled() => break,
}
}
let _ = client_writer.shutdown().await;
Ok(())
}
async fn handle_response_logging_static(
buffer: &[u8], state: &mut ResponseState, logger: &Arc<Mutex<Option<Arc<Logger>>>>,
request_id: &Arc<Mutex<Option<String>>>,
) -> bool {
let req_id_guard = request_id.lock().await;
let current_req_id = req_id_guard.clone();
drop(req_id_guard);
let is_new_response = match (&state.current_response_id, ¤t_req_id) {
(Some(current_id), Some(req_id)) => current_id != req_id,
(None, Some(_)) => true,
_ => false,
};
if is_new_response {
state.reset_for_new_response(current_req_id);
}
if state.buffer.is_empty() && !buffer.is_empty() {
if state.first_chunk_time.is_none() {
state.first_chunk_time = Some(tokio::time::Instant::now());
}
state.is_chunked = kftray_http_logs::http_response_analyzer::HttpResponseAnalyzer::detect_chunked_encoding(buffer);
}
kftray_http_logs::http_response_analyzer::HttpResponseAnalyzer::process_chunk(
buffer,
state.is_chunked,
&mut state.found_end_marker,
&mut state.total_chunks_received,
);
state.buffer.extend_from_slice(buffer);
let is_websocket_upgrade =
kftray_http_logs::http_response_analyzer::HttpResponseAnalyzer::is_websocket_upgrade(
&state.buffer,
);
if is_websocket_upgrade {
if let Some(headers_end) = state.buffer.windows(4).position(|w| w == b"\r\n\r\n") {
state.buffer.truncate(headers_end + 4);
}
}
if !state.current_response_logged {
let maybe_logger = {
let guard = logger.lock().await;
guard.clone()
};
if let Some(log) = maybe_logger {
let should_log = Self::should_log_response_static(state);
if should_log && state.current_response_id.is_some() {
let response_id = state.current_response_id.as_ref().unwrap().clone();
state.current_response_logged = true;
log.log_response(state.buffer.clone().into(), response_id)
.await;
if Self::can_clear_response_buffer_static(state) {
state.reset_for_next_response();
}
}
}
}
is_websocket_upgrade
}
fn should_log_response_static(state: &mut ResponseState) -> bool {
if state.is_chunked {
if state.found_end_marker && state.force_log_time.is_none() {
let delay = if !state.buffer.is_empty()
&& state.buffer.len() < 40_000
&& state.total_chunks_received < 20
{
tokio::time::Instant::now() + tokio::time::Duration::from_millis(10)
} else {
tokio::time::Instant::now()
};
state.force_log_time = Some(delay);
return delay <= tokio::time::Instant::now();
}
if let Some(force_time) = state.force_log_time {
let now = tokio::time::Instant::now();
if now >= force_time {
let needs_more_time =
state.buffer.windows(18).any(|w| {
w == b"Content-Encoding: gzip" || w == b"content-encoding: gzip"
}) && state.total_chunks_received < 20
&& now.saturating_duration_since(force_time)
< tokio::time::Duration::from_millis(10);
return !needs_more_time;
}
} else if let Some(first_time) = state.first_chunk_time {
return first_time.elapsed().as_secs() > 5;
}
} else {
return kftray_http_logs::http_response_analyzer::HttpResponseAnalyzer::check_content_length_match(&state.buffer)
|| kftray_http_logs::http_response_analyzer::HttpResponseAnalyzer::is_websocket_upgrade(&state.buffer);
}
false
}
fn can_clear_response_buffer_static(state: &ResponseState) -> bool {
if state.is_chunked {
state.found_end_marker
} else {
kftray_http_logs::http_response_analyzer::HttpResponseAnalyzer::check_content_length_match(&state.buffer)
}
}
}
struct ResponseState {
buffer: Vec<u8>,
is_chunked: bool,
found_end_marker: bool,
total_chunks_received: usize,
current_response_id: Option<String>,
current_response_logged: bool,
first_chunk_time: Option<tokio::time::Instant>,
force_log_time: Option<tokio::time::Instant>,
}
impl ResponseState {
fn reset_for_new_response(&mut self, request_id: Option<String>) {
self.buffer.clear();
self.is_chunked = false;
self.found_end_marker = false;
self.total_chunks_received = 0;
self.current_response_logged = false;
self.current_response_id = request_id;
self.first_chunk_time = Some(tokio::time::Instant::now());
self.force_log_time = None;
}
fn reset_for_next_response(&mut self) {
self.buffer.clear();
self.is_chunked = false;
self.found_end_marker = false;
self.total_chunks_received = 0;
self.first_chunk_time = None;
self.force_log_time = None;
}
}
#[cfg(test)]
mod tests {
use super::*;
const WEBSOCKET_HANDSHAKE: &[u8] = b"HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: test-accept\r\n\r\n";
async fn check_websocket_response_forwarding(
response: &[u8], split_at: usize, logging_enabled: bool, expect_tunnel: bool,
) {
let temp_dir = tempfile::tempdir().unwrap();
let log_path = temp_dir.path().join("websocket.log");
let http_logger = Arc::new(
Logger::new(
kftray_http_logs::LogConfig::new(temp_dir.path().to_path_buf()),
log_path.clone(),
)
.await
.unwrap(),
);
let req_id = http_logger
.log_request(bytes::Bytes::from_static(
b"GET /websocket HTTP/1.1\r\nHost: example.test\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n",
))
.await;
let logger = Arc::new(Mutex::new(logging_enabled.then(|| http_logger.clone())));
let request_id = Arc::new(Mutex::new(Some(req_id)));
let tunnel = Arc::new(AtomicBool::new(false));
let forward_tunnel = tunnel.clone();
let (mut upstream_reader, mut upstream_writer) = tokio::io::duplex(4096);
let (mut client_writer, mut client_reader) = tokio::io::duplex(4096);
let (events, subscriber) = tokio::sync::broadcast::channel(8);
let forward = tokio::spawn(async move {
TcpForwarder::forward_upstream_to_client(
logger,
1,
&mut upstream_reader,
&mut client_writer,
request_id,
CancellationToken::new(),
subscriber,
8080,
forward_tunnel,
)
.await
.unwrap();
});
tokio::time::timeout(Duration::from_secs(5), async {
upstream_writer
.write_all(&response[..split_at])
.await
.unwrap();
let mut first = vec![0; split_at];
client_reader.read_exact(&mut first).await.unwrap();
assert_eq!(first, response[..split_at]);
assert!(
!tunnel.load(Ordering::SeqCst),
"incomplete headers must not switch modes"
);
let frame = b"\x81\x0dframe-payload";
let mut last = response[split_at..].to_vec();
if expect_tunnel || !logging_enabled {
last.extend_from_slice(frame);
}
upstream_writer.write_all(&last).await.unwrap();
let mut received = vec![0; last.len()];
client_reader.read_exact(&mut received).await.unwrap();
assert_eq!(received, last, "coalesced frame bytes must be preserved");
assert_eq!(tunnel.load(Ordering::SeqCst), expect_tunnel);
if expect_tunnel || !logging_enabled {
events
.send(crate::kube::http_log_watcher::HttpLogStateEvent::new(
1, true,
))
.unwrap();
tokio::task::yield_now().await;
upstream_writer.write_all(frame).await.unwrap();
let mut received = vec![0; frame.len()];
client_reader.read_exact(&mut received).await.unwrap();
assert_eq!(received, frame, "later frame bytes must be preserved");
}
upstream_writer.shutdown().await.unwrap();
forward.await.unwrap();
})
.await
.expect("forwarding should not stall");
let status = std::str::from_utf8(response)
.unwrap()
.lines()
.next()
.unwrap();
if logging_enabled {
tokio::time::timeout(Duration::from_secs(5), async {
loop {
http_logger.flush().await.unwrap();
let contents = tokio::fs::read_to_string(&log_path).await.unwrap();
if contents.contains(status) {
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("the HTTP response must be logged before leaving HTTP mode");
}
http_logger.shutdown().await;
let contents = tokio::fs::read_to_string(&log_path).await.unwrap();
assert_eq!(
contents.matches(status).count(),
usize::from(logging_enabled)
);
assert!(
!contents.contains("frame-payload"),
"WebSocket frames are not HTTP log entries"
);
}
#[tokio::test]
async fn test_websocket_handshake_logged_before_tunnel() {
check_websocket_response_forwarding(WEBSOCKET_HANDSHAKE, 30, true, true).await;
}
#[tokio::test]
async fn test_websocket_rejection_stays_in_http_mode() {
check_websocket_response_forwarding(
b"HTTP/1.1 403 Forbidden\r\nContent-Length: 0\r\n\r\n",
15,
true,
false,
)
.await;
}
#[tokio::test]
async fn test_websocket_without_logging_preserves_bytes() {
check_websocket_response_forwarding(WEBSOCKET_HANDSHAKE, 30, false, true).await;
}
#[test]
fn test_upgrade_detector_is_bounded() {
let mut detector = UpgradeDetector::default();
assert!(!detector.observe(b"H"));
assert!(detector.observe(&WEBSOCKET_HANDSHAKE[1..]));
assert!(detector.headers.is_empty());
let mut oversized = b"HTTP/1.1 101 Switching Protocols\r\nX-Long: ".to_vec();
oversized.resize(BUFFER_SIZE * 2, b'x');
assert!(!detector.observe(&oversized));
assert!(detector.headers.is_empty());
assert!(!detector.observe(&vec![0x81; BUFFER_SIZE * 2]));
assert!(detector.headers.is_empty());
}
#[test]
fn test_upgrade_detector_does_not_classify_body_bytes() {
let mut detector = UpgradeDetector::default();
let response = format!(
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n",
WEBSOCKET_HANDSHAKE.len()
);
assert!(!detector.observe(response.as_bytes()));
assert!(!detector.observe(WEBSOCKET_HANDSHAKE));
assert!(detector.observe(WEBSOCKET_HANDSHAKE));
}
#[tokio::test]
async fn test_raw_tunnel_blocked_writes_are_cancellable() {
for upstream_to_client in [false, true] {
let (mut reader, mut source) = tokio::io::duplex(1024);
let (mut writer, _unread_peer) = tokio::io::duplex(1);
source.write_all(&[0x81; 64]).await.unwrap();
let token = CancellationToken::new();
let cancel = token.clone();
let (_events, receiver) = tokio::sync::broadcast::channel(1);
let mut forward = tokio::spawn(async move {
let logger = Arc::new(Mutex::new(None));
let request_id = Arc::new(Mutex::new(None));
let tunnel = Arc::new(AtomicBool::new(true));
if upstream_to_client {
TcpForwarder::forward_upstream_to_client(
logger,
1,
&mut reader,
&mut writer,
request_id,
token,
receiver,
8080,
tunnel,
)
.await
} else {
TcpForwarder::forward_client_to_upstream(
logger,
1,
&mut reader,
&mut writer,
request_id,
token,
receiver,
8080,
tunnel,
)
.await
}
});
tokio::time::sleep(Duration::from_millis(20)).await;
assert!(!forward.is_finished());
cancel.cancel();
let result = tokio::time::timeout(Duration::from_secs(1), &mut forward).await;
if result.is_err() {
forward.abort();
}
result
.expect("cancellation must settle a blocked raw write")
.unwrap()
.unwrap();
}
}
#[tokio::test]
async fn test_new() {
let forwarder = TcpForwarder::new(1, "pod".to_string());
assert_eq!(forwarder.config_id, 1);
assert_eq!(forwarder.workload_type, "pod");
assert!(forwarder.logger.is_none());
}
#[tokio::test]
async fn test_initialize_logger_non_service_workload() {
let mut forwarder = TcpForwarder::new(1, "deployment".to_string());
let result = forwarder.initialize_logger(8080).await;
assert!(result.is_ok());
assert!(forwarder.logger.is_none());
}
#[tokio::test]
async fn test_initialize_logger_state() {
let _db = kftray_commons::test_utils::test_db().await;
let mut forwarder = TcpForwarder::new(1, "pod".to_string());
let result = forwarder.initialize_logger(8080).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_forward_connection_basic() {
let (mut client_to_upstream, mut client_reader) = tokio::io::duplex(1024);
let (mut upstream_to_client, mut upstream_reader) = tokio::io::duplex(1024);
let test_request = b"GET /test HTTP/1.1\r\n\r\n";
let test_response = b"HTTP/1.1 200 OK\r\nContent-Length: 11\r\n\r\nHello World";
let (tx, rx) = tokio::sync::oneshot::channel::<Vec<u8>>();
let _forwarder = TcpForwarder::new(1, "pod".to_string());
let cancellation_token = CancellationToken::new();
let cancel_clone = cancellation_token.clone();
let upstream_task = tokio::spawn(async move {
let mut buf = vec![0u8; 1024];
let n = client_reader
.read(&mut buf)
.await
.expect("Read should succeed");
buf.truncate(n);
tx.send(buf).expect("Channel send should succeed");
upstream_to_client
.write_all(test_response)
.await
.expect("Write should succeed");
tokio::time::sleep(Duration::from_millis(50)).await;
cancel_clone.cancel();
});
let client_task = tokio::spawn(async move {
client_to_upstream
.write_all(test_request)
.await
.expect("Write should succeed");
client_to_upstream
.flush()
.await
.expect("Flush should succeed");
let mut response_buf = vec![0u8; 1024];
let n = upstream_reader
.read(&mut response_buf)
.await
.expect("Read should succeed");
response_buf.truncate(n);
assert_eq!(
response_buf, test_response,
"Response should match expected"
);
});
let received_request = tokio::time::timeout(Duration::from_millis(300), rx)
.await
.expect("Should complete in time")
.expect("Channel should return data");
assert_eq!(
received_request, test_request,
"Request should match expected"
);
let _ = tokio::join!(upstream_task, client_task);
}
#[tokio::test]
async fn test_handle_client_to_upstream_early_cancel() {
let (client_read, mut client_write) = tokio::io::duplex(1024);
let (upstream_write, mut upstream_read) = tokio::io::duplex(1024);
let forwarder = TcpForwarder::new(1, "pod".to_string());
let request_id = Arc::new(Mutex::new(None));
let cancellation_token = CancellationToken::new();
let cancel_clone = cancellation_token.clone();
tokio::spawn(async move {
client_write.write_all(b"test data").await.unwrap();
client_write.flush().await.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
cancel_clone.cancel();
});
tokio::spawn(async move {
let mut buf = [0u8; 1024];
let _ = upstream_read.read(&mut buf).await;
});
let mut client_reader = client_read;
let mut upstream_writer = upstream_write;
let (_, log_subscriber) = tokio::sync::broadcast::channel(1);
let logger = Arc::new(Mutex::new(None));
let websocket_tunnel_mode = Arc::new(AtomicBool::new(false));
let result = TcpForwarder::forward_client_to_upstream(
logger,
forwarder.config_id,
&mut client_reader,
&mut upstream_writer,
request_id,
cancellation_token,
log_subscriber,
8080,
websocket_tunnel_mode,
)
.await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_handle_upstream_to_client_early_cancel() {
let (upstream_read, mut upstream_write) = tokio::io::duplex(1024);
let (client_write, mut client_read) = tokio::io::duplex(1024);
let forwarder = TcpForwarder::new(1, "pod".to_string());
let request_id = Arc::new(Mutex::new(Some("req123".to_string())));
let cancellation_token = CancellationToken::new();
let cancel_clone = cancellation_token.clone();
tokio::spawn(async move {
upstream_write
.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 9\r\n\r\ntest data")
.await
.unwrap();
upstream_write.flush().await.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
cancel_clone.cancel();
});
tokio::spawn(async move {
let mut buf = [0u8; 1024];
let _ = client_read.read(&mut buf).await;
});
let mut upstream_reader = upstream_read;
let mut client_writer = client_write;
let (_, log_subscriber) = tokio::sync::broadcast::channel(1);
let logger = Arc::new(Mutex::new(None));
let websocket_tunnel_mode = Arc::new(AtomicBool::new(false));
let result = TcpForwarder::forward_upstream_to_client(
logger,
forwarder.config_id,
&mut upstream_reader,
&mut client_writer,
request_id,
cancellation_token,
log_subscriber,
8080,
websocket_tunnel_mode,
)
.await;
assert!(result.is_ok());
}
}