peer: update endpoint on handshake response messages too

2 files changed, 4 insertions, 3 deletions

diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index c807bd8..509d8c2 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -152,7 +152,7 @@ impl PeerServer {
     }
 
     // TODO use the address to update endpoint if it changes i suppose
-    fn handle_ingress_handshake_resp(&mut self, _addr: SocketAddr, packet: &Response) -> Result<(), Error> {
+    fn handle_ingress_handshake_resp(&mut self, addr: SocketAddr, packet: &Response) -> Result<(), Error> {
         ensure!(packet.len() == 92, "handshake resp packet length is incorrect");
         let mut state = self.shared_state.borrow_mut();
         {
@@ -166,7 +166,7 @@ impl PeerServer {
             .ok_or_else(|| format_err!("unknown our_index ({})", our_index))?
             .clone();
         let mut peer = peer_ref.borrow_mut();
-        let dead_index = peer.process_incoming_handshake_response(packet)?;
+        let dead_index = peer.process_incoming_handshake_response(addr, packet)?;
         if let Some(index) = dead_index {
             let _ = state.index_map.remove(&index);
         }
diff --git a/src/peer.rs b/src/peer.rs
index 2331a09..9aaf422 100644
--- a/src/peer.rs
+++ b/src/peer.rs
@@ -312,13 +312,14 @@ impl Peer {
         self.cookie.consume_reply(reply)
     }
 
-    pub fn process_incoming_handshake_response(&mut self, packet: &Response) -> Result<Option<u32>, Error> {
+    pub fn process_incoming_handshake_response(&mut self, addr: SocketAddr, packet: &Response) -> Result<Option<u32>, Error> {
         let mut session     = mem::replace(&mut self.sessions.next, None).ok_or_else(|| err_msg("no next session"))?;
         let     _           = session.noise.read_message(packet.noise_bytes(), &mut [])?;
         session             = session.into_transport_mode()?;
         session.their_index = packet.their_index();
         session.birthday    = Timestamp::now();
         self.last_handshake = Timestamp::now();
+        self.info.endpoint  = Some(addr);
 
         let current = mem::replace(&mut self.sessions.current, Some(session));
         let dead    = mem::replace(&mut self.sessions.past,    current);