Add peers and psks to device

4 files changed, 118 insertions, 55 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 33f2d81..42f8449 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,12 +1,4 @@
+mod types;
+mod noise;
 mod messages;
 mod machine;
-mod noise;
-
-#[cfg(test)]
-mod tests {
-    #[test]
-    fn it_works() {
-        assert_eq!(2 + 2, 4);
-    }
-}
-
diff --git a/src/machine.rs b/src/machine.rs
index a8d4d07..8cb6fba 100644
--- a/src/machine.rs
+++ b/src/machine.rs
@@ -1,61 +1,57 @@
+use std::sync::Mutex;
+use std::collections::HashMap;
+
+use rand::prelude::*;
+use rand::rngs::OsRng;
+
 use x25519_dalek::PublicKey;
 use x25519_dalek::StaticSecret;
 use x25519_dalek::SharedSecret;
 
-use std::sync::Mutex;
-use std::sync::mpsc::channel;
-use std::collections::HashMap;
+use crate::noise;
+use crate::types;
+
+pub struct Output (
+    Option<types::KeyPair>, // resulting key-pair of successful handshake
+    Option<Vec<u8>>         // message to send
+);
 
 pub struct Peer {
-    m   : StateMutable,
-    sk  : StaticSecret,
-    pk  : PublicKey,
-    ss  : SharedSecret,
-    psk : [u8; 32]
+    // mutable state
+    m   : Mutex<State>,
+
+    // constant state
+    pk  : PublicKey,     // public key of peer
+    ss  : SharedSecret,  // precomputed DH(static, static)
+    psk : [u8; 32]       // psk of peer
 }
 
-/* Mutable part of handshake state */
-enum StateMutable {
+enum State {
     Reset,
     InitiationSent,
-    InitiationProcessed,
-    ReponseSent
-}
-
-/* Immutable part of the handshake state */
-struct StateFixed {
 }
 
-struct StateMachine {
-    peers : Vec<Mutex<Peer>>,          // peer index  -> state
+struct Device {
+    sk    : StaticSecret,              // static secret key
+    pk    : PublicKey,                 // static public key
+    peers : Vec<Peer>,                 // peer index  -> state
     pkmap : HashMap<[u8; 32], usize>,  // public key  -> peer index
     ids   : Mutex<HashMap<u32, usize>> // receive ids -> peer index
 }
 
-struct Key {
-    key : [u8; 32],
-    id  : u32
-}
-
-struct KeyPair {
-    confimed : bool, // has the key-pair been confirmed
-    send     : Key,  // key for outbound messages
-    recv     : Key   // key for inbound messages
-}
-
-struct Output (
-    Option<KeyPair>, // resulting key-pair of successful handshake
-    Option<u32>      // id to be released
-);
-
-impl StateMachine {
+/* A mutable reference to the state machine needs to be held,
+ * during configuration.
+ */
+impl Device {
     /// Initialize a new handshake state machine
     ///
     /// # Arguments
     ///
     /// * `sk` - x25519 scalar representing the local private key
-    pub fn new(sk : StaticSecret) -> StateMachine {
-        StateMachine {
+    pub fn new(sk : StaticSecret) -> Device {
+        Device {
+            pk    : PublicKey::from(&sk),
+            sk    : sk,
             peers : vec![],
             pkmap : HashMap::new(),
             ids   : Mutex::new(HashMap::new())
@@ -73,8 +69,56 @@ impl StateMachine {
     ///
     /// The call might fail if the public key corresponds to the secret key of the machine
     pub fn add(&mut self, pk : PublicKey) -> Result<(), ()> {
-        // let ss = sk.diffie_hellman(&pk);
-        Err(())
+        // check that the pk is not added twice
+
+        if let Some(_) = self.pkmap.get(pk.as_bytes()) {
+            return Err(());
+        };
+
+        // check that the pk is not that of the device
+
+        if *self.pk.as_bytes() == *pk.as_bytes() {
+            return Err(());
+        }
+
+        // map : pk -> new index
+
+        self.pkmap.insert(*pk.as_bytes(), self.peers.len());
+
+        // map : new index -> peer
+
+        self.peers.push(Peer {
+            m   : Mutex::new(State::Reset),
+            pk  : pk,
+            ss  : self.sk.diffie_hellman(&pk),
+            psk : [0u8; 32]
+        });
+
+        Ok(())
+    }
+
+    /// Add a psk to the peer
+    ///
+    /// # Arguments
+    ///
+    /// * `pk` - The public key of the peer
+    /// * `psk` - The psk to set / unset
+    ///
+    /// # Returns
+    ///
+    /// The call might fail if the public key is not found
+    pub fn psk(&mut self, pk : PublicKey, psk : Option<[u8; 32]>) -> Result<(), ()> {
+        match self.pkmap.get(pk.as_bytes()) {
+            Some(&idx) => {
+                let peer = &mut self.peers[idx];
+                peer.psk = match psk {
+                    Some(v) => v,
+                    None => [0u8; 32],
+                };
+                Ok(())
+            },
+            _ => Err(())
+        }
     }
 
     /// Release an id back to the pool
@@ -91,12 +135,13 @@ impl StateMachine {
     /// # Arguments
     ///
     /// * `pk` - Public key of peer to initiate handshake for
-    pub fn begin(&self, pk : PublicKey) -> Result<Output, ()> {
+    pub fn begin(&self, pk : PublicKey) -> Result<Vec<u8>, ()> {
         match self.pkmap.get(pk.as_bytes()) {
             None => Err(()),
             Some(&idx) => {
-                let mut peer = self.peers.get(idx).unwrap().lock().unwrap();
-                Err(())
+                let peer = &self.peers[idx];
+                let id = self.allocate(idx);
+                noise::create_initiation(peer, id)
             }
         }
     }
@@ -106,10 +151,25 @@ impl StateMachine {
     /// # Arguments
     ///
     /// * `msg` - Byte slice containing the message (untrusted input)
-    fn process(&self, msg : &[u8]) -> Result<Output, ()> {
+    pub fn process(&self, msg : &[u8]) -> Result<Output, ()> {
         // inspect type field
         match msg.get(0) {
             _ => Err(())
         }
     }
 }
+
+impl Device {
+    // allocate a new index (id), for peer with idx
+    fn allocate(&self, idx : usize) -> u32 {
+        let mut rng = OsRng;
+        let mut table = self.ids.lock().unwrap();
+        loop {
+            let id = rng.gen();
+            if !table.contains_key(&id) {
+                table.insert(id, idx);
+                return id;
+            }
+        }
+    }
+}
diff --git a/src/noise.rs b/src/noise.rs
index 189ada0..2d4c122 100644
--- a/src/noise.rs
+++ b/src/noise.rs
@@ -1,5 +1,6 @@
 use crate::machine::Peer;
+use crate::machine::Output;
 
-fn create_initiation(st : &mut Peer) {
-
+pub fn create_initiation(peer : &Peer, id : u32) -> Result<Vec<u8>, ()> {
+    Ok(vec![])
 }
diff --git a/src/types.rs b/src/types.rs
new file mode 100644
index 0000000..3729555
--- /dev/null
+++ b/src/types.rs
@@ -0,0 +1,10 @@
+struct Key {
+    key : [u8; 32],
+    id  : u32
+}
+
+pub struct KeyPair {
+    confimed : bool, // has the key-pair been confirmed?
+    send     : Key,  // key for outbound messages
+    recv     : Key   // key for inbound messages
+}