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

//! Runs hardware devices in child processes.

use std::ffi::CString;
use std::time::Duration;

use base::error;
use base::AsRawDescriptor;
use base::RawDescriptor;
use base::Tube;
use base::TubeError;
use libc::pid_t;
use minijail::Minijail;
use remain::sorted;
use serde::Deserialize;
use serde::Serialize;
use thiserror::Error;

use crate::bus::ConfigWriteResult;
use crate::pci::CrosvmDeviceId;
use crate::pci::PciAddress;
use crate::BusAccessInfo;
use crate::BusDevice;
use crate::BusRange;
use crate::BusType;
use crate::DeviceId;

/// Errors for proxy devices.
#[sorted]
#[derive(Error, Debug)]
pub enum Error {
    #[error("Failed to fork jail process: {0}")]
    ForkingJail(minijail::Error),
    #[error("Failed to configure tube: {0}")]
    Tube(TubeError),
}

pub type Result<T> = std::result::Result<T, Error>;

const SOCKET_TIMEOUT_MS: u64 = 2000;

#[derive(Debug, Serialize, Deserialize)]
enum Command {
    Read {
        len: u32,
        info: BusAccessInfo,
    },
    Write {
        len: u32,
        info: BusAccessInfo,
        data: [u8; 8],
    },
    ReadConfig(u32),
    WriteConfig {
        reg_idx: u32,
        offset: u32,
        len: u32,
        data: [u8; 4],
    },
    ReadVirtualConfig(u32),
    WriteVirtualConfig {
        reg_idx: u32,
        value: u32,
    },
    Shutdown,
    GetRanges,
}
#[derive(Debug, Serialize, Deserialize)]
enum CommandResult {
    Ok,
    ReadResult([u8; 8]),
    ReadConfigResult(u32),
    WriteConfigResult {
        mmio_remove: Vec<BusRange>,
        mmio_add: Vec<BusRange>,
        io_remove: Vec<BusRange>,
        io_add: Vec<BusRange>,
        removed_pci_devices: Vec<PciAddress>,
    },
    ReadVirtualConfigResult(u32),
    GetRangesResult(Vec<(BusRange, BusType)>),
}

fn child_proc<D: BusDevice>(tube: Tube, device: &mut D) {
    let mut running = true;

    while running {
        let cmd = match tube.recv() {
            Ok(cmd) => cmd,
            Err(err) => {
                error!("child device process failed recv: {}", err);
                break;
            }
        };

        let res = match cmd {
            Command::Read { len, info } => {
                let mut buffer = [0u8; 8];
                device.read(info, &mut buffer[0..len as usize]);
                tube.send(&CommandResult::ReadResult(buffer))
            }
            Command::Write { len, info, data } => {
                let len = len as usize;
                device.write(info, &data[0..len]);
                // Command::Write does not have a result.
                Ok(())
            }
            Command::ReadConfig(idx) => {
                let val = device.config_register_read(idx as usize);
                tube.send(&CommandResult::ReadConfigResult(val))
            }
            Command::WriteConfig {
                reg_idx,
                offset,
                len,
                data,
            } => {
                let len = len as usize;
                let res =
                    device.config_register_write(reg_idx as usize, offset as u64, &data[0..len]);
                tube.send(&CommandResult::WriteConfigResult {
                    mmio_remove: res.mmio_remove,
                    mmio_add: res.mmio_add,
                    io_remove: res.io_remove,
                    io_add: res.io_add,
                    removed_pci_devices: res.removed_pci_devices,
                })
            }
            Command::ReadVirtualConfig(idx) => {
                let val = device.virtual_config_register_read(idx as usize);
                tube.send(&CommandResult::ReadVirtualConfigResult(val))
            }
            Command::WriteVirtualConfig { reg_idx, value } => {
                device.virtual_config_register_write(reg_idx as usize, value);
                // Command::WriteVirtualConfig does not have a result.
                Ok(())
            }
            Command::Shutdown => {
                running = false;
                tube.send(&CommandResult::Ok)
            }
            Command::GetRanges => {
                let ranges = device.get_ranges();
                tube.send(&CommandResult::GetRangesResult(ranges))
            }
        };
        if let Err(e) = res {
            error!("child device process failed send: {}", e);
        }
    }
}

/// Wraps an inner `BusDevice` that is run inside a child process via fork.
///
/// Because forks are very unfriendly to destructors and all memory mappings and file descriptors
/// are inherited, this should be used as early as possible in the main process.
pub struct ProxyDevice {
    tube: Tube,
    pid: pid_t,
    debug_label: String,
}

impl ProxyDevice {
    /// Takes the given device and isolates it into another process via fork before returning.
    ///
    /// The forked process will automatically be terminated when this is dropped, so be sure to keep
    /// a reference.
    ///
    /// # Arguments
    /// * `device` - The device to isolate to another process.
    /// * `jail` - The jail to use for isolating the given device.
    /// * `keep_rds` - File descriptors that will be kept open in the child.
    pub fn new<D: BusDevice>(
        mut device: D,
        jail: Minijail,
        mut keep_rds: Vec<RawDescriptor>,
    ) -> Result<ProxyDevice> {
        let debug_label = device.debug_label();
        let (child_tube, parent_tube) = Tube::pair().map_err(Error::Tube)?;

        keep_rds.push(child_tube.as_raw_descriptor());

        // Deduplicate the FDs since minijail expects this.
        keep_rds.sort_unstable();
        keep_rds.dedup();

        // Forking here is safe as long as the program is still single threaded.
        // We own the jail object and nobody else will try to reuse it.
        let pid = match unsafe { jail.fork(Some(&keep_rds)) }.map_err(Error::ForkingJail)? {
            0 => {
                let max_len = 15; // pthread_setname_np() limit on Linux
                let debug_label_trimmed =
                    &debug_label.as_bytes()[..std::cmp::min(max_len, debug_label.len())];
                let thread_name = CString::new(debug_label_trimmed).unwrap();
                // thread_name is a valid pointer and setting name of this thread should be safe.
                let _ =
                    unsafe { libc::pthread_setname_np(libc::pthread_self(), thread_name.as_ptr()) };
                device.on_sandboxed();
                child_proc(child_tube, &mut device);

                // We're explicitly not using std::process::exit here to avoid the cleanup of
                // stdout/stderr globals. This can cause cascading panics and SIGILL if a worker
                // thread attempts to log to stderr after at_exit handlers have been run.
                // TODO(crbug.com/992494): Remove this once device shutdown ordering is clearly
                // defined.
                //
                // exit() is trivially safe.
                // ! Never returns
                unsafe { libc::exit(0) };
            }
            p => p,
        };

        parent_tube
            .set_send_timeout(Some(Duration::from_millis(SOCKET_TIMEOUT_MS)))
            .map_err(Error::Tube)?;
        parent_tube
            .set_recv_timeout(Some(Duration::from_millis(SOCKET_TIMEOUT_MS)))
            .map_err(Error::Tube)?;
        Ok(ProxyDevice {
            tube: parent_tube,
            pid,
            debug_label,
        })
    }

    pub fn pid(&self) -> pid_t {
        self.pid
    }

    /// Send a command that does not expect a response from the child device process.
    fn send_no_result(&self, cmd: &Command) {
        let res = self.tube.send(cmd);
        if let Err(e) = res {
            error!(
                "failed write to child device process {}: {}",
                self.debug_label, e,
            );
        }
    }

    /// Send a command and read its response from the child device process.
    fn sync_send(&self, cmd: &Command) -> Option<CommandResult> {
        self.send_no_result(cmd);
        match self.tube.recv() {
            Err(e) => {
                error!(
                    "failed to read result of {:?} from child device process {}: {}",
                    cmd, self.debug_label, e,
                );
                None
            }
            Ok(r) => Some(r),
        }
    }
}

impl BusDevice for ProxyDevice {
    fn device_id(&self) -> DeviceId {
        CrosvmDeviceId::ProxyDevice.into()
    }

    fn debug_label(&self) -> String {
        self.debug_label.clone()
    }

    fn config_register_write(
        &mut self,
        reg_idx: usize,
        offset: u64,
        data: &[u8],
    ) -> ConfigWriteResult {
        let len = data.len() as u32;
        let mut buffer = [0u8; 4];
        buffer[0..data.len()].clone_from_slice(data);
        let reg_idx = reg_idx as u32;
        let offset = offset as u32;
        if let Some(CommandResult::WriteConfigResult {
            mmio_remove,
            mmio_add,
            io_remove,
            io_add,
            removed_pci_devices,
        }) = self.sync_send(&Command::WriteConfig {
            reg_idx,
            offset,
            len,
            data: buffer,
        }) {
            ConfigWriteResult {
                mmio_remove,
                mmio_add,
                io_remove,
                io_add,
                removed_pci_devices,
            }
        } else {
            Default::default()
        }
    }

    fn config_register_read(&self, reg_idx: usize) -> u32 {
        let res = self.sync_send(&Command::ReadConfig(reg_idx as u32));
        if let Some(CommandResult::ReadConfigResult(val)) = res {
            val
        } else {
            0
        }
    }

    fn virtual_config_register_write(&mut self, reg_idx: usize, value: u32) {
        let reg_idx = reg_idx as u32;
        self.send_no_result(&Command::WriteVirtualConfig { reg_idx, value });
    }

    fn virtual_config_register_read(&self, reg_idx: usize) -> u32 {
        let res = self.sync_send(&Command::ReadVirtualConfig(reg_idx as u32));
        if let Some(CommandResult::ReadVirtualConfigResult(val)) = res {
            val
        } else {
            0
        }
    }

    fn read(&mut self, info: BusAccessInfo, data: &mut [u8]) {
        let len = data.len() as u32;
        if let Some(CommandResult::ReadResult(buffer)) =
            self.sync_send(&Command::Read { len, info })
        {
            let len = data.len();
            data.clone_from_slice(&buffer[0..len]);
        }
    }

    fn write(&mut self, info: BusAccessInfo, data: &[u8]) {
        let mut buffer = [0u8; 8];
        let len = data.len() as u32;
        buffer[0..data.len()].clone_from_slice(data);
        self.send_no_result(&Command::Write {
            len,
            info,
            data: buffer,
        });
    }

    fn get_ranges(&self) -> Vec<(BusRange, BusType)> {
        if let Some(CommandResult::GetRangesResult(ranges)) = self.sync_send(&Command::GetRanges) {
            ranges
        } else {
            Default::default()
        }
    }
}

impl Drop for ProxyDevice {
    fn drop(&mut self) {
        self.sync_send(&Command::Shutdown);
    }
}

/// Note: These tests must be run with --test-threads=1 to allow minijail to fork
/// the process.
#[cfg(test)]
mod tests {
    use super::*;
    use crate::pci::PciId;

    /// A simple test echo device that outputs the same u8 that was written to it.
    struct EchoDevice {
        data: u8,
        config: u8,
    }
    impl EchoDevice {
        fn new() -> EchoDevice {
            EchoDevice { data: 0, config: 0 }
        }
    }
    impl BusDevice for EchoDevice {
        fn device_id(&self) -> DeviceId {
            PciId::new(0, 0).into()
        }

        fn debug_label(&self) -> String {
            "EchoDevice".to_owned()
        }

        fn write(&mut self, _info: BusAccessInfo, data: &[u8]) {
            assert!(data.len() == 1);
            self.data = data[0];
        }

        fn read(&mut self, _info: BusAccessInfo, data: &mut [u8]) {
            assert!(data.len() == 1);
            data[0] = self.data;
        }

        fn config_register_write(
            &mut self,
            _reg_idx: usize,
            _offset: u64,
            data: &[u8],
        ) -> ConfigWriteResult {
            let result = ConfigWriteResult {
                ..Default::default()
            };
            assert!(data.len() == 1);
            self.config = data[0];
            result
        }

        fn config_register_read(&self, _reg_idx: usize) -> u32 {
            self.config as u32
        }
    }

    fn new_proxied_echo_device() -> ProxyDevice {
        let device = EchoDevice::new();
        let keep_fds: Vec<RawDescriptor> = Vec::new();
        let minijail = Minijail::new().unwrap();
        ProxyDevice::new(device, minijail, keep_fds).unwrap()
    }

    // TODO(b/173833661): Find a way to ensure these tests are run single-threaded.
    #[test]
    #[ignore]
    fn test_debug_label() {
        let proxy_device = new_proxied_echo_device();
        assert_eq!(proxy_device.debug_label(), "EchoDevice");
    }

    #[test]
    #[ignore]
    fn test_proxied_read_write() {
        let mut proxy_device = new_proxied_echo_device();
        let address = BusAccessInfo {
            offset: 0,
            address: 0,
            id: 0,
        };
        proxy_device.write(address, &[42]);
        let mut read_buffer = [0];
        proxy_device.read(address, &mut read_buffer);
        assert_eq!(read_buffer, [42]);
    }

    #[test]
    #[ignore]
    fn test_proxied_config() {
        let mut proxy_device = new_proxied_echo_device();
        proxy_device.config_register_write(0, 0, &[42]);
        assert_eq!(proxy_device.config_register_read(0), 42);
    }
}
<commit_msg>devices: proxy: ensure device drop gets called<commit_after>// Copyright 2017 The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Runs hardware devices in child processes.

use std::ffi::CString;
use std::time::Duration;

use base::error;
use base::AsRawDescriptor;
use base::RawDescriptor;
use base::Tube;
use base::TubeError;
use libc::pid_t;
use minijail::Minijail;
use remain::sorted;
use serde::Deserialize;
use serde::Serialize;
use thiserror::Error;

use crate::bus::ConfigWriteResult;
use crate::pci::CrosvmDeviceId;
use crate::pci::PciAddress;
use crate::BusAccessInfo;
use crate::BusDevice;
use crate::BusRange;
use crate::BusType;
use crate::DeviceId;

/// Errors for proxy devices.
#[sorted]
#[derive(Error, Debug)]
pub enum Error {
    #[error("Failed to fork jail process: {0}")]
    ForkingJail(minijail::Error),
    #[error("Failed to configure tube: {0}")]
    Tube(TubeError),
}

pub type Result<T> = std::result::Result<T, Error>;

const SOCKET_TIMEOUT_MS: u64 = 2000;

#[derive(Debug, Serialize, Deserialize)]
enum Command {
    Read {
        len: u32,
        info: BusAccessInfo,
    },
    Write {
        len: u32,
        info: BusAccessInfo,
        data: [u8; 8],
    },
    ReadConfig(u32),
    WriteConfig {
        reg_idx: u32,
        offset: u32,
        len: u32,
        data: [u8; 4],
    },
    ReadVirtualConfig(u32),
    WriteVirtualConfig {
        reg_idx: u32,
        value: u32,
    },
    Shutdown,
    GetRanges,
}
#[derive(Debug, Serialize, Deserialize)]
enum CommandResult {
    Ok,
    ReadResult([u8; 8]),
    ReadConfigResult(u32),
    WriteConfigResult {
        mmio_remove: Vec<BusRange>,
        mmio_add: Vec<BusRange>,
        io_remove: Vec<BusRange>,
        io_add: Vec<BusRange>,
        removed_pci_devices: Vec<PciAddress>,
    },
    ReadVirtualConfigResult(u32),
    GetRangesResult(Vec<(BusRange, BusType)>),
}

fn child_proc<D: BusDevice>(tube: Tube, device: &mut D) {
    let mut running = true;

    while running {
        let cmd = match tube.recv() {
            Ok(cmd) => cmd,
            Err(err) => {
                error!("child device process failed recv: {}", err);
                break;
            }
        };

        let res = match cmd {
            Command::Read { len, info } => {
                let mut buffer = [0u8; 8];
                device.read(info, &mut buffer[0..len as usize]);
                tube.send(&CommandResult::ReadResult(buffer))
            }
            Command::Write { len, info, data } => {
                let len = len as usize;
                device.write(info, &data[0..len]);
                // Command::Write does not have a result.
                Ok(())
            }
            Command::ReadConfig(idx) => {
                let val = device.config_register_read(idx as usize);
                tube.send(&CommandResult::ReadConfigResult(val))
            }
            Command::WriteConfig {
                reg_idx,
                offset,
                len,
                data,
            } => {
                let len = len as usize;
                let res =
                    device.config_register_write(reg_idx as usize, offset as u64, &data[0..len]);
                tube.send(&CommandResult::WriteConfigResult {
                    mmio_remove: res.mmio_remove,
                    mmio_add: res.mmio_add,
                    io_remove: res.io_remove,
                    io_add: res.io_add,
                    removed_pci_devices: res.removed_pci_devices,
                })
            }
            Command::ReadVirtualConfig(idx) => {
                let val = device.virtual_config_register_read(idx as usize);
                tube.send(&CommandResult::ReadVirtualConfigResult(val))
            }
            Command::WriteVirtualConfig { reg_idx, value } => {
                device.virtual_config_register_write(reg_idx as usize, value);
                // Command::WriteVirtualConfig does not have a result.
                Ok(())
            }
            Command::Shutdown => {
                running = false;
                tube.send(&CommandResult::Ok)
            }
            Command::GetRanges => {
                let ranges = device.get_ranges();
                tube.send(&CommandResult::GetRangesResult(ranges))
            }
        };
        if let Err(e) = res {
            error!("child device process failed send: {}", e);
        }
    }
}

/// Wraps an inner `BusDevice` that is run inside a child process via fork.
///
/// Because forks are very unfriendly to destructors and all memory mappings and file descriptors
/// are inherited, this should be used as early as possible in the main process.
pub struct ProxyDevice {
    tube: Tube,
    pid: pid_t,
    debug_label: String,
}

impl ProxyDevice {
    /// Takes the given device and isolates it into another process via fork before returning.
    ///
    /// The forked process will automatically be terminated when this is dropped, so be sure to keep
    /// a reference.
    ///
    /// # Arguments
    /// * `device` - The device to isolate to another process.
    /// * `jail` - The jail to use for isolating the given device.
    /// * `keep_rds` - File descriptors that will be kept open in the child.
    pub fn new<D: BusDevice>(
        mut device: D,
        jail: Minijail,
        mut keep_rds: Vec<RawDescriptor>,
    ) -> Result<ProxyDevice> {
        let debug_label = device.debug_label();
        let (child_tube, parent_tube) = Tube::pair().map_err(Error::Tube)?;

        keep_rds.push(child_tube.as_raw_descriptor());

        // Deduplicate the FDs since minijail expects this.
        keep_rds.sort_unstable();
        keep_rds.dedup();

        // Forking here is safe as long as the program is still single threaded.
        // We own the jail object and nobody else will try to reuse it.
        let pid = match unsafe { jail.fork(Some(&keep_rds)) }.map_err(Error::ForkingJail)? {
            0 => {
                let max_len = 15; // pthread_setname_np() limit on Linux
                let debug_label_trimmed =
                    &debug_label.as_bytes()[..std::cmp::min(max_len, debug_label.len())];
                let thread_name = CString::new(debug_label_trimmed).unwrap();
                // thread_name is a valid pointer and setting name of this thread should be safe.
                let _ =
                    unsafe { libc::pthread_setname_np(libc::pthread_self(), thread_name.as_ptr()) };
                device.on_sandboxed();
                child_proc(child_tube, &mut device);

                // Explicitly drop the device so that its Drop implementation has a chance to run
                // before the call to `libc::exit()`.
                std::mem::drop(device);

                // We're explicitly not using std::process::exit here to avoid the cleanup of
                // stdout/stderr globals. This can cause cascading panics and SIGILL if a worker
                // thread attempts to log to stderr after at_exit handlers have been run.
                // TODO(crbug.com/992494): Remove this once device shutdown ordering is clearly
                // defined.
                //
                // exit() is trivially safe.
                // ! Never returns
                unsafe { libc::exit(0) };
            }
            p => p,
        };

        parent_tube
            .set_send_timeout(Some(Duration::from_millis(SOCKET_TIMEOUT_MS)))
            .map_err(Error::Tube)?;
        parent_tube
            .set_recv_timeout(Some(Duration::from_millis(SOCKET_TIMEOUT_MS)))
            .map_err(Error::Tube)?;
        Ok(ProxyDevice {
            tube: parent_tube,
            pid,
            debug_label,
        })
    }

    pub fn pid(&self) -> pid_t {
        self.pid
    }

    /// Send a command that does not expect a response from the child device process.
    fn send_no_result(&self, cmd: &Command) {
        let res = self.tube.send(cmd);
        if let Err(e) = res {
            error!(
                "failed write to child device process {}: {}",
                self.debug_label, e,
            );
        }
    }

    /// Send a command and read its response from the child device process.
    fn sync_send(&self, cmd: &Command) -> Option<CommandResult> {
        self.send_no_result(cmd);
        match self.tube.recv() {
            Err(e) => {
                error!(
                    "failed to read result of {:?} from child device process {}: {}",
                    cmd, self.debug_label, e,
                );
                None
            }
            Ok(r) => Some(r),
        }
    }
}

impl BusDevice for ProxyDevice {
    fn device_id(&self) -> DeviceId {
        CrosvmDeviceId::ProxyDevice.into()
    }

    fn debug_label(&self) -> String {
        self.debug_label.clone()
    }

    fn config_register_write(
        &mut self,
        reg_idx: usize,
        offset: u64,
        data: &[u8],
    ) -> ConfigWriteResult {
        let len = data.len() as u32;
        let mut buffer = [0u8; 4];
        buffer[0..data.len()].clone_from_slice(data);
        let reg_idx = reg_idx as u32;
        let offset = offset as u32;
        if let Some(CommandResult::WriteConfigResult {
            mmio_remove,
            mmio_add,
            io_remove,
            io_add,
            removed_pci_devices,
        }) = self.sync_send(&Command::WriteConfig {
            reg_idx,
            offset,
            len,
            data: buffer,
        }) {
            ConfigWriteResult {
                mmio_remove,
                mmio_add,
                io_remove,
                io_add,
                removed_pci_devices,
            }
        } else {
            Default::default()
        }
    }

    fn config_register_read(&self, reg_idx: usize) -> u32 {
        let res = self.sync_send(&Command::ReadConfig(reg_idx as u32));
        if let Some(CommandResult::ReadConfigResult(val)) = res {
            val
        } else {
            0
        }
    }

    fn virtual_config_register_write(&mut self, reg_idx: usize, value: u32) {
        let reg_idx = reg_idx as u32;
        self.send_no_result(&Command::WriteVirtualConfig { reg_idx, value });
    }

    fn virtual_config_register_read(&self, reg_idx: usize) -> u32 {
        let res = self.sync_send(&Command::ReadVirtualConfig(reg_idx as u32));
        if let Some(CommandResult::ReadVirtualConfigResult(val)) = res {
            val
        } else {
            0
        }
    }

    fn read(&mut self, info: BusAccessInfo, data: &mut [u8]) {
        let len = data.len() as u32;
        if let Some(CommandResult::ReadResult(buffer)) =
            self.sync_send(&Command::Read { len, info })
        {
            let len = data.len();
            data.clone_from_slice(&buffer[0..len]);
        }
    }

    fn write(&mut self, info: BusAccessInfo, data: &[u8]) {
        let mut buffer = [0u8; 8];
        let len = data.len() as u32;
        buffer[0..data.len()].clone_from_slice(data);
        self.send_no_result(&Command::Write {
            len,
            info,
            data: buffer,
        });
    }

    fn get_ranges(&self) -> Vec<(BusRange, BusType)> {
        if let Some(CommandResult::GetRangesResult(ranges)) = self.sync_send(&Command::GetRanges) {
            ranges
        } else {
            Default::default()
        }
    }
}

impl Drop for ProxyDevice {
    fn drop(&mut self) {
        self.sync_send(&Command::Shutdown);
    }
}

/// Note: These tests must be run with --test-threads=1 to allow minijail to fork
/// the process.
#[cfg(test)]
mod tests {
    use super::*;
    use crate::pci::PciId;

    /// A simple test echo device that outputs the same u8 that was written to it.
    struct EchoDevice {
        data: u8,
        config: u8,
    }
    impl EchoDevice {
        fn new() -> EchoDevice {
            EchoDevice { data: 0, config: 0 }
        }
    }
    impl BusDevice for EchoDevice {
        fn device_id(&self) -> DeviceId {
            PciId::new(0, 0).into()
        }

        fn debug_label(&self) -> String {
            "EchoDevice".to_owned()
        }

        fn write(&mut self, _info: BusAccessInfo, data: &[u8]) {
            assert!(data.len() == 1);
            self.data = data[0];
        }

        fn read(&mut self, _info: BusAccessInfo, data: &mut [u8]) {
            assert!(data.len() == 1);
            data[0] = self.data;
        }

        fn config_register_write(
            &mut self,
            _reg_idx: usize,
            _offset: u64,
            data: &[u8],
        ) -> ConfigWriteResult {
            let result = ConfigWriteResult {
                ..Default::default()
            };
            assert!(data.len() == 1);
            self.config = data[0];
            result
        }

        fn config_register_read(&self, _reg_idx: usize) -> u32 {
            self.config as u32
        }
    }

    fn new_proxied_echo_device() -> ProxyDevice {
        let device = EchoDevice::new();
        let keep_fds: Vec<RawDescriptor> = Vec::new();
        let minijail = Minijail::new().unwrap();
        ProxyDevice::new(device, minijail, keep_fds).unwrap()
    }

    // TODO(b/173833661): Find a way to ensure these tests are run single-threaded.
    #[test]
    #[ignore]
    fn test_debug_label() {
        let proxy_device = new_proxied_echo_device();
        assert_eq!(proxy_device.debug_label(), "EchoDevice");
    }

    #[test]
    #[ignore]
    fn test_proxied_read_write() {
        let mut proxy_device = new_proxied_echo_device();
        let address = BusAccessInfo {
            offset: 0,
            address: 0,
            id: 0,
        };
        proxy_device.write(address, &[42]);
        let mut read_buffer = [0];
        proxy_device.read(address, &mut read_buffer);
        assert_eq!(read_buffer, [42]);
    }

    #[test]
    #[ignore]
    fn test_proxied_config() {
        let mut proxy_device = new_proxied_echo_device();
        proxy_device.config_register_write(0, 0, &[42]);
        assert_eq!(proxy_device.config_register_read(0), 42);
    }
}
<|endoftext|>


<commit_before>use crate::{
    prelude::{Properties, SignedBlock},
    serializer::*,
};
use std::convert::{TryFrom, TryInto};
use std::io::{self, Cursor, Error};

#[repr(u8)]
pub enum MsgType {
    Error = 0,
    GetProperties = 1,
    GetBlock = 2,
}

pub enum MsgRequest {
    GetProperties,
    GetBlock(u64), // height
}

impl MsgRequest {
    pub fn serialize(self) -> Vec<u8> {
        match self {
            MsgRequest::GetProperties => vec![MsgType::GetProperties as u8],
            MsgRequest::GetBlock(height) => {
                let mut buf = Vec::with_capacity(9);
                buf.push(MsgType::GetBlock as u8);
                buf.push_u64(height);
                buf
            }
        }
    }

    pub fn deserialize(cursor: &mut Cursor<&[u8]>) -> io::Result<Self> {
        let tag = cursor.take_u8()?;
        match tag {
            t if t == MsgType::GetProperties as u8 => Ok(MsgRequest::GetProperties),
            t if t == MsgType::GetBlock as u8 => {
                let height = cursor.take_u64()?;
                Ok(MsgRequest::GetBlock(height))
            }
            _ => Err(Error::new(
                io::ErrorKind::InvalidData,
                "invalid msg request",
            )),
        }
    }
}

#[derive(Copy, Clone, Debug)]
#[repr(u16)]
pub enum ErrorKind {
    UnknownError = 0,
    InvalidHeight = 1,
}

impl TryFrom<u16> for ErrorKind {
    type Error = Error;

    fn try_from(value: u16) -> Result<ErrorKind, Error> {
        match value {
            e if e == ErrorKind::UnknownError as u16 => Ok(ErrorKind::UnknownError),
            e if e == ErrorKind::InvalidHeight as u16 => Ok(ErrorKind::InvalidHeight),
            _ => Err(Error::new(io::ErrorKind::InvalidData, "unknown error kind")),
        }
    }
}

#[derive(Clone, Debug)]
pub enum MsgResponse {
    Error(ErrorKind, Option<String>), // code, message
    GetProperties(Properties),
    GetBlock(SignedBlock),
}

impl MsgResponse {
    pub fn serialize(self) -> Vec<u8> {
        use std::mem;

        match self {
            MsgResponse::Error(code, msg) => match msg {
                Some(msg) => {
                    let mut buf = Vec::with_capacity(3 + msg.len());
                    buf.push(MsgType::Error as u8);
                    buf.push_u16(code as u16);
                    buf.push_bytes(msg.as_bytes());
                    buf
                }
                None => {
                    let mut buf = Vec::with_capacity(7);
                    buf.push(MsgType::Error as u8);
                    buf.push_u16(code as u16);
                    buf.push_bytes(&[]);
                    buf
                }
            },
            MsgResponse::GetProperties(props) => {
                let mut buf = Vec::with_capacity(1 + mem::size_of::<Properties>());
                buf.push(MsgType::GetProperties as u8);
                buf.push_u64(props.height);
                buf.push_balance(&props.network_fee);
                buf.push_balance(&props.token_supply);
                buf
            }
            MsgResponse::GetBlock(block) => {
                let mut buf = Vec::with_capacity(1_048_576);
                buf.push(MsgType::GetBlock as u8);
                block.encode_with_tx(&mut buf);
                buf
            }
        }
    }

    pub fn deserialize(cursor: &mut Cursor<&[u8]>) -> io::Result<Self> {
        let tag = cursor.take_u8()?;
        match tag {
            t if t == MsgType::Error as u8 => {
                let kind = cursor.take_u16()?.try_into()?;
                let msg = {
                    let buf = cursor.take_bytes()?;
                    if buf.is_empty() {
                        None
                    } else {
                        Some(String::from_utf8_lossy(&buf).into_owned())
                    }
                };
                Ok(MsgResponse::Error(kind, msg))
            }
            t if t == MsgType::GetProperties as u8 => {
                let height = cursor.take_u64()?;
                let network_fee = cursor.take_balance()?;
                let token_supply = cursor.take_balance()?;
                Ok(MsgResponse::GetProperties(Properties {
                    height,
                    network_fee,
                    token_supply,
                }))
            }
            t if t == MsgType::GetBlock as u8 => {
                let block = SignedBlock::decode_with_tx(cursor)
                    .ok_or_else(|| Error::from(io::ErrorKind::UnexpectedEof))?;
                Ok(MsgResponse::GetBlock(block))
            }
            _ => Err(Error::new(
                io::ErrorKind::InvalidData,
                "invalid msg response",
            )),
        }
    }
}
<commit_msg>Derive PartialEq on net::ErrorKind<commit_after>use crate::{
    prelude::{Properties, SignedBlock},
    serializer::*,
};
use std::convert::{TryFrom, TryInto};
use std::io::{self, Cursor, Error};

#[repr(u8)]
pub enum MsgType {
    Error = 0,
    GetProperties = 1,
    GetBlock = 2,
}

pub enum MsgRequest {
    GetProperties,
    GetBlock(u64), // height
}

impl MsgRequest {
    pub fn serialize(self) -> Vec<u8> {
        match self {
            MsgRequest::GetProperties => vec![MsgType::GetProperties as u8],
            MsgRequest::GetBlock(height) => {
                let mut buf = Vec::with_capacity(9);
                buf.push(MsgType::GetBlock as u8);
                buf.push_u64(height);
                buf
            }
        }
    }

    pub fn deserialize(cursor: &mut Cursor<&[u8]>) -> io::Result<Self> {
        let tag = cursor.take_u8()?;
        match tag {
            t if t == MsgType::GetProperties as u8 => Ok(MsgRequest::GetProperties),
            t if t == MsgType::GetBlock as u8 => {
                let height = cursor.take_u64()?;
                Ok(MsgRequest::GetBlock(height))
            }
            _ => Err(Error::new(
                io::ErrorKind::InvalidData,
                "invalid msg request",
            )),
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
#[repr(u16)]
pub enum ErrorKind {
    UnknownError = 0,
    InvalidHeight = 1,
}

impl TryFrom<u16> for ErrorKind {
    type Error = Error;

    fn try_from(value: u16) -> Result<ErrorKind, Error> {
        match value {
            e if e == ErrorKind::UnknownError as u16 => Ok(ErrorKind::UnknownError),
            e if e == ErrorKind::InvalidHeight as u16 => Ok(ErrorKind::InvalidHeight),
            _ => Err(Error::new(io::ErrorKind::InvalidData, "unknown error kind")),
        }
    }
}

#[derive(Clone, Debug)]
pub enum MsgResponse {
    Error(ErrorKind, Option<String>), // code, message
    GetProperties(Properties),
    GetBlock(SignedBlock),
}

impl MsgResponse {
    pub fn serialize(self) -> Vec<u8> {
        use std::mem;

        match self {
            MsgResponse::Error(code, msg) => match msg {
                Some(msg) => {
                    let mut buf = Vec::with_capacity(3 + msg.len());
                    buf.push(MsgType::Error as u8);
                    buf.push_u16(code as u16);
                    buf.push_bytes(msg.as_bytes());
                    buf
                }
                None => {
                    let mut buf = Vec::with_capacity(7);
                    buf.push(MsgType::Error as u8);
                    buf.push_u16(code as u16);
                    buf.push_bytes(&[]);
                    buf
                }
            },
            MsgResponse::GetProperties(props) => {
                let mut buf = Vec::with_capacity(1 + mem::size_of::<Properties>());
                buf.push(MsgType::GetProperties as u8);
                buf.push_u64(props.height);
                buf.push_balance(&props.network_fee);
                buf.push_balance(&props.token_supply);
                buf
            }
            MsgResponse::GetBlock(block) => {
                let mut buf = Vec::with_capacity(1_048_576);
                buf.push(MsgType::GetBlock as u8);
                block.encode_with_tx(&mut buf);
                buf
            }
        }
    }

    pub fn deserialize(cursor: &mut Cursor<&[u8]>) -> io::Result<Self> {
        let tag = cursor.take_u8()?;
        match tag {
            t if t == MsgType::Error as u8 => {
                let kind = cursor.take_u16()?.try_into()?;
                let msg = {
                    let buf = cursor.take_bytes()?;
                    if buf.is_empty() {
                        None
                    } else {
                        Some(String::from_utf8_lossy(&buf).into_owned())
                    }
                };
                Ok(MsgResponse::Error(kind, msg))
            }
            t if t == MsgType::GetProperties as u8 => {
                let height = cursor.take_u64()?;
                let network_fee = cursor.take_balance()?;
                let token_supply = cursor.take_balance()?;
                Ok(MsgResponse::GetProperties(Properties {
                    height,
                    network_fee,
                    token_supply,
                }))
            }
            t if t == MsgType::GetBlock as u8 => {
                let block = SignedBlock::decode_with_tx(cursor)
                    .ok_or_else(|| Error::from(io::ErrorKind::UnexpectedEof))?;
                Ok(MsgResponse::GetBlock(block))
            }
            _ => Err(Error::new(
                io::ErrorKind::InvalidData,
                "invalid msg response",
            )),
        }
    }
}
<|endoftext|>


<commit_before>mod terminal;
mod ansi;

use std::io::Write;
use self::terminal::{Terminal, Event, Key};

pub struct Interface {
    choices: Vec<String>,
    terminal: Terminal
}

impl Interface {
    /// Creates a new Interface with the provided input choices.
    pub fn with_choices(choices: Vec<String>) -> Interface {
        let mut term = Terminal::from("/dev/tty").unwrap();
        term.set_raw_mode().unwrap();

        Interface {
            choices: choices,
            terminal: term
        }
    }

    // Starts the interface
    pub fn run(&mut self) {
        let ref mut term = self.terminal;

        write!(term, "{} choices to search through\r\n", self.choices.len());

        for event in term.events().unwrap() {
            match event {
                Ok(Event::Key(Key::Ctrl('c'))) => {
                    return;
                },

                Ok(event) => write!(term, "{:?}\r\n", event),
                Err(err) => write!(term, "{:?}\r\n", err),
            };
        }
    }
}
<commit_msg>Add search query, split out render method<commit_after>mod terminal;
mod ansi;

use std::io::{self, Write};
use self::terminal::{Terminal, Event, Key};
use self::ansi::{clear, cursor};

pub struct Interface {
    choices: Vec<String>,
    search: String,
    terminal: Terminal
}

impl Interface {
    /// Creates a new Interface with the provided input choices.
    pub fn with_choices(choices: Vec<String>) -> Interface {
        let mut term = Terminal::from("/dev/tty").unwrap();
        term.set_raw_mode().unwrap();

        Interface {
            choices: choices,
            search: String::new(),
            terminal: term
        }
    }

    // Starts the interface
    pub fn run(&mut self) {
        self.render().expect("Unable to render");

        for event in self.terminal.events().unwrap() {
            match event {
                Ok(Event::Key(Key::Ctrl('c'))) => {
                    return;
                },

                Ok(Event::Key(Key::Char(ch))) => {
                    self.search.push(ch);
                    self.render().expect("Unable to render");
                },

                Ok(Event::Key(Key::Backspace)) => {
                    self.search.pop();
                    self.render().expect("Unable to render");
                },

                Ok(_) => {}
                Err(err) => write!(self.terminal, "{:?}", err).unwrap(),
            };
        }
    }

    fn render(&mut self) -> io::Result<()> {
        let ref mut term = self.terminal;
        write!(term, "{}{}", clear::Line, cursor::Column(1))?;
        write!(term, "> {}", self.search)?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>use std::collections::VecDeque;
use std::fmt::Debug;
use std::time::Instant;

use actix_net::codec::Framed;
use actix_net::service::Service;

use futures::{Async, AsyncSink, Future, Poll, Sink, Stream};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_timer::Delay;

use error::{ParseError, PayloadError};
use payload::{Payload, PayloadSender, PayloadStatus, PayloadWriter};

use body::Body;
use config::ServiceConfig;
use error::DispatchError;
use request::Request;
use response::Response;

use super::codec::{Codec, InMessage, OutMessage};

const MAX_PIPELINED_MESSAGES: usize = 16;

bitflags! {
    pub struct Flags: u8 {
        const STARTED            = 0b0000_0001;
        const KEEPALIVE_ENABLED  = 0b0000_0010;
        const KEEPALIVE          = 0b0000_0100;
        const POLLED             = 0b0000_1000;
        const FLUSHED            = 0b0001_0000;
        const SHUTDOWN           = 0b0010_0000;
        const DISCONNECTED       = 0b0100_0000;
    }
}

/// Dispatcher for HTTP/1.1 protocol
pub struct Dispatcher<T, S: Service>
where
    S::Error: Debug,
{
    service: S,
    flags: Flags,
    framed: Framed<T, Codec>,
    error: Option<DispatchError<S::Error>>,
    config: ServiceConfig,

    state: State<S>,
    payload: Option<PayloadSender>,
    messages: VecDeque<Message>,

    ka_expire: Instant,
    ka_timer: Option<Delay>,
}

enum Message {
    Item(Request),
    Error(Response),
}

enum State<S: Service> {
    None,
    ServiceCall(S::Future),
    SendResponse(Option<OutMessage>),
    SendResponseWithPayload(Option<(OutMessage, Body)>),
    Payload(Body),
}

impl<S: Service> State<S> {
    fn is_empty(&self) -> bool {
        if let State::None = self {
            true
        } else {
            false
        }
    }
}

impl<T, S> Dispatcher<T, S>
where
    T: AsyncRead + AsyncWrite,
    S: Service<Request = Request, Response = Response>,
    S::Error: Debug,
{
    /// Create http/1 dispatcher.
    pub fn new(stream: T, config: ServiceConfig, service: S) -> Self {
        Dispatcher::with_timeout(stream, config, None, service)
    }

    /// Create http/1 dispatcher with slow request timeout.
    pub fn with_timeout(
        stream: T, config: ServiceConfig, timeout: Option<Delay>, service: S,
    ) -> Self {
        let keepalive = config.keep_alive_enabled();
        let flags = if keepalive {
            Flags::KEEPALIVE | Flags::KEEPALIVE_ENABLED | Flags::FLUSHED
        } else {
            Flags::FLUSHED
        };
        let framed = Framed::new(stream, Codec::new(keepalive));

        let (ka_expire, ka_timer) = if let Some(delay) = timeout {
            (delay.deadline(), Some(delay))
        } else if let Some(delay) = config.keep_alive_timer() {
            (delay.deadline(), Some(delay))
        } else {
            (config.now(), None)
        };

        Dispatcher {
            payload: None,
            state: State::None,
            error: None,
            messages: VecDeque::new(),
            service,
            flags,
            framed,
            config,
            ka_expire,
            ka_timer,
        }
    }

    fn can_read(&self) -> bool {
        if self.flags.contains(Flags::DISCONNECTED) {
            return false;
        }

        if let Some(ref info) = self.payload {
            info.need_read() == PayloadStatus::Read
        } else {
            true
        }
    }

    // if checked is set to true, delay disconnect until all tasks have finished.
    fn client_disconnected(&mut self) {
        self.flags.insert(Flags::DISCONNECTED);
        if let Some(mut payload) = self.payload.take() {
            payload.set_error(PayloadError::Incomplete);
        }
    }

    /// Flush stream
    fn poll_flush(&mut self) -> Poll<(), DispatchError<S::Error>> {
        if !self.flags.contains(Flags::FLUSHED) {
            match self.framed.poll_complete() {
                Ok(Async::NotReady) => Ok(Async::NotReady),
                Err(err) => {
                    debug!("Error sending data: {}", err);
                    Err(err.into())
                }
                Ok(Async::Ready(_)) => {
                    // if payload is not consumed we can not use connection
                    if self.payload.is_some() && self.state.is_empty() {
                        return Err(DispatchError::PayloadIsNotConsumed);
                    }
                    self.flags.insert(Flags::FLUSHED);
                    Ok(Async::Ready(()))
                }
            }
        } else {
            Ok(Async::Ready(()))
        }
    }

    fn poll_response(&mut self) -> Result<(), DispatchError<S::Error>> {
        let mut retry = self.can_read();

        // process
        loop {
            let state = match self.state {
                State::None => loop {
                    break if let Some(msg) = self.messages.pop_front() {
                        match msg {
                            Message::Item(req) => Some(self.handle_request(req)?),
                            Message::Error(res) => Some(State::SendResponse(Some(
                                OutMessage::Response(res),
                            ))),
                        }
                    } else {
                        None
                    };
                },
                State::Payload(ref mut _body) => unimplemented!(),
                State::ServiceCall(ref mut fut) => match fut
                    .poll()
                    .map_err(DispatchError::Service)?
                {
                    Async::Ready(mut res) => {
                        self.framed.get_codec_mut().prepare_te(&mut res);
                        let body = res.replace_body(Body::Empty);
                        if body.is_empty() {
                            Some(State::SendResponse(Some(OutMessage::Response(res))))
                        } else {
                            Some(State::SendResponseWithPayload(Some((
                                OutMessage::Response(res),
                                body,
                            ))))
                        }
                    }
                    Async::NotReady => None,
                },
                State::SendResponse(ref mut item) => {
                    let msg = item.take().expect("SendResponse is empty");
                    match self.framed.start_send(msg) {
                        Ok(AsyncSink::Ready) => {
                            self.flags.set(
                                Flags::KEEPALIVE,
                                self.framed.get_codec().keepalive(),
                            );
                            self.flags.remove(Flags::FLUSHED);
                            Some(State::None)
                        }
                        Ok(AsyncSink::NotReady(msg)) => {
                            *item = Some(msg);
                            return Ok(());
                        }
                        Err(err) => {
                            if let Some(mut payload) = self.payload.take() {
                                payload.set_error(PayloadError::Incomplete);
                            }
                            return Err(DispatchError::Io(err));
                        }
                    }
                }
                State::SendResponseWithPayload(ref mut item) => {
                    let (msg, body) =
                        item.take().expect("SendResponseWithPayload is empty");
                    match self.framed.start_send(msg) {
                        Ok(AsyncSink::Ready) => {
                            self.flags.set(
                                Flags::KEEPALIVE,
                                self.framed.get_codec().keepalive(),
                            );
                            self.flags.remove(Flags::FLUSHED);
                            Some(State::Payload(body))
                        }
                        Ok(AsyncSink::NotReady(msg)) => {
                            *item = Some((msg, body));
                            return Ok(());
                        }
                        Err(err) => {
                            if let Some(mut payload) = self.payload.take() {
                                payload.set_error(PayloadError::Incomplete);
                            }
                            return Err(DispatchError::Io(err));
                        }
                    }
                }
            };

            match state {
                Some(state) => self.state = state,
                None => {
                    // if read-backpressure is enabled and we consumed some data.
                    // we may read more dataand retry
                    if !retry && self.can_read() && self.poll_request()? {
                        retry = self.can_read();
                        continue;
                    }
                    break;
                }
            }
        }

        Ok(())
    }

    fn handle_request(
        &mut self, req: Request,
    ) -> Result<State<S>, DispatchError<S::Error>> {
        let mut task = self.service.call(req);
        match task.poll().map_err(DispatchError::Service)? {
            Async::Ready(mut res) => {
                self.framed.get_codec_mut().prepare_te(&mut res);
                let body = res.replace_body(Body::Empty);
                if body.is_empty() {
                    Ok(State::SendResponse(Some(OutMessage::Response(res))))
                } else {
                    Ok(State::SendResponseWithPayload(Some((
                        OutMessage::Response(res),
                        body,
                    ))))
                }
            }
            Async::NotReady => Ok(State::ServiceCall(task)),
        }
    }

    fn one_message(&mut self, msg: InMessage) -> Result<(), DispatchError<S::Error>> {
        self.flags.insert(Flags::STARTED);

        match msg {
            InMessage::Message(msg) => {
                // handle request early
                if self.state.is_empty() {
                    self.state = self.handle_request(msg)?;
                } else {
                    self.messages.push_back(Message::Item(msg));
                }
            }
            InMessage::MessageWithPayload(msg) => {
                // payload
                let (ps, pl) = Payload::new(false);
                *msg.inner.payload.borrow_mut() = Some(pl);
                self.payload = Some(ps);

                self.messages.push_back(Message::Item(msg));
            }
            InMessage::Chunk(chunk) => {
                if let Some(ref mut payload) = self.payload {
                    payload.feed_data(chunk);
                } else {
                    error!("Internal server error: unexpected payload chunk");
                    self.flags.insert(Flags::DISCONNECTED);
                    self.messages.push_back(Message::Error(
                        Response::InternalServerError().finish(),
                    ));
                    self.error = Some(DispatchError::InternalError);
                }
            }
            InMessage::Eof => {
                if let Some(mut payload) = self.payload.take() {
                    payload.feed_eof();
                } else {
                    error!("Internal server error: unexpected eof");
                    self.flags.insert(Flags::DISCONNECTED);
                    self.messages.push_back(Message::Error(
                        Response::InternalServerError().finish(),
                    ));
                    self.error = Some(DispatchError::InternalError);
                }
            }
        }

        Ok(())
    }

    pub(self) fn poll_request(&mut self) -> Result<bool, DispatchError<S::Error>> {
        let mut updated = false;

        if self.messages.len() < MAX_PIPELINED_MESSAGES {
            'outer: loop {
                match self.framed.poll() {
                    Ok(Async::Ready(Some(msg))) => {
                        updated = true;
                        self.one_message(msg)?;
                    }
                    Ok(Async::Ready(None)) => {
                        self.client_disconnected();
                        break;
                    }
                    Ok(Async::NotReady) => break,
                    Err(ParseError::Io(e)) => {
                        self.client_disconnected();
                        self.error = Some(DispatchError::Io(e));
                        break;
                    }
                    Err(e) => {
                        if let Some(mut payload) = self.payload.take() {
                            payload.set_error(PayloadError::EncodingCorrupted);
                        }

                        // Malformed requests should be responded with 400
                        self.messages
                            .push_back(Message::Error(Response::BadRequest().finish()));
                        self.flags.insert(Flags::DISCONNECTED);
                        self.error = Some(e.into());
                        break;
                    }
                }
            }
        }

        if self.ka_timer.is_some() && updated {
            if let Some(expire) = self.config.keep_alive_expire() {
                self.ka_expire = expire;
            }
        }
        Ok(updated)
    }

    /// keep-alive timer
    fn poll_keepalive(&mut self) -> Result<(), DispatchError<S::Error>> {
        if let Some(ref mut timer) = self.ka_timer {
            match timer.poll() {
                Ok(Async::Ready(_)) => {
                    if timer.deadline() >= self.ka_expire {
                        // check for any outstanding request handling
                        if self.state.is_empty() && self.messages.is_empty() {
                            // if we get timer during shutdown, just drop connection
                            if self.flags.contains(Flags::SHUTDOWN) {
                                return Err(DispatchError::DisconnectTimeout);
                            } else if !self.flags.contains(Flags::STARTED) {
                                // timeout on first request (slow request) return 408
                                trace!("Slow request timeout");
                                self.flags.insert(Flags::STARTED | Flags::DISCONNECTED);
                                self.state =
                                    State::SendResponse(Some(OutMessage::Response(
                                        Response::RequestTimeout().finish(),
                                    )));
                            } else {
                                trace!("Keep-alive timeout, close connection");
                                self.flags.insert(Flags::SHUTDOWN);

                                // start shutdown timer
                                if let Some(deadline) =
                                    self.config.client_disconnect_timer()
                                {
                                    timer.reset(deadline)
                                } else {
                                    return Ok(());
                                }
                            }
                        } else if let Some(deadline) = self.config.keep_alive_expire() {
                            timer.reset(deadline)
                        }
                    } else {
                        timer.reset(self.ka_expire)
                    }
                }
                Ok(Async::NotReady) => (),
                Err(e) => {
                    error!("Timer error {:?}", e);
                    return Err(DispatchError::Unknown);
                }
            }
        }

        Ok(())
    }
}

impl<T, S> Future for Dispatcher<T, S>
where
    T: AsyncRead + AsyncWrite,
    S: Service<Request = Request, Response = Response>,
    S::Error: Debug,
{
    type Item = ();
    type Error = DispatchError<S::Error>;

    #[inline]
    fn poll(&mut self) -> Poll<(), Self::Error> {
        if self.flags.contains(Flags::SHUTDOWN) {
            self.poll_keepalive()?;
            try_ready!(self.poll_flush());
            Ok(AsyncWrite::shutdown(self.framed.get_mut())?)
        } else {
            self.poll_keepalive()?;
            self.poll_request()?;
            self.poll_response()?;
            self.poll_flush()?;

            // keep-alive and stream errors
            if self.state.is_empty() && self.flags.contains(Flags::FLUSHED) {
                if let Some(err) = self.error.take() {
                    Err(err)
                } else if self.flags.contains(Flags::DISCONNECTED) {
                    Ok(Async::Ready(()))
                }
                // disconnect if keep-alive is not enabled
                else if self.flags.contains(Flags::STARTED) && !self
                    .flags
                    .intersects(Flags::KEEPALIVE | Flags::KEEPALIVE_ENABLED)
                {
                    self.flags.insert(Flags::SHUTDOWN);
                    self.poll()
                } else {
                    Ok(Async::NotReady)
                }
            } else {
                Ok(Async::NotReady)
            }
        }
    }
}
<commit_msg>better keep-alive handling<commit_after>use std::collections::VecDeque;
use std::fmt::Debug;
use std::time::Instant;

use actix_net::codec::Framed;
use actix_net::service::Service;

use futures::{Async, AsyncSink, Future, Poll, Sink, Stream};
use tokio_io::{AsyncRead, AsyncWrite};
use tokio_timer::Delay;

use error::{ParseError, PayloadError};
use payload::{Payload, PayloadSender, PayloadStatus, PayloadWriter};

use body::Body;
use config::ServiceConfig;
use error::DispatchError;
use request::Request;
use response::Response;

use super::codec::{Codec, InMessage, OutMessage};

const MAX_PIPELINED_MESSAGES: usize = 16;

bitflags! {
    pub struct Flags: u8 {
        const STARTED            = 0b0000_0001;
        const KEEPALIVE_ENABLED  = 0b0000_0010;
        const KEEPALIVE          = 0b0000_0100;
        const POLLED             = 0b0000_1000;
        const FLUSHED            = 0b0001_0000;
        const SHUTDOWN           = 0b0010_0000;
        const DISCONNECTED       = 0b0100_0000;
    }
}

/// Dispatcher for HTTP/1.1 protocol
pub struct Dispatcher<T, S: Service>
where
    S::Error: Debug,
{
    service: S,
    flags: Flags,
    framed: Framed<T, Codec>,
    error: Option<DispatchError<S::Error>>,
    config: ServiceConfig,

    state: State<S>,
    payload: Option<PayloadSender>,
    messages: VecDeque<Message>,

    ka_expire: Instant,
    ka_timer: Option<Delay>,
}

enum Message {
    Item(Request),
    Error(Response),
}

enum State<S: Service> {
    None,
    ServiceCall(S::Future),
    SendResponse(Option<OutMessage>),
    SendResponseWithPayload(Option<(OutMessage, Body)>),
    Payload(Body),
}

impl<S: Service> State<S> {
    fn is_empty(&self) -> bool {
        if let State::None = self {
            true
        } else {
            false
        }
    }
}

impl<T, S> Dispatcher<T, S>
where
    T: AsyncRead + AsyncWrite,
    S: Service<Request = Request, Response = Response>,
    S::Error: Debug,
{
    /// Create http/1 dispatcher.
    pub fn new(stream: T, config: ServiceConfig, service: S) -> Self {
        Dispatcher::with_timeout(stream, config, None, service)
    }

    /// Create http/1 dispatcher with slow request timeout.
    pub fn with_timeout(
        stream: T, config: ServiceConfig, timeout: Option<Delay>, service: S,
    ) -> Self {
        let keepalive = config.keep_alive_enabled();
        let flags = if keepalive {
            Flags::KEEPALIVE | Flags::KEEPALIVE_ENABLED | Flags::FLUSHED
        } else {
            Flags::FLUSHED
        };
        let framed = Framed::new(stream, Codec::new(keepalive));

        let (ka_expire, ka_timer) = if let Some(delay) = timeout {
            (delay.deadline(), Some(delay))
        } else if let Some(delay) = config.keep_alive_timer() {
            (delay.deadline(), Some(delay))
        } else {
            (config.now(), None)
        };

        Dispatcher {
            payload: None,
            state: State::None,
            error: None,
            messages: VecDeque::new(),
            service,
            flags,
            framed,
            config,
            ka_expire,
            ka_timer,
        }
    }

    fn can_read(&self) -> bool {
        if self.flags.contains(Flags::DISCONNECTED) {
            return false;
        }

        if let Some(ref info) = self.payload {
            info.need_read() == PayloadStatus::Read
        } else {
            true
        }
    }

    // if checked is set to true, delay disconnect until all tasks have finished.
    fn client_disconnected(&mut self) {
        self.flags.insert(Flags::DISCONNECTED);
        if let Some(mut payload) = self.payload.take() {
            payload.set_error(PayloadError::Incomplete);
        }
    }

    /// Flush stream
    fn poll_flush(&mut self) -> Poll<(), DispatchError<S::Error>> {
        if !self.flags.contains(Flags::FLUSHED) {
            match self.framed.poll_complete() {
                Ok(Async::NotReady) => Ok(Async::NotReady),
                Err(err) => {
                    debug!("Error sending data: {}", err);
                    Err(err.into())
                }
                Ok(Async::Ready(_)) => {
                    // if payload is not consumed we can not use connection
                    if self.payload.is_some() && self.state.is_empty() {
                        return Err(DispatchError::PayloadIsNotConsumed);
                    }
                    self.flags.insert(Flags::FLUSHED);
                    Ok(Async::Ready(()))
                }
            }
        } else {
            Ok(Async::Ready(()))
        }
    }

    fn poll_response(&mut self) -> Result<(), DispatchError<S::Error>> {
        let mut retry = self.can_read();

        // process
        loop {
            let state = match self.state {
                State::None => loop {
                    break if let Some(msg) = self.messages.pop_front() {
                        match msg {
                            Message::Item(req) => Some(self.handle_request(req)?),
                            Message::Error(res) => Some(State::SendResponse(Some(
                                OutMessage::Response(res),
                            ))),
                        }
                    } else {
                        None
                    };
                },
                State::Payload(ref mut _body) => unimplemented!(),
                State::ServiceCall(ref mut fut) => match fut
                    .poll()
                    .map_err(DispatchError::Service)?
                {
                    Async::Ready(mut res) => {
                        self.framed.get_codec_mut().prepare_te(&mut res);
                        let body = res.replace_body(Body::Empty);
                        if body.is_empty() {
                            Some(State::SendResponse(Some(OutMessage::Response(res))))
                        } else {
                            Some(State::SendResponseWithPayload(Some((
                                OutMessage::Response(res),
                                body,
                            ))))
                        }
                    }
                    Async::NotReady => None,
                },
                State::SendResponse(ref mut item) => {
                    let msg = item.take().expect("SendResponse is empty");
                    match self.framed.start_send(msg) {
                        Ok(AsyncSink::Ready) => {
                            self.flags.set(
                                Flags::KEEPALIVE,
                                self.framed.get_codec().keepalive(),
                            );
                            self.flags.remove(Flags::FLUSHED);
                            Some(State::None)
                        }
                        Ok(AsyncSink::NotReady(msg)) => {
                            *item = Some(msg);
                            return Ok(());
                        }
                        Err(err) => {
                            if let Some(mut payload) = self.payload.take() {
                                payload.set_error(PayloadError::Incomplete);
                            }
                            return Err(DispatchError::Io(err));
                        }
                    }
                }
                State::SendResponseWithPayload(ref mut item) => {
                    let (msg, body) =
                        item.take().expect("SendResponseWithPayload is empty");
                    match self.framed.start_send(msg) {
                        Ok(AsyncSink::Ready) => {
                            self.flags.set(
                                Flags::KEEPALIVE,
                                self.framed.get_codec().keepalive(),
                            );
                            self.flags.remove(Flags::FLUSHED);
                            Some(State::Payload(body))
                        }
                        Ok(AsyncSink::NotReady(msg)) => {
                            *item = Some((msg, body));
                            return Ok(());
                        }
                        Err(err) => {
                            if let Some(mut payload) = self.payload.take() {
                                payload.set_error(PayloadError::Incomplete);
                            }
                            return Err(DispatchError::Io(err));
                        }
                    }
                }
            };

            match state {
                Some(state) => self.state = state,
                None => {
                    // if read-backpressure is enabled and we consumed some data.
                    // we may read more dataand retry
                    if !retry && self.can_read() && self.poll_request()? {
                        retry = self.can_read();
                        continue;
                    }
                    break;
                }
            }
        }

        Ok(())
    }

    fn handle_request(
        &mut self, req: Request,
    ) -> Result<State<S>, DispatchError<S::Error>> {
        let mut task = self.service.call(req);
        match task.poll().map_err(DispatchError::Service)? {
            Async::Ready(mut res) => {
                self.framed.get_codec_mut().prepare_te(&mut res);
                let body = res.replace_body(Body::Empty);
                if body.is_empty() {
                    Ok(State::SendResponse(Some(OutMessage::Response(res))))
                } else {
                    Ok(State::SendResponseWithPayload(Some((
                        OutMessage::Response(res),
                        body,
                    ))))
                }
            }
            Async::NotReady => Ok(State::ServiceCall(task)),
        }
    }

    fn one_message(&mut self, msg: InMessage) -> Result<(), DispatchError<S::Error>> {
        self.flags.insert(Flags::STARTED);

        match msg {
            InMessage::Message(msg) => {
                // handle request early
                if self.state.is_empty() {
                    self.state = self.handle_request(msg)?;
                } else {
                    self.messages.push_back(Message::Item(msg));
                }
            }
            InMessage::MessageWithPayload(msg) => {
                // payload
                let (ps, pl) = Payload::new(false);
                *msg.inner.payload.borrow_mut() = Some(pl);
                self.payload = Some(ps);

                self.messages.push_back(Message::Item(msg));
            }
            InMessage::Chunk(chunk) => {
                if let Some(ref mut payload) = self.payload {
                    payload.feed_data(chunk);
                } else {
                    error!("Internal server error: unexpected payload chunk");
                    self.flags.insert(Flags::DISCONNECTED);
                    self.messages.push_back(Message::Error(
                        Response::InternalServerError().finish(),
                    ));
                    self.error = Some(DispatchError::InternalError);
                }
            }
            InMessage::Eof => {
                if let Some(mut payload) = self.payload.take() {
                    payload.feed_eof();
                } else {
                    error!("Internal server error: unexpected eof");
                    self.flags.insert(Flags::DISCONNECTED);
                    self.messages.push_back(Message::Error(
                        Response::InternalServerError().finish(),
                    ));
                    self.error = Some(DispatchError::InternalError);
                }
            }
        }

        Ok(())
    }

    pub(self) fn poll_request(&mut self) -> Result<bool, DispatchError<S::Error>> {
        let mut updated = false;

        if self.messages.len() < MAX_PIPELINED_MESSAGES {
            'outer: loop {
                match self.framed.poll() {
                    Ok(Async::Ready(Some(msg))) => {
                        updated = true;
                        self.one_message(msg)?;
                    }
                    Ok(Async::Ready(None)) => {
                        self.client_disconnected();
                        break;
                    }
                    Ok(Async::NotReady) => break,
                    Err(ParseError::Io(e)) => {
                        self.client_disconnected();
                        self.error = Some(DispatchError::Io(e));
                        break;
                    }
                    Err(e) => {
                        if let Some(mut payload) = self.payload.take() {
                            payload.set_error(PayloadError::EncodingCorrupted);
                        }

                        // Malformed requests should be responded with 400
                        self.messages
                            .push_back(Message::Error(Response::BadRequest().finish()));
                        self.flags.insert(Flags::DISCONNECTED);
                        self.error = Some(e.into());
                        break;
                    }
                }
            }
        }

        if self.ka_timer.is_some() && updated {
            if let Some(expire) = self.config.keep_alive_expire() {
                self.ka_expire = expire;
            }
        }
        Ok(updated)
    }

    /// keep-alive timer
    fn poll_keepalive(&mut self) -> Result<(), DispatchError<S::Error>> {
        if let Some(ref mut timer) = self.ka_timer {
            match timer.poll() {
                Ok(Async::Ready(_)) => {
                    // if we get timer during shutdown, just drop connection
                    if self.flags.contains(Flags::SHUTDOWN) {
                        return Err(DispatchError::DisconnectTimeout);
                    } else if timer.deadline() >= self.ka_expire {
                        // check for any outstanding response processing
                        if self.state.is_empty() {
                            if self.flags.contains(Flags::STARTED) {
                                trace!("Keep-alive timeout, close connection");
                                self.flags.insert(Flags::SHUTDOWN);

                                // start shutdown timer
                                if let Some(deadline) =
                                    self.config.client_disconnect_timer()
                                {
                                    timer.reset(deadline)
                                } else {
                                    return Ok(());
                                }
                            } else {
                                // timeout on first request (slow request) return 408
                                trace!("Slow request timeout");
                                self.flags.insert(Flags::STARTED | Flags::DISCONNECTED);
                                self.state =
                                    State::SendResponse(Some(OutMessage::Response(
                                        Response::RequestTimeout().finish(),
                                    )));
                            }
                        } else if let Some(deadline) = self.config.keep_alive_expire() {
                            timer.reset(deadline)
                        }
                    } else {
                        timer.reset(self.ka_expire)
                    }
                }
                Ok(Async::NotReady) => (),
                Err(e) => {
                    error!("Timer error {:?}", e);
                    return Err(DispatchError::Unknown);
                }
            }
        }

        Ok(())
    }
}

impl<T, S> Future for Dispatcher<T, S>
where
    T: AsyncRead + AsyncWrite,
    S: Service<Request = Request, Response = Response>,
    S::Error: Debug,
{
    type Item = ();
    type Error = DispatchError<S::Error>;

    #[inline]
    fn poll(&mut self) -> Poll<(), Self::Error> {
        if self.flags.contains(Flags::SHUTDOWN) {
            self.poll_keepalive()?;
            try_ready!(self.poll_flush());
            Ok(AsyncWrite::shutdown(self.framed.get_mut())?)
        } else {
            self.poll_keepalive()?;
            self.poll_request()?;
            self.poll_response()?;
            self.poll_flush()?;

            // keep-alive and stream errors
            if self.state.is_empty() && self.flags.contains(Flags::FLUSHED) {
                if let Some(err) = self.error.take() {
                    Err(err)
                } else if self.flags.contains(Flags::DISCONNECTED) {
                    Ok(Async::Ready(()))
                }
                // disconnect if keep-alive is not enabled
                else if self.flags.contains(Flags::STARTED) && !self
                    .flags
                    .intersects(Flags::KEEPALIVE | Flags::KEEPALIVE_ENABLED)
                {
                    self.flags.insert(Flags::SHUTDOWN);
                    self.poll()
                } else {
                    Ok(Async::NotReady)
                }
            } else {
                Ok(Async::NotReady)
            }
        }
    }
}
<|endoftext|>


<commit_before>/// Decodes an integer encoded with a given prefix size (in bits).
/// Assumes that the buffer `buf` contains the integer to be decoded,
/// with the first byte representing the octet that contains the
/// prefix.
///
/// Returns a tuple representing the decoded integer and the number
/// of bytes from the buffer that were used.
fn decode_integer(buf: &[u8], prefix_size: u8) -> (usize, usize) {
    assert!(prefix_size >= 1 && prefix_size <= 8);
    assert!(buf.len() >= 1);
    let mask = (1u8 << prefix_size) - 1;
    let mut value: usize = (buf[0] & mask) as usize;
    if value < (mask as usize) {
        return (value, 1);
    }

    // Already one byte used (the prefix)
    let mut total = 1;
    let mut m = 0;
    for &b in buf[1..].iter() {
        total += 1;
        value += ((b & 127) as usize) * (1us << m);
        m += 7;
        if b & 128 != 128 {
            break;
        }
    }

    // TODO Error situation -- what happens if we reach here w/o having
    //      `b` byte-sized...
    (value, total)
}

mod tests {
    use super::{decode_integer};

    #[test]
    fn test_decode_integer() {
        assert_eq!((10us, 1), decode_integer(&[10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 32, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 64, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10, 342, 22], 5));

        assert_eq!((127us, 2), decode_integer(&[255, 0], 7));
        assert_eq!((127us, 2), decode_integer(&[127, 0], 7));
    }

}
<commit_msg>Implement an integer encoder according to HPACK<commit_after>/// Decodes an integer encoded with a given prefix size (in bits).
/// Assumes that the buffer `buf` contains the integer to be decoded,
/// with the first byte representing the octet that contains the
/// prefix.
///
/// Returns a tuple representing the decoded integer and the number
/// of bytes from the buffer that were used.
fn decode_integer(buf: &[u8], prefix_size: u8) -> (usize, usize) {
    assert!(prefix_size >= 1 && prefix_size <= 8);
    assert!(buf.len() >= 1);
    let mask = (1u8 << prefix_size) - 1;
    let mut value: usize = (buf[0] & mask) as usize;
    if value < (mask as usize) {
        return (value, 1);
    }

    // Already one byte used (the prefix)
    let mut total = 1;
    let mut m = 0;
    for &b in buf[1..].iter() {
        total += 1;
        value += ((b & 127) as usize) * (1us << m);
        m += 7;
        if b & 128 != 128 {
            break;
        }
    }

    // TODO Error situation -- what happens if we reach here w/o having
    //      `b` byte-sized...
    (value, total)
}

/// Encode an integer to the representation defined by HPACK.
///
/// Returns a newly allocated `Vec` containing the encoded bytes.
/// Only `prefix_size` lowest-order bits of the first byte in the
/// array are guaranteed to be used.
fn encode_integer(mut value: usize, prefix_size: u8) -> Vec<u8> {
    let mask: usize = ((1u8 << prefix_size) - 1) as usize;
    if value < mask {
        // Right now, the caller would need to be the one to combine
        // the other part of the prefix byte (but would know that it's
        // safe to do so using the bit-wise or).
        return vec![value as u8];
    }

    let mut res: Vec<u8> = Vec::new();
    res.push(mask as u8);
    value -= mask;
    while value >= 128 {
        res.push(((value % 128) + 128) as u8);
        value = value / 128;
    }
    res.push(value as u8);

    res
}

mod tests {
    use super::{decode_integer};
    use super::{encode_integer};

    #[test]
    fn test_decode_integer() {
        assert_eq!((10us, 1), decode_integer(&[10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 32, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31 + 64, 154, 10], 5));
        assert_eq!((1337us, 3), decode_integer(&[31, 154, 10, 342, 22], 5));

        assert_eq!((127us, 2), decode_integer(&[255, 0], 7));
        assert_eq!((127us, 2), decode_integer(&[127, 0], 7));
    }

    #[test]
    fn test_encode_integer() {
        assert_eq!(encode_integer(10, 5), [10]);
        assert_eq!(encode_integer(1337, 5), [31, 154, 10]);
        assert_eq!(encode_integer(127, 7), [127, 0]);
    }
}
<|endoftext|>


<commit_before>use std::io::{BufRead, Error};

pub fn read_and_trim_lines<T: BufRead>(buffered_reader: T) -> Result<Vec<String>, Error> {
    buffered_reader
        .lines()
        .map(|r| Ok(crate::utils::trim_end_inplace(r?)))
        .collect()
}
<commit_msg>Implement unit tests for the trim lines fn<commit_after>use std::io::{BufRead, Error};

pub fn read_and_trim_lines<T: BufRead>(buffered_reader: T) -> Result<Vec<String>, Error> {
    buffered_reader
        .lines()
        .map(|r| Ok(crate::utils::trim_end_inplace(r?)))
        .collect()
}

// LCOV_EXCL_START
#[cfg(test)]
mod tests {
    #[cfg(test)]
    mod read_and_trim_lines_test {
        use super::super::read_and_trim_lines;
        use std::io::BufReader;

        #[test]
        fn valid_string() {
            let input = "\tfoo\nbar \n qux\t\n \t\n";

            let res = read_and_trim_lines(BufReader::new(input.as_bytes()));
            assert!(res.is_ok());

            let res_unwrapped = res.unwrap();
            assert_eq!(res_unwrapped.len(), 4);

            for (actual, expected) in res_unwrapped.iter().zip(vec!["\tfoo", "bar", " qux", ""]) {
                assert_eq!(actual, expected);
            }
        }

        #[test]
        fn invalid_utf8_string() {
            let input = b"\xe2\x28\xa1";
            let res = read_and_trim_lines(BufReader::new(input.as_ref()));
            assert!(res.is_err());
        }
    }
}
// LCOV_EXCL_STOP
<|endoftext|>


<commit_before>use Proxy;

use event_queue::{EventQueue, create_event_queue};
use generated::client::wl_display::WlDisplay;
use std::ffi::{CStr, CString, OsStr};

use std::io;
use std::os::unix::ffi::OsStrExt;

use wayland_sys::client::*;

/// Enum representing the possible reasons why connecting to the wayland server failed
#[derive(Debug)]
pub enum ConnectError {
    /// The library was compiled with the `dlopen` feature, and the `libwayland-client.so`
    /// library could not be found at runtime
    NoWaylandLib,
    /// Any needed library was found, but the listening socket of the server could not be
    /// found.
    ///
    /// Most of the time, this means that the program was not started from a wayland session.
    NoCompositorListening,
}

/// Enum representing possible errors fatal to a wayland session
///
/// These errors are fatal, so there is no way to recover the session, you
/// must create a new one (or report failure to your user). But recovering
/// this error can provide usefull debug information and/or help provide
/// a sensible error message to the user.
#[derive(Debug)]
pub enum FatalError {
    /// Session aborted after an I/O error
    Io(io::Error),
    /// Session aborted after a protocol error
    Protocol {
        ///  name of the interface of the proxy that generated this error
        interface: String,
        /// internal id of the proxy that generated this error
        proxy_id: u32,
        /// code of the error, as defined by the `Error` enum of the interface of the proxy.
        /// It can directly be fed to the `from_raw` static method of this enum.
        error_code: u32,
    },
}

/// Connect to the compositor socket
///
/// Attempt to connect to a Wayland compositor according to the environment variables.
///
/// On success, returns the display object, as well as the default event iterator associated with it.
pub fn default_connect() -> Result<(WlDisplay, EventQueue), ConnectError> {
    if !::wayland_sys::client::is_lib_available() {
        return Err(ConnectError::NoWaylandLib);
    }
    let ptr = unsafe {
        ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                      wl_display_connect,
                      ::std::ptr::null())
    };
    if ptr.is_null() {
        Err(ConnectError::NoCompositorListening)
    } else {
        let display = unsafe { WlDisplay::from_ptr_new(ptr as *mut _) };
        let eventiter = unsafe { create_event_queue(display.ptr() as *mut wl_display, None) };
        Ok((display, eventiter))
    }
}

/// Connect to the compositor socket
///
/// Attempt to connect to a Wayland compositor on a given socket name
///
/// On success, returns the display object, as well as the default event iterator associated with it.
pub fn connect_to(name: &OsStr) -> Result<(WlDisplay, EventQueue), ConnectError> {
    if !::wayland_sys::client::is_lib_available() {
        return Err(ConnectError::NoWaylandLib);
    }
    // Only possible error is interior null, and in this case, no compositor will be listening to a socket
    // with null in its name.
    let name = CString::new(name.as_bytes().to_owned())
        .map_err(|_| ConnectError::NoCompositorListening)?;
    let ptr = unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_connect, name.as_ptr()) };
    if ptr.is_null() {
        Err(ConnectError::NoCompositorListening)
    } else {
        let display = unsafe { WlDisplay::from_ptr_new(ptr as *mut _) };
        let eventiter = unsafe { create_event_queue(display.ptr() as *mut wl_display, None) };
        Ok((display, eventiter))
    }
}

impl WlDisplay {
    /// Non-blocking write to the server
    ///
    /// Will write as many pending requests as possible to the server socket. Never blocks: if not all
    /// requests coul be written, will return an io error `WouldBlock`.
    ///
    /// On success returns the number of written requests.
    pub fn flush(&self) -> io::Result<i32> {
        let ret = unsafe {
            ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                          wl_display_flush,
                          self.ptr() as *mut _)
        };
        if ret >= 0 {
            Ok(ret)
        } else {
            Err(io::Error::last_os_error())
        }
    }

    /// Create a new EventQueue
    ///
    /// No object is by default attached to it.
    pub fn create_event_queue(&self) -> EventQueue {
        let evq = unsafe {
            ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                          wl_display_create_queue,
                          self.ptr() as *mut _)
        };
        unsafe { create_event_queue(self.ptr() as *mut _, Some(evq)) }
    }

    /// Get the last error that occured on the session
    ///
    /// Such errors are *fatal*, meaning that if this function does not
    /// return `None`, your session is not usable any longer.
    ///
    /// As such, this function mostly provide diagnistics information. You can have a hint
    /// an error might have been generated if I/O methods of EventQueue start returning errors.
    pub fn last_error(&self) -> Option<FatalError> {
        let err = unsafe {
            ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                          wl_display_get_error,
                          self.ptr() as *mut _)
        };
        if err == 0 {
            None
        } else if err == ::libc::EPROTO {
            let mut interface = ::std::ptr::null_mut();
            let mut id = 0;
            let code = unsafe {
                ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                              wl_display_get_protocol_error,
                              self.ptr() as *mut _,
                              &mut interface,
                              &mut id)
            };
            let interface = if interface.is_null() {
                "<unkown interface>".to_owned()
            } else {
                unsafe { CStr::from_ptr((*interface).name) }
                    .to_string_lossy()
                    .into_owned()
            };
            Some(FatalError::Protocol {
                     interface: interface,
                     proxy_id: id,
                     error_code: code,
                 })
        } else {
            Some(FatalError::Io(io::Error::from_raw_os_error(err)))
        }
    }

    /// Get the raw File Descriptor associated with the connection
    ///
    /// This is provided to be used in conjunction with some polling mecanism,
    /// if you want to manually control the flow with something like `epoll`.
    /// In this case, you'll most likely want to use `EventQueue::prepare_read()` and
    /// `EventQueue::dispatch_pending()` rather than `EventQueue::dispatch()`.
    ///
    /// Reading or writing anything to this FD will corrupt the internal state of
    /// the lib.
    pub fn get_fd(&self) -> ::std::os::unix::io::RawFd {
        unsafe {
            ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                          wl_display_get_fd,
                          self.ptr() as *mut _)
        }
    }
}

impl Drop for WlDisplay {
    fn drop(&mut self) {
        unsafe {
            ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                          wl_display_disconnect,
                          self.ptr() as *mut _)
        }
    }
}
<commit_msg>client: display.get_fd should really be unsafe<commit_after>use Proxy;

use event_queue::{EventQueue, create_event_queue};
use generated::client::wl_display::WlDisplay;
use std::ffi::{CStr, CString, OsStr};

use std::io;
use std::os::unix::ffi::OsStrExt;

use wayland_sys::client::*;

/// Enum representing the possible reasons why connecting to the wayland server failed
#[derive(Debug)]
pub enum ConnectError {
    /// The library was compiled with the `dlopen` feature, and the `libwayland-client.so`
    /// library could not be found at runtime
    NoWaylandLib,
    /// Any needed library was found, but the listening socket of the server could not be
    /// found.
    ///
    /// Most of the time, this means that the program was not started from a wayland session.
    NoCompositorListening,
}

/// Enum representing possible errors fatal to a wayland session
///
/// These errors are fatal, so there is no way to recover the session, you
/// must create a new one (or report failure to your user). But recovering
/// this error can provide usefull debug information and/or help provide
/// a sensible error message to the user.
#[derive(Debug)]
pub enum FatalError {
    /// Session aborted after an I/O error
    Io(io::Error),
    /// Session aborted after a protocol error
    Protocol {
        ///  name of the interface of the proxy that generated this error
        interface: String,
        /// internal id of the proxy that generated this error
        proxy_id: u32,
        /// code of the error, as defined by the `Error` enum of the interface of the proxy.
        /// It can directly be fed to the `from_raw` static method of this enum.
        error_code: u32,
    },
}

/// Connect to the compositor socket
///
/// Attempt to connect to a Wayland compositor according to the environment variables.
///
/// On success, returns the display object, as well as the default event iterator associated with it.
pub fn default_connect() -> Result<(WlDisplay, EventQueue), ConnectError> {
    if !::wayland_sys::client::is_lib_available() {
        return Err(ConnectError::NoWaylandLib);
    }
    let ptr = unsafe {
        ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                      wl_display_connect,
                      ::std::ptr::null())
    };
    if ptr.is_null() {
        Err(ConnectError::NoCompositorListening)
    } else {
        let display = unsafe { WlDisplay::from_ptr_new(ptr as *mut _) };
        let eventiter = unsafe { create_event_queue(display.ptr() as *mut wl_display, None) };
        Ok((display, eventiter))
    }
}

/// Connect to the compositor socket
///
/// Attempt to connect to a Wayland compositor on a given socket name
///
/// On success, returns the display object, as well as the default event iterator associated with it.
pub fn connect_to(name: &OsStr) -> Result<(WlDisplay, EventQueue), ConnectError> {
    if !::wayland_sys::client::is_lib_available() {
        return Err(ConnectError::NoWaylandLib);
    }
    // Only possible error is interior null, and in this case, no compositor will be listening to a socket
    // with null in its name.
    let name = CString::new(name.as_bytes().to_owned())
        .map_err(|_| ConnectError::NoCompositorListening)?;
    let ptr = unsafe { ffi_dispatch!(WAYLAND_CLIENT_HANDLE, wl_display_connect, name.as_ptr()) };
    if ptr.is_null() {
        Err(ConnectError::NoCompositorListening)
    } else {
        let display = unsafe { WlDisplay::from_ptr_new(ptr as *mut _) };
        let eventiter = unsafe { create_event_queue(display.ptr() as *mut wl_display, None) };
        Ok((display, eventiter))
    }
}

impl WlDisplay {
    /// Non-blocking write to the server
    ///
    /// Will write as many pending requests as possible to the server socket. Never blocks: if not all
    /// requests coul be written, will return an io error `WouldBlock`.
    ///
    /// On success returns the number of written requests.
    pub fn flush(&self) -> io::Result<i32> {
        let ret = unsafe {
            ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                          wl_display_flush,
                          self.ptr() as *mut _)
        };
        if ret >= 0 {
            Ok(ret)
        } else {
            Err(io::Error::last_os_error())
        }
    }

    /// Create a new EventQueue
    ///
    /// No object is by default attached to it.
    pub fn create_event_queue(&self) -> EventQueue {
        let evq = unsafe {
            ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                          wl_display_create_queue,
                          self.ptr() as *mut _)
        };
        unsafe { create_event_queue(self.ptr() as *mut _, Some(evq)) }
    }

    /// Get the last error that occured on the session
    ///
    /// Such errors are *fatal*, meaning that if this function does not
    /// return `None`, your session is not usable any longer.
    ///
    /// As such, this function mostly provide diagnistics information. You can have a hint
    /// an error might have been generated if I/O methods of EventQueue start returning errors.
    pub fn last_error(&self) -> Option<FatalError> {
        let err = unsafe {
            ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                          wl_display_get_error,
                          self.ptr() as *mut _)
        };
        if err == 0 {
            None
        } else if err == ::libc::EPROTO {
            let mut interface = ::std::ptr::null_mut();
            let mut id = 0;
            let code = unsafe {
                ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                              wl_display_get_protocol_error,
                              self.ptr() as *mut _,
                              &mut interface,
                              &mut id)
            };
            let interface = if interface.is_null() {
                "<unkown interface>".to_owned()
            } else {
                unsafe { CStr::from_ptr((*interface).name) }
                    .to_string_lossy()
                    .into_owned()
            };
            Some(FatalError::Protocol {
                     interface: interface,
                     proxy_id: id,
                     error_code: code,
                 })
        } else {
            Some(FatalError::Io(io::Error::from_raw_os_error(err)))
        }
    }

    /// Get the raw File Descriptor associated with the connection
    ///
    /// This is provided to be used in conjunction with some polling mecanism,
    /// if you want to manually control the flow with something like `epoll`.
    /// In this case, you'll most likely want to use `EventQueue::prepare_read()` and
    /// `EventQueue::dispatch_pending()` rather than `EventQueue::dispatch()`.
    ///
    /// Reading or writing anything to this FD will corrupt the internal state of
    /// the lib.
    pub unsafe fn get_fd(&self) -> ::std::os::unix::io::RawFd {
        ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                      wl_display_get_fd,
                      self.ptr() as *mut _)
    }
}

impl Drop for WlDisplay {
    fn drop(&mut self) {
        unsafe {
            ffi_dispatch!(WAYLAND_CLIENT_HANDLE,
                          wl_display_disconnect,
                          self.ptr() as *mut _)
        }
    }
}
<|endoftext|>