Skip to content

Commit

Permalink
save_kdbx4: Fix write corruption for large Databases (#218)
Browse files Browse the repository at this point in the history
* chore: add test to determine regressions for large file saves

* fix: remove block size chunking in hmac_block_stream
  • Loading branch information
baranyildirim authored Apr 28, 2024
1 parent b50d395 commit 3e99ccc
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 3 deletions.
4 changes: 1 addition & 3 deletions src/hmac_block_stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,8 +6,6 @@ use crate::error::{BlockStreamError, CryptographyError};

pub const HMAC_KEY_END: [u8; 1] = hex!("01");

const DEFAULT_BLOCK_SIZE: u32 = 1024 * 1024;

/// Read from a HMAC block stream into a raw buffer
pub(crate) fn read_hmac_block_stream(
data: &[u8],
Expand Down Expand Up @@ -62,7 +60,7 @@ pub(crate) fn write_hmac_block_stream(
let mut block_index = 0;

while pos < data.len() {
let size = std::cmp::min(DEFAULT_BLOCK_SIZE as usize, data.len() - pos);
let size = data.len() - pos;

let block = &data[pos..(pos + size)];

Expand Down
75 changes: 75 additions & 0 deletions tests/large_database_roundtrip_tests.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
mod large_file_roundtrip_tests {
use std::fs::File;

use keepass::{
db::{Database, Entry, NodeRef, Value},
DatabaseKey,
};

/// This can be tuned based on how "large" we expect databases to realistically be.
const LARGE_DATABASE_ENTRY_COUNT: usize = 100000;

/// Constants for the test database.
const TEST_DATABASE_FILE_NAME: &str = "demo.kdbx";
const TEST_DATABASE_PASSWORD: &str = "demopass";

/// Writing and reading back a large databack should function as expected.
/// This tests guards against issues that might affect large databases.
#[test]
fn write_and_read_large_database() -> Result<(), Box<dyn std::error::Error>> {
let mut db = Database::new(Default::default());

db.meta.database_name = Some("Demo database".to_string());

for i in 0..LARGE_DATABASE_ENTRY_COUNT {
let mut entry = Entry::new();
entry
.fields
.insert("Title".to_string(), Value::Unprotected(format!("Entry_{i}")));
entry.fields.insert(
"UserName".to_string(),
Value::Unprotected(format!("UserName_{i}")),
);
entry.fields.insert(
"Password".to_string(),
Value::Protected(format!("Password_{i}").as_bytes().into()),
);
db.root.add_child(entry);
}

// Define database key.
let key = DatabaseKey::new().with_password(TEST_DATABASE_PASSWORD);
#[cfg(feature = "save_kdbx4")]
{
db.save(&mut File::create(TEST_DATABASE_FILE_NAME)?, key.clone())?;
}
// Read the database that was written in the previous block.
let db = Database::open(&mut File::open(TEST_DATABASE_FILE_NAME)?, key)?;
// Validate that the data is what we expect.
let mut entry_counter = 0;
for node in &db.root {
match node {
NodeRef::Group(g) => {
println!("Saw group '{0}'", g.name);
}
NodeRef::Entry(e) => {
assert_eq!(
format!("Entry_{entry_counter}"),
e.get_title().expect("Title should be defined")
);
assert_eq!(
format!("UserName_{entry_counter}"),
e.get_username().expect("Username should be defined")
);
assert_eq!(
format!("Password_{entry_counter}"),
e.get_password().expect("Password should be defined")
);
entry_counter += 1;
}
}
}
assert_eq!(entry_counter, LARGE_DATABASE_ENTRY_COUNT);
Ok(())
}
}

0 comments on commit 3e99ccc

Please sign in to comment.