add and strip padding to transport messages

4 files changed, 32 insertions, 14 deletions

diff --git a/benches/criterion.rs b/benches/criterion.rs
index d425840..9b179d3 100644
--- a/benches/criterion.rs
+++ b/benches/criterion.rs
@@ -4,13 +4,15 @@ extern crate wireguard;
 extern crate x25519_dalek;
 extern crate rand;
 extern crate snow;
+extern crate pnet;
 
 use criterion::{Benchmark, Criterion, Throughput};
 use wireguard::protocol::{Peer, Session};
-use wireguard::noise::Noise;
+use wireguard::noise;
 use x25519_dalek::{generate_secret, generate_public};
 use rand::OsRng;
 use std::time::Duration;
+use pnet::packet::{Packet, ipv4::MutableIpv4Packet};
 
 struct Keypair {
     pub private: [u8; 32],
@@ -33,8 +35,8 @@ fn connected_peers() -> (Peer, [u8; 32], Peer, [u8; 32]) {
     let mut peer_resp = Peer::default();
     let     init_keys = Keypair::new();
     let     resp_keys = Keypair::new();
-    let mut initiator = Noise::build_initiator(&init_keys.private, &resp_keys.public, &None).unwrap();
-    let mut responder = Noise::build_responder(&resp_keys.private).unwrap();
+    let mut initiator = noise::build_initiator(&init_keys.private, &resp_keys.public, &None).unwrap();
+    let mut responder = noise::build_responder(&resp_keys.private).unwrap();
     let mut buf       = [0u8; 500];
 
     match responder {
@@ -92,8 +94,10 @@ fn benchmarks(c: &mut Criterion) {
 
     c.bench("peer_transport_incoming", Benchmark::new("peer_transport_incoming", |b| {
         let (mut peer_init, _, mut peer_resp, _) = connected_peers();
+        let mut packet = MutableIpv4Packet::owned(vec![0u8; 1420]).unwrap();
+        packet.set_version(4);
         b.iter_with_setup(move || {
-            peer_init.handle_outgoing_transport(&[1u8; 1420]).expect("SETUP handle_outgoing_transport")
+            peer_init.handle_outgoing_transport(packet.packet()).expect("SETUP handle_outgoing_transport")
         }, move |(addr, packet)| {
             peer_resp.handle_incoming_transport(addr, &packet).expect("handle_incoming_transport")
         });
diff --git a/src/ip_packet.rs b/src/ip_packet.rs
index 67ec7d4..ad67f3f 100644
--- a/src/ip_packet.rs
+++ b/src/ip_packet.rs
@@ -20,17 +20,25 @@ impl<'a> IpPacket<'a> {
         }
     }
 
-    pub fn get_source(&self) -> IpAddr {
+    pub fn source(&self) -> IpAddr {
         match *self {
             IpPacket::V4(ref packet) => packet.get_source().into(),
             IpPacket::V6(ref packet) => packet.get_source().into(),
         }
     }
 
-    pub fn get_destination(&self) -> IpAddr {
+    pub fn destination(&self) -> IpAddr {
         match *self {
             IpPacket::V4(ref packet) => packet.get_destination().into(),
             IpPacket::V6(ref packet) => packet.get_destination().into(),
         }
     }
+
+    pub fn length(&self) -> u16 {
+        match *self {
+            IpPacket::V4(ref packet) => packet.get_total_length(),
+            IpPacket::V6(ref packet) => packet.get_payload_length(),
+        }
+
+    }
 }
diff --git a/src/protocol/peer.rs b/src/protocol/peer.rs
index bf130db..cf06a89 100644
--- a/src/protocol/peer.rs
+++ b/src/protocol/peer.rs
@@ -1,8 +1,9 @@
 use anti_replay::AntiReplay;
 use byteorder::{ByteOrder, LittleEndian};
-use consts::{TRANSPORT_OVERHEAD, TRANSPORT_HEADER_SIZE, MAX_SEGMENT_SIZE, REJECT_AFTER_MESSAGES};
+use consts::{TRANSPORT_OVERHEAD, TRANSPORT_HEADER_SIZE, MAX_SEGMENT_SIZE, REJECT_AFTER_MESSAGES, PADDING_MULTIPLE};
 use cookie;
 use failure::{Error, SyncFailure, err_msg};
+use ip_packet::IpPacket;
 use noise;
 use std::{self, mem};
 use std::fmt::{self, Debug, Display, Formatter};
@@ -231,9 +232,10 @@ impl Peer {
     pub fn handle_incoming_transport(&mut self, addr: SocketAddr, packet: &[u8])
         -> Result<(Vec<u8>, Option<u32>), Error> {
 
-        let our_index = LittleEndian::read_u32(&packet[4..]);
-        let nonce = LittleEndian::read_u64(&packet[8..]);
+        let     our_index  = LittleEndian::read_u32(&packet[4..]);
+        let     nonce      = LittleEndian::read_u64(&packet[8..]);
         let mut raw_packet = vec![0u8; MAX_SEGMENT_SIZE];
+
         let session_type = {
             let (session, session_type) = self.find_session(our_index).ok_or_else(|| err_msg("no session with index"))?;
             ensure!(session.noise.is_handshake_finished(), "session is not ready for transport packets");
@@ -241,7 +243,8 @@ impl Peer {
             session.anti_replay.update(nonce)?;
             session.noise.set_receiving_nonce(nonce).map_err(SyncFailure::new)?;
             let len = session.noise.read_message(&packet[16..], &mut raw_packet).map_err(SyncFailure::new)?;
-            raw_packet.truncate(len);
+            let len = IpPacket::new(&raw_packet[..len]).ok_or_else(||err_msg("invalid IP packet"))?.length();
+            raw_packet.truncate(len as usize);
 
             session.last_received = Some(Instant::now());
 
@@ -268,7 +271,9 @@ impl Peer {
     pub fn handle_outgoing_transport(&mut self, packet: &[u8]) -> Result<(SocketAddr, Vec<u8>), Error> {
         let session        = self.sessions.current.as_mut().ok_or_else(|| err_msg("no current noise session"))?;
         let endpoint       = self.info.endpoint.ok_or_else(|| err_msg("no known peer endpoint"))?;
-        let mut out_packet = vec![0u8; packet.len() + TRANSPORT_OVERHEAD];
+        let padding        = PADDING_MULTIPLE - (packet.len() % PADDING_MULTIPLE);
+        let padded_len     = packet.len() + padding;
+        let mut out_packet = vec![0u8; padded_len + TRANSPORT_OVERHEAD];
 
         let nonce = session.noise.sending_nonce().map_err(SyncFailure::new)?;
         ensure!(nonce < REJECT_AFTER_MESSAGES, "exceeded maximum message count");
@@ -276,7 +281,8 @@ impl Peer {
         out_packet[0] = 4;
         LittleEndian::write_u32(&mut out_packet[4..], session.their_index);
         LittleEndian::write_u64(&mut out_packet[8..], nonce);
-        let len = session.noise.write_message(packet, &mut out_packet[16..])
+        let padded_packet = &[packet, &vec![0u8; padding]].concat();
+        let len = session.noise.write_message(padded_packet, &mut out_packet[16..])
             .map_err(SyncFailure::new)?;
         self.tx_bytes += len as u64;
         session.last_sent = Some(Instant::now());
diff --git a/src/router.rs b/src/router.rs
index 85cc539..113b15f 100644
--- a/src/router.rs
+++ b/src/router.rs
@@ -42,14 +42,14 @@ impl Router {
 
     pub fn route_to_peer(&self, packet: &[u8]) -> Option<SharedPeer> {
         match IpPacket::new(packet) {
-            Some(packet) => self.get_peer_from_ip(packet.get_destination()),
+            Some(packet) => self.get_peer_from_ip(packet.destination()),
             _ => None
         }
     }
 
     pub fn validate_source(&self, packet: &[u8], peer: &SharedPeer) -> Result<(), Error> {
         let routed_peer = match IpPacket::new(packet) {
-            Some(packet) => self.get_peer_from_ip(packet.get_source()),
+            Some(packet) => self.get_peer_from_ip(packet.source()),
             _ => None
         }.ok_or_else(|| err_msg("no peer found on route"))?;