<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use num_bigint_dig::BigUint;
use num_traits::Num;
use std::cmp::Ordering;
use std::fmt;
use thiserror::Error;

use crate::util::parse_int::ParseInt;

#[derive(Error, Debug, Clone, PartialEq, Eq)]
pub enum ParseBigIntError {
    #[error("integer is too large")]
    Overflow,
    #[error("integer is too small")]
    Underflow,
    #[error(transparent)]
    ParseBigIntError(#[from] num_bigint_dig::ParseBigIntError),
}

/// A fixed-size unsigned big integer.
///
/// This struct wraps a `BigUint` to facilitate defining new fixed-size unsigned integer types for
/// better type safety.
///
/// An integer stored in this type is fixed-size in the sense that the minimum number of bits
/// required to represent it, i.e. its bit length, is at most `BIT_LEN`. This size can be specified
/// using the const parameters `BIT_LEN` and `EXACT_LEN` as follows:
///   - When `EXACT_LEN` is `false`, the bit length of the integer can be at most `BIT_LEN` bits,
///     e.g. SHA-256 digests (at most 256 bits) or RSA-3072 signatures (at most 3072 bits),
///   - When `EXACT_LEN` is `true`, the number of bits required to represent the integer must be
///     exactly `BIT_LEN` bits, e.g. RSA-3072 moduli (exactly 3072 bits).
/// Note that while the type encapsulates the size information, the actual check is performed at
/// runtime when an instance is created (see `check_len()`).
///
/// This struct is not meant to be used directly, please see the `fixed_size_bigint` macro which
/// also generates the required boilerplate code for new types.
#[derive(Debug, Clone, Eq, PartialEq)]
pub(crate) struct FixedSizeBigInt<const BIT_LEN: usize, const EXACT_LEN: bool>(BigUint);

impl<const BIT_LEN: usize, const EXACT_LEN: bool> FixedSizeBigInt<BIT_LEN, EXACT_LEN> {
    const BYTE_LEN: usize = BIT_LEN.saturating_add(u8::BITS as usize - 1) / u8::BITS as usize;

    /// Checks the bit length of the `FixedSizeBigInt`.
    ///
    /// Bit length of a `FixedSizeBigInt` can be at most `BIT_LEN` if `EXACT_LEN` is `false`, must
    /// be exactly `BIT_LEN` otherwise.
    fn new_from_biguint(biguint: BigUint) -> Result<Self, ParseBigIntError> {
        match (biguint.bits().cmp(&BIT_LEN), EXACT_LEN) {
            (Ordering::Greater, _) => Err(ParseBigIntError::Overflow),
            (Ordering::Equal, _) => Ok(Self(biguint)),
            (Ordering::Less, true) => Err(ParseBigIntError::Underflow),
            (Ordering::Less, false) => Ok(Self(biguint)),
        }
    }

    /// Creates a `FixedSizeBigInt` from little-endian bytes.
    pub(crate) fn from_le_bytes(bytes: impl AsRef<[u8]>) -> Result<Self, ParseBigIntError> {
        Self::new_from_biguint(BigUint::from_bytes_le(bytes.as_ref()))
    }

    /// Creates a `FixedSizeBigInt` from big-endian bytes.
    pub(crate) fn from_be_bytes(bytes: impl AsRef<[u8]>) -> Result<Self, ParseBigIntError> {
        Self::new_from_biguint(BigUint::from_bytes_be(bytes.as_ref()))
    }

    /// Returns the bit length.
    ///
    /// Bit length of `FixedSizeBigInt` is the minimum number of bits required to represent its
    /// value. The underlying storage may be larger.
    pub(crate) fn bit_len(&self) -> usize {
        self.0.bits()
    }

    /// Returns the byte representation in little-endian order.
    pub(crate) fn to_le_bytes(&self) -> Vec<u8> {
        self.0.to_bytes_le()
    }

    /// Returns the byte representation in big-endian order.
    pub(crate) fn to_be_bytes(&self) -> Vec<u8> {
        self.0.to_bytes_be()
    }

    /// Returns the underlying `BigUint`.
    // FIXME: remove this `dead_code` if this function is ever used.
    #[allow(dead_code)]
    pub(crate) fn as_biguint(&self) -> &BigUint {
        &self.0
    }
}

impl<const BIT_LEN: usize, const EXACT_LEN: bool> ParseInt for FixedSizeBigInt<BIT_LEN, EXACT_LEN> {
    type FromStrRadixErr = ParseBigIntError;

    fn from_str_radix(src: &str, radix: u32) -> Result<Self, Self::FromStrRadixErr> {
        Self::new_from_biguint(
            BigUint::from_str_radix(src, radix).map_err(ParseBigIntError::ParseBigIntError)?,
        )
    }
}

impl<const BIT_LEN: usize, const EXACT_LEN: bool> fmt::Display
    for FixedSizeBigInt<BIT_LEN, EXACT_LEN>
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(
            &format_args!("{:#0width$x}", self.0, width = Self::BYTE_LEN * 2 + 2),
            f,
        )
    }
}

/// Helper macro for the `fixed_size_bigint` macro.
macro_rules! fixed_size_bigint_impl {
    ($struct_name:ident, $bit_len:expr, $exact_len:expr) => {
        #[derive(serde::Serialize, Debug, Clone, Eq, PartialEq)]
        #[serde(into = "String")]
        pub struct $struct_name($crate::util::bigint::FixedSizeBigInt<$bit_len, $exact_len>);

        const _: () = {
            use num_bigint_dig::BigUint;
            use std::fmt;
            use std::result::Result;

            use $crate::util::bigint::{FixedSizeBigInt, ParseBigIntError};
            use $crate::util::parse_int::ParseInt;

            impl $struct_name {
                pub fn from_le_bytes(bytes: impl AsRef<[u8]>) -> Result<Self, ParseBigIntError> {
                    Ok($struct_name(
                        FixedSizeBigInt::<$bit_len, $exact_len>::from_le_bytes(bytes)?,
                    ))
                }

                pub fn from_be_bytes(bytes: impl AsRef<[u8]>) -> Result<Self, ParseBigIntError> {
                    Ok($struct_name(
                        FixedSizeBigInt::<$bit_len, $exact_len>::from_be_bytes(bytes)?,
                    ))
                }

                pub fn bit_len(&self) -> usize {
                    self.0.bit_len()
                }

                pub fn to_le_bytes(&self) -> Vec<u8> {
                    self.0.to_le_bytes()
                }

                pub fn to_be_bytes(&self) -> Vec<u8> {
                    self.0.to_be_bytes()
                }

                pub fn as_biguint(&self) -> &BigUint {
                    self.0.as_biguint()
                }
            }

            impl ParseInt for $struct_name {
                type FromStrRadixErr = ParseBigIntError;

                fn from_str_radix(src: &str, radix: u32) -> Result<Self, Self::FromStrRadixErr> {
                    Ok($struct_name(
                        FixedSizeBigInt::<$bit_len, $exact_len>::from_str_radix(src, radix)?,
                    ))
                }
            }

            impl fmt::Display for $struct_name {
                fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                    fmt::Display::fmt(&self.0, f)
                }
            }

            impl From<$struct_name> for String {
                fn from(s: $struct_name) -> String {
                    s.0.to_string()
                }
            }
        };
    };
}

pub(crate) use fixed_size_bigint_impl;

/// Macro for defining a new fixed-size unsigned big integer type.
///
/// Defines a new type that wraps a `FixedSizeBigInt`. This macro is intended to be used within this
/// crate to define types which can then be exported as needed:
///
/// ```
/// use crate::util::bigint::fixed_size_bigint;
///
/// // Define a type for RSA-3072 moduli (exactly 3072 bits long):
/// fixed_size_bigint!(Rsa3072Modulus, 3072);
///
/// // Define a type for SHA-256 digests (at most 256 bits long):
/// fixed_size_bigint!(Sha256Digest, at_most 256);
/// ```
macro_rules! fixed_size_bigint {
    ($struct_name:ident, $bit_len:expr) => {
        $crate::util::bigint::fixed_size_bigint_impl!($struct_name, $bit_len, true);
    };
    ($struct_name:ident, at_most $bit_len:expr) => {
        $crate::util::bigint::fixed_size_bigint_impl!($struct_name, $bit_len, false);
    };
}

pub(crate) use fixed_size_bigint;

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

    fixed_size_bigint!(TestArray, at_most 16);
    fixed_size_bigint!(TestArrayExact, 16);

    #[test]
    fn test_from_to_le_bytes() {
        fn check(slice: &[u8], data: &[u8]) {
            assert_eq!(TestArray::from_le_bytes(slice).unwrap().to_le_bytes(), data);
        }
        check(&[], &[0]);
        check(&[1], &[1]);
        check(&[0, 1], &[0, 1]);
        check(&[1, 0], &[1]);

        assert!(TestArray::from_le_bytes([1, 2, 3]).is_err());
    }

    #[test]
    fn test_from_to_le_bytes_exact_len() {
        fn check(slice: &[u8], data: &[u8]) {
            assert_eq!(
                TestArrayExact::from_le_bytes(slice).unwrap().to_le_bytes(),
                data
            );
        }
        check(&[0, 128], &[0, 128]);
        check(&[255, 255, 0], &[255, 255]);

        assert!(TestArrayExact::from_le_bytes([1]).is_err());
        assert!(TestArrayExact::from_le_bytes([255, 127]).is_err());
        assert!(TestArrayExact::from_le_bytes([0, 0, 1]).is_err());
    }

    #[test]
    fn test_from_to_be_bytes() {
        fn check(slice: &[u8], data: &[u8]) {
            assert_eq!(TestArray::from_be_bytes(slice).unwrap().to_be_bytes(), data);
        }
        check(&[1], &[1]);
        check(&[1, 0], &[1, 0]);
        check(&[0, 1], &[1]);

        assert!(TestArray::from_be_bytes([1, 2, 1]).is_err());
    }

    #[test]
    fn test_from_to_be_bytes_exact_len() {
        fn check(slice: &[u8], data: &[u8]) {
            assert_eq!(
                TestArrayExact::from_be_bytes(slice).unwrap().to_be_bytes(),
                data
            );
        }
        check(&[128, 1], &[128, 1]);
        check(&[0, 255, 255], &[255, 255]);

        assert!(TestArrayExact::from_be_bytes([1]).is_err());
        assert!(TestArrayExact::from_be_bytes([127, 1]).is_err());
        assert!(TestArrayExact::from_be_bytes([1, 0, 0]).is_err());
    }

    #[test]
    fn test_bit_len() {
        fn check(slice: &[u8], bit_len: usize) {
            assert_eq!(TestArray::from_le_bytes(slice).unwrap().bit_len(), bit_len);
        }
        check(&[1], 1);
        check(&[1, 0], 1);
        check(&[255], 8);
        check(&[0, 1], 9);
        check(&[0, 128], 16);
    }

    #[test]
    fn test_from_str() {
        assert_eq!(TestArray::from_str("0x01").unwrap().to_le_bytes(), [1]);
        assert_eq!(
            TestArray::from_str("0x00201").unwrap().to_le_bytes(),
            [1, 2]
        );
        assert!(TestArray::from_str("0x030201").is_err());
    }

    #[test]
    fn test_from_str_exact_len() {
        assert_eq!(
            TestArrayExact::from_str("0x08001").unwrap().to_le_bytes(),
            [1, 128]
        );

        assert!(TestArrayExact::from_str("0x01").is_err());
        assert!(TestArrayExact::from_str("0x0201").is_err());
        assert!(TestArrayExact::from_str("0x030201").is_err());
    }

    #[test]
    fn test_fmt() {
        let exact = TestArrayExact::from_str("0xabcd").unwrap();
        assert_eq!(exact.to_string(), "0xabcd");

        let at_most = TestArray::from_str("0xab").unwrap();
        assert_eq!(at_most.to_string(), "0x00ab");

        let at_most_full = TestArray::from_str("0xabcd").unwrap();
        assert_eq!(at_most_full.to_string(), "0xabcd");
    }
}
<commit_msg>[opentitanlib] Update FixedSizeBigInt to_bytes methods to return BYTE_LEN long vectors<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use num_bigint_dig::BigUint;
use num_traits::Num;
use std::cmp::Ordering;
use std::fmt;
use std::iter;
use thiserror::Error;

use crate::util::parse_int::ParseInt;

#[derive(Error, Debug, Clone, PartialEq, Eq)]
pub enum ParseBigIntError {
    #[error("integer is too large")]
    Overflow,
    #[error("integer is too small")]
    Underflow,
    #[error(transparent)]
    ParseBigIntError(#[from] num_bigint_dig::ParseBigIntError),
}

/// A fixed-size unsigned big integer.
///
/// This struct wraps a `BigUint` to facilitate defining new fixed-size unsigned integer types for
/// better type safety.
///
/// An integer stored in this type is fixed-size in the sense that the minimum number of bits
/// required to represent it, i.e. its bit length, is at most `BIT_LEN`. This size can be specified
/// using the const parameters `BIT_LEN` and `EXACT_LEN` as follows:
///   - When `EXACT_LEN` is `false`, the bit length of the integer can be at most `BIT_LEN` bits,
///     e.g. SHA-256 digests (at most 256 bits) or RSA-3072 signatures (at most 3072 bits),
///   - When `EXACT_LEN` is `true`, the number of bits required to represent the integer must be
///     exactly `BIT_LEN` bits, e.g. RSA-3072 moduli (exactly 3072 bits).
/// Note that while the type encapsulates the size information, the actual check is performed at
/// runtime when an instance is created (see `check_len()`).
///
/// This struct is not meant to be used directly, please see the `fixed_size_bigint` macro which
/// also generates the required boilerplate code for new types.
#[derive(Debug, Clone, Eq, PartialEq)]
pub(crate) struct FixedSizeBigInt<const BIT_LEN: usize, const EXACT_LEN: bool>(BigUint);

impl<const BIT_LEN: usize, const EXACT_LEN: bool> FixedSizeBigInt<BIT_LEN, EXACT_LEN> {
    const BYTE_LEN: usize = BIT_LEN.saturating_add(u8::BITS as usize - 1) / u8::BITS as usize;

    /// Checks the bit length of the `FixedSizeBigInt`.
    ///
    /// Bit length of a `FixedSizeBigInt` can be at most `BIT_LEN` if `EXACT_LEN` is `false`, must
    /// be exactly `BIT_LEN` otherwise.
    fn new_from_biguint(biguint: BigUint) -> Result<Self, ParseBigIntError> {
        match (biguint.bits().cmp(&BIT_LEN), EXACT_LEN) {
            (Ordering::Greater, _) => Err(ParseBigIntError::Overflow),
            (Ordering::Equal, _) => Ok(Self(biguint)),
            (Ordering::Less, true) => Err(ParseBigIntError::Underflow),
            (Ordering::Less, false) => Ok(Self(biguint)),
        }
    }

    /// Creates a `FixedSizeBigInt` from little-endian bytes.
    pub(crate) fn from_le_bytes(bytes: impl AsRef<[u8]>) -> Result<Self, ParseBigIntError> {
        Self::new_from_biguint(BigUint::from_bytes_le(bytes.as_ref()))
    }

    /// Creates a `FixedSizeBigInt` from big-endian bytes.
    pub(crate) fn from_be_bytes(bytes: impl AsRef<[u8]>) -> Result<Self, ParseBigIntError> {
        Self::new_from_biguint(BigUint::from_bytes_be(bytes.as_ref()))
    }

    /// Returns the bit length.
    ///
    /// Bit length of `FixedSizeBigInt` is the minimum number of bits required to represent its
    /// value. The underlying storage may be larger.
    pub(crate) fn bit_len(&self) -> usize {
        self.0.bits()
    }

    /// Returns the byte representation in little-endian order.
    pub(crate) fn to_le_bytes(&self) -> Vec<u8> {
        let mut v = self.0.to_bytes_le();
        assert!(Self::BYTE_LEN >= v.len());
        // Append since `v` is little-endian.
        v.resize(Self::BYTE_LEN, 0);
        v
    }

    /// Returns the byte representation in big-endian order.
    pub(crate) fn to_be_bytes(&self) -> Vec<u8> {
        let mut v = self.0.to_bytes_be();
        assert!(Self::BYTE_LEN >= v.len());
        // Prepend since `v` is big-endian.
        v.splice(0..0, iter::repeat(0).take(Self::BYTE_LEN - v.len()));
        v
    }

    /// Returns the underlying `BigUint`.
    pub(crate) fn as_biguint(&self) -> &BigUint {
        &self.0
    }
}

impl<const BIT_LEN: usize, const EXACT_LEN: bool> ParseInt for FixedSizeBigInt<BIT_LEN, EXACT_LEN> {
    type FromStrRadixErr = ParseBigIntError;

    fn from_str_radix(src: &str, radix: u32) -> Result<Self, Self::FromStrRadixErr> {
        Self::new_from_biguint(
            BigUint::from_str_radix(src, radix).map_err(ParseBigIntError::ParseBigIntError)?,
        )
    }
}

impl<const BIT_LEN: usize, const EXACT_LEN: bool> fmt::Display
    for FixedSizeBigInt<BIT_LEN, EXACT_LEN>
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Display::fmt(
            &format_args!("{:#0width$x}", self.0, width = Self::BYTE_LEN * 2 + 2),
            f,
        )
    }
}

/// Helper macro for the `fixed_size_bigint` macro.
macro_rules! fixed_size_bigint_impl {
    ($struct_name:ident, $bit_len:expr, $exact_len:expr) => {
        #[derive(serde::Serialize, Debug, Clone, Eq, PartialEq)]
        #[serde(into = "String")]
        pub struct $struct_name($crate::util::bigint::FixedSizeBigInt<$bit_len, $exact_len>);

        const _: () = {
            use num_bigint_dig::BigUint;
            use std::fmt;
            use std::result::Result;

            use $crate::util::bigint::{FixedSizeBigInt, ParseBigIntError};
            use $crate::util::parse_int::ParseInt;

            impl $struct_name {
                pub fn from_le_bytes(bytes: impl AsRef<[u8]>) -> Result<Self, ParseBigIntError> {
                    Ok($struct_name(
                        FixedSizeBigInt::<$bit_len, $exact_len>::from_le_bytes(bytes)?,
                    ))
                }

                pub fn from_be_bytes(bytes: impl AsRef<[u8]>) -> Result<Self, ParseBigIntError> {
                    Ok($struct_name(
                        FixedSizeBigInt::<$bit_len, $exact_len>::from_be_bytes(bytes)?,
                    ))
                }

                pub fn bit_len(&self) -> usize {
                    self.0.bit_len()
                }

                pub fn to_le_bytes(&self) -> Vec<u8> {
                    self.0.to_le_bytes()
                }

                pub fn to_be_bytes(&self) -> Vec<u8> {
                    self.0.to_be_bytes()
                }

                pub fn as_biguint(&self) -> &BigUint {
                    self.0.as_biguint()
                }
            }

            impl ParseInt for $struct_name {
                type FromStrRadixErr = ParseBigIntError;

                fn from_str_radix(src: &str, radix: u32) -> Result<Self, Self::FromStrRadixErr> {
                    Ok($struct_name(
                        FixedSizeBigInt::<$bit_len, $exact_len>::from_str_radix(src, radix)?,
                    ))
                }
            }

            impl fmt::Display for $struct_name {
                fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                    fmt::Display::fmt(&self.0, f)
                }
            }

            impl From<$struct_name> for String {
                fn from(s: $struct_name) -> String {
                    s.0.to_string()
                }
            }
        };
    };
}

pub(crate) use fixed_size_bigint_impl;

/// Macro for defining a new fixed-size unsigned big integer type.
///
/// Defines a new type that wraps a `FixedSizeBigInt`. This macro is intended to be used within this
/// crate to define types which can then be exported as needed:
///
/// ```
/// use crate::util::bigint::fixed_size_bigint;
///
/// // Define a type for RSA-3072 moduli (exactly 3072 bits long):
/// fixed_size_bigint!(Rsa3072Modulus, 3072);
///
/// // Define a type for SHA-256 digests (at most 256 bits long):
/// fixed_size_bigint!(Sha256Digest, at_most 256);
/// ```
macro_rules! fixed_size_bigint {
    ($struct_name:ident, $bit_len:expr) => {
        $crate::util::bigint::fixed_size_bigint_impl!($struct_name, $bit_len, true);
    };
    ($struct_name:ident, at_most $bit_len:expr) => {
        $crate::util::bigint::fixed_size_bigint_impl!($struct_name, $bit_len, false);
    };
}

pub(crate) use fixed_size_bigint;

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

    fixed_size_bigint!(TestArray, at_most 16);
    fixed_size_bigint!(TestArrayExact, 16);

    #[test]
    fn test_from_to_le_bytes() {
        fn check(slice: &[u8], data: &[u8]) {
            assert_eq!(TestArray::from_le_bytes(slice).unwrap().to_le_bytes(), data);
        }
        check(&[], &[0, 0]);
        check(&[1], &[1, 0]);
        check(&[0, 1], &[0, 1]);
        check(&[1, 0], &[1, 0]);

        assert!(TestArray::from_le_bytes([1, 2, 3]).is_err());
    }

    #[test]
    fn test_from_to_le_bytes_exact_len() {
        fn check(slice: &[u8], data: &[u8]) {
            assert_eq!(
                TestArrayExact::from_le_bytes(slice).unwrap().to_le_bytes(),
                data
            );
        }
        check(&[0, 128], &[0, 128]);
        check(&[255, 255, 0], &[255, 255]);

        assert!(TestArrayExact::from_le_bytes([1]).is_err());
        assert!(TestArrayExact::from_le_bytes([255, 127]).is_err());
        assert!(TestArrayExact::from_le_bytes([0, 0, 1]).is_err());
    }

    #[test]
    fn test_from_to_be_bytes() {
        fn check(slice: &[u8], data: &[u8]) {
            assert_eq!(TestArray::from_be_bytes(slice).unwrap().to_be_bytes(), data);
        }
        check(&[1], &[0, 1]);
        check(&[1, 0], &[1, 0]);
        check(&[0, 1], &[0, 1]);

        assert!(TestArray::from_be_bytes([1, 2, 1]).is_err());
    }

    #[test]
    fn test_from_to_be_bytes_exact_len() {
        fn check(slice: &[u8], data: &[u8]) {
            assert_eq!(
                TestArrayExact::from_be_bytes(slice).unwrap().to_be_bytes(),
                data
            );
        }
        check(&[128, 1], &[128, 1]);
        check(&[0, 255, 255], &[255, 255]);

        assert!(TestArrayExact::from_be_bytes([1]).is_err());
        assert!(TestArrayExact::from_be_bytes([127, 1]).is_err());
        assert!(TestArrayExact::from_be_bytes([1, 0, 0]).is_err());
    }

    #[test]
    fn test_bit_len() {
        fn check(slice: &[u8], bit_len: usize) {
            assert_eq!(TestArray::from_le_bytes(slice).unwrap().bit_len(), bit_len);
        }
        check(&[1], 1);
        check(&[1, 0], 1);
        check(&[255], 8);
        check(&[0, 1], 9);
        check(&[0, 128], 16);
    }

    #[test]
    fn test_from_str() {
        assert_eq!(TestArray::from_str("0x01").unwrap().to_le_bytes(), [1, 0]);
        assert_eq!(
            TestArray::from_str("0x00201").unwrap().to_le_bytes(),
            [1, 2]
        );
        assert!(TestArray::from_str("0x030201").is_err());
    }

    #[test]
    fn test_from_str_exact_len() {
        assert_eq!(
            TestArrayExact::from_str("0x08001").unwrap().to_le_bytes(),
            [1, 128]
        );

        assert!(TestArrayExact::from_str("0x01").is_err());
        assert!(TestArrayExact::from_str("0x0201").is_err());
        assert!(TestArrayExact::from_str("0x030201").is_err());
    }

    #[test]
    fn test_fmt() {
        let exact = TestArrayExact::from_str("0xabcd").unwrap();
        assert_eq!(exact.to_string(), "0xabcd");

        let at_most = TestArray::from_str("0xab").unwrap();
        assert_eq!(at_most.to_string(), "0x00ab");

        let at_most_full = TestArray::from_str("0xabcd").unwrap();
        assert_eq!(at_most_full.to_string(), "0xabcd");
    }
}
<|endoftext|>


<commit_before>import std::run;
import std::run::spawn_process;
import std::io;
import std::os;
import io::writer_util;
import ctypes::{pid_t, fd_t};

export run;

#[cfg(target_os = "win32")]
fn target_env(lib_path: str, prog: str) -> option<[(str,str)]> {

    let env = std::generic_os::env();

    env = vec::map(env) {|pair|
        let (k,v) = pair;
        if k == "PATH" { ("PATH", v + ";" + lib_path) }
        else { (k,v) }
    };
    if str::ends_with(prog, "rustc.exe") {
        env += [("RUST_THREADS", "1")]
    }
    ret some(env);
}

#[cfg(target_os = "linux")]
#[cfg(target_os = "macos")]
#[cfg(target_os = "freebsd")]
fn target_env(_lib_path: str, _prog: str) -> option<[(str,str)]> {
    none
}

fn run(lib_path: str, prog: str, args: [str],
       input: option<str>) -> {status: int, out: str, err: str} {

    let pipe_in = os::pipe();
    let pipe_out = os::pipe();
    let pipe_err = os::pipe();
    let pid = spawn_process(prog, args, target_env(lib_path, prog), none,
                            pipe_in.in, pipe_out.out, pipe_err.out);

    os::close(pipe_in.in);
    os::close(pipe_out.out);
    os::close(pipe_err.out);
    if pid == -1i32 {
        os::close(pipe_in.out);
        os::close(pipe_out.in);
        os::close(pipe_err.in);
        fail;
    }

    writeclose(pipe_in.out, input);
    let errput = readclose(pipe_err.in);
    let output = readclose(pipe_out.in);
    let status = run::waitpid(pid);
    ret {status: status, out: output, err: errput};
}

fn writeclose(fd: fd_t, s: option<str>) {
    if option::is_some(s) {
        let writer = io::fd_writer(fd, false);
        writer.write_str(option::get(s));
    }

    os::close(fd);
}

fn readclose(fd: fd_t) -> str {
    // Copied from run::program_output
    let file = os::fd_FILE(fd);
    let reader = io::FILE_reader(file, false);
    let buf = "";
    while !reader.eof() {
        let bytes = reader.read_bytes(4096u);
        buf += str::from_bytes(bytes);
    }
    os::fclose(file);
    ret buf;
}
<commit_msg>use task::spawn_sched to read stdout and stderr at the same time<commit_after>import std::run;
import std::run::spawn_process;
import std::io;
import std::os;
import io::writer_util;
import ctypes::{pid_t, fd_t};

export run;

#[cfg(target_os = "win32")]
fn target_env(lib_path: str, prog: str) -> option<[(str,str)]> {

    let env = std::generic_os::env();

    env = vec::map(env) {|pair|
        let (k,v) = pair;
        if k == "PATH" { ("PATH", v + ";" + lib_path) }
        else { (k,v) }
    };
    if str::ends_with(prog, "rustc.exe") {
        env += [("RUST_THREADS", "1")]
    }
    ret some(env);
}

#[cfg(target_os = "linux")]
#[cfg(target_os = "macos")]
#[cfg(target_os = "freebsd")]
fn target_env(_lib_path: str, _prog: str) -> option<[(str,str)]> {
    none
}

fn run(lib_path: str, prog: str, args: [str],
       input: option<str>) -> {status: int, out: str, err: str} {

    let pipe_in = os::pipe();
    let pipe_out = os::pipe();
    let pipe_err = os::pipe();
    let pid = spawn_process(prog, args, target_env(lib_path, prog), none,
                            pipe_in.in, pipe_out.out, pipe_err.out);

    os::close(pipe_in.in);
    os::close(pipe_out.out);
    os::close(pipe_err.out);
    if pid == -1i32 {
        os::close(pipe_in.out);
        os::close(pipe_out.in);
        os::close(pipe_err.in);
        fail;
    }

    writeclose(pipe_in.out, input);
    let p = comm::port();
    let ch = comm::chan(p);
    task::spawn_sched(1u) {||
        let errput = readclose(pipe_err.in);
        comm::send(ch, (2, errput));
    };
    task::spawn_sched(1u) {||
        let output = readclose(pipe_out.in);
        comm::send(ch, (1, output));
    };
    let status = run::waitpid(pid);
    let errs = "";
    let outs = "";
    let count = 2;
    while count > 0 {
        let stream = comm::recv(p);
        alt stream {
            (1, s) {
                outs = s;
            }
            (2, s) {
                errs = s;
            }
        };
        count -= 1;
    };
    ret {status: status, out: outs, err: errs};
}

fn writeclose(fd: fd_t, s: option<str>) {
    if option::is_some(s) {
        let writer = io::fd_writer(fd, false);
        writer.write_str(option::get(s));
    }

    os::close(fd);
}

fn readclose(fd: fd_t) -> str {
    // Copied from run::program_output
    let file = os::fd_FILE(fd);
    let reader = io::FILE_reader(file, false);
    let buf = "";
    while !reader.eof() {
        let bytes = reader.read_bytes(4096u);
        buf += str::from_bytes(bytes);
    }
    os::fclose(file);
    ret buf;
}
<|endoftext|>


<commit_before>// https://leetcode.com/problems/number-of-visible-people-in-a-queue/
pub fn can_see_persons_count(heights: Vec<i32>) -> Vec<i32> {
    todo!()
}

fn main() {
    println!("{:?}", can_see_persons_count(vec![10, 6, 8, 5, 11, 9])); // [3,1,2,1,1,0]
    println!("{:?}", can_see_persons_count(vec![5, 1, 2, 3, 10])); // [4,1,1,1,0]
}
<commit_msg>Solve problem 1944<commit_after>use std::collections::VecDeque;

// https://leetcode.com/problems/number-of-visible-people-in-a-queue/
pub fn can_see_persons_count(heights: Vec<i32>) -> Vec<i32> {
    let mut answer: Vec<i32> = vec![0; heights.len()];
    let mut deque: VecDeque<i32> = VecDeque::new();
    for i in (0..heights.len()).rev() {
        let height = &heights[i];
        let mut count = 0;
        loop {
            let front = deque.front();
            match front {
                Some(value) => {
                    if height > value {
                        deque.pop_front();
                        count += 1;
                    } else {
                        count += 1;
                        break;
                    }
                }
                None => break,
            }
        }
        answer[i] = count;
        deque.push_front(*height);
    }
    answer
}

fn main() {
    println!("{:?}", can_see_persons_count(vec![10, 6, 8, 5, 11, 9])); // [3,1,2,1,1,0]
    println!("{:?}", can_see_persons_count(vec![5, 1, 2, 3, 10])); // [4,1,1,1,0]
}
<|endoftext|>


<commit_before>// Copyright 2018 Developers of the Rand project.
//
// 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 implementations of the `Standard` distribution for other built-in types.

use core::char;
use core::num::Wrapping;

use crate::distributions::{Distribution, Standard, Uniform};
use crate::Rng;

#[cfg(feature = "serde1")]
use serde::{Serialize, Deserialize};

// ----- Sampling distributions -----

/// Sample a `char`, uniformly distributed over ASCII letters and numbers:
/// a-z, A-Z and 0-9.
///
/// # Example
///
/// ```
/// use std::iter;
/// use rand::{Rng, thread_rng};
/// use rand::distributions::Alphanumeric;
///
/// let mut rng = thread_rng();
/// let chars: String = iter::repeat(())
///         .map(|()| rng.sample(Alphanumeric))
///         .take(7)
///         .collect();
/// println!("Random chars: {}", chars);
/// ```
///
/// # Passwords
///
/// Users sometimes ask whether it is safe to use a string of random characters
/// as a password. `Alphanumeric` generates from an alphabet of 62 symbols, thus
/// each character can provide `log2(62) = 5.95...` bits of entropy. We suggest
/// consulting external sources for more. One may start with the
/// [Wikipedia article on Password Strength](https://en.wikipedia.org/wiki/Password_strength).
///
/// We caution that strings produced by sampling `Alphanumeric` tend not
/// to be particularly memorable when used as passwords by humans.
/// Drawing words from a specially-curated word-list such as
/// [Diceware](https://en.wikipedia.org/wiki/Diceware) may be a better option
/// for memorable passwords.
/// Each random word contributes `log2(wordlist_length)` bits of entropy.
#[derive(Debug)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub struct Alphanumeric;

// ----- Implementations of distributions -----

impl Distribution<char> for Standard {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> char {
        // A valid `char` is either in the interval `[0, 0xD800)` or
        // `(0xDFFF, 0x11_0000)`. All `char`s must therefore be in
        // `[0, 0x11_0000)` but not in the "gap" `[0xD800, 0xDFFF]` which is
        // reserved for surrogates. This is the size of that gap.
        const GAP_SIZE: u32 = 0xDFFF - 0xD800 + 1;

        // Uniform::new(0, 0x11_0000 - GAP_SIZE) can also be used but it
        // seemed slower.
        let range = Uniform::new(GAP_SIZE, 0x11_0000);

        let mut n = range.sample(rng);
        if n <= 0xDFFF {
            n -= GAP_SIZE;
        }
        unsafe { char::from_u32_unchecked(n) }
    }
}

impl Distribution<char> for Alphanumeric {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> char {
        const RANGE: u32 = 26 + 26 + 10;
        const GEN_ASCII_STR_CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ\
                abcdefghijklmnopqrstuvwxyz\
                0123456789";
        // We can pick from 62 characters. This is so close to a power of 2, 64,
        // that we can do better than `Uniform`. Use a simple bitshift and
        // rejection sampling. We do not use a bitmask, because for small RNGs
        // the most significant bits are usually of higher quality.
        loop {
            let var = rng.next_u32() >> (32 - 6);
            if var < RANGE {
                return GEN_ASCII_STR_CHARSET[var as usize] as char;
            }
        }
    }
}

impl Distribution<bool> for Standard {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool {
        // We can compare against an arbitrary bit of an u32 to get a bool.
        // Because the least significant bits of a lower quality RNG can have
        // simple patterns, we compare against the most significant bit. This is
        // easiest done using a sign test.
        (rng.next_u32() as i32) < 0
    }
}

macro_rules! tuple_impl {
    // use variables to indicate the arity of the tuple
    ($($tyvar:ident),* ) => {
        // the trailing commas are for the 1 tuple
        impl< $( $tyvar ),* >
            Distribution<( $( $tyvar ),* , )>
            for Standard
            where $( Standard: Distribution<$tyvar> ),*
        {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> ( $( $tyvar ),* , ) {
                (
                    // use the $tyvar's to get the appropriate number of
                    // repeats (they're not actually needed)
                    $(
                        _rng.gen::<$tyvar>()
                    ),*
                    ,
                )
            }
        }
    }
}

impl Distribution<()> for Standard {
    #[allow(clippy::unused_unit)]
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, _: &mut R) -> () {
        ()
    }
}
tuple_impl! {A}
tuple_impl! {A, B}
tuple_impl! {A, B, C}
tuple_impl! {A, B, C, D}
tuple_impl! {A, B, C, D, E}
tuple_impl! {A, B, C, D, E, F}
tuple_impl! {A, B, C, D, E, F, G}
tuple_impl! {A, B, C, D, E, F, G, H}
tuple_impl! {A, B, C, D, E, F, G, H, I}
tuple_impl! {A, B, C, D, E, F, G, H, I, J}
tuple_impl! {A, B, C, D, E, F, G, H, I, J, K}
tuple_impl! {A, B, C, D, E, F, G, H, I, J, K, L}

macro_rules! array_impl {
    // recursive, given at least one type parameter:
    {$n:expr, $t:ident, $($ts:ident,)*} => {
        array_impl!{($n - 1), $($ts,)*}

        impl<T> Distribution<[T; $n]> for Standard where Standard: Distribution<T> {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] {
                [_rng.gen::<$t>(), $(_rng.gen::<$ts>()),*]
            }
        }
    };
    // empty case:
    {$n:expr,} => {
        impl<T> Distribution<[T; $n]> for Standard {
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] { [] }
        }
    };
}

array_impl! {32, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T,}

impl<T> Distribution<Option<T>> for Standard
where Standard: Distribution<T>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Option<T> {
        // UFCS is needed here: https://github.com/rust-lang/rust/issues/24066
        if rng.gen::<bool>() {
            Some(rng.gen())
        } else {
            None
        }
    }
}

impl<T> Distribution<Wrapping<T>> for Standard
where Standard: Distribution<T>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Wrapping<T> {
        Wrapping(rng.gen())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::RngCore;
    #[cfg(all(not(feature = "std"), feature = "alloc"))] use alloc::string::String;

    #[test]
    fn test_misc() {
        let rng: &mut dyn RngCore = &mut crate::test::rng(820);

        rng.sample::<char, _>(Standard);
        rng.sample::<bool, _>(Standard);
    }

    #[cfg(feature = "alloc")]
    #[test]
    fn test_chars() {
        use core::iter;
        let mut rng = crate::test::rng(805);

        // Test by generating a relatively large number of chars, so we also
        // take the rejection sampling path.
        let word: String = iter::repeat(())
            .map(|()| rng.gen::<char>())
            .take(1000)
            .collect();
        assert!(word.len() != 0);
    }

    #[test]
    fn test_alphanumeric() {
        let mut rng = crate::test::rng(806);

        // Test by generating a relatively large number of chars, so we also
        // take the rejection sampling path.
        let mut incorrect = false;
        for _ in 0..100 {
            let c = rng.sample(Alphanumeric);
            incorrect |= !((c >= '0' && c <= '9') ||
                           (c >= 'A' && c <= 'Z') ||
                           (c >= 'a' && c <= 'z') );
        }
        assert!(incorrect == false);
    }

    #[test]
    fn value_stability() {
        fn test_samples<T: Copy + core::fmt::Debug + PartialEq, D: Distribution<T>>(
            distr: &D, zero: T, expected: &[T],
        ) {
            let mut rng = crate::test::rng(807);
            let mut buf = [zero; 5];
            for x in &mut buf {
                *x = rng.sample(&distr);
            }
            assert_eq!(&buf, expected);
        }

        test_samples(&Standard, 'a', &[
            '\u{8cdac}',
            '\u{a346a}',
            '\u{80120}',
            '\u{ed692}',
            '\u{35888}',
        ]);
        test_samples(&Alphanumeric, 'a', &['h', 'm', 'e', '3', 'M']);
        test_samples(&Standard, false, &[true, true, false, true, false]);
        test_samples(&Standard, None as Option<bool>, &[
            Some(true),
            None,
            Some(false),
            None,
            Some(false),
        ]);
        test_samples(&Standard, Wrapping(0i32), &[
            Wrapping(-2074640887),
            Wrapping(-1719949321),
            Wrapping(2018088303),
            Wrapping(-547181756),
            Wrapping(838957336),
        ]);

        // We test only sub-sets of tuple and array impls
        test_samples(&Standard, (), &[(), (), (), (), ()]);
        test_samples(&Standard, (false,), &[
            (true,),
            (true,),
            (false,),
            (true,),
            (false,),
        ]);
        test_samples(&Standard, (false, false), &[
            (true, true),
            (false, true),
            (false, false),
            (true, false),
            (false, false),
        ]);

        test_samples(&Standard, [0u8; 0], &[[], [], [], [], []]);
        test_samples(&Standard, [0u8; 3], &[
            [9, 247, 111],
            [68, 24, 13],
            [174, 19, 194],
            [172, 69, 213],
            [149, 207, 29],
        ]);
    }
}
<commit_msg>Avoid to give advice on password generation<commit_after>// Copyright 2018 Developers of the Rand project.
//
// 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 implementations of the `Standard` distribution for other built-in types.

use core::char;
use core::num::Wrapping;

use crate::distributions::{Distribution, Standard, Uniform};
use crate::Rng;

#[cfg(feature = "serde1")]
use serde::{Serialize, Deserialize};

// ----- Sampling distributions -----

/// Sample a `char`, uniformly distributed over ASCII letters and numbers:
/// a-z, A-Z and 0-9.
///
/// # Example
///
/// ```
/// use std::iter;
/// use rand::{Rng, thread_rng};
/// use rand::distributions::Alphanumeric;
///
/// let mut rng = thread_rng();
/// let chars: String = iter::repeat(())
///         .map(|()| rng.sample(Alphanumeric))
///         .take(7)
///         .collect();
/// println!("Random chars: {}", chars);
/// ```
///
/// # Passwords
///
/// Users sometimes ask whether it is safe to use a string of random characters
/// as a password. In principle, all RNGs in Rand implementing `CryptoRng` are
/// suitable as a source of randomness for generating passwords (if they are
/// properly seeded), but it is more conservative to only use randomness
/// directly from the operating system via the `getrandom` crate, or the
/// corresponding bindings of a crypto library.
///
/// When generating passwords or keys, it is important to consider the threat
/// model and in some cases the memorability of the password. This is out of
/// scope of the Rand project, and therefore we defer to the following
/// references:
///
/// - [Wikipedia article on Password Strength](https://en.wikipedia.org/wiki/Password_strength)
/// - [Diceware for generating memorable passwords](https://en.wikipedia.org/wiki/Diceware)
#[derive(Debug)]
#[cfg_attr(feature = "serde1", derive(Serialize, Deserialize))]
pub struct Alphanumeric;

// ----- Implementations of distributions -----

impl Distribution<char> for Standard {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> char {
        // A valid `char` is either in the interval `[0, 0xD800)` or
        // `(0xDFFF, 0x11_0000)`. All `char`s must therefore be in
        // `[0, 0x11_0000)` but not in the "gap" `[0xD800, 0xDFFF]` which is
        // reserved for surrogates. This is the size of that gap.
        const GAP_SIZE: u32 = 0xDFFF - 0xD800 + 1;

        // Uniform::new(0, 0x11_0000 - GAP_SIZE) can also be used but it
        // seemed slower.
        let range = Uniform::new(GAP_SIZE, 0x11_0000);

        let mut n = range.sample(rng);
        if n <= 0xDFFF {
            n -= GAP_SIZE;
        }
        unsafe { char::from_u32_unchecked(n) }
    }
}

impl Distribution<char> for Alphanumeric {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> char {
        const RANGE: u32 = 26 + 26 + 10;
        const GEN_ASCII_STR_CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZ\
                abcdefghijklmnopqrstuvwxyz\
                0123456789";
        // We can pick from 62 characters. This is so close to a power of 2, 64,
        // that we can do better than `Uniform`. Use a simple bitshift and
        // rejection sampling. We do not use a bitmask, because for small RNGs
        // the most significant bits are usually of higher quality.
        loop {
            let var = rng.next_u32() >> (32 - 6);
            if var < RANGE {
                return GEN_ASCII_STR_CHARSET[var as usize] as char;
            }
        }
    }
}

impl Distribution<bool> for Standard {
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> bool {
        // We can compare against an arbitrary bit of an u32 to get a bool.
        // Because the least significant bits of a lower quality RNG can have
        // simple patterns, we compare against the most significant bit. This is
        // easiest done using a sign test.
        (rng.next_u32() as i32) < 0
    }
}

macro_rules! tuple_impl {
    // use variables to indicate the arity of the tuple
    ($($tyvar:ident),* ) => {
        // the trailing commas are for the 1 tuple
        impl< $( $tyvar ),* >
            Distribution<( $( $tyvar ),* , )>
            for Standard
            where $( Standard: Distribution<$tyvar> ),*
        {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> ( $( $tyvar ),* , ) {
                (
                    // use the $tyvar's to get the appropriate number of
                    // repeats (they're not actually needed)
                    $(
                        _rng.gen::<$tyvar>()
                    ),*
                    ,
                )
            }
        }
    }
}

impl Distribution<()> for Standard {
    #[allow(clippy::unused_unit)]
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, _: &mut R) -> () {
        ()
    }
}
tuple_impl! {A}
tuple_impl! {A, B}
tuple_impl! {A, B, C}
tuple_impl! {A, B, C, D}
tuple_impl! {A, B, C, D, E}
tuple_impl! {A, B, C, D, E, F}
tuple_impl! {A, B, C, D, E, F, G}
tuple_impl! {A, B, C, D, E, F, G, H}
tuple_impl! {A, B, C, D, E, F, G, H, I}
tuple_impl! {A, B, C, D, E, F, G, H, I, J}
tuple_impl! {A, B, C, D, E, F, G, H, I, J, K}
tuple_impl! {A, B, C, D, E, F, G, H, I, J, K, L}

macro_rules! array_impl {
    // recursive, given at least one type parameter:
    {$n:expr, $t:ident, $($ts:ident,)*} => {
        array_impl!{($n - 1), $($ts,)*}

        impl<T> Distribution<[T; $n]> for Standard where Standard: Distribution<T> {
            #[inline]
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] {
                [_rng.gen::<$t>(), $(_rng.gen::<$ts>()),*]
            }
        }
    };
    // empty case:
    {$n:expr,} => {
        impl<T> Distribution<[T; $n]> for Standard {
            fn sample<R: Rng + ?Sized>(&self, _rng: &mut R) -> [T; $n] { [] }
        }
    };
}

array_impl! {32, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T, T,}

impl<T> Distribution<Option<T>> for Standard
where Standard: Distribution<T>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Option<T> {
        // UFCS is needed here: https://github.com/rust-lang/rust/issues/24066
        if rng.gen::<bool>() {
            Some(rng.gen())
        } else {
            None
        }
    }
}

impl<T> Distribution<Wrapping<T>> for Standard
where Standard: Distribution<T>
{
    #[inline]
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Wrapping<T> {
        Wrapping(rng.gen())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::RngCore;
    #[cfg(all(not(feature = "std"), feature = "alloc"))] use alloc::string::String;

    #[test]
    fn test_misc() {
        let rng: &mut dyn RngCore = &mut crate::test::rng(820);

        rng.sample::<char, _>(Standard);
        rng.sample::<bool, _>(Standard);
    }

    #[cfg(feature = "alloc")]
    #[test]
    fn test_chars() {
        use core::iter;
        let mut rng = crate::test::rng(805);

        // Test by generating a relatively large number of chars, so we also
        // take the rejection sampling path.
        let word: String = iter::repeat(())
            .map(|()| rng.gen::<char>())
            .take(1000)
            .collect();
        assert!(word.len() != 0);
    }

    #[test]
    fn test_alphanumeric() {
        let mut rng = crate::test::rng(806);

        // Test by generating a relatively large number of chars, so we also
        // take the rejection sampling path.
        let mut incorrect = false;
        for _ in 0..100 {
            let c = rng.sample(Alphanumeric);
            incorrect |= !((c >= '0' && c <= '9') ||
                           (c >= 'A' && c <= 'Z') ||
                           (c >= 'a' && c <= 'z') );
        }
        assert!(incorrect == false);
    }

    #[test]
    fn value_stability() {
        fn test_samples<T: Copy + core::fmt::Debug + PartialEq, D: Distribution<T>>(
            distr: &D, zero: T, expected: &[T],
        ) {
            let mut rng = crate::test::rng(807);
            let mut buf = [zero; 5];
            for x in &mut buf {
                *x = rng.sample(&distr);
            }
            assert_eq!(&buf, expected);
        }

        test_samples(&Standard, 'a', &[
            '\u{8cdac}',
            '\u{a346a}',
            '\u{80120}',
            '\u{ed692}',
            '\u{35888}',
        ]);
        test_samples(&Alphanumeric, 'a', &['h', 'm', 'e', '3', 'M']);
        test_samples(&Standard, false, &[true, true, false, true, false]);
        test_samples(&Standard, None as Option<bool>, &[
            Some(true),
            None,
            Some(false),
            None,
            Some(false),
        ]);
        test_samples(&Standard, Wrapping(0i32), &[
            Wrapping(-2074640887),
            Wrapping(-1719949321),
            Wrapping(2018088303),
            Wrapping(-547181756),
            Wrapping(838957336),
        ]);

        // We test only sub-sets of tuple and array impls
        test_samples(&Standard, (), &[(), (), (), (), ()]);
        test_samples(&Standard, (false,), &[
            (true,),
            (true,),
            (false,),
            (true,),
            (false,),
        ]);
        test_samples(&Standard, (false, false), &[
            (true, true),
            (false, true),
            (false, false),
            (true, false),
            (false, false),
        ]);

        test_samples(&Standard, [0u8; 0], &[[], [], [], [], []]);
        test_samples(&Standard, [0u8; 3], &[
            [9, 247, 111],
            [68, 24, 13],
            [174, 19, 194],
            [172, 69, 213],
            [149, 207, 29],
        ]);
    }
}
<|endoftext|>


<commit_before>use crate::util::{process_error, read2, CargoResult, CargoResultExt};
use anyhow::bail;
use jobserver::Client;
use shell_escape::escape;
use std::collections::HashMap;
use std::env;
use std::ffi::{OsStr, OsString};
use std::fmt;
use std::path::Path;
use std::process::{Command, Output, Stdio};

/// A builder object for an external process, similar to `std::process::Command`.
#[derive(Clone, Debug)]
pub struct ProcessBuilder {
    /// The program to execute.
    program: OsString,
    /// A list of arguments to pass to the program.
    args: Vec<OsString>,
    /// Any environment variables that should be set for the program.
    env: HashMap<String, Option<OsString>>,
    /// The directory to run the program from.
    cwd: Option<OsString>,
    /// The `make` jobserver. See the [jobserver crate][jobserver_docs] for
    /// more information.
    ///
    /// [jobserver_docs]: https://docs.rs/jobserver/0.1.6/jobserver/
    jobserver: Option<Client>,
    /// `true` to include environment variable in display.
    display_env_vars: bool,
}

impl fmt::Display for ProcessBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "`")?;

        if self.display_env_vars {
            for (key, val) in self.env.iter() {
                if let Some(val) = val {
                    let val = escape(val.to_string_lossy());
                    if cfg!(windows) {
                        write!(f, "set {}={}&& ", key, val)?;
                    } else {
                        write!(f, "{}={} ", key, val)?;
                    }
                }
            }
        }

        write!(f, "{}", self.program.to_string_lossy())?;

        for arg in &self.args {
            write!(f, " {}", escape(arg.to_string_lossy()))?;
        }

        write!(f, "`")
    }
}

impl ProcessBuilder {
    /// (chainable) Sets the executable for the process.
    pub fn program<T: AsRef<OsStr>>(&mut self, program: T) -> &mut ProcessBuilder {
        self.program = program.as_ref().to_os_string();
        self
    }

    /// (chainable) Adds `arg` to the args list.
    pub fn arg<T: AsRef<OsStr>>(&mut self, arg: T) -> &mut ProcessBuilder {
        self.args.push(arg.as_ref().to_os_string());
        self
    }

    /// (chainable) Adds multiple `args` to the args list.
    pub fn args<T: AsRef<OsStr>>(&mut self, args: &[T]) -> &mut ProcessBuilder {
        self.args
            .extend(args.iter().map(|t| t.as_ref().to_os_string()));
        self
    }

    /// (chainable) Replaces the args list with the given `args`.
    pub fn args_replace<T: AsRef<OsStr>>(&mut self, args: &[T]) -> &mut ProcessBuilder {
        self.args = args.iter().map(|t| t.as_ref().to_os_string()).collect();
        self
    }

    /// (chainable) Sets the current working directory of the process.
    pub fn cwd<T: AsRef<OsStr>>(&mut self, path: T) -> &mut ProcessBuilder {
        self.cwd = Some(path.as_ref().to_os_string());
        self
    }

    /// (chainable) Sets an environment variable for the process.
    pub fn env<T: AsRef<OsStr>>(&mut self, key: &str, val: T) -> &mut ProcessBuilder {
        self.env
            .insert(key.to_string(), Some(val.as_ref().to_os_string()));
        self
    }

    /// (chainable) Unsets an environment variable for the process.
    pub fn env_remove(&mut self, key: &str) -> &mut ProcessBuilder {
        self.env.insert(key.to_string(), None);
        self
    }

    /// Gets the executable name.
    pub fn get_program(&self) -> &OsString {
        &self.program
    }

    /// Gets the program arguments.
    pub fn get_args(&self) -> &[OsString] {
        &self.args
    }

    /// Gets the current working directory for the process.
    pub fn get_cwd(&self) -> Option<&Path> {
        self.cwd.as_ref().map(Path::new)
    }

    /// Gets an environment variable as the process will see it (will inherit from environment
    /// unless explicitally unset).
    pub fn get_env(&self, var: &str) -> Option<OsString> {
        self.env
            .get(var)
            .cloned()
            .or_else(|| Some(env::var_os(var)))
            .and_then(|s| s)
    }

    /// Gets all environment variables explicitly set or unset for the process (not inherited
    /// vars).
    pub fn get_envs(&self) -> &HashMap<String, Option<OsString>> {
        &self.env
    }

    /// Sets the `make` jobserver. See the [jobserver crate][jobserver_docs] for
    /// more information.
    ///
    /// [jobserver_docs]: https://docs.rs/jobserver/0.1.6/jobserver/
    pub fn inherit_jobserver(&mut self, jobserver: &Client) -> &mut Self {
        self.jobserver = Some(jobserver.clone());
        self
    }

    /// Enables environment variable display.
    pub fn display_env_vars(&mut self) -> &mut Self {
        self.display_env_vars = true;
        self
    }

    /// Runs the process, waiting for completion, and mapping non-success exit codes to an error.
    pub fn exec(&self) -> CargoResult<()> {
        let mut command = self.build_command();
        let exit = command.status().chain_err(|| {
            process_error(&format!("could not execute process {}", self), None, None)
        })?;

        if exit.success() {
            Ok(())
        } else {
            Err(process_error(
                &format!("process didn't exit successfully: {}", self),
                Some(exit),
                None,
            )
            .into())
        }
    }

    /// Replaces the current process with the target process.
    ///
    /// On Unix, this executes the process using the Unix syscall `execvp`, which will block
    /// this process, and will only return if there is an error.
    ///
    /// On Windows this isn't technically possible. Instead we emulate it to the best of our
    /// ability. One aspect we fix here is that we specify a handler for the Ctrl-C handler.
    /// In doing so (and by effectively ignoring it) we should emulate proxying Ctrl-C
    /// handling to the application at hand, which will either terminate or handle it itself.
    /// According to Microsoft's documentation at
    /// <https://docs.microsoft.com/en-us/windows/console/ctrl-c-and-ctrl-break-signals>.
    /// the Ctrl-C signal is sent to all processes attached to a terminal, which should
    /// include our child process. If the child terminates then we'll reap them in Cargo
    /// pretty quickly, and if the child handles the signal then we won't terminate
    /// (and we shouldn't!) until the process itself later exits.
    pub fn exec_replace(&self) -> CargoResult<()> {
        imp::exec_replace(self)
    }

    /// Executes the process, returning the stdio output, or an error if non-zero exit status.
    pub fn exec_with_output(&self) -> CargoResult<Output> {
        let mut command = self.build_command();

        let output = command.output().chain_err(|| {
            process_error(&format!("could not execute process {}", self), None, None)
        })?;

        if output.status.success() {
            Ok(output)
        } else {
            Err(process_error(
                &format!("process didn't exit successfully: {}", self),
                Some(output.status),
                Some(&output),
            )
            .into())
        }
    }

    /// Executes a command, passing each line of stdout and stderr to the supplied callbacks, which
    /// can mutate the string data.
    ///
    /// If any invocations of these function return an error, it will be propagated.
    ///
    /// If `capture_output` is true, then all the output will also be buffered
    /// and stored in the returned `Output` object. If it is false, no caching
    /// is done, and the callbacks are solely responsible for handling the
    /// output.
    pub fn exec_with_streaming(
        &self,
        on_stdout_line: &mut dyn FnMut(&str) -> CargoResult<()>,
        on_stderr_line: &mut dyn FnMut(&str) -> CargoResult<()>,
        capture_output: bool,
    ) -> CargoResult<Output> {
        let mut stdout = Vec::new();
        let mut stderr = Vec::new();

        let mut cmd = self.build_command();
        cmd.stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .stdin(Stdio::null());

        let mut callback_error = None;
        let status = (|| {
            let mut child = cmd.spawn()?;
            let out = child.stdout.take().unwrap();
            let err = child.stderr.take().unwrap();
            read2(out, err, &mut |is_out, data, eof| {
                let idx = if eof {
                    data.len()
                } else {
                    match data.iter().rposition(|b| *b == b'\n') {
                        Some(i) => i + 1,
                        None => return,
                    }
                };
                {
                    // scope for new_lines
                    let new_lines = if capture_output {
                        let dst = if is_out { &mut stdout } else { &mut stderr };
                        let start = dst.len();
                        let data = data.drain(..idx);
                        dst.extend(data);
                        &dst[start..]
                    } else {
                        &data[..idx]
                    };
                    for line in String::from_utf8_lossy(new_lines).lines() {
                        if callback_error.is_some() {
                            break;
                        }
                        let callback_result = if is_out {
                            on_stdout_line(line)
                        } else {
                            on_stderr_line(line)
                        };
                        if let Err(e) = callback_result {
                            callback_error = Some(e);
                        }
                    }
                }
                if !capture_output {
                    data.drain(..idx);
                }
            })?;
            child.wait()
        })()
        .chain_err(|| process_error(&format!("could not execute process {}", self), None, None))?;
        let output = Output {
            stdout,
            stderr,
            status,
        };

        {
            let to_print = if capture_output { Some(&output) } else { None };
            if let Some(e) = callback_error {
                let cx = process_error(
                    &format!("failed to parse process output: {}", self),
                    Some(output.status),
                    to_print,
                );
                bail!(anyhow::Error::new(cx).context(e));
            } else if !output.status.success() {
                bail!(process_error(
                    &format!("process didn't exit successfully: {}", self),
                    Some(output.status),
                    to_print,
                ));
            }
        }

        Ok(output)
    }

    /// Converts `ProcessBuilder` into a `std::process::Command`, and handles the jobserver, if
    /// present.
    pub fn build_command(&self) -> Command {
        let mut command = Command::new(&self.program);
        if let Some(cwd) = self.get_cwd() {
            command.current_dir(cwd);
        }
        for arg in &self.args {
            command.arg(arg);
        }
        for (k, v) in &self.env {
            match *v {
                Some(ref v) => {
                    command.env(k, v);
                }
                None => {
                    command.env_remove(k);
                }
            }
        }
        if let Some(ref c) = self.jobserver {
            c.configure(&mut command);
        }
        command
    }
}

/// A helper function to create a `ProcessBuilder`.
pub fn process<T: AsRef<OsStr>>(cmd: T) -> ProcessBuilder {
    ProcessBuilder {
        program: cmd.as_ref().to_os_string(),
        args: Vec::new(),
        cwd: None,
        env: HashMap::new(),
        jobserver: None,
        display_env_vars: false,
    }
}

#[cfg(unix)]
mod imp {
    use crate::util::{process_error, ProcessBuilder};
    use crate::CargoResult;
    use std::os::unix::process::CommandExt;

    pub fn exec_replace(process_builder: &ProcessBuilder) -> CargoResult<()> {
        let mut command = process_builder.build_command();
        let error = command.exec();
        Err(anyhow::Error::from(error)
            .context(process_error(
                &format!("could not execute process {}", process_builder),
                None,
                None,
            ))
            .into())
    }
}

#[cfg(windows)]
mod imp {
    use crate::util::{process_error, ProcessBuilder};
    use crate::CargoResult;
    use winapi::shared::minwindef::{BOOL, DWORD, FALSE, TRUE};
    use winapi::um::consoleapi::SetConsoleCtrlHandler;

    unsafe extern "system" fn ctrlc_handler(_: DWORD) -> BOOL {
        // Do nothing; let the child process handle it.
        TRUE
    }

    pub fn exec_replace(process_builder: &ProcessBuilder) -> CargoResult<()> {
        unsafe {
            if SetConsoleCtrlHandler(Some(ctrlc_handler), TRUE) == FALSE {
                return Err(process_error("Could not set Ctrl-C handler.", None, None).into());
            }
        }

        // Just execute the process as normal.
        process_builder.exec()
    }
}
<commit_msg>Keep environment variables in a BTreeMap to preserve sort order<commit_after>use crate::util::{process_error, read2, CargoResult, CargoResultExt};
use anyhow::bail;
use jobserver::Client;
use shell_escape::escape;
use std::collections::BTreeMap;
use std::env;
use std::ffi::{OsStr, OsString};
use std::fmt;
use std::path::Path;
use std::process::{Command, Output, Stdio};

/// A builder object for an external process, similar to `std::process::Command`.
#[derive(Clone, Debug)]
pub struct ProcessBuilder {
    /// The program to execute.
    program: OsString,
    /// A list of arguments to pass to the program.
    args: Vec<OsString>,
    /// Any environment variables that should be set for the program.
    env: BTreeMap<String, Option<OsString>>,
    /// The directory to run the program from.
    cwd: Option<OsString>,
    /// The `make` jobserver. See the [jobserver crate][jobserver_docs] for
    /// more information.
    ///
    /// [jobserver_docs]: https://docs.rs/jobserver/0.1.6/jobserver/
    jobserver: Option<Client>,
    /// `true` to include environment variable in display.
    display_env_vars: bool,
}

impl fmt::Display for ProcessBuilder {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "`")?;

        if self.display_env_vars {
            for (key, val) in self.env.iter() {
                if let Some(val) = val {
                    let val = escape(val.to_string_lossy());
                    if cfg!(windows) {
                        write!(f, "set {}={}&& ", key, val)?;
                    } else {
                        write!(f, "{}={} ", key, val)?;
                    }
                }
            }
        }

        write!(f, "{}", self.program.to_string_lossy())?;

        for arg in &self.args {
            write!(f, " {}", escape(arg.to_string_lossy()))?;
        }

        write!(f, "`")
    }
}

impl ProcessBuilder {
    /// (chainable) Sets the executable for the process.
    pub fn program<T: AsRef<OsStr>>(&mut self, program: T) -> &mut ProcessBuilder {
        self.program = program.as_ref().to_os_string();
        self
    }

    /// (chainable) Adds `arg` to the args list.
    pub fn arg<T: AsRef<OsStr>>(&mut self, arg: T) -> &mut ProcessBuilder {
        self.args.push(arg.as_ref().to_os_string());
        self
    }

    /// (chainable) Adds multiple `args` to the args list.
    pub fn args<T: AsRef<OsStr>>(&mut self, args: &[T]) -> &mut ProcessBuilder {
        self.args
            .extend(args.iter().map(|t| t.as_ref().to_os_string()));
        self
    }

    /// (chainable) Replaces the args list with the given `args`.
    pub fn args_replace<T: AsRef<OsStr>>(&mut self, args: &[T]) -> &mut ProcessBuilder {
        self.args = args.iter().map(|t| t.as_ref().to_os_string()).collect();
        self
    }

    /// (chainable) Sets the current working directory of the process.
    pub fn cwd<T: AsRef<OsStr>>(&mut self, path: T) -> &mut ProcessBuilder {
        self.cwd = Some(path.as_ref().to_os_string());
        self
    }

    /// (chainable) Sets an environment variable for the process.
    pub fn env<T: AsRef<OsStr>>(&mut self, key: &str, val: T) -> &mut ProcessBuilder {
        self.env
            .insert(key.to_string(), Some(val.as_ref().to_os_string()));
        self
    }

    /// (chainable) Unsets an environment variable for the process.
    pub fn env_remove(&mut self, key: &str) -> &mut ProcessBuilder {
        self.env.insert(key.to_string(), None);
        self
    }

    /// Gets the executable name.
    pub fn get_program(&self) -> &OsString {
        &self.program
    }

    /// Gets the program arguments.
    pub fn get_args(&self) -> &[OsString] {
        &self.args
    }

    /// Gets the current working directory for the process.
    pub fn get_cwd(&self) -> Option<&Path> {
        self.cwd.as_ref().map(Path::new)
    }

    /// Gets an environment variable as the process will see it (will inherit from environment
    /// unless explicitally unset).
    pub fn get_env(&self, var: &str) -> Option<OsString> {
        self.env
            .get(var)
            .cloned()
            .or_else(|| Some(env::var_os(var)))
            .and_then(|s| s)
    }

    /// Gets all environment variables explicitly set or unset for the process (not inherited
    /// vars).
    pub fn get_envs(&self) -> &BTreeMap<String, Option<OsString>> {
        &self.env
    }

    /// Sets the `make` jobserver. See the [jobserver crate][jobserver_docs] for
    /// more information.
    ///
    /// [jobserver_docs]: https://docs.rs/jobserver/0.1.6/jobserver/
    pub fn inherit_jobserver(&mut self, jobserver: &Client) -> &mut Self {
        self.jobserver = Some(jobserver.clone());
        self
    }

    /// Enables environment variable display.
    pub fn display_env_vars(&mut self) -> &mut Self {
        self.display_env_vars = true;
        self
    }

    /// Runs the process, waiting for completion, and mapping non-success exit codes to an error.
    pub fn exec(&self) -> CargoResult<()> {
        let mut command = self.build_command();
        let exit = command.status().chain_err(|| {
            process_error(&format!("could not execute process {}", self), None, None)
        })?;

        if exit.success() {
            Ok(())
        } else {
            Err(process_error(
                &format!("process didn't exit successfully: {}", self),
                Some(exit),
                None,
            )
            .into())
        }
    }

    /// Replaces the current process with the target process.
    ///
    /// On Unix, this executes the process using the Unix syscall `execvp`, which will block
    /// this process, and will only return if there is an error.
    ///
    /// On Windows this isn't technically possible. Instead we emulate it to the best of our
    /// ability. One aspect we fix here is that we specify a handler for the Ctrl-C handler.
    /// In doing so (and by effectively ignoring it) we should emulate proxying Ctrl-C
    /// handling to the application at hand, which will either terminate or handle it itself.
    /// According to Microsoft's documentation at
    /// <https://docs.microsoft.com/en-us/windows/console/ctrl-c-and-ctrl-break-signals>.
    /// the Ctrl-C signal is sent to all processes attached to a terminal, which should
    /// include our child process. If the child terminates then we'll reap them in Cargo
    /// pretty quickly, and if the child handles the signal then we won't terminate
    /// (and we shouldn't!) until the process itself later exits.
    pub fn exec_replace(&self) -> CargoResult<()> {
        imp::exec_replace(self)
    }

    /// Executes the process, returning the stdio output, or an error if non-zero exit status.
    pub fn exec_with_output(&self) -> CargoResult<Output> {
        let mut command = self.build_command();

        let output = command.output().chain_err(|| {
            process_error(&format!("could not execute process {}", self), None, None)
        })?;

        if output.status.success() {
            Ok(output)
        } else {
            Err(process_error(
                &format!("process didn't exit successfully: {}", self),
                Some(output.status),
                Some(&output),
            )
            .into())
        }
    }

    /// Executes a command, passing each line of stdout and stderr to the supplied callbacks, which
    /// can mutate the string data.
    ///
    /// If any invocations of these function return an error, it will be propagated.
    ///
    /// If `capture_output` is true, then all the output will also be buffered
    /// and stored in the returned `Output` object. If it is false, no caching
    /// is done, and the callbacks are solely responsible for handling the
    /// output.
    pub fn exec_with_streaming(
        &self,
        on_stdout_line: &mut dyn FnMut(&str) -> CargoResult<()>,
        on_stderr_line: &mut dyn FnMut(&str) -> CargoResult<()>,
        capture_output: bool,
    ) -> CargoResult<Output> {
        let mut stdout = Vec::new();
        let mut stderr = Vec::new();

        let mut cmd = self.build_command();
        cmd.stdout(Stdio::piped())
            .stderr(Stdio::piped())
            .stdin(Stdio::null());

        let mut callback_error = None;
        let status = (|| {
            let mut child = cmd.spawn()?;
            let out = child.stdout.take().unwrap();
            let err = child.stderr.take().unwrap();
            read2(out, err, &mut |is_out, data, eof| {
                let idx = if eof {
                    data.len()
                } else {
                    match data.iter().rposition(|b| *b == b'\n') {
                        Some(i) => i + 1,
                        None => return,
                    }
                };
                {
                    // scope for new_lines
                    let new_lines = if capture_output {
                        let dst = if is_out { &mut stdout } else { &mut stderr };
                        let start = dst.len();
                        let data = data.drain(..idx);
                        dst.extend(data);
                        &dst[start..]
                    } else {
                        &data[..idx]
                    };
                    for line in String::from_utf8_lossy(new_lines).lines() {
                        if callback_error.is_some() {
                            break;
                        }
                        let callback_result = if is_out {
                            on_stdout_line(line)
                        } else {
                            on_stderr_line(line)
                        };
                        if let Err(e) = callback_result {
                            callback_error = Some(e);
                        }
                    }
                }
                if !capture_output {
                    data.drain(..idx);
                }
            })?;
            child.wait()
        })()
        .chain_err(|| process_error(&format!("could not execute process {}", self), None, None))?;
        let output = Output {
            stdout,
            stderr,
            status,
        };

        {
            let to_print = if capture_output { Some(&output) } else { None };
            if let Some(e) = callback_error {
                let cx = process_error(
                    &format!("failed to parse process output: {}", self),
                    Some(output.status),
                    to_print,
                );
                bail!(anyhow::Error::new(cx).context(e));
            } else if !output.status.success() {
                bail!(process_error(
                    &format!("process didn't exit successfully: {}", self),
                    Some(output.status),
                    to_print,
                ));
            }
        }

        Ok(output)
    }

    /// Converts `ProcessBuilder` into a `std::process::Command`, and handles the jobserver, if
    /// present.
    pub fn build_command(&self) -> Command {
        let mut command = Command::new(&self.program);
        if let Some(cwd) = self.get_cwd() {
            command.current_dir(cwd);
        }
        for arg in &self.args {
            command.arg(arg);
        }
        for (k, v) in &self.env {
            match *v {
                Some(ref v) => {
                    command.env(k, v);
                }
                None => {
                    command.env_remove(k);
                }
            }
        }
        if let Some(ref c) = self.jobserver {
            c.configure(&mut command);
        }
        command
    }
}

/// A helper function to create a `ProcessBuilder`.
pub fn process<T: AsRef<OsStr>>(cmd: T) -> ProcessBuilder {
    ProcessBuilder {
        program: cmd.as_ref().to_os_string(),
        args: Vec::new(),
        cwd: None,
        env: BTreeMap::new(),
        jobserver: None,
        display_env_vars: false,
    }
}

#[cfg(unix)]
mod imp {
    use crate::util::{process_error, ProcessBuilder};
    use crate::CargoResult;
    use std::os::unix::process::CommandExt;

    pub fn exec_replace(process_builder: &ProcessBuilder) -> CargoResult<()> {
        let mut command = process_builder.build_command();
        let error = command.exec();
        Err(anyhow::Error::from(error)
            .context(process_error(
                &format!("could not execute process {}", process_builder),
                None,
                None,
            ))
            .into())
    }
}

#[cfg(windows)]
mod imp {
    use crate::util::{process_error, ProcessBuilder};
    use crate::CargoResult;
    use winapi::shared::minwindef::{BOOL, DWORD, FALSE, TRUE};
    use winapi::um::consoleapi::SetConsoleCtrlHandler;

    unsafe extern "system" fn ctrlc_handler(_: DWORD) -> BOOL {
        // Do nothing; let the child process handle it.
        TRUE
    }

    pub fn exec_replace(process_builder: &ProcessBuilder) -> CargoResult<()> {
        unsafe {
            if SetConsoleCtrlHandler(Some(ctrlc_handler), TRUE) == FALSE {
                return Err(process_error("Could not set Ctrl-C handler.", None, None).into());
            }
        }

        // Just execute the process as normal.
        process_builder.exec()
    }
}
<|endoftext|>


<commit_before>use std::io::{stdout, stderr};

use argparse::{ArgumentParser};
use argparse::{StoreTrue, List, StoreOption, Store};
use unshare::{Command, Namespace};

use options::build_mode::build_mode;
use options::version_hash;
use container::nsutil::{set_namespace};
use process_util::{run_and_wait, convert_status};

use super::network;
use super::build::{build_container};
use super::wrap::Wrapper;
use launcher::Context;
use launcher::volumes::prepare_volumes;

pub fn run_command(context: &Context, mut args: Vec<String>)
    -> Result<i32, String>
{
    let mut cmdargs = Vec::<String>::new();
    let mut container = "".to_string();
    let mut command = "".to_string();
    let mut copy = false;
    let mut bmode = context.build_mode;
    {
        args.insert(0, "vagga _run".to_string());
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Runs arbitrary command inside the container
            ");
        ap.refer(&mut copy)
            .add_option(&["-W", "--writeable"], StoreTrue,
                "Create translient writeable container for running the command.
                 Currently we use hard-linked copy of the container, so it's
                 dangerous for some operations. Still it's ok for installing
                 packages or similar tasks");
        build_mode(&mut ap, &mut bmode);
        ap.refer(&mut container)
            .add_argument("container", Store,
                "Container to run command in")
            .required();
        ap.refer(&mut command)
            .add_argument("command", Store,
                "Command to run inside the container")
            .required();
        ap.refer(&mut cmdargs)
            .add_argument("arg", List, "Arguments to the command");

        ap.stop_on_first_argument(true);
        match ap.parse(args.clone(), &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    }
    let cinfo = try!(context.config.get_container(&container));
    let ver = try!(build_container(context, &container, bmode));
    try!(prepare_volumes(cinfo.volumes.values(), context));

    if context.isolate_network {
        try_msg!(network::isolate_network(),
            "Cannot setup isolated network: {err}");
    }

    let mut cmd: Command = Wrapper::new(Some(&ver), &context.settings);
    cmd.workdir(&context.workdir);
    cmd.arg("_run");
    cmd.args(&args[1..]);
    try!(cmd.map_users_for(cinfo, &context.settings));
    cmd.gid(0);
    cmd.groups(Vec::new());
    let res = run_and_wait(&mut cmd).map(convert_status);

    if copy {
        let mut cmd: Command = Wrapper::new(None, &context.settings);
        cmd.workdir(&context.workdir);  // TODO(tailhook) why is it needed?
        cmd.max_uidmap();
        cmd.gid(0);
        cmd.groups(Vec::new());
        cmd.arg("_clean").arg("--transient");
        match cmd.status() {
            Ok(s) if s.success() => {}
            Ok(s) => warn!("The `vagga _clean --transient` {}", s),
            Err(e) => warn!("Failed to run `vagga _clean --transient`: {}", e),
        }

    }
    return res;
}

pub fn run_in_netns(context: &Context, cname: String, mut args: Vec<String>)
    -> Result<i32, String>
{
    let mut cmdargs: Vec<String> = vec!();
    let mut container = "".to_string();
    let mut pid = None;
    let mut bmode = context.build_mode;
    {
        args.insert(0, "vagga ".to_string() + &cname);
        let mut ap = ArgumentParser::new();
        ap.set_description(
            "Run command (or shell) in one of the vagga's network namespaces");
        ap.refer(&mut pid)
            .add_option(&["--pid"], StoreOption, "
                Run in the namespace of the process with PID.
                By default you get shell in the \"gateway\" namespace.
                ");
        build_mode(&mut ap, &mut bmode);
        ap.refer(&mut container)
            .add_argument("container", Store,
                "Container to run command in")
            .required();
        ap.refer(&mut cmdargs)
            .add_argument("command", List,
                "Command (with arguments) to run inside container")
            .required();

        ap.stop_on_first_argument(true);
        match ap.parse(args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    }
    let ver = try!(build_container(context, &container, bmode));
    try!(network::join_gateway_namespaces());
    if let Some::<i32>(pid) = pid {
        try!(set_namespace(format!("/proc/{}/ns/net", pid), Namespace::Net)
            .map_err(|e| format!("Error setting networkns: {}", e)));
    }
    let mut cmd: Command = Wrapper::new(Some(&ver), &context.settings);
    cmd.workdir(&context.workdir);
    cmd.arg(cname);
    cmd.arg(container.clone());
    cmd.args(&cmdargs);
    run_and_wait(&mut cmd).map(convert_status)
}

pub fn version_hash(ctx: &Context, cname: &str, mut args: Vec<String>)
    -> Result<i32, String>
{
    args.insert(0, "vagga _version_hash".to_string());
    let opt = match version_hash::Options::parse(&args, false) {
        Ok(x) => x,
        Err(e) => return Ok(e),
    };
    let mut cmd: Command = Wrapper::new(None, &ctx.settings);
    cmd.workdir(&ctx.workdir);
    try!(cmd.map_users_for(
        &ctx.config.get_container(&opt.container).unwrap(),
        &ctx.settings));
    cmd.gid(0);
    cmd.groups(Vec::new());
    cmd.arg(&cname).args(&args[1..]);
    cmd.status()
    .map(convert_status)
    .map_err(|e| format!("Error running `vagga_wrapper {}`: {}",
                         cname, e))
}
pub fn passthrough(ctx: &Context, cname: &str, args: Vec<String>)
    -> Result<i32, String>
{
    let mut cmd: Command = Wrapper::new(None, &ctx.settings);
    cmd.workdir(&ctx.workdir);
    cmd.max_uidmap();
    cmd.gid(0);
    cmd.groups(Vec::new());
    cmd.arg(&cname).args(&args);
    cmd.status()
    .map(convert_status)
    .map_err(|e| format!("Error running `vagga_wrapper {}`: {}",
                         cname, e))
}
<commit_msg>'vagga _version_hash' shouldn't panic when container is not found<commit_after>use std::io::{stdout, stderr};

use argparse::{ArgumentParser};
use argparse::{StoreTrue, List, StoreOption, Store};
use unshare::{Command, Namespace};

use options::build_mode::build_mode;
use options::version_hash;
use container::nsutil::{set_namespace};
use process_util::{run_and_wait, convert_status};

use super::network;
use super::build::{build_container};
use super::wrap::Wrapper;
use launcher::Context;
use launcher::volumes::prepare_volumes;

pub fn run_command(context: &Context, mut args: Vec<String>)
    -> Result<i32, String>
{
    let mut cmdargs = Vec::<String>::new();
    let mut container = "".to_string();
    let mut command = "".to_string();
    let mut copy = false;
    let mut bmode = context.build_mode;
    {
        args.insert(0, "vagga _run".to_string());
        let mut ap = ArgumentParser::new();
        ap.set_description("
            Runs arbitrary command inside the container
            ");
        ap.refer(&mut copy)
            .add_option(&["-W", "--writeable"], StoreTrue,
                "Create translient writeable container for running the command.
                 Currently we use hard-linked copy of the container, so it's
                 dangerous for some operations. Still it's ok for installing
                 packages or similar tasks");
        build_mode(&mut ap, &mut bmode);
        ap.refer(&mut container)
            .add_argument("container", Store,
                "Container to run command in")
            .required();
        ap.refer(&mut command)
            .add_argument("command", Store,
                "Command to run inside the container")
            .required();
        ap.refer(&mut cmdargs)
            .add_argument("arg", List, "Arguments to the command");

        ap.stop_on_first_argument(true);
        match ap.parse(args.clone(), &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    }
    let cinfo = try!(context.config.get_container(&container));
    let ver = try!(build_container(context, &container, bmode));
    try!(prepare_volumes(cinfo.volumes.values(), context));

    if context.isolate_network {
        try_msg!(network::isolate_network(),
            "Cannot setup isolated network: {err}");
    }

    let mut cmd: Command = Wrapper::new(Some(&ver), &context.settings);
    cmd.workdir(&context.workdir);
    cmd.arg("_run");
    cmd.args(&args[1..]);
    try!(cmd.map_users_for(cinfo, &context.settings));
    cmd.gid(0);
    cmd.groups(Vec::new());
    let res = run_and_wait(&mut cmd).map(convert_status);

    if copy {
        let mut cmd: Command = Wrapper::new(None, &context.settings);
        cmd.workdir(&context.workdir);  // TODO(tailhook) why is it needed?
        cmd.max_uidmap();
        cmd.gid(0);
        cmd.groups(Vec::new());
        cmd.arg("_clean").arg("--transient");
        match cmd.status() {
            Ok(s) if s.success() => {}
            Ok(s) => warn!("The `vagga _clean --transient` {}", s),
            Err(e) => warn!("Failed to run `vagga _clean --transient`: {}", e),
        }

    }
    return res;
}

pub fn run_in_netns(context: &Context, cname: String, mut args: Vec<String>)
    -> Result<i32, String>
{
    let mut cmdargs: Vec<String> = vec!();
    let mut container = "".to_string();
    let mut pid = None;
    let mut bmode = context.build_mode;
    {
        args.insert(0, "vagga ".to_string() + &cname);
        let mut ap = ArgumentParser::new();
        ap.set_description(
            "Run command (or shell) in one of the vagga's network namespaces");
        ap.refer(&mut pid)
            .add_option(&["--pid"], StoreOption, "
                Run in the namespace of the process with PID.
                By default you get shell in the \"gateway\" namespace.
                ");
        build_mode(&mut ap, &mut bmode);
        ap.refer(&mut container)
            .add_argument("container", Store,
                "Container to run command in")
            .required();
        ap.refer(&mut cmdargs)
            .add_argument("command", List,
                "Command (with arguments) to run inside container")
            .required();

        ap.stop_on_first_argument(true);
        match ap.parse(args, &mut stdout(), &mut stderr()) {
            Ok(()) => {}
            Err(0) => return Ok(0),
            Err(_) => {
                return Ok(122);
            }
        }
    }
    let ver = try!(build_container(context, &container, bmode));
    try!(network::join_gateway_namespaces());
    if let Some::<i32>(pid) = pid {
        try!(set_namespace(format!("/proc/{}/ns/net", pid), Namespace::Net)
            .map_err(|e| format!("Error setting networkns: {}", e)));
    }
    let mut cmd: Command = Wrapper::new(Some(&ver), &context.settings);
    cmd.workdir(&context.workdir);
    cmd.arg(cname);
    cmd.arg(container.clone());
    cmd.args(&cmdargs);
    run_and_wait(&mut cmd).map(convert_status)
}

pub fn version_hash(ctx: &Context, cname: &str, mut args: Vec<String>)
    -> Result<i32, String>
{
    args.insert(0, "vagga _version_hash".to_string());
    let opt = match version_hash::Options::parse(&args, false) {
        Ok(x) => x,
        Err(e) => return Ok(e),
    };
    let mut cmd: Command = Wrapper::new(None, &ctx.settings);
    cmd.workdir(&ctx.workdir);
    try!(cmd.map_users_for(
        try!(ctx.config.get_container(&opt.container)),
        &ctx.settings));
    cmd.gid(0);
    cmd.groups(Vec::new());
    cmd.arg(&cname).args(&args[1..]);
    cmd.status()
    .map(convert_status)
    .map_err(|e| format!("Error running `vagga_wrapper {}`: {}",
                         cname, e))
}
pub fn passthrough(ctx: &Context, cname: &str, args: Vec<String>)
    -> Result<i32, String>
{
    let mut cmd: Command = Wrapper::new(None, &ctx.settings);
    cmd.workdir(&ctx.workdir);
    cmd.max_uidmap();
    cmd.gid(0);
    cmd.groups(Vec::new());
    cmd.arg(&cname).args(&args);
    cmd.status()
    .map(convert_status)
    .map_err(|e| format!("Error running `vagga_wrapper {}`: {}",
                         cname, e))
}
<|endoftext|>


<commit_before>use std::io::Read;

use rustc_serialize::json::Json;
use oauthcli::SignatureMethod::HmacSha1;
use oauthcli::OAuthAuthorizationHeaderBuilder;
use url::Url;
use url::form_urlencoded::Serializer;
use hyper::{Client, Error};
use hyper::client::Response;
use hyper::method::Method;
use hyper::status::StatusCode;
use hyper::header::{Headers, Authorization, ContentType};

pub struct TwitterClient {
	consumer_key: String,
	consumer_secret: String,
	access_key: String,
	access_secret: String,
	client: Client
}

#[derive(Debug, Clone)]
pub enum StatusUpdateError {
	Duplicated,
	RateLimitExceeded,
	InvalidTweet,
	InvalidResponse,
	Network,
	Unauthorized,
	Unknown(String)
}

impl TwitterClient {

	pub fn new(consumer_key: String, consumer_secret: String, access_key: String, access_secret: String) -> TwitterClient
	{
		TwitterClient {
			consumer_key: consumer_key,
			consumer_secret: consumer_secret,
			access_key: access_key,
			access_secret: access_secret,
			client: Client::new()
		}
	}

	pub fn is_valid(&self) -> bool
	{
		// TODO: implement token validation
		return true;
	}

	pub fn update_status(&self, message: &str, in_reply_to: Option<u64>)
	 -> Result<u64, StatusUpdateError>
	{
		let api_url = "https://api.twitter.com/1.1/statuses/update.json";
		let prev_str_opt = in_reply_to.map(|i| i.to_string());
		let mut args = vec![("status", message)];
		if prev_str_opt.is_some() {
			args.push(("in_reply_to_status_id", prev_str_opt.as_ref().unwrap()));
		}
		let result = self.request(Method::Post, api_url, args);

		let mut res = try!(result.map_err(|_| StatusUpdateError::Network));

		match res.status {

			StatusCode::Forbidden => {

				let mut body = String::new();

				try!(res.read_to_string(&mut body).map_err(|_| StatusUpdateError::Network));

				if body.contains("140") {
					return Err(StatusUpdateError::InvalidTweet);
				} else if body.contains("duplicate") {
					return Err(StatusUpdateError::Duplicated);
				}

				return Err(StatusUpdateError::Unknown(body));
			},

			StatusCode::Ok => {

				let mut body = String::new();

				try!(res.read_to_string(&mut body).map_err(|_| StatusUpdateError::Network));

				let json = try!(Json::from_str(&body).map_err(|_| StatusUpdateError::InvalidResponse));
				let id_obj = try!(json.find("id").ok_or(StatusUpdateError::InvalidResponse));
				let id = try!(id_obj.as_u64().ok_or(StatusUpdateError::InvalidResponse));

				return Ok(id);
			},

			StatusCode::TooManyRequests => return Err(StatusUpdateError::RateLimitExceeded),
			StatusCode::Unauthorized => return Err(StatusUpdateError::Unauthorized),
			_ => return Err(StatusUpdateError::Unknown(format!("Unknown status: {}", res.status)))
		};
	}

	fn request(&self, method: Method, api_url: &str, args: Vec<(&str, &str)>)
	 -> Result<Response, Error>
	{
		match method {

			Method::Get => {

				let mut headers = Headers::new();
				let oauth_header = self.construct_oauth_header("GET", api_url, args);
				headers.set(Authorization(oauth_header));

				let result = self.client.get(api_url).headers(headers).send();

				return result;

			},

			Method::Post => {

				let content_type = ContentType("application/x-www-form-urlencoded".parse().unwrap());

				let mut args_serializer = Serializer::new(String::new());
				for &(ref k, ref v) in &args {
					args_serializer.append_pair(k, v);
				}
				let args_serialized = args_serializer.finish();

				let mut headers = Headers::new();
				let oauth_header = self.construct_oauth_header("POST", api_url, args);
				headers.set(Authorization(oauth_header));
				headers.set(content_type);

				let result = self.client.post(api_url).body(&args_serialized).headers(headers).send();

				return result;

			},

			_ => {
				panic!();
			}
		}
	}

	fn construct_oauth_header(&self, method: &str, api_url: &str, args: Vec<(&str, &str)>)
	 -> String
	{
		let oauth_header = OAuthAuthorizationHeaderBuilder::new(
			method,
			&Url::parse(api_url).unwrap(),
			self.consumer_key.as_ref(),
			self.consumer_secret.as_ref(),
			HmacSha1)
			.token(self.access_key.as_ref(), self.access_secret.as_ref())
			.request_parameters(args.into_iter())
			.finish_for_twitter();

		return oauth_header.to_string();
	}
}
<commit_msg>TwitterClient: Fix the network error caused by reuse of connections<commit_after>use std::io::Read;

use rustc_serialize::json::Json;
use oauthcli::SignatureMethod::HmacSha1;
use oauthcli::OAuthAuthorizationHeaderBuilder;
use url::Url;
use url::form_urlencoded::Serializer;
use hyper::Url as HyperUrl;
use hyper::{Client, Error};
use hyper::client::Response;
use hyper::method::Method;
use hyper::status::StatusCode;
use hyper::header::{Headers, Authorization, ContentType};

pub struct TwitterClient {
	consumer_key: String,
	consumer_secret: String,
	access_key: String,
	access_secret: String,
	client: Client
}

#[derive(Debug, Clone)]
pub enum StatusUpdateError {
	Duplicated,
	RateLimitExceeded,
	InvalidTweet,
	InvalidResponse,
	Network,
	Unauthorized,
	Unknown(String)
}

impl TwitterClient {

	pub fn new(consumer_key: String, consumer_secret: String, access_key: String, access_secret: String) -> TwitterClient
	{
		TwitterClient {
			consumer_key: consumer_key,
			consumer_secret: consumer_secret,
			access_key: access_key,
			access_secret: access_secret,
			client: Client::new()
		}
	}

	pub fn is_valid(&self) -> bool
	{
		// TODO: implement token validation
		return true;
	}

	pub fn update_status(&self, message: &str, in_reply_to: Option<u64>)
	 -> Result<u64, StatusUpdateError>
	{
		let api_url = "https://api.twitter.com/1.1/statuses/update.json";
		let prev_str_opt = in_reply_to.map(|i| i.to_string());
		let mut args = vec![("status", message)];
		if prev_str_opt.is_some() {
			args.push(("in_reply_to_status_id", prev_str_opt.as_ref().unwrap()));
		}
		let result = self.request(Method::Post, api_url, args);

		let mut res = try!(result.map_err(|_| StatusUpdateError::Network));

		match res.status {

			StatusCode::Forbidden => {

				let mut body = String::new();

				try!(res.read_to_string(&mut body).map_err(|_| StatusUpdateError::Network));

				if body.contains("140") {
					return Err(StatusUpdateError::InvalidTweet);
				} else if body.contains("duplicate") {
					return Err(StatusUpdateError::Duplicated);
				}

				return Err(StatusUpdateError::Unknown(body));
			},

			StatusCode::Ok => {

				let mut body = String::new();

				try!(res.read_to_string(&mut body).map_err(|_| StatusUpdateError::Network));

				let json = try!(Json::from_str(&body).map_err(|_| StatusUpdateError::InvalidResponse));
				let id_obj = try!(json.find("id").ok_or(StatusUpdateError::InvalidResponse));
				let id = try!(id_obj.as_u64().ok_or(StatusUpdateError::InvalidResponse));

				return Ok(id);
			},

			StatusCode::TooManyRequests => return Err(StatusUpdateError::RateLimitExceeded),
			StatusCode::Unauthorized => return Err(StatusUpdateError::Unauthorized),
			_ => return Err(StatusUpdateError::Unknown(format!("Unknown status: {}", res.status)))
		};
	}

	fn request(&self, method: Method, api_url: &str, args: Vec<(&str, &str)>)
	 -> Result<Response, Error>
	{
		let mut body = String::new();
		let mut headers = Headers::new();
		let mut url_obj = HyperUrl::parse(api_url).unwrap();

		if method == Method::Post || method == Method::Get {

			let mut args_serializer = Serializer::new(String::new());
			for &(ref k, ref v) in args.iter() {
				args_serializer.append_pair(k, v);
			}
			let args_serialized = args_serializer.finish();

			if method == Method::Post {
				body.push_str(&args_serialized);
				let content_type = ContentType("application/x-www-form-urlencoded".parse().unwrap());
				headers.set(content_type);
			} else {
				url_obj.set_query(Some(&args_serialized));
			}
		}

		let oauth_header = self.construct_oauth_header(method.as_ref(), api_url, args);
		headers.set(Authorization(oauth_header));

		let m = method.clone();
		let h = headers.clone();
		let u = url_obj.clone();
		if let Ok(r) = self.client.request(m, u).headers(h).body(&body).send() {
			return Ok(r);
		}

		// retry (because the first attempt may fail due to reuse of closed sockets)
		// detail: https://github.com/hyperium/hyper/issues/796
		return self.client.request(method, url_obj).headers(headers).body(&body).send();
	}

	fn construct_oauth_header(&self, method: &str, api_url: &str, args: Vec<(&str, &str)>)
	 -> String
	{
		let oauth_header = OAuthAuthorizationHeaderBuilder::new(
			method,
			&Url::parse(api_url).unwrap(),
			self.consumer_key.as_ref(),
			self.consumer_secret.as_ref(),
			HmacSha1)
			.token(self.access_key.as_ref(), self.access_secret.as_ref())
			.request_parameters(args.into_iter())
			.finish_for_twitter();

		return oauth_header.to_string();
	}
}
<|endoftext|>