update peer endpoint after receiving authenticated transport packet

2 files changed, 5 insertions, 2 deletions

diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index 1aa4daa..b0cec2f 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -219,7 +219,7 @@ impl PeerServer {
                 if let Some(ref peer) = lookup {
                     let raw_packet = {
                         let mut peer = peer.borrow_mut();
-                        peer.handle_incoming_transport(our_index_received, nonce, &packet[16..])?
+                        peer.handle_incoming_transport(our_index_received, nonce, addr, &packet[16..])?
                     };
 
                     if raw_packet.len() == 0 {
diff --git a/src/protocol/peer.rs b/src/protocol/peer.rs
index 1a87b66..bea3a52 100644
--- a/src/protocol/peer.rs
+++ b/src/protocol/peer.rs
@@ -128,7 +128,7 @@ impl Peer {
         Ok(dead)
     }
 
-    pub fn handle_incoming_transport(&mut self, our_index: u32, nonce: u64, packet: &[u8]) -> Result<Vec<u8>, Error> {
+    pub fn handle_incoming_transport(&mut self, our_index: u32, nonce: u64, addr: SocketAddr, packet: &[u8]) -> Result<Vec<u8>, Error> {
         self.rx_bytes += packet.len() as u64;
 
         let session = self.sessions.current.as_mut().filter(|session| session.our_index == our_index)
@@ -142,6 +142,9 @@ impl Peer {
             .map_err(SyncFailure::new)?;
         let len = session.noise.read_message(packet, &mut raw_packet)
             .map_err(SyncFailure::new)?;
+
+        self.info.endpoint = Some(addr); // update peer endpoint after successful authentication
+
         raw_packet.truncate(len);
         Ok(raw_packet)
     }