enigma_protocol/
crypto.rs

1use hkdf::Hkdf;
2use sha2::{Digest, Sha256};
3use uuid::Uuid;
4use zeroize::Zeroize;
5
6use crate::error::{EnigmaProtocolError, Result};
7use crate::types::{InitiatorOrResponder, SessionBootstrap};
8
9const INFO_SEED: &[u8] = b"enigma-protocol-seed";
10const INFO_NEXT: &[u8] = b"enigma-protocol-next";
11const INFO_MSG: &[u8] = b"enigma-protocol-msg";
12
13pub fn conversation_ad_bytes(conversation_id: &Uuid, sender: &str, receiver: &str) -> Vec<u8> {
14    let mut out = Vec::with_capacity(16 + sender.len() + receiver.len());
15    out.extend_from_slice(conversation_id.as_bytes());
16    out.extend_from_slice(sender.as_bytes());
17    out.extend_from_slice(receiver.as_bytes());
18    out
19}
20
21pub fn derive_conversation_id(secret: &[u8; 32], local: &str, remote: &str) -> Uuid {
22    let mut hasher = Sha256::new();
23    let (first, second) = if local <= remote {
24        (local, remote)
25    } else {
26        (remote, local)
27    };
28    hasher.update(secret);
29    hasher.update(first.as_bytes());
30    hasher.update(second.as_bytes());
31    let digest = hasher.finalize();
32    let mut bytes = [0u8; 16];
33    bytes.copy_from_slice(&digest[..16]);
34    Uuid::from_bytes(bytes)
35}
36
37pub struct SessionRatchet {
38    conversation_id: Uuid,
39    send_chain: [u8; 32],
40    recv_chain: [u8; 32],
41    send_counter: u64,
42    recv_counter: u64,
43}
44
45impl SessionRatchet {
46    pub fn from_bootstrap(bootstrap: &SessionBootstrap, local: &str, remote: &str) -> Result<Self> {
47        match bootstrap {
48            SessionBootstrap::PreSharedSecret {
49                secret32,
50                role,
51                remote_dh_pub,
52            } => {
53                let salt = derive_salt(secret32, local, remote, remote_dh_pub.as_ref());
54                let (send_chain, recv_chain) = derive_initial_chains(secret32, &salt, *role)?;
55                let conversation_id = derive_conversation_id(secret32, local, remote);
56                Ok(Self {
57                    conversation_id,
58                    send_chain,
59                    recv_chain,
60                    send_counter: 0,
61                    recv_counter: 0,
62                })
63            }
64        }
65    }
66
67    pub fn conversation_id(&self) -> Uuid {
68        self.conversation_id
69    }
70
71    pub fn next_send_key(&mut self) -> Result<[u8; 32]> {
72        let key = next_key(&mut self.send_chain, self.send_counter)?;
73        self.send_counter = self
74            .send_counter
75            .checked_add(1)
76            .ok_or(EnigmaProtocolError::InvalidState)?;
77        Ok(key)
78    }
79
80    pub fn next_recv_key(&mut self) -> Result<[u8; 32]> {
81        let key = next_key(&mut self.recv_chain, self.recv_counter)?;
82        self.recv_counter = self
83            .recv_counter
84            .checked_add(1)
85            .ok_or(EnigmaProtocolError::InvalidState)?;
86        Ok(key)
87    }
88}
89
90impl Drop for SessionRatchet {
91    fn drop(&mut self) {
92        self.send_chain.zeroize();
93        self.recv_chain.zeroize();
94    }
95}
96
97fn derive_salt(
98    secret: &[u8; 32],
99    local: &str,
100    remote: &str,
101    remote_dh: Option<&[u8; 32]>,
102) -> [u8; 32] {
103    let mut hasher = Sha256::new();
104    let (first, second) = if local <= remote {
105        (local, remote)
106    } else {
107        (remote, local)
108    };
109    hasher.update(secret);
110    hasher.update(first.as_bytes());
111    hasher.update(second.as_bytes());
112    if let Some(key) = remote_dh {
113        hasher.update(key);
114    }
115    let digest = hasher.finalize();
116    let mut salt = [0u8; 32];
117    salt.copy_from_slice(&digest);
118    salt
119}
120
121fn derive_initial_chains(
122    seed: &[u8; 32],
123    salt: &[u8; 32],
124    role: InitiatorOrResponder,
125) -> Result<([u8; 32], [u8; 32])> {
126    let hk = Hkdf::<Sha256>::new(Some(salt), seed);
127    let mut okm = [0u8; 64];
128    hk.expand(INFO_SEED, &mut okm)
129        .map_err(|_| EnigmaProtocolError::Crypto)?;
130    let mut first = [0u8; 32];
131    first.copy_from_slice(&okm[..32]);
132    let mut second = [0u8; 32];
133    second.copy_from_slice(&okm[32..]);
134    match role {
135        InitiatorOrResponder::Initiator => Ok((first, second)),
136        InitiatorOrResponder::Responder => Ok((second, first)),
137    }
138}
139
140fn next_key(chain: &mut [u8; 32], counter: u64) -> Result<[u8; 32]> {
141    let mut counter_bytes = [0u8; 8];
142    counter_bytes.copy_from_slice(&counter.to_be_bytes());
143    let hk = Hkdf::<Sha256>::new(Some(chain), &counter_bytes);
144    let mut message_key = [0u8; 32];
145    hk.expand(INFO_MSG, &mut message_key)
146        .map_err(|_| EnigmaProtocolError::Crypto)?;
147    let mut next_chain = [0u8; 32];
148    hk.expand(INFO_NEXT, &mut next_chain)
149        .map_err(|_| EnigmaProtocolError::Crypto)?;
150    chain.copy_from_slice(&next_chain);
151    Ok(message_key)
152}