<commit_before>/*
 * Steps of XWayland server creation
 *
 * Sockets to create:
 * - a pair for XWayland to connect to smithay as a wayland client, we use our
 *   end to insert the XWayland client in the display
 * - a pair for smithay to connect to XWayland as a WM, we give our end to the
 *   WM and it deals with it
 * - 2 listening sockets on which the XWayland server will listen. We need to
 *   bind them ourselves so we know what value put in the $DISPLAY env variable.
 *   This involves some dance with a lockfile to ensure there is no collision with
 *   an other starting xserver
 *   if we listen on display $D, their paths are respectively:
 *   - /tmp/.X11-unix/X$D
 *   - @/tmp/.X11-unix/X$D (abstract socket)
 *
 * The XWayland server is spawned via an intermediate shell
 * -> wlroots does a double-fork while weston a single one, why ??
 *    -> https://stackoverflow.com/questions/881388/
 * -> once it is started, it will check if SIGUSR1 is set to ignored. If so,
 *    if will consider its parent as "smart", and send a SIGUSR1 signal when
 *    startup completes. We want to catch this so we can launch the VM.
 * -> we need to track if the XWayland crashes, to restart it
 *
 * cf https://github.com/swaywm/wlroots/blob/master/xwayland/xwayland.c
 *
 * Setting SIGUSR1 handler is complicated in multithreaded program, because
 * Xwayland will send SIGUSR1 to the process, and if a thread cannot handle
 * SIGUSR1, that thread will be killed.
 *
 * Double-fork can tackle this issue, but this is also very complex in a
 * a multithread program, after forking only signal-safe functions can be used.
 * The only workaround is to fork early before any other thread starts, but
 * doing so will expose an unsafe interface.
 *
 * We use an intermediate shell to translate the signal to simple fd IO.
 * We ask sh to setup SIGUSR1 handler, and in a subshell mute SIGUSR1 and exec
 * Xwayland. When the SIGUSR1 is received, it can communicate to us via redirected
 * STDOUT.
 */
use std::{
    cell::RefCell,
    env,
    io::{self, Read},
    os::unix::{
        io::{AsRawFd, RawFd},
        net::UnixStream,
        process::CommandExt,
    },
    process::{ChildStdout, Command, Stdio},
    rc::Rc,
};

use calloop::{
    channel::{self, sync_channel, Channel, SyncSender},
    generic::Generic,
    Interest, LoopHandle, Mode,
};

use slog::{error, info, o};

use super::x11_sockets::{prepare_x11_sockets, X11Lock};

/// The XWayland handle
#[derive(Debug)]
pub struct XWayland {
    inner: Rc<RefCell<Inner>>,
}

/// Events generated by the XWayland manager
///
/// This is a very low-level interface, only notifying you when the connection
/// with XWayland is up, or when it terminates.
///
/// Your WM code must be able to handle the XWayland server connecting then
/// disconnecting several time in a row, but only a single connection will
/// be active at any given time.
#[derive(Debug)]
#[must_use = "Connection events must be handled to prevent fd leaking"]
pub enum XWaylandEvent {
    /// The XWayland server is ready
    Ready {
        /// Privileged X11 connection to XWayland
        connection: UnixStream,

        /// Wayland client representing XWayland
        ///
        /// This should be used to create a Wayland client using [`Display::insert_client`](wayland_server::Display::insert_client).
        ///
        /// When the client is disconnected, you should [`shutdown`](XWayland::shutdown) XWayland.
        client: UnixStream,

        /// The display number the XWayland server is available at.
        ///
        /// This may be used if you wish to set the `DISPLAY` variable manually when spawning processes that
        /// may use XWayland.
        display: u32,
    },

    /// The XWayland server exited
    ///
    /// This event is sent when the [`XWayland`] handle is dropped.
    Exited,
}

impl XWayland {
    /// Create a new XWayland manager
    ///
    /// This function returns both the [`XWayland`] handle and an [`XWaylandSource`] that needs to be inserted
    /// into the [`calloop`] event loop, producing the Xwayland startup and shutdown events.
    pub fn new<L>(logger: L) -> (XWayland, XWaylandSource)
    where
        L: Into<Option<::slog::Logger>>,
    {
        let log = crate::slog_or_fallback(logger);
        // We don't expect to ever have more than 2 messages in flight, if XWayland got ready and then died right away
        let (sender, channel) = sync_channel(2);
        let inner = Rc::new(RefCell::new(Inner {
            instance: None,
            sender,
            log: log.new(o!("smithay_module" => "XWayland")),
        }));
        (XWayland { inner }, XWaylandSource { channel })
    }

    /// Attempt to start the XWayland instance
    ///
    /// If it succeeds, you'll eventually receive an `XWaylandEvent::Ready`
    /// through the source provided by `XWayland::new()` containing an
    /// `UnixStream` representing your WM connection to XWayland, and the
    /// wayland `Client` for XWayland.
    ///
    /// Does nothing if XWayland is already started or starting.
    pub fn start<D>(&self, loop_handle: LoopHandle<'_, D>) -> io::Result<()> {
        launch(&self.inner, loop_handle)
    }

    /// Shutdown XWayland
    ///
    /// Does nothing if it was not already running, otherwise kills it and you will
    /// later receive a `XWaylandEvent::Exited` event.
    pub fn shutdown(&self) {
        self.inner.borrow_mut().shutdown();
    }
}

impl Drop for XWayland {
    fn drop(&mut self) {
        self.inner.borrow_mut().shutdown();
    }
}

#[derive(Debug)]
struct XWaylandInstance {
    display_lock: X11Lock,
    wayland_client_fd: Option<UnixStream>,
    wm_fd: Option<UnixStream>,
    child_stdout: Option<ChildStdout>,
}

// Inner implementation of the XWayland manager
#[derive(Debug)]
struct Inner {
    sender: SyncSender<XWaylandEvent>,
    instance: Option<XWaylandInstance>,
    log: ::slog::Logger,
}

// Launch an XWayland server
//
// Does nothing if there is already a launched instance
fn launch<D>(inner: &Rc<RefCell<Inner>>, loop_handle: LoopHandle<'_, D>) -> io::Result<()> {
    let mut guard = inner.borrow_mut();
    if guard.instance.is_some() {
        return Ok(());
    }

    info!(guard.log, "Starting XWayland");

    let (x_wm_x11, x_wm_me) = UnixStream::pair()?;
    let (wl_x11, wl_me) = UnixStream::pair()?;

    let (lock, x_fds) = prepare_x11_sockets(guard.log.clone())?;

    // we have now created all the required sockets

    // Setup the associated wayland client to be created in an idle callback, so that we don't need
    // to access the dispatch_data *right now*
    let idle_inner = inner.clone();
    loop_handle.insert_idle(move |_| {
        let mut guard = idle_inner.borrow_mut();
        let guard = &mut *guard;

        if let Some(ref mut instance) = guard.instance {
            instance.wayland_client_fd = Some(wl_me);
        }
    });

    // all is ready, we can do the fork dance
    let child_stdout = match spawn_xwayland(lock.display(), wl_x11, x_wm_x11, &x_fds) {
        Ok(child_stdout) => child_stdout,
        Err(e) => {
            error!(guard.log, "XWayland failed to spawn"; "err" => format!("{:?}", e));
            return Err(e);
        }
    };

    let inner = inner.clone();
    loop_handle
        .insert_source(
            Generic::<_, io::Error>::new(child_stdout.as_raw_fd(), Interest::READ, Mode::Level),
            move |_, _, _| {
                // the closure must be called exactly one time, this cannot panic
                xwayland_ready(&inner);
                Ok(calloop::PostAction::Remove)
            },
        )
        .map_err(|err| io::Error::new(io::ErrorKind::Other, err))?;

    guard.instance = Some(XWaylandInstance {
        display_lock: lock,
        wayland_client_fd: None,
        wm_fd: Some(x_wm_me),
        child_stdout: Some(child_stdout),
    });

    Ok(())
}

/// An event source for monitoring XWayland status
///
/// You need to insert it in a [`calloop`] event loop to handle the events it produces,
/// of type [`XWaylandEvent`], which notify you about startup and shutdown of the Xwayland
/// instance.
#[derive(Debug)]
#[must_use]
pub struct XWaylandSource {
    channel: Channel<XWaylandEvent>,
}

impl calloop::EventSource for XWaylandSource {
    type Event = XWaylandEvent;
    type Metadata = ();
    type Ret = ();
    type Error = io::Error;

    fn process_events<F>(
        &mut self,
        readiness: calloop::Readiness,
        token: calloop::Token,
        mut callback: F,
    ) -> io::Result<calloop::PostAction>
    where
        F: FnMut(Self::Event, &mut Self::Metadata) -> Self::Ret,
    {
        self.channel
            .process_events(readiness, token, |event, &mut ()| match event {
                channel::Event::Msg(msg) => callback(msg, &mut ()),
                channel::Event::Closed => {}
            })
            .map_err(|err| io::Error::new(io::ErrorKind::BrokenPipe, err))
    }

    fn register(
        &mut self,
        poll: &mut calloop::Poll,
        factory: &mut calloop::TokenFactory,
    ) -> calloop::Result<()> {
        self.channel.register(poll, factory)
    }

    fn reregister(
        &mut self,
        poll: &mut calloop::Poll,
        factory: &mut calloop::TokenFactory,
    ) -> calloop::Result<()> {
        self.channel.reregister(poll, factory)
    }

    fn unregister(&mut self, poll: &mut calloop::Poll) -> calloop::Result<()> {
        self.channel.unregister(poll)
    }
}

impl Inner {
    // Shutdown the XWayland server and cleanup everything
    fn shutdown(&mut self) {
        // don't do anything if not running
        if self.instance.take().is_some() {
            info!(self.log, "Shutting down XWayland.");

            // send error occurs if the user dropped the channel... We cannot do much except ignore.
            let _ = self.sender.send(XWaylandEvent::Exited);

            // All connections and lockfiles are cleaned by their destructors

            // Remove DISPLAY from the env
            ::std::env::remove_var("DISPLAY");
            // We do like wlroots:
            // > We do not kill the XWayland process, it dies to broken pipe
            // > after we close our side of the wm/wl fds. This is more reliable
            // > than trying to kill something that might no longer be XWayland.
        }
    }
}

fn xwayland_ready(inner: &Rc<RefCell<Inner>>) {
    // Lots of re-borrowing to please the borrow-checker
    let mut guard = inner.borrow_mut();
    let guard = &mut *guard;
    // instance should never be None at this point
    let instance = guard.instance.as_mut().unwrap();
    // neither the child_stdout
    let child_stdout = instance.child_stdout.as_mut().unwrap();

    // This reads the one byte that is written when sh receives SIGUSR1
    let mut buffer = [0];
    let success = match child_stdout.read(&mut buffer) {
        Ok(len) => len > 0 && buffer[0] == b'S',
        Err(e) => {
            error!(guard.log, "Checking launch status failed"; "err" => format!("{:?}", e));
            false
        }
    };

    if success {
        // setup the environment
        ::std::env::set_var("DISPLAY", format!(":{}", instance.display_lock.display()));

        // signal the WM
        info!(
            guard.log,
            "XWayland is ready on DISPLAY \":{}\", signaling the WM.",
            instance.display_lock.display()
        );
        // send error occurs if the user dropped the channel... We cannot do much except ignore.
        let _ = guard.sender.send(XWaylandEvent::Ready {
            connection: instance.wm_fd.take().unwrap(), // This is a bug if None
            client: instance.wayland_client_fd.take().unwrap(),
            display: instance.display_lock.display(),
        });
    } else {
        error!(
            guard.log,
            "XWayland crashed at startup, will not try to restart it."
        );
    }
}

/// Spawn XWayland with given sockets on given display
///
/// Returns a pipe that outputs 'S' upon successful launch.
fn spawn_xwayland(
    display: u32,
    wayland_socket: UnixStream,
    wm_socket: UnixStream,
    listen_sockets: &[UnixStream],
) -> io::Result<ChildStdout> {
    let mut command = Command::new("sh");

    // We use output stream to communicate because FD is easier to handle than exit code.
    command.stdout(Stdio::piped());

    let mut xwayland_args = format!(":{} -rootless -terminate -wm {}", display, wm_socket.as_raw_fd());
    for socket in listen_sockets {
        xwayland_args.push_str(&format!(" -listen {}", socket.as_raw_fd()));
    }
    // This command let sh to:
    // * Set up signal handler for USR1
    // * Launch Xwayland with USR1 ignored so Xwayland will signal us when it is ready (also redirect
    //   Xwayland's STDOUT to STDERR so its output, if any, won't distract us)
    // * Print "S" and exit if USR1 is received
    command.arg("-c").arg(format!(
        "trap 'echo S' USR1; (trap '' USR1; exec Xwayland {}) 1>&2 & wait",
        xwayland_args
    ));

    // Setup the environment: clear everything except PATH and XDG_RUNTIME_DIR
    command.env_clear();
    for (key, value) in env::vars_os() {
        if key.to_str() == Some("PATH") || key.to_str() == Some("XDG_RUNTIME_DIR") {
            command.env(key, value);
            continue;
        }
    }
    command.env("WAYLAND_SOCKET", format!("{}", wayland_socket.as_raw_fd()));

    unsafe {
        let wayland_socket_fd = wayland_socket.as_raw_fd();
        let wm_socket_fd = wm_socket.as_raw_fd();
        let socket_fds: Vec<_> = listen_sockets.iter().map(|socket| socket.as_raw_fd()).collect();
        command.pre_exec(move || {
            // unset the CLOEXEC flag from the sockets we need to pass
            // to xwayland
            unset_cloexec(wayland_socket_fd)?;
            unset_cloexec(wm_socket_fd)?;
            for &socket in socket_fds.iter() {
                unset_cloexec(socket)?;
            }
            Ok(())
        });
    }

    let mut child = command.spawn()?;
    Ok(child.stdout.take().expect("stdout should be piped"))
}

/// Remove the `O_CLOEXEC` flag from this `Fd`
///
/// This means that the `Fd` will *not* be automatically
/// closed when we `exec()` into XWayland
fn unset_cloexec(fd: RawFd) -> io::Result<()> {
    use nix::fcntl::{fcntl, FcntlArg, FdFlag};
    fcntl(fd, FcntlArg::F_SETFD(FdFlag::empty()))?;
    Ok(())
}
<commit_msg>xwayland: handle Client creation<commit_after>/*
 * Steps of XWayland server creation
 *
 * Sockets to create:
 * - a pair for XWayland to connect to smithay as a wayland client, we use our
 *   end to insert the XWayland client in the display
 * - a pair for smithay to connect to XWayland as a WM, we give our end to the
 *   WM and it deals with it
 * - 2 listening sockets on which the XWayland server will listen. We need to
 *   bind them ourselves so we know what value put in the $DISPLAY env variable.
 *   This involves some dance with a lockfile to ensure there is no collision with
 *   an other starting xserver
 *   if we listen on display $D, their paths are respectively:
 *   - /tmp/.X11-unix/X$D
 *   - @/tmp/.X11-unix/X$D (abstract socket)
 *
 * The XWayland server is spawned via an intermediate shell
 * -> wlroots does a double-fork while weston a single one, why ??
 *    -> https://stackoverflow.com/questions/881388/
 * -> once it is started, it will check if SIGUSR1 is set to ignored. If so,
 *    if will consider its parent as "smart", and send a SIGUSR1 signal when
 *    startup completes. We want to catch this so we can launch the VM.
 * -> we need to track if the XWayland crashes, to restart it
 *
 * cf https://github.com/swaywm/wlroots/blob/master/xwayland/xwayland.c
 *
 * Setting SIGUSR1 handler is complicated in multithreaded program, because
 * Xwayland will send SIGUSR1 to the process, and if a thread cannot handle
 * SIGUSR1, that thread will be killed.
 *
 * Double-fork can tackle this issue, but this is also very complex in a
 * a multithread program, after forking only signal-safe functions can be used.
 * The only workaround is to fork early before any other thread starts, but
 * doing so will expose an unsafe interface.
 *
 * We use an intermediate shell to translate the signal to simple fd IO.
 * We ask sh to setup SIGUSR1 handler, and in a subshell mute SIGUSR1 and exec
 * Xwayland. When the SIGUSR1 is received, it can communicate to us via redirected
 * STDOUT.
 */
use std::{
    env,
    io::{self, Read},
    os::unix::{
        io::{AsRawFd, RawFd},
        net::UnixStream,
        process::CommandExt,
    },
    process::{ChildStdout, Command, Stdio},
    sync::{Arc, Mutex},
};

use calloop::{
    channel::{self, sync_channel, Channel, SyncSender},
    generic::Generic,
    Interest, LoopHandle, Mode,
};
use wayland_server::{Client, DisplayHandle, backend::{ClientData, ClientId, DisconnectReason}};

use slog::{error, info, o};

use super::x11_sockets::{prepare_x11_sockets, X11Lock};

/// The XWayland handle
#[derive(Debug)]
pub struct XWayland {
    inner: Arc<Mutex<Inner>>,
}

/// Events generated by the XWayland manager
///
/// This is a very low-level interface, only notifying you when the connection
/// with XWayland is up, or when it terminates.
///
/// Your WM code must be able to handle the XWayland server connecting then
/// disconnecting several time in a row, but only a single connection will
/// be active at any given time.
#[derive(Debug)]
#[must_use = "Connection events must be handled to prevent fd leaking"]
pub enum XWaylandEvent {
    /// The XWayland server is ready
    Ready {
        /// Privileged X11 connection to XWayland
        connection: UnixStream,

        /// Wayland client representing XWayland
        client: Client,

        /// The display number the XWayland server is available at.
        ///
        /// This may be used if you wish to set the `DISPLAY` variable manually when spawning processes that
        /// may use XWayland.
        display: u32,
    },

    /// The XWayland server exited
    ///
    /// This event is sent when the [`XWayland`] handle is dropped.
    Exited,
}

impl XWayland {
    /// Create a new XWayland manager
    ///
    /// This function returns both the [`XWayland`] handle and an [`XWaylandSource`] that needs to be inserted
    /// into the [`calloop`] event loop, producing the Xwayland startup and shutdown events.
    pub fn new<L>(logger: L) -> (XWayland, XWaylandSource)
    where
        L: Into<Option<::slog::Logger>>,
    {
        let log = crate::slog_or_fallback(logger);
        // We don't expect to ever have more than 2 messages in flight, if XWayland got ready and then died right away
        let (sender, channel) = sync_channel(2);
        let inner = Arc::new(Mutex::new(Inner {
            instance: None,
            sender,
            log: log.new(o!("smithay_module" => "XWayland")),
        }));
        (XWayland { inner }, XWaylandSource { channel })
    }

    /// Attempt to start the XWayland instance
    ///
    /// If it succeeds, you'll eventually receive an `XWaylandEvent::Ready`
    /// through the source provided by `XWayland::new()` containing an
    /// `UnixStream` representing your WM connection to XWayland, and the
    /// wayland `Client` for XWayland.
    ///
    /// Does nothing if XWayland is already started or starting.
    pub fn start<D>(&self, loop_handle: LoopHandle<'_, D>, dh: DisplayHandle) -> io::Result<()> {
        launch(&self.inner, loop_handle, dh)
    }

    /// Shutdown XWayland
    ///
    /// Does nothing if it was not already running, otherwise kills it and you will
    /// later receive a `XWaylandEvent::Exited` event.
    pub fn shutdown(&self) {
        self.inner.lock().unwrap().shutdown();
    }
}

impl Drop for XWayland {
    fn drop(&mut self) {
        self.inner.lock().unwrap().shutdown();
    }
}

#[derive(Debug)]
struct XWaylandInstance {
    display_lock: X11Lock,
    wayland_client: Option<Client>,
    wm_fd: Option<UnixStream>,
    child_stdout: ChildStdout,
}

// Inner implementation of the XWayland manager
#[derive(Debug)]
struct Inner {
    sender: SyncSender<XWaylandEvent>,
    instance: Option<XWaylandInstance>,
    log: ::slog::Logger,
}

struct XWaylandClientData {
    inner: Arc<Mutex<Inner>>,
}

impl ClientData for XWaylandClientData {
    fn initialized(&self, _client_id: ClientId) {}
    fn disconnected(&self, _client_id: ClientId, _reason: DisconnectReason) {
        // If we are unable to take a lock we are most likely called during
        // a shutdown. This will definitely be the case when the compositor exits
        // and the XWayland instance is dropped.
        if let Ok(mut guard) = self.inner.try_lock() {
            guard.shutdown();
        }
    }
}

// Launch an XWayland server
//
// Does nothing if there is already a launched instance
fn launch<D>(inner: &Arc<Mutex<Inner>>, loop_handle: LoopHandle<'_, D>, mut dh: DisplayHandle) -> io::Result<()> {
    let mut guard = inner.lock().unwrap();
    if guard.instance.is_some() {
        return Ok(());
    }

    info!(guard.log, "Starting XWayland");

    let (x_wm_x11, x_wm_me) = UnixStream::pair()?;
    let (wl_x11, wl_me) = UnixStream::pair()?;

    let (lock, x_fds) = prepare_x11_sockets(guard.log.clone())?;

    // we have now created all the required sockets

    // all is ready, we can do the fork dance
    let child_stdout = match spawn_xwayland(lock.display(), wl_x11, x_wm_x11, &x_fds) {
        Ok(child_stdout) => child_stdout,
        Err(e) => {
            error!(guard.log, "XWayland failed to spawn"; "err" => format!("{:?}", e));
            return Err(e);
        }
    };

    let loop_inner = inner.clone();
    loop_handle
        .insert_source(
            Generic::<_, io::Error>::new(child_stdout.as_raw_fd(), Interest::READ, Mode::Level),
            move |_, _, _| {
                // the closure must be called exactly one time, this cannot panic
                xwayland_ready(&loop_inner);
                Ok(calloop::PostAction::Remove)
            },
        )
        .map_err(|err| io::Error::new(io::ErrorKind::Other, err))?;

    let client = dh.insert_client(wl_me, Arc::new(XWaylandClientData { inner: inner.clone() }))?;
    guard.instance = Some(XWaylandInstance {
        display_lock: lock,
        wayland_client: Some(client),
        wm_fd: Some(x_wm_me),
        child_stdout,
    });

    Ok(())
}

/// An event source for monitoring XWayland status
///
/// You need to insert it in a [`calloop`] event loop to handle the events it produces,
/// of type [`XWaylandEvent`], which notify you about startup and shutdown of the Xwayland
/// instance.
#[derive(Debug)]
#[must_use]
pub struct XWaylandSource {
    channel: Channel<XWaylandEvent>,
}

impl calloop::EventSource for XWaylandSource {
    type Event = XWaylandEvent;
    type Metadata = ();
    type Ret = ();
    type Error = io::Error;

    fn process_events<F>(
        &mut self,
        readiness: calloop::Readiness,
        token: calloop::Token,
        mut callback: F,
    ) -> io::Result<calloop::PostAction>
    where
        F: FnMut(Self::Event, &mut Self::Metadata) -> Self::Ret,
    {
        self.channel
            .process_events(readiness, token, |event, &mut ()| match event {
                channel::Event::Msg(msg) => callback(msg, &mut ()),
                channel::Event::Closed => {}
            })
            .map_err(|err| io::Error::new(io::ErrorKind::BrokenPipe, err))
    }

    fn register(
        &mut self,
        poll: &mut calloop::Poll,
        factory: &mut calloop::TokenFactory,
    ) -> calloop::Result<()> {
        self.channel.register(poll, factory)
    }

    fn reregister(
        &mut self,
        poll: &mut calloop::Poll,
        factory: &mut calloop::TokenFactory,
    ) -> calloop::Result<()> {
        self.channel.reregister(poll, factory)
    }

    fn unregister(&mut self, poll: &mut calloop::Poll) -> calloop::Result<()> {
        self.channel.unregister(poll)
    }
}

impl Inner {
    // Shutdown the XWayland server and cleanup everything
    fn shutdown(&mut self) {
        // don't do anything if not running
        if self.instance.take().is_some() {
            info!(self.log, "Shutting down XWayland.");

            // send error occurs if the user dropped the channel... We cannot do much except ignore.
            let _ = self.sender.send(XWaylandEvent::Exited);

            // All connections and lockfiles are cleaned by their destructors

            // Remove DISPLAY from the env
            ::std::env::remove_var("DISPLAY");
            // We do like wlroots:
            // > We do not kill the XWayland process, it dies to broken pipe
            // > after we close our side of the wm/wl fds. This is more reliable
            // > than trying to kill something that might no longer be XWayland.
        }
    }
}

fn xwayland_ready(inner: &Arc<Mutex<Inner>>) {
    // Lots of re-borrowing to please the borrow-checker
    let mut guard = inner.lock().unwrap();
    let guard = &mut *guard;
    // instance should never be None at this point
    let instance = guard.instance.as_mut().unwrap();
    // neither the child_stdout
    let child_stdout = &mut instance.child_stdout;

    // This reads the one byte that is written when sh receives SIGUSR1
    let mut buffer = [0];
    let success = match child_stdout.read(&mut buffer) {
        Ok(len) => len > 0 && buffer[0] == b'S',
        Err(e) => {
            error!(guard.log, "Checking launch status failed"; "err" => format!("{:?}", e));
            false
        }
    };

    if success {
        // setup the environment
        ::std::env::set_var("DISPLAY", format!(":{}", instance.display_lock.display()));

        // signal the WM
        info!(
            guard.log,
            "XWayland is ready on DISPLAY \":{}\", signaling the WM.",
            instance.display_lock.display()
        );
        // send error occurs if the user dropped the channel... We cannot do much except ignore.
        let _ = guard.sender.send(XWaylandEvent::Ready {
            connection: instance.wm_fd.take().unwrap(), // This is a bug if None
            client: instance.wayland_client.take().unwrap(),
            display: instance.display_lock.display(),
        });
    } else {
        error!(
            guard.log,
            "XWayland crashed at startup, will not try to restart it."
        );
    }
}

/// Spawn XWayland with given sockets on given display
///
/// Returns a pipe that outputs 'S' upon successful launch.
fn spawn_xwayland(
    display: u32,
    wayland_socket: UnixStream,
    wm_socket: UnixStream,
    listen_sockets: &[UnixStream],
) -> io::Result<ChildStdout> {
    let mut command = Command::new("sh");

    // We use output stream to communicate because FD is easier to handle than exit code.
    command.stdout(Stdio::piped());

    let mut xwayland_args = format!(":{} -rootless -terminate -wm {}", display, wm_socket.as_raw_fd());
    for socket in listen_sockets {
        xwayland_args.push_str(&format!(" -listen {}", socket.as_raw_fd()));
    }
    // This command let sh to:
    // * Set up signal handler for USR1
    // * Launch Xwayland with USR1 ignored so Xwayland will signal us when it is ready (also redirect
    //   Xwayland's STDOUT to STDERR so its output, if any, won't distract us)
    // * Print "S" and exit if USR1 is received
    command.arg("-c").arg(format!(
        "trap 'echo S' USR1; (trap '' USR1; exec Xwayland {}) 1>&2 & wait",
        xwayland_args
    ));

    // Setup the environment: clear everything except PATH and XDG_RUNTIME_DIR
    command.env_clear();
    for (key, value) in env::vars_os() {
        if key.to_str() == Some("PATH") || key.to_str() == Some("XDG_RUNTIME_DIR") {
            command.env(key, value);
            continue;
        }
    }
    command.env("WAYLAND_SOCKET", format!("{}", wayland_socket.as_raw_fd()));

    unsafe {
        let wayland_socket_fd = wayland_socket.as_raw_fd();
        let wm_socket_fd = wm_socket.as_raw_fd();
        let socket_fds: Vec<_> = listen_sockets.iter().map(|socket| socket.as_raw_fd()).collect();
        command.pre_exec(move || {
            // unset the CLOEXEC flag from the sockets we need to pass
            // to xwayland
            unset_cloexec(wayland_socket_fd)?;
            unset_cloexec(wm_socket_fd)?;
            for &socket in socket_fds.iter() {
                unset_cloexec(socket)?;
            }
            Ok(())
        });
    }

    let mut child = command.spawn()?;
    Ok(child.stdout.take().expect("stdout should be piped"))
}

/// Remove the `O_CLOEXEC` flag from this `Fd`
///
/// This means that the `Fd` will *not* be automatically
/// closed when we `exec()` into XWayland
fn unset_cloexec(fd: RawFd) -> io::Result<()> {
    use nix::fcntl::{fcntl, FcntlArg, FdFlag};
    fcntl(fd, FcntlArg::F_SETFD(FdFlag::empty()))?;
    Ok(())
}
<|endoftext|>


<commit_before>use netinfo::{ConnToInodeMap, InodeToPidMap, Inode, Pid, Connection, ConnectionType, TransportType, PacketInfo};
use netinfo::error::*;
use std::collections::HashMap;

/// `PacketMatcher` provides a function which allows matching
/// source/destination adresses for packet to a process.
#[derive(Debug)]
pub struct PacketMatcher {
    /// Matches a source/destination adress to a socket inode of a process.
    conn_to_inode_map: ConnToInodeMap,

    /// Matches a socket inode to its process.
    inode_to_pid_map: InodeToPidMap,

    /// Store all already-handled connections in HashMap so we can look them up
    /// without having to refresh the other maps (less cpu intensive)
    ///
    /// If the resulting value is None, the connection could not be associated with
    /// process and it is in most cases pointless to try again.
    known_connections: HashMap<(TransportType, Connection), Option<(Inode, Pid)>>,
}

impl PacketMatcher {
    /// Constructor for `PacketMatcher`.
    pub fn new() -> PacketMatcher {
        PacketMatcher {
            conn_to_inode_map: ConnToInodeMap::new(),
            inode_to_pid_map: InodeToPidMap::new(),
            known_connections: HashMap::new(),
        }
    }

    /// This function updates the tables that are used for the matching.
    pub fn refresh(&mut self) -> Result<()> {
        self.conn_to_inode_map.refresh()?;
        self.inode_to_pid_map.refresh()?;
        Ok(())
    }

    /// Find the process to which a packet belongs.
    pub fn find_pid(&mut self, packet_info: PacketInfo) -> Result<Option<Pid>> {
        let tt = packet_info.transport_type;
        let mut c = Connection::from(packet_info.clone());
        match packet_info.inout_type {
            // Incoming packets have a their address ordered as "remote addr -> local addr" but
            // our connection-to-inode table only contains "local addr -> remote addr" entries.
            // So we have to reverse the connection.
            Some(ConnectionType::Incoming) => { c = c.get_reverse(); }
            Some(ConnectionType::Outgoing) => {}

            // TODO: use heuristics to guess whether packet is incoming or outgoing - this is just a corner case though
            None => { return Ok(None); }
        }

        self.find_pid_in_table(tt, c)
    }

    /// Find the process by transport type and connection. The connection has to be
    /// already reversed for incoming packets.
    /// This functions uses the "self.known_connections"-HashMap or adds the new connection
    /// to it.
    fn find_pid_in_table(&mut self, tt: TransportType, c: Connection) -> Result<Option<Pid>> {
        if let Some(&res) = self.known_connections.get(&(tt, c)) {
            // Known connection! Does this connection have a process?
            Ok(res.map(|(_, pid)| pid))
        } else {
            // Unknown connection!
            let inode_pid_opt = self.get_new_connection(tt, c)?;
            self.known_connections.insert((tt, c), inode_pid_opt);
            Ok(inode_pid_opt.map(|(_, pid)| pid))
        }
    }

    /// Find the process by transport type and connection. The connection has to be
    /// already reversed for incoming packets.
    /// At the beginning the internal conn->inode->pid tables are refreshed, then the
    /// it tries to assign the connection to a pid.
    fn get_new_connection(&mut self, tt: TransportType, c: Connection) -> Result<Option<(Inode, Pid)>> {
        self.refresh()?;

        Ok(
            self.conn_to_inode_map
                .find_inode(tt, c)
                .and_then(|inode| self.inode_to_pid_map.find_pid(inode)
                                                        .map(|pid| (inode, pid)))
        )
    }
}
<commit_msg>Update cached connections on refresh<commit_after>use netinfo::{ConnToInodeMap, InodeToPidMap, Inode, Pid, Connection, ConnectionType, TransportType, PacketInfo};
use netinfo::error::*;
use std::collections::HashMap;

/// Provides the tables for PacketMatcher.
#[derive(Debug)]
struct PacketMatcherTables {
    /// Matches a source/destination adress to a socket inode of a process.
    conn_to_inode_map: ConnToInodeMap,

    /// Matches a socket inode to its process.
    inode_to_pid_map: InodeToPidMap,
}

/// `PacketMatcher` provides a function which allows matching
/// source/destination adresses for packet to a process.
#[derive(Debug)]
pub struct PacketMatcher {
    /// conn->inode->pid Tables
    tables: PacketMatcherTables,

    /// Store all already-handled connections in HashMap so we can look them up
    /// without having to refresh the other maps (less cpu intensive/caching)
    ///
    /// If the resulting value is None, the connection could not be associated with
    /// process and it is in most cases pointless to try again.
    known_connections: HashMap<(TransportType, Connection), Option<(Inode, Pid)>>,
}

impl PacketMatcherTables {
    /// Constructor for `PacketMatcherTables`.
    pub fn new() -> PacketMatcherTables {
        PacketMatcherTables {
            conn_to_inode_map: ConnToInodeMap::new(),
            inode_to_pid_map: InodeToPidMap::new(),
        }
    }

    /// This function updates the tables that are used for the matching.
    pub fn refresh(&mut self) -> Result<()> {
        self.conn_to_inode_map.refresh()?;
        self.inode_to_pid_map.refresh()?;
        Ok(())
    }

    /// Find the process by transport type and connection. The connection has to be
    /// already reversed for incoming packets.
    /// At the beginning the internal conn->inode->pid tables are refreshed, then the
    /// it tries to assign the connection to a pid.
    fn get_new_connection(&mut self, tt: TransportType, c: Connection) -> Result<Option<(Inode, Pid)>> {
        self.refresh()?;

        Ok(self.map_connection(tt, c))
    }

    /// Maps connection from conn->inode->pid with internal tables.
    fn map_connection(&self, tt: TransportType, c: Connection) -> Option<(Inode, Pid)> {
        self.conn_to_inode_map
        .find_inode(tt, c)
        .and_then(|inode| self.inode_to_pid_map.find_pid(inode)
                                                .map(|pid| (inode, pid)))
    }
}

impl PacketMatcher {
    /// Constructor for `PacketMatcher`.
    pub fn new() -> PacketMatcher {
        PacketMatcher {
            tables: PacketMatcherTables::new(),
            known_connections: HashMap::new(),
        }
    }

    /// This function updates the tables that are used for the matching.
    pub fn refresh(&mut self) -> Result<()> {
        self.tables.refresh()?;
        self.update_known_connections()?;
        Ok(())
    }

    /// A process might end a connection and another might open the same connection.
    /// To prevent wrong assignment we have to update/purge the "self.known_connections"
    /// when connection is closed/reopened.
    fn update_known_connections(&mut self) -> Result<()> {
        self.known_connections =
                self.known_connections.iter()
                .filter_map(|(&(tt, c), &old)| {
                    let new = self.tables.map_connection(tt, c);
                    match (old, new) {
                        // Connection used to exist, but not anymore -> remove/drop so
                        // if reopened it will be assinged to a new program
                        (Some(_), None) => { None }

                        // in these cases the connection was changed/did not change/was created
                        (_, new) => { Some(((tt, c), new)) }
                    }

                })
                .collect();
        Ok(())
    }

    /// Find the process to which a packet belongs.
    pub fn find_pid(&mut self, packet_info: PacketInfo) -> Result<Option<Pid>> {
        let tt = packet_info.transport_type;
        let mut c = Connection::from(packet_info.clone());
        match packet_info.inout_type {
            // Incoming packets have a their address ordered as "remote addr -> local addr" but
            // our connection-to-inode table only contains "local addr -> remote addr" entries.
            // So we have to reverse the connection.
            Some(ConnectionType::Incoming) => { c = c.get_reverse(); }
            Some(ConnectionType::Outgoing) => {}

            // TODO: use heuristics to guess whether packet is incoming or outgoing - this is just a corner case though
            None => { return Ok(None); }
        }

        self.find_pid_cached(tt, c)
    }

    /// Find the process by transport type and connection. The connection has to be
    /// already reversed for incoming packets.
    /// This functions uses the "self.known_connections"-HashMap or adds the new connection
    /// to it.
    fn find_pid_cached(&mut self, tt: TransportType, c: Connection) -> Result<Option<Pid>> {
        if let Some(&res) = self.known_connections.get(&(tt, c)) {
            // Known connection! Does this connection have a process?
            Ok(res.map(|(_, pid)| pid))
        } else {
            // Unknown connection!
            let inode_pid_opt = self.tables.get_new_connection(tt, c)?;
            self.known_connections.insert((tt, c), inode_pid_opt);
            Ok(inode_pid_opt.map(|(_, pid)| pid))
        }
    }
}
<|endoftext|>


<commit_before>//
// Copyright 2021 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::server::{create_and_start_server, WasmHandler};
use hyper::{
    client::Client,
    service::{make_service_fn, service_fn},
    Body, Response,
};
use prost::Message;
use std::{
    convert::Infallible,
    fs,
    net::{Ipv6Addr, SocketAddr},
    sync::{Arc, Mutex},
};

extern crate test;
use test::Bencher;

const TEST_WASM_MODULE_PATH: &str = "testdata/non-oak-minimal.wasm";
const OAK_FUNCTIONS_SERVER_PORT: u16 = 9001;
const STATIC_SERVER_PORT: u16 = 9002;

#[tokio::test]
async fn test_server() {
    let _ = env_logger::builder().is_test(true).try_init();
    let address = SocketAddr::from((Ipv6Addr::UNSPECIFIED, OAK_FUNCTIONS_SERVER_PORT));
    let (notify_sender, notify_receiver) = tokio::sync::oneshot::channel::<()>();

    let wasm_module_bytes =
        fs::read(TEST_WASM_MODULE_PATH).expect("Couldn't read test Wasm module");

    let server_fut = create_and_start_server(&address, &wasm_module_bytes, notify_receiver);
    let client_fut = start_client(OAK_FUNCTIONS_SERVER_PORT, notify_sender);

    let (res, _) = tokio::join!(server_fut, client_fut);
    assert!(res.is_ok());
}

async fn start_client(port: u16, notify_sender: tokio::sync::oneshot::Sender<()>) {
    let client = Client::new();
    let request = hyper::Request::builder()
        .method(http::Method::POST)
        .uri(format!("http://localhost:{}/invoke", port))
        .body(Body::empty())
        .unwrap();
    let resp = client
        .request(request)
        .await
        .expect("Error while awaiting response");

    assert_eq!(resp.status(), hyper::StatusCode::OK);
    assert_eq!(hyper::body::to_bytes(resp.into_body()).await.unwrap(), "");

    notify_sender
        .send(())
        .expect("Couldn't send completion signal.");
}

// TODO(#1933): Currently there is no support for running benchmark tests in the runner.
// Run this with: `cargo bench --manifest-path=oak_functions/loader/Cargo.toml`
#[bench]
fn bench_wasm_handler(bencher: &mut Bencher) {
    let summary = bencher.bench(|bencher| {
        let wasm_module_bytes =
            fs::read(TEST_WASM_MODULE_PATH).expect("Couldn't read test Wasm module");
        let wasm_handler =
            WasmHandler::create(&wasm_module_bytes).expect("Couldn't create the server");
        let rt = tokio::runtime::Runtime::new().unwrap();
        bencher.iter(|| {
            let request = hyper::Request::builder()
                .method(http::Method::POST)
                // We don't really send the request. So the hostname can be anything. It is only
                // needed for building a valid request.
                .uri("http://example.com/invoke")
                .body(Body::empty())
                .unwrap();
            let resp = rt.block_on(wasm_handler.handle_request(request)).unwrap();
            assert_eq!(resp.status(), hyper::StatusCode::OK);
        });
    });

    // When running `cargo test` this benchmark test gets executed too, but `summary` will be `None`
    // in that case. So, here we first check that `summary` is not empty.
    if let Some(summary) = summary {
        // We expect the `mean` time for loading the test Wasm module and running its main function
        // to be less than a fixed threshold.
        // Note: `summary.mean` is in nanoseconds, even though it is not explicitly documented in
        // https://doc.rust-lang.org/test/stats/struct.Summary.html.
        assert!(summary.mean < 50_000.0, "elapsed time: {}ns", summary.mean);
    }
}

#[test]
fn load_lookup_entries_empty() {
    let _ = env_logger::builder().is_test(true).try_init();
    let empty = vec![];
    let entries = crate::load_lookup_entries(empty.as_ref()).unwrap();
    assert!(entries.is_empty());
}

// Fix the serialized representation for testing by manually annotating individual bytes.
//
// See https://developers.google.com/protocol-buffers/docs/encoding#structure.
const ENTRY_0_LENGTH_DELIMITED: &[u8] = &[
    8,  // Message total length.
    10, // Field 1 key: (1<<3) | 2
    2,  // Field 1 length.
    14, 12, // Field 1 value.
    18, // Field 2 key: (2<<3) | 2
    2,  // Field 2 length.
    19, 88, // Field 2 value.
];

const ENTRY_1_LENGTH_DELIMITED: &[u8] = &[
    15, // Message total length.
    10, // Field 1 key: (1<<3) | 2
    5,  // Field 1 length.
    b'H', b'a', b'r', b'r', b'y', // Field 1 value.
    18,   // Field 2 key: (2<<3) | 2
    6,    // Field 2 length.
    b'P', b'o', b't', b't', b'e', b'r', // Field 2 value.
];

// Ensure that the serialized representation is correct.
#[test]
fn check_serialized_lookup_entries() {
    {
        let mut buf = vec![];
        let entry = crate::proto::Entry {
            key: vec![14, 12],
            value: vec![19, 88],
        };
        entry.encode_length_delimited(&mut buf).unwrap();
        assert_eq!(buf, ENTRY_0_LENGTH_DELIMITED);
    }
    {
        let mut buf = vec![];
        let entry = crate::proto::Entry {
            key: b"Harry".to_vec(),
            value: b"Potter".to_vec(),
        };
        entry.encode_length_delimited(&mut buf).unwrap();
        assert_eq!(buf, ENTRY_1_LENGTH_DELIMITED);
    }
}

#[test]
fn load_lookup_entries_multiple_entries() {
    let _ = env_logger::builder().is_test(true).try_init();
    let mut buf = vec![];
    buf.append(&mut ENTRY_0_LENGTH_DELIMITED.to_vec());
    buf.append(&mut ENTRY_1_LENGTH_DELIMITED.to_vec());
    let entries = crate::load_lookup_entries(buf.as_ref()).unwrap();
    assert_eq!(entries.len(), 2);
    assert_eq!(entries.get(&[14, 12].to_vec()), Some(&vec![19, 88]));
    assert_eq!(entries.get(&b"Harry".to_vec()), Some(&b"Potter".to_vec()));
}

#[test]
fn load_lookup_entries_multiple_entries_trailing() {
    let _ = env_logger::builder().is_test(true).try_init();
    let mut buf = vec![];
    buf.append(&mut ENTRY_0_LENGTH_DELIMITED.to_vec());
    buf.append(&mut ENTRY_1_LENGTH_DELIMITED.to_vec());
    // Add invalid trailing bytes.
    buf.append(&mut vec![1, 2, 3]);
    let res = crate::load_lookup_entries(buf.as_ref());
    assert!(res.is_err());
}

#[test]
fn load_lookup_entries_invalid() {
    let _ = env_logger::builder().is_test(true).try_init();
    // Invalid bytes.
    let buf = vec![1, 2, 3];
    let res = crate::load_lookup_entries(buf.as_ref());
    assert!(res.is_err());
}

/// A mock implementation of a static server that always returns the same configurable response for
/// any incoming HTTP request.
#[derive(Default)]
struct MockStaticServer {
    response_body: Arc<Mutex<Vec<u8>>>,
}

impl MockStaticServer {
    /// Sets the content of the response body to return for any request.
    fn set_response_body(&self, response_body: Vec<u8>) {
        *self
            .response_body
            .lock()
            .expect("could not lock response body mutex") = response_body;
    }

    /// Starts serving, listening on the provided port.
    async fn serve(&self, port: u16) {
        let address = SocketAddr::from((Ipv6Addr::UNSPECIFIED, port));
        let response_body = self.response_body.clone();
        let server = hyper::Server::bind(&address).serve(make_service_fn(|_conn| {
            let response_body = response_body.clone();
            async {
                Ok::<_, Infallible>(service_fn(move |_req| {
                    let response_body = response_body.clone();
                    async move {
                        let response_body: Vec<u8> = response_body
                            .lock()
                            .expect("could not lock response body mutex")
                            .clone();
                        Ok::<_, Infallible>(Response::new(Body::from(response_body)))
                    }
                }))
            }
        }));
        server.await.unwrap();
    }
}

#[tokio::test]
async fn lookup_data_refresh() {
    let mock_static_server = Arc::new(MockStaticServer::default());

    let mock_static_server_clone = mock_static_server.clone();
    tokio::spawn(async move { mock_static_server_clone.serve(STATIC_SERVER_PORT).await });

    let lookup_data =
        crate::LookupData::new_empty(&format!("http://localhost:{}", STATIC_SERVER_PORT));
    assert_eq!(lookup_data.len(), 0);

    // Initially empty file, no entries.
    lookup_data.refresh().await.unwrap();
    assert_eq!(lookup_data.len(), 0);

    // Single entry.
    mock_static_server.set_response_body(ENTRY_0_LENGTH_DELIMITED.to_vec());
    lookup_data.refresh().await.unwrap();
    assert_eq!(lookup_data.len(), 1);
    assert_eq!(lookup_data.get(&[14, 12]), Some(vec![19, 88]));
    assert_eq!(lookup_data.get(b"Harry"), None);

    // Empty file again.
    mock_static_server.set_response_body(vec![]);
    lookup_data.refresh().await.unwrap();
    assert_eq!(lookup_data.len(), 0);

    // A different entry.
    mock_static_server.set_response_body(ENTRY_1_LENGTH_DELIMITED.to_vec());
    lookup_data.refresh().await.unwrap();
    assert_eq!(lookup_data.len(), 1);
    assert_eq!(lookup_data.get(&[14, 12]), None);
    assert_eq!(lookup_data.get(b"Harry"), Some(b"Potter".to_vec()));

    // Two entries.
    let mut buf = ENTRY_0_LENGTH_DELIMITED.to_vec();
    buf.append(&mut ENTRY_1_LENGTH_DELIMITED.to_vec());
    mock_static_server.set_response_body(buf);
    lookup_data.refresh().await.unwrap();
    assert_eq!(lookup_data.len(), 2);
    assert_eq!(lookup_data.get(&[14, 12]), Some(vec![19, 88]));
    assert_eq!(lookup_data.get(b"Harry"), Some(b"Potter".to_vec()));
}
<commit_msg>Use Duration time for benchmark (#1975)<commit_after>//
// Copyright 2021 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::server::{create_and_start_server, WasmHandler};
use hyper::{
    client::Client,
    service::{make_service_fn, service_fn},
    Body, Response,
};
use prost::Message;
use std::{
    convert::Infallible,
    fs,
    net::{Ipv6Addr, SocketAddr},
    sync::{Arc, Mutex},
    time::Duration,
};

extern crate test;
use test::Bencher;

const TEST_WASM_MODULE_PATH: &str = "testdata/non-oak-minimal.wasm";
const OAK_FUNCTIONS_SERVER_PORT: u16 = 9001;
const STATIC_SERVER_PORT: u16 = 9002;

#[tokio::test]
async fn test_server() {
    let _ = env_logger::builder().is_test(true).try_init();
    let address = SocketAddr::from((Ipv6Addr::UNSPECIFIED, OAK_FUNCTIONS_SERVER_PORT));
    let (notify_sender, notify_receiver) = tokio::sync::oneshot::channel::<()>();

    let wasm_module_bytes =
        fs::read(TEST_WASM_MODULE_PATH).expect("Couldn't read test Wasm module");

    let server_fut = create_and_start_server(&address, &wasm_module_bytes, notify_receiver);
    let client_fut = start_client(OAK_FUNCTIONS_SERVER_PORT, notify_sender);

    let (res, _) = tokio::join!(server_fut, client_fut);
    assert!(res.is_ok());
}

async fn start_client(port: u16, notify_sender: tokio::sync::oneshot::Sender<()>) {
    let client = Client::new();
    let request = hyper::Request::builder()
        .method(http::Method::POST)
        .uri(format!("http://localhost:{}/invoke", port))
        .body(Body::empty())
        .unwrap();
    let resp = client
        .request(request)
        .await
        .expect("Error while awaiting response");

    assert_eq!(resp.status(), hyper::StatusCode::OK);
    assert_eq!(hyper::body::to_bytes(resp.into_body()).await.unwrap(), "");

    notify_sender
        .send(())
        .expect("Couldn't send completion signal.");
}

// TODO(#1933): Currently there is no support for running benchmark tests in the runner.
// Run this with: `cargo bench --manifest-path=oak_functions/loader/Cargo.toml`
#[bench]
fn bench_wasm_handler(bencher: &mut Bencher) {
    let summary = bencher.bench(|bencher| {
        let wasm_module_bytes =
            fs::read(TEST_WASM_MODULE_PATH).expect("Couldn't read test Wasm module");
        let wasm_handler =
            WasmHandler::create(&wasm_module_bytes).expect("Couldn't create the server");
        let rt = tokio::runtime::Runtime::new().unwrap();
        bencher.iter(|| {
            let request = hyper::Request::builder()
                .method(http::Method::POST)
                // We don't really send the request. So the hostname can be anything. It is only
                // needed for building a valid request.
                .uri("http://example.com/invoke")
                .body(Body::empty())
                .unwrap();
            let resp = rt.block_on(wasm_handler.handle_request(request)).unwrap();
            assert_eq!(resp.status(), hyper::StatusCode::OK);
        });
    });

    // When running `cargo test` this benchmark test gets executed too, but `summary` will be `None`
    // in that case. So, here we first check that `summary` is not empty.
    if let Some(summary) = summary {
        // `summary.mean` is in nanoseconds, even though it is not explicitly documented in
        // https://doc.rust-lang.org/test/stats/struct.Summary.html.
        let elapsed = Duration::from_nanos(summary.mean as u64);
        // We expect the `mean` time for loading the test Wasm module and running its main function
        // to be less than a fixed threshold.
        assert!(
            elapsed < Duration::from_millis(1),
            "elapsed time: {:?}",
            elapsed
        );
    }
}

#[test]
fn load_lookup_entries_empty() {
    let _ = env_logger::builder().is_test(true).try_init();
    let empty = vec![];
    let entries = crate::load_lookup_entries(empty.as_ref()).unwrap();
    assert!(entries.is_empty());
}

// Fix the serialized representation for testing by manually annotating individual bytes.
//
// See https://developers.google.com/protocol-buffers/docs/encoding#structure.
const ENTRY_0_LENGTH_DELIMITED: &[u8] = &[
    8,  // Message total length.
    10, // Field 1 key: (1<<3) | 2
    2,  // Field 1 length.
    14, 12, // Field 1 value.
    18, // Field 2 key: (2<<3) | 2
    2,  // Field 2 length.
    19, 88, // Field 2 value.
];

const ENTRY_1_LENGTH_DELIMITED: &[u8] = &[
    15, // Message total length.
    10, // Field 1 key: (1<<3) | 2
    5,  // Field 1 length.
    b'H', b'a', b'r', b'r', b'y', // Field 1 value.
    18,   // Field 2 key: (2<<3) | 2
    6,    // Field 2 length.
    b'P', b'o', b't', b't', b'e', b'r', // Field 2 value.
];

// Ensure that the serialized representation is correct.
#[test]
fn check_serialized_lookup_entries() {
    {
        let mut buf = vec![];
        let entry = crate::proto::Entry {
            key: vec![14, 12],
            value: vec![19, 88],
        };
        entry.encode_length_delimited(&mut buf).unwrap();
        assert_eq!(buf, ENTRY_0_LENGTH_DELIMITED);
    }
    {
        let mut buf = vec![];
        let entry = crate::proto::Entry {
            key: b"Harry".to_vec(),
            value: b"Potter".to_vec(),
        };
        entry.encode_length_delimited(&mut buf).unwrap();
        assert_eq!(buf, ENTRY_1_LENGTH_DELIMITED);
    }
}

#[test]
fn load_lookup_entries_multiple_entries() {
    let _ = env_logger::builder().is_test(true).try_init();
    let mut buf = vec![];
    buf.append(&mut ENTRY_0_LENGTH_DELIMITED.to_vec());
    buf.append(&mut ENTRY_1_LENGTH_DELIMITED.to_vec());
    let entries = crate::load_lookup_entries(buf.as_ref()).unwrap();
    assert_eq!(entries.len(), 2);
    assert_eq!(entries.get(&[14, 12].to_vec()), Some(&vec![19, 88]));
    assert_eq!(entries.get(&b"Harry".to_vec()), Some(&b"Potter".to_vec()));
}

#[test]
fn load_lookup_entries_multiple_entries_trailing() {
    let _ = env_logger::builder().is_test(true).try_init();
    let mut buf = vec![];
    buf.append(&mut ENTRY_0_LENGTH_DELIMITED.to_vec());
    buf.append(&mut ENTRY_1_LENGTH_DELIMITED.to_vec());
    // Add invalid trailing bytes.
    buf.append(&mut vec![1, 2, 3]);
    let res = crate::load_lookup_entries(buf.as_ref());
    assert!(res.is_err());
}

#[test]
fn load_lookup_entries_invalid() {
    let _ = env_logger::builder().is_test(true).try_init();
    // Invalid bytes.
    let buf = vec![1, 2, 3];
    let res = crate::load_lookup_entries(buf.as_ref());
    assert!(res.is_err());
}

/// A mock implementation of a static server that always returns the same configurable response for
/// any incoming HTTP request.
#[derive(Default)]
struct MockStaticServer {
    response_body: Arc<Mutex<Vec<u8>>>,
}

impl MockStaticServer {
    /// Sets the content of the response body to return for any request.
    fn set_response_body(&self, response_body: Vec<u8>) {
        *self
            .response_body
            .lock()
            .expect("could not lock response body mutex") = response_body;
    }

    /// Starts serving, listening on the provided port.
    async fn serve(&self, port: u16) {
        let address = SocketAddr::from((Ipv6Addr::UNSPECIFIED, port));
        let response_body = self.response_body.clone();
        let server = hyper::Server::bind(&address).serve(make_service_fn(|_conn| {
            let response_body = response_body.clone();
            async {
                Ok::<_, Infallible>(service_fn(move |_req| {
                    let response_body = response_body.clone();
                    async move {
                        let response_body: Vec<u8> = response_body
                            .lock()
                            .expect("could not lock response body mutex")
                            .clone();
                        Ok::<_, Infallible>(Response::new(Body::from(response_body)))
                    }
                }))
            }
        }));
        server.await.unwrap();
    }
}

#[tokio::test]
async fn lookup_data_refresh() {
    let mock_static_server = Arc::new(MockStaticServer::default());

    let mock_static_server_clone = mock_static_server.clone();
    tokio::spawn(async move { mock_static_server_clone.serve(STATIC_SERVER_PORT).await });

    let lookup_data =
        crate::LookupData::new_empty(&format!("http://localhost:{}", STATIC_SERVER_PORT));
    assert_eq!(lookup_data.len(), 0);

    // Initially empty file, no entries.
    lookup_data.refresh().await.unwrap();
    assert_eq!(lookup_data.len(), 0);

    // Single entry.
    mock_static_server.set_response_body(ENTRY_0_LENGTH_DELIMITED.to_vec());
    lookup_data.refresh().await.unwrap();
    assert_eq!(lookup_data.len(), 1);
    assert_eq!(lookup_data.get(&[14, 12]), Some(vec![19, 88]));
    assert_eq!(lookup_data.get(b"Harry"), None);

    // Empty file again.
    mock_static_server.set_response_body(vec![]);
    lookup_data.refresh().await.unwrap();
    assert_eq!(lookup_data.len(), 0);

    // A different entry.
    mock_static_server.set_response_body(ENTRY_1_LENGTH_DELIMITED.to_vec());
    lookup_data.refresh().await.unwrap();
    assert_eq!(lookup_data.len(), 1);
    assert_eq!(lookup_data.get(&[14, 12]), None);
    assert_eq!(lookup_data.get(b"Harry"), Some(b"Potter".to_vec()));

    // Two entries.
    let mut buf = ENTRY_0_LENGTH_DELIMITED.to_vec();
    buf.append(&mut ENTRY_1_LENGTH_DELIMITED.to_vec());
    mock_static_server.set_response_body(buf);
    lookup_data.refresh().await.unwrap();
    assert_eq!(lookup_data.len(), 2);
    assert_eq!(lookup_data.get(&[14, 12]), Some(vec![19, 88]));
    assert_eq!(lookup_data.get(b"Harry"), Some(b"Potter".to_vec()));
}
<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
use mimir::Admin;
use mimir::AdminType;
use geo::contains::Contains;
use geo;
use std::iter::FromIterator;
use std::rc::Rc;
use gst::rtree::{RTree, Rect};
use std;

/// We want to strip the admin's boundary for the objects referencing it (for performance purpose)
/// thus in the `AdminGeoFinder` we store an Admin without the boundary (the option is emptied)
/// and we store the boundary aside
struct BoundaryAndAdmin(Option<geo::MultiPolygon<f64>>, Rc<Admin>);

impl BoundaryAndAdmin {
    fn new(mut admin: Admin) -> BoundaryAndAdmin {
        let b = std::mem::replace(&mut admin.boundary, None);
        let minimal_admin = Rc::new(admin);
        BoundaryAndAdmin(b, minimal_admin)
    }
}

pub struct AdminGeoFinder {
    admins: RTree<BoundaryAndAdmin>,
}

impl AdminGeoFinder {
    pub fn insert(&mut self, admin: Admin) {
        use ordered_float::OrderedFloat;
        fn min(a: OrderedFloat<f32>, b: f64) -> f32 {
            a.0.min(down(b as f32))
        }
        fn max(a: OrderedFloat<f32>, b: f64) -> f32 {
            a.0.max(up(b as f32))
        }

        let rect = {
            let mut coords = match admin.boundary {
                Some(ref b) => (*b).0.iter().flat_map(|poly| (poly.exterior).0.iter()),
                None => return,
            };
            let first_coord = match coords.next() {
                Some(c) => c,
                None => return,
            };
            let first_rect: Rect = {
                let (x, y) = (first_coord.x() as f32, first_coord.y() as f32);
                Rect::from_float(down(x), up(x), down(y), up(y))
            };
            coords.fold(first_rect, |accu, p| {
                Rect::from_float(
                    min(accu.xmin, p.x()),
                    max(accu.xmax, p.x()),
                    min(accu.ymin, p.y()),
                    max(accu.ymax, p.y()),
                )
            })
        };
        self.admins.insert(rect, BoundaryAndAdmin::new(admin));
    }

    /// Get all Admins overlapping the coordinate
    pub fn get(&self, coord: &geo::Coordinate<f64>) -> Vec<Rc<Admin>> {
        let (x, y) = (coord.x as f32, coord.y as f32);
        let search = Rect::from_float(down(x), up(x), down(y), up(y));
        self.admins
            .get(&search)
            .into_iter()
            .map(|(_, a)| a)
            .filter(|a| {
                a.0
                    .as_ref()
                    .map_or(false, |b| (*b).contains(&geo::Point(*coord)))
            })
            .map(|admin_and_boundary| admin_and_boundary.1.clone())
            .collect()
    }

    /// Iterates on all the admins with a not None boundary.
    pub fn admins<'a>(&'a self) -> Box<Iterator<Item = Admin> + 'a> {
        let iter = self.admins
            .get(&Rect::from_float(
                std::f32::NEG_INFINITY,
                std::f32::INFINITY,
                std::f32::NEG_INFINITY,
                std::f32::INFINITY,
            ))
            .into_iter()
            .map(|(_, a)| {
                let mut admin = (*a.1).clone();
                admin.boundary = a.0.clone();
                admin
            });
        Box::new(iter)
    }

    /// Iterates on all the `Rc<Admin>` in the structure as returned by `get`.
    pub fn admins_without_boundary<'a>(&'a self) -> Box<Iterator<Item = Rc<Admin>> + 'a> {
        let iter = self.admins
            .get(&Rect::from_float(
                std::f32::NEG_INFINITY,
                std::f32::INFINITY,
                std::f32::NEG_INFINITY,
                std::f32::INFINITY,
            ))
            .into_iter()
            .map(|(_, a)| a.1.clone());
        Box::new(iter)
    }
}

impl Default for AdminGeoFinder {
    fn default() -> Self {
        AdminGeoFinder {
            admins: RTree::new(),
        }
    }
}

impl FromIterator<Admin> for AdminGeoFinder {
    fn from_iter<I: IntoIterator<Item = Admin>>(admins: I) -> Self {
        let mut geofinder = AdminGeoFinder::default();

        for admin in admins {
            geofinder.insert(admin);
        }

        geofinder
    }
}

// the goal is that f in [down(f as f32) as f64, up(f as f32) as f64]
fn down(f: f32) -> f32 {
    f - (f * ::std::f32::EPSILON).abs()
}
fn up(f: f32) -> f32 {
    f + (f * ::std::f32::EPSILON).abs()
}

#[test]
fn test_up_down() {
    for &f in [1.0f64, 0., -0., -1., 0.1, -0.1, 0.9, -0.9, 42., -42.].iter() {
        let small_f = f as f32;
        assert!(
            down(small_f) as f64 <= f,
            format!("{} <= {}", down(small_f) as f64, f)
        );
        assert!(
            f <= up(small_f) as f64,
            format!("{} <= {}", f, up(small_f) as f64)
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn p(x: f64, y: f64) -> ::geo::Point<f64> {
        ::geo::Point(::geo::Coordinate { x: x, y: y })
    }

    fn make_admin(offset: f64) -> ::mimir::Admin {
        // the boundary is a big octogon
        let shape = ::geo::Polygon::new(
            ::geo::LineString(vec![
                p(3. + offset, 0. + offset),
                p(6. + offset, 0. + offset),
                p(9. + offset, 3. + offset),
                p(9. + offset, 6. + offset),
                p(6. + offset, 9. + offset),
                p(3. + offset, 9. + offset),
                p(0. + offset, 6. + offset),
                p(0. + offset, 3. + offset),
                p(3. + offset, 0. + offset),
            ]),
            vec![],
        );
        let boundary = ::geo::MultiPolygon(vec![shape]);

        ::mimir::Admin {
            id: format!("admin:offset:{}", offset),
            level: 8,
            name: "city".to_string(),
            label: format!("city {}", offset),
            zip_codes: vec!["421337".to_string()],
            weight: ::std::cell::Cell::new(0.),
            coord: ::mimir::Coord::new(4.0 + offset, 4.0 + offset),
            boundary: Some(boundary),
            insee: "outlook".to_string(),
            admin_type: AdminType::City,
        }
    }

    #[test]
    fn test_two_fake_admins() {
        let mut finder = AdminGeoFinder::default();
        finder.insert(make_admin(40.));
        finder.insert(make_admin(43.));

        // outside
        for coord in [p(48., 41.), p(411., 41.), p(51., 54.), p(53., 53.)].iter() {
            assert!(finder.get(&coord.0).is_empty());
        }

        // inside one
        let admins = finder.get(&p(44., 44.).0);
        assert_eq!(admins.len(), 1);
        assert_eq!(admins[0].id, "admin:offset:40");
        let admins = finder.get(&p(48., 48.).0);
        assert_eq!(admins.len(), 1);
        assert_eq!(admins[0].id, "admin:offset:43");

        // inside two
        let mut admins = finder.get(&p(46., 46.).0);
        admins.sort_by(|a, b| a.id.cmp(&b.id));
        assert_eq!(admins.len(), 2);
        assert_eq!(admins[0].id, "admin:offset:40");
        assert_eq!(admins[1].id, "admin:offset:43");
    }
}
<commit_msg>remove warning<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
use mimir::Admin;
use geo::contains::Contains;
use geo;
use std::iter::FromIterator;
use std::rc::Rc;
use gst::rtree::{RTree, Rect};
use std;

/// We want to strip the admin's boundary for the objects referencing it (for performance purpose)
/// thus in the `AdminGeoFinder` we store an Admin without the boundary (the option is emptied)
/// and we store the boundary aside
struct BoundaryAndAdmin(Option<geo::MultiPolygon<f64>>, Rc<Admin>);

impl BoundaryAndAdmin {
    fn new(mut admin: Admin) -> BoundaryAndAdmin {
        let b = std::mem::replace(&mut admin.boundary, None);
        let minimal_admin = Rc::new(admin);
        BoundaryAndAdmin(b, minimal_admin)
    }
}

pub struct AdminGeoFinder {
    admins: RTree<BoundaryAndAdmin>,
}

impl AdminGeoFinder {
    pub fn insert(&mut self, admin: Admin) {
        use ordered_float::OrderedFloat;
        fn min(a: OrderedFloat<f32>, b: f64) -> f32 {
            a.0.min(down(b as f32))
        }
        fn max(a: OrderedFloat<f32>, b: f64) -> f32 {
            a.0.max(up(b as f32))
        }

        let rect = {
            let mut coords = match admin.boundary {
                Some(ref b) => (*b).0.iter().flat_map(|poly| (poly.exterior).0.iter()),
                None => return,
            };
            let first_coord = match coords.next() {
                Some(c) => c,
                None => return,
            };
            let first_rect: Rect = {
                let (x, y) = (first_coord.x() as f32, first_coord.y() as f32);
                Rect::from_float(down(x), up(x), down(y), up(y))
            };
            coords.fold(first_rect, |accu, p| {
                Rect::from_float(
                    min(accu.xmin, p.x()),
                    max(accu.xmax, p.x()),
                    min(accu.ymin, p.y()),
                    max(accu.ymax, p.y()),
                )
            })
        };
        self.admins.insert(rect, BoundaryAndAdmin::new(admin));
    }

    /// Get all Admins overlapping the coordinate
    pub fn get(&self, coord: &geo::Coordinate<f64>) -> Vec<Rc<Admin>> {
        let (x, y) = (coord.x as f32, coord.y as f32);
        let search = Rect::from_float(down(x), up(x), down(y), up(y));
        self.admins
            .get(&search)
            .into_iter()
            .map(|(_, a)| a)
            .filter(|a| {
                a.0
                    .as_ref()
                    .map_or(false, |b| (*b).contains(&geo::Point(*coord)))
            })
            .map(|admin_and_boundary| admin_and_boundary.1.clone())
            .collect()
    }

    /// Iterates on all the admins with a not None boundary.
    pub fn admins<'a>(&'a self) -> Box<Iterator<Item = Admin> + 'a> {
        let iter = self.admins
            .get(&Rect::from_float(
                std::f32::NEG_INFINITY,
                std::f32::INFINITY,
                std::f32::NEG_INFINITY,
                std::f32::INFINITY,
            ))
            .into_iter()
            .map(|(_, a)| {
                let mut admin = (*a.1).clone();
                admin.boundary = a.0.clone();
                admin
            });
        Box::new(iter)
    }

    /// Iterates on all the `Rc<Admin>` in the structure as returned by `get`.
    pub fn admins_without_boundary<'a>(&'a self) -> Box<Iterator<Item = Rc<Admin>> + 'a> {
        let iter = self.admins
            .get(&Rect::from_float(
                std::f32::NEG_INFINITY,
                std::f32::INFINITY,
                std::f32::NEG_INFINITY,
                std::f32::INFINITY,
            ))
            .into_iter()
            .map(|(_, a)| a.1.clone());
        Box::new(iter)
    }
}

impl Default for AdminGeoFinder {
    fn default() -> Self {
        AdminGeoFinder {
            admins: RTree::new(),
        }
    }
}

impl FromIterator<Admin> for AdminGeoFinder {
    fn from_iter<I: IntoIterator<Item = Admin>>(admins: I) -> Self {
        let mut geofinder = AdminGeoFinder::default();

        for admin in admins {
            geofinder.insert(admin);
        }

        geofinder
    }
}

// the goal is that f in [down(f as f32) as f64, up(f as f32) as f64]
fn down(f: f32) -> f32 {
    f - (f * ::std::f32::EPSILON).abs()
}
fn up(f: f32) -> f32 {
    f + (f * ::std::f32::EPSILON).abs()
}

#[test]
fn test_up_down() {
    for &f in [1.0f64, 0., -0., -1., 0.1, -0.1, 0.9, -0.9, 42., -42.].iter() {
        let small_f = f as f32;
        assert!(
            down(small_f) as f64 <= f,
            format!("{} <= {}", down(small_f) as f64, f)
        );
        assert!(
            f <= up(small_f) as f64,
            format!("{} <= {}", f, up(small_f) as f64)
        );
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use mimir::AdminType;

    fn p(x: f64, y: f64) -> ::geo::Point<f64> {
        ::geo::Point(::geo::Coordinate { x: x, y: y })
    }

    fn make_admin(offset: f64) -> ::mimir::Admin {
        // the boundary is a big octogon
        let shape = ::geo::Polygon::new(
            ::geo::LineString(vec![
                p(3. + offset, 0. + offset),
                p(6. + offset, 0. + offset),
                p(9. + offset, 3. + offset),
                p(9. + offset, 6. + offset),
                p(6. + offset, 9. + offset),
                p(3. + offset, 9. + offset),
                p(0. + offset, 6. + offset),
                p(0. + offset, 3. + offset),
                p(3. + offset, 0. + offset),
            ]),
            vec![],
        );
        let boundary = ::geo::MultiPolygon(vec![shape]);

        ::mimir::Admin {
            id: format!("admin:offset:{}", offset),
            level: 8,
            name: "city".to_string(),
            label: format!("city {}", offset),
            zip_codes: vec!["421337".to_string()],
            weight: ::std::cell::Cell::new(0.),
            coord: ::mimir::Coord::new(4.0 + offset, 4.0 + offset),
            boundary: Some(boundary),
            insee: "outlook".to_string(),
            admin_type: AdminType::City,
        }
    }

    #[test]
    fn test_two_fake_admins() {
        let mut finder = AdminGeoFinder::default();
        finder.insert(make_admin(40.));
        finder.insert(make_admin(43.));

        // outside
        for coord in [p(48., 41.), p(411., 41.), p(51., 54.), p(53., 53.)].iter() {
            assert!(finder.get(&coord.0).is_empty());
        }

        // inside one
        let admins = finder.get(&p(44., 44.).0);
        assert_eq!(admins.len(), 1);
        assert_eq!(admins[0].id, "admin:offset:40");
        let admins = finder.get(&p(48., 48.).0);
        assert_eq!(admins.len(), 1);
        assert_eq!(admins[0].id, "admin:offset:43");

        // inside two
        let mut admins = finder.get(&p(46., 46.).0);
        admins.sort_by(|a, b| a.id.cmp(&b.id));
        assert_eq!(admins.len(), 2);
        assert_eq!(admins[0].id, "admin:offset:40");
        assert_eq!(admins[1].id, "admin:offset:43");
    }
}
<|endoftext|>


<commit_before>//! Compute dominators of a control-flow graph.
//!
//! # The Dominance Relation
//!
//! In a directed graph with a root node **R**, a node **A** is said to *dominate* a
//! node **B** iff every path from **R** to **B** contains **A**.
//!
//! The node **A** is said to *strictly dominate* the node **B** iff **A** dominates
//! **B** and **A ≠ B**.
//!
//! The node **A** is said to be the *immediate dominator* of a node **B** iff it
//! strictly dominates **B** and there does not exist any node **C** where **A**
//! dominates **C** and **C** dominates **B**.

use std::collections::{HashMap, HashSet};
use std::hash::Hash;

use visit::{DfsPostOrder, GraphBase, IntoNeighbors, Visitable, Walker};

/// The dominance relation for some graph and root.
#[derive(Debug, Clone)]
pub struct Dominators<N>
    where N: Copy + Eq + Hash
{
    root: N,
    dominators: HashMap<N, N>,
}

impl<N> Dominators<N>
    where N: Copy + Eq + Hash
{
    /// Get the root node used to construct these dominance relations.
    pub fn root(&self) -> N {
        self.root
    }

    /// Get the immediate dominator of the given node.
    ///
    /// Returns `None` for any node that is not reachable from the root, and for
    /// the root itself.
    pub fn immediate_dominator(&self, node: N) -> Option<N> {
        if node == self.root {
            None
        } else {
            self.dominators.get(&node).cloned()
        }
    }

    /// Iterate over the given node's that strict dominators.
    ///
    /// If the given node is not reachable from the root, then `None` is
    /// returned.
    pub fn strict_dominators(&self, node: N) -> Option<DominatorsIter<N>> {
        if self.dominators.contains_key(&node) {
            Some(DominatorsIter {
                dominators: self,
                node: self.immediate_dominator(node),
            })
        } else {
            None
        }
    }

    /// Iterate over all of the given node's dominators (including the given
    /// node itself).
    ///
    /// If the given node is not reachable from the root, then `None` is
    /// returned.
    pub fn dominators(&self, node: N) -> Option<DominatorsIter<N>> {
        if self.dominators.contains_key(&node) {
            Some(DominatorsIter {
                dominators: self,
                node: Some(node),
            })
        } else {
            None
        }
    }
}

/// Iterator for a node's dominators.
pub struct DominatorsIter<'a, N>
    where N: 'a + Copy + Eq + Hash
{
    dominators: &'a Dominators<N>,
    node: Option<N>,
}

impl<'a, N> Iterator for DominatorsIter<'a, N>
    where N: 'a + Copy + Eq + Hash
{
    type Item = N;

    fn next(&mut self) -> Option<Self::Item> {
        let next = self.node.take();
        if let Some(next) = next {
            self.node = self.dominators.immediate_dominator(next);
        }
        next
    }
}

/// The undefined dominator sentinel, for when we have not yet discovered a
/// node's dominator.
const UNDEFINED: usize = ::std::usize::MAX;

/// This is an implementation of the engineered ["Simple, Fast Dominance
/// Algorithm"][0] discovered by Cooper et al.
///
/// This algorithm is **O(|V|²)**, and therefore has slower theoretical running time
/// than the Lenguaer-Tarjan algorithm (which is **O(|E| log |V|)**. However,
/// Cooper et al found it to be faster in practice on control flow graphs of up
/// to ~30,000 vertices.
///
/// [0]: http://www.cs.rice.edu/~keith/EMBED/dom.pdf
pub fn simple_fast<G>(graph: G, root: G::NodeId) -> Dominators<G::NodeId>
    where G: IntoNeighbors + Visitable,
          <G as GraphBase>::NodeId: Eq + Hash
{
    let (post_order, predecessor_sets) = simple_fast_post_order(graph, root);
    let length = post_order.len();
    debug_assert!(length > 0);
    debug_assert!(post_order.last() == Some(&root));

    // From here on out we use indices into `post_order` instead of actual
    // `NodeId`s wherever possible. This greatly improves the performance of
    // this implementation, but we have to pay a little bit of upfront cost to
    // convert our data structures to play along first.

    // Maps a node to its index into `post_order`.
    let node_to_post_order_idx: HashMap<_, _> = post_order.iter()
        .enumerate()
        .map(|(idx, &node)| (node, idx))
        .collect();

    // Maps a node's `post_order` index to its set of predecessors's indices
    // into `post_order` (as a vec).
    let idx_to_predecessor_vec =
        predecessor_sets_to_idx_vecs(&post_order, &node_to_post_order_idx, predecessor_sets);

    let mut dominators = vec![UNDEFINED; length];
    dominators[length - 1] = length - 1;

    let mut changed = true;
    while changed {
        changed = false;

        // Iterate in reverse post order, skipping the root.

        for idx in (0..length - 1).rev() {
            debug_assert!(post_order[idx] != root);

            // Take the intersection of every predecessor's dominator set; that
            // is the current best guess at the immediate dominator for this
            // node.

            let new_idom_idx = {
                let mut predecessors =
                    idx_to_predecessor_vec[idx].iter().filter(|&&p| dominators[p] != UNDEFINED);
                let new_idom_idx = predecessors.next()
                    .expect("Because the root is initialized to dominate itself, and is the \
                             first node in every path, there must exist a predecessor to this \
                             node that also has a dominator");
                predecessors.fold(*new_idom_idx, |new_idom_idx, &predecessor_idx| {
                    intersect(&dominators, new_idom_idx, predecessor_idx)
                })
            };

            debug_assert!(new_idom_idx < length);

            if new_idom_idx != dominators[idx] {
                dominators[idx] = new_idom_idx;
                changed = true;
            }
        }
    }

    // All done! Translate the indices back into proper `G::NodeId`s.

    debug_assert!(!dominators.iter().any(|&dom| dom == UNDEFINED));

    Dominators {
        root: root,
        dominators: dominators.into_iter()
            .enumerate()
            .map(|(idx, dom_idx)| (post_order[idx], post_order[dom_idx]))
            .collect(),
    }
}

fn intersect(dominators: &[usize], mut finger1: usize, mut finger2: usize) -> usize {
    while finger1 != finger2 {
        if finger1 < finger2 {
            finger1 = dominators[finger1];
        } else if finger2 < finger1 {
            finger2 = dominators[finger2];
        }
    }
    finger1
}

fn predecessor_sets_to_idx_vecs<N>(post_order: &[N],
                                   node_to_post_order_idx: &HashMap<N, usize>,
                                   mut predecessor_sets: HashMap<N, HashSet<N>>)
                                   -> Vec<Vec<usize>>
    where N: Copy + Eq + Hash
{
    post_order.iter()
        .map(|node| {
            predecessor_sets.remove(node)
                .map(|predecessors| {
                    predecessors.into_iter()
                        .map(|p| *node_to_post_order_idx.get(&p).unwrap())
                        .collect()
                })
                .unwrap_or(vec![])
        })
        .collect()
}

fn simple_fast_post_order<G>(graph: G,
                             root: G::NodeId)
                             -> (Vec<G::NodeId>, HashMap<G::NodeId, HashSet<G::NodeId>>)
    where G: IntoNeighbors + Visitable,
          <G as GraphBase>::NodeId: Eq + Hash
{
    let mut post_order = vec![];
    let mut predecessor_sets = HashMap::new();

    for node in DfsPostOrder::new(graph, root).iter(graph) {
        post_order.push(node);

        for successor in graph.neighbors(node) {
            let mut predecessors = predecessor_sets.entry(successor)
                .or_insert_with(HashSet::new);
            predecessors.insert(node);
        }
    }

    (post_order, predecessor_sets)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_iter_dominators() {
        let doms: Dominators<u32> = Dominators {
            root: 0,
            dominators: [(2, 1), (1, 0), (0, 0)]
                .iter()
                .cloned()
                .collect(),
        };

        let all_doms: Vec<_> = doms.dominators(2).unwrap().collect();
        assert_eq!(vec![2, 1, 0], all_doms);

        assert_eq!(None::<()>, doms.dominators(99).map(|_| unreachable!()));

        let strict_doms: Vec<_> = doms.strict_dominators(2).unwrap().collect();
        assert_eq!(vec![1, 0], strict_doms);

        assert_eq!(None::<()>, doms.strict_dominators(99).map(|_| unreachable!()));
    }
}
<commit_msg>MAINT: Fix unused mut warning<commit_after>//! Compute dominators of a control-flow graph.
//!
//! # The Dominance Relation
//!
//! In a directed graph with a root node **R**, a node **A** is said to *dominate* a
//! node **B** iff every path from **R** to **B** contains **A**.
//!
//! The node **A** is said to *strictly dominate* the node **B** iff **A** dominates
//! **B** and **A ≠ B**.
//!
//! The node **A** is said to be the *immediate dominator* of a node **B** iff it
//! strictly dominates **B** and there does not exist any node **C** where **A**
//! dominates **C** and **C** dominates **B**.

use std::collections::{HashMap, HashSet};
use std::hash::Hash;

use visit::{DfsPostOrder, GraphBase, IntoNeighbors, Visitable, Walker};

/// The dominance relation for some graph and root.
#[derive(Debug, Clone)]
pub struct Dominators<N>
    where N: Copy + Eq + Hash
{
    root: N,
    dominators: HashMap<N, N>,
}

impl<N> Dominators<N>
    where N: Copy + Eq + Hash
{
    /// Get the root node used to construct these dominance relations.
    pub fn root(&self) -> N {
        self.root
    }

    /// Get the immediate dominator of the given node.
    ///
    /// Returns `None` for any node that is not reachable from the root, and for
    /// the root itself.
    pub fn immediate_dominator(&self, node: N) -> Option<N> {
        if node == self.root {
            None
        } else {
            self.dominators.get(&node).cloned()
        }
    }

    /// Iterate over the given node's that strict dominators.
    ///
    /// If the given node is not reachable from the root, then `None` is
    /// returned.
    pub fn strict_dominators(&self, node: N) -> Option<DominatorsIter<N>> {
        if self.dominators.contains_key(&node) {
            Some(DominatorsIter {
                dominators: self,
                node: self.immediate_dominator(node),
            })
        } else {
            None
        }
    }

    /// Iterate over all of the given node's dominators (including the given
    /// node itself).
    ///
    /// If the given node is not reachable from the root, then `None` is
    /// returned.
    pub fn dominators(&self, node: N) -> Option<DominatorsIter<N>> {
        if self.dominators.contains_key(&node) {
            Some(DominatorsIter {
                dominators: self,
                node: Some(node),
            })
        } else {
            None
        }
    }
}

/// Iterator for a node's dominators.
pub struct DominatorsIter<'a, N>
    where N: 'a + Copy + Eq + Hash
{
    dominators: &'a Dominators<N>,
    node: Option<N>,
}

impl<'a, N> Iterator for DominatorsIter<'a, N>
    where N: 'a + Copy + Eq + Hash
{
    type Item = N;

    fn next(&mut self) -> Option<Self::Item> {
        let next = self.node.take();
        if let Some(next) = next {
            self.node = self.dominators.immediate_dominator(next);
        }
        next
    }
}

/// The undefined dominator sentinel, for when we have not yet discovered a
/// node's dominator.
const UNDEFINED: usize = ::std::usize::MAX;

/// This is an implementation of the engineered ["Simple, Fast Dominance
/// Algorithm"][0] discovered by Cooper et al.
///
/// This algorithm is **O(|V|²)**, and therefore has slower theoretical running time
/// than the Lenguaer-Tarjan algorithm (which is **O(|E| log |V|)**. However,
/// Cooper et al found it to be faster in practice on control flow graphs of up
/// to ~30,000 vertices.
///
/// [0]: http://www.cs.rice.edu/~keith/EMBED/dom.pdf
pub fn simple_fast<G>(graph: G, root: G::NodeId) -> Dominators<G::NodeId>
    where G: IntoNeighbors + Visitable,
          <G as GraphBase>::NodeId: Eq + Hash
{
    let (post_order, predecessor_sets) = simple_fast_post_order(graph, root);
    let length = post_order.len();
    debug_assert!(length > 0);
    debug_assert!(post_order.last() == Some(&root));

    // From here on out we use indices into `post_order` instead of actual
    // `NodeId`s wherever possible. This greatly improves the performance of
    // this implementation, but we have to pay a little bit of upfront cost to
    // convert our data structures to play along first.

    // Maps a node to its index into `post_order`.
    let node_to_post_order_idx: HashMap<_, _> = post_order.iter()
        .enumerate()
        .map(|(idx, &node)| (node, idx))
        .collect();

    // Maps a node's `post_order` index to its set of predecessors's indices
    // into `post_order` (as a vec).
    let idx_to_predecessor_vec =
        predecessor_sets_to_idx_vecs(&post_order, &node_to_post_order_idx, predecessor_sets);

    let mut dominators = vec![UNDEFINED; length];
    dominators[length - 1] = length - 1;

    let mut changed = true;
    while changed {
        changed = false;

        // Iterate in reverse post order, skipping the root.

        for idx in (0..length - 1).rev() {
            debug_assert!(post_order[idx] != root);

            // Take the intersection of every predecessor's dominator set; that
            // is the current best guess at the immediate dominator for this
            // node.

            let new_idom_idx = {
                let mut predecessors =
                    idx_to_predecessor_vec[idx].iter().filter(|&&p| dominators[p] != UNDEFINED);
                let new_idom_idx = predecessors.next()
                    .expect("Because the root is initialized to dominate itself, and is the \
                             first node in every path, there must exist a predecessor to this \
                             node that also has a dominator");
                predecessors.fold(*new_idom_idx, |new_idom_idx, &predecessor_idx| {
                    intersect(&dominators, new_idom_idx, predecessor_idx)
                })
            };

            debug_assert!(new_idom_idx < length);

            if new_idom_idx != dominators[idx] {
                dominators[idx] = new_idom_idx;
                changed = true;
            }
        }
    }

    // All done! Translate the indices back into proper `G::NodeId`s.

    debug_assert!(!dominators.iter().any(|&dom| dom == UNDEFINED));

    Dominators {
        root: root,
        dominators: dominators.into_iter()
            .enumerate()
            .map(|(idx, dom_idx)| (post_order[idx], post_order[dom_idx]))
            .collect(),
    }
}

fn intersect(dominators: &[usize], mut finger1: usize, mut finger2: usize) -> usize {
    while finger1 != finger2 {
        if finger1 < finger2 {
            finger1 = dominators[finger1];
        } else if finger2 < finger1 {
            finger2 = dominators[finger2];
        }
    }
    finger1
}

fn predecessor_sets_to_idx_vecs<N>(post_order: &[N],
                                   node_to_post_order_idx: &HashMap<N, usize>,
                                   mut predecessor_sets: HashMap<N, HashSet<N>>)
                                   -> Vec<Vec<usize>>
    where N: Copy + Eq + Hash
{
    post_order.iter()
        .map(|node| {
            predecessor_sets.remove(node)
                .map(|predecessors| {
                    predecessors.into_iter()
                        .map(|p| *node_to_post_order_idx.get(&p).unwrap())
                        .collect()
                })
                .unwrap_or(vec![])
        })
        .collect()
}

fn simple_fast_post_order<G>(graph: G,
                             root: G::NodeId)
                             -> (Vec<G::NodeId>, HashMap<G::NodeId, HashSet<G::NodeId>>)
    where G: IntoNeighbors + Visitable,
          <G as GraphBase>::NodeId: Eq + Hash
{
    let mut post_order = vec![];
    let mut predecessor_sets = HashMap::new();

    for node in DfsPostOrder::new(graph, root).iter(graph) {
        post_order.push(node);

        for successor in graph.neighbors(node) {
            predecessor_sets.entry(successor)
                .or_insert_with(HashSet::new)
                .insert(node);
        }
    }

    (post_order, predecessor_sets)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_iter_dominators() {
        let doms: Dominators<u32> = Dominators {
            root: 0,
            dominators: [(2, 1), (1, 0), (0, 0)]
                .iter()
                .cloned()
                .collect(),
        };

        let all_doms: Vec<_> = doms.dominators(2).unwrap().collect();
        assert_eq!(vec![2, 1, 0], all_doms);

        assert_eq!(None::<()>, doms.dominators(99).map(|_| unreachable!()));

        let strict_doms: Vec<_> = doms.strict_dominators(2).unwrap().collect();
        assert_eq!(vec![1, 0], strict_doms);

        assert_eq!(None::<()>, doms.strict_dominators(99).map(|_| unreachable!()));
    }
}
<|endoftext|>


<commit_before>extern crate tcod;
use tcod::RootConsole;

const SCREEN_WIDTH: i32 = 80;
const SCREEN_HIGHT: i32 = 50;

fn main() {
    let mut root = RootConsole::initializer()
                    .size(SCREEN_WIDTH, SCREEN_HIGHT)
                    .title("Roguelike Tutorial")
                    .fullscreen(false)
                    .init();

    while !root.window_closed() {
        root.flush();
        let key = root.wait_for_keypress(true);
    }
}
<commit_msg>Added rendering of the @ symbol.<commit_after>extern crate tcod;
use tcod::RootConsole;
use tcod::{Console};
use tcod::console::{FontLayout, FontType, BackgroundFlag};

const SCREEN_WIDTH: i32 = 80;
const SCREEN_HIGHT: i32 = 50;

fn main() {
    let mut root = RootConsole::initializer()
                    .size(SCREEN_WIDTH, SCREEN_HIGHT)
                    .title("Roguelike Tutorial")
                    .fullscreen(false)
                    .font("arial10x10.png", FontLayout::Tcod)
                    .font_type(FontType::Greyscale)
                    .init();

    root.set_default_foreground(tcod::colors::WHITE);

    while !root.window_closed() {
        root.put_char(1, 1, '@', BackgroundFlag::None);
        root.flush();

        let key = root.wait_for_keypress(true);
    }
}
<|endoftext|>