Restructured for wireguard-rs

13 files changed, 119 insertions, 63 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 1ff70c1..b7f93d8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -343,22 +343,6 @@ version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 
 [[package]]
-name = "wg-handshake"
-version = "0.1.0"
-dependencies = [
- "blake2 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)",
- "byteorder 1.3.1 (registry+https://github.com/rust-lang/crates.io-index)",
- "generic-array 0.12.3 (registry+https://github.com/rust-lang/crates.io-index)",
- "hex 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)",
- "hmac 0.7.1 (registry+https://github.com/rust-lang/crates.io-index)",
- "rand 0.6.5 (registry+https://github.com/rust-lang/crates.io-index)",
- "rust-crypto 0.2.36 (registry+https://github.com/rust-lang/crates.io-index)",
- "spin 0.5.0 (registry+https://github.com/rust-lang/crates.io-index)",
- "x25519-dalek 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)",
- "zerocopy 0.2.7 (registry+https://github.com/rust-lang/crates.io-index)",
-]
-
-[[package]]
 name = "winapi"
 version = "0.3.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -378,6 +362,22 @@ version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 
 [[package]]
+name = "wireguard-rs"
+version = "0.1.0"
+dependencies = [
+ "blake2 0.8.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "byteorder 1.3.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "generic-array 0.12.3 (registry+https://github.com/rust-lang/crates.io-index)",
+ "hex 0.3.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "hmac 0.7.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "rand 0.6.5 (registry+https://github.com/rust-lang/crates.io-index)",
+ "rust-crypto 0.2.36 (registry+https://github.com/rust-lang/crates.io-index)",
+ "spin 0.5.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "x25519-dalek 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "zerocopy 0.2.7 (registry+https://github.com/rust-lang/crates.io-index)",
+]
+
+[[package]]
 name = "x25519-dalek"
 version = "0.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index af24051..0b4ce0e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,5 +1,5 @@
 [package]
-name = "wg-handshake"
+name = "wireguard-rs"
 version = "0.1.0"
 authors = ["Mathias Hall-Andersen <mathias@hall-andersen.dk>"]
 edition = "2018"
diff --git a/src/lib.rs b/src/lib.rs
deleted file mode 100644
index d5a9c31..0000000
--- a/src/lib.rs
+++ /dev/null
@@ -1,10 +0,0 @@
-mod device;
-mod messages;
-mod noise;
-mod peer;
-mod timestamp;
-mod types;
-
-// publicly exposed interface
-
-pub use device::Device;
diff --git a/src/main.rs b/src/main.rs
new file mode 100644
index 0000000..b2995e7
--- /dev/null
+++ b/src/main.rs
@@ -0,0 +1,7 @@
+mod noise;
+mod types;
+
+use noise::Device;
+use types::KeyPair;
+
+fn main() {}
diff --git a/src/mod.rs b/src/mod.rs
new file mode 100644
index 0000000..dd3d99f
--- /dev/null
+++ b/src/mod.rs
@@ -0,0 +1,2 @@
+mod noise;
+mod types;
diff --git a/src/device.rs b/src/noise/device.rs
index a6081aa..04e00f9 100644
--- a/src/device.rs
+++ b/src/noise/device.rs
@@ -7,10 +7,10 @@ use rand::rngs::OsRng;
 use x25519_dalek::PublicKey;
 use x25519_dalek::StaticSecret;
 
-use crate::messages;
-use crate::noise;
-use crate::peer::Peer;
-use crate::types::*;
+use super::messages;
+use super::noise;
+use super::peer::Peer;
+use super::types::*;
 
 pub struct Device<T> {
     pub sk: StaticSecret,                   // static secret key
@@ -106,7 +106,7 @@ where
     /// # Returns
     ///
     /// The call might fail if the public key is not found
-    pub fn psk(&mut self, pk: PublicKey, psk: Option<Psk>) -> Result<(), ConfigError> {
+    pub fn set_psk(&mut self, pk: PublicKey, psk: Option<Psk>) -> Result<(), ConfigError> {
         match self.pk_map.get_mut(pk.as_bytes()) {
             Some(mut peer) => {
                 peer.psk = match psk {
@@ -119,6 +119,24 @@ where
         }
     }
 
+    /// Return the psk for the peer
+    ///
+    /// # Arguments
+    ///
+    /// * `pk` - The public key of the peer
+    ///
+    /// # Returns
+    ///
+    /// A 32 byte array holding the PSK
+    ///
+    /// The call might fail if the public key is not found
+    pub fn get_psk(&self, pk: PublicKey) -> Result<Psk, ConfigError> {
+        match self.pk_map.get(pk.as_bytes()) {
+            Some(peer) => Ok(peer.psk),
+            _ => Err(ConfigError::new("No such public key")),
+        }
+    }
+
     /// Release an id back to the pool
     ///
     /// # Arguments
@@ -233,6 +251,11 @@ mod tests {
         let sk2 = StaticSecret::new(&mut rng);
         let pk2 = PublicKey::from(&sk2);
 
+        // pick random psk
+
+        let mut psk = [0u8; 32];
+        rng.fill_bytes(&mut psk[..]);
+
         // intialize devices on both ends
 
         let mut dev1 = Device::new(sk1);
@@ -241,6 +264,9 @@ mod tests {
         dev1.add(pk2, 1337).unwrap();
         dev2.add(pk1, 2600).unwrap();
 
+        dev1.set_psk(pk2, Some(psk)).unwrap();
+        dev2.set_psk(pk1, Some(psk)).unwrap();
+
         // do a few handshakes
 
         for i in 0..10 {
@@ -279,5 +305,11 @@ mod tests {
             dev1.release(ks_i.send.id);
             dev2.release(ks_r.send.id);
         }
+
+        assert_eq!(dev1.get_psk(pk2).unwrap(), psk);
+        assert_eq!(dev2.get_psk(pk1).unwrap(), psk);
+
+        dev1.remove(pk2).unwrap();
+        dev2.remove(pk1).unwrap();
     }
 }
diff --git a/src/messages.rs b/src/noise/messages.rs
index 78f0838..dca49b9 100644
--- a/src/messages.rs
+++ b/src/noise/messages.rs
@@ -8,7 +8,7 @@ use byteorder::LittleEndian;
 use zerocopy::byteorder::U32;
 use zerocopy::{AsBytes, ByteSlice, FromBytes, LayoutVerified};
 
-use crate::types::*;
+use super::types::*;
 
 const SIZE_TAG: usize = 16;
 const SIZE_X25519_POINT: usize = 32;
diff --git a/src/noise/mod.rs b/src/noise/mod.rs
new file mode 100644
index 0000000..d48b5e0
--- /dev/null
+++ b/src/noise/mod.rs
@@ -0,0 +1,18 @@
+/* Implementation of the:
+ *
+ * Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s
+ *
+ * Protocol pattern, see: http://www.noiseprotocol.org/noise.html.
+ * For documentation.
+ */
+
+mod device;
+mod messages;
+mod noise;
+mod peer;
+mod timestamp;
+mod types;
+
+// publicly exposed interface
+
+pub use device::Device;
diff --git a/src/noise.rs b/src/noise/noise.rs
index 08935e0..980f1db 100644
--- a/src/noise.rs
+++ b/src/noise/noise.rs
@@ -17,11 +17,13 @@ use generic_array::GenericArray;
 
 use zerocopy::AsBytes;
 
-use crate::device::Device;
-use crate::messages::{Initiation, Response};
-use crate::peer::{Peer, State};
-use crate::timestamp;
-use crate::types::*;
+use super::device::Device;
+use super::messages::{Initiation, Response};
+use super::peer::{Peer, State};
+use super::timestamp;
+use super::types::*;
+
+use crate::types::{Key, KeyPair};
 
 // HMAC hasher (generic construction)
 
diff --git a/src/peer.rs b/src/noise/peer.rs
index 2dff10e..5b01d75 100644
--- a/src/peer.rs
+++ b/src/noise/peer.rs
@@ -7,9 +7,9 @@ use x25519_dalek::PublicKey;
 use x25519_dalek::SharedSecret;
 use x25519_dalek::StaticSecret;
 
-use crate::device::Device;
-use crate::timestamp;
-use crate::types::*;
+use super::device::Device;
+use super::timestamp;
+use super::types::*;
 
 /* Represents the recomputation and state of a peer.
  *
diff --git a/src/timestamp.rs b/src/noise/timestamp.rs
index 0996f8b..0996f8b 100644
--- a/src/timestamp.rs
+++ b/src/noise/timestamp.rs
diff --git a/src/types.rs b/src/noise/types.rs
index a18057f..0d9a5d3 100644
--- a/src/types.rs
+++ b/src/noise/types.rs
@@ -1,6 +1,10 @@
 use std::error::Error;
 use std::fmt;
 
+use crate::types::KeyPair;
+
+/* Internal types for the noise IKpsk2 implementation */
+
 // config error
 
 #[derive(Debug)]
@@ -65,28 +69,6 @@ impl Error for HandshakeError {
     }
 }
 
-// types for resulting key-material
-
-#[derive(Debug)]
-pub struct Key {
-    pub key: [u8; 32],
-    pub id: u32,
-}
-
-#[cfg(test)]
-impl PartialEq for Key {
-    fn eq(&self, other: &Self) -> bool {
-        self.id == other.id && self.key[..] == other.key[..]
-    }
-}
-
-#[derive(Debug)]
-pub struct KeyPair {
-    pub confirmed: bool, // has the key-pair been confirmed?
-    pub send: Key,       // key for outbound messages
-    pub recv: Key,       // key for inbound messages
-}
-
 pub type Output<T> = (
     T, // external identifier associated with peer
     // (e.g. a reference or vector index)
diff --git a/src/types/mod.rs b/src/types/mod.rs
new file mode 100644
index 0000000..ac6a307
--- /dev/null
+++ b/src/types/mod.rs
@@ -0,0 +1,23 @@
+/* This file holds types passed between components.
+ * Whenever a type cannot be held local to a single module.
+ */
+
+#[derive(Debug)]
+pub struct Key {
+    pub key: [u8; 32],
+    pub id: u32,
+}
+
+#[cfg(test)]
+impl PartialEq for Key {
+    fn eq(&self, other: &Self) -> bool {
+        self.id == other.id && self.key[..] == other.key[..]
+    }
+}
+
+#[derive(Debug)]
+pub struct KeyPair {
+    pub confirmed: bool, // has the key-pair been confirmed?
+    pub send: Key,       // key for outbound messages
+    pub recv: Key,       // key for inbound messages
+}