<commit_before>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod common;

use anyhow::Result;
use common::{
    create_partition_entry, create_primary_header, create_secondary_header,
};
#[cfg(feature = "std")]
use gpt_disk_io::StdBlockIo;
use gpt_disk_io::{BlockIo, Disk, DiskError, MutSliceBlockIo};
use gpt_disk_types::{BlockSize, GptPartitionEntryArray};
use std::io::{Cursor, Read};

fn load_test_disk() -> Vec<u8> {
    // Test data generated as follows:
    //
    // truncate --size 4MiB disk.bin
    // sgdisk disk.bin \
    //   --disk-guid=57a7feb6-8cd5-4922-b7bd-c78b0914e870 \
    //   --new=1:2048:4096 \
    //   --change-name='1:hello world!' \
    //   --partition-guid=1:37c75ffd-8932-467a-9c56-8cf1f0456b12 \
    //   --typecode=1:ccf0994f-f7e0-4e26-a011-843e38aa2eac
    // bzip2 disk.bin
    // mv disk.bin.bz2 gpt_disk_io/tests/
    let compressed_data = Cursor::new(include_bytes!("disk.bin.bz2"));

    let mut reader = bzip2_rs::DecoderReader::new(compressed_data);
    let mut disk = Vec::new();
    reader.read_to_end(&mut disk).unwrap();
    disk
}

fn test_disk_read<Io>(block_io: Io) -> Result<(), DiskError<Io::Error>>
where
    Io: BlockIo,
{
    let mut block_buf = vec![0u8; 512];
    let mut disk = Disk::new(block_io)?;

    let primary_header = disk.read_primary_gpt_header(&mut block_buf)?;
    assert_eq!(primary_header, create_primary_header());

    let secondary_header = disk.read_secondary_gpt_header(&mut block_buf)?;
    assert_eq!(secondary_header, create_secondary_header());

    let expected_partition_entry = create_partition_entry();

    // Check the primary partition entry array.
    let primary_partition_entry = disk
        .gpt_partition_entry_array_iter(
            primary_header.get_partition_entry_array_layout().unwrap(),
            &mut block_buf,
        )?
        .next()
        .unwrap()?;
    assert_eq!(primary_partition_entry, expected_partition_entry);

    // Check the secondary partition entry array.
    let second_partition_entry = disk
        .gpt_partition_entry_array_iter(
            primary_header.get_partition_entry_array_layout().unwrap(),
            &mut block_buf,
        )?
        .next()
        .unwrap()?;
    assert_eq!(second_partition_entry, expected_partition_entry);

    Ok(())
}

fn test_disk_write<Io>(block_io: Io) -> Result<(), DiskError<Io::Error>>
where
    Io: BlockIo,
{
    let bs = BlockSize::BS_512;
    let mut block_buf = vec![0u8; bs.to_usize().unwrap()];
    let mut disk = Disk::new(block_io)?;

    let primary_header = create_primary_header();
    let secondary_header = create_secondary_header();
    let partition_entry = create_partition_entry();

    disk.write_protective_mbr(&mut block_buf)?;
    disk.write_primary_gpt_header(&primary_header, &mut block_buf)?;
    disk.write_secondary_gpt_header(&secondary_header, &mut block_buf)?;

    let layout = primary_header.get_partition_entry_array_layout().unwrap();
    let mut bytes =
        vec![0; layout.num_bytes_rounded_to_block_as_usize(bs).unwrap()];
    let mut entry_array =
        GptPartitionEntryArray::new(layout, bs, &mut bytes).unwrap();
    *entry_array.get_partition_entry_mut(0).unwrap() = partition_entry;
    disk.write_gpt_partition_entry_array(&entry_array)?;<commit_msg>Improve test coverage of partition entry array reading<commit_after>// Copyright 2022 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod common;

use anyhow::Result;
use common::{
    create_partition_entry, create_primary_header, create_secondary_header,
};
#[cfg(feature = "std")]
use gpt_disk_io::StdBlockIo;
use gpt_disk_io::{BlockIo, Disk, DiskError, MutSliceBlockIo};
use gpt_disk_types::{BlockSize, GptPartitionEntryArray};
use std::io::{Cursor, Read};

fn load_test_disk() -> Vec<u8> {
    // Test data generated as follows:
    //
    // truncate --size 4MiB disk.bin
    // sgdisk disk.bin \
    //   --disk-guid=57a7feb6-8cd5-4922-b7bd-c78b0914e870 \
    //   --new=1:2048:4096 \
    //   --change-name='1:hello world!' \
    //   --partition-guid=1:37c75ffd-8932-467a-9c56-8cf1f0456b12 \
    //   --typecode=1:ccf0994f-f7e0-4e26-a011-843e38aa2eac
    // bzip2 disk.bin
    // mv disk.bin.bz2 gpt_disk_io/tests/
    let compressed_data = Cursor::new(include_bytes!("disk.bin.bz2"));

    let mut reader = bzip2_rs::DecoderReader::new(compressed_data);
    let mut disk = Vec::new();
    reader.read_to_end(&mut disk).unwrap();
    disk
}

fn test_disk_read<Io>(block_io: Io) -> Result<(), DiskError<Io::Error>>
where
    Io: BlockIo,
{
    let bs = BlockSize::BS_512;
    let mut block_buf = vec![0u8; bs.to_usize().unwrap()];
    let mut disk = Disk::new(block_io)?;

    let primary_header = disk.read_primary_gpt_header(&mut block_buf)?;
    assert_eq!(primary_header, create_primary_header());

    let secondary_header = disk.read_secondary_gpt_header(&mut block_buf)?;
    assert_eq!(secondary_header, create_secondary_header());

    let expected_partition_entry = create_partition_entry();

    let check_partition_entry_array = |disk: &mut Disk<Io>, layout| {
        // First use the iter interface.
        {
            let mut block_buf = vec![0u8; bs.to_usize().unwrap()];
            let mut iter = disk
                .gpt_partition_entry_array_iter(layout, &mut block_buf)
                .unwrap();
            let entry = iter.next().unwrap().unwrap();
            assert_eq!(entry, expected_partition_entry);
            assert!(entry.is_used());

            let entry = iter.next().unwrap().unwrap();
            assert!(!entry.is_used());
        }

        // Then check the whole array.
        let mut array_buf = vec![0u8; bs.to_usize().unwrap() * 34];
        let array = disk
            .read_gpt_partition_entry_array(layout, &mut array_buf)
            .unwrap();
        let entry = *array.get_partition_entry(0).unwrap();
        assert_eq!(entry, expected_partition_entry);
        assert!(entry.is_used());

        let entry = *array.get_partition_entry(1).unwrap();
        assert!(!entry.is_used());
    };

    // Check the primary partition entry array.
    check_partition_entry_array(
        &mut disk,
        primary_header.get_partition_entry_array_layout().unwrap(),
    );

    // Check the secondary partition entry array.
    check_partition_entry_array(
        &mut disk,
        secondary_header.get_partition_entry_array_layout().unwrap(),
    );

    Ok(())
}

fn test_disk_write<Io>(block_io: Io) -> Result<(), DiskError<Io::Error>>
where
    Io: BlockIo,
{
    let bs = BlockSize::BS_512;
    let mut block_buf = vec![0u8; bs.to_usize().unwrap()];
    let mut disk = Disk::new(block_io)?;

    let primary_header = create_primary_header();
    let secondary_header = create_secondary_header();
    let partition_entry = create_partition_entry();

    disk.write_protective_mbr(&mut block_buf)?;
    disk.write_primary_gpt_header(&primary_header, &mut block_buf)?;
    disk.write_secondary_gpt_header(&secondary_header, &mut block_buf)?;

    let layout = primary_header.get_partition_entry_array_layout().unwrap();
    let mut bytes =
        vec![0; layout.num_bytes_rounded_to_block_as_usize(bs).unwrap()];
    let mut entry_array =
        GptPartitionEntryArray::new(layout, bs, &mut bytes).unwrap();
    *entry_array.get_partition_entry_mut(0).unwrap() = partition_entry;
    disk.write_gpt_partition_entry_array(&entry_array)?;
{'fraction_non_alphanumeric': 0.12554779106952504, 'fraction_numerical': 0.024162027715267084, 'mean_word_length': 3.4772004241781547, 'pattern_counts': {'":': 0, '<': 18, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '642397', 'n_tokens_mistral': 2817, 'n_tokens_neox': 2723, 'n_words': 593}