72 lines
1.9 KiB
Rust
72 lines
1.9 KiB
Rust
use crate::tls::{ChannelBinding, TlsStream};
|
|
use std::io;
|
|
use std::pin::Pin;
|
|
use std::task::{Context, Poll};
|
|
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
|
|
|
pub enum MaybeTlsStream<S, T> {
|
|
Raw(S),
|
|
Tls(T),
|
|
}
|
|
|
|
impl<S, T> AsyncRead for MaybeTlsStream<S, T>
|
|
where
|
|
S: AsyncRead + Unpin,
|
|
T: AsyncRead + Unpin,
|
|
{
|
|
fn poll_read(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
buf: &mut ReadBuf<'_>,
|
|
) -> Poll<io::Result<()>> {
|
|
match &mut *self {
|
|
MaybeTlsStream::Raw(s) => Pin::new(s).poll_read(cx, buf),
|
|
MaybeTlsStream::Tls(s) => Pin::new(s).poll_read(cx, buf),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<S, T> AsyncWrite for MaybeTlsStream<S, T>
|
|
where
|
|
S: AsyncWrite + Unpin,
|
|
T: AsyncWrite + Unpin,
|
|
{
|
|
fn poll_write(
|
|
mut self: Pin<&mut Self>,
|
|
cx: &mut Context<'_>,
|
|
buf: &[u8],
|
|
) -> Poll<io::Result<usize>> {
|
|
match &mut *self {
|
|
MaybeTlsStream::Raw(s) => Pin::new(s).poll_write(cx, buf),
|
|
MaybeTlsStream::Tls(s) => Pin::new(s).poll_write(cx, buf),
|
|
}
|
|
}
|
|
|
|
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
match &mut *self {
|
|
MaybeTlsStream::Raw(s) => Pin::new(s).poll_flush(cx),
|
|
MaybeTlsStream::Tls(s) => Pin::new(s).poll_flush(cx),
|
|
}
|
|
}
|
|
|
|
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
match &mut *self {
|
|
MaybeTlsStream::Raw(s) => Pin::new(s).poll_shutdown(cx),
|
|
MaybeTlsStream::Tls(s) => Pin::new(s).poll_shutdown(cx),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<S, T> TlsStream for MaybeTlsStream<S, T>
|
|
where
|
|
S: AsyncRead + AsyncWrite + Unpin,
|
|
T: TlsStream + Unpin,
|
|
{
|
|
fn channel_binding(&self) -> ChannelBinding {
|
|
match self {
|
|
MaybeTlsStream::Raw(_) => ChannelBinding::none(),
|
|
MaybeTlsStream::Tls(s) => s.channel_binding(),
|
|
}
|
|
}
|
|
}
|