only count authenticated packets in tx/rx numbers

2 files changed, 3 insertions, 7 deletions

diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index 2ad363a..a142d31 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -117,7 +117,6 @@ impl PeerServer {
         self.handle.spawn(self.tunnel_tx.clone().send(packet).then(|_| Ok(())));
     }
 
-    // TODO: create a transport packet (type 0x4) queue until a handshake has been completed
     fn handle_incoming_packet(&mut self, addr: SocketAddr, packet: Vec<u8>) -> Result<(), Error> {
         debug!("got a UDP packet from {:?} of length {}, packet type {}", &addr, packet.len(), packet[0]);
         let mut state = self.shared_state.borrow_mut();
@@ -161,9 +160,7 @@ impl PeerServer {
                     .read_message(&packet[12..60], &mut [])
                     .map_err(SyncFailure::new)?;
 
-                if payload_len != 0 {
-                    bail!("non-zero payload length in handshake response");
-                }
+                ensure!(payload_len == 0, "non-zero payload length in handshake response");
 
                 peer.ratchet_session()?;
                 info!("got handshake response, ratcheted session.");
diff --git a/src/protocol/peer.rs b/src/protocol/peer.rs
index d8e0433..ad6d361 100644
--- a/src/protocol/peer.rs
+++ b/src/protocol/peer.rs
@@ -131,7 +131,6 @@ impl Peer {
     }
 
     pub fn handle_incoming_transport(&mut self, our_index: u32, nonce: u64, addr: SocketAddr, packet: &[u8]) -> Result<Vec<u8>, Error> {
-        self.rx_bytes += packet.len() as u64;
 
         let session = self.sessions.current.as_mut().filter(|session| session.our_index == our_index)
             .or(self.sessions.past.as_mut().filter(|session| session.our_index == our_index))
@@ -145,6 +144,7 @@ impl Peer {
         let len = session.noise.read_message(packet, &mut raw_packet)
             .map_err(SyncFailure::new)?;
 
+        self.rx_bytes += packet.len() as u64;
         self.info.endpoint = Some(addr); // update peer endpoint after successful authentication
 
         raw_packet.truncate(len);
@@ -159,9 +159,9 @@ impl Peer {
         out_packet[0] = 4;
         LittleEndian::write_u32(&mut out_packet[4..], session.their_index);
         LittleEndian::write_u64(&mut out_packet[8..], session.noise.sending_nonce().map_err(SyncFailure::new)?);
-        self.tx_bytes += packet.len() as u64;
         let len = session.noise.write_message(packet, &mut out_packet[16..])
             .map_err(SyncFailure::new)?;
+        self.tx_bytes += len as u64;
         out_packet.truncate(TRANSPORT_HEADER_SIZE + len);
         Ok((endpoint, out_packet))
     }
@@ -236,7 +236,6 @@ impl Peer {
             ensure!(&timestamp > last_tai64n, "handshake timestamp earlier than last handshake's timestamp");
         }
 
-        // TODO: verify timestamp
         // TODO: hacked up API until it's officially supported in snow.
         match noise {
             snow::Session::Handshake(ref mut handshake_state) => {