refactor timestamp code a bit

5 files changed, 127 insertions, 100 deletions

diff --git a/src/interface/peer_server.rs b/src/interface/peer_server.rs
index fe97d1c..b849475 100644
--- a/src/interface/peer_server.rs
+++ b/src/interface/peer_server.rs
@@ -3,11 +3,12 @@ use consts::{REKEY_TIMEOUT, REKEY_AFTER_TIME, REJECT_AFTER_TIME, REKEY_ATTEMPT_T
 use cookie;
 use interface::{SharedPeer, SharedState, UtunPacket, config};
 use peer::{Peer, SessionType};
+use time::Timestamp;
 use timer::{Timer, TimerMessage};
 
 use std::io;
 use std::net::{IpAddr, Ipv6Addr, SocketAddr};
-use std::time::{Duration, Instant};
+use std::time::Duration;
 
 use byteorder::{ByteOrder, LittleEndian};
 use failure::{Error, err_msg};
@@ -289,8 +290,7 @@ impl PeerServer {
         }
 
         self.send_to_peer((endpoint, init_packet))?;
-        peer.last_sent_init = Some(Instant::now());
-        peer.last_tun_queue = peer.last_tun_queue.or_else(|| Some(Instant::now()));
+        peer.last_sent_init = Timestamp::now();
         let when = *REKEY_TIMEOUT + *TIMER_TICK_DURATION * 2;
         self.timer.spawn_delayed(&self.handle,
                                  when,
@@ -303,38 +303,29 @@ impl PeerServer {
             TimerMessage::Rekey(peer_ref, our_index) => {
                 {
                     let mut peer = peer_ref.borrow_mut();
-                    let     now  = Instant::now();
 
                     match peer.find_session(our_index) {
                         Some((_, SessionType::Next)) => {
-                            if let Some(sent_init) = peer.last_sent_init {
-                                let since_last_init = now.duration_since(sent_init);
-                                if since_last_init < *REKEY_TIMEOUT {
-                                    let wait = *REKEY_TIMEOUT - since_last_init + *TIMER_TICK_DURATION * 2;
-                                    self.timer.spawn_delayed(&self.handle,
-                                                             wait,
-                                                             TimerMessage::Rekey(peer_ref.clone(), our_index));
-                                    bail!("too soon since last init sent, waiting {:?} ({})", wait, our_index);
-                                }
+                            if peer.last_sent_init.elapsed() < *REKEY_TIMEOUT {
+                                let wait = *REKEY_TIMEOUT - peer.last_sent_init.elapsed() + *TIMER_TICK_DURATION * 2;
+                                self.timer.spawn_delayed(&self.handle,
+                                                         wait,
+                                                         TimerMessage::Rekey(peer_ref.clone(), our_index));
+                                bail!("too soon since last init sent, waiting {:?} ({})", wait, our_index);
                             }
-                            if let Some(init_attempt_epoch) = peer.last_tun_queue {
-                                let since_attempt_epoch = now.duration_since(init_attempt_epoch);
-                                if since_attempt_epoch > *REKEY_ATTEMPT_TIME {
-                                    peer.last_tun_queue = None;
-                                    bail!("REKEY_ATTEMPT_TIME exceeded ({})", our_index);
-                                }
+                            if peer.last_tun_queue.elapsed() > *REKEY_ATTEMPT_TIME {
+                                peer.last_tun_queue = Timestamp::unset();
+                                bail!("REKEY_ATTEMPT_TIME exceeded ({})", our_index);
                             }
                         },
                         Some((_, SessionType::Current)) => {
-                            if let Some(last_handshake) = peer.last_handshake {
-                                let since_last_handshake = now.duration_since(last_handshake);
-                                if since_last_handshake <= *REKEY_AFTER_TIME {
-                                    let wait = *REKEY_AFTER_TIME - since_last_handshake + *TIMER_TICK_DURATION * 2;
-                                    self.timer.spawn_delayed(&self.handle,
-                                                             wait,
-                                                             TimerMessage::Rekey(peer_ref.clone(), our_index));
-                                    bail!("recent last complete handshake - waiting {:?} ({})", wait, our_index);
-                                }
+                            let elapsed = peer.last_handshake.elapsed();
+                            if elapsed <= *REKEY_AFTER_TIME {
+                                let wait = *REKEY_AFTER_TIME - elapsed + *TIMER_TICK_DURATION * 2;
+                                self.timer.spawn_delayed(&self.handle,
+                                                         wait,
+                                                         TimerMessage::Rekey(peer_ref.clone(), our_index));
+                                bail!("recent last complete handshake - waiting {:?} ({})", wait, our_index);
                             }
                         },
                         _ => bail!("index is linked to a dead session, bailing.")
@@ -364,16 +355,13 @@ impl PeerServer {
                     let (session, session_type) = peer.find_session(our_index).ok_or_else(|| err_msg("missing session for timer"))?;
                     ensure!(session_type == SessionType::Current, "expired session for passive keepalive timer");
 
-                    if let Some(last_sent) = session.last_sent {
-                        let last_sent_packet = Instant::now().duration_since(last_sent);
-                        if last_sent_packet < *KEEPALIVE_TIMEOUT {
-                            self.timer.spawn_delayed(&self.handle,
-                                                     *KEEPALIVE_TIMEOUT - last_sent_packet + *TIMER_TICK_DURATION,
-                                                     TimerMessage::PassiveKeepAlive(peer_ref.clone(), our_index));
-                            bail!("passive keepalive tick (waiting {:?})", *KEEPALIVE_TIMEOUT - last_sent_packet);
-                        } else { // if we're going to send a keepalive reset last_sent
-                            session.last_sent = None;
-                        }
+                    if session.last_sent.elapsed() < *KEEPALIVE_TIMEOUT {
+                        self.timer.spawn_delayed(&self.handle,
+                                                 *KEEPALIVE_TIMEOUT - session.last_sent.elapsed() + *TIMER_TICK_DURATION,
+                                                 TimerMessage::PassiveKeepAlive(peer_ref.clone(), our_index));
+                        bail!("passive keepalive tick (waiting ~{}s)", (*KEEPALIVE_TIMEOUT - session.last_sent.elapsed()).as_secs());
+                    } else { // if we're going to send a keepalive reset last_sent
+                        session.last_sent = Timestamp::unset();
                     }
                 }
 
diff --git a/src/lib.rs b/src/lib.rs
index 14f8890..9cd3dd6 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -42,7 +42,7 @@ pub mod peer;
 pub mod types;
 pub mod anti_replay;
 pub mod router;
-pub mod tai64n;
+pub mod time;
 pub mod timer;
 pub mod ip_packet;
 pub mod xchacha20poly1305;
diff --git a/src/peer.rs b/src/peer.rs
index 6d3b59b..95daea5 100644
--- a/src/peer.rs
+++ b/src/peer.rs
@@ -10,9 +10,9 @@ use std::{self, mem};
 use std::collections::VecDeque;
 use std::fmt::{self, Debug, Display, Formatter};
 use std::net::SocketAddr;
-use std::time::{SystemTime, Instant, UNIX_EPOCH};
+use std::time::{SystemTime, UNIX_EPOCH};
 use hex;
-use tai64n::TAI64N;
+use time::{Tai64n, Timestamp};
 use rand::{self, Rng};
 use snow;
 use types::PeerInfo;
@@ -22,10 +22,10 @@ pub struct Peer {
     pub sessions              : Sessions,
     pub tx_bytes              : u64,
     pub rx_bytes              : u64,
-    pub last_sent_init        : Option<Instant>,
-    pub last_tun_queue        : Option<Instant>,
-    pub last_handshake        : Option<Instant>,
-    pub last_handshake_tai64n : Option<TAI64N>,
+    pub last_sent_init        : Timestamp,
+    pub last_tun_queue        : Timestamp,
+    pub last_handshake        : Timestamp,
+    pub last_handshake_tai64n : Option<Tai64n>,
     pub outgoing_queue        : VecDeque<UtunPacket>,
     pub cookie                : cookie::Generator,
 }
@@ -46,8 +46,8 @@ pub struct Session {
     pub our_index     : u32,
     pub their_index   : u32,
     pub anti_replay   : AntiReplay,
-    pub last_sent     : Option<Instant>,
-    pub last_received : Option<Instant>,
+    pub last_sent     : Timestamp,
+    pub last_received : Timestamp,
 }
 
 impl Session {
@@ -58,8 +58,8 @@ impl Session {
             our_index     : rand::thread_rng().gen::<u32>(),
             their_index   : their_index,
             anti_replay   : AntiReplay::default(),
-            last_sent     : None,
-            last_received : None,
+            last_sent     : Timestamp::default(),
+            last_received : Timestamp::default(),
         }
     }
 
@@ -82,15 +82,15 @@ impl From<snow::Session> for Session {
             our_index     : rand::thread_rng().gen::<u32>(),
             their_index   : 0,
             anti_replay   : AntiReplay::default(),
-            last_sent     : None,
-            last_received : None,
+            last_sent     : Timestamp::default(),
+            last_received : Timestamp::default(),
         }
     }
 }
 
 pub struct IncompleteIncomingHandshake {
     their_index : u32,
-    timestamp   : TAI64N,
+    timestamp   : Tai64n,
     noise       : snow::Session,
 }
 
@@ -149,7 +149,7 @@ impl Peer {
 
     pub fn queue_egress(&mut self, packet: UtunPacket) {
         self.outgoing_queue.push_back(packet);
-        self.last_tun_queue = Some(Instant::now());
+        self.last_tun_queue = Timestamp::now();
     }
 
     pub fn needs_new_handshake(&self) -> bool {
@@ -174,7 +174,7 @@ impl Peer {
         let     endpoint = self.info.endpoint.ok_or_else(|| err_msg("no known peer endpoint"))?;
         let mut packet   = vec![0; 148];
 
-        let tai64n = TAI64N::now();
+        let tai64n = Tai64n::now();
         packet[0] = 1;
 
         LittleEndian::write_u32(&mut packet[4..], session.our_index);
@@ -265,8 +265,7 @@ impl Peer {
         let current = mem::replace(&mut self.sessions.current, Some(session));
         let dead    = mem::replace(&mut self.sessions.past,    current);
 
-        self.last_handshake = Some(Instant::now());
-        self.last_tun_queue = None;
+        self.last_handshake = Timestamp::now();
         Ok(dead.map(|session| session.our_index))
     }
 
@@ -289,7 +288,7 @@ impl Peer {
                 .length();
             raw_packet.truncate(len as usize);
 
-            session.last_received = Some(Instant::now());
+            session.last_received = Timestamp::now();
 
             session_type
         };
@@ -299,8 +298,7 @@ impl Peer {
             let next    = std::mem::replace(&mut self.sessions.next, None);
             let current = std::mem::replace(&mut self.sessions.current, next);
             let dead    = std::mem::replace(&mut self.sessions.past, current);
-            self.last_handshake = Some(Instant::now());
-            self.last_tun_queue = None;
+            self.last_handshake = Timestamp::now();
             Some(dead.map(|session| session.our_index))
         } else {
             None
@@ -328,7 +326,7 @@ impl Peer {
         let padded_packet = &[packet, &vec![0u8; padding]].concat();
         let len = session.noise.write_message(padded_packet, &mut out_packet[16..])?;
         self.tx_bytes += len as u64;
-        session.last_sent = Some(Instant::now());
+        session.last_sent = Timestamp::now();
         out_packet.truncate(TRANSPORT_HEADER_SIZE + len);
         Ok((endpoint, out_packet))
     }
@@ -349,12 +347,10 @@ impl Peer {
         }
         s.push_str(&format!("tx_bytes={}\nrx_bytes={}\n", self.tx_bytes, self.rx_bytes));
 
-        if let Some(ref last_handshake) = self.last_handshake {
-            let system_now = SystemTime::now();
-            let time_passed = Instant::now().duration_since(*last_handshake);
-            if let Ok(time) = (system_now - time_passed).duration_since(UNIX_EPOCH) {
+        if self.last_handshake.is_set() {
+            if let Ok(time) = (SystemTime::now() - self.last_handshake.elapsed()).duration_since(UNIX_EPOCH) {
                 s.push_str(&format!("last_handshake_time_sec={}\nlast_handshake_time_nsec={}\n",
-                                    time.as_secs(), time.subsec_nanos()))
+                                    time.as_secs(), time.subsec_nanos()));
             } else {
                 warn!("SystemTime Duration error");
             }
diff --git a/src/tai64n.rs b/src/tai64n.rs
deleted file mode 100644
index 9e15ed1..0000000
--- a/src/tai64n.rs
+++ /dev/null
@@ -1,35 +0,0 @@
-use byteorder::{ByteOrder, BigEndian};
-use std::ops::Deref;
-use std::time::{SystemTime, UNIX_EPOCH};
-
-const TAI64N_BASE: i64 = 4611686018427387914;
-
-#[derive(PartialEq, PartialOrd)]
-pub struct TAI64N {
-    tai64n: [u8; 12]
-}
-
-impl TAI64N {
-    pub fn now() -> TAI64N {
-        let mut tai64n = [0u8; 12];
-        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
-        BigEndian::write_i64(&mut tai64n[0..], TAI64N_BASE + now.as_secs() as i64);
-        BigEndian::write_i32(&mut tai64n[8..], (now.subsec_nanos() as i32));
-
-        TAI64N { tai64n }
-    }
-}
-
-impl Deref for TAI64N {
-    type Target = [u8; 12];
-
-    fn deref(&self) -> &[u8; 12] {
-        &self.tai64n
-    }
-}
-
-impl From<[u8; 12]> for TAI64N {
-    fn from(tai64n: [u8; 12]) -> Self {
-        TAI64N { tai64n }
-    }
-}
diff --git a/src/time.rs b/src/time.rs
new file mode 100644
index 0000000..6261573
--- /dev/null
+++ b/src/time.rs
@@ -0,0 +1,78 @@
+use byteorder::{ByteOrder, BigEndian};
+use std::ops::Deref;
+use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
+use std::u64;
+
+const TAI64N_BASE: i64 = 4611686018427387914;
+
+#[derive(PartialEq, PartialOrd)]
+pub struct Tai64n {
+    tai64n: [u8; 12]
+}
+
+impl Tai64n {
+    pub fn now() -> Tai64n {
+        let mut tai64n = [0u8; 12];
+        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
+        BigEndian::write_i64(&mut tai64n[0..], TAI64N_BASE + now.as_secs() as i64);
+        BigEndian::write_i32(&mut tai64n[8..], now.subsec_nanos() as i32);
+
+        Tai64n { tai64n }
+    }
+}
+
+impl Deref for Tai64n {
+    type Target = [u8; 12];
+
+    fn deref(&self) -> &[u8; 12] {
+        &self.tai64n
+    }
+}
+
+impl From<[u8; 12]> for Tai64n {
+    fn from(tai64n: [u8; 12]) -> Self {
+        Tai64n { tai64n }
+    }
+}
+
+lazy_static! {
+    pub static ref FOREVER:     Duration = Duration::from_secs(u64::MAX);
+    pub static ref FOREVER_AGO: Instant  = Instant::now() - *FOREVER;
+}
+
+pub struct Timestamp(Option<Instant>);
+
+impl Default for Timestamp {
+    fn default() -> Self {
+        Timestamp(None)
+    }
+}
+
+impl Deref for Timestamp {
+    type Target = Instant;
+
+    fn deref(&self) -> &Self::Target {
+        match self.0 {
+            Some(ref time) => time,
+            None           => &*FOREVER_AGO,
+        }
+    }
+}
+
+impl Timestamp {
+    pub fn now() -> Self {
+        Timestamp(Some(Instant::now()))
+    }
+
+    pub fn unset() -> Self {
+        Timestamp(None)
+    }
+
+    pub fn is_set(&self) -> bool {
+        self.0.is_some()
+    }
+
+    pub fn elapsed(&self) -> Duration {
+        Instant::now().duration_since(**self)
+    }
+}