#![deny(missing_docs, missing_debug_implementations, unsafe_code)]
#![warn(unreachable_pub, unused_qualifications, unused_lifetimes)]
#![warn(
clippy::must_use_candidate,
clippy::unwrap_in_result,
clippy::panic_in_result_fn
)]
use futures_io::{AsyncRead, AsyncWrite};
use openssl::{
error::ErrorStack,
ssl::{self, ErrorCode, ShutdownState, Ssl, SslRef},
};
use std::{
any::Any,
fmt, future,
io::{self, Read, Write},
mem,
panic::{self, AssertUnwindSafe},
pin::Pin,
sync::{
Arc, Mutex, MutexGuard, PoisonError,
atomic::{AtomicUsize, Ordering},
},
task::{Context, Poll, Wake, Waker},
};
#[cfg(test)]
mod test;
#[cfg(doctest)]
#[doc = include_str!("../README.md")]
struct ReadmeDoctests;
const WRITE_CHUNK_SIZE: usize = 16 * 1024;
const MAX_UNSENT: usize = 4 * WRITE_CHUNK_SIZE;
const MAX_WAITERS: usize = 32;
fn chunk(buf: &[u8]) -> &[u8] {
&buf[..buf.len().min(WRITE_CHUNK_SIZE)]
}
#[derive(Default)]
struct WakerSet {
waiters: Mutex<Vec<Waiter>>,
generation: AtomicUsize,
}
struct Waiter {
waker: Waker,
records: bool,
}
impl WakerSet {
fn lock(&self) -> MutexGuard<'_, Vec<Waiter>> {
self.waiters.lock().unwrap_or_else(PoisonError::into_inner)
}
fn generation(&self) -> usize {
self.generation.load(Ordering::Acquire)
}
fn join(&self, waker: &Waker, records: bool, generation: usize) {
let mut waiters = self.lock();
if self.generation.load(Ordering::Relaxed) != generation {
drop(waiters);
waker.wake_by_ref();
return;
}
if let Some(waiter) = waiters.iter_mut().find(|w| w.waker.will_wake(waker)) {
waiter.records |= records;
return;
}
let waiter = Waiter {
waker: waker.clone(),
records,
};
if waiters.len() < MAX_WAITERS {
waiters.push(waiter);
} else {
let evicted = mem::replace(&mut *waiters, vec![waiter]);
drop(waiters);
self.wake_all(evicted);
}
}
fn wake_tasks(&self, by_transport: bool, records_only: bool) {
let woken = {
let mut waiters = self.lock();
if by_transport {
self.generation.fetch_add(1, Ordering::Release);
}
if waiters.is_empty() || (records_only && !waiters.iter().any(|w| w.records)) {
return;
}
let all = mem::take(&mut *waiters);
if records_only {
let (woken, kept) = all.into_iter().partition(|waiter| waiter.records);
*waiters = kept;
woken
} else {
all
}
};
self.wake_all(woken);
}
fn wake_all(&self, wakers: Vec<Waiter>) {
struct Requeue<'a> {
set: &'a WakerSet,
wakers: Vec<Waiter>,
next: usize,
}
impl Drop for Requeue<'_> {
fn drop(&mut self) {
if self.next < self.wakers.len() {
let rest = self.wakers.drain(self.next..);
self.set.lock().extend(rest);
}
}
}
let mut requeue = Requeue {
set: self,
wakers,
next: 0,
};
while let Some(waiter) = requeue.wakers.get(requeue.next) {
requeue.next += 1;
waiter.waker.wake_by_ref();
}
let mut wakers = mem::take(&mut requeue.wakers);
wakers.clear();
let mut waiters = self.lock();
if waiters.capacity() == 0 {
*waiters = wakers;
}
}
}
impl Wake for WakerSet {
fn wake(self: Arc<Self>) {
self.wake_by_ref();
}
fn wake_by_ref(self: &Arc<Self>) {
self.wake_tasks(true, false);
}
}
fn keeping_foreign_errors<T>(f: impl FnOnce() -> T) -> T {
let foreign = ErrorStack::get();
let result = f();
foreign.put();
result
}
fn handshake_finished(ssl: &SslRef) -> bool {
ssl.state_string_long() == "SSL negotiation finished successfully"
}
fn poll_write_all<S: AsyncWrite + ?Sized>(
mut stream: Pin<&mut S>,
cx: &mut Context<'_>,
buf: &[u8],
written: &mut usize,
) -> Poll<io::Result<()>> {
while *written < buf.len() {
let left = buf.len() - *written;
match std::task::ready!(stream.as_mut().poll_write(cx, &buf[*written..]))? {
0 => return Poll::Ready(Err(io::ErrorKind::WriteZero.into())),
n if n > left => {
let error = "the transport reported writing more than it was given";
return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, error)));
}
n => *written += n,
}
}
Poll::Ready(Ok(()))
}
#[derive(Clone, Copy, Default)]
enum ErrorPath {
#[default]
ThroughOpenSsl,
ToCaller,
}
#[derive(Clone, Copy, Default)]
enum LastRead {
#[default]
Data,
Eof,
Waiting,
Failed,
}
#[derive(Default)]
struct Call {
path: ErrorPath,
read: LastRead,
send_blocked: bool,
write_error: Option<AssertUnwindSafe<io::Error>>,
read_error: Option<AssertUnwindSafe<io::Error>>,
inject: Option<AssertUnwindSafe<io::Error>>,
}
#[derive(Default)]
struct Op {
read_progress: bool,
send_failed: bool,
read_wait: bool,
}
#[derive(Clone, Copy)]
enum Flush {
Requested,
Unflushed,
Always,
}
struct StreamWrapper<S: Unpin> {
stream: S,
wakers: Arc<WakerSet>,
waker: Waker,
out: Vec<u8>,
sent: usize,
flush_needed: bool,
unflushed: bool,
refused: Option<usize>,
poisoned: bool,
panic: Mutex<Option<Box<dyn Any + Send>>>,
close: Close,
op: Op,
call: Call,
}
impl<S> fmt::Debug for StreamWrapper<S>
where
S: fmt::Debug + Unpin,
{
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("StreamWrapper")
.field("stream", &self.stream)
.field("unsent", &self.unsent())
.field("flush_needed", &self.flush_needed)
.field("unflushed", &self.unflushed)
.field("poisoned", &self.poisoned)
.field("close", &self.close)
.finish_non_exhaustive()
}
}
impl<S: Unpin> StreamWrapper<S> {
fn unsent(&self) -> usize {
self.out.len() - self.sent
}
}
impl<S: AsyncWrite + Unpin> StreamWrapper<S> {
fn new(stream: S) -> Self {
let wakers = Arc::<WakerSet>::default();
Self {
stream,
waker: Waker::from(wakers.clone()),
wakers,
out: Vec::new(),
sent: 0,
flush_needed: false,
unflushed: false,
refused: None,
poisoned: false,
panic: Mutex::new(None),
close: Close::Open,
op: Op::default(),
call: Call::default(),
}
}
fn check_poisoned(&self) -> io::Result<()> {
if self.poisoned {
let error = "the transport panicked during an earlier call";
return Err(io::Error::other(error));
}
Ok(())
}
fn poll_transport<T>(
&mut self,
poll: impl FnOnce(Pin<&mut S>, &mut Context<'_>) -> Poll<io::Result<T>>,
) -> Poll<io::Result<T>> {
self.check_poisoned()?;
let stream = Pin::new(&mut self.stream);
let mut cx = Context::from_waker(&self.waker);
match panic::catch_unwind(AssertUnwindSafe(|| poll(stream, &mut cx))) {
Ok(result) => result,
Err(payload) => {
self.poisoned = true;
*self.panic.get_mut().unwrap_or_else(PoisonError::into_inner) = Some(payload);
Poll::Ready(Err(io::Error::other("the transport panicked")))
}
}
}
fn poll_output<T>(
&mut self,
poll: impl FnOnce(Pin<&mut S>, &mut Context<'_>) -> Poll<io::Result<T>>,
) -> Poll<io::Result<T>> {
let generation = self.wakers.generation();
let result = self.poll_transport(poll);
self.refused = result.is_pending().then_some(generation);
if let Poll::Ready(Err(_)) = result {
self.op.send_failed = true;
}
result
}
fn poll_send(&mut self, flush: Flush) -> Poll<io::Result<()>> {
if let Err(error) = self.check_poisoned() {
self.op.send_failed = true;
return Poll::Ready(Err(error));
}
if self.unsent() > 0 {
let out = mem::take(&mut self.out);
let mut sent = self.sent;
let result = self.poll_output(|stream, cx| poll_write_all(stream, cx, &out, &mut sent));
(self.out, self.sent) = (out, sent);
std::task::ready!(result)?;
self.out.clear();
self.out.shrink_to(2 * WRITE_CHUNK_SIZE);
self.sent = 0;
}
let needed = match flush {
Flush::Requested => self.flush_needed,
Flush::Unflushed => self.flush_needed || self.unflushed,
Flush::Always => true,
};
if needed {
std::task::ready!(self.poll_output(|stream, cx| stream.poll_flush(cx)))?;
self.flush_needed = false;
self.unflushed = false;
}
Poll::Ready(Ok(()))
}
fn poll_close(&mut self) -> Poll<io::Result<()>> {
self.poll_transport(|stream, cx| stream.poll_close(cx))
}
fn may_send(&self) -> bool {
!self.call.send_blocked && self.refused != Some(self.wakers.generation())
}
fn start_call(&mut self, path: ErrorPath) {
self.call = Call {
path,
..Call::default()
};
}
fn send(&mut self, flush: Flush) {
if !self.may_send() {
return;
}
let result = self.poll_send(flush);
self.note_send(result);
}
fn note_send(&mut self, result: Poll<io::Result<()>>) {
match result {
Poll::Ready(Ok(())) => {}
Poll::Ready(Err(error)) => {
self.call.send_blocked = true;
self.call.write_error = Some(AssertUnwindSafe(error));
}
Poll::Pending => self.call.send_blocked = true,
}
}
fn wait(&mut self) -> io::Result<usize> {
self.call.read = LastRead::Waiting;
Err(io::ErrorKind::WouldBlock.into())
}
fn read_failed(&mut self, error: io::Error) -> io::Result<usize> {
match self.call.path {
ErrorPath::ThroughOpenSsl => {
self.call.read = LastRead::Failed;
Err(error)
}
ErrorPath::ToCaller => {
self.call.read_error = Some(AssertUnwindSafe(error));
self.wait()
}
}
}
fn fail_or_wait(&mut self) -> io::Result<usize> {
match self.call.write_error.take() {
Some(AssertUnwindSafe(error)) => self.read_failed(error),
None => self.wait(),
}
}
fn kept_error(&mut self) -> Option<io::Error> {
(self.call.read_error.take())
.or_else(|| self.call.write_error.take())
.map(|AssertUnwindSafe(error)| error)
}
}
impl<S> Read for StreamWrapper<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if let Some(AssertUnwindSafe(error)) = self.call.inject.take() {
self.call.read = LastRead::Failed;
return Err(error);
}
self.send(Flush::Requested);
if self.unsent() > MAX_UNSENT {
return self.fail_or_wait();
}
let len = buf.len();
match self.poll_transport(|stream, cx| stream.poll_read(cx, buf)) {
Poll::Ready(Ok(0)) => {
self.call.read = LastRead::Eof;
Ok(0)
}
Poll::Ready(Ok(n)) if n > len => {
let error = "the transport reported reading more than it was given";
self.read_failed(io::Error::new(io::ErrorKind::InvalidData, error))
}
Poll::Ready(Ok(n)) => {
self.call.read = LastRead::Data;
self.op.read_progress = true;
Ok(n)
}
Poll::Ready(Err(error)) => self.read_failed(error),
Poll::Pending => {
self.op.read_wait = true;
self.send(Flush::Unflushed);
self.fail_or_wait()
}
}
}
}
impl<S> Write for StreamWrapper<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.check_poisoned()?;
if !matches!(self.close, Close::Open) {
return Ok(buf.len());
}
let mut written = 0;
if self.out.is_empty() && self.may_send() {
let result =
self.poll_output(|stream, cx| poll_write_all(stream, cx, buf, &mut written));
self.note_send(result);
} else if self.sent > 0 && self.sent >= self.unsent() {
self.out.drain(..self.sent);
self.sent = 0;
}
self.out.extend_from_slice(&buf[written..]);
self.unflushed = true;
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
if !matches!(self.close, Close::Open) {
return Ok(());
}
self.flush_needed = true;
self.send(Flush::Requested);
Ok(())
}
}
#[derive(Debug)]
enum Close {
Open,
Sent,
NoSession(Option<AssertUnwindSafe<io::Error>>),
}
pub struct SslStream<S: Unpin> {
inner: ssl::SslStream<StreamWrapper<S>>,
}
impl<S> fmt::Debug for SslStream<S>
where
S: fmt::Debug + Unpin,
{
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt.debug_struct("SslStream")
.field("stream", self.inner.get_ref())
.field("ssl", self.inner.ssl())
.finish()
}
}
impl<S> SslStream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
pub fn new(ssl: Ssl, stream: S) -> Result<Self, ErrorStack> {
ssl::SslStream::new(ssl, StreamWrapper::new(stream)).map(|inner| SslStream { inner })
}
pub fn poll_connect(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), ssl::Error>> {
self.poll_ssl_op(cx, |s| s.connect())
}
pub async fn connect(mut self: Pin<&mut Self>) -> Result<(), ssl::Error> {
future::poll_fn(|cx| self.as_mut().poll_connect(cx)).await
}
pub fn poll_accept(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), ssl::Error>> {
self.poll_ssl_op(cx, |s| s.accept())
}
pub async fn accept(mut self: Pin<&mut Self>) -> Result<(), ssl::Error> {
future::poll_fn(|cx| self.as_mut().poll_accept(cx)).await
}
pub fn poll_do_handshake(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<(), ssl::Error>> {
self.poll_ssl_op(cx, |s| s.do_handshake())
}
pub async fn do_handshake(mut self: Pin<&mut Self>) -> Result<(), ssl::Error> {
future::poll_fn(|cx| self.as_mut().poll_do_handshake(cx)).await
}
pub fn poll_peek(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<usize, ssl::Error>> {
if buf.is_empty() {
return Poll::Ready(Ok(0));
}
self.poll_ssl_op(cx, |s| s.ssl_peek(buf))
}
pub async fn peek(mut self: Pin<&mut Self>, buf: &mut [u8]) -> Result<usize, ssl::Error> {
future::poll_fn(|cx| self.as_mut().poll_peek(cx, buf)).await
}
#[cfg(ossl111)]
pub fn poll_read_early_data(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<usize, ssl::Error>> {
self.poll_op(cx, |this| {
this.poll_ssl(ErrorPath::ToCaller, |s| {
let early = s.read_early_data(buf);
let over = matches!(early, Ok(0));
let result = match early {
Ok(0) => s.do_handshake().map(|()| 0),
result => result,
};
let Err(error) = result else {
return result;
};
let waited = matches!(s.get_ref().call.read, LastRead::Waiting);
match waited.then(|| s.get_mut().kept_error()).flatten() {
Some(kept) => {
s.get_mut().call.inject = Some(AssertUnwindSafe(kept));
if over {
s.do_handshake().map(|()| 0)
} else {
s.read_early_data(buf)
}
}
None if over
&& matches!(error.code(), ErrorCode::WANT_READ | ErrorCode::WANT_WRITE) =>
{
Ok(0)
}
None => Err(error),
}
})
})
}
#[cfg(ossl111)]
pub async fn read_early_data(
mut self: Pin<&mut Self>,
buf: &mut [u8],
) -> Result<usize, ssl::Error> {
future::poll_fn(|cx| self.as_mut().poll_read_early_data(cx, buf)).await
}
#[cfg(ossl111)]
pub fn poll_write_early_data(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, ssl::Error>> {
self.poll_op(cx, |this| {
let _ = std::task::ready!(this.wrapper().poll_send(Flush::Requested));
let buf = chunk(buf);
this.poll_ssl(ErrorPath::ThroughOpenSsl, |s| s.write_early_data(buf))
})
}
#[cfg(ossl111)]
pub async fn write_early_data(
mut self: Pin<&mut Self>,
buf: &[u8],
) -> Result<usize, ssl::Error> {
future::poll_fn(|cx| self.as_mut().poll_write_early_data(cx, buf)).await
}
}
impl<S: Unpin> SslStream<S> {
#[must_use]
pub fn ssl(&self) -> &SslRef {
self.inner.ssl()
}
#[must_use]
pub fn get_ref(&self) -> &S {
&self.inner.get_ref().stream
}
pub fn get_mut(&mut self) -> &mut S {
&mut self.inner.get_mut().stream
}
#[must_use]
pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut S> {
Pin::new(Pin::into_inner(self).get_mut())
}
fn wrapper(&mut self) -> &mut StreamWrapper<S> {
self.inner.get_mut()
}
}
impl<S> SslStream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn poll_op<T>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
op: impl FnOnce(&mut Self) -> Poll<T>,
) -> Poll<T> {
let this = self.get_mut();
let wrapper = this.wrapper();
let generation = wrapper.wakers.generation();
wrapper.op = Op::default();
let result = op(this);
if !matches!(this.inner.get_ref().close, Close::Open) {
this.refuse_writes();
}
let wrapper = this.wrapper();
if wrapper.op.send_failed || wrapper.op.read_progress {
wrapper.wakers.wake_tasks(false, !wrapper.op.send_failed);
}
if result.is_pending() {
let records = wrapper.op.read_wait;
wrapper.wakers.join(cx.waker(), records, generation);
}
let panic = wrapper.panic.get_mut();
if let Some(panic) = panic.unwrap_or_else(PoisonError::into_inner).take() {
panic::resume_unwind(panic);
}
result
}
fn poll_ssl<T>(
&mut self,
path: ErrorPath,
mut f: impl FnMut(&mut ssl::SslStream<StreamWrapper<S>>) -> Result<T, ssl::Error>,
) -> Poll<Result<T, ssl::Error>> {
keeping_foreign_errors(|| {
loop {
self.inner.get_mut().start_call(path);
let error = match f(&mut self.inner) {
Ok(value) => break Poll::Ready(Ok(value)),
Err(error) => error,
};
match (self.inner.get_ref().call.read, error.code()) {
(LastRead::Waiting, _) => break Poll::Pending,
(LastRead::Failed | LastRead::Eof, _) => break Poll::Ready(Err(error)),
(LastRead::Data, ErrorCode::WANT_READ | ErrorCode::WANT_WRITE) => {}
(LastRead::Data, _) => break Poll::Ready(Err(error)),
}
}
})
}
fn refuse_writes(&mut self) {
let state = self.inner.get_shutdown();
self.inner.set_shutdown(state | ShutdownState::SENT);
}
fn poll_ssl_op<T>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
f: impl FnMut(&mut ssl::SslStream<StreamWrapper<S>>) -> Result<T, ssl::Error>,
) -> Poll<Result<T, ssl::Error>> {
self.poll_op(cx, |this| this.poll_ssl(ErrorPath::ThroughOpenSsl, f))
}
fn poll_ssl_io<T>(
&mut self,
f: impl FnMut(&mut ssl::SslStream<StreamWrapper<S>>) -> Result<T, ssl::Error>,
) -> Poll<io::Result<T>> {
match self.poll_ssl(ErrorPath::ToCaller, f) {
Poll::Ready(result) => Poll::Ready(result.map_err(io::Error::other)),
Poll::Pending => match self.wrapper().kept_error() {
Some(error) => Poll::Ready(Err(error)),
None => Poll::Pending,
},
}
}
}
impl<S> AsyncRead for SslStream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
if buf.is_empty() {
return Poll::Ready(Ok(0));
}
self.poll_op(cx, |this| {
this.poll_ssl_io(|s| match s.ssl_read(buf) {
Err(error)
if (error.code() == ErrorCode::ZERO_RETURN
|| (error.code() == ErrorCode::SYSCALL
&& matches!(s.get_ref().call.read, LastRead::Eof)))
&& handshake_finished(s.ssl()) =>
{
Ok(0)
}
result => result,
})
})
}
}
impl<S> AsyncWrite for SslStream<S>
where
S: AsyncRead + AsyncWrite + Unpin,
{
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
if buf.is_empty() {
return Poll::Ready(Ok(0));
}
self.poll_op(cx, |this| {
std::task::ready!(this.wrapper().poll_send(Flush::Requested))?;
let buf = chunk(buf);
this.poll_ssl_io(|s| s.ssl_write(buf))
})
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.poll_op(cx, |this| this.wrapper().poll_send(Flush::Always))
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.poll_op(cx, |this| {
if matches!(this.inner.get_ref().close, Close::Open) {
std::task::ready!(this.wrapper().poll_send(Flush::Requested))?;
this.inner.get_mut().start_call(ErrorPath::ThroughOpenSsl);
let shutdown = keeping_foreign_errors(|| this.inner.shutdown());
let sent = this.inner.get_shutdown().contains(ShutdownState::SENT);
this.inner.get_mut().close = match shutdown {
Ok(_) if sent => Close::Sent,
Ok(_) => Close::NoSession(None),
Err(error) => {
let error = error.into_io_error().unwrap_or_else(io::Error::other);
Close::NoSession(Some(AssertUnwindSafe(error)))
}
};
}
if matches!(this.inner.get_ref().close, Close::NoSession(_)) {
std::task::ready!(this.wrapper().poll_close())?;
let reason = match &mut this.wrapper().close {
Close::NoSession(reason) => reason.take(),
Close::Open | Close::Sent => None,
};
let error = reason.map_or_else(
|| io::Error::other("no TLS session to close"),
|AssertUnwindSafe(error)| error,
);
return Poll::Ready(Err(error));
}
match std::task::ready!(this.wrapper().poll_send(Flush::Unflushed)) {
Ok(()) => {}
Err(_) if this.inner.get_shutdown().contains(ShutdownState::RECEIVED) => {}
Err(error) => return Poll::Ready(Err(error)),
}
this.wrapper().poll_close()
})
}
}