Add an implementation of mutual message exchange
[rust-lightning] / mutual-message-exchange / src / lib.rs
1 //#![cfg_attr(not(test), no_std)]
2 #![cfg_attr(not(test), deny(missing_docs))]
3 #![forbid(unsafe_code)]
4
5 #![deny(rustdoc::broken_intra_doc_links)]
6 #![deny(rustdoc::private_intra_doc_links)]
7
8 //! A simple mutual-authentication protocol which allows two parties to maintain a set of public
9 //! keys which they're willing to exchange messages with and exchange a message with an extra
10 //! half-round-trip.
11 //!
12 //! The protocol contains one party wishing to send a message to another. The message recipient is
13 //! the `initiator` in the protocol, and speaks first. Most of the CPU cost is born by the message
14 //! sender.
15 //!
16 //! Both parties first create a [`TrustedSet`] listing the public keys which they are willing to
17 //! exchange messages with.
18 //!
19 //! In order to exchange a message, the message recipient calls [`get_init_bytes`] and sends the
20 //! resulting message bytes to the message sender. That message sender then uses
21 //! [`respond_with_message`] to determine if both sides are mutually in each others' [`TrustedSet`]
22 //! and encrypt the message if so. Finally, the initiator uses [`decode_msg`]
23 //!
24 //! If the message sender is in the initiator's trusted set and the message sender has the public
25 //! key for the initiator, the message sender will learn who the initiator is upon receipt of the
26 //! init message (without any response). The initiator will only learn who the message sender is
27 //! (and the message sender will only respond) if both sides are mutually-trusting.
28 //!
29 //! In any other case, neither party learns anything about the other, apart from a rough estimate
30 //! of the trusted set size of the initiator.
31
32 extern crate alloc;
33
34 use bitcoin_hashes::cmp::fixed_time_eq;
35
36 #[allow(dead_code)]
37 mod chacha20;
38 #[allow(dead_code)]
39 mod poly1305;
40 #[allow(dead_code)]
41 mod chacha20poly1305rfc;
42
43 use alloc::vec;
44 use alloc::vec::Vec;
45
46 use secp256k1::{SecretKey, PublicKey};
47 use secp256k1::ecdh::SharedSecret;
48
49 use bitcoin_hashes::{Hash, HashEngine};
50 use bitcoin_hashes::sha256::Hash as Sha256;
51
52 use chacha20::ChaCha20;
53 use chacha20poly1305rfc::ChaCha20Poly1305RFC;
54
55 /// The maximum number of trusted counterparties which is allowed to be in a single [`TrustedSet`].
56 pub const MAX_TRUSTED_KEYS: usize = 1024;
57
58 /// A `TrustedSet` stores the set of peers which we are willing to talk to.
59 pub struct TrustedSet {
60         trusted_ecdhs: Vec<[u8; 32]>,
61         state_key: [u8; 32],
62 }
63
64 impl TrustedSet {
65         /// Constructs a new [`TrustedSet`] given a list of trusted counterparties. The keys are not
66         /// stored, only ECDH results are.
67         ///
68         /// `trusted_counterparties` must not exceed [`MAX_TRUSTED_KEYS`] entries or the construction
69         /// will fail. In all other cases construction succeeds.
70         pub fn new(our_key: &SecretKey, trusted_counterparties: &[PublicKey]) -> Result<Self, ()> {
71                 if trusted_counterparties.len() > MAX_TRUSTED_KEYS { return Err(()) }
72
73                 let mut trusted_ecdhs = Vec::with_capacity(trusted_counterparties.len());
74                 for counterparty in trusted_counterparties.iter() {
75                         let mut ecdh_hash = Sha256::engine();
76                         ecdh_hash.input(b"Mutual Message Exchange ECDH Result");
77                         ecdh_hash.input(&SharedSecret::new(counterparty, &our_key).secret_bytes());
78                         trusted_ecdhs.push(Sha256::from_engine(ecdh_hash).to_byte_array());
79                 }
80                 let mut state_key_hash = Sha256::engine();
81                 state_key_hash.input(b"Mutual Private Auth State Key Generation");
82                 state_key_hash.input(&Sha256::hash(&our_key[..]).to_byte_array());
83                 Ok(Self {
84                         trusted_ecdhs,
85                         state_key: Sha256::from_engine(state_key_hash).to_byte_array(),
86                 })
87         }
88
89         fn get_cover_trusted_count(&self) -> usize {
90                 // In order to avoid giving away exactly how many keys we trust, we include some fake
91                 // entries in our message. To avoid too much overhead we only round the trusted set up a
92                 // bit.
93                 debug_assert!(self.trusted_ecdhs.len() <= MAX_TRUSTED_KEYS);
94                 match self.trusted_ecdhs.len() {
95                         0..=16 => 16,
96                         17..=32 => 32,
97                         33..=128 => 128,
98                         129..=512 => 512,
99                         _ => MAX_TRUSTED_KEYS,
100                 }
101         }
102 }
103
104 /// The per-trusted-peer length in the intitial bytes.
105 ///
106 /// Fixed by the protocol.
107 const PER_PEER_LEN: usize = 32 + 16;
108 /// The length of the repeated data we sent in the init bytes and expect to be repeated in the
109 /// response. This is floating in the protocol, but we fix it for ourselves.
110 const OUR_REPEATED_DATA_LEN: usize = 64 + 32 + 8 + 16;
111
112 /// In order to avoid a message recipient having any guess as to the size our trusted set is, we
113 /// shuffle the entries in our init deterministically, using the permutation calculated here.
114 fn get_idx_permutation(cover_trusted_set_len: usize, rng_seed: &[u8]) -> [u16; MAX_TRUSTED_KEYS] {
115         debug_assert!(cover_trusted_set_len <= MAX_TRUSTED_KEYS);
116         debug_assert!(MAX_TRUSTED_KEYS <= u16::MAX.into());
117         debug_assert_eq!(rng_seed.len(), 32);
118
119         let mut perm = [0; MAX_TRUSTED_KEYS];
120         for i in 0..MAX_TRUSTED_KEYS {
121                 perm[i] = i as u16;
122         }
123         let mut rng = ChaCha20::new(rng_seed, b"MPA PERM RNG");
124         for i in 0..cover_trusted_set_len {
125                 let mut pos;
126                 let max_pos = (cover_trusted_set_len - i) as u64;
127                 loop {
128                         let mut rand = [0; 8];
129                         rng.process_in_place(&mut rand);
130                         pos = u64::from_le_bytes(rand);
131                         if pos < u64::MAX / max_pos * max_pos{
132                                 pos %= max_pos;
133                                 break;
134                         }
135                 }
136                 perm.swap(i, pos as usize);
137         }
138         perm
139 }
140
141 /// Gets the initial bytes this initiator should send to a (potential) message sender.
142 ///
143 /// It requires 64 secure random bytes, a reference to a [`TrustedSet`], and a `salt` and `aad`.
144 ///
145 /// The `salt` should uniquely describe this protocol the protocol being built using this mutual
146 /// authentication handshake. The `aad` should describe the particular message type being sent
147 /// (which the sender expects).
148 pub fn get_init_bytes(secure_random_nonce: [u8; 64], trusted_set: &TrustedSet, salt: [u8; 8], aad: &[u8]) -> Vec<u8> {
149         let mut local_nonce = [0; 32];
150         local_nonce.copy_from_slice(&secure_random_nonce[..32]);
151
152         let mut chacha_salt = [0; 12];
153         chacha_salt[4..].copy_from_slice(&salt);
154
155         let mut rng = ChaCha20::new(&secure_random_nonce[32..], b"MPA INIT RNG");
156
157         // Init message format is
158         // 2 byte handshake count
159         // PER_PEER_LEN * handshake count:
160         //   32-byte encrypted initiator nonce
161         //   16-byte poly1305 tag
162         // 2 byte repeated data len
163         // repeated data len bytes of data to be repeated
164         // any further bytes uninterpreted (for extensibility)
165         //
166         // Our repeated data is a 40 byte IV (XOR'd into `state_key` and "NONCE KY" to form the ChaCha
167         // key and nonce to encrypt remaining bytes), followed by 64 bytes containing the
168         // ChaCha-encrypted `secure_random_nonce` and the 16 byte Poly1305 MAC tag for the same.
169
170         let vec_cnt = trusted_set.get_cover_trusted_count();
171         let repeated_data_offs = 2 + vec_cnt * PER_PEER_LEN;
172         let mut res = Vec::with_capacity(repeated_data_offs + 2 + OUR_REPEATED_DATA_LEN);
173         res.resize(2 + vec_cnt * PER_PEER_LEN + 2 + OUR_REPEATED_DATA_LEN, 0);
174         res[..2].copy_from_slice(&(vec_cnt as u16).to_be_bytes());
175
176         // First fill in the encrypted slots for our trusted peers and fill remaining slots with noise.
177         let idx_permutation = get_idx_permutation(vec_cnt, &secure_random_nonce[32..]);
178         for (pos, idx) in idx_permutation.iter().take(vec_cnt).enumerate() {
179                 let idx_slice = &mut res[2 + pos * PER_PEER_LEN..8 + (pos + 1) * PER_PEER_LEN];
180                 if *idx as usize >= trusted_set.trusted_ecdhs.len() {
181                         rng.process_in_place(idx_slice);
182                 } else {
183                         let ecdh = &trusted_set.trusted_ecdhs[*idx as usize];
184                         let mut cryptoor = ChaCha20Poly1305RFC::new(ecdh, &chacha_salt, aad);
185                         let (crypted, tag) = idx_slice.split_at_mut(32);
186                         cryptoor.encrypt(&local_nonce, crypted, tag);
187                 }
188         }
189
190         // Pick a random nonce to XOR into the ChaCha key and nonce. Note that we reuse the `state_key`
191         // here repeatedly, so need the IV here to make the ChaCha pad unique.
192         let mut repeated_data_key_mask = [0; 32 + 8];
193         rng.process_in_place(&mut repeated_data_key_mask);
194         let mut repeated_data_key = [0; 32];
195         for i in 0..32 {
196                 repeated_data_key[i] = trusted_set.state_key[i] ^ repeated_data_key_mask[i];
197         }
198         let mut repeated_data_nonce = [0; 12];
199         for i in 0..8 {
200                 repeated_data_nonce[4 + i] = b"NONCE KY"[i] ^ repeated_data_key_mask[32 + i];
201         }
202
203         res[repeated_data_offs..repeated_data_offs + 2]
204                 .copy_from_slice(&(OUR_REPEATED_DATA_LEN as u16).to_be_bytes());
205         res[repeated_data_offs + 2..repeated_data_offs + 2 + 32 + 8]
206                 .copy_from_slice(&repeated_data_key_mask);
207
208         let (crypted_nonce, tag) = res[repeated_data_offs + 2 + 32 + 8..].split_at_mut(64);
209         let mut state_store = ChaCha20Poly1305RFC::new(&repeated_data_key, &repeated_data_nonce, &[]);
210         state_store.encrypt(&secure_random_nonce, crypted_nonce, tag);
211         res
212 }
213
214 /// Decode the message our counterparty sent us. The `salt` and `aad` provided must match the one
215 /// set in [`get_init_bytes`] and the one used by the message sender in [`respond_with_message`].
216 ///
217 /// Returns both the message sent to us by the counterparty (if any) and a shared key which can be
218 /// used to en/decrypt future messages with the message-sender.
219 pub fn decode_msg(trusted_set: &TrustedSet, salt: [u8; 8], aad: &[u8], wire_msg: &[u8]) -> Result<(Vec<u8>, [u8; 32]), ()> {
220         // Message format is:
221         // 2 byte selected challenge index
222         // 32 + 16 byte encrypted + MAC'd nonce
223         // 2 byte repeated data len
224         // repeated data len bytes of repeated data
225         // 2 byte message length (not counting mac)
226         // message length of encrypted message data
227         // 16 byte poly1305 message MAC
228         //
229         // Our repeated data is a 40 byte IV (XOR'd into `state_key` and "NONCE KY" to form the ChaCha
230         // key and nonce to encrypt remaining bytes), followed by 64 bytes containing the
231         // ChaCha-encrypted `secure_random_nonce` and the 16 byte Poly1305 MAC tag for the same.
232         const REPEATED_DATA_OFFS: usize = 2 + 32 + 16;
233         const CONST_OVERHEAD: usize = 2 + 32 + 16 + 2 + 2 + 16;
234
235         let mut chacha_salt = [0; 12];
236         chacha_salt[4..].copy_from_slice(&salt);
237         for b in chacha_salt.iter_mut().skip(4) {
238                 *b ^= 0xff;
239         }
240         if wire_msg.len() < CONST_OVERHEAD { return Err(()); }
241
242         // Read our state storage (i.e. the "secure_random_nonce" parameter from `get_init_bytes`.
243         let mut secure_random_nonce = [0; 64];
244         let repeated_part_len;
245         let msg_offs: usize;
246         {
247                 let mut repeated_part_len_bytes = [0; 2];
248                 repeated_part_len_bytes.copy_from_slice(&wire_msg[REPEATED_DATA_OFFS..REPEATED_DATA_OFFS + 2]);
249                 repeated_part_len = u16::from_be_bytes(repeated_part_len_bytes);
250                 if repeated_part_len as usize != OUR_REPEATED_DATA_LEN { return Err(()); }
251                 if wire_msg.len() < CONST_OVERHEAD + OUR_REPEATED_DATA_LEN { return Err(()); }
252                 msg_offs = REPEATED_DATA_OFFS + 2 + OUR_REPEATED_DATA_LEN;
253
254                 let mut repeated_data_key = [0; 32];
255                 for i in 0..32 {
256                         repeated_data_key[i] = trusted_set.state_key[i] ^ wire_msg[REPEATED_DATA_OFFS + 2 + i];
257                 }
258                 let mut repeated_data_nonce = [0; 12];
259                 for i in 0..8 {
260                         repeated_data_nonce[4 + i] = b"NONCE KY"[i] ^ wire_msg[REPEATED_DATA_OFFS + 2 + 32 + i];
261                 }
262
263                 let mut state_store = ChaCha20Poly1305RFC::new(&repeated_data_key, &repeated_data_nonce, &[]);
264                 let ciphertext = &wire_msg[REPEATED_DATA_OFFS + 2 + 32 + 8..REPEATED_DATA_OFFS + 2 + OUR_REPEATED_DATA_LEN - 16];
265                 let mac = &wire_msg[REPEATED_DATA_OFFS + 2 + OUR_REPEATED_DATA_LEN - 16..REPEATED_DATA_OFFS + 2 + OUR_REPEATED_DATA_LEN];
266                 // The message sender (presumably) knows if they modified the repeated data, so there's no
267                 // need to be constant-time wrt failures here (and thus we also return early).
268                 state_store.variable_time_decrypt(ciphertext, &mut secure_random_nonce, mac)?;
269         }
270
271         let mut local_nonce = [0; 32];
272         local_nonce.copy_from_slice(&secure_random_nonce[..32]);
273
274         let mut remote_nonce = [0; 32];
275
276         // Decrypt and validate the remote nonce
277         {
278                 let mut peer_idx_bytes = [0; 2];
279                 peer_idx_bytes.copy_from_slice(&wire_msg[..2]);
280                 let peer_idx = u16::from_be_bytes(peer_idx_bytes);
281                 let idx_permutation = get_idx_permutation(trusted_set.get_cover_trusted_count(), &secure_random_nonce[32..]);
282
283                 // The message sender has already learned if they're in our trusted peer list and if we're
284                 // in theirs. Same goes for any third-party observers who would detect the same by the fact
285                 // that some response was made. Thus, there's no need to worry about timing differences
286                 // giving that away here - if the idx is bogus we can simply return and we can do a
287                 // variable time decryption (and early return if it fails).
288                 let ecdh_idx = idx_permutation.get(peer_idx as usize).ok_or(())?;
289                 let ecdh = trusted_set.trusted_ecdhs.get(*ecdh_idx as usize).ok_or(())?;
290
291                 let mut remote_nonce_key = Sha256::engine();
292                 remote_nonce_key.input(ecdh);
293                 remote_nonce_key.input(&local_nonce);
294
295                 let mut cryptoor = ChaCha20Poly1305RFC::new(&Sha256::from_engine(remote_nonce_key)[..], &chacha_salt, aad);
296                 cryptoor.variable_time_decrypt(&wire_msg[2..2+32], &mut remote_nonce, &wire_msg[2+32..2+32+16])?;
297         }
298
299         for b in chacha_salt.iter_mut().skip(4) {
300                 *b ^= 0x0f;
301         }
302
303         let mut msg_key = local_nonce;
304         for (out, remote) in msg_key.iter_mut().zip(remote_nonce.iter()) {
305                 *out ^= *remote;
306         }
307
308         let separated_msg_keys = ChaCha20::get_single_block(&msg_key, b"INLINE KEY STRCH");
309         let mut oob_msg_key = [0; 32];
310         oob_msg_key.copy_from_slice(&separated_msg_keys[32..]);
311
312         let mut msg_len_bytes = [0; 2];
313         msg_len_bytes.copy_from_slice(&wire_msg[msg_offs..msg_offs + 2]);
314         let msg_len = u16::from_be_bytes(msg_len_bytes);
315
316         let mut res = Vec::new();
317         if msg_len != 0 {
318                 if msg_len < 16 { return Err(()); }
319                 if wire_msg.len() < msg_offs + msg_len as usize { return Err(()); }
320                 res = vec![0; msg_len as usize - 16];
321                 let ciphertext = &wire_msg[msg_offs + 2..msg_offs + 2 + msg_len as usize - 16];
322                 let mac = &wire_msg[msg_offs + 2 + msg_len as usize - 16..msg_offs + 2 + msg_len as usize];
323
324                 let mut msg_cryptoor = ChaCha20Poly1305RFC::new(&separated_msg_keys[..32], &chacha_salt, aad);
325                 if msg_cryptoor.variable_time_decrypt(ciphertext, &mut res, mac).is_err() {
326                         return Err(());
327                 }
328         }
329
330         Ok((res, oob_msg_key))
331 }
332
333 /// Processes the initial message sent by the initiator and generates an encrypted response,
334 /// containing the given `msg`. Also returns a negotiated shared key which can be used to encrypt
335 /// further messages to the initiator.
336 ///
337 /// Requires a random 64 bytes, a [`TrustedSet`] of peers, the `peer_init` message sent to us by
338 /// the initiator (via [`get_init_bytes`]), and a `salt` and `aad` which match those used by the
339 /// initiator.
340 ///
341 /// The `salt` should uniquely describe this protocol the protocol being built using this mutual
342 /// authentication handshake. The `aad` should describe the particular message type being sent
343 /// (which the recipient expects).
344 pub fn respond_with_message(secure_random_nonce: [u8; 64], trusted_set: &TrustedSet, peer_init: &[u8], salt: &[u8; 8], aad: &[u8], msg: &[u8])
345 -> Result<(Vec<u8>, [u8; 32]), ()> {
346         // Init message format is
347         // 2 byte handshake count
348         // PER_PEER_LEN * handshake count:
349         //   32-byte encrypted initiator nonce
350         //   16-byte poly1305 tag
351         // 2 byte repeated data len
352         // repeated data len bytes of data to be repeated
353         // any further bytes uninterpreted (for extensibility)
354
355         if peer_init.len() < 4 { return Err(()); }
356
357         let mut handshake_count_bytes = [0; 2];
358         handshake_count_bytes.copy_from_slice(&peer_init[..2]);
359         let handshake_count = u16::from_be_bytes(handshake_count_bytes);
360         if peer_init.len() < 4 + handshake_count as usize * PER_PEER_LEN { return Err(()); }
361
362         let mut repeated_data_len_bytes = [0; 2];
363         repeated_data_len_bytes.copy_from_slice(&peer_init[2 + handshake_count as usize * PER_PEER_LEN..2 + handshake_count as usize * PER_PEER_LEN + 2]);
364         let repeated_data_len = u16::from_be_bytes(repeated_data_len_bytes);
365         if peer_init.len() < 4 + handshake_count as usize * PER_PEER_LEN + repeated_data_len as usize { return Err(()); }
366
367         let mut chacha_salt = [0; 12];
368         chacha_salt[4..].copy_from_slice(salt);
369
370         let mut local_nonce = [0; 32];
371         local_nonce.copy_from_slice(&secure_random_nonce[..32]);
372         let mut rng = ChaCha20::new(&secure_random_nonce[32..], b"MPA Key Salt");
373
374         let mut default_peer_bytes = [0; 8];
375         rng.process_in_place(&mut default_peer_bytes);
376         let mut peer_match_idx =
377                 (u64::from_le_bytes(default_peer_bytes) as usize) %
378                 core::cmp::max(trusted_set.trusted_ecdhs.len(), 1);
379         let mut remote_nonce = [0; 32];
380         rng.process_in_place(&mut remote_nonce);
381         let mut peer_ecdh = &remote_nonce;
382
383         'match_search: for (idx, peer_enc) in peer_init[2..2 + handshake_count as usize * PER_PEER_LEN].chunks(PER_PEER_LEN).enumerate() {
384                 for ecdh in trusted_set.trusted_ecdhs.iter() {
385                         let mut cryptoor = ChaCha20Poly1305RFC::new(ecdh, &chacha_salt, aad);
386
387                         let mut peer_nonce = [0; 32];
388                         // Because the sender (should have) randomized the order of their trusted-peers list,
389                         // the time taken to find a matching ECDH entry shouldn't give away who they were to a
390                         // third-party observer. Thus, variable-time decryption (and an early return) should be
391                         // fine.
392                         if cryptoor.variable_time_decrypt(&peer_enc[..32], &mut peer_nonce, &peer_enc[32..]).is_ok() {
393                                 peer_ecdh = ecdh;
394                                 remote_nonce = peer_nonce;
395                                 peer_match_idx = idx;
396                                 break 'match_search;
397                         }
398                 }
399         }
400
401         for b in chacha_salt.iter_mut().skip(4) {
402                 *b ^= 0xff;
403         }
404
405         // Message format is:
406         // 2 byte selected challenge index
407         // 32 + 16 byte encrypted + MAC'd nonce
408         // 2 byte repeated data len
409         // repeated data len bytes of repeated data
410         // 2 byte message length (not counting mac)
411         // message length of encrypted message data
412         // 16 byte poly1305 message MAC
413
414         let mut res = Vec::with_capacity(2 + 32 + 16 + 2 + repeated_data_len as usize + 2 + msg.len() + 16);
415         res.resize(2 + 32 + 16 + 2 + repeated_data_len as usize + 2 + msg.len() + 16, 0);
416         res[0..2].copy_from_slice(&(peer_match_idx as u16).to_be_bytes());
417         let mut res_write_pos = 2;
418
419         let mut noise = [0; 32];
420         rng.process_in_place(&mut noise);
421
422         let mut local_nonce_key = Sha256::engine();
423         local_nonce_key.input(peer_ecdh);
424         local_nonce_key.input(&remote_nonce);
425         {
426                 let mut cryptoor = ChaCha20Poly1305RFC::new(&Sha256::from_engine(local_nonce_key)[..], &chacha_salt, aad);
427                 let (crypted, tag) = res.split_at_mut(res_write_pos + 32);
428                 cryptoor.encrypt(&local_nonce, &mut crypted[res_write_pos..], &mut tag[..16]);
429                 res_write_pos += 32 + 16;
430         }
431
432         res[res_write_pos..res_write_pos + 2].copy_from_slice(&repeated_data_len_bytes);
433         res_write_pos += 2;
434         res[res_write_pos..res_write_pos + repeated_data_len as usize]
435                 .copy_from_slice(&peer_init[2 + handshake_count as usize * PER_PEER_LEN + 2..2 + handshake_count as usize * PER_PEER_LEN + 2 + repeated_data_len as usize]);
436         res_write_pos += repeated_data_len as usize;
437
438         for b in chacha_salt.iter_mut().skip(4) {
439                 *b ^= 0x0f;
440         }
441
442         let mut msg_key = local_nonce;
443         for (out, remote) in msg_key.iter_mut().zip(remote_nonce.iter()) {
444                 *out ^= *remote;
445         }
446
447         let separated_msg_keys = ChaCha20::get_single_block(&msg_key, b"INLINE KEY STRCH");
448         let mut oob_msg_key = [0; 32];
449         oob_msg_key.copy_from_slice(&separated_msg_keys[32..]);
450
451         let proto_msg_len = if msg.is_empty() { 0 } else { msg.len() as u16 + 16 };
452         res[res_write_pos..res_write_pos + 2].copy_from_slice(&proto_msg_len.to_be_bytes());
453         res_write_pos += 2;
454
455         let mut msg_cryptoor = ChaCha20Poly1305RFC::new(&separated_msg_keys[0..32], &chacha_salt, aad);
456         let (crypted, tag) = res.split_at_mut(res_write_pos + msg.len());
457         debug_assert_eq!(tag.len(), 16);
458         msg_cryptoor.encrypt(msg, &mut crypted[res_write_pos..], tag);
459         res_write_pos += msg.len() + 16;
460         debug_assert_eq!(res_write_pos, res.len());
461
462         Ok((res, oob_msg_key))
463 }
464
465 #[cfg(test)]
466 mod tests {
467         use super::*;
468
469         use secp256k1::{SecretKey, PublicKey, Secp256k1};
470
471         use std::hash::{BuildHasher, Hasher};
472
473         fn rand_bytes() -> [u8; 32] {
474                 let random_number = std::collections::hash_map::RandomState::new().build_hasher().finish();
475                 [random_number as u8; 32]
476         }
477         fn rand_64_bytes() -> [u8; 64] {
478                 let mut res = [0; 64];
479                 res[..32].copy_from_slice(&rand_bytes());
480                 res[32..].copy_from_slice(&rand_bytes());
481                 res
482         }
483
484         #[test]
485         fn simple_test() {
486                 let secp_ctx = Secp256k1::new();
487
488                 let initiator_key = SecretKey::from_slice(&rand_bytes()).unwrap();
489                 let initiator_pk = PublicKey::from_secret_key(&secp_ctx, &initiator_key);
490                 let receiver_key = SecretKey::from_slice(&rand_bytes()).unwrap();
491                 let receiver_pk = PublicKey::from_secret_key(&secp_ctx, &receiver_key);
492
493                 let initiator_state = TrustedSet::new(&initiator_key, &[receiver_pk]).unwrap();
494                 let init_msg = get_init_bytes(rand_64_bytes(), &initiator_state, *b"SALTSALT", b"42");
495
496                 let receiver_state = TrustedSet::new(&receiver_key, &[initiator_pk]).unwrap();
497                 let sent_msg = b"Hello Initiator!";
498                 let (receiver_wire, receiver_shared_key) =
499                         respond_with_message(rand_64_bytes(), &receiver_state, &init_msg, b"SALTSALT", b"42", sent_msg).unwrap();
500
501                 let (initiator_msg, initiator_shared_key) =
502                         decode_msg(&initiator_state, *b"SALTSALT", b"42", &receiver_wire).unwrap();
503                 assert_eq!(&initiator_msg[..], sent_msg);
504                 assert_eq!(receiver_shared_key, initiator_shared_key);
505         }
506 }