ssec_core/
encrypt.rs

1use futures_core::Stream;
2use bytes::{Bytes, BytesMut, BufMut};
3use rand_core::TryCryptoRng;
4use ctr::cipher::{KeyIvInit, StreamCipher};
5use hmac::{Mac, KeyInit};
6use core::pin::Pin;
7use core::task::{Context, Poll, ready};
8use crate::util::{HmacSha3_256, new_arr, kdf, compute_verification_hash};
9use crate::{BYTES_PER_POLL, Aes256Ctr};
10
11enum EncryptState {
12	PreHeader,
13	PostHeader,
14	Finalizing,
15	Finished
16}
17
18pin_project_lite::pin_project! {
19	pub struct Encrypt<R> {
20		#[pin]
21		read: R,
22		aes: Aes256Ctr,
23		password_salt: Box<[u8; 32]>,
24		password_verification_hash: Box<[u8; 64]>,
25		integrity_code: Option<HmacSha3_256>,
26		state: EncryptState,
27		block_buffer: BytesMut,
28		iv: [u8; 16]
29	}
30}
31
32impl<R> Encrypt<R> {
33	/// This method is *very* blocking.
34	/// If you're using Tokio I advise that you wrap this call in a `spawn_blocking`.
35	///
36	/// SECURITY: It is advisable to zero out the memory containing the password after this method returns.
37	pub fn new_uncompressed<RNG: TryCryptoRng>(read: R, password: &[u8], rng: &mut RNG) -> Result<Self, RNG::Error> {
38		let mut password_salt = new_arr::<32>();
39		rng.try_fill_bytes(password_salt.as_mut())?;
40
41		let aes_key = kdf(password, password_salt.as_ref());
42		let password_verification_hash = compute_verification_hash(&aes_key);
43
44		let mut iv = [0; 16];
45		rng.try_fill_bytes(&mut iv)?;
46
47		let aes = Aes256Ctr::new(aes_key.as_ref().get_ref().into(), (&iv).into());
48
49		Ok(Self {
50			read,
51			aes,
52			password_salt,
53			password_verification_hash,
54			integrity_code: Some(HmacSha3_256::new_from_slice(aes_key.as_ref().get_ref()).unwrap()),
55			state: EncryptState::PreHeader,
56			block_buffer: BytesMut::new(),
57			iv
58		})
59	}
60}
61
62impl<E, R: Stream<Item = Result<Bytes, E>>> Stream for Encrypt<R> {
63	type Item = Result<Bytes, E>;
64
65	fn poll_next(
66		self: Pin<&mut Self>,
67		cx: &mut Context<'_>
68	) -> Poll<Option<Self::Item>> {
69		let mut this = self.project();
70
71		loop {
72			match this.state {
73				EncryptState::PreHeader => {
74					let mut buf = Vec::with_capacity(
75						4 + // magic
76						1 + // version number
77						1 + // compression algo
78						32 + // password salt
79						64 + // password verification hash
80						16 // IV
81					);
82
83					buf.extend_from_slice(b"SSEC");
84					buf.push(0x01);
85					buf.push(0x6e);
86					buf.extend_from_slice(this.password_salt.as_ref());
87					buf.extend_from_slice(this.password_verification_hash.as_ref());
88					buf.extend_from_slice(this.iv.as_ref());
89
90					// as per spec: first we add the version byte, compression algo, then iv before the data
91					let integrity_code = this.integrity_code.as_mut().unwrap();
92					integrity_code.update(&[0x01, 0x6e]);
93					integrity_code.update(this.iv.as_ref());
94
95					match this.read.poll_next(cx) {
96						Poll::Pending => *this.state = EncryptState::PostHeader,
97						Poll::Ready(None) => *this.state = EncryptState::Finalizing,
98						Poll::Ready(Some(Err(e))) => {
99							*this.state = EncryptState::Finished;
100							return Poll::Ready(Some(Err(e)));
101						},
102						Poll::Ready(Some(Ok(bytes))) => {
103							*this.state = EncryptState::PostHeader;
104							this.block_buffer.put(bytes);
105						}
106					}
107
108					return Poll::Ready(Some(Ok(Bytes::from_owner(buf))));
109				},
110				EncryptState::PostHeader => {
111					if this.block_buffer.len() >= BYTES_PER_POLL {
112						let mut data = this.block_buffer.split_to(BYTES_PER_POLL);
113						this.aes.apply_keystream(&mut data);
114						this.integrity_code.as_mut().unwrap().update(&data);
115
116						return Poll::Ready(Some(Ok(data.freeze())));
117					} else {
118						match ready!(this.read.as_mut().poll_next(cx)) {
119							Some(Ok(bytes)) => {
120								this.block_buffer.put(bytes);
121								continue;
122							},
123							Some(Err(e)) => {
124								*this.state = EncryptState::Finished;
125								return Poll::Ready(Some(Err(e)));
126							},
127							None => {
128								*this.state = EncryptState::Finalizing;
129								continue;
130							}
131						}
132					}
133				},
134				EncryptState::Finalizing => {
135					debug_assert!(this.block_buffer.len() < BYTES_PER_POLL);
136
137					let mut final_data = this.block_buffer.split();
138
139					let mut hmac = this.integrity_code.take()
140						.expect("integrity_code only taken here");
141
142					this.aes.apply_keystream(&mut final_data);
143
144					hmac.update(&final_data);
145					final_data.put(Bytes::from_owner(hmac.finalize().into_bytes()));
146
147					*this.state = EncryptState::Finished;
148
149					return Poll::Ready(Some(Ok(final_data.freeze())));
150				},
151				EncryptState::Finished => return Poll::Ready(None)
152			}
153		}
154	}
155}
ssec_core/encrypt.rs

ssec_core/
encrypt.rs
