Skip to content

Commit

Permalink
[PM-8874] Set attachment size when encrypting (#846)
Browse files Browse the repository at this point in the history
Set the attachment size after we encrypt the contents
  • Loading branch information
dani-garcia authored Jun 14, 2024
1 parent 56d5001 commit 32ac1e4
Showing 1 changed file with 84 additions and 3 deletions.
87 changes: 84 additions & 3 deletions crates/bitwarden-vault/src/cipher/attachment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -73,13 +73,29 @@ impl<'a> KeyEncryptable<SymmetricCryptoKey, AttachmentEncryptResult> for Attachm
let encrypted_contents = self.contents.encrypt_with_key(&attachment_key)?;
attachment.key = Some(attachment_key.to_vec().encrypt_with_key(ciphers_key)?);

let contents = encrypted_contents.to_buffer()?;

// Once we have the encrypted contents, we can set the size of the attachment
attachment.size = Some(contents.len().to_string());
attachment.size_name = Some(size_name(contents.len()));

Ok(AttachmentEncryptResult {
attachment: attachment.encrypt_with_key(ciphers_key)?,
contents: encrypted_contents.to_buffer()?,
contents,
})
}
}

fn size_name(size: usize) -> String {
let units = ["Bytes", "KB", "MB", "GB", "TB"];
let size = size as f64;
let unit = (size.ln() / 1024_f64.ln()).floor() as usize;
let size = size / 1024_f64.powi(unit as i32);

let size_round = (size * 10.0_f64).round() as usize as f64 / 10.0_f64;
format!("{} {}", size_round, units[unit])
}

impl KeyDecryptable<SymmetricCryptoKey, Vec<u8>> for AttachmentFile {
fn decrypt_with_key(&self, key: &SymmetricCryptoKey) -> Result<Vec<u8>, CryptoError> {
let ciphers_key = Cipher::get_cipher_key(key, &self.cipher.key)?;
Expand Down Expand Up @@ -144,13 +160,78 @@ impl TryFrom<bitwarden_api_api::models::AttachmentResponseModel> for Attachment
#[cfg(test)]
mod tests {
use base64::{engine::general_purpose::STANDARD, Engine};
use bitwarden_crypto::{EncString, KeyDecryptable, SymmetricCryptoKey};
use bitwarden_crypto::{EncString, KeyDecryptable, KeyEncryptable, SymmetricCryptoKey};

use crate::{
cipher::cipher::{CipherRepromptType, CipherType},
Attachment, AttachmentFile, Cipher,
Attachment, AttachmentFile, AttachmentFileView, AttachmentView, Cipher,
};

#[test]
fn test_size_name_conversions() {
assert_eq!(super::size_name(0), "0 Bytes");
assert_eq!(super::size_name(19), "19 Bytes");
assert_eq!(super::size_name(1024), "1 KB");
assert_eq!(super::size_name(1570), "1.5 KB");
assert_eq!(super::size_name(1024 * 1024), "1 MB");
assert_eq!(super::size_name(1024 * 18999), "18.6 MB");
assert_eq!(super::size_name(1024 * 1024 * 1024), "1 GB");
assert_eq!(super::size_name(1024 * 1024 * 1024 * 1024), "1 TB");
}

#[test]
fn test_encrypt_attachment() {
let user_key: SymmetricCryptoKey = "w2LO+nwV4oxwswVYCxlOfRUseXfvU03VzvKQHrqeklPgiMZrspUe6sOBToCnDn9Ay0tuCBn8ykVVRb7PWhub2Q==".to_string().try_into().unwrap();

let attachment = AttachmentView {
id: None,
url: None,
size: Some("100".into()),
size_name: Some("100 Bytes".into()),
file_name: Some("Test.txt".into()),
key: None,
};

let contents = b"This is a test file that we will encrypt. It's 100 bytes long, the encrypted version will be longer!";

let attachment_file = AttachmentFileView {
cipher: Cipher {
id: None,
organization_id: None,
folder_id: None,
collection_ids: Vec::new(),
key: Some("2.Gg8yCM4IIgykCZyq0O4+cA==|GJLBtfvSJTDJh/F7X4cJPkzI6ccnzJm5DYl3yxOW2iUn7DgkkmzoOe61sUhC5dgVdV0kFqsZPcQ0yehlN1DDsFIFtrb4x7LwzJNIkMgxNyg=|1rGkGJ8zcM5o5D0aIIwAyLsjMLrPsP3EWm3CctBO3Fw=".parse().unwrap()),
name: "2.d24xECyEdMZ3MG9s6SrGNw==|XvJlTeu5KJ22M3jKosy6iw==|8xGiQty4X61cDMx6PVqkJfSQ0ZTdA/5L9TpG7QfovoM=".parse().unwrap(),
notes: None,
r#type: CipherType::Login,
login: None,
identity: None,
card: None,
secure_note: None,
favorite: false,
reprompt: CipherRepromptType::None,
organization_use_totp: false,
edit: true,
view_password: true,
local_data: None,
attachments: None,
fields: None,
password_history: None,
creation_date: "2023-07-24T12:05:09.466666700Z".parse().unwrap(),
deleted_date: None,
revision_date: "2023-07-27T19:28:05.240Z".parse().unwrap(),
},
attachment,
contents: contents.as_slice(),
};

let result = attachment_file.encrypt_with_key(&user_key).unwrap();

assert_eq!(result.contents.len(), 161);
assert_eq!(result.attachment.size, Some("161".into()));
assert_eq!(result.attachment.size_name, Some("161 Bytes".into()));
}

#[test]
fn test_attachment_key() {
let user_key: SymmetricCryptoKey = "w2LO+nwV4oxwswVYCxlOfRUseXfvU03VzvKQHrqeklPgiMZrspUe6sOBToCnDn9Ay0tuCBn8ykVVRb7PWhub2Q==".to_string().try_into().unwrap();
Expand Down

0 comments on commit 32ac1e4

Please sign in to comment.