Port unified TLS stream type to tokio-0.2

This commit is contained in:
Dirkjan Ochtman 2019-11-06 11:41:59 +01:00 committed by quininer
parent 7cccd9c3b3
commit 3e2c0446a4
2 changed files with 113 additions and 1 deletions

View File

@ -18,6 +18,7 @@ github-actions = { repository = "quininer/tokio-rustls", workflow = "ci" }
smallvec = "0.6"
tokio-io = "=0.2.0-alpha.6"
futures-core-preview = "=0.3.0-alpha.19"
pin-project = "0.4"
rustls = "0.16"
webpki = "0.21"

View File

@ -6,7 +6,8 @@ pub mod server;
use common::Stream;
use futures_core as futures;
use rustls::{ClientConfig, ClientSession, ServerConfig, ServerSession};
use pin_project::{pin_project, project};
use rustls::{ClientConfig, ClientSession, ServerConfig, ServerSession, Session};
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
@ -195,3 +196,113 @@ impl<IO: AsyncRead + AsyncWrite + Unpin> Future for Accept<IO> {
Pin::new(&mut self.0).poll(cx)
}
}
/// Unified TLS stream type
///
/// This abstracts over the inner `client::TlsStream` and `server::TlsStream`, so you can use
/// a single type to keep both client- and server-initiated TLS-encrypted connections.
#[pin_project]
pub enum TlsStream<T> {
Client(#[pin] client::TlsStream<T>),
Server(#[pin] server::TlsStream<T>),
}
impl<T> TlsStream<T> {
pub fn get_ref(&self) -> (&T, &dyn Session) {
use TlsStream::*;
match self {
Client(io) => {
let (io, session) = io.get_ref();
(io, &*session)
}
Server(io) => {
let (io, session) = io.get_ref();
(io, &*session)
}
}
}
pub fn get_mut(&mut self) -> (&mut T, &mut dyn Session) {
use TlsStream::*;
match self {
Client(io) => {
let (io, session) = io.get_mut();
(io, &mut *session)
}
Server(io) => {
let (io, session) = io.get_mut();
(io, &mut *session)
}
}
}
}
impl<T> From<client::TlsStream<T>> for TlsStream<T> {
fn from(s: client::TlsStream<T>) -> Self {
Self::Client(s)
}
}
impl<T> From<server::TlsStream<T>> for TlsStream<T> {
fn from(s: server::TlsStream<T>) -> Self {
Self::Server(s)
}
}
impl<T> AsyncRead for TlsStream<T>
where
T: AsyncRead + AsyncWrite + Unpin,
{
#[project]
#[inline]
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
#[project]
match self.project() {
TlsStream::Client(x) => x.poll_read(cx, buf),
TlsStream::Server(x) => x.poll_read(cx, buf),
}
}
}
impl<T> AsyncWrite for TlsStream<T>
where
T: AsyncRead + AsyncWrite + Unpin,
{
#[project]
#[inline]
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
#[project]
match self.project() {
TlsStream::Client(x) => x.poll_write(cx, buf),
TlsStream::Server(x) => x.poll_write(cx, buf),
}
}
#[project]
#[inline]
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
#[project]
match self.project() {
TlsStream::Client(x) => x.poll_flush(cx),
TlsStream::Server(x) => x.poll_flush(cx),
}
}
#[project]
#[inline]
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
#[project]
match self.project() {
TlsStream::Client(x) => x.poll_shutdown(cx),
TlsStream::Server(x) => x.poll_shutdown(cx),
}
}
}