let Peer process its incoming handshake

2 files changed, 37 insertions, 31 deletions

diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index b0cec2f..5fe6bc3 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -126,41 +126,28 @@ impl PeerServer {
                 let their_index = LittleEndian::read_u32(&packet[4..]);
 
                 let mut noise = Noise::build_responder(
-                    &state.interface_info.private_key.expect("no privatekey!"))?;
+                    &state.interface_info.private_key.ok_or_else(|| format_err!("no private key!"))?)?;
 
-                let mut timestamp = [0u8; 116];
+                let mut timestamp = [0u8; 12];
                 let _ = noise.read_message(&packet[8..116], &mut timestamp)
                     .map_err(SyncFailure::new)?;
 
-                let peer_ref = {
+                let mut peer_ref = {
                     let their_pubkey = noise.get_remote_static().expect("must have remote static key");
 
                     debug!("their_pubkey: {}", base64::encode(&their_pubkey[..]));
-                    let peer_ref = state.pubkey_map.get(&their_pubkey[..])
-                        .ok_or_else(|| format_err!("unknown peer pubkey"))?;
-                    peer_ref.clone()
+                    state.pubkey_map.get(&their_pubkey[..])
+                        .ok_or_else(|| format_err!("unknown peer pubkey"))?.clone()
                 };
-
                 let mut peer = peer_ref.borrow_mut();
 
-                // TODO: hacked up API until it's officially supported in snow.
-                match noise {
-                    snow::Session::Handshake(ref mut handshake_state) => {
-                        handshake_state.set_psk(2, &peer.info.psk.unwrap_or_else(|| [0u8; 32]));
-                    },
-                    _ => unreachable!()
+                let (response, next_index, dead_index) = peer.process_incoming_handshake(addr, their_index, &timestamp, noise)?;
+                let _ = state.index_map.insert(next_index, peer_ref.clone());
+                if let Some(index) = dead_index {
+                    let _ = state.index_map.remove(&index);
                 }
 
-                peer.set_next_session(Session::with_their_index(noise, their_index));
-                let _ = state.index_map.insert(peer.our_next_index().unwrap(), peer_ref.clone());
-
-                let response_packet = peer.get_response_packet()?;
-
-                self.send_to_peer((addr.clone(), response_packet));
-                let dead_session = peer.ratchet_session()?;
-                if let Some(session) = dead_session {
-                    let _ = state.index_map.remove(&session.our_index);
-                }
+                self.send_to_peer((addr, response));
                 info!("sent handshake response, ratcheted session.");
             },
             2 => {
diff --git a/src/protocol/peer.rs b/src/protocol/peer.rs
index bea3a52..6c2fe9f 100644
--- a/src/protocol/peer.rs
+++ b/src/protocol/peer.rs
@@ -229,18 +229,37 @@ impl Peer {
     /// Takes a new handshake packet (type 0x01), updates the internal peer state,
     /// and generates a response.
     ///
-    /// Returns: the response packet (type 0x02).
-    pub fn process_incoming_handshake(&mut self) -> Vec<u8> {
-        unimplemented!()
+    /// Returns: the response packet (type 0x02), and an optional dead session index that was removed.
+    pub fn process_incoming_handshake(&mut self, addr: SocketAddr, their_index: u32, timestamp: &[u8], mut noise: snow::Session)
+            -> Result<(Vec<u8>, u32, Option<u32>), Error> {
+
+        // TODO: verify timestamp
+        // TODO: hacked up API until it's officially supported in snow.
+        match noise {
+            snow::Session::Handshake(ref mut handshake_state) => {
+                handshake_state.set_psk(2, &self.info.psk.unwrap_or_else(|| [0u8; 32]));
+            },
+            _ => unreachable!()
+        }
+
+        let mut next_session = Session::with_their_index(noise, their_index);
+        let next_index = next_session.our_index;
+        let response_packet = self.get_response_packet(&mut next_session)?;
+        self.set_next_session(next_session);
+
+        let dead_index = self.ratchet_session()?.map(|session| session.our_index);
+
+        self.info.endpoint = Some(addr); // update peer endpoint after successful authentication
+
+        Ok((response_packet, next_index, dead_index))
     }
 
-    pub fn get_response_packet(&mut self) -> Result<Vec<u8>, Error> {
+    fn get_response_packet(&mut self, next_session: &mut Session) -> Result<Vec<u8>, Error> {
         let mut packet = vec![0; 76];
         packet[0] = 2; /* Type: Response */
-        let next = self.sessions.next.as_mut().ok_or_else(|| format_err!("missing next session"))?;
-        LittleEndian::write_u32(&mut packet[4..], next.our_index);
-        LittleEndian::write_u32(&mut packet[8..], next.their_index);
-        next.noise.write_message(&[], &mut packet[12..]).map_err(SyncFailure::new)?;
+        LittleEndian::write_u32(&mut packet[4..], next_session.our_index);
+        LittleEndian::write_u32(&mut packet[8..], next_session.their_index);
+        next_session.noise.write_message(&[], &mut packet[12..]).map_err(SyncFailure::new)?;
         let mut mac_key_input = [0; 40];
         memcpy(&mut mac_key_input, b"mac1----");
         memcpy(&mut mac_key_input[8..], &self.info.pub_key);