92 lines
2.6 KiB
Rust
92 lines
2.6 KiB
Rust
use std::io::SeekFrom;
|
|
|
|
use tokio::io::{AsyncSeek, AsyncSeekExt, AsyncWrite, AsyncWriteExt};
|
|
|
|
use crate::chacha::{ChaChaConfigv1, ChaChaHeaderv1};
|
|
|
|
pub struct ChaChaWriterAsync<W: AsyncWrite + AsyncSeek + Unpin + Send> {
|
|
inner: W,
|
|
header: ChaChaHeaderv1,
|
|
|
|
encryption_key: [u8; 32],
|
|
buffer: Vec<u8>,
|
|
plaintext_bytes_written: u64,
|
|
}
|
|
|
|
impl<W: AsyncWrite + AsyncSeek + Unpin + Send> ChaChaWriterAsync<W> {
|
|
pub async fn new(mut inner: W, encryption_key: [u8; 32]) -> Result<Self, std::io::Error> {
|
|
let header = ChaChaHeaderv1 {
|
|
config: ChaChaConfigv1::default(),
|
|
plaintext_size: 0,
|
|
};
|
|
inner.write_all(&serialize_header(header)?).await?;
|
|
|
|
Ok(Self {
|
|
inner,
|
|
header,
|
|
encryption_key,
|
|
buffer: Vec::new(),
|
|
plaintext_bytes_written: 0,
|
|
})
|
|
}
|
|
|
|
pub async fn write(&mut self, buf: &[u8]) -> Result<(), std::io::Error> {
|
|
self.buffer.extend_from_slice(buf);
|
|
self.plaintext_bytes_written += buf.len() as u64;
|
|
|
|
let chunk_size = self.header.config.chunk_size as usize;
|
|
while self.buffer.len() >= chunk_size {
|
|
let encrypted = encrypt_chunk(&self.encryption_key, &self.buffer[..chunk_size])?;
|
|
self.inner.write_all(&encrypted).await?;
|
|
self.buffer.drain(..chunk_size);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Encrypt and write any buffered plaintext, patch the header with the
|
|
/// final `plaintext_size`, then return the inner writer.
|
|
pub async fn finish(mut self) -> Result<W, std::io::Error> {
|
|
if !self.buffer.is_empty() {
|
|
let encrypted = encrypt_chunk(&self.encryption_key, &self.buffer)?;
|
|
self.inner.write_all(&encrypted).await?;
|
|
}
|
|
|
|
self.inner.seek(SeekFrom::Start(0)).await?;
|
|
let header_bytes = serialize_header(ChaChaHeaderv1 {
|
|
config: self.header.config,
|
|
plaintext_size: self.plaintext_bytes_written,
|
|
})?;
|
|
self.inner.write_all(&header_bytes).await?;
|
|
|
|
Ok(self.inner)
|
|
}
|
|
}
|
|
|
|
fn encrypt_chunk(key: &[u8; 32], plaintext: &[u8]) -> Result<Vec<u8>, std::io::Error> {
|
|
use chacha20poly1305::{
|
|
XChaCha20Poly1305,
|
|
aead::{Aead, AeadCore, KeyInit, OsRng},
|
|
};
|
|
|
|
let nonce = XChaCha20Poly1305::generate_nonce(&mut OsRng);
|
|
let cipher = XChaCha20Poly1305::new(chacha20poly1305::Key::from_slice(key));
|
|
let ciphertext = cipher
|
|
.encrypt(&nonce, plaintext)
|
|
.map_err(|_| std::io::Error::other("encryption failed"))?;
|
|
|
|
let mut output = Vec::with_capacity(nonce.len() + ciphertext.len());
|
|
output.extend_from_slice(&nonce);
|
|
output.extend_from_slice(&ciphertext);
|
|
Ok(output)
|
|
}
|
|
|
|
fn serialize_header(header: ChaChaHeaderv1) -> Result<Vec<u8>, std::io::Error> {
|
|
use binrw::BinWriterExt;
|
|
use std::io::Cursor;
|
|
|
|
let mut buf = Cursor::new(Vec::new());
|
|
buf.write_le(&header).map_err(std::io::Error::other)?;
|
|
Ok(buf.into_inner())
|
|
}
|