[Add] Pipes (experimental)

host/src/lib.rs

@@ -3,13 +3,22 @@ use mtp_codec::{
     registry::{Registry, VersionedCodec},
 };
 use mtp_common::CommunicationError;
+#[cfg(feature = "pipes")]
+pub use mtp_common::PipeError;
 use std::net::IpAddr;
 #[cfg(feature = "crypto")]
 use std::pin::Pin;
 use std::{error::Error, fmt};
+#[cfg(feature = "pipes")]
+use std::collections::HashMap;
+use std::sync::Arc;
+use tokio::sync::{Mutex, mpsc};
 #[cfg(feature = "crypto")]
 use tokio::time::Duration;
 
+#[cfg(feature = "pipes")]
+use mtp_transport::PipeReader;
+
 pub use MTPConnection as Connection;
 pub use MTPHost as Host;
 pub use mtp_transport::Policy;
@@ -17,6 +26,9 @@ pub use mtp_transport::Receiver;
 pub use mtp_transport::SendMode;
 pub use mtp_transport::Sender;
 
+#[cfg(feature = "pipes")]
+pub use mtp_transport::PipeWriter;
+
 /* ---- async callback type aliases ---- */
 #[cfg(feature = "crypto")]
 pub type GetExistingClient = Box<
@@ -47,6 +59,162 @@ pub enum AuthenticationPolicy {
     Unauthenticated,
 }
 
+#[cfg(feature = "pipes")]
+pub struct PipeHandle {
+    pipe_id: u32,
+    description: String,
+    sender: Sender,
+    response_rx: tokio::sync::oneshot::Receiver<Result<bool, PipeError>>,
+}
+
+#[cfg(feature = "pipes")]
+impl PipeHandle {
+    pub fn pipe_id(&self) -> u32 {
+        self.pipe_id
+    }
+
+    pub fn description(&self) -> &str {
+        &self.description
+    }
+
+    pub async fn wait(self) -> Result<Option<PipeWriter>, PipeError> {
+        match self.response_rx.await {
+            Ok(Ok(true)) => {
+                let writer = self
+                    .sender
+                    .open_pipe(self.pipe_id, &self.description)
+                    .await
+                    .map_err(PipeError::from)?;
+                Ok(Some(writer))
+            }
+            Ok(Ok(false)) => Ok(None),
+            Ok(Err(e)) => Err(e),
+            Err(_) => Err(PipeError::StreamClosed),
+        }
+    }
+}
+
+#[cfg(feature = "pipes")]
+pub struct PipeRequest {
+    pipe_id: u32,
+    description: String,
+    sender: Sender,
+    dispatcher: Arc<PipeDispatcher>,
+}
+
+#[cfg(feature = "pipes")]
+impl PipeRequest {
+    pub fn id(&self) -> u32 {
+        self.pipe_id
+    }
+
+    pub fn description(&self) -> &str {
+        &self.description
+    }
+
+    pub async fn accept(self) -> Result<PipeReader, PipeError> {
+        let (pipe_tx, pipe_rx) = tokio::sync::oneshot::channel();
+        {
+            let mut pending = self.dispatcher.pending_pipes.lock().await;
+            pending.insert(self.pipe_id, pipe_tx);
+        }
+
+        let resp = CommunicationValue::new(mtp_codec::CommunicationType::PipeResponse)
+            .with_id(self.pipe_id)
+            .add_typed_default(DataType::Accepted, DataValue::BoolTrue);
+        self.sender.send(&resp).await.map_err(PipeError::from)?;
+
+        let timeout = self.dispatcher.policy.read_timeout;
+        tokio::time::timeout(timeout, pipe_rx)
+            .await
+            .map_err(|_| PipeError::HandshakeTimeout)?
+            .map_err(|_| PipeError::StreamClosed)
+    }
+
+    pub async fn deny(self) -> Result<(), PipeError> {
+        let resp = CommunicationValue::new(mtp_codec::CommunicationType::PipeResponse)
+            .with_id(self.pipe_id)
+            .add_typed_default(DataType::Accepted, DataValue::BoolFalse);
+        self.sender.send(&resp).await.map_err(PipeError::from)?;
+        Ok(())
+    }
+}
+
+#[cfg(feature = "pipes")]
+struct PipeDispatcher {
+    pending_creations: Mutex<HashMap<u32, tokio::sync::oneshot::Sender<Result<bool, PipeError>>>>,
+    pending_pipes: Mutex<HashMap<u32, tokio::sync::oneshot::Sender<PipeReader>>>,
+    policy: Arc<mtp_transport::Policy>,
+}
+
+#[cfg(not(feature = "pipes"))]
+struct PipeDispatcher;
+
+#[cfg(not(feature = "pipes"))]
+pub(crate) struct PipeRequest;
+
+#[cfg(feature = "pipes")]
+async fn run_dispatcher(
+    receiver: Receiver,
+    sender: Sender,
+    app_tx: mpsc::Sender<Result<CommunicationValue, CommunicationError>>,
+    pipe_req_tx: mpsc::Sender<PipeRequest>,
+    dispatcher: Arc<PipeDispatcher>,
+) {
+    let pipe_req_type =
+        mtp_codec::CommunicationType::PipeRequest.to_id(&mtp_codec::TypeMap::latest());
+    let pipe_resp_type =
+        mtp_codec::CommunicationType::PipeResponse.to_id(&mtp_codec::TypeMap::latest());
+
+    loop {
+        match receiver.receive_event().await {
+            Ok(mtp_transport::TransportEvent::Message(msg)) => {
+                if msg.get_type() == pipe_req_type {
+                    let pipe_id = msg.get_id();
+                    let description = msg
+                        .get_str(DataType::Description)
+                        .unwrap_or("")
+                        .to_string();
+                    let req = PipeRequest {
+                        pipe_id,
+                        description,
+                        sender: sender.clone(),
+                        dispatcher: dispatcher.clone(),
+                    };
+                    let _ = pipe_req_tx.send(req).await;
+                    continue;
+                }
+
+                if msg.get_type() == pipe_resp_type {
+                    let pipe_id = msg.get_id();
+                    let accepted = msg.get_bool(DataType::Accepted).unwrap_or(false);
+                    let mut pending = dispatcher.pending_creations.lock().await;
+                    if let Some(tx) = pending.remove(&pipe_id) {
+                        let _ = tx.send(Ok(accepted));
+                    }
+                    continue;
+                }
+
+                if app_tx.send(Ok(msg)).await.is_err() {
+                    break;
+                }
+            }
+            Ok(mtp_transport::TransportEvent::Pipe(reader)) => {
+                let pipe_id = reader.pipe_id();
+                let mut pending = dispatcher.pending_pipes.lock().await;
+                if let Some(tx) = pending.remove(&pipe_id) {
+                    let _ = tx.send(reader);
+                }
+            }
+            Err(e) => {
+                if app_tx.send(Err(e)).await.is_err() {
+                    break;
+                }
+            }
+        }
+    }
+}
+
 /* Host configuration. */
 pub struct HostConfig {
     pub ip: IpAddr,
@@ -180,7 +348,11 @@ pub struct MTPConnection {
     pub codec: VersionedCodec,
     pub sender: Sender,
     pub receiver: Receiver,
+    app_rx: Mutex<mpsc::Receiver<Result<CommunicationValue, CommunicationError>>>,
+    pipe_req_rx: Mutex<mpsc::Receiver<PipeRequest>>,
+    pipe_dispatcher: Arc<PipeDispatcher>,
     pub description: Option<String>,
+    _dispatcher_task: tokio::task::JoinHandle<()>,
     #[cfg(feature = "crypto")]
     pub auth_state: AuthState,
     #[cfg(feature = "crypto")]
@@ -242,11 +414,10 @@ impl MTPHost {
                     Ok(result) => result,
                     Err(_) => Err(AcceptError::AuthenticationTimedOut),
                 };
-                return Ok(self.configure_pongs(connection?));
+                return connection;
             }
             AuthenticationPolicy::AllowAuthentication => {
-                let connection = self.accept_allow_auth(sender, receiver).await?;
-                return Ok(self.configure_pongs(connection));
+                return self.accept_allow_auth(sender, receiver).await;
             }
             AuthenticationPolicy::Unauthenticated => {
                 let first_msg = match receiver.receive().await {
@@ -278,19 +449,16 @@ impl MTPHost {
                     DataValue::Str(s) => Some(s.clone()),
                     _ => None,
                 };
-                Ok(self.configure_pongs(Some(MTPConnection {
-                    version: negotiated,
-                    codec,
+                Ok(Some(self.connection_from_parts(
                     sender,
                     receiver,
+                    negotiated,
+                    codec,
                     description,
-                    #[cfg(feature = "crypto")]
-                    auth_state: AuthState::Unauthenticated,
-                    #[cfg(feature = "crypto")]
-                    client_id: rand::random(),
-                    #[cfg(feature = "crypto")]
-                    client_public_key: None,
-                })))
+                    AuthState::Unauthenticated,
+                    rand::random(),
+                    None,
+                )))
             }
         }
 
@@ -318,13 +486,13 @@ impl MTPHost {
                 DataValue::Str(s) => Some(s.clone()),
                 _ => None,
             };
-            return Ok(self.configure_pongs(Some(MTPConnection {
-                version: negotiated,
-                codec,
+            return Ok(Some(self.connection_from_parts(
                 sender,
                 receiver,
+                negotiated,
+                codec,
                 description,
-            })));
+            )));
         }
     }
 
@@ -336,16 +504,204 @@ impl MTPHost {
         &self.registry
     }
 
-    fn configure_pongs(&self, connection: Option<MTPConnection>) -> Option<MTPConnection> {
-        if let Some(connection) = connection {
+    #[cfg(not(feature = "crypto"))]
+    fn connection_from_parts(
+        &self,
+        sender: Sender,
+        receiver: Receiver,
+        version: Version,
+        codec: VersionedCodec,
+        description: Option<String>,
+    ) -> MTPConnection {
+        #[cfg(feature = "pipes")]
+        {
             if self.config.send_pongs {
-                connection
-                    .receiver
-                    .respond_to_pings(connection.sender.clone());
+                receiver.respond_to_pings(sender.clone());
             }
-            Some(connection)
-        } else {
-            None
+
+            let (app_tx, app_rx) =
+                mpsc::channel(self.config.policy.receiver_queue_capacity);
+            let (pipe_req_tx, pipe_req_rx) =
+                mpsc::channel(self.config.policy.receiver_queue_capacity);
+
+            let dispatcher = Arc::new(PipeDispatcher {
+                pending_creations: Mutex::new(HashMap::new()),
+                pending_pipes: Mutex::new(HashMap::new()),
+                policy: Arc::new(self.config.policy),
+            });
+
+            let dispatcher_clone = dispatcher.clone();
+            let receiver_clone = receiver.clone();
+            let sender_clone = sender.clone();
+            let task = tokio::spawn(run_dispatcher(
+                receiver_clone,
+                sender_clone,
+                app_tx,
+                pipe_req_tx,
+                dispatcher_clone,
+            ));
+
+            MTPConnection {
+                version,
+                codec,
+                sender,
+                receiver,
+                app_rx: Mutex::new(app_rx),
+                pipe_req_rx: Mutex::new(pipe_req_rx),
+                pipe_dispatcher: dispatcher,
+                description,
+                _dispatcher_task: task,
+            }
+        }
+
+        #[cfg(not(feature = "pipes"))]
+        {
+            if self.config.send_pongs {
+                receiver.respond_to_pings(sender.clone());
+            }
+
+            let (_, app_rx) = mpsc::channel::<Result<CommunicationValue, CommunicationError>>(1);
+            let (_, pipe_req_rx) = mpsc::channel::<PipeRequest>(1);
+            let dispatcher = Arc::new(PipeDispatcher);
+            let task = tokio::spawn(async {});
+
+            MTPConnection {
+                version,
+                codec,
+                sender,
+                receiver,
+                app_rx: Mutex::new(app_rx),
+                pipe_req_rx: Mutex::new(pipe_req_rx),
+                pipe_dispatcher: dispatcher,
+                description,
+                _dispatcher_task: task,
+            }
+        }
+    }
+
+    #[cfg(feature = "crypto")]
+    fn connection_from_parts(
+        &self,
+        sender: Sender,
+        receiver: Receiver,
+        version: Version,
+        codec: VersionedCodec,
+        description: Option<String>,
+        auth_state: AuthState,
+        client_id: u64,
+        client_public_key: Option<mtp_crypto::PublicKeyBundle>,
+    ) -> MTPConnection {
+        #[cfg(feature = "pipes")]
+        {
+            if self.config.send_pongs {
+                receiver.respond_to_pings(sender.clone());
+            }
+
+            let (app_tx, app_rx) =
+                mpsc::channel(self.config.policy.receiver_queue_capacity);
+            let (pipe_req_tx, pipe_req_rx) =
+                mpsc::channel(self.config.policy.receiver_queue_capacity);
+
+            let dispatcher = Arc::new(PipeDispatcher {
+                pending_creations: Mutex::new(HashMap::new()),
+                pending_pipes: Mutex::new(HashMap::new()),
+                policy: Arc::new(self.config.policy),
+            });
+
+            let dispatcher_clone = dispatcher.clone();
+            let receiver_clone = receiver.clone();
+            let sender_clone = sender.clone();
+            let task = tokio::spawn(run_dispatcher(
+                receiver_clone,
+                sender_clone,
+                app_tx,
+                pipe_req_tx,
+                dispatcher_clone,
+            ));
+
+            MTPConnection {
+                version,
+                codec,
+                sender,
+                receiver,
+                app_rx: Mutex::new(app_rx),
+                pipe_req_rx: Mutex::new(pipe_req_rx),
+                pipe_dispatcher: dispatcher,
+                description,
+                _dispatcher_task: task,
+                auth_state,
+                client_id,
+                client_public_key,
+            }
+        }
+
+        #[cfg(not(feature = "pipes"))]
+        {
+            if self.config.send_pongs {
+                receiver.respond_to_pings(sender.clone());
+            }
+
+            let (_, app_rx) = mpsc::channel::<Result<CommunicationValue, CommunicationError>>(1);
+            let (_, pipe_req_rx) = mpsc::channel::<PipeRequest>(1);
+            let dispatcher = Arc::new(PipeDispatcher);
+            let task = tokio::spawn(async {});
+
+            MTPConnection {
+                version,
+                codec,
+                sender,
+                receiver,
+                app_rx: Mutex::new(app_rx),
+                pipe_req_rx: Mutex::new(pipe_req_rx),
+                pipe_dispatcher: dispatcher,
+                description,
+                _dispatcher_task: task,
+                auth_state,
+                client_id,
+                client_public_key,
+            }
+        }
+    }
+}
+
+#[cfg(feature = "pipes")]
+impl MTPConnection {
+    pub async fn receive(&self) -> Result<CommunicationValue, CommunicationError> {
+        let mut rx = self.app_rx.lock().await;
+        match rx.recv().await {
+            Some(result) => result,
+            None => Err(CommunicationError::StreamClosed),
+        }
+    }
+
+    pub async fn create_pipe(&self, description: &str) -> Result<PipeHandle, PipeError> {
+        let pipe_id = rand::random::<u32>();
+        let (tx, rx) = tokio::sync::oneshot::channel();
+
+        {
+            let mut pending = self.pipe_dispatcher.pending_creations.lock().await;
+            pending.insert(pipe_id, tx);
+        }
+
+        let request = CommunicationValue::new(mtp_codec::CommunicationType::PipeRequest)
+            .with_id(pipe_id)
+            .add_typed_default(DataType::Description, DataValue::Str(description.into()));
+
+        self.sender.send(&request).await.map_err(PipeError::from)?;
+
+        Ok(PipeHandle {
+            pipe_id,
+            description: description.to_string(),
+            sender: self.sender.clone(),
+            response_rx: rx,
+        })
+    }
+
+    pub async fn receive_pipe(&self) -> Result<PipeRequest, CommunicationError> {
+        let mut rx = self.pipe_req_rx.lock().await;
+        match rx.recv().await {
+            Some(req) => Ok(req),
+            None => Err(CommunicationError::StreamClosed),
         }
     }
 }
@@ -668,16 +1024,16 @@ impl MTPHost {
         let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone())
             .expect("negotiated version must be registered");
 
-        Ok(Some(MTPConnection {
-            version: negotiated,
-            codec,
+        Ok(Some(self.connection_from_parts(
             sender,
             receiver,
+            negotiated,
+            codec,
             description,
-            auth_state: AuthState::Authenticated,
-            client_id: assigned_id,
-            client_public_key: Some(client_bundle),
-        }))
+            AuthState::Authenticated,
+            assigned_id,
+            Some(client_bundle),
+        )))
     }
 
     /*
@@ -782,16 +1138,16 @@ impl MTPHost {
             };
             let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone())
                 .expect("negotiated version must be registered");
-            return Ok(Some(MTPConnection {
-                version: negotiated,
-                codec,
+            return Ok(Some(self.connection_from_parts(
                 sender,
                 receiver,
+                negotiated,
+                codec,
                 description,
-                auth_state: AuthState::Unauthenticated,
-                client_id: rand::random(),
-                client_public_key: None,
-            }));
+                AuthState::Unauthenticated,
+                rand::random(),
+                None,
+            )));
         }
 
         sender.close();