panorama/imap/src/client/inner.rs

381 lines
13 KiB
Rust
Raw Normal View History

2021-03-01 07:48:28 +00:00
use std::collections::{HashSet, VecDeque};
2021-02-20 05:03:33 +00:00
use std::pin::Pin;
use std::sync::Arc;
2021-02-21 13:42:40 +00:00
use std::task::{Context, Poll, Waker};
2021-02-20 05:03:33 +00:00
use anyhow::{Context as AnyhowContext, Error, Result};
use futures::{
2021-03-01 07:48:28 +00:00
future::{self, Either, Future, FutureExt, TryFutureExt},
stream::StreamExt,
};
2021-03-01 07:48:28 +00:00
use parking_lot::RwLock;
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},
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};
2021-02-22 20:43:09 +00:00
use crate::response::{Capability, Response, ResponseCode, Status};
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-02-22 20:36:08 +00:00
pub type CapsLock = Arc<RwLock<Option<HashSet<Capability>>>>;
pub type ResponseFuture = Box<dyn Future<Output = Result<Response>> + Send + Unpin>;
pub type ResponseSender = mpsc::UnboundedSender<Response>;
pub type ResponseStream = mpsc::UnboundedReceiver<Response>;
type ResultQueue = Arc<RwLock<VecDeque<HandlerResult>>>;
2021-02-23 02:47:00 +00:00
pub type GreetingState = Arc<RwLock<(Option<Response>, Option<Waker>)>>;
2021-03-01 07:48:28 +00:00
pub const TAG_PREFIX: &str = "ptag";
2021-02-20 05:03:33 +00:00
struct HandlerResult {
id: usize,
end: Option<oneshot::Sender<Response>>,
sender: ResponseSender,
waker: Option<Waker>,
}
2021-02-21 13:42:40 +00:00
/// The lower-level Client struct, that is shared by all of the exported structs in the state machine.
2021-02-20 05:03:33 +00:00
pub struct Client<C> {
2021-02-21 13:54:46 +00:00
config: ClientConfig,
/// write half of the connection
2021-02-20 05:03:33 +00:00
conn: WriteHalf<C>,
/// counter for monotonically incrementing unique ids
2021-02-20 05:03:33 +00:00
id: usize,
results: ResultQueue,
2021-02-20 05:03:33 +00:00
2021-02-21 13:42:40 +00:00
/// cached set of capabilities
caps: CapsLock,
2021-02-21 13:42:40 +00:00
/// join handle for the listener thread
listener_handle: JoinHandle<Result<ReadHalf<C>>>,
/// used for telling the listener thread to stop and return the read half
exit_tx: mpsc::Sender<()>,
/// used for receiving the greeting
2021-02-21 13:54:46 +00:00
greeting: GreetingState,
2021-02-20 05:03:33 +00:00
}
impl<C> Client<C>
where
C: AsyncRead + AsyncWrite + Unpin + Send + 'static,
{
/// Creates a new client that wraps a connection
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);
2021-02-23 02:47:00 +00:00
let results = Arc::new(RwLock::new(VecDeque::new()));
2021-02-21 13:42:40 +00:00
let (exit_tx, exit_rx) = mpsc::channel(1);
2021-02-23 02:47:00 +00:00
let greeting = Arc::new(RwLock::new((None, None)));
let caps: CapsLock = Arc::new(RwLock::new(None));
2021-03-02 22:47:33 +00:00
let listener_handle = tokio::spawn(
listen(
read_half,
caps.clone(),
results.clone(),
exit_rx,
greeting.clone(),
)
.map_err(|err| {
error!("Help, the listener loop died: {}", err);
err
}),
);
2021-02-20 05:03:33 +00:00
Client {
2021-02-21 13:42:40 +00:00
config,
2021-02-20 05:03:33 +00:00
conn: write_half,
id: 0,
2021-02-20 07:30:58 +00:00
results,
listener_handle,
caps,
2021-02-21 13:42:40 +00:00
exit_tx,
greeting,
2021-02-20 05:03:33 +00:00
}
}
2021-02-21 13:54:46 +00:00
/// Returns a future that doesn't resolve until we receive a greeting from the server.
pub fn wait_for_greeting(&self) -> GreetingWaiter {
2021-02-21 13:42:40 +00:00
debug!("waiting for greeting");
2021-02-21 13:54:46 +00:00
GreetingWaiter(self.greeting.clone())
2021-02-21 13:42:40 +00:00
}
2021-02-20 05:03:33 +00:00
/// Sends a command to the server and returns a handle to retrieve the result
pub async fn execute(&mut self, cmd: Command) -> Result<(ResponseFuture, ResponseStream)> {
// debug!("executing command {:?}", cmd);
2021-02-20 05:03:33 +00:00
let id = self.id;
self.id += 1;
// create a channel for sending the final response
let (end_tx, end_rx) = oneshot::channel();
// create a channel for sending responses for this particular client call
// this should queue up responses
let (tx, rx) = mpsc::unbounded_channel();
2021-03-02 22:47:33 +00:00
debug!("EX[{}]: adding handler result to the handlers queue", id);
2021-02-20 05:03:33 +00:00
{
2021-02-20 07:30:58 +00:00
let mut handlers = self.results.write();
handlers.push_back(HandlerResult {
id,
end: Some(end_tx),
sender: tx,
waker: None,
});
2021-02-20 05:03:33 +00:00
}
2021-03-02 22:47:33 +00:00
debug!("EX[{}]: send the command to the server", id);
2021-02-21 01:13:10 +00:00
let cmd_str = format!("{}{} {}\r\n", TAG_PREFIX, id, cmd);
2021-02-24 12:43:50 +00:00
// debug!("[{}] writing to socket: {:?}", id, cmd_str);
2021-02-20 05:03:33 +00:00
self.conn.write_all(cmd_str.as_bytes()).await?;
self.conn.flush().await?;
2021-02-24 12:43:50 +00:00
// debug!("[{}] written.", id);
// let resp = ExecWaiter(self, id, false).await;
2021-02-23 02:47:00 +00:00
// let resp = {
// let mut handlers = self.results.write();
// handlers.remove(&id).unwrap().0.unwrap()
// };
// let resp = end_rx.await?;
2021-03-02 22:47:33 +00:00
debug!("EX[{}]: hellosu", id);
let q = self.results.clone();
// let end = Box::new(end_rx.map_err(|err| Error::from).map(move |resp| resp));
let end = Box::new(end_rx.map_err(Error::from).map(move |resp| {
2021-03-02 22:47:33 +00:00
debug!("EX[{}]: -end result- {:?}", id, resp);
// pop the first entry from the list
let mut results = q.write();
results.pop_front();
resp
}));
Ok((end, rx))
2021-02-20 05:03:33 +00:00
}
/// Executes the CAPABILITY command
2021-02-23 02:47:00 +00:00
pub async fn capabilities(&mut self, force: bool) -> Result<()> {
{
2021-02-23 04:01:39 +00:00
let caps = self.caps.read();
2021-02-23 02:47:00 +00:00
if caps.is_some() && !force {
return Ok(());
}
}
2021-02-20 05:03:33 +00:00
let cmd = Command::Capability;
// debug!("sending: {:?} {:?}", cmd, cmd.to_string());
2021-02-23 02:47:00 +00:00
let (result, intermediate) = self
2021-02-21 13:42:40 +00:00
.execute(cmd)
.await
.context("error executing CAPABILITY command")?;
let result = result.await?;
debug!("cap resp: {:?}", result);
2021-02-23 02:47:00 +00:00
if let Some(Response::Capabilities(new_caps)) = UnboundedReceiverStream::new(intermediate)
.filter(|resp| future::ready(matches!(resp, Response::Capabilities(_))))
.next()
.await
2021-02-23 02:47:00 +00:00
{
2021-02-23 04:01:39 +00:00
let mut caps = self.caps.write();
2021-02-23 02:47:00 +00:00
*caps = Some(new_caps.iter().cloned().collect());
}
2021-02-20 07:30:58 +00:00
Ok(())
2021-02-20 05:03:33 +00:00
}
2021-02-21 13:42:40 +00:00
/// Attempts to upgrade this connection using STARTTLS
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 (resp, _) = self.execute(Command::Starttls).await?;
let resp = resp.await?;
2021-02-21 13:42:40 +00:00
debug!("server response to starttls: {:?}", resp);
debug!("sending exit for upgrade");
2021-02-21 13:42:40 +00:00
self.exit_tx.send(()).await?;
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))
}
2021-02-22 20:36:08 +00:00
/// Check if this client has a particular capability
2021-02-23 02:47:00 +00:00
pub async fn has_capability(&mut self, cap: impl AsRef<str>) -> Result<bool> {
2021-02-22 20:36:08 +00:00
let cap = cap.as_ref().to_owned();
debug!("checking for the capability: {:?}", cap);
2021-02-23 02:47:00 +00:00
let cap = parse_capability(cap)?;
2021-02-22 20:36:08 +00:00
2021-02-23 02:47:00 +00:00
self.capabilities(false).await?;
2021-02-23 04:01:39 +00:00
let caps = self.caps.read();
2021-02-22 20:36:08 +00:00
// TODO: refresh caps
let caps = caps.as_ref().unwrap();
let result = caps.contains(&cap);
debug!("cap result: {:?}", result);
Ok(result)
}
2021-02-21 13:42:40 +00:00
}
2021-02-21 13:54:46 +00:00
pub struct GreetingWaiter(GreetingState);
2021-02-21 13:42:40 +00:00
2021-02-21 13:54:46 +00:00
impl Future for GreetingWaiter {
2021-02-23 02:47:00 +00:00
type Output = Response;
2021-02-21 13:42:40 +00:00
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
let (state, waker) = &mut *self.0.write();
debug!("g {:?}", state);
2021-02-21 13:42:40 +00:00
if waker.is_none() {
*waker = Some(cx.waker().clone());
}
2021-02-23 02:47:00 +00:00
match state.take() {
Some(v) => Poll::Ready(v),
None => Poll::Pending,
2021-02-21 13:42:40 +00:00
}
}
2021-02-20 05:03:33 +00:00
}
2021-02-21 13:42:40 +00:00
/// Main listen loop for the application
async fn listen<C>(
conn: C,
caps: CapsLock,
results: ResultQueue,
2021-02-21 13:42:40 +00:00
mut exit: mpsc::Receiver<()>,
2021-02-21 13:54:46 +00:00
greeting: GreetingState,
2021-02-21 13:42:40 +00:00
) -> Result<C>
where
C: AsyncRead + Unpin,
{
2021-02-24 12:43:50 +00:00
// debug!("amogus");
2021-02-20 05:03:33 +00:00
let mut reader = BufReader::new(conn);
2021-02-21 13:42:40 +00:00
let mut greeting = Some(greeting);
2021-02-20 05:03:33 +00:00
loop {
let mut next_line = String::new();
2021-02-21 13:42:40 +00:00
let fut = reader.read_line(&mut next_line).fuse();
pin_mut!(fut);
let fut2 = exit.recv().fuse();
pin_mut!(fut2);
match future::select(fut, fut2).await {
2021-02-22 07:37:19 +00:00
Either::Left((res, _)) => {
let bytes = res.context("read failed")?;
if bytes == 0 {
bail!("connection probably died");
}
2021-03-02 22:47:33 +00:00
debug!("[LISTEN] got a new line {:?}", next_line);
2021-02-23 02:47:00 +00:00
let resp = parse_response(next_line)?;
2021-02-23 04:01:39 +00:00
// if this is the very first message, treat it as a greeting
if let Some(greeting) = greeting.take() {
let (greeting, waker) = &mut *greeting.write();
2021-03-02 22:47:33 +00:00
debug!("[LISTEN] received greeting!");
2021-02-23 02:47:00 +00:00
*greeting = Some(resp.clone());
if let Some(waker) = waker.take() {
waker.wake();
}
}
2021-02-23 04:01:39 +00:00
// update capabilities list
// TODO: probably not really necessary here (done somewhere else)?
if let Response::Capabilities(new_caps)
| Response::Data {
status: Status::Ok,
code: Some(ResponseCode::Capabilities(new_caps)),
..
} = &resp
{
let caps = &mut *caps.write();
*caps = Some(new_caps.iter().cloned().collect());
debug!("new caps: {:?}", caps);
}
match &resp {
2021-02-23 04:01:39 +00:00
Response::Data {
status: Status::Ok, ..
} => {
2021-02-23 04:01:39 +00:00
let mut results = results.write();
if let Some(HandlerResult { id, sender, .. }) = results.iter_mut().next() {
2021-03-02 22:47:33 +00:00
// we don't really care if it fails to send
// this just means that the other side has dropped the channel
//
// which is fine since that just means they don't care about
// intermediate messages
let _ = sender.send(resp);
debug!("[LISTEN] pushed to intermediate for id {}", id);
2021-02-23 04:01:39 +00:00
}
2021-03-02 22:47:33 +00:00
std::mem::drop(results);
debug!("[LISTEN] << unlocked self.results");
2021-02-21 13:42:40 +00:00
}
2021-02-22 20:36:08 +00:00
// bye
Response::Data {
status: Status::Bye,
..
} => {
bail!("disconnected");
}
Response::Done { tag, .. } => {
2021-02-22 20:43:09 +00:00
if tag.starts_with(TAG_PREFIX) {
2021-02-23 04:01:39 +00:00
// let id = tag.trim_start_matches(TAG_PREFIX).parse::<usize>()?;
let mut results = results.write();
if let Some(HandlerResult {
end: ref mut opt,
waker,
..
}) = results.iter_mut().next()
{
if let Some(opt) = opt.take() {
opt.send(resp).unwrap();
}
// *opt = Some(resp);
if let Some(waker) = waker.take() {
waker.wake();
}
}
2021-02-21 13:42:40 +00:00
}
}
2021-02-23 02:47:00 +00:00
_ => {}
2021-02-21 13:42:40 +00:00
}
}
2021-02-21 13:42:40 +00:00
Either::Right((_, _)) => {
debug!("exiting read loop");
break;
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
}