<commit_before>macro_rules! hty {
    ($ty:ty) => {
        <$ty as LargeInt>::HighHalf
    }
}

macro_rules! os_ty {
    ($ty:ty) => {
        <$ty as Int>::OtherSign
    }
}

pub mod mul;
pub mod sdiv;
pub mod shift;
pub mod udiv;

/// Trait for some basic operations on integers
pub trait Int {
    /// Type with the same width but other signedness
    type OtherSign;
    /// Unsigned version of Self
    type UnsignedInt;

    /// Returns the bitwidth of the int type
    fn bits() -> u32;

    /// Extracts the sign from self and returns a tuple.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let i = -25_i32;
    /// let (sign, u) = i.extract_sign();
    /// assert_eq!(sign, true);
    /// assert_eq!(u, 25_u32);
    /// ```
    fn extract_sign(self) -> (bool, Self::UnsignedInt);
}

macro_rules! int_impl {
    ($ity:ty, $uty:ty, $bits:expr) => {
        impl Int for $uty {
            type OtherSign = $ity;
            type UnsignedInt = $uty;

            fn bits() -> u32 {
                $bits
            }

            fn extract_sign(self) -> (bool, $uty) {
                (false, self)
            }
        }

        impl Int for $ity {
            type OtherSign = $uty;
            type UnsignedInt = $uty;

            fn bits() -> u32 {
                $bits
            }

            fn extract_sign(self) -> (bool, $uty) {
                if self < 0 {
                    (true, !(self as $uty) + 1)
                } else {
                    (false, self as $uty)
                }
            }
        }
    }
}

int_impl!(i32, u32, 32);
int_impl!(i64, u64, 64);
int_impl!(i128, u128, 128);

/// Trait to convert an integer to/from smaller parts
pub trait LargeInt {
    type LowHalf;
    type HighHalf;

    fn low(self) -> Self::LowHalf;
    fn high(self) -> Self::HighHalf;
    fn from_parts(low: Self::LowHalf, high: Self::HighHalf) -> Self;
}

macro_rules! large_int {
    ($ty:ty, $tylow:ty, $tyhigh:ty, $halfbits:expr) => {
        impl LargeInt for $ty {
            type LowHalf = $tylow;
            type HighHalf = $tyhigh;

            fn low(self) -> $tylow {
                self as $tylow
            }
            fn high(self) -> $tyhigh {
                (self >> $halfbits) as $tyhigh
            }
            fn from_parts(low: $tylow, high: $tyhigh) -> $ty {
                low as $ty | ((high as $ty) << $halfbits)
            }
        }
    }
}

large_int!(u64, u32, u32, 32);
large_int!(i64, u32, i32, 32);
large_int!(u128, u64, u64, 64);
large_int!(i128, u64, i64, 64);
<commit_msg>Fix overflow bug when creating the absolute value<commit_after>macro_rules! hty {
    ($ty:ty) => {
        <$ty as LargeInt>::HighHalf
    }
}

macro_rules! os_ty {
    ($ty:ty) => {
        <$ty as Int>::OtherSign
    }
}

pub mod mul;
pub mod sdiv;
pub mod shift;
pub mod udiv;

/// Trait for some basic operations on integers
pub trait Int {
    /// Type with the same width but other signedness
    type OtherSign;
    /// Unsigned version of Self
    type UnsignedInt;

    /// Returns the bitwidth of the int type
    fn bits() -> u32;

    /// Extracts the sign from self and returns a tuple.
    ///
    /// # Examples
    ///
    /// ```rust,ignore
    /// let i = -25_i32;
    /// let (sign, u) = i.extract_sign();
    /// assert_eq!(sign, true);
    /// assert_eq!(u, 25_u32);
    /// ```
    fn extract_sign(self) -> (bool, Self::UnsignedInt);
}

macro_rules! int_impl {
    ($ity:ty, $uty:ty, $bits:expr) => {
        impl Int for $uty {
            type OtherSign = $ity;
            type UnsignedInt = $uty;

            fn bits() -> u32 {
                $bits
            }

            fn extract_sign(self) -> (bool, $uty) {
                (false, self)
            }
        }

        impl Int for $ity {
            type OtherSign = $uty;
            type UnsignedInt = $uty;

            fn bits() -> u32 {
                $bits
            }

            fn extract_sign(self) -> (bool, $uty) {
                if self < 0 {
                    (true, (!(self as $uty)).wrapping_add(1))
                } else {
                    (false, self as $uty)
                }
            }
        }
    }
}

int_impl!(i32, u32, 32);
int_impl!(i64, u64, 64);
int_impl!(i128, u128, 128);

/// Trait to convert an integer to/from smaller parts
pub trait LargeInt {
    type LowHalf;
    type HighHalf;

    fn low(self) -> Self::LowHalf;
    fn high(self) -> Self::HighHalf;
    fn from_parts(low: Self::LowHalf, high: Self::HighHalf) -> Self;
}

macro_rules! large_int {
    ($ty:ty, $tylow:ty, $tyhigh:ty, $halfbits:expr) => {
        impl LargeInt for $ty {
            type LowHalf = $tylow;
            type HighHalf = $tyhigh;

            fn low(self) -> $tylow {
                self as $tylow
            }
            fn high(self) -> $tyhigh {
                (self >> $halfbits) as $tyhigh
            }
            fn from_parts(low: $tylow, high: $tyhigh) -> $ty {
                low as $ty | ((high as $ty) << $halfbits)
            }
        }
    }
}

large_int!(u64, u32, u32, 32);
large_int!(i64, u32, i32, 32);
large_int!(u128, u64, u64, 64);
large_int!(i128, u64, i64, 64);
<|endoftext|>


<commit_before>use super::*;

use std::iter;

/// Doc-comments are promoted to attributes that have `is_sugared_doc` = true
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Attribute {
    pub style: AttrStyle,
    pub value: MetaItem,
    pub is_sugared_doc: bool,
}

/// Distinguishes between Attributes that decorate items and Attributes that
/// are contained as statements within items. These two cases need to be
/// distinguished for pretty-printing.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum AttrStyle {
    Outer,
    Inner,
}

/// A compile-time attribute item.
///
/// E.g. `#[test]`, `#[derive(..)]` or `#[feature = "foo"]`
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum MetaItem {
    /// Word meta item.
    ///
    /// E.g. `test` as in `#[test]`
    Word(Ident),
    /// List meta item.
    ///
    /// E.g. `derive(..)` as in `#[derive(..)]`
    List(Ident, Vec<MetaItem>),
    /// Name value meta item.
    ///
    /// E.g. `feature = "foo"` as in `#[feature = "foo"]`
    NameValue(Ident, Lit),
}

impl MetaItem {
    pub fn name(&self) -> &str {
        match *self {
            MetaItem::Word(ref name) |
            MetaItem::List(ref name, _) |
            MetaItem::NameValue(ref name, _) => name.as_ref(),
        }
    }
}

pub trait FilterAttrs<'a> {
    type Ret: Iterator<Item = &'a Attribute>;

    fn outer(self) -> Self::Ret;
    fn inner(self) -> Self::Ret;
}

impl<'a, T> FilterAttrs<'a> for T where T: IntoIterator<Item = &'a Attribute> {
    type Ret = iter::Filter<T::IntoIter, fn(&&Attribute) -> bool>;

    fn outer(self) -> Self::Ret {
        fn is_outer(attr: &&Attribute) -> bool {
            attr.style == AttrStyle::Outer
        }
        self.into_iter().filter(is_outer)
    }

    fn inner(self) -> Self::Ret {
        fn is_inner(attr: &&Attribute) -> bool {
            attr.style == AttrStyle::Inner
        }
        self.into_iter().filter(is_inner)
    }
}

#[cfg(feature = "parsing")]
pub mod parsing {
    use super::*;
    use ident::parsing::ident;
    use lit::{Lit, StrStyle};
    use lit::parsing::lit;

    named!(pub inner_attr -> Attribute, do_parse!(
        punct!("#") >>
        punct!("!") >>
        punct!("[") >>
        meta_item: meta_item >>
        punct!("]") >>
        (Attribute {
            style: AttrStyle::Inner,
            value: meta_item,
            is_sugared_doc: false,
        })
    ));

    named!(pub outer_attr -> Attribute, alt!(
        do_parse!(
            punct!("#") >>
            punct!("[") >>
            meta_item: meta_item >>
            punct!("]") >>
            (Attribute {
                style: AttrStyle::Outer,
                value: meta_item,
                is_sugared_doc: false,
            })
        )
        |
        do_parse!(
            punct!("///") >>
            not!(peek!(tag!("/"))) >>
            content: take_until!("\n") >>
            (Attribute {
                style: AttrStyle::Outer,
                value: MetaItem::NameValue(
                    "doc".into(),
                    Lit::Str(
                        format!("///{}", content),
                        StrStyle::Cooked,
                    ),
                ),
                is_sugared_doc: true,
            })
        )
    ));

    named!(meta_item -> MetaItem, alt!(
        do_parse!(
            id: ident >>
            punct!("(") >>
            inner: separated_list!(punct!(","), meta_item) >>
            punct!(")") >>
            (MetaItem::List(id, inner))
        )
        |
        do_parse!(
            name: ident >>
            punct!("=") >>
            value: lit >>
            (MetaItem::NameValue(name, value))
        )
        |
        map!(ident, MetaItem::Word)
    ));
}

#[cfg(feature = "printing")]
mod printing {
    use super::*;
    use lit::{Lit, StrStyle};
    use quote::{Tokens, ToTokens};

    impl ToTokens for Attribute {
        fn to_tokens(&self, tokens: &mut Tokens) {
            match *self {
                Attribute {
                    style: AttrStyle::Outer,
                    value: MetaItem::NameValue(
                        ref name,
                        Lit::Str(ref value, StrStyle::Cooked),
                    ),
                    is_sugared_doc: true,
                } if name == "doc" && value.starts_with("///") => {
                    tokens.append(&format!("{}\n", value));
                }
                _ => {
                    tokens.append("#");
                    if let AttrStyle::Inner = self.style {
                        tokens.append("!");
                    }
                    tokens.append("[");
                    self.value.to_tokens(tokens);
                    tokens.append("]");
                }
            }
        }
    }

    impl ToTokens for MetaItem {
        fn to_tokens(&self, tokens: &mut Tokens) {
            match *self {
                MetaItem::Word(ref ident) => {
                    ident.to_tokens(tokens);
                }
                MetaItem::List(ref ident, ref inner) => {
                    ident.to_tokens(tokens);
                    tokens.append("(");
                    tokens.append_separated(inner, ",");
                    tokens.append(")");
                }
                MetaItem::NameValue(ref name, ref value) => {
                    name.to_tokens(tokens);
                    tokens.append("=");
                    value.to_tokens(tokens);
                }
            }
        }
    }
}
<commit_msg>Fix unused warning on inner_attr<commit_after>use super::*;

use std::iter;

/// Doc-comments are promoted to attributes that have `is_sugared_doc` = true
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Attribute {
    pub style: AttrStyle,
    pub value: MetaItem,
    pub is_sugared_doc: bool,
}

/// Distinguishes between Attributes that decorate items and Attributes that
/// are contained as statements within items. These two cases need to be
/// distinguished for pretty-printing.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum AttrStyle {
    Outer,
    Inner,
}

/// A compile-time attribute item.
///
/// E.g. `#[test]`, `#[derive(..)]` or `#[feature = "foo"]`
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum MetaItem {
    /// Word meta item.
    ///
    /// E.g. `test` as in `#[test]`
    Word(Ident),
    /// List meta item.
    ///
    /// E.g. `derive(..)` as in `#[derive(..)]`
    List(Ident, Vec<MetaItem>),
    /// Name value meta item.
    ///
    /// E.g. `feature = "foo"` as in `#[feature = "foo"]`
    NameValue(Ident, Lit),
}

impl MetaItem {
    pub fn name(&self) -> &str {
        match *self {
            MetaItem::Word(ref name) |
            MetaItem::List(ref name, _) |
            MetaItem::NameValue(ref name, _) => name.as_ref(),
        }
    }
}

pub trait FilterAttrs<'a> {
    type Ret: Iterator<Item = &'a Attribute>;

    fn outer(self) -> Self::Ret;
    fn inner(self) -> Self::Ret;
}

impl<'a, T> FilterAttrs<'a> for T where T: IntoIterator<Item = &'a Attribute> {
    type Ret = iter::Filter<T::IntoIter, fn(&&Attribute) -> bool>;

    fn outer(self) -> Self::Ret {
        fn is_outer(attr: &&Attribute) -> bool {
            attr.style == AttrStyle::Outer
        }
        self.into_iter().filter(is_outer)
    }

    fn inner(self) -> Self::Ret {
        fn is_inner(attr: &&Attribute) -> bool {
            attr.style == AttrStyle::Inner
        }
        self.into_iter().filter(is_inner)
    }
}

#[cfg(feature = "parsing")]
pub mod parsing {
    use super::*;
    use ident::parsing::ident;
    use lit::{Lit, StrStyle};
    use lit::parsing::lit;

    #[cfg(feature = "full")]
    named!(pub inner_attr -> Attribute, do_parse!(
        punct!("#") >>
        punct!("!") >>
        punct!("[") >>
        meta_item: meta_item >>
        punct!("]") >>
        (Attribute {
            style: AttrStyle::Inner,
            value: meta_item,
            is_sugared_doc: false,
        })
    ));

    named!(pub outer_attr -> Attribute, alt!(
        do_parse!(
            punct!("#") >>
            punct!("[") >>
            meta_item: meta_item >>
            punct!("]") >>
            (Attribute {
                style: AttrStyle::Outer,
                value: meta_item,
                is_sugared_doc: false,
            })
        )
        |
        do_parse!(
            punct!("///") >>
            not!(peek!(tag!("/"))) >>
            content: take_until!("\n") >>
            (Attribute {
                style: AttrStyle::Outer,
                value: MetaItem::NameValue(
                    "doc".into(),
                    Lit::Str(
                        format!("///{}", content),
                        StrStyle::Cooked,
                    ),
                ),
                is_sugared_doc: true,
            })
        )
    ));

    named!(meta_item -> MetaItem, alt!(
        do_parse!(
            id: ident >>
            punct!("(") >>
            inner: separated_list!(punct!(","), meta_item) >>
            punct!(")") >>
            (MetaItem::List(id, inner))
        )
        |
        do_parse!(
            name: ident >>
            punct!("=") >>
            value: lit >>
            (MetaItem::NameValue(name, value))
        )
        |
        map!(ident, MetaItem::Word)
    ));
}

#[cfg(feature = "printing")]
mod printing {
    use super::*;
    use lit::{Lit, StrStyle};
    use quote::{Tokens, ToTokens};

    impl ToTokens for Attribute {
        fn to_tokens(&self, tokens: &mut Tokens) {
            match *self {
                Attribute {
                    style: AttrStyle::Outer,
                    value: MetaItem::NameValue(
                        ref name,
                        Lit::Str(ref value, StrStyle::Cooked),
                    ),
                    is_sugared_doc: true,
                } if name == "doc" && value.starts_with("///") => {
                    tokens.append(&format!("{}\n", value));
                }
                _ => {
                    tokens.append("#");
                    if let AttrStyle::Inner = self.style {
                        tokens.append("!");
                    }
                    tokens.append("[");
                    self.value.to_tokens(tokens);
                    tokens.append("]");
                }
            }
        }
    }

    impl ToTokens for MetaItem {
        fn to_tokens(&self, tokens: &mut Tokens) {
            match *self {
                MetaItem::Word(ref ident) => {
                    ident.to_tokens(tokens);
                }
                MetaItem::List(ref ident, ref inner) => {
                    ident.to_tokens(tokens);
                    tokens.append("(");
                    tokens.append_separated(inner, ",");
                    tokens.append(")");
                }
                MetaItem::NameValue(ref name, ref value) => {
                    name.to_tokens(tokens);
                    tokens.append("=");
                    value.to_tokens(tokens);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::{HashMap, HashSet};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::Duration;

#[cfg(feature = "tls")]
use crate::exop_impl::StartTLS;
use crate::ldap::Ldap;
use crate::protocol::{ItemSender, LdapCodec, LdapOp, MaybeControls, ResultSender};
use crate::result::{LdapError, Result};
use crate::search::SearchItem;
use crate::RequestId;

use lber::structures::{Null, Tag};

#[cfg(feature = "tls")]
use futures_util::future::TryFutureExt;
use futures_util::sink::SinkExt;
#[cfg(feature = "tls")]
use native_tls::TlsConnector;
#[cfg(unix)]
use percent_encoding::percent_decode;
use tokio::io::{self, AsyncRead, AsyncWrite};
use tokio::net::TcpStream;
#[cfg(unix)]
use tokio::net::UnixStream;
use tokio::stream::StreamExt;
use tokio::sync::mpsc;
#[cfg(feature = "tls")]
use tokio::sync::oneshot;
#[cfg(feature = "tls")]
use tokio_tls::{TlsConnector as TokioTlsConnector, TlsStream};
use tokio_util::codec::{Decoder, Framed};
use url::{self, Url};

#[derive(Debug)]
enum ConnType {
    Tcp(TcpStream),
    #[cfg(feature = "tls")]
    Tls(TlsStream<TcpStream>),
    #[cfg(unix)]
    Unix(UnixStream),
}

impl AsyncRead for ConnType {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        match self.get_mut() {
            ConnType::Tcp(ts) => Pin::new(ts).poll_read(cx, buf),
            #[cfg(feature = "tls")]
            ConnType::Tls(tls) => Pin::new(tls).poll_read(cx, buf),
            #[cfg(unix)]
            ConnType::Unix(us) => Pin::new(us).poll_read(cx, buf),
        }
    }
}

impl AsyncWrite for ConnType {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
        match self.get_mut() {
            ConnType::Tcp(ts) => Pin::new(ts).poll_write(cx, buf),
            #[cfg(feature = "tls")]
            ConnType::Tls(tls) => Pin::new(tls).poll_write(cx, buf),
            #[cfg(unix)]
            ConnType::Unix(us) => Pin::new(us).poll_write(cx, buf),
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
        match self.get_mut() {
            ConnType::Tcp(ts) => Pin::new(ts).poll_flush(cx),
            #[cfg(feature = "tls")]
            ConnType::Tls(tls) => Pin::new(tls).poll_flush(cx),
            #[cfg(unix)]
            ConnType::Unix(us) => Pin::new(us).poll_flush(cx),
        }
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
        match self.get_mut() {
            ConnType::Tcp(ts) => Pin::new(ts).poll_shutdown(cx),
            #[cfg(feature = "tls")]
            ConnType::Tls(tls) => Pin::new(tls).poll_shutdown(cx),
            #[cfg(unix)]
            ConnType::Unix(us) => Pin::new(us).poll_shutdown(cx),
        }
    }
}

/// Additional settings for an LDAP connection.
///
/// The structure is opaque for better extensibility. An instance with
/// default values is constructed by [`new()`](#method.new), and all
/// available settings can be replaced through a builder-like interface,
/// by calling the appropriate functions.
#[derive(Clone, Default)]
pub struct LdapConnSettings {
    conn_timeout: Option<Duration>,
    #[cfg(feature = "tls")]
    connector: Option<TlsConnector>,
    #[cfg(feature = "tls")]
    starttls: bool,
    #[cfg(feature = "tls")]
    no_tls_verify: bool,
}

impl LdapConnSettings {
    /// Create an instance of the structure with default settings.
    pub fn new() -> LdapConnSettings {
        LdapConnSettings {
            ..Default::default()
        }
    }

    /// Set the connection timeout. If a connetion to the server can't
    /// be established before the timeout expires, an error will be
    /// returned to the user. Defaults to `None`, meaning an infinite
    /// timeout.
    pub fn set_conn_timeout(mut self, timeout: Duration) -> Self {
        self.conn_timeout = Some(timeout);
        self
    }

    #[cfg(feature = "tls")]
    /// Set a custom TLS connector, which enables setting various options
    /// when establishing a secure connection. See the documentation for
    /// [native_tls](https://docs.rs/native-tls/0.1.4/native_tls/).
    /// Defaults to `None`, which will use a connector with default
    /// settings.
    pub fn set_connector(mut self, connector: TlsConnector) -> Self {
        self.connector = Some(connector);
        self
    }

    #[cfg(feature = "tls")]
    /// If `true`, use the StartTLS extended operation to establish a
    /// secure connection. Defaults to `false`.
    pub fn set_starttls(mut self, starttls: bool) -> Self {
        self.starttls = starttls;
        self
    }

    #[cfg(feature = "tls")]
    /// The `starttls` settings indicates whether the StartTLS extended
    /// operation will be used to establish a secure connection.
    pub fn starttls(&self) -> bool {
        self.starttls
    }

    #[cfg(not(feature = "tls"))]
    /// Always `false` when no TLS support is compiled in.
    pub fn starttls(&self) -> bool {
        false
    }

    #[cfg(feature = "tls")]
    /// If `true`, try to establish a TLS connection without hostname
    /// verification. Defaults to `false`.
    ///
    /// The connection can still fail if the server certificate is
    /// considered invalid for other reasons (e.g., chain of trust or
    /// expiration date.) Depending on the platform, using a
    /// custom connector with backend-specific options _and_ setting
    /// this option to `true` may enable connections to servers with
    /// invalid certificates. One tested combination is OpenSSL with
    /// a connector for which `SSL_VERIFY_NONE` has been set.
    pub fn set_no_tls_verify(mut self, no_tls_verify: bool) -> Self {
        self.no_tls_verify = no_tls_verify;
        self
    }
}

enum LoopMode {
    #[allow(dead_code)]
    SingleOp,
    Continuous,
}

pub struct LdapConnAsync {
    msgmap: Arc<Mutex<(i32, HashSet<i32>)>>,
    resultmap: HashMap<i32, ResultSender>,
    searchmap: HashMap<i32, ItemSender>,
    rx: mpsc::UnboundedReceiver<(RequestId, LdapOp, Tag, MaybeControls, ResultSender)>,
    stream: Framed<ConnType, LdapCodec>,
}

#[macro_export]
macro_rules! drive {
    ($conn:expr) => {
        tokio::spawn(async move {
            if let Err(e) = $conn.drive().await {
                $crate::log::warn!("LDAP connection error: {}", e);
            }
        });
    };
}

impl LdapConnAsync {
    pub async fn with_settings(settings: LdapConnSettings, url: &str) -> Result<(Self, Ldap)> {
        if url.starts_with("ldapi://") {
            Ok(LdapConnAsync::new_unix(url, settings).await?)
        } else {
            Ok(LdapConnAsync::new_tcp(url, settings).await?)
        }
    }

    pub async fn new(url: &str) -> Result<(Self, Ldap)> {
        if url.starts_with("ldapi://") {
            Ok(LdapConnAsync::new_unix(url, LdapConnSettings::new()).await?)
        } else {
            Ok(LdapConnAsync::new_tcp(url, LdapConnSettings::new()).await?)
        }
    }

    #[cfg(unix)]
    async fn new_unix(url: &str, _settings: LdapConnSettings) -> Result<(Self, Ldap)> {
        let path = url.split('/').nth(2).unwrap();
        if path.is_empty() {
            return Err(LdapError::EmptyUnixPath);
        }
        if path.contains(':') {
            return Err(LdapError::PortInUnixPath);
        }
        let dec_path = percent_decode(path.as_bytes()).decode_utf8_lossy();
        let stream = UnixStream::connect(dec_path.as_ref()).await?;
        Ok(Self::conn_pair(ConnType::Unix(stream)))
    }

    #[cfg(not(unix))]
    async fn new_unix(_url: &str, _settings: LdapConnSettings) -> Result<(Self, Ldap)> {
        unimplemented!("no Unix domain sockets on non-Unix platforms");
    }

    #[allow(unused_mut)]
    async fn new_tcp(url: &str, mut settings: LdapConnSettings) -> Result<(Self, Ldap)> {
        let url = Url::parse(url)?;
        let mut port = 389;
        let scheme = match url.scheme() {
            s @ "ldap" => {
                if settings.starttls() {
                    "starttls"
                } else {
                    s
                }
            }
            #[cfg(feature = "tls")]
            s @ "ldaps" => {
                settings = settings.set_starttls(false);
                port = 636;
                s
            }
            s => return Err(LdapError::UnknownScheme(String::from(s))),
        };
        if let Some(url_port) = url.port() {
            port = url_port;
        }
        let (_hostname, host_port) = match url.host_str() {
            Some(h) if h != "" => (h, format!("{}:{}", h, port)),
            Some(h) if h == "" => ("localhost", format!("localhost:{}", port)),
            _ => panic!("unexpected None from url.host_str()"),
        };
        let stream = TcpStream::connect(host_port.as_str()).await?;
        let (mut conn, mut ldap) = Self::conn_pair(ConnType::Tcp(stream));
        match scheme {
            "ldap" => (),
            #[cfg(feature = "tls")]
            s @ "ldaps" | s @ "starttls" => {
                let connector = match settings.connector {
                    Some(connector) => connector,
                    None => {
                        let mut builder = TlsConnector::builder();
                        if settings.no_tls_verify {
                            builder.danger_accept_invalid_certs(true);
                        }
                        builder.build().expect("connector")
                    }
                };
                if s == "starttls" {
                    let (tx, rx) = oneshot::channel();
                    tokio::spawn(async move {
                        conn.single_op(tx).await;
                    });
                    let res = tokio::try_join!(
                        rx.map_err(|e| LdapError::from(e)),
                        ldap.extended(StartTLS)
                    );
                    match res {
                        Ok((conn_res, res)) => {
                            conn = conn_res?;
                            res.success()?;
                        }
                        Err(e) => return Err(e),
                    }
                }
                let parts = conn.stream.into_parts();
                let tls_stream = if let ConnType::Tcp(stream) = parts.io {
                    TokioTlsConnector::from(connector)
                        .connect(_hostname, stream)
                        .await?
                } else {
                    panic!("underlying stream not TCP");
                };
                conn.stream = parts.codec.framed(ConnType::Tls(tls_stream));
            }
            _ => unimplemented!(),
        }
        Ok((conn, ldap))
    }

    fn conn_pair(ctype: ConnType) -> (Self, Ldap) {
        let codec = LdapCodec;
        let (tx, rx) = mpsc::unbounded_channel();
        let conn = LdapConnAsync {
            msgmap: Arc::new(Mutex::new((0, HashSet::new()))),
            resultmap: HashMap::new(),
            searchmap: HashMap::new(),
            rx: rx,
            stream: codec.framed(ctype),
        };
        let ldap = Ldap {
            msgmap: conn.msgmap.clone(),
            tx: tx,
            last_id: 0,
            timeout: None,
            controls: None,
            search_opts: None,
        };
        (conn, ldap)
    }

    pub async fn drive(self) -> Result<()> {
        self.turn(LoopMode::Continuous).await.map(|_| ())
    }

    #[cfg(feature = "tls")]
    pub(crate) async fn single_op(self, tx: oneshot::Sender<Result<Self>>) {
        if let Err(_) = tx.send(self.turn(LoopMode::SingleOp).await) {
            warn!("single op send error");
        }
    }

    async fn turn(mut self, mode: LoopMode) -> Result<Self> {
        loop {
            tokio::select! {
                op_tuple = self.rx.recv() => {
                    if let Some((id, op, tag, controls, tx)) = op_tuple {
                        if let &LdapOp::Search(ref search_tx) = &op {
                            self.searchmap.insert(id, search_tx.clone());
                        }
                        if let Err(e) = self.stream.send((id, tag, controls)).await {
                            warn!("socket send error: {}", e);
                            return Err(LdapError::from(e));
                        } else {
                            match op {
                                LdapOp::Single => {
                                    self.resultmap.insert(id, tx);
                                    continue;
                                },
                                LdapOp::Search(_) => (),
                                LdapOp::Abandon(msgid) => {
                                    let mut msgmap = self.msgmap.lock().expect("msgmap mutex (abandon)");
                                    msgmap.1.remove(&id);
                                    self.resultmap.remove(&msgid);
                                    self.searchmap.remove(&msgid);
                                },
                                LdapOp::Unbind => {
                                    if let Err(e) = self.stream.close().await {
                                        warn!("socket shutdown error: {}", e);
                                        return Err(LdapError::from(e));
                                    }
                                },
                            }
                            if let Err(e) = tx.send((Tag::Null(Null { ..Default::default() }), vec![])) {
                                warn!("ldap null result send error: {:?}", e);
                            }
                        }
                    }
                },
                resp = self.stream.next() => {
                    let (id, (tag, controls)) = match resp {
                        None => break,
                        Some(Err(e)) => {
                            warn!("socket receive error: {}", e);
                            return Err(LdapError::from(e));
                        },
                        Some(Ok(resp)) => resp,
                    };
                    if let Some(tx) = self.searchmap.get(&id) {
                        let protoop = if let Tag::StructureTag(protoop) = tag {
                            protoop
                        } else {
                            panic!("unmatched tag structure: {:?}", tag);
                        };
                        let (item, mut remove) = match protoop.id {
                            4 | 25 => (SearchItem::Entry(protoop), false),
                            5 => (SearchItem::Done(Tag::StructureTag(protoop).into(), controls), true),
                            19 => (SearchItem::Referral(protoop), false),
                            _ => panic!("unrecognized op id: {}", protoop.id),
                        };
                        if let Err(e) = tx.send(item) {
                            warn!("ldap search item send error, op={}: {:?}", id, e);
                            remove = true;
                        }
                        if remove {
                            self.searchmap.remove(&id);
                        }
                    } else if let Some(tx) = self.resultmap.remove(&id) {
                        let mut msgmap = self.msgmap.lock().expect("msgmap mutex (stream rx)");
                        msgmap.1.remove(&id);
                        if let Err(e) = tx.send((tag, controls)) {
                            warn!("ldap result send error: {:?}", e);
                        }
                    } else {
                        warn!("unmatched id: {}", id);
                    }
                },
            };
            if let LoopMode::SingleOp = mode {
                break;
            }
        }
        Ok(self)
    }
}
<commit_msg>Add timeout handling to connection establishment<commit_after>use std::collections::{HashMap, HashSet};
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use std::time::Duration;

#[cfg(feature = "tls")]
use crate::exop_impl::StartTLS;
use crate::ldap::Ldap;
use crate::protocol::{ItemSender, LdapCodec, LdapOp, MaybeControls, ResultSender};
use crate::result::{LdapError, Result};
use crate::search::SearchItem;
use crate::RequestId;

use lber::structures::{Null, Tag};

#[cfg(feature = "tls")]
use futures_util::future::TryFutureExt;
use futures_util::sink::SinkExt;
#[cfg(feature = "tls")]
use native_tls::TlsConnector;
#[cfg(unix)]
use percent_encoding::percent_decode;
use tokio::io::{self, AsyncRead, AsyncWrite};
use tokio::net::TcpStream;
#[cfg(unix)]
use tokio::net::UnixStream;
use tokio::stream::StreamExt;
use tokio::sync::mpsc;
#[cfg(feature = "tls")]
use tokio::sync::oneshot;
#[cfg(feature = "tls")]
use tokio::time;
use tokio_tls::{TlsConnector as TokioTlsConnector, TlsStream};
use tokio_util::codec::{Decoder, Framed};
use url::{self, Url};

#[derive(Debug)]
enum ConnType {
    Tcp(TcpStream),
    #[cfg(feature = "tls")]
    Tls(TlsStream<TcpStream>),
    #[cfg(unix)]
    Unix(UnixStream),
}

impl AsyncRead for ConnType {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context,
        buf: &mut [u8],
    ) -> Poll<io::Result<usize>> {
        match self.get_mut() {
            ConnType::Tcp(ts) => Pin::new(ts).poll_read(cx, buf),
            #[cfg(feature = "tls")]
            ConnType::Tls(tls) => Pin::new(tls).poll_read(cx, buf),
            #[cfg(unix)]
            ConnType::Unix(us) => Pin::new(us).poll_read(cx, buf),
        }
    }
}

impl AsyncWrite for ConnType {
    fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
        match self.get_mut() {
            ConnType::Tcp(ts) => Pin::new(ts).poll_write(cx, buf),
            #[cfg(feature = "tls")]
            ConnType::Tls(tls) => Pin::new(tls).poll_write(cx, buf),
            #[cfg(unix)]
            ConnType::Unix(us) => Pin::new(us).poll_write(cx, buf),
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
        match self.get_mut() {
            ConnType::Tcp(ts) => Pin::new(ts).poll_flush(cx),
            #[cfg(feature = "tls")]
            ConnType::Tls(tls) => Pin::new(tls).poll_flush(cx),
            #[cfg(unix)]
            ConnType::Unix(us) => Pin::new(us).poll_flush(cx),
        }
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
        match self.get_mut() {
            ConnType::Tcp(ts) => Pin::new(ts).poll_shutdown(cx),
            #[cfg(feature = "tls")]
            ConnType::Tls(tls) => Pin::new(tls).poll_shutdown(cx),
            #[cfg(unix)]
            ConnType::Unix(us) => Pin::new(us).poll_shutdown(cx),
        }
    }
}

/// Additional settings for an LDAP connection.
///
/// The structure is opaque for better extensibility. An instance with
/// default values is constructed by [`new()`](#method.new), and all
/// available settings can be replaced through a builder-like interface,
/// by calling the appropriate functions.
#[derive(Clone, Default)]
pub struct LdapConnSettings {
    conn_timeout: Option<Duration>,
    #[cfg(feature = "tls")]
    connector: Option<TlsConnector>,
    #[cfg(feature = "tls")]
    starttls: bool,
    #[cfg(feature = "tls")]
    no_tls_verify: bool,
}

impl LdapConnSettings {
    /// Create an instance of the structure with default settings.
    pub fn new() -> LdapConnSettings {
        LdapConnSettings {
            ..Default::default()
        }
    }

    /// Set the connection timeout. If a connetion to the server can't
    /// be established before the timeout expires, an error will be
    /// returned to the user. Defaults to `None`, meaning an infinite
    /// timeout.
    pub fn set_conn_timeout(mut self, timeout: Duration) -> Self {
        self.conn_timeout = Some(timeout);
        self
    }

    #[cfg(feature = "tls")]
    /// Set a custom TLS connector, which enables setting various options
    /// when establishing a secure connection. See the documentation for
    /// [native_tls](https://docs.rs/native-tls/0.1.4/native_tls/).
    /// Defaults to `None`, which will use a connector with default
    /// settings.
    pub fn set_connector(mut self, connector: TlsConnector) -> Self {
        self.connector = Some(connector);
        self
    }

    #[cfg(feature = "tls")]
    /// If `true`, use the StartTLS extended operation to establish a
    /// secure connection. Defaults to `false`.
    pub fn set_starttls(mut self, starttls: bool) -> Self {
        self.starttls = starttls;
        self
    }

    #[cfg(feature = "tls")]
    /// The `starttls` settings indicates whether the StartTLS extended
    /// operation will be used to establish a secure connection.
    pub fn starttls(&self) -> bool {
        self.starttls
    }

    #[cfg(not(feature = "tls"))]
    /// Always `false` when no TLS support is compiled in.
    pub fn starttls(&self) -> bool {
        false
    }

    #[cfg(feature = "tls")]
    /// If `true`, try to establish a TLS connection without hostname
    /// verification. Defaults to `false`.
    ///
    /// The connection can still fail if the server certificate is
    /// considered invalid for other reasons (e.g., chain of trust or
    /// expiration date.) Depending on the platform, using a
    /// custom connector with backend-specific options _and_ setting
    /// this option to `true` may enable connections to servers with
    /// invalid certificates. One tested combination is OpenSSL with
    /// a connector for which `SSL_VERIFY_NONE` has been set.
    pub fn set_no_tls_verify(mut self, no_tls_verify: bool) -> Self {
        self.no_tls_verify = no_tls_verify;
        self
    }
}

enum LoopMode {
    #[allow(dead_code)]
    SingleOp,
    Continuous,
}

pub struct LdapConnAsync {
    msgmap: Arc<Mutex<(i32, HashSet<i32>)>>,
    resultmap: HashMap<i32, ResultSender>,
    searchmap: HashMap<i32, ItemSender>,
    rx: mpsc::UnboundedReceiver<(RequestId, LdapOp, Tag, MaybeControls, ResultSender)>,
    stream: Framed<ConnType, LdapCodec>,
}

#[macro_export]
macro_rules! drive {
    ($conn:expr) => {
        tokio::spawn(async move {
            if let Err(e) = $conn.drive().await {
                $crate::log::warn!("LDAP connection error: {}", e);
            }
        });
    };
}

impl LdapConnAsync {
    pub async fn with_settings(mut settings: LdapConnSettings, url: &str) -> Result<(Self, Ldap)> {
        if url.starts_with("ldapi://") {
            Ok(LdapConnAsync::new_unix(url, settings).await?)
        } else {
            let timeout = settings.conn_timeout.take();
            let conn_future = LdapConnAsync::new_tcp(url, settings);
            Ok(if let Some(timeout) = timeout {
                time::timeout(timeout, conn_future).await?
            } else {
                conn_future.await
            }?)
        }
    }

    pub async fn new(url: &str) -> Result<(Self, Ldap)> {
        Self::with_settings(LdapConnSettings::new(), url).await
    }

    #[cfg(unix)]
    async fn new_unix(url: &str, _settings: LdapConnSettings) -> Result<(Self, Ldap)> {
        let path = url.split('/').nth(2).unwrap();
        if path.is_empty() {
            return Err(LdapError::EmptyUnixPath);
        }
        if path.contains(':') {
            return Err(LdapError::PortInUnixPath);
        }
        let dec_path = percent_decode(path.as_bytes()).decode_utf8_lossy();
        let stream = UnixStream::connect(dec_path.as_ref()).await?;
        Ok(Self::conn_pair(ConnType::Unix(stream)))
    }

    #[cfg(not(unix))]
    async fn new_unix(_url: &str, _settings: LdapConnSettings) -> Result<(Self, Ldap)> {
        unimplemented!("no Unix domain sockets on non-Unix platforms");
    }

    #[allow(unused_mut)]
    async fn new_tcp(url: &str, mut settings: LdapConnSettings) -> Result<(Self, Ldap)> {
        let url = Url::parse(url)?;
        let mut port = 389;
        let scheme = match url.scheme() {
            s @ "ldap" => {
                if settings.starttls() {
                    "starttls"
                } else {
                    s
                }
            }
            #[cfg(feature = "tls")]
            s @ "ldaps" => {
                settings = settings.set_starttls(false);
                port = 636;
                s
            }
            s => return Err(LdapError::UnknownScheme(String::from(s))),
        };
        if let Some(url_port) = url.port() {
            port = url_port;
        }
        let (_hostname, host_port) = match url.host_str() {
            Some(h) if h != "" => (h, format!("{}:{}", h, port)),
            Some(h) if h == "" => ("localhost", format!("localhost:{}", port)),
            _ => panic!("unexpected None from url.host_str()"),
        };
        let stream = TcpStream::connect(host_port.as_str()).await?;
        let (mut conn, mut ldap) = Self::conn_pair(ConnType::Tcp(stream));
        match scheme {
            "ldap" => (),
            #[cfg(feature = "tls")]
            s @ "ldaps" | s @ "starttls" => {
                let connector = match settings.connector {
                    Some(connector) => connector,
                    None => {
                        let mut builder = TlsConnector::builder();
                        if settings.no_tls_verify {
                            builder.danger_accept_invalid_certs(true);
                        }
                        builder.build().expect("connector")
                    }
                };
                if s == "starttls" {
                    let (tx, rx) = oneshot::channel();
                    tokio::spawn(async move {
                        conn.single_op(tx).await;
                    });
                    let res = tokio::try_join!(
                        rx.map_err(|e| LdapError::from(e)),
                        ldap.extended(StartTLS)
                    );
                    match res {
                        Ok((conn_res, res)) => {
                            conn = conn_res?;
                            res.success()?;
                        }
                        Err(e) => return Err(e),
                    }
                }
                let parts = conn.stream.into_parts();
                let tls_stream = if let ConnType::Tcp(stream) = parts.io {
                    TokioTlsConnector::from(connector)
                        .connect(_hostname, stream)
                        .await?
                } else {
                    panic!("underlying stream not TCP");
                };
                conn.stream = parts.codec.framed(ConnType::Tls(tls_stream));
            }
            _ => unimplemented!(),
        }
        Ok((conn, ldap))
    }

    fn conn_pair(ctype: ConnType) -> (Self, Ldap) {
        let codec = LdapCodec;
        let (tx, rx) = mpsc::unbounded_channel();
        let conn = LdapConnAsync {
            msgmap: Arc::new(Mutex::new((0, HashSet::new()))),
            resultmap: HashMap::new(),
            searchmap: HashMap::new(),
            rx: rx,
            stream: codec.framed(ctype),
        };
        let ldap = Ldap {
            msgmap: conn.msgmap.clone(),
            tx: tx,
            last_id: 0,
            timeout: None,
            controls: None,
            search_opts: None,
        };
        (conn, ldap)
    }

    pub async fn drive(self) -> Result<()> {
        self.turn(LoopMode::Continuous).await.map(|_| ())
    }

    #[cfg(feature = "tls")]
    pub(crate) async fn single_op(self, tx: oneshot::Sender<Result<Self>>) {
        if let Err(_) = tx.send(self.turn(LoopMode::SingleOp).await) {
            warn!("single op send error");
        }
    }

    async fn turn(mut self, mode: LoopMode) -> Result<Self> {
        loop {
            tokio::select! {
                op_tuple = self.rx.recv() => {
                    if let Some((id, op, tag, controls, tx)) = op_tuple {
                        if let &LdapOp::Search(ref search_tx) = &op {
                            self.searchmap.insert(id, search_tx.clone());
                        }
                        if let Err(e) = self.stream.send((id, tag, controls)).await {
                            warn!("socket send error: {}", e);
                            return Err(LdapError::from(e));
                        } else {
                            match op {
                                LdapOp::Single => {
                                    self.resultmap.insert(id, tx);
                                    continue;
                                },
                                LdapOp::Search(_) => (),
                                LdapOp::Abandon(msgid) => {
                                    let mut msgmap = self.msgmap.lock().expect("msgmap mutex (abandon)");
                                    msgmap.1.remove(&id);
                                    self.resultmap.remove(&msgid);
                                    self.searchmap.remove(&msgid);
                                },
                                LdapOp::Unbind => {
                                    if let Err(e) = self.stream.close().await {
                                        warn!("socket shutdown error: {}", e);
                                        return Err(LdapError::from(e));
                                    }
                                },
                            }
                            if let Err(e) = tx.send((Tag::Null(Null { ..Default::default() }), vec![])) {
                                warn!("ldap null result send error: {:?}", e);
                            }
                        }
                    }
                },
                resp = self.stream.next() => {
                    let (id, (tag, controls)) = match resp {
                        None => break,
                        Some(Err(e)) => {
                            warn!("socket receive error: {}", e);
                            return Err(LdapError::from(e));
                        },
                        Some(Ok(resp)) => resp,
                    };
                    if let Some(tx) = self.searchmap.get(&id) {
                        let protoop = if let Tag::StructureTag(protoop) = tag {
                            protoop
                        } else {
                            panic!("unmatched tag structure: {:?}", tag);
                        };
                        let (item, mut remove) = match protoop.id {
                            4 | 25 => (SearchItem::Entry(protoop), false),
                            5 => (SearchItem::Done(Tag::StructureTag(protoop).into(), controls), true),
                            19 => (SearchItem::Referral(protoop), false),
                            _ => panic!("unrecognized op id: {}", protoop.id),
                        };
                        if let Err(e) = tx.send(item) {
                            warn!("ldap search item send error, op={}: {:?}", id, e);
                            remove = true;
                        }
                        if remove {
                            self.searchmap.remove(&id);
                        }
                    } else if let Some(tx) = self.resultmap.remove(&id) {
                        let mut msgmap = self.msgmap.lock().expect("msgmap mutex (stream rx)");
                        msgmap.1.remove(&id);
                        if let Err(e) = tx.send((tag, controls)) {
                            warn!("ldap result send error: {:?}", e);
                        }
                    } else {
                        warn!("unmatched id: {}", id);
                    }
                },
            };
            if let LoopMode::SingleOp = mode {
                break;
            }
        }
        Ok(self)
    }
}
<|endoftext|>


<commit_before>use std::str::FromStr;

/// A TIS-100 port.
#[derive(Debug, PartialEq, Copy, Clone, Hash)]
pub enum Port {
    Up,
    Down,
    Left,
    Right,
}

/// An error which can be returned when parsing a port.
#[derive(Debug, PartialEq)]
pub struct ParsePortError;

impl FromStr for Port {
    type Err = ParsePortError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "UP" => Ok(Port::Up),
            "DOWN" => Ok(Port::Down),
            "LEFT" => Ok(Port::Left),
            "RIGHT" => Ok(Port::Right),
            _ => Err(ParsePortError)
        }
    }
}

/// A TIS-100 port or pseudo-port.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum PseudoPort {
    Any,
    Last,
}

/// An error which can be returned when parsing a psedo-port.
#[derive(Debug, PartialEq)]
pub struct ParsePseudoPortError;

impl FromStr for PseudoPort {
    type Err = ParsePseudoPortError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "ANY" => Ok(PseudoPort::Any),
            "LAST" => Ok(PseudoPort::Last),
            _ => Err(ParsePseudoPortError)
        }
    }
}

/// A TIS-100 port or pseudo-port.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum IoRegister {
    Port(Port),
    Pseudo(PseudoPort),
}

/// An error which can be returned when parsing an IO register.
#[derive(Debug, PartialEq)]
pub struct ParseIoRegisterError;

impl FromStr for IoRegister {
    type Err = ParseIoRegisterError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if let Ok(port) = str::parse::<Port>(s) {
            Ok(IoRegister::Port(port))
        } else if let Ok(pseudo) = str::parse::<PseudoPort>(s) {
            Ok(IoRegister::Pseudo(pseudo))
        } else {
            Err(ParseIoRegisterError)
        }
    }
}

/// A TIS-100 register.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Register {
    Acc,
    Nil,
    Io(IoRegister),
}

/// An error which can be returned when parsing a register.
#[derive(Debug, PartialEq)]
pub struct ParseRegisterError;

impl FromStr for Register {
    type Err = ParseRegisterError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "ACC" => Ok(Register::Acc),
            "NIL" => Ok(Register::Nil),
            _ => {
                if let Ok(reg) = str::parse::<IoRegister>(s) {
                    Ok(Register::Io(reg))
                } else {
                    Err(ParseRegisterError)
                }
            }
        }
    }
}

/// The source component of a TIS-100 instruction.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Source {
    Val(isize),
    Reg(Register),
}

/// An error which can be returned when parsing an source.
#[derive(Debug, PartialEq)]
pub struct ParseSourceError;

impl FromStr for Source {
    type Err = ParseSourceError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if let Ok(val) = str::parse::<isize>(s) {
            Ok(Source::Val(val))
        } else if let Ok(register) = str::parse::<Register>(s) {
            Ok(Source::Reg(register))
        } else {
            Err(ParseSourceError)
        }
    }
}

/// The opcode component of a TIS-100 instruction.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Opcode {
    Nop,
    Mov,
    Swp,
    Sav,
    Add,
    Sub,
    Neg,
    Jmp,
    Jez,
    Jnz,
    Jgz,
    Jlz,
    Jro,
}

/// An error which can be returned when parsing an opcode.
#[derive(Debug, PartialEq)]
pub struct ParseOpcodeError;

impl FromStr for Opcode {
    type Err = ParseOpcodeError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "NOP" => Ok(Opcode::Nop),
            "MOV" => Ok(Opcode::Mov),
            "SWP" => Ok(Opcode::Swp),
            "SAV" => Ok(Opcode::Sav),
            "ADD" => Ok(Opcode::Add),
            "SUB" => Ok(Opcode::Sub),
            "NEG" => Ok(Opcode::Neg),
            "JMP" => Ok(Opcode::Jmp),
            "JEZ" => Ok(Opcode::Jez),
            "JNZ" => Ok(Opcode::Jnz),
            "JGZ" => Ok(Opcode::Jgz),
            "JLZ" => Ok(Opcode::Jlz),
            "JRO" => Ok(Opcode::Jro),
            _ => Err(ParseOpcodeError),
        }
    }
}

/// A valid TIS-100 instruction.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Instruction {
    Nop,
    Mov(Source, Register),
    Swp,
    Sav,
    Add(Source),
    Sub(Source),
    Neg,
    Jmp(usize),
    Jez(usize),
    Jnz(usize),
    Jgz(usize),
    Jlz(usize),
    Jro(Source),
}

/// The list of instructions created by parsing the program source code. The
/// instructions can then be evaluated by a basic execution node.
pub type Program = Vec<Instruction>;

#[test]
fn test_parse_port() {
    assert_eq!(str::parse::<Port>("UP"), Ok(Port::Up));
    assert_eq!(str::parse::<Port>("up"), Ok(Port::Up));
    assert_eq!(str::parse::<Port>("DOWN"), Ok(Port::Down));
    assert_eq!(str::parse::<Port>("LEFT"), Ok(Port::Left));
    assert_eq!(str::parse::<Port>("RIGHT"), Ok(Port::Right));
    assert_eq!(str::parse::<Port>("bad"), Err(ParsePortError));
}

#[test]
fn test_parse_psedo_port() {
    assert_eq!(str::parse::<PseudoPort>("ANY"), Ok(PseudoPort::Any));
    assert_eq!(str::parse::<PseudoPort>("any"), Ok(PseudoPort::Any));
    assert_eq!(str::parse::<PseudoPort>("LAST"), Ok(PseudoPort::Last));
    assert_eq!(str::parse::<PseudoPort>("bad"), Err(ParsePseudoPortError));
}

#[test]
fn test_parse_io_register() {
    assert_eq!(str::parse::<IoRegister>("UP"), Ok(IoRegister::Port(Port::Up)));
    assert_eq!(str::parse::<IoRegister>("ANY"), Ok(IoRegister::Pseudo(PseudoPort::Any)));
    assert_eq!(str::parse::<IoRegister>("bad"), Err(ParsePortError));
}

#[test]
fn test_parse_register() {
    assert_eq!(str::parse::<Register>("ACC"), Ok(Register::Acc));
    assert_eq!(str::parse::<Register>("acc"), Ok(Register::Acc));
    assert_eq!(str::parse::<Register>("NIl"), Ok(Register::Nil));
    assert_eq!(str::parse::<Register>("UP"), Ok(Register::Io(IoRegister::Port(Port::Up))));
    assert_eq!(str::parse::<Register>("bad"), Err(ParseRegisterError));
}

#[test]
fn test_parse_source() {
    assert_eq!(str::parse::<Source>("ACC"), Ok(Source::Reg(Register::Acc)));
    assert_eq!(str::parse::<Source>("1"), Ok(Source::Val(1)));
    assert_eq!(str::parse::<Source>("bad"), Err(ParseSourceError));
}

#[test]
fn test_parse_opcode() {
    assert_eq!(str::parse::<Opcode>("NOP"), Ok(Opcode::Nop));
    assert_eq!(str::parse::<Opcode>("nop"), Ok(Opcode::Nop));
    assert_eq!(str::parse::<Opcode>("MOV"), Ok(Opcode::Mov));
    assert_eq!(str::parse::<Opcode>("SWP"), Ok(Opcode::Swp));
    assert_eq!(str::parse::<Opcode>("SAV"), Ok(Opcode::Sav));
    assert_eq!(str::parse::<Opcode>("ADD"), Ok(Opcode::Add));
    assert_eq!(str::parse::<Opcode>("SUB"), Ok(Opcode::Sub));
    assert_eq!(str::parse::<Opcode>("NEG"), Ok(Opcode::Neg));
    assert_eq!(str::parse::<Opcode>("JMP"), Ok(Opcode::Jmp));
    assert_eq!(str::parse::<Opcode>("JEZ"), Ok(Opcode::Jez));
    assert_eq!(str::parse::<Opcode>("JNZ"), Ok(Opcode::Jnz));
    assert_eq!(str::parse::<Opcode>("JGZ"), Ok(Opcode::Jgz));
    assert_eq!(str::parse::<Opcode>("JLZ"), Ok(Opcode::Jlz));
    assert_eq!(str::parse::<Opcode>("JRO"), Ok(Opcode::Jro));
    assert_eq!(str::parse::<Opcode>("bad"), Err(ParseOpcodeError));
}
<commit_msg>Rename IoRegister::Port and Pseudo to PortReg and PseudoReg<commit_after>use std::str::FromStr;

/// A TIS-100 port.
#[derive(Debug, PartialEq, Copy, Clone, Hash)]
pub enum Port {
    Up,
    Down,
    Left,
    Right,
}

/// An error which can be returned when parsing a port.
#[derive(Debug, PartialEq)]
pub struct ParsePortError;

impl FromStr for Port {
    type Err = ParsePortError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "UP" => Ok(Port::Up),
            "DOWN" => Ok(Port::Down),
            "LEFT" => Ok(Port::Left),
            "RIGHT" => Ok(Port::Right),
            _ => Err(ParsePortError)
        }
    }
}

/// A TIS-100 port or pseudo-port.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum PseudoPort {
    Any,
    Last,
}

/// An error which can be returned when parsing a psedo-port.
#[derive(Debug, PartialEq)]
pub struct ParsePseudoPortError;

impl FromStr for PseudoPort {
    type Err = ParsePseudoPortError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "ANY" => Ok(PseudoPort::Any),
            "LAST" => Ok(PseudoPort::Last),
            _ => Err(ParsePseudoPortError)
        }
    }
}

/// A TIS-100 port or pseudo-port.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum IoRegister {
    PortReg(Port),
    PseudoReg(PseudoPort),
}

/// An error which can be returned when parsing an IO register.
#[derive(Debug, PartialEq)]
pub struct ParseIoRegisterError;

impl FromStr for IoRegister {
    type Err = ParseIoRegisterError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if let Ok(port) = str::parse::<Port>(s) {
            Ok(IoRegister::PortReg(port))
        } else if let Ok(pseudo) = str::parse::<PseudoPort>(s) {
            Ok(IoRegister::PseudoReg(pseudo))
        } else {
            Err(ParseIoRegisterError)
        }
    }
}

/// A TIS-100 register.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Register {
    Acc,
    Nil,
    Io(IoRegister),
}

/// An error which can be returned when parsing a register.
#[derive(Debug, PartialEq)]
pub struct ParseRegisterError;

impl FromStr for Register {
    type Err = ParseRegisterError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "ACC" => Ok(Register::Acc),
            "NIL" => Ok(Register::Nil),
            _ => {
                if let Ok(reg) = str::parse::<IoRegister>(s) {
                    Ok(Register::Io(reg))
                } else {
                    Err(ParseRegisterError)
                }
            }
        }
    }
}

/// The source component of a TIS-100 instruction.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Source {
    Val(isize),
    Reg(Register),
}

/// An error which can be returned when parsing an source.
#[derive(Debug, PartialEq)]
pub struct ParseSourceError;

impl FromStr for Source {
    type Err = ParseSourceError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if let Ok(val) = str::parse::<isize>(s) {
            Ok(Source::Val(val))
        } else if let Ok(register) = str::parse::<Register>(s) {
            Ok(Source::Reg(register))
        } else {
            Err(ParseSourceError)
        }
    }
}

/// The opcode component of a TIS-100 instruction.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Opcode {
    Nop,
    Mov,
    Swp,
    Sav,
    Add,
    Sub,
    Neg,
    Jmp,
    Jez,
    Jnz,
    Jgz,
    Jlz,
    Jro,
}

/// An error which can be returned when parsing an opcode.
#[derive(Debug, PartialEq)]
pub struct ParseOpcodeError;

impl FromStr for Opcode {
    type Err = ParseOpcodeError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_uppercase().as_str() {
            "NOP" => Ok(Opcode::Nop),
            "MOV" => Ok(Opcode::Mov),
            "SWP" => Ok(Opcode::Swp),
            "SAV" => Ok(Opcode::Sav),
            "ADD" => Ok(Opcode::Add),
            "SUB" => Ok(Opcode::Sub),
            "NEG" => Ok(Opcode::Neg),
            "JMP" => Ok(Opcode::Jmp),
            "JEZ" => Ok(Opcode::Jez),
            "JNZ" => Ok(Opcode::Jnz),
            "JGZ" => Ok(Opcode::Jgz),
            "JLZ" => Ok(Opcode::Jlz),
            "JRO" => Ok(Opcode::Jro),
            _ => Err(ParseOpcodeError),
        }
    }
}

/// A valid TIS-100 instruction.
#[derive(Debug, PartialEq, Copy, Clone)]
pub enum Instruction {
    Nop,
    Mov(Source, Register),
    Swp,
    Sav,
    Add(Source),
    Sub(Source),
    Neg,
    Jmp(usize),
    Jez(usize),
    Jnz(usize),
    Jgz(usize),
    Jlz(usize),
    Jro(Source),
}

/// The list of instructions created by parsing the program source code. The
/// instructions can then be evaluated by a basic execution node.
pub type Program = Vec<Instruction>;

#[test]
fn test_parse_port() {
    assert_eq!(str::parse::<Port>("UP"), Ok(Port::Up));
    assert_eq!(str::parse::<Port>("up"), Ok(Port::Up));
    assert_eq!(str::parse::<Port>("DOWN"), Ok(Port::Down));
    assert_eq!(str::parse::<Port>("LEFT"), Ok(Port::Left));
    assert_eq!(str::parse::<Port>("RIGHT"), Ok(Port::Right));
    assert_eq!(str::parse::<Port>("bad"), Err(ParsePortError));
}

#[test]
fn test_parse_psedo_port() {
    assert_eq!(str::parse::<PseudoPort>("ANY"), Ok(PseudoPort::Any));
    assert_eq!(str::parse::<PseudoPort>("any"), Ok(PseudoPort::Any));
    assert_eq!(str::parse::<PseudoPort>("LAST"), Ok(PseudoPort::Last));
    assert_eq!(str::parse::<PseudoPort>("bad"), Err(ParsePseudoPortError));
}

#[test]
fn test_parse_io_register() {
    assert_eq!(str::parse::<IoRegister>("UP"), Ok(IoRegister::PortReg(Port::Up)));
    assert_eq!(str::parse::<IoRegister>("ANY"), Ok(IoRegister::PseudoReg(PseudoPort::Any)));
    assert_eq!(str::parse::<IoRegister>("bad"), Err(ParsePortError));
}

#[test]
fn test_parse_register() {
    assert_eq!(str::parse::<Register>("ACC"), Ok(Register::Acc));
    assert_eq!(str::parse::<Register>("acc"), Ok(Register::Acc));
    assert_eq!(str::parse::<Register>("NIl"), Ok(Register::Nil));
    assert_eq!(str::parse::<Register>("UP"), Ok(Register::Io(IoRegister::Port(Port::Up))));
    assert_eq!(str::parse::<Register>("bad"), Err(ParseRegisterError));
}

#[test]
fn test_parse_source() {
    assert_eq!(str::parse::<Source>("ACC"), Ok(Source::Reg(Register::Acc)));
    assert_eq!(str::parse::<Source>("1"), Ok(Source::Val(1)));
    assert_eq!(str::parse::<Source>("bad"), Err(ParseSourceError));
}

#[test]
fn test_parse_opcode() {
    assert_eq!(str::parse::<Opcode>("NOP"), Ok(Opcode::Nop));
    assert_eq!(str::parse::<Opcode>("nop"), Ok(Opcode::Nop));
    assert_eq!(str::parse::<Opcode>("MOV"), Ok(Opcode::Mov));
    assert_eq!(str::parse::<Opcode>("SWP"), Ok(Opcode::Swp));
    assert_eq!(str::parse::<Opcode>("SAV"), Ok(Opcode::Sav));
    assert_eq!(str::parse::<Opcode>("ADD"), Ok(Opcode::Add));
    assert_eq!(str::parse::<Opcode>("SUB"), Ok(Opcode::Sub));
    assert_eq!(str::parse::<Opcode>("NEG"), Ok(Opcode::Neg));
    assert_eq!(str::parse::<Opcode>("JMP"), Ok(Opcode::Jmp));
    assert_eq!(str::parse::<Opcode>("JEZ"), Ok(Opcode::Jez));
    assert_eq!(str::parse::<Opcode>("JNZ"), Ok(Opcode::Jnz));
    assert_eq!(str::parse::<Opcode>("JGZ"), Ok(Opcode::Jgz));
    assert_eq!(str::parse::<Opcode>("JLZ"), Ok(Opcode::Jlz));
    assert_eq!(str::parse::<Opcode>("JRO"), Ok(Opcode::Jro));
    assert_eq!(str::parse::<Opcode>("bad"), Err(ParseOpcodeError));
}
<|endoftext|>


<commit_before>#![feature(core)]

#[macro_use]
extern crate itertools;

use std::iter::AdditiveIterator;
use itertools::Itertools;

fn main() {
    println!("24 days of Rust - itertools (day 7)");
    let mut words = "hello supercalifragilisticexpialidocious programmer".words();
    words.apply(|word| println!("{} is {} characters long", word, word.len()));
    let even = (1us..10).map(|x| x * 2);
    let odd = (1us..5).map(|x| x * 2 + 1);
    println!("{:?}", even.interleave(odd).collect::<Vec<_>>());
    println!("{:?}", (1..10).intersperse(15).collect::<Vec<_>>());
    let numbers = 1..4;
    let chars = vec!['a', 'b', 'c'];
    for (i, c) in iproduct!(numbers, chars.iter()) {
        println!("{}: {}", i, c);
    }
    /*
    let log = "GET / 4096\nGET /home/ 16301\nPOST /home/ 49\nGET / 4096\n";
    let lines = log.lines();
    let rows = icompr!(line.words().collect::<Vec<_>>() for line in lines);
    let bytes = icompr!(row[2] for row in rows if row[0] != "POST");
    let total = icompr!(b.parse::<uint>().unwrap() for b in bytes).sum();
    println!("Total GET throughput: {} bytes", total);
    */
}
<commit_msg>Added str_words feature to day 7.<commit_after>#![feature(core, str_words)]

#[macro_use]
extern crate itertools;

use std::iter::AdditiveIterator;
use itertools::Itertools;

fn main() {
    println!("24 days of Rust - itertools (day 7)");
    let mut words = "hello supercalifragilisticexpialidocious programmer".words();
    words.apply(|word| println!("{} is {} characters long", word, word.len()));
    let even = (1us..10).map(|x| x * 2);
    let odd = (1us..5).map(|x| x * 2 + 1);
    println!("{:?}", even.interleave(odd).collect::<Vec<_>>());
    println!("{:?}", (1..10).intersperse(15).collect::<Vec<_>>());
    let numbers = 1..4;
    let chars = vec!['a', 'b', 'c'];
    for (i, c) in iproduct!(numbers, chars.iter()) {
        println!("{}: {}", i, c);
    }
    /*
    let log = "GET / 4096\nGET /home/ 16301\nPOST /home/ 49\nGET / 4096\n";
    let lines = log.lines();
    let rows = icompr!(line.words().collect::<Vec<_>>() for line in lines);
    let bytes = icompr!(row[2] for row in rows if row[0] != "POST");
    let total = icompr!(b.parse::<uint>().unwrap() for b in bytes).sum();
    println!("Total GET throughput: {} bytes", total);
    */
}
<|endoftext|>


<commit_before>use super::ConcurrentDevice;

extern crate memmap;
use z80e_core_rust::{ IoDevice };
use self::memmap::{ Mmap };
pub use self::memmap::{ MmapView, Protection };

use std::sync::{ Arc, Condvar, Mutex };
use std::sync::atomic::{ AtomicBool, AtomicUsize, Ordering };
use std::io::{ self, ErrorKind, Write };
use std::path::Path;
use std::{ str, mem, ptr };

// Sector size must be a power of two
const SECTOR_SIZE: u16 = 128;

// Max disk must be a power of two <= 16. If changed, you must insert/delete Nones in DiskController::run()
const MAX_DISK: u8 = 16;

// Status port bitflags.
//     8-bit values, but "must" be usize to avoid overly-verbose casts for AtomicUsize calls.
const COMMAND_READY: usize = 1 << 0;
const DATA_READY: usize = 1 << 1;
const ERROR: usize = 1 << 7;

// Commands.
const NOP: u8 = 0;
const SEL_DSK: u8 = 1;
const SEL_TRK: u8 = 2;
const SEL_SEC: u8 = 3;
const READ: u8 = 4;
const WRITE: u8 = 5;
const RESET: u8 = 6;
const OPEN: u8 = 7;
const CLOSE: u8 = 8;
const DPB: u8 = 9;

#[derive(Clone)]
pub struct DiskController {
    pub status: Arc<AtomicUsize>,
    command_cond: Arc<Condvar>,
    buffer: Arc<Mutex<Buffer>>,
    parameters: Arc<Mutex<Parameters>>,
}

impl DiskController {
    fn new() -> DiskController {
        DiskController {
            status: Arc::new(AtomicUsize::new(0)),
            command_cond: Arc::new(Condvar::new()),
            buffer: Arc::new(Mutex::new(Buffer::new())),
            parameters: Arc::new(Mutex::new(Parameters::new())),
        }
    }
}

struct Buffer {
    bytes: [u8; SECTOR_SIZE as usize],
    i: u16,
}

impl Buffer {
    fn new() -> Buffer {
        Buffer {
            bytes: [0; SECTOR_SIZE as usize],
            i: 0,
        }
    }
}

struct Parameters {
    disk: u8,
    track: u16,
    sector: u16,
    command: u8,
    do_command: bool,
}

impl Parameters {
    fn new() -> Parameters {
        Parameters {
            disk: 0,
            track: 0,
            sector: 0,
            command: NOP,
            do_command: false,
        }
    }
}

struct StatusPort {
    controller: DiskController,
}

impl StatusPort {
    fn new(controller: DiskController) -> StatusPort {
        StatusPort {
            controller: controller,
        }
    }
}

impl IoDevice for StatusPort {
    fn read_in(&self) -> u8 {
        (self.controller.status.load(Ordering::SeqCst) & 0xff) as u8
    }
    fn write_out(&mut self, value: u8) {
        if (self.controller.status.fetch_and(!COMMAND_READY, Ordering::SeqCst) & COMMAND_READY) != 0 {
            let mut params = self.controller.parameters.lock().unwrap();
            params.command = value;
            params.do_command = true;
            self.controller.command_cond.notify_one();
        } else {
            self.controller.status.fetch_or(COMMAND_READY | ERROR, Ordering::SeqCst);
            let _ = writeln!(io::stderr(), "disk: Attempted to write command register when not ready.");
        }
    }
}

struct DataPort {
    controller: DiskController,
}

impl DataPort {
    fn new(controller: DiskController) -> DataPort {
        DataPort {
            controller: controller,
        }
    }
}

impl IoDevice for DataPort {
    fn read_in (&self) -> u8 {
        if (self.controller.status.fetch_and(!DATA_READY, Ordering::SeqCst) & DATA_READY) != 0 {
            let byte;
            {
                let mut buffer = self.controller.buffer.lock().unwrap();
                buffer.i = ((buffer.i as usize + 1) & (SECTOR_SIZE - 1) as usize) as u16;
                byte = buffer.bytes[buffer.i as usize];
            }
            self.controller.status.fetch_or(DATA_READY, Ordering::SeqCst);
            byte
        } else {
            self.controller.status.fetch_or(ERROR, Ordering::SeqCst);
            let _ = writeln!(io::stderr(), "disk: Attempted to read data register when not ready.");
            0
        }
    }
    fn write_out(&mut self, value: u8) {
        if (self.controller.status.fetch_and(!DATA_READY, Ordering::SeqCst) & DATA_READY) != 0 {
            {
                let mut buffer = self.controller.buffer.lock().unwrap();
                buffer.i = ((buffer.i as usize + 1) & (SECTOR_SIZE - 1) as usize) as u16;
                buffer.bytes[buffer.i as usize] = value;
            }
            self.controller.status.fetch_or(DATA_READY, Ordering::SeqCst);
        } else {
            self.controller.status.fetch_or(ERROR, Ordering::SeqCst);
            let _ = writeln!(io::stderr(), "disk: Attempted to write data register when not ready.");
        }
    }
}

pub struct Disk {
    view: MmapView,
    pub tracks: u16,
    pub spt: u16,
    dpb: [u8; 16],
}

impl Disk {
    pub fn open<T: AsRef<Path>>(path: &T, protection: Protection) -> io::Result<Disk> {
        let mut file = try!(Mmap::open_path(path, protection)).into_view();
        let (header, image) = try!(file.split_at(128));
        let header = unsafe { header.as_slice() };
        if match str::from_utf8(&header[0..10]) {
            Ok(x) => x,
            Err(err) => return Err(io::Error::new(ErrorKind::InvalidData, "Invalid image header encoding.")),
        }!= "<CPM_Disk>" {
            return Err(io::Error::new(ErrorKind::InvalidData, "Not a valid disk image."));
        }
        let mut dpb: [u8; 16] = [0; 16];
        for i in 0..17 {
            dpb[i] = header[32 + i];
        }
        let spt: u16 = (dpb[0] as u16) | ((dpb[1] as u16) << 8);
        let bsh: u16 = (dpb[2] as u16);
        let dsm: u16 = (dpb[5] as u16) | ((dpb[6] as u16) << 8);
        let off: u16 = (dpb[13] as u16) | ((dpb[14] as u16) << 8);
        let tracks: u16 = (dsm + 1) * (1 << bsh) / spt + off;
        Ok(Disk {
            view: image,
            tracks: tracks,
            spt: spt,
            dpb: dpb
        })
    }
}

impl ConcurrentDevice for DiskController {
    fn run(&mut self, die: Arc<AtomicBool>) {
        let mut disks = unsafe {
            let mut disks: [Option<Disk>; MAX_DISK as usize] = mem::uninitialized();
            for disk in disks.iter_mut() {
                ptr::write(disk, None);
            };
            disks
        };
        let mut parameters = self.parameters.lock().unwrap();
        loop {
            while !(*parameters).do_command {
                parameters = self.command_cond.wait(parameters).unwrap();
            }
            if die.load(Ordering::Acquire) { break; }
            let status = self.status.fetch_and(!DATA_READY, Ordering::SeqCst);
            if (status & ERROR) != 0 && parameters.command != RESET {
                self.status.fetch_or(DATA_READY, Ordering::SeqCst);
                continue;
            };
            let mut buffer = self.buffer.lock().unwrap();
            match parameters.command {
                NOP => (),
                SEL_DSK => {
                    if buffer.bytes[0] < MAX_DISK {
                        parameters.disk = buffer.bytes[0];
                    } else {
                        self.status.fetch_or(ERROR, Ordering::SeqCst);
                    }
                },
                SEL_TRK => {
                    let track = buffer.bytes[0] as u16 & ((buffer.bytes[1] as u16) << 8);
                    match disks[parameters.disk as usize] {
                        Some(ref disk) => {
                            if track < disk.tracks {
                                parameters.track = track;
                            } else {
                                self.status.fetch_or(ERROR, Ordering::SeqCst);
                            }
                        },
                        None => {
                            self.status.fetch_or(ERROR, Ordering::SeqCst);
                        }
                    }
                },
                SEL_SEC => {
                    match disks[parameters.disk as usize] {
                        Some(ref disk) => {
                            let sector = buffer.bytes[0] as u16 & ((buffer.bytes[1] as u16) << 8);
                            if sector < disk.spt {
                                parameters.sector = sector;
                            } else {
                                self.status.fetch_or(ERROR, Ordering::SeqCst);
                            }
                        },
                        None => {
                            self.status.fetch_or(ERROR, Ordering::SeqCst);
                        },
                    }
                },
                READ => {
                    match disks[parameters.disk as usize] {
                        Some(ref disk) => {
                            let off = ((parameters.track as usize * disk.spt as usize)
                                       + parameters.sector as usize) * SECTOR_SIZE as usize;
                            let bytes = unsafe { disk.view.as_slice() };
                            for (i, byte) in buffer.bytes.iter_mut().enumerate() {
                                *byte = bytes[off + i];
                            }
                        },
                        None => {
                            self.status.fetch_or(ERROR, Ordering::SeqCst);
                        },
                    }
                },
                WRITE => {
                    match disks[parameters.disk as usize] {
                        Some(ref mut disk) => {
                            let off = ((parameters.track as usize * disk.spt as usize)
                                       + parameters.sector as usize) * SECTOR_SIZE as usize;
                            let bytes = unsafe { disk.view.as_mut_slice() };
                            for (i, byte) in buffer.bytes.iter().enumerate() {
                                bytes[i] = *byte;
                            };
                        },
                        None => {
                            self.status.fetch_or(ERROR, Ordering::SeqCst);
                        },
                    }

                }
                RESET => {
                    parameters.disk = 0;
                    parameters.track = 0;
                    parameters.sector = 0;
                    parameters.command = NOP;
                    self.status.fetch_and(!ERROR, Ordering::SeqCst);
                },
                OPEN => {
                    match str::from_utf8(&buffer.bytes[0..SECTOR_SIZE as usize]) {
                        Ok(file_name) => {
                            match Disk::open(&file_name, Protection::ReadWrite) {
                                Ok(disk) => {
                                    disks[parameters.disk as usize] = Some(disk);
                                },
                                Err(err) => {
                                    let _ = write!(io::stderr(), "disk: Failed to open file.\nError:\n\t{}\n", err);
                                    self.status.fetch_or(ERROR, Ordering::SeqCst);
                                },
                            }
                        },
                        Err(err) => {
                            let _ = write!(io::stderr(), "disk: Bad UTF-8 in file name.\nError:\n\t{}\n", err);
                            self.status.fetch_or(ERROR, Ordering::SeqCst);
                        }
                    }
                },
                CLOSE => {
                    disks[parameters.disk as usize] = None;
                },
                DPB => {
                    match disks[parameters.disk as usize] {
                        Some(ref disk) => {
                            for (a, b) in disk.dpb.iter().zip(buffer.bytes.iter_mut()) {
                                *b = *a;
                            }
                        },
                        None => {
                            self.status.fetch_or(ERROR, Ordering::SeqCst);
                        }
                    }
                },
                _ => (),
            }
            self.status.fetch_or(COMMAND_READY, Ordering::SeqCst);
        }
    }
}
<commit_msg>Finish implementing ConcurrentDevice<commit_after>use super::ConcurrentDevice;

extern crate memmap;
use z80e_core_rust::{ IoDevice };
use self::memmap::{ Mmap };
pub use self::memmap::{ MmapView, Protection };

use std::sync::{ Arc, Condvar, Mutex };
use std::sync::atomic::{ AtomicBool, AtomicUsize, Ordering };
use std::io::{ self, ErrorKind, Write };
use std::path::Path;
use std::{ str, mem, ptr };

// Sector size must be a power of two
const SECTOR_SIZE: u16 = 128;

// Max disk must be a power of two <= 16. If changed, you must insert/delete Nones in DiskController::run()
const MAX_DISK: u8 = 16;

// Status port bitflags.
//     8-bit values, but "must" be usize to avoid overly-verbose casts for AtomicUsize calls.
const COMMAND_READY: usize = 1 << 0;
const DATA_READY: usize = 1 << 1;
const ERROR: usize = 1 << 7;

// Commands.
const NOP: u8 = 0;
const SEL_DSK: u8 = 1;
const SEL_TRK: u8 = 2;
const SEL_SEC: u8 = 3;
const READ: u8 = 4;
const WRITE: u8 = 5;
const RESET: u8 = 6;
const OPEN: u8 = 7;
const CLOSE: u8 = 8;
const DPB: u8 = 9;

#[derive(Clone)]
pub struct DiskController {
    pub status: Arc<AtomicUsize>,
    command_cond: Arc<Condvar>,
    buffer: Arc<Mutex<Buffer>>,
    parameters: Arc<Mutex<Parameters>>,
}

impl DiskController {
    fn new() -> DiskController {
        DiskController {
            status: Arc::new(AtomicUsize::new(DATA_READY)),
            command_cond: Arc::new(Condvar::new()),
            buffer: Arc::new(Mutex::new(Buffer::new())),
            parameters: Arc::new(Mutex::new(Parameters::new())),
        }
    }
}

struct Buffer {
    bytes: [u8; SECTOR_SIZE as usize],
    i: u16,
}

impl Buffer {
    fn new() -> Buffer {
        Buffer {
            bytes: [0; SECTOR_SIZE as usize],
            i: 0,
        }
    }
}

struct Parameters {
    disk: u8,
    track: u16,
    sector: u16,
    command: u8,
    do_command: bool,
}

impl Parameters {
    fn new() -> Parameters {
        Parameters {
            disk: 0,
            track: 0,
            sector: 0,
            command: NOP,
            do_command: false,
        }
    }
}

struct StatusPort {
    controller: DiskController,
}

impl StatusPort {
    fn new(controller: DiskController) -> StatusPort {
        StatusPort {
            controller: controller,
        }
    }
}

impl IoDevice for StatusPort {
    fn read_in(&self) -> u8 {
        (self.controller.status.load(Ordering::SeqCst) & 0xff) as u8
    }
    fn write_out(&mut self, value: u8) {
        if (self.controller.status.fetch_and(!COMMAND_READY, Ordering::SeqCst) & COMMAND_READY) != 0 {
            let mut params = self.controller.parameters.lock().unwrap();
            params.command = value;
            params.do_command = true;
            self.controller.command_cond.notify_one();
        } else {
            self.controller.status.fetch_or(COMMAND_READY | ERROR, Ordering::SeqCst);
            let _ = writeln!(io::stderr(), "disk: Attempted to write command register when not ready.");
        }
    }
}

struct DataPort {
    controller: DiskController,
}

impl DataPort {
    fn new(controller: DiskController) -> DataPort {
        DataPort {
            controller: controller,
        }
    }
}

impl IoDevice for DataPort {
    fn read_in (&self) -> u8 {
        if (self.controller.status.fetch_and(!DATA_READY, Ordering::SeqCst) & DATA_READY) != 0 {
            let byte;
            {
                let mut buffer = self.controller.buffer.lock().unwrap();
                buffer.i = ((buffer.i as usize + 1) & (SECTOR_SIZE - 1) as usize) as u16;
                byte = buffer.bytes[buffer.i as usize];
            }
            self.controller.status.fetch_or(DATA_READY, Ordering::SeqCst);
            byte
        } else {
            self.controller.status.fetch_or(ERROR, Ordering::SeqCst);
            let _ = writeln!(io::stderr(), "disk: Attempted to read data register when not ready.");
            0
        }
    }
    fn write_out(&mut self, value: u8) {
        if (self.controller.status.fetch_and(!DATA_READY, Ordering::SeqCst) & DATA_READY) != 0 {
            {
                let mut buffer = self.controller.buffer.lock().unwrap();
                buffer.i = ((buffer.i as usize + 1) & (SECTOR_SIZE - 1) as usize) as u16;
                buffer.bytes[buffer.i as usize] = value;
            }
            self.controller.status.fetch_or(DATA_READY, Ordering::SeqCst);
        } else {
            self.controller.status.fetch_or(ERROR, Ordering::SeqCst);
            let _ = writeln!(io::stderr(), "disk: Attempted to write data register when not ready.");
        }
    }
}

pub struct Disk {
    view: MmapView,
    pub tracks: u16,
    pub spt: u16,
    dpb: [u8; 16],
}

impl Disk {
    pub fn open<T: AsRef<Path>>(path: &T, protection: Protection) -> io::Result<Disk> {
        let mut file = try!(Mmap::open_path(path, protection)).into_view();
        let (header, image) = try!(file.split_at(128));
        let header = unsafe { header.as_slice() };
        if match str::from_utf8(&header[0..10]) {
            Ok(x) => x,
            Err(err) => return Err(io::Error::new(ErrorKind::InvalidData, "Invalid image header encoding.")),
        }!= "<CPM_Disk>" {
            return Err(io::Error::new(ErrorKind::InvalidData, "Not a valid disk image."));
        }
        let mut dpb: [u8; 16] = [0; 16];
        for i in 0..17 {
            dpb[i] = header[32 + i];
        }
        let spt: u16 = (dpb[0] as u16) | ((dpb[1] as u16) << 8);
        let bsh: u16 = (dpb[2] as u16);
        let dsm: u16 = (dpb[5] as u16) | ((dpb[6] as u16) << 8);
        let off: u16 = (dpb[13] as u16) | ((dpb[14] as u16) << 8);
        let tracks: u16 = (dsm + 1) * (1 << bsh) / spt + off;
        Ok(Disk {
            view: image,
            tracks: tracks,
            spt: spt,
            dpb: dpb
        })
    }
}

impl ConcurrentDevice for DiskController {
    fn run(&mut self, die: Arc<AtomicBool>) {
        let mut disks = unsafe {
            let mut disks: [Option<Disk>; MAX_DISK as usize] = mem::uninitialized();
            for disk in disks.iter_mut() {
                ptr::write(disk, None);
            };
            disks
        };
        let mut parameters = self.parameters.lock().unwrap();
        loop {
            self.status.fetch_or(COMMAND_READY, Ordering::SeqCst);
            while !(*parameters).do_command {
                parameters = self.command_cond.wait(parameters).unwrap();
            }
            if die.load(Ordering::Acquire) { break; }
            let status = self.status.fetch_and(!DATA_READY, Ordering::SeqCst);
            if (status & ERROR) != 0 && parameters.command != RESET {
                self.status.fetch_or(DATA_READY, Ordering::SeqCst);
                continue;
            };
            {
                let mut buffer = self.buffer.lock().unwrap();
                match parameters.command {
                    NOP => (),
                    SEL_DSK => {
                        if buffer.bytes[0] < MAX_DISK {
                            parameters.disk = buffer.bytes[0];
                        } else {
                            self.status.fetch_or(ERROR, Ordering::SeqCst);
                        }
                    },
                    SEL_TRK => {
                        let track = buffer.bytes[0] as u16 & ((buffer.bytes[1] as u16) << 8);
                        match disks[parameters.disk as usize] {
                            Some(ref disk) => {
                                if track < disk.tracks {
                                    parameters.track = track;
                                } else {
                                    self.status.fetch_or(ERROR, Ordering::SeqCst);
                                }
                            },
                            None => {
                                self.status.fetch_or(ERROR, Ordering::SeqCst);
                            }
                        }
                    },
                    SEL_SEC => {
                        match disks[parameters.disk as usize] {
                            Some(ref disk) => {
                                let sector = buffer.bytes[0] as u16 & ((buffer.bytes[1] as u16) << 8);
                                if sector < disk.spt {
                                    parameters.sector = sector;
                                } else {
                                    self.status.fetch_or(ERROR, Ordering::SeqCst);
                                }
                            },
                            None => {
                                self.status.fetch_or(ERROR, Ordering::SeqCst);
                            },
                        }
                    },
                    READ => {
                        match disks[parameters.disk as usize] {
                            Some(ref disk) => {
                                let off = ((parameters.track as usize * disk.spt as usize)
                                           + parameters.sector as usize) * SECTOR_SIZE as usize;
                                let bytes = unsafe { disk.view.as_slice() };
                                for (i, byte) in buffer.bytes.iter_mut().enumerate() {
                                    *byte = bytes[off + i];
                                }
                            },
                            None => {
                                self.status.fetch_or(ERROR, Ordering::SeqCst);
                            },
                        }
                    },
                    WRITE => {
                        match disks[parameters.disk as usize] {
                            Some(ref mut disk) => {
                                let off = ((parameters.track as usize * disk.spt as usize)
                                           + parameters.sector as usize) * SECTOR_SIZE as usize;
                                let bytes = unsafe { disk.view.as_mut_slice() };
                                for (i, byte) in buffer.bytes.iter().enumerate() {
                                    bytes[i] = *byte;
                                };
                            },
                            None => {
                                self.status.fetch_or(ERROR, Ordering::SeqCst);
                            },
                        }

                    }
                    RESET => {
                        parameters.disk = 0;
                        parameters.track = 0;
                        parameters.sector = 0;
                        parameters.command = NOP;
                        self.status.fetch_and(!ERROR, Ordering::SeqCst);
                    },
                    OPEN => {
                        match str::from_utf8(&buffer.bytes[0..SECTOR_SIZE as usize]) {
                            Ok(file_name) => {
                                match Disk::open(&file_name, Protection::ReadWrite) {
                                    Ok(disk) => {
                                        disks[parameters.disk as usize] = Some(disk);
                                    },
                                    Err(err) => {
                                        let _ = write!(io::stderr(), "disk: Failed to open file.\nError:\n\t{}\n", err);
                                        self.status.fetch_or(ERROR, Ordering::SeqCst);
                                    },
                                }
                            },
                            Err(err) => {
                                let _ = write!(io::stderr(), "disk: Bad UTF-8 in file name.\nError:\n\t{}\n", err);
                                self.status.fetch_or(ERROR, Ordering::SeqCst);
                            }
                        }
                    },
                    CLOSE => {
                        disks[parameters.disk as usize] = None;
                    },
                    DPB => {
                        match disks[parameters.disk as usize] {
                            Some(ref disk) => {
                                for (a, b) in disk.dpb.iter().zip(buffer.bytes.iter_mut()) {
                                    *b = *a;
                                }
                            },
                            None => {
                                self.status.fetch_or(ERROR, Ordering::SeqCst);
                            }
                        }
                    },
                    _ => {
                        self.status.fetch_or(ERROR, Ordering::SeqCst);
                        let _ = write!(io::stderr(), "disk: System sent bad command: {:02X}\n", parameters.command);
                    },
                }
            }
            self.status.fetch_or(DATA_READY, Ordering::SeqCst);
        }
    }
}
<|endoftext|>


<commit_before>extern crate ini;
use self::ini::Ini;

extern crate time;

#[cfg(target_os = "macos")]
use super::osx;

use std::env;
use std::fs;
use std::path;

const INIFILE: &'static str = ".connectr.ini";

pub struct Settings {
    pub port: u32,
    pub secret: String,
    pub client_id: String,
    pub access_token: Option<String>,
    pub refresh_token: Option<String>,
    pub expire_utc: Option<u64>,
    pub presets: Vec<(String,String)>,
}

#[cfg(target_os = "macos")]
fn bundled_ini() -> String {
    match osx::bundled_resource_path("connectr", "ini") {
        Some(path) => path,
        None => String::new(),
    }
}

#[cfg(not(target_os = "macos"))]
fn bundled_ini() -> String {
    String::new()
}

fn inifile() -> String {
    // Try to load INI file from home directory
    let path = format!("{}/{}", env::home_dir().unwrap().display(), INIFILE);
    if path::Path::new(&path).exists() {
        info!("Found config: {}", path);
        return path
    }

    // If it doesn't exist, try to copy the template from the app bundle, if
    // such a thing exists.
    let bundle_ini = bundled_ini();
    if path::Path::new(&bundle_ini).exists() {
        info!("Copied config: {}", bundle_ini);
        let _ = fs::copy(bundle_ini, path.clone());
    }
    path
}

pub fn read_settings() -> Option<Settings> {
    info!("Attempting to read config file.");
    let conf = match Ini::load_from_file(&inifile()) {
        Ok(c) => c,
        Err(_) => {
            // No connectr.ini found.  Generate a junk one in-memory, which
            // will fail shortly after with the nice error message.
            let mut c = Ini::new();
            info!("No config file found.");
            c.with_section(Some("connectr".to_owned()))
                .set("port", 5657.to_string());
            c.with_section(Some("application".to_owned()))
                .set("secret", "<PLACEHOLDER>".to_string())
                .set("client_id", "<PLACEHOLDER>".to_string());
            c
        }
    };

    let section = conf.section(Some("connectr".to_owned())).unwrap();
    let port = section.get("port").unwrap().parse().unwrap();

    let section = conf.section(Some("application".to_owned())).unwrap();
    let secret = section.get("secret").unwrap();
    let client_id = section.get("client_id").unwrap();
    if client_id.starts_with('<') || secret.starts_with('<') {
        error!("No config file found.  Exiting.");
        println!("");
        println!("ERROR: Spotify Client ID or Secret not set in connectr.ini!");
        println!("");
        println!("Create a Spotify application at https://developer.spotify.com/my-applications/ and");
        println!("add the client ID and secret to connectr.ini.");
        println!("");
        println!("Be sure to add a redirect URI of http://127.0.0.1:<PORT> to your Spotify application,");
        println!("and make sure the port matches in connectr.ini.");
        println!("");
        return None;
    }

    let mut access = None;
    let mut refresh = None;
    let mut expire_utc = None;
    if let Some(section) = conf.section(Some("tokens".to_owned())) {
        access = Some(section.get("access").unwrap().clone());
        refresh = Some(section.get("refresh").unwrap().clone());
        expire_utc = Some(section.get("expire").unwrap().parse().unwrap());
        println!("Read access token from INI!");
    }

    let mut presets = Vec::<(String,String)>::new();
    if let Some(section) = conf.section(Some("presets".to_owned())) {
        for (key, value) in section {
            presets.push((key.to_owned(), value.to_owned()));
        }
    }

    Some(Settings { secret: secret.to_string(), client_id: client_id.to_string(), port: port,
                    access_token: access, refresh_token: refresh, expire_utc: expire_utc,
                    presets: presets})
}

pub type SettingsError = String;
pub fn save_tokens(access: &str, refresh: &str, expire_utc: u64) -> Result<(), SettingsError> {
    let mut conf = Ini::load_from_file(&inifile()).unwrap();
    conf.with_section(Some("tokens".to_owned()))
        .set("access", access)
        .set("refresh", refresh)
        .set("expire", expire_utc.to_string());
    conf.write_to_file(&inifile()).unwrap();
    Ok(())
}
<commit_msg>Also use connectr.ini in cwd if it exists.<commit_after>extern crate ini;
use self::ini::Ini;

extern crate time;

#[cfg(target_os = "macos")]
use super::osx;

use std::env;
use std::fs;
use std::path;

const INIFILE: &'static str = "connectr.ini";

pub struct Settings {
    pub port: u32,
    pub secret: String,
    pub client_id: String,
    pub access_token: Option<String>,
    pub refresh_token: Option<String>,
    pub expire_utc: Option<u64>,
    pub presets: Vec<(String,String)>,
}

#[cfg(target_os = "macos")]
fn bundled_ini() -> String {
    match osx::bundled_resource_path("connectr", "ini") {
        Some(path) => path,
        None => String::new(),
    }
}

#[cfg(not(target_os = "macos"))]
fn bundled_ini() -> String {
    String::new()
}

fn inifile() -> String {
    // Try to load INI file from home directory
    let path = format!("{}/.{}", env::home_dir().unwrap().display(), INIFILE);
    if path::Path::new(&path).exists() {
        info!("Found config: {}", path);
        return path;
    }

    // If it doesn't exist, try to copy the template from the app bundle, if
    // such a thing exists.
    let bundle_ini = bundled_ini();
    if path::Path::new(&bundle_ini).exists() {
        info!("Copied config: {}", bundle_ini);
        let _ = fs::copy(bundle_ini, path.clone());
        return path;
    }

    let path = INIFILE.to_string();
    if path::Path::new(&path).exists() {
        info!("Found config: {}", path);
        return path;
    }

    // Default to looking in current working directory
    String::new()
}

pub fn read_settings() -> Option<Settings> {
    info!("Attempting to read config file.");
    let conf = match Ini::load_from_file(&inifile()) {
        Ok(c) => c,
        Err(e) => {
            info!("Load file error: {}", e);
            // No connectr.ini found.  Generate a junk one in-memory, which
            // will fail shortly after with the nice error message.
            let mut c = Ini::new();
            info!("No config file found.");
            c.with_section(Some("connectr".to_owned()))
                .set("port", 5657.to_string());
            c.with_section(Some("application".to_owned()))
                .set("secret", "<PLACEHOLDER>".to_string())
                .set("client_id", "<PLACEHOLDER>".to_string());
            c
        }
    };

    let section = conf.section(Some("connectr".to_owned())).unwrap();
    let port = section.get("port").unwrap().parse().unwrap();

    let section = conf.section(Some("application".to_owned())).unwrap();
    let secret = section.get("secret").unwrap();
    let client_id = section.get("client_id").unwrap();
    if client_id.starts_with('<') || secret.starts_with('<') {
        error!("Invalid or missing configuration.  Cannot continue.");
        println!("");
        println!("ERROR: Spotify Client ID or Secret not set in connectr.ini!");
        println!("");
        println!("Create a Spotify application at https://developer.spotify.com/my-applications/ and");
        println!("add the client ID and secret to connectr.ini.");
        println!("");
        println!("Be sure to add a redirect URI of http://127.0.0.1:<PORT> to your Spotify application,");
        println!("and make sure the port matches in connectr.ini.");
        println!("");
        return None;
    }

    let mut access = None;
    let mut refresh = None;
    let mut expire_utc = None;
    if let Some(section) = conf.section(Some("tokens".to_owned())) {
        access = Some(section.get("access").unwrap().clone());
        refresh = Some(section.get("refresh").unwrap().clone());
        expire_utc = Some(section.get("expire").unwrap().parse().unwrap());
        println!("Read access token from INI!");
    }

    let mut presets = Vec::<(String,String)>::new();
    if let Some(section) = conf.section(Some("presets".to_owned())) {
        for (key, value) in section {
            presets.push((key.to_owned(), value.to_owned()));
        }
    }

    Some(Settings { secret: secret.to_string(), client_id: client_id.to_string(), port: port,
                    access_token: access, refresh_token: refresh, expire_utc: expire_utc,
                    presets: presets})
}

pub type SettingsError = String;
pub fn save_tokens(access: &str, refresh: &str, expire_utc: u64) -> Result<(), SettingsError> {
    let mut conf = Ini::load_from_file(&inifile()).unwrap();
    conf.with_section(Some("tokens".to_owned()))
        .set("access", access)
        .set("refresh", refresh)
        .set("expire", expire_utc.to_string());
    conf.write_to_file(&inifile()).unwrap();
    Ok(())
}
<|endoftext|>