<commit_before>use lazy_static::lazy_static;
use miow::iocp::CompletionPort;
use ntapi::ntioapi::FILE_OPEN;
use ntapi::ntioapi::{IO_STATUS_BLOCK_u, IO_STATUS_BLOCK};
use ntapi::ntioapi::{NtCancelIoFileEx, NtCreateFile, NtDeviceIoControlFile};
use ntapi::ntrtl::RtlNtStatusToDosError;
use std::ffi::OsStr;
use std::fmt;
use std::fs::File;
use std::io;
use std::mem::{size_of, zeroed};
use std::os::windows::ffi::OsStrExt;
use std::os::windows::io::{AsRawHandle, FromRawHandle, RawHandle};
use std::ptr::null_mut;
use std::sync::atomic::{AtomicUsize, Ordering};
use winapi::shared::ntdef::{
    HANDLE, LARGE_INTEGER, NTSTATUS, OBJECT_ATTRIBUTES, PVOID, ULONG, UNICODE_STRING,
};
use winapi::shared::ntstatus::{STATUS_NOT_FOUND, STATUS_PENDING, STATUS_SUCCESS};
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
use winapi::um::winbase::{SetFileCompletionNotificationModes, FILE_SKIP_SET_EVENT_ON_HANDLE};
use winapi::um::winnt::SYNCHRONIZE;
use winapi::um::winnt::{FILE_SHARE_READ, FILE_SHARE_WRITE};

const IOCTL_AFD_POLL: ULONG = 0x00012024;

static NEXT_TOKEN: AtomicUsize = AtomicUsize::new(0);

lazy_static! {
    static ref AFD_HELPER_NAME: Vec<u16> = {
        OsStr::new("\\Device\\Afd\\Mio")
            .encode_wide()
            .collect::<Vec<_>>()
    };
}

struct UnicodeString(UNICODE_STRING);
unsafe impl Send for UnicodeString {}
unsafe impl Sync for UnicodeString {}

struct ObjectAttributes(OBJECT_ATTRIBUTES);
unsafe impl Send for ObjectAttributes {}
unsafe impl Sync for ObjectAttributes {}

lazy_static! {
    static ref AFD_OBJ_NAME: UnicodeString = UnicodeString(UNICODE_STRING {
        // Lengths are calced in bytes
        Length: (AFD_HELPER_NAME.len() * 2) as u16,
        MaximumLength: (AFD_HELPER_NAME.len() * 2) as u16,
        Buffer: AFD_HELPER_NAME.as_ptr() as *mut _,
    });
    static ref AFD_HELPER_ATTRIBUTES: ObjectAttributes = ObjectAttributes(OBJECT_ATTRIBUTES {
        Length: size_of::<OBJECT_ATTRIBUTES>() as ULONG,
        RootDirectory: null_mut() as HANDLE,
        ObjectName: &AFD_OBJ_NAME.0 as *const _ as *mut _,
        Attributes: 0 as ULONG,
        SecurityDescriptor: null_mut() as PVOID,
        SecurityQualityOfService: null_mut() as PVOID,
    });
}

/// Winsock2 AFD driver instance.
///
/// All operations are unsafe due to IO_STATUS_BLOCK parameter are being used by Afd driver during STATUS_PENDING before I/O Completion Port returns its result.
#[derive(Debug)]
pub struct Afd {
    fd: File,
}

#[repr(C)]
#[derive(Debug)]
pub struct AfdPollHandleInfo {
    pub handle: HANDLE,
    pub events: ULONG,
    pub status: NTSTATUS,
}

unsafe impl Send for AfdPollHandleInfo {}

#[repr(C)]
pub struct AfdPollInfo {
    pub timeout: LARGE_INTEGER,
    // Can have only value 1.
    pub number_of_handles: ULONG,
    pub exclusive: ULONG,
    pub handles: [AfdPollHandleInfo; 1],
}

impl AfdPollInfo {
    pub fn zeroed() -> AfdPollInfo {
        unsafe { zeroed() }
    }
}

impl fmt::Debug for AfdPollInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AfdPollInfo").finish()
    }
}

impl Afd {
    /// Create new Afd instance.
    pub fn new(cp: &CompletionPort) -> io::Result<Afd> {
        let mut afd_helper_handle: HANDLE = INVALID_HANDLE_VALUE;
        let mut iosb = IO_STATUS_BLOCK {
            u: IO_STATUS_BLOCK_u { Status: 0 },
            Information: 0,
        };

        unsafe {
            let status = NtCreateFile(
                &mut afd_helper_handle as *mut _,
                SYNCHRONIZE,
                &AFD_HELPER_ATTRIBUTES.0 as *const _ as *mut _,
                &mut iosb,
                null_mut(),
                0 as ULONG,
                FILE_SHARE_READ | FILE_SHARE_WRITE,
                FILE_OPEN,
                0 as ULONG,
                null_mut(),
                0 as ULONG,
            );
            if status != STATUS_SUCCESS {
                return Err(io::Error::from_raw_os_error(
                    RtlNtStatusToDosError(status) as i32
                ));
            }
            let fd = File::from_raw_handle(afd_helper_handle as RawHandle);
            let token = NEXT_TOKEN.fetch_add(1, Ordering::Relaxed) + 1;
            let afd = Afd { fd };
            cp.add_handle(token, &afd.fd)?;
            match SetFileCompletionNotificationModes(
                afd_helper_handle,
                FILE_SKIP_SET_EVENT_ON_HANDLE,
            ) {
                0 => Err(io::Error::last_os_error()),
                _ => Ok(afd),
            }
        }
    }

    /// Poll `Afd` instance with `AfdPollInfo`.
    ///
    /// # Unsafety
    ///
    /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`).
    /// `iosb` needs to be untouched after the call while operation is in effective at ALL TIME except for `cancel` method.
    /// So be careful not to `poll` twice while polling.
    /// User should deallocate there overlapped value when error to prevent memory leak.
    pub unsafe fn poll(
        &self,
        info: &mut AfdPollInfo,
        iosb: *mut IO_STATUS_BLOCK,
        overlapped: PVOID,
    ) -> io::Result<bool> {
        let info_ptr: PVOID = info as *mut _ as PVOID;
        (*iosb).u.Status = STATUS_PENDING;
        let status = NtDeviceIoControlFile(
            self.fd.as_raw_handle(),
            null_mut(),
            None,
            overlapped,
            iosb,
            IOCTL_AFD_POLL,
            info_ptr,
            size_of::<AfdPollInfo>() as u32,
            info_ptr,
            size_of::<AfdPollInfo>() as u32,
        );
        match status {
            STATUS_SUCCESS => Ok(true),
            STATUS_PENDING => Ok(false),
            _ => Err(io::Error::from_raw_os_error(
                RtlNtStatusToDosError(status) as i32
            )),
        }
    }

    /// Cancel previous polled request of `Afd`.
    ///
    /// iosb needs to be used by `poll` first for valid `cancel`.
    ///
    /// # Unsafety
    ///
    /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`).
    /// Use it only with request is still being polled so that you have valid `IO_STATUS_BLOCK` to use.
    /// User should NOT deallocate there overlapped value after the `cancel` to prevent double free.
    pub unsafe fn cancel(&self, iosb: *mut IO_STATUS_BLOCK) -> io::Result<()> {
        if (*iosb).u.Status != STATUS_PENDING {
            return Ok(());
        }

        let mut cancel_iosb = IO_STATUS_BLOCK {
            u: IO_STATUS_BLOCK_u { Status: 0 },
            Information: 0,
        };
        let status = NtCancelIoFileEx(self.fd.as_raw_handle(), iosb, &mut cancel_iosb);
        if status == STATUS_SUCCESS || status == STATUS_NOT_FOUND {
            return Ok(());
        }
        Err(io::Error::from_raw_os_error(
            RtlNtStatusToDosError(status) as i32
        ))
    }
}

pub const AFD_POLL_RECEIVE: u32 = 0x0001;
pub const AFD_POLL_RECEIVE_EXPEDITED: u32 = 0x0002;
pub const AFD_POLL_SEND: u32 = 0x0004;
pub const AFD_POLL_DISCONNECT: u32 = 0x0008;
pub const AFD_POLL_ABORT: u32 = 0x0010;
pub const AFD_POLL_LOCAL_CLOSE: u32 = 0x0020;
pub const AFD_POLL_ACCEPT: u32 = 0x0080;
pub const AFD_POLL_CONNECT_FAIL: u32 = 0x0100;
pub const KNOWN_AFD_EVENTS: u32 = AFD_POLL_RECEIVE
    | AFD_POLL_RECEIVE_EXPEDITED
    | AFD_POLL_SEND
    | AFD_POLL_DISCONNECT
    | AFD_POLL_ABORT
    | AFD_POLL_ACCEPT
    | AFD_POLL_CONNECT_FAIL;
<commit_msg>Change AFD constants to use bit patterns<commit_after>use lazy_static::lazy_static;
use miow::iocp::CompletionPort;
use ntapi::ntioapi::FILE_OPEN;
use ntapi::ntioapi::{IO_STATUS_BLOCK_u, IO_STATUS_BLOCK};
use ntapi::ntioapi::{NtCancelIoFileEx, NtCreateFile, NtDeviceIoControlFile};
use ntapi::ntrtl::RtlNtStatusToDosError;
use std::ffi::OsStr;
use std::fmt;
use std::fs::File;
use std::io;
use std::mem::{size_of, zeroed};
use std::os::windows::ffi::OsStrExt;
use std::os::windows::io::{AsRawHandle, FromRawHandle, RawHandle};
use std::ptr::null_mut;
use std::sync::atomic::{AtomicUsize, Ordering};
use winapi::shared::ntdef::{
    HANDLE, LARGE_INTEGER, NTSTATUS, OBJECT_ATTRIBUTES, PVOID, ULONG, UNICODE_STRING,
};
use winapi::shared::ntstatus::{STATUS_NOT_FOUND, STATUS_PENDING, STATUS_SUCCESS};
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
use winapi::um::winbase::{SetFileCompletionNotificationModes, FILE_SKIP_SET_EVENT_ON_HANDLE};
use winapi::um::winnt::SYNCHRONIZE;
use winapi::um::winnt::{FILE_SHARE_READ, FILE_SHARE_WRITE};

const IOCTL_AFD_POLL: ULONG = 0x00012024;

static NEXT_TOKEN: AtomicUsize = AtomicUsize::new(0);

lazy_static! {
    static ref AFD_HELPER_NAME: Vec<u16> = {
        OsStr::new("\\Device\\Afd\\Mio")
            .encode_wide()
            .collect::<Vec<_>>()
    };
}

struct UnicodeString(UNICODE_STRING);
unsafe impl Send for UnicodeString {}
unsafe impl Sync for UnicodeString {}

struct ObjectAttributes(OBJECT_ATTRIBUTES);
unsafe impl Send for ObjectAttributes {}
unsafe impl Sync for ObjectAttributes {}

lazy_static! {
    static ref AFD_OBJ_NAME: UnicodeString = UnicodeString(UNICODE_STRING {
        // Lengths are calced in bytes
        Length: (AFD_HELPER_NAME.len() * 2) as u16,
        MaximumLength: (AFD_HELPER_NAME.len() * 2) as u16,
        Buffer: AFD_HELPER_NAME.as_ptr() as *mut _,
    });
    static ref AFD_HELPER_ATTRIBUTES: ObjectAttributes = ObjectAttributes(OBJECT_ATTRIBUTES {
        Length: size_of::<OBJECT_ATTRIBUTES>() as ULONG,
        RootDirectory: null_mut() as HANDLE,
        ObjectName: &AFD_OBJ_NAME.0 as *const _ as *mut _,
        Attributes: 0 as ULONG,
        SecurityDescriptor: null_mut() as PVOID,
        SecurityQualityOfService: null_mut() as PVOID,
    });
}

/// Winsock2 AFD driver instance.
///
/// All operations are unsafe due to IO_STATUS_BLOCK parameter are being used by Afd driver during STATUS_PENDING before I/O Completion Port returns its result.
#[derive(Debug)]
pub struct Afd {
    fd: File,
}

#[repr(C)]
#[derive(Debug)]
pub struct AfdPollHandleInfo {
    pub handle: HANDLE,
    pub events: ULONG,
    pub status: NTSTATUS,
}

unsafe impl Send for AfdPollHandleInfo {}

#[repr(C)]
pub struct AfdPollInfo {
    pub timeout: LARGE_INTEGER,
    // Can have only value 1.
    pub number_of_handles: ULONG,
    pub exclusive: ULONG,
    pub handles: [AfdPollHandleInfo; 1],
}

impl AfdPollInfo {
    pub fn zeroed() -> AfdPollInfo {
        unsafe { zeroed() }
    }
}

impl fmt::Debug for AfdPollInfo {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AfdPollInfo").finish()
    }
}

impl Afd {
    /// Create new Afd instance.
    pub fn new(cp: &CompletionPort) -> io::Result<Afd> {
        let mut afd_helper_handle: HANDLE = INVALID_HANDLE_VALUE;
        let mut iosb = IO_STATUS_BLOCK {
            u: IO_STATUS_BLOCK_u { Status: 0 },
            Information: 0,
        };

        unsafe {
            let status = NtCreateFile(
                &mut afd_helper_handle as *mut _,
                SYNCHRONIZE,
                &AFD_HELPER_ATTRIBUTES.0 as *const _ as *mut _,
                &mut iosb,
                null_mut(),
                0 as ULONG,
                FILE_SHARE_READ | FILE_SHARE_WRITE,
                FILE_OPEN,
                0 as ULONG,
                null_mut(),
                0 as ULONG,
            );
            if status != STATUS_SUCCESS {
                return Err(io::Error::from_raw_os_error(
                    RtlNtStatusToDosError(status) as i32
                ));
            }
            let fd = File::from_raw_handle(afd_helper_handle as RawHandle);
            let token = NEXT_TOKEN.fetch_add(1, Ordering::Relaxed) + 1;
            let afd = Afd { fd };
            cp.add_handle(token, &afd.fd)?;
            match SetFileCompletionNotificationModes(
                afd_helper_handle,
                FILE_SKIP_SET_EVENT_ON_HANDLE,
            ) {
                0 => Err(io::Error::last_os_error()),
                _ => Ok(afd),
            }
        }
    }

    /// Poll `Afd` instance with `AfdPollInfo`.
    ///
    /// # Unsafety
    ///
    /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`).
    /// `iosb` needs to be untouched after the call while operation is in effective at ALL TIME except for `cancel` method.
    /// So be careful not to `poll` twice while polling.
    /// User should deallocate there overlapped value when error to prevent memory leak.
    pub unsafe fn poll(
        &self,
        info: &mut AfdPollInfo,
        iosb: *mut IO_STATUS_BLOCK,
        overlapped: PVOID,
    ) -> io::Result<bool> {
        let info_ptr: PVOID = info as *mut _ as PVOID;
        (*iosb).u.Status = STATUS_PENDING;
        let status = NtDeviceIoControlFile(
            self.fd.as_raw_handle(),
            null_mut(),
            None,
            overlapped,
            iosb,
            IOCTL_AFD_POLL,
            info_ptr,
            size_of::<AfdPollInfo>() as u32,
            info_ptr,
            size_of::<AfdPollInfo>() as u32,
        );
        match status {
            STATUS_SUCCESS => Ok(true),
            STATUS_PENDING => Ok(false),
            _ => Err(io::Error::from_raw_os_error(
                RtlNtStatusToDosError(status) as i32
            )),
        }
    }

    /// Cancel previous polled request of `Afd`.
    ///
    /// iosb needs to be used by `poll` first for valid `cancel`.
    ///
    /// # Unsafety
    ///
    /// This function is unsafe due to memory of `IO_STATUS_BLOCK` still being used by `Afd` instance while `Ok(false)` (`STATUS_PENDING`).
    /// Use it only with request is still being polled so that you have valid `IO_STATUS_BLOCK` to use.
    /// User should NOT deallocate there overlapped value after the `cancel` to prevent double free.
    pub unsafe fn cancel(&self, iosb: *mut IO_STATUS_BLOCK) -> io::Result<()> {
        if (*iosb).u.Status != STATUS_PENDING {
            return Ok(());
        }

        let mut cancel_iosb = IO_STATUS_BLOCK {
            u: IO_STATUS_BLOCK_u { Status: 0 },
            Information: 0,
        };
        let status = NtCancelIoFileEx(self.fd.as_raw_handle(), iosb, &mut cancel_iosb);
        if status == STATUS_SUCCESS || status == STATUS_NOT_FOUND {
            return Ok(());
        }
        Err(io::Error::from_raw_os_error(
            RtlNtStatusToDosError(status) as i32
        ))
    }
}

pub const AFD_POLL_RECEIVE: u32 = 0b000_000_001;
pub const AFD_POLL_RECEIVE_EXPEDITED: u32 = 0b000_000_010;
pub const AFD_POLL_SEND: u32 = 0b000_000_100;
pub const AFD_POLL_DISCONNECT: u32 = 0b000_001_000;
pub const AFD_POLL_ABORT: u32 = 0b000_010_000;
pub const AFD_POLL_LOCAL_CLOSE: u32 = 0b000_100_000;
//pub const AFD_POLL_CONNECT: u32 = 0b001_000_000; Currently unused.
pub const AFD_POLL_ACCEPT: u32 = 0b010_000_000;
pub const AFD_POLL_CONNECT_FAIL: u32 = 0b100_000_000;

pub const KNOWN_AFD_EVENTS: u32 = AFD_POLL_RECEIVE
    | AFD_POLL_RECEIVE_EXPEDITED
    | AFD_POLL_SEND
    | AFD_POLL_DISCONNECT
    | AFD_POLL_ABORT
    | AFD_POLL_LOCAL_CLOSE
    | AFD_POLL_ACCEPT
    | AFD_POLL_CONNECT_FAIL;
<|endoftext|>


<commit_before>pub mod config;
pub mod ignore;

use constant::{WORKING_DIR_NAME, STORAGE_DIR_NAME};
use std::fs::{self, File};
use std::io::{Error as IoError, Write};
use std::path::{Path, PathBuf};

pub use self::config::{Config, ConfigKey};
pub use self::ignore::Ignore;

pub fn working_dir() -> PathBuf {
    path_buf![WORKING_DIR_NAME]
}

pub fn storage_dir() -> PathBuf {
    path_buf![working_dir(), STORAGE_DIR_NAME]
}

pub fn working_dir_exists() -> bool {
    working_dir().is_dir()
}

pub fn storage_dir_exists() -> bool {
    storage_dir().is_dir()
}

pub fn create_working_dir() -> Result<(), IoError> {
    create_essential_dir(working_dir())
}

pub fn create_storage_dir() -> Result<(), IoError> {
    create_essential_dir(storage_dir())
}

fn create_essential_dir<P: AsRef<Path>>(path_to_dir: P) -> Result<(), IoError> {
    if !path_to_dir.as_ref().is_dir() {
        try!(fs::create_dir(path_to_dir));
    }

    Ok(())
}

pub fn create_essential_dir_all<P: AsRef<Path>>(path_to_dir: P) -> Result<(), IoError> {
    try!(fs::create_dir_all(path_to_dir));

    Ok(())
}

fn create_setting_file<P: AsRef<Path>, U: AsRef<[u8]>>(path_to_file: P, contents: U) -> Result<(), IoError> {
    let mut f = try!(File::create(path_to_file));
    try!(f.write(contents.as_ref()));

    Ok(())
}

pub fn delete_all_setting_files() -> Result<(), IoError> {
    delete_dir_all(working_dir())
}

pub fn delete_dir_all<P: AsRef<Path>>(path: P) -> Result<(), IoError> {
    try!(fs::remove_dir_all(path));

    Ok(())
}
<commit_msg>[Refactoring] Remove unnecessary "pub"s<commit_after>mod config;
mod ignore;

use constant::{WORKING_DIR_NAME, STORAGE_DIR_NAME};
use std::fs::{self, File};
use std::io::{Error as IoError, Write};
use std::path::{Path, PathBuf};

pub use self::config::{Config, ConfigKey};
pub use self::ignore::Ignore;

pub fn working_dir() -> PathBuf {
    path_buf![WORKING_DIR_NAME]
}

pub fn storage_dir() -> PathBuf {
    path_buf![working_dir(), STORAGE_DIR_NAME]
}

pub fn working_dir_exists() -> bool {
    working_dir().is_dir()
}

pub fn storage_dir_exists() -> bool {
    storage_dir().is_dir()
}

pub fn create_working_dir() -> Result<(), IoError> {
    create_essential_dir(working_dir())
}

pub fn create_storage_dir() -> Result<(), IoError> {
    create_essential_dir(storage_dir())
}

fn create_essential_dir<P: AsRef<Path>>(path_to_dir: P) -> Result<(), IoError> {
    if !path_to_dir.as_ref().is_dir() {
        try!(fs::create_dir(path_to_dir));
    }

    Ok(())
}

pub fn create_essential_dir_all<P: AsRef<Path>>(path_to_dir: P) -> Result<(), IoError> {
    try!(fs::create_dir_all(path_to_dir));

    Ok(())
}

fn create_setting_file<P: AsRef<Path>, U: AsRef<[u8]>>(path_to_file: P, contents: U) -> Result<(), IoError> {
    let mut f = try!(File::create(path_to_file));
    try!(f.write(contents.as_ref()));

    Ok(())
}

pub fn delete_all_setting_files() -> Result<(), IoError> {
    delete_dir_all(working_dir())
}

pub fn delete_dir_all<P: AsRef<Path>>(path: P) -> Result<(), IoError> {
    try!(fs::remove_dir_all(path));

    Ok(())
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
const END_TAG_LEN: u32 = 8;

use core::mem::transmute;

#[macro_use]
macro_rules! ptr {
    ($it:expr) => { &$it as *const _ }
}

pub mod elf;
pub mod elf64;

#[repr(C)]
pub struct Info { pub length: u32
                , _pad: u32
                , tag_start: Tag
                }

impl Info {

    pub unsafe fn from(addr: usize) -> &'static Self {
        let info = &*(addr as *const Info);
        assert!(info.has_end());
        info
    }

    /// Finds the tag with the given tag type.
    ///
    /// This is actually safe since the tag types are constrained by The
    /// `TagType` enum
    pub fn get_tag(&self, tag_type: TagType) -> Option<&'static Tag> {
        self.tags()
            .find(|t| t.ty == tag_type)
    }

    #[inline]
    pub fn mem_map(&self) -> Option<&'static MemMapTag> {
        self.get_tag(TagType::MemoryMap)
            .map(|tag| unsafe { &*(ptr!(tag) as *const MemMapTag) })
    }

    #[inline]
    pub fn elf64_sections(&self) -> Option<&'static elf64::SectionsTag> {
        self.get_tag(TagType::ELFSections)
            .map(|tag| unsafe { &*(ptr!(tag) as *const elf64::SectionsTag) })
    }

    #[inline]
    fn tags(&self) -> Tags { Tags(ptr!(self.tag_start)) }

    fn has_end(&self) -> bool {
        let end_tag_addr
            = (self as *const _) as usize +
              (self.length - END_TAG_LEN) as usize;
        let end_tag = unsafe {&*(end_tag_addr as *const Tag)};
        end_tag.ty == TagType::End && end_tag.length == 8
    }
}

/// A Multiboot tag.
///
/// From the specification:
///
/// Boot information consists of a fixed part and a series of tags.
/// Its start is 8-bytes aligned. Fixed part is as following:
///
///<rawtext>
///             +-------------------+
///     u32     | total_size        |
///     u32     | reserved          |
///             +-------------------+
///</rawtext>
/// `total_size` contains the total size of boot information including this
/// field and terminating tag in bytes.
/// `reserved` is always set to zero and must be ignored by OS image.
///
///  Every tag begins with following fields:
///<rawtext>
///
///             +-------------------+
///     u32     | type              |
///     u32     | size              |
///             +-------------------+
///</rawtext>
/// `type` contains an identifier of contents of the rest of the tag. `size`
/// contains the size of tag including header fields but not including padding.
/// Tags follow one another padded when necessary in order for each tag to
/// start at 8-bytes aligned address. Tags are terminated by a tag of type `0`
/// and size `8`.
#[repr(C)]
#[derive(Debug)]
pub struct Tag { ty: TagType
               , length: u32
               }

/// Types of Multiboot tags
///
/// Refer to Chapter 3 of the Multiboot 2 spec
#[repr(u32)]
#[derive(Debug, Eq, PartialEq)]
pub enum TagType { /// Tag that indicates the end of multiboot tags
                   End              = 0
                 , /// Command line passed to the bootloader
                   CommandLine      = 1
                 , BootloaderName   = 2
                 , Modules          = 3
                 , BasicMemInfo     = 4
                 , BIOSBootDev      = 5
                 , MemoryMap        = 6
                 , VBEInfo          = 7
                 , FramebufferInfo  = 8
                 , ELFSections      = 9
                 , APMTable         = 10
                 }

struct Tags(*const Tag);

impl Tags {
    #[inline] fn advance(&mut self, size: u32) {
        let next_addr = self.0 as usize + size as usize;
        self.0 = (((next_addr-1) & !0x7) + 0x8) as *const _;
    }
}

impl Iterator for Tags {
    type Item = &'static Tag;

    fn next(&mut self) -> Option<Self::Item> {
        match unsafe { &*self.0 } {
            &Tag{ ty: TagType::End, length: END_TAG_LEN } => None
          , tag => { self.advance(tag.length); Some(tag) }
        }
    }
}

#[repr(C)]
pub struct MemMapTag { tag: Tag
                     , entry_size: u32
                     , entry_version: u32
                     , first_entry: MemArea
                     }
impl MemMapTag {
    pub fn entries(&self) -> Entries {
        Entries { curr: (&self.first_entry) as *const MemArea
                , last: ((self as *const MemMapTag as u32)
                            + self.tag.length
                            - self.entry_size) as *const MemArea
                , size: self.entry_size
                }
    }
}

#[repr(u32)]
#[derive(Debug, Eq, PartialEq, Ord, PartialOrd)]
pub enum MemAreaType { Available = 1
                     , ACPI      = 3
                     , Preserve  = 4
                     }

#[repr(C)]
pub struct MemArea { pub base: u64
                   , pub length: u64
                   , ty: MemAreaType
                   , _pad: u32
                   }

pub struct Entries { curr: *const MemArea
                   , last: *const MemArea
                   , size: u32
                   }

impl Iterator for Entries {
    type Item = &'static MemArea;
    fn next(&mut self) -> Option<&'static MemArea> {
        if self.curr > self.last {
            None
        } else {
            let current = unsafe { *self.curr };
            self.curr = (self.curr as u32 + self.size) as *const MemArea;
            if current.ty == MemAreaType::Available {
                Some(&current)
            } else {
                self.next()
            }
        }
    }
}
<commit_msg>Fix deref issue in multiboot<commit_after>//
//  SOS: the Stupid Operating System
//  by Hawk Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2015 Hawk Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
const END_TAG_LEN: u32 = 8;

use core::mem::transmute;

#[macro_use]
macro_rules! ptr {
    ($it:expr) => { &$it as *const _ }
}

pub mod elf;
pub mod elf64;

#[repr(C)]
pub struct Info { pub length: u32
                , _pad: u32
                , tag_start: Tag
                }

impl Info {

    pub unsafe fn from(addr: usize) -> &'static Self {
        let info = &*(addr as *const Info);
        assert!(info.has_end());
        info
    }

    /// Finds the tag with the given tag type.
    ///
    /// This is actually safe since the tag types are constrained by The
    /// `TagType` enum
    pub fn get_tag(&self, tag_type: TagType) -> Option<&'static Tag> {
        self.tags()
            .find(|t| t.ty == tag_type)
    }

    #[inline]
    pub fn mem_map(&self) -> Option<&'static MemMapTag> {
        self.get_tag(TagType::MemoryMap)
            .map(|tag| unsafe { &*(ptr!(tag) as *const MemMapTag) })
    }

    #[inline]
    pub fn elf64_sections(&self) -> Option<&'static elf64::SectionsTag> {
        self.get_tag(TagType::ELFSections)
            .map(|tag| unsafe { &*(ptr!(tag) as *const elf64::SectionsTag) })
    }

    #[inline]
    fn tags(&self) -> Tags { Tags(ptr!(self.tag_start)) }

    fn has_end(&self) -> bool {
        let end_tag_addr
            = (self as *const _) as usize +
              (self.length - END_TAG_LEN) as usize;
        let end_tag = unsafe {&*(end_tag_addr as *const Tag)};
        end_tag.ty == TagType::End && end_tag.length == 8
    }
}

/// A Multiboot tag.
///
/// From the specification:
///
/// Boot information consists of a fixed part and a series of tags.
/// Its start is 8-bytes aligned. Fixed part is as following:
///
///<rawtext>
///             +-------------------+
///     u32     | total_size        |
///     u32     | reserved          |
///             +-------------------+
///</rawtext>
/// `total_size` contains the total size of boot information including this
/// field and terminating tag in bytes.
/// `reserved` is always set to zero and must be ignored by OS image.
///
///  Every tag begins with following fields:
///<rawtext>
///
///             +-------------------+
///     u32     | type              |
///     u32     | size              |
///             +-------------------+
///</rawtext>
/// `type` contains an identifier of contents of the rest of the tag. `size`
/// contains the size of tag including header fields but not including padding.
/// Tags follow one another padded when necessary in order for each tag to
/// start at 8-bytes aligned address. Tags are terminated by a tag of type `0`
/// and size `8`.
#[repr(C)]
#[derive(Debug)]
pub struct Tag { ty: TagType
               , length: u32
               }

/// Types of Multiboot tags
///
/// Refer to Chapter 3 of the Multiboot 2 spec
#[repr(u32)]
#[derive(Debug, Eq, PartialEq)]
pub enum TagType { /// Tag that indicates the end of multiboot tags
                   End              = 0
                 , /// Command line passed to the bootloader
                   CommandLine      = 1
                 , BootloaderName   = 2
                 , Modules          = 3
                 , BasicMemInfo     = 4
                 , BIOSBootDev      = 5
                 , MemoryMap        = 6
                 , VBEInfo          = 7
                 , FramebufferInfo  = 8
                 , ELFSections      = 9
                 , APMTable         = 10
                 }

struct Tags(*const Tag);

impl Tags {
    #[inline] fn advance(&mut self, size: u32) {
        let next_addr = self.0 as usize + size as usize;
        self.0 = (((next_addr-1) & !0x7) + 0x8) as *const _;
    }
}

impl Iterator for Tags {
    type Item = &'static Tag;

    fn next(&mut self) -> Option<Self::Item> {
        match unsafe { &*self.0 } {
            &Tag{ ty: TagType::End, length: END_TAG_LEN } => None
          , tag => { self.advance(tag.length); Some(tag) }
        }
    }
}

#[repr(C)]
pub struct MemMapTag { tag: Tag
                     , entry_size: u32
                     , entry_version: u32
                     , first_entry: MemArea
                     }
impl MemMapTag {
    pub fn entries(&self) -> Entries {
        Entries { curr: (&self.first_entry) as *const MemArea
                , last: ((self as *const MemMapTag as u32)
                            + self.tag.length
                            - self.entry_size) as *const MemArea
                , size: self.entry_size
                }
    }
}

#[repr(u32)]
#[derive(Debug, Eq, PartialEq, Ord, PartialOrd)]
pub enum MemAreaType { Available = 1
                     , ACPI      = 3
                     , Preserve  = 4
                     }

#[repr(C)]
pub struct MemArea { pub base: u64
                   , pub length: u64
                   , ty: MemAreaType
                   , _pad: u32
                   }

pub struct Entries { curr: *const MemArea
                   , last: *const MemArea
                   , size: u32
                   }

impl Iterator for Entries {
    type Item = &'static MemArea;
    fn next(&mut self) -> Option<&'static MemArea> {
        if self.curr > self.last {
            None
        } else {
            let current = unsafe { &*self.curr };
            self.curr = (self.curr as u32 + self.size) as *const MemArea;
            if current.ty == MemAreaType::Available {
                Some(current)
            } else {
                self.next()
            }
        }
    }
}
<|endoftext|>


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

use std::cmp::Ordering;

use super::symbol_range::*;
use super::countable::*;

///
/// A symbol map maps from one set of symbol ranges to another
///
pub struct SymbolMap<Symbol: PartialOrd+Clone+Countable> {
    // Ranges in this symbol map
    ranges: Vec<SymbolRange<Symbol>>,
}

impl<Symbol: PartialOrd+Clone+Countable> SymbolMap<Symbol> {
    ///
    /// Creates a new symbol map
    ///
    pub fn new() -> SymbolMap<Symbol> {
        SymbolMap { ranges: vec![] }
    }

    ///
    /// Orders two symbols (makes unordered symbols equal)
    ///
    #[inline]
    fn order_symbols(a: &Symbol, b: &Symbol) -> Ordering {
        if let Some(order) = a.partial_cmp(b) {
            order
        } else {
            Ordering::Equal
        }
    }

    ///
    /// Orders two symbol ranges
    ///
    #[inline]
    fn order_ranges(a: &SymbolRange<Symbol>, b: &SymbolRange<Symbol>) -> Ordering {
        let ordering = SymbolMap::order_symbols(&a.lowest, &b.lowest);

        if ordering == Ordering::Equal {
            SymbolMap::order_symbols(&a.highest, &b.highest)
        } else {
            ordering
        }
    }

    ///
    /// Adds a range to those that are known about by this object
    ///
    pub fn add_range(&mut self, range: &SymbolRange<Symbol>) {
        let existing = self.ranges.binary_search_by(|test_range| { SymbolMap::order_ranges(test_range, range) });

        // Insert the range if it is not already in the map
        if let Err(insertion_pos) = existing {
            self.ranges.insert(insertion_pos, range.clone());
        }
    }

    ///
    /// Finds the ranges in this map that overlap the target ranges
    ///
    pub fn find_overlapping_ranges(&self, range: &SymbolRange<Symbol>) -> Vec<&SymbolRange<Symbol>> {
        let mut result = vec![];

        // Find the first range that matches (or the insertion position, which should be the first range wit a lowest value higher than the target range)
        let existing = self.ranges.binary_search_by(|test_range| { SymbolMap::order_ranges(test_range, range) });

        // Start returning values from here
        let mut pos = match existing {
            Ok(found_position) => found_position,
            Err(insert_position) => insert_position
        };

        // Move backwards if the previous position overlaps this one
        if pos > 0 && self.ranges[pos-1].highest >= range.lowest {
            pos -= 1;
        }

        // TODO: can we construct a set of ranges such that one is missed here? Think maybe we can
        while pos < self.ranges.len() && self.ranges[pos].lowest <= range.highest {
            result.push(&self.ranges[pos]);
            pos += 1;
        }

        result
    }

    ///
    /// Creates a non-overlapping range from an overlapping one
    ///
    pub fn to_non_overlapping_map(&self) -> SymbolMap<Symbol> {
        // Create a vector of all the range symbols
        let mut range_symbols: Vec<Symbol> = self.ranges.iter()
            .map(|r| r.lowest.clone())
            .chain(self.ranges.iter().map(|r| r.highest.clone()))
            .collect();

        // Sort and dedupe (we'll get pairs of ranges)
        range_symbols.sort_by(SymbolMap::order_symbols);
        range_symbols.dedup();

        // Generate a new symbol map with these ranges
        // They'll overlap with the last value, which should be OK if following the rules for adjacent ranges in symbol_range.rs
        let mut result = vec![];
        for index in 0..range_symbols.len()-2 {
            result.push(SymbolRange::new(range_symbols[index].clone(), range_symbols[index+1].prev()));
        }

        if range_symbols.len() > 1 {
            // Last item is always inclusive
            result.push(SymbolRange::new(range_symbols[range_symbols.len()-2].clone(), range_symbols[range_symbols.len()-1].clone()));
        }

        // We already sorted everything, so bypass the usual 'add' method (which sorts as it goes)
        SymbolMap { ranges: result }
    }
}

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

    #[test]
    fn can_lookup_overlapping_ranges() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 4));
        map.add_range(&SymbolRange::new(2, 5));
        map.add_range(&SymbolRange::new(3, 6));

        let bottom = map.find_overlapping_ranges(&SymbolRange::new(0, 1));
        let all    = map.find_overlapping_ranges(&SymbolRange::new(1, 3));
        let top    = map.find_overlapping_ranges(&SymbolRange::new(6, 6));

        assert!(bottom == vec![&SymbolRange::new(0, 4)]);
        assert!(all == vec![&SymbolRange::new(0, 4), &SymbolRange::new(2, 5), &SymbolRange::new(3, 6)]);
        assert!(top == vec![&SymbolRange::new(3, 6)]);
    }

    #[test]
    fn can_lookup_mid_range() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 4));
        map.add_range(&SymbolRange::new(5, 10));
        map.add_range(&SymbolRange::new(11, 15));

        let mid = map.find_overlapping_ranges(&SymbolRange::new(1, 3));

        assert!(mid == vec![&SymbolRange::new(0, 4)]);
    }

    #[test]
    fn works_with_duplicate_lower_values() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 2));
        map.add_range(&SymbolRange::new(0, 3));

        let all    = map.find_overlapping_ranges(&SymbolRange::new(0, 1));

        assert!(all == vec![&SymbolRange::new(0, 2), &SymbolRange::new(0, 3)]);
    }

    #[test]
    fn can_get_non_overlapping_map() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 4));
        map.add_range(&SymbolRange::new(2, 5));
        map.add_range(&SymbolRange::new(3, 6));

        let non_overlapping = map.to_non_overlapping_map();

        let all = non_overlapping.find_overlapping_ranges(&SymbolRange::new(0, 6));

        assert!(all == vec![&SymbolRange::new(0, 1), &SymbolRange::new(2, 2), &SymbolRange::new(3, 3), &SymbolRange::new(4, 4), &SymbolRange::new(5, 6)]);
    }

    #[test]
    fn can_get_non_overlapping_map_with_single_symbols() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 5));
        map.add_range(&SymbolRange::new(2, 2));
        map.add_range(&SymbolRange::new(3, 6));

        let non_overlapping = map.to_non_overlapping_map();

        let all = non_overlapping.find_overlapping_ranges(&SymbolRange::new(0, 6));

        assert!(all == vec![&SymbolRange::new(0, 1), &SymbolRange::new(2, 2), &SymbolRange::new(3, 4), &SymbolRange::new(5, 6)]);
    }

    #[test]
    fn can_get_non_overlapping_map_with_single_symbols_at_start() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 0));
        map.add_range(&SymbolRange::new(0, 1));

        let non_overlapping = map.to_non_overlapping_map();

        let all = non_overlapping.find_overlapping_ranges(&SymbolRange::new(0, 6));

        assert!(all == vec![&SymbolRange::new(0, 0), &SymbolRange::new(1, 1)]);
    }

    #[test]
    fn can_get_non_overlapping_map_with_single_symbols_at_start_and_gap() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 0));
        map.add_range(&SymbolRange::new(0, 1));
        map.add_range(&SymbolRange::new(3, 6));

        let non_overlapping = map.to_non_overlapping_map();

        let all = non_overlapping.find_overlapping_ranges(&SymbolRange::new(0, 6));

        assert!(all == vec![&SymbolRange::new(0, 0), &SymbolRange::new(1, 1), &SymbolRange::new(2, 2), &SymbolRange::new(3, 6)]);
    }
}
<commit_msg>It’s important we deal with the case where symbol ranges overlap by a single symbol too<commit_after>//
//   Copyright 2016 Andrew Hunter
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//

use std::cmp::Ordering;

use super::symbol_range::*;
use super::countable::*;

///
/// A symbol map maps from one set of symbol ranges to another
///
pub struct SymbolMap<Symbol: PartialOrd+Clone+Countable> {
    // Ranges in this symbol map
    ranges: Vec<SymbolRange<Symbol>>,
}

impl<Symbol: PartialOrd+Clone+Countable> SymbolMap<Symbol> {
    ///
    /// Creates a new symbol map
    ///
    pub fn new() -> SymbolMap<Symbol> {
        SymbolMap { ranges: vec![] }
    }

    ///
    /// Orders two symbols (makes unordered symbols equal)
    ///
    #[inline]
    fn order_symbols(a: &Symbol, b: &Symbol) -> Ordering {
        if let Some(order) = a.partial_cmp(b) {
            order
        } else {
            Ordering::Equal
        }
    }

    ///
    /// Orders two symbol ranges
    ///
    #[inline]
    fn order_ranges(a: &SymbolRange<Symbol>, b: &SymbolRange<Symbol>) -> Ordering {
        let ordering = SymbolMap::order_symbols(&a.lowest, &b.lowest);

        if ordering == Ordering::Equal {
            SymbolMap::order_symbols(&a.highest, &b.highest)
        } else {
            ordering
        }
    }

    ///
    /// Adds a range to those that are known about by this object
    ///
    pub fn add_range(&mut self, range: &SymbolRange<Symbol>) {
        let existing = self.ranges.binary_search_by(|test_range| { SymbolMap::order_ranges(test_range, range) });

        // Insert the range if it is not already in the map
        if let Err(insertion_pos) = existing {
            self.ranges.insert(insertion_pos, range.clone());
        }
    }

    ///
    /// Finds the ranges in this map that overlap the target ranges
    ///
    pub fn find_overlapping_ranges(&self, range: &SymbolRange<Symbol>) -> Vec<&SymbolRange<Symbol>> {
        let mut result = vec![];

        // Find the first range that matches (or the insertion position, which should be the first range wit a lowest value higher than the target range)
        let existing = self.ranges.binary_search_by(|test_range| { SymbolMap::order_ranges(test_range, range) });

        // Start returning values from here
        let mut pos = match existing {
            Ok(found_position) => found_position,
            Err(insert_position) => insert_position
        };

        // Move backwards if the previous position overlaps this one
        if pos > 0 && self.ranges[pos-1].highest >= range.lowest {
            pos -= 1;
        }

        // TODO: can we construct a set of ranges such that one is missed here? Think maybe we can
        while pos < self.ranges.len() && self.ranges[pos].lowest <= range.highest {
            result.push(&self.ranges[pos]);
            pos += 1;
        }

        result
    }

    ///
    /// Creates a non-overlapping range from an overlapping one
    ///
    pub fn to_non_overlapping_map(&self) -> SymbolMap<Symbol> {
        // Create a vector of all the range symbols
        let mut range_symbols: Vec<Symbol> = self.ranges.iter()
            .map(|r| r.lowest.clone())
            .chain(self.ranges.iter().map(|r| r.highest.clone()))
            .collect();

        // Sort and dedupe (we'll get pairs of ranges)
        range_symbols.sort_by(SymbolMap::order_symbols);
        range_symbols.dedup();

        // Generate a new symbol map with these ranges
        // They'll overlap with the last value, which should be OK if following the rules for adjacent ranges in symbol_range.rs
        let mut result = vec![];
        for index in 0..range_symbols.len()-2 {
            result.push(SymbolRange::new(range_symbols[index].clone(), range_symbols[index+1].prev()));
        }

        if range_symbols.len() > 1 {
            // Last item is always inclusive
            result.push(SymbolRange::new(range_symbols[range_symbols.len()-2].clone(), range_symbols[range_symbols.len()-1].clone()));
        }

        // We already sorted everything, so bypass the usual 'add' method (which sorts as it goes)
        SymbolMap { ranges: result }
    }
}

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

    #[test]
    fn can_lookup_overlapping_ranges() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 4));
        map.add_range(&SymbolRange::new(2, 5));
        map.add_range(&SymbolRange::new(3, 6));

        let bottom = map.find_overlapping_ranges(&SymbolRange::new(0, 1));
        let all    = map.find_overlapping_ranges(&SymbolRange::new(1, 3));
        let top    = map.find_overlapping_ranges(&SymbolRange::new(6, 6));

        assert!(bottom == vec![&SymbolRange::new(0, 4)]);
        assert!(all == vec![&SymbolRange::new(0, 4), &SymbolRange::new(2, 5), &SymbolRange::new(3, 6)]);
        assert!(top == vec![&SymbolRange::new(3, 6)]);
    }

    #[test]
    fn can_lookup_mid_range() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 4));
        map.add_range(&SymbolRange::new(5, 10));
        map.add_range(&SymbolRange::new(11, 15));

        let mid = map.find_overlapping_ranges(&SymbolRange::new(1, 3));

        assert!(mid == vec![&SymbolRange::new(0, 4)]);
    }

    #[test]
    fn works_with_duplicate_lower_values() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 2));
        map.add_range(&SymbolRange::new(0, 3));

        let all    = map.find_overlapping_ranges(&SymbolRange::new(0, 1));

        assert!(all == vec![&SymbolRange::new(0, 2), &SymbolRange::new(0, 3)]);
    }

    #[test]
    fn can_get_non_overlapping_map() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 4));
        map.add_range(&SymbolRange::new(2, 5));
        map.add_range(&SymbolRange::new(3, 6));

        let non_overlapping = map.to_non_overlapping_map();

        let all = non_overlapping.find_overlapping_ranges(&SymbolRange::new(0, 6));

        assert!(all == vec![&SymbolRange::new(0, 1), &SymbolRange::new(2, 2), &SymbolRange::new(3, 3), &SymbolRange::new(4, 4), &SymbolRange::new(5, 6)]);
    }

    #[test]
    fn can_get_non_overlapping_map_with_single_symbols() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 5));
        map.add_range(&SymbolRange::new(2, 2));
        map.add_range(&SymbolRange::new(3, 6));

        let non_overlapping = map.to_non_overlapping_map();

        let all = non_overlapping.find_overlapping_ranges(&SymbolRange::new(0, 6));

        assert!(all == vec![&SymbolRange::new(0, 1), &SymbolRange::new(2, 2), &SymbolRange::new(3, 4), &SymbolRange::new(5, 6)]);
    }

    #[test]
    fn generate_correctly_for_single_overlap() {
        let mut map = SymbolMap::new();

        // Here the symbol '5' is in both ranges, so we should generate it as a seperate range in the non-overlapping version
        map.add_range(&SymbolRange::new(0, 5));
        map.add_range(&SymbolRange::new(5, 10));

        let non_overlapping = map.to_non_overlapping_map();

        let all = non_overlapping.find_overlapping_ranges(&SymbolRange::new(0, 10));

        assert!(all == vec![&SymbolRange::new(0, 4), &SymbolRange::new(5,5), &SymbolRange::new(6, 10)]);
    }

    #[test]
    fn can_get_non_overlapping_map_with_single_symbols_at_start() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 0));
        map.add_range(&SymbolRange::new(0, 1));

        let non_overlapping = map.to_non_overlapping_map();

        let all = non_overlapping.find_overlapping_ranges(&SymbolRange::new(0, 6));

        assert!(all == vec![&SymbolRange::new(0, 0), &SymbolRange::new(1, 1)]);
    }

    #[test]
    fn can_get_non_overlapping_map_with_single_symbols_at_start_and_gap() {
        let mut map = SymbolMap::new();

        map.add_range(&SymbolRange::new(0, 0));
        map.add_range(&SymbolRange::new(0, 1));
        map.add_range(&SymbolRange::new(3, 6));

        let non_overlapping = map.to_non_overlapping_map();

        let all = non_overlapping.find_overlapping_ranges(&SymbolRange::new(0, 6));

        assert!(all == vec![&SymbolRange::new(0, 0), &SymbolRange::new(1, 1), &SymbolRange::new(2, 2), &SymbolRange::new(3, 6)]);
    }
}
<|endoftext|>


<commit_before>use std::error::Error;

use conduit_middleware;
use conduit::Request;
use conduit_cookie::RequestSession;

use Model;
use db::RequestTransaction;
use super::User;
use util::errors::{CargoResult, Unauthorized, ChainError, std_error};

pub struct Middleware;

impl conduit_middleware::Middleware for Middleware {
    fn before(&self, req: &mut Request) -> Result<(), Box<Error+Send>> {
        let id = { req.session().get("user_id").and_then(|s| s.parse().ok()) };
        let user = match id {
            Some(id) => {
                match User::find(try!(req.tx().map_err(std_error)), id) {
                    Ok(user) => user,
                    Err(..) => return Ok(()),
                }
            }
            None => {
                let tx = try!(req.tx().map_err(std_error));
                match req.headers().find("Authorization") {
                    Some(headers) => {
                        match User::find_by_api_token(tx, &headers[0]) {
                            Ok(user) => user,
                            Err(..) => return Ok(())
                        }
                    }
                    None => return Ok(())
                }
            }
        };

        req.mut_extensions().insert(user);
        Ok(())
    }
}

pub trait RequestUser {
    fn user(&self) -> CargoResult<&User>;
}

impl<'a> RequestUser for Request + 'a {
    fn user(&self) -> CargoResult<&User> {
        self.extensions().find::<User>().chain_error(|| Unauthorized)
    }
}
<commit_msg>user/middleware: Add comments to Middleware implementation<commit_after>use std::error::Error;

use conduit_middleware;
use conduit::Request;
use conduit_cookie::RequestSession;

use Model;
use db::RequestTransaction;
use super::User;
use util::errors::{CargoResult, Unauthorized, ChainError, std_error};

pub struct Middleware;

impl conduit_middleware::Middleware for Middleware {
    fn before(&self, req: &mut Request) -> Result<(), Box<Error+Send>> {
        // Check if the request has a session cookie with a `user_id` property inside
        let id = { req.session().get("user_id").and_then(|s| s.parse().ok()) };

        let user = match id {

            // `user_id` was found on the session
            Some(id) => {

                // Look for a user in the database with the given `user_id`
                match User::find(try!(req.tx().map_err(std_error)), id) {
                    Ok(user) => user,
                    Err(..) => return Ok(()),
                }
            }

            // `user_id` was *not* found on the session
            None => {

                // Look for an `Authorization` header on the request
                let tx = try!(req.tx().map_err(std_error));
                match req.headers().find("Authorization") {
                    Some(headers) => {

                        // Look for a user in the database with a matching API token
                        match User::find_by_api_token(tx, &headers[0]) {
                            Ok(user) => user,
                            Err(..) => return Ok(())
                        }
                    }
                    None => return Ok(())
                }
            }
        };

        // Attach the `User` model from the database to the request
        req.mut_extensions().insert(user);
        Ok(())
    }
}

pub trait RequestUser {
    fn user(&self) -> CargoResult<&User>;
}

impl<'a> RequestUser for Request + 'a {
    fn user(&self) -> CargoResult<&User> {
        self.extensions().find::<User>().chain_error(|| Unauthorized)
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 Walter Kuppens.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use nes::cpu::CPU;
use nes::opcode::Opcode;
use num::FromPrimitive;

/// All 6502 instructions are a maximum size of 3 bytes. The first byte is the
/// opcode which is determines the action of the instruction. The following 2
/// bytes are the arguments and are present depending on the opcode.
#[derive(Debug)]
pub struct Instruction(u8, u8, u8);

impl Instruction {
    /// The decoder will store all bytes in the instruction regardless if they
    /// are needed or not.
    ///
    /// TODO: Add page boundary flag for instructions that need extra cycles to
    /// decode across pages.
    pub fn decode(instr: &[u8], boundary: bool) -> (Instruction, u8, u16) {
        use nes::opcode::decode_opcode;

        // Determine the length of the instruction based on the opcode.
        let opcode = decode_opcode(instr[0]);
        let (len, mut cycles, boundary_sensitive) = match opcode {
            Opcode::JMPA => (3, 3, false)
        };

        // Add an additional cycle when page boundaries are crossed while
        // parsing arguments (only if the opcode is susceptible to it).
        if boundary_sensitive && boundary {
            cycles += 1;
        }

        // Return the instruction with it's arguments filled along with the
        // amount of cycles it will take for it to execute.
        (match len {
            1 => Instruction(instr[0], 0, 0),
            2 => Instruction(instr[0], instr[1], 0),
            3 => Instruction(instr[0], instr[1], instr[2]),
            _ => { panic!("Invalid length calculated"); }
        }, len, cycles)
    }

    pub fn log(&self, cpu: &CPU) {
        // Disassemble the instruction based on the opcode.
        let disassembled = match self.opcode() {
            Opcode::JMPA => format!("JMP ${:2X}{:2X}", self.2, self.1)
        };

        // TODO: Add CPU state information.
        println!("{:04X}  {:02X} {:02X} {:02X}  {:30}  A:{:02X} X:{:02X} Y:{:02X} P:{:02X} SP:{:02X} CYC:{:3}",
                 cpu.pc, self.0, self.1, self.2, disassembled, cpu.a, cpu.x,
                 cpu.y, cpu.p, cpu.sp, cpu.cycles);
    }

    #[inline(always)]
    pub fn opcode(&self) -> Opcode {
        use nes::opcode::decode_opcode;
        decode_opcode(self.0)
    }
}
<commit_msg>Moved disassembly code to it's own function<commit_after>// Copyright 2016 Walter Kuppens.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use nes::cpu::CPU;
use nes::opcode::Opcode;
use num::FromPrimitive;

/// All 6502 instructions are a maximum size of 3 bytes. The first byte is the
/// opcode which is determines the action of the instruction. The following 2
/// bytes are the arguments and are present depending on the opcode.
#[derive(Debug)]
pub struct Instruction(u8, u8, u8);

impl Instruction {
    /// The decoder will store all bytes in the instruction regardless if they
    /// are needed or not.
    ///
    /// TODO: Add page boundary flag for instructions that need extra cycles to
    /// decode across pages.
    pub fn decode(instr: &[u8], boundary: bool) -> (Instruction, u8, u16) {
        use nes::opcode::decode_opcode;

        // Determine the length of the instruction based on the opcode.
        let opcode = decode_opcode(instr[0]);
        let (len, mut cycles, boundary_sensitive) = match opcode {
            Opcode::JMPA => (3, 3, false)
        };

        // Add an additional cycle when page boundaries are crossed while
        // parsing arguments (only if the opcode is susceptible to it).
        if boundary_sensitive && boundary {
            cycles += 1;
        }

        // Return the instruction with it's arguments filled along with the
        // amount of cycles it will take for it to execute.
        (match len {
            1 => Instruction(instr[0], 0, 0),
            2 => Instruction(instr[0], instr[1], 0),
            3 => Instruction(instr[0], instr[1], instr[2]),
            _ => { panic!("Invalid length calculated"); }
        }, len, cycles)
    }

    /// Disassembles the instruction into human readable assembly.
    pub fn disassemble(&self) -> String {
        match self.opcode() {
            Opcode::JMPA => format!("JMP ${:2X}{:2X}", self.2, self.1)
        }
    }

    /// Logs a human-readable representation of the instruction along with the
    /// CPU state in an easy to parse format.
    ///
    /// TODO: Return a string for the test suite so CPU correctness can be
    /// checked.
    pub fn log(&self, cpu: &CPU) {
        // Disassemble the instruction based on the opcode.
        let disassembled = self.disassemble();

        // Prints the CPU state in a nice parsable format. In the future this
        // output will be used for automatically testing the CPU's accuracy.
        println!("{:04X}  {:02X} {:02X} {:02X}  {:30}  A:{:02X} X:{:02X} Y:{:02X} P:{:02X} SP:{:02X} CYC:{:3}",
                 cpu.pc, self.0, self.1, self.2, disassembled, cpu.a, cpu.x,
                 cpu.y, cpu.p, cpu.sp, cpu.cycles);
    }

    /// Obtain the opcode of the instruction.
    #[inline(always)]
    pub fn opcode(&self) -> Opcode {
        use nes::opcode::decode_opcode;
        decode_opcode(self.0)
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

use std::cmp;

#[derive(Debug, PartialEq)]
enum ProgressState {
    Probe,
    Replicate,
    Snapshot,
}

struct Progress {
    matched: u64,
    next_idx: u64,
    // When in ProgressStateProbe, leader sends at most one replication message
    // per heartbeat interval. It also probes actual progress of the follower.
    //
    // When in ProgressStateReplicate, leader optimistically increases next
    // to the latest entry sent after sending replication message. This is
    // an optimized state for fast replicating log entries to the follower.
    //
    // When in ProgressStateSnapshot, leader should have sent out snapshot
    // before and stops sending any replication message.
    state: ProgressState,
    // Paused is used in ProgressStateProbe.
    // When Paused is true, raft should pause sending replication message to this peer.
    paused: bool,
    // pending_snapshot is used in ProgressStateSnapshot.
    // If there is a pending snapshot, the pendingSnapshot will be set to the
    // index of the snapshot. If pendingSnapshot is set, the replication process of
    // this Progress will be paused. raft will not resend snapshot until the pending one
    // is reported to be failed.
    pending_snapshot: u64,

    // recent_active is true if the progress is recently active. Receiving any messages
    // from the corresponding follower indicates the progress is active.
    // RecentActive can be reset to false after an election timeout.
    recent_active: bool,

    // Inflights is a sliding window for the inflight messages.
    // When inflights is full, no more message should be sent.
    // When a leader sends out a message, the index of the last
    // entry should be added to inflights. The index MUST be added
    // into inflights in order.
    // When a leader receives a reply, the previous inflights should
    // be freed by calling inflights.freeTo.
    ins: Inflights,
}

impl Progress {
    pub fn reset_state(&mut self, state: ProgressState) {
        self.paused = false;
        self.recent_active = false;
        self.pending_snapshot = 0;
        self.state = state;
        self.ins.reset();
    }

    pub fn become_probe(&mut self) {
        // If the original state is ProgressStateSnapshot, progress knows that
        // the pending snapshot has been sent to this peer successfully, then
        // probes from pendingSnapshot + 1.
        if self.state == ProgressState::Snapshot {
            let pending_snapshot = self.pending_snapshot;
            self.reset_state(ProgressState::Probe);
            self.next_idx = cmp::max(self.matched + 1, pending_snapshot + 1);
        } else {
            self.reset_state(ProgressState::Probe);
            self.next_idx = self.matched + 1;
        }
    }

    pub fn become_replicate(&mut self) {
        self.reset_state(ProgressState::Replicate);
        self.next_idx = self.matched + 1;
    }

    pub fn become_snapshot(&mut self, snapshot_idx: u64) {
        self.reset_state(ProgressState::Snapshot);
        self.pending_snapshot = snapshot_idx;
    }

    pub fn snapshot_failure(&self) -> bool {
        self.pending_snapshot == 0
    }

    // maybe_snapshot_abort unsets pendingSnapshot if Match is equal or higher than
    // the pendingSnapshot
    pub fn maybe_snapshot_abort(&self) -> bool {
        self.state == ProgressState::Snapshot && self.matched >= self.pending_snapshot
    }

    // maybe_update returns false if the given n index comes from an outdated message.
    // Otherwise it updates the progress and returns true.
    pub fn maybe_update(&mut self, n: u64) -> bool {
        let mut updated = false;
        if self.matched < n {
            self.matched = n;
            updated = true;
            self.resume()
        }

        if self.next_idx < n + 1 {
            self.next_idx = n + 1
        }

        updated
    }

    pub fn optimistic_update(&mut self, n: u64) {
        self.next_idx = n + 1;
    }

    // maybe_decr_to returns false if the given to index comes from an out of order message.
    // Otherwise it decreases the progress next index to min(rejected, last) and returns true.
    pub fn maybe_decr_to(&mut self, rejected: u64, last: u64) -> bool {
        if self.state == ProgressState::Replicate {
            // the rejection must be stale if the progress has matched and "rejected"
            // is smaller than "match".
            if rejected <= self.matched {
                return false;
            }
            self.next_idx = self.matched + 1;
            return true;
        }

        // the rejection must be stale if "rejected" does not match next - 1
        if self.next_idx - 1 != rejected {
            return false;
        }

        self.next_idx = cmp::min(rejected, last + 1);
        if self.next_idx < 1 {
            self.next_idx = 1;
        }
        self.resume();
        true
    }

    pub fn is_paused(&self) -> bool {
        match self.state {
            ProgressState::Probe => self.paused,
            ProgressState::Replicate => self.ins.full(),
            ProgressState::Snapshot => true,
        }
    }

    pub fn resume(&mut self) {
        self.paused = false;
    }

    pub fn pause(&mut self) {
        self.paused = true;
    }
}

#[derive(Debug, Default)]
struct Inflights {
    // the starting index in the buffer
    start: usize,
    // number of inflights in the buffer
    count: usize,

    // ring buffer
    buffer: Vec<u64>,
}

impl Inflights {
    pub fn new(cap: usize) -> Inflights {
        Inflights { buffer: Vec::with_capacity(cap), ..Default::default() }
    }

    // full returns true if the inflights is full.
    pub fn full(&self) -> bool {
        self.count == self.cap()
    }

    pub fn cap(&self) -> usize {
        self.buffer.capacity()
    }

    // add adds an inflight into inflights
    pub fn add(&mut self, inflight: u64) {
        if self.full() {
            panic!("cannot add into a full inflights")
        }

        let mut next = self.start + self.count;
        if next >= self.cap() {
            next -= self.cap();
        }
        self.buffer[next] = inflight;
        self.count += 1;
    }

    // free_to frees the inflights smaller or equal to the given `to` flight.
    pub fn free_to(&mut self, to: u64) {
        if self.count == 0 || to < self.buffer[self.start] {
            // out of the left side of the window
            return;
        }

        let mut i = 0usize;
        let mut idx = self.start;
        while i < self.count {
            if to < self.buffer[idx] {
                // found the first large inflight
                break;
            }

            // increase index and maybe rotate
            idx += 1;
            if idx >= self.cap() {
                idx -= self.cap();
            }

            i += 1;
        }

        // free i inflights and set new start index
        self.count -= i;
        self.start = idx;
    }

    pub fn free_first_one(&mut self) {
        let start = self.start;
        self.free_to(start);
    }

    // resets frees all inflights.
    pub fn reset(&mut self) {
        self.count = 0;
        self.start = 0;
    }
}
<commit_msg>Fix free_first_one<commit_after>#![allow(dead_code)]

use std::cmp;

#[derive(Debug, PartialEq)]
enum ProgressState {
    Probe,
    Replicate,
    Snapshot,
}

struct Progress {
    matched: u64,
    next_idx: u64,
    // When in ProgressStateProbe, leader sends at most one replication message
    // per heartbeat interval. It also probes actual progress of the follower.
    //
    // When in ProgressStateReplicate, leader optimistically increases next
    // to the latest entry sent after sending replication message. This is
    // an optimized state for fast replicating log entries to the follower.
    //
    // When in ProgressStateSnapshot, leader should have sent out snapshot
    // before and stops sending any replication message.
    state: ProgressState,
    // Paused is used in ProgressStateProbe.
    // When Paused is true, raft should pause sending replication message to this peer.
    paused: bool,
    // pending_snapshot is used in ProgressStateSnapshot.
    // If there is a pending snapshot, the pendingSnapshot will be set to the
    // index of the snapshot. If pendingSnapshot is set, the replication process of
    // this Progress will be paused. raft will not resend snapshot until the pending one
    // is reported to be failed.
    pending_snapshot: u64,

    // recent_active is true if the progress is recently active. Receiving any messages
    // from the corresponding follower indicates the progress is active.
    // RecentActive can be reset to false after an election timeout.
    recent_active: bool,

    // Inflights is a sliding window for the inflight messages.
    // When inflights is full, no more message should be sent.
    // When a leader sends out a message, the index of the last
    // entry should be added to inflights. The index MUST be added
    // into inflights in order.
    // When a leader receives a reply, the previous inflights should
    // be freed by calling inflights.freeTo.
    ins: Inflights,
}

impl Progress {
    pub fn reset_state(&mut self, state: ProgressState) {
        self.paused = false;
        self.recent_active = false;
        self.pending_snapshot = 0;
        self.state = state;
        self.ins.reset();
    }

    pub fn become_probe(&mut self) {
        // If the original state is ProgressStateSnapshot, progress knows that
        // the pending snapshot has been sent to this peer successfully, then
        // probes from pendingSnapshot + 1.
        if self.state == ProgressState::Snapshot {
            let pending_snapshot = self.pending_snapshot;
            self.reset_state(ProgressState::Probe);
            self.next_idx = cmp::max(self.matched + 1, pending_snapshot + 1);
        } else {
            self.reset_state(ProgressState::Probe);
            self.next_idx = self.matched + 1;
        }
    }

    pub fn become_replicate(&mut self) {
        self.reset_state(ProgressState::Replicate);
        self.next_idx = self.matched + 1;
    }

    pub fn become_snapshot(&mut self, snapshot_idx: u64) {
        self.reset_state(ProgressState::Snapshot);
        self.pending_snapshot = snapshot_idx;
    }

    pub fn snapshot_failure(&self) -> bool {
        self.pending_snapshot == 0
    }

    // maybe_snapshot_abort unsets pendingSnapshot if Match is equal or higher than
    // the pendingSnapshot
    pub fn maybe_snapshot_abort(&self) -> bool {
        self.state == ProgressState::Snapshot && self.matched >= self.pending_snapshot
    }

    // maybe_update returns false if the given n index comes from an outdated message.
    // Otherwise it updates the progress and returns true.
    pub fn maybe_update(&mut self, n: u64) -> bool {
        let mut updated = false;
        if self.matched < n {
            self.matched = n;
            updated = true;
            self.resume()
        }

        if self.next_idx < n + 1 {
            self.next_idx = n + 1
        }

        updated
    }

    pub fn optimistic_update(&mut self, n: u64) {
        self.next_idx = n + 1;
    }

    // maybe_decr_to returns false if the given to index comes from an out of order message.
    // Otherwise it decreases the progress next index to min(rejected, last) and returns true.
    pub fn maybe_decr_to(&mut self, rejected: u64, last: u64) -> bool {
        if self.state == ProgressState::Replicate {
            // the rejection must be stale if the progress has matched and "rejected"
            // is smaller than "match".
            if rejected <= self.matched {
                return false;
            }
            self.next_idx = self.matched + 1;
            return true;
        }

        // the rejection must be stale if "rejected" does not match next - 1
        if self.next_idx - 1 != rejected {
            return false;
        }

        self.next_idx = cmp::min(rejected, last + 1);
        if self.next_idx < 1 {
            self.next_idx = 1;
        }
        self.resume();
        true
    }

    pub fn is_paused(&self) -> bool {
        match self.state {
            ProgressState::Probe => self.paused,
            ProgressState::Replicate => self.ins.full(),
            ProgressState::Snapshot => true,
        }
    }

    pub fn resume(&mut self) {
        self.paused = false;
    }

    pub fn pause(&mut self) {
        self.paused = true;
    }
}

#[derive(Debug, Default)]
struct Inflights {
    // the starting index in the buffer
    start: usize,
    // number of inflights in the buffer
    count: usize,

    // ring buffer
    buffer: Vec<u64>,
}

impl Inflights {
    pub fn new(cap: usize) -> Inflights {
        Inflights { buffer: Vec::with_capacity(cap), ..Default::default() }
    }

    // full returns true if the inflights is full.
    pub fn full(&self) -> bool {
        self.count == self.cap()
    }

    pub fn cap(&self) -> usize {
        self.buffer.capacity()
    }

    // add adds an inflight into inflights
    pub fn add(&mut self, inflight: u64) {
        if self.full() {
            panic!("cannot add into a full inflights")
        }

        let mut next = self.start + self.count;
        if next >= self.cap() {
            next -= self.cap();
        }
        self.buffer[next] = inflight;
        self.count += 1;
    }

    // free_to frees the inflights smaller or equal to the given `to` flight.
    pub fn free_to(&mut self, to: u64) {
        if self.count == 0 || to < self.buffer[self.start] {
            // out of the left side of the window
            return;
        }

        let mut i = 0usize;
        let mut idx = self.start;
        while i < self.count {
            if to < self.buffer[idx] {
                // found the first large inflight
                break;
            }

            // increase index and maybe rotate
            idx += 1;
            if idx >= self.cap() {
                idx -= self.cap();
            }

            i += 1;
        }

        // free i inflights and set new start index
        self.count -= i;
        self.start = idx;
    }

    pub fn free_first_one(&mut self) {
        let start = self.buffer[self.start];
        self.free_to(start);
    }

    // resets frees all inflights.
    pub fn reset(&mut self) {
        self.count = 0;
        self.start = 0;
    }
}
<|endoftext|>


<commit_before>use std::thread;
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::sync::mpsc::{Sender,Receiver};
use std::collections::HashMap;
use consts::{WatchedEventType, ZkError, ZkState};
use proto::WatchedEvent;
use zookeeper::{ZkResult, ZooKeeper};
use zookeeper_ext::ZooKeeperExt;
use watch;
use listeners::{Subscription};

pub type Data = HashMap<String, Arc<Vec<u8>>>;

#[derive(Debug)]
pub struct ChildData;

#[derive(Debug)]
pub enum PathChildrenCacheEvent {
    Initialized,
    ConnectionSuspended,
    ConnectionLost,
    ConnectionReconnected,
    ChildRemoved(ChildData),
    ChildAdded(ChildData),
    ChildUpdated(ChildData),
}

#[derive(Debug)]
pub enum RefreshMode {
    Standard,
    ForceGetDataAndStat,
}

pub enum Operation {
    Shutdown(Option<Sender<ZkResult<bool>>>),
    Refresh(RefreshMode, Option<Sender<ZkResult<()>>>),
    Event(PathChildrenCacheEvent, Option<Sender<ZkResult<bool>>>),
    GetData(String /* path */, Option<Sender<ZkResult<Vec<u8>>>>),
}

use std::fmt;
impl fmt::Debug for Operation {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Operation::Shutdown(ref _sender) => {
                write!(f, "Operation::Shutdown")
            },
            Operation::Refresh(ref mode, ref _sender) => {
                write!(f, "Operation::Refresh({:?})", mode)
            },
            Operation::Event(ref ev, ref _sender) => {
                write!(f, "Operation::Event({:?})", ev)
            },
            Operation::GetData(ref path, ref _sender) => {
                write!(f, "Operation::GetData({})", path)
            }
        }
    }
}

pub struct PathChildrenCache {
    path: Arc<String>,
    zk: Arc<ZooKeeper>,
    data: Arc<Mutex<Data>>,
    worker_thread: Option<thread::JoinHandle<()>>,
    channel: Option<Sender<Operation>>,
    listener_subscription: Option<Subscription>,
}

impl PathChildrenCache {

    fn get_children(zk: Arc<ZooKeeper>, path: &str, data: Arc<Mutex<Data>>, ops_chan: Sender<Operation>) -> ZkResult<()> {

        let zk1 = zk.clone();
        let data1 = data.clone();
        let ops_chan1 = ops_chan.clone();

        let watcher = move |event: &WatchedEvent| {
            match event.event_type {
                WatchedEventType::NodeChildrenChanged => {
                    let path = event.path.as_ref().expect("Path absent");

                    // Subscribe to new changes recursively
                    ops_chan1.send(Operation::Refresh(RefreshMode::Standard, None));
                },
                _ => error!("Unexpected: {:?}", event)
            };
            // Send change event
        };

        let children = try!(zk.get_children_w(&path, watcher));

        let mut data_locked = data.lock().unwrap();

        for child in children.iter() {
            let child_path = join_path(path, child);

            if !data_locked.contains_key(&child_path) {

                let child_data = try!(Self::get_data(zk.clone(), &child_path, data.clone(), ops_chan.clone()));

                data_locked.insert(child_path, Arc::new(child_data));
            }
        }

        debug!("New data: {:?}", *data_locked);

        Ok(())
    }

    fn get_data(zk: Arc<ZooKeeper>, path: &str, data: Arc<Mutex<Data>>, ops_chan: Sender<Operation>) -> ZkResult<Vec<u8>> {
        let zk1 = zk.clone();
        let path1 = path.to_owned();

        let data_watcher = move |event: &WatchedEvent| {
            let mut data_locked = data.lock().unwrap();
            match event.event_type {
                WatchedEventType::NodeDeleted => {
                    data_locked.remove(&path1);
                },
                WatchedEventType::NodeDataChanged => {
                    // Subscribe to new changes recursively
                    ops_chan.send(Operation::GetData(path1, None));
                },
                _ => error!("Unexpected: {:?}", event)
            };

            debug!("New data: {:?}", *data_locked);

            //TODO Send change event
        };

        zk.get_data_w(path, data_watcher).map(|stuff| { stuff.0 })
    }

    pub fn new(zk: Arc<ZooKeeper>, path: &str) -> ZkResult<PathChildrenCache> {

        let data = Arc::new(Mutex::new(HashMap::new()));

        try!(zk.ensure_path(path));

        Ok(PathChildrenCache{
            path: Arc::new(path.to_string()),
            zk: zk,
            data: data,
            worker_thread: None,
            channel: None,
            listener_subscription: None,
        })
    }

    pub fn get_current_data(&self) -> Data {
        self.data.lock().unwrap().clone()
    }

    pub fn clear(&self) {
        self.data.lock().unwrap().clear()
    }

    pub fn start(&mut self) {
        let (ops_chan_tx, ops_chan_rx) = mpsc::channel();
        let (listener_chan_tx, listener_chan_rx) = mpsc::channel();

        let sub = self.zk.add_listener(listener_chan_tx.clone());
        self.listener_subscription = Some(sub);

        let zk = self.zk.clone();
        let path = self.path.clone();
        let data = self.data.clone();
        self.channel = Some(ops_chan_tx.clone());

        self.worker_thread = Some(thread::spawn(move || {
            let mut done = false;

            while !done {
                select! (
                    state = listener_chan_rx.recv() => {
                        match state {
                            Ok(state) => {
                                info!("zk state change {:?}", state);
                            },
                            Err(err) => {
                                info!("zk state chan err. shutting down. {:?}", err);
                                done = true;
                            }
                        }
                    },
                    op = ops_chan_rx.recv() => {

                        debug!("handling op {:?}", op);
                        match op {
                            Ok(Operation::Shutdown(maybe_chan)) => {
                                info!("shutting down worker thread");
                                break;
                            },
                            Ok(Operation::Refresh(mode, maybe_chan)) => {
                                debug!("getting children");
                                let result = Self::get_children(zk.clone(), &*path, data.clone(), ops_chan_tx.clone());
                                info!("got children {:?}", result);

                                maybe_chan.map_or((), |chan| {
                                    chan.send(result);
                                });
                            },
                            Ok(Operation::GetData(path, maybe_chan)) => {
                                debug!("getting data");
                                let result = Self::get_data(zk.clone(), &*path, data.clone(), ops_chan_tx.clone());
                                info!("got data {:?}", result);

                                maybe_chan.map_or((), |chan| {
                                    chan.send(result);
                                });
                            },
                            Ok(Operation::Event(event, maybe_chan)) => {
                            },
                            Err(err) => {
                                info!("error receiving from operations channel. shutting down worker thread");
                                done = true;
                            }
                        }

                    });
            }
        }));

        self.offer_operation(Operation::Refresh(RefreshMode::ForceGetDataAndStat, None));
    }

    fn offer_operation(&self, op: Operation) -> ZkResult<()> {
        match self.channel {
            Some(ref chan) => {
                chan.send(op);
                Ok(())
            },
            None => Err(ZkError::APIError)
        }
    }

}

pub fn join_path(dir: &str, child: &str) -> String {
    let dir_bytes = dir.as_bytes();
    let mut result = dir.to_string();
    if dir_bytes[dir_bytes.len() - 1] != ('/' as u8) {
        result.push_str("/");
    }
    result.push_str(child);
    result
}
<commit_msg>handle reconnect with a full refresh<commit_after>use std::thread;
use std::sync::{Arc, Mutex};
use std::sync::mpsc;
use std::sync::mpsc::{Sender,Receiver};
use std::collections::HashMap;
use consts::{WatchedEventType, ZkError, ZkState};
use proto::WatchedEvent;
use zookeeper::{ZkResult, ZooKeeper};
use zookeeper_ext::ZooKeeperExt;
use watch;
use listeners::{Subscription};

pub type Data = HashMap<String, Arc<Vec<u8>>>;

#[derive(Debug)]
pub struct ChildData;

#[derive(Debug)]
pub enum PathChildrenCacheEvent {
    Initialized,
    ConnectionSuspended,
    ConnectionLost,
    ConnectionReconnected,
    ChildRemoved(ChildData),
    ChildAdded(ChildData),
    ChildUpdated(ChildData),
}

#[derive(Debug)]
pub enum RefreshMode {
    Standard,
    ForceGetDataAndStat,
}

pub enum Operation {
    Shutdown(Option<Sender<ZkResult<bool>>>),
    Refresh(RefreshMode, Option<Sender<ZkResult<()>>>),
    Event(PathChildrenCacheEvent, Option<Sender<ZkResult<bool>>>),
    GetData(String /* path */, Option<Sender<ZkResult<Vec<u8>>>>),
}

use std::fmt;
impl fmt::Debug for Operation {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Operation::Shutdown(ref _sender) => {
                write!(f, "Operation::Shutdown")
            },
            Operation::Refresh(ref mode, ref _sender) => {
                write!(f, "Operation::Refresh({:?})", mode)
            },
            Operation::Event(ref ev, ref _sender) => {
                write!(f, "Operation::Event({:?})", ev)
            },
            Operation::GetData(ref path, ref _sender) => {
                write!(f, "Operation::GetData({})", path)
            }
        }
    }
}

pub struct PathChildrenCache {
    path: Arc<String>,
    zk: Arc<ZooKeeper>,
    data: Arc<Mutex<Data>>,
    worker_thread: Option<thread::JoinHandle<()>>,
    channel: Option<Sender<Operation>>,
    listener_subscription: Option<Subscription>,
}

impl PathChildrenCache {

    fn get_children(zk: Arc<ZooKeeper>, path: &str, data: Arc<Mutex<Data>>, ops_chan: Sender<Operation>) -> ZkResult<()> {

        let zk1 = zk.clone();
        let data1 = data.clone();
        let ops_chan1 = ops_chan.clone();

        let watcher = move |event: &WatchedEvent| {
            match event.event_type {
                WatchedEventType::NodeChildrenChanged => {
                    let path = event.path.as_ref().expect("Path absent");

                    // Subscribe to new changes recursively
                    ops_chan1.send(Operation::Refresh(RefreshMode::Standard, None));
                },
                _ => error!("Unexpected: {:?}", event)
            };
            // Send change event
        };

        let children = try!(zk.get_children_w(&path, watcher));

        let mut data_locked = data.lock().unwrap();

        for child in children.iter() {
            let child_path = join_path(path, child);

            if !data_locked.contains_key(&child_path) {

                let child_data = try!(Self::get_data(zk.clone(), &child_path, data.clone(), ops_chan.clone()));

                data_locked.insert(child_path, Arc::new(child_data));
            }
        }

        debug!("New data: {:?}", *data_locked);

        Ok(())
    }

    fn get_data(zk: Arc<ZooKeeper>, path: &str, data: Arc<Mutex<Data>>, ops_chan: Sender<Operation>) -> ZkResult<Vec<u8>> {
        let zk1 = zk.clone();
        let path1 = path.to_owned();

        let data_watcher = move |event: &WatchedEvent| {
            let mut data_locked = data.lock().unwrap();
            match event.event_type {
                WatchedEventType::NodeDeleted => {
                    data_locked.remove(&path1);
                },
                WatchedEventType::NodeDataChanged => {
                    // Subscribe to new changes recursively
                    ops_chan.send(Operation::GetData(path1, None));
                },
                _ => error!("Unexpected: {:?}", event)
            };

            debug!("New data: {:?}", *data_locked);

            //TODO Send change event
        };

        zk.get_data_w(path, data_watcher).map(|stuff| { stuff.0 })
    }

    pub fn new(zk: Arc<ZooKeeper>, path: &str) -> ZkResult<PathChildrenCache> {

        let data = Arc::new(Mutex::new(HashMap::new()));

        try!(zk.ensure_path(path));

        Ok(PathChildrenCache{
            path: Arc::new(path.to_string()),
            zk: zk,
            data: data,
            worker_thread: None,
            channel: None,
            listener_subscription: None,
        })
    }

    pub fn get_current_data(&self) -> Data {
        self.data.lock().unwrap().clone()
    }

    pub fn clear(&self) {
        self.data.lock().unwrap().clear()
    }

    pub fn start(&mut self) {
        let (ops_chan_tx, ops_chan_rx) = mpsc::channel();
        let (listener_chan_tx, listener_chan_rx) = mpsc::channel();

        let sub = self.zk.add_listener(listener_chan_tx.clone());
        self.listener_subscription = Some(sub);

        let zk = self.zk.clone();
        let path = self.path.clone();
        let data = self.data.clone();
        self.channel = Some(ops_chan_tx.clone());

        self.worker_thread = Some(thread::spawn(move || {
            let mut done = false;

            while !done {
                select! (
                    state = listener_chan_rx.recv() => {
                        match state {
                            Ok(state) => {
                                info!("zk state change {:?}", state);
                                match state {
                                    ZkState::Connected => {
                                        ops_chan_tx.send(Operation::Refresh(RefreshMode::ForceGetDataAndStat, None));
                                    },
                                    _ => {
                                    }
                                }
                            },
                            Err(err) => {
                                info!("zk state chan err. shutting down. {:?}", err);
                                done = true;
                            }
                        }
                    },
                    op = ops_chan_rx.recv() => {

                        debug!("handling op {:?}", op);
                        match op {
                            Ok(Operation::Shutdown(maybe_chan)) => {
                                info!("shutting down worker thread");
                                break;
                            },
                            Ok(Operation::Refresh(mode, maybe_chan)) => {
                                debug!("getting children");
                                let result = Self::get_children(zk.clone(), &*path, data.clone(), ops_chan_tx.clone());
                                info!("got children {:?}", result);

                                maybe_chan.map_or((), |chan| {
                                    chan.send(result);
                                });
                            },
                            Ok(Operation::GetData(path, maybe_chan)) => {
                                debug!("getting data");
                                let result = Self::get_data(zk.clone(), &*path, data.clone(), ops_chan_tx.clone());
                                info!("got data {:?}", result);

                                maybe_chan.map_or((), |chan| {
                                    chan.send(result);
                                });
                            },
                            Ok(Operation::Event(event, maybe_chan)) => {
                            },
                            Err(err) => {
                                info!("error receiving from operations channel. shutting down worker thread");
                                done = true;
                            }
                        }

                    });
            }
        }));

        self.offer_operation(Operation::Refresh(RefreshMode::ForceGetDataAndStat, None));
    }

    fn offer_operation(&self, op: Operation) -> ZkResult<()> {
        match self.channel {
            Some(ref chan) => {
                chan.send(op);
                Ok(())
            },
            None => Err(ZkError::APIError)
        }
    }

}

pub fn join_path(dir: &str, child: &str) -> String {
    let dir_bytes = dir.as_bytes();
    let mut result = dir.to_string();
    if dir_bytes[dir_bytes.len() - 1] != ('/' as u8) {
        result.push_str("/");
    }
    result.push_str(child);
    result
}
<|endoftext|>


<commit_before>/*
    Copyright 2013 Jesse 'Jeaye' Wilkerson
    See licensing in LICENSE file, or at:
        http://www.opensource.org/licenses/BSD-3-Clause

    File: obj/voxel/chunk.rs
    Author: Jesse 'Jeaye' Wilkerson
    Description:
      A pageable 3D chunk of voxel data.
*/

use math::{ Vec3f, Vec3i8 };
use primitive::{ Cube, Cube_Index };
use super::Behavior;

#[path = "../../gl/mod.rs"]
mod gl;
#[path = "../../gl/util.rs"]
mod util;
#[macro_escape]
#[path = "../../gl/check.rs"]
mod check;

pub struct Chunk
{
  vbo: gl::GLuint,
  dimensions: Vec3i8,
  voxels: ~[Cube],
  indices: ~[Cube_Index],
  behaviors: ~[Behavior],
}

impl Chunk
{
  pub fn new() -> Chunk
  {
    let chunk = Chunk
    {
      vbo: check!(gl::gen_buffers(1))[0],
      dimensions: Vec3i8::new(16, 16, 16),
      voxels: ~[],
      indices: ~[],
      behaviors: ~[],
    };

    chunk
  }

  //pub fn get(&self, x: i8, y: i8, z: i8) -> Behavior
  //{
  //  voxels[(z * self.dimensions.z * self.dimensions.z) + (y * self.dimensions.y) + x].get_behavior()
  //}

  pub fn draw(&self)
  {

  }
}

<commit_msg> #12 Removed warning in chunk<commit_after>/*
    Copyright 2013 Jesse 'Jeaye' Wilkerson
    See licensing in LICENSE file, or at:
        http://www.opensource.org/licenses/BSD-3-Clause

    File: obj/voxel/chunk.rs
    Author: Jesse 'Jeaye' Wilkerson
    Description:
      A pageable 3D chunk of voxel data.
*/

use math::{ Vec3i8 };
use primitive::{ Cube, Cube_Index };
use super::Behavior;

#[path = "../../gl/mod.rs"]
mod gl;
#[path = "../../gl/util.rs"]
mod util;
#[macro_escape]
#[path = "../../gl/check.rs"]
mod check;

pub struct Chunk
{
  vbo: gl::GLuint,
  dimensions: Vec3i8,
  voxels: ~[Cube],
  indices: ~[Cube_Index],
  behaviors: ~[Behavior],
}

impl Chunk
{
  pub fn new() -> Chunk
  {
    let chunk = Chunk
    {
      vbo: check!(gl::gen_buffers(1))[0],
      dimensions: Vec3i8::new(16, 16, 16),
      voxels: ~[],
      indices: ~[],
      behaviors: ~[],
    };

    chunk
  }

  //pub fn get(&self, x: i8, y: i8, z: i8) -> Behavior
  //{
  //  voxels[(z * self.dimensions.z * self.dimensions.z) + (y * self.dimensions.y) + x].get_behavior()
  //}

  pub fn draw(&self)
  {

  }
}

<|endoftext|>


<commit_before>use std::mem::MaybeUninit;

/// Encode u64 as varint.
/// Panics if buffer length is less than 10.
#[inline]
pub fn encode_varint64(mut value: u64, buf: &mut [MaybeUninit<u8>]) -> usize {
    assert!(buf.len() >= 10);

    fn iter(value: &mut u64, byte: &mut MaybeUninit<u8>) -> bool {
        if (*value & !0x7F) > 0 {
            byte.write(((*value & 0x7F) | 0x80) as u8);
            *value >>= 7;
            true
        } else {
            byte.write(*value as u8);
            false
        }
    }

    // Explicitly unroll loop to avoid either
    // unsafe code or bound checking when writing to `buf`

    if !iter(&mut value, &mut buf[0]) {
        return 1;
    };
    if !iter(&mut value, &mut buf[1]) {
        return 2;
    };
    if !iter(&mut value, &mut buf[2]) {
        return 3;
    };
    if !iter(&mut value, &mut buf[3]) {
        return 4;
    };
    if !iter(&mut value, &mut buf[4]) {
        return 5;
    };
    if !iter(&mut value, &mut buf[5]) {
        return 6;
    };
    if !iter(&mut value, &mut buf[6]) {
        return 7;
    };
    if !iter(&mut value, &mut buf[7]) {
        return 8;
    };
    if !iter(&mut value, &mut buf[8]) {
        return 9;
    };
    buf[9].write(value as u8);
    10
}

/// Encode u32 value as varint.
/// Panics if buffer length is less than 5.
#[inline]
pub fn encode_varint32(mut value: u32, buf: &mut [MaybeUninit<u8>]) -> usize {
    assert!(buf.len() >= 5);

    fn iter(value: &mut u32, byte: &mut MaybeUninit<u8>) -> bool {
        if (*value & !0x7F) > 0 {
            byte.write(((*value & 0x7F) | 0x80) as u8);
            *value >>= 7;
            true
        } else {
            byte.write(*value as u8);
            false
        }
    }

    // Explicitly unroll loop to avoid either
    // unsafe code or bound checking when writing to `buf`

    if !iter(&mut value, &mut buf[0]) {
        return 1;
    };
    if !iter(&mut value, &mut buf[1]) {
        return 2;
    };
    if !iter(&mut value, &mut buf[2]) {
        return 3;
    };
    if !iter(&mut value, &mut buf[3]) {
        return 4;
    };
    buf[4].write(value as u8);
    5
}

/// Encoded size of u64 value.
pub(crate) fn encoded_varint64_len(value: u64) -> usize {
    if value == 0 {
        1
    } else {
        let significant_bits = 64 - value.leading_zeros();
        (significant_bits + 6) as usize / 7
    }
}

#[cfg(test)]
mod test {
    use std::mem::MaybeUninit;

    use crate::varint::encode_varint64;
    use crate::varint::encoded_varint64_len;

    #[test]
    fn test_encoded_varint64_len() {
        fn test(n: u64) {
            let mut buf = [MaybeUninit::uninit(); 10];
            let expected = encode_varint64(n, &mut buf);
            assert_eq!(expected, encoded_varint64_len(n), "n={}", n);
        }

        for n in 0..1000 {
            test(n);
        }

        for p in 0.. {
            match 2u64.checked_pow(p) {
                Some(n) => test(n),
                None => break,
            }
        }

        for p in 0.. {
            match 3u64.checked_pow(p) {
                Some(n) => test(n),
                None => break,
            }
        }

        test(u64::MAX);
        test(u64::MAX - 1);
        test((i64::MAX as u64) + 1);
        test(i64::MAX as u64);
        test((i64::MAX as u64) - 1);
        test((u32::MAX as u64) + 1);
        test(u32::MAX as u64);
        test((u32::MAX as u64) - 1);
        test((i32::MAX as u64) + 1);
        test(i32::MAX as u64);
        test((i32::MAX as u64) - 1);
    }
}
<commit_msg>Reduce function visibility for clarity<commit_after>use std::mem::MaybeUninit;

/// Encode u64 as varint.
/// Panics if buffer length is less than 10.
#[inline]
pub(crate) fn encode_varint64(mut value: u64, buf: &mut [MaybeUninit<u8>]) -> usize {
    assert!(buf.len() >= 10);

    fn iter(value: &mut u64, byte: &mut MaybeUninit<u8>) -> bool {
        if (*value & !0x7F) > 0 {
            byte.write(((*value & 0x7F) | 0x80) as u8);
            *value >>= 7;
            true
        } else {
            byte.write(*value as u8);
            false
        }
    }

    // Explicitly unroll loop to avoid either
    // unsafe code or bound checking when writing to `buf`

    if !iter(&mut value, &mut buf[0]) {
        return 1;
    };
    if !iter(&mut value, &mut buf[1]) {
        return 2;
    };
    if !iter(&mut value, &mut buf[2]) {
        return 3;
    };
    if !iter(&mut value, &mut buf[3]) {
        return 4;
    };
    if !iter(&mut value, &mut buf[4]) {
        return 5;
    };
    if !iter(&mut value, &mut buf[5]) {
        return 6;
    };
    if !iter(&mut value, &mut buf[6]) {
        return 7;
    };
    if !iter(&mut value, &mut buf[7]) {
        return 8;
    };
    if !iter(&mut value, &mut buf[8]) {
        return 9;
    };
    buf[9].write(value as u8);
    10
}

/// Encode u32 value as varint.
/// Panics if buffer length is less than 5.
#[inline]
pub(crate) fn encode_varint32(mut value: u32, buf: &mut [MaybeUninit<u8>]) -> usize {
    assert!(buf.len() >= 5);

    fn iter(value: &mut u32, byte: &mut MaybeUninit<u8>) -> bool {
        if (*value & !0x7F) > 0 {
            byte.write(((*value & 0x7F) | 0x80) as u8);
            *value >>= 7;
            true
        } else {
            byte.write(*value as u8);
            false
        }
    }

    // Explicitly unroll loop to avoid either
    // unsafe code or bound checking when writing to `buf`

    if !iter(&mut value, &mut buf[0]) {
        return 1;
    };
    if !iter(&mut value, &mut buf[1]) {
        return 2;
    };
    if !iter(&mut value, &mut buf[2]) {
        return 3;
    };
    if !iter(&mut value, &mut buf[3]) {
        return 4;
    };
    buf[4].write(value as u8);
    5
}

/// Encoded size of u64 value.
pub(crate) fn encoded_varint64_len(value: u64) -> usize {
    if value == 0 {
        1
    } else {
        let significant_bits = 64 - value.leading_zeros();
        (significant_bits + 6) as usize / 7
    }
}

#[cfg(test)]
mod test {
    use std::mem::MaybeUninit;

    use crate::varint::encode_varint64;
    use crate::varint::encoded_varint64_len;

    #[test]
    fn test_encoded_varint64_len() {
        fn test(n: u64) {
            let mut buf = [MaybeUninit::uninit(); 10];
            let expected = encode_varint64(n, &mut buf);
            assert_eq!(expected, encoded_varint64_len(n), "n={}", n);
        }

        for n in 0..1000 {
            test(n);
        }

        for p in 0.. {
            match 2u64.checked_pow(p) {
                Some(n) => test(n),
                None => break,
            }
        }

        for p in 0.. {
            match 3u64.checked_pow(p) {
                Some(n) => test(n),
                None => break,
            }
        }

        test(u64::MAX);
        test(u64::MAX - 1);
        test((i64::MAX as u64) + 1);
        test(i64::MAX as u64);
        test((i64::MAX as u64) - 1);
        test((u32::MAX as u64) + 1);
        test(u32::MAX as u64);
        test((u32::MAX as u64) - 1);
        test((i32::MAX as u64) + 1);
        test(i32::MAX as u64);
        test((i32::MAX as u64) - 1);
    }
}
<|endoftext|>