use std::any::TypeId;
use std::error::Error as StdError;
use std::fmt;
use std::io;
use std::marker::Unpin;
use bytes::{Buf, Bytes};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::sync::oneshot;
use crate::common::io::Rewind;
use crate::common::{task, Future, Pin, Poll};
pub struct Upgraded {
io: Rewind<Box<dyn Io + Send>>,
}
pub struct OnUpgrade {
rx: Option<oneshot::Receiver<crate::Result<Upgraded>>>,
}
#[derive(Debug)]
pub struct Parts<T> {
pub io: T,
pub read_buf: Bytes,
_inner: (),
}
pub(crate) struct Pending {
tx: oneshot::Sender<crate::Result<Upgraded>>,
}
#[derive(Debug)]
struct UpgradeExpected(());
pub(crate) fn pending() -> (Pending, OnUpgrade) {
let (tx, rx) = oneshot::channel();
(Pending { tx }, OnUpgrade { rx: Some(rx) })
}
impl Upgraded {
pub(crate) fn new<T>(io: T, read_buf: Bytes) -> Self
where
T: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
Upgraded {
io: Rewind::new_buffered(Box::new(ForwardsWriteBuf(io)), read_buf),
}
}
pub fn downcast<T: AsyncRead + AsyncWrite + Unpin + 'static>(self) -> Result<Parts<T>, Self> {
let (io, buf) = self.io.into_inner();
match io.__hyper_downcast::<ForwardsWriteBuf<T>>() {
Ok(t) => Ok(Parts {
io: t.0,
read_buf: buf,
_inner: (),
}),
Err(io) => Err(Upgraded {
io: Rewind::new_buffered(io, buf),
}),
}
}
}
impl AsyncRead for Upgraded {
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [std::mem::MaybeUninit<u8>]) -> bool {
self.io.prepare_uninitialized_buffer(buf)
}
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.io).poll_read(cx, buf)
}
}
impl AsyncWrite for Upgraded {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.io).poll_write(cx, buf)
}
fn poll_write_buf<B: Buf>(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
Pin::new(self.io.get_mut()).poll_write_dyn_buf(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.io).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.io).poll_shutdown(cx)
}
}
impl fmt::Debug for Upgraded {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Upgraded").finish()
}
}
impl OnUpgrade {
pub(crate) fn none() -> Self {
OnUpgrade { rx: None }
}
pub(crate) fn is_none(&self) -> bool {
self.rx.is_none()
}
}
impl Future for OnUpgrade {
type Output = Result<Upgraded, crate::Error>;
fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
match self.rx {
Some(ref mut rx) => Pin::new(rx).poll(cx).map(|res| match res {
Ok(Ok(upgraded)) => Ok(upgraded),
Ok(Err(err)) => Err(err),
Err(_oneshot_canceled) => {
Err(crate::Error::new_canceled().with(UpgradeExpected(())))
}
}),
None => Poll::Ready(Err(crate::Error::new_user_no_upgrade())),
}
}
}
impl fmt::Debug for OnUpgrade {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("OnUpgrade").finish()
}
}
impl Pending {
pub(crate) fn fulfill(self, upgraded: Upgraded) {
trace!("pending upgrade fulfill");
let _ = self.tx.send(Ok(upgraded));
}
pub(crate) fn manual(self) {
trace!("pending upgrade handled manually");
let _ = self.tx.send(Err(crate::Error::new_user_manual_upgrade()));
}
}
impl fmt::Display for UpgradeExpected {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "upgrade expected but not completed")
}
}
impl StdError for UpgradeExpected {}
struct ForwardsWriteBuf<T>(T);
pub(crate) trait Io: AsyncRead + AsyncWrite + Unpin + 'static {
fn poll_write_dyn_buf(
&mut self,
cx: &mut task::Context<'_>,
buf: &mut dyn Buf,
) -> Poll<io::Result<usize>>;
fn __hyper_type_id(&self) -> TypeId {
TypeId::of::<Self>()
}
}
impl dyn Io + Send {
fn __hyper_is<T: Io>(&self) -> bool {
let t = TypeId::of::<T>();
self.__hyper_type_id() == t
}
fn __hyper_downcast<T: Io>(self: Box<Self>) -> Result<Box<T>, Box<Self>> {
if self.__hyper_is::<T>() {
unsafe {
let raw: *mut dyn Io = Box::into_raw(self);
Ok(Box::from_raw(raw as *mut T))
}
} else {
Err(self)
}
}
}
impl<T: AsyncRead + Unpin> AsyncRead for ForwardsWriteBuf<T> {
unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [std::mem::MaybeUninit<u8>]) -> bool {
self.0.prepare_uninitialized_buffer(buf)
}
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.0).poll_read(cx, buf)
}
}
impl<T: AsyncWrite + Unpin> AsyncWrite for ForwardsWriteBuf<T> {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.0).poll_write(cx, buf)
}
fn poll_write_buf<B: Buf>(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.0).poll_write_buf(cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.0).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.0).poll_shutdown(cx)
}
}
impl<T: AsyncRead + AsyncWrite + Unpin + 'static> Io for ForwardsWriteBuf<T> {
fn poll_write_dyn_buf(
&mut self,
cx: &mut task::Context<'_>,
mut buf: &mut dyn Buf,
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.0).poll_write_buf(cx, &mut buf)
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::AsyncWriteExt;
#[test]
fn upgraded_downcast() {
let upgraded = Upgraded::new(Mock, Bytes::new());
let upgraded = upgraded.downcast::<std::io::Cursor<Vec<u8>>>().unwrap_err();
upgraded.downcast::<Mock>().unwrap();
}
#[tokio::test]
async fn upgraded_forwards_write_buf() {
Mock.write_buf(&mut "hello".as_bytes()).await.unwrap();
let mut upgraded = Upgraded::new(Mock, Bytes::new());
upgraded.write_buf(&mut "hello".as_bytes()).await.unwrap();
}
struct Mock;
impl AsyncRead for Mock {
fn poll_read(
self: Pin<&mut Self>,
_cx: &mut task::Context<'_>,
_buf: &mut [u8],
) -> Poll<io::Result<usize>> {
unreachable!("Mock::poll_read")
}
}
impl AsyncWrite for Mock {
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut task::Context<'_>,
_buf: &[u8],
) -> Poll<io::Result<usize>> {
panic!("poll_write shouldn't be called");
}
fn poll_write_buf<B: Buf>(
self: Pin<&mut Self>,
_cx: &mut task::Context<'_>,
buf: &mut B,
) -> Poll<io::Result<usize>> {
let n = buf.remaining();
buf.advance(n);
Poll::Ready(Ok(n))
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
unreachable!("Mock::poll_flush")
}
fn poll_shutdown(
self: Pin<&mut Self>,
_cx: &mut task::Context<'_>,
) -> Poll<io::Result<()>> {
unreachable!("Mock::poll_shutdown")
}
}
}