Added opaque identity to output

4 files changed, 49 insertions, 39 deletions

diff --git a/src/device.rs b/src/device.rs
index 4579f3b..66d5644 100644
--- a/src/device.rs
+++ b/src/device.rs
@@ -12,10 +12,10 @@ use crate::messages;
 use crate::types::*;
 use crate::peer::Peer;
 
-pub struct Device {
+pub struct Device<T> {
     pub sk : StaticSecret,               // static secret key
     pub pk : PublicKey,                  // static public key
-    peers  : Vec<Peer>,                  // peer index  -> state
+    peers  : Vec<Peer<T>>,               // peer index  -> state
     pk_map : HashMap<[u8; 32], usize>,   // public key  -> peer index
     id_map : RwLock<HashMap<u32, usize>> // receive ids -> peer index
 }
@@ -23,13 +23,13 @@ pub struct Device {
 /* A mutable reference to the device needs to be held during configuration.
  * Wrapping the device in a RwLock enables peer config after "configuration time"
  */
-impl Device {
+impl <T>Device<T> where T : Copy {
     /// Initialize a new handshake state machine
     ///
     /// # Arguments
     ///
     /// * `sk` - x25519 scalar representing the local private key
-    pub fn new(sk : StaticSecret) -> Device {
+    pub fn new(sk : StaticSecret) -> Device<T> {
         Device {
             pk     : PublicKey::from(&sk),
             sk     : sk,
@@ -45,7 +45,8 @@ impl Device {
     /// # Arguments
     ///
     /// * `pk` - The public key to add
-    pub fn add(&mut self, pk : PublicKey) -> Result<(), ConfigError> {
+    /// * `identifier` - Associated identifier which can be used to distinguish the peers
+    pub fn add(&mut self, pk : PublicKey, identifier : T) -> Result<(), ConfigError> {
         // check that the pk is not added twice
 
         if let Some(_) = self.pk_map.get(pk.as_bytes()) {
@@ -66,7 +67,7 @@ impl Device {
         // map : new index -> peer
 
         self.peers.push(Peer::new(
-            idx, pk, self.sk.diffie_hellman(&pk)
+            idx, identifier, pk, self.sk.diffie_hellman(&pk)
         ));
 
         Ok(())
@@ -128,7 +129,7 @@ impl Device {
     /// # Arguments
     ///
     /// * `msg` - Byte slice containing the message (untrusted input)
-    pub fn process(&self, msg : &[u8]) -> Result<Output, HandshakeError> {
+    pub fn process(&self, msg : &[u8]) -> Result<Output<T>, HandshakeError> {
         match msg.get(0) {
             Some(&messages::TYPE_INITIATION) => {
                 // consume the initiation
@@ -152,7 +153,7 @@ impl Device {
     // Internal function
     //
     // Return the peer associated with the public key
-    pub(crate) fn lookup_pk(&self, pk : &PublicKey) -> Result<&Peer, HandshakeError> {
+    pub(crate) fn lookup_pk(&self, pk : &PublicKey) -> Result<&Peer<T>, HandshakeError> {
         match self.pk_map.get(pk.as_bytes()) {
             Some(&idx) => Ok(&self.peers[idx]),
             _ => Err(HandshakeError::UnknownPublicKey)
@@ -162,7 +163,7 @@ impl Device {
     // Internal function
     //
     // Return the peer currently associated with the receiver identifier
-    pub(crate) fn lookup_id(&self, id : u32) -> Result<&Peer, HandshakeError> {
+    pub(crate) fn lookup_id(&self, id : u32) -> Result<&Peer<T>, HandshakeError> {
         match self.id_map.read().get(&id) {
             Some(&idx) => Ok(&self.peers[idx]),
             _ => Err(HandshakeError::UnknownReceiverId)
@@ -217,8 +218,8 @@ mod tests {
         let mut dev1 = Device::new(sk1);
         let mut dev2 = Device::new(sk2);
 
-        dev1.add(pk2).unwrap();
-        dev2.add(pk1).unwrap();
+        dev1.add(pk2, 1337).unwrap();
+        dev2.add(pk1, 2600).unwrap();
 
         // do a few handshakes
 
@@ -235,7 +236,7 @@ mod tests {
 
             // process initiation and create response
 
-            let (msg2, ks_r) = dev2.process(&msg1).unwrap();
+            let (_, msg2, ks_r) = dev2.process(&msg1).unwrap();
 
             let ks_r = ks_r.unwrap();
             let msg2 = msg2.unwrap();
@@ -247,7 +248,7 @@ mod tests {
 
             // process response and obtain confirmed key-pair
 
-            let (msg3, ks_i) = dev1.process(&msg2).unwrap();
+            let (_, msg3, ks_i) = dev1.process(&msg2).unwrap();
             let ks_i = ks_i.unwrap();
 
             assert!(msg3.is_none(), "Returned message after response");
diff --git a/src/noise.rs b/src/noise.rs
index f26a7df..d5495b8 100644
--- a/src/noise.rs
+++ b/src/noise.rs
@@ -184,11 +184,11 @@ mod tests {
     }
 }
 
-pub fn create_initiation(
-    device : &Device,
-    peer : &Peer,
+pub fn create_initiation<T>(
+    device : &Device<T>,
+    peer : &Peer<T>,
     sender : u32
-) -> Result<Vec<u8>, HandshakeError> {
+) -> Result<Vec<u8>, HandshakeError> where T : Copy {
 
     let mut rng = OsRng::new().unwrap();
     let mut msg : Initiation = Default::default();
@@ -263,10 +263,10 @@ pub fn create_initiation(
     Ok(Initiation::into(msg))
 }
 
-pub fn consume_initiation<'a>(
-    device : &'a Device,
+pub fn consume_initiation<'a, T>(
+    device : &'a Device<T>,
     msg : &[u8]
-) -> Result<(&'a Peer, TemporaryState), HandshakeError> {
+) -> Result<(&'a Peer<T>, TemporaryState), HandshakeError> where T : Copy {
 
     // parse message
 
@@ -341,11 +341,11 @@ pub fn consume_initiation<'a>(
     Ok((peer, (msg.f_sender, eph_r_pk, hs, ck)))
 }
 
-pub fn create_response(
-    peer     : &Peer,
+pub fn create_response<T>(
+    peer     : &Peer<T>,
     sender   : u32,           // sending identifier
     state    : TemporaryState // state from "consume_initiation"
-) -> Result<Output, HandshakeError> {
+) -> Result<Output<T>, HandshakeError> where T : Copy {
 
     let mut rng = OsRng::new().unwrap();
     let mut msg : Response = Default::default();
@@ -417,6 +417,7 @@ pub fn create_response(
     // return response and unconfirmed key-pair
 
     Ok((
+        peer.identifier,
         Some(Response::into(msg)),
         Some(KeyPair{
             confirmed : false,
@@ -432,10 +433,10 @@ pub fn create_response(
     ))
 }
 
-pub fn consume_response(
-    device : &Device,
+pub fn consume_response<T>(
+    device : &Device<T>,
     msg : &[u8]
-) -> Result<Output, HandshakeError> {
+) -> Result<Output<T>, HandshakeError> where T : Copy {
 
     // parse message
 
@@ -497,6 +498,7 @@ pub fn consume_response(
     // return response and unconfirmed key-pair
 
     Ok((
+        peer.identifier,
         None,
         Some(KeyPair{
             confirmed : true,
diff --git a/src/peer.rs b/src/peer.rs
index d846a21..dd7d303 100644
--- a/src/peer.rs
+++ b/src/peer.rs
@@ -16,7 +16,10 @@ use crate::device::Device;
  * This type is only for internal use and not exposed.
  */
 
-pub struct Peer {
+pub struct Peer<T> {
+    // external identifier
+    pub(crate) identifier : T,
+
     // internal identifier
     pub(crate) idx : usize,
 
@@ -55,19 +58,21 @@ impl Clone for State {
     }
 }
 
-impl Peer {
+impl <T>Peer<T> where T : Copy {
     pub fn new(
-        idx    : usize,
-        pk     : PublicKey,    // public key of peer
-        ss     : SharedSecret  // precomputed DH(static, static)
+        idx        : usize,
+        identifier : T,            // external identifier
+        pk         : PublicKey,    // public key of peer
+        ss         : SharedSecret  // precomputed DH(static, static)
     ) -> Self {
         Self {
-            idx       : idx,
-            state     : Mutex::new(State::Reset),
-            timestamp : Mutex::new(None),
-            pk        : pk,
-            ss        : ss,
-            psk       : [0u8; 32]
+            idx        : idx,
+            identifier : identifier,
+            state      : Mutex::new(State::Reset),
+            timestamp  : Mutex::new(None),
+            pk         : pk,
+            ss         : ss,
+            psk        : [0u8; 32]
         }
     }
 
@@ -97,7 +102,7 @@ impl Peer {
     /// * ts_new - The associated timestamp
     pub fn check_timestamp(
         &self,
-        device : &Device,
+        device : &Device<T>,
         timestamp_new : &timestamp::TAI64N
     ) -> Result<(), HandshakeError> {
 
diff --git a/src/types.rs b/src/types.rs
index 27807f2..593500f 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -91,7 +91,9 @@ pub struct KeyPair {
     pub recv      : Key   // key for inbound messages
 }
 
-pub type Output = (
+pub type Output<T> = (
+    T,               // external identifier associated with peer
+                     // (e.g. a reference or vector index)
     Option<Vec<u8>>, // message to send
     Option<KeyPair>  // resulting key-pair of successful handshake
 );