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	/// It's very important that you provide the correct length of the `read` stream otherwise you'll get a corrupted stream.
36	///
37	/// SECURITY: It is advisable to zero out the memory containing the password after this method returns.
38	pub fn new_uncompressed<RNG: TryCryptoRng>(read: R, password: &[u8], rng: &mut RNG) -> Result<Self, RNG::Error> {
39		let mut password_salt = new_arr::<32>();
40		rng.try_fill_bytes(password_salt.as_mut())?;
41
42		let aes_key = kdf(password, password_salt.as_ref());
43		let password_verification_hash = compute_verification_hash(&aes_key);
44
45		let mut iv = [0; 16];
46		rng.try_fill_bytes(&mut iv)?;
47
48		let aes = Aes256Ctr::new(aes_key.as_ref().get_ref().into(), (&iv).into());
49
50		Ok(Self {
51			read,
52			aes,
53			password_salt,
54			password_verification_hash,
55			integrity_code: Some(HmacSha3_256::new_from_slice(aes_key.as_ref().get_ref()).unwrap()),
56			state: EncryptState::PreHeader,
57			block_buffer: BytesMut::new(),
58			iv
59		})
60	}
61}
62
63impl<E, R: Stream<Item = Result<Bytes, E>>> Stream for Encrypt<R> {
64	type Item = Result<Bytes, E>;
65
66	fn poll_next(
67		self: Pin<&mut Self>,
68		cx: &mut Context<'_>
69	) -> Poll<Option<Self::Item>> {
70		let mut this = self.project();
71
72		loop {
73			match this.state {
74				EncryptState::PreHeader => {
75					let mut buf = Vec::with_capacity(
76						4 + // magic
77						1 + // version number
78						1 + // compression algo
79						32 + // password salt
80						64 + // password verification hash
81						16 // IV
82					);
83
84					buf.extend_from_slice(b"SSEC");
85					buf.push(0x01);
86					buf.push(0x6e);
87					buf.extend_from_slice(this.password_salt.as_ref());
88					buf.extend_from_slice(this.password_verification_hash.as_ref());
89					buf.extend_from_slice(this.iv.as_ref());
90
91					// as per spec: first we add the version byte, compression algo, then iv before the data
92					let integrity_code = this.integrity_code.as_mut().unwrap();
93					integrity_code.update(&[0x01, 0x6e]);
94					integrity_code.update(this.iv.as_ref());
95
96					match this.read.poll_next(cx) {
97						Poll::Pending => *this.state = EncryptState::PostHeader,
98						Poll::Ready(None) => *this.state = EncryptState::Finalizing,
99						Poll::Ready(Some(Err(e))) => {
100							*this.state = EncryptState::Finished;
101							return Poll::Ready(Some(Err(e)));
102						},
103						Poll::Ready(Some(Ok(bytes))) => {
104							*this.state = EncryptState::PostHeader;
105							this.block_buffer.put(bytes);
106						}
107					}
108
109					return Poll::Ready(Some(Ok(Bytes::from_owner(buf))));
110				},
111				EncryptState::PostHeader => {
112					if this.block_buffer.len() >= BYTES_PER_POLL {
113						let mut data = this.block_buffer.split_to(BYTES_PER_POLL);
114						this.aes.apply_keystream(&mut data);
115						this.integrity_code.as_mut().unwrap().update(&data);
116
117						return Poll::Ready(Some(Ok(data.freeze())));
118					} else {
119						match ready!(this.read.as_mut().poll_next(cx)) {
120							Some(Ok(bytes)) => {
121								this.block_buffer.put(bytes);
122								continue;
123							},
124							Some(Err(e)) => {
125								*this.state = EncryptState::Finished;
126								return Poll::Ready(Some(Err(e)));
127							},
128							None => {
129								*this.state = EncryptState::Finalizing;
130								continue;
131							}
132						}
133					}
134				},
135				EncryptState::Finalizing => {
136					debug_assert!(this.block_buffer.len() < BYTES_PER_POLL);
137
138					let mut final_data = this.block_buffer.split();
139
140					let mut hmac = this.integrity_code.take()
141						.expect("integrity_code only taken here");
142
143					this.aes.apply_keystream(&mut final_data);
144
145					hmac.update(&final_data);
146					final_data.put(Bytes::from_owner(hmac.finalize().into_bytes()));
147
148					*this.state = EncryptState::Finished;
149
150					return Poll::Ready(Some(Ok(final_data.freeze())));
151				},
152				EncryptState::Finished => return Poll::Ready(None)
153			}
154		}
155	}
156}
ssec_core/encrypt.rs

ssec_core/
encrypt.rs
