fire_stream/packet/
encrypted_bytes.rs

1use super::{BodyBytes, BodyBytesMut, PacketBytes, PacketError};
2
3use crypto::cipher::{Key, Mac};
4
5use bytes::{Bytes, BytesMut, BytesRead, BytesWrite, BytesSeek};
6
7
8const OFFSET: usize = Mac::LEN;
9
10/// Encrypted bytes
11/// 
12/// Data Layout
13/// +-----------+-----------+
14/// |   Header  |   Body    |
15/// +-----+-----+-----+-----+
16/// | Mac | Ctn | Mac | Ctn |
17/// +-----+-----+-----+-----+
18///
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct EncryptedBytes {
21	bytes: Vec<u8>,
22	header_len: usize
23}
24
25impl PacketBytes for EncryptedBytes {
26	fn new(header_len: usize) -> Self {
27		Self {
28			// 1 offset for header
29			// 1 offset for body
30			bytes: vec![0; OFFSET * 2 + header_len],
31			header_len
32		}
33	}
34
35	fn header(&self) -> Bytes<'_> {
36		self.bytes[OFFSET..][..self.header_len].into()
37	}
38
39	fn header_mut(&mut self) -> BytesMut<'_> {
40		(&mut self.bytes[OFFSET..][..self.header_len]).into()
41	}
42
43	fn full_header_mut(&mut self) -> BytesMut<'_> {
44		(&mut self.bytes[..OFFSET + self.header_len]).into()
45	}
46
47	fn body(&self) -> BodyBytes<'_> {
48		BodyBytes::new(&self.bytes[OFFSET * 2..][self.header_len..])
49	}
50
51	fn body_mut(&mut self) -> BodyBytesMut<'_> {
52		BodyBytesMut::new(self.header_len + OFFSET * 2, &mut self.bytes)
53	}
54
55	fn full_body_mut(&mut self) -> BytesMut<'_> {
56		(&mut self.bytes[OFFSET + self.header_len..]).into()
57	}
58}
59
60impl EncryptedBytes {
61	pub(crate) fn has_body(&self) -> bool {
62		self.bytes.len() > OFFSET * 2 + self.header_len
63	}
64
65	// slice without body mac is returned
66	pub(crate) fn as_slice(&self) -> &[u8] {
67		if self.has_body() {
68			&*self.bytes
69		} else {
70			// if we don't have any body
71			// remove the space saved for it
72			&self.bytes[..self.bytes.len() - OFFSET]
73		}
74	}
75
76	pub(crate) fn encrypt(&mut self, key: &mut Key) {
77
78		// encrypt header
79		let mut header: BytesMut = self.full_header_mut().into();
80		header.seek(OFFSET);
81		let mac = key.encrypt(header.remaining_mut());
82		header.seek(0);
83		header.write(&mac.into_bytes());
84
85		if !self.has_body() {
86			return
87		}
88
89		// encrypt body
90		let mut body: BytesMut = self.full_body_mut().into();
91		body.seek(OFFSET);
92		let mac = key.encrypt(body.remaining_mut());
93		body.seek(0);
94		body.write(&mac.into_bytes());
95
96	}
97
98	pub(crate) fn decrypt_header(
99		&mut self,
100		key: &mut Key
101	) -> Result<(), PacketError> {
102		let mut header: BytesMut = self.full_header_mut().into();
103		let mac = Mac::from_slice(header.read(OFFSET));
104		key.decrypt(header.remaining_mut(), &mac)
105			.map_err(|_| PacketError::MacNotEqual)
106	}
107
108	pub(crate) fn decrypt_body(
109		&mut self,
110		key: &mut Key
111	) -> Result<(), PacketError> {
112		let mut body: BytesMut = self.full_body_mut().into();
113		let mac = Mac::from_slice(body.read(OFFSET));
114		key.decrypt(body.remaining_mut(), &mac)
115			.map_err(|_| PacketError::MacNotEqual)
116	}
117}
118
119#[cfg(test)]
120mod tests {
121
122	use super::*;
123	use crate::packet::bytes::tests::test_gen_msg;
124	use crypto::cipher::{Keypair, Nonce};
125
126	#[test]
127	fn encrypted_bytes() {
128		test_gen_msg::<EncryptedBytes>();
129	}
130
131	fn generate_key() -> Key {
132		let alice = Keypair::new();
133		let bob = Keypair::new();
134
135		let ssk = alice.diffie_hellman(bob.public());
136
137		let nonce = Nonce::from([1u8; 24]);
138		ssk.to_key(nonce)
139	}
140
141	#[test]
142	fn crypto() {
143
144		let header = [10u8; 30];
145		let mut bytes = EncryptedBytes::new(header.len());
146		bytes.header_mut().write(&header);
147
148		let body_buf = [20u8; 30];
149		let mut body = bytes.body_mut();
150		body.write(&body_buf);
151
152		let mut alice_key = generate_key();
153		let mut alice_key_2 = alice_key.dublicate();
154		let mut bob_key = alice_key.dublicate();
155
156		// should encrypt header
157		bytes.encrypt(&mut alice_key);
158
159		// validate header mac
160		let mut mut_header = header.clone();
161		let mac_r = alice_key_2.encrypt(&mut mut_header);
162		assert_eq!(mac_r.into_bytes(), bytes.full_header_mut().as_mut()[..16]);
163		assert_eq!(mut_header, bytes.header().as_slice());
164
165		// validate body mac
166		let mut mut_body = body_buf.clone();
167		let mac_r = alice_key_2.encrypt(&mut mut_body);
168		assert_eq!(mac_r.into_bytes()[..], bytes.full_body_mut().as_mut()[..16]);
169		assert_eq!(mut_body, bytes.body().as_slice());
170
171		// now decrypt everything
172		bytes.decrypt_header(&mut bob_key).unwrap();
173		bytes.decrypt_body(&mut bob_key).unwrap();
174
175		assert_eq!(bytes.header().as_slice(), header);
176		assert_eq!(bytes.body().as_slice(), body_buf);
177
178
179	}
180
181	#[test]
182	fn empty_msg() {
183
184		let header = [10u8; 30];
185		let mut bytes = EncryptedBytes::new(header.len());
186
187		assert!(!bytes.has_body());
188		assert_eq!(bytes.body().len(), 0);
189		// assert_eq!(bytes.as_slice().len(), 16 + 30);
190
191		bytes.body_mut().write_u8(10);
192		assert!(bytes.has_body());
193		assert_eq!(bytes.body().len(), 1);
194		// assert_eq!(bytes.as_slice().len(), 16 + 30 + 16 + 1);
195		assert_eq!(bytes.body().read_u8(), 10);
196
197		bytes.body_mut().resize(0);
198		assert!(!bytes.has_body());
199		assert_eq!(bytes.body().len(), 0);
200		// assert_eq!(bytes.as_slice().len(), 16 + 30);
201
202		bytes.body_mut().resize(1);
203		assert!(bytes.has_body());
204		assert_eq!(bytes.body().len(), 1);
205		// assert_eq!(bytes.as_slice().len(), 16 + 30 + 16 + 1);
206		assert_eq!(bytes.body().read_u8(), 0);
207
208	}
209
210}
fire_stream/packet/encrypted_bytes.rs

fire_stream/packet/
encrypted_bytes.rs
