Implemented AsyncReadCapture and AsyncReadWrite

These are utility types which will help capture the ClientHello as it is
being read and parsed, as well as any data following it which might get
buffered in the acceptore. This way we can get that data back out in
case we want to switch into transparent TCP mode again.
This commit is contained in:
Brian Picciano 2023-07-22 13:40:19 +02:00
parent b7289d7e7e
commit 077565b908
2 changed files with 140 additions and 1 deletions

View File

@ -359,5 +359,106 @@ impl<'a, 'b, T: AsyncRead + Unpin> Read for SyncReadAdapter<'a, 'b, T> {
}
}
/// Wraps an AsyncRead and AsyncWrite instance together to produce a single type which implements
/// AsyncRead + AsyncWrite.
pub struct AsyncReadWrite<R, W> {
r: Pin<Box<R>>,
w: Pin<Box<W>>,
}
impl<R, W> AsyncReadWrite<R, W>
where
R: Unpin,
W: Unpin,
{
pub fn new(r: R, w: W) -> Self {
Self {
r: Box::pin(r),
w: Box::pin(w),
}
}
pub fn into_inner(self) -> (R, W) {
(*Pin::into_inner(self.r), *Pin::into_inner(self.w))
}
}
impl<R, W> AsyncRead for AsyncReadWrite<R, W>
where
R: AsyncRead + Unpin,
{
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
self.r.as_mut().poll_read(cx, buf)
}
}
impl<R, W> AsyncWrite for AsyncReadWrite<R, W>
where
W: AsyncWrite + Unpin,
{
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.w.as_mut().poll_write(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.w.as_mut().poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.w.as_mut().poll_shutdown(cx)
}
}
/// Wraps an AsyncRead in order to capture all bytes which have been read by it into an internal
/// buffer.
pub struct AsyncReadCapture<R> {
r: Pin<Box<R>>,
buf: Vec<u8>,
}
impl<R> AsyncReadCapture<R>
where
R: AsyncRead + Unpin,
{
/// Initializes an AsyncReadCapture with an empty internal buffer of the given size.
pub fn with_capacity(r: R, cap: usize) -> Self {
Self {
r: Box::pin(r),
buf: Vec::with_capacity(cap),
}
}
pub fn into_inner(self) -> (R, Vec<u8>) {
(*Pin::into_inner(self.r), self.buf)
}
}
impl<R> AsyncRead for AsyncReadCapture<R>
where
R: AsyncRead,
{
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let res = self.r.as_mut().poll_read(cx, buf);
if let Poll::Ready(Ok(())) = res {
self.buf.extend_from_slice(buf.filled());
}
res
}
}
#[cfg(test)]
mod test_stream;

View File

@ -49,7 +49,7 @@ pub mod client;
mod common;
pub mod server;
use common::{MidHandshake, Stream, TlsState};
use common::{AsyncReadCapture, AsyncReadWrite, MidHandshake, Stream, TlsState};
use rustls::{ClientConfig, ClientConnection, CommonState, ServerConfig, ServerConnection};
use std::future::Future;
use std::io;
@ -333,6 +333,44 @@ where
}
}
type IOWithCapture<IO> = AsyncReadWrite<AsyncReadCapture<IO>, IO>;
pub struct TransparentConfigAcceptor<IO> {
acceptor: LazyConfigAcceptor<IOWithCapture<IO>>,
}
impl<IO> TransparentConfigAcceptor<IO>
where
IO: AsyncRead + AsyncWrite + Unpin,
{
pub fn new(acceptor: rustls::server::Acceptor, io: IO) -> Self {
let r = AsyncReadCapture::with_capacity(io, 512);
let rw = AsyncReadWrite::new(r, io);
Self {
acceptor: LazyConfigAcceptor::new(acceptor, rw),
}
}
//pub fn into_lazy_config_acceptor(self) -> LazyConfigAcceptor<IO> {
// self.acceptor
//}
}
//impl<IO> Future for TransparentConfigAcceptor<IO>
//where
// IO: AsyncRead + AsyncWrite + Unpin,
//{
// type Output = io::Result<TransparentStartHandshake<IO>>;
//
// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
// let this = self.get_mut();
// }
//}
pub struct TransparentStartHandshake<IO> {
h: StartHandshake<IO>,
}
/// Future returned from `TlsConnector::connect` which will resolve
/// once the connection handshake has finished.
pub struct Connect<IO>(MidHandshake<client::TlsStream<IO>>);