.

src/crypto/init.rs

@@ -53,7 +53,9 @@
 
 use super::{
     core::{CryptoCore, EXTRA_LEN},
-    Algorithms, EcdhPrivateKey, EcdhPublicKey, Ed25519PublicKey, Error, MsgBuffer, Payload
+    rotate::RotationState,
+    Algorithms, EcdhPrivateKey, EcdhPublicKey, Ed25519PublicKey, Error, MsgBuffer, Payload, PeerCrypto,
+    MESSAGE_TYPE_ROTATION
 };
 use crate::types::NodeId;
 use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
@@ -72,6 +74,7 @@ use std::{
     sync::Arc
 };
 
+pub const INIT_MESSAGE_FIRST_BYTE: u8 = 0xff;
 
 pub const STAGE_PING: u8 = 1;
 pub const STAGE_PONG: u8 = 2;
@@ -85,6 +88,11 @@ pub const SALTED_NODE_ID_HASH_LEN: usize = 20;
 pub type SaltedNodeIdHash = [u8; SALTED_NODE_ID_HASH_LEN];
 
 
+pub fn is_init_message(msg: &[u8]) -> bool {
+    !msg.is_empty() && msg[0] == INIT_MESSAGE_FIRST_BYTE
+}
+
+
 #[allow(clippy::large_enum_variant)]
 pub enum InitMsg {
     Ping {
@@ -141,6 +149,9 @@ impl InitMsg {
     fn read_from(buffer: &[u8], trusted_keys: &[Ed25519PublicKey]) -> Result<(Self, Ed25519PublicKey), Error> {
         let mut r = Cursor::new(buffer);
 
+        if r.read_u8().map_err(|_| Error::Parse("Init message too short"))? != INIT_MESSAGE_FIRST_BYTE {
+            return Err(Error::Parse("Init message has invalid first byte"))
+        }
         let mut public_key_salt = [0; 4];
         r.read_exact(&mut public_key_salt).map_err(|_| Error::Parse("Init message too short"))?;
         let mut public_key_hash = [0; 4];
@@ -290,6 +301,7 @@ impl InitMsg {
     fn write_to(&self, buffer: &mut [u8], key: &Ed25519KeyPair) -> Result<usize, io::Error> {
         let mut w = Cursor::new(buffer);
 
+        w.write_u8(INIT_MESSAGE_FIRST_BYTE)?;
         let rand = SystemRandom::new();
         let mut salt = [0; 4];
         rand.fill(&mut salt).unwrap();
@@ -378,9 +390,9 @@ pub enum InitResult<P: Payload> {
     Success { peer_payload: P, is_initiator: bool }
 }
 
-
 pub struct InitState<P: Payload> {
     node_id: NodeId,
+    is_initiator: bool,
     salted_node_id_hash: SaltedNodeIdHash,
     payload: P,
     key_pair: Arc<Ed25519KeyPair>,
@@ -389,7 +401,7 @@ pub struct InitState<P: Payload> {
     next_stage: u8,
     close_time: usize,
     last_message: Option<Vec<u8>>,
-    crypto: Option<CryptoCore>,
+    crypto: Option<Arc<CryptoCore>>,
     algorithms: Algorithms,
     selected_algorithm: Option<&'static Algorithm>,
     failed_retries: usize
@@ -409,6 +421,7 @@ impl<P: Payload> InitState<P> {
         hash[4..].clone_from_slice(&d.as_ref()[..16]);
         Self {
             node_id,
+            is_initiator: false,
             salted_node_id_hash: hash,
             payload,
             key_pair,
@@ -431,7 +444,7 @@ impl<P: Payload> InitState<P> {
 
         // create stage 1 msg
         self.send_message(STAGE_PING, Some(ecdh_public_key), out);
-
+        self.is_initiator = true;
         self.next_stage = STAGE_PONG;
     }
 
@@ -588,6 +601,7 @@ impl<P: Payload> InitState<P> {
                 // the node with the higher node_id "wins" and gets to initialize the connection
                 if salted_node_id_hash > self.salted_node_id_hash {
                     // reset to initial state
+                    self.is_initiator = false;
                     self.next_stage = STAGE_PING;
                     self.last_message = None;
                     self.ecdh_private_key = None;
@@ -614,7 +628,8 @@ impl<P: Payload> InitState<P> {
                 self.selected_algorithm = algorithm.map(|a| a.0);
                 if let Some((algorithm, _speed)) = algorithm {
                     let master_key = self.derive_master_key(algorithm, my_ecdh_private_key, &ecdh_public_key);
-                    self.crypto = Some(CryptoCore::new(master_key, self.salted_node_id_hash > salted_node_id_hash));
+                    self.crypto =
+                        Some(Arc::new(CryptoCore::new(master_key, self.salted_node_id_hash > salted_node_id_hash)));
                 }
 
                 // create and send stage 2 reply
@@ -630,7 +645,8 @@ impl<P: Payload> InitState<P> {
                 self.selected_algorithm = algorithm.map(|a| a.0);
                 if let Some((algorithm, _speed)) = algorithm {
                     let master_key = self.derive_master_key(algorithm, ecdh_private_key, &ecdh_public_key);
-                    self.crypto = Some(CryptoCore::new(master_key, self.salted_node_id_hash > salted_node_id_hash));
+                    self.crypto =
+                        Some(Arc::new(CryptoCore::new(master_key, self.salted_node_id_hash > salted_node_id_hash)));
                 }
 
                 // decrypt the payload
@@ -657,8 +673,19 @@ impl<P: Payload> InitState<P> {
         }
     }
 
-    pub fn take_core(&mut self) -> Option<CryptoCore> {
-        self.crypto.take()
+    pub fn finish(self, buffer: &mut MsgBuffer) -> PeerCrypto {
+        assert!(buffer.is_empty());
+        let rotation = if self.crypto.is_some() { Some(RotationState::new(!self.is_initiator, buffer)) } else { None };
+        if !buffer.is_empty() {
+            buffer.prepend_byte(MESSAGE_TYPE_ROTATION);
+        }
+        PeerCrypto {
+            algorithm: self.crypto.map(|c| c.algorithm()),
+            core: self.crypto,
+            rotation,
+            rotate_counter: 0,
+            last_init_message: self.last_message
+        }
     }
 }
 

src/crypto/mod.rs

@@ -2,11 +2,10 @@ mod core;
 mod init;
 mod rotate;
 
-pub use self::core::{EXTRA_LEN, TAG_LEN, CryptoCore};
-use self::{
-    core::{test_speed, CryptoCore},
-    init::{InitResult, InitState, CLOSING},
-    rotate::RotationState
+use self::{core::test_speed, init::InitState, rotate::RotationState};
+pub use self::{
+    core::{CryptoCore, EXTRA_LEN, TAG_LEN},
+    init::{is_init_message, INIT_MESSAGE_FIRST_BYTE}
 };
 use crate::{
     error::Error,
@@ -26,7 +25,7 @@ use thiserror::Error;
 
 
 const SALT: &[u8; 32] = b"vpncloudVPNCLOUDvpncl0udVpnCloud";
-const INIT_MESSAGE_FIRST_BYTE: u8 = 0xff;
+
 const MESSAGE_TYPE_ROTATION: u8 = 0x10;
 
 pub type Ed25519PublicKey = [u8; ED25519_PUBLIC_KEY_LEN];
@@ -186,100 +185,35 @@ impl Crypto {
         Ok(result)
     }
 
-    pub fn peer_instance<P: Payload>(&self, payload: P) -> PeerCrypto<P> {
-        PeerCrypto::new(
-            self.node_id,
-            payload,
-            self.key_pair.clone(),
-            self.trusted_keys.clone(),
-            self.algorithms.clone()
-        )
+    pub fn peer_instance<P: Payload>(&self, payload: P) -> InitState<P> {
+        InitState::new(self.node_id, payload, self.key_pair.clone(), self.trusted_keys.clone(), self.algorithms.clone())
     }
 }
 
 
 #[derive(Debug, PartialEq)]
-pub enum MessageResult<P: Payload> {
+pub enum MessageResult {
     Message(u8),
-    Initialized(P),
-    InitializedWithReply(P),
     Reply,
     None
 }
 
 
-pub struct PeerCrypto<P: Payload> {
+// TODO: completely rewrite
+// PeerCrypto is only for initialized crypto
+// Init is consumed and generates PeerCrypto
+pub struct PeerCrypto {
     #[allow(dead_code)]
-    node_id: NodeId,
-    init: Option<InitState<P>>,
+    last_init_message: Option<Vec<u8>>,
+    algorithm: Option<&'static Algorithm>,
     rotation: Option<RotationState>,
-    unencrypted: bool,
-    core: Option<CryptoCore>,
+    core: Option<Arc<CryptoCore>>,
     rotate_counter: usize
 }
 
-impl<P: Payload> PeerCrypto<P> {
-    pub fn new(
-        node_id: NodeId, init_payload: P, key_pair: Arc<Ed25519KeyPair>, trusted_keys: Arc<[Ed25519PublicKey]>,
-        algorithms: Algorithms
-    ) -> Self
-    {
-        Self {
-            node_id,
-            init: Some(InitState::new(node_id, init_payload, key_pair, trusted_keys, algorithms)),
-            rotation: None,
-            unencrypted: false,
-            core: None,
-            rotate_counter: 0
-        }
-    }
-
-    fn get_init(&mut self) -> Result<&mut InitState<P>, Error> {
-        if let Some(init) = &mut self.init {
-            Ok(init)
-        } else {
-            Err(Error::InvalidCryptoState("Initialization already finished"))
-        }
-    }
-
-    fn get_core(&mut self) -> Result<&mut CryptoCore, Error> {
-        if let Some(core) = &mut self.core {
-            Ok(core)
-        } else {
-            Err(Error::InvalidCryptoState("Crypto core not ready yet"))
-        }
-    }
-
-    fn get_rotation(&mut self) -> Result<&mut RotationState, Error> {
-        if let Some(rotation) = &mut self.rotation {
-            Ok(rotation)
-        } else {
-            Err(Error::InvalidCryptoState("Key rotation not initialized"))
-        }
-    }
-
-    pub fn initialize(&mut self, out: &mut MsgBuffer) -> Result<(), Error> {
-        let init = self.get_init()?;
-        if init.stage() != init::STAGE_PING {
-            Err(Error::InvalidCryptoState("Initialization already ongoing"))
-        } else {
-            init.send_ping(out);
-            out.prepend_byte(INIT_MESSAGE_FIRST_BYTE);
-            Ok(())
-        }
-    }
-
-    pub fn has_init(&self) -> bool {
-        self.init.is_some()
-    }
-
-    pub fn is_ready(&self) -> bool {
-        self.core.is_some()
-    }
-
+impl PeerCrypto {
     pub fn algorithm_name(&self) -> &'static str {
-        if let Some(ref core) = self.core {
-            let algo = core.algorithm();
+        if let Some(algo) = self.algorithm {
             if algo == &aead::CHACHA20_POLY1305 {
                 "CHACHA20"
             } else if algo == &aead::AES_128_GCM {
@@ -294,72 +228,43 @@ impl<P: Payload> PeerCrypto<P> {
         }
     }
 
-    fn handle_init_message(&mut self, buffer: &mut MsgBuffer) -> Result<MessageResult<P>, Error> {
-        let result = self.get_init()?.handle_init(buffer)?;
-        if !buffer.is_empty() {
-            buffer.prepend_byte(INIT_MESSAGE_FIRST_BYTE);
-        }
-        match result {
-            InitResult::Continue => Ok(MessageResult::Reply),
-            InitResult::Success { peer_payload, is_initiator } => {
-                self.core = self.get_init()?.take_core();
-                if self.core.is_none() {
-                    self.unencrypted = true;
-                }
-                if self.get_init()?.stage() == init::CLOSING {
-                    self.init = None
-                }
-                if self.core.is_some() {
-                    self.rotation = Some(RotationState::new(!is_initiator, buffer));
-                }
-                if !is_initiator {
-                    if self.unencrypted {
-                        return Ok(MessageResult::Initialized(peer_payload))
-                    }
-                    assert!(!buffer.is_empty());
-                    buffer.prepend_byte(MESSAGE_TYPE_ROTATION);
-                    self.encrypt_message(buffer)?;
-                }
-                Ok(MessageResult::InitializedWithReply(peer_payload))
-            }
-        }
+    fn handle_init_message(&mut self, buffer: &mut MsgBuffer) -> Result<MessageResult, Error> {
+        // TODO: parse message stage
+        // TODO: depending on stage resend last message
+        Ok(MessageResult::None)
     }
 
     fn handle_rotate_message(&mut self, data: &[u8]) -> Result<(), Error> {
-        if self.unencrypted {
-            return Ok(())
-        }
-        if let Some(rot) = self.get_rotation()?.handle_message(data)? {
-            let core = self.get_core()?;
-            let algo = core.algorithm();
-            let key = LessSafeKey::new(UnboundKey::new(algo, &rot.key[..algo.key_len()]).unwrap());
-            core.rotate_key(key, rot.id, rot.use_for_sending);
+        if let Some(rotation) = &mut self.rotation {
+            if let Some(rot) = rotation.handle_message(data)? {
+                let algo = self.algorithm.unwrap();
+                let key = LessSafeKey::new(UnboundKey::new(algo, &rot.key[..algo.key_len()]).unwrap());
+                self.core.unwrap().rotate_key(key, rot.id, rot.use_for_sending);
+            }
         }
         Ok(())
     }
 
-    fn encrypt_message(&mut self, buffer: &mut MsgBuffer) -> Result<(), Error> {
-        if self.unencrypted {
-            return Ok(())
+    fn encrypt_message(&mut self, buffer: &mut MsgBuffer) {
+        if let Some(core) = &mut self.core {
+            core.encrypt(buffer)
         }
-        self.get_core()?.encrypt(buffer);
-        Ok(())
     }
 
     fn decrypt_message(&mut self, buffer: &mut MsgBuffer) -> Result<(), Error> {
-        if self.unencrypted {
-            return Ok(())
+        if let Some(core) = &mut self.core {
+            core.decrypt(buffer)
+        } else {
+            Ok(())
         }
-        self.get_core()?.decrypt(buffer)
     }
 
-    pub fn handle_message(&mut self, buffer: &mut MsgBuffer) -> Result<MessageResult<P>, Error> {
+    pub fn handle_message(&mut self, buffer: &mut MsgBuffer) -> Result<MessageResult, Error> {
         if buffer.is_empty() {
             return Err(Error::InvalidCryptoState("No message in buffer"))
         }
         if is_init_message(buffer.buffer()) {
             debug!("Received init message");
-            buffer.take_prefix();
             self.handle_init_message(buffer)
         } else {
             debug!("Received encrypted message");
@@ -376,52 +281,37 @@ impl<P: Payload> PeerCrypto<P> {
         }
     }
 
-    pub fn send_message(&mut self, type_: u8, buffer: &mut MsgBuffer) -> Result<(), Error> {
+    pub fn send_message(&mut self, type_: u8, buffer: &mut MsgBuffer) {
         assert_ne!(type_, MESSAGE_TYPE_ROTATION);
         buffer.prepend_byte(type_);
-        self.encrypt_message(buffer)
+        self.encrypt_message(buffer);
     }
 
-    pub fn every_second(&mut self, out: &mut MsgBuffer) -> Result<MessageResult<P>, Error> {
+    pub fn every_second(&mut self, out: &mut MsgBuffer) -> MessageResult {
         out.clear();
         if let Some(ref mut core) = self.core {
             core.every_second()
         }
-        if let Some(ref mut init) = self.init {
-            init.every_second(out)?;
-        }
-        if self.init.as_ref().map(|i| i.stage()).unwrap_or(CLOSING) == CLOSING {
-            self.init = None
-        }
-        if !out.is_empty() {
-            out.prepend_byte(INIT_MESSAGE_FIRST_BYTE);
-            return Ok(MessageResult::Reply)
-        }
         if let Some(ref mut rotate) = self.rotation {
             self.rotate_counter += 1;
             if self.rotate_counter >= ROTATE_INTERVAL {
                 self.rotate_counter = 0;
                 if let Some(rot) = rotate.cycle(out) {
-                    let core = self.get_core()?;
-                    let algo = core.algorithm();
+                    let algo = self.algorithm.unwrap();
                     let key = LessSafeKey::new(UnboundKey::new(algo, &rot.key[..algo.key_len()]).unwrap());
-                    core.rotate_key(key, rot.id, rot.use_for_sending);
+                    self.core.unwrap().rotate_key(key, rot.id, rot.use_for_sending);
                 }
                 if !out.is_empty() {
                     out.prepend_byte(MESSAGE_TYPE_ROTATION);
-                    self.encrypt_message(out)?;
-                    return Ok(MessageResult::Reply)
+                    self.encrypt_message(out);
+                    return MessageResult::Reply
                 }
             }
         }
-        Ok(MessageResult::None)
+        MessageResult::None
     }
 }
 
-pub fn is_init_message(msg: &[u8]) -> bool {
-    !msg.is_empty() && msg[0] == INIT_MESSAGE_FIRST_BYTE
-}
-
 
 #[cfg(test)]
 mod tests {

src/engine/shared.rs

@@ -12,7 +12,7 @@ use parking_lot::Mutex;
 use std::{collections::HashMap, net::SocketAddr, sync::Arc};
 
 pub struct SharedPeerCrypto {
-    peers: Arc<Mutex<HashMap<SocketAddr, CryptoCore, Hash>>>
+    peers: Arc<Mutex<HashMap<SocketAddr, Arc<CryptoCore>, Hash>>>
 }
 
 impl SharedPeerCrypto {