<commit_before>use std::path::PathBuf;

use chrono::prelude::*;
use chrono::Duration;
use chrono_english::*;
use clap::Clap;

#[derive(Clap, Debug)]
pub enum SubCommand {
    /// Enqueue a task for execution.
    Add {
        /// The command that should be added.
        #[clap(required = true)]
        command: Vec<String>,

        /// Start the task immediately.
        #[clap(name = "immediate", short, long, conflicts_with = "stashed")]
        start_immediately: bool,

        /// Create the task in stashed state.
        /// Useful to avoid immediate execution if the queue is empty.
        #[clap(name = "stashed", short, long, conflicts_with = "immediate")]
        stashed: bool,

        /// Delays enqueueing the task until <delay> elapses. See "enqueue" for accepted formats.
        #[clap(name = "delay", short, long, conflicts_with = "immediate", parse(try_from_str=parse_delay_until))]
        delay_until: Option<DateTime<Local>>,

        /// Assign the task to a group. Groups kind of act as separate queues.
        /// I.e. all groups run in parallel and you can specify the amount of parallel tasks for each group.
        /// If no group is specified, the default group will be used.
        #[clap(name = "group", short, long)]
        group: Option<String>,

        /// Start the task once all specified tasks have successfully finished.
        /// As soon as one of the dependencies fails, this task will fail as well.
        #[clap(name = "after", short, long)]
        dependencies: Vec<usize>,
    },
    /// Remove tasks from the list.
    /// Running or paused tasks need to be killed first.
    Remove {
        /// The task ids to be removed.
        #[clap(required = true)]
        task_ids: Vec<usize>,
    },
    /// Switches the queue position of two commands. Only works on queued and stashed commands.
    Switch {
        /// The first task id.
        task_id_1: usize,
        /// The second task id.
        task_id_2: usize,
    },
    /// Stashed tasks won't be automatically started.
    /// Either enqueue them, to be normally handled or explicitly start them.
    Stash {
        /// The id(s) of the tasks you want to stash.
        #[clap(required = true)]
        task_ids: Vec<usize>,
    },
    /// Enqueue stashed tasks. They'll be handled normally afterwards.
    #[clap(after_help = "DELAY FORMAT:

    The --delay argument must be either a number of seconds or a \"date expression\" similar to GNU \
    `date -d` with some extensions. It does not attempt to parse all natural language, but is \
    incredibly flexible. Here are some supported examples.

    2020-04-01T18:30:00   // RFC 3339 timestamp
    2020-4-1 18:2:30      // Optional leading zeros
    2020-4-1 5:30pm       // Informal am/pm time
    2020-4-1 5pm          // Optional minutes and seconds
    April 1 2020 18:30:00 // English months
    1 Apr 8:30pm          // Implies current year
    4/1                   // American form date
    wednesday 10:30pm     // The closest wednesday in the future at 22:30
    wednesday             // The closest wednesday in the future
    4 months              // 4 months from today at 00:00:00
    1 week                // 1 week at the current time
    1days                 // 1 day from today at the current time
    1d 03:00              // The closest 3:00 after 1 day (24 hours)
    3h                    // 3 hours from now
    3600s                 // 3600 seconds from now
")]
    Enqueue {
        /// The id(s) of the tasks you want to enqueue.
        task_ids: Vec<usize>,

        /// Delay enqueuing the tasks until <delay> elapses. See DELAY FORMAT below.
        #[clap(name = "delay", short, long, parse(try_from_str=parse_delay_until))]
        delay_until: Option<DateTime<Local>>,
    },

    /// Resume operation of specific tasks or groups of tasks.
    /// By default, this resumes the default queue and all its tasks.
    /// Can also be used force-start specific tasks.
    #[clap(verbatim_doc_comment)]
    Start {
        /// Enforce starting these tasks. Paused tasks will be started again.
        /// This does not affect anything other than these tasks.
        task_ids: Vec<usize>,

        /// Start a specific group and all paused tasks in it.
        #[clap(short, long, group("start"))]
        group: Option<String>,

        /// Start a everything (Default queue and all groups)!
        /// All groups will be set to `running` and all paused tasks will be resumed.
        #[clap(short, long, group("start"))]
        all: bool,

        /// Also resume direct child processes of your paused tasks.
        /// By default only the main process will get a SIGSTART.
        #[clap(short, long, group("start"))]
        children: bool,
    },

    /// Restart task(s).
    /// Identical tasks will be created and by default enqueued.
    Restart {
        /// The tasks you want to restart.
        #[clap(required = true)]
        task_ids: Vec<usize>,

        /// Immediately start the task(s).
        #[clap(name = "immediate", short, long)]
        start_immediately: bool,

        /// Create the task in stashed state.
        /// Useful to avoid immediate execution.
        #[clap(short, long, conflicts_with = "immediate")]
        stashed: bool,

        /// Edit the command of the task before restarting
        #[clap(short, long)]
        edit: bool,

        /// Edit the path of the task before restarting
        #[clap(short, long)]
        path: bool,
    },

    /// Pause either running tasks or specific groups of tasks.
    /// By default, pauses the default queue and all its tasks.
    /// A paused queue (group) won't start any new tasks.
    #[clap(verbatim_doc_comment)]
    Pause {
        /// Pause these specific tasks.
        /// Does not affect the default queue, groups or any other tasks.
        task_ids: Vec<usize>,

        /// Pause a specific group.
        #[clap(short, long, group("pause"))]
        group: Option<String>,

        /// Pause everything (Default queue and all groups)!
        #[clap(short, long, group("pause"))]
        all: bool,

        /// Don not pause already running tasks and let them finish by themselves,
        /// when pausing with `default`, `all` or `group`.
        #[clap(short, long, group("pause"))]
        wait: bool,

        /// Also pause direct child processes of a task's main process.
        /// By default only the main process will get a SIGSTOP.
        /// This is useful when calling bash scripts, which start other processes themselves.
        /// This operation is not recursive!
        #[clap(short, long, group("pause"))]
        children: bool,
    },

    /// Kill specific running tasks or various groups of tasks.
    Kill {
        /// The tasks that should be killed.
        task_ids: Vec<usize>,

        /// Kill all running tasks in the default queue. Pause the default queue.
        #[clap(short, long, group("kill"))]
        default: bool,

        /// Kill all running in a group. Pauses the group.
        #[clap(short, long, group("kill"))]
        group: Option<String>,

        /// Kill ALL running tasks. This also pauses everything
        #[clap(short, long, group("kill"))]
        all: bool,

        /// Send the SIGTERM signal to all children as well.
        /// Useful when working with shell scripts.
        #[clap(short, long, group("kill"))]
        children: bool,
    },

    /// Send something to a task. Useful for sending confirmations such as 'y\n'.
    Send {
        /// The id of the task.
        task_id: usize,

        /// The input that should be sent to the process.
        input: String,
    },

    /// Edit the command or path of a stashed or queued task.
    /// This edits the command of the task by default.
    #[clap(verbatim_doc_comment)]
    Edit {
        /// The id of the task.
        task_id: usize,

        /// Edit the path of the task.
        #[clap(short, long)]
        path: bool,
    },

    /// Manage groups.
    /// By default, this will simply display all known groups.
    Group {
        /// Add a group
        #[clap(short, long)]
        add: Option<String>,

        /// Remove a group.
        /// This will move all tasks in this group to the default group!
        #[clap(short, long)]
        remove: Option<String>,
    },

    /// Display the current status of all tasks.
    Status {
        /// Print the current state as json to stdout.
        /// This does not include stdout/stderr of tasks.
        /// Use `log -j` if you want everything.
        #[clap(short, long)]
        json: bool,

        #[clap(short, long)]
        /// Only show tasks of a specific group
        group: Option<String>,
    },

    /// Display the log output of finished tasks.
    /// Prints either all logs or only the logs of specified tasks.
    Log {
        /// Specify for which specific tasks you want to see the output.
        task_ids: Vec<usize>,
        /// Print the current state as json.
        /// Includes EVERYTHING.
        #[clap(short, long)]
        json: bool,
    },

    /// Follow the output of a currently running task.
    /// This command works like tail -f.
    Follow {
        /// The id of the task you want to watch.
        /// If no or multiple tasks are running, you have to specify the id.
        /// If only a single task is running, you can omit the id.
        task_id: Option<usize>,

        /// Show stderr instead of stdout.
        #[clap(short, long)]
        err: bool,
    },

    /// Wait until tasks are finished. This can be quite useful for scripting.
    /// By default, this will wait for all tasks in the default queue to finish.
    /// Note: This will also wait for all tasks that aren't somehow 'Done'.
    /// Includes: [Paused, Stashed, Locked, Queued, ...]
    Wait {
        /// This allows you to wait for specific tasks to finish.
        #[clap(short, long)]
        task_ids: Option<Vec<usize>>,

        /// Wait for all tasks in a specific group
        #[clap(short, long)]
        group: Option<String>,

        /// Wait for all tasks across all groups and the default queue.
        #[clap(short, long)]
        all: bool,

        /// Don't show any log output while waiting
        #[clap(short, long)]
        quiet: bool,
    },

    /// Remove all finished tasks from the list (also clears logs).
    Clean,

    /// Kill all running tasks on user behalf, remove all tasks and reset max_task_id.
    Reset {
        /// Send the SIGTERM signal to all children as well.
        /// Useful when working with shell scripts.
        #[clap(short, long)]
        children: bool,
        /// Force killing all the running tasks without confirmation.
        #[clap(short, long)]
        force: bool,
    },

    /// Remotely shut down the daemon. Should only be used if the daemon isn't started by a service manager.
    Shutdown,

    /// Set the amount of allowed parallel tasks.
    Parallel {
        /// The amount of allowed parallel tasks.
        #[clap(validator=min_one)]
        parallel_tasks: usize,

        /// Specify the amount of parallel tasks for a group.
        #[clap(name = "group", short, long)]
        group: Option<String>,
    },

    /// Generates shell completion files.
    /// This can be ignored during normal operations.
    Completions {
        /// The target shell. Can be `bash`, `fish`, `powershell`, `elvish` and `zsh`.
        #[clap(arg_enum)]
        shell: Shell,
        /// The output directory to which the file should be written.
        output_directory: PathBuf,
    },
}

#[derive(Clap, Debug, PartialEq)]
pub enum Shell {
    Bash,
    Elvish,
    Fish,
    PowerShell,
    Zsh,
}

#[derive(Clap, Debug)]
#[clap(
    name = "Pueue client",
    about = "Interact with the Pueue daemon",
    author = env!("CARGO_PKG_AUTHORS"),
    version = env!("CARGO_PKG_VERSION")
)]
pub struct CliArguments {
    /// Verbose mode (-v, -vv, -vvv)
    #[clap(short, long, parse(from_occurrences))]
    pub verbose: u8,

    /// The port for the daemon. Overwrites the port in the config file.
    /// Will force TCP mode.
    #[clap(short, long)]
    pub port: Option<String>,

    /// The path to the unix socket.
    /// Overwrites the path in the config file.
    /// Will force Unix-socket mode.
    #[clap(short, long, conflicts_with = "port")]
    pub unix_socket_path: Option<String>,

    /// Path to a specific pueue config daemon, that should be used.
    /// This ignores all other config files.
    #[clap(short, long)]
    pub config: Option<PathBuf>,

    #[clap(subcommand)]
    pub cmd: SubCommand,
}

fn parse_delay_until(src: &str) -> Result<DateTime<Local>, String> {
    if let Ok(seconds) = src.parse::<i64>() {
        let delay_until = Local::now() + Duration::seconds(seconds);
        return Ok(delay_until);
    }

    if let Ok(date_time) = parse_date_string(src, Local::now(), Dialect::Us) {
        return Ok(date_time);
    }

    Err(String::from(
        "could not parse as seconds or date expression",
    ))
}

/// Validator function. The input string has to be parsable as int and bigger than 0
fn min_one(value: &str) -> Result<(), String> {
    match value.parse::<usize>() {
        Ok(value) => {
            if value < 1 {
                return Err("You must provide a value that's bigger than 0".into());
            }
            Ok(())
        }
        Err(_) => Err("Failed to parse integer".into()),
    }
}
<commit_msg>Add client CLI conflicts and remove unnecessary groups<commit_after>use std::path::PathBuf;

use chrono::prelude::*;
use chrono::Duration;
use chrono_english::*;
use clap::Clap;

#[derive(Clap, Debug)]
pub enum SubCommand {
    /// Enqueue a task for execution.
    Add {
        /// The command that should be added.
        #[clap(required = true)]
        command: Vec<String>,

        /// Start the task immediately.
        #[clap(name = "immediate", short, long, conflicts_with = "stashed")]
        start_immediately: bool,

        /// Create the task in stashed state.
        /// Useful to avoid immediate execution if the queue is empty.
        #[clap(name = "stashed", short, long, conflicts_with = "immediate")]
        stashed: bool,

        /// Delays enqueueing the task until <delay> elapses. See "enqueue" for accepted formats.
        #[clap(name = "delay", short, long, conflicts_with = "immediate", parse(try_from_str=parse_delay_until))]
        delay_until: Option<DateTime<Local>>,

        /// Assign the task to a group. Groups kind of act as separate queues.
        /// I.e. all groups run in parallel and you can specify the amount of parallel tasks for each group.
        /// If no group is specified, the default group will be used.
        #[clap(name = "group", short, long)]
        group: Option<String>,

        /// Start the task once all specified tasks have successfully finished.
        /// As soon as one of the dependencies fails, this task will fail as well.
        #[clap(name = "after", short, long)]
        dependencies: Vec<usize>,
    },
    /// Remove tasks from the list.
    /// Running or paused tasks need to be killed first.
    Remove {
        /// The task ids to be removed.
        #[clap(required = true)]
        task_ids: Vec<usize>,
    },
    /// Switches the queue position of two commands. Only works on queued and stashed commands.
    Switch {
        /// The first task id.
        task_id_1: usize,
        /// The second task id.
        task_id_2: usize,
    },
    /// Stashed tasks won't be automatically started.
    /// Either enqueue them, to be normally handled or explicitly start them.
    Stash {
        /// The id(s) of the tasks you want to stash.
        #[clap(required = true)]
        task_ids: Vec<usize>,
    },
    /// Enqueue stashed tasks. They'll be handled normally afterwards.
    #[clap(after_help = "DELAY FORMAT:

    The --delay argument must be either a number of seconds or a \"date expression\" similar to GNU \
    `date -d` with some extensions. It does not attempt to parse all natural language, but is \
    incredibly flexible. Here are some supported examples.

    2020-04-01T18:30:00   // RFC 3339 timestamp
    2020-4-1 18:2:30      // Optional leading zeros
    2020-4-1 5:30pm       // Informal am/pm time
    2020-4-1 5pm          // Optional minutes and seconds
    April 1 2020 18:30:00 // English months
    1 Apr 8:30pm          // Implies current year
    4/1                   // American form date
    wednesday 10:30pm     // The closest wednesday in the future at 22:30
    wednesday             // The closest wednesday in the future
    4 months              // 4 months from today at 00:00:00
    1 week                // 1 week at the current time
    1days                 // 1 day from today at the current time
    1d 03:00              // The closest 3:00 after 1 day (24 hours)
    3h                    // 3 hours from now
    3600s                 // 3600 seconds from now
")]
    Enqueue {
        /// The id(s) of the tasks you want to enqueue.
        task_ids: Vec<usize>,

        /// Delay enqueuing the tasks until <delay> elapses. See DELAY FORMAT below.
        #[clap(name = "delay", short, long, parse(try_from_str=parse_delay_until))]
        delay_until: Option<DateTime<Local>>,
    },

    /// Resume operation of specific tasks or groups of tasks.
    /// By default, this resumes the default queue and all its tasks.
    /// Can also be used force-start specific tasks.
    #[clap(verbatim_doc_comment)]
    Start {
        /// Enforce starting these tasks. Paused tasks will be started again.
        /// This does not affect anything other than these tasks.
        task_ids: Vec<usize>,

        /// Start a specific group and all paused tasks in it.
        #[clap(short, long, conflicts_with = "all")]
        group: Option<String>,

        /// Start a everything (Default queue and all groups)!
        /// All groups will be set to `running` and all paused tasks will be resumed.
        #[clap(short, long)]
        all: bool,

        /// Also resume direct child processes of your paused tasks.
        /// By default only the main process will get a SIGSTART.
        #[clap(short, long)]
        children: bool,
    },

    /// Restart task(s).
    /// Identical tasks will be created and by default enqueued.
    Restart {
        /// The tasks you want to restart.
        #[clap(required = true)]
        task_ids: Vec<usize>,

        /// Immediately start the task(s).
        #[clap(short, long, name = "immediate", conflicts_with = "stashed")]
        start_immediately: bool,

        /// Create the task in stashed state.
        /// Useful to avoid immediate execution.
        #[clap(short, long)]
        stashed: bool,

        /// Edit the command of the task before restarting
        #[clap(short, long)]
        edit: bool,

        /// Edit the path of the task before restarting
        #[clap(short, long)]
        path: bool,
    },

    /// Pause either running tasks or specific groups of tasks.
    /// By default, pauses the default queue and all its tasks.
    /// A paused queue (group) won't start any new tasks.
    #[clap(verbatim_doc_comment)]
    Pause {
        /// Pause these specific tasks.
        /// Does not affect the default queue, groups or any other tasks.
        task_ids: Vec<usize>,

        /// Pause a specific group.
        #[clap(short, long, conflicts_with = "all")]
        group: Option<String>,

        /// Pause everything (Default queue and all groups)!
        #[clap(short, long)]
        all: bool,

        /// Don not pause already running tasks and let them finish by themselves,
        /// when pausing with `default`, `all` or `group`.
        #[clap(short, long)]
        wait: bool,

        /// Also pause direct child processes of a task's main process.
        /// By default only the main process will get a SIGSTOP.
        /// This is useful when calling bash scripts, which start other processes themselves.
        /// This operation is not recursive!
        #[clap(short, long)]
        children: bool,
    },

    /// Kill specific running tasks or various groups of tasks.
    Kill {
        /// The tasks that should be killed.
        #[clap(short, long)]
        task_ids: Vec<usize>,

        /// Kill all running tasks in the default queue. Pause the default queue.
        #[clap(short, long, conflicts_with = "group", conflicts_with = "all")]
        default: bool,

        /// Kill all running in a group. Pauses the group.
        #[clap(short, long, conflicts_with = "all")]
        group: Option<String>,

        /// Kill ALL running tasks. This also pauses everything
        #[clap(short, long)]
        all: bool,

        /// Send the SIGTERM signal to all children as well.
        /// Useful when working with shell scripts.
        #[clap(short, long)]
        children: bool,
    },

    /// Send something to a task. Useful for sending confirmations such as 'y\n'.
    Send {
        /// The id of the task.
        task_id: usize,

        /// The input that should be sent to the process.
        input: String,
    },

    /// Edit the command or path of a stashed or queued task.
    /// This edits the command of the task by default.
    #[clap(verbatim_doc_comment)]
    Edit {
        /// The id of the task.
        task_id: usize,

        /// Edit the path of the task.
        #[clap(short, long)]
        path: bool,
    },

    /// Manage groups.
    /// By default, this will simply display all known groups.
    Group {
        /// Add a group
        #[clap(short, long, conflicts_with = "remove")]
        add: Option<String>,

        /// Remove a group.
        /// This will move all tasks in this group to the default group!
        #[clap(short, long)]
        remove: Option<String>,
    },

    /// Display the current status of all tasks.
    Status {
        /// Print the current state as json to stdout.
        /// This does not include stdout/stderr of tasks.
        /// Use `log -j` if you want everything.
        #[clap(short, long)]
        json: bool,

        #[clap(short, long)]
        /// Only show tasks of a specific group
        group: Option<String>,
    },

    /// Display the log output of finished tasks.
    /// Prints either all logs or only the logs of specified tasks.
    Log {
        /// Specify for which specific tasks you want to see the output.
        task_ids: Vec<usize>,
        /// Print the current state as json.
        /// Includes EVERYTHING.
        #[clap(short, long)]
        json: bool,
    },

    /// Follow the output of a currently running task.
    /// This command works like tail -f.
    Follow {
        /// The id of the task you want to watch.
        /// If no or multiple tasks are running, you have to specify the id.
        /// If only a single task is running, you can omit the id.
        task_id: Option<usize>,

        /// Show stderr instead of stdout.
        #[clap(short, long)]
        err: bool,
    },

    /// Wait until tasks are finished. This can be quite useful for scripting.
    /// By default, this will wait for all tasks in the default queue to finish.
    /// Note: This will also wait for all tasks that aren't somehow 'Done'.
    /// Includes: [Paused, Stashed, Locked, Queued, ...]
    Wait {
        /// This allows you to wait for specific tasks to finish.
        #[clap(short, long, conflicts_with = "group", conflicts_with = "all")]
        task_ids: Option<Vec<usize>>,

        /// Wait for all tasks in a specific group
        #[clap(short, long, conflicts_with = "all")]
        group: Option<String>,

        /// Wait for all tasks across all groups and the default queue.
        #[clap(short, long)]
        all: bool,

        /// Don't show any log output while waiting
        #[clap(short, long)]
        quiet: bool,
    },

    /// Remove all finished tasks from the list (also clears logs).
    Clean,

    /// Kill all running tasks on user behalf, remove all tasks and reset max_task_id.
    Reset {
        /// Send the SIGTERM signal to all children as well.
        /// Useful when working with shell scripts.
        #[clap(short, long)]
        children: bool,
        /// Force killing all the running tasks without confirmation.
        #[clap(short, long)]
        force: bool,
    },

    /// Remotely shut down the daemon. Should only be used if the daemon isn't started by a service manager.
    Shutdown,

    /// Set the amount of allowed parallel tasks.
    Parallel {
        /// The amount of allowed parallel tasks.
        #[clap(validator=min_one)]
        parallel_tasks: usize,

        /// Specify the amount of parallel tasks for a group.
        #[clap(name = "group", short, long)]
        group: Option<String>,
    },

    /// Generates shell completion files.
    /// This can be ignored during normal operations.
    Completions {
        /// The target shell. Can be `bash`, `fish`, `powershell`, `elvish` and `zsh`.
        #[clap(arg_enum)]
        shell: Shell,
        /// The output directory to which the file should be written.
        output_directory: PathBuf,
    },
}

#[derive(Clap, Debug, PartialEq)]
pub enum Shell {
    Bash,
    Elvish,
    Fish,
    PowerShell,
    Zsh,
}

#[derive(Clap, Debug)]
#[clap(
    name = "Pueue client",
    about = "Interact with the Pueue daemon",
    author = env!("CARGO_PKG_AUTHORS"),
    version = env!("CARGO_PKG_VERSION")
)]
pub struct CliArguments {
    /// Verbose mode (-v, -vv, -vvv)
    #[clap(short, long, parse(from_occurrences))]
    pub verbose: u8,

    /// The port for the daemon. Overwrites the port in the config file.
    /// Will force TCP mode.
    #[clap(short, long)]
    pub port: Option<String>,

    /// The path to the unix socket.
    /// Overwrites the path in the config file.
    /// Will force Unix-socket mode.
    #[clap(short, long, conflicts_with = "port")]
    pub unix_socket_path: Option<String>,

    /// Path to a specific pueue config daemon, that should be used.
    /// This ignores all other config files.
    #[clap(short, long)]
    pub config: Option<PathBuf>,

    #[clap(subcommand)]
    pub cmd: SubCommand,
}

fn parse_delay_until(src: &str) -> Result<DateTime<Local>, String> {
    if let Ok(seconds) = src.parse::<i64>() {
        let delay_until = Local::now() + Duration::seconds(seconds);
        return Ok(delay_until);
    }

    if let Ok(date_time) = parse_date_string(src, Local::now(), Dialect::Us) {
        return Ok(date_time);
    }

    Err(String::from(
        "could not parse as seconds or date expression",
    ))
}

/// Validator function. The input string has to be parsable as int and bigger than 0
fn min_one(value: &str) -> Result<(), String> {
    match value.parse::<usize>() {
        Ok(value) => {
            if value < 1 {
                return Err("You must provide a value that's bigger than 0".into());
            }
            Ok(())
        }
        Err(_) => Err("Failed to parse integer".into()),
    }
}
<|endoftext|>


<commit_before>#![feature(used, lang_items, const_fn)]

extern crate kvm;
extern crate memmap;
extern crate x86;

use kvm::{Capability, Exit, IoDirection, Segment, System, Vcpu, VirtualMachine};
use memmap::{Mmap, Protection};
use std::fs::File;
use std::io::{BufRead, BufReader};

use x86::shared::control_regs::*;
use x86::shared::paging::*;
use x86::bits64::paging::*;

struct PageTable {
    backing_memory: Mmap
}
type PageTableMemoryLayout = (PML4, [PDPT; 512]);
static PAGE_TABLE_P: PAddr = PAddr::from_u64(0x1000); // XXX:

impl PageTable {

    /// Allocated a chunk of memory to contain a basic page-table setup that covers the whole
    /// address space
    fn new() -> PageTable {
        let mut anon_mmap = Mmap::anonymous(3 * (1 << 20), Protection::ReadWrite).unwrap(); // Map 1 MiB
        PageTable { backing_memory: anon_mmap }
    }

    fn as_mut_slice<'a>(&'a mut self) -> &'a mut [u8] {
        unsafe { self.backing_memory.as_mut_slice() }
    }

    /// Constructs a simple page-table that identity maps
    /// the whole address space (guest virtual <-> guest physical).
    fn setup_identity_mapping(&mut self) {
        let page_table_memory = self.as_mut_slice();

        // "physical" layout of PageTable is:
        // PML4Table at 0x1000
        // 1st PDPT at 0x2000
        // 2nd PDPT at 0x3000 ...
        // 512th PDPT at 0x1000+512*0x1000
        // XXX: Can this be simpler why the offset?
        let page_table: &mut PageTableMemoryLayout =
            unsafe { ::std::mem::transmute(&mut page_table_memory[PAGE_TABLE_P.as_u64() as usize]) };

        // Identity map everything in our PML4:
        let (ref mut pml4, ref mut pdpts) = *page_table;
        for (i, mut pdpt) in pdpts.iter_mut().enumerate() {
            let offset = 0x2000 + 0x1000 * i;
            pml4[i] = PML4Entry::new(PAddr::from_u64(offset as _), PML4_P | PML4_RW);

            for j in 0..512 {
                // Set-up 1 GiB page-mappings:
                pdpt[j] = PDPTEntry::new(PAddr::from_u64(((512 * i + j) as u64) << 30),
                                         PDPT_P | PDPT_RW | PDPT_PS);
            }
        }
    }
}

struct Stack {
     backing_memory: Mmap
}
static STACK_BASE_T: PAddr = PAddr::from_u64(0x2000000);

impl Stack {

    /// Allocated a chunk of memory to contain a basic page-table setup that covers the whole
    /// address space
    fn new() -> Stack {
        let mut stack_mmap = Mmap::anonymous(65536, Protection::ReadWrite).unwrap();
        Stack { backing_memory: stack_mmap }
    }

    fn as_mut_slice<'a>(&'a mut self) -> &'a mut [u8]{
        unsafe { self.backing_memory.as_mut_slice() }
    }

    fn size(&self) -> usize {
        self.backing_memory.len()
    }

    fn base(&self) -> VAddr {
        VAddr::from_usize(self.backing_memory.ptr() as usize)
    }
}

struct TestEnvironment<'a> {
    sys: &'a System,
    pt: &'a mut PageTable,
    st: &'a mut Stack,
    vm: VirtualMachine<'a>,
}

impl<'a> TestEnvironment<'a> {

    fn new(sys: &'a System, st: &'a mut Stack, pt: &'a mut PageTable) -> TestEnvironment<'a> {
        let mut vm = VirtualMachine::create(sys).unwrap();
        // Ensure that the VM supports memory backing with user memory
        assert!(vm.check_capability(Capability::UserMemory) > 0);

        /*let mut pt = PageTable::new();
        pt.setup_identity_mapping();
        let mut st = Stack::new();*/

        TestEnvironment { pt: pt, st: st, sys: sys, vm: vm }
    }

    /// Map the page table memory and stack memory
    fn map_memory(mut self) {
        self.vm.set_user_memory_region(0, self.pt.as_mut_slice(), 0).unwrap();
        self.vm.set_user_memory_region(STACK_BASE_T.as_u64(), self.st.as_mut_slice(), 0).unwrap();
    }
}

#[derive(Debug)]
pub struct KvmTestMetaData {
    pub mbz: u64,
    pub meta: &'static str,
    pub identity_map: bool,
    pub physical_memory: (u64, u64),
}

/// Linker generates symbols that are inserted at the start and end of the kvm section.
extern "C" {
	static __start_kvm: std::os::raw::c_void;
	static __stop_kvm: std::os::raw::c_void;
}

/// Walks the kvm section to find see if there is a metha-data struct
/// lying around for the given test name.
fn find_meta_data(name: &str) -> Option<&KvmTestMetaData> {
    let (baseptr, size);

    // Safe: the linker will take care of initializing these symbols
	unsafe {
		baseptr = &__start_kvm as *const _ as *const KvmTestMetaData;
		size = &__stop_kvm as *const _ as usize - baseptr as usize;
	}

	let count = size / std::mem::size_of::<KvmTestMetaData>();

    // Safe: The section points to (static) KvmTestMetaData descriptions
	unsafe {
		let mods = std::slice::from_raw_parts(baseptr, count);
        for m in mods.iter() {
            if m.meta == name {
                return Some(&m);
            }
        }
	}

    None
}

/// Start the test harness.
pub fn test_start(ntests: usize) {
    println!("running {} tests (using KVM support)", ntests)
}

/// Signals that given test is ignored.
pub fn test_ignored(name: &str) {
    println!("test {} ... ignored", name);
}

/*
fn io_example() {

    // Initialize the KVM system
    let sys = System::initialize().unwrap();

    // Create a Virtual Machine
    let mut vm = VirtualMachine::create(&sys).unwrap();

    // Ensure that the VM supports memory backing with user memory
    assert!(vm.check_capability(Capability::UserMemory) > 0);

    // Once the memory is set we can't even call length.
    let page_table_memory_limit = page_table_memory.len() - 1;

    // Map the page table memory
    vm.set_user_memory_region(0, page_table_memory, 0).unwrap();
    // Map stack space
    vm.set_user_memory_region(STACK_BASE_T.as_u64(), stack_memory, 0).unwrap();

    // Map the process
    let f = File::open("/proc/self/maps").unwrap();
    let reader = BufReader::new(f);

    for line in reader.lines() {
        let line = line.unwrap();
        println!("{}", line);
        let mut s = line.split(' ');
        let mut s2 = s.next().unwrap().split('-');
        let begin = usize::from_str_radix(s2.next().unwrap(), 16).unwrap();
        let end = usize::from_str_radix(s2.next().unwrap(), 16).unwrap();
        if end < 0x800000000000 {
            let perm = s.next().unwrap();
            //println!("{:#X}-{:#X} {}", begin, end, perm);
            let slice = {
                let begin_ptr: *mut u8 = begin as *const u8 as _;
                unsafe { ::std::slice::from_raw_parts_mut(begin_ptr, end - begin) }
            };
            // Make sure process doesn't overlap with page table
            assert!(begin > page_table_memory_limit);
            vm.set_user_memory_region(begin as _, slice, 0).unwrap();
        }
    }

    let mut vcpu = Vcpu::create(&mut vm).unwrap();
    // Set supported CPUID (KVM fails without doing this)
    let mut cpuid = sys.get_supported_cpuid().unwrap();
    vcpu.set_cpuid2(&mut cpuid).unwrap();

    // Setup the special registers
    let mut sregs = vcpu.get_sregs().unwrap();
    // Set the code segment to have base 0, limit 4GB (flat segmentation)
    let segment_template = Segment {
        base: 0x0,
        limit: 0xffffffff,
        selector: 0,
        _type: 0,
        present: 0,
        dpl: 0,
        db: 1,
        s: 0,
        l: 0,
        g: 1,
        avl: 0,
        ..Default::default()
    };
    sregs.cs = Segment {
        selector: 0x8,
        _type: 0xb,
        present: 1,
        db: 0,
        s: 1,
        l: 1,
        ..segment_template
    };
    sregs.ss = Segment { ..segment_template };
    sregs.ds = Segment { ..segment_template };
    sregs.es = Segment { ..segment_template };
    sregs.fs = Segment { ..segment_template };
    sregs.gs = Segment { ..segment_template };
    // We don't need to populate the GDT if we have our segments setup
    // cr0 - protected mode on, paging enabled
    sregs.cr0 = (CR0_PROTECTED_MODE | CR0_MONITOR_COPROCESSOR | CR0_EXTENSION_TYPE |
                 CR0_ENABLE_PAGING | CR0_NUMERIC_ERROR | CR0_WRITE_PROTECT |
                 CR0_ALIGNMENT_MASK | CR0_ENABLE_PAGING)
        .bits() as u64;
    sregs.cr3 = PAGE_TABLE_P.as_u64();
    sregs.cr4 = (CR4_ENABLE_PSE | CR4_ENABLE_PAE | CR4_ENABLE_GLOBAL_PAGES | CR4_ENABLE_SSE |
                 CR4_UNMASKED_SSE |
                 CR4_ENABLE_OS_XSAVE | CR4_ENABLE_SMEP | CR4_ENABLE_VME)
        .bits() as u64;
    sregs.efer = 0xd01;

    // Set the special registers
    vcpu.set_sregs(&sregs).unwrap();

    let mut regs = vcpu.get_regs().unwrap();
    // set the instruction pointer to 1 MB
    regs.rip = vaddr.as_usize() as u64;
    println!("regs.rip = 0x{:x}", regs.rip); // but is at: 0x40cd60
    //println!("regs.rip = 0x{:x}", unsafe { *(regs.rip as *const u64) }); // but is at: 0x40cd60
    //regs.rip = 0x40cd60;
    regs.rflags = 0x246;
    regs.rsp = STACK_BASE_T.as_u64() + stack_size as u64;
    regs.rbp = regs.rsp;
    vcpu.set_regs(&regs).unwrap();

    // Actually run the VCPU

    println!("size of run {:?}", std::mem::size_of::<kvm::Run>());
    let mut vm_is_done = false;
    let mut new_regs = kvm::Regs::default();
    while !vm_is_done {
        {
            let (run, mut regs) = unsafe { vcpu.run_regs() }.unwrap();
            match run.exit_reason {
                Exit::Io => {
                    let io = unsafe { *run.io() };
                    match io.direction {
                        IoDirection::In => {
                            if io.port == 0x3fd {
                                regs.rax = 0x20; // Mark serial line as ready to write
                            } else {
                                println!("IO on unknown port: {}", io.port);
                            }
                        }
                        IoDirection::Out => {
                            if io.port == 0x3f8 {
                                println!("got char {:#?}", regs.rax as u8 as char);
                            }
                        }
                    }
                }
                Exit::Shutdown => {
                    println!("Shutting down");
                    vm_is_done = true;
                }
                _ => {
                    println!("Unknown exit reason: {:?}", run.exit_reason);
                }
            }

            new_regs = regs;
        }
        vcpu.set_regs(&new_regs).unwrap();
    }

}*/

pub fn test_before_run(name: &str) {
    print!("test {} ... ", name);

    let meta_data = find_meta_data(name).unwrap();

}

pub fn test_panic_fmt(args: std::fmt::Arguments, file: &'static str, line: u32) {
    print!("\npanicked at '");
    use std::io::Write;
    std::io::stderr().write_fmt(args);
    println!("', {}:{}", file, line);
}

pub fn test_failed(_name: &str) {
    println!("FAILED");
}

pub fn test_success(_name: &str) {
    println!("OK");
}

pub fn test_summary(passed: usize, failed: usize, ignored: usize) {
    println!("\ntest result: {} {} passed; {} failed; {} ignored",
             if failed == 0 { "OK" } else { "FAILED" },
             passed,
             failed,
             ignored);

    if failed != 0 {
        std::process::exit(101);
    }
}

#[no_mangle]
#[used]
pub static mut __TEST_PANICKED: bool = false;

pub fn test_main_static(tests: &[TestDescAndFn]) {
    unsafe {
        test_start(tests.len());

        let mut failed = 0;
        let mut ignored = 0;
        let mut passed = 0;
        for test in tests {
            if test.desc.ignore {
                ignored += 1;
                test_ignored(test.desc.name.0);
            } else {
                test_before_run(test.desc.name.0);

                __TEST_PANICKED = false;

                test.testfn.0();

                if __TEST_PANICKED == (test.desc.should_panic == ShouldPanic::Yes) {
                    passed += 1;
                    test_success(test.desc.name.0);
                } else {
                    failed += 1;
                    test_failed(test.desc.name.0);
                }
            }

        }

        test_summary(passed, failed, ignored);
    }
}

// required for compatibility with the `rustc --test` interface
pub struct TestDescAndFn {
    pub desc: TestDesc,
    pub testfn: StaticTestFn,
}

pub struct TestDesc {
    pub ignore: bool,
    pub name: StaticTestName,
    pub should_panic: ShouldPanic,
}

pub struct StaticTestName(pub &'static str);
pub struct StaticTestFn(pub fn());

#[derive(PartialEq)]
pub enum ShouldPanic {
    No,
    Yes,
}
<commit_msg>Merging KVM with test harness functionality #20.<commit_after>#![feature(used, lang_items, const_fn)]

extern crate kvm;
extern crate memmap;
extern crate x86;

use kvm::{Capability, Exit, IoDirection, Segment, System, Vcpu, VirtualMachine};
use memmap::{Mmap, Protection};
use std::fs::File;
use std::io::{BufRead, BufReader};

use x86::shared::control_regs::*;
use x86::shared::paging::*;
use x86::bits64::paging::*;

struct PageTable {
    backing_memory: Mmap
}
type PageTableMemoryLayout = (PML4, [PDPT; 512]);
static PAGE_TABLE_P: PAddr = PAddr::from_u64(0x1000); // XXX:

impl PageTable {

    /// Allocated a chunk of memory to contain a basic page-table setup that covers the whole
    /// address space
    fn new() -> PageTable {
        let mut anon_mmap = Mmap::anonymous(3 * (1 << 20), Protection::ReadWrite).unwrap(); // Map 1 MiB
        PageTable { backing_memory: anon_mmap }
    }

    fn as_mut_slice<'a>(&'a mut self) -> &'a mut [u8] {
        unsafe { self.backing_memory.as_mut_slice() }
    }

    fn as_slice(&mut self) -> &[u8] {
        unsafe { self.backing_memory.as_slice() }
    }

    fn len(&mut self) -> usize {
        self.as_slice().len()
    }

    /// Constructs a simple page-table that identity maps
    /// the whole address space (guest virtual <-> guest physical).
    fn setup_identity_mapping(&mut self) {
        let page_table_memory = self.as_mut_slice();

        // "physical" layout of PageTable is:
        // PML4Table at 0x1000
        // 1st PDPT at 0x2000
        // 2nd PDPT at 0x3000 ...
        // 512th PDPT at 0x1000+512*0x1000
        // XXX: Can this be simpler why the offset?
        let page_table: &mut PageTableMemoryLayout =
            unsafe { ::std::mem::transmute(&mut page_table_memory[PAGE_TABLE_P.as_u64() as usize]) };

        // Identity map everything in our PML4:
        let (ref mut pml4, ref mut pdpts) = *page_table;
        for (i, mut pdpt) in pdpts.iter_mut().enumerate() {
            let offset = 0x2000 + 0x1000 * i;
            pml4[i] = PML4Entry::new(PAddr::from_u64(offset as _), PML4_P | PML4_RW);

            for j in 0..512 {
                // Set-up 1 GiB page-mappings:
                pdpt[j] = PDPTEntry::new(PAddr::from_u64(((512 * i + j) as u64) << 30),
                                         PDPT_P | PDPT_RW | PDPT_PS);
            }
        }
    }
}

struct Stack {
     backing_memory: Mmap
}
static STACK_BASE_T: PAddr = PAddr::from_u64(0x2000000);

impl Stack {

    /// Allocated a chunk of memory to contain a basic page-table setup that covers the whole
    /// address space
    fn new() -> Stack {
        let mut stack_mmap = Mmap::anonymous(65536, Protection::ReadWrite).unwrap();
        Stack { backing_memory: stack_mmap }
    }

    fn len(&self) -> usize {
        self.backing_memory.len()
    }

    fn as_mut_slice<'a>(&'a mut self) -> &'a mut [u8]{
        unsafe { self.backing_memory.as_mut_slice() }
    }

    fn size(&self) -> usize {
        self.backing_memory.len()
    }

    fn base(&self) -> VAddr {
        VAddr::from_usize(self.backing_memory.ptr() as usize)
    }
}

struct TestEnvironment<'a> {
    sys: &'a System,
    pt: &'a mut PageTable,
    st: &'a mut Stack,
    vm: VirtualMachine<'a>,
}

impl<'a> TestEnvironment<'a> {

    fn new(sys: &'a System, st: &'a mut Stack, pt: &'a mut PageTable) -> TestEnvironment<'a> {
        let mut vm = VirtualMachine::create(sys).unwrap();
        // Ensure that the VM supports memory backing with user memory
        assert!(vm.check_capability(Capability::UserMemory) > 0);

        /*let mut pt = PageTable::new();
        pt.setup_identity_mapping();
        let mut st = Stack::new();*/

        TestEnvironment { pt: pt, st: st, sys: sys, vm: vm }
    }

    /// Map the page table memory and stack memory
    fn create_vcpu(&'a mut self, init_fn: VAddr) -> kvm::Vcpu {
        // Once the memory is set we can't even call length.
        let page_table_memory_limit = self.pt.len() - 1;
        let stack_size = self.st.len();

        self.vm.set_user_memory_region(0, self.pt.as_mut_slice(), 0).unwrap();
        self.vm.set_user_memory_region(STACK_BASE_T.as_u64(), self.st.as_mut_slice(), 0).unwrap();

        // Map the process
        let f = File::open("/proc/self/maps").unwrap();
        let reader = BufReader::new(f);

        for line in reader.lines() {
            let line = line.unwrap();
            let mut s = line.split(' ');
            let mut s2 = s.next().unwrap().split('-');
            let begin = usize::from_str_radix(s2.next().unwrap(), 16).unwrap();
            let end = usize::from_str_radix(s2.next().unwrap(), 16).unwrap();
            if end < 0x800000000000 {
                let perm = s.next().unwrap();
                //println!("{:#X}-{:#X} {}", begin, end, perm);
                let slice = {
                    let begin_ptr: *mut u8 = begin as *const u8 as _;
                    unsafe { ::std::slice::from_raw_parts_mut(begin_ptr, end - begin) }
                };
                // Make sure process doesn't overlap with page table
                //assert!(begin > page_table_memory_limit);
                self.vm.set_user_memory_region(begin as _, slice, 0).unwrap();
            }
        }

        let mut vcpu = Vcpu::create(&mut self.vm).unwrap();
        // Set supported CPUID (KVM fails without doing this)
        let mut cpuid = self.sys.get_supported_cpuid().unwrap();
        vcpu.set_cpuid2(&mut cpuid).unwrap();

        // Setup the special registers
        let mut sregs = vcpu.get_sregs().unwrap();

        // Set the code segment to have base 0, limit 4GB (flat segmentation)
        let segment_template = Segment {
            base: 0x0,
            limit: 0xffffffff,
            selector: 0,
            _type: 0,
            present: 0,
            dpl: 0,
            db: 1,
            s: 0,
            l: 0,
            g: 1,
            avl: 0,
            ..Default::default()
        };

        sregs.cs = Segment {
            selector: 0x8,
            _type: 0xb,
            present: 1,
            db: 0,
            s: 1,
            l: 1,
            ..segment_template
        };
        sregs.ss = Segment { ..segment_template };
        sregs.ds = Segment { ..segment_template };
        sregs.es = Segment { ..segment_template };
        sregs.fs = Segment { ..segment_template };
        sregs.gs = Segment { ..segment_template };

        // We don't need to populate the GDT if we have our segments setup
        // cr0 - protected mode on, paging enabled
        sregs.cr0 = (CR0_PROTECTED_MODE | CR0_MONITOR_COPROCESSOR | CR0_EXTENSION_TYPE |
                     CR0_ENABLE_PAGING | CR0_NUMERIC_ERROR | CR0_WRITE_PROTECT |
                     CR0_ALIGNMENT_MASK | CR0_ENABLE_PAGING).bits() as u64;
        sregs.cr3 = PAGE_TABLE_P.as_u64();
        sregs.cr4 = (CR4_ENABLE_PSE | CR4_ENABLE_PAE | CR4_ENABLE_GLOBAL_PAGES | CR4_ENABLE_SSE |
                     CR4_UNMASKED_SSE | CR4_ENABLE_OS_XSAVE | CR4_ENABLE_SMEP | CR4_ENABLE_VME)
                     .bits() as u64;
        sregs.efer = 0xd01; // XXX

        // Set the special registers
        vcpu.set_sregs(&sregs).unwrap();

        let mut regs = vcpu.get_regs().unwrap();

        // Set the instruction pointer to 1 MB
        regs.rip = init_fn.as_usize() as u64;
        regs.rflags = 0x246; // XXX
        regs.rsp = STACK_BASE_T.as_u64() + stack_size as u64;
        regs.rbp = regs.rsp;

        vcpu.set_regs(&regs).unwrap();

        vcpu
    }
}

#[derive(Debug)]
pub struct KvmTestMetaData {
    pub mbz: u64,
    pub meta: &'static str,
    pub identity_map: bool,
    pub physical_memory: (u64, u64),
}

/// Linker generates symbols that are inserted at the start and end of the kvm section.
extern "C" {
	static __start_kvm: std::os::raw::c_void;
	static __stop_kvm: std::os::raw::c_void;
}

/// Walks the kvm section to find see if there is a metha-data struct
/// lying around for the given test name.
fn find_meta_data(name: &str) -> Option<&KvmTestMetaData> {
    let (baseptr, size);

    // Safe: the linker will take care of initializing these symbols
	unsafe {
		baseptr = &__start_kvm as *const _ as *const KvmTestMetaData;
		size = &__stop_kvm as *const _ as usize - baseptr as usize;
	}

	let count = size / std::mem::size_of::<KvmTestMetaData>();

    // Safe: The section points to (static) KvmTestMetaData descriptions
	unsafe {
		let mods = std::slice::from_raw_parts(baseptr, count);
        for m in mods.iter() {
            if m.meta == name {
                return Some(&m);
            }
        }
	}

    None
}

/// Start the test harness.
pub fn test_start(ntests: usize) {
    println!("running {} tests (using KVM support)", ntests)
}

/// Signals that given test is ignored.
pub fn test_ignored(name: &str) {
    println!("test {} ... ignored", name);
}

pub fn test_before_run(name: &str) -> Option<&KvmTestMetaData> {
    print!("test {} ... ", name);
    find_meta_data(name)
}

pub fn test_panic_fmt(args: std::fmt::Arguments, file: &'static str, line: u32) {
    print!("\npanicked at '");
    use std::io::Write;
    std::io::stderr().write_fmt(args);
    println!("', {}:{}", file, line);
}

pub fn test_failed(_name: &str) {
    println!("FAILED");
}

pub fn test_success(_name: &str) {
    println!("OK");
}

pub fn test_summary(passed: usize, failed: usize, ignored: usize) {
    println!("\ntest result: {} {} passed; {} failed; {} ignored",
             if failed == 0 { "OK" } else { "FAILED" },
             passed,
             failed,
             ignored);

    if failed != 0 {
        std::process::exit(101);
    }
}

#[no_mangle]
#[used]
pub static mut __TEST_PANICKED: bool = false;

pub fn test_main_static(tests: &[TestDescAndFn]) {
    unsafe {
        test_start(tests.len());

        let mut failed = 0;
        let mut ignored = 0;
        let mut passed = 0;
        for test in tests {
            if test.desc.ignore {
                ignored += 1;
                test_ignored(test.desc.name.0);
            } else {
                let meta_data = test_before_run(test.desc.name.0);

                __TEST_PANICKED = false;
                match meta_data {
                    Some(mtd) => {
                        let sys = System::initialize().unwrap();
                        let mut st = Stack::new();
                        let mut pt = PageTable::new();
                        pt.setup_identity_mapping();

                        let mut test_environment = TestEnvironment::new(&sys, &mut st, &mut pt);
                        let test_fn_vaddr = VAddr::from_usize(test.testfn.0 as *const () as usize);
                        let mut vcpu = test_environment.create_vcpu(test_fn_vaddr);

                        let mut vm_is_done = false;
                        let mut new_regs = kvm::Regs::default();
                        while !vm_is_done {
                            {
                                let (run, mut regs) = unsafe { vcpu.run_regs() }.unwrap();
                                match run.exit_reason {
                                    Exit::Io => {
                                        let io = unsafe { *run.io() };
                                        match io.direction {
                                            IoDirection::In => {
                                                if io.port == 0x3fd {
                                                    regs.rax = 0x20; // Mark serial line as ready to write
                                                } else {
                                                    println!("IO on unknown port: {}", io.port);
                                                }
                                            }
                                            IoDirection::Out => {
                                                if io.port == 0x3f8 {
                                                    println!("got char {:#?}", regs.rax as u8 as char);
                                                }
                                            }
                                        }
                                    }
                                    Exit::Shutdown => {
                                        println!("Shutting down");
                                        vm_is_done = true;
                                    }
                                    _ => {
                                        println!("Unknown exit reason: {:?}", run.exit_reason);
                                    }
                                }

                                new_regs = regs;
                            }
                            vcpu.set_regs(&new_regs).unwrap();
                        }

                    },
                    _ => test.testfn.0() // Regular test, not running inside virtual machine
                }

                if __TEST_PANICKED == (test.desc.should_panic == ShouldPanic::Yes) {
                    passed += 1;
                    test_success(test.desc.name.0);
                } else {
                    failed += 1;
                    test_failed(test.desc.name.0);
                }
            }

        }

        test_summary(passed, failed, ignored);
    }
}

// required for compatibility with the `rustc --test` interface
pub struct TestDescAndFn {
    pub desc: TestDesc,
    pub testfn: StaticTestFn,
}

pub struct TestDesc {
    pub ignore: bool,
    pub name: StaticTestName,
    pub should_panic: ShouldPanic,
}

pub struct StaticTestName(pub &'static str);
pub struct StaticTestFn(pub fn());

#[derive(PartialEq)]
pub enum ShouldPanic {
    No,
    Yes,
}
<|endoftext|>


<commit_before>use async_trait::async_trait;
use cucumber::{t, StepContext, Steps};
use futures::stream::{self, TryStreamExt};
use snafu::ResultExt;
use std::path::{Path, PathBuf};
use tokio::fs;
use tokio::io::AsyncWriteExt;

use crate::error;
use crate::error::Error;
use crate::state::{State, Step, StepStatus};
use crate::utils::{create_dir_if_not_exists_rec, file_exists};

const GEOFABRIK_URL: &str = "https://download.geofabrik.de";
const BANO_URL: &str = "http://bano.openstreetmap.fr";

pub fn steps() -> Steps<State> {
    let mut steps: Steps<State> = Steps::new();

    steps.given_regex_async(
        "osm file has been downloaded for (.*)",
        t!(|mut state, ctx| {
            let region = ctx.matches[1].clone();

            state
                .execute_once(DownloadOsm(region), &ctx)
                .await
                .expect("failed to download OSM file");

            state
        }),
    );

    steps.given_regex_async(
        "bano files have been downloaded for (.*) into (.*)",
        t!(|mut state, ctx| {
            let departments = ctx.matches[1]
                .split(',')
                .map(|s| s.trim().to_string())
                .collect();

            let region = ctx.matches[2].clone();

            state
                .execute_once(
                    DownloadBano {
                        departments,
                        region,
                    },
                    &ctx,
                )
                .await
                .expect("failed to download OSM file");

            state
        }),
    );

    steps
}

/// Downloads the file identified by the url and saves it to the given path.
/// If a file is already present, it will append to that file.
async fn download_to_file(path: &Path, url: &str) -> Result<(), Error> {
    let mut file = tokio::io::BufWriter::new({
        fs::OpenOptions::new()
            .append(true)
            .create(true)
            .open(&path)
            .await
            .context(error::InvalidIO {
                details: format!("could no create file for download {}", path.display()),
            })?
    });

    let mut resp = reqwest::get(url)
        .await
        .context(error::Download {
            details: format!("could not download url {}", url),
        })?
        .error_for_status()
        .context(error::Download {
            details: format!("download response error for {}", url),
        })?;

    while let Some(chunk) = resp.chunk().await.context(error::Download {
        details: format!("read chunk error during download of {}", url),
    })? {
        file.write_all(&chunk).await.context(error::InvalidIO {
            details: format!("write chunk error during download of {}", url),
        })?;
    }

    file.flush().await.context(error::InvalidIO {
        details: format!("flush error during download of {}", url),
    })?;

    Ok(())
}

/// Given the name of a french region, it will download the matching OSM file
/// If that file is already in the local filesystem, then we skip the download.
///
/// This makes several assumptions:
///  1. The name of the region is one found in http://download.geofabrik.de/europe/france.html
///  2. The file will be downloaded to `tests/data/osm` under the project's root (identified
///     by the CARGO_MANIFEST_DIR environment variable
#[derive(PartialEq)]
pub struct DownloadOsm(pub String);

#[async_trait(?Send)]
impl Step for DownloadOsm {
    async fn execute(&mut self, _state: &State, _ctx: &StepContext) -> Result<StepStatus, Error> {
        let Self(region) = self;

        // Try to see if there is already a file with the expected name in tests/data/osm, in which
        // case we skip the actual download, to save time.
        let dir_path: PathBuf = [env!("CARGO_MANIFEST_DIR"), "tests", "data", "osm"]
            .iter()
            .collect();

        let file_path = dir_path.join(format!("{}-latest.osm.pbf", region));

        if file_exists(&file_path).await {
            return Ok(StepStatus::Skipped);
        }

        // No file in 'cache', so we download it
        create_dir_if_not_exists_rec(&dir_path).await?;
        let url = format!("{}/europe/france/{}-latest.osm.pbf", GEOFABRIK_URL, region);
        download_to_file(&file_path, &url).await?;

        Ok(StepStatus::Done)
    }
}

/// Given a list of French departments, it will download the matching BANO files.
/// If these files are already in the local filesystem, then we skip the download.
/// Then we concatenate these files into a single file with the name of the region.
/// The reason for this is that we want at the indexing stage to check that admins
/// have been indexed prior to indexing addresses, and so we need the same name for
/// the bano region and the osm region.
///
/// This makes several assumptions:
///  1. The file will be downloaded to `tests/data/bano` under the project's root (identified
///     by the CARGO_MANIFEST_DIR environment variable
#[derive(PartialEq)]
pub struct DownloadBano {
    pub departments: Vec<String>,
    pub region: String,
}

#[async_trait(?Send)]
impl Step for DownloadBano {
    // FIXME: The only StepStatus returned is StepStatus::Done. Not handling very well if there
    // are files downloaded before.
    // This function will not stop if one of the download fails, but it will report an error.
    async fn execute(&mut self, _state: &State, _ctx: &StepContext) -> Result<StepStatus, Error> {
        let Self {
            departments,
            region,
        } = self;
        let dir_path: PathBuf = [env!("CARGO_MANIFEST_DIR"), "tests", "data", "bano"]
            .iter()
            .collect();

        create_dir_if_not_exists_rec(&dir_path).await?;

        // this is the path for the concatenated departments
        let file_path = dir_path.join(format!("{}.csv", region));

        if file_exists(&file_path).await {
            return Ok(StepStatus::Skipped);
        }

        let res: Result<Vec<()>, Error> = stream::try_unfold(
            BanoState {
                departments: departments.to_vec(),
                index: 0,
            },
            |state| async {
                if state.departments.len() > state.index {
                    let department = state.departments[state.index].clone();
                    let next_state = BanoState {
                        departments: state.departments,
                        index: state.index + 1,
                    };
                    let url = format!("{}/data/bano-{:02}.csv", BANO_URL, department);
                    download_to_file(&file_path, &url).await?;
                    Ok(Some(((), next_state)))
                } else {
                    Ok(None)
                }
            },
        )
        .try_collect()
        .await;

        res.map(|_| StepStatus::Done)
    }
}

struct BanoState {
    departments: Vec<String>,
    index: usize, // index into the departments
}
<commit_msg>Update tests/steps/download.rs<commit_after>use async_trait::async_trait;
use cucumber::{t, StepContext, Steps};
use futures::stream::{self, TryStreamExt};
use snafu::ResultExt;
use std::path::{Path, PathBuf};
use tokio::fs;
use tokio::io::AsyncWriteExt;

use crate::error;
use crate::error::Error;
use crate::state::{State, Step, StepStatus};
use crate::utils::{create_dir_if_not_exists_rec, file_exists};

const GEOFABRIK_URL: &str = "https://download.geofabrik.de";
const BANO_URL: &str = "http://bano.openstreetmap.fr";

pub fn steps() -> Steps<State> {
    let mut steps: Steps<State> = Steps::new();

    steps.given_regex_async(
        "osm file has been downloaded for (.*)",
        t!(|mut state, ctx| {
            let region = ctx.matches[1].clone();

            state
                .execute_once(DownloadOsm(region), &ctx)
                .await
                .expect("failed to download OSM file");

            state
        }),
    );

    steps.given_regex_async(
        "bano files have been downloaded for (.*) into (.*)",
        t!(|mut state, ctx| {
            let departments = ctx.matches[1]
                .split(',')
                .map(|s| s.trim().to_string())
                .collect();

            let region = ctx.matches[2].clone();

            state
                .execute_once(
                    DownloadBano {
                        departments,
                        region,
                    },
                    &ctx,
                )
                .await
                .expect("failed to download OSM file");

            state
        }),
    );

    steps
}

/// Downloads the file identified by the url and saves it to the given path.
/// If a file is already present, it will append to that file.
async fn download_to_file(path: &Path, url: &str) -> Result<(), Error> {
    let mut file = tokio::io::BufWriter::new({
        fs::OpenOptions::new()
            .append(true)
            .create(true)
            .open(&path)
            .await
            .context(error::InvalidIO {
                details: format!("could no create file for download {}", path.display()),
            })?
    });

    let mut resp = reqwest::get(url)
        .await
        .context(error::Download {
            details: format!("could not download url {}", url),
        })?
        .error_for_status()
        .context(error::Download {
            details: format!("download response error for {}", url),
        })?;

    while let Some(chunk) = resp.chunk().await.context(error::Download {
        details: format!("read chunk error during download of {}", url),
    })? {
        file.write_all(&chunk).await.context(error::InvalidIO {
            details: format!("write chunk error during download of {}", url),
        })?;
    }

    file.flush().await.context(error::InvalidIO {
        details: format!("flush error during download of {}", url),
    })?;

    Ok(())
}

/// Given the name of a french region, it will download the matching OSM file
/// If that file is already in the local filesystem, then we skip the download.
///
/// This makes several assumptions:
///  1. The name of the region is one found in http://download.geofabrik.de/europe/france.html
///  2. The file will be downloaded to `tests/data/osm` under the project's root (identified
///     by the CARGO_MANIFEST_DIR environment variable
#[derive(PartialEq)]
pub struct DownloadOsm(pub String);

#[async_trait(?Send)]
impl Step for DownloadOsm {
    async fn execute(&mut self, _state: &State, _ctx: &StepContext) -> Result<StepStatus, Error> {
        let Self(region) = self;

        // Try to see if there is already a file with the expected name in tests/data/osm, in which
        // case we skip the actual download, to save time.
        let dir_path: PathBuf = [env!("CARGO_MANIFEST_DIR"), "tests", "data", "osm"]
            .iter()
            .collect();

        let file_path = dir_path.join(format!("{}-latest.osm.pbf", region));

        if file_exists(&file_path).await {
            return Ok(StepStatus::Skipped);
        }

        // No file in 'cache', so we download it
        create_dir_if_not_exists_rec(&dir_path).await?;
        let url = format!("{}/europe/france/{}-latest.osm.pbf", GEOFABRIK_URL, region);
        download_to_file(&file_path, &url).await?;

        Ok(StepStatus::Done)
    }
}

/// Given a list of French departments, it will download the matching BANO files.
/// If these files are already in the local filesystem, then we skip the download.
/// Then we concatenate these files into a single file with the name of the region.
/// The reason for this is that we want at the indexing stage to check that admins
/// have been indexed prior to indexing addresses, and so we need the same name for
/// the bano region and the osm region.
///
/// This makes several assumptions:
///  1. The file will be downloaded to `tests/data/bano` under the project's root (identified
///     by the CARGO_MANIFEST_DIR environment variable
#[derive(PartialEq)]
pub struct DownloadBano {
    pub departments: Vec<String>,
    pub region: String,
}

#[async_trait(?Send)]
impl Step for DownloadBano {
    // FIXME: The only StepStatus returned is StepStatus::Done. Not handling very well if there
    // are files downloaded before.
    // This function will not stop if one of the download fails, but it will report an error.
    async fn execute(&mut self, _state: &State, _ctx: &StepContext) -> Result<StepStatus, Error> {
        let Self {
            departments,
            region,
        } = self;
        let dir_path: PathBuf = [env!("CARGO_MANIFEST_DIR"), "tests", "data", "bano"]
            .iter()
            .collect();

        create_dir_if_not_exists_rec(&dir_path).await?;

        // this is the path for the concatenated departments
        let file_path = &dir_path.join(format!("{}.csv", region));

        if file_exists(file_path).await {
            return Ok(StepStatus::Skipped);
        }

        stream::iter(departments.iter().map(Ok))
            .try_for_each(|department| async move {
                let url = format!("{}/data/bano-{:02}.csv", BANO_URL, department);
                download_to_file(file_path, &url).await
            })
            .await?;

        Ok(StepStatus::Done)
    }
}
<|endoftext|>


<commit_before>extern crate futures;

use futures::Future;
use futures::unsync::oneshot::{channel, Canceled};

#[test]
fn smoke() {
    let (tx, rx) = channel();
    tx.complete(33);
    assert_eq!(rx.wait().unwrap(), 33);
}

#[test]
fn canceled() {
    let (_, rx) = channel::<()>();
    assert_eq!(rx.wait().unwrap_err(), Canceled);
}

#[test]
fn poll_cancel() {
    let (mut tx, _) = channel::<()>();
    assert!(tx.poll_cancel().unwrap().is_ready());
}

#[test]
fn tx_complete_rx_unparked() {
    let (tx, rx) = channel();

    let res = rx.join(::futures::lazy(move || {
        tx.complete(55);
        Ok(11)
    }));
    assert_eq!(res.wait().unwrap(), (55, 11));
}

#[test]
fn tx_dropped_rx_unparked() {
    let (tx, rx) = channel::<i32>();

    let res = rx.join(::futures::lazy(move || {
        let _tx = tx;
        Ok(11)
    }));
    assert_eq!(res.wait().unwrap_err(), Canceled);
}
<commit_msg>Fix unsync-oneshot tests<commit_after>extern crate futures;

use futures::Future;
use futures::future;
use futures::unsync::oneshot::{channel, Canceled};

#[test]
fn smoke() {
    let (tx, rx) = channel();
    tx.complete(33);
    assert_eq!(rx.wait().unwrap(), 33);
}

#[test]
fn canceled() {
    let (_, rx) = channel::<()>();
    assert_eq!(rx.wait().unwrap_err(), Canceled);
}

#[test]
fn poll_cancel() {
    let (mut tx, _) = channel::<()>();
    assert!(tx.poll_cancel().unwrap().is_ready());
}

#[test]
fn tx_complete_rx_unparked() {
    let (tx, rx) = channel();

    let res = rx.join(future::lazy(move || {
        tx.complete(55);
        Ok(11)
    }));
    assert_eq!(res.wait().unwrap(), (55, 11));
}

#[test]
fn tx_dropped_rx_unparked() {
    let (tx, rx) = channel::<i32>();

    let res = rx.join(future::lazy(move || {
        let _tx = tx;
        Ok(11)
    }));
    assert_eq!(res.wait().unwrap_err(), Canceled);
}
<|endoftext|>


<commit_before>extern crate nom;
extern crate classfile_parser;

use nom::IResult;

use classfile_parser::class_parser;
use classfile_parser::constant_info::ConstantInfo;

#[test]
fn test_attribute_bootstrap_methods() {
    let stack_map_class = include_bytes!("../java-assets/compiled-classes/BootstrapMethods.class");
    let res = class_parser(stack_map_class);
    match res {
        IResult::Done(_, c) => {
            println!("Valid class file, version {},{} const_pool({}), this=const[{}], super=const[{}], interfaces({}), fields({}), methods({}), attributes({}), access({:?})", c.major_version, c.minor_version, c.const_pool_size, c.this_class, c.super_class, c.interfaces_count, c.fields_count, c.methods_count, c.attributes_count, c.access_flags);

            let mut bootstrap_method_const_index = 0;

            println!("Constant pool:");
            for (const_index, const_item) in c.const_pool.iter().enumerate() {
                println!("\t[{}] = {:?}", (const_index + 1), const_item);
                match *const_item {
                    ConstantInfo::Utf8(ref c) => {
                        if c.utf8_string == "BootstrapMethods" {
                            bootstrap_method_const_index = (const_index + 1) as u16;
                        }
                    },
                    _ => {},
                }
            }
            println!("Bootstrap Methods constant index = {}", bootstrap_method_const_index);

            println!("Attributes:");
            for (attribute_index, attribute_item) in c.attributes.iter().enumerate() {
                println!("\t[{}] = {:?}", (attribute_index + 1), attribute_item);
                if attribute_item.attribute_name_index == bootstrap_method_const_index {
                    use classfile_parser::attribute_info::bootstrap_methods_attribute_parser;
                    let bsm = bootstrap_methods_attribute_parser(&attribute_item.info);
                    println!("\t\t{:?}", bsm);
                }
            }
        },
        _ => panic!("Not a class file"),
    }
}
<commit_msg>Cleaned up Bootstrap Methods attribute test for issue #9<commit_after>extern crate nom;
extern crate classfile_parser;

use nom::IResult;

use classfile_parser::class_parser;
use classfile_parser::constant_info::ConstantInfo;
use classfile_parser::attribute_info::bootstrap_methods_attribute_parser;

#[test]
fn test_attribute_bootstrap_methods() {
    match class_parser(include_bytes!("../java-assets/compiled-classes/BootstrapMethods.class")) {
        IResult::Done(_, c) => {
            println!("Valid class file, version {},{} const_pool({}), this=const[{}], super=const[{}], interfaces({}), fields({}), methods({}), attributes({}), access({:?})", c.major_version, c.minor_version, c.const_pool_size, c.this_class, c.super_class, c.interfaces_count, c.fields_count, c.methods_count, c.attributes_count, c.access_flags);

            let mut bootstrap_method_const_index = 0;

            println!("Constant pool:");
            for (const_index, const_item) in c.const_pool.iter().enumerate() {
                println!("\t[{}] = {:?}", (const_index + 1), const_item);
                match *const_item {
                    ConstantInfo::Utf8(ref c) => {
                        if c.utf8_string == "BootstrapMethods" {
                            if bootstrap_method_const_index != 0 {
                                assert!(false, "Should not find more than one BootstrapMethods constant");
                            }
                            bootstrap_method_const_index = (const_index + 1) as u16;
                        }
                    }
                    _ => {}
                }
            }
            assert_ne!(bootstrap_method_const_index, 0);

            println!("Bootstrap Methods constant index = {}", bootstrap_method_const_index);

            for (_, attribute_item) in c.attributes.iter().enumerate() {
                if attribute_item.attribute_name_index == bootstrap_method_const_index {
                    match bootstrap_methods_attribute_parser(&attribute_item.info) {
                        IResult::Done(_, bsma) => {
                            assert_eq!(bsma.num_bootstrap_methods, 1);
                            let bsm = &bsma.bootstrap_methods[0];
                            assert_eq!(bsm.bootstrap_method_ref, 36);

                            println!("{:?}", bsm);
                            println!("\tmethod ref: {:?}", c.const_pool[36]);
                            println!("\t\tdescriptor: {:?}", c.const_pool[53]);
                            println!("\t\t\tclass_index: {:?}", c.const_pool[9]);
                            println!("\t\t\t\tname_index: {:?}", c.const_pool[51]);
                            println!("\t\t\t\t\tclass_index: {:?}", c.const_pool[64]);
                            println!("\t\t\t\t\tname_and_type_index: {:?}", c.const_pool[65]);
                            println!("\t\t\t\t\t\tname_index: {:?}", c.const_pool[30]);
                            println!("\t\t\t\t\t\tdescriptor_index: {:?}", c.const_pool[31]);
                            println!("\t\t\tname_and_type_index: {:?}", c.const_pool[66]);
                            return;
                        }
                        _ => panic!("Failed to parse bootstrap method attribute"),
                    }
                }
            }

            assert!(false, "Should not get to here");
        }
        _ => panic!("Not a valid class file"),
    }
}

#[test]
fn should_have_no_bootstrap_method_attr_if_no_invoke_dynamic() {
    match class_parser(include_bytes!("../java-assets/compiled-classes/BasicClass.class")) {
        IResult::Done(_, c) => {
            for (_, const_item) in c.const_pool.iter().enumerate() {
                match *const_item {
                    ConstantInfo::Utf8(ref c) => {
                        if c.utf8_string == "BootstrapMethods" {
                            assert!(false, "Should not have found a BootstrapMethods constant in a class not requiring it")
                        }
                    }
                    _ => {}
                }
            }
        }
        _ => panic!("Not a valid class file"),
    }
}
<|endoftext|>


<commit_before>/// Pins a value on the stack.
///
/// Calls to `async fn` return anonymous [`Future`] values that are `!Unpin`.
/// These values must be pinned before they can be polled. Calling `.await` will
/// handle this, but consumes the future. If it is required to call `.await` on
/// a `&mut _` reference, the caller is responsible for pinning the future.
///
/// Pinning may be done by allocating with [`Box::pin`] or by using the stack
/// with the `pin!` macro.
///
/// The following will **fail to compile**:
///
/// ```compile_fail
/// async fn my_async_fn() {
///     // async logic here
/// }
///
/// #[tokio::main]
/// async fn main() {
///     let mut future = my_async_fn();
///     (&mut future).await;
/// }
/// ```
///
/// To make this work requires pinning:
///
/// ```
/// use tokio::pin;
///
/// async fn my_async_fn() {
///     // async logic here
/// }
///
/// #[tokio::main]
/// async fn main() {
///     let future = my_async_fn();
///     pin!(future);
///
///     (&mut future).await;
/// }
/// ```
///
/// Pinning is useful when using `select!` and stream operators that require `T:
/// Stream + Unpin`.
///
/// [`Future`]: trait@std::future::Future
/// [`Box::pin`]: #
///
/// # Usage
///
/// The `pin!` macro takes **identifiers** as arguments. It does **not** work
/// with expressions.
///
/// The following does not compile as an expression is passed to `pin!`.
///
/// ```compile_fail
/// async fn my_async_fn() {
///     // async logic here
/// }
///
/// #[tokio::main]
/// async fn main() {
///     let mut future = pin!(my_async_fn());
///     (&mut future).await;
/// }
/// ```
///
/// # Examples
///
/// Using with select:
///
/// ```
/// use tokio::{pin, select};
/// use tokio_stream::{self as stream, StreamExt};
///
/// async fn my_async_fn() {
///     // async logic here
/// }
///
/// #[tokio::main]
/// async fn main() {
///     let mut stream = stream::iter(vec![1, 2, 3, 4]);
///
///     let future = my_async_fn();
///     pin!(future);
///
///     loop {
///         select! {
///             _ = &mut future => {
///                 // Stop looping `future` will be polled after completion
///                 break;
///             }
///             Some(val) = stream.next() => {
///                 println!("got value = {}", val);
///             }
///         }
///     }
/// }
/// ```
///
/// Because assigning to a variable followed by pinning is common, there is also
/// a variant of the macro that supports doing both in one go.
///
/// ```
/// use tokio::{pin, select};
///
/// async fn my_async_fn() {
///     // async logic here
/// }
///
/// #[tokio::main]
/// async fn main() {
///     pin! {
///         let future1 = my_async_fn();
///         let future2 = my_async_fn();
///     }
///
///     select! {
///         _ = &mut future1 => {}
///         _ = &mut future2 => {}
///     }
/// }
/// ```
#[macro_export]
macro_rules! pin {
    ($($x:ident),*) => { $(
        // Move the value to ensure that it is owned
        let mut $x = $x;
        // Shadow the original binding so that it can't be directly accessed
        // ever again.
        #[allow(unused_mut)]
        let mut $x = unsafe {
            $crate::macros::support::Pin::new_unchecked(&mut $x)
        };
    )* };
    ($(
            let $x:ident = $init:expr;
    )*) => {
        $(
            let $x = $init;
            $crate::pin!($x);
        )*
    };
}
<commit_msg>doc: fix incorrect link in doc (#3670)<commit_after>/// Pins a value on the stack.
///
/// Calls to `async fn` return anonymous [`Future`] values that are `!Unpin`.
/// These values must be pinned before they can be polled. Calling `.await` will
/// handle this, but consumes the future. If it is required to call `.await` on
/// a `&mut _` reference, the caller is responsible for pinning the future.
///
/// Pinning may be done by allocating with [`Box::pin`] or by using the stack
/// with the `pin!` macro.
///
/// The following will **fail to compile**:
///
/// ```compile_fail
/// async fn my_async_fn() {
///     // async logic here
/// }
///
/// #[tokio::main]
/// async fn main() {
///     let mut future = my_async_fn();
///     (&mut future).await;
/// }
/// ```
///
/// To make this work requires pinning:
///
/// ```
/// use tokio::pin;
///
/// async fn my_async_fn() {
///     // async logic here
/// }
///
/// #[tokio::main]
/// async fn main() {
///     let future = my_async_fn();
///     pin!(future);
///
///     (&mut future).await;
/// }
/// ```
///
/// Pinning is useful when using `select!` and stream operators that require `T:
/// Stream + Unpin`.
///
/// [`Future`]: trait@std::future::Future
/// [`Box::pin`]: std::boxed::Box::pin
///
/// # Usage
///
/// The `pin!` macro takes **identifiers** as arguments. It does **not** work
/// with expressions.
///
/// The following does not compile as an expression is passed to `pin!`.
///
/// ```compile_fail
/// async fn my_async_fn() {
///     // async logic here
/// }
///
/// #[tokio::main]
/// async fn main() {
///     let mut future = pin!(my_async_fn());
///     (&mut future).await;
/// }
/// ```
///
/// # Examples
///
/// Using with select:
///
/// ```
/// use tokio::{pin, select};
/// use tokio_stream::{self as stream, StreamExt};
///
/// async fn my_async_fn() {
///     // async logic here
/// }
///
/// #[tokio::main]
/// async fn main() {
///     let mut stream = stream::iter(vec![1, 2, 3, 4]);
///
///     let future = my_async_fn();
///     pin!(future);
///
///     loop {
///         select! {
///             _ = &mut future => {
///                 // Stop looping `future` will be polled after completion
///                 break;
///             }
///             Some(val) = stream.next() => {
///                 println!("got value = {}", val);
///             }
///         }
///     }
/// }
/// ```
///
/// Because assigning to a variable followed by pinning is common, there is also
/// a variant of the macro that supports doing both in one go.
///
/// ```
/// use tokio::{pin, select};
///
/// async fn my_async_fn() {
///     // async logic here
/// }
///
/// #[tokio::main]
/// async fn main() {
///     pin! {
///         let future1 = my_async_fn();
///         let future2 = my_async_fn();
///     }
///
///     select! {
///         _ = &mut future1 => {}
///         _ = &mut future2 => {}
///     }
/// }
/// ```
#[macro_export]
macro_rules! pin {
    ($($x:ident),*) => { $(
        // Move the value to ensure that it is owned
        let mut $x = $x;
        // Shadow the original binding so that it can't be directly accessed
        // ever again.
        #[allow(unused_mut)]
        let mut $x = unsafe {
            $crate::macros::support::Pin::new_unchecked(&mut $x)
        };
    )* };
    ($(
            let $x:ident = $init:expr;
    )*) => {
        $(
            let $x = $init;
            $crate::pin!($x);
        )*
    };
}
<|endoftext|>


<commit_before><commit_msg>add args type for view()<commit_after><|endoftext|>


<commit_before>use std::io::Write;

/// A P-Data value writer.
///
/// This exposes an API to iteratively construct and send Data messages
/// to another node.
/// Using this as a [standard writer](std::io::Write)
/// will automatically split the incoming bytes
/// into separate PDUs if they do not fit in a single one.
///
/// Use an association's `send_pdata` method
/// to create a new P-Data value writer.
#[must_use]
pub struct PDataWriter<W: Write> {
    buffer: Vec<u8>,
    stream: W,
    presentation_context_id: u8,
    max_data_length: u32,
}

impl<W> PDataWriter<W>
where
    W: Write,
{
    /// Construct a new P-Data value writer.
    pub(crate) fn new(stream: W, presentation_context_id: u8, max_pdu_length: u32) -> Self {
        let max_data_length = calculate_max_data_len_single(max_pdu_length);
        PDataWriter {
            stream,
            presentation_context_id,
            max_data_length,
            buffer: Vec::with_capacity(max_data_length as usize),
        }
    }

    /// Send the header of a single P-Data PDU,
    /// containing a single data fragment.
    fn send_pdata_header(&mut self, data_len: u32, is_last: bool) -> std::io::Result<()> {
        let mut message_header = 0x00;
        if is_last {
            message_header |= 0x02;
        }

        let pdu_len_bytes = (data_len + 2 + 4).to_be_bytes();
        let data_len_bytes = (data_len + 2).to_be_bytes();
        let header = [
            // PDU-type + reserved byte
            0x04,
            0x00,
            // full PDU length
            pdu_len_bytes[0],
            pdu_len_bytes[1],
            pdu_len_bytes[2],
            pdu_len_bytes[3],
            // presentation data length (data + 2 properties below)
            data_len_bytes[0],
            data_len_bytes[1],
            data_len_bytes[2],
            data_len_bytes[3],
            // presentation context id
            self.presentation_context_id,
            // message control header
            message_header,
        ];

        self.stream.write_all(&header)
    }

    fn finish(&mut self) -> std::io::Result<()> {
        self.dispatch_excess_data()?;
        if !self.buffer.is_empty() {
            // send last PDU
            self.send_pdata_header(self.buffer.len() as u32, true)?;
            self.stream.write_all(&self.buffer[..])?;
            self.buffer.clear();
        }
        Ok(())
    }

    fn dispatch_excess_data(&mut self) -> std::io::Result<()> {
        while self.buffer.len() > self.max_data_length as usize {
            // send PDU now
            self.send_pdata_header(self.max_data_length, false)?;
            let data = &self.buffer[0..self.max_data_length as usize];
            self.stream.write_all(data)?;

            // shift the remaining contents to the beginning of the buffer
            let (p1, p2) = (&mut self.buffer[..]).split_at_mut(self.max_data_length as usize);
            for (e1, e2) in std::iter::Iterator::zip(p1.iter_mut(), p2.iter()) {
                *e1 = *e2;
            }
            self.buffer
                .truncate(self.buffer.len() - self.max_data_length as usize);
        }
        Ok(())
    }
}

impl<W> Write for PDataWriter<W>
where
    W: Write,
{
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.buffer.extend(buf);
        self.dispatch_excess_data()?;
        Ok(buf.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        // do nothing
        Ok(())
    }
}

impl<W> Drop for PDataWriter<W>
where
    W: Write,
{
    fn drop(&mut self) {
        let _ = self.finish();
    }
}

/// Determine the maximum length of actual data
/// when encapsulated in a PDU with the given length property.
/// Does not account for the first 2 bytes (type + reserved).
#[inline]
fn calculate_max_data_len_single(pdu_len: u32) -> u32 {
    pdu_len
    // data length
    - 4
    // control header
    - 2
}

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

    use crate::pdu::PDataValueType;
    use crate::pdu::Pdu;
    use crate::pdu::{read_pdu, reader::MINIMUM_PDU_SIZE};

    use super::PDataWriter;

    #[test]
    fn test_write_pdata_and_finish() {
        let presentation_context_id = 12;

        let mut buf = Vec::new();
        {
            let mut writer = PDataWriter::new(&mut buf, presentation_context_id, MINIMUM_PDU_SIZE);
            writer.write_all(&(0..64).collect::<Vec<u8>>()).unwrap();
            writer.finish().unwrap();
        }

        let mut cursor = &buf[..];
        let same_pdu = read_pdu(&mut cursor, MINIMUM_PDU_SIZE, true).unwrap();

        // concatenate data chunks, compare with all data

        match same_pdu {
            Pdu::PData { data: data_1 } => {
                let data_1 = &data_1[0];

                // check that this PDU is consistent
                assert_eq!(data_1.value_type, PDataValueType::Data);
                assert_eq!(data_1.presentation_context_id, presentation_context_id);
                assert_eq!(data_1.data.len(), 64);
                assert_eq!(data_1.data, (0..64).collect::<Vec<u8>>());
            }
            pdu => panic!("Expected PData, got {:?}", pdu),
        }

        assert_eq!(cursor.len(), 0);
    }

    #[test]
    fn test_write_large_pdata_and_finish() {
        let presentation_context_id = 32;

        let my_data: Vec<_> = (0..6000).map(|x| x as u8).collect();
        assert_eq!(my_data.len(), 6000);

        let mut buf = Vec::new();
        {
            let mut writer = PDataWriter::new(&mut buf, presentation_context_id, MINIMUM_PDU_SIZE);
            writer.write_all(&my_data).unwrap();
            writer.finish().unwrap();
        }

        let mut cursor = &buf[..];
        let same_pdu_1 = read_pdu(&mut cursor, MINIMUM_PDU_SIZE, true).unwrap();
        let same_pdu_2 = read_pdu(&mut cursor, MINIMUM_PDU_SIZE, true).unwrap();

        // concatenate data chunks, compare with all data

        match (same_pdu_1, same_pdu_2) {
            (Pdu::PData { data: data_1 }, Pdu::PData { data: data_2 }) => {
                let data_1 = &data_1[0];
                let data_2 = &data_2[0];

                // check that these two PDUs are consistent
                assert_eq!(data_1.value_type, PDataValueType::Data);
                assert_eq!(data_2.value_type, PDataValueType::Data);
                assert_eq!(data_1.presentation_context_id, presentation_context_id);
                assert_eq!(data_2.presentation_context_id, presentation_context_id);
                assert_eq!(data_1.data.len() + data_2.data.len(), 6000);

                let data_1 = &data_1.data;
                let data_2 = &data_2.data;

                let mut all_data: Vec<u8> = Vec::new();
                all_data.extend(data_1);
                all_data.extend(data_2);
                assert_eq!(all_data, my_data);
            }
            x => panic!("Expected two PDatas, got {:?}", x),
        }

        assert_eq!(cursor.len(), 0);
    }
}
<commit_msg>Fix PDataWriter for large payloads (#164)<commit_after>use std::io::Write;

use crate::pdu::reader::PDU_HEADER_SIZE;

/// A P-Data value writer.
///
/// This exposes an API to iteratively construct and send Data messages
/// to another node.
/// Using this as a [standard writer](std::io::Write)
/// will automatically split the incoming bytes
/// into separate PDUs if they do not fit in a single one.
///
/// # Example
///
/// Use an association's `send_pdata` method
/// to create a new P-Data value writer.
///
/// ```no_run
/// # use std::io::Write;
/// # use dicom_ul::association::{ClientAssociationOptions, PDataWriter};
/// # use dicom_ul::pdu::{Pdu, PDataValue, PDataValueType};
/// # fn command_data() -> Vec<u8> { unimplemented!() }
/// # fn dicom_data() -> &'static [u8] { unimplemented!() }
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let mut association = ClientAssociationOptions::new()
///    .establish("129.168.0.5:104")?;
///
/// let presentation_context_id = association.presentation_contexts()[0].id;
///
/// // send a command first
/// association.send(&Pdu::PData {
///     data: vec![PDataValue {
///     presentation_context_id,
///     value_type: PDataValueType::Command,
///         is_last: true,
///         data: command_data(),
///     }],
/// });
///
/// // then send a DICOM object which may be split into multiple PDUs
/// let mut pdata = association.send_pdata(presentation_context_id);
/// pdata.write_all(dicom_data())?;
/// pdata.finish()?;
///
/// let pdu_ac = association.receive()?;
/// # Ok(())
/// # }
#[must_use]
pub struct PDataWriter<W: Write> {
    buffer: Vec<u8>,
    stream: W,
    presentation_context_id: u8,
    max_data_len: u32,
}

impl<W> PDataWriter<W>
where
    W: Write,
{
    /// Construct a new P-Data value writer.
    ///
    /// `max_pdu_length` is the maximum value of the PDU-length property.
    pub(crate) fn new(stream: W, presentation_context_id: u8, max_pdu_length: u32) -> Self {
        let max_data_length = calculate_max_data_len_single(max_pdu_length);
        PDataWriter {
            stream,
            presentation_context_id,
            max_data_len: max_data_length,
            buffer: Vec::with_capacity(max_data_length as usize),
        }
    }

    /// Declare to have finished sending P-Data fragments,
    /// thus emitting the last P-Data fragment PDU.
    ///
    /// This is also done automatically once the P-Data writer is dropped.
    pub fn finish(mut self) -> std::io::Result<()> {
        self.finish_impl()?;
        Ok(())
    }

    fn create_pdata_header(&self, data_len: u32, is_last: bool) -> [u8; 12] {
        let pdu_len = data_len + 4 + 2;
        let pdv_data_len = data_len + 2;
        let pdu_len_bytes = pdu_len.to_be_bytes();
        let data_len_bytes = pdv_data_len.to_be_bytes();
        [
            // PDU-type + reserved byte
            0x04,
            0x00,
            // full PDU length (minus PDU type and reserved byte)
            pdu_len_bytes[0],
            pdu_len_bytes[1],
            pdu_len_bytes[2],
            pdu_len_bytes[3],
            // presentation data length (data + 2 properties below)
            data_len_bytes[0],
            data_len_bytes[1],
            data_len_bytes[2],
            data_len_bytes[3],
            // presentation context id
            self.presentation_context_id,
            // message control header
            if is_last { 0x02 } else { 0x00 },
        ]
    }

    /// Send the header of a single P-Data PDU,
    /// containing a single data fragment.
    fn send_pdata_header(&mut self, data_len: u32, is_last: bool) -> std::io::Result<()> {
        self.stream
            .write_all(&self.create_pdata_header(data_len, is_last))
    }

    fn finish_impl(&mut self) -> std::io::Result<()> {
        self.dispatch_excess_data()?;
        if !self.buffer.is_empty() {
            // send last PDU
            self.send_pdata_header(self.buffer.len() as u32, true)?;
            self.stream.write_all(&self.buffer[..])?;
            // clear buffer so that subsequent calls to `finish_impl`
            // do not send any more PDUs
            self.buffer.clear();
        }
        Ok(())
    }

    /// Use the current state of the buffer to send new PDUs
    ///
    /// Pre-condition:
    /// buffer must have enough data for one full P-Data-tf PDU
    fn dispatch_pdus(&mut self) -> std::io::Result<()> {
        debug_assert!(self.buffer.len() >= self.max_data_len as usize);
        // send PDU now

        let mut acc_len = 0;
        for chunk in self.buffer.chunks(self.max_data_len as usize) {
            if chunk.len() < self.max_data_len as usize {
                break;
            }
            let header = self.create_pdata_header(self.max_data_len, false);
            self.stream.write_all(&header)?;
            self.stream.write_all(chunk)?;
            acc_len += chunk.len();
        }

        // shift the remaining contents to the beginning of the buffer
        // (guaranteed to be shorter than `self.max_data_len`)
        let (p1, p2) = (&mut self.buffer[..]).split_at_mut(acc_len);
        for (e1, e2) in std::iter::Iterator::zip(p1.iter_mut(), p2.iter()) {
            *e1 = *e2;
        }
        self.buffer.truncate(self.buffer.len() - acc_len as usize);

        Ok(())
    }

    fn dispatch_excess_data(&mut self) -> std::io::Result<()> {
        while self.buffer.len() > self.max_data_len as usize {
            self.dispatch_pdus()?;
        }
        Ok(())
    }
}

impl<W> Write for PDataWriter<W>
where
    W: Write,
{
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.buffer.extend(buf);
        self.dispatch_excess_data()?;
        Ok(buf.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        // do nothing
        Ok(())
    }
}

/// With the P-Data writer dropped,
/// this `Drop` implementation
/// will construct and emit the last P-Data fragment PDU
/// if there is any data left to send.
impl<W> Drop for PDataWriter<W>
where
    W: Write,
{
    fn drop(&mut self) {
        let _ = self.finish_impl();
    }
}

/// Determine the maximum length of actual PDV data
/// when encapsulated in a PDU with the given length property.
/// Does not account for the first 2 bytes (type + reserved).
#[inline]
fn calculate_max_data_len_single(pdu_len: u32) -> u32 {
    // data length: 4 bytes
    // control header: 2 bytes
    pdu_len - PDU_HEADER_SIZE
}

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

    use crate::pdu::reader::{read_pdu, MINIMUM_PDU_SIZE, PDU_HEADER_SIZE};
    use crate::pdu::PDataValueType;
    use crate::pdu::Pdu;

    use super::PDataWriter;

    #[test]
    fn test_write_pdata_and_finish() {
        let presentation_context_id = 12;

        let mut buf = Vec::new();
        {
            let mut writer = PDataWriter::new(&mut buf, presentation_context_id, MINIMUM_PDU_SIZE);
            writer.write_all(&(0..64).collect::<Vec<u8>>()).unwrap();
            writer.finish().unwrap();
        }

        let mut cursor = &buf[..];
        let same_pdu = read_pdu(&mut cursor, MINIMUM_PDU_SIZE, true).unwrap();

        // concatenate data chunks, compare with all data

        match same_pdu {
            Pdu::PData { data: data_1 } => {
                let data_1 = &data_1[0];

                // check that this PDU is consistent
                assert_eq!(data_1.value_type, PDataValueType::Data);
                assert_eq!(data_1.presentation_context_id, presentation_context_id);
                assert_eq!(data_1.data.len(), 64);
                assert_eq!(data_1.data, (0..64).collect::<Vec<u8>>());
            }
            pdu => panic!("Expected PData, got {:?}", pdu),
        }

        assert_eq!(cursor.len(), 0);
    }

    #[test]
    fn test_write_large_pdata_and_finish() {
        let presentation_context_id = 32;

        let my_data: Vec<_> = (0..9000).map(|x: u32| x as u8).collect();

        let mut buf = Vec::new();
        {
            let mut writer = PDataWriter::new(&mut buf, presentation_context_id, MINIMUM_PDU_SIZE);
            writer.write_all(&my_data).unwrap();
            writer.finish().unwrap();
        }

        let mut cursor = &buf[..];
        let pdu_1 = read_pdu(&mut cursor, MINIMUM_PDU_SIZE, true).unwrap();
        let pdu_2 = read_pdu(&mut cursor, MINIMUM_PDU_SIZE, true).unwrap();
        let pdu_3 = read_pdu(&mut cursor, MINIMUM_PDU_SIZE, true).unwrap();

        // concatenate data chunks, compare with all data

        match (pdu_1, pdu_2, pdu_3) {
            (
                Pdu::PData { data: data_1 },
                Pdu::PData { data: data_2 },
                Pdu::PData { data: data_3 },
            ) => {
                assert_eq!(data_1.len(), 1);
                let data_1 = &data_1[0];
                assert_eq!(data_2.len(), 1);
                let data_2 = &data_2[0];
                assert_eq!(data_3.len(), 1);
                let data_3 = &data_3[0];

                // check that these two PDUs are consistent
                assert_eq!(data_1.value_type, PDataValueType::Data);
                assert_eq!(data_2.value_type, PDataValueType::Data);
                assert_eq!(data_1.presentation_context_id, presentation_context_id);
                assert_eq!(data_2.presentation_context_id, presentation_context_id);

                // check expected lengths
                assert_eq!(
                    data_1.data.len(),
                    (MINIMUM_PDU_SIZE - PDU_HEADER_SIZE) as usize
                );
                assert_eq!(
                    data_2.data.len(),
                    (MINIMUM_PDU_SIZE - PDU_HEADER_SIZE) as usize
                );
                assert_eq!(data_3.data.len(), 820);

                // check data consistency
                assert_eq!(
                    &data_1.data[..],
                    (0..MINIMUM_PDU_SIZE - PDU_HEADER_SIZE).map(|x| x as u8).collect::<Vec<_>>()
                );
                assert_eq!(
                    data_1.data.len() + data_2.data.len() + data_3.data.len(),
                    9000
                );

                let data_1 = &data_1.data;
                let data_2 = &data_2.data;
                let data_3 = &data_3.data;

                let mut all_data: Vec<u8> = Vec::new();
                all_data.extend(data_1);
                all_data.extend(data_2);
                all_data.extend(data_3);
                assert_eq!(all_data, my_data);
            }
            x => panic!("Expected 3 PDatas, got {:?}", x),
        }

        assert_eq!(cursor.len(), 0);
    }
}
<|endoftext|>