diff --git a/wg-event-gen/src/controller.rs b/wg-event-gen/src/controller.rs index 2cff8fb..1946159 100644 --- a/wg-event-gen/src/controller.rs +++ b/wg-event-gen/src/controller.rs @@ -1,5 +1,5 @@ use crate::model::{ECCKey, Peer, PeerBuilder, WireguardController}; -use std::io::{BufRead, BufReader, Result, Write}; +use std::io::{BufRead, BufReader, Error, ErrorKind, Result, Write}; use std::net::{IpAddr, SocketAddr}; use std::os::unix::net::UnixStream; use std::path::PathBuf; @@ -18,43 +18,48 @@ impl WireguardController for Userspace { let mut stream = UnixStream::connect(&self.0)?; stream.write_all(b"get=1\n")?; - fn build_peer(builder: &mut PeerBuilder, line: io::Result) -> Option> { - let line = line?; + fn build_peer(builder: &mut PeerBuilder, line: Result) -> Option> { + let line = match line { + Ok(line) => line, + Err(e) => return Some(Err(e)), + }; + fn parse_err(res: std::error::Result) -> Result { + res.map_err(|err: String| Error::new(ErrorKind::InvalidData, err.into())) + } let mut iter = line.chars(); let key = iter.by_ref().take_while(|c| c != &'=').collect::(); let value = iter.collect::(); - let mut ips: Vec<(IpAddr, u8)> = Vec::with_capacity(0); - if builder.allowed_ips.len() > 1 { - ips.append(&builder.allowed_ips); - } - let update_handshake = |d: Duration| { - builder.last_handshake(last_handshake.map(|c| c + d)); - }; - let value_as_num = || value.parse::(); + let value_as_num = || parse_err(value.parse::()); let mut peer: Option = None; let add_key = |key| { - if builder.key.is_some() { - //If next entry begins - peer = Some(builder.build().map_err(|err: String| { - io::Error::new(ErrorKind::InvalidData, err.into()) - })?); + if builder.is_whole() { + peer = Some(parse_err(builder.build())?); *builder = PeerBuilder::default(); + } else { + peer = None } builder.key(key); + Ok(peer) }; match key.as_ref() { - "public_key" => add_key(ECCKey::from_base64(value)?), - "private_key" => add_key(ECCKey::from_base64(value)?), - "endpoint" => builder.endpoint(value.parse::()?), + "public_key" => { + add_key(parse_err(ECCKey::from_base64(value))?)?; + } + /*"private_key" => { + add_key(ECCKey::from_base64(value)?)?; + } + "endpoint" => { + builder.endpoint(Some(parse_err(value.parse::())?)); + } "last_handshake_time_sec" => { - update_handshake(Duration::from_secs(value_as_num().into())); + builder.add_last_handshake(Duration::from_secs(value_as_num()?)); } "last_handshake_time_nsec" => { - update_handshake(Duration::from_nanos(value_as_num().into())); - } + builder.add_last_handshake(Duration::from_nanos(value_as_num()?.into())); + }, "persistent_keepalive" => { - builder.persistent_keepalive(Some(Duration::from_secs(value_as_num().into()))); - } + builder.persistent_keepalive(Some(Duration::from_secs(value_as_num()?.into()))); + },*/ "allowed_ip" => { let mut parts = value.split("/").into_iter(); let net = match ( @@ -67,7 +72,7 @@ impl WireguardController for Userspace { _ => None, }; if let Some(net) = net { - builder.allowed_ips.map(|ips| ips.push(net)); + builder.add_allowed_ip(net); } } } diff --git a/wg-event-gen/src/model.rs b/wg-event-gen/src/model.rs index e01d7db..86e9257 100644 --- a/wg-event-gen/src/model.rs +++ b/wg-event-gen/src/model.rs @@ -32,7 +32,6 @@ impl ECCKey { struct SharedKey([u8; 32]); #[derive(Debug, Builder, PartialEq, Eq, Hash, Clone)] -#[builder(public)] pub struct Peer { key: ECCKey, #[builder(default = "None")] @@ -49,6 +48,37 @@ pub struct Peer { parsed: Instant, } +impl PeerBuilder { + fn validate(&self) -> Result<(), String> { + if let Some(ref key) = self.key { + Ok(()) + } else { + Err("No key supplied".into()) + } + } + + pub fn is_whole(&self) -> bool { + self.validate().is_ok() + } + + pub fn add_allowed_ip(&mut self, ip: (IpAddr, u8)) { + if !self.allowed_ips.is_some() { + self.allowed_ips = Some(Vec::new()); + } + self.allowed_ips.map(|ips| ips.push(ip)); + } + + pub fn add_last_handshake(&mut self, d: Duration) { + if !self.last_handshake.is_some() { + self.last_handshake = Some(Some(d)); + } else { + self.last_handshake = self + .last_handshake + .map(|shake| shake.map(|shake| shake + d)); + } + } +} + pub trait WireguardController { fn peers<'a>(&'a mut self) -> io::Result> + 'a>>;