<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// 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.

//! A wrapper around another RNG that reseeds it after it
//! generates a certain number of random bytes.

use core::default::Default;

use {Rng, SeedableRng};

/// How many bytes of entropy the underling RNG is allowed to generate
/// before it is reseeded
const DEFAULT_GENERATION_THRESHOLD: u64 = 32 * 1024;

/// A wrapper around any RNG which reseeds the underlying RNG after it
/// has generated a certain number of random bytes.
#[derive(Debug)]
pub struct ReseedingRng<R, Rsdr> {
    rng: R,
    generation_threshold: u64,
    bytes_generated: u64,
    /// Controls the behaviour when reseeding the RNG.
    pub reseeder: Rsdr,
}

impl<R: Rng, Rsdr: Reseeder<R>> ReseedingRng<R, Rsdr> {
    /// Create a new `ReseedingRng` with the given parameters.
    ///
    /// # Arguments
    ///
    /// * `rng`: the random number generator to use.
    /// * `generation_threshold`: the number of bytes of entropy at which to reseed the RNG.
    /// * `reseeder`: the reseeding object to use.
    pub fn new(rng: R, generation_threshold: u64, reseeder: Rsdr) -> ReseedingRng<R,Rsdr> {
        ReseedingRng {
            rng: rng,
            generation_threshold: generation_threshold,
            bytes_generated: 0,
            reseeder: reseeder
        }
    }

    /// Reseed the internal RNG if the number of bytes that have been
    /// generated exceed the threshold.
    pub fn reseed_if_necessary(&mut self) {
        if self.bytes_generated >= self.generation_threshold {
            self.reseeder.reseed(&mut self.rng);
            self.bytes_generated = 0;
        }
    }
}

impl<R: Rng, Rsdr: Reseeder<R>> Rng for ReseedingRng<R, Rsdr> {
    fn next_u32(&mut self) -> u32 {
        self.reseed_if_necessary();
        self.bytes_generated += 4;
        self.rng.next_u32()
    }

    fn next_u64(&mut self) -> u64 {
        self.reseed_if_necessary();
        self.bytes_generated += 8;
        self.rng.next_u64()
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.reseed_if_necessary();
        self.bytes_generated += dest.len() as u64;
        self.rng.fill_bytes(dest)
    }
}

impl<S, R: SeedableRng<S>, Rsdr: Reseeder<R> + Default>
     SeedableRng<(Rsdr, S)> for ReseedingRng<R, Rsdr> {
    fn reseed(&mut self, (rsdr, seed): (Rsdr, S)) {
        self.rng.reseed(seed);
        self.reseeder = rsdr;
        self.bytes_generated = 0;
    }

    /// Create a new `ReseedingRng` from the given reseeder and
    /// seed. This uses a default value for `generation_threshold`.
    fn from_seed((rsdr, seed): (Rsdr, S)) -> ReseedingRng<R, Rsdr> {
        ReseedingRng {
            rng: SeedableRng::from_seed(seed),
            generation_threshold: DEFAULT_GENERATION_THRESHOLD,
            bytes_generated: 0,
            reseeder: rsdr
        }
    }
}

/// Something that can be used to reseed an RNG via `ReseedingRng`.
///
/// # Example
///
/// ```rust
/// use rand::{Rng, SeedableRng, StdRng};
/// use rand::reseeding::{Reseeder, ReseedingRng};
///
/// struct TickTockReseeder { tick: bool }
/// impl Reseeder<StdRng> for TickTockReseeder {
///     fn reseed(&mut self, rng: &mut StdRng) {
///         let val = if self.tick {0} else {1};
///         rng.reseed(&[val]);
///         self.tick = !self.tick;
///     }
/// }
/// fn main() {
///     let rsdr = TickTockReseeder { tick: true };
///
///     let inner = StdRng::new().unwrap();
///     let mut rng = ReseedingRng::new(inner, 10, rsdr);
///
///     // this will repeat, because it gets reseeded very regularly.
///     let s: String = rng.gen_ascii_chars().take(100).collect();
///     println!("{}", s);
/// }
///
/// ```
pub trait Reseeder<R> {
    /// Reseed the given RNG.
    fn reseed(&mut self, rng: &mut R);
}

/// Reseed an RNG using a `Default` instance. This reseeds by
/// replacing the RNG with the result of a `Default::default` call.
#[derive(Clone, Copy, Debug)]
pub struct ReseedWithDefault;

impl<R: Rng + Default> Reseeder<R> for ReseedWithDefault {
    fn reseed(&mut self, rng: &mut R) {
        *rng = Default::default();
    }
}
impl Default for ReseedWithDefault {
    fn default() -> ReseedWithDefault { ReseedWithDefault }
}

#[cfg(test)]
mod test {
    use impls;
    use std::default::Default;
    use std::iter::repeat;
    use super::{ReseedingRng, ReseedWithDefault};
    use {SeedableRng, Rng};

    struct Counter {
        i: u32
    }

    impl Rng for Counter {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            // very random
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            impls::next_u64_via_u32(self)
        }

        fn fill_bytes(&mut self, dest: &mut [u8]) {
            impls::fill_bytes_via_u64(self, dest)
        }
    }
    impl Default for Counter {
        fn default() -> Counter {
            Counter { i: 0 }
        }
    }
    impl SeedableRng<u32> for Counter {
        fn reseed(&mut self, seed: u32) {
            self.i = seed;
        }
        fn from_seed(seed: u32) -> Counter {
            Counter { i: seed }
        }
    }
    type MyRng = ReseedingRng<Counter, ReseedWithDefault>;

    #[test]
    fn test_reseeding() {
        let mut rs = ReseedingRng::new(Counter {i:0}, 400, ReseedWithDefault);

        let mut i = 0;
        for _ in 0..1000 {
            assert_eq!(rs.next_u32(), i % 100);
            i += 1;
        }
    }

    #[test]
    fn test_rng_seeded() {
        let mut ra: MyRng = SeedableRng::from_seed((ReseedWithDefault, 2));
        let mut rb: MyRng = SeedableRng::from_seed((ReseedWithDefault, 2));
        assert!(::test::iter_eq(ra.gen_ascii_chars().take(100),
                                rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_reseed() {
        let mut r: MyRng = SeedableRng::from_seed((ReseedWithDefault, 3));
        let string1: String = r.gen_ascii_chars().take(100).collect();

        r.reseed((ReseedWithDefault, 3));

        let string2: String = r.gen_ascii_chars().take(100).collect();
        assert_eq!(string1, string2);
    }

    const FILL_BYTES_V_LEN: usize = 13579;
    #[test]
    fn test_rng_fill_bytes() {
        let mut v = repeat(0u8).take(FILL_BYTES_V_LEN).collect::<Vec<_>>();
        ::test::rng().fill_bytes(&mut v);

        // Sanity test: if we've gotten here, `fill_bytes` has not infinitely
        // recursed.
        assert_eq!(v.len(), FILL_BYTES_V_LEN);

        // To test that `fill_bytes` actually did something, check that the
        // average of `v` is not 0.
        let mut sum = 0.0;
        for &x in v.iter() {
            sum += x as f64;
        }
        assert!(sum / v.len() as f64 != 0.0);
    }
}
<commit_msg>Test reseeding without allocating<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// 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.

//! A wrapper around another RNG that reseeds it after it
//! generates a certain number of random bytes.

use core::default::Default;

use {Rng, SeedableRng};

/// How many bytes of entropy the underling RNG is allowed to generate
/// before it is reseeded
const DEFAULT_GENERATION_THRESHOLD: u64 = 32 * 1024;

/// A wrapper around any RNG which reseeds the underlying RNG after it
/// has generated a certain number of random bytes.
#[derive(Debug)]
pub struct ReseedingRng<R, Rsdr> {
    rng: R,
    generation_threshold: u64,
    bytes_generated: u64,
    /// Controls the behaviour when reseeding the RNG.
    pub reseeder: Rsdr,
}

impl<R: Rng, Rsdr: Reseeder<R>> ReseedingRng<R, Rsdr> {
    /// Create a new `ReseedingRng` with the given parameters.
    ///
    /// # Arguments
    ///
    /// * `rng`: the random number generator to use.
    /// * `generation_threshold`: the number of bytes of entropy at which to reseed the RNG.
    /// * `reseeder`: the reseeding object to use.
    pub fn new(rng: R, generation_threshold: u64, reseeder: Rsdr) -> ReseedingRng<R,Rsdr> {
        ReseedingRng {
            rng: rng,
            generation_threshold: generation_threshold,
            bytes_generated: 0,
            reseeder: reseeder
        }
    }

    /// Reseed the internal RNG if the number of bytes that have been
    /// generated exceed the threshold.
    pub fn reseed_if_necessary(&mut self) {
        if self.bytes_generated >= self.generation_threshold {
            self.reseeder.reseed(&mut self.rng);
            self.bytes_generated = 0;
        }
    }
}

impl<R: Rng, Rsdr: Reseeder<R>> Rng for ReseedingRng<R, Rsdr> {
    fn next_u32(&mut self) -> u32 {
        self.reseed_if_necessary();
        self.bytes_generated += 4;
        self.rng.next_u32()
    }

    fn next_u64(&mut self) -> u64 {
        self.reseed_if_necessary();
        self.bytes_generated += 8;
        self.rng.next_u64()
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.reseed_if_necessary();
        self.bytes_generated += dest.len() as u64;
        self.rng.fill_bytes(dest)
    }
}

impl<S, R: SeedableRng<S>, Rsdr: Reseeder<R> + Default>
     SeedableRng<(Rsdr, S)> for ReseedingRng<R, Rsdr> {
    fn reseed(&mut self, (rsdr, seed): (Rsdr, S)) {
        self.rng.reseed(seed);
        self.reseeder = rsdr;
        self.bytes_generated = 0;
    }

    /// Create a new `ReseedingRng` from the given reseeder and
    /// seed. This uses a default value for `generation_threshold`.
    fn from_seed((rsdr, seed): (Rsdr, S)) -> ReseedingRng<R, Rsdr> {
        ReseedingRng {
            rng: SeedableRng::from_seed(seed),
            generation_threshold: DEFAULT_GENERATION_THRESHOLD,
            bytes_generated: 0,
            reseeder: rsdr
        }
    }
}

/// Something that can be used to reseed an RNG via `ReseedingRng`.
///
/// # Example
///
/// ```rust
/// use rand::{Rng, SeedableRng, StdRng};
/// use rand::reseeding::{Reseeder, ReseedingRng};
///
/// struct TickTockReseeder { tick: bool }
/// impl Reseeder<StdRng> for TickTockReseeder {
///     fn reseed(&mut self, rng: &mut StdRng) {
///         let val = if self.tick {0} else {1};
///         rng.reseed(&[val]);
///         self.tick = !self.tick;
///     }
/// }
/// fn main() {
///     let rsdr = TickTockReseeder { tick: true };
///
///     let inner = StdRng::new().unwrap();
///     let mut rng = ReseedingRng::new(inner, 10, rsdr);
///
///     // this will repeat, because it gets reseeded very regularly.
///     let s: String = rng.gen_ascii_chars().take(100).collect();
///     println!("{}", s);
/// }
///
/// ```
pub trait Reseeder<R> {
    /// Reseed the given RNG.
    fn reseed(&mut self, rng: &mut R);
}

/// Reseed an RNG using a `Default` instance. This reseeds by
/// replacing the RNG with the result of a `Default::default` call.
#[derive(Clone, Copy, Debug)]
pub struct ReseedWithDefault;

impl<R: Rng + Default> Reseeder<R> for ReseedWithDefault {
    fn reseed(&mut self, rng: &mut R) {
        *rng = Default::default();
    }
}
impl Default for ReseedWithDefault {
    fn default() -> ReseedWithDefault { ReseedWithDefault }
}

#[cfg(test)]
mod test {
    use impls;
    use core::default::Default;
    use super::{ReseedingRng, ReseedWithDefault};
    use {SeedableRng, Rng};

    struct Counter {
        i: u32
    }

    impl Rng for Counter {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            // very random
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            impls::next_u64_via_u32(self)
        }

        fn fill_bytes(&mut self, dest: &mut [u8]) {
            impls::fill_bytes_via_u64(self, dest)
        }
    }
    impl Default for Counter {
        fn default() -> Counter {
            Counter { i: 0 }
        }
    }
    impl SeedableRng<u32> for Counter {
        fn reseed(&mut self, seed: u32) {
            self.i = seed;
        }
        fn from_seed(seed: u32) -> Counter {
            Counter { i: seed }
        }
    }
    type MyRng = ReseedingRng<Counter, ReseedWithDefault>;

    #[test]
    fn test_reseeding() {
        let mut rs = ReseedingRng::new(Counter {i:0}, 400, ReseedWithDefault);

        let mut i = 0;
        for _ in 0..1000 {
            assert_eq!(rs.next_u32(), i % 100);
            i += 1;
        }
    }

    #[test]
    fn test_rng_seeded() {
        let mut ra: MyRng = SeedableRng::from_seed((ReseedWithDefault, 2));
        let mut rb: MyRng = SeedableRng::from_seed((ReseedWithDefault, 2));
        assert!(::test::iter_eq(ra.gen_ascii_chars().take(100),
                                rb.gen_ascii_chars().take(100)));
    }

    #[test]
    fn test_rng_reseed() {
        let mut rng: MyRng = SeedableRng::from_seed((ReseedWithDefault, 3));
        let mut results1 = [0u32; 32];
        for i in results1.iter_mut() { *i = rng.next_u32(); }

        rng.reseed((ReseedWithDefault, 3));

        let mut results2 = [0u32; 32];
        for i in results2.iter_mut() { *i = rng.next_u32(); }
        assert_eq!(results1, results2);
    }

    const FILL_BYTES_V_LEN: usize = 13579;
    #[test]
    fn test_rng_fill_bytes() {
        let mut v = [0u8; FILL_BYTES_V_LEN];
        ::test::rng().fill_bytes(&mut v);

        // To test that `fill_bytes` actually did something, check that the
        // average of `v` is not 0.
        let mut sum = 0.0;
        for &x in v.iter() {
            sum += x as f64;
        }
        assert!(sum / v.len() as f64 != 0.0);
    }
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// 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.

//! A wrapper around another PRNG that reseeds it after it
//! generates a certain number of random bytes.

use {Rng, SeedableRng, Error, ErrorKind};

/// A wrapper around any PRNG which reseeds the underlying PRNG after it has
/// generated a certain number of random bytes.
///
/// Reseeding is never strictly *necessary*. Cryptographic PRNGs don't have a
/// limited number of bytes they can output, or at least not a limit reachable
/// in any practical way. There is no such thing as 'running out of entropy'.
///
/// Some small non-cryptographic PRNGs can have very small periods, for
/// example less than 264. Would reseeding help to ensure that you do
/// not wrap around at the end of the period? A period of 264 still
/// takes several centuries of CPU-years on current hardware. Reseeding will
/// actually make things worse, because the reseeded PRNG will just continue
/// somewhere else *in the same period*, with a high chance of overlapping with
/// previously used parts of it.
///
/// # When should you use `ReseedingRng`?
///
/// - Reseeding can be seen as some form of 'security in depth'. Even if in the
///   future a cryptographic weakness is found in the CSPRNG being used,
///   occasionally reseeding should make exploiting it much more difficult or
///   even impossible.
/// - It can be used as a poor man's cryptography (not recommended, just use a
///   good CSPRNG). Previous implementations of `thread_rng` for example used
///   `ReseedingRng` with the ISAAC RNG. That algorithm, although apparently
///   strong and with no known attack, does not come with any proof of security
///   and does not meet the current standards for a cryptographically secure
///   PRNG. By reseeding it frequently (every 32 MiB) it seems safe to assume
///   there is no attack that can operate on the tiny window between reseeds.
///
/// # Error handling
///
/// If reseeding fails, `try_fill_bytes` is the only `Rng` method to report it.
/// For all other `Rng` methods, `ReseedingRng` will not panic but try to
/// handle the error intelligently; if handling the source error fails these
/// methods will continue generating data from the wrapped PRNG without
/// reseeding.
///
/// It is usually best to use the infallible methods `next_u32`, `next_u64` and
/// `fill_bytes` because they can make use of this error handling strategy.
/// Use `try_fill_bytes` and possibly `try_reseed` if you want to handle
/// reseeding errors explicitly.
#[derive(Debug)]
pub struct ReseedingRng<R, Rsdr> {
    rng: R,
    reseeder: Rsdr,
    threshold: i64,
    bytes_until_reseed: i64,
}

impl<R: Rng+SeedableRng, Rsdr: Rng> ReseedingRng<R, Rsdr> {
    /// Create a new `ReseedingRng` with the given parameters.
    ///
    /// # Arguments
    ///
    /// * `rng`: the random number generator to use.
    /// * `threshold`: the number of generated bytes after which to reseed the RNG.
    /// * `reseeder`: the RNG to use for reseeding.
    pub fn new(rng: R, threshold: u64, reseeder: Rsdr) -> ReseedingRng<R,Rsdr> {
        assert!(threshold <= ::core::i64::MAX as u64);
        ReseedingRng {
            rng: rng,
            reseeder: reseeder,
            threshold: threshold as i64,
            bytes_until_reseed: threshold as i64,
        }
    }

    /// Reseed the internal PRNG.
    ///
    /// This will try to work around errors in the RNG used for reseeding
    /// intelligently. If the error kind indicates retrying might help, it will
    /// immediately retry a couple of times. If the error kind indicates the
    /// seeding RNG is not ready, it will retry later, after `threshold / 256`
    /// generated bytes. On other errors in the source RNG, this will skip
    /// reseeding and continue using the internal PRNG, until another
    /// `threshold` bytes have been generated (at which point it will try
    /// reseeding again).
    #[inline(never)]
    pub fn reseed(&mut self) {
        trace!("Reseeding RNG after generating {} bytes",
               self.threshold - self.bytes_until_reseed);
        self.bytes_until_reseed = self.threshold;
        let mut err_count = 0;
        loop {
            if let Err(e) = R::from_rng(&mut self.reseeder)
                            .map(|result| self.rng = result) {
                let kind = e.kind();
                if kind.should_wait() {
                    self.bytes_until_reseed = self.threshold >> 8;
                    warn!("Reseeding RNG delayed for {} bytes",
                           self.bytes_until_reseed);
                } else if kind.should_retry() {
                    err_count += 1;
                    // Retry immediately for 5 times (arbitrary limit)
                    if err_count <= 5 { continue; }
                }
                warn!("Reseeding RNG failed; continuing without reseeding. Error: {:?}", e);
            }
            break; // Successfully reseeded, delayed, or given up.
        }
    }

    /// Reseed the internal RNG if the number of bytes that have been
    /// generated exceed the threshold.
    ///
    /// If reseeding fails, return an error with the original cause. Note that
    /// if the cause has a permanent failure, we report a transient error and
    /// skip reseeding; this means that only two error kinds can be reported
    /// from this method: `ErrorKind::Transient` and `ErrorKind::NotReady`.
    #[inline(never)]
    pub fn try_reseed(&mut self) -> Result<(), Error> {
        trace!("Reseeding RNG after {} generated bytes",
               self.threshold - self.bytes_until_reseed);
        if let Err(err) = R::from_rng(&mut self.reseeder)
                          .map(|result| self.rng = result) {
            let newkind = match err.kind() {
                a @ ErrorKind::NotReady => a,
                b @ ErrorKind::Transient => b,
                _ => {
                    self.bytes_until_reseed = self.threshold; // skip reseeding
                    ErrorKind::Transient
                }
            };
            return Err(Error::with_cause(newkind, "reseeding failed", err));
        }
        self.bytes_until_reseed = self.threshold;
        Ok(())
    }
}

impl<R: Rng+SeedableRng, Rsdr: Rng> Rng for ReseedingRng<R, Rsdr> {
    fn next_u32(&mut self) -> u32 {
        let value = self.rng.next_u32();
        self.bytes_until_reseed -= 4;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
        value
    }

    fn next_u64(&mut self) -> u64 {
        let value = self.rng.next_u64();
        self.bytes_until_reseed -= 8;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
        value
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.rng.fill_bytes(dest);
        self.bytes_until_reseed -= dest.len() as i64;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.rng.try_fill_bytes(dest)?;
        self.bytes_until_reseed -= dest.len() as i64;
        if self.bytes_until_reseed <= 0 {
            self.try_reseed()?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use {impls, le};
    use super::{ReseedingRng};
    use {SeedableRng, Rng, Error};

    struct Counter {
        i: u32
    }

    impl Rng for Counter {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            // very random
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            impls::next_u64_via_u32(self)
        }

        fn fill_bytes(&mut self, dest: &mut [u8]) {
            impls::fill_bytes_via_u64(self, dest)
        }
    }
    impl SeedableRng for Counter {
        type Seed = [u8; 4];
        fn from_seed(seed: Self::Seed) -> Self {
            let mut seed_u32 = [0u32; 1];
            le::read_u32_into(&seed, &mut seed_u32);
            Counter { i: seed_u32[0] }
        }
    }

    #[derive(Debug, Clone)]
    struct ResetCounter;
    impl Rng for ResetCounter {
        fn next_u32(&mut self) -> u32 { unimplemented!() }
        fn next_u64(&mut self) -> u64 { unimplemented!() }
        fn fill_bytes(&mut self, _dest: &mut [u8]) { unimplemented!() }
        fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
            for i in dest.iter_mut() { *i = 0; }
            Ok(())
        }
    }

    #[test]
    fn test_reseeding() {
        let mut rs = ReseedingRng::new(Counter {i:0}, 400, ResetCounter);

        let mut i = 0;
        for _ in 0..1000 {
            assert_eq!(rs.next_u32(), i % 100);
            i += 1;
        }
    }
}
<commit_msg>Use Display instead of Debug for error logging<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// https://rust-lang.org/COPYRIGHT.
//
// 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.

//! A wrapper around another PRNG that reseeds it after it
//! generates a certain number of random bytes.

use {Rng, SeedableRng, Error, ErrorKind};

/// A wrapper around any PRNG which reseeds the underlying PRNG after it has
/// generated a certain number of random bytes.
///
/// Reseeding is never strictly *necessary*. Cryptographic PRNGs don't have a
/// limited number of bytes they can output, or at least not a limit reachable
/// in any practical way. There is no such thing as 'running out of entropy'.
///
/// Some small non-cryptographic PRNGs can have very small periods, for
/// example less than 264. Would reseeding help to ensure that you do
/// not wrap around at the end of the period? A period of 264 still
/// takes several centuries of CPU-years on current hardware. Reseeding will
/// actually make things worse, because the reseeded PRNG will just continue
/// somewhere else *in the same period*, with a high chance of overlapping with
/// previously used parts of it.
///
/// # When should you use `ReseedingRng`?
///
/// - Reseeding can be seen as some form of 'security in depth'. Even if in the
///   future a cryptographic weakness is found in the CSPRNG being used,
///   occasionally reseeding should make exploiting it much more difficult or
///   even impossible.
/// - It can be used as a poor man's cryptography (not recommended, just use a
///   good CSPRNG). Previous implementations of `thread_rng` for example used
///   `ReseedingRng` with the ISAAC RNG. That algorithm, although apparently
///   strong and with no known attack, does not come with any proof of security
///   and does not meet the current standards for a cryptographically secure
///   PRNG. By reseeding it frequently (every 32 MiB) it seems safe to assume
///   there is no attack that can operate on the tiny window between reseeds.
///
/// # Error handling
///
/// If reseeding fails, `try_fill_bytes` is the only `Rng` method to report it.
/// For all other `Rng` methods, `ReseedingRng` will not panic but try to
/// handle the error intelligently; if handling the source error fails these
/// methods will continue generating data from the wrapped PRNG without
/// reseeding.
///
/// It is usually best to use the infallible methods `next_u32`, `next_u64` and
/// `fill_bytes` because they can make use of this error handling strategy.
/// Use `try_fill_bytes` and possibly `try_reseed` if you want to handle
/// reseeding errors explicitly.
#[derive(Debug)]
pub struct ReseedingRng<R, Rsdr> {
    rng: R,
    reseeder: Rsdr,
    threshold: i64,
    bytes_until_reseed: i64,
}

impl<R: Rng+SeedableRng, Rsdr: Rng> ReseedingRng<R, Rsdr> {
    /// Create a new `ReseedingRng` with the given parameters.
    ///
    /// # Arguments
    ///
    /// * `rng`: the random number generator to use.
    /// * `threshold`: the number of generated bytes after which to reseed the RNG.
    /// * `reseeder`: the RNG to use for reseeding.
    pub fn new(rng: R, threshold: u64, reseeder: Rsdr) -> ReseedingRng<R,Rsdr> {
        assert!(threshold <= ::core::i64::MAX as u64);
        ReseedingRng {
            rng: rng,
            reseeder: reseeder,
            threshold: threshold as i64,
            bytes_until_reseed: threshold as i64,
        }
    }

    /// Reseed the internal PRNG.
    ///
    /// This will try to work around errors in the RNG used for reseeding
    /// intelligently. If the error kind indicates retrying might help, it will
    /// immediately retry a couple of times. If the error kind indicates the
    /// seeding RNG is not ready, it will retry later, after `threshold / 256`
    /// generated bytes. On other errors in the source RNG, this will skip
    /// reseeding and continue using the internal PRNG, until another
    /// `threshold` bytes have been generated (at which point it will try
    /// reseeding again).
    #[inline(never)]
    pub fn reseed(&mut self) {
        trace!("Reseeding RNG after generating {} bytes",
               self.threshold - self.bytes_until_reseed);
        self.bytes_until_reseed = self.threshold;
        let mut err_count = 0;
        loop {
            if let Err(e) = R::from_rng(&mut self.reseeder)
                            .map(|result| self.rng = result) {
                let kind = e.kind();
                if kind.should_wait() {
                    self.bytes_until_reseed = self.threshold >> 8;
                    warn!("Reseeding RNG delayed for {} bytes",
                           self.bytes_until_reseed);
                } else if kind.should_retry() {
                    err_count += 1;
                    // Retry immediately for 5 times (arbitrary limit)
                    if err_count <= 5 { continue; }
                }
                warn!("Reseeding RNG failed; continuing without reseeding. Error: {}", e);
            }
            break; // Successfully reseeded, delayed, or given up.
        }
    }

    /// Reseed the internal RNG if the number of bytes that have been
    /// generated exceed the threshold.
    ///
    /// If reseeding fails, return an error with the original cause. Note that
    /// if the cause has a permanent failure, we report a transient error and
    /// skip reseeding; this means that only two error kinds can be reported
    /// from this method: `ErrorKind::Transient` and `ErrorKind::NotReady`.
    #[inline(never)]
    pub fn try_reseed(&mut self) -> Result<(), Error> {
        trace!("Reseeding RNG after {} generated bytes",
               self.threshold - self.bytes_until_reseed);
        if let Err(err) = R::from_rng(&mut self.reseeder)
                          .map(|result| self.rng = result) {
            let newkind = match err.kind() {
                a @ ErrorKind::NotReady => a,
                b @ ErrorKind::Transient => b,
                _ => {
                    self.bytes_until_reseed = self.threshold; // skip reseeding
                    ErrorKind::Transient
                }
            };
            return Err(Error::with_cause(newkind, "reseeding failed", err));
        }
        self.bytes_until_reseed = self.threshold;
        Ok(())
    }
}

impl<R: Rng+SeedableRng, Rsdr: Rng> Rng for ReseedingRng<R, Rsdr> {
    fn next_u32(&mut self) -> u32 {
        let value = self.rng.next_u32();
        self.bytes_until_reseed -= 4;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
        value
    }

    fn next_u64(&mut self) -> u64 {
        let value = self.rng.next_u64();
        self.bytes_until_reseed -= 8;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
        value
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.rng.fill_bytes(dest);
        self.bytes_until_reseed -= dest.len() as i64;
        if self.bytes_until_reseed <= 0 {
            self.reseed();
        }
    }

    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
        self.rng.try_fill_bytes(dest)?;
        self.bytes_until_reseed -= dest.len() as i64;
        if self.bytes_until_reseed <= 0 {
            self.try_reseed()?;
        }
        Ok(())
    }
}

#[cfg(test)]
mod test {
    use {impls, le};
    use super::{ReseedingRng};
    use {SeedableRng, Rng, Error};

    struct Counter {
        i: u32
    }

    impl Rng for Counter {
        fn next_u32(&mut self) -> u32 {
            self.i += 1;
            // very random
            self.i - 1
        }
        fn next_u64(&mut self) -> u64 {
            impls::next_u64_via_u32(self)
        }

        fn fill_bytes(&mut self, dest: &mut [u8]) {
            impls::fill_bytes_via_u64(self, dest)
        }
    }
    impl SeedableRng for Counter {
        type Seed = [u8; 4];
        fn from_seed(seed: Self::Seed) -> Self {
            let mut seed_u32 = [0u32; 1];
            le::read_u32_into(&seed, &mut seed_u32);
            Counter { i: seed_u32[0] }
        }
    }

    #[derive(Debug, Clone)]
    struct ResetCounter;
    impl Rng for ResetCounter {
        fn next_u32(&mut self) -> u32 { unimplemented!() }
        fn next_u64(&mut self) -> u64 { unimplemented!() }
        fn fill_bytes(&mut self, _dest: &mut [u8]) { unimplemented!() }
        fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), Error> {
            for i in dest.iter_mut() { *i = 0; }
            Ok(())
        }
    }

    #[test]
    fn test_reseeding() {
        let mut rs = ReseedingRng::new(Counter {i:0}, 400, ResetCounter);

        let mut i = 0;
        for _ in 0..1000 {
            assert_eq!(rs.next_u32(), i % 100);
            i += 1;
        }
    }
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{
    fs::File,
    io::Read,
    path::{Path, PathBuf},
    thread::sleep,
    time::Duration,
};

use nix::mount::{umount2, MntFlags};

use devicemapper::{DevId, DmFlags, DmName, DmNameBuf, DmOptions, DM};

use crate::{
    engine::strat_engine::{
        cmd::udev_settle,
        device::blkdev_size,
        dm::{get_dm, get_dm_init},
    },
    stratis::StratisResult,
};

mod cleanup_errors {
    use std::fmt;

    use devicemapper::DmError;

    #[derive(Debug)]
    pub enum Error {
        Ioe(std::io::Error),
        Nix(nix::Error),
        Msg(String),
        Chained(String, Box<Error>),
        Dm(DmError),
    }

    pub type Result<T> = std::result::Result<T, Error>;

    impl From<nix::Error> for Error {
        fn from(err: nix::Error) -> Error {
            Error::Nix(err)
        }
    }

    impl From<std::io::Error> for Error {
        fn from(err: std::io::Error) -> Error {
            Error::Ioe(err)
        }
    }

    impl From<String> for Error {
        fn from(err: String) -> Error {
            Error::Msg(err)
        }
    }

    impl From<DmError> for Error {
        fn from(err: DmError) -> Error {
            Error::Dm(err)
        }
    }

    impl fmt::Display for Error {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            match self {
                Error::Ioe(err) => write!(f, "IO error: {}", err),
                Error::Nix(err) => write!(f, "Nix error: {}", err),
                Error::Msg(err) => write!(f, "{}", err),
                Error::Chained(msg, err) => write!(f, "{}: {}", msg, err),
                Error::Dm(err) => write!(f, "DM error: {}", err),
            }
        }
    }

    impl std::error::Error for Error {}
}

use self::cleanup_errors::{Error, Result};

/// Attempt to remove all device mapper devices which match the stratis naming convention.
/// FIXME: Current implementation complicated by https://bugzilla.redhat.com/show_bug.cgi?id=1506287
pub fn dm_stratis_devices_remove() -> Result<()> {
    /// One iteration of removing devicemapper devices
    fn one_iteration() -> Result<(bool, Vec<DmNameBuf>)> {
        let mut progress_made = false;
        let mut remain = get_dm()
            .list_devices()
            .map_err(|e| {
                Error::Chained(
                    "failed while listing DM devices, giving up".into(),
                    Box::new(e.into()),
                )
            })?
            .iter()
            .map(|d| &d.0)
            .filter_map(|n| {
                if !n.to_string().starts_with("stratis-1") {
                    None
                } else {
                    match get_dm().device_remove(&DevId::Name(n), DmOptions::default()) {
                        Ok(_) => {
                            progress_made = true;
                            None
                        }
                        Err(_) => Some(n.to_owned()),
                    }
                }
            })
            .collect::<Vec<_>>();

        // Retries if no progress has been made.
        if !remain.is_empty() && !progress_made {
            remain = remain
                .into_iter()
                .filter(|name| {
                    for _ in 0..3 {
                        match get_dm().device_remove(&DevId::Name(name), DmOptions::default()) {
                            Ok(_) => {
                                progress_made = true;
                                return false;
                            }
                            Err(e) => {
                                debug!(
                                    "Failed to remove device {} on retry: {}",
                                    name.to_string(),
                                    e
                                );
                                sleep(Duration::from_secs(1));
                            }
                        }
                    }
                    true
                })
                .collect();
        }

        Ok((progress_made, remain))
    }

    /// Do one iteration of removals until progress stops. Return remaining
    /// dm devices.
    fn do_while_progress() -> Result<Vec<DmNameBuf>> {
        let mut result = one_iteration()?;
        while result.0 {
            result = one_iteration()?;
        }
        Ok(result.1)
    }

    || -> Result<()> {
        udev_settle().unwrap();
        get_dm_init().map_err(|_| Error::Msg("Unable to initialize DM".into()))?;
        do_while_progress().and_then(|remain| {
            if !remain.is_empty() {
                Err(format!("Some Stratis DM devices remaining: {:?}", remain).into())
            } else {
                Ok(())
            }
        })
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure removal of all test-generated DM devices".into(),
            Box::new(e),
        )
    })
}

/// Try and un-mount any filesystems that have the name stratis in the mount point, returning
/// immediately on the first one we are unable to unmount.
fn stratis_filesystems_unmount() -> Result<()> {
    || -> Result<()> {
        let mut mount_data = String::new();
        File::open("/proc/self/mountinfo")?.read_to_string(&mut mount_data)?;
        let parser = libmount::mountinfo::Parser::new(mount_data.as_bytes());

        for mount_point in parser
            .filter_map(|x| x.ok())
            .filter_map(|m| m.mount_point.into_owned().into_string().ok())
            .filter(|mp| mp.contains("stratis"))
        {
            umount2(&PathBuf::from(mount_point), MntFlags::MNT_DETACH)?;
        }

        Ok(())
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure all Stratis filesystems were unmounted".into(),
            Box::new(e),
        )
    })
}

/// When a unit test panics we can leave the system in an inconsistent state.  This function
/// tries to clean up by un-mounting any mounted file systems which contain the string
/// "stratis_testing" and then it tries to remove any device mapper tables which are also stratis
/// created.
pub fn clean_up() -> Result<()> {
    stratis_filesystems_unmount().and_then(|_| dm_stratis_devices_remove())
}

pub struct FailDevice {
    backing_device: PathBuf,
    test_device_name: String,
    dm_context: DM,
    size: u64,
}

impl FailDevice {
    pub fn new(backing_device: &Path, test_device_name: &str) -> StratisResult<Self> {
        let dm = DM::new()?;
        let dm_name = DmName::new(test_device_name)?;
        let dev_id = DevId::Name(dm_name);

        let size = {
            let file = File::open(backing_device)?;
            blkdev_size(&file)?
        };

        dm.device_create(dm_name, None, DmOptions::default())?;
        dm.table_load(
            &dev_id,
            &[(
                0,
                *size.sectors(),
                "linear".to_string(),
                format!("{} 0", backing_device.display()),
            )],
            DmOptions::default(),
        )?;
        dm.device_suspend(&dev_id, DmOptions::default())?;

        Ok(FailDevice {
            backing_device: backing_device.to_owned(),
            test_device_name: test_device_name.to_owned(),
            dm_context: dm,
            size: *size.sectors(),
        })
    }

    pub fn as_path(&self) -> PathBuf {
        vec!["/dev/mapper", self.test_device_name.as_str()]
            .into_iter()
            .collect::<PathBuf>()
    }

    pub fn start_failing(&self, num_sectors_after_start: u64) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::default().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[
                (
                    0,
                    num_sectors_after_start,
                    "error".to_string(),
                    String::new(),
                ),
                (
                    num_sectors_after_start,
                    self.size - num_sectors_after_start,
                    "linear".to_string(),
                    format!(
                        "{} {}",
                        self.backing_device.display(),
                        num_sectors_after_start
                    ),
                ),
            ],
            DmOptions::default(),
        )?;
        self.dm_context
            .device_suspend(&dev_id, DmOptions::default())?;

        Ok(())
    }

    pub fn stop_failing(&self) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::default().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[(
                0,
                self.size,
                "linear".to_string(),
                format!("{} 0", self.backing_device.display()),
            )],
            DmOptions::default(),
        )?;
        self.dm_context
            .device_suspend(&dev_id, DmOptions::default())?;

        Ok(())
    }
}

impl Drop for FailDevice {
    fn drop(&mut self) {
        fn drop_fail(dev: &mut FailDevice) -> StratisResult<()> {
            let dev_id = DevId::Name(DmName::new(dev.test_device_name.as_str())?);

            let (dev_info, _) = dev.dm_context.table_status(&dev_id, DmOptions::default())?;
            if dev_info.flags() & DmFlags::DM_SUSPEND == DmFlags::DM_SUSPEND {
                dev.dm_context
                    .device_suspend(&dev_id, DmOptions::default())?;
            }
            dev.dm_context
                .device_remove(&dev_id, DmOptions::default())?;

            Ok(())
        }

        if let Err(e) = drop_fail(self) {
            warn!(
                "Teardown of test device /dev/mapper/{} failed: {}",
                self.test_device_name, e
            );
        }
    }
}
<commit_msg>Fix manual-retain clippy error<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::{
    fs::File,
    io::Read,
    path::{Path, PathBuf},
    thread::sleep,
    time::Duration,
};

use nix::mount::{umount2, MntFlags};

use devicemapper::{DevId, DmFlags, DmName, DmNameBuf, DmOptions, DM};

use crate::{
    engine::strat_engine::{
        cmd::udev_settle,
        device::blkdev_size,
        dm::{get_dm, get_dm_init},
    },
    stratis::StratisResult,
};

mod cleanup_errors {
    use std::fmt;

    use devicemapper::DmError;

    #[derive(Debug)]
    pub enum Error {
        Ioe(std::io::Error),
        Nix(nix::Error),
        Msg(String),
        Chained(String, Box<Error>),
        Dm(DmError),
    }

    pub type Result<T> = std::result::Result<T, Error>;

    impl From<nix::Error> for Error {
        fn from(err: nix::Error) -> Error {
            Error::Nix(err)
        }
    }

    impl From<std::io::Error> for Error {
        fn from(err: std::io::Error) -> Error {
            Error::Ioe(err)
        }
    }

    impl From<String> for Error {
        fn from(err: String) -> Error {
            Error::Msg(err)
        }
    }

    impl From<DmError> for Error {
        fn from(err: DmError) -> Error {
            Error::Dm(err)
        }
    }

    impl fmt::Display for Error {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            match self {
                Error::Ioe(err) => write!(f, "IO error: {}", err),
                Error::Nix(err) => write!(f, "Nix error: {}", err),
                Error::Msg(err) => write!(f, "{}", err),
                Error::Chained(msg, err) => write!(f, "{}: {}", msg, err),
                Error::Dm(err) => write!(f, "DM error: {}", err),
            }
        }
    }

    impl std::error::Error for Error {}
}

use self::cleanup_errors::{Error, Result};

/// Attempt to remove all device mapper devices which match the stratis naming convention.
/// FIXME: Current implementation complicated by https://bugzilla.redhat.com/show_bug.cgi?id=1506287
pub fn dm_stratis_devices_remove() -> Result<()> {
    /// One iteration of removing devicemapper devices
    fn one_iteration() -> Result<(bool, Vec<DmNameBuf>)> {
        let mut progress_made = false;
        let mut remain = get_dm()
            .list_devices()
            .map_err(|e| {
                Error::Chained(
                    "failed while listing DM devices, giving up".into(),
                    Box::new(e.into()),
                )
            })?
            .iter()
            .map(|d| &d.0)
            .filter_map(|n| {
                if !n.to_string().starts_with("stratis-1") {
                    None
                } else {
                    match get_dm().device_remove(&DevId::Name(n), DmOptions::default()) {
                        Ok(_) => {
                            progress_made = true;
                            None
                        }
                        Err(_) => Some(n.to_owned()),
                    }
                }
            })
            .collect::<Vec<_>>();

        // Retries if no progress has been made.
        if !remain.is_empty() && !progress_made {
            remain.retain(|name| {
                for _ in 0..3 {
                    match get_dm().device_remove(&DevId::Name(name), DmOptions::default()) {
                        Ok(_) => {
                            progress_made = true;
                            return false;
                        }
                        Err(e) => {
                            debug!(
                                "Failed to remove device {} on retry: {}",
                                name.to_string(),
                                e
                            );
                            sleep(Duration::from_secs(1));
                        }
                    }
                }
                true
            });
        }

        Ok((progress_made, remain))
    }

    /// Do one iteration of removals until progress stops. Return remaining
    /// dm devices.
    fn do_while_progress() -> Result<Vec<DmNameBuf>> {
        let mut result = one_iteration()?;
        while result.0 {
            result = one_iteration()?;
        }
        Ok(result.1)
    }

    || -> Result<()> {
        udev_settle().unwrap();
        get_dm_init().map_err(|_| Error::Msg("Unable to initialize DM".into()))?;
        do_while_progress().and_then(|remain| {
            if !remain.is_empty() {
                Err(format!("Some Stratis DM devices remaining: {:?}", remain).into())
            } else {
                Ok(())
            }
        })
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure removal of all test-generated DM devices".into(),
            Box::new(e),
        )
    })
}

/// Try and un-mount any filesystems that have the name stratis in the mount point, returning
/// immediately on the first one we are unable to unmount.
fn stratis_filesystems_unmount() -> Result<()> {
    || -> Result<()> {
        let mut mount_data = String::new();
        File::open("/proc/self/mountinfo")?.read_to_string(&mut mount_data)?;
        let parser = libmount::mountinfo::Parser::new(mount_data.as_bytes());

        for mount_point in parser
            .filter_map(|x| x.ok())
            .filter_map(|m| m.mount_point.into_owned().into_string().ok())
            .filter(|mp| mp.contains("stratis"))
        {
            umount2(&PathBuf::from(mount_point), MntFlags::MNT_DETACH)?;
        }

        Ok(())
    }()
    .map_err(|e| {
        Error::Chained(
            "Failed to ensure all Stratis filesystems were unmounted".into(),
            Box::new(e),
        )
    })
}

/// When a unit test panics we can leave the system in an inconsistent state.  This function
/// tries to clean up by un-mounting any mounted file systems which contain the string
/// "stratis_testing" and then it tries to remove any device mapper tables which are also stratis
/// created.
pub fn clean_up() -> Result<()> {
    stratis_filesystems_unmount().and_then(|_| dm_stratis_devices_remove())
}

pub struct FailDevice {
    backing_device: PathBuf,
    test_device_name: String,
    dm_context: DM,
    size: u64,
}

impl FailDevice {
    pub fn new(backing_device: &Path, test_device_name: &str) -> StratisResult<Self> {
        let dm = DM::new()?;
        let dm_name = DmName::new(test_device_name)?;
        let dev_id = DevId::Name(dm_name);

        let size = {
            let file = File::open(backing_device)?;
            blkdev_size(&file)?
        };

        dm.device_create(dm_name, None, DmOptions::default())?;
        dm.table_load(
            &dev_id,
            &[(
                0,
                *size.sectors(),
                "linear".to_string(),
                format!("{} 0", backing_device.display()),
            )],
            DmOptions::default(),
        )?;
        dm.device_suspend(&dev_id, DmOptions::default())?;

        Ok(FailDevice {
            backing_device: backing_device.to_owned(),
            test_device_name: test_device_name.to_owned(),
            dm_context: dm,
            size: *size.sectors(),
        })
    }

    pub fn as_path(&self) -> PathBuf {
        vec!["/dev/mapper", self.test_device_name.as_str()]
            .into_iter()
            .collect::<PathBuf>()
    }

    pub fn start_failing(&self, num_sectors_after_start: u64) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::default().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[
                (
                    0,
                    num_sectors_after_start,
                    "error".to_string(),
                    String::new(),
                ),
                (
                    num_sectors_after_start,
                    self.size - num_sectors_after_start,
                    "linear".to_string(),
                    format!(
                        "{} {}",
                        self.backing_device.display(),
                        num_sectors_after_start
                    ),
                ),
            ],
            DmOptions::default(),
        )?;
        self.dm_context
            .device_suspend(&dev_id, DmOptions::default())?;

        Ok(())
    }

    pub fn stop_failing(&self) -> StratisResult<()> {
        let dm_name = DmName::new(self.test_device_name.as_str())?;
        let dev_id = DevId::Name(dm_name);

        self.dm_context
            .device_suspend(&dev_id, DmOptions::default().set_flags(DmFlags::DM_SUSPEND))?;
        self.dm_context.table_load(
            &dev_id,
            &[(
                0,
                self.size,
                "linear".to_string(),
                format!("{} 0", self.backing_device.display()),
            )],
            DmOptions::default(),
        )?;
        self.dm_context
            .device_suspend(&dev_id, DmOptions::default())?;

        Ok(())
    }
}

impl Drop for FailDevice {
    fn drop(&mut self) {
        fn drop_fail(dev: &mut FailDevice) -> StratisResult<()> {
            let dev_id = DevId::Name(DmName::new(dev.test_device_name.as_str())?);

            let (dev_info, _) = dev.dm_context.table_status(&dev_id, DmOptions::default())?;
            if dev_info.flags() & DmFlags::DM_SUSPEND == DmFlags::DM_SUSPEND {
                dev.dm_context
                    .device_suspend(&dev_id, DmOptions::default())?;
            }
            dev.dm_context
                .device_remove(&dev_id, DmOptions::default())?;

            Ok(())
        }

        if let Err(e) = drop_fail(self) {
            warn!(
                "Teardown of test device /dev/mapper/{} failed: {}",
                self.test_device_name, e
            );
        }
    }
}
<|endoftext|>


<commit_before><commit_msg>Sketch out client API.<commit_after>use std::net::{UdpSocket, SocketAddr};
use std::io::{Error};

/// Client socket for statsd servers.
///
/// After creating a metric you can use `Client`
/// to send metrics to the configured statsd server
///
/// # Example
///
/// Creating a client and sending metrics is easy.
///
/// ```ignore
/// use statsd::client::Client;
///
/// let client = Client::new("127.0.0.1:8125", "myapp.");
/// client.incr("some.metric.completed");
/// ```
pub struct Client {
    socket: UdpSocket,
    server_addr: SocketAddr,
    prefix: String
}

impl Client {
    /// Construct a new statsd client given an host/port & prefix
    pub fn new(host: &str, prefix: &str) -> Result<Client, Error> {
        let addr = SocketAddr.from_str(host);
        let mut socket = try!(UdpSocket::bind(addr));
        Ok(Client {socket: socket, prefix: prefix.to_string(), server_addr: addr})
    }

    /// Increment a metric by 1
    ///
    /// ```ignore
    /// # Increment a given metric by 1.
    /// client.incr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn incr(&mut self, metric: &str) {
        self.count(metric, 1.0);
    }

    /// Decrement a metric by -1
    ///
    /// ```ignore
    /// # Decrement a given metric by 1
    /// client.decr("metric.completed");
    /// ```
    ///
    /// This modifies a counter with an effective sampling
    /// rate of 1.0.
    pub fn decr(&mut self, metric: &str) {
        self.count(metric, -1.0);
    }

    /// Modify a counter by `value`.
    ///
    /// Will increment or decrement a counter by `value` with
    /// a sampling rate of 1.0.
    ///
    /// ```ignore
    /// # Increment by 12
    /// client.count("metric.completed", 12);
    /// ```
    pub fn count(&mut self, metric: &str, value: usize) {
        let data = format!("{}.{}:{}|c", self.prefix, metric, value);
        self.send(data);
    }

    /// Data goes in, data comes out.
    fn send(&mut self, data: &str) {
        let _ = self.socket.send_to(data.as_bytes(), self.dest);
    }
}
<|endoftext|>


<commit_before>use std::thread;
use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::sync::mpsc;
use std::sync::Arc;
use std::str::from_utf8;
use std::mem;

use futures;
use futures::Future;
use futures::stream::Stream;

use tokio_core::io as tokio_io;
use tokio_core::io::TaskIo;
use tokio_core::io::TaskIoRead;
use tokio_core::io::TaskIoWrite;

use tokio_core;
use tokio_core::Loop;
use tokio_core::LoopHandle;
use tokio_core::TcpStream;

use solicit::http::client::RequestStream;
use solicit::http::client::ClientSession;
use solicit::http::session::Client;
use solicit::http::session::SessionState;
use solicit::http::session::DefaultSessionState;
use solicit::http::session::DefaultStream;
use solicit::http::session::StreamState;
use solicit::http::session::StreamDataError;
use solicit::http::session::StreamDataChunk;
use solicit::http::session::Stream as solicit_Stream;
use solicit::http::connection::HttpConnection;
use solicit::http::connection::SendStatus;
use solicit::http::connection::EndStream;
use solicit::http::connection::DataChunk;
use solicit::http::priority::SimplePrioritizer;
use solicit::http::StreamId;
use solicit::http::HttpScheme;
use solicit::http::HttpError;
use solicit::http::HttpResult;
use solicit::http::Header;
use solicit::http::StaticHeader;

use method::MethodDescriptor;

use error::*;
use result::*;

use futures_misc::*;
use futures_grpc::*;

use grpc::*;
use solicit_async::*;
use solicit_misc::*;

use http_client::*;

trait GrpcResponseHandlerTrait : Send + 'static + HttpResponseHandler {
}

struct GrpcResponseHandlerTyped<Req : Send + 'static, Resp : Send + 'static> {
    method: Arc<MethodDescriptor<Req, Resp>>,
    complete: tokio_core::Sender<ResultOrEof<Resp, GrpcError>>,
    remaining_response: Vec<u8>,
}

impl<Req : Send + 'static, Resp : Send + 'static> GrpcResponseHandlerTrait for GrpcResponseHandlerTyped<Req, Resp> {
}

impl<Req : Send + 'static, Resp : Send + 'static> HttpResponseHandler for GrpcResponseHandlerTyped<Req, Resp> {
    fn headers(&mut self, headers: Vec<StaticHeader>) -> HttpResult<()> {
        println!("client: received headers");
        if slice_get_header(&headers, ":status") != Some("200") {
            if let Some(message) = slice_get_header(&headers, HEADER_GRPC_MESSAGE) {
                self.complete.send(ResultOrEof::Error(GrpcError::GrpcMessage(GrpcMessageError { grpc_message: message.to_owned() })));
                Err(HttpError::Other(Box::new(GrpcError::GrpcMessage(GrpcMessageError { grpc_message: message.to_owned() }))))
            } else {
                self.complete.send(ResultOrEof::Error(GrpcError::Other("not 200")));
                Err(HttpError::Other(Box::new(GrpcError::Other("not 200"))))
            }
        } else {
            Ok(())
        }
    }

    fn data_frame(&mut self, chunk: Vec<u8>) -> HttpResult<()> {
        self.remaining_response.extend(&chunk);
        loop {
            let len = match parse_grpc_frame(&self.remaining_response) {
                Err(e) => {
                    self.complete.send(ResultOrEof::Error(e));
                    break;
                }
                Ok(None) => break,
                Ok(Some((message, len))) => {
                    let resp = self.method.resp_marshaller.read(&message);
                    self.complete.send(From::from(resp)).ok();
                    len
                }
            };
            self.remaining_response.drain(..len);
        }
        Ok(())
    }

    fn trailers(&mut self, headers: Vec<StaticHeader>) -> HttpResult<()> {
        let status_200 = slice_get_header(&headers, ":status") == Some("200");
        let grpc_status_0 = slice_get_header(&headers, HEADER_GRPC_STATUS) == Some("0");
        if /* status_200 && */ grpc_status_0 {
            Ok(())
        } else {
            if let Some(message) = slice_get_header(&headers, HEADER_GRPC_MESSAGE) {
                self.complete.send(ResultOrEof::Error(GrpcError::GrpcMessage(GrpcMessageError { grpc_message: message.to_owned() })));
                Err(HttpError::Other(Box::new(GrpcError::GrpcMessage(GrpcMessageError { grpc_message: message.to_owned() }))))
            } else {
                self.complete.send(ResultOrEof::Error(GrpcError::Other("not xxx")));
                Err(HttpError::Other(Box::new(GrpcError::Other("not xxx"))))
            }
        }
    }

    fn end(&mut self) -> HttpResult<()> {
        self.complete.send(ResultOrEof::Eof);
        Ok(())
    }
}

struct GrpcResponseHandler {
    tr: Box<GrpcResponseHandlerTrait>,
}

impl HttpResponseHandler for GrpcResponseHandler {
    fn headers(&mut self, headers: Vec<StaticHeader>) -> HttpResult<()> {
        self.tr.headers(headers)
    }

    fn data_frame(&mut self, chunk: Vec<u8>) -> HttpResult<()> {
        self.tr.data_frame(chunk)
    }

    fn trailers(&mut self, headers: Vec<StaticHeader>) -> HttpResult<()> {
        self.tr.trailers(headers)
    }

    fn end(&mut self) -> HttpResult<()> {
        self.tr.end()
    }
}

// Data sent from event loop to GrpcClient
struct LoopToClient {
    // used only once to send shutdown signal
    shutdown_tx: tokio_core::Sender<()>,
    loop_handle: LoopHandle,
    http_conn: HttpConnectionAsync<GrpcResponseHandler>,
}

/// gRPC client implementation.
/// Used by generated code.
pub struct GrpcClient {
    loop_to_client: LoopToClient,
    thread_join_handle: Option<thread::JoinHandle<()>>,
    host: String,
    http_scheme: HttpScheme,
}

trait CallRequest : Send {
    // serialize message contained in this request
    fn write_req(&self) -> GrpcResult<Vec<u8>>;
    // gRPC method name (path)
    fn method_name(&self) -> &str;
    // called when server receives response from server
    // to send it back method called
    fn process_response(&mut self, message: ResultOrEof<&[u8], GrpcError>);
}

struct CallRequestHolder {
    request: Box<CallRequest>,
}

struct CallRequestTyped<Req, Resp> {
    // this call method descriptor
    method: Arc<MethodDescriptor<Req, Resp>>,
    // the request
    req: Req,
    // channel to send response back to called
    complete: tokio_core::Sender<ResultOrEof<Resp, GrpcError>>,
}

impl<Req : Send + 'static, Resp : Send + 'static> CallRequest for CallRequestTyped<Req, Resp> {
    fn write_req(&self) -> GrpcResult<Vec<u8>> {
        self.method.req_marshaller.write(&self.req)
    }

    fn method_name(&self) -> &str {
        &self.method.name
    }

    fn process_response(&mut self, message: ResultOrEof<&[u8], GrpcError>) {
        let resp = message.and_then(|message| self.method.resp_marshaller.read(&message));
        self.complete.send(resp).ok();
    }
}

impl GrpcClient {
    /// Create a client connected to specified host and port.
    pub fn new(host: &str, port: u16) -> GrpcResult<GrpcClient> {

        // TODO: sync
        // TODO: try connect to all addrs
        let socket_addr = try!((host, port).to_socket_addrs()).next().unwrap();

        // We need some data back from event loop.
        // This channel is used to exchange that data
        let (get_from_loop_tx, get_from_loop_rx) = mpsc::channel();

        // Start event loop.
        let join_handle = thread::spawn(move || {
            run_client_event_loop(socket_addr, get_from_loop_tx);
        });

        // Get back call channel and shutdown channel.
        let loop_to_client = try!(get_from_loop_rx.recv()
            .map_err(|_| GrpcError::Other("get response from loop")));

        Ok(GrpcClient {
            loop_to_client: loop_to_client,
            thread_join_handle: Some(join_handle),
            host: host.to_owned(),
            http_scheme: HttpScheme::Http,
        })
    }

    pub fn new_resp_channel<Resp : Send + 'static>(&self)
        -> (tokio_core::Sender<ResultOrEof<Resp, GrpcError>>, GrpcStream<Resp>)
    {
        let (sender, receiver) = self.loop_to_client.loop_handle.clone().channel::<ResultOrEof<Resp, GrpcError>>();
        let receiver: GrpcStream<ResultOrEof<Resp, GrpcError>> = future_flatten_to_stream(receiver)
            .map_err(GrpcError::from)
            .boxed();

        let receiver: GrpcStream<Resp> = stream_with_eof_and_error(receiver).boxed();

        (sender, receiver)
    }

    pub fn call_unary<Req : Send + 'static, Resp : Send + 'static>(&self, req: Req, method: Arc<MethodDescriptor<Req, Resp>>)
        -> GrpcFuture<Resp>
    {
        stream_single(self.call_server_streaming(req, method))
    }

    pub fn call_server_streaming<Req : Send + 'static, Resp : Send + 'static>(&self, req: Req, method: Arc<MethodDescriptor<Req, Resp>>)
        -> GrpcStream<Resp>
    {
        // A channel to send response back to caller
        let (complete, receiver) = self.new_resp_channel();

        fn header(name: &str, value: &str) -> StaticHeader {
            Header::new(name.as_bytes().to_owned(), value.as_bytes().to_owned())
        }

        let headers = vec![
            header(":method", "POST"),
            header(":path", &method.name),
            header(":authority", &self.host),
            header(":scheme", from_utf8(self.http_scheme.as_bytes()).unwrap()),
        ];

        let request_frames = {
            let method = method.clone();
            stream_once::<_, GrpcError>(req)
                .and_then(move |req| {
                    let grpc_frame = try!(method.req_marshaller.write(&req));
                    Ok(write_grpc_frame_to_vec(&grpc_frame))
                })
                .map_err(|e| HttpError::Other(Box::new(e)))
        };

        let start_request = self.loop_to_client.http_conn.start_request(
            headers,
            request_frames.boxed(),
            GrpcResponseHandler {
                tr: Box::new(GrpcResponseHandlerTyped {
                    method: method.clone(),
                    complete: complete,
                    remaining_response: Vec::new(),
                }),
            }
        ).map_err(GrpcError::from);

        future_flatten_to_stream(start_request.and_then(|()| Ok(receiver))).boxed()
    }

    pub fn call_client_streaming<Req : Send + 'static, Resp : Send + 'static>(&self, req: GrpcStream<Req>, method: Arc<MethodDescriptor<Req, Resp>>)
        -> GrpcFuture<Resp>
    {
        unimplemented!();
    }

    pub fn call_bidi<Req : Send + 'static, Resp : Send + 'static>(&self, req: GrpcStream<Req>, method: Arc<MethodDescriptor<Req, Resp>>)
        -> GrpcStream<Resp>
    {
        unimplemented!();
    }
}

// We shutdown client in destructor.
impl Drop for GrpcClient {
    fn drop(&mut self) {
        // ignore error because even loop may be already dead
        self.loop_to_client.shutdown_tx.send(()).ok();

        // do not ignore errors because we own event loop thread
        self.thread_join_handle.take().expect("handle.take")
            .join().expect("join thread");
    }
}

fn slice_get_header<'a>(headers: &'a [Header<'a, 'a>], name: &str) -> Option<&'a str> {
    headers.iter()
        .find(|h| h.name() == name.as_bytes())
        .and_then(|h| from_utf8(h.value()).ok())
}

fn default_stream_get_header<'a>(stream: &'a DefaultStream, name: &str) -> Option<&'a str> {
    match stream.headers {
        Some(ref v) => slice_get_header(v, name),
        None => None,
    }
}

// Event loop entry point
fn run_client_event_loop(
    socket_addr: SocketAddr,
    send_to_back: mpsc::Sender<LoopToClient>)
{
    // Create an event loop.
    let mut lp = Loop::new().unwrap();

    // Create a channel to receive shutdown signal.
    let (shutdown_tx, shutdown_rx) = lp.handle().channel();

    let (http_conn, http_conn_future) = HttpConnectionAsync::new(lp.handle(), &socket_addr);
    let http_conn_future = http_conn_future.map_err(GrpcError::from);

    // Send channels back to GrpcClient
    send_to_back
        .send(LoopToClient {
            shutdown_tx: shutdown_tx,
            loop_handle: lp.handle(),
            http_conn: http_conn,
        })
        .expect("send back");

    let shutdown_rx = shutdown_rx.map_err(GrpcError::from);

    let shutdown = shutdown_rx.and_then(|shutdown_rx| {
        shutdown_rx.into_future()
            .map_err(|_| GrpcError::Other("error in shutdown channel"))
            .and_then(|_| {
                // Must complete with error,
                // so `join` with this future cancels another future.
                let result: Result<(), _> = Err(GrpcError::Other("shutdown"));
                result
            })
    });

    // Wait for either completion of connection (i. e. error)
    // or shutdown signal.
    let done = http_conn_future.join(shutdown);

    // TODO: do not ignore error
    lp.run(done).ok();
}
<commit_msg>more streaming<commit_after>use std::thread;
use std::net::SocketAddr;
use std::net::ToSocketAddrs;
use std::sync::mpsc;
use std::sync::Arc;
use std::str::from_utf8;
use std::mem;

use futures;
use futures::Future;
use futures::stream::Stream;

use tokio_core::io as tokio_io;
use tokio_core::io::TaskIo;
use tokio_core::io::TaskIoRead;
use tokio_core::io::TaskIoWrite;

use tokio_core;
use tokio_core::Loop;
use tokio_core::LoopHandle;
use tokio_core::TcpStream;

use solicit::http::client::RequestStream;
use solicit::http::client::ClientSession;
use solicit::http::session::Client;
use solicit::http::session::SessionState;
use solicit::http::session::DefaultSessionState;
use solicit::http::session::DefaultStream;
use solicit::http::session::StreamState;
use solicit::http::session::StreamDataError;
use solicit::http::session::StreamDataChunk;
use solicit::http::session::Stream as solicit_Stream;
use solicit::http::connection::HttpConnection;
use solicit::http::connection::SendStatus;
use solicit::http::connection::EndStream;
use solicit::http::connection::DataChunk;
use solicit::http::priority::SimplePrioritizer;
use solicit::http::StreamId;
use solicit::http::HttpScheme;
use solicit::http::HttpError;
use solicit::http::HttpResult;
use solicit::http::Header;
use solicit::http::StaticHeader;

use method::MethodDescriptor;

use error::*;
use result::*;

use futures_misc::*;
use futures_grpc::*;

use grpc::*;
use solicit_async::*;
use solicit_misc::*;

use http_client::*;

trait GrpcResponseHandlerTrait : Send + 'static + HttpResponseHandler {
}

struct GrpcResponseHandlerTyped<Req : Send + 'static, Resp : Send + 'static> {
    method: Arc<MethodDescriptor<Req, Resp>>,
    complete: tokio_core::Sender<ResultOrEof<Resp, GrpcError>>,
    remaining_response: Vec<u8>,
}

impl<Req : Send + 'static, Resp : Send + 'static> GrpcResponseHandlerTrait for GrpcResponseHandlerTyped<Req, Resp> {
}

impl<Req : Send + 'static, Resp : Send + 'static> HttpResponseHandler for GrpcResponseHandlerTyped<Req, Resp> {
    fn headers(&mut self, headers: Vec<StaticHeader>) -> HttpResult<()> {
        println!("client: received headers");
        if slice_get_header(&headers, ":status") != Some("200") {
            if let Some(message) = slice_get_header(&headers, HEADER_GRPC_MESSAGE) {
                self.complete.send(ResultOrEof::Error(GrpcError::GrpcMessage(GrpcMessageError { grpc_message: message.to_owned() })));
                Err(HttpError::Other(Box::new(GrpcError::GrpcMessage(GrpcMessageError { grpc_message: message.to_owned() }))))
            } else {
                self.complete.send(ResultOrEof::Error(GrpcError::Other("not 200")));
                Err(HttpError::Other(Box::new(GrpcError::Other("not 200"))))
            }
        } else {
            Ok(())
        }
    }

    fn data_frame(&mut self, chunk: Vec<u8>) -> HttpResult<()> {
        self.remaining_response.extend(&chunk);
        loop {
            let len = match parse_grpc_frame(&self.remaining_response) {
                Err(e) => {
                    self.complete.send(ResultOrEof::Error(e));
                    break;
                }
                Ok(None) => break,
                Ok(Some((message, len))) => {
                    let resp = self.method.resp_marshaller.read(&message);
                    self.complete.send(From::from(resp)).ok();
                    len
                }
            };
            self.remaining_response.drain(..len);
        }
        Ok(())
    }

    fn trailers(&mut self, headers: Vec<StaticHeader>) -> HttpResult<()> {
        let status_200 = slice_get_header(&headers, ":status") == Some("200");
        let grpc_status_0 = slice_get_header(&headers, HEADER_GRPC_STATUS) == Some("0");
        if /* status_200 && */ grpc_status_0 {
            Ok(())
        } else {
            if let Some(message) = slice_get_header(&headers, HEADER_GRPC_MESSAGE) {
                self.complete.send(ResultOrEof::Error(GrpcError::GrpcMessage(GrpcMessageError { grpc_message: message.to_owned() })));
                Err(HttpError::Other(Box::new(GrpcError::GrpcMessage(GrpcMessageError { grpc_message: message.to_owned() }))))
            } else {
                self.complete.send(ResultOrEof::Error(GrpcError::Other("not xxx")));
                Err(HttpError::Other(Box::new(GrpcError::Other("not xxx"))))
            }
        }
    }

    fn end(&mut self) -> HttpResult<()> {
        self.complete.send(ResultOrEof::Eof);
        Ok(())
    }
}

struct GrpcResponseHandler {
    tr: Box<GrpcResponseHandlerTrait>,
}

impl HttpResponseHandler for GrpcResponseHandler {
    fn headers(&mut self, headers: Vec<StaticHeader>) -> HttpResult<()> {
        self.tr.headers(headers)
    }

    fn data_frame(&mut self, chunk: Vec<u8>) -> HttpResult<()> {
        self.tr.data_frame(chunk)
    }

    fn trailers(&mut self, headers: Vec<StaticHeader>) -> HttpResult<()> {
        self.tr.trailers(headers)
    }

    fn end(&mut self) -> HttpResult<()> {
        self.tr.end()
    }
}

// Data sent from event loop to GrpcClient
struct LoopToClient {
    // used only once to send shutdown signal
    shutdown_tx: tokio_core::Sender<()>,
    loop_handle: LoopHandle,
    http_conn: HttpConnectionAsync<GrpcResponseHandler>,
}

/// gRPC client implementation.
/// Used by generated code.
pub struct GrpcClient {
    loop_to_client: LoopToClient,
    thread_join_handle: Option<thread::JoinHandle<()>>,
    host: String,
    http_scheme: HttpScheme,
}

trait CallRequest : Send {
    // serialize message contained in this request
    fn write_req(&self) -> GrpcResult<Vec<u8>>;
    // gRPC method name (path)
    fn method_name(&self) -> &str;
    // called when server receives response from server
    // to send it back method called
    fn process_response(&mut self, message: ResultOrEof<&[u8], GrpcError>);
}

struct CallRequestHolder {
    request: Box<CallRequest>,
}

struct CallRequestTyped<Req, Resp> {
    // this call method descriptor
    method: Arc<MethodDescriptor<Req, Resp>>,
    // the request
    req: Req,
    // channel to send response back to called
    complete: tokio_core::Sender<ResultOrEof<Resp, GrpcError>>,
}

impl<Req : Send + 'static, Resp : Send + 'static> CallRequest for CallRequestTyped<Req, Resp> {
    fn write_req(&self) -> GrpcResult<Vec<u8>> {
        self.method.req_marshaller.write(&self.req)
    }

    fn method_name(&self) -> &str {
        &self.method.name
    }

    fn process_response(&mut self, message: ResultOrEof<&[u8], GrpcError>) {
        let resp = message.and_then(|message| self.method.resp_marshaller.read(&message));
        self.complete.send(resp).ok();
    }
}

impl GrpcClient {
    /// Create a client connected to specified host and port.
    pub fn new(host: &str, port: u16) -> GrpcResult<GrpcClient> {

        // TODO: sync
        // TODO: try connect to all addrs
        let socket_addr = try!((host, port).to_socket_addrs()).next().unwrap();

        // We need some data back from event loop.
        // This channel is used to exchange that data
        let (get_from_loop_tx, get_from_loop_rx) = mpsc::channel();

        // Start event loop.
        let join_handle = thread::spawn(move || {
            run_client_event_loop(socket_addr, get_from_loop_tx);
        });

        // Get back call channel and shutdown channel.
        let loop_to_client = try!(get_from_loop_rx.recv()
            .map_err(|_| GrpcError::Other("get response from loop")));

        Ok(GrpcClient {
            loop_to_client: loop_to_client,
            thread_join_handle: Some(join_handle),
            host: host.to_owned(),
            http_scheme: HttpScheme::Http,
        })
    }

    pub fn new_resp_channel<Resp : Send + 'static>(&self)
        -> (tokio_core::Sender<ResultOrEof<Resp, GrpcError>>, GrpcStream<Resp>)
    {
        let (sender, receiver) = self.loop_to_client.loop_handle.clone().channel::<ResultOrEof<Resp, GrpcError>>();
        let receiver: GrpcStream<ResultOrEof<Resp, GrpcError>> = future_flatten_to_stream(receiver)
            .map_err(GrpcError::from)
            .boxed();

        let receiver: GrpcStream<Resp> = stream_with_eof_and_error(receiver).boxed();

        (sender, receiver)
    }

    pub fn call_impl<Req : Send + 'static, Resp : Send + 'static>(&self, req: GrpcStream<Req>, method: Arc<MethodDescriptor<Req, Resp>>)
        -> GrpcStream<Resp>
    {
        // A channel to send response back to caller
        let (complete, receiver) = self.new_resp_channel();

        fn header(name: &str, value: &str) -> StaticHeader {
            Header::new(name.as_bytes().to_owned(), value.as_bytes().to_owned())
        }

        let headers = vec![
            header(":method", "POST"),
            header(":path", &method.name),
            header(":authority", &self.host),
            header(":scheme", from_utf8(self.http_scheme.as_bytes()).unwrap()),
        ];

        let request_frames = {
            let method = method.clone();
            req
                .and_then(move |req| {
                    let grpc_frame = try!(method.req_marshaller.write(&req));
                    Ok(write_grpc_frame_to_vec(&grpc_frame))
                })
                .map_err(|e| HttpError::Other(Box::new(e)))
        };

        let start_request = self.loop_to_client.http_conn.start_request(
            headers,
            request_frames.boxed(),
            GrpcResponseHandler {
                tr: Box::new(GrpcResponseHandlerTyped {
                    method: method.clone(),
                    complete: complete,
                    remaining_response: Vec::new(),
                }),
            }
        ).map_err(GrpcError::from);

        future_flatten_to_stream(start_request.and_then(|()| Ok(receiver))).boxed()
    }

    pub fn call_unary<Req : Send + 'static, Resp : Send + 'static>(&self, req: Req, method: Arc<MethodDescriptor<Req, Resp>>)
        -> GrpcFuture<Resp>
    {
        stream_single(self.call_impl(stream_once(req).boxed(), method))
    }

    pub fn call_server_streaming<Req : Send + 'static, Resp : Send + 'static>(&self, req: Req, method: Arc<MethodDescriptor<Req, Resp>>)
        -> GrpcStream<Resp>
    {
        self.call_impl(stream_once(req).boxed(), method)
    }

    pub fn call_client_streaming<Req : Send + 'static, Resp : Send + 'static>(&self, req: GrpcStream<Req>, method: Arc<MethodDescriptor<Req, Resp>>)
        -> GrpcFuture<Resp>
    {
        stream_single(self.call_impl(req, method))
    }

    pub fn call_bidi<Req : Send + 'static, Resp : Send + 'static>(&self, req: GrpcStream<Req>, method: Arc<MethodDescriptor<Req, Resp>>)
        -> GrpcStream<Resp>
    {
        self.call_impl(req, method)
    }
}

// We shutdown client in destructor.
impl Drop for GrpcClient {
    fn drop(&mut self) {
        // ignore error because even loop may be already dead
        self.loop_to_client.shutdown_tx.send(()).ok();

        // do not ignore errors because we own event loop thread
        self.thread_join_handle.take().expect("handle.take")
            .join().expect("join thread");
    }
}

fn slice_get_header<'a>(headers: &'a [Header<'a, 'a>], name: &str) -> Option<&'a str> {
    headers.iter()
        .find(|h| h.name() == name.as_bytes())
        .and_then(|h| from_utf8(h.value()).ok())
}

fn default_stream_get_header<'a>(stream: &'a DefaultStream, name: &str) -> Option<&'a str> {
    match stream.headers {
        Some(ref v) => slice_get_header(v, name),
        None => None,
    }
}

// Event loop entry point
fn run_client_event_loop(
    socket_addr: SocketAddr,
    send_to_back: mpsc::Sender<LoopToClient>)
{
    // Create an event loop.
    let mut lp = Loop::new().unwrap();

    // Create a channel to receive shutdown signal.
    let (shutdown_tx, shutdown_rx) = lp.handle().channel();

    let (http_conn, http_conn_future) = HttpConnectionAsync::new(lp.handle(), &socket_addr);
    let http_conn_future = http_conn_future.map_err(GrpcError::from);

    // Send channels back to GrpcClient
    send_to_back
        .send(LoopToClient {
            shutdown_tx: shutdown_tx,
            loop_handle: lp.handle(),
            http_conn: http_conn,
        })
        .expect("send back");

    let shutdown_rx = shutdown_rx.map_err(GrpcError::from);

    let shutdown = shutdown_rx.and_then(|shutdown_rx| {
        shutdown_rx.into_future()
            .map_err(|_| GrpcError::Other("error in shutdown channel"))
            .and_then(|_| {
                // Must complete with error,
                // so `join` with this future cancels another future.
                let result: Result<(), _> = Err(GrpcError::Other("shutdown"));
                result
            })
    });

    // Wait for either completion of connection (i. e. error)
    // or shutdown signal.
    let done = http_conn_future.join(shutdown);

    // TODO: do not ignore error
    lp.run(done).ok();
}
<|endoftext|>


<commit_before>//! Internal client
//!
//! Do not use directly - use the [`Algorithmia`](../struct.Algorithmia.html) struct instead
pub use hyper::client::response::Response;
use hyper::{Client, Url};
use hyper::client::RequestBuilder;
use hyper::header::{Authorization, UserAgent};
use hyper::method::Method;

use std::sync::Arc;
use std::clone;

pub use hyper::client::Body;

/// Represent the different ways to auth with the API
#[derive(Clone)]
pub enum ApiAuth {
    SimpleAuth(String),
    NoAuth
}

/// Internal HttpClient to build requests: wraps `hyper` client
pub struct HttpClient{
    pub base_url: String,
    api_auth: ApiAuth,
    hyper_client: Arc<Client>,
    user_agent: String,
}

impl HttpClient {
    /// Instantiate an HttpClient - creates a new `hyper` client
    pub fn new(api_auth: ApiAuth, base_url: String) -> HttpClient {
        HttpClient {
            api_auth: api_auth,
            base_url: base_url,
            hyper_client: Arc::new(Client::new()),
            user_agent: format!("algorithmia-rust/{} (Rust {}", option_env!("CARGO_PKG_VERSION").unwrap_or("unknown"), option_env!("CFG_RELEASE").unwrap_or("unknown")),
        }
    }

    /// Helper to make Algorithmia GET requests with the API key
    pub fn get(&self, url: Url) -> RequestBuilder {
        self.build_request(Method::Get, url)
    }

    /// Helper to make Algorithmia GET requests with the API key
    pub fn head(&self, url: Url) -> RequestBuilder {
        self.build_request(Method::Head, url)
    }

    /// Helper to make Algorithmia POST requests with the API key
    pub fn post(&self, url: Url) -> RequestBuilder {
        self.build_request(Method::Post, url)
    }

    /// Helper to make Algorithmia PUT requests with the API key
    pub fn put(&self, url: Url) -> RequestBuilder {
        self.build_request(Method::Put, url)
    }

    /// Helper to make Algorithmia POST requests with the API key
    pub fn delete(&self, url: Url) -> RequestBuilder {
        self.build_request(Method::Delete, url)
    }

    fn build_request(&self, verb: Method, url: Url) -> RequestBuilder {
        let mut req = self.hyper_client.request(verb, url);

        // TODO: Support Secure Auth
        req = req.header(UserAgent(self.user_agent.clone()));
        if let ApiAuth::SimpleAuth(ref api_key) = self.api_auth {
           req = req.header(Authorization(format!("Simple {}", api_key)))
        }
        req
    }
}

/// Allow cloning in order to reuse http client (and API key) for multiple connections
impl clone::Clone for HttpClient {
    fn clone(&self) -> HttpClient {
        HttpClient {
            api_auth: self.api_auth.clone(),
            base_url: self.base_url.clone(),
            hyper_client: self.hyper_client.clone(),
            user_agent: self.user_agent.clone(),
        }
    }
}

impl <'a> From<&'a str> for ApiAuth {
    fn from(api_key: &'a str) -> Self {
        match api_key.len() {
            0 => ApiAuth::NoAuth,
            _ => ApiAuth::SimpleAuth(api_key.into()),
        }
    }
}

impl From<()> for ApiAuth {
    fn from(_: ()) -> Self {
        ApiAuth::NoAuth
    }
}<commit_msg>Fix alt_client to stip any trailing slash appended by Url<commit_after>//! Internal client
//!
//! Do not use directly - use the [`Algorithmia`](../struct.Algorithmia.html) struct instead
pub use hyper::client::response::Response;
use hyper::{Client, Url};
use hyper::client::RequestBuilder;
use hyper::header::{Authorization, UserAgent};
use hyper::method::Method;

use std::sync::Arc;
use std::clone;

pub use hyper::client::Body;

/// Represent the different ways to auth with the API
#[derive(Clone)]
pub enum ApiAuth {
    SimpleAuth(String),
    NoAuth
}

/// Internal HttpClient to build requests: wraps `hyper` client
pub struct HttpClient{
    pub base_url: String,
    api_auth: ApiAuth,
    hyper_client: Arc<Client>,
    user_agent: String,
}

impl HttpClient {
    /// Instantiate an HttpClient - creates a new `hyper` client
    pub fn new(api_auth: ApiAuth, base_url: String) -> HttpClient {
        let checked_url = match base_url.chars().rev().next() {
            Some(c) if c == '/' => base_url[0..(base_url.len()-1)].to_string(),
            _ => base_url.clone(),
        };
        HttpClient {
            api_auth: api_auth,
            base_url: checked_url,
            hyper_client: Arc::new(Client::new()),
            user_agent: format!("algorithmia-rust/{} (Rust {}", option_env!("CARGO_PKG_VERSION").unwrap_or("unknown"), option_env!("CFG_RELEASE").unwrap_or("unknown")),
        }
    }

    /// Helper to make Algorithmia GET requests with the API key
    pub fn get(&self, url: Url) -> RequestBuilder {
        self.build_request(Method::Get, url)
    }

    /// Helper to make Algorithmia GET requests with the API key
    pub fn head(&self, url: Url) -> RequestBuilder {
        self.build_request(Method::Head, url)
    }

    /// Helper to make Algorithmia POST requests with the API key
    pub fn post(&self, url: Url) -> RequestBuilder {
        self.build_request(Method::Post, url)
    }

    /// Helper to make Algorithmia PUT requests with the API key
    pub fn put(&self, url: Url) -> RequestBuilder {
        self.build_request(Method::Put, url)
    }

    /// Helper to make Algorithmia POST requests with the API key
    pub fn delete(&self, url: Url) -> RequestBuilder {
        self.build_request(Method::Delete, url)
    }

    fn build_request(&self, verb: Method, url: Url) -> RequestBuilder {
        let mut req = self.hyper_client.request(verb, url);

        // TODO: Support Secure Auth
        req = req.header(UserAgent(self.user_agent.clone()));
        if let ApiAuth::SimpleAuth(ref api_key) = self.api_auth {
           req = req.header(Authorization(format!("Simple {}", api_key)))
        }
        req
    }
}

/// Allow cloning in order to reuse http client (and API key) for multiple connections
impl clone::Clone for HttpClient {
    fn clone(&self) -> HttpClient {
        HttpClient {
            api_auth: self.api_auth.clone(),
            base_url: self.base_url.clone(),
            hyper_client: self.hyper_client.clone(),
            user_agent: self.user_agent.clone(),
        }
    }
}

impl <'a> From<&'a str> for ApiAuth {
    fn from(api_key: &'a str) -> Self {
        match api_key.len() {
            0 => ApiAuth::NoAuth,
            _ => ApiAuth::SimpleAuth(api_key.into()),
        }
    }
}

impl From<()> for ApiAuth {
    fn from(_: ()) -> Self {
        ApiAuth::NoAuth
    }
}<|endoftext|>


<commit_before>use std::fs::File;
use std::io::{ BufReader, BufRead };

pub mod face;
pub mod point3d;

use self::face::Face;
use self::point3d::Point3D;

pub struct WavefrontModel {
	vertices: Vec<Point3D>,
	normals: Vec<Point3D>,
	faces: Vec<Face>
}

impl WavefrontModel {
	pub fn new(path: &'static str) -> WavefrontModel {
		let file = BufReader::new(File::open(&path).unwrap());

		let mut vertices = Vec::new();
		let mut normals = Vec::new();
		let mut faces = Vec::new();

		for wrapped_line in file.lines()  {
			let line: String = wrapped_line.unwrap();
			let mut tokens: Vec<&str> = line.split(' ').collect();

			let line_type = tokens.remove(0);

			match line_type {
				"v" => vertices.push(WavefrontModel::parse_point(tokens)),
				"vn" => normals.push(WavefrontModel::parse_point(tokens)),
				"f" => faces.push(WavefrontModel::parse_face(tokens)),
				_ => (),
			}
		}

		WavefrontModel {
			vertices: vertices,
			normals: normals,
			faces: faces
		}
	}

	fn parse_point(tokens: Vec<&str>) -> Point3D {
		Point3D { x: tokens[0].parse::<f32>().unwrap(), y: tokens[1].parse::<f32>().unwrap(), z: tokens[2].parse::<f32>().unwrap() }
	}

	fn parse_face(tokens: Vec<&str>) -> Face {
		let v1: Vec<&str> = tokens[0].split('/').collect();
		let v2: Vec<&str> = tokens[1].split('/').collect();
		let v3: Vec<&str> = tokens[2].split('/').collect();

		// Obj files index from `1` hence the `-1` offset
		Face {
			v1: v1[0].parse::<usize>().unwrap() - 1,
			vn1: v1[2].parse::<usize>().unwrap() - 1,

			v2: v2[0].parse::<usize>().unwrap() - 1,
			vn2: v2[2].parse::<usize>().unwrap() - 1,

			v3: v3[0].parse::<usize>().unwrap() - 1,
			vn3: v3[2].parse::<usize>().unwrap() - 1
		}
	}

	pub fn draw(&self) {
		for face in self.faces.iter() {
			// println!("Face [ {}, {}, {} ]", self.vertices[face.v1], self.vertices[face.v2], self.vertices[face.v3]);
			println!("{}", self.vertices[face.v1]);
		}
	}
}<commit_msg>Output all 3 vertices for face<commit_after>use std::fs::File;
use std::io::{ BufReader, BufRead };

pub mod face;
pub mod point3d;

use self::face::Face;
use self::point3d::Point3D;

pub struct WavefrontModel {
	vertices: Vec<Point3D>,
	normals: Vec<Point3D>,
	faces: Vec<Face>
}

impl WavefrontModel {
	pub fn new(path: &'static str) -> WavefrontModel {
		let file = BufReader::new(File::open(&path).unwrap());

		let mut vertices = Vec::new();
		let mut normals = Vec::new();
		let mut faces = Vec::new();

		for wrapped_line in file.lines()  {
			let line: String = wrapped_line.unwrap();
			let mut tokens: Vec<&str> = line.split(' ').collect();

			let line_type = tokens.remove(0);

			match line_type {
				"v" => vertices.push(WavefrontModel::parse_point(tokens)),
				"vn" => normals.push(WavefrontModel::parse_point(tokens)),
				"f" => faces.push(WavefrontModel::parse_face(tokens)),
				_ => (),
			}
		}

		WavefrontModel {
			vertices: vertices,
			normals: normals,
			faces: faces
		}
	}

	fn parse_point(tokens: Vec<&str>) -> Point3D {
		Point3D { x: tokens[0].parse::<f32>().unwrap(), y: tokens[1].parse::<f32>().unwrap(), z: tokens[2].parse::<f32>().unwrap() }
	}

	fn parse_face(tokens: Vec<&str>) -> Face {
		let v1: Vec<&str> = tokens[0].split('/').collect();
		let v2: Vec<&str> = tokens[1].split('/').collect();
		let v3: Vec<&str> = tokens[2].split('/').collect();

		// Obj files index from `1` hence the `-1` offset
		Face {
			v1: v1[0].parse::<usize>().unwrap() - 1,
			vn1: v1[2].parse::<usize>().unwrap() - 1,

			v2: v2[0].parse::<usize>().unwrap() - 1,
			vn2: v2[2].parse::<usize>().unwrap() - 1,

			v3: v3[0].parse::<usize>().unwrap() - 1,
			vn3: v3[2].parse::<usize>().unwrap() - 1
		}
	}

	pub fn draw(&self) {
		for face in self.faces.iter() {
			println!("Face [ {}, {}, {} ]", self.vertices[face.v1], self.vertices[face.v2], self.vertices[face.v3]);
		}
	}
}<|endoftext|>