clear out dead indices from ratcheted-away sessions

2 files changed, 7 insertions, 4 deletions

diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index 0ca7f3d..d9883d9 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -160,7 +160,10 @@ impl PeerServer {
                 let response_packet = peer.get_response_packet();
 
                 self.handle.spawn(self.udp_tx.clone().send((addr.clone(), response_packet)).then(|_| Ok(())));
-                peer.ratchet_session().unwrap();
+                let dead_session = peer.ratchet_session().unwrap();
+                if let Some(session) = dead_session {
+                    let _ = state.index_map.remove(&session.our_index);
+                }
                 info!("sent handshake response, ratcheted session.");
             },
             2 => {
diff --git a/src/protocol/peer.rs b/src/protocol/peer.rs
index a44dffd..789ea8c 100644
--- a/src/protocol/peer.rs
+++ b/src/protocol/peer.rs
@@ -104,15 +104,15 @@ impl Peer {
         let _ = std::mem::replace(&mut self.sessions.next, Some(session));
     }
 
-    pub fn ratchet_session(&mut self) -> Result<(), ()> {
+    pub fn ratchet_session(&mut self) -> Result<Option<Session>, ()> {
         let next = std::mem::replace(&mut self.sessions.next, None).ok_or(())?;
         let next = next.into_transport_mode();
 
         let current = std::mem::replace(&mut self.sessions.current, Some(next));
-        let _       = std::mem::replace(&mut self.sessions.past,    current);
+        let dead    = std::mem::replace(&mut self.sessions.past,    current);
 
         self.last_handshake = Some(SystemTime::now());
-        Ok(())
+        Ok(dead)
     }
 
     pub fn decrypt_transport_packet(&mut self, our_index: u32, nonce: u64, packet: &[u8]) -> Result<Vec<u8>, ()> {