use crate::{
s3::RequestOptions,
s3::S3,
stream::{
InitialStreamParams, STDIN_BUFFER_SIZE, Stream, complete_multipart_upload,
create_initial_stream, create_nonce_header, encrypt_chunk, get_key, init_encryption,
initiate_multipart_upload, maybe_upload_part, setup_stream_progress, upload_final_part,
write_to_stream,
},
};
use anyhow::{Result, anyhow};
use chacha20poly1305::aead::stream::EncryptorBE32;
use futures::stream::TryStreamExt;
use std::{collections::BTreeMap, path::PathBuf};
use tokio::io::stdin;
use tokio_util::codec::{BytesCodec, FramedRead};
pub async fn stream_stdin_encrypted(
s3: &S3,
object_key: &str,
acl: Option<String>,
meta: Option<BTreeMap<String, String>>,
quiet: bool,
tmp_dir: PathBuf,
globals: RequestOptions,
) -> Result<String> {
let encryption_key = globals
.enc_key
.as_ref()
.ok_or_else(|| anyhow!("Encryption key is required"))?;
let key = get_key(object_key, globals.compress, globals.encrypt);
let mut meta = meta.unwrap_or_default();
meta.insert(
"Content-Type".to_string(),
"application/vnd.s3m.encrypted".to_string(),
);
let upload_id = initiate_multipart_upload(s3, &key, acl, meta).await?;
let progress_sender = setup_stream_progress(quiet).await;
let (cipher, nonce_bytes) = init_encryption(encryption_key);
let encryptor = EncryptorBE32::from_aead(cipher, (&nonce_bytes).into());
let nonce_header = create_nonce_header(&nonce_bytes);
let first_stream: Stream = create_initial_stream(InitialStreamParams {
upload_id: &upload_id,
tmp_dir: &tmp_dir,
key: &key,
s3,
progress_sender,
globals: &globals,
header_data: Some(&nonce_header),
})?;
let mut stream = FramedRead::new(stdin(), BytesCodec::new())
.map_err(|e| anyhow!("Error reading STDIN chunk: {e}"))
.try_fold(
(first_stream, encryptor),
|(mut current_upload_state_acc, mut current_encryptor_acc), chunk| async move {
let encrypted_data = encrypt_chunk(&mut current_encryptor_acc, &chunk)
.map_err(|e| anyhow!("Failed to encrypt chunk: {e}"))?;
write_to_stream(&mut current_upload_state_acc, &encrypted_data).map_err(|e| {
anyhow!("Failed to write encrypted chunk to upload stream: {e}")
})?;
maybe_upload_part(&mut current_upload_state_acc, STDIN_BUFFER_SIZE).await?;
Ok((current_upload_state_acc, current_encryptor_acc))
},
)
.await
.map(|(final_stream_state, _)| final_stream_state)?;
let final_etag = upload_final_part(&mut stream, &key, &upload_id, s3, &globals).await?;
stream.etags.push(final_etag);
if let Some(sender) = stream.channel.take() {
drop(sender);
}
complete_multipart_upload(s3, &key, &upload_id, stream.etags).await
}