use crate::protocol::{HeaderLine, Message, MessageIO, Protocol, ProtocolError};
use crate::{Negotiated, NegotiationError, Version};
use futures::prelude::*;
use std::{
convert::TryFrom as _,
iter, mem,
pin::Pin,
task::{Context, Poll},
};
pub fn dialer_select_proto<R, I>(
inner: R,
protocols: I,
version: Version,
) -> DialerSelectFuture<R, I::IntoIter>
where
R: AsyncRead + AsyncWrite,
I: IntoIterator,
I::Item: AsRef<str>,
{
let protocols = protocols.into_iter().peekable();
DialerSelectFuture {
version,
protocols,
state: State::SendHeader {
io: MessageIO::new(inner),
},
}
}
#[pin_project::pin_project]
pub struct DialerSelectFuture<R, I: Iterator> {
protocols: iter::Peekable<I>,
state: State<R, I::Item>,
version: Version,
}
enum State<R, N> {
SendHeader { io: MessageIO<R> },
SendProtocol { io: MessageIO<R>, protocol: N },
FlushProtocol { io: MessageIO<R>, protocol: N },
AwaitProtocol { io: MessageIO<R>, protocol: N },
Done,
}
impl<R, I> Future for DialerSelectFuture<R, I>
where
R: AsyncRead + AsyncWrite + Unpin,
I: Iterator,
I::Item: AsRef<str>,
{
type Output = Result<(I::Item, Negotiated<R>), NegotiationError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
loop {
match mem::replace(this.state, State::Done) {
State::SendHeader { mut io } => {
match Pin::new(&mut io).poll_ready(cx)? {
Poll::Ready(()) => {}
Poll::Pending => {
*this.state = State::SendHeader { io };
return Poll::Pending;
}
}
let h = HeaderLine::from(*this.version);
if let Err(err) = Pin::new(&mut io).start_send(Message::Header(h)) {
return Poll::Ready(Err(From::from(err)));
}
let protocol = this.protocols.next().ok_or(NegotiationError::Failed)?;
*this.state = State::SendProtocol { io, protocol };
}
State::SendProtocol { mut io, protocol } => {
match Pin::new(&mut io).poll_ready(cx)? {
Poll::Ready(()) => {}
Poll::Pending => {
*this.state = State::SendProtocol { io, protocol };
return Poll::Pending;
}
}
let p = Protocol::try_from(protocol.as_ref())?;
if let Err(err) = Pin::new(&mut io).start_send(Message::Protocol(p.clone())) {
return Poll::Ready(Err(From::from(err)));
}
log::debug!("Dialer: Proposed protocol: {}", p);
if this.protocols.peek().is_some() {
*this.state = State::FlushProtocol { io, protocol }
} else {
match this.version {
Version::V1 => *this.state = State::FlushProtocol { io, protocol },
Version::V1Lazy => {
log::debug!("Dialer: Expecting proposed protocol: {}", p);
let hl = HeaderLine::from(Version::V1Lazy);
let io = Negotiated::expecting(io.into_reader(), p, Some(hl));
return Poll::Ready(Ok((protocol, io)));
}
}
}
}
State::FlushProtocol { mut io, protocol } => {
match Pin::new(&mut io).poll_flush(cx)? {
Poll::Ready(()) => *this.state = State::AwaitProtocol { io, protocol },
Poll::Pending => {
*this.state = State::FlushProtocol { io, protocol };
return Poll::Pending;
}
}
}
State::AwaitProtocol { mut io, protocol } => {
let msg = match Pin::new(&mut io).poll_next(cx)? {
Poll::Ready(Some(msg)) => msg,
Poll::Pending => {
*this.state = State::AwaitProtocol { io, protocol };
return Poll::Pending;
}
Poll::Ready(None) => return Poll::Ready(Err(NegotiationError::Failed)),
};
match msg {
Message::Header(v) if v == HeaderLine::from(*this.version) => {
*this.state = State::AwaitProtocol { io, protocol };
}
Message::Protocol(ref p) if p.as_ref() == protocol.as_ref() => {
log::debug!("Dialer: Received confirmation for protocol: {}", p);
let io = Negotiated::completed(io.into_inner());
return Poll::Ready(Ok((protocol, io)));
}
Message::NotAvailable => {
log::debug!(
"Dialer: Received rejection of protocol: {}",
protocol.as_ref()
);
let protocol = this.protocols.next().ok_or(NegotiationError::Failed)?;
*this.state = State::SendProtocol { io, protocol }
}
_ => return Poll::Ready(Err(ProtocolError::InvalidMessage.into())),
}
}
State::Done => panic!("State::poll called after completion"),
}
}
}
}