diff options
-rw-r--r-- | benches/criterion.rs | 12 | ||||
-rw-r--r-- | src/ip_packet.rs | 12 | ||||
-rw-r--r-- | src/protocol/peer.rs | 18 | ||||
-rw-r--r-- | src/router.rs | 4 |
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"))?; |