<commit_before>use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{error, info, warn};
use winapi::shared::minwindef::FALSE;
use winapi::shared::ntdef::NULL;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::processthreadsapi::{
    OpenProcess, OpenThread, ResumeThread, SuspendThread, TerminateProcess,
};
use winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Process32First, Process32Next, Thread32First, Thread32Next,
    PROCESSENTRY32, TH32CS_SNAPPROCESS, TH32CS_SNAPTHREAD, THREADENTRY32,
};
use winapi::um::winnt::{PROCESS_TERMINATE, THREAD_SUSPEND_RESUME};

use crate::task_handler::ProcessAction;
use pueue_lib::network::message::Signal as InternalSignal;

pub fn compile_shell_command(command_string: &str) -> Command {
    // Chain two `powershell` commands, one that sets the output encoding to utf8 and then the user provided one.
    let mut command = Command::new("powershell");
    command.arg("-c").arg(format!(
        "[Console]::OutputEncoding = [Text.UTF8Encoding]::UTF8; {}",
        command_string
    ));

    command
}

pub fn send_internal_signal_to_child(
    child: &Child,
    signal: InternalSignal,
    send_to_children: bool,
) -> Result<bool> {
    bail!("Trying to send unix signal on a windows machine. This isn't supported.");
}

/// Send a signal to a windows process.
pub fn run_action_on_child(child: &Child, action: &ProcessAction, _children: bool) -> Result<bool> {
    let pids = get_cur_task_processes(child.id());
    if pids.is_empty() {
        bail!("Process has just gone away");
    }

    match action {
        ProcessAction::Pause => {
            for pid in pids {
                for thread in get_threads(pid) {
                    suspend_thread(thread);
                }
            }
        }
        ProcessAction::Resume => {
            for pid in pids {
                for thread in get_threads(pid) {
                    resume_thread(thread);
                }
            }
        }
    }

    Ok(true)
}

/// Kill a child process
pub fn kill_child(task_id: usize, child: &mut Child, _kill_children: bool) -> bool {
    match child.kill() {
        Err(_) => {
            info!("Task {} has already finished by itself", task_id);
            false
        }
        Ok(_) => {
            let pids = get_cur_task_processes(child.id());

            for pid in pids {
                terminate_process(pid);
            }
            true
        }
    }
}

/// Get current task pid, all child pid and all children's children
fn get_cur_task_processes(task_pid: u32) -> Vec<u32> {
    let mut all_pids = Vec::new();

    // Get all pids by BFS
    let mut parent_pids = vec![task_pid];
    while let Some(pid) = parent_pids.pop() {
        all_pids.push(pid);

        get_child_pids(pid, &mut parent_pids);
    }

    // Keep parent pid ahead of child. We need execute action for parent process first.
    all_pids.reverse();
    all_pids
}

/// Get child pids of a specific process.
fn get_child_pids(target_pid: u32, pid_list: &mut Vec<u32>) {
    unsafe {
        // Take a snapshot of all processes in the system.
        // While enumerating the set of processes, new processes can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, target_pid);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return;
        }

        // Walk the list of processes.
        let mut process_entry = PROCESSENTRY32 {
            dwSize: std::mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };
        if Process32First(snapshot_handle, &mut process_entry) == FALSE {
            error!("Couldn't get first process.");
            CloseHandle(snapshot_handle);
            return;
        }

        loop {
            if process_entry.th32ParentProcessID == target_pid {
                pid_list.push(process_entry.th32ProcessID);
            }

            if Process32Next(snapshot_handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }
}

/// Get all thread id of a specific process
fn get_threads(target_pid: u32) -> Vec<u32> {
    let mut threads = Vec::new();

    unsafe {
        // Take a snapshot of all threads in the system.
        // While enumerating the set of threads, new threads can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return threads;
        }

        // Walk the list of threads.
        let mut thread_entry = THREADENTRY32 {
            dwSize: std::mem::size_of::<THREADENTRY32>() as u32,
            ..Default::default()
        };
        if Thread32First(snapshot_handle, &mut thread_entry) == FALSE {
            error!("Couldn't get first thread.");
            CloseHandle(snapshot_handle);
            return threads;
        }

        loop {
            if thread_entry.th32OwnerProcessID == target_pid {
                threads.push(thread_entry.th32ThreadID);
            }

            if Thread32Next(snapshot_handle, &mut thread_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }

    threads
}

/// Suspend a thread
/// Each thread has a suspend count (with a maximum value of `MAXIMUM_SUSPEND_COUNT`).
/// If the suspend count is greater than zero, the thread is suspended; otherwise, the thread is not suspended and is eligible for execution.
/// Calling `SuspendThread` causes the target thread's suspend count to be incremented.
/// Attempting to increment past the maximum suspend count causes an error without incrementing the count.
/// [SuspendThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-suspendthread)
fn suspend_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If SuspendThread fails, the return value is (DWORD) -1
            if u32::max_value() == SuspendThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to suspend thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Resume a thread
/// ResumeThread checks the suspend count of the subject thread.
/// If the suspend count is zero, the thread is not currently suspended. Otherwise, the subject thread's suspend count is decremented.
/// If the resulting value is zero, then the execution of the subject thread is resumed.
/// [ResumeThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-resumethread)
fn resume_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If ResumeThread fails, the return value is (DWORD) -1
            if u32::max_value() == ResumeThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to resume thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Terminate a process
/// [TerminateProcess](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-terminateprocess)
fn terminate_process(pid: u32) {
    unsafe {
        // Get a handle for the target process
        let process_handle = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
        // If TerminateProcess fails, the return value is zero.
        if 0 == TerminateProcess(process_handle, 1) {
            let err_code = GetLastError();
            warn!(
                "Failed to terminate process {} with error code {}",
                pid, err_code
            );
        }

        CloseHandle(process_handle);
    }
}

/// Assert that certain process id no longer exists
pub fn process_exists(pid: u32) -> bool {
    unsafe {
        let handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

        let mut process_entry = PROCESSENTRY32::default();
        process_entry.dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;

        loop {
            if process_entry.th32ProcessID == pid {
                CloseHandle(handle);
                return true;
            }

            if Process32Next(handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(handle);
    }

    false
}

#[cfg(test)]
mod test {
    use std::thread::sleep;
    use std::time::Duration;

    use pretty_assertions::assert_eq;

    use super::*;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: u32) -> bool {
        !process_exists(pid)
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("sleep 0.1")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[test]
    /// Ensure a `powershell -c` command will be properly killed without detached processes.
    fn test_shell_command_is_killed() {
        let mut child = compile_shell_command("sleep 60; sleep 60; echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(1000));

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }

    #[test]
    /// Ensure that a `powershell -c` process with a child process that has children of it's own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() {
        let mut child = compile_shell_command("powershell -c 'sleep 60; sleep 60'; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(1000));

        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Assert that the direct child (powershell -c) has been killed.
        sleep(Duration::from_millis(500));
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }

    #[test]
    /// Ensure a normal command without `powershell -c` will be killed.
    fn test_normal_command_is_killed() {
        let mut child = Command::new("ping")
            .arg("localhost")
            .arg("-t")
            .spawn()
            .expect("Failed to spawn ping");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));
    }

    #[test]
    /// Ensure a normal command and all it's children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() {
        let mut child = Command::new("powershell")
            .arg("-c")
            .arg("sleep 60; sleep 60; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = get_cur_task_processes(pid);
        assert_eq!(process_ids.len(), 1);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }
    }
}
<commit_msg>fix: Incremental wait logic for windows process tests<commit_after>use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{error, info, warn};
use winapi::shared::minwindef::FALSE;
use winapi::shared::ntdef::NULL;
use winapi::um::errhandlingapi::GetLastError;
use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
use winapi::um::processthreadsapi::{
    OpenProcess, OpenThread, ResumeThread, SuspendThread, TerminateProcess,
};
use winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Process32First, Process32Next, Thread32First, Thread32Next,
    PROCESSENTRY32, TH32CS_SNAPPROCESS, TH32CS_SNAPTHREAD, THREADENTRY32,
};
use winapi::um::winnt::{PROCESS_TERMINATE, THREAD_SUSPEND_RESUME};

use crate::task_handler::ProcessAction;
use pueue_lib::network::message::Signal as InternalSignal;

pub fn compile_shell_command(command_string: &str) -> Command {
    // Chain two `powershell` commands, one that sets the output encoding to utf8 and then the user provided one.
    let mut command = Command::new("powershell");
    command.arg("-c").arg(format!(
        "[Console]::OutputEncoding = [Text.UTF8Encoding]::UTF8; {}",
        command_string
    ));

    command
}

pub fn send_internal_signal_to_child(
    child: &Child,
    signal: InternalSignal,
    send_to_children: bool,
) -> Result<bool> {
    bail!("Trying to send unix signal on a windows machine. This isn't supported.");
}

/// Send a signal to a windows process.
pub fn run_action_on_child(child: &Child, action: &ProcessAction, _children: bool) -> Result<bool> {
    let pids = get_cur_task_processes(child.id());
    if pids.is_empty() {
        bail!("Process has just gone away");
    }

    match action {
        ProcessAction::Pause => {
            for pid in pids {
                for thread in get_threads(pid) {
                    suspend_thread(thread);
                }
            }
        }
        ProcessAction::Resume => {
            for pid in pids {
                for thread in get_threads(pid) {
                    resume_thread(thread);
                }
            }
        }
    }

    Ok(true)
}

/// Kill a child process
pub fn kill_child(task_id: usize, child: &mut Child, _kill_children: bool) -> bool {
    match child.kill() {
        Err(_) => {
            info!("Task {} has already finished by itself", task_id);
            false
        }
        Ok(_) => {
            let pids = get_cur_task_processes(child.id());

            for pid in pids {
                terminate_process(pid);
            }
            true
        }
    }
}

/// Get current task pid, all child pid and all children's children
fn get_cur_task_processes(task_pid: u32) -> Vec<u32> {
    let mut all_pids = Vec::new();

    // Get all pids by BFS
    let mut parent_pids = vec![task_pid];
    while let Some(pid) = parent_pids.pop() {
        all_pids.push(pid);

        get_child_pids(pid, &mut parent_pids);
    }

    // Keep parent pid ahead of child. We need execute action for parent process first.
    all_pids.reverse();
    all_pids
}

/// Get child pids of a specific process.
fn get_child_pids(target_pid: u32, pid_list: &mut Vec<u32>) {
    unsafe {
        // Take a snapshot of all processes in the system.
        // While enumerating the set of processes, new processes can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, target_pid);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return;
        }

        // Walk the list of processes.
        let mut process_entry = PROCESSENTRY32 {
            dwSize: std::mem::size_of::<PROCESSENTRY32>() as u32,
            ..Default::default()
        };
        if Process32First(snapshot_handle, &mut process_entry) == FALSE {
            error!("Couldn't get first process.");
            CloseHandle(snapshot_handle);
            return;
        }

        loop {
            if process_entry.th32ParentProcessID == target_pid {
                pid_list.push(process_entry.th32ProcessID);
            }

            if Process32Next(snapshot_handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }
}

/// Get all thread id of a specific process
fn get_threads(target_pid: u32) -> Vec<u32> {
    let mut threads = Vec::new();

    unsafe {
        // Take a snapshot of all threads in the system.
        // While enumerating the set of threads, new threads can be created and destroyed.
        let snapshot_handle = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
        if snapshot_handle == INVALID_HANDLE_VALUE {
            error!("Failed to get process {} snapShot", target_pid);
            return threads;
        }

        // Walk the list of threads.
        let mut thread_entry = THREADENTRY32 {
            dwSize: std::mem::size_of::<THREADENTRY32>() as u32,
            ..Default::default()
        };
        if Thread32First(snapshot_handle, &mut thread_entry) == FALSE {
            error!("Couldn't get first thread.");
            CloseHandle(snapshot_handle);
            return threads;
        }

        loop {
            if thread_entry.th32OwnerProcessID == target_pid {
                threads.push(thread_entry.th32ThreadID);
            }

            if Thread32Next(snapshot_handle, &mut thread_entry) == FALSE {
                break;
            }
        }

        CloseHandle(snapshot_handle);
    }

    threads
}

/// Suspend a thread
/// Each thread has a suspend count (with a maximum value of `MAXIMUM_SUSPEND_COUNT`).
/// If the suspend count is greater than zero, the thread is suspended; otherwise, the thread is not suspended and is eligible for execution.
/// Calling `SuspendThread` causes the target thread's suspend count to be incremented.
/// Attempting to increment past the maximum suspend count causes an error without incrementing the count.
/// [SuspendThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-suspendthread)
fn suspend_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If SuspendThread fails, the return value is (DWORD) -1
            if u32::max_value() == SuspendThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to suspend thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Resume a thread
/// ResumeThread checks the suspend count of the subject thread.
/// If the suspend count is zero, the thread is not currently suspended. Otherwise, the subject thread's suspend count is decremented.
/// If the resulting value is zero, then the execution of the subject thread is resumed.
/// [ResumeThread](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-resumethread)
fn resume_thread(tid: u32) {
    unsafe {
        // Attempt to convert the thread ID into a handle
        let thread_handle = OpenThread(THREAD_SUSPEND_RESUME, FALSE, tid);
        if thread_handle != NULL {
            // If ResumeThread fails, the return value is (DWORD) -1
            if u32::max_value() == ResumeThread(thread_handle) {
                let err_code = GetLastError();
                warn!(
                    "Failed to resume thread {} with error code {}",
                    tid, err_code
                );
            }
        }

        CloseHandle(thread_handle);
    }
}

/// Terminate a process
/// [TerminateProcess](https://docs.microsoft.com/en-us/windows/win32/api/processthreadsapi/nf-processthreadsapi-terminateprocess)
fn terminate_process(pid: u32) {
    unsafe {
        // Get a handle for the target process
        let process_handle = OpenProcess(PROCESS_TERMINATE, FALSE, pid);
        // If TerminateProcess fails, the return value is zero.
        if 0 == TerminateProcess(process_handle, 1) {
            let err_code = GetLastError();
            warn!(
                "Failed to terminate process {} with error code {}",
                pid, err_code
            );
        }

        CloseHandle(process_handle);
    }
}

/// Assert that certain process id no longer exists
pub fn process_exists(pid: u32) -> bool {
    unsafe {
        let handle = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);

        let mut process_entry = PROCESSENTRY32::default();
        process_entry.dwSize = std::mem::size_of::<PROCESSENTRY32>() as u32;

        loop {
            if process_entry.th32ProcessID == pid {
                CloseHandle(handle);
                return true;
            }

            if Process32Next(handle, &mut process_entry) == FALSE {
                break;
            }
        }

        CloseHandle(handle);
    }

    false
}

#[cfg(test)]
mod test {
    use std::thread::sleep;
    use std::time::Duration;

    use pretty_assertions::assert_eq;

    use super::*;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: u32) -> bool {
        !process_exists(pid)
    }

    /// A test helper function, which ensures that a specific amount of subprocesses can be
    /// observed for a given PID in a given time window.
    /// If the correct amount can be observed, the process ids are then returned.
    ///
    /// The process count is checked every few milliseconds for the given duration.
    fn assert_process_ids(pid: u32, expected_processes: usize, millis: usize) -> Result<Vec<u32>> {
        // Check every 50 milliseconds.
        let interval = 50;
        let tries = millis / interval;
        let mut current_try = 0;

        while current_try <= tries {
            // Continue waiting if the count doesn't match.
            let process_ids = get_cur_task_processes(pid);
            if process_ids.len() != expected_processes {
                current_try += 1;
                sleep(Duration::from_millis(interval as u64));
                continue;
            }

            return Ok(process_ids);
        }

        bail!(
            "{} processes were expected. Last process count was {}",
            expected_processes,
            get_cur_task_processes(pid).len()
        )
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("sleep 0.1")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[test]
    /// Ensure a `powershell -c` command will be properly killed without detached processes.
    fn test_shell_command_is_killed() -> Result<()> {
        let mut child = compile_shell_command("sleep 60; sleep 60; echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = assert_process_ids(pid, 1, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }

        Ok(())
    }

    #[test]
    /// Ensure that a `powershell -c` process with a child process that has children of it's own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() -> Result<()> {
        let mut child = compile_shell_command("powershell -c 'sleep 60; sleep 60'; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();
        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = assert_process_ids(pid, 2, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Assert that the direct child (powershell -c) has been killed.
        sleep(Duration::from_millis(500));
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }

        Ok(())
    }

    #[test]
    /// Ensure a normal command without `powershell -c` will be killed.
    fn test_normal_command_is_killed() -> Result<()> {
        let mut child = Command::new("ping")
            .arg("localhost")
            .arg("-t")
            .spawn()
            .expect("Failed to spawn ping");
        let pid = child.id();

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = assert_process_ids(pid, 1, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));

        Ok(())
    }

    #[test]
    /// Ensure a normal command and all it's children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() -> Result<()> {
        let mut child = Command::new("powershell")
            .arg("-c")
            .arg("sleep 60; sleep 60; sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid = child.id();

        // Get all processes, so we can make sure they no longer exist afterwards.
        let process_ids = assert_process_ids(pid, 1, 5000)?;

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for pid in process_ids {
            assert!(process_is_gone(pid));
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>use glut;
use azure;
use geom;

use geom::point::Point2D;
use geom::rect::Rect;
use geom::size::Size2D;
use geom::matrix::Matrix4;
use layers::*;
use scene::*;
use rendergl::*;
use util::convert_rgb32_to_rgb24;

use glut::{DOUBLE, check_loop, create_window, destroy_window, display_func, init};
use glut::{init_display_mode, post_redisplay, swap_buffers, timer_func};

use azure::cairo::CAIRO_FORMAT_RGB24;
use CairoContext = azure::cairo_hl::Context;
use azure::azure_hl::{Color, ColorPattern, DrawTarget};
use azure::cairo_hl::ImageSurface;

use comm::{peek, recv, send, Port, Chan};
use libc::c_uint;
use os::{getenv, setenv};
use task::TaskBuilder;

struct Renderer {
    mut layer: @TiledImageLayer,
    //mut layer: @ImageLayer,
    mut t: f32,
    mut delta: f32,
    mut render_context: Option<RenderContext>,
    mut image: Option<@Image>,
}

impl Renderer {
    fn get_display_callback(this: @Renderer) -> fn@() {
        fn@() {
            (*this).display_callback();
        }
    }

    fn display_callback() {
        match self.render_context {
            None => {
                self.render_context = Some(init_render_context());
            }
            Some(_) => {
                // Nothing to do.
            }
        }
        let context = match self.render_context {
            None => fail,
            Some(ctx) => ctx
        };

        let t = self.t;
        self.layer.common.transform = Matrix4(400.0f32 * t, 0.0f32,       0.0f32, 0.0f32,
                                              0.0f32,       300.0f32 * t, 0.0f32, 0.0f32,
                                              0.0f32,       0.0f32,       1.0f32, 0.0f32,
                                              0.0f32,       0.0f32,       0.0f32, 1.0f32);

        let mut scene = Scene(TiledImageLayerKind(self.layer), Size2D(400.0f32, 300.0f32));
        //let mut scene = Scene(ImageLayerKind(self.layer), Size2D(400.0f32, 300.0f32));
        render_scene(context, scene);

        //self.t += self.delta;
        if self.t < 0.0f32 || self.t > 1.0f32 {
            self.delta = -self.delta;
        }

        swap_buffers();

        post_redisplay();
    }
}

fn Renderer() -> Renderer {
        let cairo_image = ImageSurface(CAIRO_FORMAT_RGB24, 500, 704);

        let image : Image;

        let draw_target = DrawTarget(cairo_image);
        draw_target.fill_rect(Rect(Point2D(50.0f32, 50.0f32), Size2D(300.0f32, 284.0f32)),
                              ColorPattern(Color(1.0f32, 1.0f32, 0.0f32, 1.0f32)));
        draw_target.flush();

        let (width, height) = (cairo_image.width() as uint, cairo_image.height() as uint);
        let (tile_width, tile_height) = (width / 4, height / 4);
        let cairo_data = cairo_image.data();

        let tiles = DVec();
        for 4.timesi |y| {
            for 4.timesi |x| {
                // Extract the relevant part of the image.
                let data = DVec();

                let mut scanline_start = (y * tile_height * width + x * tile_width) * 4;
                for tile_height.times {
                    for (tile_width * 4).timesi |offset| {
                        data.push(cairo_data[scanline_start + offset]);
                    }

                    scanline_start += width * 4;
                }

                let data = convert_rgb32_to_rgb24(vec::from_mut(dvec::unwrap(data)));
                let image = @Image(tile_width, tile_height, RGB24Format, data);
                tiles.push(image);
            }
        }

        let tiles = dvec::unwrap(tiles);

    Renderer {
        image : Some(image),
        layer : @TiledImageLayer(tiles, 4),
        t : 1.0f32,
        delta : -0.001f32,
        render_context : None
    }
}

#[test]
fn test_triangle_and_square() unsafe {
    let builder = task::task().sched_mode(task::PlatformThread);

    let po: Port<()> = Port();
    let ch = Chan(po);
    let _result_ch: Chan<()> = do builder.spawn_listener |_po| {
        let renderer = @Renderer();

        init();
        init_display_mode(DOUBLE as c_uint);
        let window = create_window(~"Rust Layers");
        display_func(renderer.get_display_callback(renderer));

        let wakeup = Port();
        let wakeup_chan = Chan(wakeup);
        do timer_func(30000) {
            send(wakeup_chan, ());
        }

        loop {
            check_loop();

            if peek(wakeup) {
                recv(wakeup);
                send(ch, ());
                destroy_window(window);
                break;
            }
        }

        send(ch, ());
    };

    recv(po);
}

<commit_msg>Fix test breakage<commit_after>use glut;
use azure;
use geom;

use geom::point::Point2D;
use geom::rect::Rect;
use geom::size::Size2D;
use geom::matrix::Matrix4;
use layers::*;
use scene::*;
use rendergl::*;
use util::convert_rgb32_to_rgb24;

use glut::{DOUBLE, check_loop, create_window, destroy_window, display_func, init};
use glut::{init_display_mode, post_redisplay, swap_buffers, timer_func};

use azure::cairo::CAIRO_FORMAT_RGB24;
use CairoContext = azure::cairo_hl::Context;
use azure::azure_hl::{Color, ColorPattern, DrawTarget};
use azure::cairo_hl::ImageSurface;

use comm::{peek, recv, send, Port, Chan};
use libc::c_uint;
use os::{getenv, setenv};
use task::TaskBuilder;

struct Renderer {
    mut layer: @TiledImageLayer,
    //mut layer: @ImageLayer,
    mut t: f32,
    mut delta: f32,
    mut render_context: Option<RenderContext>,
    mut image: Option<@Image>,
}

impl Renderer {
    fn get_display_callback(this: @Renderer) -> fn@() {
        fn@() {
            (*this).display_callback();
        }
    }

    fn display_callback() {
        match self.render_context {
            None => {
                self.render_context = Some(init_render_context());
            }
            Some(_) => {
                // Nothing to do.
            }
        }
        let context = match self.render_context {
            None => fail,
            Some(ctx) => ctx
        };

        let t = self.t;
        self.layer.common.transform = Matrix4(400.0f32 * t, 0.0f32,       0.0f32, 0.0f32,
                                              0.0f32,       300.0f32 * t, 0.0f32, 0.0f32,
                                              0.0f32,       0.0f32,       1.0f32, 0.0f32,
                                              0.0f32,       0.0f32,       0.0f32, 1.0f32);

        let mut scene = Scene(TiledImageLayerKind(self.layer), Size2D(400.0f32, 300.0f32));
        //let mut scene = Scene(ImageLayerKind(self.layer), Size2D(400.0f32, 300.0f32));
        render_scene(context, scene);

        //self.t += self.delta;
        if self.t < 0.0f32 || self.t > 1.0f32 {
            self.delta = -self.delta;
        }

        swap_buffers();

        post_redisplay();
    }
}

fn Renderer() -> Renderer {
        let cairo_image = ImageSurface(CAIRO_FORMAT_RGB24, 500, 704);

        let image : Image;

        let draw_target = DrawTarget(cairo_image);
        draw_target.fill_rect(Rect(Point2D(50.0f32, 50.0f32), Size2D(300.0f32, 284.0f32)),
                              ColorPattern(Color(1.0f32, 1.0f32, 0.0f32, 1.0f32)));
        draw_target.flush();

        let (width, height) = (cairo_image.width() as uint, cairo_image.height() as uint);
        let (tile_width, tile_height) = (width / 4, height / 4);
        let cairo_data = cairo_image.data();

        let tiles = DVec();
        for 4.timesi |y| {
            for 4.timesi |x| {
                // Extract the relevant part of the image.
                let data = DVec();

                let mut scanline_start = (y * tile_height * width + x * tile_width) * 4;
                for tile_height.times {
                    for (tile_width * 4).timesi |offset| {
                        data.push(cairo_data[scanline_start + offset]);
                    }

                    scanline_start += width * 4;
                }

                let data = convert_rgb32_to_rgb24(vec::from_mut(dvec::unwrap(data)));
                let image = @Image(tile_width, tile_height, RGB24Format, data);
                tiles.push(image);
            }
        }

        let tiles = dvec::unwrap(tiles);

    Renderer {
        image : Some(tiles[0]),
        layer : @TiledImageLayer(tiles, 4),
        t : 1.0f32,
        delta : -0.001f32,
        render_context : None
    }
}

#[test]
fn test_triangle_and_square() unsafe {
    let builder = task::task().sched_mode(task::PlatformThread);

    let po: Port<()> = Port();
    let ch = Chan(po);
    let _result_ch: Chan<()> = do builder.spawn_listener |_po| {
        let renderer = @Renderer();

        init();
        init_display_mode(DOUBLE as c_uint);
        let window = create_window(~"Rust Layers");
        display_func(renderer.get_display_callback(renderer));

        let wakeup = Port();
        let wakeup_chan = Chan(wakeup);
        do timer_func(30000) {
            send(wakeup_chan, ());
        }

        loop {
            check_loop();

            if peek(wakeup) {
                recv(wakeup);
                send(ch, ());
                destroy_window(window);
                break;
            }
        }

        send(ch, ());
    };

    recv(po);
}

<|endoftext|>


<commit_before>//! Driver for proc macro server

use crate::{expand_task, list_macros};
use ra_proc_macro::msg::{self, Message};
use std::io;

pub fn run() -> io::Result<()> {
    loop {
        // bubble up the error for read request,
        // as the stdin pipe may be closed.
        let req = match read_request()? {
            None => continue,
            Some(req) => req,
        };

        let res = match req {
            msg::Request::ListMacro(task) => Ok(msg::Response::ListMacro(list_macros(&task))),
            msg::Request::ExpansionMacro(task) => {
                expand_task(&task).map(msg::Response::ExpansionMacro)
            }
        };

        let msg = res.unwrap_or_else(|err| {
            msg::Response::Error(msg::ResponseError {
                code: msg::ErrorCode::ExpansionError,
                message: err,
            })
        });

        if let Err(err) = write_response(msg) {
            eprintln!("Write message error: {}", err);
        }
    }
}

fn read_request() -> io::Result<Option<msg::Request>> {
    msg::Request::read(&mut io::stdin().lock())
}

fn write_response(msg: msg::Response) -> io::Result<()> {
    msg.write(&mut io::stdout().lock())
}
<commit_msg>Refactor a bit<commit_after>//! Driver for proc macro server

use crate::{expand_task, list_macros};
use ra_proc_macro::msg::{self, Message};
use std::io;

pub fn run() -> io::Result<()> {
    while let Some(req) = read_request()? {
        let res = match req {
            msg::Request::ListMacro(task) => Ok(msg::Response::ListMacro(list_macros(&task))),
            msg::Request::ExpansionMacro(task) => {
                expand_task(&task).map(msg::Response::ExpansionMacro)
            }
        };

        let msg = res.unwrap_or_else(|err| {
            msg::Response::Error(msg::ResponseError {
                code: msg::ErrorCode::ExpansionError,
                message: err,
            })
        });

        if let Err(err) = write_response(msg) {
            eprintln!("Write message error: {}", err);
        }
    }

    Ok(())
}

fn read_request() -> io::Result<Option<msg::Request>> {
    msg::Request::read(&mut io::stdin().lock())
}

fn write_response(msg: msg::Response) -> io::Result<()> {
    msg.write(&mut io::stdout().lock())
}
<|endoftext|>


<commit_before>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

extern crate sys_util;

use std::mem;
use std::ffi::CStr;
use std::io::{Read, Seek, SeekFrom};

use sys_util::{GuestAddress, GuestMemory};

#[allow(dead_code)]
#[allow(non_camel_case_types)]
#[allow(non_snake_case)]
#[allow(non_upper_case_globals)]
mod elf;

#[derive(Debug, PartialEq)]
pub enum Error {
    BigEndianElfOnLittle,
    CommandLineCopy,
    CommandLineOverflow,
    InvalidElfMagicNumber,
    InvalidProgramHeaderSize,
    InvalidProgramHeaderOffset,
    InvalidProgramHeaderAddress,
    ReadElfHeader,
    ReadKernelImage,
    ReadProgramHeader,
    SeekKernelStart,
    SeekElfStart,
    SeekProgramHeader,
}
pub type Result<T> = std::result::Result<T, Error>;

/// Loads a kernel from a vmlinux elf image to a slice
///
/// # Arguments
///
/// * `guest_mem` - The guest memory region the kernel is written to.
/// * `kernel_start` - The offset into `guest_mem` at which to load the kernel.
/// * `kernel_image` - Input vmlinux image.
pub fn load_kernel<F>(guest_mem: &GuestMemory, kernel_start: GuestAddress, kernel_image: &mut F) -> Result<()>
    where F: Read + Seek
{
    let mut ehdr: elf::Elf64_Ehdr = Default::default();
    kernel_image.seek(SeekFrom::Start(0))
        .map_err(|_| Error::SeekElfStart)?;
    unsafe {
        // read_struct is safe when reading a POD struct.  It can be used and dropped without issue.
        sys_util::read_struct(kernel_image, &mut ehdr).map_err(|_| Error::ReadElfHeader)?;
    }

    // Sanity checks
    if ehdr.e_ident[elf::EI_MAG0 as usize] != elf::ELFMAG0 as u8 ||
       ehdr.e_ident[elf::EI_MAG1 as usize] != elf::ELFMAG1 ||
       ehdr.e_ident[elf::EI_MAG2 as usize] != elf::ELFMAG2 ||
       ehdr.e_ident[elf::EI_MAG3 as usize] != elf::ELFMAG3 {
        return Err(Error::InvalidElfMagicNumber);
    }
    if ehdr.e_ident[elf::EI_DATA as usize] != elf::ELFDATA2LSB as u8 {
        return Err(Error::BigEndianElfOnLittle);
    }
    if ehdr.e_phentsize as usize != mem::size_of::<elf::Elf64_Phdr>() {
        return Err(Error::InvalidProgramHeaderSize);
    }
    if (ehdr.e_phoff as usize) < mem::size_of::<elf::Elf64_Ehdr>() {
        // If the program header is backwards, bail.
        return Err(Error::InvalidProgramHeaderOffset);
    }

    kernel_image.seek(SeekFrom::Start(ehdr.e_phoff))
        .map_err(|_| Error::SeekProgramHeader)?;
    let phdrs: Vec<elf::Elf64_Phdr> = unsafe {
        // Reading the structs is safe for a slice of POD structs.
        sys_util::read_struct_slice(kernel_image, ehdr.e_phnum as usize)
            .map_err(|_| Error::ReadProgramHeader)?
    };

    // Read in each section pointed to by the program headers.
    for phdr in phdrs.iter() {
        if (phdr.p_type & elf::PT_LOAD) == 0 || phdr.p_filesz == 0 {
            continue;
        }

        kernel_image.seek(SeekFrom::Start(phdr.p_offset))
            .map_err(|_| Error::SeekKernelStart)?;

        let mem_offset = kernel_start.checked_add(phdr.p_paddr as usize)
            .ok_or(Error::InvalidProgramHeaderAddress)?;
        guest_mem.read_to_memory(mem_offset, kernel_image, phdr.p_filesz as usize)
            .map_err(|_| Error::ReadKernelImage)?;
    }

    Ok(())
}

/// Writes the command line string to the given memory slice.
///
/// # Arguments
///
/// * `guest_mem` - A u8 slice that will be partially overwritten by the command line.
/// * `guest_addr` - The address in `guest_mem` at which to load the command line.
/// * `cmdline` - The kernel command line.
pub fn load_cmdline(guest_mem: &GuestMemory, guest_addr: GuestAddress, cmdline: &CStr) -> Result<()> {
    let len = cmdline.to_bytes().len();
    if len <= 0 {
        return Ok(());
    }

    let end = guest_addr.checked_add(len + 1)
        .ok_or(Error::CommandLineOverflow)?; // Extra for null termination.
    if end > guest_mem.end_addr() {
        return Err(Error::CommandLineOverflow)?;
    }

    guest_mem.write_slice_at_addr(cmdline.to_bytes_with_nul(), guest_addr)
        .map_err(|_| Error::CommandLineCopy)?;

    Ok(())
}

#[cfg(test)]
mod test {
    use std::io::Cursor;
    use super::*;
    use sys_util::{GuestAddress, GuestMemory};

    const MEM_SIZE: usize = 0x8000;

    fn create_guest_mem() -> GuestMemory {
        GuestMemory::new(&vec![(GuestAddress(0x0), MEM_SIZE)]).unwrap()
    }

    #[test]
    fn cmdline_overflow() {
        let gm = create_guest_mem();
        let cmdline_address = GuestAddress(MEM_SIZE - 5);
        assert_eq!(Err(Error::CommandLineOverflow),
                   load_cmdline(&gm,
                                cmdline_address,
                                CStr::from_bytes_with_nul(b"12345\0").unwrap()));
    }

    #[test]
    fn cmdline_write_end() {
        let gm = create_guest_mem();
        let mut cmdline_address = GuestAddress(45);
        assert_eq!(Ok(()),
                   load_cmdline(&gm,
                                cmdline_address,
                                CStr::from_bytes_with_nul(b"1234\0").unwrap()));
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '1' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '2' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '3' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '4' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '\0' as u8);
    }

    // Elf64 image that prints hello world on x86_64.
    fn make_elf_bin() -> Vec<u8> {
        let mut v = Vec::new();
        v.extend_from_slice(include_bytes!("test_elf.bin"));
        v
    }

    #[test]
    fn load_elf() {
        let gm = create_guest_mem();
        let kernel_addr = GuestAddress(0x0);
        let image = make_elf_bin();
        assert_eq!(Ok(()),
                   load_kernel(&gm, kernel_addr, &mut Cursor::new(&image)));
    }

    #[test]
    fn bad_magic() {
        let gm = create_guest_mem();
        let kernel_addr = GuestAddress(0x0);
        let mut bad_image = make_elf_bin();
        bad_image[0x1] = 0x33;
        assert_eq!(Err(Error::InvalidElfMagicNumber),
                   load_kernel(&gm, kernel_addr, &mut Cursor::new(&bad_image)));
    }

    #[test]
    fn bad_endian() {
        // Only little endian is supported
        let gm = create_guest_mem();
        let kernel_addr = GuestAddress(0x0);
        let mut bad_image = make_elf_bin();
        bad_image[0x5] = 2;
        assert_eq!(Err(Error::BigEndianElfOnLittle),
                   load_kernel(&gm, kernel_addr, &mut Cursor::new(&bad_image)));
    }

    #[test]
    fn bad_phoff() {
        // program header has to be past the end of the elf header
        let gm = create_guest_mem();
        let kernel_addr = GuestAddress(0x0);
        let mut bad_image = make_elf_bin();
        bad_image[0x20] = 0x10;
        assert_eq!(Err(Error::InvalidProgramHeaderOffset),
                   load_kernel(&gm, kernel_addr, &mut Cursor::new(&bad_image)));
    }
}
<commit_msg>kernel_loader: Fix clippy warnings.<commit_after>// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

extern crate sys_util;

use std::mem;
use std::ffi::CStr;
use std::io::{Read, Seek, SeekFrom};

use sys_util::{GuestAddress, GuestMemory};

#[allow(dead_code)]
#[allow(non_camel_case_types)]
#[allow(non_snake_case)]
#[allow(non_upper_case_globals)]
mod elf;

#[derive(Debug, PartialEq)]
pub enum Error {
    BigEndianElfOnLittle,
    CommandLineCopy,
    CommandLineOverflow,
    InvalidElfMagicNumber,
    InvalidProgramHeaderSize,
    InvalidProgramHeaderOffset,
    InvalidProgramHeaderAddress,
    ReadElfHeader,
    ReadKernelImage,
    ReadProgramHeader,
    SeekKernelStart,
    SeekElfStart,
    SeekProgramHeader,
}
pub type Result<T> = std::result::Result<T, Error>;

/// Loads a kernel from a vmlinux elf image to a slice
///
/// # Arguments
///
/// * `guest_mem` - The guest memory region the kernel is written to.
/// * `kernel_start` - The offset into `guest_mem` at which to load the kernel.
/// * `kernel_image` - Input vmlinux image.
pub fn load_kernel<F>(guest_mem: &GuestMemory, kernel_start: GuestAddress, kernel_image: &mut F) -> Result<()>
    where F: Read + Seek
{
    let mut ehdr: elf::Elf64_Ehdr = Default::default();
    kernel_image.seek(SeekFrom::Start(0))
        .map_err(|_| Error::SeekElfStart)?;
    unsafe {
        // read_struct is safe when reading a POD struct.  It can be used and dropped without issue.
        sys_util::read_struct(kernel_image, &mut ehdr).map_err(|_| Error::ReadElfHeader)?;
    }

    // Sanity checks
    if ehdr.e_ident[elf::EI_MAG0 as usize] != elf::ELFMAG0 as u8 ||
       ehdr.e_ident[elf::EI_MAG1 as usize] != elf::ELFMAG1 ||
       ehdr.e_ident[elf::EI_MAG2 as usize] != elf::ELFMAG2 ||
       ehdr.e_ident[elf::EI_MAG3 as usize] != elf::ELFMAG3 {
        return Err(Error::InvalidElfMagicNumber);
    }
    if ehdr.e_ident[elf::EI_DATA as usize] != elf::ELFDATA2LSB as u8 {
        return Err(Error::BigEndianElfOnLittle);
    }
    if ehdr.e_phentsize as usize != mem::size_of::<elf::Elf64_Phdr>() {
        return Err(Error::InvalidProgramHeaderSize);
    }
    if (ehdr.e_phoff as usize) < mem::size_of::<elf::Elf64_Ehdr>() {
        // If the program header is backwards, bail.
        return Err(Error::InvalidProgramHeaderOffset);
    }

    kernel_image.seek(SeekFrom::Start(ehdr.e_phoff))
        .map_err(|_| Error::SeekProgramHeader)?;
    let phdrs: Vec<elf::Elf64_Phdr> = unsafe {
        // Reading the structs is safe for a slice of POD structs.
        sys_util::read_struct_slice(kernel_image, ehdr.e_phnum as usize)
            .map_err(|_| Error::ReadProgramHeader)?
    };

    // Read in each section pointed to by the program headers.
    for phdr in &phdrs {
        if (phdr.p_type & elf::PT_LOAD) == 0 || phdr.p_filesz == 0 {
            continue;
        }

        kernel_image.seek(SeekFrom::Start(phdr.p_offset))
            .map_err(|_| Error::SeekKernelStart)?;

        let mem_offset = kernel_start.checked_add(phdr.p_paddr as usize)
            .ok_or(Error::InvalidProgramHeaderAddress)?;
        guest_mem.read_to_memory(mem_offset, kernel_image, phdr.p_filesz as usize)
            .map_err(|_| Error::ReadKernelImage)?;
    }

    Ok(())
}

/// Writes the command line string to the given memory slice.
///
/// # Arguments
///
/// * `guest_mem` - A u8 slice that will be partially overwritten by the command line.
/// * `guest_addr` - The address in `guest_mem` at which to load the command line.
/// * `cmdline` - The kernel command line.
pub fn load_cmdline(guest_mem: &GuestMemory, guest_addr: GuestAddress, cmdline: &CStr) -> Result<()> {
    let len = cmdline.to_bytes().len();
    if len == 0 {
        return Ok(());
    }

    let end = guest_addr.checked_add(len + 1)
        .ok_or(Error::CommandLineOverflow)?; // Extra for null termination.
    if end > guest_mem.end_addr() {
        return Err(Error::CommandLineOverflow)?;
    }

    guest_mem.write_slice_at_addr(cmdline.to_bytes_with_nul(), guest_addr)
        .map_err(|_| Error::CommandLineCopy)?;

    Ok(())
}

#[cfg(test)]
mod test {
    use std::io::Cursor;
    use super::*;
    use sys_util::{GuestAddress, GuestMemory};

    const MEM_SIZE: usize = 0x8000;

    fn create_guest_mem() -> GuestMemory {
        GuestMemory::new(&vec![(GuestAddress(0x0), MEM_SIZE)]).unwrap()
    }

    #[test]
    fn cmdline_overflow() {
        let gm = create_guest_mem();
        let cmdline_address = GuestAddress(MEM_SIZE - 5);
        assert_eq!(Err(Error::CommandLineOverflow),
                   load_cmdline(&gm,
                                cmdline_address,
                                CStr::from_bytes_with_nul(b"12345\0").unwrap()));
    }

    #[test]
    fn cmdline_write_end() {
        let gm = create_guest_mem();
        let mut cmdline_address = GuestAddress(45);
        assert_eq!(Ok(()),
                   load_cmdline(&gm,
                                cmdline_address,
                                CStr::from_bytes_with_nul(b"1234\0").unwrap()));
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '1' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '2' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '3' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '4' as u8);
        cmdline_address = cmdline_address.unchecked_add(1);
        let val: u8 = gm.read_obj_from_addr(cmdline_address).unwrap();
        assert_eq!(val, '\0' as u8);
    }

    // Elf64 image that prints hello world on x86_64.
    fn make_elf_bin() -> Vec<u8> {
        let mut v = Vec::new();
        v.extend_from_slice(include_bytes!("test_elf.bin"));
        v
    }

    #[test]
    fn load_elf() {
        let gm = create_guest_mem();
        let kernel_addr = GuestAddress(0x0);
        let image = make_elf_bin();
        assert_eq!(Ok(()),
                   load_kernel(&gm, kernel_addr, &mut Cursor::new(&image)));
    }

    #[test]
    fn bad_magic() {
        let gm = create_guest_mem();
        let kernel_addr = GuestAddress(0x0);
        let mut bad_image = make_elf_bin();
        bad_image[0x1] = 0x33;
        assert_eq!(Err(Error::InvalidElfMagicNumber),
                   load_kernel(&gm, kernel_addr, &mut Cursor::new(&bad_image)));
    }

    #[test]
    fn bad_endian() {
        // Only little endian is supported
        let gm = create_guest_mem();
        let kernel_addr = GuestAddress(0x0);
        let mut bad_image = make_elf_bin();
        bad_image[0x5] = 2;
        assert_eq!(Err(Error::BigEndianElfOnLittle),
                   load_kernel(&gm, kernel_addr, &mut Cursor::new(&bad_image)));
    }

    #[test]
    fn bad_phoff() {
        // program header has to be past the end of the elf header
        let gm = create_guest_mem();
        let kernel_addr = GuestAddress(0x0);
        let mut bad_image = make_elf_bin();
        bad_image[0x20] = 0x10;
        assert_eq!(Err(Error::InvalidProgramHeaderOffset),
                   load_kernel(&gm, kernel_addr, &mut Cursor::new(&bad_image)));
    }
}
<|endoftext|>


<commit_before>/// A TOML [1] configuration file parser
///
/// [1]: https://github.com/mojombo/toml

use std::io::mem::MemReader;
use std::io::File;
use std::hashmap::HashMap;

fn read_char(rd: &mut MemReader) -> char {
  match rd.read_byte() {
    Some(b) => b as char,
    None => fail!()
  }
}

#[deriving(ToStr)]
enum Value {
    True,
    False,
    Unsigned(u64),
    Integer(i64),
    Float(f64),
    String(~str),
    Array(~[Value]),
    Datetime, // XXX
    Map(HashMap<~str, Value>) // XXX: This is no value
}

// parse values recursivly
fn parse_value(rd: &mut MemReader) -> Value {
    enum State {
        st_value,
        st_number,
        st_got_number
    }

    let mut state: State = st_value;
    let mut val: ~str = ~"";

    loop {
        match state {
            st_value => {
                if rd.eof() { fail!() }
                let ch = read_char(rd);
                match ch {
                    '\r' | '\n' | ' ' | '\t' => { }
                    '0' .. '9' => {
                        val.push_char(ch);
                        state = st_number;
                    }
                    't' => {
                        assert!(read_char(rd) == 'r');
                        assert!(read_char(rd) == 'u');
                        assert!(read_char(rd) == 'e');
                        return True
                    }
                    'f' => {
                        assert!(read_char(rd) == 'a');
                        assert!(read_char(rd) == 'l');
                        assert!(read_char(rd) == 's');
                        assert!(read_char(rd) == 'e');
                        return False
                    }
                    '"' => {
                        loop {
                            let ch = read_char(rd);
                            if ch == '"' {
                                break;
                            }
                            val.push_char(ch);
                        }
                        return String(val)
                    }
                    _ => { fail!() }
                }
            }
            st_got_number => {
                return Unsigned(from_str(val).unwrap())
            }
            st_number => {
                if rd.eof() {
                    state = st_got_number;
                } else {
                    let ch = read_char(rd);
                    match ch {
                        '\r' | '\n' | ' ' | '\t' => { state = st_got_number }
                        '#' => { skip_comment(rd); state = st_got_number }
                        '0' .. '9' => {
                            val.push_char(ch);
                        }
                        _ => { fail!() }
                    }
                }
            }
        }
    }
}

// We must be already within the '#"
fn skip_comment(rd: &mut MemReader) {
    loop {
        if rd.eof() { return }
        match read_char(rd) {
            '\n' => { return }
            _ => { }
        }
    }
}

trait Visitor {
    fn section(&mut self, name: &str, is_array: bool);
    fn pair(&mut self, key: &str, val: Value);
}

struct MyVisitor {
    root: HashMap<~str, Value>,
    current_path: ~str
}

impl MyVisitor {
    fn new() -> MyVisitor {
        MyVisitor { root: HashMap::new(), current_path: ~"" }
    }
    fn get_root<'a>(&'a self) -> &'a HashMap<~str, Value> {
        return &self.root;
    }
}

impl Visitor for MyVisitor {
    fn section(&mut self, name: &str, is_array: bool) {
        assert!(is_array == false);
        debug!("Section: {}", name);
        self.current_path = name.to_owned(); //clone();
    }
    fn pair(&mut self, key: &str, val: Value) {
        debug!("Pair: {} {:s}", key, val.to_str());
        let mut m = self.root.find_or_insert(self.current_path.clone(), Map(HashMap::new())); // XXX: remove clone
        match *m {
            Map(ref mut map) => {
                let fresh = map.insert(key.to_owned(), val);
                assert!(fresh == true);
            }
            _ => { fail!("Invalid TOML") }
        }
    }
}

fn parse<V: Visitor>(rd: &mut MemReader, visitor: &mut V) {
    enum State {
        st_toplevel,
        st_comment,
        st_section,
        st_ident,
        st_assign_wanted,
        st_value_wanted
    }

    let mut state = st_toplevel;
    let mut buf: ~str = ~"";

    loop {
        match state {
            st_toplevel => {
                if rd.eof() { break }
                let ch = read_char(rd);
                match ch {
                    // ignore whitespace
                    '\r' | '\n' | ' ' | '\t' => { }

                    // comment
                    '#' => { state = st_comment }

                    // section
                    '[' => { state = st_section }

                    // identifier
                    'a' .. 'z' | 'A' .. 'Z' | '_' => {
                        state = st_ident;
                        buf.push_char(ch)
                    }

                    _ => { fail!() }
                }
            }
            st_comment => {
                if rd.eof() { break }
                match read_char(rd) {
                    '\n' => { state = st_toplevel }
                    _ => { }
                }
            }
            st_section => {
                if rd.eof() { fail!() }
                let ch = read_char(rd);
                match ch {
                    'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' | '.' | '_'=> {
                        buf.push_char(ch);
                    }
                    ']' => {
                        visitor.section(buf, false);
                        buf.truncate(0);
                        state = st_toplevel;
                    }
                    _ => { fail!() }
                }
            }
            st_ident => {
                if rd.eof() { fail!() }
                let ch = read_char(rd);
                match ch {
                    'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' | '_'=> {
                        buf.push_char(ch);
                    }

                    '\r' | '\n' | ' ' | '\t' => { state = st_assign_wanted }

                    '=' => { state = st_value_wanted }

                    _ => { fail!() }
                }
            }
            st_assign_wanted => {
                if rd.eof() { fail!() }
                match read_char(rd) {
                    '\r' | '\n' | ' ' | '\t' => { }
                    '=' => { state = st_value_wanted }
                    _ => { fail!() }
                }
            }
            st_value_wanted => {
                visitor.pair(buf, parse_value(rd));
                buf.truncate(0);
                state = st_toplevel;
            }
        }
    }
}

fn main() {
  let contents = File::open(&Path::new(std::os::args()[1])).read_to_end();
  let mut rd = MemReader::new(contents);
  let mut visitor = MyVisitor::new();
  parse(&mut rd, &mut visitor);
  println!("{:s}", visitor.get_root().to_str());
}
<commit_msg>Visitor returns false on error<commit_after>/// A TOML [1] configuration file parser
///
/// [1]: https://github.com/mojombo/toml

use std::io::mem::MemReader;
use std::io::File;
use std::hashmap::HashMap;

fn read_char(rd: &mut MemReader) -> char {
  match rd.read_byte() {
    Some(b) => b as char,
    None => fail!()
  }
}

#[deriving(ToStr)]
enum Value {
    True,
    False,
    Unsigned(u64),
    Integer(i64),
    Float(f64),
    String(~str),
    Array(~[Value]),
    Datetime, // XXX
    Map(HashMap<~str, Value>) // XXX: This is no value
}

// parse values recursivly
fn parse_value(rd: &mut MemReader) -> Value {
    enum State {
        st_value,
        st_number,
        st_got_number
    }

    let mut state: State = st_value;
    let mut val: ~str = ~"";

    loop {
        match state {
            st_value => {
                if rd.eof() { fail!() }
                let ch = read_char(rd);
                match ch {
                    '\r' | '\n' | ' ' | '\t' => { }
                    '0' .. '9' => {
                        val.push_char(ch);
                        state = st_number;
                    }
                    't' => {
                        assert!(read_char(rd) == 'r');
                        assert!(read_char(rd) == 'u');
                        assert!(read_char(rd) == 'e');
                        return True
                    }
                    'f' => {
                        assert!(read_char(rd) == 'a');
                        assert!(read_char(rd) == 'l');
                        assert!(read_char(rd) == 's');
                        assert!(read_char(rd) == 'e');
                        return False
                    }
                    '"' => {
                        loop {
                            let ch = read_char(rd);
                            if ch == '"' {
                                break;
                            }
                            val.push_char(ch);
                        }
                        return String(val)
                    }
                    _ => { fail!() }
                }
            }
            st_got_number => {
                return Unsigned(from_str(val).unwrap())
            }
            st_number => {
                if rd.eof() {
                    state = st_got_number;
                } else {
                    let ch = read_char(rd);
                    match ch {
                        '\r' | '\n' | ' ' | '\t' => { state = st_got_number }
                        '#' => { skip_comment(rd); state = st_got_number }
                        '0' .. '9' => {
                            val.push_char(ch);
                        }
                        _ => { fail!() }
                    }
                }
            }
        }
    }
}

// We must be already within the '#"
fn skip_comment(rd: &mut MemReader) {
    loop {
        if rd.eof() { return }
        match read_char(rd) {
            '\n' => { return }
            _ => { }
        }
    }
}

trait Visitor {
    fn section(&mut self, name: &str, is_array: bool) -> bool;
    fn pair(&mut self, key: &str, val: Value) -> bool;
}

struct MyVisitor {
    root: HashMap<~str, Value>,
    current_path: ~str
}

impl MyVisitor {
    fn new() -> MyVisitor {
        MyVisitor { root: HashMap::new(), current_path: ~"" }
    }
    fn get_root<'a>(&'a self) -> &'a HashMap<~str, Value> {
        return &self.root;
    }
}

impl Visitor for MyVisitor {
    fn section(&mut self, name: &str, is_array: bool) -> bool {
        assert!(is_array == false);
        debug!("Section: {}", name);
        self.current_path = name.to_owned();
        return true
    }
    fn pair(&mut self, key: &str, val: Value) -> bool {
        debug!("Pair: {} {:s}", key, val.to_str());
        let mut m = self.root.find_or_insert(self.current_path.clone(), Map(HashMap::new())); // XXX: remove clone
        match *m {
            Map(ref mut map) => {
                let ok = map.insert(key.to_owned(), val);
                return ok
            }
            _ => { return false }
        }
        return true
    }
}

fn parse<V: Visitor>(rd: &mut MemReader, visitor: &mut V) {
    enum State {
        st_toplevel,
        st_comment,
        st_section,
        st_ident,
        st_assign_wanted,
        st_value_wanted
    }

    let mut state = st_toplevel;
    let mut buf: ~str = ~"";

    loop {
        match state {
            st_toplevel => {
                if rd.eof() { break }
                let ch = read_char(rd);
                match ch {
                    // ignore whitespace
                    '\r' | '\n' | ' ' | '\t' => { }

                    // comment
                    '#' => { state = st_comment }

                    // section
                    '[' => { state = st_section }

                    // identifier
                    'a' .. 'z' | 'A' .. 'Z' | '_' => {
                        state = st_ident;
                        buf.push_char(ch)
                    }

                    _ => { fail!() }
                }
            }
            st_comment => {
                if rd.eof() { break }
                match read_char(rd) {
                    '\n' => { state = st_toplevel }
                    _ => { }
                }
            }
            st_section => {
                if rd.eof() { fail!() }
                let ch = read_char(rd);
                match ch {
                    'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' | '.' | '_'=> {
                        buf.push_char(ch);
                    }
                    ']' => {
                        visitor.section(buf, false);
                        buf.truncate(0);
                        state = st_toplevel;
                    }
                    _ => { fail!() }
                }
            }
            st_ident => {
                if rd.eof() { fail!() }
                let ch = read_char(rd);
                match ch {
                    'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' | '_'=> {
                        buf.push_char(ch);
                    }

                    '\r' | '\n' | ' ' | '\t' => { state = st_assign_wanted }

                    '=' => { state = st_value_wanted }

                    _ => { fail!() }
                }
            }
            st_assign_wanted => {
                if rd.eof() { fail!() }
                match read_char(rd) {
                    '\r' | '\n' | ' ' | '\t' => { }
                    '=' => { state = st_value_wanted }
                    _ => { fail!() }
                }
            }
            st_value_wanted => {
                visitor.pair(buf, parse_value(rd));
                buf.truncate(0);
                state = st_toplevel;
            }
        }
    }
}

fn main() {
  let contents = File::open(&Path::new(std::os::args()[1])).read_to_end();
  let mut rd = MemReader::new(contents);
  let mut visitor = MyVisitor::new();
  parse(&mut rd, &mut visitor);
  println!("{:s}", visitor.get_root().to_str());
}
<|endoftext|>


<commit_before>/// A TOML [1] configuration file parser
///
/// Copyright (c) 2014 by Michael Neumann
///
/// [1]: https://github.com/mojombo/toml

use std::io::Buffer;
use std::hashmap::HashMap;
use std::char;

#[deriving(ToStr)]
pub enum Value {
    True,
    False,
    Unsigned(u64),
    Integer(i64),
    Float(f64),
    String(~str),
    Array(~[Value]),
    Datetime, // XXX
    Table(~[Value]),
    Map(HashMap<~str, Value>)
}

pub trait Visitor {
    fn section(&mut self, name: ~str, is_array: bool) -> bool;
    fn pair(&mut self, key: ~str, val: Value) -> bool;
}

pub struct ValueBuilder {
    root: HashMap<~str, Value>,
    current_section: ~str,
    section_is_array: bool
}

impl ValueBuilder {
    pub fn new() -> ValueBuilder {
        ValueBuilder { root: HashMap::new(), current_section: ~"", section_is_array: false }
    }

    fn insert(&mut self, path: &str, value: Value) -> bool {
        assert!(path.len() > 0);

        let p: ~[&str] = path.split_str(".").collect();
        assert!(p.len() > 0);

        ValueBuilder::ins(p, &mut self.root, value);

        return true;
    }

    fn ins(path: &[&str], ht: &mut HashMap<~str, Value>, val: Value) {
        println!("{:?}", path);
        if path.len() == 0 { return }
        let head = path.head().to_owned();

        if path.len() == 1 {
            let ok = ht.insert(head, val);
            assert!(ok);
            return;
        }
        else {
            let m = ht.find_or_insert(head, Map(HashMap::new()));
            match *m {
                Map(ref mut map) => {
                    ValueBuilder::ins(path.slice_from(1), map, val);
                }
                _ => { fail!() }
            }
        }
    }

    pub fn get_root<'a>(&'a self) -> &'a HashMap<~str, Value> {
        return &self.root;
    }
}

impl Visitor for ValueBuilder {
    fn section(&mut self, name: ~str, is_array: bool) -> bool {
        let ok = self.insert(name, Map(HashMap::new()));
        assert!(ok); // XXX

        self.current_section = name;
        self.section_is_array = is_array; // XXX: not implemented yet

        return true
    }

    fn pair(&mut self, key: ~str, val: Value) -> bool {
        if self.current_section.len() == 0 {
            let ok = self.insert(key, val);
            assert!(ok); // XXX
        } else {
            let path = self.current_section + "." + key;
            let ok = self.insert(path, val);
            assert!(ok); // XXX
        }

        return true;
    }
}

pub struct Parser<'a, BUF> {
    rd: &'a mut BUF,
    current_char: Option<char>
}

impl<'a, BUF: Buffer> Parser<'a, BUF> {
    pub fn new(rd: &'a mut BUF) -> Parser<'a, BUF> {
        let ch = rd.read_char();
        Parser { rd: rd, current_char: ch }
    }

    fn advance(&mut self) {
        self.current_char = self.rd.read_char();
    }

    fn ch(&self) -> Option<char> {
        return self.current_char;
    }

    fn eos(&self) -> bool {
        return self.current_char.is_none();
    }

    fn advance_if(&mut self, c: char) -> bool {
        match self.ch() {
            Some(ch) if ch == c => {
               self.advance();
               true
            }
            _ => {
                false
            }
        }
    }

    fn read_digit(&mut self, radix: uint) -> Option<u8> {
        if self.eos() { return None }
        match char::to_digit(self.ch().unwrap(), radix) {
            Some(n) => {
                self.advance();
                Some(n as u8)
            }
            None => { None }
        }
    }

    fn read_digits(&mut self) -> Option<u64> {
        let mut num: u64;
        match self.read_digit(10) {
            Some(n) => { num = n as u64; }
            None => { return None }
        }
        loop {
            match self.read_digit(10) {
                Some(n) => {
                    // XXX: check range
                    num = num * 10 + (n as u64);
                }
                None => {
                    return Some(num)
                }
            }
        }
    }

    // allows a single "."
    fn read_float_mantissa(&mut self) -> f64 {
        let mut num: f64 = 0.0;
        let mut div: f64 = 10.0;

        loop {
            match self.read_digit(10) {
                Some(n) => {
                    num = num + (n as f64)/div;
                    div = div * 10.0;
                }
                None => {
                    return num;
                }
            }
        }
    }

    fn parse_value(&mut self) -> Option<Value> {
        self.skip_whitespaces();

        if self.eos() { return None }
        match self.ch().unwrap() {
            '-' => {
                self.advance();
                match self.read_digits() {
                    Some(n) => {
                        if self.ch() == Some('.') {
                            // floating point
                            self.advance();
                            let num = self.read_float_mantissa();
                            let num = (n as f64) + num;
                            return Some(Float(-num));
                        }
                        else {
                            match n.to_i64() {
                                Some(i) => Some(Integer(-i)),
                                None => None // XXX: Use Result
                            }
                        }
                    }
                    None => {
                        return None
                    }
                }
            }
            '0' .. '9' => {
                match self.read_digits() {
                    Some(n) => {
                        match self.ch() {
                            Some('.') => {
                                // floating point
                                self.advance();
                                let num = self.read_float_mantissa();
                                let num = (n as f64) + num;
                                return Some(Float(num));
                            }
                            Some('-') => {
                                // XXX
                                fail!("Datetime not yet supported");
                            }
                            _ => {
                                return Some(Unsigned(n))
                            }
                        }
                    }
                    None => {
                        assert!(false);
                        return None
                    }
                }
            }
            't' => {
                self.advance();
                if self.advance_if('r') &&
                   self.advance_if('u') &&
                   self.advance_if('e') {
                    return Some(True)
                } else {
                    return None
                }
            }
            'f' => {
                self.advance();
                if self.advance_if('a') &&
                   self.advance_if('l') &&
                   self.advance_if('s') &&
                   self.advance_if('e') {
                    return Some(True)
                } else {
                    return None
                }
            }
            '[' => {
                self.advance();
                let mut arr = ~[];
                loop {
                    match self.parse_value() {
                        Some(val) => {
                            arr.push(val);
                        }
                        None => {
                            break;
                        }
                    }

                    self.skip_whitespaces_and_comments();
                    if !self.advance_if(',') { break }
                }
                self.skip_whitespaces_and_comments();
                if self.advance_if(']') {
                    return Some(Array(arr));
                } else {
                    return None;
                }
            }
            '"' => {
                match self.parse_string() {
                    Some(str) => { return Some(String(str)) }
                    None => { return None }
                }
            }
            _ => { return None }
        }
    }

    fn parse_string(&mut self) -> Option<~str> {
        if !self.advance_if('"') { return None }

        let mut str = ~"";
        loop {
            if self.ch().is_none() { return None }
            match self.ch().unwrap() {
                '\r' | '\n' | '\u000C' | '\u0008' => { return None }
                '\\' => {
                    self.advance();
                    if self.ch().is_none() { return None }
                    match self.ch().unwrap() {
                        'b' => { str.push_char('\u0008'); self.advance() },
                        't' => { str.push_char('\t'); self.advance() },
                        'n' => { str.push_char('\n'); self.advance() },
                        'f' => { str.push_char('\u000C'); self.advance() },
                        'r' => { str.push_char('\r'); self.advance() },
                        '"' => { str.push_char('"'); self.advance() },
                        '/' => { str.push_char('/'); self.advance() },
                        '\\' => { str.push_char('\\'); self.advance() },
                        'u' => {
                            self.advance();
                            let d1 = self.read_digit(16);
                            let d2 = self.read_digit(16);
                            let d3 = self.read_digit(16);
                            let d4 = self.read_digit(16);
                            match (d1, d2, d3, d4) {
                                (Some(d1), Some(d2), Some(d3), Some(d4)) => {
                                    // XXX: how to construct an UTF character
                                    let ch = (((((d1 as u32 << 8) | d2 as u32) << 8) | d3 as u32) << 8) | d4 as u32;
                                    match char::from_u32(ch) {
                                        Some(ch) => {
                                            str.push_char(ch);
                                        }
                                        None => {
                                            return None;
                                        }
                                    }
                                }
                                _ => return None
                            }
                        }
                        _ => { return None }
                    }
                }
                '"' => {
                    self.advance();
                    return Some(str);
                }
                c => {
                    str.push_char(c);
                    self.advance();
                }
            }
        }
    }

    fn read_token(&mut self, f: |char| -> bool) -> ~str {
        let mut token = ~"";
        loop {
            match self.ch() {
                Some(ch) => {
                    if f(ch) { token.push_char(ch) }
                    else { break }
                }
                None => { break }
            }
            self.advance();
        }

        return token;
    }

    fn parse_section_identifier(&mut self) -> ~str {
        self.read_token(|ch| {
            match ch {
                'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' | '.' | '_' => true,
                _ => false
            }
        })
    }

    fn skip_whitespaces(&mut self) {
        loop {
            match self.ch() {
                Some(' ') | Some('\t') | Some('\r') | Some('\n') => {
                    self.advance();
                }
                _ => { break }
            }
        }
    }

    fn skip_whitespaces_and_comments(&mut self) {
        loop {
            match self.ch() {
                Some(' ') | Some('\t') | Some('\r') | Some('\n') => {
                    self.advance();
                }
                Some('#') => {
                    self.skip_comment();
                }
                _ => { break }
            }
        }
    }

    fn skip_comment(&mut self) {
        assert!(self.ch() == Some('#'));
        // skip to end of line
        loop {
            self.advance();
            match self.ch() {
                Some('\n') => { break }
                None => { return }
                _ => { /* skip */ }
            }
        }
        self.advance();
    }

    pub fn parse<V: Visitor>(&mut self, visitor: &mut V) -> bool {
        loop {
            if self.eos() { return true }
            match self.ch().unwrap() {
                // ignore whitespace
                '\r' | '\n' | ' ' | '\t' => {
                    self.advance();
                }

                // comment
                '#' => {
                    self.skip_comment();
                }

                // section
                '[' => {
                    self.advance();
                    let mut double_section = false;
                    match self.ch() {
                        Some('[') => {
                            double_section = true;
                            self.advance();
                        }
                        _ => {}
                    }

                    let section_name = self.parse_section_identifier();

                    if !self.advance_if(']') { return false }
                    if double_section {
                        if !self.advance_if(']') { return false }
                    }

                    visitor.section(section_name, double_section);
                }

                // identifier
                'a' .. 'z' | 'A' .. 'Z' | '_' => {

                    let ident = self.read_token(|ch| {
                        match ch {
                            'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' | '_' => true,
                            _ => false
                        }
                    });

                    self.skip_whitespaces();

                    if !self.advance_if('=') { return false } // assign wanted

                    match self.parse_value() {
                        Some(val) => { visitor.pair(ident, val); }
                        None => { return false; }
                    }
                    // do not advance!
                }

                _ => { return false }
            } /* end match */
        }

        assert!(false);
    }
}
<commit_msg>Handle duplicate keys better<commit_after>/// A TOML [1] configuration file parser
///
/// Copyright (c) 2014 by Michael Neumann
///
/// [1]: https://github.com/mojombo/toml

use std::io::Buffer;
use std::hashmap::HashMap;
use std::char;

#[deriving(ToStr)]
pub enum Value {
    True,
    False,
    Unsigned(u64),
    Integer(i64),
    Float(f64),
    String(~str),
    Array(~[Value]),
    Datetime, // XXX
    Table(~[Value]),
    Map(HashMap<~str, Value>)
}

pub trait Visitor {
    fn section(&mut self, name: ~str, is_array: bool) -> bool;
    fn pair(&mut self, key: ~str, val: Value) -> bool;
}

pub struct ValueBuilder {
    root: HashMap<~str, Value>,
    current_section: ~str,
    section_is_array: bool
}

impl ValueBuilder {
    pub fn new() -> ValueBuilder {
        ValueBuilder { root: HashMap::new(), current_section: ~"", section_is_array: false }
    }

    fn insert(&mut self, path: &str, value: Value) -> bool {
        let path: ~[&str] = path.split_str(".").collect();
        assert!(path.len() > 0);

        return ValueBuilder::ins(path, &mut self.root, value);
    }

    fn ins(path: &[&str], ht: &mut HashMap<~str, Value>, val: Value) -> bool {
        assert!(path.len() > 0);

        let head = path.head().to_owned();

        if path.len() == 1 {
            let ok = ht.insert(head, val);
            if !ok {
                debug!("Duplicate key");
            }
            return ok;
        }
        else {
            let m = ht.find_or_insert(head, Map(HashMap::new())); // Optimize
            match *m {
                Map(ref mut map) => {
                    return ValueBuilder::ins(path.slice_from(1), map, val);
                }
                _ => {
                    debug!("Wrong type/duplicate key");
                    return false;
                }
            }
        }
    }

    pub fn get_root<'a>(&'a self) -> &'a HashMap<~str, Value> {
        return &self.root;
    }
}

impl Visitor for ValueBuilder {
    fn section(&mut self, name: ~str, is_array: bool) -> bool {
        let ok = self.insert(name, Map(HashMap::new()));
        if !ok {
            debug!("Duplicate key: {}", name);
        }

        self.current_section = name;
        self.section_is_array = is_array; // XXX: not implemented yet

        return ok;
    }

    fn pair(&mut self, key: ~str, val: Value) -> bool {
        if self.current_section.len() == 0 {
            let ok = self.insert(key, val);
            if !ok {
                debug!("Duplicate key: {}", key);
            }
            return ok;
        } else {
            let path = self.current_section + "." + key;
            let ok = self.insert(path, val);
            if !ok {
                debug!("Duplicate key: {}", path);
            }
            return ok;
        }
    }
}

pub struct Parser<'a, BUF> {
    rd: &'a mut BUF,
    current_char: Option<char>
}

impl<'a, BUF: Buffer> Parser<'a, BUF> {
    pub fn new(rd: &'a mut BUF) -> Parser<'a, BUF> {
        let ch = rd.read_char();
        Parser { rd: rd, current_char: ch }
    }

    fn advance(&mut self) {
        self.current_char = self.rd.read_char();
    }

    fn ch(&self) -> Option<char> {
        return self.current_char;
    }

    fn eos(&self) -> bool {
        return self.current_char.is_none();
    }

    fn advance_if(&mut self, c: char) -> bool {
        match self.ch() {
            Some(ch) if ch == c => {
               self.advance();
               true
            }
            _ => {
                false
            }
        }
    }

    fn read_digit(&mut self, radix: uint) -> Option<u8> {
        if self.eos() { return None }
        match char::to_digit(self.ch().unwrap(), radix) {
            Some(n) => {
                self.advance();
                Some(n as u8)
            }
            None => { None }
        }
    }

    fn read_digits(&mut self) -> Option<u64> {
        let mut num: u64;
        match self.read_digit(10) {
            Some(n) => { num = n as u64; }
            None => { return None }
        }
        loop {
            match self.read_digit(10) {
                Some(n) => {
                    // XXX: check range
                    num = num * 10 + (n as u64);
                }
                None => {
                    return Some(num)
                }
            }
        }
    }

    // allows a single "."
    fn read_float_mantissa(&mut self) -> f64 {
        let mut num: f64 = 0.0;
        let mut div: f64 = 10.0;

        loop {
            match self.read_digit(10) {
                Some(n) => {
                    num = num + (n as f64)/div;
                    div = div * 10.0;
                }
                None => {
                    return num;
                }
            }
        }
    }

    fn parse_value(&mut self) -> Option<Value> {
        self.skip_whitespaces();

        if self.eos() { return None }
        match self.ch().unwrap() {
            '-' => {
                self.advance();
                match self.read_digits() {
                    Some(n) => {
                        if self.ch() == Some('.') {
                            // floating point
                            self.advance();
                            let num = self.read_float_mantissa();
                            let num = (n as f64) + num;
                            return Some(Float(-num));
                        }
                        else {
                            match n.to_i64() {
                                Some(i) => Some(Integer(-i)),
                                None => None // XXX: Use Result
                            }
                        }
                    }
                    None => {
                        return None
                    }
                }
            }
            '0' .. '9' => {
                match self.read_digits() {
                    Some(n) => {
                        match self.ch() {
                            Some('.') => {
                                // floating point
                                self.advance();
                                let num = self.read_float_mantissa();
                                let num = (n as f64) + num;
                                return Some(Float(num));
                            }
                            Some('-') => {
                                // XXX
                                fail!("Datetime not yet supported");
                            }
                            _ => {
                                return Some(Unsigned(n))
                            }
                        }
                    }
                    None => {
                        assert!(false);
                        return None
                    }
                }
            }
            't' => {
                self.advance();
                if self.advance_if('r') &&
                   self.advance_if('u') &&
                   self.advance_if('e') {
                    return Some(True)
                } else {
                    return None
                }
            }
            'f' => {
                self.advance();
                if self.advance_if('a') &&
                   self.advance_if('l') &&
                   self.advance_if('s') &&
                   self.advance_if('e') {
                    return Some(True)
                } else {
                    return None
                }
            }
            '[' => {
                self.advance();
                let mut arr = ~[];
                loop {
                    match self.parse_value() {
                        Some(val) => {
                            arr.push(val);
                        }
                        None => {
                            break;
                        }
                    }

                    self.skip_whitespaces_and_comments();
                    if !self.advance_if(',') { break }
                }
                self.skip_whitespaces_and_comments();
                if self.advance_if(']') {
                    return Some(Array(arr));
                } else {
                    return None;
                }
            }
            '"' => {
                match self.parse_string() {
                    Some(str) => { return Some(String(str)) }
                    None => { return None }
                }
            }
            _ => { return None }
        }
    }

    fn parse_string(&mut self) -> Option<~str> {
        if !self.advance_if('"') { return None }

        let mut str = ~"";
        loop {
            if self.ch().is_none() { return None }
            match self.ch().unwrap() {
                '\r' | '\n' | '\u000C' | '\u0008' => { return None }
                '\\' => {
                    self.advance();
                    if self.ch().is_none() { return None }
                    match self.ch().unwrap() {
                        'b' => { str.push_char('\u0008'); self.advance() },
                        't' => { str.push_char('\t'); self.advance() },
                        'n' => { str.push_char('\n'); self.advance() },
                        'f' => { str.push_char('\u000C'); self.advance() },
                        'r' => { str.push_char('\r'); self.advance() },
                        '"' => { str.push_char('"'); self.advance() },
                        '/' => { str.push_char('/'); self.advance() },
                        '\\' => { str.push_char('\\'); self.advance() },
                        'u' => {
                            self.advance();
                            let d1 = self.read_digit(16);
                            let d2 = self.read_digit(16);
                            let d3 = self.read_digit(16);
                            let d4 = self.read_digit(16);
                            match (d1, d2, d3, d4) {
                                (Some(d1), Some(d2), Some(d3), Some(d4)) => {
                                    // XXX: how to construct an UTF character
                                    let ch = (((((d1 as u32 << 8) | d2 as u32) << 8) | d3 as u32) << 8) | d4 as u32;
                                    match char::from_u32(ch) {
                                        Some(ch) => {
                                            str.push_char(ch);
                                        }
                                        None => {
                                            return None;
                                        }
                                    }
                                }
                                _ => return None
                            }
                        }
                        _ => { return None }
                    }
                }
                '"' => {
                    self.advance();
                    return Some(str);
                }
                c => {
                    str.push_char(c);
                    self.advance();
                }
            }
        }
    }

    fn read_token(&mut self, f: |char| -> bool) -> ~str {
        let mut token = ~"";
        loop {
            match self.ch() {
                Some(ch) => {
                    if f(ch) { token.push_char(ch) }
                    else { break }
                }
                None => { break }
            }
            self.advance();
        }

        return token;
    }

    fn parse_section_identifier(&mut self) -> ~str {
        self.read_token(|ch| {
            match ch {
                'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' | '.' | '_' => true,
                _ => false
            }
        })
    }

    fn skip_whitespaces(&mut self) {
        loop {
            match self.ch() {
                Some(' ') | Some('\t') | Some('\r') | Some('\n') => {
                    self.advance();
                }
                _ => { break }
            }
        }
    }

    fn skip_whitespaces_and_comments(&mut self) {
        loop {
            match self.ch() {
                Some(' ') | Some('\t') | Some('\r') | Some('\n') => {
                    self.advance();
                }
                Some('#') => {
                    self.skip_comment();
                }
                _ => { break }
            }
        }
    }

    fn skip_comment(&mut self) {
        assert!(self.ch() == Some('#'));
        // skip to end of line
        loop {
            self.advance();
            match self.ch() {
                Some('\n') => { break }
                None => { return }
                _ => { /* skip */ }
            }
        }
        self.advance();
    }

    pub fn parse<V: Visitor>(&mut self, visitor: &mut V) -> bool {
        loop {
            if self.eos() { return true }
            match self.ch().unwrap() {
                // ignore whitespace
                '\r' | '\n' | ' ' | '\t' => {
                    self.advance();
                }

                // comment
                '#' => {
                    self.skip_comment();
                }

                // section
                '[' => {
                    self.advance();
                    let mut double_section = false;
                    match self.ch() {
                        Some('[') => {
                            double_section = true;
                            self.advance();
                        }
                        _ => {}
                    }

                    let section_name = self.parse_section_identifier();

                    if !self.advance_if(']') { return false }
                    if double_section {
                        if !self.advance_if(']') { return false }
                    }

                    visitor.section(section_name, double_section);
                }

                // identifier
                'a' .. 'z' | 'A' .. 'Z' | '_' => {

                    let ident = self.read_token(|ch| {
                        match ch {
                            'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' | '_' => true,
                            _ => false
                        }
                    });

                    self.skip_whitespaces();

                    if !self.advance_if('=') { return false } // assign wanted

                    match self.parse_value() {
                        Some(val) => { visitor.pair(ident, val); }
                        None => { return false; }
                    }
                    // do not advance!
                }

                _ => { return false }
            } /* end match */
        }

        assert!(false);
    }
}
<|endoftext|>