Maintain session challenge (#52)

* Maintains sessions for invalid auth header sigs

* fmt

src/error.rs

@@ -1,3 +1,4 @@
+use crate::handler::Challenge;
 use rlp::DecoderError;
 
 #[derive(Debug)]
@@ -15,8 +16,8 @@ pub enum Discv5Error {
     InvalidRemotePublicKey,
     /// The secret key does not match the provided ENR.
     InvalidSecretKey,
-    /// An invalid signature was received.
-    InvalidSignature,
+    /// An invalid signature was received for a challenge.
+    InvalidChallengeSignature(Challenge),
     /// The Service channel has been closed early.
     ServiceChannelClosed,
     /// The discv5 service is not running.

src/handler/mod.rs

@@ -118,13 +118,14 @@ pub enum HandlerResponse {
 #[derive(Debug, Clone, PartialEq)]
 pub struct WhoAreYouRef(pub NodeAddress, MessageNonce);
 
-pub(crate) struct Challenge {
+#[derive(Debug)]
+pub struct Challenge {
     data: ChallengeData,
     remote_enr: Option<Enr>,
 }
 
-#[derive(Debug)]
 /// A request to a node that we are waiting for a response.
+#[derive(Debug)]
 pub(crate) struct RequestCall {
     contact: NodeContact,
     /// The raw discv5 packet sent.
@@ -759,6 +760,14 @@ impl Handler {
                             .await;
                     }
                 }
+                Err(Discv5Error::InvalidChallengeSignature(challenge)) => {
+                    warn!(
+                        "Authentication header contained invalid signature. Ignoring packet from: {}",
+                        node_address
+                    );
+                    // insert back the challenge
+                    self.active_challenges.insert(node_address, challenge);
+                }
                 Err(e) => {
                     warn!(
                         "Invalid Authentication header. Dropping session. Error: {:?}",

src/handler/session.rs

@@ -128,38 +128,31 @@ impl Session {
         ephem_pubkey: &[u8],
         enr_record: Option<Enr>,
     ) -> Result<(Session, Enr), Discv5Error> {
-        // generate session keys
-        let (decryption_key, encryption_key) = crypto::derive_keys_from_pubkey(
-            &local_key.read(),
-            local_id,
-            remote_id,
-            &challenge.data,
-            ephem_pubkey,
-        )?;
-
         // check and verify a potential ENR update
-        let session_enr = match (enr_record, challenge.remote_enr) {
-            (Some(new_enr), Some(known_enr)) => {
-                if new_enr.seq() > known_enr.seq() {
-                    new_enr
-                } else {
-                    known_enr
+
+        // Duplicate code here to avoid cloning an ENR
+        let remote_public_key = {
+            let enr = match (enr_record.as_ref(), challenge.remote_enr.as_ref()) {
+                (Some(new_enr), Some(known_enr)) => {
+                    if new_enr.seq() > known_enr.seq() {
+                        new_enr
+                    } else {
+                        known_enr
+                    }
                 }
-            }
-            (Some(new_enr), None) => new_enr,
-            (None, Some(known_enr)) => known_enr,
-            (None, None) => {
-                warn!(
+                (Some(new_enr), None) => new_enr,
+                (None, Some(known_enr)) => known_enr,
+                (None, None) => {
+                    warn!(
                 "Peer did not respond with their ENR. Session could not be established. Node: {}",
                 remote_id
             );
-                return Err(Discv5Error::SessionNotEstablished);
-            }
+                    return Err(Discv5Error::SessionNotEstablished);
+                }
+            };
+            enr.public_key()
         };
 
-        // ENR must exist here
-        let remote_public_key = session_enr.public_key();
-
         // verify the auth header nonce
         if !crypto::verify_authentication_nonce(
             &remote_public_key,
@@ -168,14 +161,41 @@ impl Session {
             &local_id,
             id_nonce_sig,
         ) {
-            return Err(Discv5Error::InvalidSignature);
+            return Err(Discv5Error::InvalidChallengeSignature(challenge));
         }
 
+        // The keys are derived after the message has been verified to prevent potential extra work
+        // for invalid messages.
+
+        // generate session keys
+        let (decryption_key, encryption_key) = crypto::derive_keys_from_pubkey(
+            &local_key.read(),
+            local_id,
+            remote_id,
+            &challenge.data,
+            ephem_pubkey,
+        )?;
+
         let keys = Keys {
             encryption_key,
             decryption_key,
         };
 
+        // Takes ownership of the provided ENRs - Slightly annoying code duplication, but avoids
+        // cloning ENRs
+        let session_enr = match (enr_record, challenge.remote_enr) {
+            (Some(new_enr), Some(known_enr)) => {
+                if new_enr.seq() > known_enr.seq() {
+                    new_enr
+                } else {
+                    known_enr
+                }
+            }
+            (Some(new_enr), None) => new_enr,
+            (None, Some(known_enr)) => known_enr,
+            (None, None) => unreachable!("Checked in the first match above"),
+        };
+
         Ok((Session::new(keys), session_enr))
     }
 

src/packet/mod.rs

@@ -47,6 +47,12 @@ pub type IdNonce = [u8; ID_NONCE_LENGTH];
 // This is the WHOAREYOU authenticated data.
 pub struct ChallengeData([u8; 63]);
 
+impl std::fmt::Debug for ChallengeData {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", hex::encode(&self.0))
+    }
+}
+
 impl std::convert::TryFrom<Vec<u8>> for ChallengeData {
     type Error = ();