From a1c26ff62d2dfd0bab12e35b266ce46eee024b77 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adam=20Cig=C3=A1nek?= <adam.ciganek@gmail.com>
Date: Mon, 14 Dec 2020 14:36:51 +0100
Subject: [PATCH] fix: handle message send to self

---
 src/routing/approved.rs  | 13 +++++++++-
 src/routing/tests/mod.rs | 56 ++++++++++++++++++++++++++++++++++++----
 2 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/src/routing/approved.rs b/src/routing/approved.rs
index 5e82d99a05..65b393e12a 100644
--- a/src/routing/approved.rs
+++ b/src/routing/approved.rs
@@ -1793,7 +1793,18 @@ impl Approved {
             SrcLocation::Node(_) => {
                 // If the source is a single node, we don't even need to vote, so let's cut this short.
                 let msg = Message::single_src(&self.node, dst, variant, None, None)?;
-                Ok(self.relay_message(&msg)?.into_iter().collect())
+                let mut commands = vec![];
+
+                if dst.contains(&self.node.name(), self.section.prefix()) {
+                    commands.push(Command::HandleMessage {
+                        sender: Some(self.node.addr),
+                        message: msg.clone(),
+                    });
+                }
+
+                commands.extend(self.relay_message(&msg)?);
+
+                Ok(commands)
             }
             SrcLocation::Section(_) => {
                 let vote = self.create_send_message_vote(dst, variant, None)?;
diff --git a/src/routing/tests/mod.rs b/src/routing/tests/mod.rs
index 1cbb83ee0a..2b7e6889f4 100644
--- a/src/routing/tests/mod.rs
+++ b/src/routing/tests/mod.rs
@@ -14,7 +14,7 @@ use crate::{
     consensus::{test_utils::*, Proven, Vote},
     crypto,
     event::Event,
-    location::DstLocation,
+    location::{DstLocation, SrcLocation},
     majority,
     messages::{
         BootstrapResponse, JoinRequest, Message, PlainMessage, ResourceProofResponse, Variant,
@@ -22,10 +22,7 @@ use crate::{
     network::Network,
     node::Node,
     peer::Peer,
-    relocation,
-    relocation::RelocateDetails,
-    relocation::RelocatePayload,
-    relocation::SignedRelocateDetails,
+    relocation::{self, RelocateDetails, RelocatePayload, SignedRelocateDetails},
     section::{
         test_utils::*, EldersInfo, MemberInfo, PeerState, Section, SectionKeyShare,
         SectionProofChain, MIN_AGE,
@@ -1388,6 +1385,55 @@ async fn relocation(relocated_peer_role: RelocatedPeerRole) -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn node_message_to_self() -> Result<()> {
+    message_to_self(MessageDst::Node).await
+}
+
+#[tokio::test]
+async fn section_message_to_self() -> Result<()> {
+    message_to_self(MessageDst::Section).await
+}
+
+enum MessageDst {
+    Node,
+    Section,
+}
+
+async fn message_to_self(dst: MessageDst) -> Result<()> {
+    let node = create_node();
+    let peer = node.peer();
+    let state = Approved::first_node(node, mpsc::unbounded_channel().0)?;
+    let stage = Stage::new(state, create_comm()?);
+
+    let src = SrcLocation::Node(*peer.name());
+    let dst = match dst {
+        MessageDst::Node => DstLocation::Node(*peer.name()),
+        MessageDst::Section => DstLocation::Section(rand::random()),
+    };
+    let content = Bytes::from_static(b"hello");
+
+    let commands = stage
+        .handle_command(Command::SendUserMessage {
+            src,
+            dst,
+            content: content.clone(),
+        })
+        .await?;
+
+    assert_matches!(&commands[..], [Command::HandleMessage { sender, message }] => {
+        assert_eq!(sender.as_ref(), Some(peer.addr()));
+        assert_eq!(message.src().src_location(), src);
+        assert_eq!(message.dst(), &dst);
+        assert_matches!(
+            message.variant(),
+            Variant::UserMessage(actual_content) if actual_content == &content
+        );
+    });
+
+    Ok(())
+}
+
 // TODO: add more tests here
 
 fn create_peer() -> Peer {