diff options
-rw-r--r-- | src/interface/mod.rs | 2 | ||||
-rw-r--r-- | src/interface/peer_server.rs | 4 | ||||
-rw-r--r-- | src/protocol/peer.rs | 26 |
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 { |