// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{
    fmt,
    io::{self, Read, Seek, SeekFrom},
    str::from_utf8,
};

use byteorder::{ByteOrder, LittleEndian};
use chrono::{DateTime, Utc};
use crc::crc32;
use uuid::Uuid;

use devicemapper::{Sectors, IEC, SECTOR_SIZE};

use crate::{
    engine::{
        strat_engine::{
            backstore::metadata::{
                mda,
                sizes::{
                    static_header_size, BDAExtendedSize, BlockdevSize, MDADataSize, MDASize,
                    ReservedSize, STATIC_HEADER_SIZE,
                },
            },
            device::SyncAll,
        },
        types::{DevUuid, PoolUuid},
    },
    stratis::{ErrorEnum, StratisError, StratisResult},
};

const RESERVED_SECTORS: Sectors = Sectors(3 * IEC::Mi / (SECTOR_SIZE as u64)); // = 3 MiB

const STRAT_MAGIC: &[u8] = b"!Stra0tis\x86\xff\x02^\x41rh";

const STRAT_SIGBLOCK_VERSION: u8 = 1;

/// Get a Stratis pool UUID and device UUID from any device.
/// If there is an error while obtaining these values return the error.
/// If the device does not appear to be a Stratis device, return None.
pub fn device_identifiers<F>(f: &mut F) -> StratisResult<Option<((PoolUuid, DevUuid))>>
where
    F: Read + Seek + SyncAll,
{
    StaticHeader::setup(f).map(|sh| sh.map(|sh| (sh.pool_uuid, sh.dev_uuid)))
}

/// Remove Stratis identifying information from device.
pub fn disown_device<F>(f: &mut F) -> StratisResult<()>
where
    F: Seek + SyncAll,
{
    StaticHeader::wipe(f)
}

// Transform a constant in sectors to a constant in bytes
macro_rules! bytes {
    ($number:expr) => {
        $number * SECTOR_SIZE
    };
}

#[derive(Debug)]
pub struct BDA {
    header: StaticHeader,
    regions: mda::MDARegions,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum MetadataLocation {
    Both,
    First,
    Second,
}

impl BDA {
    /// Initialize a blockdev with a Stratis BDA.
    pub fn initialize<F>(
        f: &mut F,
        pool_uuid: Uuid,
        dev_uuid: Uuid,
        mda_data_size: MDADataSize,
        blkdev_size: BlockdevSize,
        initialization_time: u64,
    ) -> StratisResult<BDA>
    where
        F: Seek + SyncAll,
    {
        let header = StaticHeader::new(
            pool_uuid,
            dev_uuid,
            mda_data_size.region_size().mda_size(),
            blkdev_size,
            initialization_time,
        );

        header.write(f, MetadataLocation::Both)?;

        let regions =
            mda::MDARegions::initialize(STATIC_HEADER_SIZE.sectors().bytes(), header.mda_size, f)?;

        Ok(BDA { header, regions })
    }

    /// Load a BDA on initial setup of a device.
    /// Returns None if no BDA appears to exist.
    pub fn load<F>(f: &mut F) -> StratisResult<Option<BDA>>
    where
        F: Read + Seek + SyncAll,
    {
        let header = match StaticHeader::setup(f)? {
            Some(header) => header,
            None => return Ok(None),
        };

        // Assume that, since a valid StaticHeader was found on the device,
        // that this implies that BDA::initialize() was succesfully executed
        // sometime in the past. Since that is the case, valid MDA headers
        // were written to the device. Returns an error if there is an error
        // when loading the MDARegions, which can only be caused by an I/O
        // error or invalid MDA headers.
        let regions =
            mda::MDARegions::load(STATIC_HEADER_SIZE.sectors().bytes(), header.mda_size, f)?;

        Ok(Some(BDA { header, regions }))
    }

    /// Save metadata to the disk
    pub fn save_state<F>(
        &mut self,
        time: &DateTime<Utc>,
        metadata: &[u8],
        f: &mut F,
    ) -> StratisResult<()>
    where
        F: Seek + SyncAll,
    {
        self.regions
            .save_state(STATIC_HEADER_SIZE.sectors().bytes(), time, metadata, f)
    }

    /// Read latest metadata from the disk
    pub fn load_state<F>(&self, mut f: &mut F) -> StratisResult<Option<Vec<u8>>>
    where
        F: Read + Seek,
    {
        self.regions
            .load_state(STATIC_HEADER_SIZE.sectors().bytes(), &mut f)
    }

    /// The time when the most recent metadata was written to the BDA,
    /// if any.
    pub fn last_update_time(&self) -> Option<&DateTime<Utc>> {
        self.regions.last_update_time()
    }

    /// The UUID of the device.
    pub fn dev_uuid(&self) -> DevUuid {
        self.header.dev_uuid
    }

    /// The UUID of the device's pool.
    pub fn pool_uuid(&self) -> PoolUuid {
        self.header.pool_uuid
    }

    /// The size of the device.
    pub fn dev_size(&self) -> BlockdevSize {
        self.header.blkdev_size
    }

    /// The number of sectors the BDA itself occupies.
    pub fn extended_size(&self) -> BDAExtendedSize {
        self.header.bda_extended_size()
    }

    /// The maximum size of variable length metadata that can be accommodated.
    pub fn max_data_size(&self) -> MDADataSize {
        self.regions.max_data_size()
    }

    /// Timestamp when the device was initialized.
    pub fn initialization_time(&self) -> u64 {
        self.header.initialization_time
    }
}

#[derive(Eq, PartialEq)]
pub struct StaticHeader {
    blkdev_size: BlockdevSize,
    pool_uuid: PoolUuid,
    dev_uuid: DevUuid,
    mda_size: MDASize,
    reserved_size: ReservedSize,
    flags: u64,
    /// Seconds portion of DateTime<Utc> value.
    initialization_time: u64,
}

impl StaticHeader {
    fn new(
        pool_uuid: PoolUuid,
        dev_uuid: DevUuid,
        mda_size: MDASize,
        blkdev_size: BlockdevSize,
        initialization_time: u64,
    ) -> StaticHeader {
        StaticHeader {
            blkdev_size,
            pool_uuid,
            dev_uuid,
            mda_size,
            reserved_size: ReservedSize::new(RESERVED_SECTORS),
            flags: 0,
            initialization_time,
        }
    }

    /// Read the data at both signature block locations.
    ///
    /// Return the data from each location as an array of bytes
    /// or an error if the read fails. The values are returned
    /// in the same order in which they occur on the device.
    ///
    /// Read the contents of each signature block separately,
    /// as this increases the probability that at least one read
    /// will not fail.
    fn read<F>(
        f: &mut F,
    ) -> (
        io::Result<[u8; bytes!(static_header_size::SIGBLOCK_SECTORS)]>,
        io::Result<[u8; bytes!(static_header_size::SIGBLOCK_SECTORS)]>,
    )
    where
        F: Read + Seek,
    {
        let mut buf_loc_1 = [0u8; bytes!(static_header_size::SIGBLOCK_SECTORS)];
        let mut buf_loc_2 = [0u8; bytes!(static_header_size::SIGBLOCK_SECTORS)];

        fn read_sector_at_offset<F>(f: &mut F, offset: usize, mut buf: &mut [u8]) -> io::Result<()>
        where
            F: Read + Seek,
        {
            f.seek(SeekFrom::Start(offset as u64))
                .and_then(|_| f.read_exact(&mut buf))
        }

        (
            read_sector_at_offset(
                f,
                bytes!(static_header_size::FIRST_SIGBLOCK_START_SECTORS),
                &mut buf_loc_1,
            )
            .map(|_| buf_loc_1),
            read_sector_at_offset(
                f,
                bytes!(static_header_size::SECOND_SIGBLOCK_START_SECTORS),
                &mut buf_loc_2,
            )
            .map(|_| buf_loc_2),
        )
    }

    // Writes signature_block according to the value of which.
    // If first location is specified, write zeroes to empty regions in the
    // first 8 sectors. If the second location is specified, writes zeroes to empty
    // regions in the second 8 sectors.
    fn write<F>(&self, f: &mut F, which: MetadataLocation) -> io::Result<()>
    where
        F: Seek + SyncAll,
    {
        let signature_block = self.sigblock_to_buf();
        let zeroed = [0u8; bytes!(static_header_size::POST_SIGBLOCK_PADDING_SECTORS)];
        f.seek(SeekFrom::Start(0))?;

        // Write to a static header region in the static header.
        fn write_region<F>(f: &mut F, signature_block: &[u8], zeroed: &[u8]) -> io::Result<()>
        where
            F: Seek + SyncAll,
        {
            f.write_all(&zeroed[..bytes!(static_header_size::PRE_SIGBLOCK_PADDING_SECTORS)])?;
            f.write_all(signature_block)?;
            f.write_all(&zeroed[..bytes!(static_header_size::POST_SIGBLOCK_PADDING_SECTORS)])?;
            f.sync_all()?;
            Ok(())
        };

        if which == MetadataLocation::Both || which == MetadataLocation::First {
            write_region(f, &signature_block, &zeroed)?;
        } else {
            f.seek(SeekFrom::Start(
                bytes!(static_header_size::SIGBLOCK_REGION_SECTORS) as u64,
            ))?;
        }

        if which == MetadataLocation::Both || which == MetadataLocation::Second {
            write_region(f, &signature_block, &zeroed)?;
        }
        Ok(())
    }

    pub fn bda_extended_size(&self) -> BDAExtendedSize {
        BDAExtendedSize::new(self.mda_size.bda_size().sectors() + self.reserved_size.sectors())
    }

    /// Try to find a valid StaticHeader on a device.
    /// Return the latest copy that validates as a Stratis BDA, however verify both
    /// copies and if one validates but one does not, re-write the one that is incorrect.  If both
    /// copies are valid, but one is newer than the other, rewrite the older one to match.
    /// Return None if it's not a Stratis device.
    /// Return an error if the metadata seems to indicate that the device is
    /// a Stratis device, but no well-formed signature block could be read.
    /// Return an error if neither sigblock location can be read.
    /// Return an error if the sigblocks differ in some unaccountable way.
    /// Returns an error if a write intended to repair an ill-formed,
    /// unreadable, or stale signature block failed.
    fn setup<F>(f: &mut F) -> StratisResult<Option<StaticHeader>>
    where
        F: Read + Seek + SyncAll,
    {
        // Action taken when one sigblock is interpreted as invalid.
        //
        // If the other sigblock is interpreted as a Stratis header, attempts repair
        // of the invalid sigblock, returning an error if that fails, otherwise returning
        // the valid sigblock.
        //
        // In all other cases, return the error associated with the invalid sigblock.
        fn ok_err_static_header_handling<F>(
            f: &mut F,
            maybe_sh: Option<StaticHeader>,
            sh_error: StratisError,
            repair_location: MetadataLocation,
        ) -> StratisResult<Option<StaticHeader>>
        where
            F: Read + Seek + SyncAll,
        {
            if let Some(sh) = maybe_sh {
                write_header(f, sh, repair_location)
            } else {
                Err(sh_error)
            }
        }

        // Action taken when there was an I/O error reading the other sigblock.
        //
        // * If this sigblock region is interpreted as having no siglblock, it returns None.
        // * If this sigblock region has a valid sigblock, attempts repair of the other
        //   sigblock region with the valid sigblock, returning the valid sigblock
        //   if the repair succeeds, otherwise returning an error.
        // * If this sigblock appears to be invalid, return the error encountered when
        //   reading the sigblock.
        fn copy_ok_err_handling<F>(
            f: &mut F,
            maybe_sh: StratisResult<Option<StaticHeader>>,
            repair_location: MetadataLocation,
        ) -> StratisResult<Option<StaticHeader>>
        where
            F: Read + Seek + SyncAll,
        {
            match maybe_sh {
                Ok(loc) => {
                    if let Some(ref sh) = loc {
                        sh.write(f, repair_location)?;
                    }
                    Ok(loc)
                }
                Err(e) => Err(e),
            }
        }

        fn write_header<F>(
            f: &mut F,
            sh: StaticHeader,
            repair_location: MetadataLocation,
        ) -> StratisResult<Option<StaticHeader>>
        where
            F: Read + Seek + SyncAll,
        {
            sh.write(f, repair_location)?;
            Ok(Some(sh))
        }

        let (maybe_buf_1, maybe_buf_2) = StaticHeader::read(f);
        match (
            maybe_buf_1.map(|buf| StaticHeader::sigblock_from_buf(&buf)),
            maybe_buf_2.map(|buf| StaticHeader::sigblock_from_buf(&buf)),
        ) {
            // We read both copies without an IO error.
            (Ok(buf_loc_1), Ok(buf_loc_2)) => match (buf_loc_1, buf_loc_2) {
                (Ok(loc_1), Ok(loc_2)) => match (loc_1, loc_2) {
                    (Some(loc_1), Some(loc_2)) => {
                        if loc_1 == loc_2 {
                            Ok(Some(loc_1))
                        } else if loc_1.initialization_time == loc_2.initialization_time {
                            // Inexplicable disagreement among static headers
                            let err_str =
                                "Appeared to be a Stratis device, but signature blocks disagree.";
                            Err(StratisError::Engine(ErrorEnum::Invalid, err_str.into()))
                        } else if loc_1.initialization_time > loc_2.initialization_time {
                            // If the first header block is newer, overwrite second with
                            // contents of first.
                            write_header(f, loc_1, MetadataLocation::Second)
                        } else {
                            // The second header block must be newer, so overwrite first
                            // with contents of second.
                            write_header(f, loc_2, MetadataLocation::First)
                        }
                    }
                    (None, None) => Ok(None),
                    (Some(loc_1), None) => write_header(f, loc_1, MetadataLocation::Second),
                    (None, Some(loc_2)) => write_header(f, loc_2, MetadataLocation::First),
                },
                (Ok(loc_1), Err(loc_2)) => {
                    ok_err_static_header_handling(f, loc_1, loc_2, MetadataLocation::Second)
                }
                (Err(loc_1), Ok(loc_2)) => {
                    ok_err_static_header_handling(f, loc_2, loc_1, MetadataLocation::First)
                }
                (Err(_), Err(_)) => {
                    let err_str = "Appeared to be a Stratis device, but no valid sigblock found";
                    Err(StratisError::Engine(ErrorEnum::Invalid, err_str.into()))
                }
            },
            // Copy 1 read OK, 2 resulted in an IO error
            (Ok(buf_loc_1), Err(_)) => copy_ok_err_handling(f, buf_loc_1, MetadataLocation::Second),
            // Copy 2 read OK, 1 resulted in IO Error
            (Err(_), Ok(buf_loc_2)) => copy_ok_err_handling(f, buf_loc_2, MetadataLocation::First),
            (Err(_), Err(_)) => {
                // Unable to read the device at all.
                let err_str = "Unable to read data at sigblock locations.";
                Err(StratisError::Engine(ErrorEnum::Invalid, err_str.into()))
            }
        }
    }

    /// Generate a buf suitable for writing to blockdev
    fn sigblock_to_buf(&self) -> [u8; bytes!(static_header_size::SIGBLOCK_SECTORS)] {
        let mut buf = [0u8; bytes!(static_header_size::SIGBLOCK_SECTORS)];
        buf[4..20].clone_from_slice(STRAT_MAGIC);
        LittleEndian::write_u64(&mut buf[20..28], *self.blkdev_size.sectors());
        buf[28] = STRAT_SIGBLOCK_VERSION;
        buf[32..64].clone_from_slice(self.pool_uuid.to_simple_ref().to_string().as_bytes());
        buf[64..96].clone_from_slice(self.dev_uuid.to_simple_ref().to_string().as_bytes());
        LittleEndian::write_u64(&mut buf[96..104], *self.mda_size.sectors());
        LittleEndian::write_u64(&mut buf[104..112], *self.reserved_size.sectors());
        LittleEndian::write_u64(&mut buf[120..128], self.initialization_time);

        let hdr_crc =
            crc32::checksum_castagnoli(&buf[4..bytes!(static_header_size::SIGBLOCK_SECTORS)]);
        LittleEndian::write_u32(&mut buf[..4], hdr_crc);
        buf
    }

    /// Parse a buffer to a StaticHeader.
    /// Return None if no stratis magic number found.
    /// Return an error if stored checksum and calculated checksum do not
    /// match.
    /// Return an error if the version number is not expected.
    fn sigblock_from_buf(buf: &[u8]) -> StratisResult<Option<StaticHeader>> {
        assert_eq!(buf.len(), bytes!(static_header_size::SIGBLOCK_SECTORS));

        if &buf[4..20] != STRAT_MAGIC {
            return Ok(None);
        }

        let crc = crc32::checksum_castagnoli(&buf[4..bytes!(static_header_size::SIGBLOCK_SECTORS)]);
        if crc != LittleEndian::read_u32(&buf[..4]) {
            return Err(StratisError::Engine(
                ErrorEnum::Invalid,
                "header CRC invalid".into(),
            ));
        }

        let blkdev_size = BlockdevSize::new(Sectors(LittleEndian::read_u64(&buf[20..28])));

        let version = buf[28];
        if version != STRAT_SIGBLOCK_VERSION {
            return Err(StratisError::Engine(
                ErrorEnum::Invalid,
                format!("Unknown sigblock version: {}", version),
            ));
        }

        let pool_uuid = Uuid::parse_str(from_utf8(&buf[32..64])?)?;
        let dev_uuid = Uuid::parse_str(from_utf8(&buf[64..96])?)?;

        let mda_size = MDASize(Sectors(LittleEndian::read_u64(&buf[96..104])));

        Ok(Some(StaticHeader {
            pool_uuid,
            dev_uuid,
            blkdev_size,
            mda_size,
            reserved_size: ReservedSize::new(Sectors(LittleEndian::read_u64(&buf[104..112]))),
            flags: 0,
            initialization_time: LittleEndian::read_u64(&buf[120..128]),
        }))
    }

    /// Zero out the entire static header region on the designated file.
    pub fn wipe<F>(f: &mut F) -> StratisResult<()>
    where
        F: Seek + SyncAll,
    {
        let zeroed = [0u8; bytes!(static_header_size::STATIC_HEADER_SECTORS)];
        f.seek(SeekFrom::Start(0))?;
        f.write_all(&zeroed)?;
        f.sync_all()?;
        Ok(())
    }
}

impl fmt::Debug for StaticHeader {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("StaticHeader")
            .field("blkdev_size", &self.blkdev_size)
            .field("pool_uuid", &self.pool_uuid.to_simple_ref())
            .field("dev_uuid", &self.dev_uuid.to_simple_ref())
            .field("mda_size", &self.mda_size)
            .field("reserved_size", &self.reserved_size)
            .field("flags", &self.flags)
            .field("initialization_time", &self.initialization_time)
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use proptest::{collection::vec, num, option, prelude::BoxedStrategy, strategy::Strategy};
    use uuid::Uuid;

    use devicemapper::{Bytes, Sectors, IEC};

    use super::*;

    /// Return a static header with random block device and MDA size.
    /// The block device is less than the minimum, for efficiency in testing.
    fn random_static_header(blkdev_size: u64, mda_size_factor: u32) -> StaticHeader {
        let pool_uuid = Uuid::new_v4();
        let dev_uuid = Uuid::new_v4();
        let mda_size = MDADataSize::new(
            MDADataSize::default().bytes() + Bytes(u64::from(mda_size_factor * 4)),
        )
        .region_size()
        .mda_size();
        let blkdev_size = (Bytes(IEC::Mi) + Sectors(blkdev_size).bytes()).sectors();
        StaticHeader::new(
            pool_uuid,
            dev_uuid,
            mda_size,
            BlockdevSize::new(blkdev_size),
            Utc::now().timestamp() as u64,
        )
    }

    /// Make a static header strategy
    fn static_header_strategy() -> BoxedStrategy<StaticHeader> {
        (0..64u64, 0..64u32)
            .prop_map(|(b, m)| random_static_header(b, m))
            .boxed()
    }

    proptest! {
        #[test]
        /// Construct an arbitrary StaticHeader object.
        /// Verify that the "memory buffer" is unowned.
        /// Initialize a static header.
        /// Verify that Stratis buffer validates.
        /// Wipe the static header.
        /// Verify that the buffer is again unowned.
        fn test_ownership(ref sh in static_header_strategy()) {
            let buf_size = *sh.mda_size.sectors().bytes() as usize + bytes!(static_header_size::STATIC_HEADER_SECTORS);
            let mut buf = Cursor::new(vec![0; buf_size]);
            prop_assert!(StaticHeader::setup(&mut buf).unwrap().is_none());

            sh.write(&mut buf, MetadataLocation::Both).unwrap();

            prop_assert!(StaticHeader::setup(&mut buf)
                         .unwrap()
                         .map(|new_sh| new_sh.pool_uuid == sh.pool_uuid && new_sh.dev_uuid == sh.dev_uuid)
                         .unwrap_or(false));

            StaticHeader::wipe(&mut buf).unwrap();
            prop_assert!(StaticHeader::setup(&mut buf).unwrap().is_none());
        }
    }

    proptest! {
        #[test]
        /// Construct an arbitrary StaticHeader object.
        /// Initialize a BDA.
        /// Verify that the last update time is None.
        fn empty_bda(ref sh in static_header_strategy()) {
            let buf_size = *sh.mda_size.sectors().bytes() as usize + bytes!(static_header_size::STATIC_HEADER_SECTORS);
            let mut buf = Cursor::new(vec![0; buf_size]);
            let bda = BDA::initialize(
                &mut buf,
                sh.pool_uuid,
                sh.dev_uuid,
                sh.mda_size.region_size().data_size(),
                sh.blkdev_size,
                Utc::now().timestamp() as u64,
            ).unwrap();
            prop_assert!(bda.last_update_time().is_none());
        }
    }

    #[test]
    /// Construct a BDA and verify that an error is returned if timestamp
    /// of saved data is older than timestamp of most recently written data.
    fn test_early_times_err() {
        let data = [0u8; 3];

        // Construct a BDA.
        let sh = random_static_header(0, 0);
        let mut buf = Cursor::new(vec![0; *sh.blkdev_size.sectors().bytes() as usize]);
        let mut bda = BDA::initialize(
            &mut buf,
            sh.pool_uuid,
            sh.dev_uuid,
            sh.mda_size.region_size().data_size(),
            sh.blkdev_size,
            Utc::now().timestamp() as u64,
        )
        .unwrap();

        let timestamp0 = Utc::now();
        let timestamp1 = Utc::now();
        assert_ne!(timestamp0, timestamp1);

        let mut buf = Cursor::new(vec![0; *sh.blkdev_size.sectors().bytes() as usize]);
        bda.save_state(&timestamp1, &data, &mut buf).unwrap();

        // Error, because current timestamp is older than written to newer.
        assert_matches!(bda.save_state(&timestamp0, &data, &mut buf), Err(_));

        let timestamp2 = Utc::now();
        let timestamp3 = Utc::now();
        assert_ne!(timestamp2, timestamp3);

        bda.save_state(&timestamp3, &data, &mut buf).unwrap();

        // Error, because current timestamp is older than written to newer.
        assert_matches!(bda.save_state(&timestamp2, &data, &mut buf), Err(_));
    }

    proptest! {
        #[test]
        /// Construct an arbitrary StaticHeader object.
        /// Initialize a BDA.
        /// Save metadata and verify correct update time and state.
        /// Reload BDA and verify that new BDA has correct update time.
        /// Load state using new BDA and verify correct state.
        /// Save metadata again, and reload one more time, verifying new timestamp.
        fn check_state(
            ref sh in static_header_strategy(),
            ref state in vec(num::u8::ANY, 1..100),
            ref next_state in vec(num::u8::ANY, 1..100)
        ) {
            let buf_size = *sh.mda_size.sectors().bytes() as usize + bytes!(static_header_size::STATIC_HEADER_SECTORS);
            let mut buf = Cursor::new(vec![0; buf_size]);
            let mut bda = BDA::initialize(
                &mut buf,
                sh.pool_uuid,
                sh.dev_uuid,
                sh.mda_size.region_size().data_size(),
                sh.blkdev_size,
                Utc::now().timestamp() as u64,
            ).unwrap();
            let current_time = Utc::now();
            bda.save_state(&current_time, state, &mut buf).unwrap();
            let loaded_state = bda.load_state(&mut buf).unwrap();
            prop_assert!(bda.last_update_time().map(|t| t == &current_time).unwrap_or(false));
            prop_assert!(loaded_state.map(|s| &s == state).unwrap_or(false));

            let mut bda = BDA::load(&mut buf).unwrap().unwrap();
            let loaded_state = bda.load_state(&mut buf).unwrap();
            prop_assert!(loaded_state.map(|s| &s == state).unwrap_or(false));
            prop_assert!(bda.last_update_time().map(|t| t == &current_time).unwrap_or(false));

            let current_time = Utc::now();
            bda.save_state(&current_time, next_state, &mut buf)
                .unwrap();
            let loaded_state = bda.load_state(&mut buf).unwrap();
            prop_assert!(loaded_state.map(|s| &s == next_state).unwrap_or(false));
            prop_assert!(bda.last_update_time().map(|t| t == &current_time).unwrap_or(false));

        }
    }

    proptest! {
        #[test]
        /// Construct an arbitrary StaticHeader object.
        /// Write it to a buffer, read it out and make sure you get the same thing.
        fn static_header(ref sh1 in static_header_strategy()) {
            let buf = sh1.sigblock_to_buf();
            let sh2 = StaticHeader::sigblock_from_buf(&buf).unwrap().unwrap();
            prop_assert_eq!(sh1.pool_uuid, sh2.pool_uuid);
            prop_assert_eq!(sh1.dev_uuid, sh2.dev_uuid);
            prop_assert_eq!(sh1.blkdev_size, sh2.blkdev_size);
            prop_assert_eq!(sh1.mda_size, sh2.mda_size);
            prop_assert_eq!(sh1.reserved_size, sh2.reserved_size);
            prop_assert_eq!(sh1.flags, sh2.flags);
            prop_assert_eq!(sh1.initialization_time, sh2.initialization_time);
        }
    }

    proptest! {
        #[test]
        /// Verify correct reading of the static header if only one of
        /// the two static headers is corrupted. Verify expected behavior
        /// if both are corrupted, which varies depending on whether the
        /// Stratis magic number or some other part of the header is corrupted.
        fn test_corrupted_sigblock_recovery(primary in option::of(0..bytes!(static_header_size::SIGBLOCK_SECTORS)),
                             secondary in option::of(0..bytes!(static_header_size::SIGBLOCK_SECTORS))) {

            // Corrupt a byte at the specified position.
            fn corrupt_byte<F>(f: &mut F, position: u64) -> io::Result<()>
            where
                F: Read + Seek + SyncAll,
            {
                let mut byte_to_corrupt = [0; 1];
                f.seek(SeekFrom::Start(position))
                    .and_then(|_| f.read_exact(&mut byte_to_corrupt))?;

                byte_to_corrupt[0] = !byte_to_corrupt[0];

                f.seek(SeekFrom::Start(position))
                    .and_then(|_| f.write_all(&byte_to_corrupt))
                    .and_then(|_| f.sync_all())
            }

            let sh = random_static_header(10000, 4);
            let buf_size = bytes!(static_header_size::STATIC_HEADER_SECTORS);

            let mut reference_buf = Cursor::new(vec![0; buf_size]);
            sh.write(&mut reference_buf, MetadataLocation::Both).unwrap();

            let mut buf = Cursor::new(vec![0; buf_size]);
            sh.write(&mut buf, MetadataLocation::Both).unwrap();

            if let Some(index) = primary {
                corrupt_byte(
                    &mut buf,
                    (bytes!(static_header_size::FIRST_SIGBLOCK_START_SECTORS) + index) as u64,
                )
                .unwrap();
            }

            if let Some(index) = secondary {
                corrupt_byte(
                    &mut buf,
                    (bytes!(static_header_size::SECOND_SIGBLOCK_START_SECTORS) + index) as u64,
                )
                .unwrap();
            }

            let setup_result = StaticHeader::setup(&mut buf);

            match (primary, secondary) {
                (Some(p_index), Some(s_index)) => {
                    match (p_index, s_index) {
                        // Both magic locations are corrupted, conclusion is
                        // that this is not a Stratis static header.
                        (4..=19, 4..=19) => {
                            prop_assert!(setup_result.is_ok());
                            prop_assert_eq!(setup_result.unwrap(), None);
                        }
                        // Both sigblocks were corrupted, but at least one
                        // was recognized as a Stratis sigblock.
                        _ => {
                            prop_assert!(setup_result.is_err());
                        }
                    }
                    // No healing was attempted.
                    prop_assert_ne!(reference_buf.get_ref(), buf.get_ref());

                }
                // Only one header was corrupted, so the other was healed.
                _ => {
                    prop_assert!(setup_result.is_ok());
                    prop_assert_eq!(setup_result.unwrap(), Some(sh));
                    prop_assert_eq!(reference_buf.get_ref(), buf.get_ref());
                }
            }
        }
    }

    #[test]
    /// Test that the newer sigblock is copied to the older sigblock's location
    /// if the sigblock's timestamps don't match but they are otherwise
    /// identical.
    fn test_rewrite_older_sigblock() {
        let sh = random_static_header(10000, 4);

        let ts = Utc::now().timestamp() as u64;
        let sh_older =
            StaticHeader::new(sh.pool_uuid, sh.dev_uuid, sh.mda_size, sh.blkdev_size, ts);
        let sh_newer = StaticHeader::new(
            sh.pool_uuid,
            sh.dev_uuid,
            sh.mda_size,
            sh.blkdev_size,
            ts + 1,
        );
        assert_ne!(sh_older, sh_newer);

        let buf_size = bytes!(static_header_size::STATIC_HEADER_SECTORS);

        let mut reference_buf = Cursor::new(vec![0; buf_size]);
        sh_newer
            .write(&mut reference_buf, MetadataLocation::Both)
            .unwrap();

        // Test that StaticHeader::setup succeeds by writing the older
        // signature block to the specified older location and the newer
        // sigblock to the specified newer location and then calling
        // StaticHeader::setup. StaticHeader::setup should return without
        // error with the newer sigblock. As a side-effect, it should have
        // overwritten the location of the older sigblock with the value of
        // the newer sigblock.
        let test_rewrite = |sh_older: &StaticHeader,
                            sh_newer: &StaticHeader,
                            older_location: MetadataLocation,
                            newer_location: MetadataLocation| {
            let mut buf = Cursor::new(vec![0; buf_size]);
            sh_older.write(&mut buf, older_location).unwrap();
            sh_newer.write(&mut buf, newer_location).unwrap();
            assert_ne!(buf.get_ref(), reference_buf.get_ref());
            assert_eq!(
                StaticHeader::setup(&mut buf).unwrap().as_ref(),
                Some(sh_newer)
            );
            assert_eq!(buf.get_ref(), reference_buf.get_ref());
        };

        test_rewrite(
            &sh_older,
            &sh_newer,
            MetadataLocation::First,
            MetadataLocation::Second,
        );

        test_rewrite(
            &sh_older,
            &sh_newer,
            MetadataLocation::Second,
            MetadataLocation::First,
        );
    }
}
