<commit_before>// Copyright 2015 Joe Neeman.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!
This module provides functions for quickly skipping through the haystack, looking for places that
might conceivably be the start of a match. Just about everything in this module is an iterator over
`(usize, usize, usize)` triples.

  - The first `usize` is the index where the match begun. If this does turn out to be a match, the
    DFA should report the match as beginning here. This should always be at a character boundary.
  - The second `usize` is the index that the DFA should begin matching from. This could be
    different from the first index because we might already know what state the DFA would be in
    if it encountered the prefix we found. In that case, there is no need for the DFA to go back
    and re-examine the prefix. This should always be at a character boundary.
  - The third `usize` is the state that the DFA should start in.
 */

use aho_corasick::{Automaton, FullAcAutomaton};
use bytes::ByteSet;
use memchr::memchr;
use memmem::{Searcher, TwoWaySearcher};
use program::InitStates;
use std;

pub trait Skipper {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)>;
}

/* TODO: see whether specialization gives better performance than boxing. If it doesn't
 * go back to this implementation instead of the one in threaded.
impl Skipper {
    pub fn from_prefix<'a, 'b>(p: &'a Prefix, init: &'b InitStates) -> Box<Skipper + 'a + 'b> {
        use prefix::Prefix::*;

        match *p {
            Empty => Box::new(NoSkipper(init)),
            AsciiChar(ref set, _) => Box::new(SkipToAsciiSet(set.clone(), state)),
            Byte(b, _) => Box::new(SkipToByte(b)),
            Lit(ref s, _) => Box::new(SkipToStr(&s)),
            Ac(ref ac, _) => Box::new(AcSkipper(ac)),
            LoopUntil(ref set, _) => Box::new(LoopSkipper(set.clone())),
        }
    }
}
*/

/// An iterator that searchest for a given byte. The second position is position of the byte.
pub struct ByteIter<'a> {
    input: &'a [u8],
    byte: u8,
    pos: usize,
    state: usize,
}

impl<'a> ByteIter<'a> {
    pub fn new(s: &'a [u8], b: u8, state: usize) -> ByteIter<'a> {
        ByteIter {
            input: s,
            byte: b,
            pos: 0,
            state: state,
        }
    }
}

impl<'a> Iterator for ByteIter<'a> {
    type Item = (usize, usize, usize);

    fn next(&mut self) -> Option<(usize, usize, usize)> {
        let ret =
        memchr(self.byte, &self.input[self.pos..])
            .map(|pos| {
                self.pos += pos + 1;
                (self.pos - 1, self.pos - 1, self.state)
            });
        ret
    }
}

/// An iterator over (possibly overlapping) matches of a string. The second position is the one
/// at the start of the match.
pub struct StrIter<'hay, 'needle> {
    input: &'hay [u8],
    searcher: TwoWaySearcher<'needle>,
    pos: usize,
    state: usize,
}

impl<'hay, 'needle> StrIter<'hay, 'needle> {
    pub fn new(hay: &'hay [u8], needle: &'needle [u8], state: usize) -> StrIter<'hay, 'needle> {
        StrIter {
            input: hay,
            searcher: TwoWaySearcher::new(needle),
            pos: 0,
            state: state,
        }
    }
}

impl<'hay, 'needle> Iterator for StrIter<'hay, 'needle> {
    type Item = (usize, usize, usize);

    fn next(&mut self) -> Option<(usize, usize, usize)> {
        if let Some(pos) = self.searcher.search_in(&self.input[self.pos..]) {
            let ret = Some((self.pos + pos, self.pos + pos, self.state));
            self.pos += pos + 1;
            ret
        } else {
            None
        }
    }
}

/// An iterator over all non-overlapping (but possibly empty) strings of chars belonging to a given
/// set. The second position is the one after the end of the match.
pub struct LoopIter<'a, 'b> {
    input: &'a [u8],
    bytes: &'b ByteSet,
    pos: usize,
    state: usize,
}

impl<'a, 'b> LoopIter<'a, 'b> {
    pub fn new(input: &'a [u8], bytes: &'b ByteSet, state: usize) -> LoopIter<'a, 'b> {
        LoopIter {
            input: input,
            bytes: bytes,
            pos: 0,
            state: state,
        }
    }
}

impl<'a, 'b> Iterator for LoopIter<'a, 'b> {
    type Item = (usize, usize, usize);

    fn next(&mut self) -> Option<(usize, usize, usize)> {
        if let Some(pos) = self.input[self.pos..].iter()
                .position(|c| !self.bytes.0[*c as usize]) {
            let ret = Some((self.pos, self.pos + pos, self.state));
            self.pos += pos + 1;
            ret
        } else {
            None
        }
    }
}

/// An iterator over all characters belonging to a certain set of bytes. The second position is the
/// position of the match.
pub struct ByteSetIter<'a, 'b> {
    input: &'a [u8],
    bytes: &'b ByteSet,
    pos: usize,
    state: usize,
}

impl<'a, 'b> ByteSetIter<'a, 'b> {
    pub fn new(input: &'a [u8], bytes: &'b ByteSet, state: usize) -> ByteSetIter<'a, 'b> {
        ByteSetIter {
            input: input,
            bytes: bytes,
            pos: 0,
            state: state,
        }
    }
}

impl<'a, 'b> Iterator for ByteSetIter<'a, 'b> {
    type Item = (usize, usize, usize);

    fn next(&mut self) -> Option<(usize, usize, usize)> {
        if let Some(pos) = self.input[self.pos..].iter()
                .position(|c| self.bytes.0[*c as usize]) {
            self.pos += pos + 1;
            Some((self.pos - 1, self.pos - 1, self.state))
        } else {
            None
        }
    }
}

pub struct NoSkipper<'a>(pub &'a InitStates);
impl<'a> Skipper for NoSkipper<'a> {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        for pos in pos..s.len() {
            if let Some(state) = self.0.state_at_pos(s, pos) {
                return Some((pos, pos, state));
            }
        }
        None
    }
}

pub struct SkipToByte(pub u8, pub usize);
impl Skipper for SkipToByte {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        if let Some(offset) = memchr(self.0, &s[pos..]) {
            Some((pos + offset, pos + offset, self.1))
        } else {
            None
        }
    }
}

pub struct SkipToStr<'a>(pub TwoWaySearcher<'a>, pub usize);
impl<'a> SkipToStr<'a> {
    pub fn new(needle: &'a [u8], state: usize) -> SkipToStr<'a> {
        SkipToStr(TwoWaySearcher::new(needle), state)
    }
}
impl<'a> Skipper for SkipToStr<'a> {
    fn skip(&self, hay: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        if let Some(offset) = self.0.search_in(&hay[pos..]) {
            Some((pos + offset, pos + offset, self.1))
        } else {
            None
        }
    }
}

pub struct SkipToByteSet<'a>(pub &'a ByteSet, pub usize);
impl<'a> Skipper for SkipToByteSet<'a> {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        if let Some(offset) = s[pos..].iter().position(|c| (self.0).0[*c as usize]) {
            Some((pos + offset, pos + offset, self.1))
        } else {
            None
        }
    }
}

pub struct LoopSkipper<'a>(pub &'a ByteSet, pub usize);
impl<'a> Skipper for LoopSkipper<'a> {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        if let Some(offset) = s[pos..].iter().position(|c| !(self.0).0[*c as usize]) {
            Some((pos, pos + offset, self.1))
        } else {
            None
        }
    }
}

pub struct AcSkipper<'a>(pub &'a FullAcAutomaton<Vec<u8>>, pub usize);
impl<'a> Skipper for AcSkipper<'a> {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        let ac_input = unsafe { std::str::from_utf8_unchecked(&s[pos..]) };
        if let Some(mat) = self.0.find(ac_input).next() {
            Some((pos + mat.start, pos + mat.start, self.1))
        } else {
            None
        }
    }
}

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

    #[test]
    fn test_byte_iter() {
        let bi = ByteIter::new("abcaba".as_bytes(), 'a' as u8, 5);
        assert_eq!(bi.collect::<Vec<_>>(),
            vec![(0, 0, 5), (3, 3, 5), (5, 5, 5)]);
    }

    /*
    #[test]
    fn test_str_iter() {
        let si = StrIter::new("abcaba", "ab", 5);
        assert_eq!(si.collect::<Vec<_>>(),
            vec![(0, 0, 5), (3, 3, 5)]);

        let si = StrIter::new("aaaa", "aa", 5);
        assert_eq!(si.collect::<Vec<_>>(),
            vec![(0, 0, 5), (1, 1, 5), (2, 2, 5)]);
    }

    #[test]
    fn test_loop_iter() {
        let cs = ExtAsciiSet {
            set: AsciiSet::from_chars("b".chars()),
            contains_non_ascii: false,
        };
        let li = LoopIter::new("baaababaa", cs, 5);
        assert_eq!(li.collect::<Vec<_>>(),
            vec![(0, 0, 5), (1, 4, 5), (5, 6, 5)]);
    }

    #[test]
    fn test_ascii_set_iter() {
        let cs = AsciiSet::from_chars("ac".chars());
        let asi = AsciiSetIter::new("abcba", cs, 5);
        assert_eq!(asi.collect::<Vec<_>>(),
            vec![(0, 0, 5), (2, 2, 5), (4, 4, 5)]);
    }
    */
}
<commit_msg>Get rid of an unsafe, now that AC takes [u8].<commit_after>// Copyright 2015 Joe Neeman.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

/*!
This module provides functions for quickly skipping through the haystack, looking for places that
might conceivably be the start of a match. Just about everything in this module is an iterator over
`(usize, usize, usize)` triples.

  - The first `usize` is the index where the match begun. If this does turn out to be a match, the
    DFA should report the match as beginning here. This should always be at a character boundary.
  - The second `usize` is the index that the DFA should begin matching from. This could be
    different from the first index because we might already know what state the DFA would be in
    if it encountered the prefix we found. In that case, there is no need for the DFA to go back
    and re-examine the prefix. This should always be at a character boundary.
  - The third `usize` is the state that the DFA should start in.
 */

use aho_corasick::{Automaton, FullAcAutomaton};
use bytes::ByteSet;
use memchr::memchr;
use memmem::{Searcher, TwoWaySearcher};
use program::InitStates;

pub trait Skipper {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)>;
}

/* TODO: see whether specialization gives better performance than boxing. If it doesn't
 * go back to this implementation instead of the one in threaded.
impl Skipper {
    pub fn from_prefix<'a, 'b>(p: &'a Prefix, init: &'b InitStates) -> Box<Skipper + 'a + 'b> {
        use prefix::Prefix::*;

        match *p {
            Empty => Box::new(NoSkipper(init)),
            AsciiChar(ref set, _) => Box::new(SkipToAsciiSet(set.clone(), state)),
            Byte(b, _) => Box::new(SkipToByte(b)),
            Lit(ref s, _) => Box::new(SkipToStr(&s)),
            Ac(ref ac, _) => Box::new(AcSkipper(ac)),
            LoopUntil(ref set, _) => Box::new(LoopSkipper(set.clone())),
        }
    }
}
*/

/// An iterator that searchest for a given byte. The second position is position of the byte.
pub struct ByteIter<'a> {
    input: &'a [u8],
    byte: u8,
    pos: usize,
    state: usize,
}

impl<'a> ByteIter<'a> {
    pub fn new(s: &'a [u8], b: u8, state: usize) -> ByteIter<'a> {
        ByteIter {
            input: s,
            byte: b,
            pos: 0,
            state: state,
        }
    }
}

impl<'a> Iterator for ByteIter<'a> {
    type Item = (usize, usize, usize);

    fn next(&mut self) -> Option<(usize, usize, usize)> {
        let ret =
        memchr(self.byte, &self.input[self.pos..])
            .map(|pos| {
                self.pos += pos + 1;
                (self.pos - 1, self.pos - 1, self.state)
            });
        ret
    }
}

/// An iterator over (possibly overlapping) matches of a string. The second position is the one
/// at the start of the match.
pub struct StrIter<'hay, 'needle> {
    input: &'hay [u8],
    searcher: TwoWaySearcher<'needle>,
    pos: usize,
    state: usize,
}

impl<'hay, 'needle> StrIter<'hay, 'needle> {
    pub fn new(hay: &'hay [u8], needle: &'needle [u8], state: usize) -> StrIter<'hay, 'needle> {
        StrIter {
            input: hay,
            searcher: TwoWaySearcher::new(needle),
            pos: 0,
            state: state,
        }
    }
}

impl<'hay, 'needle> Iterator for StrIter<'hay, 'needle> {
    type Item = (usize, usize, usize);

    fn next(&mut self) -> Option<(usize, usize, usize)> {
        if let Some(pos) = self.searcher.search_in(&self.input[self.pos..]) {
            let ret = Some((self.pos + pos, self.pos + pos, self.state));
            self.pos += pos + 1;
            ret
        } else {
            None
        }
    }
}

/// An iterator over all non-overlapping (but possibly empty) strings of chars belonging to a given
/// set. The second position is the one after the end of the match.
pub struct LoopIter<'a, 'b> {
    input: &'a [u8],
    bytes: &'b ByteSet,
    pos: usize,
    state: usize,
}

impl<'a, 'b> LoopIter<'a, 'b> {
    pub fn new(input: &'a [u8], bytes: &'b ByteSet, state: usize) -> LoopIter<'a, 'b> {
        LoopIter {
            input: input,
            bytes: bytes,
            pos: 0,
            state: state,
        }
    }
}

impl<'a, 'b> Iterator for LoopIter<'a, 'b> {
    type Item = (usize, usize, usize);

    fn next(&mut self) -> Option<(usize, usize, usize)> {
        if let Some(pos) = self.input[self.pos..].iter()
                .position(|c| !self.bytes.0[*c as usize]) {
            let ret = Some((self.pos, self.pos + pos, self.state));
            self.pos += pos + 1;
            ret
        } else {
            None
        }
    }
}

/// An iterator over all characters belonging to a certain set of bytes. The second position is the
/// position of the match.
pub struct ByteSetIter<'a, 'b> {
    input: &'a [u8],
    bytes: &'b ByteSet,
    pos: usize,
    state: usize,
}

impl<'a, 'b> ByteSetIter<'a, 'b> {
    pub fn new(input: &'a [u8], bytes: &'b ByteSet, state: usize) -> ByteSetIter<'a, 'b> {
        ByteSetIter {
            input: input,
            bytes: bytes,
            pos: 0,
            state: state,
        }
    }
}

impl<'a, 'b> Iterator for ByteSetIter<'a, 'b> {
    type Item = (usize, usize, usize);

    fn next(&mut self) -> Option<(usize, usize, usize)> {
        if let Some(pos) = self.input[self.pos..].iter()
                .position(|c| self.bytes.0[*c as usize]) {
            self.pos += pos + 1;
            Some((self.pos - 1, self.pos - 1, self.state))
        } else {
            None
        }
    }
}

pub struct NoSkipper<'a>(pub &'a InitStates);
impl<'a> Skipper for NoSkipper<'a> {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        for pos in pos..s.len() {
            if let Some(state) = self.0.state_at_pos(s, pos) {
                return Some((pos, pos, state));
            }
        }
        None
    }
}

pub struct SkipToByte(pub u8, pub usize);
impl Skipper for SkipToByte {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        if let Some(offset) = memchr(self.0, &s[pos..]) {
            Some((pos + offset, pos + offset, self.1))
        } else {
            None
        }
    }
}

pub struct SkipToStr<'a>(pub TwoWaySearcher<'a>, pub usize);
impl<'a> SkipToStr<'a> {
    pub fn new(needle: &'a [u8], state: usize) -> SkipToStr<'a> {
        SkipToStr(TwoWaySearcher::new(needle), state)
    }
}
impl<'a> Skipper for SkipToStr<'a> {
    fn skip(&self, hay: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        if let Some(offset) = self.0.search_in(&hay[pos..]) {
            Some((pos + offset, pos + offset, self.1))
        } else {
            None
        }
    }
}

pub struct SkipToByteSet<'a>(pub &'a ByteSet, pub usize);
impl<'a> Skipper for SkipToByteSet<'a> {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        if let Some(offset) = s[pos..].iter().position(|c| (self.0).0[*c as usize]) {
            Some((pos + offset, pos + offset, self.1))
        } else {
            None
        }
    }
}

pub struct LoopSkipper<'a>(pub &'a ByteSet, pub usize);
impl<'a> Skipper for LoopSkipper<'a> {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        if let Some(offset) = s[pos..].iter().position(|c| !(self.0).0[*c as usize]) {
            Some((pos, pos + offset, self.1))
        } else {
            None
        }
    }
}

pub struct AcSkipper<'a>(pub &'a FullAcAutomaton<Vec<u8>>, pub usize);
impl<'a> Skipper for AcSkipper<'a> {
    fn skip(&self, s: &[u8], pos: usize) -> Option<(usize, usize, usize)> {
        if let Some(mat) = self.0.find(&s[pos..]).next() {
            Some((pos + mat.start, pos + mat.start, self.1))
        } else {
            None
        }
    }
}

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

    #[test]
    fn test_byte_iter() {
        let bi = ByteIter::new("abcaba".as_bytes(), 'a' as u8, 5);
        assert_eq!(bi.collect::<Vec<_>>(),
            vec![(0, 0, 5), (3, 3, 5), (5, 5, 5)]);
    }

    /*
    #[test]
    fn test_str_iter() {
        let si = StrIter::new("abcaba", "ab", 5);
        assert_eq!(si.collect::<Vec<_>>(),
            vec![(0, 0, 5), (3, 3, 5)]);

        let si = StrIter::new("aaaa", "aa", 5);
        assert_eq!(si.collect::<Vec<_>>(),
            vec![(0, 0, 5), (1, 1, 5), (2, 2, 5)]);
    }

    #[test]
    fn test_loop_iter() {
        let cs = ExtAsciiSet {
            set: AsciiSet::from_chars("b".chars()),
            contains_non_ascii: false,
        };
        let li = LoopIter::new("baaababaa", cs, 5);
        assert_eq!(li.collect::<Vec<_>>(),
            vec![(0, 0, 5), (1, 4, 5), (5, 6, 5)]);
    }

    #[test]
    fn test_ascii_set_iter() {
        let cs = AsciiSet::from_chars("ac".chars());
        let asi = AsciiSetIter::new("abcba", cs, 5);
        assert_eq!(asi.collect::<Vec<_>>(),
            vec![(0, 0, 5), (2, 2, 5), (4, 4, 5)]);
    }
    */
}
<|endoftext|>


<commit_before>use set1;

//challenge 9
pub fn pkcs_7(buf: &mut Vec<u8>, block_len: usize) {
    assert!(block_len < 256);
    let mut remain = block_len as isize - buf.len() as isize;
    let mut new_len = block_len;
    if remain < 0 {
        panic!("buf is larger than block size");
    } else if remain == 0  {
        new_len = block_len + block_len;
        remain = block_len as isize;
    }

    let to_pad = remain as u8;
    unsafe {
        let orig_len = buf.len();
        buf.set_len(new_len);
        let b_ptr = buf.as_mut_ptr().offset(orig_len as isize);
        for x in 0..remain {
            *b_ptr.offset(x as isize) = to_pad;
        }
    }
}

//challenge 10 and 15

#[derive(Debug)]
pub struct FormatError(());

pub fn pkcs_7_remove(buf: &mut Vec<u8>) -> Result<(), FormatError> {
    let last_byte = buf[buf.len() - 1];
    if last_byte as usize > buf.len() {
        return Err(FormatError(()));
    } else {
        let new_len = buf.len()- last_byte as usize;
        for &x in &buf[new_len..buf.len()-2] {
            if x != last_byte {
                return Err(FormatError(()));
            }
        }
        buf.resize(new_len, 0);
        return Ok(());
    }
}

fn shift_suffix<F>(mut f: F, attacker_len: usize, suffix_len: usize) -> Vec<u8>
        where F: FnMut(&[u8]) -> Vec<u8> {
    let mut attacker: Vec<u8> = vec![0; attacker_len];
    let attacker_len = attacker.len();
    let mut known_suffix: Vec<u8> = vec![];

    while known_suffix.len() != suffix_len {
        let mut ans = f(&vec![0; attacker_len - known_suffix.len() - 1]);
        ans.resize(attacker_len, 0);
        attacker[attacker_len-known_suffix.len()-1 .. attacker_len-1].clone_from_slice(&known_suffix);

        for x in 0u16..256u16 {
            attacker[attacker_len-1] = x as u8;
            // ans and trial result should only differs in 1 block, so it is a waste to compare
            // the whole array, but performance is not too important here.
            if f(&attacker)[..attacker_len] == ans[..attacker_len] {
                known_suffix.push(x as u8);
                break;
            }
        }
    }

    return known_suffix;
}

// challenge 12
pub fn ecb_decypt_appended<F>(mut f: F) -> Vec<u8>
        where F: FnMut(&[u8]) -> Vec<u8> {
    let p0_len = f(&[0u8; 0]).len();

    let mut block_len = 0;
    let mut suffix_len = 0;
    for x in 1usize.. {
        let l = f(&vec![0u8; x]).len();
        if l != p0_len {
            suffix_len = p0_len - x;
            block_len = l - p0_len;
            break;
        }
    }

    let attacker_len = suffix_len - (suffix_len % block_len) + block_len;
    shift_suffix(f, attacker_len, suffix_len)
}

//challenge14
pub fn ecb_wrapped_decypt_appended<F>(mut f: F) -> Vec<u8>
        where F: FnMut(&[u8]) -> Vec<u8> {
    let p0_len = f(&[0u8; 0]).len();

    /*
     * Symbols Meaning
     *<--pre--><--attacker_payload--><--post-->
     *^ len:p ^^       len:k        ^^ len:z  ^,  whete p+k+o mod 16 = 0
     *         ^     len:a        ^  ,where p+a mod 16 = 0
     */
    let mut block_len = 0;
    //let mut suffix_len = 0;
    let mut k = 0;
    for x in 1usize.. {
        let l = f(&vec![0u8; x]).len();
        if l != p0_len {
            block_len = l - p0_len;
            k = x + l;  // +l just make k sufficiently large, not always necessary.
            break;
        }
    }

    let mut attacker: Vec<u8> = vec![0; k];
    let enc_all0 = f(&attacker);
    attacker[k-1] = 1;
    let enc_f1 = f(&attacker);
    // pa means p + a
    let pa = enc_all0.iter().zip(enc_f1.iter()).position(|(a,b)| a!= b).unwrap() - 1;
    let mut a = k - 2;  // the index to change from 0 to 1 of attacker

    loop {
        attacker[a] = 1;
        // check if the a-th byte is at block ending boundary by diffing.  Check more than 1 byte
        // to avoid collision.
        if f(&attacker)[pa-1..pa+1] != enc_f1[pa-1..pa+1] {
            break;
        }
        a = a - 1;
    }

    // enc_all0.len()-block_len is equivalent to p+k+z
    let z = (enc_all0.len()-block_len) - pa - (k-a);
    a = a+1;  // convert c-style index to length

    shift_suffix(f, a, z)
}

//challenge16, 128 bits block size is asumed
pub fn cbc_flip(ciphertext: &mut [u8], plaintext: &[u8], desired: &[u8]) {
    assert!(ciphertext.len() >= 32 && plaintext.len() >= 32);
    assert!(desired.len() == 16);
    // Need to xor 3 buffers togeether and then safe it into first block of ciphertext.
    // It is easier to allocate a tmp
    // let tmp = vec![0; 16];
    // xor to there first and then xor again to the first block.
    // But good C programmer will not do this.  We will take the first block as tmp.
    // However it means borrow it twice, one as mut.  That is unsafe in rust.
    // But think about it this operation is actually valid operation.
    // So we should not be afraid of `unsafe {}`
    unsafe {
        let block: *mut [u8] = &mut ciphertext[0..16];
        set1::xor_buffers_buf(&ciphertext[0..16], &plaintext[16..32], &mut *block);
        set1::xor_buffers_buf(&ciphertext[0..16], &desired[..], &mut *block);
    }
}

#[cfg(test)]
mod test {
    use super::*;
    extern crate base64;
    extern crate openssl;
    use self::openssl::symm::{self, Cipher};

    struct MyEncrypter {
        suffix: Vec<u8>,
        key: Vec<u8>
    }

    impl MyEncrypter {
        pub fn new() -> Self {
            let secret = b"Um9sbGluJyBpbiBteSA1LjAKV2l0aCBteSByYWctdG9wIGRvd24gc28gbXkg
aGFpciBjYW4gYmxvdwpUaGUgZ2lybGllcyBvbiBzdGFuZGJ5IHdhdmluZyBq
dXN0IHRvIHNheSBoaQpEaWQgeW91IHN0b3A/IE5vLCBJIGp1c3QgZHJvdmUg
YnkK".to_vec();
            let secret = base64::decode_config(&secret, base64::MIME).expect("Base64 malformated");
            MyEncrypter {
                suffix: secret,
                key: b"LorentzTransform".to_vec()
            }
        }

        pub fn encrypt(&mut self, buf: &[u8]) -> Vec<u8> {
            let mut plaintext = vec![];
            plaintext.extend_from_slice(buf);
            plaintext.extend_from_slice(&self.suffix);

            //println!("encrypting {:?}", plaintext);
            let cipher = Cipher::aes_128_ecb();
            return symm::encrypt(
                cipher,
                &self.key,
                None,
                &plaintext).unwrap();
        }
    }

    #[test]
    fn test_ecb_decypt_appended() {
        let mut encrypter = MyEncrypter::new();
        let res = ecb_decypt_appended(|buf| encrypter.encrypt(buf));
        // The answer starts with "Rollin' in my", not showing the full one here as a spoiler
        assert!(res.starts_with(&b"Rollin' in my"[..]));
    }

    #[test]
    fn test_invalid_pkcs_7_remove() {
        let mut block = b"ICEI ICE BABY\x05\x05\x05\x05".to_vec();
        assert!(pkcs_7_remove(&mut block).is_err());
    }

    #[test]
    fn test_pkcs_7_not_multiple() {
        let mut block = b"Minkowski".to_vec();
        pkcs_7(&mut block, 16);
        assert_eq!(&block, b"Minkowski\x07\x07\x07\x07\x07\x07\x07");
    }

    #[test]
    fn test_pkcs_7_multiple() {
        let mut block = b"Euclid".to_vec();
        pkcs_7(&mut block, 6);
        assert_eq!(&block, b"Euclid\x06\x06\x06\x06\x06\x06");
    }
}
<commit_msg>Bugfix on pkcs7 removal<commit_after>use set1;

//challenge 9
pub fn pkcs_7(buf: &mut Vec<u8>, block_len: usize) {
    assert!(block_len < 256);
    let mut remain = block_len as isize - buf.len() as isize;
    let mut new_len = block_len;
    if remain < 0 {
        panic!("buf is larger than block size");
    } else if remain == 0  {
        new_len = block_len + block_len;
        remain = block_len as isize;
    }

    let to_pad = remain as u8;
    unsafe {
        let orig_len = buf.len();
        buf.set_len(new_len);
        let b_ptr = buf.as_mut_ptr().offset(orig_len as isize);
        for x in 0..remain {
            *b_ptr.offset(x as isize) = to_pad;
        }
    }
}

//challenge 10 and 15

#[derive(Debug)]
pub struct FormatError(());

pub fn pkcs_7_remove(buf: &mut Vec<u8>) -> Result<(), FormatError> {
    let last_byte = buf[buf.len() - 1];
    if last_byte as usize > buf.len() || last_byte == 0 {
        return Err(FormatError(()));
    } else {
        let new_len = buf.len() - last_byte as usize;
        for &x in &buf[new_len..buf.len()-1] {
            if x != last_byte {
                return Err(FormatError(()));
            }
        }
        buf.resize(new_len, 0);
        return Ok(());
    }
}

fn shift_suffix<F>(mut f: F, attacker_len: usize, suffix_len: usize) -> Vec<u8>
        where F: FnMut(&[u8]) -> Vec<u8> {
    let mut attacker: Vec<u8> = vec![0; attacker_len];
    let attacker_len = attacker.len();
    let mut known_suffix: Vec<u8> = vec![];

    while known_suffix.len() != suffix_len {
        let mut ans = f(&vec![0; attacker_len - known_suffix.len() - 1]);
        ans.resize(attacker_len, 0);
        attacker[attacker_len-known_suffix.len()-1 .. attacker_len-1].clone_from_slice(&known_suffix);

        for x in 0u16..256u16 {
            attacker[attacker_len-1] = x as u8;
            // ans and trial result should only differs in 1 block, so it is a waste to compare
            // the whole array, but performance is not too important here.
            if f(&attacker)[..attacker_len] == ans[..attacker_len] {
                known_suffix.push(x as u8);
                break;
            }
        }
    }

    return known_suffix;
}

// challenge 12
pub fn ecb_decypt_appended<F>(mut f: F) -> Vec<u8>
        where F: FnMut(&[u8]) -> Vec<u8> {
    let p0_len = f(&[0u8; 0]).len();

    let mut block_len = 0;
    let mut suffix_len = 0;
    for x in 1usize.. {
        let l = f(&vec![0u8; x]).len();
        if l != p0_len {
            suffix_len = p0_len - x;
            block_len = l - p0_len;
            break;
        }
    }

    let attacker_len = suffix_len - (suffix_len % block_len) + block_len;
    shift_suffix(f, attacker_len, suffix_len)
}

//challenge14
pub fn ecb_wrapped_decypt_appended<F>(mut f: F) -> Vec<u8>
        where F: FnMut(&[u8]) -> Vec<u8> {
    let p0_len = f(&[0u8; 0]).len();

    /*
     * Symbols Meaning
     *<--pre--><--attacker_payload--><--post-->
     *^ len:p ^^       len:k        ^^ len:z  ^,  whete p+k+o mod 16 = 0
     *         ^     len:a        ^  ,where p+a mod 16 = 0
     */
    let mut block_len = 0;
    //let mut suffix_len = 0;
    let mut k = 0;
    for x in 1usize.. {
        let l = f(&vec![0u8; x]).len();
        if l != p0_len {
            block_len = l - p0_len;
            k = x + l;  // +l just make k sufficiently large, not always necessary.
            break;
        }
    }

    let mut attacker: Vec<u8> = vec![0; k];
    let enc_all0 = f(&attacker);
    attacker[k-1] = 1;
    let enc_f1 = f(&attacker);
    // pa means p + a
    let pa = enc_all0.iter().zip(enc_f1.iter()).position(|(a,b)| a!= b).unwrap() - 1;
    let mut a = k - 2;  // the index to change from 0 to 1 of attacker

    loop {
        attacker[a] = 1;
        // check if the a-th byte is at block ending boundary by diffing.  Check more than 1 byte
        // to avoid collision.
        if f(&attacker)[pa-1..pa+1] != enc_f1[pa-1..pa+1] {
            break;
        }
        a = a - 1;
    }

    // enc_all0.len()-block_len is equivalent to p+k+z
    let z = (enc_all0.len()-block_len) - pa - (k-a);
    a = a+1;  // convert c-style index to length

    shift_suffix(f, a, z)
}

//challenge16, 128 bits block size is asumed
pub fn cbc_flip(ciphertext: &mut [u8], plaintext: &[u8], desired: &[u8]) {
    assert!(ciphertext.len() >= 32 && plaintext.len() >= 32);
    assert!(desired.len() == 16);
    // Need to xor 3 buffers together and then save it into first block of ciphertext.
    // It is easier to allocate a tmp
    // let tmp = vec![0; 16];
    // xor into tmp first and then xor again to the first block.
    // But good C programmer will not do this.  We will take the first block as tmp.
    // However it means borrow it twice, one as mut.  That is unsafe in rust.
    // But think about it this operation is actually valid operation.
    // So we should not be afraid of `unsafe {}`
    unsafe {
        let block: *mut [u8] = &mut ciphertext[0..16];
        set1::xor_buffers_buf(&ciphertext[0..16], &plaintext[16..32], &mut *block);
        set1::xor_buffers_buf(&ciphertext[0..16], &desired[..], &mut *block);
    }
}

#[cfg(test)]
mod test {
    use super::*;
    extern crate base64;
    extern crate openssl;
    use self::openssl::symm::{self, Cipher};

    struct MyEncrypter {
        suffix: Vec<u8>,
        key: Vec<u8>
    }

    impl MyEncrypter {
        pub fn new() -> Self {
            let secret = b"Um9sbGluJyBpbiBteSA1LjAKV2l0aCBteSByYWctdG9wIGRvd24gc28gbXkg
aGFpciBjYW4gYmxvdwpUaGUgZ2lybGllcyBvbiBzdGFuZGJ5IHdhdmluZyBq
dXN0IHRvIHNheSBoaQpEaWQgeW91IHN0b3A/IE5vLCBJIGp1c3QgZHJvdmUg
YnkK".to_vec();
            let secret = base64::decode_config(&secret, base64::MIME).expect("Base64 malformated");
            MyEncrypter {
                suffix: secret,
                key: b"LorentzTransform".to_vec()
            }
        }

        pub fn encrypt(&mut self, buf: &[u8]) -> Vec<u8> {
            let mut plaintext = vec![];
            plaintext.extend_from_slice(buf);
            plaintext.extend_from_slice(&self.suffix);

            //println!("encrypting {:?}", plaintext);
            let cipher = Cipher::aes_128_ecb();
            return symm::encrypt(
                cipher,
                &self.key,
                None,
                &plaintext).unwrap();
        }
    }

    #[test]
    fn test_ecb_decypt_appended() {
        let mut encrypter = MyEncrypter::new();
        let res = ecb_decypt_appended(|buf| encrypter.encrypt(buf));
        // The answer starts with "Rollin' in my", not showing the full one here as a spoiler
        assert!(res.starts_with(&b"Rollin' in my"[..]));
    }

    #[test]
    fn test_invalid_pkcs_7_remove() {
        let mut block = b"ICEI ICE BABY\x05\x05\x05\x05".to_vec();
        assert!(pkcs_7_remove(&mut block).is_err());
    }

    #[test]
    fn test_pkcs_7_not_multiple() {
        let mut block = b"Minkowski".to_vec();
        pkcs_7(&mut block, 16);
        assert_eq!(&block, b"Minkowski\x07\x07\x07\x07\x07\x07\x07");
    }

    #[test]
    fn test_pkcs_7_multiple() {
        let mut block = b"Euclid".to_vec();
        pkcs_7(&mut block, 6);
        assert_eq!(&block, b"Euclid\x06\x06\x06\x06\x06\x06");
    }
}
<|endoftext|>


<commit_before>
use imp_prelude::*;
use error::{ShapeError, ErrorKind, from_kind};

/// Stack arrays along the given axis.
///
/// ```
/// use ndarray::{arr2, Axis, stack};
///
/// let a = arr2(&[[2., 2.],
///                [3., 3.]]);
/// assert!(
///     stack(Axis(0), &[a.view(), a.view()])
///     == Ok(arr2(&[[2., 2.],
///                  [3., 3.],
///                  [2., 2.],
///                  [3., 3.]]))
/// );
/// ```
pub fn stack<'a, A, D>(axis: Axis, arrays: &[ArrayView<'a, A, D>])
    -> Result<OwnedArray<A, D>, ShapeError>
    where A: Copy,
          D: RemoveAxis
{
    if arrays.len() == 0 {
        return Err(from_kind(ErrorKind::Unsupported));
    }
    let mut res_dim = arrays[0].dim();
    if axis.axis() >= res_dim.ndim() {
        return Err(from_kind(ErrorKind::OutOfBounds));
    }
    let common_dim = res_dim.remove_axis(axis);
    if arrays.iter().any(|a| a.dim().remove_axis(axis) != common_dim) {
        return Err(from_kind(ErrorKind::IncompatibleShape));
    }

    let stacked_dim = arrays.iter()
                            .fold(0, |acc, a| acc + a.shape().axis(axis));
    res_dim.set_axis(axis, stacked_dim);

    // we can safely use uninitialized values here because they are Copy
    // and we will only ever write to them
    let size = res_dim.size();
    let mut v = Vec::with_capacity(size);
    unsafe {
        v.set_len(size);
    }
    let mut res = try!(OwnedArray::from_vec_dim(res_dim, v));

    {
        let mut assign_view = res.view_mut();
        for array in arrays {
            let len = array.shape().axis(axis);
            let (mut front, rest) = assign_view.split_at(axis, len);
            front.assign(array);
            assign_view = rest;
        }
    }
    Ok(res)
}

/// Stack arrays along the given axis.
///
/// Uses the [`stack`][1] function, calling `ArrayView::from(&a)` on each
/// argument `a`.
///
/// [1]: fn.stack.html
///
/// ```
/// #[macro_use(stack)]
/// extern crate ndarray;
///
/// use ndarray::{arr2, Axis, stack};
///
/// # fn main() {
///
/// let a = arr2(&[[2., 2.],
///                [3., 3.]]);
/// assert!(
///     stack![Axis(0), a, a]
///     == arr2(&[[2., 2.],
///               [3., 3.],
///               [2., 2.],
///               [3., 3.]])
/// );
/// # }
/// ```
#[macro_export]
macro_rules! stack {
    ($axis:expr, $( $array:expr ),+ ) => {
        ndarray::stack($axis, &[ $(ndarray::ArrayView::from(&$array) ),* ]).unwrap()
    }
}
<commit_msg>Edit docs for stack![]<commit_after>
use imp_prelude::*;
use error::{ShapeError, ErrorKind, from_kind};

/// Stack arrays along the given axis.
///
/// ```
/// use ndarray::{arr2, Axis, stack};
///
/// let a = arr2(&[[2., 2.],
///                [3., 3.]]);
/// assert!(
///     stack(Axis(0), &[a.view(), a.view()])
///     == Ok(arr2(&[[2., 2.],
///                  [3., 3.],
///                  [2., 2.],
///                  [3., 3.]]))
/// );
/// ```
pub fn stack<'a, A, D>(axis: Axis, arrays: &[ArrayView<'a, A, D>])
    -> Result<OwnedArray<A, D>, ShapeError>
    where A: Copy,
          D: RemoveAxis
{
    if arrays.len() == 0 {
        return Err(from_kind(ErrorKind::Unsupported));
    }
    let mut res_dim = arrays[0].dim();
    if axis.axis() >= res_dim.ndim() {
        return Err(from_kind(ErrorKind::OutOfBounds));
    }
    let common_dim = res_dim.remove_axis(axis);
    if arrays.iter().any(|a| a.dim().remove_axis(axis) != common_dim) {
        return Err(from_kind(ErrorKind::IncompatibleShape));
    }

    let stacked_dim = arrays.iter()
                            .fold(0, |acc, a| acc + a.shape().axis(axis));
    res_dim.set_axis(axis, stacked_dim);

    // we can safely use uninitialized values here because they are Copy
    // and we will only ever write to them
    let size = res_dim.size();
    let mut v = Vec::with_capacity(size);
    unsafe {
        v.set_len(size);
    }
    let mut res = try!(OwnedArray::from_vec_dim(res_dim, v));

    {
        let mut assign_view = res.view_mut();
        for array in arrays {
            let len = array.shape().axis(axis);
            let (mut front, rest) = assign_view.split_at(axis, len);
            front.assign(array);
            assign_view = rest;
        }
    }
    Ok(res)
}

/// Stack arrays along the given axis.
///
/// Uses the [`stack`][1] function, calling `ArrayView::from(&a)` on each
/// argument `a`.
///
/// [1]: fn.stack.html
///
/// ***Panics*** if the `stack` function would return an error.
///
/// ```
/// #[macro_use(stack)]
/// extern crate ndarray;
///
/// use ndarray::{arr2, Axis, stack};
///
/// # fn main() {
///
/// let a = arr2(&[[2., 2.],
///                [3., 3.]]);
/// assert!(
///     stack![Axis(0), a, a]
///     == arr2(&[[2., 2.],
///               [3., 3.],
///               [2., 2.],
///               [3., 3.]])
/// );
/// # }
/// ```
#[macro_export]
macro_rules! stack {
    ($axis:expr, $( $array:expr ),+ ) => {
        ndarray::stack($axis, &[ $(ndarray::ArrayView::from(&$array) ),* ]).unwrap()
    }
}
<|endoftext|>


<commit_before>use cargo_test_support::registry::{Dependency, Package};
use cargo_test_support::ProjectBuilder;
use cargo_test_support::{is_nightly, paths, project, rustc_host, Execs};
use std::path::PathBuf;

struct Setup {
    rustc_wrapper: PathBuf,
    real_sysroot: String,
}

fn setup() -> Option<Setup> {
    if !is_nightly() {
        // -Zbuild-std is nightly
        // We don't want these tests to run on rust-lang/rust.
        return None;
    }

    // Our mock sysroot requires a few packages from crates.io, so make sure
    // they're "published" to crates.io. Also edit their code a bit to make sure
    // that they have access to our custom crates with custom apis.
    Package::new("registry-dep-using-core", "1.0.0")
        .file(
            "src/lib.rs",
            "
                #![no_std]

                #[cfg(feature = \"mockbuild\")]
                pub fn custom_api() {
                }

                #[cfg(not(feature = \"mockbuild\"))]
                pub fn non_sysroot_api() {
                    core::custom_api();
                }
            ",
        )
        .add_dep(Dependency::new("rustc-std-workspace-core", "*").optional(true))
        .feature("mockbuild", &["rustc-std-workspace-core"])
        .publish();
    Package::new("registry-dep-using-alloc", "1.0.0")
        .file(
            "src/lib.rs",
            "
                #![no_std]

                extern crate alloc;

                #[cfg(feature = \"mockbuild\")]
                pub fn custom_api() {
                }

                #[cfg(not(feature = \"mockbuild\"))]
                pub fn non_sysroot_api() {
                    core::custom_api();
                    alloc::custom_api();
                }
            ",
        )
        .add_dep(Dependency::new("rustc-std-workspace-core", "*").optional(true))
        .add_dep(Dependency::new("rustc-std-workspace-alloc", "*").optional(true))
        .feature(
            "mockbuild",
            &["rustc-std-workspace-core", "rustc-std-workspace-alloc"],
        )
        .publish();
    Package::new("registry-dep-using-std", "1.0.0")
        .file(
            "src/lib.rs",
            "
                #[cfg(feature = \"mockbuild\")]
                pub fn custom_api() {
                }

                #[cfg(not(feature = \"mockbuild\"))]
                pub fn non_sysroot_api() {
                    std::custom_api();
                }
            ",
        )
        .add_dep(Dependency::new("rustc-std-workspace-std", "*").optional(true))
        .feature("mockbuild", &["rustc-std-workspace-std"])
        .publish();

    let p = ProjectBuilder::new(paths::root().join("rustc-wrapper"))
        .file(
            "src/main.rs",
            r#"
                use std::process::Command;
                use std::env;
                fn main() {
                    let mut args = env::args().skip(1).collect::<Vec<_>>();

                    let crates = [
                        "alloc",
                        "compiler_builtins",
                        "core",
                        "panic_unwind",
                        "proc_macro",
                        "std",
                        "test",
                    ];
                    let is_sysroot_crate = args.iter()
                        .any(|a| a.starts_with("rustc_std_workspace") || crates.iter().any(|b| a == b));

                    if is_sysroot_crate {
                        let arg = args.iter().position(|a| a == "--sysroot").unwrap();
                        args[arg + 1] = env::var("REAL_SYSROOT").unwrap();
                    }

                    let ret = Command::new(&args[0]).args(&args[1..]).status().unwrap();
                    std::process::exit(ret.code().unwrap_or(1));
                }
            "#,
        )
        .build();
    p.cargo("build").run();

    return Some(Setup {
        rustc_wrapper: p.bin("foo"),
        real_sysroot: paths::sysroot(),
    });
}

fn enable_build_std(e: &mut Execs, setup: &Setup, arg: Option<&str>) {
    // First up, force Cargo to use our "mock sysroot" which mimics what
    // libstd looks like upstream.
    let root = paths::root();
    let root = root
        .parent() // chop off test name
        .unwrap()
        .parent() // chop off `citN`
        .unwrap()
        .parent() // chop off `target`
        .unwrap()
        .join("tests/testsuite/mock-std");
    e.env("__CARGO_TESTS_ONLY_SRC_ROOT", &root);

    // Actually enable `-Zbuild-std` for now
    let arg = match arg {
        Some(s) => format!("-Zbuild-std={}", s),
        None => "-Zbuild-std".to_string(),
    };
    e.arg(arg);
    e.masquerade_as_nightly_cargo();

    // We do various shenanigans to ensure our "mock sysroot" actually links
    // with the real sysroot, so we don't have to actually recompile std for
    // each test. Perform all that logic here, namely:
    //
    // * RUSTC_WRAPPER - uses our shim executable built above to control rustc
    // * REAL_SYSROOT - used by the shim executable to swap out to the real
    //   sysroot temporarily for some compilations
    // * RUST{,DOC}FLAGS - an extra `-L` argument to ensure we can always load
    //   crates from the sysroot, but only indirectly through other crates.
    e.env("RUSTC_WRAPPER", &setup.rustc_wrapper);
    e.env("REAL_SYSROOT", &setup.real_sysroot);
    let libdir = format!("/lib/rustlib/{}/lib", rustc_host());
    e.env(
        "RUSTFLAGS",
        format!("-Ldependency={}{}", setup.real_sysroot, libdir),
    );
    e.env(
        "RUSTDOCFLAGS",
        format!("-Ldependency={}{}", setup.real_sysroot, libdir),
    );
}

// Helper methods used in the tests below
trait BuildStd: Sized {
    fn build_std(&mut self, setup: &Setup) -> &mut Self;
    fn build_std_arg(&mut self, setup: &Setup, arg: &str) -> &mut Self;
    fn target_host(&mut self) -> &mut Self;
}

impl BuildStd for Execs {
    fn build_std(&mut self, setup: &Setup) -> &mut Self {
        enable_build_std(self, setup, None);
        self
    }

    fn build_std_arg(&mut self, setup: &Setup, arg: &str) -> &mut Self {
        enable_build_std(self, setup, Some(arg));
        self
    }

    fn target_host(&mut self) -> &mut Self {
        self.arg("--target").arg(rustc_host());
        self
    }
}

#[cargo_test]
fn basic() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };

    let p = project()
        .file(
            "src/main.rs",
            "
                fn main() {
                    std::custom_api();
                    foo::f();
                }

                #[test]
                fn smoke_bin_unit() {
                    std::custom_api();
                    foo::f();
                }
            ",
        )
        .file(
            "src/lib.rs",
            "
                extern crate alloc;
                extern crate proc_macro;

                /// ```
                /// foo::f();
                /// ```
                pub fn f() {
                    core::custom_api();
                    std::custom_api();
                    alloc::custom_api();
                    proc_macro::custom_api();
                }

                #[test]
                fn smoke_lib_unit() {
                    std::custom_api();
                    f();
                }
            ",
        )
        .file(
            "tests/smoke.rs",
            "
                #[test]
                fn smoke_integration() {
                    std::custom_api();
                    foo::f();
                }
            ",
        )
        .build();

    p.cargo("check").build_std(&setup).target_host().run();
    p.cargo("build").build_std(&setup).target_host().run();
    p.cargo("run").build_std(&setup).target_host().run();
    p.cargo("test").build_std(&setup).target_host().run();
}

#[cargo_test]
fn simple_lib_std() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project().file("src/lib.rs", "").build();
    p.cargo("build -v")
        .build_std(&setup)
        .target_host()
        .with_stderr_contains("[RUNNING] `[..]--crate-name std [..]`")
        .run();
    // Check freshness.
    p.change_file("src/lib.rs", " ");
    p.cargo("build -v")
        .build_std(&setup)
        .target_host()
        .with_stderr_contains("[FRESH] std[..]")
        .run();
}

#[cargo_test]
fn simple_bin_std() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project().file("src/main.rs", "fn main() {}").build();
    p.cargo("run -v").build_std(&setup).target_host().run();
}

#[cargo_test]
fn lib_nostd() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                #![no_std]
                pub fn foo() {
                    assert_eq!(core::u8::MIN, 0);
                }
            "#,
        )
        .build();
    p.cargo("build -v --lib")
        .build_std_arg(&setup, "core")
        .target_host()
        .with_stderr_does_not_contain("[..]libstd[..]")
        .run();
}

#[cargo_test]
fn check_core() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file("src/lib.rs", "#![no_std] fn unused_fn() {}")
        .build();

    p.cargo("check -v")
        .build_std_arg(&setup, "core")
        .target_host()
        .with_stderr_contains("[WARNING] [..]unused_fn[..]`")
        .run();
}

#[cargo_test]
fn depend_same_as_std() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };

    let p = project()
        .file(
            "src/lib.rs",
            r#"
                pub fn f() {
                    registry_dep_using_core::non_sysroot_api();
                    registry_dep_using_alloc::non_sysroot_api();
                    registry_dep_using_std::non_sysroot_api();
                }
            "#,
        )
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.1.0"
                edition = "2018"

                [dependencies]
                registry-dep-using-core = "1.0"
                registry-dep-using-alloc = "1.0"
                registry-dep-using-std = "1.0"
            "#,
        )
        .build();

    p.cargo("build -v").build_std(&setup).target_host().run();
}

#[cargo_test]
fn test() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                #[cfg(test)]
                mod tests {
                    #[test]
                    fn it_works() {
                        assert_eq!(2 + 2, 4);
                    }
                }
            "#,
        )
        .build();

    p.cargo("test -v")
        .build_std(&setup)
        .target_host()
        .with_stdout_contains("test tests::it_works ... ok")
        .run();
}

#[cargo_test]
fn target_proc_macro() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                extern crate proc_macro;
                pub fn f() {
                    let _ts = proc_macro::TokenStream::new();
                }
            "#,
        )
        .build();

    p.cargo("build -v").build_std(&setup).target_host().run();
}

#[cargo_test]
fn bench() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                #![feature(test)]
                extern crate test;

                #[bench]
                fn b1(b: &mut test::Bencher) {
                    b.iter(|| ())
                }
            "#,
        )
        .build();

    p.cargo("bench -v").build_std(&setup).target_host().run();
}

#[cargo_test]
fn doc() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                /// Doc
                pub fn f() -> Result<(), ()> {Ok(())}
            "#,
        )
        .build();

    p.cargo("doc -v").build_std(&setup).target_host().run();
}

#[cargo_test]
fn check_std() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            "
                extern crate core;
                extern crate alloc;
                extern crate proc_macro;
                pub fn f() {}
            ",
        )
        .file("src/main.rs", "fn main() {}")
        .file(
            "tests/t1.rs",
            r#"
                #[test]
                fn t1() {
                    assert_eq!(1, 2);
                }
            "#,
        )
        .build();

    p.cargo("check -v --all-targets")
        .build_std(&setup)
        .target_host()
        .run();
    p.cargo("check -v --all-targets --profile=test")
        .build_std(&setup)
        .target_host()
        .run();
}

#[cargo_test]
fn doctest() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                /// Doc
                /// ```
                /// std::custom_api();
                /// ```
                pub fn f() {}
            "#,
        )
        .build();

    p.cargo("test --doc -v")
        .build_std(&setup)
        .with_stdout_contains("test src/lib.rs - f [..] ... ok")
        .target_host()
        .run();
}

#[cargo_test]
fn no_implicit_alloc() {
    // Demonstrate that alloc is not implicitly in scope.
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                pub fn f() {
                    let _: Vec<i32> = alloc::vec::Vec::new();
                }
            "#,
        )
        .build();

    p.cargo("build -v")
        .build_std(&setup)
        .target_host()
        .with_stderr_contains("[..]use of undeclared [..]`alloc`")
        .with_status(101)
        .run();
}

#[cargo_test]
fn macro_expanded_shadow() {
    // This tests a bug caused by the previous use of `--extern` to directly
    // load sysroot crates. This necessitated the switch to `--sysroot` to
    // retain existing behavior. See
    // https://github.com/rust-lang/wg-cargo-std-aware/issues/40 for more
    // detail.
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                macro_rules! a {
                    () => (extern crate std as alloc;)
                }
                a!();
            "#,
        )
        .build();

    p.cargo("build -v").build_std(&setup).target_host().run();
}
<commit_msg>Simplify mock rustc wrapper.<commit_after>use cargo_test_support::registry::{Dependency, Package};
use cargo_test_support::ProjectBuilder;
use cargo_test_support::{is_nightly, paths, project, rustc_host, Execs};
use std::path::PathBuf;

struct Setup {
    rustc_wrapper: PathBuf,
    real_sysroot: String,
}

fn setup() -> Option<Setup> {
    if !is_nightly() {
        // -Zbuild-std is nightly
        // We don't want these tests to run on rust-lang/rust.
        return None;
    }

    // Our mock sysroot requires a few packages from crates.io, so make sure
    // they're "published" to crates.io. Also edit their code a bit to make sure
    // that they have access to our custom crates with custom apis.
    Package::new("registry-dep-using-core", "1.0.0")
        .file(
            "src/lib.rs",
            "
                #![no_std]

                #[cfg(feature = \"mockbuild\")]
                pub fn custom_api() {
                }

                #[cfg(not(feature = \"mockbuild\"))]
                pub fn non_sysroot_api() {
                    core::custom_api();
                }
            ",
        )
        .add_dep(Dependency::new("rustc-std-workspace-core", "*").optional(true))
        .feature("mockbuild", &["rustc-std-workspace-core"])
        .publish();
    Package::new("registry-dep-using-alloc", "1.0.0")
        .file(
            "src/lib.rs",
            "
                #![no_std]

                extern crate alloc;

                #[cfg(feature = \"mockbuild\")]
                pub fn custom_api() {
                }

                #[cfg(not(feature = \"mockbuild\"))]
                pub fn non_sysroot_api() {
                    core::custom_api();
                    alloc::custom_api();
                }
            ",
        )
        .add_dep(Dependency::new("rustc-std-workspace-core", "*").optional(true))
        .add_dep(Dependency::new("rustc-std-workspace-alloc", "*").optional(true))
        .feature(
            "mockbuild",
            &["rustc-std-workspace-core", "rustc-std-workspace-alloc"],
        )
        .publish();
    Package::new("registry-dep-using-std", "1.0.0")
        .file(
            "src/lib.rs",
            "
                #[cfg(feature = \"mockbuild\")]
                pub fn custom_api() {
                }

                #[cfg(not(feature = \"mockbuild\"))]
                pub fn non_sysroot_api() {
                    std::custom_api();
                }
            ",
        )
        .add_dep(Dependency::new("rustc-std-workspace-std", "*").optional(true))
        .feature("mockbuild", &["rustc-std-workspace-std"])
        .publish();

    let p = ProjectBuilder::new(paths::root().join("rustc-wrapper"))
        .file(
            "src/main.rs",
            r#"
                use std::process::Command;
                use std::env;
                fn main() {
                    let mut args = env::args().skip(1).collect::<Vec<_>>();

                    let is_sysroot_crate = env::var_os("RUSTC_BOOTSTRAP").is_some();
                    if is_sysroot_crate {
                        let arg = args.iter().position(|a| a == "--sysroot").unwrap();
                        args[arg + 1] = env::var("REAL_SYSROOT").unwrap();
                    }

                    let ret = Command::new(&args[0]).args(&args[1..]).status().unwrap();
                    std::process::exit(ret.code().unwrap_or(1));
                }
            "#,
        )
        .build();
    p.cargo("build").run();

    return Some(Setup {
        rustc_wrapper: p.bin("foo"),
        real_sysroot: paths::sysroot(),
    });
}

fn enable_build_std(e: &mut Execs, setup: &Setup, arg: Option<&str>) {
    // First up, force Cargo to use our "mock sysroot" which mimics what
    // libstd looks like upstream.
    let root = paths::root();
    let root = root
        .parent() // chop off test name
        .unwrap()
        .parent() // chop off `citN`
        .unwrap()
        .parent() // chop off `target`
        .unwrap()
        .join("tests/testsuite/mock-std");
    e.env("__CARGO_TESTS_ONLY_SRC_ROOT", &root);

    // Actually enable `-Zbuild-std` for now
    let arg = match arg {
        Some(s) => format!("-Zbuild-std={}", s),
        None => "-Zbuild-std".to_string(),
    };
    e.arg(arg);
    e.masquerade_as_nightly_cargo();

    // We do various shenanigans to ensure our "mock sysroot" actually links
    // with the real sysroot, so we don't have to actually recompile std for
    // each test. Perform all that logic here, namely:
    //
    // * RUSTC_WRAPPER - uses our shim executable built above to control rustc
    // * REAL_SYSROOT - used by the shim executable to swap out to the real
    //   sysroot temporarily for some compilations
    // * RUST{,DOC}FLAGS - an extra `-L` argument to ensure we can always load
    //   crates from the sysroot, but only indirectly through other crates.
    e.env("RUSTC_WRAPPER", &setup.rustc_wrapper);
    e.env("REAL_SYSROOT", &setup.real_sysroot);
    let libdir = format!("/lib/rustlib/{}/lib", rustc_host());
    e.env(
        "RUSTFLAGS",
        format!("-Ldependency={}{}", setup.real_sysroot, libdir),
    );
    e.env(
        "RUSTDOCFLAGS",
        format!("-Ldependency={}{}", setup.real_sysroot, libdir),
    );
}

// Helper methods used in the tests below
trait BuildStd: Sized {
    fn build_std(&mut self, setup: &Setup) -> &mut Self;
    fn build_std_arg(&mut self, setup: &Setup, arg: &str) -> &mut Self;
    fn target_host(&mut self) -> &mut Self;
}

impl BuildStd for Execs {
    fn build_std(&mut self, setup: &Setup) -> &mut Self {
        enable_build_std(self, setup, None);
        self
    }

    fn build_std_arg(&mut self, setup: &Setup, arg: &str) -> &mut Self {
        enable_build_std(self, setup, Some(arg));
        self
    }

    fn target_host(&mut self) -> &mut Self {
        self.arg("--target").arg(rustc_host());
        self
    }
}

#[cargo_test]
fn basic() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };

    let p = project()
        .file(
            "src/main.rs",
            "
                fn main() {
                    std::custom_api();
                    foo::f();
                }

                #[test]
                fn smoke_bin_unit() {
                    std::custom_api();
                    foo::f();
                }
            ",
        )
        .file(
            "src/lib.rs",
            "
                extern crate alloc;
                extern crate proc_macro;

                /// ```
                /// foo::f();
                /// ```
                pub fn f() {
                    core::custom_api();
                    std::custom_api();
                    alloc::custom_api();
                    proc_macro::custom_api();
                }

                #[test]
                fn smoke_lib_unit() {
                    std::custom_api();
                    f();
                }
            ",
        )
        .file(
            "tests/smoke.rs",
            "
                #[test]
                fn smoke_integration() {
                    std::custom_api();
                    foo::f();
                }
            ",
        )
        .build();

    p.cargo("check").build_std(&setup).target_host().run();
    p.cargo("build").build_std(&setup).target_host().run();
    p.cargo("run").build_std(&setup).target_host().run();
    p.cargo("test").build_std(&setup).target_host().run();
}

#[cargo_test]
fn simple_lib_std() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project().file("src/lib.rs", "").build();
    p.cargo("build -v")
        .build_std(&setup)
        .target_host()
        .with_stderr_contains("[RUNNING] `[..]--crate-name std [..]`")
        .run();
    // Check freshness.
    p.change_file("src/lib.rs", " ");
    p.cargo("build -v")
        .build_std(&setup)
        .target_host()
        .with_stderr_contains("[FRESH] std[..]")
        .run();
}

#[cargo_test]
fn simple_bin_std() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project().file("src/main.rs", "fn main() {}").build();
    p.cargo("run -v").build_std(&setup).target_host().run();
}

#[cargo_test]
fn lib_nostd() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                #![no_std]
                pub fn foo() {
                    assert_eq!(core::u8::MIN, 0);
                }
            "#,
        )
        .build();
    p.cargo("build -v --lib")
        .build_std_arg(&setup, "core")
        .target_host()
        .with_stderr_does_not_contain("[..]libstd[..]")
        .run();
}

#[cargo_test]
fn check_core() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file("src/lib.rs", "#![no_std] fn unused_fn() {}")
        .build();

    p.cargo("check -v")
        .build_std_arg(&setup, "core")
        .target_host()
        .with_stderr_contains("[WARNING] [..]unused_fn[..]`")
        .run();
}

#[cargo_test]
fn depend_same_as_std() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };

    let p = project()
        .file(
            "src/lib.rs",
            r#"
                pub fn f() {
                    registry_dep_using_core::non_sysroot_api();
                    registry_dep_using_alloc::non_sysroot_api();
                    registry_dep_using_std::non_sysroot_api();
                }
            "#,
        )
        .file(
            "Cargo.toml",
            r#"
                [package]
                name = "foo"
                version = "0.1.0"
                edition = "2018"

                [dependencies]
                registry-dep-using-core = "1.0"
                registry-dep-using-alloc = "1.0"
                registry-dep-using-std = "1.0"
            "#,
        )
        .build();

    p.cargo("build -v").build_std(&setup).target_host().run();
}

#[cargo_test]
fn test() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                #[cfg(test)]
                mod tests {
                    #[test]
                    fn it_works() {
                        assert_eq!(2 + 2, 4);
                    }
                }
            "#,
        )
        .build();

    p.cargo("test -v")
        .build_std(&setup)
        .target_host()
        .with_stdout_contains("test tests::it_works ... ok")
        .run();
}

#[cargo_test]
fn target_proc_macro() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                extern crate proc_macro;
                pub fn f() {
                    let _ts = proc_macro::TokenStream::new();
                }
            "#,
        )
        .build();

    p.cargo("build -v").build_std(&setup).target_host().run();
}

#[cargo_test]
fn bench() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                #![feature(test)]
                extern crate test;

                #[bench]
                fn b1(b: &mut test::Bencher) {
                    b.iter(|| ())
                }
            "#,
        )
        .build();

    p.cargo("bench -v").build_std(&setup).target_host().run();
}

#[cargo_test]
fn doc() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                /// Doc
                pub fn f() -> Result<(), ()> {Ok(())}
            "#,
        )
        .build();

    p.cargo("doc -v").build_std(&setup).target_host().run();
}

#[cargo_test]
fn check_std() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            "
                extern crate core;
                extern crate alloc;
                extern crate proc_macro;
                pub fn f() {}
            ",
        )
        .file("src/main.rs", "fn main() {}")
        .file(
            "tests/t1.rs",
            r#"
                #[test]
                fn t1() {
                    assert_eq!(1, 2);
                }
            "#,
        )
        .build();

    p.cargo("check -v --all-targets")
        .build_std(&setup)
        .target_host()
        .run();
    p.cargo("check -v --all-targets --profile=test")
        .build_std(&setup)
        .target_host()
        .run();
}

#[cargo_test]
fn doctest() {
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                /// Doc
                /// ```
                /// std::custom_api();
                /// ```
                pub fn f() {}
            "#,
        )
        .build();

    p.cargo("test --doc -v")
        .build_std(&setup)
        .with_stdout_contains("test src/lib.rs - f [..] ... ok")
        .target_host()
        .run();
}

#[cargo_test]
fn no_implicit_alloc() {
    // Demonstrate that alloc is not implicitly in scope.
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                pub fn f() {
                    let _: Vec<i32> = alloc::vec::Vec::new();
                }
            "#,
        )
        .build();

    p.cargo("build -v")
        .build_std(&setup)
        .target_host()
        .with_stderr_contains("[..]use of undeclared [..]`alloc`")
        .with_status(101)
        .run();
}

#[cargo_test]
fn macro_expanded_shadow() {
    // This tests a bug caused by the previous use of `--extern` to directly
    // load sysroot crates. This necessitated the switch to `--sysroot` to
    // retain existing behavior. See
    // https://github.com/rust-lang/wg-cargo-std-aware/issues/40 for more
    // detail.
    let setup = match setup() {
        Some(s) => s,
        None => return,
    };
    let p = project()
        .file(
            "src/lib.rs",
            r#"
                macro_rules! a {
                    () => (extern crate std as alloc;)
                }
                a!();
            "#,
        )
        .build();

    p.cargo("build -v").build_std(&setup).target_host().run();
}
<|endoftext|>


<commit_before>//! Synchronized types that allow access to a `Histogram` from multiple threads.

use crate::errors::*;
use crate::{Counter, Histogram};
use std::borrow::Borrow;
use std::ops::{AddAssign, Deref, DerefMut};
use std::sync::{atomic, Arc, Mutex};
use std::time;

/// A write-only handle to a [`SyncHistogram`].
///
/// This handle allows you to record samples from multiple threads, each with its own `Recorder`,
/// concurrently. Writes to a `Recorder` are wait-free and scalable except for when the
/// [`SyncHistogram`] initiates a _phase shift_. During a phase shift, the next write on each
/// associated `Recorder` merges its results into a shared [`Histogram`] that is then made
/// available to the [`SyncHistogram`] once the phase shift completes.
///
/// An idle `Recorder` will hold up a phase shift indefinitely, or until it times out (is using
/// [`SyncHistogram::refresh_timeout`]. If a `Recorder` will remain idle for extended periods of
/// time, it should call [`Recorder::idle`], which will tell the reader not to wait for this
/// particular writer.
///
/// When a `Recorder` is dropped, all samples are made visible to the next
/// [`SyncHistogram::refresh`]. If you wish to ensure that all currently recorded samples are
/// communicated to the reader at a particular point in time, use [`Recorder::synchronize`].
#[derive(Debug)]
pub struct Recorder<C: Counter> {
    local: Histogram<C>,
    shared: Arc<Shared<C>>,
    last_phase: usize,
}

// make it more ergonomic to record samples
impl<C: Counter> AddAssign<u64> for Recorder<C> {
    fn add_assign(&mut self, value: u64) {
        self.record(value).unwrap();
    }
}

impl<C: Counter> Clone for Recorder<C> {
    fn clone(&self) -> Self {
        // reader will have to wait for one more recorder
        {
            let mut truth = self.shared.truth.lock().unwrap();
            truth.recorders += 1;
        }

        // new recorder starts at the same phase as we do with an empty histogram
        Recorder {
            local: Histogram::new_from(&self.local),
            shared: self.shared.clone(),
            last_phase: self.last_phase,
        }
    }
}

impl<C: Counter> Drop for Recorder<C> {
    fn drop(&mut self) {
        // we'll need to decrement the # of recorders
        let mut truth = self.shared.truth.lock().unwrap();
        truth.recorders -= 1;

        // we also want to communicate the remainder of our samples to the reader
        // we do this under the lock so that if the reader reads .recorders after we update it
        // above, it is guaranteed to see the samples from this recorder.
        // we also _have_ to do it at some point during drop as the reader _may_ have read
        // .recorders _before_ we decremented it above, in which case it's blocking on us!
        // note that we cannot call self.update() here as it would drop the mutex guard
        let h = Histogram::new_from(&self.local);
        let h = std::mem::replace(&mut self.local, h);
        let _ = self.shared.sender.send(h).is_ok(); // if this is err, the reader went away

        // explicitly drop guard to ensure we don't accidentally drop it above
        drop(truth);
    }
}

#[derive(Debug)]
struct Critical {
    recorders: usize,
}

#[derive(Debug)]
struct Shared<C: Counter> {
    truth: Mutex<Critical>,
    sender: crossbeam_channel::Sender<Histogram<C>>,
    phase: atomic::AtomicUsize,
}

/// This guard denotes that a [`Recorder`] is currently idle, and should not be waited on by a
/// [`SyncHistogram`] phase-shift.
pub struct IdleRecorderGuard<'a, C: Counter>(&'a mut Recorder<C>);

impl<'a, C: Counter> Drop for IdleRecorderGuard<'a, C> {
    fn drop(&mut self) {
        // the Recorder is no longer idle, so the reader has to wait for us again
        // this basically means re-incrementing .recorders
        let mut crit = self.0.shared.truth.lock().unwrap();
        crit.recorders += 1;

        // we need to figure out what phase we're joining
        // the easiest way to do that is to adopt the current phase
        //
        // note that we have to load the phase while holding the lock.
        // if we did not, the reader could come along, read our ++'d .recorders (and so wait for us
        // to send), and bump the phase, all before we read it, which would lead us to believe that
        // we were already synchronized when in reality we were not, which would stall the reader
        // even if we issued more writes.
        self.0.last_phase = self.0.shared.phase.load(atomic::Ordering::Acquire);

        // explicitly drop guard to ensure we don't accidentally drop it above
        drop(crit);
    }
}

impl<C: Counter> Recorder<C> {
    fn with_hist<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut Histogram<C>) -> R,
    {
        let r = f(&mut self.local);
        let phase = self.shared.phase.load(atomic::Ordering::Acquire);
        if phase != self.last_phase {
            self.update();
            self.last_phase = phase;
        }
        r
    }

    // return our current histogram and leave a cleared one in its place
    fn shed(&mut self) -> Histogram<C> {
        let h = Histogram::new_from(&self.local);
        std::mem::replace(&mut self.local, h)
    }

    fn update(&mut self) {
        let h = self.shed();
        let _ = self.shared.sender.send(h).is_ok(); // if this is err, the reader went away
    }

    /// Call this method if the Recorder will be idle for a while.
    ///
    /// Until the returned guard is dropped, the associated [`SyncHistogram`] will not wait for
    /// this recorder on a phase shift.
    pub fn idle(&mut self) -> IdleRecorderGuard<C> {
        let phase;
        {
            // we're leaving rotation, so we need to decrement .recorders
            let mut crit = self.shared.truth.lock().unwrap();
            crit.recorders -= 1;

            // make sure we don't hold up the current phase shift (if any)
            phase = self.shared.phase.load(atomic::Ordering::Acquire);
            if phase != self.last_phase {
                // can't call self.update() due to borrow of self.shared above
                let h = Histogram::new_from(&self.local);
                let h = std::mem::replace(&mut self.local, h);
                let _ = self.shared.sender.send(h).is_ok(); // if this is err, the reader went away
            }
        }
        self.last_phase = phase;

        IdleRecorderGuard(self)
    }

    /// See [`Histogram::add`].
    pub fn add<B: Borrow<Histogram<C>>>(&mut self, source: B) -> Result<(), AdditionError> {
        self.with_hist(move |h| h.add(source))
    }

    /// See [`Histogram::add_correct`].
    pub fn add_correct<B: Borrow<Histogram<C>>>(
        &mut self,
        source: B,
        interval: u64,
    ) -> Result<(), RecordError> {
        self.with_hist(move |h| h.add_correct(source, interval))
    }

    /// See [`Histogram::subtract`].
    pub fn subtract<B: Borrow<Histogram<C>>>(
        &mut self,
        subtrahend: B,
    ) -> Result<(), SubtractionError> {
        self.with_hist(move |h| h.subtract(subtrahend))
    }

    /// See [`Histogram::record`].
    pub fn record(&mut self, value: u64) -> Result<(), RecordError> {
        self.with_hist(move |h| h.record(value))
    }

    /// See [`Histogram::saturating_record`].
    pub fn saturating_record(&mut self, value: u64) {
        self.with_hist(move |h| h.saturating_record(value))
    }

    /// See [`Histogram::record_n`].
    pub fn record_n(&mut self, value: u64, count: C) -> Result<(), RecordError> {
        self.with_hist(move |h| h.record_n(value, count))
    }

    /// See [`Histogram::saturating_record_n`].
    pub fn saturating_record_n(&mut self, value: u64, count: C) {
        self.with_hist(move |h| h.saturating_record_n(value, count))
    }

    /// See [`Histogram::record_correct`].
    pub fn record_correct(&mut self, value: u64, interval: u64) -> Result<(), RecordError> {
        self.with_hist(move |h| h.record_correct(value, interval))
    }

    /// See [`Histogram::record_n_correct`].
    pub fn record_n_correct(
        &mut self,
        value: u64,
        count: C,
        interval: u64,
    ) -> Result<(), RecordError> {
        self.with_hist(move |h| h.record_n_correct(value, count, interval))
    }
}

/// A `Histogram` that can be written to by multiple threads concurrently.
///
/// Each writer thread should have a [`Recorder`], which allows it to record new samples without
/// synchronization. New recorded samples are made available through this histogram by calling
/// [`SyncHistogram::refresh`], which blocks until it has synchronized with every recorder.
pub struct SyncHistogram<C: Counter> {
    merged: Histogram<C>,
    shared: Arc<Shared<C>>,
    receiver: crossbeam_channel::Receiver<Histogram<C>>,
}

impl<C: Counter> SyncHistogram<C> {
    fn refresh_inner(&mut self, timeout: Option<time::Duration>) {
        let end = timeout.map(|dur| time::Instant::now() + dur);

        // time to start a phase change
        // we first want to drain any histograms left over by dropped recorders
        // note that we do this _before_ incrementing the phase, so we know they're "old"
        while let Ok(h) = self.receiver.try_recv() {
            self.merged
                .add(&h)
                .expect("TODO: failed to merge histogram");
        }

        // make sure no recorders can join or leave in the middle of this
        let recorders = self.shared.truth.lock().unwrap().recorders;

        // then, we tell writers to phase
        let _ = self.shared.phase.fetch_add(1, atomic::Ordering::AcqRel);

        // we want to wait for writers to all have phased
        let mut phased = 0;

        // at this point, we expect to get at least truth.recorders histograms
        while phased < recorders {
            let h = if let Some(end) = end {
                let now = time::Instant::now();
                if now > end {
                    break;
                }

                match self.receiver.recv_timeout(end - now) {
                    Ok(h) => h,
                    Err(crossbeam_channel::RecvTimeoutError::Timeout) => break,
                    Err(crossbeam_channel::RecvTimeoutError::Disconnected) => unreachable!(),
                }
            } else {
                self.receiver
                    .recv()
                    .expect("SyncHistogram has an Arc<Shared> with a Receiver")
            };

            self.merged
                .add(&h)
                .expect("TODO: failed to merge histogram");
            phased += 1;
        }

        // we also gobble up extra histograms we may have been sent from more dropped writers
        while let Ok(h) = self.receiver.try_recv() {
            self.merged
                .add(&h)
                .expect("TODO: failed to merge histogram");
        }
    }

    /// Block until writes from all [`Recorder`] instances for this histogram have been
    /// incorporated.
    pub fn refresh(&mut self) {
        self.refresh_inner(None)
    }

    /// Block until writes from all [`Recorder`] instances for this histogram have been
    /// incorporated, or until the given amount of time has passed.
    pub fn refresh_timeout(&mut self, timeout: time::Duration) {
        self.refresh_inner(Some(timeout))
    }

    /// Obtain another multi-threaded writer for this histogram.
    ///
    /// Note that writes made to the `Recorder` will not be visible until the next call to
    /// [`SyncHistogram::refresh`].
    pub fn recorder(&self) -> Recorder<C> {
        // we will have to wait for one more recorder
        {
            let mut truth = self.shared.truth.lock().unwrap();
            truth.recorders += 1;
        }

        // new recorder starts at the current phase with an empty histogram
        Recorder {
            local: Histogram::new_from(&self.merged),
            shared: self.shared.clone(),
            last_phase: self.shared.phase.load(atomic::Ordering::Acquire),
        }
    }
}

impl<C: Counter> From<Histogram<C>> for SyncHistogram<C> {
    fn from(h: Histogram<C>) -> Self {
        let (tx, rx) = crossbeam_channel::unbounded();
        SyncHistogram {
            merged: h,
            receiver: rx,
            shared: Arc::new(Shared {
                truth: Mutex::new(Critical { recorders: 0 }),
                sender: tx,
                phase: atomic::AtomicUsize::new(0),
            }),
        }
    }
}

impl<C: Counter> Deref for SyncHistogram<C> {
    type Target = Histogram<C>;
    fn deref(&self) -> &Self::Target {
        &self.merged
    }
}

impl<C: Counter> DerefMut for SyncHistogram<C> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.merged
    }
}
<commit_msg>Fix some errors in sync docs<commit_after>//! Synchronized types that allow access to a `Histogram` from multiple threads.

use crate::errors::*;
use crate::{Counter, Histogram};
use std::borrow::Borrow;
use std::ops::{AddAssign, Deref, DerefMut};
use std::sync::{atomic, Arc, Mutex};
use std::time;

/// A write-only handle to a [`SyncHistogram`].
///
/// This handle allows you to record samples from multiple threads, each with its own `Recorder`,
/// concurrently. Writes to a `Recorder` are wait-free and scalable except for when the
/// [`SyncHistogram`] initiates a _phase shift_. During a phase shift, the next write on each
/// associated `Recorder` merges its results into a shared [`Histogram`] that is then made
/// available to the [`SyncHistogram`] once the phase shift completes. Phase shifts should also be
/// relatively cheap for writers, as they mainly need to perform a channel send on an unbounded,
/// lock-free channel.
///
/// An idle `Recorder` will hold up a phase shift indefinitely, or until it times out (is using
/// [`SyncHistogram::refresh_timeout`]. If a `Recorder` will remain idle for extended periods of
/// time, it should call [`Recorder::idle`], which will tell the reader not to wait for this
/// particular writer.
///
/// When a `Recorder` is dropped, all samples are made visible to the next
/// [`SyncHistogram::refresh`].
#[derive(Debug)]
pub struct Recorder<C: Counter> {
    local: Histogram<C>,
    shared: Arc<Shared<C>>,
    last_phase: usize,
}

// make it more ergonomic to record samples
impl<C: Counter> AddAssign<u64> for Recorder<C> {
    fn add_assign(&mut self, value: u64) {
        self.record(value).unwrap();
    }
}

impl<C: Counter> Clone for Recorder<C> {
    fn clone(&self) -> Self {
        // reader will have to wait for one more recorder
        {
            let mut truth = self.shared.truth.lock().unwrap();
            truth.recorders += 1;
        }

        // new recorder starts at the same phase as we do with an empty histogram
        Recorder {
            local: Histogram::new_from(&self.local),
            shared: self.shared.clone(),
            last_phase: self.last_phase,
        }
    }
}

impl<C: Counter> Drop for Recorder<C> {
    fn drop(&mut self) {
        // we'll need to decrement the # of recorders
        let mut truth = self.shared.truth.lock().unwrap();
        truth.recorders -= 1;

        // we also want to communicate the remainder of our samples to the reader
        // we do this under the lock so that if the reader reads .recorders after we update it
        // above, it is guaranteed to see the samples from this recorder.
        // we also _have_ to do it at some point during drop as the reader _may_ have read
        // .recorders _before_ we decremented it above, in which case it's blocking on us!
        // note that we cannot call self.update() here as it would drop the mutex guard
        let h = Histogram::new_from(&self.local);
        let h = std::mem::replace(&mut self.local, h);
        let _ = self.shared.sender.send(h).is_ok(); // if this is err, the reader went away

        // explicitly drop guard to ensure we don't accidentally drop it above
        drop(truth);
    }
}

#[derive(Debug)]
struct Critical {
    recorders: usize,
}

#[derive(Debug)]
struct Shared<C: Counter> {
    truth: Mutex<Critical>,
    sender: crossbeam_channel::Sender<Histogram<C>>,
    phase: atomic::AtomicUsize,
}

/// This guard denotes that a [`Recorder`] is currently idle, and should not be waited on by a
/// [`SyncHistogram`] phase-shift.
pub struct IdleRecorderGuard<'a, C: Counter>(&'a mut Recorder<C>);

impl<'a, C: Counter> Drop for IdleRecorderGuard<'a, C> {
    fn drop(&mut self) {
        // the Recorder is no longer idle, so the reader has to wait for us again
        // this basically means re-incrementing .recorders
        let mut crit = self.0.shared.truth.lock().unwrap();
        crit.recorders += 1;

        // we need to figure out what phase we're joining
        // the easiest way to do that is to adopt the current phase
        //
        // note that we have to load the phase while holding the lock.
        // if we did not, the reader could come along, read our ++'d .recorders (and so wait for us
        // to send), and bump the phase, all before we read it, which would lead us to believe that
        // we were already synchronized when in reality we were not, which would stall the reader
        // even if we issued more writes.
        self.0.last_phase = self.0.shared.phase.load(atomic::Ordering::Acquire);

        // explicitly drop guard to ensure we don't accidentally drop it above
        drop(crit);
    }
}

impl<C: Counter> Recorder<C> {
    fn with_hist<F, R>(&mut self, f: F) -> R
    where
        F: FnOnce(&mut Histogram<C>) -> R,
    {
        let r = f(&mut self.local);
        let phase = self.shared.phase.load(atomic::Ordering::Acquire);
        if phase != self.last_phase {
            self.update();
            self.last_phase = phase;
        }
        r
    }

    // return our current histogram and leave a cleared one in its place
    fn shed(&mut self) -> Histogram<C> {
        let h = Histogram::new_from(&self.local);
        std::mem::replace(&mut self.local, h)
    }

    fn update(&mut self) {
        let h = self.shed();
        let _ = self.shared.sender.send(h).is_ok(); // if this is err, the reader went away
    }

    /// Call this method if the Recorder will be idle for a while.
    ///
    /// Until the returned guard is dropped, the associated [`SyncHistogram`] will not wait for
    /// this recorder on a phase shift.
    pub fn idle(&mut self) -> IdleRecorderGuard<C> {
        let phase;
        {
            // we're leaving rotation, so we need to decrement .recorders
            let mut crit = self.shared.truth.lock().unwrap();
            crit.recorders -= 1;

            // make sure we don't hold up the current phase shift (if any)
            phase = self.shared.phase.load(atomic::Ordering::Acquire);
            if phase != self.last_phase {
                // can't call self.update() due to borrow of self.shared above
                let h = Histogram::new_from(&self.local);
                let h = std::mem::replace(&mut self.local, h);
                let _ = self.shared.sender.send(h).is_ok(); // if this is err, the reader went away
            }
        }
        self.last_phase = phase;

        IdleRecorderGuard(self)
    }

    /// See [`Histogram::add`].
    pub fn add<B: Borrow<Histogram<C>>>(&mut self, source: B) -> Result<(), AdditionError> {
        self.with_hist(move |h| h.add(source))
    }

    /// See [`Histogram::add_correct`].
    pub fn add_correct<B: Borrow<Histogram<C>>>(
        &mut self,
        source: B,
        interval: u64,
    ) -> Result<(), RecordError> {
        self.with_hist(move |h| h.add_correct(source, interval))
    }

    /// See [`Histogram::subtract`].
    pub fn subtract<B: Borrow<Histogram<C>>>(
        &mut self,
        subtrahend: B,
    ) -> Result<(), SubtractionError> {
        self.with_hist(move |h| h.subtract(subtrahend))
    }

    /// See [`Histogram::record`].
    pub fn record(&mut self, value: u64) -> Result<(), RecordError> {
        self.with_hist(move |h| h.record(value))
    }

    /// See [`Histogram::saturating_record`].
    pub fn saturating_record(&mut self, value: u64) {
        self.with_hist(move |h| h.saturating_record(value))
    }

    /// See [`Histogram::record_n`].
    pub fn record_n(&mut self, value: u64, count: C) -> Result<(), RecordError> {
        self.with_hist(move |h| h.record_n(value, count))
    }

    /// See [`Histogram::saturating_record_n`].
    pub fn saturating_record_n(&mut self, value: u64, count: C) {
        self.with_hist(move |h| h.saturating_record_n(value, count))
    }

    /// See [`Histogram::record_correct`].
    pub fn record_correct(&mut self, value: u64, interval: u64) -> Result<(), RecordError> {
        self.with_hist(move |h| h.record_correct(value, interval))
    }

    /// See [`Histogram::record_n_correct`].
    pub fn record_n_correct(
        &mut self,
        value: u64,
        count: C,
        interval: u64,
    ) -> Result<(), RecordError> {
        self.with_hist(move |h| h.record_n_correct(value, count, interval))
    }
}

/// A `Histogram` that can be written to by multiple threads concurrently.
///
/// Each writer thread should have a [`Recorder`], which allows it to record new samples without
/// synchronization. New recorded samples are made available through this histogram by calling
/// [`SyncHistogram::refresh`], which blocks until it has synchronized with every recorder.
pub struct SyncHistogram<C: Counter> {
    merged: Histogram<C>,
    shared: Arc<Shared<C>>,
    receiver: crossbeam_channel::Receiver<Histogram<C>>,
}

impl<C: Counter> SyncHistogram<C> {
    fn refresh_inner(&mut self, timeout: Option<time::Duration>) {
        let end = timeout.map(|dur| time::Instant::now() + dur);

        // time to start a phase change
        // we first want to drain any histograms left over by dropped recorders
        // note that we do this _before_ incrementing the phase, so we know they're "old"
        while let Ok(h) = self.receiver.try_recv() {
            self.merged
                .add(&h)
                .expect("TODO: failed to merge histogram");
        }

        // make sure no recorders can join or leave in the middle of this
        let recorders = self.shared.truth.lock().unwrap().recorders;

        // then, we tell writers to phase
        let _ = self.shared.phase.fetch_add(1, atomic::Ordering::AcqRel);

        // we want to wait for writers to all have phased
        let mut phased = 0;

        // at this point, we expect to get at least truth.recorders histograms
        while phased < recorders {
            let h = if let Some(end) = end {
                let now = time::Instant::now();
                if now > end {
                    break;
                }

                match self.receiver.recv_timeout(end - now) {
                    Ok(h) => h,
                    Err(crossbeam_channel::RecvTimeoutError::Timeout) => break,
                    Err(crossbeam_channel::RecvTimeoutError::Disconnected) => unreachable!(),
                }
            } else {
                self.receiver
                    .recv()
                    .expect("SyncHistogram has an Arc<Shared> with a Receiver")
            };

            self.merged
                .add(&h)
                .expect("TODO: failed to merge histogram");
            phased += 1;
        }

        // we also gobble up extra histograms we may have been sent from more dropped writers
        while let Ok(h) = self.receiver.try_recv() {
            self.merged
                .add(&h)
                .expect("TODO: failed to merge histogram");
        }
    }

    /// Block until writes from all [`Recorder`] instances for this histogram have been
    /// incorporated.
    pub fn refresh(&mut self) {
        self.refresh_inner(None)
    }

    /// Block until writes from all [`Recorder`] instances for this histogram have been
    /// incorporated, or until the given amount of time has passed.
    pub fn refresh_timeout(&mut self, timeout: time::Duration) {
        self.refresh_inner(Some(timeout))
    }

    /// Obtain another multi-threaded writer for this histogram.
    ///
    /// Note that writes made to the `Recorder` will not be visible until the next call to
    /// [`SyncHistogram::refresh`].
    pub fn recorder(&self) -> Recorder<C> {
        // we will have to wait for one more recorder
        {
            let mut truth = self.shared.truth.lock().unwrap();
            truth.recorders += 1;
        }

        // new recorder starts at the current phase with an empty histogram
        Recorder {
            local: Histogram::new_from(&self.merged),
            shared: self.shared.clone(),
            last_phase: self.shared.phase.load(atomic::Ordering::Acquire),
        }
    }
}

impl<C: Counter> From<Histogram<C>> for SyncHistogram<C> {
    fn from(h: Histogram<C>) -> Self {
        let (tx, rx) = crossbeam_channel::unbounded();
        SyncHistogram {
            merged: h,
            receiver: rx,
            shared: Arc::new(Shared {
                truth: Mutex::new(Critical { recorders: 0 }),
                sender: tx,
                phase: atomic::AtomicUsize::new(0),
            }),
        }
    }
}

impl<C: Counter> Deref for SyncHistogram<C> {
    type Target = Histogram<C>;
    fn deref(&self) -> &Self::Target {
        &self.merged
    }
}

impl<C: Counter> DerefMut for SyncHistogram<C> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.merged
    }
}
<|endoftext|>