use std::sync::Arc;
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(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 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,
);
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,
);
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,
) -> 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 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 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,
) -> 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
};
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
&& 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 should_log
&& let Some(ref mut state) = response_state.as_mut() {
Self::handle_response_logging_static(&buffer[..n], state, &logger, &request_id).await;
}
if let Err(e) = client_writer.write_all(&buffer[..n]).await {
return Err(e.into());
}
},
log_event = log_subscriber.recv() => {
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>>>,
) {
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);
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();
}
}
}
}
}
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::*;
#[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 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 result = TcpForwarder::forward_client_to_upstream(
logger,
forwarder.config_id,
&mut client_reader,
&mut upstream_writer,
request_id,
cancellation_token,
log_subscriber,
8080,
)
.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 result = TcpForwarder::forward_upstream_to_client(
logger,
forwarder.config_id,
&mut upstream_reader,
&mut client_writer,
request_id,
cancellation_token,
log_subscriber,
8080,
)
.await;
assert!(result.is_ok());
}
}