From cc23bd74410003a83a94fa2287ce502b4a82a5ac Mon Sep 17 00:00:00 2001
From: Aaron Feickert <66188213+AaronFeickert@users.noreply.github.com>
Date: Thu, 19 Jan 2023 15:25:43 -0600
Subject: [PATCH] Do not propagate unsigned encrypted messages

---
 comms/dht/src/inbound/decryption.rs | 52 +++++++++++++++++++----------
 1 file changed, 35 insertions(+), 17 deletions(-)

diff --git a/comms/dht/src/inbound/decryption.rs b/comms/dht/src/inbound/decryption.rs
index 439c8698fae..9f87be9bb95 100644
--- a/comms/dht/src/inbound/decryption.rs
+++ b/comms/dht/src/inbound/decryption.rs
@@ -224,8 +224,10 @@ where S: Service<DecryptedDhtMessage, Response = (), Error = PipelineError>
         node_identity: Arc<NodeIdentity>,
         message: DhtInboundMessage,
     ) -> Result<DecryptedDhtMessage, DecryptionError> {
+        // Perform initial checks on message validity
         let validated_msg = Self::initial_validation(message)?;
 
+        // The message is unencrypted and valid
         if !validated_msg.header().flags.is_encrypted() {
             return Self::success_not_encrypted(validated_msg).await;
         }
@@ -237,22 +239,14 @@ where S: Service<DecryptedDhtMessage, Response = (), Error = PipelineError>
             validated_msg.message().dht_header.message_tag
         );
 
-        let dht_header = validated_msg.header();
-
-        let e_pk = dht_header
-            .ephemeral_public_key
-            .as_ref()
-            // No ephemeral key with ENCRYPTED flag set
-            .ok_or( DecryptionError::EphemeralKeyNotProvidedForEncryptedMessage)?;
-
-        if !validated_msg.message().dht_header.destination.is_unknown() &&
-            validated_msg
-                .message()
-                .dht_header
-                .destination
-                .public_key()
-                .map(|pk| pk != node_identity.public_key())
-                .unwrap_or(false)
+        // The message is encrypted, so see if it is for us
+        if validated_msg
+            .message()
+            .dht_header
+            .destination
+            .public_key()
+            .map(|pk| pk != node_identity.public_key())
+            .unwrap_or(false)
         {
             debug!(
                 target: LOG_TARGET,
@@ -264,6 +258,13 @@ where S: Service<DecryptedDhtMessage, Response = (), Error = PipelineError>
             return Ok(DecryptedDhtMessage::failed(validated_msg.into_message()));
         }
 
+        // The message is encrypted and for us, so derive its encryption key
+        let dht_header = validated_msg.header();
+        let e_pk = dht_header
+            .ephemeral_public_key
+            .as_ref()
+            // This has already been checked, but we need it to avoid an unwrap
+            .ok_or( DecryptionError::EphemeralKeyNotProvidedForEncryptedMessage)?;
         let shared_secret = CommsDHKE::new(node_identity.secret_key(), e_pk);
         let message = validated_msg.message();
 
@@ -310,6 +311,7 @@ where S: Service<DecryptedDhtMessage, Response = (), Error = PipelineError>
             message.tag,
             message.dht_header.message_tag
         );
+        // Decrypt and verify the message
         match Self::attempt_decrypt_message_body(&shared_secret, &message.body) {
             Ok(message_body) => {
                 debug!(
@@ -346,13 +348,21 @@ where S: Service<DecryptedDhtMessage, Response = (), Error = PipelineError>
     }
 
     /// Performs message validation that should be performed by all nodes. If an error is encountered, the message is
-    /// invalid and should never have been sent.
+    /// invalid and should never have been propagated.
+    /// 
+    /// These failure modes are detectable by any node, so it is generally safe to ban an offending peer.
     fn initial_validation(message: DhtInboundMessage) -> Result<ValidatedDhtInboundMessage, DecryptionError> {
+        // Messages must not be empty
         if message.body.is_empty() {
             return Err(DecryptionError::EncryptedMessageEmptyBody);
         }
 
         if message.dht_header.flags.is_encrypted() {
+            // An encrypted message needs:
+            // - a destination
+            // - an ephemeral public key used for DHKE
+            // - an encrypted message signature
+
             // Check if there is no destination specified and discard
             if message.dht_header.destination.is_unknown() {
                 return Err(DecryptionError::EncryptedMessageNoDestination);
@@ -363,10 +373,17 @@ where S: Service<DecryptedDhtMessage, Response = (), Error = PipelineError>
                 return Err(DecryptionError::EphemeralKeyNotProvidedForEncryptedMessage);
             }
 
+            // An encrypted message signature is required
+            if message.dht_header.message_signature.is_empty() {
+                return Err(DecryptionError::MessageSignatureNotProvidedForEncryptedMessage);
+            }
+
             Ok(ValidatedDhtInboundMessage::new(message, None))
         } else if message.dht_header.message_signature.is_empty() {
+            // An unencrypted message does not require a message signature
             Ok(ValidatedDhtInboundMessage::new(message, None))
         } else {
+            // But if it has one, it must be valid!
             let message_signature: MessageSignature =
                 ProtoMessageSignature::decode(message.dht_header.message_signature.as_slice())
                     .map_err(|_| DecryptionError::MessageSignatureClearTextDecodeFailed)?
@@ -393,6 +410,7 @@ where S: Service<DecryptedDhtMessage, Response = (), Error = PipelineError>
         let encrypted_message_signature = Some(&dht_header.message_signature)
             .filter(|b| !b.is_empty())
             // This should not have been sent/propagated
+            // This is already checked elsewhere, but we need it to avoid an unwrap
             .ok_or( DecryptionError::MessageSignatureNotProvidedForEncryptedMessage)?;
 
         // obtain key signature for authenticated decrypt signature