use std::time::Duration;
use rama_core::{
Service,
error::{BoxError, ErrorContext},
io::{
HeapReader, PeekIoProvider, PrefixedIo,
peek::{PeekOutput, peek_input_until, peek_input_until_with_offset},
},
service::RejectService,
telemetry::tracing,
};
use rama_utils::octets::kib;
use tokio::time::Instant;
use crate::client::{ClientHello, parse_client_hello_handshake};
use super::NoTlsRejectError;
#[derive(Debug, Clone)]
pub struct PeekTlsClientHelloService<S, F = RejectService<(), NoTlsRejectError>> {
service: S,
fallback: F,
peek_timeout: Option<Duration>,
}
impl<S> PeekTlsClientHelloService<S> {
pub fn new(service: S) -> Self {
Self {
service,
fallback: RejectService::new(NoTlsRejectError),
peek_timeout: None,
}
}
pub fn with_fallback<F>(self, fallback: F) -> PeekTlsClientHelloService<S, F> {
PeekTlsClientHelloService {
service: self.service,
fallback,
peek_timeout: self.peek_timeout,
}
}
}
impl<S, F> PeekTlsClientHelloService<S, F> {
rama_utils::macros::generate_set_and_with! {
pub fn peek_timeout(mut self, peek_timeout: Option<Duration>) -> Self {
self.peek_timeout = peek_timeout;
self
}
}
}
pub async fn peek_client_hello_from_input<PeekableInput>(
mut input: PeekableInput,
timeout: Option<Duration>,
) -> Result<
(
PeekableInput::Mapped<TlsClientHelloPrefixedIo<PeekableInput::PeekIo>>,
Option<ClientHello>,
),
std::io::Error,
>
where
PeekableInput: PeekIoProvider<PeekIo: Unpin>,
{
let mut peek_buf = [0u8; TLS_HEADER_PEEK_LEN];
let peekable_io = input.peek_io_mut();
let start = Instant::now();
let PeekOutput { data, peek_size } =
peek_input_until(peekable_io, &mut peek_buf, timeout, |buffer| {
if buffer.len() == TLS_HEADER_PEEK_LEN
&& matches!(buffer, [0x16, 0x03, 0x00..=0x04, ..])
{
Some(())
} else {
None
}
})
.await;
let is_tls = data.is_some();
tracing::trace!("tls prefix header read (is tls: {is_tls})");
if !is_tls {
if TLS_HEADER_PEEK_LEN.saturating_sub(peek_size) > 0 {
tracing::trace!(
"move tls peek buffer cursor due to reading not enough (read: {peek_size})"
);
}
let prefix_data = HeapReader::from(&peek_buf[..peek_size.min(TLS_HEADER_PEEK_LEN)]);
let peeked_input = input.map_peek_io(|io| PrefixedIo::new(prefix_data, io));
tracing::trace!("return early for non-tls traffic: missing peek header");
return Ok((peeked_input, None));
}
let n = ((peek_buf[3] as usize) << 8) | (peek_buf[4] as usize);
const MAX_TLS_RECORD_BODY: usize = kib(16);
let record_size = (n + TLS_HEADER_PEEK_LEN).min(MAX_TLS_RECORD_BODY + TLS_HEADER_PEEK_LEN);
let mut v = vec![0u8; record_size];
v[..TLS_HEADER_PEEK_LEN].copy_from_slice(&peek_buf[..]);
let new_timeout = timeout.map(|t| t.saturating_sub(start.elapsed()));
let PeekOutput {
data: maybe_client_hello,
peek_size,
} = peek_input_until_with_offset(
peekable_io,
&mut v,
TLS_HEADER_PEEK_LEN,
new_timeout,
|buffer| {
let n = buffer.len();
parse_client_hello_handshake(buffer)
.inspect_err(|err| {
tracing::debug!("failed parse client hello handshake ({n}) byte(s): {err}",)
})
.ok()
},
)
.await;
let new_peek_size = peek_size.saturating_sub(TLS_HEADER_PEEK_LEN);
if new_peek_size != n {
tracing::trace!(
peek_size = new_peek_size,
expected_peek_size = n,
"unexpected read size for client hello handshake data: try regardless..."
);
}
let prefix_data = HeapReader::from(v);
let peeked_input = input.map_peek_io(|io| PrefixedIo::new(prefix_data, io));
Ok((peeked_input, maybe_client_hello))
}
impl<PeekableInput, Output, S, F> Service<PeekableInput> for PeekTlsClientHelloService<S, F>
where
PeekableInput: PeekIoProvider<PeekIo: Unpin>,
Output: Send + 'static,
S: Service<
InputWithClientHello<
PeekableInput::Mapped<TlsClientHelloPrefixedIo<PeekableInput::PeekIo>>,
>,
Output = Output,
Error: Into<BoxError>,
>,
F: Service<
PeekableInput::Mapped<TlsClientHelloPrefixedIo<PeekableInput::PeekIo>>,
Output = Output,
Error: Into<BoxError>,
>,
{
type Output = Output;
type Error = BoxError;
async fn serve(&self, input: PeekableInput) -> Result<Self::Output, Self::Error> {
let (peeked_input, maybe_client_hello) =
peek_client_hello_from_input(input, self.peek_timeout)
.await
.context("I/O error while peeking TLS:CH from existing input")?;
if let Some(client_hello) = maybe_client_hello {
self.service
.serve(InputWithClientHello {
input: peeked_input,
client_hello,
})
.await
.map_err(Into::into)
} else {
self.fallback.serve(peeked_input).await.map_err(Into::into)
}
}
}
const TLS_HEADER_PEEK_LEN: usize = 5;
pub type TlsClientHelloPrefixedIo<S> = PrefixedIo<HeapReader, S>;
#[derive(Debug, Clone)]
pub struct InputWithClientHello<Input> {
pub input: Input,
pub client_hello: ClientHello,
}
#[cfg(test)]
mod test {
use rama_core::{
ServiceInput,
service::{RejectError, service_fn},
};
use std::convert::Infallible;
use tokio::io::AsyncReadExt as _;
use rama_core::io::Io;
use super::*;
const CH_ONE_ONE_ONE_ONE: &[u8] = &[
0x16, 0x03, 0x01, 0x02, 0x00, 0x01, 0x00, 0x01, 0xfc, 0x03, 0x03, 0x02, 0x15, 0xfd, 0xe2,
0x92, 0xc0, 0x46, 0x9f, 0x92, 0xbe, 0xd7, 0xe9, 0x1a, 0x3c, 0x50, 0x5e, 0x55, 0x49, 0x17,
0xa6, 0xf8, 0xa5, 0xca, 0xa4, 0x6d, 0x60, 0xcc, 0xea, 0xf7, 0x25, 0xf0, 0x6e, 0x20, 0x41,
0x20, 0x18, 0x66, 0x5c, 0xae, 0x08, 0xb0, 0x10, 0x96, 0x3c, 0xad, 0xb4, 0x13, 0xe1, 0x92,
0xce, 0x96, 0xad, 0x9d, 0x45, 0x05, 0xb7, 0xa6, 0x4c, 0x01, 0x71, 0x08, 0x74, 0x0d, 0x1f,
0x35, 0x00, 0x2a, 0x3a, 0x3a, 0x13, 0x01, 0x13, 0x02, 0x13, 0x03, 0xc0, 0x2c, 0xc0, 0x2b,
0xcc, 0xa9, 0xc0, 0x30, 0xc0, 0x2f, 0xcc, 0xa8, 0xc0, 0x0a, 0xc0, 0x09, 0xc0, 0x14, 0xc0,
0x13, 0x00, 0x9d, 0x00, 0x9c, 0x00, 0x35, 0x00, 0x2f, 0xc0, 0x08, 0xc0, 0x12, 0x00, 0x0a,
0x01, 0x00, 0x01, 0x89, 0xda, 0xda, 0x00, 0x00, 0x00, 0x00, 0x00, 0x14, 0x00, 0x12, 0x00,
0x00, 0x0f, 0x6f, 0x6e, 0x65, 0x2e, 0x6f, 0x6e, 0x65, 0x2e, 0x6f, 0x6e, 0x65, 0x2e, 0x6f,
0x6e, 0x65, 0x00, 0x17, 0x00, 0x00, 0xff, 0x01, 0x00, 0x01, 0x00, 0x00, 0x0a, 0x00, 0x0c,
0x00, 0x0a, 0xfa, 0xfa, 0x00, 0x1d, 0x00, 0x17, 0x00, 0x18, 0x00, 0x19, 0x00, 0x0b, 0x00,
0x02, 0x01, 0x00, 0x00, 0x05, 0x00, 0x05, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x0d, 0x00,
0x16, 0x00, 0x14, 0x04, 0x03, 0x08, 0x04, 0x04, 0x01, 0x05, 0x03, 0x08, 0x05, 0x08, 0x05,
0x05, 0x01, 0x08, 0x06, 0x06, 0x01, 0x02, 0x01, 0x00, 0x12, 0x00, 0x00, 0x00, 0x33, 0x00,
0x2b, 0x00, 0x29, 0xfa, 0xfa, 0x00, 0x01, 0x00, 0x00, 0x1d, 0x00, 0x20, 0x7c, 0xe1, 0xc6,
0xc2, 0x01, 0x69, 0x42, 0xba, 0x2b, 0xec, 0x07, 0x2f, 0x04, 0xbd, 0xb6, 0x2a, 0x7e, 0x04,
0x6b, 0x96, 0x98, 0x51, 0x4e, 0x80, 0xb3, 0x2a, 0x4c, 0x4f, 0x1f, 0x39, 0x82, 0x2b, 0x00,
0x2d, 0x00, 0x02, 0x01, 0x01, 0x00, 0x2b, 0x00, 0x0b, 0x0a, 0x6a, 0x6a, 0x03, 0x04, 0x03,
0x03, 0x03, 0x02, 0x03, 0x01, 0x00, 0x1b, 0x00, 0x03, 0x02, 0x00, 0x01, 0x3a, 0x3a, 0x00,
0x01, 0x00, 0x00, 0x15, 0x00, 0xd3, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
];
const TLS_BUT_NO_SNI: &[u8] = &[
0x16, 0x03, 0x01, 0x02, 0x00, 0x01, 0x00, 0x01, 0xfc, 0x03, 0x03, 0x28, 0x5b, 0x8f, 0x90,
0x22, 0x2a, 0x90, 0x95, 0x89, 0xa9, 0x62, 0x1f, 0xdb, 0x68, 0xbe, 0x4c, 0x0e, 0xdf, 0xe4,
0x76, 0x50, 0x48, 0xa5, 0x40, 0x56, 0x5f, 0x9a, 0xba, 0x19, 0x29, 0x66, 0xdd, 0x20, 0x7a,
0x7f, 0x7e, 0xc7, 0xbd, 0xfb, 0x88, 0x07, 0xd9, 0xf5, 0x99, 0xfa, 0xf3, 0x0d, 0x37, 0x30,
0x52, 0x4d, 0x44, 0xe4, 0x26, 0xc0, 0xd1, 0x9a, 0xcd, 0x78, 0xf6, 0x7a, 0xf1, 0x7a, 0x66,
0xe1, 0x00, 0x3e, 0x13, 0x02, 0x13, 0x03, 0x13, 0x01, 0xc0, 0x2c, 0xc0, 0x30, 0x00, 0x9f,
0xcc, 0xa9, 0xcc, 0xa8, 0xcc, 0xaa, 0xc0, 0x2b, 0xc0, 0x2f, 0x00, 0x9e, 0xc0, 0x24, 0xc0,
0x28, 0x00, 0x6b, 0xc0, 0x23, 0xc0, 0x27, 0x00, 0x67, 0xc0, 0x0a, 0xc0, 0x14, 0x00, 0x39,
0xc0, 0x09, 0xc0, 0x13, 0x00, 0x33, 0x00, 0x9d, 0x00, 0x9c, 0x00, 0x3d, 0x00, 0x3c, 0x00,
0x35, 0x00, 0x2f, 0x00, 0xff, 0x01, 0x00, 0x01, 0x75, 0x00, 0x0b, 0x00, 0x04, 0x03, 0x00,
0x01, 0x02, 0x00, 0x0a, 0x00, 0x16, 0x00, 0x14, 0x00, 0x1d, 0x00, 0x17, 0x00, 0x1e, 0x00,
0x19, 0x00, 0x18, 0x01, 0x00, 0x01, 0x01, 0x01, 0x02, 0x01, 0x03, 0x01, 0x04, 0x00, 0x10,
0x00, 0x0e, 0x00, 0x0c, 0x02, 0x68, 0x32, 0x08, 0x68, 0x74, 0x74, 0x70, 0x2f, 0x31, 0x2e,
0x31, 0x00, 0x16, 0x00, 0x00, 0x00, 0x17, 0x00, 0x00, 0x00, 0x31, 0x00, 0x00, 0x00, 0x0d,
0x00, 0x2a, 0x00, 0x28, 0x04, 0x03, 0x05, 0x03, 0x06, 0x03, 0x08, 0x07, 0x08, 0x08, 0x08,
0x09, 0x08, 0x0a, 0x08, 0x0b, 0x08, 0x04, 0x08, 0x05, 0x08, 0x06, 0x04, 0x01, 0x05, 0x01,
0x06, 0x01, 0x03, 0x03, 0x03, 0x01, 0x03, 0x02, 0x04, 0x02, 0x05, 0x02, 0x06, 0x02, 0x00,
0x2b, 0x00, 0x05, 0x04, 0x03, 0x04, 0x03, 0x03, 0x00, 0x2d, 0x00, 0x02, 0x01, 0x01, 0x00,
0x33, 0x00, 0x26, 0x00, 0x24, 0x00, 0x1d, 0x00, 0x20, 0xe0, 0xb9, 0xfb, 0x5a, 0xd5, 0x60,
0x30, 0x39, 0xad, 0xfb, 0xd3, 0x94, 0xa2, 0xff, 0x08, 0x71, 0x9b, 0xcc, 0x6f, 0xbe, 0x9e,
0xcc, 0x7b, 0xad, 0x3c, 0xd0, 0xde, 0xe8, 0x3e, 0x5d, 0xba, 0x6b, 0x00, 0x15, 0x00, 0xca,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
];
#[tokio::test]
async fn test_client_hello_peek_service() {
let tls_service = service_fn(async |input: InputWithClientHello<_>| {
let sni = input
.client_hello
.ext_server_name()
.map(ToString::to_string);
Ok::<_, Infallible>(sni)
});
let plain_service = service_fn(async || Ok::<_, Infallible>(Some("plain".to_owned())));
let peek_tls_svc = PeekTlsClientHelloService::new(tls_service).with_fallback(plain_service);
let response = peek_tls_svc
.serve(ServiceInput::new(std::io::Cursor::new(b"".to_vec())))
.await
.unwrap();
assert_eq!(Some("plain".to_owned()), response);
let response = peek_tls_svc
.serve(ServiceInput::new(std::io::Cursor::new(
CH_ONE_ONE_ONE_ONE.to_vec(),
)))
.await
.unwrap();
assert_eq!(Some("one.one.one.one".to_owned()), response);
let response = peek_tls_svc
.serve(ServiceInput::new(std::io::Cursor::new(b"foo".to_vec())))
.await
.unwrap();
assert_eq!(Some("plain".to_owned()), response);
let response = peek_tls_svc
.serve(ServiceInput::new(std::io::Cursor::new(b"foobar".to_vec())))
.await
.unwrap();
assert_eq!(Some("plain".to_owned()), response);
let response = peek_tls_svc
.serve(ServiceInput::new(std::io::Cursor::new(
TLS_BUT_NO_SNI.to_vec(),
)))
.await
.unwrap();
assert_eq!(None, response);
}
#[tokio::test]
async fn test_peek_router_read_eof() {
async fn tls_service_fn(
InputWithClientHello {
mut input,
client_hello,
}: InputWithClientHello<impl Io + Unpin>,
) -> Result<&'static str, BoxError> {
let mut v = Vec::default();
_ = input.read_to_end(&mut v).await?;
assert_eq!(CH_ONE_ONE_ONE_ONE, v);
assert!(client_hello.ext_server_name().is_some());
assert_eq!(
"one.one.one.one",
client_hello.ext_server_name().unwrap().to_string()
);
Ok("ok")
}
let tls_service = service_fn(tls_service_fn);
let peek_tls_svc =
PeekTlsClientHelloService::new(tls_service).with_fallback(RejectService::<
&'static str,
RejectError,
>::new(
RejectError::default()
));
let response = peek_tls_svc
.serve(ServiceInput::new(std::io::Cursor::new(
CH_ONE_ONE_ONE_ONE.to_vec(),
)))
.await
.unwrap();
assert_eq!("ok", response);
}
#[tokio::test]
async fn test_peek_router_read_no_tls_eof() {
let cases = ["", "foo", "abcd", "abcde", "foobarbazbananas"];
for content in cases {
async fn tls_service_fn() -> Result<Vec<u8>, BoxError> {
Ok("tls".as_bytes().to_vec())
}
let tls_service = service_fn(tls_service_fn);
async fn plain_service_fn(mut stream: impl Io + Unpin) -> Result<Vec<u8>, BoxError> {
let mut v = Vec::default();
_ = stream.read_to_end(&mut v).await?;
Ok(v)
}
let plain_service = service_fn(plain_service_fn);
let peek_tls_svc =
PeekTlsClientHelloService::new(tls_service).with_fallback(plain_service);
let response = peek_tls_svc
.serve(ServiceInput::new(std::io::Cursor::new(
content.as_bytes().to_vec(),
)))
.await
.unwrap();
assert_eq!(content.as_bytes(), &response[..]);
}
}
#[tokio::test]
async fn test_peek_router_read_tls_no_sni_eof() {
async fn tls_service_fn(
InputWithClientHello {
mut input,
client_hello,
}: InputWithClientHello<impl Io + Unpin>,
) -> Result<&'static str, BoxError> {
let mut v = Vec::default();
_ = input.read_to_end(&mut v).await?;
assert_eq!(TLS_BUT_NO_SNI, v);
assert!(client_hello.ext_server_name().is_none());
Ok("ok")
}
let tls_service = service_fn(tls_service_fn);
let peek_tls_svc =
PeekTlsClientHelloService::new(tls_service).with_fallback(RejectService::<
&'static str,
RejectError,
>::new(
RejectError::default()
));
let response = peek_tls_svc
.serve(ServiceInput::new(std::io::Cursor::new(
TLS_BUT_NO_SNI.to_vec(),
)))
.await
.unwrap();
assert_eq!("ok", response);
}
}