Separate handler requests from service requests (#152)

* fix stuff

* separate handler requests vs service requests

Author: divagant-martian

src/handler/active_requests.rs

@@ -2,7 +2,7 @@ use super::*;
 use delay_map::HashMapDelay;
 use more_asserts::debug_unreachable;
 
-pub(crate) struct ActiveRequests {
+pub(super) struct ActiveRequests {
     /// A list of raw messages we are awaiting a response from the remote.
     active_requests_mapping: HashMapDelay<NodeAddress, RequestCall>,
     // WHOAREYOU messages do not include the source node id. We therefore maintain another
@@ -13,29 +13,26 @@ pub(crate) struct ActiveRequests {
 }
 
 impl ActiveRequests {
-    pub(crate) fn new(request_timeout: Duration) -> Self {
+    pub fn new(request_timeout: Duration) -> Self {
         ActiveRequests {
             active_requests_mapping: HashMapDelay::new(request_timeout),
             active_requests_nonce_mapping: HashMap::new(),
         }
     }
 
-    pub(crate) fn insert(&mut self, node_address: NodeAddress, request_call: RequestCall) {
+    pub fn insert(&mut self, node_address: NodeAddress, request_call: RequestCall) {
         let nonce = *request_call.packet().message_nonce();
         self.active_requests_mapping
             .insert(node_address.clone(), request_call);
         self.active_requests_nonce_mapping
             .insert(nonce, node_address);
     }
 
-    pub(crate) fn get(&self, node_address: &NodeAddress) -> Option<&RequestCall> {
+    pub fn get(&self, node_address: &NodeAddress) -> Option<&RequestCall> {
         self.active_requests_mapping.get(node_address)
     }
 
-    pub(crate) fn remove_by_nonce(
-        &mut self,
-        nonce: &MessageNonce,
-    ) -> Option<(NodeAddress, RequestCall)> {
+    pub fn remove_by_nonce(&mut self, nonce: &MessageNonce) -> Option<(NodeAddress, RequestCall)> {
         match self.active_requests_nonce_mapping.remove(nonce) {
             Some(node_address) => match self.active_requests_mapping.remove(&node_address) {
                 Some(request_call) => Some((node_address, request_call)),
@@ -49,7 +46,7 @@ impl ActiveRequests {
         }
     }
 
-    pub(crate) fn remove(&mut self, node_address: &NodeAddress) -> Option<RequestCall> {
+    pub fn remove(&mut self, node_address: &NodeAddress) -> Option<RequestCall> {
         match self.active_requests_mapping.remove(node_address) {
             Some(request_call) => {
                 // Remove the associated nonce mapping.
@@ -75,7 +72,7 @@ impl ActiveRequests {
     // this makes is so that if there is a panic, the error is printed in the caller of this
     // function.
     #[track_caller]
-    pub(crate) fn check_invariant(&self) {
+    pub fn check_invariant(&self) {
         // First check that for every `MessageNonce` there is an associated `NodeAddress`.
         for (nonce, address) in self.active_requests_nonce_mapping.iter() {
             if !self.active_requests_mapping.contains_key(address) {

src/handler/mod.rs

@@ -153,6 +153,22 @@ pub struct Challenge {
     remote_enr: Option<Enr>,
 }
 
+/// Request ID from the handler's perspective.
+#[derive(Debug, Clone)]
+enum HandlerReqId {
+    /// Requests made by the handler.
+    Internal(RequestId),
+    /// Requests made from outside the handler.
+    External(RequestId),
+}
+
+/// A request queued for sending.
+struct PendingRequest {
+    contact: NodeContact,
+    request_id: HandlerReqId,
+    request: RequestBody,
+}
+
 /// Process to handle handshakes and sessions established from raw RPC communications between nodes.
 pub struct Handler {
     /// Configuration for the discv5 service.
@@ -169,7 +185,7 @@ pub struct Handler {
     /// The expected responses by SocketAddr which allows packets to pass the underlying filter.
     filter_expected_responses: Arc<RwLock<HashMap<SocketAddr, usize>>>,
     /// Requests awaiting a handshake completion.
-    pending_requests: HashMap<NodeAddress, Vec<(NodeContact, Request)>>,
+    pending_requests: HashMap<NodeAddress, Vec<PendingRequest>>,
     /// Currently in-progress outbound handshakes (WHOAREYOU packets) with peers.
     active_challenges: HashMapDelay<NodeAddress, Challenge>,
     /// Established sessions with peers.
@@ -191,6 +207,7 @@ type HandlerReturn = (
     mpsc::UnboundedSender<HandlerIn>,
     mpsc::Receiver<HandlerOut>,
 );
+
 impl Handler {
     /// A new Session service which instantiates the UDP socket send/recv tasks.
     pub async fn spawn(
@@ -274,13 +291,13 @@ impl Handler {
                 Some(handler_request) = self.service_recv.recv() => {
                     match handler_request {
                         HandlerIn::Request(contact, request) => {
-                           let id = request.id.clone();
-                           if let Err(request_error) =  self.send_request(contact, *request).await {
-                               // If the sending failed report to the application
-                               if let Err(e) = self.service_send.send(HandlerOut::RequestFailed(id, request_error)).await {
-                                   warn!("Failed to inform that request failed {}", e)
-                               }
-                           }
+                            let Request { id, body: request } = *request;
+                            if let Err(request_error) =  self.send_request(contact, HandlerReqId::External(id.clone()), request).await {
+                                // If the sending failed report to the application
+                                if let Err(e) = self.service_send.send(HandlerOut::RequestFailed(id, request_error)).await {
+                                    warn!("Failed to inform that request failed {}", e)
+                                }
+                            }
                         }
                         HandlerIn::Response(dst, response) => self.send_response(dst, *response).await,
                         HandlerIn::WhoAreYou(wru_ref, enr) => self.send_challenge(wru_ref, enr).await,
@@ -395,7 +412,7 @@ impl Handler {
             // increment the request retry count and restart the timeout
             trace!(
                 "Resending message: {} to {}",
-                request_call.request(),
+                request_call.body(),
                 node_address
             );
             self.send(node_address.clone(), request_call.packet().clone())
@@ -409,7 +426,8 @@ impl Handler {
     async fn send_request(
         &mut self,
         contact: NodeContact,
-        request: Request,
+        request_id: HandlerReqId,
+        request: RequestBody,
     ) -> Result<(), RequestError> {
         let node_address = contact.node_address();
 
@@ -426,15 +444,25 @@ impl Handler {
             self.pending_requests
                 .entry(node_address)
                 .or_insert_with(Vec::new)
-                .push((contact, request));
+                .push(PendingRequest {
+                    contact,
+                    request_id,
+                    request,
+                });
             return Ok(());
         }
 
         let (packet, initiating_session) = {
             if let Some(session) = self.sessions.get_mut(&node_address) {
                 // Encrypt the message and send
+                let request = match &request_id {
+                    HandlerReqId::Internal(id) | HandlerReqId::External(id) => Request {
+                        id: id.clone(),
+                        body: request.clone(),
+                    },
+                };
                 let packet = session
-                    .encrypt_message(self.node_id, &request.clone().encode())
+                    .encrypt_message(self.node_id, &request.encode())
                     .map_err(|e| RequestError::EncryptionFailed(format!("{:?}", e)))?;
                 (packet, false)
             } else {
@@ -450,7 +478,13 @@ impl Handler {
             }
         };
 
-        let call = RequestCall::new(contact, packet.clone(), request, initiating_session);
+        let call = RequestCall::new(
+            contact,
+            packet.clone(),
+            request_id,
+            request,
+            initiating_session,
+        );
         // let the filter know we are expecting a response
         self.add_expected_response(node_address.socket_addr);
         self.send(node_address.clone(), packet).await;
@@ -593,7 +627,7 @@ impl Handler {
             updated_enr,
             &self.node_id,
             &challenge_data,
-            &(request_call.request().clone().encode()),
+            &request_call.encode(),
         ) {
             Ok(v) => v,
             Err(e) => {
@@ -625,10 +659,7 @@ impl Handler {
                 // not have a session for a request) and the packet is not a PING (we are not
                 // trying to update an old session that may have expired.
                 let connection_direction = {
-                    match (
-                        request_call.initiating_session(),
-                        &request_call.request().body,
-                    ) {
+                    match (request_call.initiating_session(), &request_call.body()) {
                         (true, RequestBody::Ping { .. }) => ConnectionDirection::Incoming,
                         (true, _) => ConnectionDirection::Outgoing,
                         (false, _) => ConnectionDirection::Incoming,
@@ -668,13 +699,12 @@ impl Handler {
                 self.send(node_address.clone(), auth_packet).await;
 
                 let id = RequestId::random();
-                let request = Request {
-                    id: id.clone(),
-                    body: RequestBody::FindNode { distances: vec![0] },
-                };
-
-                session.awaiting_enr = Some(id);
-                if let Err(e) = self.send_request(contact, request).await {
+                let request = RequestBody::FindNode { distances: vec![0] };
+                session.awaiting_enr = Some(id.clone());
+                if let Err(e) = self
+                    .send_request(contact, HandlerReqId::Internal(id), request)
+                    .await
+                {
                     warn!("Failed to send Enr request {}", e)
                 }
             }
@@ -802,21 +832,35 @@ impl Handler {
                 self.pending_requests.entry(node_address)
             {
                 // If it exists, there must be a request here
-                let (contact, request) = entry.get_mut().remove(0);
+                let PendingRequest {
+                    contact,
+                    request_id,
+                    request,
+                } = entry.get_mut().remove(0);
                 if entry.get().is_empty() {
                     entry.remove();
                 }
-                let id = request.id.clone();
                 trace!("Sending next awaiting message. Node: {}", contact);
-                if let Err(request_error) = self.send_request(contact, request).await {
+                if let Err(request_error) = self
+                    .send_request(contact, request_id.clone(), request)
+                    .await
+                {
                     warn!("Failed to send next awaiting request {}", request_error);
                     // Inform the service that the request failed
-                    if let Err(e) = self
-                        .service_send
-                        .send(HandlerOut::RequestFailed(id, request_error))
-                        .await
-                    {
-                        warn!("Failed to inform that request failed {}", e);
+                    match request_id {
+                        HandlerReqId::Internal(_) => {
+                            // An internal request could not be sent. For now we do nothing about
+                            // this.
+                        }
+                        HandlerReqId::External(id) => {
+                            if let Err(e) = self
+                                .service_send
+                                .send(HandlerOut::RequestFailed(id, request_error))
+                                .await
+                            {
+                                warn!("Failed to inform that request failed {}", e);
+                            }
+                        }
                     }
                 }
             }
@@ -954,7 +998,10 @@ impl Handler {
     async fn handle_response(&mut self, node_address: NodeAddress, response: Response) {
         // Find a matching request, if any
         if let Some(mut request_call) = self.active_requests.remove(&node_address) {
-            if request_call.id() != &response.id {
+            let id = match request_call.id() {
+                HandlerReqId::Internal(id) | HandlerReqId::External(id) => id,
+            };
+            if id != &response.id {
                 trace!(
                     "Received an RPC Response to an unknown request. Likely late response. {}",
                     node_address
@@ -1055,13 +1102,19 @@ impl Handler {
     ) {
         // The Request has expired, remove the session.
         // Fail the current request
-        let request_id = request_call.request().id.clone();
-        if let Err(e) = self
-            .service_send
-            .send(HandlerOut::RequestFailed(request_id, error.clone()))
-            .await
-        {
-            warn!("Failed to inform request failure {}", e)
+        match request_call.id() {
+            HandlerReqId::Internal(_) => {
+                // Do not report failures on requests belonging to the handler.
+            }
+            HandlerReqId::External(id) => {
+                if let Err(e) = self
+                    .service_send
+                    .send(HandlerOut::RequestFailed(id.clone(), error.clone()))
+                    .await
+                {
+                    warn!("Failed to inform request failure {}", e)
+                }
+            }
         }
 
         let node_address = request_call.contact().node_address();
@@ -1083,13 +1136,20 @@ impl Handler {
                 .store(self.sessions.len(), Ordering::Relaxed);
         }
         if let Some(to_remove) = self.pending_requests.remove(node_address) {
-            for request in to_remove {
-                if let Err(e) = self
-                    .service_send
-                    .send(HandlerOut::RequestFailed(request.1.id, error.clone()))
-                    .await
-                {
-                    warn!("Failed to inform request failure {}", e)
+            for PendingRequest { request_id, .. } in to_remove {
+                match request_id {
+                    HandlerReqId::Internal(_) => {
+                        // Do not report failures on requests belonging to the handler.
+                    }
+                    HandlerReqId::External(id) => {
+                        if let Err(e) = self
+                            .service_send
+                            .send(HandlerOut::RequestFailed(id, error.clone()))
+                            .await
+                        {
+                            warn!("Failed to inform request failure {}", e)
+                        }
+                    }
                 }
             }
         }