From ac4a27cdee61048273be49fc9c9500d4009c6192 Mon Sep 17 00:00:00 2001
From: bochaco <gabrielviganotti@gmail.com>
Date: Sun, 23 May 2021 18:35:09 -0300
Subject: [PATCH] fix(end-user): assign clients a xorname which always matches
 the section prefix so they are propoerly routed in a multi-section network

---
 src/routing/core/public_api.rs  |  3 ++-
 src/routing/enduser_registry.rs | 14 +++++++++-----
 src/routing/mod.rs              | 13 ++++++++-----
 tests/messages.rs               |  7 ++++++-
 4 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/src/routing/core/public_api.rs b/src/routing/core/public_api.rs
index 9b0ff92839..8ca1937f7b 100644
--- a/src/routing/core/public_api.rs
+++ b/src/routing/core/public_api.rs
@@ -43,7 +43,8 @@ impl Core {
     }
 
     pub fn try_add(&mut self, sender: SocketAddr) -> Result<EndUser> {
-        self.end_users.try_add(sender)
+        let section_prefix = self.section.prefix();
+        self.end_users.try_add(sender, section_prefix)
     }
 
     pub fn node(&self) -> &Node {
diff --git a/src/routing/enduser_registry.rs b/src/routing/enduser_registry.rs
index d5671ffc1e..b21aae2495 100644
--- a/src/routing/enduser_registry.rs
+++ b/src/routing/enduser_registry.rs
@@ -12,7 +12,7 @@ use std::{
     collections::{btree_map::Entry, BTreeMap},
     net::SocketAddr,
 };
-use xor_name::XorName;
+use xor_name::{Prefix, XorName};
 
 pub type SocketId = XorName;
 pub(crate) struct EndUserRegistry {
@@ -36,12 +36,16 @@ impl EndUserRegistry {
         self.socket_id_mapping.get(&socket_id)
     }
 
-    pub fn try_add(&mut self, sender: SocketAddr) -> Result<EndUser> {
+    pub fn try_add(&mut self, sender: SocketAddr, section_prefix: &Prefix) -> Result<EndUser> {
         // create a unique socket id from client socket addr
-        let user_xorname = XorName::from_content(&[
+        let socket_id = XorName::from_content(&[
             &bincode::serialize(&sender).map_err(|_| Error::FailedSignature)?
         ]);
 
+        // assign a XorName to the end user which belongs to this section's prefix
+        // so messages directed to this end user are correctly routed back through us
+        let user_xorname = section_prefix.substituted_in(socket_id);
+
         // TODO: we probably should remove the socket_id from the EndUser struct,
         // and pass the socket id separatelly as part of nodes' messages,
         // instead of it being part of the SrcLocation/DstLocation in nodes' messages.
@@ -49,10 +53,10 @@ impl EndUserRegistry {
         // aggregation of messages when sent to another section with aggregation AtDestination.
         let end_user = EndUser {
             xorname: user_xorname,
-            socket_id: user_xorname,
+            socket_id,
         };
 
-        match self.socket_id_mapping.entry(user_xorname) {
+        match self.socket_id_mapping.entry(socket_id) {
             Entry::Vacant(entry) => {
                 let _ = self.clients.insert(sender, end_user);
                 let _ = entry.insert(sender);
diff --git a/src/routing/mod.rs b/src/routing/mod.rs
index a5cba50838..89a5793d02 100644
--- a/src/routing/mod.rs
+++ b/src/routing/mod.rs
@@ -516,13 +516,16 @@ async fn handle_message(dispatcher: Arc<Dispatcher>, bytes: Bytes, sender: Socke
                 .copied();
 
             let end_user = match end_user {
-                Some(end_user) => end_user,
+                Some(end_user) => {
+                    debug!(
+                        "Message from client {}, socket id already exists: {:?}",
+                        sender, end_user
+                    );
+                    end_user
+                }
                 None => {
                     // this is the first time we receive a message from this client
-                    trace!(
-                        "First message from client {}, creating a socket id and caching it",
-                        sender
-                    );
+                    debug!("First message from client {}, creating a socket id", sender);
 
                     // TODO: remove the enduser registry and simply encrypt socket addr with
                     // this node's keypair and use that as the socket id
diff --git a/tests/messages.rs b/tests/messages.rs
index c84b3e31e3..771445cd01 100644
--- a/tests/messages.rs
+++ b/tests/messages.rs
@@ -49,6 +49,7 @@ async fn test_messages_client_node() -> Result<()> {
     config.local_ip = Some(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
 
     let node_addr = node.our_connection_info();
+    let section_prefix = node.our_prefix().await;
     let section_key = *node.section_chain().await.last_key();
 
     let client = QuicP2p::with_config(Some(config), &[node_addr], false)?;
@@ -97,8 +98,12 @@ async fn test_messages_client_node() -> Result<()> {
     node_handler.await??;
 
     if let Some((_, resp)) = incoming_messages.next().await {
-        let user_xorname =
+        // the xorname assigned to each end user is computed from
+        // the client socket addr plus the client section prefix
+        let socket_id =
             XorName::from_content(&[&bincode::serialize(&client_endpoint.socket_addr())?]);
+        let user_xorname = section_prefix.substituted_in(socket_id);
+
         let expected_bytes = query.serialize(user_xorname, section_key)?;
 
         assert_eq!(resp, expected_bytes);