<commit_before>use std::ffi::CString;

use c::FileType;
use c::XmpPacketInfo as PacketInfo;

use crate::error::Error;
use crate::Result;
use crate::xmp::Xmp;
use crate::xmpstring::XmpString;

bitflags! {
    /// Flag options for opening files.
    #[derive(Default)]
    pub struct OpenFlags: u32 {
        /// No open option
        const NONE = 0x0000_0000;
        /// Open for read-only access.
        const READ = 0x0000_0001;
        /// Open for reading and writing.
        const FOR_UPDATE = 0x0000_0002;
        /// Only the XMP is wanted, allows space/time optimizations.
        const ONLY_XMP = 0x0000_0004;
        /// Cache thumbnail if possible, GetThumbnail will be called.
        const CACHE_TNAIL = 0x0000_0008;
        /// Be strict about locating XMP and reconciling with other forms.
        const STRICTLY = 0x0000_0010;
        /// Require the use of a smart handler.
        const USE_SMART_HANDLER = 0x0000_0020;
        /// Force packet scanning, don't use a smart handler.
        const USE_PACKET_SCANNING = 0x0000_0040;
        /// Only packet scan files "known" to need scanning.
        const LIMITED_SCANNING = 0x0000_0080;
        /// Attempt to repair a file opened for update, default is to not open (throw an exception).
        const REPAIR_FILE = 0x0000_0100;
        /// Optimize MPEG4 to support stream when updating This can take some time
        const OPTIMIZE_FILE_LAYOUT = 0x0000_0200;
        /// Set if calling from background thread.
        const IN_BACKGROUND = 0x1000_0000;
    }
}

bitflags! {
    /// Flag options to close files.
    #[derive(Default)]
    pub struct CloseFlags: u32 {
        /// No close option
        const NONE = 0x0000;
        /// Write into a temporary file and swap for crash safety.
        const SAFE_UPDATE = 0x0001;
    }
}

bitflags! {
    /// Result flag for file / format infos.
    #[derive(Default)]
    pub struct FormatOptionFlags: u32 {
        const NONE = 0;
        /// Can inject first-time XMP into an existing file.
        const CAN_INJECT_XMP = 0x0000_0001;
        /// Can expand XMP or other metadata in an existing file.
        const CAN_EXPAND = 0x0000_0002;
        /// Can copy one file to another, writing new metadata.
        const CAN_REWRITE = 0x0000_0004;
        /// Can expand, but prefers in-place update.
        const PREFERS_IN_PLACE = 0x0000_0008;
        /// Supports reconciliation between XMP and other forms.
        const CAN_RECONCILE = 0x0000_0010;
        /// Allows access to just the XMP, ignoring other forms.
        const ALLOWS_ONLY_XMP = 0x0000_0020;
        /// File handler returns raw XMP packet information.
        const RETURNS_RAW_PACKET = 0x0000_0040;
        /// The file handler does the file open and close.
        const HANDLER_OWNS_FILE = 0x0000_0100;
        /// The file handler allows crash-safe file updates.
        const ALLOW_SAFE_UPDATE = 0x0000_0200;
        /// The file format needs the XMP packet to be read-only.
        const NEEDS_READONLY_PACKET = 0x0000_0400;
        /// The file handler uses a "sidecar" file for the XMP.
        const USE_SIDECAR_XMP = 0x0000_0800;
        /// The format is folder oriented, for example the P2 video format.
        const FOLDER_BASED_FORMAT = 0x0000_1000;
    }
}

pub struct XmpFile(*mut c::XmpFile);

impl Default for XmpFile {
    fn default() -> XmpFile {
        XmpFile(unsafe { c::xmp_files_new() })
    }
}

impl XmpFile {
    /// Create new XmpFile
    pub fn new() -> XmpFile {
        XmpFile::default()
    }

    /// Create and open a new XmpFile
    /// Equivalent to calling new then open.
    /// Return Err in case of failure
    pub fn open_new(p: &str, options: OpenFlags) -> Result<XmpFile> {
        let pp = CString::new(p).unwrap();
        let ptr = unsafe { c::xmp_files_open_new(pp.as_ptr(), options.bits()) };
        if ptr.is_null() {
            return Err(crate::get_error());
        }
        Ok(XmpFile(ptr))
    }

    /// Open an XmpFile. Usually called after new.
    pub fn open(&mut self, path: &str, options: OpenFlags) -> Result<()> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        let pp = CString::new(path).unwrap();
        if unsafe { c::xmp_files_open(self.0, pp.as_ptr(), options.bits()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Close the XmpFile
    pub fn close(&mut self, options: CloseFlags) -> Result<()> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_close(self.0, options.bits()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Return true if native pointer is null
    pub fn is_null(&self) -> bool {
        self.0.is_null()
    }

    /// Get a new Xmp fronm the currently open file
    pub fn get_new_xmp(&self) -> Result<Xmp> {
        let ptr = unsafe { c::xmp_files_get_new_xmp(self.0) };
        if ptr.is_null() {
            return Err(crate::get_error());
        }
        Ok(Xmp::from(ptr))
    }

    /// Get the xmp data an Xmp.
    pub fn get_xmp(&self, xmp: &mut Xmp) -> Result<()> {
        if self.is_null() || xmp.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_get_xmp(self.0, xmp.as_mut_ptr()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Get the xmp packet as a string.
    pub fn get_xmp_xmpstring(&self, packet: &mut XmpString, info: &mut PacketInfo) -> Result<()> {
        if self.is_null() || packet.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe {
            c::xmp_files_get_xmp_xmpstring(self.0, packet.as_mut_ptr(), info as *mut PacketInfo)
        } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Return true if it can put the Xmp into the XmpFile.
    pub fn can_put_xmp(&self, xmp: &Xmp) -> bool {
        if self.is_null() || xmp.is_null() {
            return false;
        }
        unsafe { c::xmp_files_can_put_xmp(self.0, xmp.as_ptr()) }
    }

    /// Return true if it can put the XmpString packet into the XmpFile.
    pub fn can_put_xmp_xmpstring(&self, xmp_packet: &XmpString) -> bool {
        if self.is_null() || xmp_packet.is_null() {
            return false;
        }
        unsafe { c::xmp_files_can_put_xmp_xmpstring(self.0, xmp_packet.as_ptr()) }
    }

    /// Return true if it can put the XmpString packet into the XmpFile.
    pub fn can_put_xmp_str(&self, xmp_packet: &str) -> bool {
        if self.is_null() {
            return false;
        }
        let pp = CString::new(xmp_packet).unwrap();
        unsafe { c::xmp_files_can_put_xmp_cstr(self.0, pp.as_ptr(), xmp_packet.len()) }
    }

    /// Put the Xmp into the XmpFile
    pub fn put_xmp(&mut self, xmp: &Xmp) -> Result<()> {
        if self.is_null() || xmp.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_put_xmp(self.0, xmp.as_ptr()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Get info from the XmpFile.
    pub fn get_file_info(
        &self,
        file_path: &mut String,
        options: &mut OpenFlags,
        format: &mut FileType,
        handler_flags: &mut FormatOptionFlags,
    ) -> bool {
        if self.is_null() {
            return false;
        }
        let mut s = XmpString::new();

        let mut raw_options: u32 = 0;
        let mut raw_handler_flags: u32 = 0;
        let result = unsafe {
            c::xmp_files_get_file_info(
                self.0,
                s.as_mut_ptr(),
                &mut raw_options,
                format,
                &mut raw_handler_flags,
            )
        };
        *options = OpenFlags::from_bits(raw_options).unwrap_or_default();
        *handler_flags = FormatOptionFlags::from_bits(raw_handler_flags).unwrap_or_default();

        file_path.push_str(s.to_str());
        result
    }

    /// Check the file format for the specified path
    pub fn check_file_format(path: &str) -> FileType {
        let pp = CString::new(path).unwrap();
        unsafe { c::xmp_files_check_file_format(pp.as_ptr()) }
    }

    /// Get FormatOptions for the FileType
    pub fn get_format_info(format: FileType) -> Result<FormatOptionFlags> {
        let mut raw_options: u32 = 0;
        if unsafe { c::xmp_files_get_format_info(format, &mut raw_options) } {
            Ok(FormatOptionFlags::from_bits(raw_options).unwrap_or_default())
        } else {
            Err(crate::get_error())
        }
    }
}

impl Drop for XmpFile {
    /// Drop the XmpFile.
    fn drop(&mut self) {
        if !self.is_null() {
            unsafe { c::xmp_files_free(self.0) };
        }
    }
}

#[test]
fn it_works() {
    use error::Error;

    let inited = crate::init();

    assert!(inited);

    let xf = XmpFile::new();
    assert!(!xf.is_null());

    assert!(crate::get_error() == Error::from(c::XmpError::Unknown));
}
<commit_msg>Fix test build<commit_after>use std::ffi::CString;

use c::FileType;
use c::XmpPacketInfo as PacketInfo;

use crate::error::Error;
use crate::Result;
use crate::xmp::Xmp;
use crate::xmpstring::XmpString;

bitflags! {
    /// Flag options for opening files.
    #[derive(Default)]
    pub struct OpenFlags: u32 {
        /// No open option
        const NONE = 0x0000_0000;
        /// Open for read-only access.
        const READ = 0x0000_0001;
        /// Open for reading and writing.
        const FOR_UPDATE = 0x0000_0002;
        /// Only the XMP is wanted, allows space/time optimizations.
        const ONLY_XMP = 0x0000_0004;
        /// Cache thumbnail if possible, GetThumbnail will be called.
        const CACHE_TNAIL = 0x0000_0008;
        /// Be strict about locating XMP and reconciling with other forms.
        const STRICTLY = 0x0000_0010;
        /// Require the use of a smart handler.
        const USE_SMART_HANDLER = 0x0000_0020;
        /// Force packet scanning, don't use a smart handler.
        const USE_PACKET_SCANNING = 0x0000_0040;
        /// Only packet scan files "known" to need scanning.
        const LIMITED_SCANNING = 0x0000_0080;
        /// Attempt to repair a file opened for update, default is to not open (throw an exception).
        const REPAIR_FILE = 0x0000_0100;
        /// Optimize MPEG4 to support stream when updating This can take some time
        const OPTIMIZE_FILE_LAYOUT = 0x0000_0200;
        /// Set if calling from background thread.
        const IN_BACKGROUND = 0x1000_0000;
    }
}

bitflags! {
    /// Flag options to close files.
    #[derive(Default)]
    pub struct CloseFlags: u32 {
        /// No close option
        const NONE = 0x0000;
        /// Write into a temporary file and swap for crash safety.
        const SAFE_UPDATE = 0x0001;
    }
}

bitflags! {
    /// Result flag for file / format infos.
    #[derive(Default)]
    pub struct FormatOptionFlags: u32 {
        const NONE = 0;
        /// Can inject first-time XMP into an existing file.
        const CAN_INJECT_XMP = 0x0000_0001;
        /// Can expand XMP or other metadata in an existing file.
        const CAN_EXPAND = 0x0000_0002;
        /// Can copy one file to another, writing new metadata.
        const CAN_REWRITE = 0x0000_0004;
        /// Can expand, but prefers in-place update.
        const PREFERS_IN_PLACE = 0x0000_0008;
        /// Supports reconciliation between XMP and other forms.
        const CAN_RECONCILE = 0x0000_0010;
        /// Allows access to just the XMP, ignoring other forms.
        const ALLOWS_ONLY_XMP = 0x0000_0020;
        /// File handler returns raw XMP packet information.
        const RETURNS_RAW_PACKET = 0x0000_0040;
        /// The file handler does the file open and close.
        const HANDLER_OWNS_FILE = 0x0000_0100;
        /// The file handler allows crash-safe file updates.
        const ALLOW_SAFE_UPDATE = 0x0000_0200;
        /// The file format needs the XMP packet to be read-only.
        const NEEDS_READONLY_PACKET = 0x0000_0400;
        /// The file handler uses a "sidecar" file for the XMP.
        const USE_SIDECAR_XMP = 0x0000_0800;
        /// The format is folder oriented, for example the P2 video format.
        const FOLDER_BASED_FORMAT = 0x0000_1000;
    }
}

pub struct XmpFile(*mut c::XmpFile);

impl Default for XmpFile {
    fn default() -> XmpFile {
        XmpFile(unsafe { c::xmp_files_new() })
    }
}

impl XmpFile {
    /// Create new XmpFile
    pub fn new() -> XmpFile {
        XmpFile::default()
    }

    /// Create and open a new XmpFile
    /// Equivalent to calling new then open.
    /// Return Err in case of failure
    pub fn open_new(p: &str, options: OpenFlags) -> Result<XmpFile> {
        let pp = CString::new(p).unwrap();
        let ptr = unsafe { c::xmp_files_open_new(pp.as_ptr(), options.bits()) };
        if ptr.is_null() {
            return Err(crate::get_error());
        }
        Ok(XmpFile(ptr))
    }

    /// Open an XmpFile. Usually called after new.
    pub fn open(&mut self, path: &str, options: OpenFlags) -> Result<()> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        let pp = CString::new(path).unwrap();
        if unsafe { c::xmp_files_open(self.0, pp.as_ptr(), options.bits()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Close the XmpFile
    pub fn close(&mut self, options: CloseFlags) -> Result<()> {
        if self.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_close(self.0, options.bits()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Return true if native pointer is null
    pub fn is_null(&self) -> bool {
        self.0.is_null()
    }

    /// Get a new Xmp fronm the currently open file
    pub fn get_new_xmp(&self) -> Result<Xmp> {
        let ptr = unsafe { c::xmp_files_get_new_xmp(self.0) };
        if ptr.is_null() {
            return Err(crate::get_error());
        }
        Ok(Xmp::from(ptr))
    }

    /// Get the xmp data an Xmp.
    pub fn get_xmp(&self, xmp: &mut Xmp) -> Result<()> {
        if self.is_null() || xmp.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_get_xmp(self.0, xmp.as_mut_ptr()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Get the xmp packet as a string.
    pub fn get_xmp_xmpstring(&self, packet: &mut XmpString, info: &mut PacketInfo) -> Result<()> {
        if self.is_null() || packet.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe {
            c::xmp_files_get_xmp_xmpstring(self.0, packet.as_mut_ptr(), info as *mut PacketInfo)
        } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Return true if it can put the Xmp into the XmpFile.
    pub fn can_put_xmp(&self, xmp: &Xmp) -> bool {
        if self.is_null() || xmp.is_null() {
            return false;
        }
        unsafe { c::xmp_files_can_put_xmp(self.0, xmp.as_ptr()) }
    }

    /// Return true if it can put the XmpString packet into the XmpFile.
    pub fn can_put_xmp_xmpstring(&self, xmp_packet: &XmpString) -> bool {
        if self.is_null() || xmp_packet.is_null() {
            return false;
        }
        unsafe { c::xmp_files_can_put_xmp_xmpstring(self.0, xmp_packet.as_ptr()) }
    }

    /// Return true if it can put the XmpString packet into the XmpFile.
    pub fn can_put_xmp_str(&self, xmp_packet: &str) -> bool {
        if self.is_null() {
            return false;
        }
        let pp = CString::new(xmp_packet).unwrap();
        unsafe { c::xmp_files_can_put_xmp_cstr(self.0, pp.as_ptr(), xmp_packet.len()) }
    }

    /// Put the Xmp into the XmpFile
    pub fn put_xmp(&mut self, xmp: &Xmp) -> Result<()> {
        if self.is_null() || xmp.is_null() {
            return Err(Error::from(c::XmpError::BadObject));
        }
        if unsafe { c::xmp_files_put_xmp(self.0, xmp.as_ptr()) } {
            Ok(())
        } else {
            Err(crate::get_error())
        }
    }

    /// Get info from the XmpFile.
    pub fn get_file_info(
        &self,
        file_path: &mut String,
        options: &mut OpenFlags,
        format: &mut FileType,
        handler_flags: &mut FormatOptionFlags,
    ) -> bool {
        if self.is_null() {
            return false;
        }
        let mut s = XmpString::new();

        let mut raw_options: u32 = 0;
        let mut raw_handler_flags: u32 = 0;
        let result = unsafe {
            c::xmp_files_get_file_info(
                self.0,
                s.as_mut_ptr(),
                &mut raw_options,
                format,
                &mut raw_handler_flags,
            )
        };
        *options = OpenFlags::from_bits(raw_options).unwrap_or_default();
        *handler_flags = FormatOptionFlags::from_bits(raw_handler_flags).unwrap_or_default();

        file_path.push_str(s.to_str());
        result
    }

    /// Check the file format for the specified path
    pub fn check_file_format(path: &str) -> FileType {
        let pp = CString::new(path).unwrap();
        unsafe { c::xmp_files_check_file_format(pp.as_ptr()) }
    }

    /// Get FormatOptions for the FileType
    pub fn get_format_info(format: FileType) -> Result<FormatOptionFlags> {
        let mut raw_options: u32 = 0;
        if unsafe { c::xmp_files_get_format_info(format, &mut raw_options) } {
            Ok(FormatOptionFlags::from_bits(raw_options).unwrap_or_default())
        } else {
            Err(crate::get_error())
        }
    }
}

impl Drop for XmpFile {
    /// Drop the XmpFile.
    fn drop(&mut self) {
        if !self.is_null() {
            unsafe { c::xmp_files_free(self.0) };
        }
    }
}

#[test]
fn it_works() {
    use crate::error::Error;

    let inited = crate::init();

    assert!(inited);

    let xf = XmpFile::new();
    assert!(!xf.is_null());

    assert!(crate::get_error() == Error::from(c::XmpError::Unknown));
}
<|endoftext|>


<commit_before>use core::fmt;
use libc::{c_int, c_uint, uint8_t};
use rustc_serialize::hex::ToHex;
use std::error::Error;
use std::fmt::{Display, Formatter};

#[repr(C)]
struct YK_KEY;

#[link(name="ykpers-1")]
extern {
    fn yk_init() -> c_int;
    fn yk_open_first_key() -> *const YK_KEY;
    fn yk_close_key(yk: *const YK_KEY) -> c_int;
    fn yk_get_serial(yk: *const YK_KEY, slot: uint8_t, flags: c_uint,
                     serial: *mut c_uint) -> c_int;
    fn yk_challenge_response(yk: *const YK_KEY, yk_cmd: uint8_t, may_block: c_int,
                             challenge_len: c_uint, challenge: *const uint8_t,
                             response_len: c_uint, response: *mut u8) -> c_int;
}

/* When a YK_KEY goes out of scope, close the handle */
impl Drop for YK_KEY {
    fn drop(&mut self) {
        unsafe { yk_close_key(self) };
    }
}

/*****************************
 *  Safe interface to ykpers *
 *****************************/

#[derive(Debug)]
pub enum YubikeyError { InvalidYubikeySlot,
                        EmptyCRChallenge,
                        UnknownError }

impl Error for YubikeyError {
    fn description(&self) -> &str {
        match *self {
            YubikeyError::InvalidYubikeySlot => "The selected Yubikey slot is invalid. Valid are 1, 2",
            YubikeyError::EmptyCRChallenge   => "The specified challenge was empty",
            YubikeyError::UnknownError       => "An unknown error occured"
        }
}

impl Display for YubikeyError {
    fn fmt(&self, formatter: &mut Formatter) -> fmt::Result {
        formatter.write_str(self.description())
    }
}

/* Opaque pointer to foreign Yubikey type. Under the hood this is just a handle
 *  for a USB device. */
pub type Yubikey = *const YK_KEY;

/* This function must be called at least once before using any of the Yubikey
 *  functionality. It initializes the Yubikey C library. */
pub fn yubikey_init() {
    unsafe { yk_init() };
}

/* Returns the first plugged in Yubikey. */
pub fn get_yubikey() -> Yubikey {
    unsafe { yk_open_first_key() }
}

/* Returns the serial number of the Yubikey */
pub fn get_serial(yk: Yubikey) -> u32 {
    unsafe {
        let mut serial: c_uint = 0;
        yk_get_serial(yk, 0, 0, &mut serial);
        serial as u32
    }
}

pub fn challenge_response(yk: Yubikey, slot: u8, challenge: &[u8], may_block: bool) -> Result<String, YubikeyError> {
    // Yubikey commands are defined in ykdef.h
    let yk_cmd = try!(match slot {
        1 => Ok(0x30), //#define SLOT_CHAL_HMAC1 0x30
        2 => Ok(0x38), //#define SLOT_CHAL_HMAC2 0x38
        _ => Err(YubikeyError::InvalidYubikeySlot)
    });

    let challenge_len = challenge.len() as c_uint;

    if challenge_len == 0 {
        return Err(YubikeyError::EmptyCRChallenge);
    }

    let response_len = 64; // Length of HMAC-SHA1 response

    let mut response = Vec::with_capacity(response_len as usize);

    let rc = unsafe {
        let cr_status = yk_challenge_response(yk, yk_cmd, may_block as c_int,
                                              challenge_len, challenge.as_ptr(),
                                              response_len, response.as_mut_ptr());

        response.set_len(response_len as usize);
        cr_status
    };

    if rc == 0 {
        // There's some way to get a better error code from the Yubikey, but
        // that's not needed right now.
        Err(YubikeyError::UnknownError)
    } else {
        // There must be a better way to do this, but I'm not sure what it is yet
        let response_str =
            String::from_str(response.as_slice().to_hex().as_str().slice_chars(0,40));
        Ok(response_str)
    }
}
<commit_msg>[yubikey] Minor syntax fix<commit_after>use core::fmt;
use libc::{c_int, c_uint, uint8_t};
use rustc_serialize::hex::ToHex;
use std::error::Error;
use std::fmt::{Display, Formatter};

#[repr(C)]
struct YK_KEY;

#[link(name="ykpers-1")]
extern {
    fn yk_init() -> c_int;
    fn yk_open_first_key() -> *const YK_KEY;
    fn yk_close_key(yk: *const YK_KEY) -> c_int;
    fn yk_get_serial(yk: *const YK_KEY, slot: uint8_t, flags: c_uint,
                     serial: *mut c_uint) -> c_int;
    fn yk_challenge_response(yk: *const YK_KEY, yk_cmd: uint8_t, may_block: c_int,
                             challenge_len: c_uint, challenge: *const uint8_t,
                             response_len: c_uint, response: *mut u8) -> c_int;
}

/* When a YK_KEY goes out of scope, close the handle */
impl Drop for YK_KEY {
    fn drop(&mut self) {
        unsafe { yk_close_key(self) };
    }
}

/*****************************
 *  Safe interface to ykpers *
 *****************************/

#[derive(Debug)]
pub enum YubikeyError { InvalidYubikeySlot,
                        EmptyCRChallenge,
                        UnknownError }

impl Error for YubikeyError {
    fn description(&self) -> &str {
        match *self {
            YubikeyError::InvalidYubikeySlot => "The selected Yubikey slot is invalid. Valid are 1, 2",
            YubikeyError::EmptyCRChallenge   => "The specified challenge was empty",
            YubikeyError::UnknownError       => "An unknown error occured"
        }
    }
}

impl Display for YubikeyError {
    fn fmt(&self, formatter: &mut Formatter) -> fmt::Result {
        formatter.write_str(self.description())
    }
}

/* Opaque pointer to foreign Yubikey type. Under the hood this is just a handle
 *  for a USB device. */
pub type Yubikey = *const YK_KEY;

/* This function must be called at least once before using any of the Yubikey
 *  functionality. It initializes the Yubikey C library. */
pub fn yubikey_init() {
    unsafe { yk_init() };
}

/* Returns the first plugged in Yubikey. */
pub fn get_yubikey() -> Yubikey {
    unsafe { yk_open_first_key() }
}

/* Returns the serial number of the Yubikey */
pub fn get_serial(yk: Yubikey) -> u32 {
    unsafe {
        let mut serial: c_uint = 0;
        yk_get_serial(yk, 0, 0, &mut serial);
        serial as u32
    }
}

pub fn challenge_response(yk: Yubikey, slot: u8, challenge: &[u8], may_block: bool) -> Result<String, YubikeyError> {
    // Yubikey commands are defined in ykdef.h
    let yk_cmd = try!(match slot {
        1 => Ok(0x30), //#define SLOT_CHAL_HMAC1 0x30
        2 => Ok(0x38), //#define SLOT_CHAL_HMAC2 0x38
        _ => Err(YubikeyError::InvalidYubikeySlot)
    });

    let challenge_len = challenge.len() as c_uint;

    if challenge_len == 0 {
        return Err(YubikeyError::EmptyCRChallenge);
    }

    let response_len = 64; // Length of HMAC-SHA1 response

    let mut response = Vec::with_capacity(response_len as usize);

    let rc = unsafe {
        let cr_status = yk_challenge_response(yk, yk_cmd, may_block as c_int,
                                              challenge_len, challenge.as_ptr(),
                                              response_len, response.as_mut_ptr());

        response.set_len(response_len as usize);
        cr_status
    };

    if rc == 0 {
        // There's some way to get a better error code from the Yubikey, but
        // that's not needed right now.
        Err(YubikeyError::UnknownError)
    } else {
        // There must be a better way to do this, but I'm not sure what it is yet
        let response_str =
            String::from_str(response.as_slice().to_hex().as_str().slice_chars(0,40));
        Ok(response_str)
    }
}
<|endoftext|>


<commit_before>use std::sync::Mutex;
use std::collections::HashMap;
use circular_queue::CircularQueue;
use regex::{Regex, RegexBuilder};

use irc::client::prelude::*;

use plugin::*;

use failure::Fail;
use failure::ResultExt;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use self::error::*;

lazy_static! {
    static ref RE: Regex = Regex::new(r"^s/((?:\\/|[^/])+)/((?:\\/|[^/])*)/(?:(\w+))?\s*$").unwrap();
}

#[derive(PluginName, Debug)]
pub struct Sed {
    per_channel: usize,
    channel_messages: Mutex<HashMap<String, CircularQueue<String>>>,
}

macro_rules! try_lock {
    ( $m:expr ) => {
        match $m.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        }
    }
}

impl Sed {
    pub fn new(per_channel: usize) -> Sed {
        Sed {
            per_channel: per_channel,
            channel_messages: Mutex::new(HashMap::new()),
        }
    }

    fn add_message(&self, channel: String, message: String) {
        let mut channel_messages = try_lock!(self.channel_messages);
        let messages = channel_messages
            .entry(channel)
            .or_insert(CircularQueue::with_capacity(self.per_channel));
        messages.push(message);
    }

    fn format_escaped(&self, input: &str) -> String {
        let mut output = String::with_capacity(input.len());
        let mut escape = false;

        for c in input.chars() {
            if !escape && c == '\\' {
                escape = true;
                continue;
            }
            escape = false;

            output.push(c);
        }

        output
    }

    fn run_regex(&self, channel: &str, message: &str) -> Result<String, SedError> {
        let mut global_match = false;
        let mut case_insens = false;
        let mut ign_whitespace = false;
        let mut swap_greed = false;
        let mut disable_unicode = false;

        let captures = RE.captures(message).unwrap();
        debug!("{:?}", captures);

        let first = self.format_escaped(captures.get(1).unwrap().as_str());
        let second = self.format_escaped(captures.get(2).unwrap().as_str());

        if let Some(flags) = captures.get(3) {
            let flags = flags.as_str();

            global_match = flags.contains('g');
            case_insens = flags.contains('i');
            ign_whitespace = flags.contains('x');
            swap_greed = flags.contains('U');
            disable_unicode = !flags.contains('u');
        }

        let user_re = RegexBuilder::new(&first)
            .case_insensitive(case_insens)
            .ignore_whitespace(ign_whitespace)
            .unicode(disable_unicode)
            .swap_greed(swap_greed)
            .build()
            .context(ErrorKind::InvalidRegex)?;

        let channel_messages = try_lock!(self.channel_messages);
        let messages = channel_messages.get(channel).ok_or(ErrorKind::NoMessages)?;

        for message in messages.iter() {
            if user_re.is_match(message) {
                let response = if global_match {
                    user_re.replace_all(message, &second[..])
                } else {
                    user_re.replace(message, &second[..])
                };

                return Ok(response.to_string());
            }
        }

        Err(ErrorKind::NoMatch)?
    }
}

impl Plugin for Sed {
    fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus {
        match message.command {
            Command::PRIVMSG(_, ref content) => {
                let channel = message.response_target().unwrap();
                if channel == message.source_nickname().unwrap() {
                    return ExecutionStatus::Done;
                }

                if RE.is_match(content) {
                    let result = match self.run_regex(channel, content) {
                        Ok(msg) => client.send_privmsg(channel, &msg),
                        Err(e) => client.send_notice(channel, &e.to_string()),
                    };

                    match result {
                        Err(e) => {
                            ExecutionStatus::Err(e.context(FrippyErrorKind::Connection).into())
                        }
                        Ok(_) => ExecutionStatus::Done,
                    }
                } else {
                    self.add_message(channel.to_string(), content.to_string());

                    ExecutionStatus::Done
                }
            }
            _ => ExecutionStatus::Done,
        }
    }

    fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> {
        panic!("Sed should not use threading")
    }

    fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> {
        Ok(client
            .send_notice(
                &command.source,
                "Currently this Plugin does not implement any commands.",
            )
            .context(FrippyErrorKind::Connection)?)
    }

    fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> {
        Err(String::from(
            "Evaluation of commands is not implemented for sed at this time",
        ))
    }
}

pub mod error {
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "SedError"]
    pub enum ErrorKind {
        /// Invalid regex error
        #[fail(display = "Invalid regex")]
        InvalidRegex,

        /// No messages found error
        #[fail(display = "No messages were found for this channel")]
        NoMessages,

        /// No match found error
        #[fail(display = "No recent messages match this regex")]
        NoMatch,
    }
}
<commit_msg>Remove the requirement to escape all backslashes<commit_after>use std::sync::Mutex;
use std::collections::HashMap;
use circular_queue::CircularQueue;
use regex::{Regex, RegexBuilder};

use irc::client::prelude::*;

use plugin::*;

use failure::Fail;
use failure::ResultExt;
use error::ErrorKind as FrippyErrorKind;
use error::FrippyError;
use self::error::*;

lazy_static! {
    static ref RE: Regex = Regex::new(r"^s/((?:\\/|[^/])+)/((?:\\/|[^/])*)/(?:(\w+))?\s*$").unwrap();
}

#[derive(PluginName, Debug)]
pub struct Sed {
    per_channel: usize,
    channel_messages: Mutex<HashMap<String, CircularQueue<String>>>,
}

macro_rules! try_lock {
    ( $m:expr ) => {
        match $m.lock() {
            Ok(guard) => guard,
            Err(poisoned) => poisoned.into_inner(),
        }
    }
}

impl Sed {
    pub fn new(per_channel: usize) -> Sed {
        Sed {
            per_channel: per_channel,
            channel_messages: Mutex::new(HashMap::new()),
        }
    }

    fn add_message(&self, channel: String, message: String) {
        let mut channel_messages = try_lock!(self.channel_messages);
        let messages = channel_messages
            .entry(channel)
            .or_insert(CircularQueue::with_capacity(self.per_channel));
        messages.push(message);
    }

    fn format_escaped(&self, input: &str) -> String {
        let mut output = String::with_capacity(input.len());
        let mut escape = false;

        for c in input.chars() {
            if escape && !r"/\".contains(c) {
                output.push('\\');

            } else if !escape && c == '\\' {
                escape = true;
                continue;
            }
            escape = false;

            output.push(c);
        }

        output
    }

    fn run_regex(&self, channel: &str, message: &str) -> Result<String, SedError> {
        let mut global_match = false;
        let mut case_insens = false;
        let mut ign_whitespace = false;
        let mut swap_greed = false;
        let mut disable_unicode = false;

        let captures = RE.captures(message).unwrap();
        debug!("{:?}", captures);

        let first = self.format_escaped(captures.get(1).unwrap().as_str());
        let second = self.format_escaped(captures.get(2).unwrap().as_str());

        if let Some(flags) = captures.get(3) {
            let flags = flags.as_str();

            global_match = flags.contains('g');
            case_insens = flags.contains('i');
            ign_whitespace = flags.contains('x');
            swap_greed = flags.contains('U');
            disable_unicode = !flags.contains('u');
        }

        let user_re = RegexBuilder::new(&first)
            .case_insensitive(case_insens)
            .ignore_whitespace(ign_whitespace)
            .unicode(disable_unicode)
            .swap_greed(swap_greed)
            .build()
            .context(ErrorKind::InvalidRegex)?;

        let channel_messages = try_lock!(self.channel_messages);
        let messages = channel_messages.get(channel).ok_or(ErrorKind::NoMessages)?;

        for message in messages.iter() {
            if user_re.is_match(message) {
                let response = if global_match {
                    user_re.replace_all(message, &second[..])
                } else {
                    user_re.replace(message, &second[..])
                };

                return Ok(response.to_string());
            }
        }

        Err(ErrorKind::NoMatch)?
    }
}

impl Plugin for Sed {
    fn execute(&self, client: &IrcClient, message: &Message) -> ExecutionStatus {
        match message.command {
            Command::PRIVMSG(_, ref content) => {
                let channel = message.response_target().unwrap();
                if channel == message.source_nickname().unwrap() {
                    return ExecutionStatus::Done;
                }

                if RE.is_match(content) {
                    let result = match self.run_regex(channel, content) {
                        Ok(msg) => client.send_privmsg(channel, &msg),
                        Err(e) => client.send_notice(channel, &e.to_string()),
                    };

                    match result {
                        Err(e) => {
                            ExecutionStatus::Err(e.context(FrippyErrorKind::Connection).into())
                        }
                        Ok(_) => ExecutionStatus::Done,
                    }
                } else {
                    self.add_message(channel.to_string(), content.to_string());

                    ExecutionStatus::Done
                }
            }
            _ => ExecutionStatus::Done,
        }
    }

    fn execute_threaded(&self, _: &IrcClient, _: &Message) -> Result<(), FrippyError> {
        panic!("Sed should not use threading")
    }

    fn command(&self, client: &IrcClient, command: PluginCommand) -> Result<(), FrippyError> {
        Ok(client
            .send_notice(
                &command.source,
                "Currently this Plugin does not implement any commands.",
            )
            .context(FrippyErrorKind::Connection)?)
    }

    fn evaluate(&self, _: &IrcClient, _: PluginCommand) -> Result<String, String> {
        Err(String::from(
            "Evaluation of commands is not implemented for sed at this time",
        ))
    }
}

pub mod error {
    #[derive(Copy, Clone, Eq, PartialEq, Debug, Fail, Error)]
    #[error = "SedError"]
    pub enum ErrorKind {
        /// Invalid regex error
        #[fail(display = "Invalid regex")]
        InvalidRegex,

        /// No messages found error
        #[fail(display = "No messages were found for this channel")]
        NoMessages,

        /// No match found error
        #[fail(display = "No recent messages match this regex")]
        NoMatch,
    }
}
<|endoftext|>


<commit_before>use docset::{DocSet, SkipResult};
use query::Scorer;
use common::TinySet;
use std::cmp::Ordering;
use DocId;
use Score;
use query::score_combiner::{DoNothingCombiner, ScoreCombiner};

const HORIZON_NUM_TINYBITSETS: usize = 32;
const HORIZON: u32 = 64u32 * HORIZON_NUM_TINYBITSETS as u32;

/// Creates a `DocSet` that iterator through the intersection of two `DocSet`s.
pub struct Union<TScorer, TScoreCombiner = DoNothingCombiner> {
    docsets: Vec<TScorer>,
    bitsets: Box<[TinySet; HORIZON_NUM_TINYBITSETS]>,
    scores: Box<[TScoreCombiner; HORIZON as usize]>,
    cursor: usize,
    offset: DocId,
    doc: DocId,
    score: Score,
}

impl<TScorer, TScoreCombiner> From<Vec<TScorer>> for Union<TScorer, TScoreCombiner>
where
    TScoreCombiner: ScoreCombiner,
    TScorer: Scorer,
{
    fn from(docsets: Vec<TScorer>) -> Union<TScorer, TScoreCombiner> {
        let non_empty_docsets: Vec<TScorer> = docsets
            .into_iter()
            .flat_map(
                |mut docset| {
                    if docset.advance() {
                        Some(docset)
                    } else {
                        None
                    }
                },
            )
            .collect();
        Union {
            docsets: non_empty_docsets,
            bitsets: Box::new([TinySet::empty(); HORIZON_NUM_TINYBITSETS]),
            scores: Box::new([TScoreCombiner::default(); HORIZON as usize]),
            cursor: HORIZON_NUM_TINYBITSETS,
            offset: 0,
            doc: 0,
            score: 0f32,
        }
    }
}

fn refill<TScorer: Scorer, TScoreCombiner: ScoreCombiner>(
    scorers: &mut Vec<TScorer>,
    bitsets: &mut [TinySet; HORIZON_NUM_TINYBITSETS],
    score_combiner: &mut [TScoreCombiner; HORIZON as usize],
    min_doc: DocId,
) {
    scorers.drain_filter(|scorer| {
        let horizon = min_doc + HORIZON as u32;
        loop {
            let doc = scorer.doc();
            if doc >= horizon {
                return false;
            }
            // add this document
            let delta = doc - min_doc;
            bitsets[(delta / 64) as usize].insert_mut(delta % 64u32);
            score_combiner[delta as usize].update(scorer);
            if !scorer.advance() {
                // remove the docset, it has been entirely consumed.
                return true;
            }
        }
    });
}

impl<TScorer: Scorer, TScoreCombiner: ScoreCombiner> Union<TScorer, TScoreCombiner> {
    fn refill(&mut self) -> bool {
        if let Some(min_doc) = self.docsets.iter_mut().map(|docset| docset.doc()).min() {
            self.offset = min_doc;
            self.cursor = 0;
            refill(
                &mut self.docsets,
                &mut *self.bitsets,
                &mut *self.scores,
                min_doc,
            );
            true
        } else {
            false
        }
    }

    fn advance_buffered(&mut self) -> bool {
        while self.cursor < HORIZON_NUM_TINYBITSETS {
            if let Some(val) = self.bitsets[self.cursor].pop_lowest() {
                let delta = val + (self.cursor as u32) * 64;
                self.doc = self.offset + delta;
                let score_combiner = &mut self.scores[delta as usize];
                self.score = score_combiner.score();
                score_combiner.clear();
                return true;
            } else {
                self.cursor += 1;
            }
        }
        false
    }
}

impl<TScorer, TScoreCombiner> DocSet for Union<TScorer, TScoreCombiner>
    where TScorer: Scorer, TScoreCombiner: ScoreCombiner
{
    fn advance(&mut self) -> bool {
        if self.advance_buffered() {
            return true;
        }
        if self.refill() {
            self.advance();
            true
        } else {
            false
        }
    }

    fn count(&mut self) -> u32 {
        let mut count = self.bitsets[self.cursor..HORIZON_NUM_TINYBITSETS]
            .iter()
            .map(|bitset| bitset.len())
            .sum::<u32>();
        for bitset in self.bitsets.iter_mut() {
            bitset.clear();
        }
        while self.refill() {
            count += self.bitsets.iter().map(|bitset| bitset.len()).sum::<u32>();
            for bitset in self.bitsets.iter_mut() {
                bitset.clear();
            }
        }
        self.cursor = HORIZON_NUM_TINYBITSETS;
        count
    }

    fn skip_next(&mut self, target: DocId) -> SkipResult {
        if !self.advance() {
            return SkipResult::End;
        }
        match self.doc.cmp(&target) {
            Ordering::Equal => {
                return SkipResult::Reached;
            }
            Ordering::Greater => {
                return SkipResult::OverStep;
            }
            Ordering::Less => {}
        }
        let gap = target - self.offset;
        if gap < HORIZON {
            // Our value is within the buffered horizon.

            // Skipping to  corresponding bucket.
            let new_cursor = gap as usize / 64;
            for obsolete_tinyset in &mut self.bitsets[self.cursor..new_cursor] {
                obsolete_tinyset.clear();
            }
            for score_combiner in &mut self.scores[self.cursor * 64..new_cursor * 64] {
                score_combiner.clear();
            }
            self.cursor = new_cursor;

            // Advancing until we reach the end of the bucket
            // or we reach a doc greater or equal to the target.
            while self.advance() {
                match self.doc().cmp(&target) {
                    Ordering::Equal => {
                        return SkipResult::Reached;
                    }
                    Ordering::Greater => {
                        return SkipResult::OverStep;
                    }
                    Ordering::Less => {}
                }
            }
            SkipResult::End
        } else {
            // clear the buffered info.
            for obsolete_tinyset in self.bitsets.iter_mut() {
                *obsolete_tinyset = TinySet::empty();
            }
            for score_combiner in self.scores.iter_mut() {
                score_combiner.clear();
            }

            // The target is outside of the buffered horizon.
            // advance all docsets to a doc >= to the target.
            self.docsets
                .drain_filter(|docset| match docset.doc().cmp(&target) {
                    Ordering::Less => match docset.skip_next(target) {
                        SkipResult::End => true,
                        SkipResult::Reached | SkipResult::OverStep => false,
                    },
                    Ordering::Equal | Ordering::Greater => false,
                });

            // at this point all of the docsets
            // are positionned on a doc >= to the target.
            if self.refill() {
                self.advance();
                if self.doc() == target {
                    SkipResult::Reached
                } else {
                    debug_assert!(self.doc() > target);
                    SkipResult::OverStep
                }
            } else {
                SkipResult::End
            }
        }
    }

    fn doc(&self) -> DocId {
        self.doc
    }

    fn size_hint(&self) -> u32 {
        0u32
    }
}

impl<TScorer, TScoreCombiner> Scorer for Union<TScorer, TScoreCombiner>
where
    TScoreCombiner: ScoreCombiner,
    TScorer: Scorer,
{
    fn score(&mut self) -> Score {
        self.score
    }
}

#[cfg(test)]
mod tests {

    use super::Union;
    use tests;
    use test::Bencher;
    use DocId;
    use std::collections::BTreeSet;
    use super::HORIZON;
    use docset::{DocSet, SkipResult};
    use postings::tests::test_skip_against_unoptimized;
    use query::VecDocSet;
    use query::ConstScorer;
    use query::score_combiner::DoNothingCombiner;

    fn aux_test_union(vals: Vec<Vec<u32>>) {
        let mut val_set: BTreeSet<u32> = BTreeSet::new();
        for vs in &vals {
            for &v in vs {
                val_set.insert(v);
            }
        }
        let union_vals: Vec<u32> = val_set.into_iter().collect();
        let mut union_expected = VecDocSet::from(union_vals);
        let make_union = || {
            Union::from(
                vals.iter()
                    .cloned()
                    .map(VecDocSet::from)
                    .map(ConstScorer::new)
                    .collect::<Vec<ConstScorer<VecDocSet>>>(),
            )
        };
        let mut union: Union<_, DoNothingCombiner> = make_union();
        let mut count = 0;
        while union.advance() {
            assert!(union_expected.advance());
            assert_eq!(union_expected.doc(), union.doc());
            count += 1;
        }
        assert!(!union_expected.advance());
        assert_eq!(count, make_union().count());
    }

    #[test]
    fn test_union() {
        aux_test_union(vec![
            vec![1, 3333, 100000000u32],
            vec![1, 2, 100000000u32],
            vec![1, 2, 100000000u32],
            vec![],
        ]);
        aux_test_union(vec![
            vec![1, 3333, 100000000u32],
            vec![1, 2, 100000000u32],
            vec![1, 2, 100000000u32],
            vec![],
        ]);
        aux_test_union(vec![
            tests::sample_with_seed(100_000, 0.01, 1),
            tests::sample_with_seed(100_000, 0.05, 2),
            tests::sample_with_seed(100_000, 0.001, 3),
        ]);
    }

    fn test_aux_union_skip(docs_list: &[Vec<DocId>], skip_targets: Vec<DocId>) {
        let mut btree_set = BTreeSet::new();
        for docs in docs_list {
            for &doc in docs.iter() {
                btree_set.insert(doc);
            }
        }
        let docset_factory = || {
            let res: Box<DocSet> = box Union::<_, DoNothingCombiner>::from(
                docs_list
                    .iter()
                    .map(|docs| docs.clone())
                    .map(VecDocSet::from)
                    .map(ConstScorer::new)
                    .collect::<Vec<_>>(),
            );
            res
        };
        let mut docset = docset_factory();
        for el in btree_set {
            assert!(docset.advance());
            assert_eq!(el, docset.doc());
        }
        assert!(!docset.advance());
        test_skip_against_unoptimized(docset_factory, skip_targets);
    }

    #[test]
    fn test_union_skip_corner_case() {
        test_aux_union_skip(&[vec![165132, 167382], vec![25029, 25091]], vec![25029]);
    }

    #[test]
    fn test_union_skip_corner_case2() {
        test_aux_union_skip(
            &[vec![1u32, 1u32 + HORIZON], vec![2u32, 1000u32, 10_000u32]],
            vec![0u32, 1u32, 2u32, 3u32, 1u32 + HORIZON, 2u32 + HORIZON],
        );
    }

    #[test]
    fn test_union_skip_corner_case3() {
        let mut docset = Union::<_, DoNothingCombiner>::from(vec![
            ConstScorer::new(VecDocSet::from(vec![0u32, 5u32])),
            ConstScorer::new(VecDocSet::from(vec![1u32, 4u32])),
        ]);
        assert!(docset.advance());
        assert_eq!(docset.doc(), 0u32);
        assert_eq!(docset.skip_next(0u32), SkipResult::OverStep);
        assert_eq!(docset.doc(), 1u32)
    }

    #[test]
    fn test_union_skip_random() {
        test_aux_union_skip(
            &[
                vec![1, 2, 3, 7],
                vec![1, 3, 9, 10000],
                vec![1, 3, 8, 9, 100],
            ],
            vec![1, 2, 3, 5, 6, 7, 8, 100],
        );
        test_aux_union_skip(
            &[
                tests::sample_with_seed(100_000, 0.001, 1),
                tests::sample_with_seed(100_000, 0.002, 2),
                tests::sample_with_seed(100_000, 0.005, 3),
            ],
            tests::sample_with_seed(100_000, 0.01, 4),
        );
    }

    #[test]
    fn test_union_skip_specific() {
        test_aux_union_skip(
            &[
                vec![1, 2, 3, 7],
                vec![1, 3, 9, 10000],
                vec![1, 3, 8, 9, 100],
            ],
            vec![1, 2, 3, 7, 8, 9, 99, 100, 101, 500, 20000],
        );
    }

    #[bench]
    fn bench_union_3_high(bench: &mut Bencher) {
        let union_docset: Vec<Vec<DocId>> = vec![
            tests::sample_with_seed(100_000, 0.1, 0),
            tests::sample_with_seed(100_000, 0.2, 1),
        ];
        bench.iter(|| {
            let mut v = Union::<_, DoNothingCombiner>::from(
                union_docset
                    .iter()
                    .map(|doc_ids| VecDocSet::from(doc_ids.clone()))
                    .map(ConstScorer::new)
                    .collect::<Vec<_>>(),
            );
            while v.advance() {}
        });
    }
    #[bench]
    fn bench_union_3_low(bench: &mut Bencher) {
        let union_docset: Vec<Vec<DocId>> = vec![
            tests::sample_with_seed(100_000, 0.01, 0),
            tests::sample_with_seed(100_000, 0.05, 1),
            tests::sample_with_seed(100_000, 0.001, 2),
        ];
        bench.iter(|| {
            let mut v = Union::<_, DoNothingCombiner>::from(
                union_docset
                    .iter()
                    .map(|doc_ids| VecDocSet::from(doc_ids.clone()))
                    .map(ConstScorer::new)
                    .collect::<Vec<_>>(),
            );
            while v.advance() {}
        });
    }

}
<commit_msg>Larger union horizon<commit_after>use docset::{DocSet, SkipResult};
use query::Scorer;
use common::TinySet;
use std::cmp::Ordering;
use DocId;
use Score;
use query::score_combiner::{DoNothingCombiner, ScoreCombiner};

const HORIZON_NUM_TINYBITSETS: usize = 64;
const HORIZON: u32 = 64u32 * HORIZON_NUM_TINYBITSETS as u32;

/// Creates a `DocSet` that iterator through the intersection of two `DocSet`s.
pub struct Union<TScorer, TScoreCombiner = DoNothingCombiner> {
    docsets: Vec<TScorer>,
    bitsets: Box<[TinySet; HORIZON_NUM_TINYBITSETS]>,
    scores: Box<[TScoreCombiner; HORIZON as usize]>,
    cursor: usize,
    offset: DocId,
    doc: DocId,
    score: Score,
}

impl<TScorer, TScoreCombiner> From<Vec<TScorer>> for Union<TScorer, TScoreCombiner>
where
    TScoreCombiner: ScoreCombiner,
    TScorer: Scorer,
{
    fn from(docsets: Vec<TScorer>) -> Union<TScorer, TScoreCombiner> {
        let non_empty_docsets: Vec<TScorer> = docsets
            .into_iter()
            .flat_map(
                |mut docset| {
                    if docset.advance() {
                        Some(docset)
                    } else {
                        None
                    }
                },
            )
            .collect();
        Union {
            docsets: non_empty_docsets,
            bitsets: Box::new([TinySet::empty(); HORIZON_NUM_TINYBITSETS]),
            scores: Box::new([TScoreCombiner::default(); HORIZON as usize]),
            cursor: HORIZON_NUM_TINYBITSETS,
            offset: 0,
            doc: 0,
            score: 0f32,
        }
    }
}

fn refill<TScorer: Scorer, TScoreCombiner: ScoreCombiner>(
    scorers: &mut Vec<TScorer>,
    bitsets: &mut [TinySet; HORIZON_NUM_TINYBITSETS],
    score_combiner: &mut [TScoreCombiner; HORIZON as usize],
    min_doc: DocId,
) {
    scorers.drain_filter(|scorer| {
        let horizon = min_doc + HORIZON as u32;
        loop {
            let doc = scorer.doc();
            if doc >= horizon {
                return false;
            }
            // add this document
            let delta = doc - min_doc;
            bitsets[(delta / 64) as usize].insert_mut(delta % 64u32);
            score_combiner[delta as usize].update(scorer);
            if !scorer.advance() {
                // remove the docset, it has been entirely consumed.
                return true;
            }
        }
    });
}

impl<TScorer: Scorer, TScoreCombiner: ScoreCombiner> Union<TScorer, TScoreCombiner> {
    fn refill(&mut self) -> bool {
        if let Some(min_doc) = self.docsets.iter_mut().map(|docset| docset.doc()).min() {
            self.offset = min_doc;
            self.cursor = 0;
            refill(
                &mut self.docsets,
                &mut *self.bitsets,
                &mut *self.scores,
                min_doc,
            );
            true
        } else {
            false
        }
    }

    fn advance_buffered(&mut self) -> bool {
        while self.cursor < HORIZON_NUM_TINYBITSETS {
            if let Some(val) = self.bitsets[self.cursor].pop_lowest() {
                let delta = val + (self.cursor as u32) * 64;
                self.doc = self.offset + delta;
                let score_combiner = &mut self.scores[delta as usize];
                self.score = score_combiner.score();
                score_combiner.clear();
                return true;
            } else {
                self.cursor += 1;
            }
        }
        false
    }
}

impl<TScorer, TScoreCombiner> DocSet for Union<TScorer, TScoreCombiner>
    where TScorer: Scorer, TScoreCombiner: ScoreCombiner
{
    fn advance(&mut self) -> bool {
        if self.advance_buffered() {
            return true;
        }
        if self.refill() {
            self.advance();
            true
        } else {
            false
        }
    }

    fn count(&mut self) -> u32 {
        let mut count = self.bitsets[self.cursor..HORIZON_NUM_TINYBITSETS]
            .iter()
            .map(|bitset| bitset.len())
            .sum::<u32>();
        for bitset in self.bitsets.iter_mut() {
            bitset.clear();
        }
        while self.refill() {
            count += self.bitsets.iter().map(|bitset| bitset.len()).sum::<u32>();
            for bitset in self.bitsets.iter_mut() {
                bitset.clear();
            }
        }
        self.cursor = HORIZON_NUM_TINYBITSETS;
        count
    }

    fn skip_next(&mut self, target: DocId) -> SkipResult {
        if !self.advance() {
            return SkipResult::End;
        }
        match self.doc.cmp(&target) {
            Ordering::Equal => {
                return SkipResult::Reached;
            }
            Ordering::Greater => {
                return SkipResult::OverStep;
            }
            Ordering::Less => {}
        }
        let gap = target - self.offset;
        if gap < HORIZON {
            // Our value is within the buffered horizon.

            // Skipping to  corresponding bucket.
            let new_cursor = gap as usize / 64;
            for obsolete_tinyset in &mut self.bitsets[self.cursor..new_cursor] {
                obsolete_tinyset.clear();
            }
            for score_combiner in &mut self.scores[self.cursor * 64..new_cursor * 64] {
                score_combiner.clear();
            }
            self.cursor = new_cursor;

            // Advancing until we reach the end of the bucket
            // or we reach a doc greater or equal to the target.
            while self.advance() {
                match self.doc().cmp(&target) {
                    Ordering::Equal => {
                        return SkipResult::Reached;
                    }
                    Ordering::Greater => {
                        return SkipResult::OverStep;
                    }
                    Ordering::Less => {}
                }
            }
            SkipResult::End
        } else {
            // clear the buffered info.
            for obsolete_tinyset in self.bitsets.iter_mut() {
                *obsolete_tinyset = TinySet::empty();
            }
            for score_combiner in self.scores.iter_mut() {
                score_combiner.clear();
            }

            // The target is outside of the buffered horizon.
            // advance all docsets to a doc >= to the target.
            self.docsets
                .drain_filter(|docset| match docset.doc().cmp(&target) {
                    Ordering::Less => match docset.skip_next(target) {
                        SkipResult::End => true,
                        SkipResult::Reached | SkipResult::OverStep => false,
                    },
                    Ordering::Equal | Ordering::Greater => false,
                });

            // at this point all of the docsets
            // are positionned on a doc >= to the target.
            if self.refill() {
                self.advance();
                if self.doc() == target {
                    SkipResult::Reached
                } else {
                    debug_assert!(self.doc() > target);
                    SkipResult::OverStep
                }
            } else {
                SkipResult::End
            }
        }
    }

    fn doc(&self) -> DocId {
        self.doc
    }

    fn size_hint(&self) -> u32 {
        0u32
    }
}

impl<TScorer, TScoreCombiner> Scorer for Union<TScorer, TScoreCombiner>
where
    TScoreCombiner: ScoreCombiner,
    TScorer: Scorer,
{
    fn score(&mut self) -> Score {
        self.score
    }
}

#[cfg(test)]
mod tests {

    use super::Union;
    use tests;
    use test::Bencher;
    use DocId;
    use std::collections::BTreeSet;
    use super::HORIZON;
    use docset::{DocSet, SkipResult};
    use postings::tests::test_skip_against_unoptimized;
    use query::VecDocSet;
    use query::ConstScorer;
    use query::score_combiner::DoNothingCombiner;

    fn aux_test_union(vals: Vec<Vec<u32>>) {
        let mut val_set: BTreeSet<u32> = BTreeSet::new();
        for vs in &vals {
            for &v in vs {
                val_set.insert(v);
            }
        }
        let union_vals: Vec<u32> = val_set.into_iter().collect();
        let mut union_expected = VecDocSet::from(union_vals);
        let make_union = || {
            Union::from(
                vals.iter()
                    .cloned()
                    .map(VecDocSet::from)
                    .map(ConstScorer::new)
                    .collect::<Vec<ConstScorer<VecDocSet>>>(),
            )
        };
        let mut union: Union<_, DoNothingCombiner> = make_union();
        let mut count = 0;
        while union.advance() {
            assert!(union_expected.advance());
            assert_eq!(union_expected.doc(), union.doc());
            count += 1;
        }
        assert!(!union_expected.advance());
        assert_eq!(count, make_union().count());
    }

    #[test]
    fn test_union() {
        aux_test_union(vec![
            vec![1, 3333, 100000000u32],
            vec![1, 2, 100000000u32],
            vec![1, 2, 100000000u32],
            vec![],
        ]);
        aux_test_union(vec![
            vec![1, 3333, 100000000u32],
            vec![1, 2, 100000000u32],
            vec![1, 2, 100000000u32],
            vec![],
        ]);
        aux_test_union(vec![
            tests::sample_with_seed(100_000, 0.01, 1),
            tests::sample_with_seed(100_000, 0.05, 2),
            tests::sample_with_seed(100_000, 0.001, 3),
        ]);
    }

    fn test_aux_union_skip(docs_list: &[Vec<DocId>], skip_targets: Vec<DocId>) {
        let mut btree_set = BTreeSet::new();
        for docs in docs_list {
            for &doc in docs.iter() {
                btree_set.insert(doc);
            }
        }
        let docset_factory = || {
            let res: Box<DocSet> = box Union::<_, DoNothingCombiner>::from(
                docs_list
                    .iter()
                    .map(|docs| docs.clone())
                    .map(VecDocSet::from)
                    .map(ConstScorer::new)
                    .collect::<Vec<_>>(),
            );
            res
        };
        let mut docset = docset_factory();
        for el in btree_set {
            assert!(docset.advance());
            assert_eq!(el, docset.doc());
        }
        assert!(!docset.advance());
        test_skip_against_unoptimized(docset_factory, skip_targets);
    }

    #[test]
    fn test_union_skip_corner_case() {
        test_aux_union_skip(&[vec![165132, 167382], vec![25029, 25091]], vec![25029]);
    }

    #[test]
    fn test_union_skip_corner_case2() {
        test_aux_union_skip(
            &[vec![1u32, 1u32 + HORIZON], vec![2u32, 1000u32, 10_000u32]],
            vec![0u32, 1u32, 2u32, 3u32, 1u32 + HORIZON, 2u32 + HORIZON],
        );
    }

    #[test]
    fn test_union_skip_corner_case3() {
        let mut docset = Union::<_, DoNothingCombiner>::from(vec![
            ConstScorer::new(VecDocSet::from(vec![0u32, 5u32])),
            ConstScorer::new(VecDocSet::from(vec![1u32, 4u32])),
        ]);
        assert!(docset.advance());
        assert_eq!(docset.doc(), 0u32);
        assert_eq!(docset.skip_next(0u32), SkipResult::OverStep);
        assert_eq!(docset.doc(), 1u32)
    }

    #[test]
    fn test_union_skip_random() {
        test_aux_union_skip(
            &[
                vec![1, 2, 3, 7],
                vec![1, 3, 9, 10000],
                vec![1, 3, 8, 9, 100],
            ],
            vec![1, 2, 3, 5, 6, 7, 8, 100],
        );
        test_aux_union_skip(
            &[
                tests::sample_with_seed(100_000, 0.001, 1),
                tests::sample_with_seed(100_000, 0.002, 2),
                tests::sample_with_seed(100_000, 0.005, 3),
            ],
            tests::sample_with_seed(100_000, 0.01, 4),
        );
    }

    #[test]
    fn test_union_skip_specific() {
        test_aux_union_skip(
            &[
                vec![1, 2, 3, 7],
                vec![1, 3, 9, 10000],
                vec![1, 3, 8, 9, 100],
            ],
            vec![1, 2, 3, 7, 8, 9, 99, 100, 101, 500, 20000],
        );
    }

    #[bench]
    fn bench_union_3_high(bench: &mut Bencher) {
        let union_docset: Vec<Vec<DocId>> = vec![
            tests::sample_with_seed(100_000, 0.1, 0),
            tests::sample_with_seed(100_000, 0.2, 1),
        ];
        bench.iter(|| {
            let mut v = Union::<_, DoNothingCombiner>::from(
                union_docset
                    .iter()
                    .map(|doc_ids| VecDocSet::from(doc_ids.clone()))
                    .map(ConstScorer::new)
                    .collect::<Vec<_>>(),
            );
            while v.advance() {}
        });
    }
    #[bench]
    fn bench_union_3_low(bench: &mut Bencher) {
        let union_docset: Vec<Vec<DocId>> = vec![
            tests::sample_with_seed(100_000, 0.01, 0),
            tests::sample_with_seed(100_000, 0.05, 1),
            tests::sample_with_seed(100_000, 0.001, 2),
        ];
        bench.iter(|| {
            let mut v = Union::<_, DoNothingCombiner>::from(
                union_docset
                    .iter()
                    .map(|doc_ids| VecDocSet::from(doc_ids.clone()))
                    .map(ConstScorer::new)
                    .collect::<Vec<_>>(),
            );
            while v.advance() {}
        });
    }

}
<|endoftext|>


<commit_before>pub(crate) mod line;

use std::{cmp::max, mem, str};

use termbox_simple::Termbox;

pub(crate) use self::line::{Line, SegStyle};
use crate::config::Colors;

pub(crate) struct MsgArea {
    lines: Vec<Line>,

    // Rendering related
    width: i32,
    height: i32,

    /// Vertical scroll: An offset from the last visible line.
    /// E.g. when this is 0, `self.lines[self.lines.len() - 1]` is drawn at the
    /// bottom of screen.
    scroll: i32,

    line_buf: Line,

    /// Cached total rendered height of all lines. Invalidate on resize, update
    /// when adding new lines.
    lines_height: Option<i32>,
}

impl MsgArea {
    pub(crate) fn new(width: i32, height: i32) -> MsgArea {
        MsgArea {
            lines: Vec::new(),
            width,
            height,
            scroll: 0,
            line_buf: Line::new(),
            lines_height: Some(0),
        }
    }

    pub(crate) fn resize(&mut self, width: i32, height: i32) {
        self.width = width;
        self.height = height;
        self.lines_height = None;
    }

    pub(crate) fn draw(&self, tb: &mut Termbox, colors: &Colors, pos_x: i32, pos_y: i32) {
        // Where to render current line
        let mut row = pos_y + self.height - 1;

        // How many visible lines to skip
        let mut skip = self.scroll;

        // Draw lines in reverse order
        let mut line_idx = (self.lines.len() as i32) - 1;
        while line_idx >= 0 && row >= pos_y {
            let line = &self.lines[line_idx as usize];
            let line_height = line.rendered_height(self.width);
            debug_assert!(line_height > 0);

            if skip >= line_height {
                // skip the whole line
                line_idx -= 1;
                skip -= line_height;
                continue;
            }

            // Rendered line height
            let height = line_height - skip;

            // Where to start rendering this line?
            let line_row = row - height + 1;

            // How many lines to skip in the `Line` before rendering
            let render_from = max(0, pos_y - line_row);

            line.draw(tb, colors, pos_x, line_row, render_from, height, self.width);
            row = line_row - 1;
            line_idx -= 1;
            skip = 0;

            if line_row < pos_y {
                break;
            }
        }
    }
}

////////////////////////////////////////////////////////////////////////////////
// Scrolling

impl MsgArea {
    fn lines_height(&mut self) -> i32 {
        match self.lines_height {
            Some(height) => height,
            None => {
                let mut total_height = 0;
                for line in &self.lines {
                    total_height += line.rendered_height(self.width);
                }
                self.lines_height = Some(total_height);
                total_height
            }
        }
    }

    pub(crate) fn scroll_up(&mut self) {
        if self.scroll < max(0, self.lines_height() - self.height) {
            self.scroll += 1;
        }
    }

    pub(crate) fn scroll_down(&mut self) {
        if self.scroll > 0 {
            self.scroll -= 1;
        }
    }

    pub(crate) fn scroll_top(&mut self) {
        self.scroll = max(0, self.lines_height() - self.height);
    }

    pub(crate) fn scroll_bottom(&mut self) {
        self.scroll = 0;
    }

    pub(crate) fn page_up(&mut self) {
        for _ in 0..10 {
            self.scroll_up();
        }
    }

    pub(crate) fn page_down(&mut self) {
        self.scroll = max(0, self.scroll - 10);
    }
}

////////////////////////////////////////////////////////////////////////////////
// Adding/removing text

impl MsgArea {
    pub(crate) fn set_style(&mut self, style: SegStyle) {
        self.line_buf.set_style(style);
    }

    pub(crate) fn add_text(&mut self, str: &str) {
        self.line_buf.add_text(str);
    }

    pub(crate) fn add_char(&mut self, char: char) {
        self.line_buf.add_char(char);
    }

    pub(crate) fn flush_line(&mut self) -> usize {
        let line_height = self.line_buf.rendered_height(self.width);
        self.lines
            .push(mem::replace(&mut self.line_buf, Line::new()));
        if self.scroll != 0 {
            self.scroll += line_height;
        }
        if let Some(ref mut total_height) = self.lines_height {
            *total_height += line_height;
        }
        self.lines.len() - 1
    }

    #[inline]
    pub(crate) fn modify_line<F>(&mut self, idx: usize, f: F)
    where
        F: Fn(&mut Line),
    {
        f(&mut self.lines[idx]);
    }

    pub(crate) fn clear(&mut self) {
        self.lines.clear();
        self.scroll = 0;
        self.lines_height = None;
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn newline_scrolling() {
        let mut msg_area = MsgArea::new(100, 1);
        // Adding a new line when scroll is 0 should not change it
        assert_eq!(msg_area.scroll, 0);
        msg_area.add_text("line1");
        msg_area.flush_line();
        assert_eq!(msg_area.scroll, 0);

        msg_area.add_text("line2");
        msg_area.flush_line();
        assert_eq!(msg_area.scroll, 0);

        msg_area.scroll_up();
        assert_eq!(msg_area.scroll, 1);
        msg_area.add_text("line3");
        msg_area.flush_line();
        assert_eq!(msg_area.scroll, 2);
    }
}
<commit_msg>tui: Remove an unused attribute<commit_after>pub(crate) mod line;

use std::{cmp::max, mem, str};

use termbox_simple::Termbox;

pub(crate) use self::line::{Line, SegStyle};
use crate::config::Colors;

pub(crate) struct MsgArea {
    lines: Vec<Line>,

    // Rendering related
    width: i32,
    height: i32,

    /// Vertical scroll: An offset from the last visible line.
    /// E.g. when this is 0, `self.lines[self.lines.len() - 1]` is drawn at the
    /// bottom of screen.
    scroll: i32,

    line_buf: Line,

    /// Cached total rendered height of all lines. Invalidate on resize, update
    /// when adding new lines.
    lines_height: Option<i32>,
}

impl MsgArea {
    pub(crate) fn new(width: i32, height: i32) -> MsgArea {
        MsgArea {
            lines: Vec::new(),
            width,
            height,
            scroll: 0,
            line_buf: Line::new(),
            lines_height: Some(0),
        }
    }

    pub(crate) fn resize(&mut self, width: i32, height: i32) {
        self.width = width;
        self.height = height;
        self.lines_height = None;
    }

    pub(crate) fn draw(&self, tb: &mut Termbox, colors: &Colors, pos_x: i32, pos_y: i32) {
        // Where to render current line
        let mut row = pos_y + self.height - 1;

        // How many visible lines to skip
        let mut skip = self.scroll;

        // Draw lines in reverse order
        let mut line_idx = (self.lines.len() as i32) - 1;
        while line_idx >= 0 && row >= pos_y {
            let line = &self.lines[line_idx as usize];
            let line_height = line.rendered_height(self.width);
            debug_assert!(line_height > 0);

            if skip >= line_height {
                // skip the whole line
                line_idx -= 1;
                skip -= line_height;
                continue;
            }

            // Rendered line height
            let height = line_height - skip;

            // Where to start rendering this line?
            let line_row = row - height + 1;

            // How many lines to skip in the `Line` before rendering
            let render_from = max(0, pos_y - line_row);

            line.draw(tb, colors, pos_x, line_row, render_from, height, self.width);
            row = line_row - 1;
            line_idx -= 1;
            skip = 0;

            if line_row < pos_y {
                break;
            }
        }
    }
}

////////////////////////////////////////////////////////////////////////////////
// Scrolling

impl MsgArea {
    fn lines_height(&mut self) -> i32 {
        match self.lines_height {
            Some(height) => height,
            None => {
                let mut total_height = 0;
                for line in &self.lines {
                    total_height += line.rendered_height(self.width);
                }
                self.lines_height = Some(total_height);
                total_height
            }
        }
    }

    pub(crate) fn scroll_up(&mut self) {
        if self.scroll < max(0, self.lines_height() - self.height) {
            self.scroll += 1;
        }
    }

    pub(crate) fn scroll_down(&mut self) {
        if self.scroll > 0 {
            self.scroll -= 1;
        }
    }

    pub(crate) fn scroll_top(&mut self) {
        self.scroll = max(0, self.lines_height() - self.height);
    }

    pub(crate) fn scroll_bottom(&mut self) {
        self.scroll = 0;
    }

    pub(crate) fn page_up(&mut self) {
        for _ in 0..10 {
            self.scroll_up();
        }
    }

    pub(crate) fn page_down(&mut self) {
        self.scroll = max(0, self.scroll - 10);
    }
}

////////////////////////////////////////////////////////////////////////////////
// Adding/removing text

impl MsgArea {
    pub(crate) fn set_style(&mut self, style: SegStyle) {
        self.line_buf.set_style(style);
    }

    pub(crate) fn add_text(&mut self, str: &str) {
        self.line_buf.add_text(str);
    }

    pub(crate) fn add_char(&mut self, char: char) {
        self.line_buf.add_char(char);
    }

    pub(crate) fn flush_line(&mut self) -> usize {
        let line_height = self.line_buf.rendered_height(self.width);
        self.lines
            .push(mem::replace(&mut self.line_buf, Line::new()));
        if self.scroll != 0 {
            self.scroll += line_height;
        }
        if let Some(ref mut total_height) = self.lines_height {
            *total_height += line_height;
        }
        self.lines.len() - 1
    }

    pub(crate) fn modify_line<F>(&mut self, idx: usize, f: F)
    where
        F: Fn(&mut Line),
    {
        f(&mut self.lines[idx]);
    }

    pub(crate) fn clear(&mut self) {
        self.lines.clear();
        self.scroll = 0;
        self.lines_height = None;
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn newline_scrolling() {
        let mut msg_area = MsgArea::new(100, 1);
        // Adding a new line when scroll is 0 should not change it
        assert_eq!(msg_area.scroll, 0);
        msg_area.add_text("line1");
        msg_area.flush_line();
        assert_eq!(msg_area.scroll, 0);

        msg_area.add_text("line2");
        msg_area.flush_line();
        assert_eq!(msg_area.scroll, 0);

        msg_area.scroll_up();
        assert_eq!(msg_area.scroll, 1);
        msg_area.add_text("line3");
        msg_area.flush_line();
        assert_eq!(msg_area.scroll, 2);
    }
}
<|endoftext|>


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

mod proto;

use log::info;
use oak::grpc;
use oak::grpc::OakNode;
use oak_derive::OakExports;
use proto::chat::{CreateRoomRequest, DestroyRoomRequest};
use proto::chat::{JoinRoomRequest, ReceivedMessage, SentMessage};
use proto::chat_grpc::{dispatch, ChatNode};
use protobuf::well_known_types::Empty;
use protobuf::ProtobufEnum;
use std::collections::HashMap;

type RoomId = Vec<u8>;
type AdminId = Vec<u8>;

#[derive(OakExports)]
struct Node {
    room_to_admin: HashMap<RoomId, AdminId>,
    admin_to_room: HashMap<AdminId, RoomId>,
    rooms: HashMap<AdminId, Room>,
}

impl OakNode for Node {
    fn new() -> Self {
        oak_log::init_default();
        Node {
            room_to_admin: HashMap::new(),
            admin_to_room: HashMap::new(),
            rooms: HashMap::new(),
        }
    }
    fn invoke(&mut self, method: &str, req: &[u8], writer: grpc::ChannelResponseWriter) {
        dispatch(self, method, req, writer)
    }
}

fn unknown_id_err(which: &str) -> grpc::Result<Empty> {
    Err(grpc::build_status(
        grpc::Code::INVALID_ARGUMENT,
        &format!("{} ID unrecognized", which),
    ))
}

fn duplicate_id_err(which: &str) -> grpc::Result<Empty> {
    Err(grpc::build_status(
        grpc::Code::ALREADY_EXISTS,
        &format!("{} ID already exists", which),
    ))
}

impl ChatNode for Node {
    fn create_room(&mut self, req: CreateRoomRequest) -> grpc::Result<Empty> {
        info!("Create new room {}", req.name);
        if self.rooms.contains_key(&req.admin_id) {
            return duplicate_id_err("Admin");
        }
        if self.room_to_admin.contains_key(&req.room_id) {
            return duplicate_id_err("Room");
        }
        self.rooms.insert(
            req.admin_id.clone(),
            Room {
                name: req.name,
                writers: Vec::new(),
            },
        );
        self.room_to_admin
            .insert(req.room_id.clone(), req.admin_id.clone());
        self.admin_to_room
            .insert(req.admin_id.clone(), req.room_id.clone());
        Ok(Empty::new())
    }

    fn destroy_room(&mut self, req: DestroyRoomRequest) -> grpc::Result<Empty> {
        if !self.rooms.contains_key(&req.admin_id) {
            return unknown_id_err("Admin");
        }
        let room = self.rooms.get_mut(&req.admin_id).unwrap();
        info!("Destroying room {}", room.name);
        room.clear();
        let user_id = self.admin_to_room.get(&req.admin_id).unwrap().clone();
        self.admin_to_room.remove(&req.admin_id);
        self.room_to_admin.remove(&user_id);
        self.rooms.remove(&req.admin_id);
        Ok(Empty::new())
    }

    fn join_room(&mut self, req: JoinRoomRequest, mut writer: grpc::ChannelResponseWriter) {
        if !self.room_to_admin.contains_key(&req.room_id) {
            writer.close(unknown_id_err("Room").map(|_| ()));
            return;
        }

        let admin_id = self.room_to_admin.get(&req.room_id).unwrap();
        let room = self.rooms.get_mut(admin_id).unwrap();
        room.writers.push(writer);

        // Send a joining announcement.
        info!("{} has joined room {}", req.user_handle, room.name);
        room.broadcast(
            "<server>",
            &format!("{} has joined room {}", req.user_handle, room.name).to_string(),
        );
    }

    fn send_message(&mut self, req: SentMessage) -> grpc::Result<Empty> {
        if !self.room_to_admin.contains_key(&req.room_id) {
            return unknown_id_err("Room");
        }
        let admin_id = self.room_to_admin.get(&req.room_id).unwrap();
        let room = self.rooms.get_mut(admin_id).unwrap();
        room.broadcast(&req.user_handle, &req.text);
        Ok(Empty::new())
    }
}

struct Room {
    name: String,
    writers: Vec<grpc::ChannelResponseWriter>,
}

impl Room {
    fn broadcast(&mut self, from: &str, text: &str) {
        info!("broadcast {}: {}", from, text);
        let mut msg = ReceivedMessage::new();
        msg.room_name = self.name.clone();
        msg.user_handle = from.to_string();
        msg.text = text.to_string();
        for writer in &mut self.writers {
            writer.write(msg.clone(), grpc::WriteMode::KeepOpen);
        }
    }
    // Close all writers associated with the room.
    fn clear(&mut self) {
        info!("clear room {}", self.name);
        for writer in &mut self.writers {
            writer.close(Ok(()));
        }
    }
}
<commit_msg>Make error functions generic (#378)<commit_after>//
// Copyright 2019 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

mod proto;

use log::info;
use oak::grpc;
use oak::grpc::OakNode;
use oak_derive::OakExports;
use proto::chat::{CreateRoomRequest, DestroyRoomRequest};
use proto::chat::{JoinRoomRequest, ReceivedMessage, SentMessage};
use proto::chat_grpc::{dispatch, ChatNode};
use protobuf::well_known_types::Empty;
use protobuf::ProtobufEnum;
use std::collections::HashMap;

type RoomId = Vec<u8>;
type AdminId = Vec<u8>;

#[derive(OakExports)]
struct Node {
    room_to_admin: HashMap<RoomId, AdminId>,
    admin_to_room: HashMap<AdminId, RoomId>,
    rooms: HashMap<AdminId, Room>,
}

impl OakNode for Node {
    fn new() -> Self {
        oak_log::init_default();
        Node {
            room_to_admin: HashMap::new(),
            admin_to_room: HashMap::new(),
            rooms: HashMap::new(),
        }
    }
    fn invoke(&mut self, method: &str, req: &[u8], writer: grpc::ChannelResponseWriter) {
        dispatch(self, method, req, writer)
    }
}

fn unknown_id_err<T>(which: &str) -> grpc::Result<T> {
    Err(grpc::build_status(
        grpc::Code::INVALID_ARGUMENT,
        &format!("{} ID unrecognized", which),
    ))
}

fn duplicate_id_err<T>(which: &str) -> grpc::Result<T> {
    Err(grpc::build_status(
        grpc::Code::ALREADY_EXISTS,
        &format!("{} ID already exists", which),
    ))
}

impl ChatNode for Node {
    fn create_room(&mut self, req: CreateRoomRequest) -> grpc::Result<Empty> {
        info!("Create new room {}", req.name);
        if self.rooms.contains_key(&req.admin_id) {
            return duplicate_id_err("Admin");
        }
        if self.room_to_admin.contains_key(&req.room_id) {
            return duplicate_id_err("Room");
        }
        self.rooms.insert(
            req.admin_id.clone(),
            Room {
                name: req.name,
                writers: Vec::new(),
            },
        );
        self.room_to_admin
            .insert(req.room_id.clone(), req.admin_id.clone());
        self.admin_to_room
            .insert(req.admin_id.clone(), req.room_id.clone());
        Ok(Empty::new())
    }

    fn destroy_room(&mut self, req: DestroyRoomRequest) -> grpc::Result<Empty> {
        if !self.rooms.contains_key(&req.admin_id) {
            return unknown_id_err("Admin");
        }
        let room = self.rooms.get_mut(&req.admin_id).unwrap();
        info!("Destroying room {}", room.name);
        room.clear();
        let user_id = self.admin_to_room.get(&req.admin_id).unwrap().clone();
        self.admin_to_room.remove(&req.admin_id);
        self.room_to_admin.remove(&user_id);
        self.rooms.remove(&req.admin_id);
        Ok(Empty::new())
    }

    fn join_room(&mut self, req: JoinRoomRequest, mut writer: grpc::ChannelResponseWriter) {
        if !self.room_to_admin.contains_key(&req.room_id) {
            writer.close(unknown_id_err("Room"));
            return;
        }

        let admin_id = self.room_to_admin.get(&req.room_id).unwrap();
        let room = self.rooms.get_mut(admin_id).unwrap();
        room.writers.push(writer);

        // Send a joining announcement.
        info!("{} has joined room {}", req.user_handle, room.name);
        room.broadcast(
            "<server>",
            &format!("{} has joined room {}", req.user_handle, room.name).to_string(),
        );
    }

    fn send_message(&mut self, req: SentMessage) -> grpc::Result<Empty> {
        if !self.room_to_admin.contains_key(&req.room_id) {
            return unknown_id_err("Room");
        }
        let admin_id = self.room_to_admin.get(&req.room_id).unwrap();
        let room = self.rooms.get_mut(admin_id).unwrap();
        room.broadcast(&req.user_handle, &req.text);
        Ok(Empty::new())
    }
}

struct Room {
    name: String,
    writers: Vec<grpc::ChannelResponseWriter>,
}

impl Room {
    fn broadcast(&mut self, from: &str, text: &str) {
        info!("broadcast {}: {}", from, text);
        let mut msg = ReceivedMessage::new();
        msg.room_name = self.name.clone();
        msg.user_handle = from.to_string();
        msg.text = text.to_string();
        for writer in &mut self.writers {
            writer.write(msg.clone(), grpc::WriteMode::KeepOpen);
        }
    }
    // Close all writers associated with the room.
    fn clear(&mut self) {
        info!("clear room {}", self.name);
        for writer in &mut self.writers {
            writer.close(Ok(()));
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate lazy_static;
extern crate typenum;
extern crate odds;

extern crate rex;

mod util;

use std::iter;
use std::fmt::Debug;
use std::path::Path;
use std::convert::{From, Into};
use std::io::{Cursor, Write, Seek, SeekFrom};

use odds::vec::VecExt;

use rex::frontend::{Event, KeyPress};

use util::mock_filesystem::{DefMockFilesystem, MockFilesystem};

fn assert_iter_eq<I, J>(one: I, other: J) where
    I: IntoIterator,
    J: IntoIterator,
    I::Item: PartialEq<J::Item> + Debug,
    J::Item: Debug,
{
    let mut one = one.into_iter();
    let mut other = other.into_iter();

    loop {
        match (one.next(), other.next()) {
            (None, None) => return,
            (Some(x), Some(y)) => assert_eq!(x, y),
            (a @ _, b @ _) => panic!("left is {:?}, right is {:?}", a, b)
        }
    }
}
<commit_msg>Add tests to the basic copy/paste<commit_after>#[macro_use]
extern crate lazy_static;
extern crate typenum;
extern crate odds;

extern crate rex;

mod util;

use std::iter;
use std::fmt::Debug;
use std::path::Path;
use std::convert::{From, Into};
use std::io::{Cursor, Write, Seek, SeekFrom};

use odds::vec::VecExt;

use rex::frontend::{Event, KeyPress};

use util::mock_filesystem::{DefMockFilesystem, MockFilesystem};

fn assert_iter_eq<I, J>(one: I, other: J) where
    I: IntoIterator,
    J: IntoIterator,
    I::Item: PartialEq<J::Item> + Debug,
    J::Item: Debug,
{
    let mut one = one.into_iter();
    let mut other = other.into_iter();

    loop {
        match (one.next(), other.next()) {
            (None, None) => return,
            (Some(x), Some(y)) => assert_eq!(x, y),
            (a @ _, b @ _) => panic!("left is {:?}, right is {:?}", a, b)
        }
    }
}

#[test]
/// Test that basic copy/paste works
fn test_copy_paste() {
    let v : Vec<u8> = (0..0xff).into_iter().collect();
    let mut v_copy = v.clone();

    // let mut result_vec = vec![];
    let (mut edit, mut frontend) = util::simple_init_with_vec(v);

    // Add some junk data in the begining
    frontend.run_keys(&mut edit, vec![KeyPress::Insert]);
    frontend.run_str(&mut edit, "AABBCCDDEE");
    v_copy.splice(..0, vec![0xAA, 0xBB, 0xCC, 0xDD, 0xEE]);

    // Copy it
    frontend.run_keys(&mut edit, vec![KeyPress::Left, KeyPress::Shortcut(' '), KeyPress::Home, KeyPress::Shortcut('c')]);

    // Paste it in the middle
    frontend.run_keys(&mut edit, vec![KeyPress::Shortcut('g')]);
    frontend.run_str(&mut edit, "50");
    frontend.run_keys(&mut edit, vec![KeyPress::Enter, KeyPress::Shortcut('v')]);
    assert_eq!(edit.get_position(), 55);
    v_copy.splice(50..50, vec![0xAA, 0xBB, 0xCC, 0xDD, 0xEE]);

    // // Paste it in the end
    frontend.run_keys(&mut edit, vec![KeyPress::PageDown, KeyPress::PageDown, KeyPress::PageDown, KeyPress::Shortcut('v')]);
    let l = v_copy.len();
    v_copy.splice(l.., vec![0xAA, 0xBB, 0xCC, 0xDD, 0xEE]);

    edit.save(Path::new("test_copy_paste"));
    println!("result_vec = {:?}", v_copy);
    assert_iter_eq(v_copy.iter(), DefMockFilesystem::get_inner("test_copy_paste").iter());
}
<|endoftext|>