Skip to content

WIP - Feat: buf polling #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Mar 19, 2025
Merged

WIP - Feat: buf polling #4

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
02bea1b
wip
junkurihara Mar 19, 2025
bdd31cb
wip
junkurihara Mar 19, 2025
a2e6e06
update tls probe buffering
junkurihara Mar 19, 2025
cc43701
update mimalloc
junkurihara Mar 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
wip
  • Loading branch information
@junkurihara
junkurihara committed Mar 19, 2025
commit 02bea1b5f984e8b0bc77c6ff94598976c5b0a37f
2 changes: 1 addition & 1 deletion proxy-l4-lib/src/constants.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ pub const TCP_PROTOCOL_DETECTION_TIMEOUT_MSEC: u64 = 100;
/// But considering the hybrid post-quantum key exchange (key_share extension is > 1KB in X25519MLKEM768),
/// the buffer size should be large, at least 2KB, to parse the Client Hello message.
/// https://datatracker.ietf.org/doc/html/rfc8446#section-5.1
pub const TCP_PROTOCOL_DETECTION_BUFFER_SIZE: usize = 4096;
pub const TCP_PROTOCOL_DETECTION_BUFFER_SIZE: usize = 16384;

/// UDP buffer size, theoretical limit is 65535 bytes in IPv4
/// But the practical limit is, due to the MTU, less than 1500 bytes.
Expand Down
76 changes: 40 additions & 36 deletions proxy-l4-lib/src/tcp_proxy.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use crate::{
};
use std::{net::SocketAddr, sync::Arc};
use tokio::{
io::copy_bidirectional,
io::{copy_bidirectional, AsyncReadExt},
net::TcpStream,
time::{timeout, Duration},
};
Expand Down Expand Up @@ -239,40 +239,43 @@ impl TcpProxyProtocol {
/// TODO: In TLS, we can get the length of ClientHello payload from its header.
/// TODO: Thus, we should use `stream.read_exact` method for the fetching.
/// TODO: This consumes the stream queue, and hence we need change the handling of the first packets of all types TCP stream.
pub(crate) async fn detect_protocol(incoming_stream: &TcpStream) -> Result<Self, ProxyError> {
let mut buf = vec![0u8; TCP_PROTOCOL_DETECTION_BUFFER_SIZE];
let read_len = peek_tcp_stream(incoming_stream, &mut buf).await?;

// TODO: Add more protocol detection
if buf.starts_with(b"SSH-") {
debug!("SSH connection detected");
return Ok(Self::Ssh);
}

// TODO: Refactor this part to get the exact length of the ClientHello payload
if let Some(res) = probe_tls_handshake(&buf.as_slice()[..read_len]) {
let read_again_len = match res {
TlsProbeResult::Success(info) => {
debug!("TLS connection detected");
return Ok(Self::Tls(info));
}
TlsProbeResult::PeekMore => {
debug!("TLS connection detected, but need more data");
peek_tcp_stream(incoming_stream, &mut buf).await?
}
};
let res = probe_tls_handshake(&buf.as_slice()[..read_again_len]);
if let Some(TlsProbeResult::Success(info)) = res {
debug!("TLS connection detected");
return Ok(Self::Tls(info));
}
debug!("Peeked again, but failed to get enough data of TLS Client Hello");
}

if buf.windows(4).any(|w| w.eq(b"HTTP")) {
debug!("HTTP connection detected");
return Ok(Self::Http);
}
pub(crate) async fn detect_protocol(incoming_stream: &mut TcpStream, buf: &mut [u8]) -> Result<Self, ProxyError> {
// Read the first several byte
let red_len = incoming_stream.read(buf).await?;

// let mut buf = vec![0u8; TCP_PROTOCOL_DETECTION_BUFFER_SIZE];
// let read_len = peek_tcp_stream(incoming_stream, &mut buf).await?;

// // TODO: Add more protocol detection
// if buf.starts_with(b"SSH-") {
// debug!("SSH connection detected");
// return Ok(Self::Ssh);
// }

// // TODO: Refactor this part to get the exact length of the ClientHello payload
// if let Some(res) = probe_tls_handshake(&buf.as_slice()[..read_len]) {
// let read_again_len = match res {
// TlsProbeResult::Success(info) => {
// debug!("TLS connection detected");
// return Ok(Self::Tls(info));
// }
// TlsProbeResult::PeekMore => {
// debug!("TLS connection detected, but need more data");
// peek_tcp_stream(incoming_stream, &mut buf).await?
// }
// };
// let res = probe_tls_handshake(&buf.as_slice()[..read_again_len]);
// if let Some(TlsProbeResult::Success(info)) = res {
// debug!("TLS connection detected");
// return Ok(Self::Tls(info));
// }
// debug!("Peeked again, but failed to get enough data of TLS Client Hello");
// }

// if buf.windows(4).any(|w| w.eq(b"HTTP")) {
// debug!("HTTP connection detected");
// return Ok(Self::Http);
// }

debug!("Untyped TCP connection");
Ok(Self::Any)
Expand Down Expand Up @@ -337,7 +340,8 @@ impl TcpProxy {
let dst_mux = Arc::clone(&self.destination_mux);
let connection_count = self.connection_count.clone();
async move {
let protocol = match TcpProxyProtocol::detect_protocol(&incoming_stream).await {
let mut buf = vec![0u8; TCP_PROTOCOL_DETECTION_BUFFER_SIZE];
let protocol = match TcpProxyProtocol::detect_protocol(&mut incoming_stream, &mut buf).await {
Ok(p) => p,
Err(e) => {
error!("Failed to detect protocol: {e}");
Expand Down