1 files changed, 30 insertions, 8 deletions

diff --git a/src/network/utilization.rs b/src/network/utilization.rs
index 3b67a0a..54537a1 100644
--- a/src/network/utilization.rs
+++ b/src/network/utilization.rs
@@ -1,32 +1,50 @@
 use crate::network::{Connection, Direction, Segment};
 
 use ::std::collections::HashMap;
+use ::std::time::SystemTime;
 
+#[derive(Clone)]
 pub struct TotalBandwidth {
     pub total_bytes_downloaded: u128,
     pub total_bytes_uploaded: u128,
 }
 
 impl TotalBandwidth {
-    pub fn increment_bytes_downloaded(&mut self, ip_length: u128) {
-        self.total_bytes_downloaded += ip_length;
+    pub fn increment_bytes_downloaded(&mut self, data_length: u128, reset_time: &SystemTime) {
+        if let Ok(elapsed) = reset_time.elapsed() {
+            if elapsed.as_millis() < 1000 {
+                self.total_bytes_downloaded += data_length;
+            }
+        }
     }
-    pub fn increment_bytes_uploaded(&mut self, ip_length: u128) {
-        self.total_bytes_uploaded += ip_length;
+    pub fn increment_bytes_uploaded(&mut self, data_length: u128, reset_time: &SystemTime) {
+        if let Ok(elapsed) = reset_time.elapsed() {
+            if elapsed.as_millis() < 1000 {
+                self.total_bytes_uploaded += data_length;
+            }
+        }
     }
 }
 
+#[derive(Clone)]
 pub struct Utilization {
     pub connections: HashMap<Connection, TotalBandwidth>,
+    reset_time: SystemTime,
 }
 
 impl Utilization {
     pub fn new() -> Self {
         let connections = HashMap::new();
-        Utilization { connections }
+        Utilization {
+            connections,
+            reset_time: SystemTime::now(),
+        }
     }
-    pub fn reset(&mut self) {
+    pub fn clone_and_reset(&mut self) -> Self {
+        let clone = self.clone();
+        self.reset_time = SystemTime::now();
         self.connections.clear();
+        clone
     }
     pub fn update(&mut self, seg: &Segment) {
         let total_bandwidth =
@@ -37,8 +55,12 @@ impl Utilization {
                     total_bytes_uploaded: 0,
                 });
         match seg.direction {
-            Direction::Download => total_bandwidth.increment_bytes_downloaded(seg.ip_length),
-            Direction::Upload => total_bandwidth.increment_bytes_uploaded(seg.ip_length),
+            Direction::Download => {
+                total_bandwidth.increment_bytes_downloaded(seg.data_length, &self.reset_time)
+            }
+            Direction::Upload => {
+                total_bandwidth.increment_bytes_uploaded(seg.data_length, &self.reset_time)
+            }
         }
     }
 }