moli/moli-leveldb-parser/src/tests/table.rs

// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
// Ports of table/table_test.cc and db/dbformat_test.cc, plus corruption cases.
// See ../../TESTS.md, ../../LICENSE-LevelDB and ../../AUTHORS-LevelDB.

use std::fs;

use anyhow::Result;

use super::support::*;
use crate::{
    MAX_SEQUENCE,
    coding::Decoder,
    read_snapshot,
    table::{read_entries, read_table},
};

fn scan(bytes: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
    let directory = tempfile::tempdir()?;
    let path = directory.path().join("table.ldb");
    fs::write(&path, bytes)?;
    let mut result = Vec::new();
    read_table(&path, bytes.len() as u64, |key, value| {
        result.push((key.to_vec(), value.to_vec()));
        Ok(())
    })?;
    Ok(result)
}

fn check(entries: &Entries) -> Result<()> {
    let expected: Vec<_> = entries
        .iter()
        .map(|(k, v)| (k.clone(), v.clone()))
        .collect();
    for restart_interval in [1, 16, 1024] {
        let mut actual = Vec::new();
        read_entries(&block(&expected, restart_interval), |k, v| {
            actual.push((k.to_vec(), v.to_vec()));
            Ok(())
        })?;
        assert_eq!(actual, expected, "restart interval {restart_interval}");
        for compression in [0, 1] {
            let bytes = table(&expected, restart_interval, 13, compression);
            assert_eq!(
                scan(&bytes)?,
                expected,
                "restart interval {restart_interval}, compression {compression}"
            );
        }
    }
    Ok(())
}

#[test]
fn empty() -> Result<()> {
    check(&Entries::new())
}

#[test]
fn zero_restart_points_in_block() -> Result<()> {
    // Upstream explicitly accepts this four-byte Java LevelDB empty block.
    read_entries(&[0; 4], |_, _| panic!("zero-restart block must be empty"))
}

#[test]
fn simple_empty_key() -> Result<()> {
    check(&[(vec![], b"v".to_vec())].into())
}

#[test]
fn simple_single() -> Result<()> {
    check(&[(b"abc".to_vec(), b"v".to_vec())].into())
}

#[test]
fn simple_multi() -> Result<()> {
    check(
        &[
            (b"abc".to_vec(), b"v".to_vec()),
            (b"abcd".to_vec(), b"v".to_vec()),
            (b"ac".to_vec(), b"v2".to_vec()),
        ]
        .into(),
    )
}

#[test]
fn simple_special_key() -> Result<()> {
    check(&[(vec![0xff, 0xff], b"v3".to_vec())].into())
}

#[test]
fn randomized() -> Result<()> {
    // Same entry-count schedule, key alphabet and skewed lengths as Harness.
    let mut random = Random::new(306);
    for count in (0..50).chain((50..2000).step_by(200)) {
        let mut entries = Entries::new();
        for _ in 0..count {
            let key_length = random.skewed(4);
            let key = random.key(key_length);
            let value_length = random.skewed(5);
            entries.insert(key, random.string(value_length));
        }
        check(&entries)?;
    }
    Ok(())
}

#[test]
fn randomized_long_db() -> Result<()> {
    // Generated by Chromium's pinned C++ engine. The Rust fixture engine panics
    // in its shortest-separator routine for empty keys during this workload.
    let directory = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("tests/fixtures/chromium/randomized-long-db");
    let mut random = Random::new(301);
    let mut expected = Entries::new();
    for _ in 0..100_000 {
        let key_length = random.skewed(4);
        let key = random.key(key_length);
        let value_length = random.skewed(5);
        let value = random.string(value_length);
        expected.insert(key, value);
    }
    assert!(
        !crate::manifest::Manifest::read(&directory)?
            .tables
            .is_empty()
    );
    assert_eq!(read_snapshot(&directory)?, expected);
    Ok(())
}

#[test]
fn approximate_offset_of_plain_read_contents() -> Result<()> {
    // Keep upstream's large-value/block-boundary workload; this parser exposes
    // a full scan, so validate contents instead of a writer's offset estimate.
    let entries = vec![
        (b"k01".to_vec(), b"hello".to_vec()),
        (b"k02".to_vec(), b"hello2".to_vec()),
        (b"k03".to_vec(), vec![b'x'; 10_000]),
        (b"k04".to_vec(), vec![b'x'; 200_000]),
        (b"k05".to_vec(), vec![b'x'; 300_000]),
        (b"k06".to_vec(), b"hello3".to_vec()),
        (b"k07".to_vec(), vec![b'x'; 100_000]),
    ];
    assert_eq!(scan(&table(&entries, 16, 2, 0))?, entries);
    Ok(())
}

#[test]
fn approximate_offset_of_compressed_read_contents() -> Result<()> {
    let mut random = Random::new(301);
    let first = big_string(&random.string(2500), 10_000);
    let second = big_string(&random.string(2500), 10_000);
    let entries = vec![
        (b"k01".to_vec(), b"hello".to_vec()),
        (b"k02".to_vec(), first),
        (b"k03".to_vec(), b"hello3".to_vec()),
        (b"k04".to_vec(), second),
    ];
    let compressed = table(&entries, 16, 2, 1);
    assert!(compressed.len() < 10_000);
    assert_eq!(scan(&compressed)?, entries);
    Ok(())
}

#[test]
fn internal_key_encode_decode() -> Result<()> {
    let directory = tempfile::tempdir()?;
    let sequences = [
        1,
        2,
        3,
        (1 << 8) - 1,
        1 << 8,
        (1 << 8) + 1,
        (1 << 16) - 1,
        1 << 16,
        (1 << 16) + 1,
        (1 << 32) - 1,
        1 << 32,
        (1 << 32) + 1,
        MAX_SEQUENCE,
    ];
    for sequence in sequences {
        let mut edit = metadata(3, sequence);
        let entries: Vec<_> = [b"".as_slice(), b"k", b"hello", b"longggggggggggggggggggggg"]
            .into_iter()
            .map(|key| (internal_key(key, sequence, 1), b"value".to_vec()))
            .collect();
        // The data reader consumes the on-disk order; public output is bytewise.
        let mut entries = entries;
        entries.sort_by(|a, b| a.0[..a.0.len() - 8].cmp(&b.0[..b.0.len() - 8]));
        save_table(directory.path(), &mut edit, 0, 4, &entries)?;
        write_manifest(directory.path(), &[edit])?;
        let expected: Entries = [b"".as_slice(), b"k", b"hello", b"longggggggggggggggggggggg"]
            .into_iter()
            .map(|key| (key.to_vec(), b"value".to_vec()))
            .collect();
        assert_eq!(read_snapshot(directory.path())?, expected);
        let mut edit = metadata(3, sequence);
        save_table(
            directory.path(),
            &mut edit,
            0,
            4,
            &[(internal_key(b"hello", sequence, 0), vec![])],
        )?;
        write_manifest(directory.path(), &[edit])?;
        assert!(read_snapshot(directory.path())?.is_empty());
    }
    Ok(())
}

#[test]
fn internal_key_decode_from_empty() -> Result<()> {
    let directory = tempfile::tempdir()?;
    for length in 0..8 {
        let bytes = table(&[(vec![b'k'; length], b"value".to_vec())], 1, 1, 0);
        fs::write(directory.path().join("000004.ldb"), &bytes)?;
        let mut edit = metadata(3, 1);
        let key = internal_key(b"key", 1, 1);
        add_table(&mut edit, 0, 4, bytes.len() as u64, &key, &key);
        write_manifest(directory.path(), &[edit])?;
        assert!(
            format!("{:#}", read_snapshot(directory.path()).unwrap_err()).contains("internal key")
        );
    }
    Ok(())
}

#[test]
fn unknown_internal_value_type() -> Result<()> {
    let directory = tempfile::tempdir()?;
    let mut edit = metadata(3, 1);
    save_table(
        directory.path(),
        &mut edit,
        0,
        4,
        &[(internal_key(b"key", 1, 2), vec![])],
    )?;
    write_manifest(directory.path(), &[edit])?;
    assert!(format!("{:#}", read_snapshot(directory.path()).unwrap_err()).contains("value type"));
    Ok(())
}

#[test]
fn truncated_footer_and_bad_magic() {
    for length in 0..48 {
        assert!(scan(&[0; 48][..length]).is_err());
    }
    let mut bytes = table(&[], 1, 1, 0);
    let length = bytes.len();
    bytes[length - 1] ^= 1;
    assert!(scan(&bytes).unwrap_err().to_string().contains("magic"));
}

#[test]
fn manifest_table_size_must_match() -> Result<()> {
    let directory = tempfile::tempdir()?;
    let path = directory.path().join("table.ldb");
    let bytes = table(&[], 1, 1, 0);
    fs::write(&path, &bytes)?;
    for size in [0, bytes.len() as u64 - 1, bytes.len() as u64 + 1, u64::MAX] {
        assert!(read_table(&path, size, |_, _| Ok(())).is_err());
    }
    Ok(())
}

#[test]
fn index_handles_cannot_point_outside_the_table() {
    for (offset, size) in [(u64::MAX, 1), (0, u64::MAX), (10_000, 1), (0, 10_000)] {
        let mut bytes = table(&[], 1, 1, 0);
        let mut footer = vec![0, 0];
        varint(offset, &mut footer);
        varint(size, &mut footer);
        footer.resize(40, 0);
        let start = bytes.len() - 48;
        bytes[start..start + 40].copy_from_slice(&footer);
        assert!(scan(&bytes).unwrap_err().to_string().contains("outside"));
    }
}

#[test]
fn data_and_index_checksums() {
    let original = table(&[(b"key".to_vec(), b"value".to_vec())], 1, 1, 0);
    let mut footer = Decoder::new(&original[original.len() - 48..original.len() - 8]);
    footer.varint().unwrap();
    footer.varint().unwrap();
    let index_offset = footer.varint().unwrap() as usize;
    for offset in [0, index_offset] {
        let mut bytes = original.clone();
        bytes[offset] ^= 1;
        assert!(scan(&bytes).unwrap_err().to_string().contains("checksum"));
    }
}

#[test]
fn corrupt_compression_with_valid_checksum() {
    let entries = [(b"key".to_vec(), vec![b'x'; 100])];
    let original = table(&entries, 1, 1, 0);
    let size = block(&entries, 1).len();
    for compression in [1, 2, 255] {
        let mut bytes = original.clone();
        bytes[..size].fill(0xff);
        bytes[size] = compression;
        let checksum = crc(&bytes[..size + 1]);
        bytes[size + 1..size + 5].copy_from_slice(&checksum.to_le_bytes());
        assert!(scan(&bytes).is_err(), "compression {compression}");
    }
}

#[test]
fn truncated_and_excessive_restart_arrays() {
    for length in 0..4 {
        assert!(read_entries(&[0; 4][..length], |_, _| Ok(())).is_err());
    }
    for count in [2u32, u32::MAX] {
        let mut bytes = vec![0; 4];
        bytes.extend_from_slice(&count.to_le_bytes());
        assert!(read_entries(&bytes, |_, _| Ok(())).is_err());
    }
}

#[test]
fn restart_offsets_must_be_ordered_entry_boundaries() {
    let original = block(
        &[
            (b"abc".to_vec(), vec![1]),
            (b"abcd".to_vec(), vec![2]),
            (b"abd".to_vec(), vec![3]),
        ],
        1,
    );
    let restarts = original.len() - 16;
    for (index, offset) in [(0, 1u32), (1, 0), (1, 1), (1, u32::MAX), (2, 0)] {
        let mut bytes = original.clone();
        let start = restarts + index * 4;
        bytes[start..start + 4].copy_from_slice(&offset.to_le_bytes());
        assert!(
            read_entries(&bytes, |_, _| Ok(())).is_err(),
            "index {index}, offset {offset}"
        );
    }
}

#[test]
fn restart_entry_cannot_share_a_prefix() {
    let mut bytes = block(&[(b"abc".to_vec(), vec![]), (b"abcd".to_vec(), vec![])], 1);
    bytes[6] = 1;
    assert!(
        read_entries(&bytes, |_, _| Ok(()))
            .unwrap_err()
            .to_string()
            .contains("shared prefix")
    );
}

#[test]
fn entries_cannot_extend_into_the_restart_array() {
    let original = block(&[(b"key".to_vec(), b"value".to_vec())], 1);
    for (index, value) in [(0, 1), (1, 100), (2, 100)] {
        let mut bytes = original.clone();
        bytes[index] = value;
        assert!(read_entries(&bytes, |_, _| Ok(())).is_err());
    }
}

#[test]
fn table_visitor_error_is_propagated() -> Result<()> {
    let directory = tempfile::tempdir()?;
    let path = directory.path().join("table.ldb");
    let bytes = table(&[(b"a".to_vec(), vec![]), (b"b".to_vec(), vec![])], 1, 1, 0);
    fs::write(&path, &bytes)?;
    let mut calls = 0;
    let error = read_table(&path, bytes.len() as u64, |_, _| {
        calls += 1;
        anyhow::bail!("visitor failed")
    })
    .unwrap_err();
    assert_eq!(calls, 1);
    assert!(error.to_string().contains("visitor failed"));
    Ok(())
}