// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE-BSD-3-Clause file.
//
// Copyright © 2020 Intel Corporation
//
// SPDX-License-Identifier: Apache-2.0 AND BSD-3-Clause

pub mod async_io;
pub mod disk_file;
pub mod error;
pub mod fcntl;
pub mod fixed_vhd;
#[cfg(feature = "io_uring")]
/// Enabled with the `"io_uring"` feature
pub mod fixed_vhd_async;
pub mod fixed_vhd_sync;
pub mod qcow;
pub mod qcow_sync;
#[cfg(feature = "io_uring")]
/// Async primitives based on `io-uring`
///
/// Enabled with the `"io_uring"` feature
pub mod raw_async;
pub mod raw_async_aio;
#[cfg(test)]
mod raw_async_io_tests;
pub mod raw_sync;
pub mod vhd;
pub mod vhdx;
pub mod vhdx_sync;

use std::alloc::{Layout, alloc_zeroed, dealloc};
use std::collections::VecDeque;
use std::fmt::{self, Debug};
use std::fs::{File, OpenOptions};
use std::io::{self, IoSlice, IoSliceMut, Read, Seek, SeekFrom, Write};
use std::os::linux::fs::MetadataExt;
use std::os::unix::fs::FileTypeExt;
use std::os::unix::io::AsRawFd;
use std::path::Path;
use std::str::FromStr;
use std::time::Instant;
use std::{cmp, mem, result};

#[cfg(feature = "io_uring")]
use io_uring::{IoUring, Probe, opcode};
use libc::{
    FALLOC_FL_KEEP_SIZE, FALLOC_FL_PUNCH_HOLE, FALLOC_FL_ZERO_RANGE, S_IFBLK, S_IFMT, ioctl,
};
use log::{debug, error, info, warn};
use serde::{Deserialize, Serialize};
use smallvec::SmallVec;
use thiserror::Error;
use virtio_bindings::virtio_blk::*;
use virtio_queue::DescriptorChain;
use vm_memory::bitmap::Bitmap;
use vm_memory::{
    Address, ByteValued, Bytes, GuestAddress, GuestMemory, GuestMemoryError, GuestMemoryLoadGuard,
};
use vm_virtio::{AccessPlatform, Translatable};
use vmm_sys_util::eventfd::EventFd;
use vmm_sys_util::{aio, ioctl_io_nr, ioctl_ior_nr};

use crate::async_io::{AsyncIo, AsyncIoError, AsyncIoResult};
use crate::error::{BlockError, BlockErrorKind, BlockResult, ErrorOp};
use crate::vhdx::VhdxError;

const SECTOR_SHIFT: u8 = 9;
pub const SECTOR_SIZE: u64 = 0x01 << SECTOR_SHIFT;

/// Maximum number of segments per DISCARD or WRITE_ZEROES request.
pub const MAX_DISCARD_WRITE_ZEROES_SEG: u32 = 1;

/// Size and field offsets within `struct virtio_blk_discard_write_zeroes`.
const DISCARD_WZ_SEG_SIZE: u32 = mem::size_of::<virtio_blk_discard_write_zeroes>() as u32;
const DISCARD_WZ_MAX_PAYLOAD: u32 = DISCARD_WZ_SEG_SIZE * MAX_DISCARD_WRITE_ZEROES_SEG;
const DISCARD_WZ_SECTOR_OFFSET: u64 =
    mem::offset_of!(virtio_blk_discard_write_zeroes, sector) as u64;
const DISCARD_WZ_NUM_SECTORS_OFFSET: u64 =
    mem::offset_of!(virtio_blk_discard_write_zeroes, num_sectors) as u64;
const DISCARD_WZ_FLAGS_OFFSET: u64 = mem::offset_of!(virtio_blk_discard_write_zeroes, flags) as u64;

#[derive(Error, Debug)]
pub enum Error {
    #[error("Guest gave us bad memory addresses")]
    GuestMemory(#[source] GuestMemoryError),
    #[error("Guest gave us offsets that would have overflowed a usize")]
    CheckedOffset(GuestAddress, usize /* sector offset */),
    #[error("Guest gave us a write only descriptor that protocol says to read from")]
    UnexpectedWriteOnlyDescriptor,
    #[error("Guest gave us a read only descriptor that protocol says to write to")]
    UnexpectedReadOnlyDescriptor,
    #[error("Guest gave us too few descriptors in a descriptor chain")]
    DescriptorChainTooShort,
    #[error("Guest gave us a descriptor that was too short to use")]
    DescriptorLengthTooSmall,
    #[error("Failed to detect image type")]
    DetectImageType(#[source] std::io::Error),
    #[error("Failure in fixed vhd")]
    FixedVhdError(#[source] std::io::Error),
    #[error("Getting a block's metadata failed")]
    GetFileMetadata(#[source] std::io::Error),
    #[error("The requested operation would cause a seek beyond disk end")]
    InvalidOffset,
    #[error("Failure in qcow")]
    QcowError(#[source] qcow::Error),
    #[error("Failure in raw file")]
    RawFileError(#[source] std::io::Error),
    #[error("The requested operation does not support multiple descriptors")]
    TooManyDescriptors,
    #[error("Request contains too many segments ({0}, max {MAX_DISCARD_WRITE_ZEROES_SEG})")]
    TooManySegments(u32),
    #[error("Failure in vhdx")]
    VhdxError(#[source] VhdxError),
}

fn build_device_id(disk_path: &Path) -> result::Result<String, Error> {
    let blk_metadata = match disk_path.metadata() {
        Err(e) => return Err(Error::GetFileMetadata(e)),
        Ok(m) => m,
    };
    // This is how kvmtool does it.
    let device_id = format!(
        "{}{}{}",
        blk_metadata.st_dev(),
        blk_metadata.st_rdev(),
        blk_metadata.st_ino()
    );
    Ok(device_id)
}

pub fn build_serial(disk_path: &Path) -> Vec<u8> {
    let mut default_serial = vec![0; VIRTIO_BLK_ID_BYTES as usize];
    match build_device_id(disk_path) {
        Err(_) => {
            warn!("Could not generate device id. We'll use a default.");
        }
        Ok(m) => {
            // The kernel only knows to read a maximum of VIRTIO_BLK_ID_BYTES.
            // This will also zero out any leftover bytes.
            let disk_id = m.as_bytes();
            let bytes_to_copy = cmp::min(disk_id.len(), VIRTIO_BLK_ID_BYTES as usize);
            default_serial[..bytes_to_copy].clone_from_slice(&disk_id[..bytes_to_copy]);
        }
    }
    default_serial
}

#[derive(Error, Debug)]
pub enum ExecuteError {
    #[error("Bad request")]
    BadRequest(#[source] Error),
    #[error("Failed to flush")]
    Flush(#[source] io::Error),
    #[error("Failed to read")]
    Read(#[source] GuestMemoryError),
    #[error("Failed to read_exact")]
    ReadExact(#[source] io::Error),
    #[error("Can't execute an operation other than `read` or `get_id` on a read-only device")]
    ReadOnly,
    #[error("Failed to seek")]
    Seek(#[source] io::Error),
    #[error("Failed to write")]
    Write(#[source] GuestMemoryError),
    #[error("Failed to write_all")]
    WriteAll(#[source] io::Error),
    #[error("Unsupported request: {0}")]
    Unsupported(u32),
    #[error("Unsupported flags {flags:#x} for request type {request_type}")]
    UnsupportedFlags { request_type: u32, flags: u32 },
    #[error("Failed to submit io uring")]
    SubmitIoUring(#[source] io::Error),
    #[error("Failed to get guest address")]
    GetHostAddress(#[source] GuestMemoryError),
    #[error("Failed to async read")]
    AsyncRead(#[source] AsyncIoError),
    #[error("Failed to async write")]
    AsyncWrite(#[source] AsyncIoError),
    #[error("failed to async flush")]
    AsyncFlush(#[source] AsyncIoError),
    #[error("Failed to async punch hole")]
    AsyncPunchHole(#[source] AsyncIoError),
    #[error("Failed to async write zeroes")]
    AsyncWriteZeroes(#[source] AsyncIoError),
    #[error("Failed allocating a temporary buffer")]
    TemporaryBufferAllocation(#[source] io::Error),
}

impl ExecuteError {
    pub fn status(&self) -> u8 {
        let status = match *self {
            ExecuteError::BadRequest(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::Flush(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::Read(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::ReadExact(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::ReadOnly => VIRTIO_BLK_S_IOERR,
            ExecuteError::Seek(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::Write(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::WriteAll(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::Unsupported(_) => VIRTIO_BLK_S_UNSUPP,
            ExecuteError::UnsupportedFlags { .. } => VIRTIO_BLK_S_UNSUPP,
            ExecuteError::SubmitIoUring(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::GetHostAddress(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::AsyncRead(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::AsyncWrite(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::AsyncFlush(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::AsyncPunchHole(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::AsyncWriteZeroes(_) => VIRTIO_BLK_S_IOERR,
            ExecuteError::TemporaryBufferAllocation(_) => VIRTIO_BLK_S_IOERR,
        };
        status as u8
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RequestType {
    In,
    Out,
    Flush,
    GetDeviceId,
    Discard,
    WriteZeroes,
    Unsupported(u32),
}

pub fn request_type<B: Bitmap + 'static>(
    mem: &vm_memory::GuestMemoryMmap<B>,
    desc_addr: GuestAddress,
) -> result::Result<RequestType, Error> {
    let type_ = mem.read_obj(desc_addr).map_err(Error::GuestMemory)?;
    match type_ {
        VIRTIO_BLK_T_IN => Ok(RequestType::In),
        VIRTIO_BLK_T_OUT => Ok(RequestType::Out),
        VIRTIO_BLK_T_FLUSH => Ok(RequestType::Flush),
        VIRTIO_BLK_T_GET_ID => Ok(RequestType::GetDeviceId),
        VIRTIO_BLK_T_DISCARD => Ok(RequestType::Discard),
        VIRTIO_BLK_T_WRITE_ZEROES => Ok(RequestType::WriteZeroes),
        t => Ok(RequestType::Unsupported(t)),
    }
}

fn sector<B: Bitmap + 'static>(
    mem: &vm_memory::GuestMemoryMmap<B>,
    desc_addr: GuestAddress,
) -> result::Result<u64, Error> {
    const SECTOR_OFFSET: usize = 8;
    let addr = match mem.checked_offset(desc_addr, SECTOR_OFFSET) {
        Some(v) => v,
        None => return Err(Error::CheckedOffset(desc_addr, SECTOR_OFFSET)),
    };

    mem.read_obj(addr).map_err(Error::GuestMemory)
}

const DEFAULT_DESCRIPTOR_VEC_SIZE: usize = 32;

#[derive(Debug)]
pub struct AlignedOperation {
    origin_ptr: u64,
    aligned_ptr: u64,
    size: usize,
    layout: Layout,
}

pub struct BatchRequest {
    pub offset: libc::off_t,
    pub iovecs: SmallVec<[libc::iovec; DEFAULT_DESCRIPTOR_VEC_SIZE]>,
    pub user_data: u64,
    pub request_type: RequestType,
}

pub struct ExecuteAsync {
    // `true` if the execution will complete asynchronously
    pub async_complete: bool,
    // request need to be batched for submission if any
    pub batch_request: Option<BatchRequest>,
}

#[derive(Debug)]
pub struct Request {
    pub request_type: RequestType,
    pub sector: u64,
    pub data_descriptors: SmallVec<[(GuestAddress, u32); DEFAULT_DESCRIPTOR_VEC_SIZE]>,
    pub status_addr: GuestAddress,
    pub writeback: bool,
    pub aligned_operations: SmallVec<[AlignedOperation; DEFAULT_DESCRIPTOR_VEC_SIZE]>,
    pub start: Instant,
}

impl Request {
    pub fn parse<B: Bitmap + 'static>(
        desc_chain: &mut DescriptorChain<GuestMemoryLoadGuard<vm_memory::GuestMemoryMmap<B>>>,
        access_platform: Option<&dyn AccessPlatform>,
    ) -> result::Result<Request, Error> {
        let hdr_desc = desc_chain
            .next()
            .ok_or(Error::DescriptorChainTooShort)
            .inspect_err(|_| {
                error!("Missing head descriptor");
            })?;

        // The head contains the request type which MUST be readable.
        if hdr_desc.is_write_only() {
            return Err(Error::UnexpectedWriteOnlyDescriptor);
        }

        let hdr_desc_addr = hdr_desc
            .addr()
            .translate_gva(access_platform, hdr_desc.len() as usize);

        let mut req = Request {
            request_type: request_type(desc_chain.memory(), hdr_desc_addr)?,
            sector: sector(desc_chain.memory(), hdr_desc_addr)?,
            data_descriptors: SmallVec::with_capacity(DEFAULT_DESCRIPTOR_VEC_SIZE),
            status_addr: GuestAddress(0),
            writeback: true,
            aligned_operations: SmallVec::with_capacity(DEFAULT_DESCRIPTOR_VEC_SIZE),
            start: Instant::now(),
        };

        let status_desc;
        let mut desc = desc_chain
            .next()
            .ok_or(Error::DescriptorChainTooShort)
            .inspect_err(|_| {
                error!("Only head descriptor present: request = {req:?}");
            })?;

        if desc.has_next() {
            req.data_descriptors.reserve_exact(1);
            while desc.has_next() {
                if desc.is_write_only() && req.request_type == RequestType::Out {
                    return Err(Error::UnexpectedWriteOnlyDescriptor);
                }
                if desc.is_write_only() && req.request_type == RequestType::Discard {
                    return Err(Error::UnexpectedWriteOnlyDescriptor);
                }
                if desc.is_write_only() && req.request_type == RequestType::WriteZeroes {
                    return Err(Error::UnexpectedWriteOnlyDescriptor);
                }
                if !desc.is_write_only() && req.request_type == RequestType::In {
                    return Err(Error::UnexpectedReadOnlyDescriptor);
                }
                if !desc.is_write_only() && req.request_type == RequestType::GetDeviceId {
                    return Err(Error::UnexpectedReadOnlyDescriptor);
                }

                req.data_descriptors.push((
                    desc.addr()
                        .translate_gva(access_platform, desc.len() as usize),
                    desc.len(),
                ));
                desc = desc_chain
                    .next()
                    .ok_or(Error::DescriptorChainTooShort)
                    .inspect_err(|_| {
                        error!("DescriptorChain corrupted: request = {req:?}");
                    })?;
            }
            status_desc = desc;
        } else {
            status_desc = desc;
            // Only flush requests are allowed to skip the data descriptor.
            if req.request_type != RequestType::Flush {
                error!("Need a data descriptor: request = {req:?}");
                return Err(Error::DescriptorChainTooShort);
            }
        }

        // The status MUST always be writable.
        if !status_desc.is_write_only() {
            return Err(Error::UnexpectedReadOnlyDescriptor);
        }

        if status_desc.len() < 1 {
            return Err(Error::DescriptorLengthTooSmall);
        }

        req.status_addr = status_desc
            .addr()
            .translate_gva(access_platform, status_desc.len() as usize);

        Ok(req)
    }

    pub fn execute<T: Seek + Read + Write, B: Bitmap + 'static>(
        &self,
        disk: &mut T,
        disk_nsectors: u64,
        mem: &vm_memory::GuestMemoryMmap<B>,
        serial: &[u8],
    ) -> result::Result<u32, ExecuteError> {
        disk.seek(SeekFrom::Start(self.sector << SECTOR_SHIFT))
            .map_err(ExecuteError::Seek)?;
        let mut len = 0;
        for (data_addr, data_len) in &self.data_descriptors {
            let mut top: u64 = u64::from(*data_len) / SECTOR_SIZE;
            if u64::from(*data_len) % SECTOR_SIZE != 0 {
                top += 1;
            }
            top = top
                .checked_add(self.sector)
                .ok_or(ExecuteError::BadRequest(Error::InvalidOffset))?;
            if top > disk_nsectors {
                return Err(ExecuteError::BadRequest(Error::InvalidOffset));
            }

            match self.request_type {
                RequestType::In => {
                    let mut buf = vec![0u8; *data_len as usize];
                    disk.read_exact(&mut buf).map_err(ExecuteError::ReadExact)?;
                    mem.read_exact_volatile_from(
                        *data_addr,
                        &mut buf.as_slice(),
                        *data_len as usize,
                    )
                    .map_err(ExecuteError::Read)?;
                    len += data_len;
                }
                RequestType::Out => {
                    let mut buf: Vec<u8> = Vec::new();
                    mem.write_all_volatile_to(*data_addr, &mut buf, *data_len as usize)
                        .map_err(ExecuteError::Write)?;
                    disk.write_all(&buf).map_err(ExecuteError::WriteAll)?;
                    if !self.writeback {
                        disk.flush().map_err(ExecuteError::Flush)?;
                    }
                }
                RequestType::Flush => disk.flush().map_err(ExecuteError::Flush)?,
                RequestType::GetDeviceId => {
                    if (*data_len as usize) < serial.len() {
                        return Err(ExecuteError::BadRequest(Error::InvalidOffset));
                    }
                    mem.write_slice(serial, *data_addr)
                        .map_err(ExecuteError::Write)?;
                }
                RequestType::Discard => {
                    return Err(ExecuteError::Unsupported(VIRTIO_BLK_T_DISCARD));
                }
                RequestType::WriteZeroes => {
                    return Err(ExecuteError::Unsupported(VIRTIO_BLK_T_WRITE_ZEROES));
                }
                RequestType::Unsupported(t) => return Err(ExecuteError::Unsupported(t)),
            }
        }
        Ok(len)
    }

    pub fn execute_async<B: Bitmap + 'static>(
        &mut self,
        mem: &vm_memory::GuestMemoryMmap<B>,
        disk_nsectors: u64,
        disk_image: &mut dyn AsyncIo,
        serial: &[u8],
        disable_sector0_writes: bool,
        user_data: u64,
    ) -> result::Result<ExecuteAsync, ExecuteError> {
        let sector = self.sector;
        let request_type = self.request_type;
        let offset = (sector << SECTOR_SHIFT) as libc::off_t;
        let alignment = disk_image.alignment();

        let mut iovecs: SmallVec<[libc::iovec; DEFAULT_DESCRIPTOR_VEC_SIZE]> =
            SmallVec::with_capacity(self.data_descriptors.len());
        for &(data_addr, data_len) in &self.data_descriptors {
            let _: u32 = data_len; // compiler-checked documentation
            const _: () = assert!(
                core::mem::size_of::<u32>() <= core::mem::size_of::<usize>(),
                "unsupported platform"
            );
            if data_len == 0 {
                continue;
            }
            let mut top: u64 = u64::from(data_len) / SECTOR_SIZE;
            if u64::from(data_len) % SECTOR_SIZE != 0 {
                top += 1;
            }
            let data_len = data_len as usize;
            top = top
                .checked_add(sector)
                .ok_or(ExecuteError::BadRequest(Error::InvalidOffset))?;
            if top > disk_nsectors {
                return Err(ExecuteError::BadRequest(Error::InvalidOffset));
            }

            let origin_ptr = mem
                .get_slice(data_addr, data_len)
                .map_err(ExecuteError::GetHostAddress)?;
            assert!(origin_ptr.len() >= data_len);
            let origin_ptr = origin_ptr.ptr_guard();

            // O_DIRECT requires buffer addresses to be aligned to the
            // backend device's logical block size. In case it's not properly
            // aligned, an intermediate buffer is created with the correct
            // alignment, and a copy from/to the origin buffer is performed,
            // depending on the type of operation.
            let iov_base = if (origin_ptr.as_ptr() as u64).is_multiple_of(alignment) {
                origin_ptr.as_ptr() as *mut libc::c_void
            } else {
                let layout = Layout::from_size_align(data_len, alignment as usize).unwrap();
                // SAFETY: layout has non-zero size
                let aligned_ptr = unsafe { alloc_zeroed(layout) };
                if aligned_ptr.is_null() {
                    return Err(ExecuteError::TemporaryBufferAllocation(
                        io::Error::last_os_error(),
                    ));
                }

                // We need to perform the copy beforehand in case we're writing
                // data out.
                if request_type == RequestType::Out {
                    // SAFETY: destination buffer has been allocated with
                    // the proper size.
                    unsafe { std::ptr::copy(origin_ptr.as_ptr(), aligned_ptr, data_len) };
                }

                // Store both origin and aligned pointers for complete_async()
                // to process them.
                self.aligned_operations.push(AlignedOperation {
                    origin_ptr: origin_ptr.as_ptr() as u64,
                    aligned_ptr: aligned_ptr as u64,
                    size: data_len,
                    layout,
                });

                aligned_ptr as *mut libc::c_void
            };

            let iovec = libc::iovec {
                iov_base,
                iov_len: data_len as libc::size_t,
            };
            iovecs.push(iovec);
        }

        let mut ret = ExecuteAsync {
            async_complete: true,
            batch_request: None,
        };
        // Queue operations expected to be submitted.
        match request_type {
            RequestType::In => {
                for (data_addr, data_len) in &self.data_descriptors {
                    mem.get_slice(*data_addr, *data_len as usize)
                        .map_err(ExecuteError::GetHostAddress)?
                        .bitmap()
                        .mark_dirty(0, *data_len as usize);
                }
                if disk_image.batch_requests_enabled() {
                    ret.batch_request = Some(BatchRequest {
                        offset,
                        iovecs,
                        user_data,
                        request_type,
                    });
                } else {
                    disk_image
                        .read_vectored(offset, &iovecs, user_data)
                        .map_err(ExecuteError::AsyncRead)?;
                }
            }
            RequestType::Out => {
                if disk_image.batch_requests_enabled() {
                    ret.batch_request = Some(BatchRequest {
                        offset,
                        iovecs,
                        user_data,
                        request_type,
                    });
                } else {
                    disk_image
                        .write_vectored(offset, &iovecs, user_data)
                        .map_err(ExecuteError::AsyncWrite)?;
                }
            }
            RequestType::Flush => {
                disk_image
                    .fsync(Some(user_data))
                    .map_err(ExecuteError::AsyncFlush)?;
            }
            RequestType::GetDeviceId => {
                let (data_addr, data_len) = if self.data_descriptors.len() == 1 {
                    (self.data_descriptors[0].0, self.data_descriptors[0].1)
                } else {
                    return Err(ExecuteError::BadRequest(Error::TooManyDescriptors));
                };
                if (data_len as usize) < serial.len() {
                    return Err(ExecuteError::BadRequest(Error::InvalidOffset));
                }
                mem.write_slice(serial, data_addr)
                    .map_err(ExecuteError::Write)?;
                ret.async_complete = false;
                return Ok(ret);
            }
            RequestType::Discard => {
                let (data_addr, data_len) = if self.data_descriptors.len() == 1 {
                    (self.data_descriptors[0].0, self.data_descriptors[0].1)
                } else {
                    return Err(ExecuteError::BadRequest(Error::TooManyDescriptors));
                };

                if data_len < DISCARD_WZ_SEG_SIZE {
                    return Err(ExecuteError::BadRequest(Error::DescriptorLengthTooSmall));
                }
                if data_len > DISCARD_WZ_MAX_PAYLOAD {
                    return Err(ExecuteError::BadRequest(Error::TooManySegments(
                        data_len.div_ceil(DISCARD_WZ_SEG_SIZE),
                    )));
                }

                let mut discard_sector = [0u8; 8];
                let mut discard_num_sectors = [0u8; 4];
                let mut discard_flags = [0u8; 4];

                let sector_addr = data_addr.checked_add(DISCARD_WZ_SECTOR_OFFSET).unwrap();
                mem.read_slice(&mut discard_sector, sector_addr)
                    .map_err(ExecuteError::Read)?;

                let num_sectors_addr = data_addr
                    .checked_add(DISCARD_WZ_NUM_SECTORS_OFFSET)
                    .unwrap();
                mem.read_slice(&mut discard_num_sectors, num_sectors_addr)
                    .map_err(ExecuteError::Read)?;

                let flags_addr = data_addr.checked_add(DISCARD_WZ_FLAGS_OFFSET).unwrap();
                mem.read_slice(&mut discard_flags, flags_addr)
                    .map_err(ExecuteError::Read)?;

                let discard_flags = u32::from_le_bytes(discard_flags);
                // Per virtio spec v1.2 reject discard if any flag is set, including unmap.
                if discard_flags != 0 {
                    warn!("Unsupported flags {discard_flags:#x} in discard request");
                    return Err(ExecuteError::UnsupportedFlags {
                        request_type: VIRTIO_BLK_T_DISCARD,
                        flags: discard_flags,
                    });
                }

                let discard_sector = u64::from_le_bytes(discard_sector);

                if discard_sector == 0 && disable_sector0_writes {
                    return Err(ExecuteError::BadRequest(Error::InvalidOffset));
                }

                let discard_num_sectors = u32::from_le_bytes(discard_num_sectors);

                let top = discard_sector
                    .checked_add(discard_num_sectors as u64)
                    .ok_or(ExecuteError::BadRequest(Error::InvalidOffset))?;
                if top > disk_nsectors {
                    return Err(ExecuteError::BadRequest(Error::InvalidOffset));
                }

                let discard_offset = discard_sector * SECTOR_SIZE;
                let discard_length = (discard_num_sectors as u64) * SECTOR_SIZE;

                disk_image
                    .punch_hole(discard_offset, discard_length, user_data)
                    .map_err(ExecuteError::AsyncPunchHole)?;
            }
            RequestType::WriteZeroes => {
                let (data_addr, data_len) = if self.data_descriptors.len() == 1 {
                    (self.data_descriptors[0].0, self.data_descriptors[0].1)
                } else {
                    return Err(ExecuteError::BadRequest(Error::TooManyDescriptors));
                };

                if data_len < DISCARD_WZ_SEG_SIZE {
                    return Err(ExecuteError::BadRequest(Error::DescriptorLengthTooSmall));
                }
                if data_len > DISCARD_WZ_MAX_PAYLOAD {
                    return Err(ExecuteError::BadRequest(Error::TooManySegments(
                        data_len.div_ceil(DISCARD_WZ_SEG_SIZE),
                    )));
                }

                let mut wz_sector = [0u8; 8];
                let mut wz_num_sectors = [0u8; 4];
                let mut wz_flags = [0u8; 4];

                let sector_addr = data_addr.checked_add(DISCARD_WZ_SECTOR_OFFSET).unwrap();
                mem.read_slice(&mut wz_sector, sector_addr)
                    .map_err(ExecuteError::Read)?;

                let num_sectors_addr = data_addr
                    .checked_add(DISCARD_WZ_NUM_SECTORS_OFFSET)
                    .unwrap();
                mem.read_slice(&mut wz_num_sectors, num_sectors_addr)
                    .map_err(ExecuteError::Read)?;

                let flags_addr = data_addr.checked_add(DISCARD_WZ_FLAGS_OFFSET).unwrap();
                mem.read_slice(&mut wz_flags, flags_addr)
                    .map_err(ExecuteError::Read)?;

                let wz_sector = u64::from_le_bytes(wz_sector);
                let wz_num_sectors = u32::from_le_bytes(wz_num_sectors);

                let wz_flags = u32::from_le_bytes(wz_flags);
                // Per virtio spec v1.2 reject write zeroes if any unknown flag is set.
                if (wz_flags & !VIRTIO_BLK_WRITE_ZEROES_FLAG_UNMAP) != 0 {
                    warn!("Unsupported flags {wz_flags:#x} in write zeroes request");
                    return Err(ExecuteError::UnsupportedFlags {
                        request_type: VIRTIO_BLK_T_WRITE_ZEROES,
                        flags: wz_flags,
                    });
                }

                let wz_offset = wz_sector * SECTOR_SIZE;
                if wz_offset == 0 && disable_sector0_writes {
                    return Err(ExecuteError::BadRequest(Error::InvalidOffset));
                }

                let top = wz_sector
                    .checked_add(wz_num_sectors as u64)
                    .ok_or(ExecuteError::BadRequest(Error::InvalidOffset))?;
                if top > disk_nsectors {
                    return Err(ExecuteError::BadRequest(Error::InvalidOffset));
                }

                let wz_length = (wz_num_sectors as u64) * SECTOR_SIZE;

                if wz_flags & VIRTIO_BLK_WRITE_ZEROES_FLAG_UNMAP != 0 {
                    disk_image
                        .punch_hole(wz_offset, wz_length, user_data)
                        .map_err(ExecuteError::AsyncPunchHole)?;
                } else {
                    disk_image
                        .write_zeroes(wz_offset, wz_length, user_data)
                        .map_err(ExecuteError::AsyncWriteZeroes)?;
                }
            }
            RequestType::Unsupported(t) => return Err(ExecuteError::Unsupported(t)),
        }

        Ok(ret)
    }

    pub fn complete_async(&mut self) -> result::Result<(), Error> {
        for aligned_operation in self.aligned_operations.drain(..) {
            // We need to perform the copy after the data has been read inside
            // the aligned buffer in case we're reading data in.
            if self.request_type == RequestType::In {
                // SAFETY: origin buffer has been allocated with the
                // proper size.
                unsafe {
                    std::ptr::copy(
                        aligned_operation.aligned_ptr as *const u8,
                        aligned_operation.origin_ptr as *mut u8,
                        aligned_operation.size,
                    );
                };
            }

            // Free the temporary aligned buffer.
            // SAFETY: aligned_ptr was allocated by alloc_zeroed with the same
            // layout
            unsafe {
                dealloc(
                    aligned_operation.aligned_ptr as *mut u8,
                    aligned_operation.layout,
                );
            };
        }

        Ok(())
    }

    pub fn set_writeback(&mut self, writeback: bool) {
        self.writeback = writeback;
    }
}

#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize)]
#[repr(C, packed)]
pub struct VirtioBlockConfig {
    pub capacity: u64,
    pub size_max: u32,
    pub seg_max: u32,
    pub geometry: VirtioBlockGeometry,
    pub blk_size: u32,
    pub physical_block_exp: u8,
    pub alignment_offset: u8,
    pub min_io_size: u16,
    pub opt_io_size: u32,
    pub writeback: u8,
    pub unused: u8,
    pub num_queues: u16,
    pub max_discard_sectors: u32,
    pub max_discard_seg: u32,
    pub discard_sector_alignment: u32,
    pub max_write_zeroes_sectors: u32,
    pub max_write_zeroes_seg: u32,
    pub write_zeroes_may_unmap: u8,
    pub unused1: [u8; 3],
}
#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize)]
#[repr(C, packed)]
pub struct VirtioBlockGeometry {
    pub cylinders: u16,
    pub heads: u8,
    pub sectors: u8,
}

// SAFETY: data structure only contain a series of integers
unsafe impl ByteValued for VirtioBlockConfig {}
// SAFETY: data structure only contain a series of integers
unsafe impl ByteValued for VirtioBlockGeometry {}

/// Check if aio can be used on the current system.
pub fn block_aio_is_supported() -> bool {
    aio::IoContext::new(1).is_ok()
}

/// Check if io_uring for block device can be used on the current system, as
/// it correctly supports the expected io_uring features.
pub fn block_io_uring_is_supported() -> bool {
    #[cfg(not(feature = "io_uring"))]
    {
        info!("io_uring is disabled by crate features");
        false
    }

    #[cfg(feature = "io_uring")]
    {
        let error_msg = "io_uring not supported:";

        // Check we can create an io_uring instance, which effectively verifies
        // that io_uring_setup() syscall is supported.
        let io_uring = match IoUring::new(1) {
            Ok(io_uring) => io_uring,
            Err(e) => {
                info!("{error_msg} failed to create io_uring instance: {e}");
                return false;
            }
        };

        let submitter = io_uring.submitter();

        let mut probe = Probe::new();

        // Check we can register a probe to validate supported operations.
        match submitter.register_probe(&mut probe) {
            Ok(_) => {}
            Err(e) => {
                info!("{error_msg} failed to register a probe: {e}");
                return false;
            }
        }

        // Check IORING_OP_FSYNC is supported
        if !probe.is_supported(opcode::Fsync::CODE) {
            info!("{error_msg} IORING_OP_FSYNC operation not supported");
            return false;
        }

        // Check IORING_OP_READV is supported
        if !probe.is_supported(opcode::Readv::CODE) {
            info!("{error_msg} IORING_OP_READV operation not supported");
            return false;
        }

        // Check IORING_OP_WRITEV is supported
        if !probe.is_supported(opcode::Writev::CODE) {
            info!("{error_msg} IORING_OP_WRITEV operation not supported");
            return false;
        }

        true
    }
}

/// Probe whether the file/device supports punch hole and zero range
pub fn probe_sparse_support(file: &File) -> bool {
    let fd = file.as_raw_fd();

    let is_block_device = {
        let mut stat = std::mem::MaybeUninit::<libc::stat>::uninit();
        // SAFETY: FFI call with valid fd and buffer
        let ret = unsafe { libc::fstat(fd, stat.as_mut_ptr()) };
        if ret != 0 {
            warn!(
                "Failed to stat file descriptor for sparse probe: {}",
                io::Error::last_os_error()
            );
            return false;
        }
        // SAFETY: stat result is valid at this point
        unsafe { (*stat.as_ptr()).st_mode & S_IFMT == S_IFBLK }
    };

    if is_block_device {
        probe_block_device_sparse_support(fd)
    } else {
        probe_file_sparse_support(fd)
    }
}

/// Probe sparse support for a regular file using fallocate().
fn probe_file_sparse_support(fd: libc::c_int) -> bool {
    // SAFETY: FFI call with valid fd
    let file_size = unsafe { libc::lseek(fd, 0, libc::SEEK_END) };
    if file_size < 0 {
        let err = io::Error::last_os_error();
        warn!("Failed to get file size for sparse probe: {err}");
        return false;
    }

    // SAFETY: FFI call with valid fd, probing past EOF is safe with KEEP_SIZE
    let punch_hole =
        unsafe { libc::fallocate(fd, FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE, file_size, 1) }
            == 0;

    if !punch_hole {
        let err = io::Error::last_os_error();
        if err.raw_os_error() == Some(libc::EOPNOTSUPP) {
            debug!("File does not support FALLOC_FL_PUNCH_HOLE: {err}");
        } else {
            debug!("PUNCH_HOLE probe returned unexpected error: {err}");
        }
    }

    // SAFETY: FFI call with valid fd, probing past EOF is safe with KEEP_SIZE
    let zero_range =
        unsafe { libc::fallocate(fd, FALLOC_FL_ZERO_RANGE | FALLOC_FL_KEEP_SIZE, file_size, 1) }
            == 0;

    if !zero_range {
        let err = io::Error::last_os_error();
        if err.raw_os_error() == Some(libc::EOPNOTSUPP) {
            debug!("File does not support FALLOC_FL_ZERO_RANGE: {err}");
        }
    }

    let supported = punch_hole || zero_range;
    info!(
        "Probed file sparse support: punch_hole={punch_hole}, zero_range={zero_range} => {supported}"
    );
    supported
}

/// Probe sparse support for a block device.
///
/// Block devices always report sparse support. `BLKZEROOUT` is guaranteed to
/// succeed as the kernel provides a software fallback writing explicit zeros
/// when the hardware lacks a native write zeroes command. `BLKDISCARD` may fail
/// at runtime with `EOPNOTSUPP` on devices without trim or discard support, but
/// Linux guests handle this gracefully by ceasing discard requests.
///
/// There is no non destructive read only ioctl to query block device discard
/// or write zeroes capabilities.
fn probe_block_device_sparse_support(_fd: libc::c_int) -> bool {
    info!("Block device: assuming sparse support");
    true
}

/// Preallocate disk space for a disk image file.
///
/// Uses `fallocate()` to allocate all disk space upfront, ensuring storage
/// availability and reducing fragmentation. Allocating all blocks upfront is
/// more likely to place them contiguously than allocating on demand during
/// random writes.
pub fn preallocate_disk<P: AsRef<Path>>(file: &File, path: P) {
    let size = match file.metadata() {
        Ok(m) => m.len(),
        Err(e) => {
            warn!("Failed to get metadata for {:?}: {}", path.as_ref(), e);
            return;
        }
    };

    if size == 0 {
        return;
    }

    // SAFETY: FFI call with valid file descriptor and size
    let ret = unsafe { libc::fallocate(file.as_raw_fd(), 0, 0, size as libc::off_t) };

    if ret != 0 {
        warn!(
            "Failed to preallocate disk space for {:?}: {}",
            path.as_ref(),
            io::Error::last_os_error()
        );
    } else {
        debug!(
            "Preallocated {size} bytes for disk image {:?}",
            path.as_ref()
        );
    }
}

pub trait AsyncAdaptor {
    fn read_vectored_sync(
        &mut self,
        offset: libc::off_t,
        iovecs: &[libc::iovec],
        user_data: u64,
        eventfd: &EventFd,
        completion_list: &mut VecDeque<(u64, i32)>,
    ) -> AsyncIoResult<()>
    where
        Self: Read + Seek,
    {
        // Convert libc::iovec into IoSliceMut
        let mut slices: SmallVec<[IoSliceMut; DEFAULT_DESCRIPTOR_VEC_SIZE]> =
            SmallVec::with_capacity(iovecs.len());
        for iovec in iovecs.iter() {
            // SAFETY: on Linux IoSliceMut wraps around libc::iovec
            slices.push(IoSliceMut::new(unsafe {
                std::mem::transmute::<libc::iovec, &mut [u8]>(*iovec)
            }));
        }

        let result = {
            // Move the cursor to the right offset
            self.seek(SeekFrom::Start(offset as u64))
                .map_err(AsyncIoError::ReadVectored)?;

            let mut r = 0;
            for b in slices.iter_mut() {
                r += self.read(b).map_err(AsyncIoError::ReadVectored)?;
            }
            r
        };

        completion_list.push_back((user_data, result as i32));
        eventfd.write(1).unwrap();

        Ok(())
    }

    fn write_vectored_sync(
        &mut self,
        offset: libc::off_t,
        iovecs: &[libc::iovec],
        user_data: u64,
        eventfd: &EventFd,
        completion_list: &mut VecDeque<(u64, i32)>,
    ) -> AsyncIoResult<()>
    where
        Self: Write + Seek,
    {
        // Convert libc::iovec into IoSlice
        let mut slices: SmallVec<[IoSlice; DEFAULT_DESCRIPTOR_VEC_SIZE]> =
            SmallVec::with_capacity(iovecs.len());
        for iovec in iovecs.iter() {
            // SAFETY: on Linux IoSlice wraps around libc::iovec
            slices.push(IoSlice::new(unsafe {
                std::mem::transmute::<libc::iovec, &mut [u8]>(*iovec)
            }));
        }

        let result = {
            // Move the cursor to the right offset
            self.seek(SeekFrom::Start(offset as u64))
                .map_err(AsyncIoError::WriteVectored)?;

            let mut r = 0;
            for b in slices.iter() {
                r += self.write(b).map_err(AsyncIoError::WriteVectored)?;
            }
            r
        };

        completion_list.push_back((user_data, result as i32));
        eventfd.write(1).unwrap();

        Ok(())
    }

    fn fsync_sync(
        &mut self,
        user_data: Option<u64>,
        eventfd: &EventFd,
        completion_list: &mut VecDeque<(u64, i32)>,
    ) -> AsyncIoResult<()>
    where
        Self: Write,
    {
        let result: i32 = {
            // Flush
            self.flush().map_err(AsyncIoError::Fsync)?;

            0
        };

        if let Some(user_data) = user_data {
            completion_list.push_back((user_data, result));
            eventfd.write(1).unwrap();
        }

        Ok(())
    }
}

#[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum ImageType {
    FixedVhd,
    Qcow2,
    Raw,
    Vhdx,
    #[default]
    Unknown,
}

impl fmt::Display for ImageType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ImageType::FixedVhd => write!(f, "vhd"),
            ImageType::Qcow2 => write!(f, "qcow2"),
            ImageType::Raw => write!(f, "raw"),
            ImageType::Vhdx => write!(f, "vhdx"),
            ImageType::Unknown => write!(f, "unknown"),
        }
    }
}

pub enum ImageTypeParseError {
    InvalidValue(String),
}

impl FromStr for ImageType {
    type Err = ImageTypeParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "vhd" => Ok(ImageType::FixedVhd),
            "qcow2" => Ok(ImageType::Qcow2),
            "raw" => Ok(ImageType::Raw),
            "vhdx" => Ok(ImageType::Vhdx),
            _ => Err(ImageTypeParseError::InvalidValue(s.to_string())),
        }
    }
}

const QCOW_MAGIC: u32 = 0x5146_49fb;
const VHDX_SIGN: u64 = 0x656C_6966_7864_6876;

/// Read a block into memory aligned by the source block size (needed for O_DIRECT)
pub fn read_aligned_block_size(f: &mut File) -> std::io::Result<Vec<u8>> {
    let blocksize = DiskTopology::probe(f)?.logical_block_size as usize;
    // SAFETY: We are allocating memory that is naturally aligned (size = alignment) and we meet
    // requirements for safety from Vec::from_raw_parts() as we are using the global allocator
    // and transferring ownership of the memory.
    let mut data = unsafe {
        Vec::from_raw_parts(
            alloc_zeroed(Layout::from_size_align_unchecked(blocksize, blocksize)),
            blocksize,
            blocksize,
        )
    };
    f.read_exact(&mut data)?;
    Ok(data)
}

/// Open a disk image file, returning a [`BlockError`] with path context
/// on failure.
pub fn open_disk_image(path: &Path, options: &OpenOptions) -> BlockResult<File> {
    options.open(path).map_err(|e| {
        BlockError::new(BlockErrorKind::Io, e)
            .with_op(ErrorOp::Open)
            .with_path(path)
    })
}

/// Determine image type through file parsing.
pub fn detect_image_type(f: &mut File) -> BlockResult<ImageType> {
    let block = read_aligned_block_size(f)
        .map_err(|e| BlockError::new(BlockErrorKind::Io, e).with_op(ErrorOp::DetectImageType))?;

    // Check 4 first bytes to get the header value and determine the image type
    let image_type = if u32::from_be_bytes(block[0..4].try_into().unwrap()) == QCOW_MAGIC {
        ImageType::Qcow2
    } else if vhd::is_fixed_vhd(f)
        .map_err(|e| BlockError::new(BlockErrorKind::Io, e).with_op(ErrorOp::DetectImageType))?
    {
        ImageType::FixedVhd
    } else if u64::from_le_bytes(block[0..8].try_into().unwrap()) == VHDX_SIGN {
        ImageType::Vhdx
    } else {
        ImageType::Raw
    };

    Ok(image_type)
}

pub trait BlockBackend: Read + Write + Seek + Send + Debug {
    /// Returns the logical disk size a guest will see.
    ///
    /// For raw formats, this is equal to [`Self::physical_size`]. For file formats
    /// that wrap disk images in a container (e.g. QCOW2), this refers to the
    /// effective size that the guest will see.
    fn logical_size(&self) -> Result<u64, Error>;
    /// Returns the physical size of the underlying file.
    fn physical_size(&self) -> Result<u64, Error>;
}

#[derive(Debug)]
pub struct DiskTopology {
    pub logical_block_size: u64,
    pub physical_block_size: u64,
    pub minimum_io_size: u64,
    pub optimal_io_size: u64,
}

impl Default for DiskTopology {
    fn default() -> Self {
        Self {
            logical_block_size: 512,
            physical_block_size: 512,
            minimum_io_size: 512,
            optimal_io_size: 0,
        }
    }
}

ioctl_io_nr!(BLKSSZGET, 0x12, 104);
ioctl_io_nr!(BLKPBSZGET, 0x12, 123);
ioctl_io_nr!(BLKIOMIN, 0x12, 120);
ioctl_io_nr!(BLKIOOPT, 0x12, 121);
ioctl_ior_nr!(BLKGETSIZE64, 0x12, 114, u64);

/// Returns `(logical_size, physical_size)` in bytes for regular files and block devices.
///
/// For regular files, logical size is `st_size` and physical size is
/// `st_blocks * 512` (actual host allocation). For block devices both
/// values equal the `BLKGETSIZE64` result.
pub fn query_device_size(file: &File) -> io::Result<(u64, u64)> {
    let m = file.metadata()?;
    if m.is_file() {
        // st_blocks is always in 512-byte units on Linux
        Ok((m.len(), m.st_blocks() * 512))
    } else if m.file_type().is_block_device() {
        let mut size: u64 = 0;
        // SAFETY: BLKGETSIZE64 reads the device size into a u64 pointer.
        let ret = unsafe { libc::ioctl(file.as_raw_fd(), BLKGETSIZE64() as _, &mut size) };
        if ret != 0 {
            return Err(io::Error::last_os_error());
        }
        Ok((size, size))
    } else {
        Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!(
                "disk image must be a regular file or block device, is: {:?}",
                m.file_type()
            ),
        ))
    }
}

#[derive(Copy, Clone)]
enum BlockSize {
    LogicalBlock,
    PhysicalBlock,
    MinimumIo,
    OptimalIo,
}

impl DiskTopology {
    fn is_block_device(f: &File) -> std::io::Result<bool> {
        let mut stat = std::mem::MaybeUninit::<libc::stat>::uninit();
        // SAFETY: FFI call with a valid fd and buffer
        let ret = unsafe { libc::fstat(f.as_raw_fd(), stat.as_mut_ptr()) };
        if ret != 0 {
            return Err(std::io::Error::last_os_error());
        }

        // SAFETY: stat is valid at this point
        let is_block = unsafe { (*stat.as_ptr()).st_mode & S_IFMT == S_IFBLK };
        Ok(is_block)
    }

    // libc::ioctl() takes different types on different architectures
    fn query_block_size(f: &File, block_size_type: BlockSize) -> std::io::Result<u64> {
        let mut block_size = 0;
        // SAFETY: FFI call with correct arguments
        let ret = unsafe {
            ioctl(
                f.as_raw_fd(),
                match block_size_type {
                    BlockSize::LogicalBlock => BLKSSZGET(),
                    BlockSize::PhysicalBlock => BLKPBSZGET(),
                    BlockSize::MinimumIo => BLKIOMIN(),
                    BlockSize::OptimalIo => BLKIOOPT(),
                } as _,
                &mut block_size,
            )
        };
        if ret != 0 {
            return Err(std::io::Error::last_os_error());
        }

        Ok(block_size)
    }

    /// Query the O_DIRECT alignment requirement for a regular file.
    ///
    /// Uses `statx(STATX_DIOALIGN)` (Linux >= 6.1) to obtain the exact
    /// memory and offset alignment the kernel requires for direct I/O on
    /// this specific file. Unlike `fstatvfs().f_bsize`, which only returns
    /// the filesystem's preferred I/O block size, `STATX_DIOALIGN` reports
    /// the true per-file DIO constraints accounting for the filesystem,
    /// underlying block device, and any stacking (loop, dm, etc.).
    fn query_file_alignment(f: &File) -> u64 {
        // The libc crate does not expose statx / STATX_DIOALIGN on all
        // targets (e.g. musl), so define the constant and a minimal repr(C)
        // struct locally and invoke the syscall directly.
        const STATX_DIOALIGN: u32 = 0x2000;

        // Minimal statx layout, only the needed fields,
        // everything else is padding.
        #[repr(C)]
        struct Statx {
            stx_mask: u32,
            _pad: [u8; 148],
            stx_dio_mem_align: u32,
            stx_dio_offset_align: u32,
            _pad2: [u8; 96],
        }

        let mut stx = mem::MaybeUninit::<Statx>::zeroed();
        // SAFETY: FFI syscall with valid fd and correctly sized buffer.
        let ret = unsafe {
            libc::syscall(
                libc::SYS_statx,
                f.as_raw_fd(),
                c"".as_ptr(),
                libc::AT_EMPTY_PATH,
                STATX_DIOALIGN,
                stx.as_mut_ptr(),
            )
        };
        if ret == 0 {
            // SAFETY: statx succeeded, the struct is fully initialized.
            let stx = unsafe { stx.assume_init() };
            if stx.stx_mask & STATX_DIOALIGN != 0 && stx.stx_dio_mem_align > 0 {
                let align = cmp::max(stx.stx_dio_mem_align, stx.stx_dio_offset_align) as u64;
                debug!("statx(STATX_DIOALIGN) returned alignment {align}");
                return align;
            }
        }

        debug!("O_DIRECT alignment query failed, falling back to default {SECTOR_SIZE}");
        SECTOR_SIZE
    }

    pub fn probe(f: &File) -> std::io::Result<Self> {
        if !Self::is_block_device(f)? {
            // For regular files opened with O_DIRECT, the logical block size
            // must reflect the filesystem DIO alignment so the guest issues
            // correctly sized I/O.
            // SAFETY: fcntl(F_GETFL) is always safe on a valid fd.
            let flags = unsafe { libc::fcntl(f.as_raw_fd(), libc::F_GETFL) };
            if flags >= 0 && (flags & libc::O_DIRECT) != 0 {
                let alignment = Self::query_file_alignment(f);
                return Ok(DiskTopology {
                    logical_block_size: alignment,
                    physical_block_size: alignment,
                    minimum_io_size: alignment,
                    optimal_io_size: 0,
                });
            }
            return Ok(DiskTopology::default());
        }

        Ok(DiskTopology {
            logical_block_size: Self::query_block_size(f, BlockSize::LogicalBlock)?,
            physical_block_size: Self::query_block_size(f, BlockSize::PhysicalBlock)?,
            minimum_io_size: Self::query_block_size(f, BlockSize::MinimumIo)?,
            optimal_io_size: Self::query_block_size(f, BlockSize::OptimalIo)?,
        })
    }
}

#[cfg(test)]
mod unit_tests {
    use std::alloc::{Layout, alloc_zeroed, dealloc};
    use std::fs::OpenOptions;
    use std::io::Write;
    use std::os::unix::fs::OpenOptionsExt;
    use std::{ptr, slice};

    use vmm_sys_util::tempfile::TempFile;

    use super::*;

    #[test]
    fn test_probe_regular_file_returns_valid_alignment() {
        let temp_file = TempFile::new().unwrap();
        let mut f = temp_file.into_file();
        f.write_all(&[0u8; 4096]).unwrap();
        f.sync_all().unwrap();

        let topo = DiskTopology::probe(&f).unwrap();

        assert_eq!(
            topo.logical_block_size, SECTOR_SIZE,
            "probe() should return {SECTOR_SIZE} for regular files without O_DIRECT, got {}",
            topo.logical_block_size
        );
    }

    #[test]
    fn test_probe_regular_file_with_direct_returns_dio_alignment() {
        let temp_file = TempFile::new().unwrap();
        let path = temp_file.as_path().to_owned();
        {
            let f = temp_file.as_file();
            f.set_len(1 << 20).unwrap(); // 1 MiB
            f.sync_all().unwrap();
        }

        let f = OpenOptions::new()
            .read(true)
            .write(true)
            .custom_flags(libc::O_DIRECT)
            .open(&path)
            .unwrap();
        let topo = DiskTopology::probe(&f).unwrap();

        assert!(
            topo.logical_block_size.is_power_of_two(),
            "logical_block_size {} is not a power of two",
            topo.logical_block_size
        );
        assert!(
            topo.logical_block_size >= SECTOR_SIZE,
            "logical_block_size {} is less than SECTOR_SIZE ({SECTOR_SIZE})",
            topo.logical_block_size
        );

        let alignment = topo.logical_block_size as usize;
        let layout = Layout::from_size_align(4096, alignment);
        assert!(
            layout.is_ok(),
            "Layout::from_size_align(4096, {alignment}) failed: {:?}",
            layout.err()
        );
    }

    #[test]
    fn test_dio_write_read_with_probed_alignment() {
        let temp_file = TempFile::new().unwrap();
        let path = temp_file.as_path().to_owned();
        {
            let f = temp_file.as_file();
            f.set_len(1 << 20).unwrap(); // 1 MiB
            f.sync_all().unwrap();
        }

        let f = OpenOptions::new()
            .read(true)
            .write(true)
            .custom_flags(libc::O_DIRECT)
            .open(&path)
            .unwrap();
        let topo = DiskTopology::probe(&f).unwrap();
        let alignment = topo.logical_block_size as usize;

        let layout = Layout::from_size_align(alignment, alignment).unwrap();
        // SAFETY: layout is valid (non-zero, power-of-two alignment).
        let buf = unsafe { alloc_zeroed(layout) };
        assert!(!buf.is_null());

        // SAFETY: buf is valid for `alignment` bytes.
        unsafe { ptr::write_bytes(buf, 0xAB, alignment) };

        // SAFETY: buf is aligned and sized for O_DIRECT; fd is valid.
        let written =
            unsafe { libc::pwrite(f.as_raw_fd(), buf as *const libc::c_void, alignment, 0) };
        assert_eq!(
            written as usize,
            alignment,
            "O_DIRECT pwrite failed: {}",
            io::Error::last_os_error()
        );

        // SAFETY: buf is valid for `alignment` bytes.
        unsafe { ptr::write_bytes(buf, 0x00, alignment) };
        // SAFETY: buf is aligned and sized for O_DIRECT; fd is valid.
        let read = unsafe { libc::pread(f.as_raw_fd(), buf as *mut libc::c_void, alignment, 0) };
        assert_eq!(
            read as usize,
            alignment,
            "O_DIRECT pread failed: {}",
            io::Error::last_os_error()
        );

        // SAFETY: buf is valid for `alignment` bytes after successful pread.
        let slice = unsafe { slice::from_raw_parts(buf, alignment) };
        assert!(
            slice.iter().all(|&b| b == 0xAB),
            "Data mismatch after O_DIRECT roundtrip"
        );

        // SAFETY: buf was allocated with this layout via alloc_zeroed.
        unsafe { dealloc(buf, layout) };
    }

    #[test]
    fn test_query_device_size_regular_file() {
        let temp_file = TempFile::new().unwrap();
        let mut f = temp_file.into_file();
        // 5 sectors + 13 extra bytes - not page aligned, not sectoraligned
        f.write_all(&[0xAB; 5 * 512 + 13]).unwrap();
        f.sync_all().unwrap();

        let (logical, physical) = query_device_size(&f).unwrap();
        assert_eq!(logical, 5 * 512 + 13);
        assert!(physical > 0);
    }

    #[test]
    fn test_query_device_size_sparse_file_punch_hole() {
        let temp_file = TempFile::new().unwrap();
        let f = temp_file.as_file();
        // Allocate 1 MiB
        let size: i64 = 1 << 20;
        f.set_len(size as u64).unwrap();
        // SAFETY: fd is valid, range is within file size.
        let ret = unsafe {
            libc::fallocate(
                f.as_raw_fd(),
                0, // allocate
                0,
                size,
            )
        };
        assert_eq!(ret, 0, "fallocate failed: {}", io::Error::last_os_error());
        f.sync_all().unwrap();

        let (log_before, phys_before) = query_device_size(f).unwrap();
        assert_eq!(log_before, size as u64);
        assert_eq!(phys_before, size as u64);

        // Punch a hole in the middle 512 KiB
        // SAFETY: fd is valid, range is within file size.
        let ret = unsafe {
            libc::fallocate(
                f.as_raw_fd(),
                libc::FALLOC_FL_PUNCH_HOLE | libc::FALLOC_FL_KEEP_SIZE,
                size / 4,
                size / 2,
            )
        };
        assert_eq!(ret, 0, "punch hole failed: {}", io::Error::last_os_error());
        f.sync_all().unwrap();

        let (logical, physical) = query_device_size(f).unwrap();
        assert_eq!(logical, size as u64, "logical size must not change");
        assert!(
            physical < logical,
            "physical ({physical}) should be less than logical ({logical}) after punch hole"
        );
    }

    #[test]
    fn test_query_device_size_rejects_char_device() {
        let f = std::fs::File::open("/dev/zero").unwrap();
        let err = query_device_size(&f).unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
    }
}