Use less hard coded lengths and slice indexes

3 files changed, 38 insertions, 22 deletions

diff --git a/src/protocol/controller.rs b/src/protocol/controller.rs
index bb9f3e9..b0d284b 100644
--- a/src/protocol/controller.rs
+++ b/src/protocol/controller.rs
@@ -173,7 +173,7 @@ struct Transport {
 }
 
 fn udp_process_handshake_init(wg: Arc<WgState>, sock: &UdpSocket, p: &[u8], addr: SocketAddr) {
-    if p.len() != 148 {
+    if p.len() != HANDSHAKE_INIT_LEN {
         return;
     }
 
@@ -185,8 +185,8 @@ fn udp_process_handshake_init(wg: Arc<WgState>, sock: &UdpSocket, p: &[u8], addr
         if !cookie_verify(p, &cookie) {
             debug!("Mac2 verify failed, send cookie reply.");
             let peer_id = Id::from_slice(&p[4..8]);
-            let mac1 = &p[116..132];
-            let reply = cookie_reply(info.psk.as_ref(), &info.pubkey, &cookie, peer_id, mac1);
+            let mac1 = get_mac1(p);
+            let reply = cookie_reply(info.psk.as_ref(), &info.pubkey, &cookie, peer_id, &mac1);
             sock.send_to(&reply, addr).unwrap();
             return;
         } else {
@@ -214,9 +214,7 @@ fn udp_process_handshake_init(wg: Arc<WgState>, sock: &UdpSocket, p: &[u8], addr
             let mut response = responde(info.deref(), &mut r, self_id);
 
             // Save mac1.
-            let mut mac1 = [0u8; 16];
-            mac1.copy_from_slice(&response[60..76]);
-            peer.last_mac1 = Some(mac1);
+            peer.last_mac1 = Some(get_mac1(&response));
 
             cookie_sign(&mut response, peer.get_cookie());
 
@@ -240,7 +238,7 @@ fn udp_process_handshake_init(wg: Arc<WgState>, sock: &UdpSocket, p: &[u8], addr
 }
 
 fn udp_process_handshake_resp(wg: &WgState, sock: &UdpSocket, p: &[u8], addr: SocketAddr) {
-    if p.len() != 92 {
+    if p.len() != HANDSHAKE_RESP_LEN {
         return;
     }
 
@@ -252,8 +250,8 @@ fn udp_process_handshake_resp(wg: &WgState, sock: &UdpSocket, p: &[u8], addr: So
         if !cookie_verify(p, &cookie) {
             debug!("Mac2 verify failed, send cookie reply.");
             let peer_id = Id::from_slice(&p[4..8]);
-            let mac1 = &p[60..76];
-            let reply = cookie_reply(info.psk.as_ref(), &info.pubkey, &cookie, peer_id, mac1);
+            let mac1 = get_mac1(p);
+            let reply = cookie_reply(info.psk.as_ref(), &info.pubkey, &cookie, peer_id, &mac1);
             sock.send_to(&reply, addr).unwrap();
             return;
         } else {
@@ -567,9 +565,7 @@ fn do_handshake(wg: Arc<WgState>, peer0: SharedPeerState, sock: Arc<UdpSocket>)
     sock.send_to(&i, endpoint).unwrap();
     peer.count_send((&i).len());
 
-    let mut mac1 = [0u8; 16];
-    mac1.copy_from_slice(&i[116..132]);
-    peer.last_mac1 = Some(mac1);
+    peer.last_mac1 = Some(get_mac1(&i));
 
     let resend = {
         let wg = wg.clone();
diff --git a/src/protocol/cookie.rs b/src/protocol/cookie.rs
index 26c031a..dd23263 100644
--- a/src/protocol/cookie.rs
+++ b/src/protocol/cookie.rs
@@ -91,6 +91,18 @@ pub fn process_cookie_reply(psk: Option<&[u8; 32]>,
     Ok(cookie)
 }
 
+/// Extract `mac1` from a message.
+///
+/// # Panics
+///
+/// If the message is not at least 32-byte long.
+pub fn get_mac1(m: &[u8]) -> [u8; 16] {
+    let mut out = [0u8; 16];
+    let len = m.len();
+    out.copy_from_slice(&m[len - 32..len-16]);
+    out
+}
+
 pub fn cookie_sign(m: &mut [u8], cookie: Option<&Cookie>) {
     if cookie.is_none() {
         return;
diff --git a/src/protocol/handshake.rs b/src/protocol/handshake.rs
index f30ece4..fd5b1ff 100644
--- a/src/protocol/handshake.rs
+++ b/src/protocol/handshake.rs
@@ -31,6 +31,9 @@ use protocol::*;
 
 const PROLOGUE: &'static [u8] = b"WireGuard v0 zx2c4 Jason@zx2c4.com";
 
+pub const HANDSHAKE_INIT_LEN: usize = 148;
+pub const HANDSHAKE_RESP_LEN: usize = 92;
+
 pub type HS = HandshakeState<X25519, ChaCha20Poly1305, NoiseBlake2s>;
 
 #[derive(Clone)]
@@ -82,8 +85,8 @@ fn mac<K>(key: Option<K>, data: &[&[u8]]) -> [u8; 16]
 /// Will generate a new ephemeral key and use current timestamp.
 ///
 /// Returns: Message, noise handshake state.
-pub fn initiate(wg: &WgInfo, peer: &PeerInfo, self_index: Id) -> ([u8; 148], HS) {
-    let mut msg = [0u8; 148];
+pub fn initiate(wg: &WgInfo, peer: &PeerInfo, self_index: Id) -> ([u8; HANDSHAKE_INIT_LEN], HS) {
+    let mut msg = [0u8; HANDSHAKE_INIT_LEN];
 
     let mut hs = {
         let mut hsbuilder = HandshakeStateBuilder::<X25519>::new();
@@ -124,10 +127,12 @@ pub struct InitProcessResult {
 /// Process a handshake initiation message.
 ///
 /// Will generate a new ephemeral key.
+///
+/// # Panics
+///
+/// If the message length is not `HANDSHAKE_INIT_LEN`.
 pub fn process_initiation(wg: &WgInfo, msg: &[u8]) -> Result<InitProcessResult, ()> {
-    if msg.len() != 148 {
-        return Err(());
-    }
+    debug_assert_eq!(msg.len(), HANDSHAKE_INIT_LEN);
 
     // Check mac1.
     let mac1 = mac(wg.psk.as_ref(), &[&wg.pubkey, &msg[..116]]);
@@ -168,8 +173,9 @@ pub fn process_initiation(wg: &WgInfo, msg: &[u8]) -> Result<InitProcessResult,
 }
 
 /// Generate handshake response message.
-pub fn responde(wg: &WgInfo, result: &mut InitProcessResult, self_id: Id) -> [u8; 92] {
-    let mut response = [0u8; 92];
+pub fn responde(wg: &WgInfo, result: &mut InitProcessResult, self_id: Id)
+        -> [u8; HANDSHAKE_RESP_LEN] {
+    let mut response = [0u8; HANDSHAKE_RESP_LEN];
 
     // Type and zeros.
     response[0..4].copy_from_slice(&[2, 0, 0, 0]);
@@ -189,10 +195,12 @@ pub fn responde(wg: &WgInfo, result: &mut InitProcessResult, self_id: Id) -> [u8
 /// Process handshake response message.
 ///
 /// Returns peer index.
+///
+/// # Panics
+///
+/// If the message length is not `HANDSHAKE_RESP_LEN`.
 pub fn process_response(wg: &WgInfo, hs: &mut HS, msg: &[u8]) -> Result<Id, ()> {
-    if msg.len() != 92 {
-        return Err(());
-    }
+    debug_assert_eq!(msg.len(), HANDSHAKE_RESP_LEN);
 
     // Check mac1.
     let mac1 = mac(wg.psk.as_ref(), &[&wg.pubkey, &msg[..60]]);