<commit_before>//! A private parser implementation of IPv4, IPv6, and socket addresses.
//!
//! This module is "publicly exported" through the `FromStr` implementations
//! below.

#[cfg(test)]
mod tests;

use crate::convert::TryInto as _;
use crate::error::Error;
use crate::fmt;
use crate::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
use crate::str::FromStr;

trait ReadNumberHelper: crate::marker::Sized + crate::cmp::PartialEq {
    const ZERO: Self;
    fn checked_mul(&self, other: u32) -> Option<Self>;
    fn checked_add(&self, other: u32) -> Option<Self>;
}

macro_rules! impl_helper {
    ($($t:ty)*) => ($(impl ReadNumberHelper for $t {
        const ZERO: Self = 0;
        #[inline]
        fn checked_mul(&self, other: u32) -> Option<Self> {
            Self::checked_mul(*self, other.try_into().ok()?)
        }
        #[inline]
        fn checked_add(&self, other: u32) -> Option<Self> {
            Self::checked_add(*self, other.try_into().ok()?)
        }
    })*)
}

impl_helper! { u8 u16 u32 }

struct Parser<'a> {
    // Parsing as ASCII, so can use byte array.
    state: &'a [u8],
}

impl<'a> Parser<'a> {
    fn new(input: &'a str) -> Parser<'a> {
        Parser { state: input.as_bytes() }
    }

    /// Run a parser, and restore the pre-parse state if it fails.
    fn read_atomically<T, F>(&mut self, inner: F) -> Option<T>
    where
        F: FnOnce(&mut Parser<'_>) -> Option<T>,
    {
        let state = self.state;
        let result = inner(self);
        if result.is_none() {
            self.state = state;
        }
        result
    }

    /// Run a parser, but fail if the entire input wasn't consumed.
    /// Doesn't run atomically.
    fn parse_with<T, F>(&mut self, inner: F) -> Result<T, AddrParseError>
    where
        F: FnOnce(&mut Parser<'_>) -> Option<T>,
    {
        let result = inner(self);
        if self.state.is_empty() { result } else { None }.ok_or(AddrParseError(()))
    }

    /// Peek the next character from the input
    fn peek_char(&self) -> Option<char> {
        self.state.first().map(|&b| char::from(b))
    }

    /// Read the next character from the input
    fn read_char(&mut self) -> Option<char> {
        self.state.split_first().map(|(&b, tail)| {
            self.state = tail;
            char::from(b)
        })
    }

    #[must_use]
    /// Read the next character from the input if it matches the target.
    fn read_given_char(&mut self, target: char) -> Option<()> {
        self.read_atomically(|p| {
            p.read_char().and_then(|c| if c == target { Some(()) } else { None })
        })
    }

    /// Helper for reading separators in an indexed loop. Reads the separator
    /// character iff index > 0, then runs the parser. When used in a loop,
    /// the separator character will only be read on index > 0 (see
    /// read_ipv4_addr for an example)
    fn read_separator<T, F>(&mut self, sep: char, index: usize, inner: F) -> Option<T>
    where
        F: FnOnce(&mut Parser<'_>) -> Option<T>,
    {
        self.read_atomically(move |p| {
            if index > 0 {
                p.read_given_char(sep)?;
            }
            inner(p)
        })
    }

    // Read a number off the front of the input in the given radix, stopping
    // at the first non-digit character or eof. Fails if the number has more
    // digits than max_digits or if there is no number.
    fn read_number<T: ReadNumberHelper>(
        &mut self,
        radix: u32,
        max_digits: Option<usize>,
        allow_zero_prefix: bool,
    ) -> Option<T> {
        self.read_atomically(move |p| {
            let mut result = T::ZERO;
            let mut digit_count = 0;
            let has_leading_zero = p.peek_char() == Some('0');

            while let Some(digit) = p.read_atomically(|p| p.read_char()?.to_digit(radix)) {
                result = result.checked_mul(radix)?;
                result = result.checked_add(digit)?;
                digit_count += 1;
                if let Some(max_digits) = max_digits {
                    if digit_count > max_digits {
                        return None;
                    }
                }
            }

            if digit_count == 0 {
                None
            } else if !allow_zero_prefix
                && has_leading_zero
                && (result != T::ZERO || result == T::ZERO && digit_count > 1)
            {
                None
            } else {
                Some(result)
            }
        })
    }

    /// Read an IPv4 address.
    fn read_ipv4_addr(&mut self) -> Option<Ipv4Addr> {
        self.read_atomically(|p| {
            let mut groups = [0; 4];

            for (i, slot) in groups.iter_mut().enumerate() {
                *slot = p.read_separator('.', i, |p| {
                    // Disallow octal number in IP string.
                    // https://tools.ietf.org/html/rfc6943#section-3.1.1
                    p.read_number(10, None, false)
                })?;
            }

            Some(groups.into())
        })
    }

    /// Read an IPv6 Address.
    fn read_ipv6_addr(&mut self) -> Option<Ipv6Addr> {
        /// Read a chunk of an IPv6 address into `groups`. Returns the number
        /// of groups read, along with a bool indicating if an embedded
        /// trailing IPv4 address was read. Specifically, read a series of
        /// colon-separated IPv6 groups (0x0000 - 0xFFFF), with an optional
        /// trailing embedded IPv4 address.
        fn read_groups(p: &mut Parser<'_>, groups: &mut [u16]) -> (usize, bool) {
            let limit = groups.len();

            for (i, slot) in groups.iter_mut().enumerate() {
                // Try to read a trailing embedded IPv4 address. There must be
                // at least two groups left.
                if i < limit - 1 {
                    let ipv4 = p.read_separator(':', i, |p| p.read_ipv4_addr());

                    if let Some(v4_addr) = ipv4 {
                        let [one, two, three, four] = v4_addr.octets();
                        groups[i + 0] = u16::from_be_bytes([one, two]);
                        groups[i + 1] = u16::from_be_bytes([three, four]);
                        return (i + 2, true);
                    }
                }

                let group = p.read_separator(':', i, |p| p.read_number(16, Some(4), true));

                match group {
                    Some(g) => *slot = g,
                    None => return (i, false),
                }
            }
            (groups.len(), false)
        }

        self.read_atomically(|p| {
            // Read the front part of the address; either the whole thing, or up
            // to the first ::
            let mut head = [0; 8];
            let (head_size, head_ipv4) = read_groups(p, &mut head);

            if head_size == 8 {
                return Some(head.into());
            }

            // IPv4 part is not allowed before `::`
            if head_ipv4 {
                return None;
            }

            // Read `::` if previous code parsed less than 8 groups.
            // `::` indicates one or more groups of 16 bits of zeros.
            p.read_given_char(':')?;
            p.read_given_char(':')?;

            // Read the back part of the address. The :: must contain at least one
            // set of zeroes, so our max length is 7.
            let mut tail = [0; 7];
            let limit = 8 - (head_size + 1);
            let (tail_size, _) = read_groups(p, &mut tail[..limit]);

            // Concat the head and tail of the IP address
            head[(8 - tail_size)..8].copy_from_slice(&tail[..tail_size]);

            Some(head.into())
        })
    }

    /// Read an IP Address, either IPv4 or IPv6.
    fn read_ip_addr(&mut self) -> Option<IpAddr> {
        self.read_ipv4_addr().map(IpAddr::V4).or_else(move || self.read_ipv6_addr().map(IpAddr::V6))
    }

    /// Read a `:` followed by a port in base 10.
    fn read_port(&mut self) -> Option<u16> {
        self.read_atomically(|p| {
            p.read_given_char(':')?;
            p.read_number(10, None, true)
        })
    }

    /// Read a `%` followed by a scope ID in base 10.
    fn read_scope_id(&mut self) -> Option<u32> {
        self.read_atomically(|p| {
            p.read_given_char('%')?;
            p.read_number(10, None, true)
        })
    }

    /// Read an IPv4 address with a port.
    fn read_socket_addr_v4(&mut self) -> Option<SocketAddrV4> {
        self.read_atomically(|p| {
            let ip = p.read_ipv4_addr()?;
            let port = p.read_port()?;
            Some(SocketAddrV4::new(ip, port))
        })
    }

    /// Read an IPv6 address with a port.
    fn read_socket_addr_v6(&mut self) -> Option<SocketAddrV6> {
        self.read_atomically(|p| {
            p.read_given_char('[')?;
            let ip = p.read_ipv6_addr()?;
            let scope_id = p.read_scope_id().unwrap_or(0);
            p.read_given_char(']')?;

            let port = p.read_port()?;
            Some(SocketAddrV6::new(ip, port, 0, scope_id))
        })
    }

    /// Read an IP address with a port
    fn read_socket_addr(&mut self) -> Option<SocketAddr> {
        self.read_socket_addr_v4()
            .map(SocketAddr::V4)
            .or_else(|| self.read_socket_addr_v6().map(SocketAddr::V6))
    }
}

#[stable(feature = "ip_addr", since = "1.7.0")]
impl FromStr for IpAddr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<IpAddr, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_ip_addr())
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for Ipv4Addr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<Ipv4Addr, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_ipv4_addr())
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for Ipv6Addr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<Ipv6Addr, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_ipv6_addr())
    }
}

#[stable(feature = "socket_addr_from_str", since = "1.5.0")]
impl FromStr for SocketAddrV4 {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<SocketAddrV4, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_socket_addr_v4())
    }
}

#[stable(feature = "socket_addr_from_str", since = "1.5.0")]
impl FromStr for SocketAddrV6 {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<SocketAddrV6, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_socket_addr_v6())
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for SocketAddr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<SocketAddr, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_socket_addr())
    }
}

/// An error which can be returned when parsing an IP address or a socket address.
///
/// This error is used as the error type for the [`FromStr`] implementation for
/// [`IpAddr`], [`Ipv4Addr`], [`Ipv6Addr`], [`SocketAddr`], [`SocketAddrV4`], and
/// [`SocketAddrV6`].
///
/// # Potential causes
///
/// `AddrParseError` may be thrown because the provided string does not parse as the given type,
/// often because it includes information only handled by a different address type.
///
/// ```should_panic
/// use std::net::IpAddr;
/// let _foo: IpAddr = "127.0.0.1:8080".parse().expect("Cannot handle the socket port");
/// ```
///
/// [`IpAddr`] doesn't handle the port. Use [`SocketAddr`] instead.
///
/// ```
/// use std::net::SocketAddr;
///
/// // No problem, the `panic!` message has disappeared.
/// let _foo: SocketAddr = "127.0.0.1:8080".parse().expect("unreachable panic");
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AddrParseError(());

#[stable(feature = "addr_parse_error_error", since = "1.4.0")]
impl fmt::Display for AddrParseError {
    #[allow(deprecated, deprecated_in_future)]
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.write_str(self.description())
    }
}

#[stable(feature = "addr_parse_error_error", since = "1.4.0")]
impl Error for AddrParseError {
    #[allow(deprecated)]
    fn description(&self) -> &str {
        "invalid IP address syntax"
    }
}
<commit_msg>Reject too-long IPs quicker<commit_after>//! A private parser implementation of IPv4, IPv6, and socket addresses.
//!
//! This module is "publicly exported" through the `FromStr` implementations
//! below.

#[cfg(test)]
mod tests;

use crate::convert::TryInto as _;
use crate::error::Error;
use crate::fmt;
use crate::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
use crate::str::FromStr;

trait ReadNumberHelper: crate::marker::Sized + crate::cmp::PartialEq {
    const ZERO: Self;
    fn checked_mul(&self, other: u32) -> Option<Self>;
    fn checked_add(&self, other: u32) -> Option<Self>;
}

macro_rules! impl_helper {
    ($($t:ty)*) => ($(impl ReadNumberHelper for $t {
        const ZERO: Self = 0;
        #[inline]
        fn checked_mul(&self, other: u32) -> Option<Self> {
            Self::checked_mul(*self, other.try_into().ok()?)
        }
        #[inline]
        fn checked_add(&self, other: u32) -> Option<Self> {
            Self::checked_add(*self, other.try_into().ok()?)
        }
    })*)
}

impl_helper! { u8 u16 u32 }

struct Parser<'a> {
    // Parsing as ASCII, so can use byte array.
    state: &'a [u8],
}

impl<'a> Parser<'a> {
    fn new(input: &'a str) -> Parser<'a> {
        Parser { state: input.as_bytes() }
    }

    /// Run a parser, and restore the pre-parse state if it fails.
    fn read_atomically<T, F>(&mut self, inner: F) -> Option<T>
    where
        F: FnOnce(&mut Parser<'_>) -> Option<T>,
    {
        let state = self.state;
        let result = inner(self);
        if result.is_none() {
            self.state = state;
        }
        result
    }

    /// Run a parser, but fail if the entire input wasn't consumed.
    /// Doesn't run atomically.
    fn parse_with<T, F>(&mut self, inner: F) -> Result<T, AddrParseError>
    where
        F: FnOnce(&mut Parser<'_>) -> Option<T>,
    {
        let result = inner(self);
        if self.state.is_empty() { result } else { None }.ok_or(AddrParseError(()))
    }

    /// Peek the next character from the input
    fn peek_char(&self) -> Option<char> {
        self.state.first().map(|&b| char::from(b))
    }

    /// Read the next character from the input
    fn read_char(&mut self) -> Option<char> {
        self.state.split_first().map(|(&b, tail)| {
            self.state = tail;
            char::from(b)
        })
    }

    #[must_use]
    /// Read the next character from the input if it matches the target.
    fn read_given_char(&mut self, target: char) -> Option<()> {
        self.read_atomically(|p| {
            p.read_char().and_then(|c| if c == target { Some(()) } else { None })
        })
    }

    /// Helper for reading separators in an indexed loop. Reads the separator
    /// character iff index > 0, then runs the parser. When used in a loop,
    /// the separator character will only be read on index > 0 (see
    /// read_ipv4_addr for an example)
    fn read_separator<T, F>(&mut self, sep: char, index: usize, inner: F) -> Option<T>
    where
        F: FnOnce(&mut Parser<'_>) -> Option<T>,
    {
        self.read_atomically(move |p| {
            if index > 0 {
                p.read_given_char(sep)?;
            }
            inner(p)
        })
    }

    // Read a number off the front of the input in the given radix, stopping
    // at the first non-digit character or eof. Fails if the number has more
    // digits than max_digits or if there is no number.
    fn read_number<T: ReadNumberHelper>(
        &mut self,
        radix: u32,
        max_digits: Option<usize>,
        allow_zero_prefix: bool,
    ) -> Option<T> {
        self.read_atomically(move |p| {
            let mut result = T::ZERO;
            let mut digit_count = 0;
            let has_leading_zero = p.peek_char() == Some('0');

            while let Some(digit) = p.read_atomically(|p| p.read_char()?.to_digit(radix)) {
                result = result.checked_mul(radix)?;
                result = result.checked_add(digit)?;
                digit_count += 1;
                if let Some(max_digits) = max_digits {
                    if digit_count > max_digits {
                        return None;
                    }
                }
            }

            if digit_count == 0 {
                None
            } else if !allow_zero_prefix
                && has_leading_zero
                && (result != T::ZERO || result == T::ZERO && digit_count > 1)
            {
                None
            } else {
                Some(result)
            }
        })
    }

    /// Read an IPv4 address.
    fn read_ipv4_addr(&mut self) -> Option<Ipv4Addr> {
        self.read_atomically(|p| {
            let mut groups = [0; 4];

            for (i, slot) in groups.iter_mut().enumerate() {
                *slot = p.read_separator('.', i, |p| {
                    // Disallow octal number in IP string.
                    // https://tools.ietf.org/html/rfc6943#section-3.1.1
                    p.read_number(10, None, false)
                })?;
            }

            Some(groups.into())
        })
    }

    /// Read an IPv6 Address.
    fn read_ipv6_addr(&mut self) -> Option<Ipv6Addr> {
        /// Read a chunk of an IPv6 address into `groups`. Returns the number
        /// of groups read, along with a bool indicating if an embedded
        /// trailing IPv4 address was read. Specifically, read a series of
        /// colon-separated IPv6 groups (0x0000 - 0xFFFF), with an optional
        /// trailing embedded IPv4 address.
        fn read_groups(p: &mut Parser<'_>, groups: &mut [u16]) -> (usize, bool) {
            let limit = groups.len();

            for (i, slot) in groups.iter_mut().enumerate() {
                // Try to read a trailing embedded IPv4 address. There must be
                // at least two groups left.
                if i < limit - 1 {
                    let ipv4 = p.read_separator(':', i, |p| p.read_ipv4_addr());

                    if let Some(v4_addr) = ipv4 {
                        let [one, two, three, four] = v4_addr.octets();
                        groups[i + 0] = u16::from_be_bytes([one, two]);
                        groups[i + 1] = u16::from_be_bytes([three, four]);
                        return (i + 2, true);
                    }
                }

                let group = p.read_separator(':', i, |p| p.read_number(16, Some(4), true));

                match group {
                    Some(g) => *slot = g,
                    None => return (i, false),
                }
            }
            (groups.len(), false)
        }

        self.read_atomically(|p| {
            // Read the front part of the address; either the whole thing, or up
            // to the first ::
            let mut head = [0; 8];
            let (head_size, head_ipv4) = read_groups(p, &mut head);

            if head_size == 8 {
                return Some(head.into());
            }

            // IPv4 part is not allowed before `::`
            if head_ipv4 {
                return None;
            }

            // Read `::` if previous code parsed less than 8 groups.
            // `::` indicates one or more groups of 16 bits of zeros.
            p.read_given_char(':')?;
            p.read_given_char(':')?;

            // Read the back part of the address. The :: must contain at least one
            // set of zeroes, so our max length is 7.
            let mut tail = [0; 7];
            let limit = 8 - (head_size + 1);
            let (tail_size, _) = read_groups(p, &mut tail[..limit]);

            // Concat the head and tail of the IP address
            head[(8 - tail_size)..8].copy_from_slice(&tail[..tail_size]);

            Some(head.into())
        })
    }

    /// Read an IP Address, either IPv4 or IPv6.
    fn read_ip_addr(&mut self) -> Option<IpAddr> {
        self.read_ipv4_addr().map(IpAddr::V4).or_else(move || self.read_ipv6_addr().map(IpAddr::V6))
    }

    /// Read a `:` followed by a port in base 10.
    fn read_port(&mut self) -> Option<u16> {
        self.read_atomically(|p| {
            p.read_given_char(':')?;
            p.read_number(10, None, true)
        })
    }

    /// Read a `%` followed by a scope ID in base 10.
    fn read_scope_id(&mut self) -> Option<u32> {
        self.read_atomically(|p| {
            p.read_given_char('%')?;
            p.read_number(10, None, true)
        })
    }

    /// Read an IPv4 address with a port.
    fn read_socket_addr_v4(&mut self) -> Option<SocketAddrV4> {
        self.read_atomically(|p| {
            let ip = p.read_ipv4_addr()?;
            let port = p.read_port()?;
            Some(SocketAddrV4::new(ip, port))
        })
    }

    /// Read an IPv6 address with a port.
    fn read_socket_addr_v6(&mut self) -> Option<SocketAddrV6> {
        self.read_atomically(|p| {
            p.read_given_char('[')?;
            let ip = p.read_ipv6_addr()?;
            let scope_id = p.read_scope_id().unwrap_or(0);
            p.read_given_char(']')?;

            let port = p.read_port()?;
            Some(SocketAddrV6::new(ip, port, 0, scope_id))
        })
    }

    /// Read an IP address with a port
    fn read_socket_addr(&mut self) -> Option<SocketAddr> {
        self.read_socket_addr_v4()
            .map(SocketAddr::V4)
            .or_else(|| self.read_socket_addr_v6().map(SocketAddr::V6))
    }
}

#[stable(feature = "ip_addr", since = "1.7.0")]
impl FromStr for IpAddr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<IpAddr, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_ip_addr())
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for Ipv4Addr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<Ipv4Addr, AddrParseError> {
        // don't try to parse if too long
        if s.len() > 15 {
            Err(AddrParseError(()))
        } else {
            Parser::new(s).parse_with(|p| p.read_ipv4_addr())
        }
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for Ipv6Addr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<Ipv6Addr, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_ipv6_addr())
    }
}

#[stable(feature = "socket_addr_from_str", since = "1.5.0")]
impl FromStr for SocketAddrV4 {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<SocketAddrV4, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_socket_addr_v4())
    }
}

#[stable(feature = "socket_addr_from_str", since = "1.5.0")]
impl FromStr for SocketAddrV6 {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<SocketAddrV6, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_socket_addr_v6())
    }
}

#[stable(feature = "rust1", since = "1.0.0")]
impl FromStr for SocketAddr {
    type Err = AddrParseError;
    fn from_str(s: &str) -> Result<SocketAddr, AddrParseError> {
        Parser::new(s).parse_with(|p| p.read_socket_addr())
    }
}

/// An error which can be returned when parsing an IP address or a socket address.
///
/// This error is used as the error type for the [`FromStr`] implementation for
/// [`IpAddr`], [`Ipv4Addr`], [`Ipv6Addr`], [`SocketAddr`], [`SocketAddrV4`], and
/// [`SocketAddrV6`].
///
/// # Potential causes
///
/// `AddrParseError` may be thrown because the provided string does not parse as the given type,
/// often because it includes information only handled by a different address type.
///
/// ```should_panic
/// use std::net::IpAddr;
/// let _foo: IpAddr = "127.0.0.1:8080".parse().expect("Cannot handle the socket port");
/// ```
///
/// [`IpAddr`] doesn't handle the port. Use [`SocketAddr`] instead.
///
/// ```
/// use std::net::SocketAddr;
///
/// // No problem, the `panic!` message has disappeared.
/// let _foo: SocketAddr = "127.0.0.1:8080".parse().expect("unreachable panic");
/// ```
#[stable(feature = "rust1", since = "1.0.0")]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AddrParseError(());

#[stable(feature = "addr_parse_error_error", since = "1.4.0")]
impl fmt::Display for AddrParseError {
    #[allow(deprecated, deprecated_in_future)]
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt.write_str(self.description())
    }
}

#[stable(feature = "addr_parse_error_error", since = "1.4.0")]
impl Error for AddrParseError {
    #[allow(deprecated)]
    fn description(&self) -> &str {
        "invalid IP address syntax"
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Error, ErrorKind, Write};

pub struct MultiLineWriter<W: Write> {
    inner: Option<W>,
    buf: Vec<u8>,
    last_newline: usize,
    panicked: bool,
}

impl<W: Write> MultiLineWriter<W> {
    pub fn new(inner: W) -> MultiLineWriter<W> {
        MultiLineWriter::with_capacity(1024, inner)
    }

    pub fn with_capacity(capacity: usize, inner: W) -> MultiLineWriter<W> {
        //MultiLineWriter { inner: BufWriter::with_capacity(capacity, inner), capacity, last_newline: usize, need_flush: false }
        MultiLineWriter {
            inner: Some(inner),
            buf: Vec::with_capacity(capacity),
            panicked: false,
            last_newline: 0,
        }
    }

    pub fn get_ref(&self) -> &W {
        self.inner.as_ref().unwrap()
    }

    pub fn get_mut(&mut self) -> &mut W {
        self.inner.as_mut().unwrap()
    }

    fn flush_buf(&mut self, flush_entire_buffer: bool) -> io::Result<()> {
        let mut written = 0;
        let len = if flush_entire_buffer {
            self.buf.len()
        } else {
            self.last_newline + 1
        };

        let mut ret = Ok(());
        while written < len {
            self.panicked = true;
            let r = self.inner.as_mut().unwrap().write(&self.buf[written..len]);
            self.panicked = false;

            match r {
                Ok(0) => {
                    ret = Err(Error::new(
                        ErrorKind::WriteZero,
                        "failed to write the buffered data",
                    ));
                    break;
                }
                Ok(n) => written += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => {
                    ret = Err(e);
                    break;
                }
            }
        }
        //println!("buf len {} last newline {}, written {}", self.buf.len(), self.last_newline, written);
        if written > 0 {
            //println!("FLUSHED: {}", ::std::str::from_utf8(&self.buf[..written]).unwrap());
            self.buf.drain(..written);
        }

        if flush_entire_buffer {
            self.last_newline = 0;
        } else if written > self.last_newline {
            self.last_newline = 0
        } else {
            self.last_newline -= written;
        }

        ret
    }
}

impl<W: Write> Write for MultiLineWriter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if self.buf.len() + buf.len() > self.buf.capacity() {
            self.flush_buf(false)?;
        }
        if buf.len() >= self.buf.capacity() {
            self.panicked = true;
            let r = self.get_mut().write(buf);
            self.panicked = false;
            r
        } else {
            if let Some(i) = memchr::memrchr(b'\n', buf) {
                self.last_newline = self.buf.len() + i;
            };

            self.buf.write(buf)
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        self.flush_buf(true).and_then(|()| self.get_mut().flush())
    }
}

impl<W: Write> Drop for MultiLineWriter<W> {
    fn drop(&mut self) {
        if self.inner.is_some() && !self.panicked {
            // dtors should not panic, so we ignore a failed flush
            let _r = self.flush_buf(true);
        }
    }
}

/*
impl<W: Write> Write for MultiLineWriter<W> {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
    let buffer_len = self.inner.buffer().len();

    let i = match memchr::memrchr(b'\n', buf) {
      Some(i) => i,
      None => buf.len(),
    };

    if buffer_len + i > self.capacity {
      self.inner.flush()?;
      self.last_newline = 0;
    }

  }
}
*/

/*
impl<W: Write> Drop for BufWriter<W> {
  fn drop(&mut self) {
    if self.inner.is_some() && !self.panicked {
      // dtors should not panic, so we ignore a failed flush
      let _r = self.flush_buf();
    }
  }
}
*/
<commit_msg>make the default multi line writer buffer larger<commit_after>use std::io::{self, Error, ErrorKind, Write};

pub struct MultiLineWriter<W: Write> {
    inner: Option<W>,
    buf: Vec<u8>,
    last_newline: usize,
    panicked: bool,
}

impl<W: Write> MultiLineWriter<W> {
    pub fn new(inner: W) -> MultiLineWriter<W> {
        MultiLineWriter::with_capacity(4096, inner)
    }

    pub fn with_capacity(capacity: usize, inner: W) -> MultiLineWriter<W> {
        //MultiLineWriter { inner: BufWriter::with_capacity(capacity, inner), capacity, last_newline: usize, need_flush: false }
        MultiLineWriter {
            inner: Some(inner),
            buf: Vec::with_capacity(capacity),
            panicked: false,
            last_newline: 0,
        }
    }

    pub fn get_ref(&self) -> &W {
        self.inner.as_ref().unwrap()
    }

    pub fn get_mut(&mut self) -> &mut W {
        self.inner.as_mut().unwrap()
    }

    fn flush_buf(&mut self, flush_entire_buffer: bool) -> io::Result<()> {
        let mut written = 0;
        let len = if flush_entire_buffer {
            self.buf.len()
        } else {
            self.last_newline + 1
        };

        let mut ret = Ok(());
        while written < len {
            self.panicked = true;
            let r = self.inner.as_mut().unwrap().write(&self.buf[written..len]);
            self.panicked = false;

            match r {
                Ok(0) => {
                    ret = Err(Error::new(
                        ErrorKind::WriteZero,
                        "failed to write the buffered data",
                    ));
                    break;
                }
                Ok(n) => written += n,
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
                Err(e) => {
                    ret = Err(e);
                    break;
                }
            }
        }
        //println!("buf len {} last newline {}, written {}", self.buf.len(), self.last_newline, written);
        if written > 0 {
            //println!("FLUSHED: {}", ::std::str::from_utf8(&self.buf[..written]).unwrap());
            self.buf.drain(..written);
        }

        if flush_entire_buffer {
            self.last_newline = 0;
        } else if written > self.last_newline {
            self.last_newline = 0
        } else {
            self.last_newline -= written;
        }

        ret
    }
}

impl<W: Write> Write for MultiLineWriter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        if self.buf.len() + buf.len() > self.buf.capacity() {
            self.flush_buf(false)?;
        }
        if buf.len() >= self.buf.capacity() {
            self.panicked = true;
            let r = self.get_mut().write(buf);
            self.panicked = false;
            r
        } else {
            if let Some(i) = memchr::memrchr(b'\n', buf) {
                self.last_newline = self.buf.len() + i;
            };

            self.buf.write(buf)
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        self.flush_buf(true).and_then(|()| self.get_mut().flush())
    }
}

impl<W: Write> Drop for MultiLineWriter<W> {
    fn drop(&mut self) {
        if self.inner.is_some() && !self.panicked {
            // dtors should not panic, so we ignore a failed flush
            let _r = self.flush_buf(true);
        }
    }
}

/*
impl<W: Write> Write for MultiLineWriter<W> {
  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
    let buffer_len = self.inner.buffer().len();

    let i = match memchr::memrchr(b'\n', buf) {
      Some(i) => i,
      None => buf.len(),
    };

    if buffer_len + i > self.capacity {
      self.inner.flush()?;
      self.last_newline = 0;
    }

  }
}
*/

/*
impl<W: Write> Drop for BufWriter<W> {
  fn drop(&mut self) {
    if self.inner.is_some() && !self.panicked {
      // dtors should not panic, so we ignore a failed flush
      let _r = self.flush_buf();
    }
  }
}
*/
<|endoftext|>


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

use exonum_merkledb::{
    access::{AsReadonly, Prefixed, RawAccess},
    Snapshot, SystemSchema,
};

use super::{DispatcherSchema, InstanceDescriptor, InstanceQuery, InstanceStatus};
use crate::blockchain::{IndexProof, Schema as CoreSchema};

/// Provides access to blockchain data for the executing service.
#[derive(Debug, Clone, Copy)]
pub struct BlockchainData<'a, T> {
    access: T,
    service_instance: InstanceDescriptor<'a>,
}

impl<'a, T: RawAccess + AsReadonly> BlockchainData<'a, T> {
    /// Creates structured access to blockchain data based on the unstructured access
    /// (e.g., a `Snapshot` or a `Fork`) and the descriptor of the executing service.
    pub fn new(access: T, service_instance: InstanceDescriptor<'a>) -> Self {
        Self {
            access,
            service_instance,
        }
    }

    /// Returns core schema.
    pub fn for_core(&self) -> CoreSchema<T::Readonly> {
        CoreSchema::new(self.access.as_readonly())
    }

    /// Returns dispatcher schema.
    pub fn for_dispatcher(&self) -> DispatcherSchema<T::Readonly> {
        DispatcherSchema::new(self.access.as_readonly())
    }

    /// Returns a mount point for another service. If the service with `id` does not exist,
    /// returns `None`.
    ///
    /// Note that this method does not check the service type; the caller is responsible
    /// for constructing a schema of a correct type around the returned access. Constructing
    /// an incorrect schema can lead to a panic or unexpected behavior.
    pub fn for_service<'q>(
        &self,
        id: impl Into<InstanceQuery<'q>>,
    ) -> Option<Prefixed<'static, T::Readonly>> {
        mount_point_for_service(self.access.as_readonly(), id)
    }

    /// Returns a mount point for the data of the executing service instance.
    /// Unlike other data, this one may be writeable provided that this `BlockchainData`
    /// wraps a `Fork`.
    pub fn for_executing_service(&self) -> Prefixed<'a, T> {
        Prefixed::new(self.service_instance.name, self.access.clone())
    }
}

impl BlockchainData<'_, &dyn Snapshot> {
    /// Returns a proof for a Merkelized index with the specified name
    /// in the currently executing service.
    ///
    /// # Return value
    ///
    /// If the index does not exist or is not Merkelized, returns `None`. Note that this may
    /// occur before the index is accessed the first time, which, depending on the service logic,
    /// may happen indefinitely after the service has been initialized. Thus, `unwrap`ping the
    /// returned value may unexpectedly lead to a panic unless the index is initialized early
    /// (e.g., during service initialization).
    pub fn proof_for_service_index(&self, index_name: &str) -> Option<IndexProof> {
        let core_schema = self.for_core();
        let height = core_schema.height();
        let block_proof = core_schema.block_and_precommits(height).unwrap();

        let full_index_name = [self.service_instance.name, ".", index_name].concat();
        let aggregator = SystemSchema::new(self.access).state_aggregator();
        aggregator.get(&full_index_name)?;
        let index_proof = aggregator.get_proof(full_index_name);
        Some(IndexProof {
            block_proof,
            index_proof,
        })
    }
}

fn mount_point_for_service<'q, T: RawAccess>(
    access: T,
    id: impl Into<InstanceQuery<'q>>,
) -> Option<Prefixed<'static, T>> {
    let state = DispatcherSchema::new(access.clone())
        .get_instance(id)
        .filter(|state| match (state.status, state.pending_status) {
            (Some(InstanceStatus::Active), _) | (None, Some(InstanceStatus::Active)) => true,
            _ => false,
        })?;
    Some(Prefixed::new(state.spec.name, access))
}

/// Extension trait for `Snapshot` allowing to access blockchain data in a more structured way.
pub trait SnapshotExt {
    /// Returns core schema.
    fn for_core(&self) -> CoreSchema<&'_ dyn Snapshot>;
    /// Returns dispatcher schema.
    fn for_dispatcher(&self) -> DispatcherSchema<&'_ dyn Snapshot>;
    /// Returns a mount point for a service. If the service does not exist, returns `None`.
    fn for_service<'q>(
        &self,
        id: impl Into<InstanceQuery<'q>>,
    ) -> Option<Prefixed<'static, &dyn Snapshot>>;
}

impl SnapshotExt for dyn Snapshot {
    fn for_core(&self) -> CoreSchema<&'_ dyn Snapshot> {
        CoreSchema::new(self)
    }

    fn for_dispatcher(&self) -> DispatcherSchema<&'_ dyn Snapshot> {
        DispatcherSchema::new(self)
    }

    fn for_service<'q>(
        &self,
        id: impl Into<InstanceQuery<'q>>,
    ) -> Option<Prefixed<'static, &dyn Snapshot>> {
        mount_point_for_service(self, id)
    }
}
<commit_msg>Allow to return `IndexProof` by full index name [ECR-4013] (#1635)<commit_after>// Copyright 2019 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use exonum_merkledb::{
    access::{AsReadonly, Prefixed, RawAccess},
    Snapshot, SystemSchema,
};

use super::{DispatcherSchema, InstanceDescriptor, InstanceQuery, InstanceStatus};
use crate::blockchain::{IndexProof, Schema as CoreSchema};

/// Provides access to blockchain data for the executing service.
#[derive(Debug, Clone, Copy)]
pub struct BlockchainData<'a, T> {
    access: T,
    service_instance: InstanceDescriptor<'a>,
}

impl<'a, T: RawAccess + AsReadonly> BlockchainData<'a, T> {
    /// Creates structured access to blockchain data based on the unstructured access
    /// (e.g., a `Snapshot` or a `Fork`) and the descriptor of the executing service.
    pub fn new(access: T, service_instance: InstanceDescriptor<'a>) -> Self {
        Self {
            access,
            service_instance,
        }
    }

    /// Returns core schema.
    pub fn for_core(&self) -> CoreSchema<T::Readonly> {
        CoreSchema::new(self.access.as_readonly())
    }

    /// Returns dispatcher schema.
    pub fn for_dispatcher(&self) -> DispatcherSchema<T::Readonly> {
        DispatcherSchema::new(self.access.as_readonly())
    }

    /// Returns a mount point for another service. If the service with `id` does not exist,
    /// returns `None`.
    ///
    /// Note that this method does not check the service type; the caller is responsible
    /// for constructing a schema of a correct type around the returned access. Constructing
    /// an incorrect schema can lead to a panic or unexpected behavior.
    pub fn for_service<'q>(
        &self,
        id: impl Into<InstanceQuery<'q>>,
    ) -> Option<Prefixed<'static, T::Readonly>> {
        mount_point_for_service(self.access.as_readonly(), id)
    }

    /// Returns a mount point for the data of the executing service instance.
    /// Unlike other data, this one may be writeable provided that this `BlockchainData`
    /// wraps a `Fork`.
    pub fn for_executing_service(&self) -> Prefixed<'a, T> {
        Prefixed::new(self.service_instance.name, self.access.clone())
    }
}

impl BlockchainData<'_, &dyn Snapshot> {
    /// Returns a proof for a Merkelized index with the specified name
    /// in the currently executing service.
    ///
    /// # Return value
    ///
    /// If the index does not exist or is not Merkelized, returns `None`. Note that this may
    /// occur before the index is accessed the first time, which, depending on the service logic,
    /// may happen indefinitely after the service has been initialized. Thus, `unwrap`ping the
    /// returned value may unexpectedly lead to a panic unless the index is initialized early
    /// (e.g., during service initialization).
    pub fn proof_for_service_index(&self, index_name: &str) -> Option<IndexProof> {
        let full_index_name = [self.service_instance.name, ".", index_name].concat();
        self.access.proof_for_index(&full_index_name)
    }
}

fn mount_point_for_service<'q, T: RawAccess>(
    access: T,
    id: impl Into<InstanceQuery<'q>>,
) -> Option<Prefixed<'static, T>> {
    let state = DispatcherSchema::new(access.clone())
        .get_instance(id)
        .filter(|state| match (state.status, state.pending_status) {
            (Some(InstanceStatus::Active), _) | (None, Some(InstanceStatus::Active)) => true,
            _ => false,
        })?;
    Some(Prefixed::new(state.spec.name, access))
}

/// Extension trait for `Snapshot` allowing to access blockchain data in a more structured way.
pub trait SnapshotExt {
    /// Returns core schema.
    fn for_core(&self) -> CoreSchema<&'_ dyn Snapshot>;
    /// Returns dispatcher schema.
    fn for_dispatcher(&self) -> DispatcherSchema<&'_ dyn Snapshot>;
    /// Returns a mount point for a service. If the service does not exist, returns `None`.
    fn for_service<'q>(
        &self,
        id: impl Into<InstanceQuery<'q>>,
    ) -> Option<Prefixed<'static, &dyn Snapshot>>;

    /// Returns a proof for a Merkelized index with the specified name.
    ///
    /// # Return value
    ///
    /// If the index does not exist or is not Merkelized, returns `None`. Note that this may
    /// occur before the index is accessed the first time, which, depending on the service logic,
    /// may happen indefinitely after the service has been initialized. Thus, `unwrap`ping the
    /// returned value may unexpectedly lead to a panic unless the index is initialized early
    /// (e.g., during service initialization).
    #[doc(hidden)]
    fn proof_for_index(&self, index_name: &str) -> Option<IndexProof>;
}

impl SnapshotExt for dyn Snapshot {
    fn for_core(&self) -> CoreSchema<&'_ dyn Snapshot> {
        CoreSchema::new(self)
    }

    fn for_dispatcher(&self) -> DispatcherSchema<&'_ dyn Snapshot> {
        DispatcherSchema::new(self)
    }

    fn for_service<'q>(
        &self,
        id: impl Into<InstanceQuery<'q>>,
    ) -> Option<Prefixed<'static, &dyn Snapshot>> {
        mount_point_for_service(self, id)
    }

    fn proof_for_index(&self, index_name: &str) -> Option<IndexProof> {
        let core_schema = self.for_core();
        let height = core_schema.height();
        let block_proof = core_schema.block_and_precommits(height).unwrap();

        let aggregator = SystemSchema::new(self).state_aggregator();
        aggregator.get(index_name)?;
        let index_proof = aggregator.get_proof(index_name.to_string());
        Some(IndexProof {
            block_proof,
            index_proof,
        })
    }
}
<|endoftext|>


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

#[macro_use]
extern crate log;
extern crate chrono;
extern crate fern;

extern crate dirs;

extern crate xi_core_lib;
extern crate xi_rpc;

use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

use xi_core_lib::XiCore;
use xi_rpc::RpcLoop;

const XI_LOG_DIR: &str = "xi-core";
const XI_LOG_FILE: &str = "xi-core.log";

fn get_logging_directory_path<P: AsRef<Path>>(directory: P) -> Result<PathBuf, io::Error> {
    match dirs::data_local_dir() {
        Some(mut log_dir) => {
            log_dir.push(directory);
            Ok(log_dir)
        }
        None => Err(io::Error::new(
            io::ErrorKind::NotFound,
            "No standard logging directory known for this platform",
        )),
    }
}

fn setup_logging(logging_path_result: Result<PathBuf, io::Error>) -> Result<(), fern::InitError> {
    let level_filter = match std::env::var("XI_LOG") {
        Ok(level) => match level.to_lowercase().as_ref() {
            "trace" => log::LevelFilter::Trace,
            "debug" => log::LevelFilter::Debug,
            _ => log::LevelFilter::Info,
        },
        // Default to info
        Err(_) => log::LevelFilter::Info,
    };

    let mut fern_dispatch = fern::Dispatch::new()
        .format(|out, message, record| {
            out.finish(format_args!(
                "{}[{}][{}] {}",
                chrono::Local::now().format("[%Y-%m-%d][%H:%M:%S]"),
                record.target(),
                record.level(),
                message,
            ))
        }).level(level_filter)
        .chain(io::stderr());

    if let Ok(logging_file_path) = &logging_path_result {
        // Ensure the logging directory is created
        let parent_path = logging_file_path.parent().ok_or(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!(
                "Unable to get the parent of the following Path: {}",
                logging_file_path.display(),
            ),
        ))?;
        fs::create_dir_all(parent_path)?;
        // Attach it to fern
        fern_dispatch = fern_dispatch.chain(fern::log_file(logging_file_path)?);
    };

    // Start fern
    fern_dispatch.apply()?;
    info!("Logging with fern is setup");

    // Log details of the logging_file_path result using fern/log
    // Either logging the path fern is outputting to or the error from obtaining the path
    match &logging_path_result {
        Ok(logging_file_path) => {
            info!(
                "Logging to the following file: {}",
                logging_file_path.display()
            );
        }
        Err(e) => {
            let message = "There was an issue getting the path for the log file";
            warn!("{}: {:?}, falling back to stderr.", message, e);
        }
    }
    Ok(())
}

fn prepare_logging_path(logfile_config: LogfileConfig) -> Result<PathBuf, io::Error> {
    // Use the file name set in logfile_config or fallback to the default
    let logfile_file_name = match logfile_config.file {
        Some(file_name) => file_name,
        None => PathBuf::from(XI_LOG_FILE),
    };
    // Use the directory name set in logfile_config or fallback to the default
    let logfile_directory_name = match logfile_config.directory {
        Some(dir) => dir,
        None => PathBuf::from(XI_LOG_DIR),
    };

    let mut logging_directory_path = get_logging_directory_path(logfile_directory_name)?;

    // Add the file name & return the full path
    logging_directory_path.push(logfile_file_name);
    Ok(logging_directory_path)
}

fn get_flags() -> HashMap<String, Option<String>> {
    let mut flags: HashMap<String, Option<String>> = HashMap::new();

    let flag_prefix = "-";
    let mut args_itterator = std::env::args().peekable();
    while let Some(arg) = args_itterator.next() {
        if arg.starts_with(flag_prefix) {
            let key = arg.trim_left_matches(flag_prefix).to_string();

            // Check the next argument doesn't start with the flag prefix
            // map_or accounts for peek returning an Option
            let next_arg_not_a_flag: bool = args_itterator
                .peek()
                .map_or(false, |val| !val.starts_with(flag_prefix));
            if next_arg_not_a_flag {
                flags.insert(key, args_itterator.next());
            }
        }
    }
    flags
}

struct EnvFlagConfig {
    env_name: &'static str,
    flag_name: &'static str,
}

fn extract_env_or_flag(flags: &HashMap<String, Option<String>>, conf: EnvFlagConfig) -> Option<String> {
    std::env::var(conf.env_name)
        .ok()
        .or(flags.get(conf.flag_name).cloned().unwrap_or(None))
}

struct LogfileConfig {
    directory: Option<PathBuf>,
    file: Option<PathBuf>,
}

fn generate_logfile_config(flags: &HashMap<String, Option<String>>) -> LogfileConfig {
    // If the key is set, get the Option within
    let log_dir_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_DIR",
        flag_name: "log-dir",
    };
    let log_file_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_FILE",
        flag_name: "log-file",
    };
    let log_dir_flag_option = extract_env_or_flag(&flags, log_dir_env_flag).map(PathBuf::from);

    let log_file_flag_option = extract_env_or_flag(&flags, log_file_env_flag).map(PathBuf::from);

    LogfileConfig {
        directory: log_dir_flag_option,
        file: log_file_flag_option,
    }
}

fn main() {
    let mut state = XiCore::new();
    let stdin = io::stdin();
    let stdout = io::stdout();
    let mut rpc_looper = RpcLoop::new(stdout);

    let flags: HashMap<String, Option<String>> = get_flags();

    let logfile_config = generate_logfile_config(&flags);

    let logging_path: Result<PathBuf, io::Error> = prepare_logging_path(logfile_config);
    if let Err(e) = setup_logging(logging_path) {
        eprintln!(
            "[ERROR] setup_logging returned error, logging disabled: {:?}",
            e
        );
    }

    match rpc_looper.mainloop(|| stdin.lock(), &mut state) {
        Ok(_) => (),
        Err(err) => error!("xi-core exited with error:\n{:?}", err),
    }
}
<commit_msg>Swapped Result for Option in setup_logging to avoid passing it the Error.<commit_after>// Copyright 2016 The xi-editor Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#[macro_use]
extern crate log;
extern crate chrono;
extern crate fern;

extern crate dirs;

extern crate xi_core_lib;
extern crate xi_rpc;

use std::collections::HashMap;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

use xi_core_lib::XiCore;
use xi_rpc::RpcLoop;

const XI_LOG_DIR: &str = "xi-core";
const XI_LOG_FILE: &str = "xi-core.log";

fn get_logging_directory_path<P: AsRef<Path>>(directory: P) -> Result<PathBuf, io::Error> {
    match dirs::data_local_dir() {
        Some(mut log_dir) => {
            log_dir.push(directory);
            Ok(log_dir)
        }
        None => Err(io::Error::new(
            io::ErrorKind::NotFound,
            "No standard logging directory known for this platform",
        )),
    }
}

fn setup_logging(logging_path_result: Option<PathBuf>) -> Result<(), fern::InitError> {
    let level_filter = match std::env::var("XI_LOG") {
        Ok(level) => match level.to_lowercase().as_ref() {
            "trace" => log::LevelFilter::Trace,
            "debug" => log::LevelFilter::Debug,
            _ => log::LevelFilter::Info,
        },
        // Default to info
        Err(_) => log::LevelFilter::Info,
    };

    let mut fern_dispatch = fern::Dispatch::new()
        .format(|out, message, record| {
            out.finish(format_args!(
                "{}[{}][{}] {}",
                chrono::Local::now().format("[%Y-%m-%d][%H:%M:%S]"),
                record.target(),
                record.level(),
                message,
            ))
        }).level(level_filter)
        .chain(io::stderr());

    if let Some(logging_file_path) = &logging_path_result {
        // Ensure the logging directory is created
        let parent_path = logging_file_path.parent().ok_or(io::Error::new(
            io::ErrorKind::InvalidInput,
            format!(
                "Unable to get the parent of the following Path: {}",
                logging_file_path.display(),
            ),
        ))?;
        fs::create_dir_all(parent_path)?;
        // Attach it to fern
        fern_dispatch = fern_dispatch.chain(fern::log_file(logging_file_path)?);
    };

    // Start fern
    fern_dispatch.apply()?;
    info!("Logging with fern is setup");

    // Log details of the logging_file_path result using fern/log
    // Either logging the path fern is outputting to or the error from obtaining the path
    match &logging_path_result {
        Some(logging_file_path) => info!(
            "Logging to the following file: {}",
            logging_file_path.display()
        ),
        None => warn!("There was no path supplied for the log file, falling back to stderr."),
    }
    Ok(())
}

fn generate_logging_path(logfile_config: LogfileConfig) -> Result<PathBuf, io::Error> {
    // Use the file name set in logfile_config or fallback to the default
    let logfile_file_name = match logfile_config.file {
        Some(file_name) => file_name,
        None => PathBuf::from(XI_LOG_FILE),
    };
    // Use the directory name set in logfile_config or fallback to the default
    let logfile_directory_name = match logfile_config.directory {
        Some(dir) => dir,
        None => PathBuf::from(XI_LOG_DIR),
    };

    let mut logging_directory_path = get_logging_directory_path(logfile_directory_name)?;

    // Add the file name & return the full path
    logging_directory_path.push(logfile_file_name);
    Ok(logging_directory_path)
}

fn get_flags() -> HashMap<String, Option<String>> {
    let mut flags: HashMap<String, Option<String>> = HashMap::new();

    let flag_prefix = "-";
    let mut args_itterator = std::env::args().peekable();
    while let Some(arg) = args_itterator.next() {
        if arg.starts_with(flag_prefix) {
            let key = arg.trim_left_matches(flag_prefix).to_string();

            // Check the next argument doesn't start with the flag prefix
            // map_or accounts for peek returning an Option
            let next_arg_not_a_flag: bool = args_itterator
                .peek()
                .map_or(false, |val| !val.starts_with(flag_prefix));
            if next_arg_not_a_flag {
                flags.insert(key, args_itterator.next());
            }
        }
    }
    flags
}

struct EnvFlagConfig {
    env_name: &'static str,
    flag_name: &'static str,
}

fn extract_env_or_flag(
    flags: &HashMap<String, Option<String>>,
    conf: EnvFlagConfig,
) -> Option<String> {
    std::env::var(conf.env_name)
        .ok()
        .or(flags.get(conf.flag_name).cloned().unwrap_or(None))
}

struct LogfileConfig {
    directory: Option<PathBuf>,
    file: Option<PathBuf>,
}

fn generate_logfile_config(flags: &HashMap<String, Option<String>>) -> LogfileConfig {
    // If the key is set, get the Option within
    let log_dir_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_DIR",
        flag_name: "log-dir",
    };
    let log_file_env_flag = EnvFlagConfig {
        env_name: "XI_LOG_FILE",
        flag_name: "log-file",
    };
    let log_dir_flag_option = extract_env_or_flag(&flags, log_dir_env_flag).map(PathBuf::from);

    let log_file_flag_option = extract_env_or_flag(&flags, log_file_env_flag).map(PathBuf::from);

    LogfileConfig {
        directory: log_dir_flag_option,
        file: log_file_flag_option,
    }
}

fn main() {
    let mut state = XiCore::new();
    let stdin = io::stdin();
    let stdout = io::stdout();
    let mut rpc_looper = RpcLoop::new(stdout);

    let flags: HashMap<String, Option<String>> = get_flags();

    let logfile_config = generate_logfile_config(&flags);

    let logging_path_result: Result<PathBuf, io::Error> = generate_logging_path(logfile_config);
    let logging_path: Option<PathBuf> = match logging_path_result {
        Ok(val) => Some(val),
        Err(e) => {
            eprintln!(
                "[WARNING] Unable to successfully generate the logging path: {:?}",
                e
            );
            None
        }
    };
    if let Err(e) = setup_logging(logging_path) {
        eprintln!(
            "[ERROR] setup_logging returned error, logging disabled: {:?}",
            e
        );
    }

    match rpc_looper.mainloop(|| stdin.lock(), &mut state) {
        Ok(_) => (),
        Err(err) => error!("xi-core exited with error:\n{:?}", err),
    }
}
<|endoftext|>


<commit_before>// Copyright 2017-2019 int08h LLC

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

// for value_t_or_exit!()
#[macro_use]
extern crate clap;

use ring::rand;
use ring::rand::SecureRandom;

use byteorder::{LittleEndian, ReadBytesExt};

use chrono::offset::Utc;
use chrono::TimeZone;

use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::iter::Iterator;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};

use clap::{App, Arg};
use roughenough::merkle::root_from_paths;
use roughenough::sign::Verifier;
use roughenough::{
    roughenough_version, RtMessage, Tag, CERTIFICATE_CONTEXT, SIGNED_RESPONSE_CONTEXT,
};

fn create_nonce() -> [u8; 64] {
    let rng = rand::SystemRandom::new();
    let mut nonce = [0u8; 64];
    rng.fill(&mut nonce).unwrap();

    nonce
}

fn make_request(nonce: &[u8]) -> Vec<u8> {
    let mut msg = RtMessage::new(1);
    msg.add_field(Tag::NONC, nonce).unwrap();
    msg.pad_to_kilobyte();

    msg.encode().unwrap()
}

fn receive_response(sock: &mut UdpSocket) -> RtMessage {
    let mut buf = [0; 744];
    let resp_len = sock.recv_from(&mut buf).unwrap().0;

    RtMessage::from_bytes(&buf[0..resp_len]).unwrap()
}

fn stress_test_forever(addr: &SocketAddr) -> ! {
    if !addr.ip().is_loopback() {
        panic!("Cannot use non-loopback address {} for stress testing", addr.ip());
    }

    println!("Stress testing!");

    let nonce = create_nonce();
    let socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
    let request = make_request(&nonce);
    loop {
        socket.send_to(&request, addr).unwrap();
    }
}

struct ResponseHandler {
    pub_key: Option<Vec<u8>>,
    msg: HashMap<Tag, Vec<u8>>,
    srep: HashMap<Tag, Vec<u8>>,
    cert: HashMap<Tag, Vec<u8>>,
    dele: HashMap<Tag, Vec<u8>>,
    nonce: [u8; 64],
}

struct ParsedResponse {
    verified: bool,
    midpoint: u64,
    radius: u32,
}

impl ResponseHandler {
    pub fn new(pub_key: Option<Vec<u8>>, response: RtMessage, nonce: [u8; 64]) -> ResponseHandler {
        let msg = response.into_hash_map();
        let srep = RtMessage::from_bytes(&msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let cert = RtMessage::from_bytes(&msg[&Tag::CERT])
            .unwrap()
            .into_hash_map();
        let dele = RtMessage::from_bytes(&cert[&Tag::DELE])
            .unwrap()
            .into_hash_map();

        ResponseHandler {
            pub_key,
            msg,
            srep,
            cert,
            dele,
            nonce,
        }
    }

    pub fn extract_time(&self) -> ParsedResponse {
        let midpoint = self.srep[&Tag::MIDP]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let radius = self.srep[&Tag::RADI]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();

        let verified = if self.pub_key.is_some() {
            self.validate_dele();
            self.validate_srep();
            self.validate_merkle();
            self.validate_midpoint(midpoint);
            true
        } else {
            false
        };

        ParsedResponse {
            verified,
            midpoint,
            radius,
        }
    }

    fn validate_dele(&self) {
        let mut full_cert = Vec::from(CERTIFICATE_CONTEXT.as_bytes());
        full_cert.extend(&self.cert[&Tag::DELE]);

        assert!(
            self.validate_sig(
                self.pub_key.as_ref().unwrap(),
                &self.cert[&Tag::SIG],
                &full_cert
            ),
            "Invalid signature on DELE tag, response may not be authentic"
        );
    }

    fn validate_srep(&self) {
        let mut full_srep = Vec::from(SIGNED_RESPONSE_CONTEXT.as_bytes());
        full_srep.extend(&self.msg[&Tag::SREP]);

        assert!(
            self.validate_sig(&self.dele[&Tag::PUBK], &self.msg[&Tag::SIG], &full_srep),
            "Invalid signature on SREP tag, response may not be authentic"
        );
    }

    fn validate_merkle(&self) {
        let srep = RtMessage::from_bytes(&self.msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let index = self.msg[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();
        let paths = &self.msg[&Tag::PATH];

        let hash = root_from_paths(index as usize, &self.nonce, paths);

        assert_eq!(
            hash, srep[&Tag::ROOT],
            "Nonce is not present in the response's merkle tree"
        );
    }

    fn validate_midpoint(&self, midpoint: u64) {
        let mint = self.dele[&Tag::MINT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let maxt = self.dele[&Tag::MAXT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();

        assert!(
            midpoint >= mint,
            "Response midpoint {} lies *before* delegation span ({}, {})",
            midpoint, mint, maxt
        );
        assert!(
            midpoint <= maxt,
            "Response midpoint {} lies *after* delegation span ({}, {})",
            midpoint, mint, maxt
        );
    }

    fn validate_sig(&self, public_key: &[u8], sig: &[u8], data: &[u8]) -> bool {
        let mut verifier = Verifier::new(public_key);
        verifier.update(data);
        verifier.verify(sig)
    }
}

fn main() {
    let matches = App::new("roughenough client")
    .version(roughenough_version().as_ref())
    .arg(Arg::with_name("host")
      .required(true)
      .help("The Roughtime server to connect to")
      .takes_value(true))
    .arg(Arg::with_name("port")
      .required(true)
      .help("The Roughtime server port to connect to")
      .takes_value(true))
    .arg(Arg::with_name("verbose")
      .short("v")
      .long("verbose")
      .help("Print more output"))
    .arg(Arg::with_name("json")
      .short("j")
      .long("json")
      .help("Print output in JSON"))
    .arg(Arg::with_name("public-key")
      .short("p")
      .long("public-key")
      .takes_value(true)
      .help("The server public key used to validate responses. If unset, no validation will be performed"))
    .arg(Arg::with_name("time-format")
      .short("f")
      .long("time-format")
      .takes_value(true)
      .help("The strftime format string used to print the time recieved from the server")
      .default_value("%b %d %Y %H:%M:%S")
    )
    .arg(Arg::with_name("num-requests")
      .short("n")
      .long("num-requests")
      .takes_value(true)
      .help("The number of requests to make to the server (each from a different source port). This is mainly useful for testing batch response handling")
      .default_value("1")
    )
    .arg(Arg::with_name("stress")
      .short("s")
      .long("stress")
      .help("Stress-tests the server by sending the same request as fast as possible. Please only use this on your own server")
    )
    .arg(Arg::with_name("output")
      .short("o")
      .long("output")
      .takes_value(true)
      .help("Writes all requests to the specified file, in addition to sending them to the server. Useful for generating fuzzer inputs")
    )
    .get_matches();

    let host = matches.value_of("host").unwrap();
    let port = value_t_or_exit!(matches.value_of("port"), u16);
    let verbose = matches.is_present("verbose");
    let json = matches.is_present("json");
    let num_requests = value_t_or_exit!(matches.value_of("num-requests"), u16) as usize;
    let time_format = matches.value_of("time-format").unwrap();
    let stress = matches.is_present("stress");
    let pub_key = matches
        .value_of("public-key")
        .map(|pkey| hex::decode(pkey).expect("Error parsing public key!"));
    let out = matches.value_of("output");

    if verbose {
        eprintln!("Requesting time from: {:?}:{:?}", host, port);
    }

    let addr = (host, port).to_socket_addrs().unwrap().next().unwrap();

    if stress {
        stress_test_forever(&addr)
    }

    let mut requests = Vec::with_capacity(num_requests);
    let mut file = out.map(|o| File::create(o).expect("Failed to create file!"));

    for _ in 0..num_requests {
        let nonce = create_nonce();
        let socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
        let request = make_request(&nonce);

        if let Some(f) = file.as_mut() {
            f.write_all(&request).expect("Failed to write to file!")
        }

        requests.push((nonce, request, socket));
    }

    for &mut (_, ref request, ref mut socket) in &mut requests {
        socket.send_to(request, addr).unwrap();
    }

    for (nonce, _, mut socket) in requests {
        let resp = receive_response(&mut socket);

        let ParsedResponse {
            verified,
            midpoint,
            radius,
        } = ResponseHandler::new(pub_key.clone(), resp.clone(), nonce).extract_time();

        let map = resp.into_hash_map();
        let index = map[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();

        let seconds = midpoint / 10_u64.pow(6);
        let nsecs = (midpoint - (seconds * 10_u64.pow(6))) * 10_u64.pow(3);
        let spec = Utc.timestamp(seconds as i64, nsecs as u32);
        let out = spec.format(time_format).to_string();
        let verify_str = if verified { "Yes" } else { "No" };

        if verbose {
            eprintln!(
                "Received time from server: midpoint={:?}, radius={:?}, verified={} (merkle_index={})",
                out, radius, verify_str, index
            );
        }

        if json {
            println!(
                r#"{{ "midpoint": {:?}, "radius": {:?}, "verified": {}, "merkle_index": {} }}"#,
                out, radius, verified, index
            );
        } else {
            println!("{}", out);
        }
    }
}

<commit_msg>Updates to the client's CLI --help text<commit_after>// Copyright 2017-2019 int08h LLC

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

// for value_t_or_exit!()
#[macro_use]
extern crate clap;

use ring::rand;
use ring::rand::SecureRandom;

use byteorder::{LittleEndian, ReadBytesExt};

use chrono::offset::Utc;
use chrono::TimeZone;

use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::iter::Iterator;
use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};

use clap::{App, Arg};
use roughenough::merkle::root_from_paths;
use roughenough::sign::Verifier;
use roughenough::{
    roughenough_version, RtMessage, Tag, CERTIFICATE_CONTEXT, SIGNED_RESPONSE_CONTEXT,
};

fn create_nonce() -> [u8; 64] {
    let rng = rand::SystemRandom::new();
    let mut nonce = [0u8; 64];
    rng.fill(&mut nonce).unwrap();

    nonce
}

fn make_request(nonce: &[u8]) -> Vec<u8> {
    let mut msg = RtMessage::new(1);
    msg.add_field(Tag::NONC, nonce).unwrap();
    msg.pad_to_kilobyte();

    msg.encode().unwrap()
}

fn receive_response(sock: &mut UdpSocket) -> RtMessage {
    let mut buf = [0; 744];
    let resp_len = sock.recv_from(&mut buf).unwrap().0;

    RtMessage::from_bytes(&buf[0..resp_len]).unwrap()
}

fn stress_test_forever(addr: &SocketAddr) -> ! {
    if !addr.ip().is_loopback() {
        panic!("Cannot use non-loopback address {} for stress testing", addr.ip());
    }

    println!("Stress testing!");

    let nonce = create_nonce();
    let socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
    let request = make_request(&nonce);
    loop {
        socket.send_to(&request, addr).unwrap();
    }
}

struct ResponseHandler {
    pub_key: Option<Vec<u8>>,
    msg: HashMap<Tag, Vec<u8>>,
    srep: HashMap<Tag, Vec<u8>>,
    cert: HashMap<Tag, Vec<u8>>,
    dele: HashMap<Tag, Vec<u8>>,
    nonce: [u8; 64],
}

struct ParsedResponse {
    verified: bool,
    midpoint: u64,
    radius: u32,
}

impl ResponseHandler {
    pub fn new(pub_key: Option<Vec<u8>>, response: RtMessage, nonce: [u8; 64]) -> ResponseHandler {
        let msg = response.into_hash_map();
        let srep = RtMessage::from_bytes(&msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let cert = RtMessage::from_bytes(&msg[&Tag::CERT])
            .unwrap()
            .into_hash_map();
        let dele = RtMessage::from_bytes(&cert[&Tag::DELE])
            .unwrap()
            .into_hash_map();

        ResponseHandler {
            pub_key,
            msg,
            srep,
            cert,
            dele,
            nonce,
        }
    }

    pub fn extract_time(&self) -> ParsedResponse {
        let midpoint = self.srep[&Tag::MIDP]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let radius = self.srep[&Tag::RADI]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();

        let verified = if self.pub_key.is_some() {
            self.validate_dele();
            self.validate_srep();
            self.validate_merkle();
            self.validate_midpoint(midpoint);
            true
        } else {
            false
        };

        ParsedResponse {
            verified,
            midpoint,
            radius,
        }
    }

    fn validate_dele(&self) {
        let mut full_cert = Vec::from(CERTIFICATE_CONTEXT.as_bytes());
        full_cert.extend(&self.cert[&Tag::DELE]);

        assert!(
            self.validate_sig(
                self.pub_key.as_ref().unwrap(),
                &self.cert[&Tag::SIG],
                &full_cert
            ),
            "Invalid signature on DELE tag, response may not be authentic"
        );
    }

    fn validate_srep(&self) {
        let mut full_srep = Vec::from(SIGNED_RESPONSE_CONTEXT.as_bytes());
        full_srep.extend(&self.msg[&Tag::SREP]);

        assert!(
            self.validate_sig(&self.dele[&Tag::PUBK], &self.msg[&Tag::SIG], &full_srep),
            "Invalid signature on SREP tag, response may not be authentic"
        );
    }

    fn validate_merkle(&self) {
        let srep = RtMessage::from_bytes(&self.msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let index = self.msg[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();
        let paths = &self.msg[&Tag::PATH];

        let hash = root_from_paths(index as usize, &self.nonce, paths);

        assert_eq!(
            hash, srep[&Tag::ROOT],
            "Nonce is not present in the response's merkle tree"
        );
    }

    fn validate_midpoint(&self, midpoint: u64) {
        let mint = self.dele[&Tag::MINT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let maxt = self.dele[&Tag::MAXT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();

        assert!(
            midpoint >= mint,
            "Response midpoint {} lies *before* delegation span ({}, {})",
            midpoint, mint, maxt
        );
        assert!(
            midpoint <= maxt,
            "Response midpoint {} lies *after* delegation span ({}, {})",
            midpoint, mint, maxt
        );
    }

    fn validate_sig(&self, public_key: &[u8], sig: &[u8], data: &[u8]) -> bool {
        let mut verifier = Verifier::new(public_key);
        verifier.update(data);
        verifier.verify(sig)
    }
}

fn main() {
    let matches = App::new("roughenough client")
    .version(roughenough_version().as_ref())
    .arg(Arg::with_name("host")
      .required(true)
      .help("The Roughtime server to connect to.")
      .takes_value(true))
    .arg(Arg::with_name("port")
      .required(true)
      .help("The Roughtime server port to connect to.")
      .takes_value(true))
    .arg(Arg::with_name("verbose")
      .short("v")
      .long("verbose")
      .help("Output additional details about the server's response."))
    .arg(Arg::with_name("json")
      .short("j")
      .long("json")
      .help("Output the server's response in JSON format."))
    .arg(Arg::with_name("public-key")
      .short("p")
      .long("public-key")
      .takes_value(true)
      .help("The server public key used to validate responses. If unset, no validation will be performed."))
    .arg(Arg::with_name("time-format")
      .short("f")
      .long("time-format")
      .takes_value(true)
      .help("The strftime format string used to print the time recieved from the server.")
      .default_value("%b %d %Y %H:%M:%S")
    )
    .arg(Arg::with_name("num-requests")
      .short("n")
      .long("num-requests")
      .takes_value(true)
      .help("The number of requests to make to the server (each from a different source port). This is mainly useful for testing batch response handling.")
      .default_value("1")
    )
    .arg(Arg::with_name("stress")
      .short("s")
      .long("stress")
      .help("Stress test the server by sending the same request as fast as possible. Please only use this on your own server.")
    )
    .arg(Arg::with_name("output")
      .short("o")
      .long("output")
      .takes_value(true)
      .help("Writes all requests to the specified file, in addition to sending them to the server. Useful for generating fuzzer inputs.")
    )
    .get_matches();

    let host = matches.value_of("host").unwrap();
    let port = value_t_or_exit!(matches.value_of("port"), u16);
    let verbose = matches.is_present("verbose");
    let json = matches.is_present("json");
    let num_requests = value_t_or_exit!(matches.value_of("num-requests"), u16) as usize;
    let time_format = matches.value_of("time-format").unwrap();
    let stress = matches.is_present("stress");
    let pub_key = matches
        .value_of("public-key")
        .map(|pkey| hex::decode(pkey).expect("Error parsing public key!"));
    let out = matches.value_of("output");

    if verbose {
        eprintln!("Requesting time from: {:?}:{:?}", host, port);
    }

    let addr = (host, port).to_socket_addrs().unwrap().next().unwrap();

    if stress {
        stress_test_forever(&addr)
    }

    let mut requests = Vec::with_capacity(num_requests);
    let mut file = out.map(|o| File::create(o).expect("Failed to create file!"));

    for _ in 0..num_requests {
        let nonce = create_nonce();
        let socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
        let request = make_request(&nonce);

        if let Some(f) = file.as_mut() {
            f.write_all(&request).expect("Failed to write to file!")
        }

        requests.push((nonce, request, socket));
    }

    for &mut (_, ref request, ref mut socket) in &mut requests {
        socket.send_to(request, addr).unwrap();
    }

    for (nonce, _, mut socket) in requests {
        let resp = receive_response(&mut socket);

        let ParsedResponse {
            verified,
            midpoint,
            radius,
        } = ResponseHandler::new(pub_key.clone(), resp.clone(), nonce).extract_time();

        let map = resp.into_hash_map();
        let index = map[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();

        let seconds = midpoint / 10_u64.pow(6);
        let nsecs = (midpoint - (seconds * 10_u64.pow(6))) * 10_u64.pow(3);
        let spec = Utc.timestamp(seconds as i64, nsecs as u32);
        let out = spec.format(time_format).to_string();
        let verify_str = if verified { "Yes" } else { "No" };

        if verbose {
            eprintln!(
                "Received time from server: midpoint={:?}, radius={:?}, verified={} (merkle_index={})",
                out, radius, verify_str, index
            );
        }

        if json {
            println!(
                r#"{{ "midpoint": {:?}, "radius": {:?}, "verified": {}, "merkle_index": {} }}"#,
                out, radius, verified, index
            );
        } else {
            println!("{}", out);
        }
    }
}

<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
// Copyright 2016-2017 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The binomial distribution.

use Rng;
use distributions::{Distribution, Cauchy};
use distributions::utils::log_gamma;

/// The binomial distribution `Binomial(n, p)`.
///
/// This distribution has density function:
/// `f(k) = n!/(k! (n-k)!) p^k (1-p)^(n-k)` for `k >= 0`.
///
/// # Example
///
/// ```
/// use rand::distributions::{Binomial, Distribution};
///
/// let bin = Binomial::new(20, 0.3);
/// let v = bin.sample(&mut rand::thread_rng());
/// println!("{} is from a binomial distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Binomial {
    /// Number of trials.
    n: u64,
    /// Probability of success.
    p: f64,
}

impl Binomial {
    /// Construct a new `Binomial` with the given shape parameters `n` (number
    /// of trials) and `p` (probability of success).
    ///
    /// Panics if `p < 0` or `p > 1`.
    pub fn new(n: u64, p: f64) -> Binomial {
        assert!(p >= 0.0, "Binomial::new called with p < 0");
        assert!(p <= 1.0, "Binomial::new called with p > 1");
        Binomial { n, p }
    }
}

impl Distribution<u64> for Binomial {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        // Handle these values directly.
        if self.p == 0.0 {
            return 0;
        } else if self.p == 1.0 {
            return self.n;
        }

        // binomial distribution is symmetrical with respect to p -> 1-p, k -> n-k
        // switch p so that it is less than 0.5 - this allows for lower expected values
        // we will just invert the result at the end
        let p = if self.p <= 0.5 {
            self.p
        } else {
            1.0 - self.p
        };

        let result;

        // For small n * min(p, 1 - p), the BINV algorithm based on the inverse
        // transformation of the binomial distribution is more efficient:
        //
        // Voratas Kachitvichyanukul and Bruce W. Schmeiser. 1988. Binomial
        // random variate generation. Commun. ACM 31, 2 (February 1988),
        // 216-222. http://dx.doi.org/10.1145/42372.42381
        if (self.n as f64) * p < 10. {
            let q = 1. - p;
            let s = p / q;
            let a = ((self.n + 1) as f64) * s;
            let mut r = q.powf(self.n as f64);
            // FIXME: Using integer exponentiation should be faster, but it is not available for
            // `u64` exponents, neither in `std` or `num`.
            let mut u: f64 = rng.gen();
            let mut x = 0;
            while u > r as f64 {
                u -= r;
                x += 1;
                r *= a / (x as f64) - s;
            }
            result = x;
        } else {
            // FIXME: Using the BTPE algorithm is probably faster.

            // prepare some cached values
            let float_n = self.n as f64;
            let ln_fact_n = log_gamma(float_n + 1.0);
            let pc = 1.0 - p;
            let log_p = p.ln();
            let log_pc = pc.ln();
            let expected = self.n as f64 * p;
            let sq = (expected * (2.0 * pc)).sqrt();
            let mut lresult;

            // we use the Cauchy distribution as the comparison distribution
            // f(x) ~ 1/(1+x^2)
            let cauchy = Cauchy::new(0.0, 1.0);
            loop {
                let mut comp_dev: f64;
                loop {
                    // draw from the Cauchy distribution
                    comp_dev = rng.sample(cauchy);
                    // shift the peak of the comparison ditribution
                    lresult = expected + sq * comp_dev;
                    // repeat the drawing until we are in the range of possible values
                    if lresult >= 0.0 && lresult < float_n + 1.0 {
                        break;
                    }
                }

                // the result should be discrete
                lresult = lresult.floor();

                let log_binomial_dist = ln_fact_n - log_gamma(lresult+1.0) -
                    log_gamma(float_n - lresult + 1.0) + lresult*log_p + (float_n - lresult)*log_pc;
                // this is the binomial probability divided by the comparison probability
                // we will generate a uniform random value and if it is larger than this,
                // we interpret it as a value falling out of the distribution and repeat
                let comparison_coeff = (log_binomial_dist.exp() * sq) * (1.2 * (1.0 + comp_dev*comp_dev));

                if comparison_coeff >= rng.gen() {
                    break;
                }
            }
            result = lresult as u64;
        }

        // invert the result for p < 0.5
        if p != self.p {
            self.n - result
        } else {
            result
        }
    }
}

#[cfg(test)]
mod test {
    use Rng;
    use distributions::Distribution;
    use super::Binomial;

    fn test_binomial_mean_and_variance<R: Rng>(n: u64, p: f64, rng: &mut R) {
        let binomial = Binomial::new(n, p);

        let expected_mean = n as f64 * p;
        let expected_variance = n as f64 * p * (1.0 - p);

        let mut results = [0.0; 1000];
        for i in results.iter_mut() { *i = binomial.sample(rng) as f64; }

        let mean = results.iter().sum::<f64>() / results.len() as f64;
        assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0);

        let variance =
            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>()
            / results.len() as f64;
        assert!((variance - expected_variance).abs() < expected_variance / 10.0);
    }

    #[test]
    fn test_binomial() {
        let mut rng = ::test::rng(351);
        test_binomial_mean_and_variance(150, 0.1, &mut rng);
        test_binomial_mean_and_variance(70, 0.6, &mut rng);
        test_binomial_mean_and_variance(40, 0.5, &mut rng);
        test_binomial_mean_and_variance(20, 0.7, &mut rng);
        test_binomial_mean_and_variance(20, 0.5, &mut rng);
    }

    #[test]
    fn test_binomial_end_points() {
        let mut rng = ::test::rng(352);
        assert_eq!(rng.sample(Binomial::new(20, 0.0)), 0);
        assert_eq!(rng.sample(Binomial::new(20, 1.0)), 20);
    }

    #[test]
    #[should_panic]
    fn test_binomial_invalid_lambda_neg() {
        Binomial::new(20, -10.0);
    }
}
<commit_msg>Binomial: Use integer exponentation<commit_after>// Copyright 2018 Developers of the Rand project.
// Copyright 2016-2017 The Rust Project Developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! The binomial distribution.

use Rng;
use distributions::{Distribution, Cauchy};
use distributions::utils::log_gamma;

/// The binomial distribution `Binomial(n, p)`.
///
/// This distribution has density function:
/// `f(k) = n!/(k! (n-k)!) p^k (1-p)^(n-k)` for `k >= 0`.
///
/// # Example
///
/// ```
/// use rand::distributions::{Binomial, Distribution};
///
/// let bin = Binomial::new(20, 0.3);
/// let v = bin.sample(&mut rand::thread_rng());
/// println!("{} is from a binomial distribution", v);
/// ```
#[derive(Clone, Copy, Debug)]
pub struct Binomial {
    /// Number of trials.
    n: u64,
    /// Probability of success.
    p: f64,
}

impl Binomial {
    /// Construct a new `Binomial` with the given shape parameters `n` (number
    /// of trials) and `p` (probability of success).
    ///
    /// Panics if `p < 0` or `p > 1`.
    pub fn new(n: u64, p: f64) -> Binomial {
        assert!(p >= 0.0, "Binomial::new called with p < 0");
        assert!(p <= 1.0, "Binomial::new called with p > 1");
        Binomial { n, p }
    }
}

/// Raise a `base` to the power of `exp`, using exponentiation by squaring.
///
/// This implementation is based on the one in the `num_traits` crate. It is
/// slightly modified to accept `u64` exponents.
fn pow(mut base: f64, mut exp: u64) -> f64 {
    if exp == 0 {
        return 1.;
    }

    while exp & 1 == 0 {
        base *= base;
        exp >>= 1;
    }

    let mut acc = base;
    while exp > 1 {
        exp >>= 1;
        base *= base;
        if exp & 1 == 1 {
            acc *= base;
        }
    }
    acc
}

impl Distribution<u64> for Binomial {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> u64 {
        // Handle these values directly.
        if self.p == 0.0 {
            return 0;
        } else if self.p == 1.0 {
            return self.n;
        }

        // binomial distribution is symmetrical with respect to p -> 1-p, k -> n-k
        // switch p so that it is less than 0.5 - this allows for lower expected values
        // we will just invert the result at the end
        let p = if self.p <= 0.5 {
            self.p
        } else {
            1.0 - self.p
        };

        let result;

        // For small n * min(p, 1 - p), the BINV algorithm based on the inverse
        // transformation of the binomial distribution is more efficient:
        //
        // Voratas Kachitvichyanukul and Bruce W. Schmeiser. 1988. Binomial
        // random variate generation. Commun. ACM 31, 2 (February 1988),
        // 216-222. http://dx.doi.org/10.1145/42372.42381
        if (self.n as f64) * p < 10. {
            let q = 1. - p;
            let s = p / q;
            let a = ((self.n + 1) as f64) * s;
            let mut r = pow(q, self.n);
            let mut u: f64 = rng.gen();
            let mut x = 0;
            while u > r as f64 {
                u -= r;
                x += 1;
                r *= a / (x as f64) - s;
            }
            result = x;
        } else {
            // FIXME: Using the BTPE algorithm is probably faster.

            // prepare some cached values
            let float_n = self.n as f64;
            let ln_fact_n = log_gamma(float_n + 1.0);
            let pc = 1.0 - p;
            let log_p = p.ln();
            let log_pc = pc.ln();
            let expected = self.n as f64 * p;
            let sq = (expected * (2.0 * pc)).sqrt();
            let mut lresult;

            // we use the Cauchy distribution as the comparison distribution
            // f(x) ~ 1/(1+x^2)
            let cauchy = Cauchy::new(0.0, 1.0);
            loop {
                let mut comp_dev: f64;
                loop {
                    // draw from the Cauchy distribution
                    comp_dev = rng.sample(cauchy);
                    // shift the peak of the comparison ditribution
                    lresult = expected + sq * comp_dev;
                    // repeat the drawing until we are in the range of possible values
                    if lresult >= 0.0 && lresult < float_n + 1.0 {
                        break;
                    }
                }

                // the result should be discrete
                lresult = lresult.floor();

                let log_binomial_dist = ln_fact_n - log_gamma(lresult+1.0) -
                    log_gamma(float_n - lresult + 1.0) + lresult*log_p + (float_n - lresult)*log_pc;
                // this is the binomial probability divided by the comparison probability
                // we will generate a uniform random value and if it is larger than this,
                // we interpret it as a value falling out of the distribution and repeat
                let comparison_coeff = (log_binomial_dist.exp() * sq) * (1.2 * (1.0 + comp_dev*comp_dev));

                if comparison_coeff >= rng.gen() {
                    break;
                }
            }
            result = lresult as u64;
        }

        // invert the result for p < 0.5
        if p != self.p {
            self.n - result
        } else {
            result
        }
    }
}

#[cfg(test)]
mod test {
    use Rng;
    use distributions::Distribution;
    use super::Binomial;

    fn test_binomial_mean_and_variance<R: Rng>(n: u64, p: f64, rng: &mut R) {
        let binomial = Binomial::new(n, p);

        let expected_mean = n as f64 * p;
        let expected_variance = n as f64 * p * (1.0 - p);

        let mut results = [0.0; 1000];
        for i in results.iter_mut() { *i = binomial.sample(rng) as f64; }

        let mean = results.iter().sum::<f64>() / results.len() as f64;
        assert!((mean as f64 - expected_mean).abs() < expected_mean / 50.0);

        let variance =
            results.iter().map(|x| (x - mean) * (x - mean)).sum::<f64>()
            / results.len() as f64;
        assert!((variance - expected_variance).abs() < expected_variance / 10.0);
    }

    #[test]
    fn test_binomial() {
        let mut rng = ::test::rng(351);
        test_binomial_mean_and_variance(150, 0.1, &mut rng);
        test_binomial_mean_and_variance(70, 0.6, &mut rng);
        test_binomial_mean_and_variance(40, 0.5, &mut rng);
        test_binomial_mean_and_variance(20, 0.7, &mut rng);
        test_binomial_mean_and_variance(20, 0.5, &mut rng);
    }

    #[test]
    fn test_binomial_end_points() {
        let mut rng = ::test::rng(352);
        assert_eq!(rng.sample(Binomial::new(20, 0.0)), 0);
        assert_eq!(rng.sample(Binomial::new(20, 1.0)), 20);
    }

    #[test]
    #[should_panic]
    fn test_binomial_invalid_lambda_neg() {
        Binomial::new(20, -10.0);
    }
}
<|endoftext|>


<commit_before>use super::{sleep_ms, spawn};
use eventual::*;
use std::sync::mpsc::channel;

#[test]
pub fn test_joining_two_tuple_futures_async() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (c2, f2) = Future::<i32, ()>::pair();
    let (tx, rx) = channel();

    join((f1, f2)).receive(move |res| {
        tx.send(res.unwrap()).unwrap();
    });

    assert!(rx.try_recv().is_err());
    c1.complete(1);

    assert!(rx.try_recv().is_err());
    c2.complete(2);

    assert_eq!(rx.recv().unwrap(), (1, 2));
}

#[test]
pub fn test_joining_two_futures_sync() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (c2, f2) = Future::<i32, ()>::pair();

    spawn(move || {
        sleep_ms(50);
        c1.complete(1);
    });

    spawn(move || {
        sleep_ms(75);
        c2.complete(2);
    });

    let vals = join((f1, f2)).await().unwrap();
    assert_eq!(vals, (1, 2));
}

#[test]
pub fn test_lazily_propagates_interest() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (_c, f2) = Future::<i32, ()>::pair();
    let (tx, rx) = channel();

    c1.receive(move |res| {
        if let Ok(c) = res {
            tx.send("fail").unwrap();
            c.complete(1);
        }
    });

    // Join but don't register a callback
    let _j = join((f1, f2));

    assert!(rx.try_recv().is_err());
}

#[test]
pub fn test_join_errors_on_failure() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (_c, f2) = Future::<i32, ()>::pair();
    let (tx, rx) = channel();

    join((f1, f2)).receive(move |res| {
        if res.is_err() {
            tx.send("win").unwrap();
        }
    });

    c1.fail(());

    assert_eq!("win", rx.try_recv().unwrap());
}

/*

    Test is blocked by a Rust bug
    https://github.com/rust-lang/rust/issues/21080

#[test]
pub fn test_joining_three_futures_async() {
    let (f1, c1) = Future::<i32, ()>::pair();
    let (f2, c2) = Future::<i32, ()>::pair();
    let (f3, c3) = Future::<i32, ()>::pair();
    let (tx, rx) = channel::<(i32, i32, i32)>();

    join((f1, f2, f3)).receive(move |res: AsyncResult<(i32, i32, i32), ()>| {
        tx.send(res.unwrap()).unwrap();
    });

    assert!(rx.try_recv().is_err());
    c1.complete(1);

    assert!(rx.try_recv().is_err());
    c2.complete(2);

    assert!(rx.try_recv().is_err());
    c3.complete(3);

    assert_eq!(rx.recv().unwrap(), (1, 2, 3));
}
*/

#[test]
pub fn test_joining_two_vec_futures_async() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (c2, f2) = Future::<i32, ()>::pair();
    let (tx, rx) = channel();

    join(vec![f1, f2]).receive(move |res| {
        tx.send(res.unwrap()).unwrap();
    });

    assert!(rx.try_recv().is_err());
    c1.complete(1);

    assert!(rx.try_recv().is_err());
    c2.complete(2);

    assert_eq!(rx.recv().unwrap(), [1, 2]);
}
<commit_msg>Uncomment a test that is no longer blocked by a rustc bug.<commit_after>use super::{sleep_ms, spawn};
use eventual::*;
use std::sync::mpsc::channel;

#[test]
pub fn test_joining_two_tuple_futures_async() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (c2, f2) = Future::<i32, ()>::pair();
    let (tx, rx) = channel();

    join((f1, f2)).receive(move |res| {
        tx.send(res.unwrap()).unwrap();
    });

    assert!(rx.try_recv().is_err());
    c1.complete(1);

    assert!(rx.try_recv().is_err());
    c2.complete(2);

    assert_eq!(rx.recv().unwrap(), (1, 2));
}

#[test]
pub fn test_joining_two_futures_sync() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (c2, f2) = Future::<i32, ()>::pair();

    spawn(move || {
        sleep_ms(50);
        c1.complete(1);
    });

    spawn(move || {
        sleep_ms(75);
        c2.complete(2);
    });

    let vals = join((f1, f2)).await().unwrap();
    assert_eq!(vals, (1, 2));
}

#[test]
pub fn test_lazily_propagates_interest() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (_c, f2) = Future::<i32, ()>::pair();
    let (tx, rx) = channel();

    c1.receive(move |res| {
        if let Ok(c) = res {
            tx.send("fail").unwrap();
            c.complete(1);
        }
    });

    // Join but don't register a callback
    let _j = join((f1, f2));

    assert!(rx.try_recv().is_err());
}

#[test]
pub fn test_join_errors_on_failure() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (_c, f2) = Future::<i32, ()>::pair();
    let (tx, rx) = channel();

    join((f1, f2)).receive(move |res| {
        if res.is_err() {
            tx.send("win").unwrap();
        }
    });

    c1.fail(());

    assert_eq!("win", rx.try_recv().unwrap());
}

#[test]
pub fn test_joining_three_futures_async() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (c2, f2) = Future::<i32, ()>::pair();
    let (c3, f3) = Future::<i32, ()>::pair();
    let (tx, rx) = channel::<(i32, i32, i32)>();

    join((f1, f2, f3)).receive(move |res: AsyncResult<(i32, i32, i32), ()>| {
        tx.send(res.unwrap()).unwrap();
    });

    assert!(rx.try_recv().is_err());
    c1.complete(1);

    assert!(rx.try_recv().is_err());
    c2.complete(2);

    assert!(rx.try_recv().is_err());
    c3.complete(3);

    assert_eq!(rx.recv().unwrap(), (1, 2, 3));
}

#[test]
pub fn test_joining_two_vec_futures_async() {
    let (c1, f1) = Future::<i32, ()>::pair();
    let (c2, f2) = Future::<i32, ()>::pair();
    let (tx, rx) = channel();

    join(vec![f1, f2]).receive(move |res| {
        tx.send(res.unwrap()).unwrap();
    });

    assert!(rx.try_recv().is_err());
    c1.complete(1);

    assert!(rx.try_recv().is_err());
    c2.complete(2);

    assert_eq!(rx.recv().unwrap(), [1, 2]);
}
<|endoftext|>