ssec_core/
chaff.rs

1use rand_core::TryRng;
2use bytes::Bytes;
3use futures_core::Stream;
4use core::pin::Pin;
5use core::task::{Context, Poll};
6use crate::HEADER_LENGTH;
7
8enum ChaffState {
9	PreHeader,
10	Data,
11	Finished
12}
13
14pin_project_lite::pin_project! {
15	pub struct ChaffStream<RNG> {
16		rng: RNG,
17		state: ChaffState,
18		remaining_bytes: usize,
19		chunk_size: usize
20	}
21}
22
23const MIN_CHUNK_SIZE: usize = HEADER_LENGTH;
24
25#[derive(Debug)]
26pub struct ChunkSizeTooSmallError;
27
28impl std::fmt::Display for ChunkSizeTooSmallError {
29	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30		write!(f, "chunk size too small, must be at least {MIN_CHUNK_SIZE} bytes")
31	}
32}
33
34impl std::error::Error for ChunkSizeTooSmallError {}
35
36#[derive(Debug, Clone, Copy)]
37pub struct ChaffStreamArgs {
38	output_length: usize,
39	chunk_size: usize
40}
41
42impl ChaffStreamArgs {
43	/// The `output_length` parameter controls the size of the hypothetical chaff input file.
44	/// In other words, the length of the stream is the length of the headers plus `output_length`.
45	pub fn with_length(output_length: usize) -> Self {
46		Self {
47			output_length,
48			chunk_size: 2048
49		}
50	}
51
52	pub fn set_chunk_size(&mut self, chunk_size: usize) -> Result<(), ChunkSizeTooSmallError> {
53		if chunk_size < MIN_CHUNK_SIZE {
54			return Err(ChunkSizeTooSmallError);
55		}
56
57		self.chunk_size = chunk_size;
58		Ok(())
59	}
60}
61
62impl<RNG:TryRng> ChaffStream<RNG> {
63	pub fn new(args: ChaffStreamArgs, rng: RNG) -> Self {
64		debug_assert!(args.chunk_size >= MIN_CHUNK_SIZE);
65
66		Self {
67			rng,
68			state: ChaffState::PreHeader,
69			remaining_bytes: args.output_length,
70			chunk_size: args.chunk_size
71		}
72	}
73}
74
75impl<RNG: TryRng> Stream for ChaffStream<RNG> {
76	type Item = Result<Bytes, RNG::Error>;
77
78	fn poll_next(
79		self: Pin<&mut Self>,
80		_cx: &mut Context<'_>
81	) -> Poll<Option<Self::Item>> {
82		let this = self.project();
83
84		match this.state {
85			ChaffState::PreHeader => {
86				let output_len: usize = (HEADER_LENGTH + *this.remaining_bytes).min(*this.chunk_size);
87				debug_assert!(output_len >= HEADER_LENGTH);
88
89				let mut output = Vec::with_capacity(output_len);
90
91				output.extend_from_slice(b"SSEC");
92				output.push(0x01);
93				output.push(0x6e);
94				output.extend_from_slice(&vec![0u8; output_len - 6]);
95
96				if let Err(err) = this.rng.try_fill_bytes(&mut output[6..]) {
97					*this.state = ChaffState::Finished;
98					return Poll::Ready(Some(Err(err)));
99				};
100
101				let bytes_left = *this.remaining_bytes - (output.len() - HEADER_LENGTH);
102				if bytes_left == 0 {
103					*this.state = ChaffState::Finished
104				} else {
105					*this.state = ChaffState::Data;
106					*this.remaining_bytes = bytes_left;
107				}
108
109				Poll::Ready(Some(Ok(Bytes::from_owner(output))))
110			},
111			ChaffState::Data => {
112				let chaff_len: usize = (*this.chunk_size).min(*this.remaining_bytes);
113				let mut chaff_data = vec![0u8; chaff_len];
114
115				if let Err(err) = this.rng.try_fill_bytes(&mut chaff_data) {
116					*this.state = ChaffState::Finished;
117					return Poll::Ready(Some(Err(err)));
118				};
119
120				let bytes_left = *this.remaining_bytes - chaff_len;
121				if bytes_left == 0 {
122					*this.state = ChaffState::Finished
123				} else {
124					*this.remaining_bytes = bytes_left;
125				}
126
127				Poll::Ready(Some(Ok(Bytes::from_owner(chaff_data))))
128			},
129			ChaffState::Finished => Poll::Ready(None)
130		}
131	}
132}
ssec_core/chaff.rs

ssec_core/
chaff.rs
