<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.

//! Small system utility modules for usage by other modules.

extern crate data_model;
extern crate libc;
extern crate syscall_defines;
#[allow(unused_imports)]
#[macro_use]
extern crate poll_token_derive;

#[macro_use]
pub mod handle_eintr;
#[macro_use]
pub mod ioctl;
#[macro_use]
pub mod syslog;
mod mmap;
mod shm;
mod eventfd;
mod errno;
mod guest_address;
mod guest_memory;
mod poll;
mod struct_util;
mod tempdir;
mod terminal;
pub mod signal;
mod fork;
mod signalfd;
mod sock_ctrl_msg;
mod passwd;
mod file_flags;
mod timerfd;
mod write_zeroes;

pub use mmap::*;
pub use shm::*;
pub use eventfd::*;
pub use errno::{Error, Result};
use errno::errno_result;
pub use guest_address::*;
pub use guest_memory::*;
pub use poll::*;
pub use struct_util::*;
pub use tempdir::*;
pub use terminal::*;
pub use signal::*;
pub use fork::*;
pub use signalfd::*;
pub use ioctl::*;
pub use sock_ctrl_msg::*;
pub use passwd::*;
pub use poll_token_derive::*;
pub use file_flags::*;
pub use timerfd::*;

pub use mmap::Error as MmapError;
pub use guest_memory::Error as GuestMemoryError;
pub use signalfd::Error as SignalFdError;
pub use write_zeroes::WriteZeroes as WriteZeroes;

use std::ffi::CStr;
use std::fs::{File, remove_file};
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::os::unix::net::UnixDatagram;
use std::ptr;

use libc::{kill, syscall, sysconf, waitpid, pipe2, c_long, pid_t, uid_t, gid_t, _SC_PAGESIZE,
           SIGKILL, WNOHANG, O_CLOEXEC};

use syscall_defines::linux::LinuxSyscall::SYS_getpid;

/// Safe wrapper for `sysconf(_SC_PAGESIZE)`.
#[inline(always)]
pub fn pagesize() -> usize {
    // Trivially safe
    unsafe { sysconf(_SC_PAGESIZE) as usize }
}

/// Uses the system's page size in bytes to round the given value up to the nearest page boundary.
#[inline(always)]
pub fn round_up_to_page_size(v: usize) -> usize {
    let page_mask = pagesize() - 1;
    (v + page_mask) & !page_mask
}

/// This bypasses `libc`'s caching `getpid(2)` wrapper which can be invalid if a raw clone was used
/// elsewhere.
#[inline(always)]
pub fn getpid() -> pid_t {
    // Safe because this syscall can never fail and we give it a valid syscall number.
    unsafe { syscall(SYS_getpid as c_long) as pid_t }
}

/// Safe wrapper for `geteuid(2)`.
#[inline(always)]
pub fn geteuid() -> uid_t {
    // trivially safe
    unsafe { libc::geteuid() }
}

/// Safe wrapper for `getegid(2)`.
#[inline(always)]
pub fn getegid() -> gid_t {
    // trivially safe
    unsafe { libc::getegid() }
}

/// Safe wrapper for chown(2).
#[inline(always)]
pub fn chown(path: &CStr, uid: uid_t, gid: gid_t) -> Result<()> {
    // Safe since we pass in a valid string pointer and check the return value.
    let ret = unsafe { libc::chown(path.as_ptr(), uid, gid) };

    if ret < 0 {
        errno_result()
    } else {
        Ok(())
    }
}

/// The operation to perform with `flock`.
pub enum FlockOperation {
    LockShared,
    LockExclusive,
    Unlock,
}

/// Safe wrapper for flock(2) with the operation `op` and optionally `nonblocking`. The lock will be
/// dropped automatically when `file` is dropped.
#[inline(always)]
pub fn flock(file: &AsRawFd, op: FlockOperation, nonblocking: bool) -> Result<()> {
    let mut operation = match op {
        FlockOperation::LockShared => libc::LOCK_SH,
        FlockOperation::LockExclusive => libc::LOCK_EX,
        FlockOperation::Unlock => libc::LOCK_UN,
    };

    if nonblocking {
        operation |= libc::LOCK_NB;
    }

    // Safe since we pass in a valid fd and flock operation, and check the return value.
    let ret = unsafe { libc::flock(file.as_raw_fd(), operation) };

    if ret < 0 {
        errno_result()
    } else {
        Ok(())
    }
}

/// The operation to perform with `fallocate`.
pub enum FallocateMode {
    PunchHole,
    ZeroRange,
}

/// Safe wrapper for `fallocate()`.
pub fn fallocate(
    file: &AsRawFd,
    mode: FallocateMode,
    keep_size: bool,
    offset: u64,
    len: u64
) -> Result<()> {
    let offset = if offset > libc::off_t::max_value() as u64 {
        return Err(Error::new(libc::EINVAL));
    } else {
        offset as libc::off_t
    };

    let len = if len > libc::off_t::max_value() as u64 {
        return Err(Error::new(libc::EINVAL));
    } else {
        len as libc::off_t
    };

    let mut mode = match mode {
        FallocateMode::PunchHole => libc::FALLOC_FL_PUNCH_HOLE,
        FallocateMode::ZeroRange => libc::FALLOC_FL_ZERO_RANGE,
    };

    if keep_size {
        mode |= libc::FALLOC_FL_KEEP_SIZE;
    }

    // Safe since we pass in a valid fd and fallocate mode, validate offset and len,
    // and check the return value.
    let ret = unsafe { libc::fallocate(file.as_raw_fd(), mode, offset, len) };
    if ret < 0 {
        errno_result()
    } else {
        Ok(())
    }
}

/// Reaps a child process that has terminated.
///
/// Returns `Ok(pid)` where `pid` is the process that was reaped or `Ok(0)` if none of the children
/// have terminated. An `Error` is with `errno == ECHILD` if there are no children left to reap.
///
/// # Examples
///
/// Reaps all child processes until there are no terminated children to reap.
///
/// ```
/// # extern crate libc;
/// # extern crate sys_util;
/// fn reap_children() {
///     loop {
///         match sys_util::reap_child() {
///             Ok(0) => println!("no children ready to reap"),
///             Ok(pid) => {
///                 println!("reaped {}", pid);
///                 continue
///             },
///             Err(e) if e.errno() == libc::ECHILD => println!("no children left"),
///             Err(e) => println!("error reaping children: {:?}", e),
///         }
///         break
///     }
/// }
/// ```
pub fn reap_child() -> Result<pid_t> {
    // Safe because we pass in no memory, prevent blocking with WNOHANG, and check for error.
    let ret = unsafe { waitpid(-1, ptr::null_mut(), WNOHANG) };
    if ret == -1 {
        errno_result()
    } else {
        Ok(ret)
    }
}

/// Kill all processes in the current process group.
///
/// On success, this kills all processes in the current process group, including the current
/// process, meaning this will not return. This is equivalent to a call to `kill(0, SIGKILL)`.
pub fn kill_process_group() -> Result<()> {
    let ret = unsafe { kill(0, SIGKILL) };
    if ret == -1 {
        errno_result()
    } else {
        // Kill succeeded, so this process never reaches here.
        unreachable!();
    }
}

/// Spawns a pipe pair where the first pipe is the read end and the second pipe is the write end.
///
/// If `close_on_exec` is true, the `O_CLOEXEC` flag will be set during pipe creation.
pub fn pipe(close_on_exec: bool) -> Result<(File, File)> {
    let flags = if close_on_exec { O_CLOEXEC } else { 0 };
    let mut pipe_fds = [-1; 2];
    // Safe because pipe2 will only write 2 element array of i32 to the given pointer, and we check
    // for error.
    let ret = unsafe { pipe2(&mut pipe_fds[0], flags) };
    if ret == -1 {
        errno_result()
    } else {
        // Safe because both fds must be valid for pipe2 to have returned sucessfully and we have
        // exclusive ownership of them.
        Ok(unsafe {
            (
                File::from_raw_fd(pipe_fds[0]),
                File::from_raw_fd(pipe_fds[1])
            )
        })
    }
}

/// Used to attempt to clean up a named pipe after it is no longer used.
pub struct UnlinkUnixDatagram(pub UnixDatagram);
impl AsRef<UnixDatagram> for UnlinkUnixDatagram {
    fn as_ref(&self) -> &UnixDatagram {
        &self.0
    }
}
impl Drop for UnlinkUnixDatagram {
    fn drop(&mut self) {
        if let Ok(addr) = self.0.local_addr() {
            if let Some(path) = addr.as_pathname() {
                if let Err(e) = remove_file(path) {
                    warn!("failed to remove control socket file: {:?}", e);
                }
            }
        }
    }
}

<commit_msg>sys_util: use fallocate64 for large file support<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.

//! Small system utility modules for usage by other modules.

extern crate data_model;
extern crate libc;
extern crate syscall_defines;
#[allow(unused_imports)]
#[macro_use]
extern crate poll_token_derive;

#[macro_use]
pub mod handle_eintr;
#[macro_use]
pub mod ioctl;
#[macro_use]
pub mod syslog;
mod mmap;
mod shm;
mod eventfd;
mod errno;
mod guest_address;
mod guest_memory;
mod poll;
mod struct_util;
mod tempdir;
mod terminal;
pub mod signal;
mod fork;
mod signalfd;
mod sock_ctrl_msg;
mod passwd;
mod file_flags;
mod timerfd;
mod write_zeroes;

pub use mmap::*;
pub use shm::*;
pub use eventfd::*;
pub use errno::{Error, Result};
use errno::errno_result;
pub use guest_address::*;
pub use guest_memory::*;
pub use poll::*;
pub use struct_util::*;
pub use tempdir::*;
pub use terminal::*;
pub use signal::*;
pub use fork::*;
pub use signalfd::*;
pub use ioctl::*;
pub use sock_ctrl_msg::*;
pub use passwd::*;
pub use poll_token_derive::*;
pub use file_flags::*;
pub use timerfd::*;

pub use mmap::Error as MmapError;
pub use guest_memory::Error as GuestMemoryError;
pub use signalfd::Error as SignalFdError;
pub use write_zeroes::WriteZeroes as WriteZeroes;

use std::ffi::CStr;
use std::fs::{File, remove_file};
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::os::unix::net::UnixDatagram;
use std::ptr;

use libc::{kill, syscall, sysconf, waitpid, pipe2, c_long, pid_t, uid_t, gid_t, _SC_PAGESIZE,
           SIGKILL, WNOHANG, O_CLOEXEC};

use syscall_defines::linux::LinuxSyscall::SYS_getpid;

/// Safe wrapper for `sysconf(_SC_PAGESIZE)`.
#[inline(always)]
pub fn pagesize() -> usize {
    // Trivially safe
    unsafe { sysconf(_SC_PAGESIZE) as usize }
}

/// Uses the system's page size in bytes to round the given value up to the nearest page boundary.
#[inline(always)]
pub fn round_up_to_page_size(v: usize) -> usize {
    let page_mask = pagesize() - 1;
    (v + page_mask) & !page_mask
}

/// This bypasses `libc`'s caching `getpid(2)` wrapper which can be invalid if a raw clone was used
/// elsewhere.
#[inline(always)]
pub fn getpid() -> pid_t {
    // Safe because this syscall can never fail and we give it a valid syscall number.
    unsafe { syscall(SYS_getpid as c_long) as pid_t }
}

/// Safe wrapper for `geteuid(2)`.
#[inline(always)]
pub fn geteuid() -> uid_t {
    // trivially safe
    unsafe { libc::geteuid() }
}

/// Safe wrapper for `getegid(2)`.
#[inline(always)]
pub fn getegid() -> gid_t {
    // trivially safe
    unsafe { libc::getegid() }
}

/// Safe wrapper for chown(2).
#[inline(always)]
pub fn chown(path: &CStr, uid: uid_t, gid: gid_t) -> Result<()> {
    // Safe since we pass in a valid string pointer and check the return value.
    let ret = unsafe { libc::chown(path.as_ptr(), uid, gid) };

    if ret < 0 {
        errno_result()
    } else {
        Ok(())
    }
}

/// The operation to perform with `flock`.
pub enum FlockOperation {
    LockShared,
    LockExclusive,
    Unlock,
}

/// Safe wrapper for flock(2) with the operation `op` and optionally `nonblocking`. The lock will be
/// dropped automatically when `file` is dropped.
#[inline(always)]
pub fn flock(file: &AsRawFd, op: FlockOperation, nonblocking: bool) -> Result<()> {
    let mut operation = match op {
        FlockOperation::LockShared => libc::LOCK_SH,
        FlockOperation::LockExclusive => libc::LOCK_EX,
        FlockOperation::Unlock => libc::LOCK_UN,
    };

    if nonblocking {
        operation |= libc::LOCK_NB;
    }

    // Safe since we pass in a valid fd and flock operation, and check the return value.
    let ret = unsafe { libc::flock(file.as_raw_fd(), operation) };

    if ret < 0 {
        errno_result()
    } else {
        Ok(())
    }
}

/// The operation to perform with `fallocate`.
pub enum FallocateMode {
    PunchHole,
    ZeroRange,
}

// TODO(dverkamp): Remove this once fallocate64 is available in libc.
extern {
    pub fn fallocate64(
        fd: libc::c_int,
        mode: libc::c_int,
        offset: libc::off64_t,
        len: libc::off64_t
    ) -> libc::c_int;
}

/// Safe wrapper for `fallocate()`.
pub fn fallocate(
    file: &AsRawFd,
    mode: FallocateMode,
    keep_size: bool,
    offset: u64,
    len: u64
) -> Result<()> {
    let offset = if offset > libc::off64_t::max_value() as u64 {
        return Err(Error::new(libc::EINVAL));
    } else {
        offset as libc::off64_t
    };

    let len = if len > libc::off64_t::max_value() as u64 {
        return Err(Error::new(libc::EINVAL));
    } else {
        len as libc::off64_t
    };

    let mut mode = match mode {
        FallocateMode::PunchHole => libc::FALLOC_FL_PUNCH_HOLE,
        FallocateMode::ZeroRange => libc::FALLOC_FL_ZERO_RANGE,
    };

    if keep_size {
        mode |= libc::FALLOC_FL_KEEP_SIZE;
    }

    // Safe since we pass in a valid fd and fallocate mode, validate offset and len,
    // and check the return value.
    // TODO(dverkamp): Replace this with libc::fallocate64 once it is available.
    let ret = unsafe { fallocate64(file.as_raw_fd(), mode, offset, len) };
    if ret < 0 {
        errno_result()
    } else {
        Ok(())
    }
}

/// Reaps a child process that has terminated.
///
/// Returns `Ok(pid)` where `pid` is the process that was reaped or `Ok(0)` if none of the children
/// have terminated. An `Error` is with `errno == ECHILD` if there are no children left to reap.
///
/// # Examples
///
/// Reaps all child processes until there are no terminated children to reap.
///
/// ```
/// # extern crate libc;
/// # extern crate sys_util;
/// fn reap_children() {
///     loop {
///         match sys_util::reap_child() {
///             Ok(0) => println!("no children ready to reap"),
///             Ok(pid) => {
///                 println!("reaped {}", pid);
///                 continue
///             },
///             Err(e) if e.errno() == libc::ECHILD => println!("no children left"),
///             Err(e) => println!("error reaping children: {:?}", e),
///         }
///         break
///     }
/// }
/// ```
pub fn reap_child() -> Result<pid_t> {
    // Safe because we pass in no memory, prevent blocking with WNOHANG, and check for error.
    let ret = unsafe { waitpid(-1, ptr::null_mut(), WNOHANG) };
    if ret == -1 {
        errno_result()
    } else {
        Ok(ret)
    }
}

/// Kill all processes in the current process group.
///
/// On success, this kills all processes in the current process group, including the current
/// process, meaning this will not return. This is equivalent to a call to `kill(0, SIGKILL)`.
pub fn kill_process_group() -> Result<()> {
    let ret = unsafe { kill(0, SIGKILL) };
    if ret == -1 {
        errno_result()
    } else {
        // Kill succeeded, so this process never reaches here.
        unreachable!();
    }
}

/// Spawns a pipe pair where the first pipe is the read end and the second pipe is the write end.
///
/// If `close_on_exec` is true, the `O_CLOEXEC` flag will be set during pipe creation.
pub fn pipe(close_on_exec: bool) -> Result<(File, File)> {
    let flags = if close_on_exec { O_CLOEXEC } else { 0 };
    let mut pipe_fds = [-1; 2];
    // Safe because pipe2 will only write 2 element array of i32 to the given pointer, and we check
    // for error.
    let ret = unsafe { pipe2(&mut pipe_fds[0], flags) };
    if ret == -1 {
        errno_result()
    } else {
        // Safe because both fds must be valid for pipe2 to have returned sucessfully and we have
        // exclusive ownership of them.
        Ok(unsafe {
            (
                File::from_raw_fd(pipe_fds[0]),
                File::from_raw_fd(pipe_fds[1])
            )
        })
    }
}

/// Used to attempt to clean up a named pipe after it is no longer used.
pub struct UnlinkUnixDatagram(pub UnixDatagram);
impl AsRef<UnixDatagram> for UnlinkUnixDatagram {
    fn as_ref(&self) -> &UnixDatagram {
        &self.0
    }
}
impl Drop for UnlinkUnixDatagram {
    fn drop(&mut self) {
        if let Ok(addr) = self.0.local_addr() {
            if let Some(path) = addr.as_pathname() {
                if let Err(e) = remove_file(path) {
                    warn!("failed to remove control socket file: {:?}", e);
                }
            }
        }
    }
}

<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

#[cfg(feature = "serde")]
#[doc(hidden)]
#[macro_export]
macro_rules! add_serde_if_enabled {
    ($(#[$attr:meta])* -- $i:item) => {
        $(#[$attr])*
        #[derive($crate::serde::Serialize, $crate::serde::Deserialize)]
        $i
    }
}

#[cfg(not(feature = "serde"))]
#[doc(hidden)]
#[macro_export]
macro_rules! add_serde_if_enabled {
    ($(#[$attr:meta])* -- $i:item) => {
        $(#[$attr])*
        $i
    }
}

/// The main macro that creates RPC services.
///
/// Rpc methods are specified, mirroring trait syntax:
///
/// ```
/// # #![feature(await_macro, pin, arbitrary_self_types, async_await, futures_api, proc_macro_hygiene)]
/// # fn main() {}
/// # tarpc::service! {
/// /// Say hello
/// rpc hello(name: String) -> String;
/// # }
/// ```
///
/// Attributes can be attached to each rpc. These attributes
/// will then be attached to the generated service traits'
/// corresponding `fn`s, as well as to the client stubs' RPCs.
///
/// The following items are expanded in the enclosing module:
///
/// * `trait Service` -- defines the RPC service.
///   * `fn serve` -- turns a service impl into a request handler.
/// * `Client` -- a client stub with a fn for each RPC.
///   * `fn new_stub` -- creates a new Client stub.
///
#[macro_export]
macro_rules! service {
// Entry point
    (
        $(
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) $(-> $out:ty)*;
        )*
    ) => {
        $crate::service! {{
            $(
                $(#[$attr])*
                rpc $fn_name( $( $arg : $in_ ),* ) $(-> $out)*;
            )*
        }}
    };
// Pattern for when the next rpc has an implicit unit return type.
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* );

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        $crate::service! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> ();
        }
    };
// Pattern for when the next rpc has an explicit return type.
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        $crate::service! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> $out;
        }
    };
// Pattern for when all return types have been expanded
    (
        { } // none left to expand
        $(
            $(#[$attr:meta])*
            rpc $fn_name:ident ( $( $arg:ident : $in_:ty ),* ) -> $out:ty;
        )*
    ) => {
        $crate::add_serde_if_enabled! {
            #[derive(Debug)]
            #[doc(hidden)]
            #[allow(non_camel_case_types, unused)]
            --
            pub enum Request__ {
                $(
                    $fn_name{ $($arg: $in_,)* }
                ),*
            }
        }

        $crate::add_serde_if_enabled! {
            #[derive(Debug)]
            #[doc(hidden)]
            #[allow(non_camel_case_types, unused)]
            --
            pub enum Response__ {
                $(
                    $fn_name($out)
                ),*
            }
        }

        // TODO: proc_macro can't currently parse $crate, so this needs to be imported for the
        // usage of snake_to_camel! to work.
        use $crate::futures::Future as Future__;

        /// Defines the RPC service. The additional trait bounds are required so that services can
        /// multiplex requests across multiple tasks, potentially on multiple threads.
        pub trait Service: Clone + Send + 'static {
            $(
                $crate::snake_to_camel! {
                    /// The type of future returned by `{}`.
                    type $fn_name: Future__<Output = $out> + Send;
                }

                $(#[$attr])*
                fn $fn_name(&self, ctx: $crate::context::Context, $($arg:$in_),*) -> $crate::ty_snake_to_camel!(Self::$fn_name);
            )*
        }

        // TODO: use an existential type instead of this when existential types work.
        #[allow(non_camel_case_types)]
        pub enum Response<S: Service> {
            $(
                $fn_name($crate::ty_snake_to_camel!(<S as Service>::$fn_name)),
            )*
        }

        impl<S: Service> ::std::fmt::Debug for Response<S> {
            fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
                fmt.debug_struct("Response").finish()
            }
        }

        impl<S: Service> ::std::future::Future for Response<S> {
            type Output = ::std::io::Result<Response__>;

            fn poll(self: ::std::pin::Pin<&mut Self>, waker: &::std::task::LocalWaker)
                -> ::std::task::Poll<::std::io::Result<Response__>>
            {
                unsafe {
                    match ::std::pin::Pin::get_mut_unchecked(self) {
                        $(
                            Response::$fn_name(resp) =>
                                ::std::pin::Pin::new_unchecked(resp)
                                    .poll(waker)
                                    .map(Response__::$fn_name)
                                    .map(Ok),
                        )*
                    }
                }
            }
        }

        /// Returns a serving function to use with rpc::server::Server.
        pub fn serve<S: Service>(service: S)
            -> impl FnMut($crate::context::Context, Request__) -> Response<S> + Send + 'static + Clone {
                move |ctx, req| {
                    match req {
                        $(
                            Request__::$fn_name{ $($arg,)* } => {
                                let resp = Service::$fn_name(&mut service.clone(), ctx, $($arg),*);
                                Response::$fn_name(resp)
                            }
                        )*
                    }
                }
            }

        #[allow(unused)]
        #[derive(Clone, Debug)]
        /// The client stub that makes RPC calls to the server. Exposes a Future interface.
        pub struct Client($crate::client::Client<Request__, Response__>);

        /// Returns a new client stub that sends requests over the given transport.
        pub async fn new_stub<T>(config: $crate::client::Config, transport: T)
            -> ::std::io::Result<Client>
        where
            T: $crate::Transport<
                    Item = $crate::Response<Response__>,
                    SinkItem = $crate::ClientMessage<Request__>> + Send,
        {
            Ok(Client(await!($crate::client::Client::new(config, transport))?))
        }

        impl Client {
            $(
                #[allow(unused)]
                $(#[$attr])*
                pub fn $fn_name(&mut self, ctx: $crate::context::Context, $($arg: $in_),*)
                    -> impl ::std::future::Future<Output = ::std::io::Result<$out>> + '_ {
                    let request__ = Request__::$fn_name { $($arg,)* };
                    let resp = self.0.call(ctx, request__);
                    async move {
                        match await!(resp)? {
                            Response__::$fn_name(msg__) => ::std::result::Result::Ok(msg__),
                            _ => unreachable!(),
                        }
                    }
                }
            )*
        }
    }
}

// allow dead code; we're just testing that the macro expansion compiles
#[allow(dead_code)]
#[cfg(test)]
mod syntax_test {
    service! {
        #[deny(warnings)]
        #[allow(non_snake_case)]
        rpc TestCamelCaseDoesntConflict();
        rpc hello() -> String;
        #[doc="attr"]
        rpc attr(s: String) -> String;
        rpc no_args_no_return();
        rpc no_args() -> ();
        rpc one_arg(foo: String) -> i32;
        rpc two_args_no_return(bar: String, baz: u64);
        rpc two_args(bar: String, baz: u64) -> String;
        rpc no_args_ret_error() -> i32;
        rpc one_arg_ret_error(foo: String) -> String;
        rpc no_arg_implicit_return_error();
        #[doc="attr"]
        rpc one_arg_implicit_return_error(foo: String);
    }
}

#[cfg(test)]
mod functional_test {
    use futures::{
        compat::TokioDefaultSpawner,
        future::{ready, Ready},
        prelude::*,
    };
    use rpc::{
        client, context,
        server::{self, Handler},
        transport::channel,
    };
    use std::io;
    use tokio::runtime::current_thread;

    service! {
        rpc add(x: i32, y: i32) -> i32;
        rpc hey(name: String) -> String;
    }

    #[derive(Clone)]
    struct Server;

    impl Service for Server {
        type AddFut = Ready<i32>;

        fn add(&self, _: context::Context, x: i32, y: i32) -> Self::AddFut {
            ready(x + y)
        }

        type HeyFut = Ready<String>;

        fn hey(&self, _: context::Context, name: String) -> Self::HeyFut {
            ready(format!("Hey, {}.", name))
        }
    }

    #[test]
    fn sequential() {
        let _ = env_logger::try_init();
        rpc::init(TokioDefaultSpawner);

        let test = async {
            let (tx, rx) = channel::unbounded();
            tokio_executor::spawn(
                rpc::Server::new(server::Config::default())
                    .incoming(stream::once(ready(Ok(rx))))
                    .respond_with(serve(Server))
                    .unit_error()
                    .boxed()
                    .compat()
            );

            let mut client = await!(new_stub(client::Config::default(), tx))?;
            assert_eq!(3, await!(client.add(context::current(), 1, 2))?);
            assert_eq!(
                "Hey, Tim.",
                await!(client.hey(context::current(), "Tim".to_string()))?
            );
            Ok::<_, io::Error>(())
        }
            .map_err(|e| panic!(e.to_string()));

        current_thread::block_on_all(test.boxed().compat()).unwrap();
    }

    #[test]
    fn concurrent() {
        let _ = env_logger::try_init();
        rpc::init(TokioDefaultSpawner);

        let test = async {
            let (tx, rx) = channel::unbounded();
            tokio_executor::spawn(
                rpc::Server::new(server::Config::default())
                    .incoming(stream::once(ready(Ok(rx))))
                    .respond_with(serve(Server))
                    .unit_error()
                    .boxed()
                    .compat()
            );

            let client = await!(new_stub(client::Config::default(), tx))?;
            let mut c = client.clone();
            let req1 = c.add(context::current(), 1, 2);
            let mut c = client.clone();
            let req2 = c.add(context::current(), 3, 4);
            let mut c = client.clone();
            let req3 = c.hey(context::current(), "Tim".to_string());

            assert_eq!(3, await!(req1)?);
            assert_eq!(7, await!(req2)?);
            assert_eq!("Hey, Tim.", await!(req3)?);
            Ok::<_, io::Error>(())
        }
            .map_err(|e| panic!("test failed: {}", e));

        current_thread::block_on_all(test.boxed().compat()).unwrap();
    }
}
<commit_msg>Make Request and Resonse enums' docs public, because they show up in the serve fn.<commit_after>// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

#[cfg(feature = "serde")]
#[doc(hidden)]
#[macro_export]
macro_rules! add_serde_if_enabled {
    ($(#[$attr:meta])* -- $i:item) => {
        $(#[$attr])*
        #[derive($crate::serde::Serialize, $crate::serde::Deserialize)]
        $i
    }
}

#[cfg(not(feature = "serde"))]
#[doc(hidden)]
#[macro_export]
macro_rules! add_serde_if_enabled {
    ($(#[$attr:meta])* -- $i:item) => {
        $(#[$attr])*
        $i
    }
}

/// The main macro that creates RPC services.
///
/// Rpc methods are specified, mirroring trait syntax:
///
/// ```
/// # #![feature(await_macro, pin, arbitrary_self_types, async_await, futures_api, proc_macro_hygiene)]
/// # fn main() {}
/// # tarpc::service! {
/// /// Say hello
/// rpc hello(name: String) -> String;
/// # }
/// ```
///
/// Attributes can be attached to each rpc. These attributes
/// will then be attached to the generated service traits'
/// corresponding `fn`s, as well as to the client stubs' RPCs.
///
/// The following items are expanded in the enclosing module:
///
/// * `trait Service` -- defines the RPC service.
///   * `fn serve` -- turns a service impl into a request handler.
/// * `Client` -- a client stub with a fn for each RPC.
///   * `fn new_stub` -- creates a new Client stub.
///
#[macro_export]
macro_rules! service {
// Entry point
    (
        $(
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) $(-> $out:ty)*;
        )*
    ) => {
        $crate::service! {{
            $(
                $(#[$attr])*
                rpc $fn_name( $( $arg : $in_ ),* ) $(-> $out)*;
            )*
        }}
    };
// Pattern for when the next rpc has an implicit unit return type.
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* );

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        $crate::service! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> ();
        }
    };
// Pattern for when the next rpc has an explicit return type.
    (
        {
            $(#[$attr:meta])*
            rpc $fn_name:ident( $( $arg:ident : $in_:ty ),* ) -> $out:ty;

            $( $unexpanded:tt )*
        }
        $( $expanded:tt )*
    ) => {
        $crate::service! {
            { $( $unexpanded )* }

            $( $expanded )*

            $(#[$attr])*
            rpc $fn_name( $( $arg : $in_ ),* ) -> $out;
        }
    };
// Pattern for when all return types have been expanded
    (
        { } // none left to expand
        $(
            $(#[$attr:meta])*
            rpc $fn_name:ident ( $( $arg:ident : $in_:ty ),* ) -> $out:ty;
        )*
    ) => {
        $crate::add_serde_if_enabled! {
            /// The request sent over the wire from the client to the server.
            #[derive(Debug)]
            #[allow(non_camel_case_types, unused)]
            --
            pub enum Request {
                $(
                    $(#[$attr])*
                    $fn_name{ $($arg: $in_,)* }
                ),*
            }
        }

        $crate::add_serde_if_enabled! {
            /// The response sent over the wire from the server to the client.
            #[derive(Debug)]
            #[allow(non_camel_case_types, unused)]
            --
            pub enum Response {
                $(
                    $(#[$attr])*
                    $fn_name($out)
                ),*
            }
        }

        // TODO: proc_macro can't currently parse $crate, so this needs to be imported for the
        // usage of snake_to_camel! to work.
        use $crate::futures::Future as Future__;

        /// Defines the RPC service. The additional trait bounds are required so that services can
        /// multiplex requests across multiple tasks, potentially on multiple threads.
        pub trait Service: Clone + Send + 'static {
            $(
                $crate::snake_to_camel! {
                    /// The type of future returned by `{}`.
                    type $fn_name: Future__<Output = $out> + Send;
                }

                $(#[$attr])*
                fn $fn_name(&self, ctx: $crate::context::Context, $($arg:$in_),*) -> $crate::ty_snake_to_camel!(Self::$fn_name);
            )*
        }

        // TODO: use an existential type instead of this when existential types work.
        /// A future resolving to a server [`Response`].
        #[allow(non_camel_case_types)]
        pub enum ResponseFut<S: Service> {
            $(
                $(#[$attr])*
                $fn_name($crate::ty_snake_to_camel!(<S as Service>::$fn_name)),
            )*
        }

        impl<S: Service> ::std::fmt::Debug for ResponseFut<S> {
            fn fmt(&self, fmt: &mut ::std::fmt::Formatter) -> ::std::fmt::Result {
                fmt.debug_struct("Response").finish()
            }
        }

        impl<S: Service> ::std::future::Future for ResponseFut<S> {
            type Output = ::std::io::Result<Response>;

            fn poll(self: ::std::pin::Pin<&mut Self>, waker: &::std::task::LocalWaker)
                -> ::std::task::Poll<::std::io::Result<Response>>
            {
                unsafe {
                    match ::std::pin::Pin::get_mut_unchecked(self) {
                        $(
                            ResponseFut::$fn_name(resp) =>
                                ::std::pin::Pin::new_unchecked(resp)
                                    .poll(waker)
                                    .map(Response::$fn_name)
                                    .map(Ok),
                        )*
                    }
                }
            }
        }

        /// Returns a serving function to use with rpc::server::Server.
        pub fn serve<S: Service>(service: S)
            -> impl FnMut($crate::context::Context, Request) -> ResponseFut<S> + Send + 'static + Clone {
                move |ctx, req| {
                    match req {
                        $(
                            Request::$fn_name{ $($arg,)* } => {
                                let resp = Service::$fn_name(&mut service.clone(), ctx, $($arg),*);
                                ResponseFut::$fn_name(resp)
                            }
                        )*
                    }
                }
            }

        #[allow(unused)]
        #[derive(Clone, Debug)]
        /// The client stub that makes RPC calls to the server. Exposes a Future interface.
        pub struct Client($crate::client::Client<Request, Response>);

        /// Returns a new client stub that sends requests over the given transport.
        pub async fn new_stub<T>(config: $crate::client::Config, transport: T)
            -> ::std::io::Result<Client>
        where
            T: $crate::Transport<
                    Item = $crate::Response<Response>,
                    SinkItem = $crate::ClientMessage<Request>> + Send,
        {
            Ok(Client(await!($crate::client::Client::new(config, transport))?))
        }

        impl Client {
            $(
                #[allow(unused)]
                $(#[$attr])*
                pub fn $fn_name(&mut self, ctx: $crate::context::Context, $($arg: $in_),*)
                    -> impl ::std::future::Future<Output = ::std::io::Result<$out>> + '_ {
                    let request__ = Request::$fn_name { $($arg,)* };
                    let resp = self.0.call(ctx, request__);
                    async move {
                        match await!(resp)? {
                            Response::$fn_name(msg__) => ::std::result::Result::Ok(msg__),
                            _ => unreachable!(),
                        }
                    }
                }
            )*
        }
    }
}

// allow dead code; we're just testing that the macro expansion compiles
#[allow(dead_code)]
#[cfg(test)]
mod syntax_test {
    service! {
        #[deny(warnings)]
        #[allow(non_snake_case)]
        rpc TestCamelCaseDoesntConflict();
        rpc hello() -> String;
        #[doc="attr"]
        rpc attr(s: String) -> String;
        rpc no_args_no_return();
        rpc no_args() -> ();
        rpc one_arg(foo: String) -> i32;
        rpc two_args_no_return(bar: String, baz: u64);
        rpc two_args(bar: String, baz: u64) -> String;
        rpc no_args_ret_error() -> i32;
        rpc one_arg_ret_error(foo: String) -> String;
        rpc no_arg_implicit_return_error();
        #[doc="attr"]
        rpc one_arg_implicit_return_error(foo: String);
    }
}

#[cfg(test)]
mod functional_test {
    use futures::{
        compat::TokioDefaultSpawner,
        future::{ready, Ready},
        prelude::*,
    };
    use rpc::{
        client, context,
        server::{self, Handler},
        transport::channel,
    };
    use std::io;
    use tokio::runtime::current_thread;

    service! {
        rpc add(x: i32, y: i32) -> i32;
        rpc hey(name: String) -> String;
    }

    #[derive(Clone)]
    struct Server;

    impl Service for Server {
        type AddFut = Ready<i32>;

        fn add(&self, _: context::Context, x: i32, y: i32) -> Self::AddFut {
            ready(x + y)
        }

        type HeyFut = Ready<String>;

        fn hey(&self, _: context::Context, name: String) -> Self::HeyFut {
            ready(format!("Hey, {}.", name))
        }
    }

    #[test]
    fn sequential() {
        let _ = env_logger::try_init();
        rpc::init(TokioDefaultSpawner);

        let test = async {
            let (tx, rx) = channel::unbounded();
            tokio_executor::spawn(
                rpc::Server::new(server::Config::default())
                    .incoming(stream::once(ready(Ok(rx))))
                    .respond_with(serve(Server))
                    .unit_error()
                    .boxed()
                    .compat()
            );

            let mut client = await!(new_stub(client::Config::default(), tx))?;
            assert_eq!(3, await!(client.add(context::current(), 1, 2))?);
            assert_eq!(
                "Hey, Tim.",
                await!(client.hey(context::current(), "Tim".to_string()))?
            );
            Ok::<_, io::Error>(())
        }
            .map_err(|e| panic!(e.to_string()));

        current_thread::block_on_all(test.boxed().compat()).unwrap();
    }

    #[test]
    fn concurrent() {
        let _ = env_logger::try_init();
        rpc::init(TokioDefaultSpawner);

        let test = async {
            let (tx, rx) = channel::unbounded();
            tokio_executor::spawn(
                rpc::Server::new(server::Config::default())
                    .incoming(stream::once(ready(Ok(rx))))
                    .respond_with(serve(Server))
                    .unit_error()
                    .boxed()
                    .compat()
            );

            let client = await!(new_stub(client::Config::default(), tx))?;
            let mut c = client.clone();
            let req1 = c.add(context::current(), 1, 2);
            let mut c = client.clone();
            let req2 = c.add(context::current(), 3, 4);
            let mut c = client.clone();
            let req3 = c.hey(context::current(), "Tim".to_string());

            assert_eq!(3, await!(req1)?);
            assert_eq!(7, await!(req2)?);
            assert_eq!("Hey, Tim.", await!(req3)?);
            Ok::<_, io::Error>(())
        }
            .map_err(|e| panic!("test failed: {}", e));

        current_thread::block_on_all(test.boxed().compat()).unwrap();
    }
}
<|endoftext|>


<commit_before>use std::{thread, time};

use winit::{
    event::{Event, KeyboardInput, WindowEvent},
    event_loop::{ControlFlow, EventLoop},
    window::WindowBuilder,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Mode {
    Wait,
    WaitUntil,
    Poll,
}

const WAIT_TIME: time::Duration = time::Duration::from_millis(100);
const POLL_SLEEP_TIME: time::Duration = time::Duration::from_millis(100);

fn main() {
    simple_logger::init().unwrap();

    println!("Press '1' to switch to Wait mode.");
    println!("Press '2' to switch to WaitUntil mode.");
    println!("Press '3' to switch to Poll mode.");
    println!("Press 'R' to toggle request_redraw() calls.");
    println!("Press 'Esc' to close the window.");

    let event_loop = EventLoop::new();
    let window = WindowBuilder::new()
        .with_title("Press 1, 2, 3 to change control flow mode. Press R to toggle redraw requests.")
        .build(&event_loop)
        .unwrap();

    let mut mode = Mode::Wait;
    let mut request_redraw = false;
    let mut wait_cancelled = false;
    let mut close_requested = false;

    event_loop.run(move |event, _, control_flow| {
        use winit::event::{ElementState, StartCause, VirtualKeyCode};
        println!("{:?}", event);
        match event {
            Event::NewEvents(start_cause) => {
                wait_cancelled = mode == Mode::WaitUntil;
                match start_cause {
                    StartCause::ResumeTimeReached {
                        start: _,
                        requested_resume: _,
                    } => {
                        wait_cancelled = false;
                    }
                    _ => (),
                }
            }
            Event::WindowEvent { event, .. } => match event {
                WindowEvent::CloseRequested => {
                    close_requested = true;
                }
                WindowEvent::KeyboardInput {
                    input:
                        KeyboardInput {
                            virtual_keycode: Some(virtual_code),
                            state: ElementState::Pressed,
                            ..
                        },
                    ..
                } => match virtual_code {
                    VirtualKeyCode::Key1 => {
                        mode = Mode::Wait;
                        println!("\nmode: {:?}\n", mode);
                    }
                    VirtualKeyCode::Key2 => {
                        mode = Mode::WaitUntil;
                        println!("\nmode: {:?}\n", mode);
                    }
                    VirtualKeyCode::Key3 => {
                        mode = Mode::Poll;
                        println!("\nmode: {:?}\n", mode);
                    }
                    VirtualKeyCode::R => {
                        request_redraw = !request_redraw;
                        println!("\nrequest_redraw: {}\n", request_redraw);
                    }
                    VirtualKeyCode::Escape => {
                        close_requested = true;
                    }
                    _ => (),
                },
                _ => (),
            },
            Event::MainEventsCleared => {
                if request_redraw && !wait_cancelled && !close_requested {
                    window.request_redraw();
                }
                if close_requested {
                    *control_flow = ControlFlow::Exit;
                }
            }
            Event::RedrawRequested(_window_id) => {}
            Event::RedrawEventsCleared => {
                *control_flow = match mode {
                    Mode::Wait => ControlFlow::Wait,
                    Mode::WaitUntil => {
                        if wait_cancelled {
                            *control_flow
                        } else {
                            ControlFlow::WaitUntil(time::Instant::now() + WAIT_TIME)
                        }
                    }
                    Mode::Poll => {
                        thread::sleep(POLL_SLEEP_TIME);
                        ControlFlow::Poll
                    }
                };
            }
            _ => (),
        }
    });
}
<commit_msg>control_flow example: fix wait_cancelled logic again (#1511)<commit_after>use std::{thread, time};

use winit::{
    event::{Event, KeyboardInput, WindowEvent},
    event_loop::{ControlFlow, EventLoop},
    window::WindowBuilder,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Mode {
    Wait,
    WaitUntil,
    Poll,
}

const WAIT_TIME: time::Duration = time::Duration::from_millis(100);
const POLL_SLEEP_TIME: time::Duration = time::Duration::from_millis(100);

fn main() {
    simple_logger::init().unwrap();

    println!("Press '1' to switch to Wait mode.");
    println!("Press '2' to switch to WaitUntil mode.");
    println!("Press '3' to switch to Poll mode.");
    println!("Press 'R' to toggle request_redraw() calls.");
    println!("Press 'Esc' to close the window.");

    let event_loop = EventLoop::new();
    let window = WindowBuilder::new()
        .with_title("Press 1, 2, 3 to change control flow mode. Press R to toggle redraw requests.")
        .build(&event_loop)
        .unwrap();

    let mut mode = Mode::Wait;
    let mut request_redraw = false;
    let mut wait_cancelled = false;
    let mut close_requested = false;

    event_loop.run(move |event, _, control_flow| {
        use winit::event::{ElementState, StartCause, VirtualKeyCode};
        println!("{:?}", event);
        match event {
            Event::NewEvents(start_cause) => {
                wait_cancelled = match start_cause {
                    StartCause::WaitCancelled { .. } => mode == Mode::WaitUntil,
                    _ => false,
                }
            }
            Event::WindowEvent { event, .. } => match event {
                WindowEvent::CloseRequested => {
                    close_requested = true;
                }
                WindowEvent::KeyboardInput {
                    input:
                        KeyboardInput {
                            virtual_keycode: Some(virtual_code),
                            state: ElementState::Pressed,
                            ..
                        },
                    ..
                } => match virtual_code {
                    VirtualKeyCode::Key1 => {
                        mode = Mode::Wait;
                        println!("\nmode: {:?}\n", mode);
                    }
                    VirtualKeyCode::Key2 => {
                        mode = Mode::WaitUntil;
                        println!("\nmode: {:?}\n", mode);
                    }
                    VirtualKeyCode::Key3 => {
                        mode = Mode::Poll;
                        println!("\nmode: {:?}\n", mode);
                    }
                    VirtualKeyCode::R => {
                        request_redraw = !request_redraw;
                        println!("\nrequest_redraw: {}\n", request_redraw);
                    }
                    VirtualKeyCode::Escape => {
                        close_requested = true;
                    }
                    _ => (),
                },
                _ => (),
            },
            Event::MainEventsCleared => {
                if request_redraw && !wait_cancelled && !close_requested {
                    window.request_redraw();
                }
                if close_requested {
                    *control_flow = ControlFlow::Exit;
                }
            }
            Event::RedrawRequested(_window_id) => {}
            Event::RedrawEventsCleared => {
                *control_flow = match mode {
                    Mode::Wait => ControlFlow::Wait,
                    Mode::WaitUntil => {
                        if wait_cancelled {
                            *control_flow
                        } else {
                            ControlFlow::WaitUntil(time::Instant::now() + WAIT_TIME)
                        }
                    }
                    Mode::Poll => {
                        thread::sleep(POLL_SLEEP_TIME);
                        ControlFlow::Poll
                    }
                };
            }
            _ => (),
        }
    });
}
<|endoftext|>


<commit_before>// A port of the simplistic benchmark from
//
//    http://github.com/PaulKeeble/ScalaVErlangAgents
//
// I *think* it's the same, more or less.

#![crate_name = "msgsend"]

extern crate time;
extern crate zmq;

use std::os;
use std::thread::Thread;
use std::sync::mpsc::{channel, Sender, Receiver};

fn server(mut pull_socket: zmq::Socket, mut push_socket: zmq::Socket, mut workers: uint) {
    let mut count = 0u;
    let mut msg = zmq::Message::new().unwrap();

    while workers != 0 {
        match pull_socket.recv(&mut msg, 0) {
            Err(e) => panic!(e),
            Ok(()) => {
                let s = msg.as_str().unwrap();
                if s.is_empty() {
                    workers -= 1;
                } else {
                    count += s.parse().unwrap();
                }
            }
        }
    }

    match push_socket.send_str(count.to_string().as_slice(), 0) {
        Ok(()) => { }
        Err(e) => panic!(e),
    }
}

fn spawn_server(ctx: &mut zmq::Context, workers: uint) -> Sender<()> {
    let mut pull_socket = ctx.socket(zmq::PULL).unwrap();
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();

    pull_socket.bind("inproc://server-pull").unwrap();
    push_socket.bind("inproc://server-push").unwrap();

    // Spawn the server.
    let (ready_tx, ready_rx) = channel();
    let (start_tx, start_rx) = channel();

    Thread::spawn(move|| {
        // Let the main thread know we're ready.
        ready_tx.send(());

        // Wait until we need to start.
        start_rx.recv();

        server(pull_socket, push_socket, workers);
    });

    // Wait for the server to start.
    ready_rx.recv();

    start_tx
}

fn worker(mut push_socket: zmq::Socket, count: uint) {
    for _ in range(0, count) {
        push_socket.send_str(100u.to_string().as_slice(), 0).unwrap();
    }

    // Let the server know we're done.
    push_socket.send_str("", 0).unwrap();
}

fn spawn_worker(ctx: &mut zmq::Context, count: uint) -> Receiver<()> {
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();

    // Spawn the worker.
    let (tx, rx) = channel();
    Thread::spawn(move|| {
        // Let the main thread we're ready.
        tx.send(());

        worker(push_socket, count);

        tx.send(());
    });

    // Wait for the worker to start.
    rx.recv();

    rx
}

fn run(ctx: &mut zmq::Context, size: uint, workers: uint) {
    let start_ch = spawn_server(ctx, workers);

    // Create some command/control sockets.
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();
    let mut pull_socket = ctx.socket(zmq::PULL).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    pull_socket.connect("inproc://server-push").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();
    //pull_socket.connect("tcp://127.0.0.1:3457").unwrap();

    // Spawn all the workers.
    let mut worker_results = Vec::new();
    for _ in range(0, workers) {
        worker_results.push(spawn_worker(ctx, size / workers));
    }

    let start = time::precise_time_s();

    start_ch.send(());

    // Block until all the workers finish.
    for po in worker_results.iter() {
        po.recv();
    }

    // Receive the final count.
    let result: i32 = match pull_socket.recv_msg(0) {
        Ok(msg) => {
            let msg_str: &str = msg.as_str().unwrap();
            msg_str.parse().unwrap()
        },
        Err(e) => panic!(e),
    };

    let end = time::precise_time_s();
    let elapsed = end - start;

    println!("Count is {}", result);
    println!("Test took {} seconds", elapsed);
    let thruput = ((size / workers * workers) as f64) / elapsed;
    println!("Throughput={} per sec", thruput);
}

fn main() {
    let args = os::args();

    let args = if os::getenv("RUST_BENCH").is_some() {
        vec!("".to_string(), "1000000".to_string(), "10000".to_string())
    } else if args.len() <= 1u {
        vec!("".to_string(), "10000".to_string(), "4".to_string())
    } else {
        args
    };

    let size: uint = args[1].as_slice().parse().unwrap();
    let workers: uint = args[2].as_slice().parse().unwrap();

    let mut ctx = zmq::Context::new();

    run(&mut ctx, size, workers);
}
<commit_msg>Unwrap messages to discard the Result w/o unused warnings...<commit_after>// A port of the simplistic benchmark from
//
//    http://github.com/PaulKeeble/ScalaVErlangAgents
//
// I *think* it's the same, more or less.

#![crate_name = "msgsend"]

extern crate time;
extern crate zmq;

use std::os;
use std::thread::Thread;
use std::sync::mpsc::{channel, Sender, Receiver};

fn server(mut pull_socket: zmq::Socket, mut push_socket: zmq::Socket, mut workers: uint) {
    let mut count = 0u;
    let mut msg = zmq::Message::new().unwrap();

    while workers != 0 {
        match pull_socket.recv(&mut msg, 0) {
            Err(e) => panic!(e),
            Ok(()) => {
                let s = msg.as_str().unwrap();
                if s.is_empty() {
                    workers -= 1;
                } else {
                    count += s.parse().unwrap();
                }
            }
        }
    }

    match push_socket.send_str(count.to_string().as_slice(), 0) {
        Ok(()) => { }
        Err(e) => panic!(e),
    }
}

fn spawn_server(ctx: &mut zmq::Context, workers: uint) -> Sender<()> {
    let mut pull_socket = ctx.socket(zmq::PULL).unwrap();
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();

    pull_socket.bind("inproc://server-pull").unwrap();
    push_socket.bind("inproc://server-push").unwrap();

    // Spawn the server.
    let (ready_tx, ready_rx) = channel();
    let (start_tx, start_rx) = channel();

    Thread::spawn(move|| {
        // Let the main thread know we're ready.
        ready_tx.send(()).unwrap();

        // Wait until we need to start.
        start_rx.recv().unwrap();

        server(pull_socket, push_socket, workers);
    });

    // Wait for the server to start.
    ready_rx.recv().unwrap();

    start_tx
}

fn worker(mut push_socket: zmq::Socket, count: uint) {
    for _ in range(0, count) {
        push_socket.send_str(100u.to_string().as_slice(), 0).unwrap();
    }

    // Let the server know we're done.
    push_socket.send_str("", 0).unwrap();
}

fn spawn_worker(ctx: &mut zmq::Context, count: uint) -> Receiver<()> {
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();

    // Spawn the worker.
    let (tx, rx) = channel();
    Thread::spawn(move|| {
        // Let the main thread we're ready.
        tx.send(()).unwrap();

        worker(push_socket, count);

        tx.send(()).unwrap();
    });

    // Wait for the worker to start.
    rx.recv().unwrap();

    rx
}

fn run(ctx: &mut zmq::Context, size: uint, workers: uint) {
    let start_ch = spawn_server(ctx, workers);

    // Create some command/control sockets.
    let mut push_socket = ctx.socket(zmq::PUSH).unwrap();
    let mut pull_socket = ctx.socket(zmq::PULL).unwrap();

    push_socket.connect("inproc://server-pull").unwrap();
    pull_socket.connect("inproc://server-push").unwrap();
    //push_socket.connect("tcp://127.0.0.1:3456").unwrap();
    //pull_socket.connect("tcp://127.0.0.1:3457").unwrap();

    // Spawn all the workers.
    let mut worker_results = Vec::new();
    for _ in range(0, workers) {
        worker_results.push(spawn_worker(ctx, size / workers));
    }

    let start = time::precise_time_s();

    start_ch.send(()).unwrap();

    // Block until all the workers finish.
    for po in worker_results.iter() {
        po.recv().unwrap();
    }

    // Receive the final count.
    let result: i32 = match pull_socket.recv_msg(0) {
        Ok(msg) => {
            let msg_str: &str = msg.as_str().unwrap();
            msg_str.parse().unwrap()
        },
        Err(e) => panic!(e),
    };

    let end = time::precise_time_s();
    let elapsed = end - start;

    println!("Count is {}", result);
    println!("Test took {} seconds", elapsed);
    let thruput = ((size / workers * workers) as f64) / elapsed;
    println!("Throughput={} per sec", thruput);
}

fn main() {
    let args = os::args();

    let args = if os::getenv("RUST_BENCH").is_some() {
        vec!("".to_string(), "1000000".to_string(), "10000".to_string())
    } else if args.len() <= 1u {
        vec!("".to_string(), "10000".to_string(), "4".to_string())
    } else {
        args
    };

    let size: uint = args[1].as_slice().parse().unwrap();
    let workers: uint = args[2].as_slice().parse().unwrap();

    let mut ctx = zmq::Context::new();

    run(&mut ctx, size, workers);
}
<|endoftext|>


<commit_before>extern crate robin_core;

mod eval_expr_tests {
    use robin_core::ast::Expression;
    use robin_core::to_javascript::ToJavaScript;
    use robin_core::error::{Error, ErrorLevel, ErrorKind};

    #[test]
    fn if_with_no_args_should_return_an_err() {
        let expr = Expression::FuncCall("if".to_string(), vec![]);

        let err = Err(Error(ErrorLevel::Error,
                            ErrorKind::TooFewArguments,
                            "Too few Arguments applied for if"));

        assert_eq!(expr.eval(), err)
    }

    #[test]
    fn if_with_no_expression_after_condition_should_return_an_err() {
        let expr = Expression::FuncCall("if".to_string(),
                                        vec![Box::new(Expression::Boolean(true))]);

        let err = Err(Error(ErrorLevel::Error,
                            ErrorKind::TooFewArguments,
                            "No expression applied for condition"));

        assert_eq!(expr.eval(), err)
    }

    #[test]
    fn if_with_only_one_branch_should_return_an_func() {
        let expr = Expression::FuncCall("if".to_string(),
                                        vec![Box::new(Expression::Boolean(true)),
                                             Box::new(Expression::Number(1.0))]);

        assert_eq!(expr.eval(), Ok(String::from("(if (true) { 1 })")))
    }

    #[test]
    fn if_with_else_branch_should_return_a_func() {
        let expr = Expression::FuncCall("if".to_string(),
                                        vec![Box::new(Expression::Boolean(true)),
                                             Box::new(Expression::Number(1.0)),
                                             Box::new(Expression::Number(1.0))]);

        assert_eq!(expr.eval(),
                   Ok(String::from("(if (true) { 1 } else { 1 })")))
    }
}
<commit_msg>Fix If-statement tests.<commit_after>extern crate robin_core;

mod eval_expr_tests {
    use robin_core::ast::Expression;
    use robin_core::to_javascript::ToJavaScript;
    use robin_core::error::{Error, ErrorLevel, ErrorKind};

    #[test]
    fn if_with_no_args_should_return_an_err() {
        let expr = Expression::FuncCall("if".to_string(), vec![]);

        let err = Err(Error(ErrorLevel::Error,
                            ErrorKind::TooFewArguments,
                            "Too few Arguments applied for if"));

        assert_eq!(expr.eval(), err)
    }

    #[test]
    fn if_with_no_expression_after_condition_should_return_an_err() {
        let expr = Expression::FuncCall("if".to_string(),
                                        vec![Box::new(Expression::Boolean(true))]);

        let err = Err(Error(ErrorLevel::Error,
                            ErrorKind::TooFewArguments,
                            "No expression applied for condition"));

        assert_eq!(expr.eval(), err)
    }

    #[test]
    fn if_with_only_one_branch_should_return_an_func() {
        let expr = Expression::FuncCall("if".to_string(),
                                        vec![Box::new(Expression::Boolean(true)),
                                             Box::new(Expression::Number(1.0))]);

        assert_eq!(expr.eval(), Ok(String::from("if (true) { 1 }")))
    }

    #[test]
    fn if_with_else_branch_should_return_a_func() {
        let expr = Expression::FuncCall("if".to_string(),
                                        vec![Box::new(Expression::Boolean(true)),
                                             Box::new(Expression::Number(1.0)),
                                             Box::new(Expression::Number(1.0))]);

        assert_eq!(expr.eval(), Ok(String::from("if (true) { 1 } else { 1 }")))
    }
}
<|endoftext|>


<commit_before>use traits::Grid;
use traits::Engine;
use traits::ReprConsumer;

pub struct Sequential<G: Grid, C: ReprConsumer> {
    grid: G,
    consumer: C
}

impl<G: Grid, C: ReprConsumer> Sequential<G, C> {

    pub fn new(grid: G, consumer: C) -> Self {
        Sequential {
            grid: grid,
            consumer: consumer
        }
    }
}

impl<G: Grid, C: ReprConsumer> Engine for Sequential<G, C> {
    fn run_times(&mut self, times: i64) {
        for _ in 0..times {
            self.grid.step();
            let repr = self.grid.repr();
            self.consumer.consume(repr);
        }
    }
}<commit_msg>Less-verbose consumer-grid interface.<commit_after>use traits::Grid;
use traits::Engine;
use traits::ReprConsumer;

pub struct Sequential<G: Grid, C: ReprConsumer> {
    grid: G,
    consumer: C
}

impl<G: Grid, C: ReprConsumer> Sequential<G, C> {

    pub fn new(grid: G, consumer: C) -> Self {
        Sequential {
            grid: grid,
            consumer: consumer
        }
    }
}

impl<G: Grid, C: ReprConsumer> Engine for Sequential<G, C> {
    fn run_times(&mut self, times: i64) {
        for _ in 0..times {
            self.grid.step();
            self.consumer.consume(&self.grid);
        }
    }
}<|endoftext|>


<commit_before>use std::io::{self, Read, Seek, SeekFrom};
use std::fs::{File, Permissions};
use std::fs::{create_dir_all, set_permissions, hard_link};
use std::os::unix::fs::{PermissionsExt, symlink};
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};

use tar::Archive;
use flate2::FlateReadExt;
use xz2::read::XzDecoder;
use bzip2::read::BzDecoder;
use libmount::BindMount;

use quire::validate as V;
use container::mount::{unmount};
use builder::context::Context;
use builder::download::{maybe_download_and_check_hashsum};
use builder::commands::generic::run_command_at;
use file_util::{read_visible_entries, create_dir, create_dir_mode,
    copy_stream, set_owner_group};
use build_step::{BuildStep, VersionError, StepError, Digest, Config, Guard};

#[derive(RustcDecodable, Debug)]
pub struct Tar {
    pub url: String,
    pub sha256: Option<String>,
    pub path: PathBuf,
    pub subdir: PathBuf,
}

impl Tar {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("url", V::Scalar::new())
        .member("sha256", V::Scalar::new().optional())
        .member("path", V::Directory::new().is_absolute(true).default("/"))
        .member("subdir", V::Directory::new().default("").is_absolute(false))
    }
}

#[derive(RustcDecodable, Debug)]
pub struct TarInstall {
    pub url: String,
    pub sha256: Option<String>,
    pub subdir: Option<PathBuf>,
    pub script: String,
}

impl TarInstall {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("url", V::Scalar::new())
        .member("sha256", V::Scalar::new().optional())
        .member("subdir", V::Directory::new().optional().is_absolute(false))
        .member("script", V::Scalar::new()
                .default("./configure --prefix=/usr\n\
                          make\n\
                          make install\n"))
    }
}

pub fn unpack_file(_ctx: &mut Context, src: &Path, tgt: &Path,
    includes: &[&Path], excludes: &[&Path], preserve_owner: bool)
    -> Result<(), String>
{

    info!("Unpacking {:?} -> {:?}", src, tgt);
    let read_err = |e| format!("Error reading {:?}: {}", src, e);

    let mut file = try!(File::open(&src).map_err(&read_err));

    let mut buf = [0u8; 8];
    let nbytes = try!(file.read(&mut buf).map_err(&read_err));
    try!(file.seek(SeekFrom::Start(0)).map_err(&read_err));
    let magic = &buf[..nbytes];
    if magic.len() >= 2 && magic[..2] == [0x1f, 0x8b] {
        return unpack_stream(
            try!(file.gz_decode().map_err(&read_err)),
            src, tgt, includes, excludes, preserve_owner);
    } else if magic.len() >= 6 && magic[..6] ==
        [ 0xFD, b'7', b'z', b'X', b'Z', 0x00]
    {
        return unpack_stream(XzDecoder::new(file),
            src, tgt, includes, excludes, preserve_owner);
    } else if magic.len() >= 3 && magic[..3] == [ b'B', b'Z', b'h'] {
        return unpack_stream(BzDecoder::new(file),
            src, tgt, includes, excludes, preserve_owner);
    } else {
        return Err(format!("unpacking {:?}: unexpected compression", src));
    }

}

fn unpack_stream<F: Read>(file: F, srcpath: &Path, tgt: &Path,
    includes: &[&Path], excludes: &[&Path], preserve_owner: bool)
    -> Result<(), String>
{
    let read_err = |e| format!("Error reading {:?}: {}", srcpath, e);
    let mut arc = Archive::new(file);
    let mut hardlinks = Vec::new();

    for item in try!(arc.entries().map_err(&read_err)) {
        let mut src = try!(item.map_err(&read_err));
        let path_ref = try!(src.path().map_err(&read_err))
            .to_path_buf();
        let mut orig_path: &Path = &path_ref;
        if orig_path.is_absolute() {
            orig_path = orig_path.strip_prefix("/").unwrap();
        }
        if includes.len() > 0 {
            if !includes.iter().any(|x| orig_path.starts_with(x)) {
                continue;
            }
        }
        if excludes.iter().any(|x| orig_path.starts_with(x)) {
            continue;
        }
        let path = tgt.join(orig_path);
        let write_err = |e| format!("Error writing {:?}: {}", path, e);
        let entry = src.header().entry_type();
        let uid = try!(src.header().uid().map_err(&read_err));
        let gid = try!(src.header().gid().map_err(&read_err));
        if entry.is_dir() {
            let mode = try!(src.header().mode().map_err(&read_err));
            try!(create_dir_mode(&path, mode).map_err(&write_err));
            if preserve_owner {
                try!(set_owner_group(&path, uid, gid).map_err(&write_err));
            }
        } else if entry.is_symlink() {
            let src = try!(try!(src.header().link_name().map_err(&read_err))
                .ok_or(format!("Error unpacking {:?}, broken symlink", path)));
            match symlink(&src, &path) {
                Ok(_) => {},
                Err(e) => {
                    if e.kind() == io::ErrorKind::NotFound {
                        if let Some(parent) = path.parent() {
                            try!(create_dir(parent, true).map_err(&write_err));
                            try!(symlink(&src, &path).map_err(&write_err))
                        } else {
                            return Err(write_err(e));
                        }
                    } else {
                        return Err(write_err(e));
                    }
                }
            };
        } else if entry.is_hard_link() {
            let link = try!(try!(src.link_name().map_err(&read_err))
                .ok_or(format!("Error unpacking {:?}, broken symlink", path)));
            let link = if link.is_absolute() {
                link.strip_prefix("/").unwrap()
            } else {
                &*link
            };
            hardlinks.push((tgt.join(link).to_path_buf(), path.to_path_buf()));
        } else if entry.is_pax_global_extensions() ||
                  entry.is_pax_local_extensions() ||
                  entry.is_gnu_longname() ||
                  entry.is_gnu_longlink()
        {
            // nothing to do
        } else {
            let mut dest = match File::create(&path) {
                Ok(x) => x,
                Err(e) => {
                    if e.kind() == io::ErrorKind::NotFound {
                        if let Some(parent) = path.parent() {
                            try!(create_dir(parent, true).map_err(&write_err));
                            try!(File::create(&path).map_err(&write_err))
                        } else {
                            return Err(write_err(e));
                        }
                    } else {
                        return Err(write_err(e));
                    }
                }
            };
            try!(copy_stream(&mut src, &mut dest).map_err(|e|
                format!("Error unpacking {:?} -> {:?}: {}",
                        srcpath, path, e)));
            let mode = try!(src.header().mode().map_err(&read_err));
            try!(set_permissions(&path, Permissions::from_mode(mode))
                .map_err(&write_err));
            if preserve_owner {
                try!(set_owner_group(&path, uid, gid).map_err(&write_err));
            }
        }
    }
    for (src, dst) in hardlinks.into_iter() {
        let write_err = |e| {
            format!("Error hardlinking {:?} - {:?}: {}", &src, &dst, e)
        };
        match hard_link(&src, &dst) {
            Ok(_) => {},
            Err(e) => {
                if e.kind() == io::ErrorKind::NotFound {
                    if let Some(parent) = dst.parent() {
                        try!(create_dir(parent, true).map_err(&write_err));
                        try!(hard_link(&src, &dst).map_err(&write_err))
                    } else {
                        return Err(write_err(e));
                    }
                } else {
                    return Err(write_err(e));
                }
            }
        };
    }
    Ok(())
}

pub fn tar_command(ctx: &mut Context, tar: &Tar) -> Result<(), String>
{
    let fpath = PathBuf::from("/vagga/root")
        .join(tar.path.strip_prefix("/").unwrap());
    let (filename, _) = try!(maybe_download_and_check_hashsum(
        ctx, &tar.url, tar.sha256.clone()));

    if &Path::new(&tar.subdir) == &Path::new(".") {
        try!(unpack_file(ctx, &filename, &fpath, &[], &[], false));
    } else {
        let tmppath = PathBuf::from("/vagga/root/tmp")
            .join(filename.file_name().unwrap());
        let tmpsub = tmppath.join(&tar.subdir);
        try_msg!(create_dir(&tmpsub, true), "Error making dir: {err}");
        if !fpath.exists() {
            try_msg!(create_dir(&fpath, true), "Error making dir: {err}");
        }
        try_msg!(BindMount::new(&fpath, &tmpsub).mount(),
            "temporary tar mount: {err}");
        let res = if tar.subdir.as_path() == Path::new("") {
            unpack_file(ctx, &filename, &tmppath, &[], &[], false)
        } else {
            unpack_file(ctx, &filename, &tmppath,
                &[&tar.subdir.clone()], &[], false)
        };
        try!(unmount(&tmpsub));
        try!(res);
    }
    Ok(())
}

pub fn tar_install(ctx: &mut Context, tar: &TarInstall)
    -> Result<(), String>
{
    let (filename, _) = try!(maybe_download_and_check_hashsum(
        ctx, &tar.url, tar.sha256.clone()));

    let tmppath = PathBuf::from("/vagga/root/tmp")
        .join(filename.file_name().unwrap());
    try!(create_dir_all(&tmppath)
         .map_err(|e| format!("Error making dir: {}", e)));
    try!(set_permissions(&tmppath, Permissions::from_mode(0o755))
         .map_err(|e| format!("Error setting permissions: {}", e)));
    try!(unpack_file(ctx, &filename, &tmppath, &[], &[], false));
    let workdir = if let Some(ref subpath) = tar.subdir {
        tmppath.join(subpath)
    } else {
        let items = try!(read_visible_entries(&tmppath)
            .map_err(|e| format!("Error reading dir: {}", e)));
        if items.len() != 1 {
            if items.len() == 0 {
                return Err("Tar archive was empty".to_string());
            } else {
                return Err("Multiple directories was unpacked. \
                    If thats expected use `subdir: \".\"` or any \
                    other directory".to_string());
            }
        }
        items.into_iter().next().unwrap()
    };
    let workdir = PathBuf::from("/").join(
        workdir.strip_prefix("/vagga/root").unwrap());
    return run_command_at(ctx, &[
        "/bin/sh".to_string(),
        "-exc".to_string(),
        tar.script.to_string()],
        &workdir);
}

impl BuildStep for Tar {
    fn hash(&self, _cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        if let Some(ref sha) = self.sha256 {
            hash.field("hash", sha);
        } else {
            hash.field("url", &self.url);
        }
        hash.field("path", self.path.as_os_str().as_bytes());
        hash.field("subdir", self.subdir.as_os_str().as_bytes());
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        if build {
            try!(tar_command(&mut guard.ctx, self));
        }
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        None
    }
}

impl BuildStep for TarInstall {
    fn hash(&self, _cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        if let Some(ref sha) = self.sha256 {
            hash.field("hash", sha);
        } else {
            hash.field("url", &self.url);
        }
        hash.opt_field("subdir",
            &self.subdir.as_ref().map(|x| x.as_os_str().as_bytes()));
        hash.field("script", &self.script);
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        if build {
            try!(tar_install(&mut guard.ctx, self));
        }
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        None
    }
}
<commit_msg>Fix unpacking archives having no uids or gids<commit_after>use std::io::{self, Read, Seek, SeekFrom};
use std::fs::{File, Permissions};
use std::fs::{create_dir_all, set_permissions, hard_link};
use std::os::unix::fs::{PermissionsExt, symlink};
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};

use tar::Archive;
use flate2::FlateReadExt;
use xz2::read::XzDecoder;
use bzip2::read::BzDecoder;
use libmount::BindMount;

use quire::validate as V;
use container::mount::{unmount};
use builder::context::Context;
use builder::download::{maybe_download_and_check_hashsum};
use builder::commands::generic::run_command_at;
use file_util::{read_visible_entries, create_dir, create_dir_mode,
    copy_stream, set_owner_group};
use build_step::{BuildStep, VersionError, StepError, Digest, Config, Guard};

#[derive(RustcDecodable, Debug)]
pub struct Tar {
    pub url: String,
    pub sha256: Option<String>,
    pub path: PathBuf,
    pub subdir: PathBuf,
}

impl Tar {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("url", V::Scalar::new())
        .member("sha256", V::Scalar::new().optional())
        .member("path", V::Directory::new().is_absolute(true).default("/"))
        .member("subdir", V::Directory::new().default("").is_absolute(false))
    }
}

#[derive(RustcDecodable, Debug)]
pub struct TarInstall {
    pub url: String,
    pub sha256: Option<String>,
    pub subdir: Option<PathBuf>,
    pub script: String,
}

impl TarInstall {
    pub fn config() -> V::Structure<'static> {
        V::Structure::new()
        .member("url", V::Scalar::new())
        .member("sha256", V::Scalar::new().optional())
        .member("subdir", V::Directory::new().optional().is_absolute(false))
        .member("script", V::Scalar::new()
                .default("./configure --prefix=/usr\n\
                          make\n\
                          make install\n"))
    }
}

pub fn unpack_file(_ctx: &mut Context, src: &Path, tgt: &Path,
    includes: &[&Path], excludes: &[&Path], preserve_owner: bool)
    -> Result<(), String>
{

    info!("Unpacking {:?} -> {:?}", src, tgt);
    let read_err = |e| format!("Error reading {:?}: {}", src, e);

    let mut file = try!(File::open(&src).map_err(&read_err));

    let mut buf = [0u8; 8];
    let nbytes = try!(file.read(&mut buf).map_err(&read_err));
    try!(file.seek(SeekFrom::Start(0)).map_err(&read_err));
    let magic = &buf[..nbytes];
    if magic.len() >= 2 && magic[..2] == [0x1f, 0x8b] {
        return unpack_stream(
            try!(file.gz_decode().map_err(&read_err)),
            src, tgt, includes, excludes, preserve_owner);
    } else if magic.len() >= 6 && magic[..6] ==
        [ 0xFD, b'7', b'z', b'X', b'Z', 0x00]
    {
        return unpack_stream(XzDecoder::new(file),
            src, tgt, includes, excludes, preserve_owner);
    } else if magic.len() >= 3 && magic[..3] == [ b'B', b'Z', b'h'] {
        return unpack_stream(BzDecoder::new(file),
            src, tgt, includes, excludes, preserve_owner);
    } else {
        return Err(format!("unpacking {:?}: unexpected compression", src));
    }

}

fn unpack_stream<F: Read>(file: F, srcpath: &Path, tgt: &Path,
    includes: &[&Path], excludes: &[&Path], preserve_owner: bool)
    -> Result<(), String>
{
    let read_err = |e| format!("Error reading {:?}: {}", srcpath, e);
    let mut arc = Archive::new(file);
    let mut hardlinks = Vec::new();

    for item in try!(arc.entries().map_err(&read_err)) {
        let mut src = try!(item.map_err(&read_err));
        let path_ref = try!(src.path().map_err(&read_err))
            .to_path_buf();
        let mut orig_path: &Path = &path_ref;
        if orig_path.is_absolute() {
            orig_path = orig_path.strip_prefix("/").unwrap();
        }
        if includes.len() > 0 {
            if !includes.iter().any(|x| orig_path.starts_with(x)) {
                continue;
            }
        }
        if excludes.iter().any(|x| orig_path.starts_with(x)) {
            continue;
        }
        let path = tgt.join(orig_path);
        let write_err = |e| format!("Error writing {:?}: {}", path, e);
        let entry = src.header().entry_type();

        // Some archives don't have uids
        // TODO(tailhook) should this be handled in tar-rs?
        let uid = src.header().uid().unwrap_or(0);
        let gid = src.header().gid().unwrap_or(0);

        if entry.is_dir() {
            let mode = try!(src.header().mode().map_err(&read_err));
            try!(create_dir_mode(&path, mode).map_err(&write_err));
            if preserve_owner {
                try!(set_owner_group(&path, uid, gid).map_err(&write_err));
            }
        } else if entry.is_symlink() {
            let src = try!(try!(src.header().link_name().map_err(&read_err))
                .ok_or(format!("Error unpacking {:?}, broken symlink", path)));
            match symlink(&src, &path) {
                Ok(_) => {},
                Err(e) => {
                    if e.kind() == io::ErrorKind::NotFound {
                        if let Some(parent) = path.parent() {
                            try!(create_dir(parent, true).map_err(&write_err));
                            try!(symlink(&src, &path).map_err(&write_err))
                        } else {
                            return Err(write_err(e));
                        }
                    } else {
                        return Err(write_err(e));
                    }
                }
            };
        } else if entry.is_hard_link() {
            let link = try!(try!(src.link_name().map_err(&read_err))
                .ok_or(format!("Error unpacking {:?}, broken symlink", path)));
            let link = if link.is_absolute() {
                link.strip_prefix("/").unwrap()
            } else {
                &*link
            };
            hardlinks.push((tgt.join(link).to_path_buf(), path.to_path_buf()));
        } else if entry.is_pax_global_extensions() ||
                  entry.is_pax_local_extensions() ||
                  entry.is_gnu_longname() ||
                  entry.is_gnu_longlink()
        {
            // nothing to do
        } else {
            let mut dest = match File::create(&path) {
                Ok(x) => x,
                Err(e) => {
                    if e.kind() == io::ErrorKind::NotFound {
                        if let Some(parent) = path.parent() {
                            try!(create_dir(parent, true).map_err(&write_err));
                            try!(File::create(&path).map_err(&write_err))
                        } else {
                            return Err(write_err(e));
                        }
                    } else {
                        return Err(write_err(e));
                    }
                }
            };
            try!(copy_stream(&mut src, &mut dest).map_err(|e|
                format!("Error unpacking {:?} -> {:?}: {}",
                        srcpath, path, e)));
            let mode = try!(src.header().mode().map_err(&read_err));
            try!(set_permissions(&path, Permissions::from_mode(mode))
                .map_err(&write_err));
            if preserve_owner {
                try!(set_owner_group(&path, uid, gid).map_err(&write_err));
            }
        }
    }
    for (src, dst) in hardlinks.into_iter() {
        let write_err = |e| {
            format!("Error hardlinking {:?} - {:?}: {}", &src, &dst, e)
        };
        match hard_link(&src, &dst) {
            Ok(_) => {},
            Err(e) => {
                if e.kind() == io::ErrorKind::NotFound {
                    if let Some(parent) = dst.parent() {
                        try!(create_dir(parent, true).map_err(&write_err));
                        try!(hard_link(&src, &dst).map_err(&write_err))
                    } else {
                        return Err(write_err(e));
                    }
                } else {
                    return Err(write_err(e));
                }
            }
        };
    }
    Ok(())
}

pub fn tar_command(ctx: &mut Context, tar: &Tar) -> Result<(), String>
{
    let fpath = PathBuf::from("/vagga/root")
        .join(tar.path.strip_prefix("/").unwrap());
    let (filename, _) = try!(maybe_download_and_check_hashsum(
        ctx, &tar.url, tar.sha256.clone()));

    if &Path::new(&tar.subdir) == &Path::new(".") {
        try!(unpack_file(ctx, &filename, &fpath, &[], &[], false));
    } else {
        let tmppath = PathBuf::from("/vagga/root/tmp")
            .join(filename.file_name().unwrap());
        let tmpsub = tmppath.join(&tar.subdir);
        try_msg!(create_dir(&tmpsub, true), "Error making dir: {err}");
        if !fpath.exists() {
            try_msg!(create_dir(&fpath, true), "Error making dir: {err}");
        }
        try_msg!(BindMount::new(&fpath, &tmpsub).mount(),
            "temporary tar mount: {err}");
        let res = if tar.subdir.as_path() == Path::new("") {
            unpack_file(ctx, &filename, &tmppath, &[], &[], false)
        } else {
            unpack_file(ctx, &filename, &tmppath,
                &[&tar.subdir.clone()], &[], false)
        };
        try!(unmount(&tmpsub));
        try!(res);
    }
    Ok(())
}

pub fn tar_install(ctx: &mut Context, tar: &TarInstall)
    -> Result<(), String>
{
    let (filename, _) = try!(maybe_download_and_check_hashsum(
        ctx, &tar.url, tar.sha256.clone()));

    let tmppath = PathBuf::from("/vagga/root/tmp")
        .join(filename.file_name().unwrap());
    try!(create_dir_all(&tmppath)
         .map_err(|e| format!("Error making dir: {}", e)));
    try!(set_permissions(&tmppath, Permissions::from_mode(0o755))
         .map_err(|e| format!("Error setting permissions: {}", e)));
    try!(unpack_file(ctx, &filename, &tmppath, &[], &[], false));
    let workdir = if let Some(ref subpath) = tar.subdir {
        tmppath.join(subpath)
    } else {
        let items = try!(read_visible_entries(&tmppath)
            .map_err(|e| format!("Error reading dir: {}", e)));
        if items.len() != 1 {
            if items.len() == 0 {
                return Err("Tar archive was empty".to_string());
            } else {
                return Err("Multiple directories was unpacked. \
                    If thats expected use `subdir: \".\"` or any \
                    other directory".to_string());
            }
        }
        items.into_iter().next().unwrap()
    };
    let workdir = PathBuf::from("/").join(
        workdir.strip_prefix("/vagga/root").unwrap());
    return run_command_at(ctx, &[
        "/bin/sh".to_string(),
        "-exc".to_string(),
        tar.script.to_string()],
        &workdir);
}

impl BuildStep for Tar {
    fn hash(&self, _cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        if let Some(ref sha) = self.sha256 {
            hash.field("hash", sha);
        } else {
            hash.field("url", &self.url);
        }
        hash.field("path", self.path.as_os_str().as_bytes());
        hash.field("subdir", self.subdir.as_os_str().as_bytes());
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        if build {
            try!(tar_command(&mut guard.ctx, self));
        }
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        None
    }
}

impl BuildStep for TarInstall {
    fn hash(&self, _cfg: &Config, hash: &mut Digest)
        -> Result<(), VersionError>
    {
        if let Some(ref sha) = self.sha256 {
            hash.field("hash", sha);
        } else {
            hash.field("url", &self.url);
        }
        hash.opt_field("subdir",
            &self.subdir.as_ref().map(|x| x.as_os_str().as_bytes()));
        hash.field("script", &self.script);
        Ok(())
    }
    fn build(&self, guard: &mut Guard, build: bool)
        -> Result<(), StepError>
    {
        if build {
            try!(tar_install(&mut guard.ctx, self));
        }
        Ok(())
    }
    fn is_dependent_on(&self) -> Option<&str> {
        None
    }
}
<|endoftext|>