bubble failures in handshake processing for peers

3 files changed, 18 insertions, 14 deletions

diff --git a/src/interface/mod.rs b/src/interface/mod.rs
index 5513ab2..06dff46 100644
--- a/src/interface/mod.rs
+++ b/src/interface/mod.rs
@@ -202,7 +202,7 @@ impl Interface {
                         let mut peer = Peer::new(info.clone());
                         peer.set_next_session(noise.into());
 
-                        let init_packet = peer.get_handshake_packet();
+                        let init_packet = peer.get_handshake_packet().unwrap();
                         let our_index = peer.our_next_index().unwrap();
                         let peer = Rc::new(RefCell::new(peer));
 
diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index dce84f6..415c78f 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -149,7 +149,7 @@ impl PeerServer {
                 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();
+                let response_packet = peer.get_response_packet().unwrap();
 
                 self.handle.spawn(self.udp_tx.clone().send((addr.clone(), response_packet)).then(|_| Ok(())));
                 let dead_session = peer.ratchet_session()?;
@@ -250,7 +250,7 @@ impl PeerServer {
 
                 let _ = state.index_map.insert(peer.our_next_index().unwrap(), peer_ref.clone());
 
-                let init_packet = peer.get_handshake_packet();
+                let init_packet = peer.get_handshake_packet().unwrap();
                 let endpoint = peer.info.endpoint.unwrap().clone();
 
                 self.handle.spawn(self.udp_tx.clone().send((endpoint, init_packet)).then(|_| Ok(())));
diff --git a/src/protocol/peer.rs b/src/protocol/peer.rs
index a2eefba..fc837e5 100644
--- a/src/protocol/peer.rs
+++ b/src/protocol/peer.rs
@@ -129,7 +129,8 @@ impl Peer {
         }
 
         let mut raw_packet = vec![0u8; 1500];
-        session.noise.set_receiving_nonce(nonce).unwrap();
+        session.noise.set_receiving_nonce(nonce)
+            .map_err(SyncFailure::new)?;
         let len = session.noise.read_message(packet, &mut raw_packet)
             .map_err(SyncFailure::new)?;
         raw_packet.truncate(len);
@@ -176,15 +177,18 @@ impl Peer {
         }
     }
 
-    pub fn get_handshake_packet(&mut self) -> Vec<u8> {
+    pub fn get_handshake_packet(&mut self) -> Result<Vec<u8>, Error> {
         let now = time::get_time();
         let mut tai64n = [0; 12];
         BigEndian::write_i64(&mut tai64n[0..], 4611686018427387914 + now.sec);
         BigEndian::write_i32(&mut tai64n[8..], now.nsec);
         let mut initiation_packet = vec![0; 148];
         initiation_packet[0] = 1; /* Type: Initiation */
-        LittleEndian::write_u32(&mut initiation_packet[4..], self.our_next_index().unwrap());
-        self.sessions.next.as_mut().unwrap().noise.write_message(&tai64n, &mut initiation_packet[8..]).unwrap();
+
+        let next = self.sessions.next.as_mut().ok_or_else(|| format_err!("missing next session"))?;
+        LittleEndian::write_u32(&mut initiation_packet[4..], next.our_index);
+        next.noise.write_message(&tai64n, &mut initiation_packet[8..]).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);
@@ -192,7 +196,7 @@ impl Peer {
         let mac = blake2s(16, mac_key.as_bytes(), &initiation_packet[0..116]);
         memcpy(&mut initiation_packet[116..], mac.as_bytes());
 
-        initiation_packet
+        Ok(initiation_packet)
     }
 
     /// Takes a new handshake packet (type 0x01), updates the internal peer state,
@@ -203,13 +207,13 @@ impl Peer {
         unimplemented!()
     }
 
-    pub fn get_response_packet(&mut self) -> Vec<u8> {
+    pub fn get_response_packet(&mut self) -> Result<Vec<u8>, Error> {
         let mut packet = vec![0; 76];
         packet[0] = 2; /* Type: Response */
-        let session = self.sessions.next.as_mut().unwrap();
-        LittleEndian::write_u32(&mut packet[4..], session.our_index);
-        LittleEndian::write_u32(&mut packet[8..], session.their_index);
-        session.noise.write_message(&[], &mut packet[12..]).unwrap();
+        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)?;
         let mut mac_key_input = [0; 40];
         memcpy(&mut mac_key_input, b"mac1----");
         memcpy(&mut mac_key_input[8..], &self.info.pub_key);
@@ -217,7 +221,7 @@ impl Peer {
         let mac = blake2s(16, mac_key.as_bytes(), &packet[0..44]);
         memcpy(&mut packet[44..], mac.as_bytes());
 
-        packet
+        Ok(packet)
     }
 
     pub fn to_config_string(&self) -> String {