<commit_before>use error::*;
use super::Command;

use constant::*;
use lib::io::*;
use lib::setting::*;
use std::collections::BTreeSet;
use std::path::{MAIN_SEPARATOR, Path};

pub struct Ignore {
    pub command: Option<String>,
    pub paths  : Vec<String>,
}

impl Command for Ignore {
    fn validation(&self) -> bool { true }

    fn usage(&self) -> Usage {
        return Usage::new(UsageKind::Ignore);
    }

    fn main(&self) -> Result<(), CliError> {
        match self.command {
            Some(ref c) => match c.as_ref() {
                "add"     => add(self),
                "remove"  => remove(self),
                "current" => ignore_current_files(),
                "clear"   => clear_ignore_file(),
                _         => Err(From::from(self.usage())),
            },
            None => Err(From::from(self.usage())),
        }
    }
}

fn append_prefix_if_need(path: &String) -> String {
    let prefix       = format!(".{}", MAIN_SEPARATOR);
    let prefix: &str = prefix.as_ref();
    format!("{}{}", if path.starts_with(prefix) { "" } else { prefix }, path)
}

fn add(ignore: &Ignore) -> Result<(), CliError> {
    let paths = &ignore.paths;

    if paths.len() == 0 {
        return Err(From::from(ignore.usage()));
    }

    print_with_tag(0, Tag::Execution, format!("Confirm files to be added to \"{}\"", IGNORE_FILE_NAME));

    let ignores_to_be_added = paths
        .iter()
        .map(append_prefix_if_need)
        .filter(|p| Path::new(p).is_file())
        .collect::<BTreeSet<String>>();

    for file in &ignores_to_be_added {
        print_with_tag(1, Tag::Info, format!("\"{}\" will be ignored", file));
    }

    let original_ignores = try!(read_ignore_file());

    let new_ignores = original_ignores
        .union(&ignores_to_be_added)
        .fold(String::new(), |c, ref f| c + f + "\n");

    print_with_tag(0, Tag::Execution, format!("Recreate \"{}\" file", IGNORE_FILE_NAME));

    try!(create_ignore_file(new_ignores));

    Ok(())
}

fn remove(ignore: &Ignore) -> Result<(), CliError> {
    let paths = &ignore.paths;

    if paths.len() == 0 {
        return Err(From::from(ignore.usage()));
    }

    let ignores_to_be_removed = paths
        .iter()
        .map(append_prefix_if_need)
        .collect::<BTreeSet<String>>();

    for file in &ignores_to_be_removed {
        print_with_tag(1, Tag::Info, format!("\"{}\" will not be ignored", file));
    }

    let original_ignores = try!(read_ignore_file());

    let new_ignores = original_ignores
        .difference(&ignores_to_be_removed)
        .fold(String::new(), |c, ref f| c + f + "\n");

    print_with_tag(0, Tag::Execution, format!("Recreate \"{}\" file", IGNORE_FILE_NAME));

    try!(create_ignore_file(new_ignores));

    Ok(())
}

fn ignore_current_files() -> Result<(), CliError> {
    println!("do current command");

    Ok(())
}

fn clear_ignore_file() -> Result<(), CliError> {
    println!("do clear command");

    Ok(())
}
<commit_msg>[Refactoring] implement "Ignore"<commit_after>use error::*;
use super::Command;

use constant::*;
use lib::io::*;
use lib::setting::*;
use std::collections::BTreeSet;
use std::path::{MAIN_SEPARATOR, Path};

pub struct Ignore {
    pub command: Option<String>,
    pub paths  : Vec<String>,
}

impl Command for Ignore {
    fn validation(&self) -> bool { true }

    fn usage(&self) -> Usage {
        return Usage::new(UsageKind::Ignore);
    }

    fn main(&self) -> Result<(), CliError> {
        match self.command {
            Some(ref c) => match c.as_ref() {
                "add"     => self.add(),
                "remove"  => self.remove(),
                "current" => self.ignore_current_files(),
                "clear"   => self.clear_ignore_file(),
                _         => Err(From::from(self.usage())),
            },
            None => Err(From::from(self.usage())),
        }
    }
}

fn append_prefix_if_need(path: &String) -> String {
    let prefix       = format!(".{}", MAIN_SEPARATOR);
    let prefix: &str = prefix.as_ref();
    format!("{}{}", if path.starts_with(prefix) { "" } else { prefix }, path)
}

impl Ignore {
    fn add(&self) -> Result<(), CliError> {
        let paths = &self.paths;

        if paths.len() == 0 {
            return Err(From::from(self.usage()));
        }

        print_with_tag(0, Tag::Execution, format!("Confirm files to be added to \"{}\"", IGNORE_FILE_NAME));

        let ignores_to_be_added = paths
            .iter()
            .map(append_prefix_if_need)
            .filter(|p| Path::new(p).is_file())
            .collect::<BTreeSet<String>>();

        for file in &ignores_to_be_added {
            print_with_tag(1, Tag::Info, format!("\"{}\" will be ignored", file));
        }

        let original_ignores = try!(read_ignore_file());

        let new_ignores = original_ignores
            .union(&ignores_to_be_added)
            .fold(String::new(), |c, ref f| c + f + "\n");

        print_with_tag(0, Tag::Execution, format!("Recreate \"{}\" file", IGNORE_FILE_NAME));

        try!(create_ignore_file(new_ignores));

        Ok(())
    }

    fn remove(&self) -> Result<(), CliError> {
        let paths = &self.paths;

        if paths.len() == 0 {
            return Err(From::from(self.usage()));
        }

        let ignores_to_be_removed = paths
            .iter()
            .map(append_prefix_if_need)
            .collect::<BTreeSet<String>>();

        for file in &ignores_to_be_removed {
            print_with_tag(1, Tag::Info, format!("\"{}\" will not be ignored", file));
        }

        let original_ignores = try!(read_ignore_file());

        let new_ignores = original_ignores
            .difference(&ignores_to_be_removed)
            .fold(String::new(), |c, ref f| c + f + "\n");

        print_with_tag(0, Tag::Execution, format!("Recreate \"{}\" file", IGNORE_FILE_NAME));

        try!(create_ignore_file(new_ignores));

        Ok(())
    }

    fn ignore_current_files(&self) -> Result<(), CliError> {
        println!("do current command");

        Ok(())
    }

    fn clear_ignore_file(&self) -> Result<(), CliError> {
        println!("do clear command");

        Ok(())
    }
}
<|endoftext|>


<commit_before>use crate::ssz_encode;
use hashing::hash;
use int_to_bytes::int_to_bytes32;

const BYTES_PER_CHUNK: usize = 32;
const HASHSIZE: usize = 32;
const MERKLE_HASH_CHUNCK: usize = 2 * BYTES_PER_CHUNK;

pub struct TreeHashCache<'a> {
    chunk_offset: usize,
    cache: &'a mut [u8],
    chunk_modified: &'a mut [bool],
}

impl<'a> TreeHashCache<'a> {
    pub fn increment(&mut self) {
        self.chunk_offset += 1
    }

    pub fn modify_current_chunk(&mut self, to: &[u8]) -> Option<()> {
        self.modify_chunk(0, to)
    }

    pub fn modify_chunk(&mut self, chunk: usize, to: &[u8]) -> Option<()> {
        let start = chunk * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK;
        self.cache.get_mut(start..end)?.copy_from_slice(to);

        self.chunk_modified[chunk] = true;

        Some(())
    }

    pub fn changed(&self, chunk: usize) -> Option<bool> {
        self.chunk_modified.get(chunk).cloned()
    }

    pub fn children_modified(&self, parent_chunk: usize) -> Option<bool> {
        let children = children(parent_chunk);

        Some(self.changed(children.0)? | self.changed(children.1)?)
    }

    pub fn hash_children(&self, parent_chunk: usize) -> Option<Vec<u8>> {
        let children = children(parent_chunk);

        let start = children.0 * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK * 2;

        Some(hash(&self.cache.get(start..end)?))
    }

    pub fn just_the_leaves(&mut self, leaves: usize) -> Option<TreeHashCache> {
        let nodes = num_nodes(leaves);
        let internal = nodes - leaves;
        let leaves_start = (self.chunk_offset + internal) * HASHSIZE;

        Some(TreeHashCache {
            chunk_offset: self.chunk_offset + internal,
            cache: self.cache.get_mut(leaves_start..leaves * HASHSIZE)?,
            chunk_modified: self
                .chunk_modified
                .get_mut(self.chunk_offset..self.chunk_offset + leaves)?,
        })
    }
}

fn children(parent: usize) -> (usize, usize) {
    ((2 * parent + 1), (2 * parent + 2))
}

fn num_nodes(num_leaves: usize) -> usize {
    2 * num_leaves - 1
}

pub trait CachedTreeHash {
    fn build_cache_bytes(&self) -> Vec<u8>;

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()>;
}

impl CachedTreeHash for u64 {
    fn build_cache_bytes(&self) -> Vec<u8> {
        merkleize(&ssz_encode(self))
    }

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> {
        if self != other {
            cache.modify_current_chunk(&hash(&ssz_encode(self)));
        }

        cache.increment();

        Some(())
    }
}

pub struct Inner {
    pub a: u64,
    pub b: u64,
    pub c: u64,
    pub d: u64,
}

impl CachedTreeHash for Inner {
    fn build_cache_bytes(&self) -> Vec<u8> {
        let mut leaves = vec![];

        leaves.append(&mut self.a.build_cache_bytes());
        leaves.append(&mut self.b.build_cache_bytes());
        leaves.append(&mut self.c.build_cache_bytes());
        leaves.append(&mut self.d.build_cache_bytes());

        merkleize(&leaves)
    }

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> {
        let num_leaves = 4;

        let mut leaf_cache = cache.just_the_leaves(num_leaves)?;
        self.a.cached_hash_tree_root(&other.a, &mut leaf_cache)?;
        self.b.cached_hash_tree_root(&other.b, &mut leaf_cache)?;
        self.c.cached_hash_tree_root(&other.c, &mut leaf_cache)?;
        self.d.cached_hash_tree_root(&other.d, &mut leaf_cache)?;

        let nodes = num_nodes(num_leaves);
        let internal_chunks = nodes - num_leaves;

        for chunk in 0..internal_chunks {
            if cache.children_modified(chunk)? {
                cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?;
            }
        }

        Some(())
    }
}

/// A reference function to test against.
pub fn merkleize(values: &[u8]) -> Vec<u8> {
    let leaves = values.len() / HASHSIZE;

    if leaves == 0 || !leaves.is_power_of_two() {
        panic!("Handle bad leaf count");
    }

    let mut o: Vec<u8> = vec![0; (num_nodes(leaves) - leaves) * HASHSIZE];
    o.append(&mut values.to_vec());

    let mut i = o.len();
    let mut j = o.len() - values.len();

    while i >= MERKLE_HASH_CHUNCK {
        i -= MERKLE_HASH_CHUNCK;
        let hash = hash(&o[i..i + MERKLE_HASH_CHUNCK]);

        j -= HASHSIZE;
        o[j..j + HASHSIZE].copy_from_slice(&hash);
    }

    o
}

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

    fn join(many: Vec<&[u8]>) -> Vec<u8> {
        let mut all = vec![];
        for one in many {
            all.extend_from_slice(&mut one.clone())
        }
        all
    }

    /*
    #[test]
    fn container() {
        let data1 = hash(&vec![1; 32]);
        let data2 = hash(&vec![2; 32]);
        let data3 = hash(&vec![3; 32]);
        let data4 = hash(&vec![4; 32]);

        let data = join(vec![&data1, &data2, &data3, &data4]);

        let cache = cache_builder(&data).unwrap();
    }
    */
    #[test]
    fn merkleize_4_leaves() {
        let data1 = hash(&int_to_bytes32(1));
        let data2 = hash(&int_to_bytes32(2));
        let data3 = hash(&int_to_bytes32(3));
        let data4 = hash(&int_to_bytes32(4));

        let data = join(vec![&data1, &data2, &data3, &data4]);

        let cache = merkleize(&data);

        let hash_12 = {
            let mut joined = vec![];
            joined.append(&mut data1.clone());
            joined.append(&mut data2.clone());
            hash(&joined)
        };
        let hash_34 = {
            let mut joined = vec![];
            joined.append(&mut data3.clone());
            joined.append(&mut data4.clone());
            hash(&joined)
        };
        let hash_hash12_hash_34 = {
            let mut joined = vec![];
            joined.append(&mut hash_12.clone());
            joined.append(&mut hash_34.clone());
            hash(&joined)
        };

        for (i, chunk) in cache.chunks(HASHSIZE).enumerate().rev() {
            let expected = match i {
                0 => hash_hash12_hash_34.clone(),
                1 => hash_12.clone(),
                2 => hash_34.clone(),
                3 => data1.clone(),
                4 => data2.clone(),
                5 => data3.clone(),
                6 => data4.clone(),
                _ => vec![],
            };

            assert_eq!(chunk, &expected[..], "failed at {}", i);
        }
    }
}
/*

/// Get merkle root of some hashed values - the input leaf nodes is expected to already be hashed
/// Outputs a `Vec<u8>` byte array of the merkle root given a set of leaf node values.
pub fn cache_builder(values: &[u8]) -> Option<Vec<u8>> {
    let leaves = values.len() / HASHSIZE;

    if leaves == 0 || !leaves.is_power_of_two() {
        return None;
    }

    let mut o: Vec<u8> = vec![0; (num_nodes(leaves) - leaves) * HASHSIZE];
    o.append(&mut values.to_vec());

    let mut i = o.len();
    let mut j = o.len() - values.len();

    while i >= MERKLE_HASH_CHUNCK {
        i -= MERKLE_HASH_CHUNCK;
        let hash = hash(&o[i..i + MERKLE_HASH_CHUNCK]);

        j -= HASHSIZE;
        o.get_mut(j..j + HASHSIZE)?.copy_from_slice(&hash);
    }

    return Some(o);
}

fn parent(child: usize) -> usize {
    (child - 1) / 2
}

fn children(parent: usize) -> (usize, usize) {
    ((2 * parent + 1), (2 * parent + 2))
}

fn num_nodes(num_leaves: usize) -> usize {
    2 * num_leaves - 1
}

pub struct Outer {
    pub a: u64,
    pub b: u64,
    pub inner: Inner,
}

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

    fn join(many: Vec<&[u8]>) -> Vec<u8> {
        let mut all = vec![];
        for one in many {
            all.extend_from_slice(&mut one.clone())
        }
        all
    }

    /*
    #[test]
    fn container() {
        let data1 = hash(&vec![1; 32]);
        let data2 = hash(&vec![2; 32]);
        let data3 = hash(&vec![3; 32]);
        let data4 = hash(&vec![4; 32]);

        let data = join(vec![&data1, &data2, &data3, &data4]);

        let cache = cache_builder(&data).unwrap();
    }
    */
#[test]
fn can_build_cache() {
let data1 = hash(&vec![1; 32]);
let data2 = hash(&vec![2; 32]);
let data3 = hash(&vec![3; 32]);
let data4 = hash(&vec![4; 32]);

let data = join(vec![&data1, &data2, &data3, &data4]);

let cache = cache_builder(&data).unwrap();

let hash_12 = {
let mut joined = vec![];
joined.append(&mut data1.clone());
joined.append(&mut data2.clone());
hash(&joined)
};
let hash_34 = {
let mut joined = vec![];
joined.append(&mut data3.clone());
joined.append(&mut data4.clone());
hash(&joined)
};
let hash_hash12_hash_34 = {
let mut joined = vec![];
joined.append(&mut hash_12.clone());
joined.append(&mut hash_34.clone());
hash(&joined)
};

for (i, chunk) in cache.chunks(HASHSIZE).enumerate().rev() {
let expected = match i {
0 => hash_hash12_hash_34.clone(),
1 => hash_12.clone(),
2 => hash_34.clone(),
3 => data1.clone(),
4 => data2.clone(),
5 => data3.clone(),
6 => data4.clone(),
_ => vec![],
};

assert_eq!(chunk, &expected[..], "failed at {}", i);
}
}
}

/*
pub trait TreeHash {
    fn hash_tree_root(&self) -> Vec<u8>;
}

/// Returns a 32 byte hash of 'list' - a vector of byte vectors.
/// Note that this will consume 'list'.
pub fn merkle_hash(list: &mut Vec<Vec<u8>>) -> Vec<u8> {
    // flatten list
    let mut chunkz = list_to_blob(list);

    // get data_len as bytes. It will hashed will the merkle root
    let mut datalen = list.len().to_le_bytes().to_vec();
    zpad(&mut datalen, 32);

    // merklelize
    while chunkz.len() > HASHSIZE {
        let mut new_chunkz: Vec<u8> = Vec::new();

        for two_chunks in chunkz.chunks(BYTES_PER_CHUNK * 2) {
            // Hash two chuncks together
            new_chunkz.append(&mut hash(two_chunks));
        }

        chunkz = new_chunkz;
    }

    chunkz.append(&mut datalen);
    hash(&chunkz)
}

fn list_to_blob(list: &mut Vec<Vec<u8>>) -> Vec<u8> {
    // pack - fit as many many items per chunk as we can and then
    // right pad to BYTES_PER_CHUNCK
    let (items_per_chunk, chunk_count) = if list.is_empty() {
        (1, 1)
    } else {
        let items_per_chunk = BYTES_PER_CHUNK / list[0].len();
        let chunk_count = list.len() / items_per_chunk;
        (items_per_chunk, chunk_count)
    };

    let mut chunkz = Vec::new();
    if list.is_empty() {
        // handle and empty list
        chunkz.append(&mut vec![0; BYTES_PER_CHUNK * 2]);
    } else if list[0].len() <= BYTES_PER_CHUNK {
        // just create a blob here; we'll divide into
        // chunked slices when we merklize
        let mut chunk = Vec::with_capacity(BYTES_PER_CHUNK);
        let mut item_count_in_chunk = 0;
        chunkz.reserve(chunk_count * BYTES_PER_CHUNK);
        for item in list.iter_mut() {
            item_count_in_chunk += 1;
            chunk.append(item);

            // completed chunk?
            if item_count_in_chunk == items_per_chunk {
                zpad(&mut chunk, BYTES_PER_CHUNK);
                chunkz.append(&mut chunk);
                item_count_in_chunk = 0;
            }
        }

        // left-over uncompleted chunk?
        if item_count_in_chunk != 0 {
            zpad(&mut chunk, BYTES_PER_CHUNK);
            chunkz.append(&mut chunk);
        }
    }

    // extend the number of chunks to a power of two if necessary
    if !chunk_count.is_power_of_two() {
        let zero_chunks_count = chunk_count.next_power_of_two() - chunk_count;
        chunkz.append(&mut vec![0; zero_chunks_count * BYTES_PER_CHUNK]);
    }

    chunkz
}

/// right pads with zeros making 'bytes' 'size' in length
fn zpad(bytes: &mut Vec<u8>, size: usize) {
    if bytes.len() < size {
        bytes.resize(size, 0);
    }
}

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

    #[test]
    fn test_merkle_hash() {
        let data1 = vec![1; 32];
        let data2 = vec![2; 32];
        let data3 = vec![3; 32];
        let mut list = vec![data1, data2, data3];
        let result = merkle_hash(&mut list);

        //note: should test againt a known test hash value
        assert_eq!(HASHSIZE, result.len());
    }
}
*/
*/
<commit_msg>Remove unused code<commit_after>use crate::ssz_encode;
use hashing::hash;
use int_to_bytes::int_to_bytes32;

const BYTES_PER_CHUNK: usize = 32;
const HASHSIZE: usize = 32;
const MERKLE_HASH_CHUNCK: usize = 2 * BYTES_PER_CHUNK;

pub struct TreeHashCache<'a> {
    chunk_offset: usize,
    cache: &'a mut [u8],
    chunk_modified: &'a mut [bool],
}

impl<'a> TreeHashCache<'a> {
    pub fn increment(&mut self) {
        self.chunk_offset += 1
    }

    pub fn modify_current_chunk(&mut self, to: &[u8]) -> Option<()> {
        self.modify_chunk(0, to)
    }

    pub fn modify_chunk(&mut self, chunk: usize, to: &[u8]) -> Option<()> {
        let start = chunk * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK;
        self.cache.get_mut(start..end)?.copy_from_slice(to);

        self.chunk_modified[chunk] = true;

        Some(())
    }

    pub fn changed(&self, chunk: usize) -> Option<bool> {
        self.chunk_modified.get(chunk).cloned()
    }

    pub fn children_modified(&self, parent_chunk: usize) -> Option<bool> {
        let children = children(parent_chunk);

        Some(self.changed(children.0)? | self.changed(children.1)?)
    }

    pub fn hash_children(&self, parent_chunk: usize) -> Option<Vec<u8>> {
        let children = children(parent_chunk);

        let start = children.0 * BYTES_PER_CHUNK;
        let end = start + BYTES_PER_CHUNK * 2;

        Some(hash(&self.cache.get(start..end)?))
    }

    pub fn just_the_leaves(&mut self, leaves: usize) -> Option<TreeHashCache> {
        let nodes = num_nodes(leaves);
        let internal = nodes - leaves;
        let leaves_start = (self.chunk_offset + internal) * HASHSIZE;

        Some(TreeHashCache {
            chunk_offset: self.chunk_offset + internal,
            cache: self.cache.get_mut(leaves_start..leaves * HASHSIZE)?,
            chunk_modified: self
                .chunk_modified
                .get_mut(self.chunk_offset..self.chunk_offset + leaves)?,
        })
    }
}

fn children(parent: usize) -> (usize, usize) {
    ((2 * parent + 1), (2 * parent + 2))
}

fn num_nodes(num_leaves: usize) -> usize {
    2 * num_leaves - 1
}

pub trait CachedTreeHash {
    fn build_cache_bytes(&self) -> Vec<u8>;

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()>;
}

impl CachedTreeHash for u64 {
    fn build_cache_bytes(&self) -> Vec<u8> {
        merkleize(&ssz_encode(self))
    }

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> {
        if self != other {
            cache.modify_current_chunk(&hash(&ssz_encode(self)));
        }

        cache.increment();

        Some(())
    }
}

pub struct Inner {
    pub a: u64,
    pub b: u64,
    pub c: u64,
    pub d: u64,
}

impl CachedTreeHash for Inner {
    fn build_cache_bytes(&self) -> Vec<u8> {
        let mut leaves = vec![];

        leaves.append(&mut self.a.build_cache_bytes());
        leaves.append(&mut self.b.build_cache_bytes());
        leaves.append(&mut self.c.build_cache_bytes());
        leaves.append(&mut self.d.build_cache_bytes());

        merkleize(&leaves)
    }

    fn cached_hash_tree_root(&self, other: &Self, cache: &mut TreeHashCache) -> Option<()> {
        let num_leaves = 4;

        let mut leaf_cache = cache.just_the_leaves(num_leaves)?;
        self.a.cached_hash_tree_root(&other.a, &mut leaf_cache)?;
        self.b.cached_hash_tree_root(&other.b, &mut leaf_cache)?;
        self.c.cached_hash_tree_root(&other.c, &mut leaf_cache)?;
        self.d.cached_hash_tree_root(&other.d, &mut leaf_cache)?;

        let nodes = num_nodes(num_leaves);
        let internal_chunks = nodes - num_leaves;

        for chunk in 0..internal_chunks {
            if cache.children_modified(chunk)? {
                cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?;
            }
        }

        Some(())
    }
}

/// A reference function to test against.
pub fn merkleize(values: &[u8]) -> Vec<u8> {
    let leaves = values.len() / HASHSIZE;

    if leaves == 0 || !leaves.is_power_of_two() {
        panic!("Handle bad leaf count");
    }

    let mut o: Vec<u8> = vec![0; (num_nodes(leaves) - leaves) * HASHSIZE];
    o.append(&mut values.to_vec());

    let mut i = o.len();
    let mut j = o.len() - values.len();

    while i >= MERKLE_HASH_CHUNCK {
        i -= MERKLE_HASH_CHUNCK;
        let hash = hash(&o[i..i + MERKLE_HASH_CHUNCK]);

        j -= HASHSIZE;
        o[j..j + HASHSIZE].copy_from_slice(&hash);
    }

    o
}

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

    fn join(many: Vec<&[u8]>) -> Vec<u8> {
        let mut all = vec![];
        for one in many {
            all.extend_from_slice(&mut one.clone())
        }
        all
    }

    #[test]
    fn merkleize_4_leaves() {
        let data1 = hash(&int_to_bytes32(1));
        let data2 = hash(&int_to_bytes32(2));
        let data3 = hash(&int_to_bytes32(3));
        let data4 = hash(&int_to_bytes32(4));

        let data = join(vec![&data1, &data2, &data3, &data4]);

        let cache = merkleize(&data);

        let hash_12 = {
            let mut joined = vec![];
            joined.append(&mut data1.clone());
            joined.append(&mut data2.clone());
            hash(&joined)
        };
        let hash_34 = {
            let mut joined = vec![];
            joined.append(&mut data3.clone());
            joined.append(&mut data4.clone());
            hash(&joined)
        };
        let hash_hash12_hash_34 = {
            let mut joined = vec![];
            joined.append(&mut hash_12.clone());
            joined.append(&mut hash_34.clone());
            hash(&joined)
        };

        for (i, chunk) in cache.chunks(HASHSIZE).enumerate().rev() {
            let expected = match i {
                0 => hash_hash12_hash_34.clone(),
                1 => hash_12.clone(),
                2 => hash_34.clone(),
                3 => data1.clone(),
                4 => data2.clone(),
                5 => data3.clone(),
                6 => data4.clone(),
                _ => vec![],
            };

            assert_eq!(chunk, &expected[..], "failed at {}", i);
        }
    }
}
<|endoftext|>


<commit_before>use super::Blob;
use failure::Error;
use rusoto_s3::{GetObjectRequest, S3Client, S3};
use std::convert::TryInto;
use std::io::Read;
use time::Timespec;

pub(crate) struct S3Backend<'a> {
    client: &'a S3Client,
    bucket: String,
}

impl<'a> S3Backend<'a> {
    pub(super) fn new(client: &'a S3Client, bucket: &str) -> Self {
        Self {
            client,
            bucket: bucket.into(),
        }
    }

    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
        let res = self
            .client
            .get_object(GetObjectRequest {
                bucket: self.bucket.clone(),
                key: path.into(),
                ..Default::default()
            })
            .sync()?;

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::with_capacity(
            res.content_length
                .and_then(|l| l.try_into().ok())
                .unwrap_or(0),
        );
        b.read_to_end(&mut content).unwrap();

        let date_updated = parse_timespec(&res.last_modified.unwrap())?;

        Ok(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated,
            content,
        })
    }
}

fn parse_timespec(raw: &str) -> Result<Timespec, Error> {
    Ok(time::strptime(raw, "%a, %d %b %Y %H:%M:%S %Z")?.to_timespec())
}

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

    #[test]
    fn test_parse_timespec() {
        crate::test::wrapper(|_| {
            // Test valid conversions
            assert_eq!(
                parse_timespec("Thu, 1 Jan 1970 00:00:00 GMT")?,
                Timespec::new(0, 0)
            );
            assert_eq!(
                parse_timespec("Mon, 16 Apr 2018 04:33:50 GMT")?,
                Timespec::new(1523853230, 0)
            );

            // Test invalid conversion
            assert!(parse_timespec("foo").is_err());

            Ok(())
        })
    }
}
<commit_msg>Remove unnecessary copy<commit_after>use super::Blob;
use failure::Error;
use rusoto_s3::{GetObjectRequest, S3Client, S3};
use std::convert::TryInto;
use std::io::Read;
use time::Timespec;

pub(crate) struct S3Backend<'a> {
    client: &'a S3Client,
    bucket: &'a str,
}

impl<'a> S3Backend<'a> {
    pub(super) fn new(client: &'a S3Client, bucket: &'a str) -> Self {
        Self {
            client,
            bucket,
        }
    }

    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
        let res = self
            .client
            .get_object(GetObjectRequest {
                bucket: self.bucket.to_string(),
                key: path.into(),
                ..Default::default()
            })
            .sync()?;

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::with_capacity(
            res.content_length
                .and_then(|l| l.try_into().ok())
                .unwrap_or(0),
        );
        b.read_to_end(&mut content).unwrap();

        let date_updated = parse_timespec(&res.last_modified.unwrap())?;

        Ok(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated,
            content,
        })
    }
}

fn parse_timespec(raw: &str) -> Result<Timespec, Error> {
    Ok(time::strptime(raw, "%a, %d %b %Y %H:%M:%S %Z")?.to_timespec())
}

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

    #[test]
    fn test_parse_timespec() {
        crate::test::wrapper(|_| {
            // Test valid conversions
            assert_eq!(
                parse_timespec("Thu, 1 Jan 1970 00:00:00 GMT")?,
                Timespec::new(0, 0)
            );
            assert_eq!(
                parse_timespec("Mon, 16 Apr 2018 04:33:50 GMT")?,
                Timespec::new(1523853230, 0)
            );

            // Test invalid conversion
            assert!(parse_timespec("foo").is_err());

            Ok(())
        })
    }
}
<|endoftext|>


<commit_before>use super::Blob;
use failure::Error;
use futures::Future;
use log::{error, warn};
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_s3::{GetObjectRequest, PutObjectRequest, S3Client, S3};
use std::convert::TryInto;
use std::io::Read;
use time::Timespec;
use tokio::runtime::Runtime;

pub(crate) struct S3Backend<'a> {
    client: S3Client,
    bucket: &'a str,
    runtime: Runtime,
}

impl<'a> S3Backend<'a> {
    pub(crate) fn new(client: S3Client, bucket: &'a str) -> Self {
        Self {
            client,
            bucket,
            runtime: Runtime::new().unwrap(),
        }
    }

    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
        let res = self
            .client
            .get_object(GetObjectRequest {
                bucket: self.bucket.to_string(),
                key: path.into(),
                ..Default::default()
            })
            .sync()?;

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::with_capacity(
            res.content_length
                .and_then(|l| l.try_into().ok())
                .unwrap_or(0),
        );
        b.read_to_end(&mut content).unwrap();

        let date_updated = parse_timespec(&res.last_modified.unwrap())?;

        Ok(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated,
            content,
        })
    }

    pub(super) fn store_batch(&mut self, batch: &[Blob]) -> Result<(), Error> {
        use futures::stream::FuturesUnordered;
        use futures::stream::Stream;
        let mut attempts = 0;

        loop {
            let mut futures = FuturesUnordered::new();
            for blob in batch {
                futures.push(
                    self.client
                        .put_object(PutObjectRequest {
                            bucket: self.bucket.to_string(),
                            key: blob.path.clone(),
                            body: Some(blob.content.clone().into()),
                            content_type: Some(blob.mime.clone()),
                            ..Default::default()
                        })
                        .inspect(|_| {
                            crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                        }),
                );
            }
            attempts += 1;

            match self.runtime.block_on(futures.map(drop).collect()) {
                // this batch was successful, start another batch if there are still more files
                Ok(_) => break,
                Err(err) => {
                    error!("failed to upload to s3: {:?}", err);
                    // if a futures error occurs, retry the batch
                    if attempts > 2 {
                        panic!("failed to upload 3 times, exiting");
                    }
                }
            }
        }
        Ok(())
    }
}

// public for testing
pub(crate) const TIME_FMT: &str = "%a, %d %b %Y %H:%M:%S %Z";

fn parse_timespec(raw: &str) -> Result<Timespec, Error> {
    Ok(time::strptime(raw, TIME_FMT)?.to_timespec())
}

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

    pub(crate) struct TestS3(S3Backend<'static>);

    use crate::storage::s3::S3Backend;
    use rusoto_core::RusotoResult;
    use rusoto_s3::{
        CreateBucketRequest, DeleteBucketRequest, DeleteObjectRequest, ListObjectsRequest,
        PutObjectError, PutObjectOutput, PutObjectRequest, S3,
    };

    impl TestS3 {
        pub(crate) fn new() -> Self {
            // A random bucket name is generated and used for the current connection.
            // This allows each test to create a fresh bucket to test with.
            let bucket = format!("docs-rs-test-bucket-{}", rand::random::<u64>());
            let client = crate::storage::s3::s3_client().unwrap();
            client
                .create_bucket(CreateBucketRequest {
                    bucket: bucket.clone(),
                    ..Default::default()
                })
                .sync()
                .expect("failed to create test bucket");
            let bucket = Box::leak(bucket.into_boxed_str());
            TestS3(S3Backend::new(client, bucket))
        }
        pub(crate) fn upload(&self, blob: Blob) -> RusotoResult<PutObjectOutput, PutObjectError> {
            self.0
                .client
                .put_object(PutObjectRequest {
                    bucket: self.0.bucket.to_owned(),
                    body: Some(blob.content.into()),
                    content_type: Some(blob.mime),
                    key: blob.path,
                    ..PutObjectRequest::default()
                })
                .sync()
        }
        fn assert_404(&self, path: &'static str) {
            use rusoto_core::RusotoError;
            use rusoto_s3::GetObjectError;

            let err = self.0.get(path).unwrap_err();
            match err
                .downcast_ref::<RusotoError<GetObjectError>>()
                .expect("wanted GetObject")
            {
                RusotoError::Unknown(http) => assert_eq!(http.status, 404),
                RusotoError::Service(GetObjectError::NoSuchKey(_)) => {}
                x => panic!("wrong error: {:?}", x),
            };
        }
        fn assert_blob(&self, blob: &Blob, path: &str) {
            let actual = self.0.get(path).unwrap();
            assert_eq!(blob.path, actual.path);
            assert_eq!(blob.content, actual.content);
            assert_eq!(blob.mime, actual.mime);
            // NOTE: this does _not_ compare the upload time since min.io doesn't allow this to be configured
        }
    }

    impl Drop for TestS3 {
        fn drop(&mut self) {
            let objects = self
                .0
                .client
                .list_objects(ListObjectsRequest {
                    bucket: self.0.bucket.to_owned(),
                    ..Default::default()
                })
                .sync()
                .unwrap();
            assert!(!objects.is_truncated.unwrap_or(false));
            for path in objects.contents.unwrap() {
                self.0
                    .client
                    .delete_object(DeleteObjectRequest {
                        bucket: self.0.bucket.to_owned(),
                        key: path.key.unwrap(),
                        ..Default::default()
                    })
                    .sync()
                    .unwrap();
            }
            let delete_req = DeleteBucketRequest {
                bucket: self.0.bucket.to_owned(),
            };
            self.0
                .client
                .delete_bucket(delete_req)
                .sync()
                .expect("failed to delete test bucket");
        }
    }

    #[test]
    fn test_parse_timespec() {
        crate::test::wrapper(|_| {
            // Test valid conversions
            assert_eq!(
                parse_timespec("Thu, 1 Jan 1970 00:00:00 GMT")?,
                Timespec::new(0, 0)
            );
            assert_eq!(
                parse_timespec("Mon, 16 Apr 2018 04:33:50 GMT")?,
                Timespec::new(1523853230, 0)
            );

            // Test invalid conversion
            assert!(parse_timespec("foo").is_err());

            Ok(())
        })
    }

    #[test]
    fn test_path_get() {
        crate::test::wrapper(|env| {
            let blob = Blob {
                path: "dir/foo.txt".into(),
                mime: "text/plain".into(),
                date_updated: Timespec::new(42, 0),
                content: "Hello world!".into(),
            };

            // Add a test file to the database
            let s3 = env.s3();
            s3.upload(blob.clone()).unwrap();

            // Test that the proper file was returned
            s3.assert_blob(&blob, "dir/foo.txt");

            // Test that other files are not returned
            s3.assert_404("dir/bar.txt");
            s3.assert_404("foo.txt");

            Ok(())
        });
    }
}

pub(crate) static S3_BUCKET_NAME: &str = "rust-docs-rs";

pub(crate) fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT")
            .ok()
            .map(|e| Region::Custom {
                name: std::env::var("S3_REGION").unwrap_or_else(|_| "us-west-1".to_owned()),
                endpoint: e,
            })
            .unwrap_or(Region::UsWest1),
    ))
}
<commit_msg>[BROKEN] start adding more tests<commit_after>use super::Blob;
use failure::Error;
use futures::Future;
use log::{error, warn};
use rusoto_core::region::Region;
use rusoto_credential::DefaultCredentialsProvider;
use rusoto_s3::{GetObjectRequest, PutObjectRequest, S3Client, S3};
use std::convert::TryInto;
use std::io::Read;
use time::Timespec;
use tokio::runtime::Runtime;

pub(crate) struct S3Backend<'a> {
    client: S3Client,
    bucket: &'a str,
    runtime: Runtime,
}

impl<'a> S3Backend<'a> {
    pub(crate) fn new(client: S3Client, bucket: &'a str) -> Self {
        Self {
            client,
            bucket,
            runtime: Runtime::new().unwrap(),
        }
    }

    pub(super) fn get(&self, path: &str) -> Result<Blob, Error> {
        let res = self
            .client
            .get_object(GetObjectRequest {
                bucket: self.bucket.to_string(),
                key: path.into(),
                ..Default::default()
            })
            .sync()?;

        let mut b = res.body.unwrap().into_blocking_read();
        let mut content = Vec::with_capacity(
            res.content_length
                .and_then(|l| l.try_into().ok())
                .unwrap_or(0),
        );
        b.read_to_end(&mut content).unwrap();

        let date_updated = parse_timespec(&res.last_modified.unwrap())?;

        Ok(Blob {
            path: path.into(),
            mime: res.content_type.unwrap(),
            date_updated,
            content,
        })
    }

    pub(super) fn store_batch(&mut self, batch: &[Blob]) -> Result<(), Error> {
        use futures::stream::FuturesUnordered;
        use futures::stream::Stream;
        let mut attempts = 0;

        loop {
            let mut futures = FuturesUnordered::new();
            for blob in batch {
                futures.push(
                    self.client
                        .put_object(PutObjectRequest {
                            bucket: self.bucket.to_string(),
                            key: blob.path.clone(),
                            body: Some(blob.content.clone().into()),
                            content_type: Some(blob.mime.clone()),
                            ..Default::default()
                        })
                        .inspect(|_| {
                            crate::web::metrics::UPLOADED_FILES_TOTAL.inc_by(1);
                        }),
                );
            }
            attempts += 1;

            match self.runtime.block_on(futures.map(drop).collect()) {
                // this batch was successful, start another batch if there are still more files
                Ok(_) => break,
                Err(err) => {
                    error!("failed to upload to s3: {:?}", err);
                    // if a futures error occurs, retry the batch
                    if attempts > 2 {
                        panic!("failed to upload 3 times, exiting");
                    }
                }
            }
        }
        Ok(())
    }
}

// public for testing
pub(crate) const TIME_FMT: &str = "%a, %d %b %Y %H:%M:%S %Z";

fn parse_timespec(raw: &str) -> Result<Timespec, Error> {
    Ok(time::strptime(raw, TIME_FMT)?.to_timespec())
}

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

    use crate::storage::s3::S3Backend;
    use rusoto_core::RusotoResult;
    use rusoto_s3::{
        CreateBucketRequest, DeleteBucketRequest, DeleteObjectRequest, ListObjectsRequest,
        PutObjectError, PutObjectOutput, PutObjectRequest, S3,
    };

    use std::cell::RefCell;

    pub(crate) struct TestS3(RefCell<S3Backend<'static>>);

    impl TestS3 {
        pub(crate) fn new() -> Self {
            // A random bucket name is generated and used for the current connection.
            // This allows each test to create a fresh bucket to test with.
            let bucket = format!("docs-rs-test-bucket-{}", rand::random::<u64>());
            let client = s3_client().unwrap();
            client
                .create_bucket(CreateBucketRequest {
                    bucket: bucket.clone(),
                    ..Default::default()
                })
                .sync()
                .expect("failed to create test bucket");
            let bucket = Box::leak(bucket.into_boxed_str());
            TestS3(S3Backend::new(client, bucket))
        }
        /// NOTE: this mocks the upload instead of using store_one
        /// TODO: I think `store_one` would be better
        pub(crate) fn upload(&self, blob: Blob) -> RusotoResult<PutObjectOutput, PutObjectError> {
            self.0
                .client
                .put_object(PutObjectRequest {
                    bucket: self.0.bucket.to_owned(),
                    body: Some(blob.content.into()),
                    content_type: Some(blob.mime),
                    key: blob.path,
                    ..PutObjectRequest::default()
                })
                .sync()
        }
        fn assert_404(&self, path: &'static str) {
            use rusoto_core::RusotoError;
            use rusoto_s3::GetObjectError;

            let err = self.0.get(path).unwrap_err();
            match err
                .downcast_ref::<RusotoError<GetObjectError>>()
                .expect("wanted GetObject")
            {
                RusotoError::Unknown(http) => assert_eq!(http.status, 404),
                RusotoError::Service(GetObjectError::NoSuchKey(_)) => {}
                x => panic!("wrong error: {:?}", x),
            };
        }
        fn assert_blob(&self, blob: &Blob, path: &str) {
            let actual = self.0.get(path).unwrap();
            assert_eq!(blob.path, actual.path);
            assert_eq!(blob.content, actual.content);
            assert_eq!(blob.mime, actual.mime);
            // NOTE: this does _not_ compare the upload time since min.io doesn't allow this to be configured
        }
    }

    impl Drop for TestS3 {
        fn drop(&mut self) {
            let list_req = ListObjectsRequest {
                bucket: self.0.bucket.to_owned(),
                ..Default::default()
            };
            let objects = self.0.client.list_objects(list_req).sync().unwrap();
            assert!(!objects.is_truncated.unwrap_or(false));
            for path in objects.contents.unwrap() {
                let delete_req = DeleteObjectRequest {
                    bucket: self.0.bucket.to_owned(),
                    key: path.key.unwrap(),
                    ..Default::default()
                };
                self.0.client.delete_object(delete_req).sync().unwrap();
            }
            let delete_req = DeleteBucketRequest {
                bucket: self.0.bucket.to_owned(),
            };
            self.0
                .client
                .delete_bucket(delete_req)
                .sync()
                .expect("failed to delete test bucket");
        }
    }

    #[test]
    fn test_parse_timespec() {
        // Test valid conversions
        assert_eq!(
            parse_timespec("Thu, 1 Jan 1970 00:00:00 GMT").unwrap(),
            Timespec::new(0, 0)
        );
        assert_eq!(
            parse_timespec("Mon, 16 Apr 2018 04:33:50 GMT").unwrap(),
            Timespec::new(1523853230, 0)
        );

        // Test invalid conversion
        assert!(parse_timespec("foo").is_err());
    }

    #[test]
    fn test_path_get() {
        wrapper(|env| {
            let blob = Blob {
                path: "dir/foo.txt".into(),
                mime: "text/plain".into(),
                date_updated: Timespec::new(42, 0),
                content: "Hello world!".into(),
            };

            // Add a test file to the database
            let s3 = env.s3();
            s3.upload(blob.clone()).unwrap();

            // Test that the proper file was returned
            s3.assert_blob(&blob, "dir/foo.txt");

            // Test that other files are not returned
            s3.assert_404("dir/bar.txt");
            s3.assert_404("foo.txt");

            Ok(())
        });
    }

    #[test]
    fn test_store() {
        wrapper(|env| {
            let s3 = env.s3();
            let names = [
                "a",
                "b",
                "a_very_long_file_name_that_has_an.extension",
                "parent/child",
                "h/i/g/h/l/y/_/n/e/s/t/e/d/_/d/i/r/e/c/t/o/r/i/e/s",
                "trailing/slash/",
            ];
            let blobs: Vec<_> = names.iter().map(|&path| Blob {
                path: path.into(),
                mime: "text/plain".into(),
                date_updated: Timespec::new(42, 0),
                content: "Hello world!".into(),
            }).collect();
            s3.0.store_batch(&blobs).unwrap();
            for blob in &blobs {
                s3.assert_blob(blob, &blob.path);
            }
            Ok(())
        })
    }
}

pub(crate) static S3_BUCKET_NAME: &str = "rust-docs-rs";

pub(crate) fn s3_client() -> Option<S3Client> {
    // If AWS keys aren't configured, then presume we should use the DB exclusively
    // for file storage.
    if std::env::var_os("AWS_ACCESS_KEY_ID").is_none() && std::env::var_os("FORCE_S3").is_none() {
        return None;
    }
    let creds = match DefaultCredentialsProvider::new() {
        Ok(creds) => creds,
        Err(err) => {
            warn!("failed to retrieve AWS credentials: {}", err);
            return None;
        }
    };
    Some(S3Client::new_with(
        rusoto_core::request::HttpClient::new().unwrap(),
        creds,
        std::env::var("S3_ENDPOINT")
            .ok()
            .map(|e| Region::Custom {
                name: std::env::var("S3_REGION").unwrap_or_else(|_| "us-west-1".to_owned()),
                endpoint: e,
            })
            .unwrap_or(Region::UsWest1),
    ))
}
<|endoftext|>


<commit_before>use super::{Icon,IconScraper};
use std::str::FromStr;

pub trait Strategy {
    fn get_guesses(self, &mut IconScraper) -> Vec<Icon>;
}

pub struct DefaultFaviconPathStrategy;
impl Strategy for DefaultFaviconPathStrategy {
    fn get_guesses(self, parser: &mut IconScraper) -> Vec<Icon> {
        let mut icon = Icon::from_url(parser.document_url.join("/favicon.ico").unwrap());
        if icon.fetch().is_ok() {
            vec![icon]
        } else {
            vec![]
        }
    }
}

pub struct LinkRelStrategy;
impl Strategy for LinkRelStrategy {
    fn get_guesses(self, parser: &mut IconScraper) -> Vec<Icon> {
        let mut rv = vec![];
        let dom = match parser.dom {
            Some(ref x) => x,
            None => return rv
        };

        for data in dom.select("link[rel*=icon]").unwrap() {
            let attrs = data.attributes.borrow();
            let href = match attrs.get(atom!("href")) {
                Some(x) => x,
                None => continue
            };

            let icon_url = match parser.document_url.join(href) {
                Ok(x) => x,
                Err(_) => continue
            };

            let mut sizes = attrs.get(atom!("sizes"))
                .unwrap_or("")
                .split('x')
                .filter_map(|d| u32::from_str(d).ok());

            let (x, y) = match (sizes.next(), sizes.next()) {
                (Some(x), Some(y)) => (Some(x), Some(y)),
                _ => (None, None)
            };

            rv.push({
                let mut icon = Icon::from_url(icon_url);
                icon.width = x;
                icon.height = y;
                icon
            });
        };

        rv
    }
}

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

    use url;
    use kuchiki;
    use kuchiki::traits::*;

    #[test]
    fn test_apple_touch_icon_without_size_attr() {
        // laverna.cc does this.
        let mut scraper = IconScraper {
            document_url: url::Url::parse("http://example.com/").unwrap(),
            dom: Some(kuchiki::parse_html().one("<!DOCTYPE html>

                    <link rel=apple-touch-icon href=apple-touch-icon.png>

            "))
        };

        let mut icons = LinkRelStrategy.get_guesses(&mut scraper);
        assert_eq!(icons.len(), 1);
        assert_eq!(icons.pop().unwrap().url, url::Url::parse("http://example.com/apple-touch-icon.png").unwrap());
    }
}
<commit_msg>Test sharesome<commit_after>use super::{Icon,IconScraper};
use std::str::FromStr;

pub trait Strategy {
    fn get_guesses(self, &mut IconScraper) -> Vec<Icon>;
}

pub struct DefaultFaviconPathStrategy;
impl Strategy for DefaultFaviconPathStrategy {
    fn get_guesses(self, parser: &mut IconScraper) -> Vec<Icon> {
        let mut icon = Icon::from_url(parser.document_url.join("/favicon.ico").unwrap());
        if icon.fetch().is_ok() {
            vec![icon]
        } else {
            vec![]
        }
    }
}

pub struct LinkRelStrategy;
impl Strategy for LinkRelStrategy {
    fn get_guesses(self, parser: &mut IconScraper) -> Vec<Icon> {
        let mut rv = vec![];
        let dom = match parser.dom {
            Some(ref x) => x,
            None => return rv
        };

        for data in dom.select("link[rel*=icon]").unwrap() {
            let attrs = data.attributes.borrow();
            let href = match attrs.get(atom!("href")) {
                Some(x) => x,
                None => continue
            };

            let icon_url = match parser.document_url.join(href) {
                Ok(x) => x,
                Err(_) => continue
            };

            let mut sizes = attrs.get(atom!("sizes"))
                .unwrap_or("")
                .split('x')
                .filter_map(|d| u32::from_str(d).ok());

            let (x, y) = match (sizes.next(), sizes.next()) {
                (Some(x), Some(y)) => (Some(x), Some(y)),
                _ => (None, None)
            };

            rv.push({
                let mut icon = Icon::from_url(icon_url);
                icon.width = x;
                icon.height = y;
                icon
            });
        };

        rv
    }
}

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

    use url;
    use kuchiki;
    use kuchiki::traits::*;

    #[test]
    fn test_apple_touch_icon_without_size_attr() {
        // laverna.cc does this.
        let mut scraper = IconScraper {
            document_url: url::Url::parse("http://example.com/").unwrap(),
            dom: Some(kuchiki::parse_html().one("<!DOCTYPE html>

                    <link rel=apple-touch-icon href=apple-touch-icon.png>

            "))
        };

        let mut icons = LinkRelStrategy.get_guesses(&mut scraper);
        assert_eq!(icons.len(), 1);
        assert_eq!(icons.pop().unwrap().url, url::Url::parse("http://example.com/apple-touch-icon.png").unwrap());
    }

    #[test]
    fn test_sharesome() {
        let mut scraper = IconScraper::from_http("https://sharesome.5apps.com/");
        assert_eq!(scraper.fetch_icons().largest().unwrap().url, url::Url::parse("https://sharesome.5apps.com/application_icon_x512.png").unwrap());
    }
}
<|endoftext|>


<commit_before>/* This file is part of Grust, GObject introspection bindings for Rust
 *
 * Copyright (C) 2013  Mikhail Zabaluev <mikhail.zabaluev@gmail.com>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301  USA
 */

extern mod grust (name="grust", vers="0.1");
extern mod gio (name="grust-Gio", vers="2.0");

extern mod std;

use core::result::{Result,Ok};
use core::comm::{Port,stream};
use core::task::TaskResult;
use core::util;
// use std::timer::recv_timeout;
// use std::uv_global_loop;
use grust::eventloop::EventLoop;

// We have to do this because of an rpath problem with crates linking to
// foreign libraries
extern mod grustna {
}

// Test timeout in milliseconds
static TEST_TIMEOUT: uint = 3000u;

fn spawn_with_future(func: ~fn()) -> Port<TaskResult> {
    let mut (port, _) = stream::<TaskResult>();
    task::task()
        .future_result(|p| {
            port = p;
        })
        .spawn(func);
    port
}

fn tcase(test: ~fn()) {
    grust::init();

    let port = spawn_with_future(test);

    // recv_timeout is broken, see https://github.com/mozilla/rust/issues/6089
    match Some(port.recv()) /* recv_timeout(&uv_global_loop::get(), TEST_TIMEOUT, &port) */ {
        Some(task::Success) => {}
        Some(task::Failure) => {
            fail!(~"test failed");
        }
        None => {
            error!("test timed out");
            unsafe { libc::abort(); }
        }
    }
}

fn tcase_result(test: ~fn() -> Result<(),()>) {
    do tcase {
        let result = test();
        assert!(result == Ok(()));
    }
}

#[test]
fn simple() {
    do tcase {
        let fobj = gio::File::new_for_path("/dev/null");
        let f = fobj.interface() as &gio::File;
        assert!(f.get_path().to_str() == ~"/dev/null");
    }
}

#[test]
fn as_interface() {
    do tcase {
        gio::File::new_for_path("/dev/null").as_interface(|fobj| {
                let f = fobj as &gio::File;
                assert!(f.get_path().to_str() == ~"/dev/null");
            });
    }
}

#[test]
fn new_ref() {
    do tcase {
        let fobj = gio::File::new_for_path("/dev/null");
        let gobj = fobj.interface().new_ref();
        let g = gobj.interface() as &gio::File;
        assert!(g.get_path().to_str() == ~"/dev/null");
    }
}

#[test]
fn clone() {
    do tcase {
        let fobj = gio::File::new_for_path("/dev/null");
        fobj.clone().as_interface(|gobj| {
                let g = gobj as &gio::File;
                assert!(g.get_path().to_str() == ~"/dev/null");
            });
    }
}

#[test]
fn off_task() {
    do tcase_result {
        let f = ~gio::File::new_for_path("/dev/null");
        do task::try {
            let f = f.interface() as &gio::File;
            let p = f.get_path();
            assert!(p.to_str() == ~"/dev/null");
        }
    }
}

#[test]
fn off_task_as_interface() {
    do tcase_result {
        let fobj = ~gio::File::new_for_path("/dev/null");
        do task::try {
            fobj.as_interface(|fifa| {
                    let f = fifa as &gio::File;
                    let p = f.get_path();
                    assert!(p.to_str() == ~"/dev/null");
                });
        }
    }
}

#[test]
fn async() {
    do tcase {
        let fobj = gio::File::new_for_path("/dev/null");
        let f = fobj.interface() as &gio::File;
        let el = EventLoop::new();
        let elo = ~el.clone();
        f.read_async(0, None,
                |obj, res| {
                    let f: &gio::File = obj.cast::<gio::raw::GFile>()
                                        as &gio::File;
                    let in = f.read_finish(res);
                    util::ignore(in);
                    elo.quit();
                });
        el.run();
    }
}

#[test]
fn async_off_stack() {
    do tcase {
        let fobj = ~gio::File::new_for_path("/dev/null");
        let el = EventLoop::new();
        let elo = ~el.clone();
        task::spawn_sched(task::ThreadPerCore, || {
            let f = fobj.interface() as &gio::File;
            let elo2 = ~elo.clone();
            f.read_async(0, None, |obj, res| {
                    let f = obj.cast::<gio::raw::GFile>() as &gio::File;
                    f.read_finish(res);
                    elo2.quit();
                });
        });
        el.run();
    }
}
<commit_msg>Added an off-task call test, ignored due to main loop blocking<commit_after>/* This file is part of Grust, GObject introspection bindings for Rust
 *
 * Copyright (C) 2013  Mikhail Zabaluev <mikhail.zabaluev@gmail.com>
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
 * 02110-1301  USA
 */

extern mod grust (name="grust", vers="0.1");
extern mod gio (name="grust-Gio", vers="2.0");

extern mod std;

use core::result::{Result,Ok};
use core::comm::{Port,stream};
use core::task::TaskResult;
use core::util;
// use std::timer::recv_timeout;
// use std::uv_global_loop;
use grust::eventloop::EventLoop;

// We have to do this because of an rpath problem with crates linking to
// foreign libraries
extern mod grustna {
}

// Test timeout in milliseconds
static TEST_TIMEOUT: uint = 3000u;

fn spawn_with_future(func: ~fn()) -> Port<TaskResult> {
    let mut (port, _) = stream::<TaskResult>();
    task::task()
        .future_result(|p| {
            port = p;
        })
        .spawn(func);
    port
}

fn tcase(test: ~fn()) {
    grust::init();

    let port = spawn_with_future(test);

    // recv_timeout is broken, see https://github.com/mozilla/rust/issues/6089
    match Some(port.recv()) /* recv_timeout(&uv_global_loop::get(), TEST_TIMEOUT, &port) */ {
        Some(task::Success) => {}
        Some(task::Failure) => {
            fail!(~"test failed");
        }
        None => {
            error!("test timed out");
            unsafe { libc::abort(); }
        }
    }
}

fn tcase_result(test: ~fn() -> Result<(),()>) {
    do tcase {
        let result = test();
        assert!(result == Ok(()));
    }
}

#[test]
fn simple() {
    do tcase {
        let fobj = gio::File::new_for_path("/dev/null");
        let f = fobj.interface() as &gio::File;
        assert!(f.get_path().to_str() == ~"/dev/null");
    }
}

#[test]
fn as_interface() {
    do tcase {
        gio::File::new_for_path("/dev/null").as_interface(|fobj| {
                let f = fobj as &gio::File;
                assert!(f.get_path().to_str() == ~"/dev/null");
            });
    }
}

#[test]
fn new_ref() {
    do tcase {
        let fobj = gio::File::new_for_path("/dev/null");
        let gobj = fobj.interface().new_ref();
        let g = gobj.interface() as &gio::File;
        assert!(g.get_path().to_str() == ~"/dev/null");
    }
}

#[test]
fn clone() {
    do tcase {
        let fobj = gio::File::new_for_path("/dev/null");
        fobj.clone().as_interface(|gobj| {
                let g = gobj as &gio::File;
                assert!(g.get_path().to_str() == ~"/dev/null");
            });
    }
}

#[test]
fn off_task() {
    do tcase_result {
        let f = ~gio::File::new_for_path("/dev/null");
        do task::try {
            let f = f.interface() as &gio::File;
            let p = f.get_path();
            assert!(p.to_str() == ~"/dev/null");
        }
    }
}

#[test]
fn off_task_as_interface() {
    do tcase_result {
        let fobj = ~gio::File::new_for_path("/dev/null");
        do task::try {
            fobj.as_interface(|fifa| {
                    let f = fifa as &gio::File;
                    let p = f.get_path();
                    assert!(p.to_str() == ~"/dev/null");
                });
        }
    }
}

#[test]
fn async() {
    do tcase {
        let fobj = gio::File::new_for_path("/dev/null");
        let f = fobj.interface() as &gio::File;
        let el = EventLoop::new();
        let elo = ~el.clone();
        f.read_async(0, None,
                |obj, res| {
                    let f: &gio::File = obj.cast::<gio::raw::GFile>()
                                        as &gio::File;
                    let in = f.read_finish(res);
                    util::ignore(in);
                    elo.quit();
                });
        el.run();
    }
}

#[test]
#[ignore]  // See https://github.com/mzabaluev/grust/issues/4
fn async_off_task() {
    do tcase {
        let fobj = ~gio::File::new_for_path("/dev/null");
        let el = EventLoop::new();
        let elo = ~el.clone();
        do spawn {
            let f = fobj.interface() as &gio::File;
            let elo2 = ~elo.clone();
            f.read_async(0, None, |obj, res| {
                    let fifa: &gio::interfaces::File = obj.cast();
                    let f = fifa as &gio::File;
                    f.read_finish(res);
                    elo2.quit();
                });
        }
        el.run();
    }
}

#[test]
fn async_off_stack() {
    do tcase {
        let fobj = ~gio::File::new_for_path("/dev/null");
        let el = EventLoop::new();
        let elo = ~el.clone();
        task::spawn_sched(task::ThreadPerCore, || {
            let f = fobj.interface() as &gio::File;
            let elo2 = ~elo.clone();
            f.read_async(0, None, |obj, res| {
                    let f = obj.cast::<gio::raw::GFile>() as &gio::File;
                    f.read_finish(res);
                    elo2.quit();
                });
        });
        el.run();
    }
}
<|endoftext|>


<commit_before>extern crate libc;
extern crate security_framework;
extern crate security_framework_sys;
extern crate tempdir;

use self::security_framework::base;
use self::security_framework::certificate::SecCertificate;
use self::security_framework::identity::SecIdentity;
use self::security_framework::import_export::{Pkcs12ImportOptions, ImportedIdentityOptions};
use self::security_framework::secure_transport::{self, SslContext, ProtocolSide, ConnectionType,
                                                 SslProtocol, ClientBuilder};
use self::security_framework_sys::base::errSecIO;
use self::tempdir::TempDir;
use std::fmt;
use std::io;
use std::error;
use std::sync::Mutex;
use std::sync::{Once, ONCE_INIT};

#[cfg(not(target_os = "ios"))]
use self::security_framework::os::macos::keychain::{self, SecKeychain, KeychainSettings};

use Protocol;

static SET_AT_EXIT: Once = ONCE_INIT;

lazy_static! {
    static ref TEMP_KEYCHAIN: Mutex<Option<(SecKeychain, TempDir)>> = Mutex::new(None);
}

fn convert_protocol(protocol: Protocol) -> SslProtocol {
    match protocol {
        Protocol::Sslv3 => SslProtocol::Ssl3,
        Protocol::Tlsv10 => SslProtocol::Tls1,
        Protocol::Tlsv11 => SslProtocol::Tls11,
        Protocol::Tlsv12 => SslProtocol::Tls12,
        Protocol::__NonExhaustive => unreachable!(),
    }
}

fn protocol_min_max(protocols: &[Protocol]) -> (SslProtocol, SslProtocol) {
    let mut min = Protocol::Tlsv12;
    let mut max = Protocol::Sslv3;
    for protocol in protocols {
        if (*protocol as usize) < (min as usize) {
            min = *protocol;
        }
        if (*protocol as usize) > (max as usize) {
            max = *protocol;
        }
    }
    (convert_protocol(min), convert_protocol(max))
}

pub struct Error(base::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

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

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<base::Error> for Error {
    fn from(error: base::Error) -> Error {
        Error(error)
    }
}

#[derive(Clone)]
pub struct Pkcs12 {
    identity: SecIdentity,
    chain: Vec<SecCertificate>,
}

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let mut imports = try!(Pkcs12::import_options(buf, pass));
        let import = imports.pop().unwrap();

        let identity = import.identity.expect(
            "Pkcs12 files must include an identity",
        );

        // FIXME: Compare the certificates for equality using CFEqual
        let identity_cert = try!(identity.certificate()).to_der();

        Ok(Pkcs12 {
            identity: identity,
            chain: import
                .cert_chain
                .unwrap_or(vec![])
                .into_iter()
                .filter(|c| c.to_der() != identity_cert)
                .collect(),
        })
    }

    #[cfg(not(target_os = "ios"))]
    fn import_options(
        buf: &[u8],
        pass: &str,
    ) -> Result<Vec<ImportedIdentityOptions>, Error> {
        SET_AT_EXIT.call_once(|| {
            extern "C" fn atexit() {
                *TEMP_KEYCHAIN.lock().unwrap() = None;
            }
            unsafe {
                libc::atexit(atexit);
            }
        });

        let keychain = match *TEMP_KEYCHAIN.lock().unwrap() {
            Some((ref keychain, _)) => keychain.clone(),
            ref mut lock @ None => {
                let dir = TempDir::new("native-tls").map_err(|_| {
                    Error(base::Error::from(errSecIO))
                })?;

                let mut keychain = keychain::CreateOptions::new().password(pass).create(
                    dir.path().join("tmp.keychain"),
                )?;
                keychain.set_settings(&KeychainSettings::new())?;

                *lock = Some((keychain, dir));
                lock.as_ref().unwrap().0.clone()
            }
        };
        let imports = try!(
            Pkcs12ImportOptions::new()
                .passphrase(pass)
                .keychain(keychain)
                .import_optional(buf)
        );
        Ok(imports)
    }

    #[cfg(target_os = "ios")]
    fn import_options(
        buf: &[u8],
        pass: &str,
    ) -> Result<Vec<ImportedIdentityOptions>, Error> {
        let imports = try!(
            Pkcs12ImportOptions::new()
                .passphrase(pass)
                .import_optional(buf)
        );
        Ok(imports)
    }
}

pub struct Certificate(SecCertificate);

impl Certificate {
    pub fn from_der(buf: &[u8]) -> Result<Certificate, Error> {
        let cert = try!(SecCertificate::from_der(buf));
        Ok(Certificate(cert))
    }
}

pub enum HandshakeError<S> {
    Interrupted(MidHandshakeTlsStream<S>),
    Failure(Error),
}

impl<S> From<secure_transport::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::HandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream::Server(s))
            }
        }
    }
}

impl<S> From<secure_transport::ClientHandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::ClientHandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::ClientHandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::ClientHandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream::Client(s))
            }
        }
    }
}

impl<S> From<base::Error> for HandshakeError<S> {
    fn from(e: base::Error) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub enum MidHandshakeTlsStream<S> {
    Server(secure_transport::MidHandshakeSslStream<S>),
    Client(secure_transport::MidHandshakeClientBuilder<S>),
}

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
where
    S: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            MidHandshakeTlsStream::Server(ref s) => s.fmt(fmt),
            MidHandshakeTlsStream::Client(ref s) => s.fmt(fmt),
        }
    }
}

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn get_ref(&self) -> &S {
        match *self {
            MidHandshakeTlsStream::Server(ref s) => s.get_ref(),
            MidHandshakeTlsStream::Client(ref s) => s.get_ref(),
        }
    }

    pub fn get_mut(&mut self) -> &mut S {
        match *self {
            MidHandshakeTlsStream::Server(ref mut s) => s.get_mut(),
            MidHandshakeTlsStream::Client(ref mut s) => s.get_mut(),
        }
    }

    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self {
            MidHandshakeTlsStream::Server(s) => {
                match s.handshake() {
                    Ok(s) => Ok(TlsStream(s)),
                    Err(e) => Err(e.into()),
                }
            }
            MidHandshakeTlsStream::Client(s) => {
                match s.handshake() {
                    Ok(s) => Ok(TlsStream(s)),
                    Err(e) => Err(e.into()),
                }
            }
        }
    }
}

pub struct TlsConnectorBuilder(TlsConnector);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        self.0.pkcs12 = Some(pkcs12);
        Ok(())
    }

    pub fn add_root_certificate(&mut self, cert: Certificate) -> Result<(), Error> {
        self.0.roots.push(cert.0);
        Ok(())
    }

    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        self.0.protocols = protocols.to_vec();
        Ok(())
    }

    pub fn build(self) -> Result<TlsConnector, Error> {
        Ok(self.0)
    }
}

#[derive(Clone)]
pub struct TlsConnector {
    pkcs12: Option<Pkcs12>,
    protocols: Vec<Protocol>,
    roots: Vec<SecCertificate>,
}

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        Ok(TlsConnectorBuilder(TlsConnector {
            pkcs12: None,
            protocols: vec![Protocol::Tlsv10, Protocol::Tlsv11, Protocol::Tlsv12],
            roots: vec![],
        }))
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        self.connect_inner(Some(domain), stream)
    }

    pub fn connect_no_domain<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        self.connect_inner(None, stream)
    }

    fn connect_inner<S>(
        &self,
        domain: Option<&str>,
        stream: S,
    ) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let mut builder = ClientBuilder::new();
        let (min, max) = protocol_min_max(&self.protocols);
        builder.protocol_min(min);
        builder.protocol_max(max);
        if let Some(pkcs12) = self.pkcs12.as_ref() {
            builder.identity(&pkcs12.identity, &pkcs12.chain);
        }
        builder.anchor_certificates(&self.roots);

        let r = match domain {
            Some(domain) => builder.handshake2(domain, stream),
            None => builder.danger_handshake_without_providing_domain_for_certificate_validation_and_server_name_indication(stream),
        };
        match r {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsAcceptorBuilder(TlsAcceptor);

impl TlsAcceptorBuilder {
    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        self.0.protocols = protocols.to_vec();
        Ok(())
    }

    pub fn build(self) -> Result<TlsAcceptor, Error> {
        Ok(self.0)
    }
}

#[derive(Clone)]
pub struct TlsAcceptor {
    pkcs12: Pkcs12,
    protocols: Vec<Protocol>,
}

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        Ok(TlsAcceptorBuilder(TlsAcceptor {
            pkcs12: pkcs12,
            protocols: vec![Protocol::Tlsv10, Protocol::Tlsv11, Protocol::Tlsv12],
        }))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let mut ctx = try!(SslContext::new(
            ProtocolSide::Server,
            ConnectionType::Stream,
        ));

        let (min, max) = protocol_min_max(&self.protocols);
        try!(ctx.set_protocol_version_min(min));
        try!(ctx.set_protocol_version_max(max));
        try!(ctx.set_certificate(
            &self.pkcs12.identity,
            &self.pkcs12.chain,
        ));
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsStream<S>(secure_transport::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(try!(self.0.context().buffered_read_size()))
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        try!(self.0.close());
        Ok(())
    }
}

impl<S: io::Read + io::Write> io::Read for TlsStream<S> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

/// Security Framework-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the Security Framework `SslStream`.
    fn raw_stream(&self) -> &secure_transport::SslStream<S>;

    /// Returns a mutable reference to the Security Framework `SslStream`.
    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &secure_transport::SslStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S> {
        &mut (self.0).0
    }
}

/// Security Framework-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Deprecated
    #[deprecated(since = "0.1.2", note = "use add_root_certificate")]
    fn anchor_certificates(&mut self, certs: &[SecCertificate]) -> &mut Self;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn anchor_certificates(&mut self, certs: &[SecCertificate]) -> &mut Self {
        (self.0).0.roots = certs.to_owned();
        self
    }
}

/// Security Framework-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying Security Framework error for inspection.
    fn security_framework_error(&self) -> &base::Error;
}

impl ErrorExt for ::Error {
    fn security_framework_error(&self) -> &base::Error {
        &(self.0).0
    }
}
<commit_msg>Don't need temp keychain on ios<commit_after>extern crate libc;
extern crate security_framework;
extern crate security_framework_sys;
extern crate tempdir;

use self::security_framework::base;
use self::security_framework::certificate::SecCertificate;
use self::security_framework::identity::SecIdentity;
use self::security_framework::import_export::{Pkcs12ImportOptions, ImportedIdentityOptions};
use self::security_framework::secure_transport::{self, SslContext, ProtocolSide, ConnectionType,
                                                 SslProtocol, ClientBuilder};
use self::security_framework_sys::base::errSecIO;
use self::tempdir::TempDir;
use std::fmt;
use std::io;
use std::error;
use std::sync::Mutex;
use std::sync::{Once, ONCE_INIT};

#[cfg(not(target_os = "ios"))]
use self::security_framework::os::macos::keychain::{self, SecKeychain, KeychainSettings};

use Protocol;

static SET_AT_EXIT: Once = ONCE_INIT;

#[cfg(not(target_os = "ios"))]
lazy_static! {
    static ref TEMP_KEYCHAIN: Mutex<Option<(SecKeychain, TempDir)>> = Mutex::new(None);
}

fn convert_protocol(protocol: Protocol) -> SslProtocol {
    match protocol {
        Protocol::Sslv3 => SslProtocol::Ssl3,
        Protocol::Tlsv10 => SslProtocol::Tls1,
        Protocol::Tlsv11 => SslProtocol::Tls11,
        Protocol::Tlsv12 => SslProtocol::Tls12,
        Protocol::__NonExhaustive => unreachable!(),
    }
}

fn protocol_min_max(protocols: &[Protocol]) -> (SslProtocol, SslProtocol) {
    let mut min = Protocol::Tlsv12;
    let mut max = Protocol::Sslv3;
    for protocol in protocols {
        if (*protocol as usize) < (min as usize) {
            min = *protocol;
        }
        if (*protocol as usize) > (max as usize) {
            max = *protocol;
        }
    }
    (convert_protocol(min), convert_protocol(max))
}

pub struct Error(base::Error);

impl error::Error for Error {
    fn description(&self) -> &str {
        error::Error::description(&self.0)
    }

    fn cause(&self) -> Option<&error::Error> {
        error::Error::cause(&self.0)
    }
}

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

impl fmt::Debug for Error {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl From<base::Error> for Error {
    fn from(error: base::Error) -> Error {
        Error(error)
    }
}

#[derive(Clone)]
pub struct Pkcs12 {
    identity: SecIdentity,
    chain: Vec<SecCertificate>,
}

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let mut imports = try!(Pkcs12::import_options(buf, pass));
        let import = imports.pop().unwrap();

        let identity = import.identity.expect(
            "Pkcs12 files must include an identity",
        );

        // FIXME: Compare the certificates for equality using CFEqual
        let identity_cert = try!(identity.certificate()).to_der();

        Ok(Pkcs12 {
            identity: identity,
            chain: import
                .cert_chain
                .unwrap_or(vec![])
                .into_iter()
                .filter(|c| c.to_der() != identity_cert)
                .collect(),
        })
    }

    #[cfg(not(target_os = "ios"))]
    fn import_options(
        buf: &[u8],
        pass: &str,
    ) -> Result<Vec<ImportedIdentityOptions>, Error> {
        SET_AT_EXIT.call_once(|| {
            extern "C" fn atexit() {
                *TEMP_KEYCHAIN.lock().unwrap() = None;
            }
            unsafe {
                libc::atexit(atexit);
            }
        });

        let keychain = match *TEMP_KEYCHAIN.lock().unwrap() {
            Some((ref keychain, _)) => keychain.clone(),
            ref mut lock @ None => {
                let dir = TempDir::new("native-tls").map_err(|_| {
                    Error(base::Error::from(errSecIO))
                })?;

                let mut keychain = keychain::CreateOptions::new().password(pass).create(
                    dir.path().join("tmp.keychain"),
                )?;
                keychain.set_settings(&KeychainSettings::new())?;

                *lock = Some((keychain, dir));
                lock.as_ref().unwrap().0.clone()
            }
        };
        let imports = try!(
            Pkcs12ImportOptions::new()
                .passphrase(pass)
                .keychain(keychain)
                .import_optional(buf)
        );
        Ok(imports)
    }

    #[cfg(target_os = "ios")]
    fn import_options(
        buf: &[u8],
        pass: &str,
    ) -> Result<Vec<ImportedIdentityOptions>, Error> {
        let imports = try!(
            Pkcs12ImportOptions::new()
                .passphrase(pass)
                .import_optional(buf)
        );
        Ok(imports)
    }
}

pub struct Certificate(SecCertificate);

impl Certificate {
    pub fn from_der(buf: &[u8]) -> Result<Certificate, Error> {
        let cert = try!(SecCertificate::from_der(buf));
        Ok(Certificate(cert))
    }
}

pub enum HandshakeError<S> {
    Interrupted(MidHandshakeTlsStream<S>),
    Failure(Error),
}

impl<S> From<secure_transport::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::HandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream::Server(s))
            }
        }
    }
}

impl<S> From<secure_transport::ClientHandshakeError<S>> for HandshakeError<S> {
    fn from(e: secure_transport::ClientHandshakeError<S>) -> HandshakeError<S> {
        match e {
            secure_transport::ClientHandshakeError::Failure(e) => HandshakeError::Failure(e.into()),
            secure_transport::ClientHandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream::Client(s))
            }
        }
    }
}

impl<S> From<base::Error> for HandshakeError<S> {
    fn from(e: base::Error) -> HandshakeError<S> {
        HandshakeError::Failure(e.into())
    }
}

pub enum MidHandshakeTlsStream<S> {
    Server(secure_transport::MidHandshakeSslStream<S>),
    Client(secure_transport::MidHandshakeClientBuilder<S>),
}

impl<S> fmt::Debug for MidHandshakeTlsStream<S>
where
    S: fmt::Debug,
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            MidHandshakeTlsStream::Server(ref s) => s.fmt(fmt),
            MidHandshakeTlsStream::Client(ref s) => s.fmt(fmt),
        }
    }
}

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn get_ref(&self) -> &S {
        match *self {
            MidHandshakeTlsStream::Server(ref s) => s.get_ref(),
            MidHandshakeTlsStream::Client(ref s) => s.get_ref(),
        }
    }

    pub fn get_mut(&mut self) -> &mut S {
        match *self {
            MidHandshakeTlsStream::Server(ref mut s) => s.get_mut(),
            MidHandshakeTlsStream::Client(ref mut s) => s.get_mut(),
        }
    }

    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self {
            MidHandshakeTlsStream::Server(s) => {
                match s.handshake() {
                    Ok(s) => Ok(TlsStream(s)),
                    Err(e) => Err(e.into()),
                }
            }
            MidHandshakeTlsStream::Client(s) => {
                match s.handshake() {
                    Ok(s) => Ok(TlsStream(s)),
                    Err(e) => Err(e.into()),
                }
            }
        }
    }
}

pub struct TlsConnectorBuilder(TlsConnector);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        self.0.pkcs12 = Some(pkcs12);
        Ok(())
    }

    pub fn add_root_certificate(&mut self, cert: Certificate) -> Result<(), Error> {
        self.0.roots.push(cert.0);
        Ok(())
    }

    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        self.0.protocols = protocols.to_vec();
        Ok(())
    }

    pub fn build(self) -> Result<TlsConnector, Error> {
        Ok(self.0)
    }
}

#[derive(Clone)]
pub struct TlsConnector {
    pkcs12: Option<Pkcs12>,
    protocols: Vec<Protocol>,
    roots: Vec<SecCertificate>,
}

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        Ok(TlsConnectorBuilder(TlsConnector {
            pkcs12: None,
            protocols: vec![Protocol::Tlsv10, Protocol::Tlsv11, Protocol::Tlsv12],
            roots: vec![],
        }))
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        self.connect_inner(Some(domain), stream)
    }

    pub fn connect_no_domain<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        self.connect_inner(None, stream)
    }

    fn connect_inner<S>(
        &self,
        domain: Option<&str>,
        stream: S,
    ) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let mut builder = ClientBuilder::new();
        let (min, max) = protocol_min_max(&self.protocols);
        builder.protocol_min(min);
        builder.protocol_max(max);
        if let Some(pkcs12) = self.pkcs12.as_ref() {
            builder.identity(&pkcs12.identity, &pkcs12.chain);
        }
        builder.anchor_certificates(&self.roots);

        let r = match domain {
            Some(domain) => builder.handshake2(domain, stream),
            None => builder.danger_handshake_without_providing_domain_for_certificate_validation_and_server_name_indication(stream),
        };
        match r {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsAcceptorBuilder(TlsAcceptor);

impl TlsAcceptorBuilder {
    pub fn supported_protocols(&mut self, protocols: &[Protocol]) -> Result<(), Error> {
        self.0.protocols = protocols.to_vec();
        Ok(())
    }

    pub fn build(self) -> Result<TlsAcceptor, Error> {
        Ok(self.0)
    }
}

#[derive(Clone)]
pub struct TlsAcceptor {
    pkcs12: Pkcs12,
    protocols: Vec<Protocol>,
}

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        Ok(TlsAcceptorBuilder(TlsAcceptor {
            pkcs12: pkcs12,
            protocols: vec![Protocol::Tlsv10, Protocol::Tlsv11, Protocol::Tlsv12],
        }))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let mut ctx = try!(SslContext::new(
            ProtocolSide::Server,
            ConnectionType::Stream,
        ));

        let (min, max) = protocol_min_max(&self.protocols);
        try!(ctx.set_protocol_version_min(min));
        try!(ctx.set_protocol_version_max(max));
        try!(ctx.set_certificate(
            &self.pkcs12.identity,
            &self.pkcs12.chain,
        ));
        match ctx.handshake(stream) {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

pub struct TlsStream<S>(secure_transport::SslStream<S>);

impl<S: fmt::Debug> fmt::Debug for TlsStream<S> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(&self.0, fmt)
    }
}

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

    pub fn get_mut(&mut self) -> &mut S {
        self.0.get_mut()
    }

    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(try!(self.0.context().buffered_read_size()))
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        try!(self.0.close());
        Ok(())
    }
}

impl<S: io::Read + io::Write> io::Read for TlsStream<S> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.0.read(buf)
    }
}

impl<S: io::Read + io::Write> io::Write for TlsStream<S> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.0.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.0.flush()
    }
}

/// Security Framework-specific extensions to `TlsStream`.
pub trait TlsStreamExt<S> {
    /// Returns a shared reference to the Security Framework `SslStream`.
    fn raw_stream(&self) -> &secure_transport::SslStream<S>;

    /// Returns a mutable reference to the Security Framework `SslStream`.
    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S>;
}

impl<S> TlsStreamExt<S> for ::TlsStream<S> {
    fn raw_stream(&self) -> &secure_transport::SslStream<S> {
        &(self.0).0
    }

    fn raw_stream_mut(&mut self) -> &mut secure_transport::SslStream<S> {
        &mut (self.0).0
    }
}

/// Security Framework-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Deprecated
    #[deprecated(since = "0.1.2", note = "use add_root_certificate")]
    fn anchor_certificates(&mut self, certs: &[SecCertificate]) -> &mut Self;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn anchor_certificates(&mut self, certs: &[SecCertificate]) -> &mut Self {
        (self.0).0.roots = certs.to_owned();
        self
    }
}

/// Security Framework-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying Security Framework error for inspection.
    fn security_framework_error(&self) -> &base::Error;
}

impl ErrorExt for ::Error {
    fn security_framework_error(&self) -> &base::Error {
        &(self.0).0
    }
}
<|endoftext|>