Handle events inside PProxy

Author: Ma233

src/error.rs

@@ -10,6 +10,8 @@ pub enum Error {
     EssentialTaskClosed,
     #[error("Litep2p error: {0}")]
     Litep2p(#[from] litep2p::Error),
+    #[error("Litep2p request response error: {0:?}")]
+    Litep2pRequestResponseError(litep2p::protocol::request_response::RequestResponseError),
     #[error("Unexpected response type")]
     UnexpectedResponseType,
 }
@@ -25,3 +27,9 @@ impl From<tokio::sync::oneshot::error::RecvError> for Error {
         Error::EssentialTaskClosed
     }
 }
+
+impl From<litep2p::protocol::request_response::RequestResponseError> for Error {
+    fn from(err: litep2p::protocol::request_response::RequestResponseError) -> Self {
+        Error::Litep2pRequestResponseError(err)
+    }
+}

src/lib.rs

@@ -1,12 +1,14 @@
+use std::collections::HashMap;
 use std::net::SocketAddr;
 
 use litep2p::protocol::request_response::DialOptions;
+use litep2p::protocol::request_response::RequestResponseEvent;
+use litep2p::types::RequestId;
 use litep2p::PeerId;
 use multiaddr::Multiaddr;
 use multiaddr::Protocol;
 use tokio::sync::mpsc;
 use tokio::sync::oneshot;
-use tracing::trace;
 use tracing::warn;
 
 use crate::command::proto::AddPeerRequest;
@@ -29,6 +31,9 @@ const DEFAULT_CHANNEL_SIZE: usize = 4096;
 /// Public result type used by the crate.
 pub type Result<T> = std::result::Result<T, error::Error>;
 
+type CommandNotification = Result<PProxyCommandResponse>;
+type CommandNotifier = oneshot::Sender<CommandNotification>;
+
 #[derive(Debug)]
 pub enum PProxyCommand {
     AddPeer(AddPeerRequest),
@@ -41,12 +46,13 @@ pub enum PProxyCommandResponse {
 }
 
 pub struct PProxy {
-    command_rx: mpsc::Receiver<(PProxyCommand, oneshot::Sender<PProxyCommandResponse>)>,
+    command_rx: mpsc::Receiver<(PProxyCommand, CommandNotifier)>,
+    tunnel_notifier: HashMap<RequestId, CommandNotifier>,
     p2p_server: P2pServer,
 }
 
 pub struct PProxyHandle {
-    command_tx: mpsc::Sender<(PProxyCommand, oneshot::Sender<PProxyCommandResponse>)>,
+    command_tx: mpsc::Sender<(PProxyCommand, CommandNotifier)>,
 }
 
 impl PProxy {
@@ -56,6 +62,7 @@ impl PProxy {
         (
             Self {
                 command_rx,
+                tunnel_notifier: HashMap::new(),
                 p2p_server: P2pServer::new(server_addr),
             },
             PProxyHandle { command_tx },
@@ -70,41 +77,75 @@ impl PProxy {
 
                 event = self.p2p_server.next_event() => match event {
                     None => return,
-                    Some(ref event) => if let Err (error) = self.handle_p2p_server_event(event).await {
+                    Some(ref event) => if let Err(error) = self.handle_p2p_server_event(event).await {
                         warn!("failed to handle event {:?}: {:?}", event, error);
 
                     }
                 },
 
                 command = self.command_rx.recv() => match command {
                     None => return,
-                    Some((ref command, ref tx)) => match command {
-                        PProxyCommand::AddPeer(request) => {
-                            if let Err(error) = self.on_add_peer(request, tx).await {
-                                warn!("failed to handle command {:?}: {:?}", command, error);
-                            }
-                        }
-                        PProxyCommand::RequestHttpServer(request) => {
-                            if let Err(error) = self.on_request_http_server(request, tx).await {
-                                warn!("failed to handle command {:?}: {:?}", command, error);
-                            }
-                        }
+                    Some((ref command, tx)) => if let Err(error) = self.handle_command(command, tx).await {
+                        warn!("failed to handle command {:?}: {:?}", command, error);
                     }
                 }
             }
         }
     }
 
     async fn handle_p2p_server_event(&mut self, event: &P2pServerEvent) -> Result<()> {
-        trace!("handle P2pServerEvent: {:?}", event);
+        match event {
+            P2pServerEvent::TunnelEvent(RequestResponseEvent::RequestReceived {
+                request_id,
+                ..
+            }) => self
+                .p2p_server
+                .tunnel_handle
+                .send_response(*request_id, vec![1, 2, 3, 4]),
+
+            P2pServerEvent::TunnelEvent(RequestResponseEvent::RequestFailed {
+                request_id,
+                error,
+                ..
+            }) => {
+                let Some(tx) = self.tunnel_notifier.remove(request_id) else {
+                    todo!();
+                };
+                tx.send(Err(Error::Litep2pRequestResponseError(error.clone())))
+                    .map_err(|_| Error::EssentialTaskClosed)?;
+            }
+
+            P2pServerEvent::TunnelEvent(RequestResponseEvent::ResponseReceived {
+                request_id,
+                response,
+                ..
+            }) => {
+                let Some(tx) = self.tunnel_notifier.remove(request_id) else {
+                    todo!();
+                };
+                tx.send(Ok(PProxyCommandResponse::RequestHttpServer(
+                    RequestHttpServerResponse {
+                        data: response.clone(),
+                    },
+                )))
+                .map_err(|_| Error::EssentialTaskClosed)?;
+            }
+            _ => {}
+        }
+
         Ok(())
     }
 
-    async fn on_add_peer(
-        &mut self,
-        request: &AddPeerRequest,
-        _rx: &oneshot::Sender<PProxyCommandResponse>,
-    ) -> Result<AddPeerResponse> {
+    async fn handle_command(&mut self, command: &PProxyCommand, tx: CommandNotifier) -> Result<()> {
+        match command {
+            PProxyCommand::AddPeer(request) => self.on_add_peer(request.clone(), tx).await,
+            PProxyCommand::RequestHttpServer(request) => {
+                self.on_request_http_server(request.clone(), tx).await
+            }
+        }
+    }
+
+    async fn on_add_peer(&mut self, request: AddPeerRequest, tx: CommandNotifier) -> Result<()> {
         let addr: Multiaddr = request
             .address
             .parse()
@@ -123,27 +164,30 @@ impl PProxy {
             .litep2p
             .add_known_address(peer_id, vec![addr].into_iter());
 
-        Ok(AddPeerResponse {
+        tx.send(Ok(PProxyCommandResponse::AddPeer(AddPeerResponse {
             peer_id: peer_id.to_string(),
-        })
+        })))
+        .map_err(|_| Error::EssentialTaskClosed)
     }
 
     async fn on_request_http_server(
         &mut self,
-        request: &RequestHttpServerRequest,
-        _rx: &oneshot::Sender<PProxyCommandResponse>,
+        request: RequestHttpServerRequest,
+        rx: CommandNotifier,
     ) -> Result<()> {
         let peer_id = request
             .peer_id
             .parse()
             .map_err(|_| Error::PeerIdParseError(request.peer_id.clone()))?;
 
-        let _request_id = self
+        let request_id = self
             .p2p_server
             .tunnel_handle
             .send_request(peer_id, request.data.clone(), DialOptions::Dial)
             .await?;
 
+        self.tunnel_notifier.insert(request_id, rx);
+
         Ok(())
     }
 }
@@ -156,7 +200,7 @@ impl PProxyHandle {
             .send((PProxyCommand::AddPeer(request), tx))
             .await?;
 
-        let response = rx.await?;
+        let response = rx.await??;
 
         match response {
             PProxyCommandResponse::AddPeer(response) => Ok(response),
@@ -174,7 +218,7 @@ impl PProxyHandle {
             .send((PProxyCommand::RequestHttpServer(request), tx))
             .await?;
 
-        let response = rx.await?;
+        let response = rx.await??;
 
         match response {
             PProxyCommandResponse::RequestHttpServer(response) => Ok(response),

src/server.rs

@@ -5,6 +5,7 @@ use futures::StreamExt;
 use litep2p::config::ConfigBuilder;
 use litep2p::protocol::request_response::Config as RequestResponseConfig;
 use litep2p::protocol::request_response::ConfigBuilder as RequestResponseConfigBuilder;
+use litep2p::protocol::request_response::RequestResponseEvent;
 use litep2p::protocol::request_response::RequestResponseHandle;
 use litep2p::transport::quic::config::Config as QuicConfig;
 use litep2p::types::protocol::ProtocolName;
@@ -14,8 +15,8 @@ use litep2p::Litep2pEvent;
 #[derive(Debug)]
 pub enum P2pServerEvent {
     Litep2p(Litep2pEvent),
-    ControlEvent,
-    TunnelEvent,
+    ControlEvent(RequestResponseEvent),
+    TunnelEvent(RequestResponseEvent),
 }
 
 pub struct P2pServer {
@@ -74,11 +75,11 @@ impl P2pServer {
             ev = self.litep2p.next_event() => {
                 ev.map(P2pServerEvent::Litep2p)
             }
-            _ev = self.control_handle.next() => {
-                Some(P2pServerEvent::ControlEvent)
+            ev = self.control_handle.next() => {
+                ev.map(P2pServerEvent::ControlEvent)
             }
-            _ev = self.tunnel_handle.next() => {
-                Some(P2pServerEvent::TunnelEvent)
+            ev = self.tunnel_handle.next() => {
+                ev.map(P2pServerEvent::TunnelEvent)
             }
         }
     }