<commit_before>#![deny(trivial_numeric_casts, unsafe_code, unstable_features)]
#![warn(
    missing_debug_implementations,
    missing_docs,
    unused_qualifications,
    unused_import_braces
)]
//! This crate contains the DICOM transfer syntax registry.
//! The transfer syntax registry maps a DICOM UID of a transfer syntax into the
//! respective transfer syntax specifier. In the default implementation, the
//! container of transfer syntaxes is populated before-main through the
//! [inventory] pattern, then making all registerd TSes readily available
//! through the [`TransferSyntaxRegistry`] type.
//!
//! The default Cargo feature `inventory-registry` can be deactivated for
//! environments which do not support `inventory`, with the downside of only
//! providing the built-in transfer syntaxes.
//!
//! This registry should not have to be used directly, except when developing
//! higher level APIs, which should learn to negotiate and resolve the expected
//! transfer syntax automatically.
//!
//! ## Transfer Syntax descriptors
//!
//! This crate encompasses the basic DICOM level of conformance:
//! _Implicit VR Little Endian_,
//! _Explicit VR Little Endian_,
//! and _Explicit VR Big Endian_ are built-in.
//! Transfer syntaxes which are not supported,
//! or which rely on encapsulated pixel data,
//! are only listed as _stubs_ to be replaced by separate libraries.
//! The full list is available in the [`entries`](entries) module.
//!
//! [inventory]: https://docs.rs/inventory/0.1.4/inventory

use byteordered::Endianness;
use dicom_encoding::transfer_syntax::{
    AdapterFreeTransferSyntax as Ts, Codec, TransferSyntaxIndex,
};
use lazy_static::lazy_static;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::fmt;

pub use dicom_encoding::TransferSyntax;
pub mod entries;

/// Data type for a registry of DICOM.
pub struct TransferSyntaxRegistryImpl {
    m: HashMap<&'static str, TransferSyntax>,
}

impl fmt::Debug for TransferSyntaxRegistryImpl {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let entries: HashMap<&str, &str> =
            self.m.iter().map(|(uid, ts)| (*uid, ts.name())).collect();
        f.debug_struct("TransferSyntaxRegistryImpl")
            .field("m", &entries)
            .finish()
    }
}

impl TransferSyntaxRegistryImpl {
    /// Obtain a DICOM codec by transfer syntax UID.
    fn get<U: AsRef<str>>(&self, uid: U) -> Option<&TransferSyntax> {
        let ts_uid = uid
            .as_ref()
            .trim_end_matches(|c: char| c.is_whitespace() || c == '\0');
        self.m.get(ts_uid)
    }

    /// Register the given transfer syntax (TS) to the system. It can override
    /// another TS with the same UID, in the only case that the TS requires
    /// certain codecs which are not supported by the previously registered
    /// TS. If no such requirements are imposed, this function returns `false`
    /// and no changes are made.
    fn register(&mut self, ts: TransferSyntax) -> bool {
        match self.m.entry(ts.uid()) {
            Entry::Occupied(mut e) => {
                let replace = match (&e.get().codec(), ts.codec()) {
                    (Codec::Unsupported, Codec::Dataset(_))
                    | (Codec::EncapsulatedPixelData, Codec::PixelData(_)) => true,
                    // weird one ahead: the two specifiers do not agree on
                    // requirements, better keep it as a separate match arm for
                    // debugging purposes
                    (Codec::Unsupported, Codec::PixelData(_)) => {
                        tracing::warn!("Inconsistent requirements for transfer syntax {}: `Unsupported` cannot be replaced with `PixelData`", ts.uid());
                        false
                    }
                    // ignoring TS with less or equal implementation
                    _ => false,
                };

                if replace {
                    e.insert(ts);
                    true
                } else {
                    false
                }
            }
            Entry::Vacant(e) => {
                e.insert(ts);
                true
            }
        }
    }
}

impl TransferSyntaxIndex for TransferSyntaxRegistryImpl {
    #[inline]
    fn get(&self, uid: &str) -> Option<&TransferSyntax> {
        Self::get(self, uid)
    }
}

/// Zero-sized representative of the main transfer syntax registry.
#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Hash)]
pub struct TransferSyntaxRegistry;

impl TransferSyntaxIndex for TransferSyntaxRegistry {
    #[inline]
    fn get(&self, uid: &str) -> Option<&TransferSyntax> {
        get_registry().get(uid)
    }
}

lazy_static! {

    static ref REGISTRY: TransferSyntaxRegistryImpl = {
        let mut registry = TransferSyntaxRegistryImpl {
            m: HashMap::with_capacity(32),
        };

        use self::entries::*;
        let built_in_ts: [TransferSyntax; 29] = [
            IMPLICIT_VR_LITTLE_ENDIAN.erased(),
            EXPLICIT_VR_LITTLE_ENDIAN.erased(),
            EXPLICIT_VR_BIG_ENDIAN.erased(),

            DEFLATED_EXPLICIT_VR_LITTLE_ENDIAN.erased(),
            JPIP_REFERENCED_DEFLATE.erased(),
            JPEG_BASELINE.erased(),
            JPEG_EXTENDED.erased(),
            JPEG_LOSSLESS_NON_HIERARCHICAL.erased(),
            JPEG_LOSSLESS_NON_HIERARCHICAL_FIRST_ORDER_PREDICTION.erased(),
            JPEG_LS_LOSSLESS_IMAGE_COMPRESSION.erased(),
            JPEG_LS_LOSSY_IMAGE_COMPRESSION.erased(),
            JPEG_2000_IMAGE_COMPRESSION_LOSSLESS_ONLY.erased(),
            JPEG_2000_IMAGE_COMPRESSION.erased(),
            JPEG_2000_PART2_MULTI_COMPONENT_IMAGE_COMPRESSION_LOSSLESS_ONLY.erased(),
            JPEG_2000_PART2_MULTI_COMPONENT_IMAGE_COMPRESSION.erased(),
            JPIP_REFERENCED.erased(),
            MPEG2_MAIN_PROFILE_MAIN_LEVEL.erased(),
            MPEG2_MAIN_PROFILE_HIGH_LEVEL.erased(),
            MPEG4_AVC_H264_HIGH_PROFILE.erased(),
            MPEG4_AVC_H264_BD_COMPATIBLE_HIGH_PROFILE.erased(),
            MPEG4_AVC_H264_HIGH_PROFILE_FOR_2D_VIDEO.erased(),
            MPEG4_AVC_H264_HIGH_PROFILE_FOR_3D_VIDEO.erased(),
            MPEG4_AVC_H264_STEREO_HIGH_PROFILE.erased(),
            HEVC_H265_MAIN_PROFILE.erased(),
            HEVC_H265_MAIN_10_PROFILE.erased(),
            RLE_LOSSLESS.erased(),
            SMPTE_ST_2110_20_UNCOMPRESSED_PROGRESSIVE.erased(),
            SMPTE_ST_2110_20_UNCOMPRESSED_INTERLACED.erased(),
            SMPTE_ST_2110_30_PCM.erased(),
        ];

        // add built-in TSes manually
        for ts in built_in_ts {
            registry.register(ts);
        }
        // add TSes from inventory, if available
        inventory_populate(&mut registry);

        registry
    };
}

#[cfg(feature = "inventory-registry")]
#[inline]
fn inventory_populate(registry: &mut TransferSyntaxRegistryImpl) {
    use dicom_encoding::transfer_syntax::TransferSyntaxFactory;

    for TransferSyntaxFactory(tsf) in inventory::iter::<TransferSyntaxFactory> {
        let ts = tsf();
        registry.register(ts);
    }
}

#[cfg(not(feature = "inventory-registry"))]
#[inline]
fn inventory_populate(_: &mut TransferSyntaxRegistryImpl) {
    // do nothing
}

/// Retrieve a reference to the global codec registry.
#[inline]
pub(crate) fn get_registry() -> &'static TransferSyntaxRegistryImpl {
    &REGISTRY
}

/// create a TS with an unsupported pixel encapsulation
pub(crate) const fn create_ts_stub(uid: &'static str, name: &'static str) -> Ts {
    TransferSyntax::new(
        uid,
        name,
        Endianness::Little,
        true,
        Codec::EncapsulatedPixelData,
    )
}

/// Retrieve the default transfer syntax.
pub fn default() -> Ts {
    entries::IMPLICIT_VR_LITTLE_ENDIAN
}
<commit_msg>[ts-registry] add TransferSyntaxRegistry::iter<commit_after>#![deny(trivial_numeric_casts, unsafe_code, unstable_features)]
#![warn(
    missing_debug_implementations,
    missing_docs,
    unused_qualifications,
    unused_import_braces
)]
//! This crate contains the DICOM transfer syntax registry.
//! The transfer syntax registry maps a DICOM UID of a transfer syntax into the
//! respective transfer syntax specifier. In the default implementation, the
//! container of transfer syntaxes is populated before-main through the
//! [inventory] pattern, then making all registerd TSes readily available
//! through the [`TransferSyntaxRegistry`] type.
//!
//! The default Cargo feature `inventory-registry` can be deactivated for
//! environments which do not support `inventory`, with the downside of only
//! providing the built-in transfer syntaxes.
//!
//! This registry should not have to be used directly, except when developing
//! higher level APIs, which should learn to negotiate and resolve the expected
//! transfer syntax automatically.
//!
//! ## Transfer Syntax descriptors
//!
//! This crate encompasses the basic DICOM level of conformance:
//! _Implicit VR Little Endian_,
//! _Explicit VR Little Endian_,
//! and _Explicit VR Big Endian_ are built-in.
//! Transfer syntaxes which are not supported,
//! or which rely on encapsulated pixel data,
//! are only listed as _stubs_ to be replaced by separate libraries.
//! The full list is available in the [`entries`](entries) module.
//!
//! [inventory]: https://docs.rs/inventory/0.1.4/inventory

use byteordered::Endianness;
use dicom_encoding::transfer_syntax::{
    AdapterFreeTransferSyntax as Ts, Codec, TransferSyntaxIndex,
};
use lazy_static::lazy_static;
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::fmt;

pub use dicom_encoding::TransferSyntax;
pub mod entries;

/// Data type for a registry of DICOM.
pub struct TransferSyntaxRegistryImpl {
    m: HashMap<&'static str, TransferSyntax>,
}

impl fmt::Debug for TransferSyntaxRegistryImpl {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let entries: HashMap<&str, &str> =
            self.m.iter().map(|(uid, ts)| (*uid, ts.name())).collect();
        f.debug_struct("TransferSyntaxRegistryImpl")
            .field("m", &entries)
            .finish()
    }
}

impl TransferSyntaxRegistryImpl {

    /// Obtain an iterator of all registered transfer syntaxes.
    pub fn iter(&self) -> impl Iterator<Item = &TransferSyntax> {
        self.m.values()
    }

    /// Obtain a DICOM codec by transfer syntax UID.
    fn get<U: AsRef<str>>(&self, uid: U) -> Option<&TransferSyntax> {
        let ts_uid = uid
            .as_ref()
            .trim_end_matches(|c: char| c.is_whitespace() || c == '\0');
        self.m.get(ts_uid)
    }

    /// Register the given transfer syntax (TS) to the system. It can override
    /// another TS with the same UID, in the only case that the TS requires
    /// certain codecs which are not supported by the previously registered
    /// TS. If no such requirements are imposed, this function returns `false`
    /// and no changes are made.
    fn register(&mut self, ts: TransferSyntax) -> bool {
        match self.m.entry(ts.uid()) {
            Entry::Occupied(mut e) => {
                let replace = match (&e.get().codec(), ts.codec()) {
                    (Codec::Unsupported, Codec::Dataset(_))
                    | (Codec::EncapsulatedPixelData, Codec::PixelData(_)) => true,
                    // weird one ahead: the two specifiers do not agree on
                    // requirements, better keep it as a separate match arm for
                    // debugging purposes
                    (Codec::Unsupported, Codec::PixelData(_)) => {
                        tracing::warn!("Inconsistent requirements for transfer syntax {}: `Unsupported` cannot be replaced with `PixelData`", ts.uid());
                        false
                    }
                    // ignoring TS with less or equal implementation
                    _ => false,
                };

                if replace {
                    e.insert(ts);
                    true
                } else {
                    false
                }
            }
            Entry::Vacant(e) => {
                e.insert(ts);
                true
            }
        }
    }
}

impl TransferSyntaxIndex for TransferSyntaxRegistryImpl {
    #[inline]
    fn get(&self, uid: &str) -> Option<&TransferSyntax> {
        Self::get(self, uid)
    }
}

impl TransferSyntaxRegistry {
    /// Obtain an iterator of all registered transfer syntaxes.
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = &TransferSyntax> {
        get_registry().iter()
    }
}

/// Zero-sized representative of the main transfer syntax registry.
#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Hash)]
pub struct TransferSyntaxRegistry;

impl TransferSyntaxIndex for TransferSyntaxRegistry {
    #[inline]
    fn get(&self, uid: &str) -> Option<&TransferSyntax> {
        get_registry().get(uid)
    }
}

lazy_static! {

    static ref REGISTRY: TransferSyntaxRegistryImpl = {
        let mut registry = TransferSyntaxRegistryImpl {
            m: HashMap::with_capacity(32),
        };

        use self::entries::*;
        let built_in_ts: [TransferSyntax; 29] = [
            IMPLICIT_VR_LITTLE_ENDIAN.erased(),
            EXPLICIT_VR_LITTLE_ENDIAN.erased(),
            EXPLICIT_VR_BIG_ENDIAN.erased(),

            DEFLATED_EXPLICIT_VR_LITTLE_ENDIAN.erased(),
            JPIP_REFERENCED_DEFLATE.erased(),
            JPEG_BASELINE.erased(),
            JPEG_EXTENDED.erased(),
            JPEG_LOSSLESS_NON_HIERARCHICAL.erased(),
            JPEG_LOSSLESS_NON_HIERARCHICAL_FIRST_ORDER_PREDICTION.erased(),
            JPEG_LS_LOSSLESS_IMAGE_COMPRESSION.erased(),
            JPEG_LS_LOSSY_IMAGE_COMPRESSION.erased(),
            JPEG_2000_IMAGE_COMPRESSION_LOSSLESS_ONLY.erased(),
            JPEG_2000_IMAGE_COMPRESSION.erased(),
            JPEG_2000_PART2_MULTI_COMPONENT_IMAGE_COMPRESSION_LOSSLESS_ONLY.erased(),
            JPEG_2000_PART2_MULTI_COMPONENT_IMAGE_COMPRESSION.erased(),
            JPIP_REFERENCED.erased(),
            MPEG2_MAIN_PROFILE_MAIN_LEVEL.erased(),
            MPEG2_MAIN_PROFILE_HIGH_LEVEL.erased(),
            MPEG4_AVC_H264_HIGH_PROFILE.erased(),
            MPEG4_AVC_H264_BD_COMPATIBLE_HIGH_PROFILE.erased(),
            MPEG4_AVC_H264_HIGH_PROFILE_FOR_2D_VIDEO.erased(),
            MPEG4_AVC_H264_HIGH_PROFILE_FOR_3D_VIDEO.erased(),
            MPEG4_AVC_H264_STEREO_HIGH_PROFILE.erased(),
            HEVC_H265_MAIN_PROFILE.erased(),
            HEVC_H265_MAIN_10_PROFILE.erased(),
            RLE_LOSSLESS.erased(),
            SMPTE_ST_2110_20_UNCOMPRESSED_PROGRESSIVE.erased(),
            SMPTE_ST_2110_20_UNCOMPRESSED_INTERLACED.erased(),
            SMPTE_ST_2110_30_PCM.erased(),
        ];

        // add built-in TSes manually
        for ts in built_in_ts {
            registry.register(ts);
        }
        // add TSes from inventory, if available
        inventory_populate(&mut registry);

        registry
    };
}

#[cfg(feature = "inventory-registry")]
#[inline]
fn inventory_populate(registry: &mut TransferSyntaxRegistryImpl) {
    use dicom_encoding::transfer_syntax::TransferSyntaxFactory;

    for TransferSyntaxFactory(tsf) in inventory::iter::<TransferSyntaxFactory> {
        let ts = tsf();
        registry.register(ts);
    }
}

#[cfg(not(feature = "inventory-registry"))]
#[inline]
fn inventory_populate(_: &mut TransferSyntaxRegistryImpl) {
    // do nothing
}

/// Retrieve a reference to the global codec registry.
#[inline]
pub(crate) fn get_registry() -> &'static TransferSyntaxRegistryImpl {
    &REGISTRY
}

/// create a TS with an unsupported pixel encapsulation
pub(crate) const fn create_ts_stub(uid: &'static str, name: &'static str) -> Ts {
    TransferSyntax::new(
        uid,
        name,
        Endianness::Little,
        true,
        Codec::EncapsulatedPixelData,
    )
}

/// Retrieve the default transfer syntax.
pub fn default() -> Ts {
    entries::IMPLICIT_VR_LITTLE_ENDIAN
}

#[cfg(test)]
mod tests {
    use dicom_encoding::TransferSyntaxIndex;

    use crate::TransferSyntaxRegistry;

    #[test]
    fn has_mandatory_tss() {
        let implicit_vr_le = TransferSyntaxRegistry.get("1.2.840.10008.1.2")
            .expect("transfer syntax registry should provide Implicit VR Little Endian");
        assert_eq!(implicit_vr_le.uid(), "1.2.840.10008.1.2");
        assert!(implicit_vr_le.fully_supported());

        // should also work with trailing null character
        let implicit_vr_le_2 = TransferSyntaxRegistry.get("1.2.840.10008.1.2\0")
            .expect("transfer syntax registry should provide Implicit VR Little Endian with padded TS UID");

        assert_eq!(implicit_vr_le_2.uid(), implicit_vr_le.uid());

        let explicit_vr_le = TransferSyntaxRegistry.get("1.2.840.10008.1.2.1")
            .expect("transfer syntax registry should provide Explicit VR Little Endian");
        assert_eq!(explicit_vr_le.uid(), "1.2.840.10008.1.2.1");
        assert!(explicit_vr_le.fully_supported());

        // should also work with trailing null character
        let explicit_vr_le_2 = TransferSyntaxRegistry.get("1.2.840.10008.1.2.1\0")
            .expect("transfer syntax registry should provide Explicit VR Little Endian with padded TS UID");

        assert_eq!(explicit_vr_le_2.uid(), explicit_vr_le.uid());
    }

    #[test]
    fn provides_iter() {
        let all_tss: Vec<_> = TransferSyntaxRegistry.iter().collect();

        assert!(all_tss.len() >= 2);

        // contains at least Implicit VR Little Endian and Explicit VR Little Endian
        assert!(all_tss.iter().any(|ts| ts.uid() == "1.2.840.10008.1.2"));
        assert!(all_tss.iter().any(|ts| ts.uid() == "1.2.840.10008.1.2.1"));
    }
}
<|endoftext|>


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

#[derive(Debug, RustcEncodable)]
pub struct LinterWarning {
    pub message: &'static str,
    pub start: usize,
    pub end: usize
}

impl fmt::Display for LinterWarning {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            LinterWarning { ref message, ref start, ref end } => {
                f.write_str(
                    format!("Warning: {} ({}, {})", message, start, end).as_str()
                )
            },
        }
    }
}

impl Error for LinterWarning {
    #[allow(unused_variables)]
    fn description(&self) -> &str {
        match *self {
            LinterWarning { ref message, ref start, ref end } => { message },
        }
    }
}
<commit_msg>Do not create a reference to a reference<commit_after>use std::error::Error;
use std::fmt;

#[derive(Debug, RustcEncodable)]
pub struct LinterWarning {
    pub message: &'static str,
    pub start: usize,
    pub end: usize
}

impl fmt::Display for LinterWarning {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            LinterWarning { message, ref start, ref end } => {
                f.write_str(
                    format!("Warning: {} ({}, {})", message, start, end).as_str()
                )
            },
        }
    }
}

impl Error for LinterWarning {
    #[allow(unused_variables)]
    fn description(&self) -> &str {
        match *self {
            LinterWarning { message, ref start, ref end } => { message },
        }
    }
}
<|endoftext|>


<commit_before>extern crate r2d2;

use std::io;

quick_error! {
    /// The most generic error message in ReQL
    #[derive(Debug)]
    pub enum Error {
        Compile(descr: String) { }
        Runtime(err: RuntimeError) { from() }
        Driver(err: DriverError) { from() }
    }
}

quick_error! {
    /// The parent class of all runtime errors
    ///
    /// All errors on the server unrelated to compilation. Programs may use this to catch any runtime
    /// error, but the server will always return a more specific error class.
    #[derive(Debug)]
    pub enum RuntimeError {
        QueryLogic(err: QueryLogicError) { from() }
        ResourceLimit(descr: String) {}
        User(descr: String) {}
        Internal(descr: String) {}
        Timeout(descr: String) {}
        Availability(err: AvailabilityError) { from() }
        Permissions(descr: String) {}
    }
}

quick_error! {
    /// The query contains a logical impossibility, such as adding a number to a string.
    #[derive(Debug)]
    pub enum QueryLogicError {
        NonExistence(descr: String) {}
    }
}

quick_error! {
    /// A server in the cluster is unavailable
    ///
    /// The parent class of `OpFailedError` and `OpIndeterminateError`. Programs may use this
    /// to catch any availability error, but the server will always return one of this class’s
    /// children.
    #[derive(Debug)]
    pub enum AvailabilityError {
        OpFailed(descr: String) {}
        OpIndeterminate(descr: String) {}
    }
}

quick_error! {
    /// An error has occurred within the driver
    ///
    /// This may be a driver bug, or it may be an unfulfillable command, such as an unserializable
    /// query.
    #[derive(Debug)]
    pub enum DriverError {
        Auth(descr: String) {}
        Initialization(err: r2d2::InitializationError) { from() }
        Connection(err: ConnectionError) { from() }
    }
}

quick_error! {
    /// Connection related errors
    #[derive(Debug)]
    pub enum ConnectionError {
        PoolWrite(descr: String) {}
        PoolRead(descr: String) {}
        Io(err: io::Error) { from() }
        Other(descr: String) {}
    }
}

/// Converts from r2d2 error
impl From<r2d2::InitializationError> for Error {
    fn from(err: r2d2::InitializationError) -> Error {
        From::from(DriverError::Initialization(err))
    }
}

impl From<ConnectionError> for Error {
    fn from(err: ConnectionError) -> Error {
        From::from(DriverError::Connection(err))
    }
}

/// Converts from IO error
impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        From::from(ConnectionError::Io(err))
    }
}
<commit_msg>Add ParseResponse error<commit_after>extern crate r2d2;

use std::io;
use std::str;

quick_error! {
    /// The most generic error message in ReQL
    #[derive(Debug)]
    pub enum Error {
        Compile(descr: String) { }
        Runtime(err: RuntimeError) { from() }
        Driver(err: DriverError) { from() }
    }
}

quick_error! {
    /// The parent class of all runtime errors
    ///
    /// All errors on the server unrelated to compilation. Programs may use this to catch any runtime
    /// error, but the server will always return a more specific error class.
    #[derive(Debug)]
    pub enum RuntimeError {
        QueryLogic(err: QueryLogicError) { from() }
        ResourceLimit(descr: String) {}
        User(descr: String) {}
        Internal(descr: String) {}
        Timeout(descr: String) {}
        Availability(err: AvailabilityError) { from() }
        Permissions(descr: String) {}
    }
}

quick_error! {
    /// The query contains a logical impossibility, such as adding a number to a string.
    #[derive(Debug)]
    pub enum QueryLogicError {
        NonExistence(descr: String) {}
    }
}

quick_error! {
    /// A server in the cluster is unavailable
    ///
    /// The parent class of `OpFailedError` and `OpIndeterminateError`. Programs may use this
    /// to catch any availability error, but the server will always return one of this class’s
    /// children.
    #[derive(Debug)]
    pub enum AvailabilityError {
        OpFailed(descr: String) {}
        OpIndeterminate(descr: String) {}
    }
}

quick_error! {
    /// An error has occurred within the driver
    ///
    /// This may be a driver bug, or it may be an unfulfillable command, such as an unserializable
    /// query.
    #[derive(Debug)]
    pub enum DriverError {
        Auth(descr: String) {}
        Initialization(err: r2d2::InitializationError) { from() }
        Connection(err: ConnectionError) { from() }
        ParseResponse(err: str::Utf8Error) { from() }
    }
}

quick_error! {
    /// Connection related errors
    #[derive(Debug)]
    pub enum ConnectionError {
        PoolWrite(descr: String) {}
        PoolRead(descr: String) {}
        Io(err: io::Error) { from() }
        Other(descr: String) {}
    }
}

/// Converts from r2d2 error
impl From<r2d2::InitializationError> for Error {
    fn from(err: r2d2::InitializationError) -> Error {
        From::from(DriverError::Initialization(err))
    }
}

impl From<ConnectionError> for Error {
    fn from(err: ConnectionError) -> Error {
        From::from(DriverError::Connection(err))
    }
}

/// Converts from IO error
impl From<io::Error> for Error {
    fn from(err: io::Error) -> Error {
        From::from(ConnectionError::Io(err))
    }
}

/// Converts from Utf8Error error
impl From<str::Utf8Error> for Error {
    fn from(err: str::Utf8Error) -> Error {
        From::from(DriverError::ParseResponse(err))
    }
}
<|endoftext|>


<commit_before>//! Manage xml character escapes

use memchr;
use std::borrow::Cow;

#[cfg_attr(feature = "failure", derive(Fail))]
#[derive(Debug)]
pub enum EscapeError {
    /// Entity with Null character
    EntityWithNull(::std::ops::Range<usize>),
    /// Unrecognized escape symbol
    UnrecognizedSymbol(
        ::std::ops::Range<usize>,
        ::std::result::Result<String, ::std::string::FromUtf8Error>,
    ),
    /// Cannot find `;` after `&`
    UnterminatedEntity(::std::ops::Range<usize>),
    /// Cannot convert Hexa to utf8
    TooLongHexadecimal,
    /// Character is not a valid hexadecimal value
    InvalidHexadecimal(char),
    /// Cannot convert decimal to hexa
    TooLongDecimal,
    /// Character is not a valid decimal value
    InvalidDecimal(char),
}

impl std::fmt::Display for EscapeError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            EscapeError::EntityWithNull(e) => write!(
                f,
                "Error while escaping character at range {:?}: Null character entity not allowed",
                e
            ),
            EscapeError::UnrecognizedSymbol(rge, res) => write!(
                f,
                "Error while escaping character at range {:?}: Unrecognized escape symbol: {:?}",
                rge, res
            ),
            EscapeError::UnterminatedEntity(e) => write!(
                f,
                "Error while escaping character at range {:?}: Cannot find ';' after '&'",
                e
            ),
            EscapeError::TooLongHexadecimal => write!(f, "Cannot convert hexadecimal to utf8"),
            EscapeError::InvalidHexadecimal(e) => {
                write!(f, "'{}' is not a valid hexadecimal character", e)
            }
            EscapeError::TooLongDecimal => write!(f, "Cannot convert decimal to utf8"),
            EscapeError::InvalidDecimal(e) => write!(f, "'{}' is not a valid decimal character", e),
        }
    }
}

// UTF-8 ranges and tags for encoding characters
const TAG_CONT: u8 = 0b1000_0000;
const TAG_TWO_B: u8 = 0b1100_0000;
const TAG_THREE_B: u8 = 0b1110_0000;
const TAG_FOUR_B: u8 = 0b1111_0000;
const MAX_ONE_B: u32 = 0x80;
const MAX_TWO_B: u32 = 0x800;
const MAX_THREE_B: u32 = 0x10000;

enum ByteOrChar {
    Byte(u8),
    Char(u32),
}

/// helper function to escape a `&[u8]` and replace all
/// xml special characters (<, >, &, ', ") with their corresponding
/// xml escaped value.
pub fn escape(raw: &[u8]) -> Cow<[u8]> {
    let mut escapes: Vec<(usize, &'static [u8])> = Vec::new();
    let mut bytes = raw.iter();
    fn to_escape(b: u8) -> bool {
        match b {
            b'<' | b'>' | b'\'' | b'&' | b'"' => true,
            _ => false,
        }
    }

    let mut loc = 0;
    while let Some(i) = bytes.position(|&b| to_escape(b)) {
        loc += i;
        match raw[loc] {
            b'<' => escapes.push((loc, b"<")),
            b'>' => escapes.push((loc, b">")),
            b'\'' => escapes.push((loc, b"'")),
            b'&' => escapes.push((loc, b"&")),
            b'"' => escapes.push((loc, b""")),
            _ => unreachable!("Only '<', '>','\', '&' and '\"' are escaped"),
        }
        loc += 1;
    }

    if escapes.is_empty() {
        Cow::Borrowed(raw)
    } else {
        let len = raw.len();
        let mut v = Vec::with_capacity(len);
        let mut start = 0;
        for (i, r) in escapes {
            v.extend_from_slice(&raw[start..i]);
            v.extend_from_slice(r);
            start = i + 1;
        }

        if start < len {
            v.extend_from_slice(&raw[start..]);
        }
        Cow::Owned(v)
    }
}

/// helper function to unescape a `&[u8]` and replace all
/// xml escaped characters ('&...;') into their corresponding value
pub fn unescape(raw: &[u8]) -> Result<Cow<[u8]>, EscapeError> {
    let mut escapes = Vec::new();

    let mut start = 0;
    while let Some(i) = memchr::memchr(b'&', &raw[start..]) {
        start += i + 1;
        if let Some(j) = memchr::memchr(b';', &raw[start..]) {
            let end = start + j;
            // search for character correctness
            let b_o_c = match &raw[start..end] {
                b"lt" => ByteOrChar::Byte(b'<'),
                b"gt" => ByteOrChar::Byte(b'>'),
                b"amp" => ByteOrChar::Byte(b'&'),
                b"apos" => ByteOrChar::Byte(b'\''),
                b"quot" => ByteOrChar::Byte(b'\"'),
                b"#x0" | b"#0" => return Err(EscapeError::EntityWithNull(start..end)),
                bytes if bytes.starts_with(b"#x") => {
                    ByteOrChar::Char(parse_hexadecimal(&bytes[2..])?)
                }
                bytes if bytes.starts_with(b"#") => ByteOrChar::Char(parse_decimal(&bytes[1..])?),
                bytes => {
                    return Err(EscapeError::UnrecognizedSymbol(
                        start..end,
                        String::from_utf8(bytes.to_vec()),
                    ));
                }
            };
            escapes.push((start - 1..end, b_o_c));
            start = end + 1;
        } else {
            return Err(EscapeError::UnterminatedEntity(start..raw.len()));
        }
    }
    if escapes.is_empty() {
        Ok(Cow::Borrowed(raw))
    } else {
        let len = raw.len();
        let mut v = Vec::with_capacity(len);
        let mut start = 0;
        for (r, b) in escapes {
            v.extend_from_slice(&raw[start..r.start]);
            match b {
                ByteOrChar::Byte(b) => v.push(b),
                ByteOrChar::Char(c) => push_utf8(&mut v, c),
            }
            start = r.end + 1;
        }
        if start < raw.len() {
            v.extend_from_slice(&raw[start..]);
        }
        Ok(Cow::Owned(v))
    }
}

fn push_utf8(buf: &mut Vec<u8>, code: u32) {
    if code < MAX_ONE_B {
        buf.push(code as u8);
    } else if code < MAX_TWO_B {
        buf.push((code >> 6 & 0x1F) as u8 | TAG_TWO_B);
        buf.push((code & 0x3F) as u8 | TAG_CONT);
    } else if code < MAX_THREE_B {
        buf.push((code >> 12 & 0x0F) as u8 | TAG_THREE_B);
        buf.push((code >> 6 & 0x3F) as u8 | TAG_CONT);
        buf.push((code & 0x3F) as u8 | TAG_CONT);
    } else {
        buf.push((code >> 18 & 0x07) as u8 | TAG_FOUR_B);
        buf.push((code >> 12 & 0x3F) as u8 | TAG_CONT);
        buf.push((code >> 6 & 0x3F) as u8 | TAG_CONT);
        buf.push((code & 0x3F) as u8 | TAG_CONT);
    }
}

fn parse_hexadecimal(bytes: &[u8]) -> Result<u32, EscapeError> {
    // maximum code is 0x10FFFF => 6 characters
    if bytes.len() > 6 {
        return Err(EscapeError::TooLongHexadecimal);
    }
    let mut code = 0;
    for &b in bytes {
        code <<= 4;
        code += match b {
            b'0'..=b'9' => b - b'0',
            b'a'..=b'f' => b - b'a' + 10,
            b'A'..=b'F' => b - b'A' + 10,
            b => return Err(EscapeError::InvalidHexadecimal(b as char)),
        } as u32;
    }
    Ok(code)
}

fn parse_decimal(bytes: &[u8]) -> Result<u32, EscapeError> {
    // maximum code is 0x10FFFF = 1114111 => 7 characters
    if bytes.len() > 7 {
        return Err(EscapeError::TooLongDecimal);
    }
    let mut code = 0;
    for &b in bytes {
        code *= 10;
        code += match b {
            b'0'..=b'9' => b - b'0',
            b => return Err(EscapeError::InvalidDecimal(b as char)),
        } as u32;
    }
    Ok(code)
}

#[test]
fn test_unescape() {
    assert_eq!(&*unescape(b"test").unwrap(), b"test");
    assert_eq!(&*unescape(b"<test>").unwrap(), b"<test>");
    assert_eq!(&*unescape(b"0").unwrap(), b"0");
    assert_eq!(&*unescape(b"0").unwrap(), b"0");
}

#[test]
fn test_escape() {
    assert_eq!(&*escape(b"test"), b"test");
    assert_eq!(&*escape(b"<test>"), b"<test>");
    assert_eq!(&*escape(b"\"a\"bc"), b""a"bc");
    assert_eq!(&*escape(b"\"a\"b&c"), b""a"b&c");
    assert_eq!(
        &*escape(b"prefix_\"a\"b&<>c"),
        "prefix_"a"b&<>c".as_bytes()
    );
}
<commit_msg>use memchr for unescape<commit_after>//! Manage xml character escapes

use memchr;
use std::borrow::Cow;

#[cfg_attr(feature = "failure", derive(Fail))]
#[derive(Debug)]
pub enum EscapeError {
    /// Entity with Null character
    EntityWithNull(::std::ops::Range<usize>),
    /// Unrecognized escape symbol
    UnrecognizedSymbol(
        ::std::ops::Range<usize>,
        ::std::result::Result<String, ::std::string::FromUtf8Error>,
    ),
    /// Cannot find `;` after `&`
    UnterminatedEntity(::std::ops::Range<usize>),
    /// Cannot convert Hexa to utf8
    TooLongHexadecimal,
    /// Character is not a valid hexadecimal value
    InvalidHexadecimal(char),
    /// Cannot convert decimal to hexa
    TooLongDecimal,
    /// Character is not a valid decimal value
    InvalidDecimal(char),
}

impl std::fmt::Display for EscapeError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            EscapeError::EntityWithNull(e) => write!(
                f,
                "Error while escaping character at range {:?}: Null character entity not allowed",
                e
            ),
            EscapeError::UnrecognizedSymbol(rge, res) => write!(
                f,
                "Error while escaping character at range {:?}: Unrecognized escape symbol: {:?}",
                rge, res
            ),
            EscapeError::UnterminatedEntity(e) => write!(
                f,
                "Error while escaping character at range {:?}: Cannot find ';' after '&'",
                e
            ),
            EscapeError::TooLongHexadecimal => write!(f, "Cannot convert hexadecimal to utf8"),
            EscapeError::InvalidHexadecimal(e) => {
                write!(f, "'{}' is not a valid hexadecimal character", e)
            }
            EscapeError::TooLongDecimal => write!(f, "Cannot convert decimal to utf8"),
            EscapeError::InvalidDecimal(e) => write!(f, "'{}' is not a valid decimal character", e),
        }
    }
}

// UTF-8 ranges and tags for encoding characters
const TAG_CONT: u8 = 0b1000_0000;
const TAG_TWO_B: u8 = 0b1100_0000;
const TAG_THREE_B: u8 = 0b1110_0000;
const TAG_FOUR_B: u8 = 0b1111_0000;
const MAX_ONE_B: u32 = 0x80;
const MAX_TWO_B: u32 = 0x800;
const MAX_THREE_B: u32 = 0x10000;

enum ByteOrChar {
    Byte(u8),
    Char(u32),
}

/// helper function to escape a `&[u8]` and replace all
/// xml special characters (<, >, &, ', ") with their corresponding
/// xml escaped value.
pub fn escape(raw: &[u8]) -> Cow<[u8]> {
    let mut escapes: Vec<(usize, &'static [u8])> = Vec::new();
    let mut bytes = raw.iter();
    fn to_escape(b: u8) -> bool {
        match b {
            b'<' | b'>' | b'\'' | b'&' | b'"' => true,
            _ => false,
        }
    }

    let mut loc = 0;
    while let Some(i) = bytes.position(|&b| to_escape(b)) {
        loc += i;
        match raw[loc] {
            b'<' => escapes.push((loc, b"<")),
            b'>' => escapes.push((loc, b">")),
            b'\'' => escapes.push((loc, b"'")),
            b'&' => escapes.push((loc, b"&")),
            b'"' => escapes.push((loc, b""")),
            _ => unreachable!("Only '<', '>','\', '&' and '\"' are escaped"),
        }
        loc += 1;
    }

    if escapes.is_empty() {
        Cow::Borrowed(raw)
    } else {
        let len = raw.len();
        let mut v = Vec::with_capacity(len);
        let mut start = 0;
        for (i, r) in escapes {
            v.extend_from_slice(&raw[start..i]);
            v.extend_from_slice(r);
            start = i + 1;
        }

        if start < len {
            v.extend_from_slice(&raw[start..]);
        }
        Cow::Owned(v)
    }
}

/// helper function to unescape a `&[u8]` and replace all
/// xml escaped characters ('&...;') into their corresponding value
pub fn unescape(raw: &[u8]) -> Result<Cow<[u8]>, EscapeError> {
    let mut escapes = Vec::new();

    let mut iter = memchr::memchr2_iter(b'&', b';', raw);
    while let Some(start) = iter.by_ref().find(|p| raw[*p] == b'&') {
        match iter.next() {
            Some(end) if raw[end] == b';' => {
                // search for character correctness
                let b_o_c = match &raw[start + 1..end] {
                    b"lt" => ByteOrChar::Byte(b'<'),
                    b"gt" => ByteOrChar::Byte(b'>'),
                    b"amp" => ByteOrChar::Byte(b'&'),
                    b"apos" => ByteOrChar::Byte(b'\''),
                    b"quot" => ByteOrChar::Byte(b'\"'),
                    b"#x0" | b"#0" => return Err(EscapeError::EntityWithNull(start..end)),
                    bytes if bytes.starts_with(b"#x") => {
                        ByteOrChar::Char(parse_hexadecimal(&bytes[2..])?)
                    }
                    bytes if bytes.starts_with(b"#") => {
                        ByteOrChar::Char(parse_decimal(&bytes[1..])?)
                    }
                    bytes => {
                        return Err(EscapeError::UnrecognizedSymbol(
                            start + 1..end,
                            String::from_utf8(bytes.to_vec()),
                        ));
                    }
                };
                escapes.push((start..end, b_o_c));
            }
            _ => return Err(EscapeError::UnterminatedEntity(start..raw.len())),
        }
    }

    if escapes.is_empty() {
        Ok(Cow::Borrowed(raw))
    } else {
        let len = raw.len();
        let mut v = Vec::with_capacity(len);
        let mut start = 0;
        for (r, b) in escapes {
            v.extend_from_slice(&raw[start..r.start]);
            match b {
                ByteOrChar::Byte(b) => v.push(b),
                ByteOrChar::Char(c) => push_utf8(&mut v, c),
            }
            start = r.end + 1;
        }
        if start < raw.len() {
            v.extend_from_slice(&raw[start..]);
        }
        Ok(Cow::Owned(v))
    }
}

fn push_utf8(buf: &mut Vec<u8>, code: u32) {
    if code < MAX_ONE_B {
        buf.push(code as u8);
    } else if code < MAX_TWO_B {
        buf.push((code >> 6 & 0x1F) as u8 | TAG_TWO_B);
        buf.push((code & 0x3F) as u8 | TAG_CONT);
    } else if code < MAX_THREE_B {
        buf.push((code >> 12 & 0x0F) as u8 | TAG_THREE_B);
        buf.push((code >> 6 & 0x3F) as u8 | TAG_CONT);
        buf.push((code & 0x3F) as u8 | TAG_CONT);
    } else {
        buf.push((code >> 18 & 0x07) as u8 | TAG_FOUR_B);
        buf.push((code >> 12 & 0x3F) as u8 | TAG_CONT);
        buf.push((code >> 6 & 0x3F) as u8 | TAG_CONT);
        buf.push((code & 0x3F) as u8 | TAG_CONT);
    }
}

fn parse_hexadecimal(bytes: &[u8]) -> Result<u32, EscapeError> {
    // maximum code is 0x10FFFF => 6 characters
    if bytes.len() > 6 {
        return Err(EscapeError::TooLongHexadecimal);
    }
    let mut code = 0;
    for &b in bytes {
        code <<= 4;
        code += match b {
            b'0'..=b'9' => b - b'0',
            b'a'..=b'f' => b - b'a' + 10,
            b'A'..=b'F' => b - b'A' + 10,
            b => return Err(EscapeError::InvalidHexadecimal(b as char)),
        } as u32;
    }
    Ok(code)
}

fn parse_decimal(bytes: &[u8]) -> Result<u32, EscapeError> {
    // maximum code is 0x10FFFF = 1114111 => 7 characters
    if bytes.len() > 7 {
        return Err(EscapeError::TooLongDecimal);
    }
    let mut code = 0;
    for &b in bytes {
        code *= 10;
        code += match b {
            b'0'..=b'9' => b - b'0',
            b => return Err(EscapeError::InvalidDecimal(b as char)),
        } as u32;
    }
    Ok(code)
}

#[test]
fn test_unescape() {
    assert_eq!(&*unescape(b"test").unwrap(), b"test");
    assert_eq!(&*unescape(b"<test>").unwrap(), b"<test>");
    assert_eq!(&*unescape(b"0").unwrap(), b"0");
    assert_eq!(&*unescape(b"0").unwrap(), b"0");
}

#[test]
fn test_escape() {
    assert_eq!(&*escape(b"test"), b"test");
    assert_eq!(&*escape(b"<test>"), b"<test>");
    assert_eq!(&*escape(b"\"a\"bc"), b""a"bc");
    assert_eq!(&*escape(b"\"a\"b&c"), b""a"b&c");
    assert_eq!(
        &*escape(b"prefix_\"a\"b&<>c"),
        "prefix_"a"b&<>c".as_bytes()
    );
}
<|endoftext|>


<commit_before>
use std::io::Read;  // Trait providing read_to_string()

use std::env;

use hyper;
use rustc_serialize;

use projects;
use groups;

pub const API_VERSION: u16 = 3;

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Version {
    pub version: String,
    pub revision: String,
}

#[derive(Debug)]
pub struct Pagination {
    pub page: u16,
    pub per_page: u16,
}

#[derive(Debug)]
pub struct GitLab {
    scheme: String,
    domain: String,
    port: u16,
    private_token: String,
    client: hyper::Client,
    pagination: Pagination,
}

impl GitLab {

    pub fn new(scheme: &str, domain: &str, port: u16, private_token: &str) -> GitLab {
        GitLab {
            scheme: scheme.to_string(),
            domain: domain.to_string(),
            port:   port,
            private_token: private_token.to_string(),
            client: match env::var("HTTP_PROXY") {
                Ok(proxy) => {
                    println!("proxy: {:?}", proxy);
                    let proxy: Vec<&str> = proxy.trim_left_matches("http://").split(':').collect();
                    let hostname = proxy[0].to_string();
                    let port = proxy[1];

                    hyper::Client::with_http_proxy(hostname, port.parse().unwrap())
                },
                Err(_) => hyper::Client::new(),
            },
            pagination: Pagination {page: 1, per_page: 20},
        }
    }

    pub fn new_http(domain: &str, private_token: &str) -> GitLab {
        GitLab::new("http", &domain, 80, &private_token)
    }

    pub fn new_https(domain: &str, private_token: &str) -> GitLab {
        GitLab::new("https", &domain, 443, &private_token)
    }

    pub fn build_url(&self, command: &str) -> String {
        format!("{}://{}:{}/api/v{}/{}?private_token={}&page={}&per_page={}",
                                self.scheme,
                                self.domain,
                                self.port,
                                API_VERSION,
                                command,
                                self.private_token,
                                self.pagination.page,
                                self.pagination.per_page)
    }

    pub fn attempt_connection(&self) -> Result<hyper::client::Response, hyper::Error> {
        // self.get("version")
        let url = self.build_url("version");
        // Close connections after each GET.
        let res = self.client.get(&url).header(hyper::header::Connection::close()).send();

        res
    }

    pub fn set_pagination(&mut self, pagination: Pagination) {
        self.pagination = pagination;
    }

    pub fn get<T>(&self, command: &str) -> Result<T, rustc_serialize::json::DecoderError>
            where T: rustc_serialize::Decodable {

        let url = self.build_url(command);
        let mut res: hyper::client::Response =
                        self.client
                        .get(&url)
                        .header(hyper::header::Connection::close())
                        .send()
                        .unwrap();

        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();

        rustc_serialize::json::decode(body.as_str())
    }

    pub fn version(&self) -> Result<Version, rustc_serialize::json::DecoderError> {
        self.get("version")
    }

    pub fn groups(&self) -> Result<groups::Groups, rustc_serialize::json::DecoderError> {
        self.get("groups")
    }

    pub fn projects(&self) -> Result<projects::Projects, rustc_serialize::json::DecoderError> {
        self.get("groups")
    }
}

#[cfg(test)]
mod tests {
    use std::env;
    use gitlab::GitLab;
    use gitlab::Pagination;

    #[test]
    fn list_groups() {
        let token = match env::var("GITLAB_TOKEN") {
            Ok(val) => val,
            Err(_)  => panic!("Please set environment variable 'GITLAB_TOKEN'"),
        };

        let mut gl = GitLab::new_https("gitlab.com", &token);
        gl.set_pagination(Pagination{page: 1, per_page: 100});
        println!("gl: {:?}", gl);

        let groups = gl.groups().unwrap();
        println!("groups: {:?}", groups);
    }

    #[test]
    fn list_projects() {
        let token = match env::var("GITLAB_TOKEN") {
            Ok(val) => val,
            Err(_)  => panic!("Please set environment variable 'GITLAB_TOKEN'"),
        };

        let mut gl = GitLab::new_https("gitlab.com", &token);
        gl.set_pagination(Pagination{page: 1, per_page: 100});
        println!("gl: {:?}", gl);

        println!("projects: {:?}", gl.projects().unwrap());
        // assert_eq!(gl.attempt_connection().unwrap().status, hyper::Ok);

    }
}
<commit_msg>Bug fix: Get the projects (not groups) in projects()<commit_after>
use std::io::Read;  // Trait providing read_to_string()

use std::env;

use hyper;
use rustc_serialize;

use projects;
use groups;

pub const API_VERSION: u16 = 3;

#[derive(Debug, RustcDecodable, RustcEncodable)]
pub struct Version {
    pub version: String,
    pub revision: String,
}

#[derive(Debug)]
pub struct Pagination {
    pub page: u16,
    pub per_page: u16,
}

#[derive(Debug)]
pub struct GitLab {
    scheme: String,
    domain: String,
    port: u16,
    private_token: String,
    client: hyper::Client,
    pagination: Pagination,
}

impl GitLab {

    pub fn new(scheme: &str, domain: &str, port: u16, private_token: &str) -> GitLab {
        GitLab {
            scheme: scheme.to_string(),
            domain: domain.to_string(),
            port:   port,
            private_token: private_token.to_string(),
            client: match env::var("HTTP_PROXY") {
                Ok(proxy) => {
                    println!("proxy: {:?}", proxy);
                    let proxy: Vec<&str> = proxy.trim_left_matches("http://").split(':').collect();
                    let hostname = proxy[0].to_string();
                    let port = proxy[1];

                    hyper::Client::with_http_proxy(hostname, port.parse().unwrap())
                },
                Err(_) => hyper::Client::new(),
            },
            pagination: Pagination {page: 1, per_page: 20},
        }
    }

    pub fn new_http(domain: &str, private_token: &str) -> GitLab {
        GitLab::new("http", &domain, 80, &private_token)
    }

    pub fn new_https(domain: &str, private_token: &str) -> GitLab {
        GitLab::new("https", &domain, 443, &private_token)
    }

    pub fn build_url(&self, command: &str) -> String {
        format!("{}://{}:{}/api/v{}/{}?private_token={}&page={}&per_page={}",
                                self.scheme,
                                self.domain,
                                self.port,
                                API_VERSION,
                                command,
                                self.private_token,
                                self.pagination.page,
                                self.pagination.per_page)
    }

    pub fn attempt_connection(&self) -> Result<hyper::client::Response, hyper::Error> {
        // self.get("version")
        let url = self.build_url("version");
        // Close connections after each GET.
        let res = self.client.get(&url).header(hyper::header::Connection::close()).send();

        res
    }

    pub fn set_pagination(&mut self, pagination: Pagination) {
        self.pagination = pagination;
    }

    pub fn get<T>(&self, command: &str) -> Result<T, rustc_serialize::json::DecoderError>
            where T: rustc_serialize::Decodable {

        let url = self.build_url(command);
        let mut res: hyper::client::Response =
                        self.client
                        .get(&url)
                        .header(hyper::header::Connection::close())
                        .send()
                        .unwrap();

        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();

        rustc_serialize::json::decode(body.as_str())
    }

    pub fn version(&self) -> Result<Version, rustc_serialize::json::DecoderError> {
        self.get("version")
    }

    pub fn groups(&self) -> Result<groups::Groups, rustc_serialize::json::DecoderError> {
        self.get("groups")
    }

    pub fn projects(&self) -> Result<projects::Projects, rustc_serialize::json::DecoderError> {
        self.get("projects")
    }
}

#[cfg(test)]
mod tests {
    use std::env;
    use gitlab::GitLab;
    use gitlab::Pagination;

    #[test]
    fn list_groups() {
        let token = match env::var("GITLAB_TOKEN") {
            Ok(val) => val,
            Err(_)  => panic!("Please set environment variable 'GITLAB_TOKEN'"),
        };

        let mut gl = GitLab::new_https("gitlab.com", &token);
        gl.set_pagination(Pagination{page: 1, per_page: 100});
        println!("gl: {:?}", gl);

        let groups = gl.groups().unwrap();
        println!("groups: {:?}", groups);
    }

    #[test]
    fn list_projects() {
        let token = match env::var("GITLAB_TOKEN") {
            Ok(val) => val,
            Err(_)  => panic!("Please set environment variable 'GITLAB_TOKEN'"),
        };

        let mut gl = GitLab::new_https("gitlab.com", &token);
        gl.set_pagination(Pagination{page: 1, per_page: 100});
        println!("gl: {:?}", gl);

        println!("projects: {:?}", gl.projects().unwrap());
        // assert_eq!(gl.attempt_connection().unwrap().status, hyper::Ok);

    }
}
<|endoftext|>


<commit_before>//! GPT-header object and helper functions.

use std::fmt;
use std::fs::{File, OpenOptions};
use std::io::{Cursor, Error, ErrorKind, Read, Result, Seek, SeekFrom, Write};
use std::path::Path;

extern crate crc;
extern crate itertools;

use self::itertools::Itertools;

use self::crc::{crc32, Hasher32};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use disk;
use partition;
use uuid;

#[derive(Debug, Eq, PartialEq)]
pub struct Header {
    /// EFI PART
    pub signature: String, // Offset  0. "EFI PART", 45h 46h 49h 20h 50h 41h 52h 54h
    /// 00 00 01 00
    pub revision: u32, // Offset  8
    /// little endian
    pub header_size_le: u32, // Offset 12
    /// CRC32 of the header with crc32 section zeroed
    pub crc32: u32, // Offset 16
    /// must be 0
    pub reserved: u32, // Offset 20
    /// For main header, 1
    pub current_lba: u64, // Offset 24
    /// LBA for backup header
    pub backup_lba: u64, // Offset 32
    /// First usable LBA for partitions (primary table last LBA + 1)
    pub first_usable: u64, // Offset 40
    /// Last usable LBA (secondary partition table first LBA - 1)
    pub last_usable: u64, // Offset 48
    /// UUID of the disk
    pub disk_guid: uuid::Uuid, // Offset 56
    /// Starting LBA of partition entries
    pub part_start: u64, // Offset 72
    /// Number of partition entries
    pub num_parts: u32, // Offset 80
    /// Size of a partition entry, usually 128
    pub part_size: u32, // Offset 84
    /// CRC32 of the partition table
    pub crc32_parts: u32, // Offset 88
}

impl Header {
    pub(crate) fn compute_new(
        primary: bool,
        pp: &[partition::Partition],
        guid: uuid::Uuid,
        backup_offset: u64,
    ) -> Result<Self> {
        let (cur, bak) = if primary {
            (1, backup_offset)
        } else {
            (backup_offset, 1)
        };
        let first = 34u64;
        let last = backup_offset
            .checked_sub(first)
            .ok_or_else(|| Error::new(ErrorKind::Other, "header underflow - last usable"))?;

        let hdr = Header {
            signature: "EFI PART".to_string(),
            revision: 65536,
            header_size_le: 92,
            crc32: 0,
            reserved: 0,
            current_lba: cur,
            backup_lba: bak,
            first_usable: first,
            last_usable: last,
            disk_guid: guid,
            part_start: 2,
            num_parts: pp.len() as u32,
            part_size: 128,
            crc32_parts: 0,
        };

        Ok(hdr)
    }

    /// Write the header to a location.  With a crc32 set to zero
    /// this will set the crc32 after writing the Header out
    pub fn write_primary(&self, file: &mut File, lb_size: disk::LogicalBlockSize) -> Result<usize> {
        let mut bytes_written: usize = 0;

        // Build up byte array in memory
        let parts_checksum = partentry_checksum(file, self, lb_size)?;
        let bytes = self.as_bytes(None, Some(parts_checksum))?;

        // Calculate the crc32 from the byte array
        let checksum = calculate_crc32(&bytes)?;

        // Write it to disk in 1 shot
        let start = self.current_lba
            .checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "primary header overflow - offset"))?;
        let _ = file.seek(SeekFrom::Start(start))?;
        let csum_len = file.write(&self.as_bytes(Some(checksum), Some(parts_checksum))?)?;
        bytes_written = bytes_written
            .checked_add(csum_len)
            .ok_or_else(|| Error::new(ErrorKind::Other, "primary header overflow - checksum"))?;

        Ok(bytes_written)
    }

    // TODO: implement writing backup header too.
    pub fn write_backup(&self, file: &mut File, lb_size: disk::LogicalBlockSize) -> Result<usize> {
        let start = self.backup_lba
            .checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "backup header overflow - offset"))?;
        let _ = file.seek(SeekFrom::Start(start))?;
        Ok(0)
    }

    fn as_bytes(&self, checksum: Option<u32>, parts_checksum: Option<u32>) -> Result<Vec<u8>> {
        let mut buff: Vec<u8> = Vec::new();

        buff.write_all(self.signature.as_bytes())?;
        buff.write_u32::<LittleEndian>(self.revision)?;
        buff.write_u32::<LittleEndian>(self.header_size_le)?;
        match checksum {
            Some(c) => buff.write_u32::<LittleEndian>(c)?,
            None => buff.write_u32::<LittleEndian>(0)?,
        };
        buff.write_u32::<LittleEndian>(0)?;
        buff.write_u64::<LittleEndian>(self.current_lba)?;
        buff.write_u64::<LittleEndian>(self.backup_lba)?;
        buff.write_u64::<LittleEndian>(self.first_usable)?;
        buff.write_u64::<LittleEndian>(self.last_usable)?;
        buff.write_all(self.disk_guid.as_bytes())?;
        buff.write_u64::<LittleEndian>(self.part_start)?;
        buff.write_u32::<LittleEndian>(self.num_parts)?;
        buff.write_u32::<LittleEndian>(self.part_size)?;
        match parts_checksum {
            Some(c) => buff.write_u32::<LittleEndian>(c)?,
            None => buff.write_u32::<LittleEndian>(0)?,
        };
        trace!("Buffer: {:02x}", buff.iter().format(","));
        Ok(buff)
    }
}

/// Parses a uuid with first 3 portions in little endian.
pub fn parse_uuid(rdr: &mut Cursor<&[u8]>) -> Result<uuid::Uuid> {
    let d1: u32 = rdr.read_u32::<LittleEndian>()?;
    let d2: u16 = rdr.read_u16::<LittleEndian>()?;
    let d3: u16 = rdr.read_u16::<LittleEndian>()?;
    let uuid = uuid::Uuid::from_fields(
        d1,
        d2,
        d3,
        &rdr.get_ref()[rdr.position() as usize..rdr.position() as usize + 8],
    );
    rdr.seek(SeekFrom::Current(8))?;

    match uuid {
        Ok(uuid) => Ok(uuid),
        Err(_) => Err(Error::new(ErrorKind::Other, "Invalid Disk UUID?")),
    }
}

impl fmt::Display for Header {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Disk:\t\t{}\nCRC32:\t\t{}\nTable CRC:\t{}",
            self.disk_guid, self.crc32, self.crc32_parts
        )
    }
}

/// Read a GPT header from a given path.
///
/// ## Example
///
/// ```rust,no_run
/// use gpt::header::read_header;
///
/// let lb_size = gpt::disk::DEFAULT_SECTOR_SIZE;
/// let diskpath = std::path::Path::new("/dev/sdz");
///
/// let h = read_header(diskpath, lb_size).unwrap();
/// ```
pub fn read_header(path: &Path, sector_size: disk::LogicalBlockSize) -> Result<Header> {
    let mut file = File::open(path)?;
    read_primary_header(&mut file, sector_size)
}

pub(crate) fn read_primary_header(
    file: &mut File,
    sector_size: disk::LogicalBlockSize,
) -> Result<Header> {
    let cur = file.seek(SeekFrom::Current(0)).unwrap_or(0);
    let offset: u64 = sector_size.into();
    let res = file_read_header(file, offset);
    let _ = file.seek(SeekFrom::Start(cur));
    res
}

pub(crate) fn read_backup_header(
    file: &mut File,
    sector_size: disk::LogicalBlockSize,
) -> Result<Header> {
    let cur = file.seek(SeekFrom::Current(0)).unwrap_or(0);
    let h2sect = find_backup_lba(file, sector_size)?;
    let offset = h2sect
        .checked_mul(sector_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "backup header overflow - offset"))?;
    let res = file_read_header(file, offset);
    let _ = file.seek(SeekFrom::Start(cur));
    res
}

pub(crate) fn file_read_header(file: &mut File, offset: u64) -> Result<Header> {
    let _ = file.seek(SeekFrom::Start(offset));
    let mut hdr: [u8; 92] = [0; 92];

    let _ = file.read_exact(&mut hdr);
    let mut reader = Cursor::new(&hdr[..]);

    let sigstr = String::from_utf8_lossy(
        &reader.get_ref()[reader.position() as usize..reader.position() as usize + 8],
    );
    reader.seek(SeekFrom::Current(8))?;

    if sigstr != "EFI PART" {
        return Err(Error::new(ErrorKind::Other, "invalid GPT signature"));
    };

    let h = Header {
        signature: sigstr.to_string(),
        revision: reader.read_u32::<LittleEndian>()?,
        header_size_le: reader.read_u32::<LittleEndian>()?,
        crc32: reader.read_u32::<LittleEndian>()?,
        reserved: reader.read_u32::<LittleEndian>()?,
        current_lba: reader.read_u64::<LittleEndian>()?,
        backup_lba: reader.read_u64::<LittleEndian>()?,
        first_usable: reader.read_u64::<LittleEndian>()?,
        last_usable: reader.read_u64::<LittleEndian>()?,
        disk_guid: parse_uuid(&mut reader)?,
        part_start: reader.read_u64::<LittleEndian>()?,
        num_parts: reader.read_u32::<LittleEndian>()?,
        part_size: reader.read_u32::<LittleEndian>()?,
        crc32_parts: reader.read_u32::<LittleEndian>()?,
    };

    let mut hdr_crc = hdr;
    for crc_byte in hdr_crc.iter_mut().skip(16).take(4) {
        *crc_byte = 0;
    }
    let c = crc32::checksum_ieee(&hdr_crc);
    trace!("hdr_crc: {:?}, h.crc32: {:?}", c, h.crc32);
    if crc32::checksum_ieee(&hdr_crc) == h.crc32 {
        Ok(h)
    } else {
        Err(Error::new(ErrorKind::Other, "invalid CRC32 checksum"))
    }
}

pub(crate) fn find_backup_lba(f: &mut File, sector_size: disk::LogicalBlockSize) -> Result<u64> {
    trace!("Querying file size to find backup header location");
    let lb_size: u64 = sector_size.into();
    let m = f.metadata()?;
    if m.len() <= lb_size {
        return Err(Error::new(
            ErrorKind::Other,
            "disk image too small for backup header",
        ));
    }
    let backup_location = (m.len().saturating_sub(lb_size)) / lb_size;
    trace!("Backup location: {:#x}", backup_location);

    Ok(backup_location)
}

fn calculate_crc32(b: &[u8]) -> Result<u32> {
    let mut digest = crc32::Digest::new(crc32::IEEE);
    trace!("Writing buffer to digest calculator");
    digest.write(b);

    Ok(digest.sum32())
}

pub(crate) fn partentry_checksum(
    file: &mut File,
    hdr: &Header,
    lb_size: disk::LogicalBlockSize,
) -> Result<u32> {
    // Seek to start of partition table.
    let start = hdr.part_start
        .checked_mul(lb_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "header overflow - partition table start"))?;
    let _ = file.seek(SeekFrom::Start(start))?;

    // Read partition table.
    let pt_len = u64::from(hdr.num_parts)
        .checked_mul(hdr.part_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "partition table - size"))?;
    let mut buf = vec![0; pt_len as usize];
    file.read_exact(&mut buf)?;

    // Compute CRC32 over all table bits.
    let mut digest = crc32::Digest::new(crc32::IEEE);
    digest.write(&buf);
    Ok(digest.sum32())
}

/// A helper function to create a new header and write it to disk.
/// If the uuid isn't given a random one will be generated.  Use
/// this in conjunction with Partition::write()
// TODO: Move this to Header::new() and Header::write to write it
// that will match the Partition::write() API
pub fn write_header(
    p: &Path,
    uuid: Option<uuid::Uuid>,
    sector_size: disk::LogicalBlockSize,
) -> Result<uuid::Uuid> {
    debug!("opening {} for writing", p.display());
    let mut file = OpenOptions::new().write(true).read(true).open(p)?;
    let bak = find_backup_lba(&mut file, sector_size)?;
    let guid = match uuid {
        Some(u) => u,
        None => {
            let u = uuid::Uuid::new_v4();
            debug!("Generated random uuid: {}", u);
            u
        }
    };

    let hdr = Header::compute_new(true, &[], guid, bak)?;
    debug!("new header: {:#?}", hdr);
    hdr.write_primary(&mut file, sector_size)?;

    Ok(guid)
}
<commit_msg>header: support writing backup header<commit_after>//! GPT-header object and helper functions.

use std::fmt;
use std::fs::{File, OpenOptions};
use std::io::{Cursor, Error, ErrorKind, Read, Result, Seek, SeekFrom, Write};
use std::path::Path;

extern crate crc;
extern crate itertools;

use self::itertools::Itertools;

use self::crc::{crc32, Hasher32};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
use disk;
use partition;
use uuid;

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Header {
    /// EFI PART
    pub signature: String, // Offset  0. "EFI PART", 45h 46h 49h 20h 50h 41h 52h 54h
    /// 00 00 01 00
    pub revision: u32, // Offset  8
    /// little endian
    pub header_size_le: u32, // Offset 12
    /// CRC32 of the header with crc32 section zeroed
    pub crc32: u32, // Offset 16
    /// must be 0
    pub reserved: u32, // Offset 20
    /// For main header, 1
    pub current_lba: u64, // Offset 24
    /// LBA for backup header
    pub backup_lba: u64, // Offset 32
    /// First usable LBA for partitions (primary table last LBA + 1)
    pub first_usable: u64, // Offset 40
    /// Last usable LBA (secondary partition table first LBA - 1)
    pub last_usable: u64, // Offset 48
    /// UUID of the disk
    pub disk_guid: uuid::Uuid, // Offset 56
    /// Starting LBA of partition entries
    pub part_start: u64, // Offset 72
    /// Number of partition entries
    pub num_parts: u32, // Offset 80
    /// Size of a partition entry, usually 128
    pub part_size: u32, // Offset 84
    /// CRC32 of the partition table
    pub crc32_parts: u32, // Offset 88
}

impl Header {
    pub(crate) fn compute_new(
        primary: bool,
        pp: &[partition::Partition],
        guid: uuid::Uuid,
        backup_offset: u64,
    ) -> Result<Self> {
        let (cur, bak) = if primary {
            (1, backup_offset)
        } else {
            (backup_offset, 1)
        };
        let first = 34u64;
        let last = backup_offset
            .checked_sub(first)
            .ok_or_else(|| Error::new(ErrorKind::Other, "header underflow - last usable"))?;

        let hdr = Header {
            signature: "EFI PART".to_string(),
            revision: 65536,
            header_size_le: 92,
            crc32: 0,
            reserved: 0,
            current_lba: cur,
            backup_lba: bak,
            first_usable: first,
            last_usable: last,
            disk_guid: guid,
            part_start: 2,
            num_parts: pp.len() as u32,
            part_size: 128,
            crc32_parts: 0,
        };

        Ok(hdr)
    }

    /// Write the primary header.
    ///
    /// With a CRC32 set to zero this will set the crc32 after
    /// writing the header out.
    pub fn write_primary(&self, file: &mut File, lb_size: disk::LogicalBlockSize) -> Result<usize> {
        // This is the primary header. It must start before the backup one.
        if self.current_lba >= self.backup_lba {
            return Err(Error::new(
                ErrorKind::Other,
                "primary header does not start before backup one",
            ));
        }
        self.file_write_header(file, self.current_lba, lb_size)
    }

    /// Write the backup header.
    ///
    /// With a CRC32 set to zero this will set the crc32 after
    /// writing the header out.
    pub fn write_backup(&self, file: &mut File, lb_size: disk::LogicalBlockSize) -> Result<usize> {
        // This is the backup header. It must start after the primary one.
        if self.current_lba <= self.backup_lba {
            return Err(Error::new(
                ErrorKind::Other,
                "backup header does not start after primary one",
            ));
        }
        self.file_write_header(file, self.current_lba, lb_size)
    }

    /// Write an header to an arbitrary LBA.
    ///
    /// With a CRC32 set to zero this will set the crc32 after
    /// writing the header out.
    fn file_write_header(
        &self,
        file: &mut File,
        lba: u64,
        lb_size: disk::LogicalBlockSize,
    ) -> Result<usize> {
        // Build up byte array in memory
        let parts_checksum = partentry_checksum(file, self, lb_size)?;
        let bytes = self.as_bytes(None, Some(parts_checksum))?;

        // Calculate the CRC32 from the byte array
        let checksum = calculate_crc32(&bytes)?;

        // Write it to disk in 1 shot
        let start = lba.checked_mul(lb_size.into())
            .ok_or_else(|| Error::new(ErrorKind::Other, "header overflow - offset"))?;
        let _ = file.seek(SeekFrom::Start(start))?;
        let len = file.write(&self.as_bytes(Some(checksum), Some(parts_checksum))?)?;

        Ok(len)
    }

    fn as_bytes(&self, checksum: Option<u32>, parts_checksum: Option<u32>) -> Result<Vec<u8>> {
        let mut buff: Vec<u8> = Vec::new();

        buff.write_all(self.signature.as_bytes())?;
        buff.write_u32::<LittleEndian>(self.revision)?;
        buff.write_u32::<LittleEndian>(self.header_size_le)?;
        match checksum {
            Some(c) => buff.write_u32::<LittleEndian>(c)?,
            None => buff.write_u32::<LittleEndian>(0)?,
        };
        buff.write_u32::<LittleEndian>(0)?;
        buff.write_u64::<LittleEndian>(self.current_lba)?;
        buff.write_u64::<LittleEndian>(self.backup_lba)?;
        buff.write_u64::<LittleEndian>(self.first_usable)?;
        buff.write_u64::<LittleEndian>(self.last_usable)?;
        buff.write_all(self.disk_guid.as_bytes())?;
        buff.write_u64::<LittleEndian>(self.part_start)?;
        buff.write_u32::<LittleEndian>(self.num_parts)?;
        buff.write_u32::<LittleEndian>(self.part_size)?;
        match parts_checksum {
            Some(c) => buff.write_u32::<LittleEndian>(c)?,
            None => buff.write_u32::<LittleEndian>(0)?,
        };
        trace!("Buffer: {:02x}", buff.iter().format(","));
        Ok(buff)
    }
}

/// Parses a uuid with first 3 portions in little endian.
pub fn parse_uuid(rdr: &mut Cursor<&[u8]>) -> Result<uuid::Uuid> {
    let d1: u32 = rdr.read_u32::<LittleEndian>()?;
    let d2: u16 = rdr.read_u16::<LittleEndian>()?;
    let d3: u16 = rdr.read_u16::<LittleEndian>()?;
    let uuid = uuid::Uuid::from_fields(
        d1,
        d2,
        d3,
        &rdr.get_ref()[rdr.position() as usize..rdr.position() as usize + 8],
    );
    rdr.seek(SeekFrom::Current(8))?;

    match uuid {
        Ok(uuid) => Ok(uuid),
        Err(_) => Err(Error::new(ErrorKind::Other, "Invalid Disk UUID?")),
    }
}

impl fmt::Display for Header {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Disk:\t\t{}\nCRC32:\t\t{}\nTable CRC:\t{}",
            self.disk_guid, self.crc32, self.crc32_parts
        )
    }
}

/// Read a GPT header from a given path.
///
/// ## Example
///
/// ```rust,no_run
/// use gpt::header::read_header;
///
/// let lb_size = gpt::disk::DEFAULT_SECTOR_SIZE;
/// let diskpath = std::path::Path::new("/dev/sdz");
///
/// let h = read_header(diskpath, lb_size).unwrap();
/// ```
pub fn read_header(path: &Path, sector_size: disk::LogicalBlockSize) -> Result<Header> {
    let mut file = File::open(path)?;
    read_primary_header(&mut file, sector_size)
}

pub(crate) fn read_primary_header(
    file: &mut File,
    sector_size: disk::LogicalBlockSize,
) -> Result<Header> {
    let cur = file.seek(SeekFrom::Current(0)).unwrap_or(0);
    let offset: u64 = sector_size.into();
    let res = file_read_header(file, offset);
    let _ = file.seek(SeekFrom::Start(cur));
    res
}

pub(crate) fn read_backup_header(
    file: &mut File,
    sector_size: disk::LogicalBlockSize,
) -> Result<Header> {
    let cur = file.seek(SeekFrom::Current(0)).unwrap_or(0);
    let h2sect = find_backup_lba(file, sector_size)?;
    let offset = h2sect
        .checked_mul(sector_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "backup header overflow - offset"))?;
    let res = file_read_header(file, offset);
    let _ = file.seek(SeekFrom::Start(cur));
    res
}

pub(crate) fn file_read_header(file: &mut File, offset: u64) -> Result<Header> {
    let _ = file.seek(SeekFrom::Start(offset));
    let mut hdr: [u8; 92] = [0; 92];

    let _ = file.read_exact(&mut hdr);
    let mut reader = Cursor::new(&hdr[..]);

    let sigstr = String::from_utf8_lossy(
        &reader.get_ref()[reader.position() as usize..reader.position() as usize + 8],
    );
    reader.seek(SeekFrom::Current(8))?;

    if sigstr != "EFI PART" {
        return Err(Error::new(ErrorKind::Other, "invalid GPT signature"));
    };

    let h = Header {
        signature: sigstr.to_string(),
        revision: reader.read_u32::<LittleEndian>()?,
        header_size_le: reader.read_u32::<LittleEndian>()?,
        crc32: reader.read_u32::<LittleEndian>()?,
        reserved: reader.read_u32::<LittleEndian>()?,
        current_lba: reader.read_u64::<LittleEndian>()?,
        backup_lba: reader.read_u64::<LittleEndian>()?,
        first_usable: reader.read_u64::<LittleEndian>()?,
        last_usable: reader.read_u64::<LittleEndian>()?,
        disk_guid: parse_uuid(&mut reader)?,
        part_start: reader.read_u64::<LittleEndian>()?,
        num_parts: reader.read_u32::<LittleEndian>()?,
        part_size: reader.read_u32::<LittleEndian>()?,
        crc32_parts: reader.read_u32::<LittleEndian>()?,
    };

    let mut hdr_crc = hdr;
    for crc_byte in hdr_crc.iter_mut().skip(16).take(4) {
        *crc_byte = 0;
    }
    let c = crc32::checksum_ieee(&hdr_crc);
    trace!("header CRC32: {:#x} - computed CRC32: {:#x}", h.crc32, c);
    if crc32::checksum_ieee(&hdr_crc) == h.crc32 {
        Ok(h)
    } else {
        Err(Error::new(ErrorKind::Other, "invalid CRC32 checksum"))
    }
}

pub(crate) fn find_backup_lba(f: &mut File, sector_size: disk::LogicalBlockSize) -> Result<u64> {
    trace!("querying file size to find backup header location");
    let lb_size: u64 = sector_size.into();
    let m = f.metadata()?;
    if m.len() <= lb_size {
        return Err(Error::new(
            ErrorKind::Other,
            "disk image too small for backup header",
        ));
    }
    let bak_offset = m.len().saturating_sub(lb_size);
    let bak_lba = bak_offset / lb_size;
    trace!(
        "backup header: LBA={}, bytes offset={}",
        bak_lba,
        bak_offset
    );

    Ok(bak_lba)
}

fn calculate_crc32(b: &[u8]) -> Result<u32> {
    let mut digest = crc32::Digest::new(crc32::IEEE);
    trace!("Writing buffer to digest calculator");
    digest.write(b);

    Ok(digest.sum32())
}

pub(crate) fn partentry_checksum(
    file: &mut File,
    hdr: &Header,
    lb_size: disk::LogicalBlockSize,
) -> Result<u32> {
    // Seek to start of partition table.
    let start = hdr.part_start
        .checked_mul(lb_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "header overflow - partition table start"))?;
    let _ = file.seek(SeekFrom::Start(start))?;

    // Read partition table.
    let pt_len = u64::from(hdr.num_parts)
        .checked_mul(hdr.part_size.into())
        .ok_or_else(|| Error::new(ErrorKind::Other, "partition table - size"))?;
    let mut buf = vec![0; pt_len as usize];
    file.read_exact(&mut buf)?;

    // Compute CRC32 over all table bits.
    let mut digest = crc32::Digest::new(crc32::IEEE);
    digest.write(&buf);
    Ok(digest.sum32())
}

/// A helper function to create a new header and write it to disk.
/// If the uuid isn't given a random one will be generated.  Use
/// this in conjunction with Partition::write()
// TODO: Move this to Header::new() and Header::write to write it
// that will match the Partition::write() API
pub fn write_header(
    p: &Path,
    uuid: Option<uuid::Uuid>,
    sector_size: disk::LogicalBlockSize,
) -> Result<uuid::Uuid> {
    debug!("opening {} for writing", p.display());
    let mut file = OpenOptions::new().write(true).read(true).open(p)?;
    let bak = find_backup_lba(&mut file, sector_size)?;
    let guid = match uuid {
        Some(u) => u,
        None => {
            let u = uuid::Uuid::new_v4();
            debug!("Generated random uuid: {}", u);
            u
        }
    };

    let hdr = Header::compute_new(true, &[], guid, bak)?;
    debug!("new header: {:#?}", hdr);
    hdr.write_primary(&mut file, sector_size)?;

    Ok(guid)
}
<|endoftext|>


<commit_before>
//! Various methods for computing with vectors.

use {Matrix2d, Ray, Vec2d, Rectangle, RoundRectangle};

/// Multiplies two matrices.
#[inline(always)]
pub fn multiply(m: &[f64, ..6], b: &[f64, ..6]) -> Matrix2d {
    [m[0]*b[0]+m[1]*b[3]+m[2]*0.0,  m[0]*b[1]+m[1]*b[4]+m[2]*0.0,  m[0]*b[2]+m[1]*b[5]+m[2]*1.0,
     m[3]*b[0]+m[4]*b[3]+m[5]*0.0,  m[3]*b[1]+m[4]*b[4]+m[5]*0.0,  m[3]*b[2]+m[4]*b[5]+m[5]*1.0]
}

/// Creates a translation matrix.
#[inline(always)]
pub fn translate(x: f64, y: f64) -> Matrix2d {
    [1.0, 0.0, x,
     0.0, 1.0, y]
}

/// Creates a rotation matrix.
#[inline(always)]
pub fn rotate_radians(angle: f64) -> Matrix2d {
    let c = angle.cos();
    let s = angle.sin();
    [c, s, 0.0,
    -s, c, 0.0]
}

/// Create a scale matrix.
#[inline(always)]
pub fn scale(sx: f64, sy: f64) -> Matrix2d {
    [sx, 0.0, 0.0,
     0.0, sy, 0.0]
}

/// Create a shear matrix.
#[inline(always)]
pub fn shear(sx: f64, sy: f64) -> Matrix2d {
    [1.0, sx, 0.0,
     sy, 1.0, 0.0]
}

/// Compute the shortest vector from point to ray.
/// A ray stores starting point and directional vector.
#[inline(always)]
pub fn separation(ray: Ray, x: f64, y: f64) -> Vec2d {
    // Get the directional vector.
    let (dir_x, dir_y) = (ray[2], ray[3]);
    // Get displacement vector from point.
    let (dx, dy) = (ray[0] - x, ray[1] - y);
    // Compute the component of position in ray direction.
    let dot = dir_x * x + dir_y * y;
    // The directional vector multiplied with the dot gives us a parallel vector.
    // When we subtract this from the displacement we get a vector normal to the ray.
    // This is the shortest vector from the point to the ray.
    [dx - dot * dir_x, dy - dot * dir_y]
}

/// Returns the least separation out of four.
/// Each seperation can be computed using `separation` function.
/// The separation returned can be used to solve collision of rectangles.
#[inline(always)]
pub fn least_separation_4(sep1: Vec2d, sep2: Vec2d, sep3: Vec2d, sep4: Vec2d) -> Vec2d {
    let dot1 = sep1[0] * sep1[0] + sep1[1] * sep1[1];
    let dot2 = sep2[0] * sep2[0] + sep2[1] * sep2[1];
    let dot3 = sep3[0] * sep3[0] + sep3[1] * sep3[1];
    let dot4 = sep4[0] * sep4[0] + sep4[1] * sep4[1];
    // Search for the smallest dot product.
    if dot1 < dot2 {
        if dot3 < dot4 {
            if dot1 < dot3 { sep1 } else { sep3 }
        } else {
            if dot1 < dot4 { sep1 } else { sep4 }
        }
    } else {
        if dot3 < dot4 {
            if dot2 < dot3 { sep2 } else { sep3 }
        } else {
            if dot2 < dot4 { sep2 } else { sep4 }
        }
    }
}

/// Shrinks a rectangle by a factor on all sides.
#[inline(always)]
pub fn margin_rectangle(rect: &Rectangle, m: f64) -> Rectangle {
    let w = rect[2] - 2.0 * m;
    let h = rect[3] - 2.0 * m;
    let (x, w) = if w < 0.0 { (rect[0] + 0.5 * rect[2], 0.0) } else { (rect[0] + m, w) };
    let (y, h) = if h < 0.0 { (rect[1] + 0.5 * rect[3], 0.0) } else { (rect[1] + m, h) };
    [x, y, w, h]
}

/// Shrinks a rectangle by a factor on all sides.
#[inline(always)]
pub fn margin_round_rectangle(rect: &RoundRectangle, m: f64) -> RoundRectangle {
    let w = rect[2] - 2.0 * m;
    let h = rect[3] - 2.0 * m;
    let (x, w) = if w < 0.0 { (rect[0] + 0.5 * rect[2], 0.0) } else { (rect[0] + m, w) };
    let (y, h) = if h < 0.0 { (rect[1] + 0.5 * rect[3], 0.0) } else { (rect[1] + m, h) };
    [x, y, w, h, rect[4]]
}

/// Computes a relative rectangle using the rectangle as a tile.
#[inline(always)]
pub fn relative_rectangle(rect: &Rectangle, x: f64, y: f64) -> Rectangle {
    [rect[0] + x * rect[2], rect[1] + y * rect[3], rect[2], rect[3]]
}

/// Computes a relative round rectangle using the round rectangle as a tile.
#[inline(always)]
pub fn relative_round_rectangle(rect: &RoundRectangle, x: f64, y: f64) -> RoundRectangle {
    [rect[0] + x * rect[2], rect[1] + y * rect[3], rect[2], rect[3], rect[4]]
}

/// Computes modular offset safely for numbers.
#[inline(always)]
pub fn modular_offset<T: Add<T, T> + Rem<T, T>>(n: &T, i: &T, off: &T) -> T {
    (*i + (*off % *n + *n)) % *n
}

#[test]
fn test_modular_offset() {
    assert_eq!(modular_offset(&3.0_f64, &0.0_f64, &-1.0_f64), 2.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &1.0_f64, &-1.0_f64), 0.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &2.0_f64, &-1.0_f64), 1.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &3.0_f64, &-1.0_f64), 2.0_f64);

    assert_eq!(modular_offset(&3.0_f64, &0.0_f64, &1.0_f64), 1.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &1.0_f64, &1.0_f64), 2.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &2.0_f64, &1.0_f64), 0.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &3.0_f64, &1.0_f64), 1.0_f64);
}
<commit_msg>Added computing area of simple polygons<commit_after>
//! Various methods for computing with vectors.

use {Matrix2d, Ray, Vec2d, Rectangle, RoundRectangle};
use modular_index::{previous};

/// Multiplies two matrices.
#[inline(always)]
pub fn multiply(m: &[f64, ..6], b: &[f64, ..6]) -> Matrix2d {
    [m[0]*b[0]+m[1]*b[3]+m[2]*0.0,  m[0]*b[1]+m[1]*b[4]+m[2]*0.0,  m[0]*b[2]+m[1]*b[5]+m[2]*1.0,
     m[3]*b[0]+m[4]*b[3]+m[5]*0.0,  m[3]*b[1]+m[4]*b[4]+m[5]*0.0,  m[3]*b[2]+m[4]*b[5]+m[5]*1.0]
}

/// Creates a translation matrix.
#[inline(always)]
pub fn translate(x: f64, y: f64) -> Matrix2d {
    [1.0, 0.0, x,
     0.0, 1.0, y]
}

/// Creates a rotation matrix.
#[inline(always)]
pub fn rotate_radians(angle: f64) -> Matrix2d {
    let c = angle.cos();
    let s = angle.sin();
    [c, s, 0.0,
    -s, c, 0.0]
}

/// Create a scale matrix.
#[inline(always)]
pub fn scale(sx: f64, sy: f64) -> Matrix2d {
    [sx, 0.0, 0.0,
     0.0, sy, 0.0]
}

/// Create a shear matrix.
#[inline(always)]
pub fn shear(sx: f64, sy: f64) -> Matrix2d {
    [1.0, sx, 0.0,
     sy, 1.0, 0.0]
}

/// Compute the shortest vector from point to ray.
/// A ray stores starting point and directional vector.
#[inline(always)]
pub fn separation(ray: Ray, x: f64, y: f64) -> Vec2d {
    // Get the directional vector.
    let (dir_x, dir_y) = (ray[2], ray[3]);
    // Get displacement vector from point.
    let (dx, dy) = (ray[0] - x, ray[1] - y);
    // Compute the component of position in ray direction.
    let dot = dir_x * x + dir_y * y;
    // The directional vector multiplied with the dot gives us a parallel vector.
    // When we subtract this from the displacement we get a vector normal to the ray.
    // This is the shortest vector from the point to the ray.
    [dx - dot * dir_x, dy - dot * dir_y]
}

/// Returns the least separation out of four.
/// Each seperation can be computed using `separation` function.
/// The separation returned can be used to solve collision of rectangles.
#[inline(always)]
pub fn least_separation_4(sep1: Vec2d, sep2: Vec2d, sep3: Vec2d, sep4: Vec2d) -> Vec2d {
    let dot1 = sep1[0] * sep1[0] + sep1[1] * sep1[1];
    let dot2 = sep2[0] * sep2[0] + sep2[1] * sep2[1];
    let dot3 = sep3[0] * sep3[0] + sep3[1] * sep3[1];
    let dot4 = sep4[0] * sep4[0] + sep4[1] * sep4[1];
    // Search for the smallest dot product.
    if dot1 < dot2 {
        if dot3 < dot4 {
            if dot1 < dot3 { sep1 } else { sep3 }
        } else {
            if dot1 < dot4 { sep1 } else { sep4 }
        }
    } else {
        if dot3 < dot4 {
            if dot2 < dot3 { sep2 } else { sep3 }
        } else {
            if dot2 < dot4 { sep2 } else { sep4 }
        }
    }
}

/// Shrinks a rectangle by a factor on all sides.
#[inline(always)]
pub fn margin_rectangle(rect: &Rectangle, m: f64) -> Rectangle {
    let w = rect[2] - 2.0 * m;
    let h = rect[3] - 2.0 * m;
    let (x, w) = if w < 0.0 { (rect[0] + 0.5 * rect[2], 0.0) } else { (rect[0] + m, w) };
    let (y, h) = if h < 0.0 { (rect[1] + 0.5 * rect[3], 0.0) } else { (rect[1] + m, h) };
    [x, y, w, h]
}

/// Shrinks a rectangle by a factor on all sides.
#[inline(always)]
pub fn margin_round_rectangle(rect: &RoundRectangle, m: f64) -> RoundRectangle {
    let w = rect[2] - 2.0 * m;
    let h = rect[3] - 2.0 * m;
    let (x, w) = if w < 0.0 { (rect[0] + 0.5 * rect[2], 0.0) } else { (rect[0] + m, w) };
    let (y, h) = if h < 0.0 { (rect[1] + 0.5 * rect[3], 0.0) } else { (rect[1] + m, h) };
    [x, y, w, h, rect[4]]
}

/// Computes a relative rectangle using the rectangle as a tile.
#[inline(always)]
pub fn relative_rectangle(rect: &Rectangle, x: f64, y: f64) -> Rectangle {
    [rect[0] + x * rect[2], rect[1] + y * rect[3], rect[2], rect[3]]
}

/// Computes a relative round rectangle using the round rectangle as a tile.
#[inline(always)]
pub fn relative_round_rectangle(rect: &RoundRectangle, x: f64, y: f64) -> RoundRectangle {
    [rect[0] + x * rect[2], rect[1] + y * rect[3], rect[2], rect[3], rect[4]]
}

/// Computes modular offset safely for numbers.
#[inline(always)]
pub fn modular_offset<T: Add<T, T> + Rem<T, T>>(n: &T, i: &T, off: &T) -> T {
    (*i + (*off % *n + *n)) % *n
}

#[test]
fn test_modular_offset() {
    assert_eq!(modular_offset(&3.0_f64, &0.0_f64, &-1.0_f64), 2.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &1.0_f64, &-1.0_f64), 0.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &2.0_f64, &-1.0_f64), 1.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &3.0_f64, &-1.0_f64), 2.0_f64);

    assert_eq!(modular_offset(&3.0_f64, &0.0_f64, &1.0_f64), 1.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &1.0_f64, &1.0_f64), 2.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &2.0_f64, &1.0_f64), 0.0_f64);
    assert_eq!(modular_offset(&3.0_f64, &3.0_f64, &1.0_f64), 1.0_f64);
}

/// Computes the area of a simple polygon.
/// A simple polygon is one that does not intersect itself.
pub fn area(polygon: &[f64]) -> f64 {
    let n = polygon.len() / 2;
    let mut sum = 0.0_f64;
    for i in range(0, n) {
        let (qx, qy) = (polygon[i * 2], polygon[i * 2 + 1]);
        let p_i = previous(n, i);
        let (px, py) = (polygon[p_i * 2], polygon[p_i * 2 + 1]);
        sum += px * qy - qx * py;
    }

    0.5 * sum
}

<|endoftext|>


<commit_before>//!
//! The **nodrop** crate has the following cargo feature flags:
//!
//! - `std`
//!   - Optional, enabled by default
//!   - Requires Rust 1.6 *to disable*
//!   - Use libstd
//! - `use_needs_drop`
//!   - Optional
//!   - Requires nightly channel.
//!   - Use `needs_drop` to skip overwriting if not necessary
//! - `use_union`
//!   - Optional
//!   - Requires nightly channel
//!   - Using untagged union, finally we have an implementation of `NoDrop` without hacks,
//!     for example the fact that `NoDrop<T>` never has a destructor anymore.
//!

#![feature(untagged_unions)]

#![cfg_attr(not(test), no_std)]
#[cfg(not(test))]
extern crate core as std;

use std::ops::{Deref, DerefMut};

#[allow(unions_with_drop_fields)]
union UnionFlag<T> {
    value: T,
}

pub struct NoDrop<T>(UnionFlag<T>);

impl<T> NoDrop<T> {
    /// Create a new **NoDrop**.
    #[inline]
    pub fn new(value: T) -> Self {
        NoDrop(UnionFlag { value: value })
    }

    /// Extract the inner value.
    ///
    /// Once extracted, the value can of course drop again.
    #[inline]
    pub fn into_inner(self) -> T {
        unsafe {
            self.0.value
        }
    }
}

impl<T> Deref for NoDrop<T> {
    type Target = T;

    // Use type invariant, always initialized
    #[inline]
    fn deref(&self) -> &T {
        unsafe {
            &self.0.value
        }
    }
}

impl<T> DerefMut for NoDrop<T> {
    // Use type invariant, always initialized
    #[inline]
    fn deref_mut(&mut self) -> &mut T {
        unsafe {
            &mut self.0.value
        }
    }
}

#[cfg(test)]
mod tests {
    use super::NoDrop;
    use std::mem;

    #[test]
    fn test_drop() {
        use std::cell::Cell;

        let flag = &Cell::new(0);

        struct Bump<'a>(&'a Cell<i32>);

        impl<'a> Drop for Bump<'a> {
            fn drop(&mut self) {
                let n = self.0.get();
                self.0.set(n + 1);
            }
        }

        {
            let _ = NoDrop::new([Bump(flag), Bump(flag)]);
        }
        assert_eq!(flag.get(), 0);

        // test something with the nullable pointer optimization
        flag.set(0);

        {
            let mut array = NoDrop::new(Vec::new());
            array.push(vec![Bump(flag)]);
            array.push(vec![Bump(flag), Bump(flag)]);
            array.push(vec![]);
            array.push(vec![Bump(flag)]);
            drop(array.pop());
            assert_eq!(flag.get(), 1);
            drop(array.pop());
            assert_eq!(flag.get(), 1);
            drop(array.pop());
            assert_eq!(flag.get(), 3);
        }

        // last one didn't drop.
        assert_eq!(flag.get(), 3);

        flag.set(0);
        {
            let array = NoDrop::new(Bump(flag));
            array.into_inner();
            assert_eq!(flag.get(), 1);
        }
        assert_eq!(flag.get(), 1);
    }

    #[test]
    fn test_size_of() {
        assert!(mem::size_of::<NoDrop<&i32>>() == mem::size_of::<&i32>());
        assert!(mem::size_of::<NoDrop<Vec<i32>>>() == mem::size_of::<Vec<i32>>());
        // No non-nullable pointer optimization!
        assert!(mem::size_of::<Option<NoDrop<&i32>>>() > mem::size_of::<NoDrop<&i32>>());
    }
}
<commit_msg>nodrop-union: Implement Copy, Clone on NoDrop<T><commit_after>//!
//! The **nodrop** crate has the following cargo feature flags:
//!
//! - `std`
//!   - Optional, enabled by default
//!   - Requires Rust 1.6 *to disable*
//!   - Use libstd
//! - `use_needs_drop`
//!   - Optional
//!   - Requires nightly channel.
//!   - Use `needs_drop` to skip overwriting if not necessary
//! - `use_union`
//!   - Optional
//!   - Requires nightly channel
//!   - Using untagged union, finally we have an implementation of `NoDrop` without hacks,
//!     for example the fact that `NoDrop<T>` never has a destructor anymore.
//!

#![feature(untagged_unions)]

#![cfg_attr(not(test), no_std)]
#[cfg(not(test))]
extern crate core as std;

use std::ops::{Deref, DerefMut};

#[allow(unions_with_drop_fields)]
#[derive(Copy)]
union UnionFlag<T> {
    value: T,
}

impl<T: Clone> Clone for UnionFlag<T> {
    fn clone(&self) -> Self {
        unsafe {
            UnionFlag { value: self.value.clone() }
        }
    }
}

#[derive(Copy, Clone)]
pub struct NoDrop<T>(UnionFlag<T>);

impl<T> NoDrop<T> {
    /// Create a new **NoDrop**.
    #[inline]
    pub fn new(value: T) -> Self {
        NoDrop(UnionFlag { value: value })
    }

    /// Extract the inner value.
    ///
    /// Once extracted, the value can of course drop again.
    #[inline]
    pub fn into_inner(self) -> T {
        unsafe {
            self.0.value
        }
    }
}

impl<T> Deref for NoDrop<T> {
    type Target = T;

    // Use type invariant, always initialized
    #[inline]
    fn deref(&self) -> &T {
        unsafe {
            &self.0.value
        }
    }
}

impl<T> DerefMut for NoDrop<T> {
    // Use type invariant, always initialized
    #[inline]
    fn deref_mut(&mut self) -> &mut T {
        unsafe {
            &mut self.0.value
        }
    }
}

#[cfg(test)]
mod tests {
    use super::NoDrop;
    use std::mem;

    #[test]
    fn test_drop() {
        use std::cell::Cell;

        let flag = &Cell::new(0);

        struct Bump<'a>(&'a Cell<i32>);

        impl<'a> Drop for Bump<'a> {
            fn drop(&mut self) {
                let n = self.0.get();
                self.0.set(n + 1);
            }
        }

        {
            let _ = NoDrop::new([Bump(flag), Bump(flag)]);
        }
        assert_eq!(flag.get(), 0);

        // test something with the nullable pointer optimization
        flag.set(0);

        {
            let mut array = NoDrop::new(Vec::new());
            array.push(vec![Bump(flag)]);
            array.push(vec![Bump(flag), Bump(flag)]);
            array.push(vec![]);
            array.push(vec![Bump(flag)]);
            drop(array.pop());
            assert_eq!(flag.get(), 1);
            drop(array.pop());
            assert_eq!(flag.get(), 1);
            drop(array.pop());
            assert_eq!(flag.get(), 3);
        }

        // last one didn't drop.
        assert_eq!(flag.get(), 3);

        flag.set(0);
        {
            let array = NoDrop::new(Bump(flag));
            array.into_inner();
            assert_eq!(flag.get(), 1);
        }
        assert_eq!(flag.get(), 1);
    }

    #[test]
    fn test_size_of() {
        assert!(mem::size_of::<NoDrop<&i32>>() == mem::size_of::<&i32>());
        assert!(mem::size_of::<NoDrop<Vec<i32>>>() == mem::size_of::<Vec<i32>>());
        // No non-nullable pointer optimization!
        assert!(mem::size_of::<Option<NoDrop<&i32>>>() > mem::size_of::<NoDrop<&i32>>());
    }
}
<|endoftext|>


<commit_before>//
// zhttpto.rs
//
// Starting code for PS1
// Running on Rust 0.9
//
// Note that this code has serious security risks!  You should not run it
// on any system with access to sensitive files.
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu and David Evans
// Version 0.3

use std::io::fs::PathExtensions;
use std::io::{Acceptor, Listener};
use std::io::{File, TcpListener, TcpStream};
use std::str;

fn extract_uri(request: &str) -> &str {
    let mut spaces: Vec<uint> = Vec::with_capacity(2);

    for (i, c) in request.chars().enumerate() {
        if c == ' ' {
            spaces.push(i)
        }

        if spaces.len() == 2 {
            return request.slice(spaces[0] + 2, spaces[1]);
        }
    }

    fail!("Bad request!");
}

fn respond(mut stream: TcpStream, code: int, body: &str) {
    let status = match code {
        200 => "200 OK",
        403 => "403 Forbidden",
        404 => "404 Not Found",
        500 => "500 Internal Server Error",
        _ => fail!(format!("Unknown status code {}", code)),
    };

    let response = format!("HTTP/1.1 {}\r\n\
                            Content-Type: text/html; charset=UTF-8\r\n\r\n\
                            {}\r\n",
                            status, body);
    stream.write(response.as_bytes());
    println!("Connection terminated.");
}

fn handle(mut stream: TcpStream) {

    match stream.peer_name() {
        Ok(peer_name) => println!("Received connection from: {}", peer_name),
        Err(e) => fail!(e)
    }

    let mut buffer = [0, ..500];
    stream.read(buffer);

    let request = str::from_utf8(buffer);
    println!("Received request:\n{}", request);

    let mut uri = match request {
        None => fail!("Bad request!"),
        Some(x) => extract_uri(x),
    };

    if uri == "" {
        uri = "index.html"
    }

    if !uri.ends_with(".html") {
        respond(stream, 403, "");
        return
    }

    let path = Path::new(uri);

    if !path.exists() || !path.is_file() {
        respond(stream, 404, "");
        return
    }

    let mut file = File::open(&path);

    let body = match file.read_to_string() {
        Ok(x) => x,
        Err(e) => fail!(e)
    };

    respond(stream, 200, body.as_slice());
}

fn main() {
    let mut counter = 0i;

    let listener = TcpListener::bind("127.0.0.1", 4414);
    let mut acceptor = listener.listen();
    println!("Listening on 127.0.0.1:4414...");

    for stream in acceptor.incoming() {
        counter += 1;
        println!("So far, the server has handled {} requests.", counter)

        // Spawn a task to handle the connection
        match stream {
            Ok(stream) => spawn(proc() { handle(stream) }),
            Err(e) => { fail!(e) }
        }
    }
}
<commit_msg>Extracts the uri in a more elegant way<commit_after>//
// zhttpto.rs
//
// Starting code for PS1
// Running on Rust 0.9
//
// Note that this code has serious security risks!  You should not run it
// on any system with access to sensitive files.
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu and David Evans
// Version 0.3

use std::io::fs::PathExtensions;
use std::io::{Acceptor, Listener};
use std::io::{File, TcpListener, TcpStream};
use std::str;

fn extract_uri(request: &str) -> &str {
    match request.splitn(2, ' ').nth(1) {
        Some(r) => r.slice_from(1),  // Strips the trailing '/'
        None => fail!("Bad request!"),
    }
}

fn respond(mut stream: TcpStream, code: int, body: &str) {
    let status = match code {
        200 => "200 OK",
        403 => "403 Forbidden",
        404 => "404 Not Found",
        500 => "500 Internal Server Error",
        _ => fail!(format!("Unknown status code {}", code)),
    };

    let response = format!("HTTP/1.1 {}\r\n\
                            Content-Type: text/html; charset=UTF-8\r\n\r\n\
                            {}\r\n",
                            status, body);
    stream.write(response.as_bytes());
    println!("Connection terminated.");
}

fn handle(mut stream: TcpStream) {

    match stream.peer_name() {
        Ok(peer_name) => println!("Received connection from: {}", peer_name),
        Err(e) => fail!(e)
    }

    let mut buffer = [0, ..500];
    stream.read(buffer);

    let request = str::from_utf8(buffer);
    println!("Received request:\n{}", request);

    let mut uri = match request {
        None => fail!("Bad request!"),
        Some(x) => extract_uri(x),
    };

    if uri == "" {
        uri = "index.html"
    }

    if !uri.ends_with(".html") {
        respond(stream, 403, "");
        return
    }

    let path = Path::new(uri);

    if !path.exists() || !path.is_file() {
        respond(stream, 404, "");
        return
    }

    let mut file = File::open(&path);

    let body = match file.read_to_string() {
        Ok(x) => x,
        Err(e) => fail!(e)
    };

    respond(stream, 200, body.as_slice());
}

fn main() {
    let mut counter = 0i;

    let listener = TcpListener::bind("127.0.0.1", 4414);
    let mut acceptor = listener.listen();
    println!("Listening on 127.0.0.1:4414...");

    for stream in acceptor.incoming() {
        counter += 1;
        println!("So far, the server has handled {} requests.", counter)

        // Spawn a task to handle the connection
        match stream {
            Ok(stream) => spawn(proc() { handle(stream) }),
            Err(e) => { fail!(e) }
        }
    }
}
<|endoftext|>


<commit_before>static DEFAULT_SOURCE_NAME: &'static str = "GraphQL";

#[derive(Debug)]
pub struct Source<'a> {
  pub name: &'a str,
  pub body: &'a str,
}

impl<'a> Source<'a> {
  pub fn new(body: &'a str) -> Source<'a> {
    Source {
      name: DEFAULT_SOURCE_NAME,
      body: body,
    }
  }

  pub fn with_name(name: &'a str, body: &'a str) -> Source<'a> {
    Source {
      name: name,
      body: body,
    }
  }
}<commit_msg>feat: add from_file<commit_after>use std::io;
use std::io::prelude::*;
use std::fs::File;
use std::path::{Path, PathBuf};

static DEFAULT_SOURCE_NAME: &'static str = "GraphQL";

#[derive(Debug)]
pub struct Source {
  pub name: String,
  pub body: String,
}

impl Source {
  pub fn new(body: String) -> Self {
    Source {
      name: DEFAULT_SOURCE_NAME.to_owned(),
      body: body,
    }
  }

  pub fn with_name(name: String, body: String) -> Self {
    Source {
      name: name,
      body: body,
    }
  }

  pub fn from_file<P: AsRef<Path>>(path: P) -> Result<Self, io::Error> {
    let mut p = PathBuf::new();
    p.push(path);
    let mut f = try!(File::open(p.clone()));
    let mut s = String::new();
    try!(f.read_to_string(&mut s));
    Ok(Source::with_name(Self::get_file_name(p), s))
  }

  fn get_file_name(path: PathBuf) -> String {
    if let Some(n) = path.as_path().file_name() {
      if let Ok(name) = n.to_os_string().into_string() {
        return name;
      }
    }
    DEFAULT_SOURCE_NAME.to_owned()
  }
}
<|endoftext|>