progress [CI SKIP]

wg-event-gen/src/model.rs

@@ -1,19 +1,22 @@
 use base64::{decode, encode};
 use std::error::Error;
+use std::hash::Hash;
 use std::io;
 use std::net::{IpAddr, SocketAddr};
 use std::time::Duration;
 use std::time::Instant;
 
-#[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)]
+const KEY_SIZE: usize = 48; //TODO: use VEC instead of array
+
+#[derive(Debug, PartialEq, Eq, Hash, Clone)]
 pub enum ECCKey {
-    PublicKey([u8; 32]),
-    PrivateKey([u8; 32]),
+    PublicKey(Vec<u8>),
+    PrivateKey(Vec<u8>),
 }
 
 pub trait Base64Backed {
-    fn from_bytes(bytes: [u8; 32]) -> Self;
-    fn bytes(&self) -> &[u8; 32];
+    fn from_bytes(bytes: Vec<u8>) -> Self;
+    fn bytes(&self) -> &Vec<u8>;
     fn from_base64<I: AsRef<str>>(key: I) -> io::Result<Self>
     where
         Self: Sized,
@@ -29,12 +32,17 @@ pub trait Base64Backed {
         }; /*.map_err(|err| {
 
            })?;*/
-        let mut bytes = [0; 32];
-        if key.len() != 32 {
-            return Err(io::Error::new(io::ErrorKind::Other, "Mismatched key size"));
+        if key.len() != KEY_SIZE {
+            return Err(io::Error::new(
+                io::ErrorKind::Other,
+                format!(
+                    "Mismatched key size. Expected: {}, Got {}",
+                    KEY_SIZE,
+                    key.len()
+                ),
+            ));
         }
-        bytes.copy_from_slice(&key);
-        Ok(Self::from_bytes(bytes))
+        Ok(Self::from_bytes(key))
     }
 
     fn as_base64(&self) -> io::Result<String> {
@@ -43,14 +51,14 @@ pub trait Base64Backed {
 }
 
 impl Base64Backed for ECCKey {
-    fn bytes(&self) -> &[u8; 32] {
+    fn bytes(&self) -> &Vec<u8> {
         match self {
             ECCKey::PublicKey(bytes) => &bytes,
             ECCKey::PrivateKey(bytes) => &bytes,
         }
     }
 
-    fn from_bytes(bytes: [u8; 32]) -> ECCKey {
+    fn from_bytes(bytes: Vec<u8>) -> ECCKey {
         ECCKey::PublicKey(bytes)
     }
 }
@@ -62,40 +70,42 @@ impl ECCKey {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, Hash, Copy, Clone)]
-pub struct SharedKey([u8; 32]);
+#[derive(Debug, PartialEq, Eq, Hash, Clone)]
+pub struct SharedKey(Vec<u8>);
 
 impl Base64Backed for SharedKey {
-    fn bytes(&self) -> &[u8; 32] {
+    fn bytes(&self) -> &Vec<u8> {
         &self.0
     }
-    fn from_bytes(bytes: [u8; 32]) -> SharedKey {
+    fn from_bytes(bytes: Vec<u8>) -> SharedKey {
         SharedKey(bytes)
     }
 }
 
 #[derive(Debug, Builder, PartialEq, Eq, Hash, Clone)]
 pub struct Interface {
-    key: ECCKey,
-    port: usize,
-    fwmark: Option<String>
+    pub key: ECCKey,
+    pub port: usize,
+    pub fwmark: Option<String>,
 }
 
 #[derive(Debug, Builder, PartialEq, Eq, Hash, Clone)]
 pub struct Peer {
-    key: ECCKey,
+    pub key: ECCKey,
     #[builder(default = "None")]
-    shared_key: Option<SharedKey>,
+    pub shared_key: Option<SharedKey>,
     #[builder(default = "None")]
-    endpoint: Option<SocketAddr>,
-    allowed_ips: Vec<(IpAddr, u8)>,
-    last_handshake: Option<Duration>,
+    pub endpoint: Option<SocketAddr>,
+    #[builder(default = "Vec::new()")]
+    pub allowed_ips: Vec<(IpAddr, u8)>,
     #[builder(default = "None")]
-    persistent_keepalive: Option<Duration>,
+    pub last_handshake: Option<Duration>,
+    #[builder(default = "None")]
+    pub persistent_keepalive: Option<Duration>,
     #[builder(default = "(0u64,0u64)")]
-    traffic: (u64, u64),
+    pub traffic: (u64, u64),
     #[builder(default = "Instant::now()")]
-    parsed: Instant,
+    pub parsed: Instant,
 }
 
 impl PeerBuilder {
@@ -128,6 +138,15 @@ impl PeerBuilder {
                 .map(|shake| shake.map(|shake| shake + d));
         }
     }
+
+    pub fn add_traffic(&mut self, txrx: (u64, u64)) {
+        if let Some(ref mut traffic) = &mut self.traffic {
+            traffic.0 += txrx.0;
+            traffic.1 += txrx.1;
+        } else {
+            self.traffic = Some(txrx);
+        }
+    }
 }
 
 pub trait WireguardController {