// Copyright 2018 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#![deny(warnings)]

use libc;
use std::fs::{self, File, OpenOptions};
use std::io::{Error, ErrorKind, Read, Seek, SeekFrom, Write};
use std::os::unix::io::AsRawFd;
use std::path::Path;
use structopt::StructOpt;

const DEV_DIR: &str = "/dev";

// Ioctl requests are comprised of 32 bits:
// [31:30]  Access mode
// [29:16]  Size of the parameter structure
// [15:8]   Request type
// [7:0]    Request number
// For the ioctls we care about here there are no parameters, so we set only a type and a number.
// See Linux:include/uapi/asm-generic/ioctl.h for more details.
const TYPE_SHIFT: usize = 8;
macro_rules! define_ioctl {
    ($name:ident, $typ:expr, $num:expr, $return_type:ty) => {
        fn $name(file: &File) -> $return_type {
            let request = ($typ << TYPE_SHIFT) | ($num & 0xff);
            let mut r: $return_type = 0;
            unsafe {
                libc::ioctl(file.as_raw_fd(), request, &mut r);
            }
            r
        }
    };
}

// Block ioctl types and number are defined in Linux:include/uapi/linux/fs.h.
define_ioctl!(block_dev_size, 0x12, 96, u64);
define_ioctl!(block_dev_sector_size, 0x12, 104, u32);

#[derive(StructOpt, Debug)]
struct Config {
    block_size: u32,
    block_count: u64,
    #[structopt(subcommand)]
    cmd: Command,
}

#[derive(StructOpt, Debug)]
enum Command {
    #[structopt(name = "check")]
    Check,
    #[structopt(name = "read")]
    Read { offset: u64, expected: u8 },
    #[structopt(name = "write")]
    Write { offset: u64, value: u8 },
}

fn find_block_device(block_size: u32, block_count: u64, write: bool) -> std::io::Result<File> {
    let dir = Path::new(DEV_DIR);
    if !dir.is_dir() {
        return Err(Error::new(
            ErrorKind::Other,
            format!("{} is not a directory", DEV_DIR),
        ));
    }
    for entry in fs::read_dir(dir)?.collect::<std::io::Result<Vec<_>>>()? {
        if !entry
            .file_name()
            .to_str()
            .map_or(false, |f| f.starts_with("vd"))
        {
            continue;
        }
        let file = OpenOptions::new()
            .read(true)
            .write(write)
            .open(entry.path())?;
        if block_size != block_dev_sector_size(&file) {
            continue;
        }
        if block_count != block_dev_size(&file) {
            continue;
        }
        return Ok(file);
    }
    Err(Error::new(ErrorKind::NotFound, "Block device not found"))
}

fn read_block(
    block_dev: &mut File,
    block_size: u32,
    offset: u64,
    expected: u8,
) -> std::io::Result<()> {
    block_dev.seek(SeekFrom::Start(offset * u64::from(block_size)))?;
    let mut data: Vec<u8> = vec![0; block_size as usize];
    block_dev.read_exact(&mut data)?;
    if !data.iter().all(|&b| b == expected) {
        return Err(Error::new(ErrorKind::Other, "Incorrect data read"));
    }
    Ok(())
}

fn write_block(
    block_dev: &mut File,
    block_size: u32,
    offset: u64,
    value: u8,
) -> std::io::Result<()> {
    block_dev.seek(SeekFrom::Start(offset * u64::from(block_size)))?;
    let data: Vec<u8> = vec![value; block_size as usize];
    block_dev.write_all(&data)?;
    block_dev.sync_all()?;
    Ok(())
}

fn main() -> std::io::Result<()> {
    let config = Config::from_args();
    let write = if let Command::Write { .. } = config.cmd {
        true
    } else {
        false
    };
    let mut block_dev = find_block_device(config.block_size, config.block_count, write)?;
    let result = match config.cmd {
        Command::Check => Ok(()),
        Command::Read { offset, expected } => {
            read_block(&mut block_dev, config.block_size, offset, expected)
        }
        Command::Write { offset, value } => {
            write_block(&mut block_dev, config.block_size, offset, value)
        }
    };
    if result.is_ok() {
        println!("PASS");
    }
    result
}