Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion crates/lance-context-wal/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ vectored-io = ["dep:bytes", "dep:tokio"]
runner = ["dep:anyhow", "dep:clap", "dep:libc", "dep:serde", "dep:serde_json", "dep:sha2"]
contract = ["dep:serde", "dep:sha2"]
arrow-contract = ["contract", "dep:arrow-schema", "arrow-schema/serde", "dep:serde_json"]
record = ["dep:bytes", "bytes/serde", "dep:serde", "dep:sha2", "dep:uuid"]
record = ["dep:bytes", "bytes/serde", "dep:serde", "dep:ring", "dep:uuid"]
placement = ["dep:anyhow", "dep:serde", "dep:serde_json", "dep:sha2", "dep:uuid"]
admission = ["dep:anyhow", "dep:serde"]
ipc = ["dep:anyhow", "dep:arrow-array", "dep:arrow-buffer", "dep:arrow-ipc", "dep:arrow-schema", "dep:bytes"]
Expand Down Expand Up @@ -52,6 +52,7 @@ sha2 = { version = "0.10", optional = true }
uuid = { version = "1", features = ["serde"], optional = true }

[dev-dependencies]
sha2 = "0.10"
libc = "0.2"
slog = "2"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "io-util"] }
Expand Down
14 changes: 11 additions & 3 deletions crates/lance-context-wal/src/record.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,17 @@

use bytes::Bytes;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use uuid::Uuid;

// Keep SHA-256 and its persisted bytes unchanged. Ring also has optimized
// implementations for CPUs that do not expose dedicated SHA instructions.
fn payload_digest(body: &[u8]) -> [u8; 32] {
let digest = ring::digest::digest(&ring::digest::SHA256, body);
let mut bytes = [0; 32];
bytes.copy_from_slice(digest.as_ref());
bytes
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BatchMeta {
#[serde(default)]
Expand Down Expand Up @@ -102,7 +110,7 @@ impl BatchMeta {
/// an envelope. Length/admission checks run before the digest scan.
pub fn validate_payload(&self, body: &[u8], limit: usize) -> Result<(), ValidationError> {
self.check_body_length(body.len(), limit)?;
if self.digest != <[u8; 32]>::from(Sha256::digest(body)) {
if self.digest != payload_digest(body) {
return Err(ValidationError::DigestMismatch);
}
Ok(())
Expand All @@ -113,7 +121,7 @@ impl Batch {
/// Retain the caller's allocation and compute its digest once. Routing and
/// sequence admission still belong to the caller/replicated state machine.
pub fn new(stream: Uuid, bucket: u32, producer: Uuid, sequence: u64, body: Bytes) -> Self {
let digest = Sha256::digest(&body).into();
let digest = payload_digest(&body);
let body_bytes = body.len() as u64;
Self {
meta: BatchMeta {
Expand Down
25 changes: 25 additions & 0 deletions crates/lance-context-wal/src/record/tests.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,30 @@
use super::*;

#[test]
fn payload_hash_backend_preserves_existing_digests_at_block_boundaries() {
use sha2::{Digest, Sha256};

let data: Vec<_> = (0..512 * 1024 + 3)
.map(|index| ((index * 29 + index / 251) % 256) as u8)
.collect();
for offset in [0, 1, 3] {
for size in [0, 1, 55, 56, 63, 64, 65, 127, 128, 129, 4096, 512 * 1024] {
let body = &data[offset..offset + size];
let old_digest: [u8; 32] = Sha256::digest(body).into();
assert_eq!(payload_digest(body), old_digest);
if !body.is_empty() {
// A record written by the previous backend remains readable.
let mut batch = record(Bytes::copy_from_slice(body));
assert_eq!(batch.meta.digest, old_digest);
batch.meta.digest = old_digest;
batch.validate(size).unwrap();
batch.meta.digest[31] ^= 1;
assert_eq!(batch.validate(size), Err(ValidationError::DigestMismatch));
}
}
}
}

fn record(body: Bytes) -> Batch {
Batch::new(
Uuid::from_bytes([0x11; 16]),
Expand Down
Loading