use std::{num::NonZeroUsize, time::Duration};
use rama_core::{
Service,
error::{BoxError, ErrorContext},
io::{
HeapReader, PeekIoProvider, PrefixedIo,
peek::{
PeekOutput, PeekStopReason, PeekTimeoutError, PeekTimeoutPolicy, PeekVerdict,
peek_input_until_verdict_growing_with_output, peek_input_until_verdict_with_options,
},
},
service::RejectService,
telemetry::tracing,
};
use rama_utils::octets::kib;
use tokio::time::Instant;
use crate::client::{ClientHello, ClientHelloHandshakePrefix, parse_client_hello_handshake_prefix};
use super::NoTlsRejectError;
#[derive(Debug, Clone)]
pub struct PeekTlsClientHelloService<S, F = RejectService<(), NoTlsRejectError>> {
service: S,
fallback: F,
peek_timeout: Option<Duration>,
peek_timeout_policy: PeekTimeoutPolicy,
}
impl<S> PeekTlsClientHelloService<S> {
pub fn new(service: S) -> Self {
Self {
service,
fallback: RejectService::new(NoTlsRejectError),
peek_timeout: None,
peek_timeout_policy: PeekTimeoutPolicy::default(),
}
}
pub fn with_fallback<F>(self, fallback: F) -> PeekTlsClientHelloService<S, F> {
PeekTlsClientHelloService {
service: self.service,
fallback,
peek_timeout: self.peek_timeout,
peek_timeout_policy: self.peek_timeout_policy,
}
}
}
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
}
}
rama_utils::macros::generate_set_and_with! {
pub fn peek_timeout_policy(mut self, peek_timeout_policy: PeekTimeoutPolicy) -> Self {
self.peek_timeout_policy = peek_timeout_policy;
self
}
}
}
pub async fn peek_client_hello_from_input<PeekableInput>(
input: PeekableInput,
timeout: Option<Duration>,
) -> Result<
(
PeekableInput::Mapped<TlsClientHelloPrefixedIo<PeekableInput::PeekIo>>,
Option<ClientHello>,
),
std::io::Error,
>
where
PeekableInput: PeekIoProvider<PeekIo: Unpin>,
{
peek_client_hello_from_input_with_timeout_policy(input, timeout, PeekTimeoutPolicy::FailOpen)
.await
}
pub async fn peek_client_hello_from_input_with_timeout_policy<PeekableInput>(
mut input: PeekableInput,
timeout: Option<Duration>,
timeout_policy: PeekTimeoutPolicy,
) -> 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,
stop_reason,
} = peek_input_until_verdict_with_options(
peekable_io,
&mut peek_buf,
0,
timeout,
None,
super::peek::tls_record_header_verdict,
)
.await;
if stop_reason == PeekStopReason::Timeout && timeout_policy == PeekTimeoutPolicy::FailClosed {
return Err(PeekTimeoutError::new().into());
}
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));
}
const MAX_TLS_RECORD: usize = kib(16) + TLS_HEADER_PEEK_LEN;
let read_chunk = NonZeroUsize::new(1024).unwrap_or(NonZeroUsize::MIN);
let new_timeout = timeout.map(|t| t.saturating_sub(start.elapsed()));
let mut v = Vec::with_capacity(read_chunk.get() + TLS_HEADER_PEEK_LEN);
v.extend_from_slice(&peek_buf[..]);
let output = peek_input_until_verdict_growing_with_output(
peekable_io,
&mut v,
MAX_TLS_RECORD,
read_chunk,
new_timeout,
|buffer| match parse_client_hello_handshake_prefix(buffer) {
ClientHelloHandshakePrefix::Complete(hello) => PeekVerdict::Match(hello),
ClientHelloHandshakePrefix::Incomplete => PeekVerdict::NeedMore,
ClientHelloHandshakePrefix::Invalid => PeekVerdict::Reject,
},
)
.await;
if output.stop_reason == PeekStopReason::Timeout
&& timeout_policy == PeekTimeoutPolicy::FailClosed
{
return Err(PeekTimeoutError::new().into());
}
let prefix_data = HeapReader::from(v);
let peeked_input = input.map_peek_io(|io| PrefixedIo::new(prefix_data, io));
Ok((peeked_input, output.data))
}
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_with_timeout_policy(
input,
self.peek_timeout,
self.peek_timeout_policy,
)
.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,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
};
use tokio::io::{AsyncReadExt as _, AsyncWriteExt 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_partial_client_hello_does_not_prefix_zero_padding() {
const SPLIT_AT: usize = 64;
const PEEK_TIMEOUT: Duration = Duration::from_millis(10);
let (mut client, server) = tokio::io::duplex(CH_ONE_ONE_ONE_ONE.len());
let (prefix_sent_tx, prefix_sent_rx) = tokio::sync::oneshot::channel();
let (continue_tx, continue_rx) = tokio::sync::oneshot::channel();
let writer = tokio::spawn(async move {
client
.write_all(&CH_ONE_ONE_ONE_ONE[..SPLIT_AT])
.await
.unwrap();
prefix_sent_tx.send(()).unwrap();
continue_rx.await.unwrap();
client
.write_all(&CH_ONE_ONE_ONE_ONE[SPLIT_AT..])
.await
.unwrap();
});
prefix_sent_rx.await.unwrap();
let (mut input, client_hello) =
peek_client_hello_from_input(ServiceInput::new(server), Some(PEEK_TIMEOUT))
.await
.unwrap();
assert!(client_hello.is_none());
continue_tx.send(()).unwrap();
let mut bytes = Vec::new();
input.read_to_end(&mut bytes).await.unwrap();
writer.await.unwrap();
assert_eq!(CH_ONE_ONE_ONE_ONE, bytes);
}
#[tokio::test]
async fn partial_client_hello_fail_closed_returns_timeout_without_fallback() {
const SPLIT_AT: usize = 64;
let fallback_calls = Arc::new(AtomicUsize::new(0));
let fallback_calls_for_service = Arc::clone(&fallback_calls);
let fallback = service_fn(move || {
fallback_calls_for_service.fetch_add(1, Ordering::SeqCst);
async { Ok::<_, Infallible>("fallback") }
});
let router =
PeekTlsClientHelloService::new(service_fn(async || Ok::<_, Infallible>("tls")))
.with_peek_timeout(Duration::from_millis(10))
.with_peek_timeout_policy(PeekTimeoutPolicy::FailClosed)
.with_fallback(fallback);
let (mut client, server) = tokio::io::duplex(CH_ONE_ONE_ONE_ONE.len());
client
.write_all(&CH_ONE_ONE_ONE_ONE[..SPLIT_AT])
.await
.unwrap();
let error = router.serve(ServiceInput::new(server)).await.unwrap_err();
assert!(error_chain_has_peek_timeout(error.as_ref()));
assert_eq!(0, fallback_calls.load(Ordering::SeqCst));
}
#[tokio::test]
async fn malformed_client_hello_falls_back_under_both_timeout_policies() {
let malformed = [
0x16, 0x03, 0x03, 0x00, 0x06, 0x02, 0x00, 0x00, 0x02, 0xaa, 0xbb,
];
for policy in [PeekTimeoutPolicy::FailOpen, PeekTimeoutPolicy::FailClosed] {
let router =
PeekTlsClientHelloService::new(service_fn(async || Ok::<_, Infallible>("tls")))
.with_peek_timeout(Duration::from_millis(10))
.with_peek_timeout_policy(policy)
.with_fallback(service_fn(async || Ok::<_, Infallible>("fallback")));
let response = router
.serve(ServiceInput::new(std::io::Cursor::new(malformed.to_vec())))
.await
.unwrap();
assert_eq!("fallback", response);
}
}
fn error_chain_has_peek_timeout(mut error: &(dyn std::error::Error + 'static)) -> bool {
loop {
if error.downcast_ref::<PeekTimeoutError>().is_some()
|| error
.downcast_ref::<std::io::Error>()
.is_some_and(|error| error.kind() == std::io::ErrorKind::TimedOut)
{
return true;
}
let Some(source) = error.source() else {
return false;
};
error = source;
}
}
#[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);
}
}