panorama/imap/src/client/inner.rs

250 lines
7.5 KiB
Rust
Raw Normal View History

2021-02-20 05:03:33 +00:00
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
2021-02-20 05:03:33 +00:00
use anyhow::Result;
use futures::{
future::{self, Either, FutureExt, TryFutureExt},
stream::{Peekable, Stream, StreamExt},
};
use parking_lot::Mutex;
2021-02-20 05:03:33 +00:00
use tokio::{
io::{
2021-02-22 20:43:09 +00:00
self, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader, ReadHalf, WriteHalf,
2021-02-20 05:03:33 +00:00
},
sync::{
mpsc,
oneshot::{self, error::TryRecvError},
},
2021-02-20 05:03:33 +00:00
task::JoinHandle,
};
2021-02-21 13:54:46 +00:00
use tokio_rustls::{
client::TlsStream, rustls::ClientConfig as RustlsConfig, webpki::DNSNameRef, TlsConnector,
};
use tokio_stream::wrappers::UnboundedReceiverStream;
2021-02-20 05:03:33 +00:00
use crate::command::Command;
2021-02-23 02:47:00 +00:00
use crate::parser::{parse_capability, parse_response};
use crate::response::{Response, ResponseDone};
2021-02-21 13:42:40 +00:00
2021-02-21 13:54:46 +00:00
use super::ClientConfig;
2021-02-20 05:03:33 +00:00
2021-03-01 07:48:28 +00:00
pub const TAG_PREFIX: &str = "ptag";
type Command2 = (Command, mpsc::UnboundedSender<Response>);
2021-02-20 05:03:33 +00:00
pub struct Client<C> {
ctr: usize,
2021-02-21 13:54:46 +00:00
config: ClientConfig,
2021-02-20 05:03:33 +00:00
conn: WriteHalf<C>,
cmd_tx: mpsc::UnboundedSender<Command2>,
greeting_rx: Option<oneshot::Receiver<()>>,
exit_tx: oneshot::Sender<()>,
2021-02-21 13:42:40 +00:00
listener_handle: JoinHandle<Result<ReadHalf<C>>>,
2021-02-20 05:03:33 +00:00
}
impl<C> Client<C>
where
C: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
2021-02-21 13:54:46 +00:00
pub fn new(conn: C, config: ClientConfig) -> Self {
2021-02-20 05:03:33 +00:00
let (read_half, write_half) = io::split(conn);
let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
let (greeting_tx, greeting_rx) = oneshot::channel();
let (exit_tx, exit_rx) = oneshot::channel();
let handle = tokio::spawn(
listen(read_half, cmd_rx, greeting_tx, exit_rx).map_err(|err| {
2021-03-02 22:47:33 +00:00
error!("Help, the listener loop died: {}", err);
err
}),
);
2021-02-20 05:03:33 +00:00
Client {
ctr: 0,
2021-02-20 05:03:33 +00:00
conn: write_half,
config,
cmd_tx,
greeting_rx: Some(greeting_rx),
2021-02-21 13:42:40 +00:00
exit_tx,
listener_handle: handle,
2021-02-20 05:03:33 +00:00
}
}
pub async fn wait_for_greeting(&mut self) -> Result<()> {
if let Some(greeting_rx) = self.greeting_rx.take() {
greeting_rx.await?;
}
Ok(())
2021-02-21 13:42:40 +00:00
}
pub async fn execute(&mut self, cmd: Command) -> Result<ResponseStream> {
let id = self.ctr;
self.ctr += 1;
2021-02-21 01:13:10 +00:00
let cmd_str = format!("{}{} {}\r\n", TAG_PREFIX, id, cmd);
2021-02-20 05:03:33 +00:00
self.conn.write_all(cmd_str.as_bytes()).await?;
self.conn.flush().await?;
let (tx, rx) = mpsc::unbounded_channel();
self.cmd_tx.send((cmd, tx))?;
Ok(ResponseStream { inner: rx })
2021-02-20 05:03:33 +00:00
}
pub async fn has_capability(&mut self, cap: impl AsRef<str>) -> Result<bool> {
// TODO: cache capabilities if needed?
let cap = cap.as_ref();
let cap = parse_capability(cap)?;
2021-02-23 02:47:00 +00:00
let resp = self.execute(Command::Capability).await?;
let (_, data) = resp.wait().await?;
for resp in data {
if let Response::Capabilities(caps) = resp {
return Ok(caps.contains(&cap));
}
// debug!("cap: {:?}", resp);
2021-02-23 02:47:00 +00:00
}
Ok(false)
2021-02-20 05:03:33 +00:00
}
2021-02-21 13:42:40 +00:00
pub async fn upgrade(mut self) -> Result<Client<TlsStream<C>>> {
// TODO: make sure STARTTLS is in the capability list
2021-02-22 20:36:08 +00:00
if !self.has_capability("STARTTLS").await? {
bail!("server doesn't support this capability");
}
2021-02-21 13:42:40 +00:00
// first, send the STARTTLS command
let mut resp = self.execute(Command::Starttls).await?;
let resp = resp.next().await.unwrap();
2021-02-21 13:42:40 +00:00
debug!("server response to starttls: {:?}", resp);
debug!("sending exit for upgrade");
// TODO: check that the channel is still open?
self.exit_tx.send(()).unwrap();
2021-02-21 13:42:40 +00:00
let reader = self.listener_handle.await??;
let writer = self.conn;
let conn = reader.unsplit(writer);
let server_name = &self.config.hostname;
2021-02-21 13:54:46 +00:00
let mut tls_config = RustlsConfig::new();
2021-02-21 13:42:40 +00:00
tls_config
.root_store
.add_server_trust_anchors(&webpki_roots::TLS_SERVER_ROOTS);
let tls_config = TlsConnector::from(Arc::new(tls_config));
let dnsname = DNSNameRef::try_from_ascii_str(server_name).unwrap();
let stream = tls_config.connect(dnsname, conn).await?;
debug!("upgraded, stream is using TLS now");
2021-02-21 13:42:40 +00:00
Ok(Client::new(stream, self.config))
}
}
pub struct ResponseStream {
inner: mpsc::UnboundedReceiver<Response>,
}
2021-02-21 13:42:40 +00:00
impl ResponseStream {
/// Retrieves just the DONE item in the stream, discarding the rest
pub async fn done(mut self) -> Result<Option<ResponseDone>> {
while let Some(resp) = self.inner.recv().await {
if let Response::Done(done) = resp {
return Ok(Some(done));
}
2021-02-21 13:42:40 +00:00
}
Ok(None)
}
2021-02-21 13:42:40 +00:00
/// Waits for the entire stream to finish, returning the DONE status and the stream
pub async fn wait(mut self) -> Result<(Option<ResponseDone>, Vec<Response>)> {
let mut done = None;
let mut vec = Vec::new();
while let Some(resp) = self.inner.recv().await {
if let Response::Done(d) = resp {
done = Some(d);
break;
} else {
vec.push(resp);
}
2021-02-21 13:42:40 +00:00
}
Ok((done, vec))
2021-02-21 13:42:40 +00:00
}
2021-02-20 05:03:33 +00:00
}
impl Stream for ResponseStream {
type Item = Response;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
self.inner.poll_recv(cx)
}
}
#[allow(unreachable_code)]
2021-02-21 13:42:40 +00:00
async fn listen<C>(
conn: ReadHalf<C>,
mut cmd_rx: mpsc::UnboundedReceiver<Command2>,
greeting_tx: oneshot::Sender<()>,
mut exit_rx: oneshot::Receiver<()>,
) -> Result<ReadHalf<C>>
2021-02-21 13:42:40 +00:00
where
C: AsyncRead + Unpin,
{
2021-02-20 05:03:33 +00:00
let mut reader = BufReader::new(conn);
let mut greeting_tx = Some(greeting_tx);
let mut curr_cmd: Option<Command2> = None;
let mut exit_rx = exit_rx.map_err(|_| ()).shared();
// let mut exit_fut = Some(exit_rx.fuse());
// let mut fut1 = None;
2021-02-21 13:42:40 +00:00
2021-02-20 05:03:33 +00:00
loop {
let mut next_line = String::new();
let read_fut = reader.read_line(&mut next_line).fuse();
pin_mut!(read_fut);
// only listen for a new command if there isn't one already
let mut cmd_fut = if let Some(_) = curr_cmd {
// if there is one, just make a future that never resolves so it'll always pick the
// other options in the select.
future::pending().boxed().fuse()
} else {
cmd_rx.recv().boxed().fuse()
};
select! {
_ = exit_rx => {
debug!("exiting the loop");
break;
}
cmd = cmd_fut => {
if curr_cmd.is_none() {
curr_cmd = cmd;
}
}
len = read_fut => {
trace!("read line {:?}", next_line);
// res should not be None here
2021-02-23 02:47:00 +00:00
let resp = parse_response(next_line)?;
// if this is the very first response, then it's a greeting
if let Some(greeting_tx) = greeting_tx.take() {
greeting_tx.send(());
2021-02-23 04:01:39 +00:00
}
if let Response::Done(_) = resp {
// since this is the DONE message, clear curr_cmd so another one can be sent
if let Some((_, cmd_tx)) = curr_cmd.take() {
cmd_tx.send(resp)?;
2021-02-22 20:36:08 +00:00
}
} else if let Some((ref cmd, ref mut cmd_tx)) = curr_cmd {
cmd_tx.send(resp)?;
2021-02-21 13:42:40 +00:00
}
}
2021-02-20 07:30:58 +00:00
}
2021-02-20 05:03:33 +00:00
}
2021-02-21 13:42:40 +00:00
let conn = reader.into_inner();
Ok(conn)
2021-02-20 05:03:33 +00:00
}