<commit_before>#[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
use crate::transport::smtp::{error, Error};
#[cfg(feature = "native-tls")]
use native_tls::{Protocol, TlsConnector};
#[cfg(feature = "rustls-tls")]
use rustls::{ClientConfig, RootCertStore, ServerCertVerified, ServerCertVerifier, TLSError};
use std::fmt::{self, Debug};
#[cfg(feature = "rustls-tls")]
use std::sync::Arc;
#[cfg(feature = "rustls-tls")]
use webpki::DNSNameRef;

/// Accepted protocols by default.
/// This removes TLS 1.0 and 1.1 compared to tls-native defaults.
// This is also rustls' default behavior
#[cfg(feature = "native-tls")]
const DEFAULT_TLS_MIN_PROTOCOL: Protocol = Protocol::Tlsv12;

/// How to apply TLS to a client connection
#[derive(Clone)]
#[allow(missing_copy_implementations)]
pub enum Tls {
    /// Insecure connection only (for testing purposes)
    None,
    /// Start with insecure connection and use `STARTTLS` when available
    #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
    Opportunistic(TlsParameters),
    /// Start with insecure connection and require `STARTTLS`
    #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
    Required(TlsParameters),
    /// Use TLS wrapped connection
    #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
    Wrapper(TlsParameters),
}

impl Debug for Tls {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &self {
            Self::None => f.pad("None"),
            #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
            Self::Opportunistic(_) => f.pad("Opportunistic"),
            #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
            Self::Required(_) => f.pad("Required"),
            #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
            Self::Wrapper(_) => f.pad("Wrapper"),
        }
    }
}

/// Parameters to use for secure clients
#[derive(Clone)]
pub struct TlsParameters {
    pub(crate) connector: InnerTlsParameters,
    /// The domain name which is expected in the TLS certificate from the server
    pub(super) domain: String,
}

/// Builder for `TlsParameters`
#[derive(Debug, Clone)]
pub struct TlsParametersBuilder {
    domain: String,
    root_certs: Vec<Certificate>,
    accept_invalid_hostnames: bool,
    accept_invalid_certs: bool,
}

impl TlsParametersBuilder {
    /// Creates a new builder for `TlsParameters`
    pub fn new(domain: String) -> Self {
        Self {
            domain,
            root_certs: Vec::new(),
            accept_invalid_hostnames: false,
            accept_invalid_certs: false,
        }
    }

    /// Add a custom root certificate
    ///
    /// Can be used to safely connect to a server using a self signed certificate, for example.
    pub fn add_root_certificate(mut self, cert: Certificate) -> Self {
        self.root_certs.push(cert);
        self
    }

    /// Controls whether certificates with an invalid hostname are accepted
    ///
    /// Defaults to `false`.
    ///
    /// # Warning
    ///
    /// You should think very carefully before using this method.
    /// If hostname verification is disabled *any* valid certificate,
    /// including those from other sites, are trusted.
    ///
    /// This method introduces significant vulnerabilities to man-in-the-middle attacks.
    ///
    /// Hostname verification can only be disabled with the `native-tls` TLS backend.
    #[cfg(feature = "native-tls")]
    #[cfg_attr(docsrs, doc(cfg(feature = "native-tls")))]
    pub fn dangerous_accept_invalid_hostnames(mut self, accept_invalid_hostnames: bool) -> Self {
        self.accept_invalid_hostnames = accept_invalid_hostnames;
        self
    }

    /// Controls whether invalid certificates are accepted
    ///
    /// Defaults to `false`.
    ///
    /// # Warning
    ///
    /// You should think very carefully before using this method.
    /// If certificate verification is disabled, *any* certificate
    /// is trusted for use, including:
    ///
    /// * Self signed certificates
    /// * Certificates from different hostnames
    /// * Expired certificates
    ///
    /// This method should only be used as a last resort, as it introduces
    /// significant vulnerabilities to man-in-the-middle attacks.
    pub fn dangerous_accept_invalid_certs(mut self, accept_invalid_certs: bool) -> Self {
        self.accept_invalid_certs = accept_invalid_certs;
        self
    }

    /// Creates a new `TlsParameters` using native-tls or rustls
    /// depending on which one is available
    #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
    #[cfg_attr(docsrs, doc(cfg(any(feature = "native-tls", feature = "rustls-tls"))))]
    // TODO: remove below line once native-tls is supported with async-std
    #[allow(unreachable_code)]
    pub fn build(self) -> Result<TlsParameters, Error> {
        // TODO: remove once native-tls is supported with async-std
        #[cfg(all(feature = "rustls-tls", feature = "async-std1"))]
        return self.build_rustls();

        #[cfg(feature = "native-tls")]
        return self.build_native();

        #[cfg(not(feature = "native-tls"))]
        return self.build_rustls();
    }

    /// Creates a new `TlsParameters` using native-tls with the provided configuration
    #[cfg(feature = "native-tls")]
    #[cfg_attr(docsrs, doc(cfg(feature = "native-tls")))]
    pub fn build_native(self) -> Result<TlsParameters, Error> {
        let mut tls_builder = TlsConnector::builder();

        for cert in self.root_certs {
            tls_builder.add_root_certificate(cert.native_tls);
        }
        tls_builder.danger_accept_invalid_hostnames(self.accept_invalid_hostnames);
        tls_builder.danger_accept_invalid_certs(self.accept_invalid_certs);

        tls_builder.min_protocol_version(Some(DEFAULT_TLS_MIN_PROTOCOL));
        let connector = tls_builder.build().map_err(error::tls)?;
        Ok(TlsParameters {
            connector: InnerTlsParameters::NativeTls(connector),
            domain: self.domain,
        })
    }

    /// Creates a new `TlsParameters` using rustls with the provided configuration
    #[cfg(feature = "rustls-tls")]
    #[cfg_attr(docsrs, doc(cfg(feature = "rustls-tls")))]
    pub fn build_rustls(self) -> Result<TlsParameters, Error> {
        use webpki_roots::TLS_SERVER_ROOTS;

        let mut tls = ClientConfig::new();

        for cert in self.root_certs {
            for rustls_cert in cert.rustls {
                tls.root_store.add(&rustls_cert).map_err(error::tls)?;
            }
        }
        if self.accept_invalid_certs {
            tls.dangerous()
                .set_certificate_verifier(Arc::new(InvalidCertsVerifier {}));
        }

        tls.root_store.add_server_trust_anchors(&TLS_SERVER_ROOTS);
        Ok(TlsParameters {
            connector: InnerTlsParameters::RustlsTls(tls),
            domain: self.domain,
        })
    }
}

#[derive(Clone)]
pub enum InnerTlsParameters {
    #[cfg(feature = "native-tls")]
    NativeTls(TlsConnector),
    #[cfg(feature = "rustls-tls")]
    RustlsTls(ClientConfig),
}

impl TlsParameters {
    /// Creates a new `TlsParameters` using native-tls or rustls
    /// depending on which one is available
    #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
    #[cfg_attr(docsrs, doc(cfg(any(feature = "native-tls", feature = "rustls-tls"))))]
    pub fn new(domain: String) -> Result<Self, Error> {
        TlsParametersBuilder::new(domain).build()
    }

    pub fn builder(domain: String) -> TlsParametersBuilder {
        TlsParametersBuilder::new(domain)
    }

    /// Creates a new `TlsParameters` using native-tls
    #[cfg(feature = "native-tls")]
    #[cfg_attr(docsrs, doc(cfg(feature = "native-tls")))]
    pub fn new_native(domain: String) -> Result<Self, Error> {
        TlsParametersBuilder::new(domain).build_native()
    }

    /// Creates a new `TlsParameters` using rustls
    #[cfg(feature = "rustls-tls")]
    #[cfg_attr(docsrs, doc(cfg(feature = "rustls-tls")))]
    pub fn new_rustls(domain: String) -> Result<Self, Error> {
        TlsParametersBuilder::new(domain).build_rustls()
    }

    pub fn domain(&self) -> &str {
        &self.domain
    }
}

/// A client certificate that can be used with [`TlsParametersBuilder::add_root_certificate`]
#[derive(Clone)]
#[allow(missing_copy_implementations)]
pub struct Certificate {
    #[cfg(feature = "native-tls")]
    native_tls: native_tls::Certificate,
    #[cfg(feature = "rustls-tls")]
    rustls: Vec<rustls::Certificate>,
}

#[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
impl Certificate {
    /// Create a `Certificate` from a DER encoded certificate
    pub fn from_der(der: Vec<u8>) -> Result<Self, Error> {
        #[cfg(feature = "native-tls")]
        let native_tls_cert = native_tls::Certificate::from_der(&der).map_err(error::tls)?;

        Ok(Self {
            #[cfg(feature = "native-tls")]
            native_tls: native_tls_cert,
            #[cfg(feature = "rustls-tls")]
            rustls: vec![rustls::Certificate(der)],
        })
    }

    /// Create a `Certificate` from a PEM encoded certificate
    pub fn from_pem(pem: &[u8]) -> Result<Self, Error> {
        #[cfg(feature = "native-tls")]
        let native_tls_cert = native_tls::Certificate::from_pem(pem).map_err(error::tls)?;

        #[cfg(feature = "rustls-tls")]
        let rustls_cert = {
            use rustls::internal::pemfile;
            use std::io::Cursor;

            let mut pem = Cursor::new(pem);
            pemfile::certs(&mut pem).map_err(|_| error::tls("invalid certificates"))?
        };

        Ok(Self {
            #[cfg(feature = "native-tls")]
            native_tls: native_tls_cert,
            #[cfg(feature = "rustls-tls")]
            rustls: rustls_cert,
        })
    }
}

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

#[cfg(feature = "rustls-tls")]
struct InvalidCertsVerifier;

#[cfg(feature = "rustls-tls")]
impl ServerCertVerifier for InvalidCertsVerifier {
    fn verify_server_cert(
        &self,
        _roots: &RootCertStore,
        _presented_certs: &[rustls::Certificate],
        _dns_name: DNSNameRef<'_>,
        _ocsp_response: &[u8],
    ) -> Result<ServerCertVerified, TLSError> {
        Ok(ServerCertVerified::assertion())
    }
}
<commit_msg>tls: use rustls if both native-tls and rustls-tls are enabled (#586)<commit_after>#[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
use crate::transport::smtp::{error, Error};
#[cfg(feature = "native-tls")]
use native_tls::{Protocol, TlsConnector};
#[cfg(feature = "rustls-tls")]
use rustls::{ClientConfig, RootCertStore, ServerCertVerified, ServerCertVerifier, TLSError};
use std::fmt::{self, Debug};
#[cfg(feature = "rustls-tls")]
use std::sync::Arc;
#[cfg(feature = "rustls-tls")]
use webpki::DNSNameRef;

/// Accepted protocols by default.
/// This removes TLS 1.0 and 1.1 compared to tls-native defaults.
// This is also rustls' default behavior
#[cfg(feature = "native-tls")]
const DEFAULT_TLS_MIN_PROTOCOL: Protocol = Protocol::Tlsv12;

/// How to apply TLS to a client connection
#[derive(Clone)]
#[allow(missing_copy_implementations)]
pub enum Tls {
    /// Insecure connection only (for testing purposes)
    None,
    /// Start with insecure connection and use `STARTTLS` when available
    #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
    Opportunistic(TlsParameters),
    /// Start with insecure connection and require `STARTTLS`
    #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
    Required(TlsParameters),
    /// Use TLS wrapped connection
    #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
    Wrapper(TlsParameters),
}

impl Debug for Tls {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &self {
            Self::None => f.pad("None"),
            #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
            Self::Opportunistic(_) => f.pad("Opportunistic"),
            #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
            Self::Required(_) => f.pad("Required"),
            #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
            Self::Wrapper(_) => f.pad("Wrapper"),
        }
    }
}

/// Parameters to use for secure clients
#[derive(Clone)]
pub struct TlsParameters {
    pub(crate) connector: InnerTlsParameters,
    /// The domain name which is expected in the TLS certificate from the server
    pub(super) domain: String,
}

/// Builder for `TlsParameters`
#[derive(Debug, Clone)]
pub struct TlsParametersBuilder {
    domain: String,
    root_certs: Vec<Certificate>,
    accept_invalid_hostnames: bool,
    accept_invalid_certs: bool,
}

impl TlsParametersBuilder {
    /// Creates a new builder for `TlsParameters`
    pub fn new(domain: String) -> Self {
        Self {
            domain,
            root_certs: Vec::new(),
            accept_invalid_hostnames: false,
            accept_invalid_certs: false,
        }
    }

    /// Add a custom root certificate
    ///
    /// Can be used to safely connect to a server using a self signed certificate, for example.
    pub fn add_root_certificate(mut self, cert: Certificate) -> Self {
        self.root_certs.push(cert);
        self
    }

    /// Controls whether certificates with an invalid hostname are accepted
    ///
    /// Defaults to `false`.
    ///
    /// # Warning
    ///
    /// You should think very carefully before using this method.
    /// If hostname verification is disabled *any* valid certificate,
    /// including those from other sites, are trusted.
    ///
    /// This method introduces significant vulnerabilities to man-in-the-middle attacks.
    ///
    /// Hostname verification can only be disabled with the `native-tls` TLS backend.
    #[cfg(feature = "native-tls")]
    #[cfg_attr(docsrs, doc(cfg(feature = "native-tls")))]
    pub fn dangerous_accept_invalid_hostnames(mut self, accept_invalid_hostnames: bool) -> Self {
        self.accept_invalid_hostnames = accept_invalid_hostnames;
        self
    }

    /// Controls whether invalid certificates are accepted
    ///
    /// Defaults to `false`.
    ///
    /// # Warning
    ///
    /// You should think very carefully before using this method.
    /// If certificate verification is disabled, *any* certificate
    /// is trusted for use, including:
    ///
    /// * Self signed certificates
    /// * Certificates from different hostnames
    /// * Expired certificates
    ///
    /// This method should only be used as a last resort, as it introduces
    /// significant vulnerabilities to man-in-the-middle attacks.
    pub fn dangerous_accept_invalid_certs(mut self, accept_invalid_certs: bool) -> Self {
        self.accept_invalid_certs = accept_invalid_certs;
        self
    }

    /// Creates a new `TlsParameters` using native-tls or rustls
    /// depending on which one is available
    #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
    #[cfg_attr(docsrs, doc(cfg(any(feature = "native-tls", feature = "rustls-tls"))))]
    pub fn build(self) -> Result<TlsParameters, Error> {
        #[cfg(feature = "rustls-tls")]
        return self.build_rustls();

        #[cfg(not(feature = "rustls-tls"))]
        return self.build_native();
    }

    /// Creates a new `TlsParameters` using native-tls with the provided configuration
    #[cfg(feature = "native-tls")]
    #[cfg_attr(docsrs, doc(cfg(feature = "native-tls")))]
    pub fn build_native(self) -> Result<TlsParameters, Error> {
        let mut tls_builder = TlsConnector::builder();

        for cert in self.root_certs {
            tls_builder.add_root_certificate(cert.native_tls);
        }
        tls_builder.danger_accept_invalid_hostnames(self.accept_invalid_hostnames);
        tls_builder.danger_accept_invalid_certs(self.accept_invalid_certs);

        tls_builder.min_protocol_version(Some(DEFAULT_TLS_MIN_PROTOCOL));
        let connector = tls_builder.build().map_err(error::tls)?;
        Ok(TlsParameters {
            connector: InnerTlsParameters::NativeTls(connector),
            domain: self.domain,
        })
    }

    /// Creates a new `TlsParameters` using rustls with the provided configuration
    #[cfg(feature = "rustls-tls")]
    #[cfg_attr(docsrs, doc(cfg(feature = "rustls-tls")))]
    pub fn build_rustls(self) -> Result<TlsParameters, Error> {
        use webpki_roots::TLS_SERVER_ROOTS;

        let mut tls = ClientConfig::new();

        for cert in self.root_certs {
            for rustls_cert in cert.rustls {
                tls.root_store.add(&rustls_cert).map_err(error::tls)?;
            }
        }
        if self.accept_invalid_certs {
            tls.dangerous()
                .set_certificate_verifier(Arc::new(InvalidCertsVerifier {}));
        }

        tls.root_store.add_server_trust_anchors(&TLS_SERVER_ROOTS);
        Ok(TlsParameters {
            connector: InnerTlsParameters::RustlsTls(tls),
            domain: self.domain,
        })
    }
}

#[derive(Clone)]
pub enum InnerTlsParameters {
    #[cfg(feature = "native-tls")]
    NativeTls(TlsConnector),
    #[cfg(feature = "rustls-tls")]
    RustlsTls(ClientConfig),
}

impl TlsParameters {
    /// Creates a new `TlsParameters` using native-tls or rustls
    /// depending on which one is available
    #[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
    #[cfg_attr(docsrs, doc(cfg(any(feature = "native-tls", feature = "rustls-tls"))))]
    pub fn new(domain: String) -> Result<Self, Error> {
        TlsParametersBuilder::new(domain).build()
    }

    pub fn builder(domain: String) -> TlsParametersBuilder {
        TlsParametersBuilder::new(domain)
    }

    /// Creates a new `TlsParameters` using native-tls
    #[cfg(feature = "native-tls")]
    #[cfg_attr(docsrs, doc(cfg(feature = "native-tls")))]
    pub fn new_native(domain: String) -> Result<Self, Error> {
        TlsParametersBuilder::new(domain).build_native()
    }

    /// Creates a new `TlsParameters` using rustls
    #[cfg(feature = "rustls-tls")]
    #[cfg_attr(docsrs, doc(cfg(feature = "rustls-tls")))]
    pub fn new_rustls(domain: String) -> Result<Self, Error> {
        TlsParametersBuilder::new(domain).build_rustls()
    }

    pub fn domain(&self) -> &str {
        &self.domain
    }
}

/// A client certificate that can be used with [`TlsParametersBuilder::add_root_certificate`]
#[derive(Clone)]
#[allow(missing_copy_implementations)]
pub struct Certificate {
    #[cfg(feature = "native-tls")]
    native_tls: native_tls::Certificate,
    #[cfg(feature = "rustls-tls")]
    rustls: Vec<rustls::Certificate>,
}

#[cfg(any(feature = "native-tls", feature = "rustls-tls"))]
impl Certificate {
    /// Create a `Certificate` from a DER encoded certificate
    pub fn from_der(der: Vec<u8>) -> Result<Self, Error> {
        #[cfg(feature = "native-tls")]
        let native_tls_cert = native_tls::Certificate::from_der(&der).map_err(error::tls)?;

        Ok(Self {
            #[cfg(feature = "native-tls")]
            native_tls: native_tls_cert,
            #[cfg(feature = "rustls-tls")]
            rustls: vec![rustls::Certificate(der)],
        })
    }

    /// Create a `Certificate` from a PEM encoded certificate
    pub fn from_pem(pem: &[u8]) -> Result<Self, Error> {
        #[cfg(feature = "native-tls")]
        let native_tls_cert = native_tls::Certificate::from_pem(pem).map_err(error::tls)?;

        #[cfg(feature = "rustls-tls")]
        let rustls_cert = {
            use rustls::internal::pemfile;
            use std::io::Cursor;

            let mut pem = Cursor::new(pem);
            pemfile::certs(&mut pem).map_err(|_| error::tls("invalid certificates"))?
        };

        Ok(Self {
            #[cfg(feature = "native-tls")]
            native_tls: native_tls_cert,
            #[cfg(feature = "rustls-tls")]
            rustls: rustls_cert,
        })
    }
}

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

#[cfg(feature = "rustls-tls")]
struct InvalidCertsVerifier;

#[cfg(feature = "rustls-tls")]
impl ServerCertVerifier for InvalidCertsVerifier {
    fn verify_server_cert(
        &self,
        _roots: &RootCertStore,
        _presented_certs: &[rustls::Certificate],
        _dns_name: DNSNameRef<'_>,
        _ocsp_response: &[u8],
    ) -> Result<ServerCertVerified, TLSError> {
        Ok(ServerCertVerified::assertion())
    }
}
<|endoftext|>


<commit_before>//! Bit level parsers and combinators
//!
//! Bit parsing is handled by tweaking the input in most macros.
//! In byte level parsing, the input is generally a `&[u8]` passed from combinator
//! to combinator until the slices are manipulated.
//!
//! Bit parsers take a `(&[u8], usize)` as input. The first part of the tuple is an byte slice,
//! the second part is a bit offset in the first byte of the slice.
//!
//! By passing a pair like this, we can leverage most of the combinators, and avoid
//! transforming the whole slice to a vector of booleans. This should make it easy
//! to see a byte slice as a bit stream, and parse code points of arbitrary bit length.

/// `bits!( parser ) => ( &[u8], (&[u8], usize) -> IResult<(&[u8], usize), T> ) -> IResult<&[u8], T>`
/// transforms its byte slice input in a bit stream for the underlying parsers
///
/// ```
/// # #[macro_use] extern crate nom;
/// # use nom::IResult::Done;
/// # fn main() {
///  named!( take_3_bits<u8>, bits!( take_bits!( u8, 3 ) ) );
///
///  let input = vec![0b10101010, 0b11110000, 0b00110011];
///  let sl    = &input[..];
///
///  assert_eq!(take_3_bits( sl ), Done(&sl[1..], 5) );
/// # }
#[macro_export]
macro_rules! bits (
  ($i:expr, $submac:ident!( $($args:tt)* )) => (
    bits_impl!($i, $submac!($($args)*));
  );
  ($i:expr, $f:expr) => (
    bits_impl!($i, call!($f));
  );
);

/// Internal parser, do not use directly
#[macro_export]
macro_rules! bits_impl (
  ($i:expr, $submac:ident!( $($args:tt)* )) => (
    {
      let input = ($i, 0usize);
      match $submac!(input, $($args)*) {
        $crate::IResult::Error(e)                            => {
          let err = match e {
            $crate::Err::Code(k) | $crate::Err::Node(k, _) | $crate::Err::Position(k, _) | $crate::Err::NodePosition(k, _, _) => {
              $crate::Err::Position(k, $i)
            }
          };
          $crate::IResult::Error(err)
        }
        $crate::IResult::Incomplete($crate::Needed::Unknown) => $crate::IResult::Incomplete($crate::Needed::Unknown),
        $crate::IResult::Incomplete($crate::Needed::Size(i)) => $crate::IResult::Incomplete($crate::Needed::Size(i)),
        $crate::IResult::Done((i, bit_index), o)             => {
          let byte_index = if bit_index == 0 { 0 } else { 1 };

          $crate::IResult::Done(&i[byte_index..], o)
        }
      }
    }
  );
);

/// `take_bits!(type, nb) => ( (&[T], usize), U, usize) -> IResult<(&[T], usize), U>`
/// generates a parser consuming the specified number of bytes
///
/// ```
/// # #[macro_use] extern crate nom;
/// # use nom::IResult::Done;
/// # fn main() {
///  named!( take_pair<(u8, u8)>, bits!( pair!( take_bits!( u8, 3 ), take_bits!(u8, 5) ) ) );
///
///  let input = vec![0b10101010, 0b11110000, 0b00110011];
///  let sl    = &input[..];
///
///  assert_eq!(take_pair( sl ),       Done(&sl[1..], (5, 10)) );
///  assert_eq!(take_pair( &sl[1..] ), Done(&sl[2..], (7, 16)) );
/// # }
/// ```
#[macro_export]
macro_rules! take_bits (
  ($i:expr, $t:ty, $count:expr) => (
    {
      use std::ops::Div;
      let (input, bit_offset) = $i;
      let res : $crate::IResult<(&[u8],usize), $t> = if $count == 0 {
        $crate::IResult::Done( (input, bit_offset), 0)
      } else {
        let cnt = ($count as usize + bit_offset).div(8);
        if input.len() * 8 < $count as usize + bit_offset {
          $crate::IResult::Incomplete($crate::Needed::Size(cnt+1))
        } else {
          let mut acc:$t            = 0;
          let mut offset: usize     = bit_offset;
          let mut remaining: usize  = $count;
          let mut end_offset: usize = 0;

          for it in 0..cnt+1 {
            if remaining == 0 {
              break;
            }
            let val: $t = if offset == 0 {
              input[it] as $t
            } else {
              ((input[it] << offset) as u8 >> offset) as $t
            };

            if remaining < 8 - offset {
              acc += val >> (8 - offset - remaining);
              end_offset = remaining + offset;
              break;
            } else {
              acc += val << remaining - (8 - offset);
              remaining -= 8 - offset;
              offset = 0;
            }
          }
          $crate::IResult::Done( (&input[cnt..], end_offset) , acc)
        }
      };
      res
    }
  );
);

#[macro_export]
macro_rules! tag_bits (
  ($i:expr, $t:ty, $count:expr, $p: pat) => (
    {
      match take_bits!($i, $t, $count) {
        $crate::IResult::Incomplete(i) => $crate::IResult::Incomplete(i),
        $crate::IResult::Done(i, o)    => {
          if let $p = o {
            let res: $crate::IResult<(&[u8],usize),$t> = $crate::IResult::Done(i, o);
            res
          } else {
            $crate::IResult::Error($crate::Err::Code($crate::ErrorKind::TagBits))
          }
        },
        _                              => {
          $crate::IResult::Error($crate::Err::Code($crate::ErrorKind::TagBits))
        }
      }
    }
  )
);

#[cfg(test)]
mod tests {
  use internal::{IResult,Needed};

  #[test]
  fn take_bits() {
    let input = vec![0b10101010, 0b11110000, 0b00110011];
    let sl    = &input[..];

    assert_eq!(take_bits!( (sl, 0), u8,   0 ), IResult::Done((sl, 0), 0));
    assert_eq!(take_bits!( (sl, 0), u8,   8 ), IResult::Done((&sl[1..], 0), 170));
    assert_eq!(take_bits!( (sl, 0), u8,   3 ), IResult::Done((&sl[0..], 3), 5));
    assert_eq!(take_bits!( (sl, 0), u8,   6 ), IResult::Done((&sl[0..], 6), 42));
    assert_eq!(take_bits!( (sl, 1), u8,   1 ), IResult::Done((&sl[0..], 2), 0));
    assert_eq!(take_bits!( (sl, 1), u8,   2 ), IResult::Done((&sl[0..], 3), 1));
    assert_eq!(take_bits!( (sl, 1), u8,   3 ), IResult::Done((&sl[0..], 4), 2));
    assert_eq!(take_bits!( (sl, 6), u8,   3 ), IResult::Done((&sl[1..], 1), 5));
    assert_eq!(take_bits!( (sl, 0), u16, 10 ), IResult::Done((&sl[1..], 2), 683));
    assert_eq!(take_bits!( (sl, 0), u16,  8 ), IResult::Done((&sl[1..], 0), 170));
    assert_eq!(take_bits!( (sl, 6), u16, 10 ), IResult::Done((&sl[2..], 0), 752));
    assert_eq!(take_bits!( (sl, 6), u16, 11 ), IResult::Done((&sl[2..], 1), 1504));
    assert_eq!(take_bits!( (sl, 0), u32, 20 ), IResult::Done((&sl[2..], 4), 700163));
    assert_eq!(take_bits!( (sl, 4), u32, 20 ), IResult::Done((&sl[3..], 0), 716851));
    assert_eq!(take_bits!( (sl, 4), u32, 22 ), IResult::Incomplete(Needed::Size(4)));
  }

  #[test]
  fn tag_bits() {
    let input = vec![0b10101010, 0b11110000, 0b00110011];
    let sl    = &input[..];

    assert_eq!(tag_bits!( (sl, 0), u8,   3, 0b101), IResult::Done((&sl[0..], 3), 5));
    assert_eq!(tag_bits!( (sl, 0), u8,   4, 0b1010), IResult::Done((&sl[0..], 4), 10));
  }
}
<commit_msg>Take_bits should return the needed number of bits in Incomplete<commit_after>//! Bit level parsers and combinators
//!
//! Bit parsing is handled by tweaking the input in most macros.
//! In byte level parsing, the input is generally a `&[u8]` passed from combinator
//! to combinator until the slices are manipulated.
//!
//! Bit parsers take a `(&[u8], usize)` as input. The first part of the tuple is an byte slice,
//! the second part is a bit offset in the first byte of the slice.
//!
//! By passing a pair like this, we can leverage most of the combinators, and avoid
//! transforming the whole slice to a vector of booleans. This should make it easy
//! to see a byte slice as a bit stream, and parse code points of arbitrary bit length.

/// `bits!( parser ) => ( &[u8], (&[u8], usize) -> IResult<(&[u8], usize), T> ) -> IResult<&[u8], T>`
/// transforms its byte slice input in a bit stream for the underlying parsers
///
/// ```
/// # #[macro_use] extern crate nom;
/// # use nom::IResult::Done;
/// # fn main() {
///  named!( take_3_bits<u8>, bits!( take_bits!( u8, 3 ) ) );
///
///  let input = vec![0b10101010, 0b11110000, 0b00110011];
///  let sl    = &input[..];
///
///  assert_eq!(take_3_bits( sl ), Done(&sl[1..], 5) );
/// # }
#[macro_export]
macro_rules! bits (
  ($i:expr, $submac:ident!( $($args:tt)* )) => (
    bits_impl!($i, $submac!($($args)*));
  );
  ($i:expr, $f:expr) => (
    bits_impl!($i, call!($f));
  );
);

/// Internal parser, do not use directly
#[macro_export]
macro_rules! bits_impl (
  ($i:expr, $submac:ident!( $($args:tt)* )) => (
    {
      let input = ($i, 0usize);
      match $submac!(input, $($args)*) {
        $crate::IResult::Error(e)                            => {
          let err = match e {
            $crate::Err::Code(k) | $crate::Err::Node(k, _) | $crate::Err::Position(k, _) | $crate::Err::NodePosition(k, _, _) => {
              $crate::Err::Position(k, $i)
            }
          };
          $crate::IResult::Error(err)
        }
        $crate::IResult::Incomplete($crate::Needed::Unknown) => $crate::IResult::Incomplete($crate::Needed::Unknown),
        $crate::IResult::Incomplete($crate::Needed::Size(i)) => $crate::IResult::Incomplete($crate::Needed::Size(i)),
        $crate::IResult::Done((i, bit_index), o)             => {
          let byte_index = if bit_index == 0 { 0 } else { 1 };

          $crate::IResult::Done(&i[byte_index..], o)
        }
      }
    }
  );
);

/// `take_bits!(type, nb) => ( (&[T], usize), U, usize) -> IResult<(&[T], usize), U>`
/// generates a parser consuming the specified number of bytes
///
/// ```
/// # #[macro_use] extern crate nom;
/// # use nom::IResult::Done;
/// # fn main() {
///  named!( take_pair<(u8, u8)>, bits!( pair!( take_bits!( u8, 3 ), take_bits!(u8, 5) ) ) );
///
///  let input = vec![0b10101010, 0b11110000, 0b00110011];
///  let sl    = &input[..];
///
///  assert_eq!(take_pair( sl ),       Done(&sl[1..], (5, 10)) );
///  assert_eq!(take_pair( &sl[1..] ), Done(&sl[2..], (7, 16)) );
/// # }
/// ```
#[macro_export]
macro_rules! take_bits (
  ($i:expr, $t:ty, $count:expr) => (
    {
      use std::ops::Div;
      //println!("taking {} bits from {:?}", $count, $i);
      let (input, bit_offset) = $i;
      let res : $crate::IResult<(&[u8],usize), $t> = if $count == 0 {
        $crate::IResult::Done( (input, bit_offset), 0)
      } else {
        let cnt = ($count as usize + bit_offset).div(8);
        if input.len() * 8 < $count as usize + bit_offset {
          //println!("returning incomplete: {}", $count as usize + bit_offset);
          $crate::IResult::Incomplete($crate::Needed::Size($count as usize))
        } else {
          let mut acc:$t            = 0;
          let mut offset: usize     = bit_offset;
          let mut remaining: usize  = $count;
          let mut end_offset: usize = 0;

          for it in 0..cnt+1 {
            if remaining == 0 {
              break;
            }
            let val: $t = if offset == 0 {
              input[it] as $t
            } else {
              ((input[it] << offset) as u8 >> offset) as $t
            };

            if remaining < 8 - offset {
              acc += val >> (8 - offset - remaining);
              end_offset = remaining + offset;
              break;
            } else {
              acc += val << remaining - (8 - offset);
              remaining -= 8 - offset;
              offset = 0;
            }
          }
          $crate::IResult::Done( (&input[cnt..], end_offset) , acc)
        }
      };
      res
    }
  );
);

#[macro_export]
macro_rules! tag_bits (
  ($i:expr, $t:ty, $count:expr, $p: pat) => (
    {
      match take_bits!($i, $t, $count) {
        $crate::IResult::Incomplete(i) => $crate::IResult::Incomplete(i),
        $crate::IResult::Done(i, o)    => {
          if let $p = o {
            let res: $crate::IResult<(&[u8],usize),$t> = $crate::IResult::Done(i, o);
            res
          } else {
            $crate::IResult::Error($crate::Err::Code($crate::ErrorKind::TagBits))
          }
        },
        _                              => {
          $crate::IResult::Error($crate::Err::Code($crate::ErrorKind::TagBits))
        }
      }
    }
  )
);

#[cfg(test)]
mod tests {
  use internal::{IResult,Needed};

  #[test]
  fn take_bits() {
    let input = vec![0b10101010, 0b11110000, 0b00110011];
    let sl    = &input[..];

    assert_eq!(take_bits!( (sl, 0), u8,   0 ), IResult::Done((sl, 0), 0));
    assert_eq!(take_bits!( (sl, 0), u8,   8 ), IResult::Done((&sl[1..], 0), 170));
    assert_eq!(take_bits!( (sl, 0), u8,   3 ), IResult::Done((&sl[0..], 3), 5));
    assert_eq!(take_bits!( (sl, 0), u8,   6 ), IResult::Done((&sl[0..], 6), 42));
    assert_eq!(take_bits!( (sl, 1), u8,   1 ), IResult::Done((&sl[0..], 2), 0));
    assert_eq!(take_bits!( (sl, 1), u8,   2 ), IResult::Done((&sl[0..], 3), 1));
    assert_eq!(take_bits!( (sl, 1), u8,   3 ), IResult::Done((&sl[0..], 4), 2));
    assert_eq!(take_bits!( (sl, 6), u8,   3 ), IResult::Done((&sl[1..], 1), 5));
    assert_eq!(take_bits!( (sl, 0), u16, 10 ), IResult::Done((&sl[1..], 2), 683));
    assert_eq!(take_bits!( (sl, 0), u16,  8 ), IResult::Done((&sl[1..], 0), 170));
    assert_eq!(take_bits!( (sl, 6), u16, 10 ), IResult::Done((&sl[2..], 0), 752));
    assert_eq!(take_bits!( (sl, 6), u16, 11 ), IResult::Done((&sl[2..], 1), 1504));
    assert_eq!(take_bits!( (sl, 0), u32, 20 ), IResult::Done((&sl[2..], 4), 700163));
    assert_eq!(take_bits!( (sl, 4), u32, 20 ), IResult::Done((&sl[3..], 0), 716851));
    assert_eq!(take_bits!( (sl, 4), u32, 22 ), IResult::Incomplete(Needed::Size(22)));
  }

  #[test]
  fn tag_bits() {
    let input = vec![0b10101010, 0b11110000, 0b00110011];
    let sl    = &input[..];

    assert_eq!(tag_bits!( (sl, 0), u8,   3, 0b101), IResult::Done((&sl[0..], 3), 5));
    assert_eq!(tag_bits!( (sl, 0), u8,   4, 0b1010), IResult::Done((&sl[0..], 4), 10));
  }
}
<|endoftext|>


<commit_before>//! Bit level parsers and combinators
//!
//! Bit parsing is handled by tweaking the input in most macros.
//! In byte level parsing, the input is generally a `&[u8]` passed from combinator
//! to combinator until the slices are manipulated.
//!
//! Bit parsers take a `(&[u8], usize)` as input. The first part of the tuple is an byte slice,
//! the second part is a bit offset in the first byte of the slice.
//!
//! By passing a pair like this, we can leverage most of the combinators, and avoid
//! transforming the whole slice to a vector of booleans. This should make it easy
//! to see a byte slice as a bit stream, and parse code points of arbitrary bit length.

/// `bits!( parser ) => ( &[u8], (&[u8], usize) -> IResult<(&[u8], usize), T> ) -> IResult<&[u8], T>`
/// transforms its byte slice input in a bit stream for the underlying parsers
///
/// ```
/// # #[macro_use] extern crate nom;
/// # use nom::IResult::Done;
/// # fn main() {
///  named!( take_3_bits<u8>, bits!( take_bits!( u8, 3 ) ) );
///
///  let input = vec![0b10101010, 0b11110000, 0b00110011];
///  let sl    = &input[..];
///
///  assert_eq!(take_3_bits( sl ), Done(&sl[1..], 5) );
/// # }
#[macro_export]
macro_rules! bits (
  ($i:expr, $submac:ident!( $($args:tt)* )) => (
    bits_impl!($i, $submac!($($args)*));
  );
  ($i:expr, $f:expr) => (
    bits_impl!($i, call!($f));
  );
);

/// Internal parser, do not use directly
#[macro_export]
macro_rules! bits_impl (
  ($i:expr, $submac:ident!( $($args:tt)* )) => (
    {
      let input = ($i, 0usize);
      match $submac!(input, $($args)*) {
        $crate::IResult::Error(e)                            => {
          let err = match e {
            $crate::Err::Code(k) | $crate::Err::Node(k, _) | $crate::Err::Position(k, _) | $crate::Err::NodePosition(k, _, _) => {
              $crate::Err::Position(k, $i)
            }
          };
          $crate::IResult::Error(err)
        }
        $crate::IResult::Incomplete($crate::Needed::Unknown) => $crate::IResult::Incomplete($crate::Needed::Unknown),
        $crate::IResult::Incomplete($crate::Needed::Size(i)) => $crate::IResult::Incomplete($crate::Needed::Size(i)),
        $crate::IResult::Done((i, bit_index), o)             => {
          let byte_index = if bit_index == 0 { 0 } else { 1 };

          $crate::IResult::Done(&i[byte_index..], o)
        }
      }
    }
  );
);

/// `take_bits!(type, nb) => ( (&[T], usize), U, usize) -> IResult<(&[T], usize), U>`
/// generates a parser consuming the specified number of bytes
///
/// ```
/// # #[macro_use] extern crate nom;
/// # use nom::IResult::Done;
/// # fn main() {
///  named!( take_pair<(u8, u8)>, bits!( pair!( take_bits!( u8, 3 ), take_bits!(u8, 5) ) ) );
///
///  let input = vec![0b10101010, 0b11110000, 0b00110011];
///  let sl    = &input[..];
///
///  assert_eq!(take_pair( sl ),       Done(&sl[1..], (5, 10)) );
///  assert_eq!(take_pair( &sl[1..] ), Done(&sl[2..], (7, 16)) );
/// # }
/// ```
#[macro_export]
macro_rules! take_bits (
  ($i:expr, $t:ty, $count:expr) => (
    {
      use std::ops::Div;
      //println!("taking {} bits from {:?}", $count, $i);
      let (input, bit_offset) = $i;
      let res : $crate::IResult<(&[u8],usize), $t> = if $count == 0 {
        $crate::IResult::Done( (input, bit_offset), 0)
      } else {
        let cnt = ($count as usize + bit_offset).div(8);
        if input.len() * 8 < $count as usize + bit_offset {
          //println!("returning incomplete: {}", $count as usize + bit_offset);
          $crate::IResult::Incomplete($crate::Needed::Size($count as usize))
        } else {
          let mut acc:$t            = 0;
          let mut offset: usize     = bit_offset;
          let mut remaining: usize  = $count;
          let mut end_offset: usize = 0;

          for it in 0..cnt+1 {
            if remaining == 0 {
              break;
            }
            let val: $t = if offset == 0 {
              input[it] as $t
            } else {
              ((input[it] << offset) as u8 >> offset) as $t
            };

            if remaining < 8 - offset {
              acc += val >> (8 - offset - remaining);
              end_offset = remaining + offset;
              break;
            } else {
              acc += val << remaining - (8 - offset);
              remaining -= 8 - offset;
              offset = 0;
            }
          }
          $crate::IResult::Done( (&input[cnt..], end_offset) , acc)
        }
      };
      res
    }
  );
);

#[macro_export]
macro_rules! tag_bits (
  ($i:expr, $t:ty, $count:expr, $p: pat) => (
    {
      match take_bits!($i, $t, $count) {
        $crate::IResult::Incomplete(i) => $crate::IResult::Incomplete(i),
        $crate::IResult::Done(i, o)    => {
          if let $p = o {
            let res: $crate::IResult<(&[u8],usize),$t> = $crate::IResult::Done(i, o);
            res
          } else {
            $crate::IResult::Error($crate::Err::Position($crate::ErrorKind::TagBits, $i))
          }
        },
        _                              => {
          $crate::IResult::Error($crate::Err::Position($crate::ErrorKind::TagBits, $i))
        }
      }
    }
  )
);

#[cfg(test)]
mod tests {
  use internal::{IResult,Needed};

  #[test]
  fn take_bits() {
    let input = vec![0b10101010, 0b11110000, 0b00110011];
    let sl    = &input[..];

    assert_eq!(take_bits!( (sl, 0), u8,   0 ), IResult::Done((sl, 0), 0));
    assert_eq!(take_bits!( (sl, 0), u8,   8 ), IResult::Done((&sl[1..], 0), 170));
    assert_eq!(take_bits!( (sl, 0), u8,   3 ), IResult::Done((&sl[0..], 3), 5));
    assert_eq!(take_bits!( (sl, 0), u8,   6 ), IResult::Done((&sl[0..], 6), 42));
    assert_eq!(take_bits!( (sl, 1), u8,   1 ), IResult::Done((&sl[0..], 2), 0));
    assert_eq!(take_bits!( (sl, 1), u8,   2 ), IResult::Done((&sl[0..], 3), 1));
    assert_eq!(take_bits!( (sl, 1), u8,   3 ), IResult::Done((&sl[0..], 4), 2));
    assert_eq!(take_bits!( (sl, 6), u8,   3 ), IResult::Done((&sl[1..], 1), 5));
    assert_eq!(take_bits!( (sl, 0), u16, 10 ), IResult::Done((&sl[1..], 2), 683));
    assert_eq!(take_bits!( (sl, 0), u16,  8 ), IResult::Done((&sl[1..], 0), 170));
    assert_eq!(take_bits!( (sl, 6), u16, 10 ), IResult::Done((&sl[2..], 0), 752));
    assert_eq!(take_bits!( (sl, 6), u16, 11 ), IResult::Done((&sl[2..], 1), 1504));
    assert_eq!(take_bits!( (sl, 0), u32, 20 ), IResult::Done((&sl[2..], 4), 700163));
    assert_eq!(take_bits!( (sl, 4), u32, 20 ), IResult::Done((&sl[3..], 0), 716851));
    assert_eq!(take_bits!( (sl, 4), u32, 22 ), IResult::Incomplete(Needed::Size(22)));
  }

  #[test]
  fn tag_bits() {
    let input = vec![0b10101010, 0b11110000, 0b00110011];
    let sl    = &input[..];

    assert_eq!(tag_bits!( (sl, 0), u8,   3, 0b101), IResult::Done((&sl[0..], 3), 5));
    assert_eq!(tag_bits!( (sl, 0), u8,   4, 0b1010), IResult::Done((&sl[0..], 4), 10));
  }
}
<commit_msg>bits! translates Error position and Incomplete number<commit_after>//! Bit level parsers and combinators
//!
//! Bit parsing is handled by tweaking the input in most macros.
//! In byte level parsing, the input is generally a `&[u8]` passed from combinator
//! to combinator until the slices are manipulated.
//!
//! Bit parsers take a `(&[u8], usize)` as input. The first part of the tuple is an byte slice,
//! the second part is a bit offset in the first byte of the slice.
//!
//! By passing a pair like this, we can leverage most of the combinators, and avoid
//! transforming the whole slice to a vector of booleans. This should make it easy
//! to see a byte slice as a bit stream, and parse code points of arbitrary bit length.

/// `bits!( parser ) => ( &[u8], (&[u8], usize) -> IResult<(&[u8], usize), T> ) -> IResult<&[u8], T>`
/// transforms its byte slice input in a bit stream for the underlying parsers
///
/// ```
/// # #[macro_use] extern crate nom;
/// # use nom::IResult::Done;
/// # fn main() {
///  named!( take_3_bits<u8>, bits!( take_bits!( u8, 3 ) ) );
///
///  let input = vec![0b10101010, 0b11110000, 0b00110011];
///  let sl    = &input[..];
///
///  assert_eq!(take_3_bits( sl ), Done(&sl[1..], 5) );
/// # }
#[macro_export]
macro_rules! bits (
  ($i:expr, $submac:ident!( $($args:tt)* )) => (
    bits_impl!($i, $submac!($($args)*));
  );
  ($i:expr, $f:expr) => (
    bits_impl!($i, call!($f));
  );
);

/// Internal parser, do not use directly
#[macro_export]
macro_rules! bits_impl (
  ($i:expr, $submac:ident!( $($args:tt)* )) => (
    {
      let input = ($i, 0usize);
      match $submac!(input, $($args)*) {
        $crate::IResult::Error(e)                            => {
          let err = match e {
            $crate::Err::Code(k) | $crate::Err::Node(k, _) => $crate::Err::Code(k),
            $crate::Err::Position(k, (i,b)) | $crate::Err::NodePosition(k, (i,b), _) => {
              $crate::Err::Position(k, &i[b/8..])
            }
          };
          $crate::IResult::Error(err)
        }
        $crate::IResult::Incomplete($crate::Needed::Unknown) => $crate::IResult::Incomplete($crate::Needed::Unknown),
        $crate::IResult::Incomplete($crate::Needed::Size(i)) => {
          //println!("bits parser returned Needed::Size({})", i);
          $crate::IResult::Incomplete($crate::Needed::Size(i / 8 + 1))
        },
        $crate::IResult::Done((i, bit_index), o)             => {
          let byte_index = bit_index / 8 + if bit_index % 8 == 0 { 0 } else { 1 } ;
          //println!("bit index=={} => byte index=={}", bit_index, byte_index);
          $crate::IResult::Done(&i[byte_index..], o)
        }
      }
    }
  );
);

/// `take_bits!(type, nb) => ( (&[T], usize), U, usize) -> IResult<(&[T], usize), U>`
/// generates a parser consuming the specified number of bytes
///
/// ```
/// # #[macro_use] extern crate nom;
/// # use nom::IResult::Done;
/// # fn main() {
///  named!( take_pair<(u8, u8)>, bits!( pair!( take_bits!( u8, 3 ), take_bits!(u8, 5) ) ) );
///
///  let input = vec![0b10101010, 0b11110000, 0b00110011];
///  let sl    = &input[..];
///
///  assert_eq!(take_pair( sl ),       Done(&sl[1..], (5, 10)) );
///  assert_eq!(take_pair( &sl[1..] ), Done(&sl[2..], (7, 16)) );
/// # }
/// ```
#[macro_export]
macro_rules! take_bits (
  ($i:expr, $t:ty, $count:expr) => (
    {
      use std::ops::Div;
      //println!("taking {} bits from {:?}", $count, $i);
      let (input, bit_offset) = $i;
      let res : $crate::IResult<(&[u8],usize), $t> = if $count == 0 {
        $crate::IResult::Done( (input, bit_offset), 0)
      } else {
        let cnt = ($count as usize + bit_offset).div(8);
        if input.len() * 8 < $count as usize + bit_offset {
          //println!("returning incomplete: {}", $count as usize + bit_offset);
          $crate::IResult::Incomplete($crate::Needed::Size($count as usize))
        } else {
          let mut acc:$t            = 0;
          let mut offset: usize     = bit_offset;
          let mut remaining: usize  = $count;
          let mut end_offset: usize = 0;

          for it in 0..cnt+1 {
            if remaining == 0 {
              break;
            }
            let val: $t = if offset == 0 {
              input[it] as $t
            } else {
              ((input[it] << offset) as u8 >> offset) as $t
            };

            if remaining < 8 - offset {
              acc += val >> (8 - offset - remaining);
              end_offset = remaining + offset;
              break;
            } else {
              acc += val << remaining - (8 - offset);
              remaining -= 8 - offset;
              offset = 0;
            }
          }
          $crate::IResult::Done( (&input[cnt..], end_offset) , acc)
        }
      };
      res
    }
  );
);

#[macro_export]
macro_rules! tag_bits (
  ($i:expr, $t:ty, $count:expr, $p: pat) => (
    {
      match take_bits!($i, $t, $count) {
        $crate::IResult::Incomplete(i) => $crate::IResult::Incomplete(i),
        $crate::IResult::Done(i, o)    => {
          if let $p = o {
            let res: $crate::IResult<(&[u8],usize),$t> = $crate::IResult::Done(i, o);
            res
          } else {
            $crate::IResult::Error($crate::Err::Position($crate::ErrorKind::TagBits, $i))
          }
        },
        _                              => {
          $crate::IResult::Error($crate::Err::Position($crate::ErrorKind::TagBits, $i))
        }
      }
    }
  )
);

#[cfg(test)]
mod tests {
  use internal::{IResult,Needed,Err};
  use ErrorKind;

  #[test]
  fn take_bits() {
    let input = vec![0b10101010, 0b11110000, 0b00110011];
    let sl    = &input[..];

    assert_eq!(take_bits!( (sl, 0), u8,   0 ), IResult::Done((sl, 0), 0));
    assert_eq!(take_bits!( (sl, 0), u8,   8 ), IResult::Done((&sl[1..], 0), 170));
    assert_eq!(take_bits!( (sl, 0), u8,   3 ), IResult::Done((&sl[0..], 3), 5));
    assert_eq!(take_bits!( (sl, 0), u8,   6 ), IResult::Done((&sl[0..], 6), 42));
    assert_eq!(take_bits!( (sl, 1), u8,   1 ), IResult::Done((&sl[0..], 2), 0));
    assert_eq!(take_bits!( (sl, 1), u8,   2 ), IResult::Done((&sl[0..], 3), 1));
    assert_eq!(take_bits!( (sl, 1), u8,   3 ), IResult::Done((&sl[0..], 4), 2));
    assert_eq!(take_bits!( (sl, 6), u8,   3 ), IResult::Done((&sl[1..], 1), 5));
    assert_eq!(take_bits!( (sl, 0), u16, 10 ), IResult::Done((&sl[1..], 2), 683));
    assert_eq!(take_bits!( (sl, 0), u16,  8 ), IResult::Done((&sl[1..], 0), 170));
    assert_eq!(take_bits!( (sl, 6), u16, 10 ), IResult::Done((&sl[2..], 0), 752));
    assert_eq!(take_bits!( (sl, 6), u16, 11 ), IResult::Done((&sl[2..], 1), 1504));
    assert_eq!(take_bits!( (sl, 0), u32, 20 ), IResult::Done((&sl[2..], 4), 700163));
    assert_eq!(take_bits!( (sl, 4), u32, 20 ), IResult::Done((&sl[3..], 0), 716851));
    assert_eq!(take_bits!( (sl, 4), u32, 22 ), IResult::Incomplete(Needed::Size(22)));
  }

  #[test]
  fn tag_bits() {
    let input = vec![0b10101010, 0b11110000, 0b00110011];
    let sl    = &input[..];

    assert_eq!(tag_bits!( (sl, 0), u8,   3, 0b101), IResult::Done((&sl[0..], 3), 5));
    assert_eq!(tag_bits!( (sl, 0), u8,   4, 0b1010), IResult::Done((&sl[0..], 4), 10));
  }

  named!(ch<(&[u8],usize),(u8,u8)>,
    chain!(
      tag_bits!(u8, 3, 0b101) ~
      x: take_bits!(u8, 4)    ~
      y: take_bits!(u8, 5)    ,
      || { (x,y) }
    )
  );

  #[test]
  fn chain_bits() {
    let input = vec![0b10101010, 0b11110000, 0b00110011];
    let sl    = &input[..];
    assert_eq!(ch((&input[..],0)), IResult::Done((&sl[1..], 4), (5,15)));
    assert_eq!(ch((&input[..],4)), IResult::Done((&sl[2..], 0), (7,16)));
    assert_eq!(ch((&input[..1],0)), IResult::Incomplete(Needed::Size(12)));
  }

  named!(ch_bytes<(u8,u8)>, bits!(ch));
  #[test]
  fn bits_to_bytes() {
    let input = vec![0b10101010, 0b11110000, 0b00110011];
    assert_eq!(ch_bytes(&input[..]), IResult::Done(&input[2..], (5,15)));
    assert_eq!(ch_bytes(&input[..1]), IResult::Incomplete(Needed::Size(2)));
    assert_eq!(ch_bytes(&input[1..]), IResult::Error(Err::Position(ErrorKind::TagBits, &input[1..])));
  }
}
<|endoftext|>


<commit_before>use std::str;

use crate::{
    Client, Result, error,
    proto::Version,
};
use futures::prelude::*;
use romio::TcpStream;
use byteorder::{ByteOrder, LittleEndian};
use serde::{Serialize, Deserialize};
use scram::client::{ScramClient, ServerFirst, ServerFinal};

const NULL_BYTE: u8 = b'\0';

#[derive(Debug)]
struct Session {
    id: u64,
    broken: bool,
    stream: TcpStream,
}

#[derive(Debug)]
pub struct Connection<'a> {
    client: Client<'a>,
    session: Session,
}

#[derive(Serialize, Deserialize, Debug)]
struct ServerInfo {
    success: bool,
    min_protocol_version: usize,
    max_protocol_version: usize,
    server_version: String,
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthRequest {
    protocol_version: i32,
    authentication_method: String,
    authentication: String,
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthResponse {
    success: bool,
    authentication: Option<String>,
    error_code: Option<usize>,
    error: Option<String>,
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthConfirmation {
    authentication: String,
}

struct HandShake<'a> {
    // this should be enough for the handshake messages
    buf: [u8; 512],
    conn: Connection<'a>,
}

impl<'a> Connection<'a> {
    pub(crate) async fn new(client: Client<'a>, id: u64) -> Result<Connection<'a>> {
        let cfg = client.config();
        let stream = await!(TcpStream::connect(cfg.server()))?;
        let session = Session { id, stream, broken: false };
        let conn = Connection { client, session };
        let handshake = HandShake { conn, buf: [0u8; 512] };
        await!(handshake.greet())
    }
}

impl<'a> HandShake<'a> {
    async fn verify_version(&mut self) -> Result<&mut HandShake<'a>> {
        LittleEndian::write_u32(&mut self.buf, Version::V1_0 as u32);
        await!(self.conn.session.stream.write_all(&self.buf[..4]))?;
        await!(self.conn.session.stream.read(&mut self.buf))?;
        let resp = self.read_buf();
        match serde_json::from_slice::<ServerInfo>(resp) {
            Ok(info) => {
                if !info.success {
                    return Err(error::Runtime::Internal(resp.to_owned()))?;
                }
            }
            Err(_) => {
                let msg = str::from_utf8(resp)?;
                return Err(error::Driver::Other(msg.to_owned()))?;
            }
        }
        Ok(self)
    }

    async fn greet(mut self) -> Result<Connection<'a>> {
        await!(self.verify_version())?;

        // Send client first message
        let cfg = &self.conn.client.config();
        let scram = ScramClient::new(cfg.user(), cfg.password(), None)?;
        let (scram, msg) = client_first(scram)?;
        await!(self.conn.session.stream.write_all(&msg))?;

        // Send client final message
        await!(self.conn.session.stream.read(&mut self.buf))?;
        let info = AuthResponse::from_slice(self.read_buf())?;
        match info.authentication {
            Some(ref auth) => {
                let (scram, msg) = client_final(scram, auth)?;
                await!(self.conn.session.stream.write_all(&msg))?;
                await!(self.conn.session.stream.read(&mut self.buf))?;
                server_final(scram, self.read_buf())?;
                await!(self.conn.session.stream.flush())?;
            }
            None => {
                let msg = String::from("Server did not send authentication info.");
                return Err(error::Driver::Other(msg))?;
            }
        }

        Ok(self.conn)
    }

    fn read_buf(&self) -> &[u8] {
        let len = self.buf.iter()
            .take_while(|x| **x != NULL_BYTE)
            .count();
        &self.buf[..len]
    }
}

impl AuthResponse {
    fn from_slice(resp: &[u8]) -> Result<Self> {
        match serde_json::from_slice::<AuthResponse>(resp) {
            Ok(info) => {
                if !info.success {
                    // If error code is between 10 and 20, this is an auth error
                    if let Some(10...20) = info.error_code {
                        if let Some(msg) = info.error {
                            return Err(error::Driver::Auth(msg))?;
                        }
                    }
                    return Err(error::Runtime::Internal(resp.to_owned()))?;
                }
                Ok(info)
            }
            Err(_) => {
                let msg = str::from_utf8(resp)?;
                Err(error::Driver::Other(msg.to_owned()))?
            }
        }
    }
}

fn client_first<'a>(scram: ScramClient<'a>) -> Result<(ServerFirst<'a>, Vec<u8>)> {
    let (scram, client_first) = scram.client_first();
    let ar = AuthRequest {
        protocol_version: 0,
        authentication_method: String::from("SCRAM-SHA-256"),
        authentication: client_first,
    };
    let mut msg = serde_json::to_vec(&ar)?;
    msg.push(NULL_BYTE);
    Ok((scram, msg))
}

fn client_final<'a>(scram: ServerFirst<'a>, auth: &str) -> Result<(ServerFinal, Vec<u8>)> {
    let scram = scram.handle_server_first(auth)?;
    let (scram, client_final) = scram.client_final();
    let conf = AuthConfirmation { authentication: client_final };
    let mut msg = serde_json::to_vec(&conf)?;
    msg.push(NULL_BYTE);
    Ok((scram, msg))
}

fn server_final(scram: ServerFinal, resp: &[u8]) -> Result<()> {
    let info = AuthResponse::from_slice(resp)?;
    if let Some(auth) = info.authentication {
        if let Err(err) = scram.handle_server_final(&auth) {
            return Err(err)?;
        }
    }
    Ok(())
}
<commit_msg>optimise message exchange<commit_after>use std::str;

use crate::{
    Client, Result, error,
    proto::Version,
};
use futures::prelude::*;
use romio::TcpStream;
use byteorder::{ByteOrder, LittleEndian};
use serde::{Serialize, Deserialize};
use scram::client::{ScramClient, ServerFirst, ServerFinal};

const NULL_BYTE: u8 = b'\0';

#[derive(Debug)]
struct Session {
    id: u64,
    broken: bool,
    stream: TcpStream,
}

#[derive(Debug)]
pub struct Connection<'a> {
    client: Client<'a>,
    session: Session,
}

#[derive(Serialize, Deserialize, Debug)]
struct ServerInfo {
    success: bool,
    min_protocol_version: usize,
    max_protocol_version: usize,
    server_version: String,
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthRequest {
    protocol_version: i32,
    authentication_method: String,
    authentication: String,
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthResponse {
    success: bool,
    authentication: Option<String>,
    error_code: Option<usize>,
    error: Option<String>,
}

#[derive(Serialize, Deserialize, Debug)]
struct AuthConfirmation {
    authentication: String,
}

struct HandShake<'a> {
    // this should be enough for the handshake messages
    buf: [u8; 512],
    conn: Connection<'a>,
}

impl<'a> Connection<'a> {
    pub(crate) async fn new(client: Client<'a>, id: u64) -> Result<Connection<'a>> {
        let cfg = client.config();
        let stream = await!(TcpStream::connect(cfg.server()))?;
        let session = Session { id, stream, broken: false };
        let conn = Connection { client, session };
        let handshake = HandShake { conn, buf: [0u8; 512] };
        await!(handshake.greet())
    }
}

impl<'a> HandShake<'a> {
    // Performs the actual handshake
    //
    // This method optimises message exchange as suggested in the RethinkDB
    // documentation by sending message 3 right after message 1, without waiting
    // for message 2 first.
    async fn greet(mut self) -> Result<Connection<'a>> {
        // Send the version we support
        LittleEndian::write_u32(&mut self.buf, Version::V1_0 as u32);
        await!(self.conn.session.stream.write_all(&self.buf[..4]))?; // message 1

        // Send client first message
        let cfg = &self.conn.client.config();
        let scram = ScramClient::new(cfg.user(), cfg.password(), None)?;
        let (scram, msg) = client_first(scram)?;
        await!(self.conn.session.stream.write_all(&msg))?; // message 3

        // Receive supported versions
        await!(self.conn.session.stream.read(&mut self.buf))?; // message 2
        ServerInfo::from_slice(self.read_buf())?;

        // Receive server first message
        await!(self.conn.session.stream.read(&mut self.buf))?; // message 4
        let info = AuthResponse::from_slice(self.read_buf())?;
        let auth = match info.authentication {
            Some(auth) => auth,
            None => {
                let msg = String::from("Server did not send authentication info.");
                return Err(error::Driver::Other(msg))?;
            }
        };

        // Send client final message
        let (scram, msg) = client_final(scram, &auth)?;
        await!(self.conn.session.stream.write_all(&msg))?; // message 5

        // Receive server final message
        await!(self.conn.session.stream.read(&mut self.buf))?; // message 6
        server_final(scram, self.read_buf())?;
        await!(self.conn.session.stream.flush())?;

        Ok(self.conn)
    }

    fn read_buf(&self) -> &[u8] {
        let len = self.buf.iter()
            .take_while(|x| **x != NULL_BYTE)
            .count();
        &self.buf[..len]
    }
}

impl ServerInfo {
    fn from_slice(resp: &[u8]) -> Result<Self> {
        match serde_json::from_slice::<ServerInfo>(resp) {
            Ok(info) => {
                if !info.success {
                    return Err(error::Runtime::Internal(resp.to_owned()))?;
                }
                Ok(info)
            }
            Err(_) => {
                let msg = str::from_utf8(resp)?;
                Err(error::Driver::Other(msg.to_owned()))?
            }
        }
    }
}

impl AuthResponse {
    fn from_slice(resp: &[u8]) -> Result<Self> {
        match serde_json::from_slice::<AuthResponse>(resp) {
            Ok(info) => {
                if !info.success {
                    // If error code is between 10 and 20, this is an auth error
                    if let Some(10...20) = info.error_code {
                        if let Some(msg) = info.error {
                            return Err(error::Driver::Auth(msg))?;
                        }
                    }
                    return Err(error::Runtime::Internal(resp.to_owned()))?;
                }
                Ok(info)
            }
            Err(_) => {
                let msg = str::from_utf8(resp)?;
                Err(error::Driver::Other(msg.to_owned()))?
            }
        }
    }
}

fn client_first<'a>(scram: ScramClient<'a>) -> Result<(ServerFirst<'a>, Vec<u8>)> {
    let (scram, client_first) = scram.client_first();
    let ar = AuthRequest {
        protocol_version: 0,
        authentication_method: String::from("SCRAM-SHA-256"),
        authentication: client_first,
    };
    let mut msg = serde_json::to_vec(&ar)?;
    msg.push(NULL_BYTE);
    Ok((scram, msg))
}

fn client_final<'a>(scram: ServerFirst<'a>, auth: &str) -> Result<(ServerFinal, Vec<u8>)> {
    let scram = scram.handle_server_first(auth)?;
    let (scram, client_final) = scram.client_final();
    let conf = AuthConfirmation { authentication: client_final };
    let mut msg = serde_json::to_vec(&conf)?;
    msg.push(NULL_BYTE);
    Ok((scram, msg))
}

fn server_final(scram: ServerFinal, resp: &[u8]) -> Result<()> {
    let info = AuthResponse::from_slice(resp)?;
    if let Some(auth) = info.authentication {
        if let Err(err) = scram.handle_server_final(&auth) {
            return Err(err)?;
        }
    }
    Ok(())
}
<|endoftext|>