<commit_before>/// A graph similarity score using neighbor matching according to [this paper][1].
///
/// [1]: http://arxiv.org/abs/1009.5290 "2010, Mladen Nikolic, Measuring Similarity
///      of Graph Nodes by Neighbor Matching"

extern crate nalgebra;
extern crate munkres;

use nalgebra::{DMat, Shape, ApproxEq};
use munkres::{WeightMatrix, solve_assignment};
use std::cmp;
use std::mem;

pub type Idx = u32;

trait Edges {
    /// The number of edges
    fn len(&self) -> usize;

    /// Returns the target node of the nth-edge
    fn nth_edge(&self, n: usize) -> usize;

    /// Returns the nth edge weight
    fn nth_edge_weight(&self, _n: usize) -> f32 {
        panic!();
    }
}

impl<'a> Edges for &'a [Idx] {
    #[inline]
    fn len(&self) -> usize {
        let x: &[Idx] = self;
        x.len()
    }
    #[inline]
    fn nth_edge(&self, n: usize) -> usize {
        self[n] as usize
    }
}

#[inline]
/// Calculates the similarity of two nodes `i` and `j`.
///
/// `n_i` contains the neighborhood of i (either in or out neighbors, not both)
/// `n_j` contains the neighborhood of j (either in or out neighbors, not both)
/// `x`   the similarity matrix.
fn s_next<T: Edges>(n_i: T, n_j: T, x: &DMat<f32>) -> f32 {
    let max_deg = cmp::max(n_i.len(), n_j.len());
    let min_deg = cmp::min(n_i.len(), n_j.len());

    if min_deg == 0 {
        // in the paper, 0/0 is defined as 1.0
        return 1.0;
    }

    assert!(min_deg > 0 && max_deg > 0);

    // map indicies from 0..min(degree) to the node indices
    let mapidx = |(a, b)| (n_i.nth_edge(a), n_j.nth_edge(b));

    let mut w = WeightMatrix::from_fn(min_deg, |ab| x[mapidx(ab)]);

    let assignment = solve_assignment(&mut w);
    assert!(assignment.len() == min_deg);

    let sum: f32 = assignment.iter().fold(0.0, |acc, &ab| acc + x[mapidx(ab)]);

    return sum / max_deg as f32;
}

#[inline]
/// Calculates x[k+1]
///
/// `node_color_scale((i,j))`: If two nodes `i` (of graph A) and `j` (of graph B)
/// are of different color, this can be set to return 0.0. Alternatively a
/// node-color distance (within 0...1) could be used to penalize.
fn next_x<F>(x: &DMat<f32>,
             new_x: &mut DMat<f32>,
             in_a: &[Vec<Idx>],
             in_b: &[Vec<Idx>],
             out_a: &[Vec<Idx>],
             out_b: &[Vec<Idx>],
             node_color_scale: F)
    where F: Fn((usize, usize)) -> f32
{
    let shape = x.shape();
    assert!(shape == new_x.shape());

    for i in 0..shape.0 {
        for j in 0..shape.1 {
            let in_i: &[Idx] = &in_a[i];
            let in_j: &[Idx] = &in_b[j];
            let out_i: &[Idx] = &out_a[i];
            let out_j: &[Idx] = &out_b[j];
            new_x[(i, j)] = node_color_scale((i, j)) *
                            (s_next(in_i, in_j, x) + s_next(out_i, out_j, x)) /
                            2.0;
        }
    }
}

#[derive(Copy, Clone, Debug)]
pub struct Graph<'a> {
    in_edges: &'a [Vec<Idx>],
    out_edges: &'a [Vec<Idx>],
}

impl<'a> Graph<'a> {
    pub fn new<'b>(in_edges: &'b [Vec<Idx>], out_edges: &'b [Vec<Idx>]) -> Graph<'b> {
        assert!(in_edges.len() == out_edges.len());
        Graph {
            in_edges: in_edges,
            out_edges: out_edges,
        }
    }

    fn num_nodes(&self) -> usize {
        let n = self.in_edges.len();
        assert!(n == self.out_edges.len());
        n
    }

    fn node_degree(&self, node_idx: usize) -> usize {
        self.in_edges[node_idx].len() + self.out_edges[node_idx].len()
    }
}

#[derive(Debug)]
pub struct GraphSimilarityMatrix<'a, F: 'a> {
    graph_a: Graph<'a>,
    graph_b: Graph<'a>,
    node_color_scale: &'a F,
    // current version of similarity matrix
    current: DMat<f32>,
    // previous version of similarity matrix
    previous: DMat<f32>,
    // current number of iterations
    num_iterations: usize,
}

impl<'a, F> GraphSimilarityMatrix<'a, F> where F: Fn((usize, usize)) -> f32
{
    pub fn new<'b>(graph_a: Graph<'b>,
                   graph_b: Graph<'b>,
                   node_color_scale: &'b F)
                   -> GraphSimilarityMatrix<'b, F> {
        // `x` is the node-similarity matrix.
        // we initialize `x`, so that x[i,j]=1 for all i in A.edges() and j in
        // B.edges().
        let x: DMat<f32> = DMat::from_fn(graph_a.num_nodes(), graph_b.num_nodes(), |i, j| {
            node_color_scale((i, j)) *
            if graph_a.node_degree(i) > 0 && graph_b.node_degree(j) > 0 {
                1.0
            } else {
                0.0
            }
        });

        let new_x: DMat<f32> = DMat::new_zeros(graph_a.num_nodes(), graph_b.num_nodes());

        GraphSimilarityMatrix {
            graph_a: graph_a,
            graph_b: graph_b,
            node_color_scale: node_color_scale,
            current: x,
            previous: new_x,
            num_iterations: 0,
        }
    }

    fn in_eps(&self, eps: f32) -> bool {
        self.previous.approx_eq_eps(&self.current, &eps)
    }

    /// Calculates the next iteration of the similarity matrix.
    pub fn next(&mut self) {
        next_x(&self.current,
               &mut self.previous,
               self.graph_a.in_edges,
               self.graph_b.in_edges,
               self.graph_a.out_edges,
               self.graph_b.out_edges,
               self.node_color_scale);
        mem::swap(&mut self.previous, &mut self.current);
        self.num_iterations += 1;
    }

    #[inline]
    /// Iteratively calculate the similarity matrix.
    ///
    /// `eps`:   When to stop the iteration
    /// `stop_after_iter`: Stop after iteration (Calculate x(stop_after_iter))
    pub fn iterate(&mut self, eps: f32, stop_after_iter: usize) {
        for _ in 0..stop_after_iter {
            if self.in_eps(eps) {
                break;
            }
            self.next();
        }
    }

    pub fn matrix(&self) -> &DMat<f32> {
        &self.current
    }

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

    fn optimal_node_assignment(&self, n: usize) -> Vec<(usize, usize)> {
        let x = &self.current;
        assert!(n > 0);
        let mut w = WeightMatrix::from_fn(n, |ij| x[ij]);
        let assignment = solve_assignment(&mut w);
        assert!(assignment.len() == n);
        assignment
    }

    fn score_optimal_sum(&self, n: usize) -> f32 {
        self.optimal_node_assignment(n).iter().fold(0.0, |acc, &ab| acc + self.current[ab])
    }

    /// Calculate a measure how good the edge weights match up.
    ///
    /// We start by calculating the optimal node assignment between nodes of graph A and
    /// graph B, then compare all outgoing edges of similar nodes by again using an
    /// assignment.
    pub fn score_outgoing_edge_weights(&self) -> f32 {
        // XXX
        0.0
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Used as default in the paper.
    pub fn score_sum_norm_min_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            self.score_optimal_sum(n) / n as f32
        } else {
            0.0
        }
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Penalizes the difference in size of graphs.
    pub fn score_sum_norm_max_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        let m = cmp::max(x.nrows(), x.ncols());

        if n > 0 {
            assert!(m > 0);
            self.score_optimal_sum(n) / m as f32
        } else {
            0.0
        }
    }

    /// Calculates the average over the whole node similarity matrix. This is faster,
    /// as no assignment has to be found. "Graphs with greater number of automorphisms
    /// would be considered to be more self-similar than graphs without automorphisms."
    pub fn score_average(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            let sum: f32 = x.as_vec().iter().fold(0.0, |acc, &v| acc + v);
            let len = x.as_vec().len();
            assert!(len > 0);
            sum / len as f32
        } else {
            0.0
        }
    }
}

#[test]
fn test_matrix() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let node_color = |_| 1.0;
    let mut s = GraphSimilarityMatrix::new(Graph::new(&in_a, &out_a),
                                           Graph::new(&in_b, &out_b),
                                           &node_color);
    s.iterate(0.1, 100);

    // println!("{:?}", s);
    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(2, mat.nrows());
    assert_eq!(2, mat.ncols());

    // A and B are isomorphic
    assert_eq!(1.0, mat[(0, 0)]);
    assert_eq!(1.0, mat[(0, 1)]);
    assert_eq!(1.0, mat[(1, 0)]);
    assert_eq!(1.0, mat[(1, 1)]);
}

#[test]
fn test_matrix_iter1() {
    let in_a = vec![vec![0, 0, 0]];
    let out_a = vec![vec![0, 0, 0]];

    let in_b = vec![vec![0, 0, 0, 0, 0]];
    let out_b = vec![vec![0, 0, 0, 0, 0]];

    let node_color = |_| 1.0;
    let mut s = GraphSimilarityMatrix::new(Graph::new(&in_a, &out_a),
                                           Graph::new(&in_b, &out_b),
                                           &node_color);
    s.iterate(0.1, 1);

    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(3.0 / 5.0, mat[(0, 0)]);
}

#[test]
fn test_score() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let node_color = |_| 1.0;
    let mut s = GraphSimilarityMatrix::new(Graph::new(&in_a, &out_a),
                                           Graph::new(&in_b, &out_b),
                                           &node_color);
    s.iterate(0.1, 100);

    assert_eq!(1, s.num_iterations());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_min_degree());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_max_degree());
}
<commit_msg>Simplify lifetimes<commit_after>/// A graph similarity score using neighbor matching according to [this paper][1].
///
/// [1]: http://arxiv.org/abs/1009.5290 "2010, Mladen Nikolic, Measuring Similarity
///      of Graph Nodes by Neighbor Matching"

extern crate nalgebra;
extern crate munkres;

use nalgebra::{DMat, Shape, ApproxEq};
use munkres::{WeightMatrix, solve_assignment};
use std::cmp;
use std::mem;

pub type Idx = u32;

trait Edges {
    /// The number of edges
    fn len(&self) -> usize;

    /// Returns the target node of the nth-edge
    fn nth_edge(&self, n: usize) -> usize;

    /// Returns the nth edge weight
    fn nth_edge_weight(&self, _n: usize) -> f32 {
        panic!();
    }
}

impl<'a> Edges for &'a [Idx] {
    #[inline]
    fn len(&self) -> usize {
        let x: &[Idx] = self;
        x.len()
    }
    #[inline]
    fn nth_edge(&self, n: usize) -> usize {
        self[n] as usize
    }
}

#[inline]
/// Calculates the similarity of two nodes `i` and `j`.
///
/// `n_i` contains the neighborhood of i (either in or out neighbors, not both)
/// `n_j` contains the neighborhood of j (either in or out neighbors, not both)
/// `x`   the similarity matrix.
fn s_next<T: Edges>(n_i: T, n_j: T, x: &DMat<f32>) -> f32 {
    let max_deg = cmp::max(n_i.len(), n_j.len());
    let min_deg = cmp::min(n_i.len(), n_j.len());

    if min_deg == 0 {
        // in the paper, 0/0 is defined as 1.0
        return 1.0;
    }

    assert!(min_deg > 0 && max_deg > 0);

    // map indicies from 0..min(degree) to the node indices
    let mapidx = |(a, b)| (n_i.nth_edge(a), n_j.nth_edge(b));

    let mut w = WeightMatrix::from_fn(min_deg, |ab| x[mapidx(ab)]);

    let assignment = solve_assignment(&mut w);
    assert!(assignment.len() == min_deg);

    let sum: f32 = assignment.iter().fold(0.0, |acc, &ab| acc + x[mapidx(ab)]);

    return sum / max_deg as f32;
}

#[inline]
/// Calculates x[k+1]
///
/// `node_color_scale((i,j))`: If two nodes `i` (of graph A) and `j` (of graph B)
/// are of different color, this can be set to return 0.0. Alternatively a
/// node-color distance (within 0...1) could be used to penalize.
fn next_x<F>(x: &DMat<f32>,
             new_x: &mut DMat<f32>,
             in_a: &[Vec<Idx>],
             in_b: &[Vec<Idx>],
             out_a: &[Vec<Idx>],
             out_b: &[Vec<Idx>],
             node_color_scale: F)
    where F: Fn((usize, usize)) -> f32
{
    let shape = x.shape();
    assert!(shape == new_x.shape());

    for i in 0..shape.0 {
        for j in 0..shape.1 {
            let in_i: &[Idx] = &in_a[i];
            let in_j: &[Idx] = &in_b[j];
            let out_i: &[Idx] = &out_a[i];
            let out_j: &[Idx] = &out_b[j];
            new_x[(i, j)] = node_color_scale((i, j)) *
                            (s_next(in_i, in_j, x) + s_next(out_i, out_j, x)) /
                            2.0;
        }
    }
}

#[derive(Copy, Clone, Debug)]
pub struct Graph<'a> {
    in_edges: &'a [Vec<Idx>],
    out_edges: &'a [Vec<Idx>],
}

impl<'a> Graph<'a> {
    pub fn new(in_edges: &'a [Vec<Idx>], out_edges: &'a [Vec<Idx>]) -> Graph<'a> {
        assert!(in_edges.len() == out_edges.len());
        Graph {
            in_edges: in_edges,
            out_edges: out_edges,
        }
    }

    fn num_nodes(&self) -> usize {
        let n = self.in_edges.len();
        assert!(n == self.out_edges.len());
        n
    }

    fn node_degree(&self, node_idx: usize) -> usize {
        self.in_edges[node_idx].len() + self.out_edges[node_idx].len()
    }
}

#[derive(Debug)]
pub struct GraphSimilarityMatrix<'a, F: 'a> {
    graph_a: Graph<'a>,
    graph_b: Graph<'a>,
    node_color_scale: &'a F,
    // current version of similarity matrix
    current: DMat<f32>,
    // previous version of similarity matrix
    previous: DMat<f32>,
    // current number of iterations
    num_iterations: usize,
}

impl<'a, F> GraphSimilarityMatrix<'a, F> where F: Fn((usize, usize)) -> f32
{
    pub fn new(graph_a: Graph<'a>,
               graph_b: Graph<'a>,
               node_color_scale: &'a F)
               -> GraphSimilarityMatrix<'a, F> {
        // `x` is the node-similarity matrix.
        // we initialize `x`, so that x[i,j]=1 for all i in A.edges() and j in
        // B.edges().
        let x: DMat<f32> = DMat::from_fn(graph_a.num_nodes(), graph_b.num_nodes(), |i, j| {
            node_color_scale((i, j)) *
            if graph_a.node_degree(i) > 0 && graph_b.node_degree(j) > 0 {
                1.0
            } else {
                0.0
            }
        });

        let new_x: DMat<f32> = DMat::new_zeros(graph_a.num_nodes(), graph_b.num_nodes());

        GraphSimilarityMatrix {
            graph_a: graph_a,
            graph_b: graph_b,
            node_color_scale: node_color_scale,
            current: x,
            previous: new_x,
            num_iterations: 0,
        }
    }

    fn in_eps(&self, eps: f32) -> bool {
        self.previous.approx_eq_eps(&self.current, &eps)
    }

    /// Calculates the next iteration of the similarity matrix.
    pub fn next(&mut self) {
        next_x(&self.current,
               &mut self.previous,
               self.graph_a.in_edges,
               self.graph_b.in_edges,
               self.graph_a.out_edges,
               self.graph_b.out_edges,
               self.node_color_scale);
        mem::swap(&mut self.previous, &mut self.current);
        self.num_iterations += 1;
    }

    #[inline]
    /// Iteratively calculate the similarity matrix.
    ///
    /// `eps`:   When to stop the iteration
    /// `stop_after_iter`: Stop after iteration (Calculate x(stop_after_iter))
    pub fn iterate(&mut self, eps: f32, stop_after_iter: usize) {
        for _ in 0..stop_after_iter {
            if self.in_eps(eps) {
                break;
            }
            self.next();
        }
    }

    pub fn matrix(&self) -> &DMat<f32> {
        &self.current
    }

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

    fn optimal_node_assignment(&self, n: usize) -> Vec<(usize, usize)> {
        let x = &self.current;
        assert!(n > 0);
        let mut w = WeightMatrix::from_fn(n, |ij| x[ij]);
        let assignment = solve_assignment(&mut w);
        assert!(assignment.len() == n);
        assignment
    }

    fn score_optimal_sum(&self, n: usize) -> f32 {
        self.optimal_node_assignment(n).iter().fold(0.0, |acc, &ab| acc + self.current[ab])
    }

    /// Calculate a measure how good the edge weights match up.
    ///
    /// We start by calculating the optimal node assignment between nodes of graph A and
    /// graph B, then compare all outgoing edges of similar nodes by again using an
    /// assignment.
    pub fn score_outgoing_edge_weights(&self) -> f32 {
        // XXX
        0.0
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Used as default in the paper.
    pub fn score_sum_norm_min_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            self.score_optimal_sum(n) / n as f32
        } else {
            0.0
        }
    }

    /// Sums the optimal assignment of the node similarities and normalizes (divides)
    /// by the min degree of both graphs.
    /// Penalizes the difference in size of graphs.
    pub fn score_sum_norm_max_degree(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        let m = cmp::max(x.nrows(), x.ncols());

        if n > 0 {
            assert!(m > 0);
            self.score_optimal_sum(n) / m as f32
        } else {
            0.0
        }
    }

    /// Calculates the average over the whole node similarity matrix. This is faster,
    /// as no assignment has to be found. "Graphs with greater number of automorphisms
    /// would be considered to be more self-similar than graphs without automorphisms."
    pub fn score_average(&self) -> f32 {
        let x = &self.current;
        let n = cmp::min(x.nrows(), x.ncols());
        if n > 0 {
            let sum: f32 = x.as_vec().iter().fold(0.0, |acc, &v| acc + v);
            let len = x.as_vec().len();
            assert!(len > 0);
            sum / len as f32
        } else {
            0.0
        }
    }
}

#[test]
fn test_matrix() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let node_color = |_| 1.0;
    let mut s = GraphSimilarityMatrix::new(Graph::new(&in_a, &out_a),
                                           Graph::new(&in_b, &out_b),
                                           &node_color);
    s.iterate(0.1, 100);

    // println!("{:?}", s);
    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(2, mat.nrows());
    assert_eq!(2, mat.ncols());

    // A and B are isomorphic
    assert_eq!(1.0, mat[(0, 0)]);
    assert_eq!(1.0, mat[(0, 1)]);
    assert_eq!(1.0, mat[(1, 0)]);
    assert_eq!(1.0, mat[(1, 1)]);
}

#[test]
fn test_matrix_iter1() {
    let in_a = vec![vec![0, 0, 0]];
    let out_a = vec![vec![0, 0, 0]];

    let in_b = vec![vec![0, 0, 0, 0, 0]];
    let out_b = vec![vec![0, 0, 0, 0, 0]];

    let node_color = |_| 1.0;
    let mut s = GraphSimilarityMatrix::new(Graph::new(&in_a, &out_a),
                                           Graph::new(&in_b, &out_b),
                                           &node_color);
    s.iterate(0.1, 1);

    assert_eq!(1, s.num_iterations());
    let mat = s.matrix();
    assert_eq!(3.0 / 5.0, mat[(0, 0)]);
}

#[test]
fn test_score() {
    // A: 0 --> 1
    let in_a = vec![vec![], vec![0]];
    let out_a = vec![vec![1], vec![]];

    // B: 0 <-- 1
    let in_b = vec![vec![1], vec![]];
    let out_b = vec![vec![], vec![0]];

    let node_color = |_| 1.0;
    let mut s = GraphSimilarityMatrix::new(Graph::new(&in_a, &out_a),
                                           Graph::new(&in_b, &out_b),
                                           &node_color);
    s.iterate(0.1, 100);

    assert_eq!(1, s.num_iterations());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_min_degree());

    // The score is 1.0 <=> A and B are isomorphic
    assert_eq!(1.0, s.score_sum_norm_max_degree());
}
<|endoftext|>


<commit_before>// Copyright 2016 oauth-client-rs Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

#![warn(bad_style)]
// #![warn(missing_docs)]
#![warn(unused)]
#![warn(unused_extern_crates)]
#![warn(unused_import_braces)]
#![warn(unused_qualifications)]
#![warn(unused_results)]

extern crate crypto;
extern crate curl;
#[macro_use]
extern crate log;
extern crate rand;
extern crate rustc_serialize;
extern crate time;
extern crate url;

use std::borrow::Cow;
use std::collections::HashMap;
use std::{error, fmt};
use rand::Rng;
use rustc_serialize::base64::{self, ToBase64};
use crypto::hmac::Hmac;
use crypto::mac::{Mac, MacResult};
use crypto::sha1::Sha1;
use curl::http::{self, Response};
use curl::http::handle::Method;
use url::percent_encoding;

#[derive(Debug)]
pub enum Error {
    Curl(curl::ErrCode),
    HttpStatus(Response),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::Curl(ref err) => write!(f, "Curl error: {}", err),
            Error::HttpStatus(ref resp) => write!(f, "HTTP status error: {}", resp),
        }
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::Curl(ref err) => err.description(),
            Error::HttpStatus(_) => "HTTP status error",
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::Curl(ref err) => Some(err),
            Error::HttpStatus(_) => None,
        }
    }
}

impl From<curl::ErrCode> for Error {
    fn from(err: curl::ErrCode) -> Error {
        Error::Curl(err)
    }
}

#[derive(Clone, Debug)]
pub struct Token<'a> {
    pub key: Cow<'a, str>,
    pub secret: Cow<'a, str>,
}

impl<'a> Token<'a> {
    pub fn new<K, S>(key: K, secret: S) -> Token<'a>
        where K: Into<Cow<'a, str>>,
              S: Into<Cow<'a, str>>
    {
        Token {
            key: key.into(),
            secret: secret.into(),
        }
    }
}

pub type ParamList<'a> = HashMap<Cow<'a, str>, Cow<'a, str>>;

fn insert_param<'a, K, V>(param: &mut ParamList<'a>, key: K, value: V) -> Option<Cow<'a, str>>
    where K: Into<Cow<'a, str>>,
          V: Into<Cow<'a, str>>
{
    param.insert(key.into(), value.into())
}

fn join_query<'a>(param: &ParamList<'a>) -> String {
    let mut pairs = param.iter()
                         .map(|(k, v)| format!("{}={}", encode(&k), encode(&v)))
                         .collect::<Vec<_>>();
    pairs.sort();
    pairs.join("&")
}

fn encode(s: &str) -> String {
    percent_encoding::utf8_percent_encode(s, percent_encoding::FORM_URLENCODED_ENCODE_SET)
}

fn hmac_sha1(key: &[u8], data: &[u8]) -> MacResult {
    let mut hmac = Hmac::new(Sha1::new(), key);
    hmac.input(data);
    hmac.result()
}

fn signature(method: &str,
             uri: &str,
             query: &str,
             consumer_secret: &str,
             token_secret: Option<&str>)
             -> String {
    let base = format!("{}&{}&{}", encode(method), encode(uri), encode(query));
    let key = format!("{}&{}",
                      encode(consumer_secret),
                      encode(token_secret.unwrap_or("")));
    let conf = base64::Config {
        char_set: base64::CharacterSet::Standard,
        newline: base64::Newline::LF,
        pad: true,
        line_length: None,
    };
    debug!("Signature base string: {}", base);
    debug!("Authorization header: Authorization: {}", base);
    hmac_sha1(key.as_bytes(), base.as_bytes()).code().to_base64(conf)
}

fn header(param: &ParamList) -> String {
    let mut pairs = param.iter()
                         .filter(|&(k, _)| k.starts_with("oauth_"))
                         .map(|(k, v)| format!("{}=\"{}\"", k, encode(&v)))
                         .collect::<Vec<_>>();
    pairs.sort();
    format!("OAuth {}", pairs.join(", "))
}

fn body(param: &ParamList) -> String {
    let mut pairs = param.iter()
                         .filter(|&(k, _)| !k.starts_with("oauth_"))
                         .map(|(k, v)| format!("{}={}", k, encode(&v)))
                         .collect::<Vec<_>>();
    pairs.sort();
    format!("{}", pairs.join("&"))
}

fn get_header(method: Method,
              uri: &str,
              consumer: &Token,
              token: Option<&Token>,
              other_param: Option<&ParamList>)
              -> (String, String) {
    let mut param = HashMap::new();
    let timestamp = format!("{}", time::now_utc().to_timespec().sec);
    let nonce = rand::thread_rng().gen_ascii_chars().take(32).collect::<String>();

    let _ = insert_param(&mut param, "oauth_consumer_key", consumer.key.to_string());
    let _ = insert_param(&mut param, "oauth_nonce", nonce);
    let _ = insert_param(&mut param, "oauth_signature_method", "HMAC-SHA1");
    let _ = insert_param(&mut param, "oauth_timestamp", timestamp);
    let _ = insert_param(&mut param, "oauth_version", "1.0");
    if let Some(tk) = token {
        let _ = insert_param(&mut param, "oauth_token", tk.key.as_ref());
    }

    if let Some(ps) = other_param {
        for (k, v) in ps.iter() {
            let _ = insert_param(&mut param, k.as_ref(), v.as_ref());
        }
    }

    let method_str = match method {
        Method::Options => "OPTIONS",
        Method::Get => "GET",
        Method::Head => "HEAD",
        Method::Post => "POST",
        Method::Put => "PUT",
        Method::Patch => "PATCH",
        Method::Delete => "DELETE",
        Method::Trace => "TRACE",
        Method::Connect => "CONNECT",
    };

    let sign = signature(method_str,
                         uri,
                         join_query(&param).as_ref(),
                         consumer.secret.as_ref(),
                         token.map(|t| t.secret.as_ref()));
    let _ = insert_param(&mut param, "oauth_signature", sign);

    (header(&param), body(&param))
}

pub fn authorization_header(method: Method,
                            uri: &str,
                            consumer: &Token,
                            token: Option<&Token>,
                            other_param: Option<&ParamList>)
                            -> String {
    get_header(method, uri, consumer, token, other_param).0
}

pub fn get(uri: &str,
           consumer: &Token,
           token: Option<&Token>,
           other_param: Option<&ParamList>)
           -> Result<Vec<u8>, Error> {
    let (header, body) = get_header(Method::Get, uri, consumer, token, other_param);
    let req_uri = if body.len() > 0 {
        format!("{}?{}", uri, body)
    } else {
        format!("{}", uri)
    };
    let resp = try!(http::handle()
                        .get(req_uri)
                        .header("Authorization", header.as_ref())
                        .exec());
    debug!("{}", resp);
    if resp.get_code() != 200 {
        return Err(Error::HttpStatus(resp));
    }
    Ok(resp.move_body())
}

pub fn post(uri: &str,
            consumer: &Token,
            token: Option<&Token>,
            other_param: Option<&ParamList>)
            -> Result<Vec<u8>, Error> {
    let (header, body) = get_header(Method::Post, uri, consumer, token, other_param);
    let resp = try!(http::handle()
                        .post(uri, &body)
                        .header("Authorization", header.as_ref())
                        .content_type("application/x-www-form-urlencoded")
                        .exec());
    debug!("{}", resp);
    if resp.get_code() != 200 {
        return Err(Error::HttpStatus(resp));
    }
    Ok(resp.move_body())
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    #[test]
    fn query() {
        let mut map = HashMap::new();
        let _ = map.insert("aaa".into(), "AAA".into());
        let _ = map.insert("bbbb".into(), "BBBB".into());
        let query = super::join_query(&map);
        assert_eq!("aaa=AAA&bbbb=BBBB", query);
    }
}
<commit_msg>Add docs<commit_after>// Copyright 2016 oauth-client-rs Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! OAuth 1.0 client library for Rust.
//!
//! Dependent on libcurl.
//!
//! [Repository](https://github.com/charlag/oauth-client-rs)
//!
//!# Examples
//!
//! Send request for request token.
//!
//! ```
//!let consumer = Token::new("key", "secret")
//!let bytes = oauth::get(api::REQUEST_TOKEN, consumer, None, None).unwrap();
//! ```
#![warn(bad_style)]
#![warn(missing_docs)]
#![warn(unused)]
#![warn(unused_extern_crates)]
#![warn(unused_import_braces)]
#![warn(unused_qualifications)]
#![warn(unused_results)]

extern crate crypto;
extern crate curl;
#[macro_use]
extern crate log;
extern crate rand;
extern crate rustc_serialize;
extern crate time;
extern crate url;

use std::borrow::Cow;
use std::collections::HashMap;
use std::{error, fmt};
use rand::Rng;
use rustc_serialize::base64::{self, ToBase64};
use crypto::hmac::Hmac;
use crypto::mac::{Mac, MacResult};
use crypto::sha1::Sha1;
use curl::http::{self, Response};
use curl::http::handle::Method;
use url::percent_encoding;

/// The `Error` type
#[derive(Debug)]
pub enum Error {
    /// Curl error
    Curl(curl::ErrCode),
    /// Http status
    HttpStatus(Response),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::Curl(ref err) => write!(f, "Curl error: {}", err),
            Error::HttpStatus(ref resp) => write!(f, "HTTP status error: {}", resp),
        }
    }
}

impl error::Error for Error {
    fn description(&self) -> &str {
        match *self {
            Error::Curl(ref err) => err.description(),
            Error::HttpStatus(_) => "HTTP status error",
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            Error::Curl(ref err) => Some(err),
            Error::HttpStatus(_) => None,
        }
    }
}

impl From<curl::ErrCode> for Error {
    fn from(err: curl::ErrCode) -> Error {
        Error::Curl(err)
    }
}

/// Token structure for the OAuth
#[derive(Clone, Debug)]
pub struct Token<'a> {
    pub key: Cow<'a, str>,
    pub secret: Cow<'a, str>,
}

impl<'a> Token<'a> {
    /// Create new token from `key` and `secret`
    ///
    ///# Examples
    ///
    /// ```
    ///let consumer = Token::new("key", "secret");
    /// ```
    pub fn new<K, S>(key: K, secret: S) -> Token<'a>
        where K: Into<Cow<'a, str>>,
              S: Into<Cow<'a, str>>
    {
        Token {
            key: key.into(),
            secret: secret.into(),
        }
    }
}

/// Alias for `HashMap<Cow<'a, str>, Cow<'a, str>>`
pub type ParamList<'a> = HashMap<Cow<'a, str>, Cow<'a, str>>;

fn insert_param<'a, K, V>(param: &mut ParamList<'a>, key: K, value: V) -> Option<Cow<'a, str>>
    where K: Into<Cow<'a, str>>,
          V: Into<Cow<'a, str>>
{
    param.insert(key.into(), value.into())
}

fn join_query<'a>(param: &ParamList<'a>) -> String {
    let mut pairs = param.iter()
                         .map(|(k, v)| format!("{}={}", encode(&k), encode(&v)))
                         .collect::<Vec<_>>();
    pairs.sort();
    pairs.join("&")
}

/// Percent encode string
fn encode(s: &str) -> String {
    percent_encoding::utf8_percent_encode(s, percent_encoding::FORM_URLENCODED_ENCODE_SET)
}

/// Wrapper function around 'crypto::Hmac'
fn hmac_sha1(key: &[u8], data: &[u8]) -> MacResult {
    let mut hmac = Hmac::new(Sha1::new(), key);
    hmac.input(data);
    hmac.result()
}

/// Create signature. See https://dev.twitter.com/oauth/overview/creating-signatures
fn signature(method: &str,
             uri: &str,
             query: &str,
             consumer_secret: &str,
             token_secret: Option<&str>)
             -> String {
    let base = format!("{}&{}&{}", encode(method), encode(uri), encode(query));
    let key = format!("{}&{}",
                      encode(consumer_secret),
                      encode(token_secret.unwrap_or("")));
    let conf = base64::Config {
        char_set: base64::CharacterSet::Standard,
        newline: base64::Newline::LF,
        pad: true,
        line_length: None,
    };
    debug!("Signature base string: {}", base);
    debug!("Authorization header: Authorization: {}", base);
    hmac_sha1(key.as_bytes(), base.as_bytes()).code().to_base64(conf)
}

/// Constuct plain-text header
fn header(param: &ParamList) -> String {
    let mut pairs = param.iter()
                         .filter(|&(k, _)| k.starts_with("oauth_"))
                         .map(|(k, v)| format!("{}=\"{}\"", k, encode(&v)))
                         .collect::<Vec<_>>();
    pairs.sort();
    format!("OAuth {}", pairs.join(", "))
}

/// Construct plain-text body from 'PaaramList'
fn body(param: &ParamList) -> String {
    let mut pairs = param.iter()
                         .filter(|&(k, _)| !k.starts_with("oauth_"))
                         .map(|(k, v)| format!("{}={}", k, encode(&v)))
                         .collect::<Vec<_>>();
    pairs.sort();
    format!("{}", pairs.join("&"))
}

/// Create header and body
fn get_header(method: Method,
              uri: &str,
              consumer: &Token,
              token: Option<&Token>,
              other_param: Option<&ParamList>)
              -> (String, String) {
    let mut param = HashMap::new();
    let timestamp = format!("{}", time::now_utc().to_timespec().sec);
    let nonce = rand::thread_rng().gen_ascii_chars().take(32).collect::<String>();

    let _ = insert_param(&mut param, "oauth_consumer_key", consumer.key.to_string());
    let _ = insert_param(&mut param, "oauth_nonce", nonce);
    let _ = insert_param(&mut param, "oauth_signature_method", "HMAC-SHA1");
    let _ = insert_param(&mut param, "oauth_timestamp", timestamp);
    let _ = insert_param(&mut param, "oauth_version", "1.0");
    if let Some(tk) = token {
        let _ = insert_param(&mut param, "oauth_token", tk.key.as_ref());
    }

    if let Some(ps) = other_param {
        for (k, v) in ps.iter() {
            let _ = insert_param(&mut param, k.as_ref(), v.as_ref());
        }
    }

    let method_str = match method {
        Method::Options => "OPTIONS",
        Method::Get => "GET",
        Method::Head => "HEAD",
        Method::Post => "POST",
        Method::Put => "PUT",
        Method::Patch => "PATCH",
        Method::Delete => "DELETE",
        Method::Trace => "TRACE",
        Method::Connect => "CONNECT",
    };

    let sign = signature(method_str,
                         uri,
                         join_query(&param).as_ref(),
                         consumer.secret.as_ref(),
                         token.map(|t| t.secret.as_ref()));
    let _ = insert_param(&mut param, "oauth_signature", sign);

    (header(&param), body(&param))
}

/// Create an authorization header.
/// See https://dev.twitter.com/oauth/overview/authorizing-requests
///
///# Examples
///
///```
///let header = oauth::authorization_header(Method::Get, api::REQUEST_TOKEN, consumer, None, None);
///```
pub fn authorization_header(method: Method,
                            uri: &str,
                            consumer: &Token,
                            token: Option<&Token>,
                            other_param: Option<&ParamList>)
                            -> String {
    get_header(method, uri, consumer, token, other_param).0
}

/// Send authorized GET request to the specified URL.
/// `consumer` is a consumer token.
///
///# Examples
///
///```
///let REQUEST_TOKEN: &'static str = "http://oauthbin.com/v1/request-token"
///let bytes = oauth::get(REQUEST_TOKEN, consumer, None, None).unwrap();
///let resp = String::from_utf8(bytes).unwrap();
///```
pub fn get(uri: &str,
           consumer: &Token,
           token: Option<&Token>,
           other_param: Option<&ParamList>)
           -> Result<Vec<u8>, Error> {
    let (header, body) = get_header(Method::Get, uri, consumer, token, other_param);
    let req_uri = if body.len() > 0 {
        format!("{}?{}", uri, body)
    } else {
        format!("{}", uri)
    };
    let resp = try!(http::handle()
                        .get(req_uri)
                        .header("Authorization", header.as_ref())
                        .exec());
    debug!("{}", resp);
    if resp.get_code() != 200 {
        return Err(Error::HttpStatus(resp));
    }
    Ok(resp.move_body())
}

/// Send authorized POST request to the specified URL.
/// `consumer` is a consumer token.
///
///# Examples
///
///```
///let ACCESS_TOKEN: &'static str = "http://oauthbin.com/v1/access-token"
///let bytes = oauth::post(ACCESS_TOKEN, consumer, Some(request), None).unwrap();
///let resp = String::from_utf8(bytes).unwrap();
///```
pub fn post(uri: &str,
            consumer: &Token,
            token: Option<&Token>,
            other_param: Option<&ParamList>)
            -> Result<Vec<u8>, Error> {
    let (header, body) = get_header(Method::Post, uri, consumer, token, other_param);
    let resp = try!(http::handle()
                        .post(uri, &body)
                        .header("Authorization", header.as_ref())
                        .content_type("application/x-www-form-urlencoded")
                        .exec());
    debug!("{}", resp);
    if resp.get_code() != 200 {
        return Err(Error::HttpStatus(resp));
    }
    Ok(resp.move_body())
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    #[test]
    fn query() {
        let mut map = HashMap::new();
        let _ = map.insert("aaa".into(), "AAA".into());
        let _ = map.insert("bbbb".into(), "BBBB".into());
        let query = super::join_query(&map);
        assert_eq!("aaa=AAA&bbbb=BBBB", query);
    }
}
<|endoftext|>


<commit_before>#![feature(vec_push_all)]
#![feature(test)]

extern crate crypto;
extern crate rustc_serialize;
extern crate test;

use crypto::mac::{Mac, MacResult};
use crypto::sha2::Sha256;
use crypto::hmac::Hmac;
use rustc_serialize::base64;

pub const ALMOND_HASH_SEED : &'static [u8; 32] = b"this_is_a_bit_of_arbitrary_data!";

pub struct Almond {
    hash: [u8; 32],
    caveats: Vec<Vec<u8>>,
    generation: u8,
    almond_type: Vec<u8>,
}

impl Almond {
    pub fn create(key: &[u8], generation: u8, almond_type: Vec<u8>) -> Almond {
        let mut almond = Almond {
            hash: *ALMOND_HASH_SEED,
            caveats: Vec::new(),
            generation: generation,
            almond_type: almond_type,
        };

        add_to_hash(&mut almond.hash, key);
        add_to_hash(&mut almond.hash, &[generation]);
        add_to_hash(&mut almond.hash, &almond.almond_type[..]);

        almond
    }

    pub fn parse_and_verify(key: &[u8], input: &[u8]) -> Result<Almond, &'static str> {
        if input.len() < 34 {
            return Err("Too small");
        }

        let hash = &input[..32];
        let generation = input[32];

        let mut split_it = input[33..].split(|c| *c == b'\n');

        let almond_type = try!(split_it.next().ok_or("No initial newline"));

        let mut almond = Almond::create(key, generation, almond_type.to_vec());

        for caveat in split_it {
            almond.add_caveat(caveat.to_vec());
        }

        // Always compare hashes using equality operators that are
        // resistent to timing attacks.
        if MacResult::new(hash) == MacResult::new(almond.hash()) {
            Ok(almond)
        } else {
            Err("Hash do not match")
        }
    }

    pub fn add_caveat(&mut self, caveat: Vec<u8>) -> &mut Self {
        add_to_hash(&mut self.hash, &caveat[..]);
        self.caveats.push(caveat);
        self
    }

    pub fn almond_type(&self) -> &[u8] {
        &self.almond_type[..]
    }

    pub fn generation(&self) -> u8 {
        self.generation
    }

    pub fn caveats(&self) -> &[Vec<u8>] {
        &self.caveats[..]
    }

    // pub fn caveat_map(&self) -> HashMap<&[u8], Option<&[u8]>> {
    //     let mut map = HashMap::with_capacity(self.caveats.len());
    //
    //     for caveat in &self.caveats() {
    //         let mut it = caveat.splitn(2, |c| *c == b' ');
    //
    //         // By definition this must have at least one item
    //         let key = it.next().unwrap();
    //         let value = it.next();
    //
    //         map.insert(key, value);
    //     }
    //
    //     map
    // }

    pub fn hash(&self) -> &[u8; 32] {
        &self.hash
    }

    pub fn serialize(&self) -> Vec<u8> {
        let mut result : Vec<u8> = Vec::new();

        result.push_all(&self.hash);
        result.push(self.generation);
        result.push_all(&self.almond_type);
        result.push(b'\n');

        for caveat in &self.caveats {
            result.push_all(&caveat[..]);
            result.push(b'\n');
        }

        result.pop();

        result
    }
}

impl base64::ToBase64 for Almond {
    fn to_base64(&self, config: base64::Config) -> String {
        let serialized = self.serialize();
        serialized[..].to_base64(config)
    }
}

fn add_to_hash(hash: &mut [u8], data: &[u8]) {
    let hasher = Sha256::new();
    let mut mac = Hmac::new(hasher, hash);
    mac.input(data);
    mac.raw_result(hash);
}

struct DeconstructedCaveatEntry<'a> {
    pub key: &'a [u8],
    pub value: Option<&'a [u8]>,
    pub accepted: Option<bool>,
}

pub struct Verifier<'a> {
    caveats: Vec<DeconstructedCaveatEntry<'a>>,
    reject: bool,
}

impl <'a> Verifier<'a> {
    pub fn new(almond: &'a Almond, generation: u8, almond_type: &[u8]) -> Verifier<'a> {
        let caveats = almond.caveats().iter()
            .map(
                |caveat| {
                    let mut it = caveat.splitn(2, |c| *c == b' ');

                    // By definition this must have at least one item
                    let key = it.next().unwrap();
                    let value = it.next();

                    DeconstructedCaveatEntry {
                        key: key,
                        value: value,
                        accepted: None,
                    }
                }
            )
            .collect();

        Verifier {
            caveats: caveats,
            reject: almond.generation() != generation || almond.almond_type() != almond_type
        }
    }

    pub fn allow(&mut self, key: &[u8]) -> &mut Self {
        for item in &mut self.caveats {
            if item.key == key {
                item.accepted = item.accepted.or(Some(true));
            }
        }

        self
    }

    pub fn satisfies<F>(&mut self, key: &[u8], mut predicate: F) -> &mut Self
        where F: FnMut(Option<&[u8]>) -> bool
    {
        for item in &mut self.caveats {
            if item.key == key {
                let res = predicate(item.value);
                item.accepted = Some(res && item.accepted.unwrap_or(true));
            }
        }

        self
    }

    pub fn satisfies_exact(&mut self, key: &[u8], value: Option<&[u8]>) {
        for item in &mut self.caveats {
            if item.key == key {
                let res = item.value.map(|x| &*x) == value;
                item.accepted = Some(res && item.accepted.unwrap_or(true));
            }
        }
    }

    pub fn verify(&self) -> bool {
        !self.reject && self.caveats.iter().all(|item| item.accepted.unwrap_or(false))
    }
}

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

    #[test]
    fn basic_test() {
        use rustc_serialize::base64::{ToBase64, URL_SAFE};

        let key = b"this_is_a_secret";

        let mut almond = Almond::create(key, 1, b"login".to_vec());
        almond.add_caveat(b"user erikj".to_vec());
        let r = almond.to_base64(URL_SAFE);
        assert_eq!(
            r,
            concat!(
                "yyTNYc-CAXTVkgXkNnl8wdMzBTMgHyLRS",
                "lXrjdf5Uw0BbG9naW4KdXNlciBlcmlrag",
            )
        );

        Almond::parse_and_verify(key, &almond.serialize()[..]).unwrap();
    }

    #[test]
    fn verify_test() {
        let key = b"this_is_a_secret";
        let mut almond = Almond::create(key, 1, b"login".to_vec());
        almond.add_caveat(b"user erikj".to_vec());
        almond.add_caveat(b"guest".to_vec());

        let mut v = Verifier::new(&almond, 1, b"login");
        v.allow(b"user");
        assert!(!v.verify());

        v.satisfies_exact(b"guest", None);
        assert!(v.verify());

        v.satisfies_exact(b"user", Some(b"noterikj"));
        assert!(!v.verify());
    }

    #[test]
    fn verify_test_gen_type() {
        let key = b"this_is_a_secret";
        let mut almond = Almond::create(key, 1, b"login".to_vec());
        almond.add_caveat(b"user erikj".to_vec());

        let mut v = Verifier::new(&almond, 2, b"login");
        v.allow(b"user");
        assert!(!v.verify());

        let mut v = Verifier::new(&almond, 1, b"notlogin");
        v.allow(b"user");
        assert!(!v.verify());
    }

    #[bench]
    fn create(b: &mut Bencher) {
        let key = b"this_is_a_secret";
        b.iter(|| {
            let mut almond = Almond::create(key, 1, b"login".to_vec());
            almond.add_caveat(b"user erikj".to_vec());
            almond.add_caveat(b"fooo bar".to_vec());
            almond.add_caveat(b"testing".to_vec());
            almond.add_caveat(b"teeeeeeeeeeeest".to_vec());
            almond
        });
    }
}
<commit_msg>Add another benchmark. Clean things up<commit_after>#![feature(vec_push_all)]
#![feature(test)]

extern crate crypto;
extern crate rustc_serialize;
extern crate test;

use crypto::mac::{Mac, MacResult};
use crypto::sha2::Sha256;
use crypto::hmac::Hmac;
use rustc_serialize::base64;

pub const ALMOND_HASH_SEED : &'static [u8; 32] = b"this_is_a_bit_of_arbitrary_data!";

pub struct Almond {
    hash: [u8; 32],
    caveats: Vec<Vec<u8>>,
    generation: u8,
    almond_type: Vec<u8>,
}

impl Almond {
    pub fn create(key: &[u8], generation: u8, almond_type: Vec<u8>) -> Almond {
        let mut almond = Almond {
            hash: *ALMOND_HASH_SEED,
            caveats: Vec::new(),
            generation: generation,
            almond_type: almond_type,
        };

        add_to_hash(&mut almond.hash, key);
        add_to_hash(&mut almond.hash, &[generation]);
        add_to_hash(&mut almond.hash, &almond.almond_type[..]);

        almond
    }

    pub fn parse_and_verify(key: &[u8], input: &[u8]) -> Result<Almond, &'static str> {
        if input.len() < 34 {
            return Err("Too small");
        }

        let hash = &input[..32];
        let generation = input[32];

        let mut split_it = input[33..].split(|c| *c == b'\n');

        let almond_type = try!(split_it.next().ok_or("No initial newline"));

        let mut almond = Almond::create(key, generation, almond_type.to_vec());

        for caveat in split_it {
            almond.add_caveat(caveat.to_vec());
        }

        // Always compare hashes using equality operators that are
        // resistent to timing attacks.
        if MacResult::new(hash) == MacResult::new(almond.hash()) {
            Ok(almond)
        } else {
            Err("Hash do not match")
        }
    }

    pub fn add_caveat(&mut self, caveat: Vec<u8>) -> &mut Self {
        add_to_hash(&mut self.hash, &caveat[..]);
        self.caveats.push(caveat);
        self
    }

    pub fn almond_type(&self) -> &[u8] {
        &self.almond_type[..]
    }

    pub fn generation(&self) -> u8 {
        self.generation
    }

    pub fn caveats(&self) -> &[Vec<u8>] {
        &self.caveats[..]
    }

    pub fn hash(&self) -> &[u8; 32] {
        &self.hash
    }

    pub fn serialize(&self) -> Vec<u8> {
        let mut result : Vec<u8> = Vec::new();

        result.push_all(&self.hash);
        result.push(self.generation);
        result.push_all(&self.almond_type);
        result.push(b'\n');

        for caveat in &self.caveats {
            result.push_all(&caveat[..]);
            result.push(b'\n');
        }

        result.pop();

        result
    }
}

impl base64::ToBase64 for Almond {
    fn to_base64(&self, config: base64::Config) -> String {
        let serialized = self.serialize();
        serialized[..].to_base64(config)
    }
}

fn add_to_hash(hash: &mut [u8], data: &[u8]) {
    let hasher = Sha256::new();
    let mut mac = Hmac::new(hasher, hash);
    mac.input(data);
    mac.raw_result(hash);
}

struct DeconstructedCaveatEntry<'a> {
    pub key: &'a [u8],
    pub value: Option<&'a [u8]>,
    pub accepted: Option<bool>,
}

pub struct Verifier<'a> {
    caveats: Vec<DeconstructedCaveatEntry<'a>>,
    reject: bool,
}

impl <'a> Verifier<'a> {
    pub fn new(almond: &'a Almond, generation: u8, almond_type: &[u8]) -> Verifier<'a> {
        let caveats = almond.caveats().iter()
            .map(
                |caveat| {
                    let mut it = caveat.splitn(2, |c| *c == b' ');

                    // By definition this must have at least one item
                    let key = it.next().unwrap();
                    let value = it.next();

                    DeconstructedCaveatEntry {
                        key: key,
                        value: value,
                        accepted: None,
                    }
                }
            )
            .collect();

        Verifier {
            caveats: caveats,
            reject: almond.generation() != generation || almond.almond_type() != almond_type
        }
    }

    pub fn allow(&mut self, key: &[u8]) -> &mut Self {
        for item in &mut self.caveats {
            if item.key == key {
                item.accepted = item.accepted.or(Some(true));
            }
        }

        self
    }

    pub fn satisfies<F>(&mut self, key: &[u8], mut predicate: F) -> &mut Self
        where F: FnMut(Option<&[u8]>) -> bool
    {
        for item in &mut self.caveats {
            if item.key == key {
                let res = predicate(item.value);
                item.accepted = Some(res && item.accepted.unwrap_or(true));
            }
        }

        self
    }

    pub fn satisfies_exact(&mut self, key: &[u8], value: Option<&[u8]>) {
        for item in &mut self.caveats {
            if item.key == key {
                let res = item.value.map(|x| &*x) == value;
                item.accepted = Some(res && item.accepted.unwrap_or(true));
            }
        }
    }

    pub fn verify(&self) -> bool {
        !self.reject && self.caveats.iter().all(|item| item.accepted.unwrap_or(false))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use test::Bencher;
    use rustc_serialize::base64::{FromBase64, ToBase64, URL_SAFE};

    #[test]
    fn basic_test() {
        let key = b"this_is_a_secret";

        let mut almond = Almond::create(key, 1, b"login".to_vec());
        almond.add_caveat(b"user erikj".to_vec());
        let r = almond.to_base64(URL_SAFE);
        assert_eq!(
            r,
            concat!(
                "yyTNYc-CAXTVkgXkNnl8wdMzBTMgHyLRS",
                "lXrjdf5Uw0BbG9naW4KdXNlciBlcmlrag",
            )
        );

        Almond::parse_and_verify(key, &almond.serialize()[..]).unwrap();
    }

    #[test]
    fn verify_test() {
        let key = b"this_is_a_secret";
        let mut almond = Almond::create(key, 1, b"login".to_vec());
        almond.add_caveat(b"user erikj".to_vec());
        almond.add_caveat(b"guest".to_vec());

        let mut v = Verifier::new(&almond, 1, b"login");
        v.allow(b"user");
        assert!(!v.verify());

        v.satisfies_exact(b"guest", None);
        assert!(v.verify());

        v.satisfies_exact(b"user", Some(b"noterikj"));
        assert!(!v.verify());
    }

    #[test]
    fn verify_test_gen_type() {
        let key = b"this_is_a_secret";
        let mut almond = Almond::create(key, 1, b"login".to_vec());
        almond.add_caveat(b"user erikj".to_vec());

        let mut v = Verifier::new(&almond, 2, b"login");
        v.allow(b"user");
        assert!(!v.verify());

        let mut v = Verifier::new(&almond, 1, b"notlogin");
        v.allow(b"user");
        assert!(!v.verify());
    }

    #[bench]
    fn create(b: &mut Bencher) {
        let key = b"this_is_a_secret";
        b.iter(|| {
            let mut almond = Almond::create(key, 1, b"login".to_vec());
            almond.add_caveat(b"user erikj".to_vec());
            almond.add_caveat(b"fooo bar".to_vec());
            almond.add_caveat(b"testing".to_vec());
            almond.add_caveat(b"teeeeeeeeeeeest".to_vec());
            almond
        });
    }

    #[bench]
    fn parse(b: &mut Bencher) {
        let key = b"this_is_a_secret";

        let b64 = concat!(
            "Rv31sT9t5d31LHPeBFjPewo0TJ1ARbDok7vOWBVNSM4BbG9naW4KdXN",
            "lciBlcmlragpmb29vIGJhcgp0ZXN0aW5nCnRlZWVlZWVlZWVlZWVzdA",
        );
        let parsed = b64.from_base64().unwrap();

        b.iter(|| {
            let almond = Almond::parse_and_verify(
                key,
                &parsed[..],
            );
            almond.unwrap();
        });
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

extern crate byteorder;
extern crate chrono;
extern crate kudu_pb;
extern crate mio;
extern crate netbuf;
extern crate parking_lot;
extern crate protobuf;
extern crate rand;
extern crate slab;
extern crate uuid;
extern crate vec_map;

#[cfg(test)] extern crate env_logger;
#[cfg(test)] extern crate tempdir;

#[cfg(any(feature="quickcheck", test))] extern crate quickcheck;

#[macro_use] extern crate log;

mod backoff;
mod bit_set;
mod client;
mod dns;
mod error;
mod key;
mod master;
mod meta_cache;
mod partition;
mod row;
mod rpc;
mod schema;
mod table;
mod tablet;
mod tablet_server;
mod util;
mod value;

#[cfg(test)]
mod mini_cluster;

pub use client::*;
pub use error::*;
pub use partition::*;
pub use row::Row;
pub use schema::*;
pub use table::*;
pub use tablet::*;
pub use tablet_server::TabletServer;
pub use value::Value;

use std::fmt;
use std::str;

use uuid::Uuid;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DataType {
    Bool,
    Int8,
    Int16,
    Int32,
    Int64,
    Timestamp,
    Float,
    Double,
    Binary,
    String,
}

impl DataType {

    fn is_var_len(self) -> bool {
        match self {
            DataType::String | DataType::Binary => true,
            _ => false,
        }
    }

    fn size(self) -> usize {
        match self {
            DataType::Bool | DataType::Int8 => 1,
            DataType::Int16 => 2,
            DataType::Int32 | DataType::Float => 4,
            DataType::Int64 | DataType::Timestamp | DataType::Double => 8,
            DataType::Binary | DataType::String => 16,
        }
    }

    fn to_pb(self) -> kudu_pb::common::DataType {
        match self {
            DataType::Bool => kudu_pb::common::DataType::BOOL,
            DataType::Int8 => kudu_pb::common::DataType::INT8,
            DataType::Int16 => kudu_pb::common::DataType::INT16,
            DataType::Int32 => kudu_pb::common::DataType::INT32,
            DataType::Int64 => kudu_pb::common::DataType::INT64,
            DataType::Timestamp => kudu_pb::common::DataType::TIMESTAMP,
            DataType::Float => kudu_pb::common::DataType::FLOAT,
            DataType::Double => kudu_pb::common::DataType::DOUBLE,
            DataType::Binary => kudu_pb::common::DataType::BINARY,
            DataType::String => kudu_pb::common::DataType::STRING,
        }
    }

    fn from_pb(pb: kudu_pb::common::DataType) -> Result<DataType> {
        match pb {
            kudu_pb::common::DataType::BOOL => Ok(DataType::Bool),
            kudu_pb::common::DataType::INT8 => Ok(DataType::Int8),
            kudu_pb::common::DataType::INT16 => Ok(DataType::Int16),
            kudu_pb::common::DataType::INT32 => Ok(DataType::Int32),
            kudu_pb::common::DataType::INT64 => Ok(DataType::Int64),
            kudu_pb::common::DataType::TIMESTAMP => Ok(DataType::Timestamp),
            kudu_pb::common::DataType::FLOAT => Ok(DataType::Float),
            kudu_pb::common::DataType::DOUBLE => Ok(DataType::Double),
            kudu_pb::common::DataType::BINARY => Ok(DataType::Binary),
            kudu_pb::common::DataType::STRING => Ok(DataType::String),
            _ => Err(Error::VersionMismatch("unknown data type".to_string())),
        }
    }

    #[cfg(any(feature="quickcheck", test))]
    pub fn arbitrary_primary_key<G>(g: &mut G) -> DataType where G: quickcheck::Gen {
        *g.choose(&[
                  DataType::Int8,
                  DataType::Int16,
                  DataType::Int32,
                  DataType::Int64,
                  DataType::Timestamp,
                  DataType::Binary,
                  DataType::String,
        ]).unwrap()
    }
}

#[cfg(any(feature="quickcheck", test))]
impl quickcheck::Arbitrary for DataType {
    fn arbitrary<G>(g: &mut G) -> DataType where G: quickcheck::Gen {
        *g.choose(&[
                  DataType::Bool,
                  DataType::Int8,
                  DataType::Int16,
                  DataType::Int32,
                  DataType::Int64,
                  DataType::Timestamp,
                  DataType::Float,
                  DataType::Double,
                  DataType::Binary,
                  DataType::String,
        ]).unwrap()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum EncodingType {
    Auto,
    Plain,
    Prefix,
    GroupVarint,
    RunLength,
    Dictionary,
    BitShuffle,
}

impl EncodingType {
    fn to_pb(self) -> kudu_pb::common::EncodingType {
        match self {
            EncodingType::Auto => kudu_pb::common::EncodingType::AUTO_ENCODING,
            EncodingType::Plain => kudu_pb::common::EncodingType::PLAIN_ENCODING,
            EncodingType::Prefix => kudu_pb::common::EncodingType::PREFIX_ENCODING,
            EncodingType::GroupVarint => kudu_pb::common::EncodingType::GROUP_VARINT,
            EncodingType::RunLength => kudu_pb::common::EncodingType::RLE,
            EncodingType::Dictionary => kudu_pb::common::EncodingType::DICT_ENCODING,
            EncodingType::BitShuffle => kudu_pb::common::EncodingType::BIT_SHUFFLE,
        }
    }

    fn from_pb(pb: kudu_pb::common::EncodingType) -> Result<EncodingType> {
        match pb {
            kudu_pb::common::EncodingType::AUTO_ENCODING => Ok(EncodingType::Auto),
            kudu_pb::common::EncodingType::PLAIN_ENCODING => Ok(EncodingType::Plain),
            kudu_pb::common::EncodingType::PREFIX_ENCODING => Ok(EncodingType::Prefix),
            kudu_pb::common::EncodingType::GROUP_VARINT => Ok(EncodingType::GroupVarint),
            kudu_pb::common::EncodingType::RLE => Ok(EncodingType::RunLength),
            kudu_pb::common::EncodingType::DICT_ENCODING => Ok(EncodingType::Dictionary),
            kudu_pb::common::EncodingType::BIT_SHUFFLE => Ok(EncodingType::BitShuffle),
            _ => Err(Error::VersionMismatch("unknown encoding type".to_string())),
        }
    }

    #[cfg(any(feature="quickcheck", test))]
    pub fn arbitrary<G>(g: &mut G, data_type: DataType) -> EncodingType where G: quickcheck::Gen {
        match data_type {
            DataType::Bool => *g.choose(&[
                EncodingType::Auto,
                EncodingType::Plain,
                EncodingType::RunLength
            ]).unwrap(),
            DataType::Int8 | DataType::Int16 |
            DataType::Int32 | DataType::Int64 | DataType::Timestamp => *g.choose(&[
                EncodingType::Auto,
                EncodingType::Plain,
                EncodingType::RunLength,
                EncodingType::BitShuffle
            ]).unwrap(),
            DataType::Float | DataType::Double => *g.choose(&[
                EncodingType::Auto,
                EncodingType::Plain,
                EncodingType::BitShuffle
            ]).unwrap(),
            DataType::Binary | DataType::String => *g.choose(&[
                EncodingType::Auto,
                EncodingType::Plain,
                EncodingType::Prefix,
                EncodingType::Dictionary
            ]).unwrap(),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum CompressionType {
    Default,
    None,
    Snappy,
    Lz4,
    Zlib,
}

impl CompressionType {
    fn to_pb(self) -> kudu_pb::common::CompressionType {
        match self {
            CompressionType::Default => kudu_pb::common::CompressionType::DEFAULT_COMPRESSION,
            CompressionType::None => kudu_pb::common::CompressionType::NO_COMPRESSION,
            CompressionType::Snappy => kudu_pb::common::CompressionType::SNAPPY,
            CompressionType::Lz4 => kudu_pb::common::CompressionType::LZ4,
            CompressionType::Zlib => kudu_pb::common::CompressionType::ZLIB,
        }
    }

    fn from_pb(pb: kudu_pb::common::CompressionType) -> Result<CompressionType> {
        match pb {
            kudu_pb::common::CompressionType::DEFAULT_COMPRESSION => Ok(CompressionType::Default),
            kudu_pb::common::CompressionType::NO_COMPRESSION => Ok(CompressionType::None),
            kudu_pb::common::CompressionType::SNAPPY => Ok(CompressionType::Snappy),
            kudu_pb::common::CompressionType::LZ4 => Ok(CompressionType::Lz4),
            kudu_pb::common::CompressionType::ZLIB => Ok(CompressionType::Zlib),
            _ => Err(Error::VersionMismatch("unknown compression type".to_string())),
        }
    }
}

#[cfg(any(feature="quickcheck", test))]
impl quickcheck::Arbitrary for CompressionType {
    fn arbitrary<G>(g: &mut G) -> CompressionType where G: quickcheck::Gen {
        *g.choose(&[
                  CompressionType::Default,
                  CompressionType::None,
                  CompressionType::Snappy,
                  CompressionType::Lz4,
                  CompressionType::Zlib,
        ]).unwrap()
    }
}

/// The role of a Kudu master or tserver in a consensus group.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RaftRole {
    Follower,
    Leader,
    Learner,
    NonParticipant,
    Unknown,
}

impl RaftRole {
    fn from_pb(pb: kudu_pb::consensus_metadata::RaftPeerPB_Role) -> RaftRole {
        match pb {
            kudu_pb::consensus_metadata::RaftPeerPB_Role::UNKNOWN_ROLE => RaftRole::Unknown,
            kudu_pb::consensus_metadata::RaftPeerPB_Role::FOLLOWER => RaftRole::Follower,
            kudu_pb::consensus_metadata::RaftPeerPB_Role::LEADER => RaftRole::Leader,
            kudu_pb::consensus_metadata::RaftPeerPB_Role::LEARNER => RaftRole::Learner,
            kudu_pb::consensus_metadata::RaftPeerPB_Role::NON_PARTICIPANT => RaftRole::NonParticipant,
        }
    }
}

macro_rules! id {
    ($id:ident) => {
        #[derive(Clone, PartialEq, Eq, Hash)]
        pub struct $id {
            id: Uuid,
        }

        impl $id {
            pub fn as_bytes(&self) -> &[u8; 16] {
                self.id.as_bytes()
            }

            fn parse(input: &str) -> Result<$id> {
                Uuid::parse_str(input).map_err(|error| Error::Serialization(format!("{}", error)))
                                      .map(|id| $id { id: id })
            }

            fn parse_bytes(input: &[u8]) -> Result<$id> {
                str::from_utf8(input)
                    .map_err(|error| Error::Serialization(format!("{}", error)))
                    .and_then($id::parse)
            }
        }

        impl fmt::Debug for $id {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                write!(f, "{}", self.id.simple())
            }
        }

        impl fmt::Display for $id {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                write!(f, "{}", self.id.simple())
            }
        }
    }
}

id!(MasterId);
id!(ReplicaId);
id!(TableId);
id!(TabletId);
id!(TabletServerId);
<commit_msg>expose Master<commit_after>#![allow(dead_code)]

extern crate byteorder;
extern crate chrono;
extern crate kudu_pb;
extern crate mio;
extern crate netbuf;
extern crate parking_lot;
extern crate protobuf;
extern crate rand;
extern crate slab;
extern crate uuid;
extern crate vec_map;

#[cfg(test)] extern crate env_logger;
#[cfg(test)] extern crate tempdir;

#[cfg(any(feature="quickcheck", test))] extern crate quickcheck;

#[macro_use] extern crate log;

mod backoff;
mod bit_set;
mod client;
mod dns;
mod error;
mod key;
mod master;
mod meta_cache;
mod partition;
mod row;
mod rpc;
mod schema;
mod table;
mod tablet;
mod tablet_server;
mod util;
mod value;

#[cfg(test)]
mod mini_cluster;

pub use client::*;
pub use error::*;
pub use master::Master;
pub use partition::*;
pub use row::Row;
pub use schema::*;
pub use table::*;
pub use tablet::*;
pub use tablet_server::TabletServer;
pub use value::Value;

use std::fmt;
use std::str;

use uuid::Uuid;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DataType {
    Bool,
    Int8,
    Int16,
    Int32,
    Int64,
    Timestamp,
    Float,
    Double,
    Binary,
    String,
}

impl DataType {

    fn is_var_len(self) -> bool {
        match self {
            DataType::String | DataType::Binary => true,
            _ => false,
        }
    }

    fn size(self) -> usize {
        match self {
            DataType::Bool | DataType::Int8 => 1,
            DataType::Int16 => 2,
            DataType::Int32 | DataType::Float => 4,
            DataType::Int64 | DataType::Timestamp | DataType::Double => 8,
            DataType::Binary | DataType::String => 16,
        }
    }

    fn to_pb(self) -> kudu_pb::common::DataType {
        match self {
            DataType::Bool => kudu_pb::common::DataType::BOOL,
            DataType::Int8 => kudu_pb::common::DataType::INT8,
            DataType::Int16 => kudu_pb::common::DataType::INT16,
            DataType::Int32 => kudu_pb::common::DataType::INT32,
            DataType::Int64 => kudu_pb::common::DataType::INT64,
            DataType::Timestamp => kudu_pb::common::DataType::TIMESTAMP,
            DataType::Float => kudu_pb::common::DataType::FLOAT,
            DataType::Double => kudu_pb::common::DataType::DOUBLE,
            DataType::Binary => kudu_pb::common::DataType::BINARY,
            DataType::String => kudu_pb::common::DataType::STRING,
        }
    }

    fn from_pb(pb: kudu_pb::common::DataType) -> Result<DataType> {
        match pb {
            kudu_pb::common::DataType::BOOL => Ok(DataType::Bool),
            kudu_pb::common::DataType::INT8 => Ok(DataType::Int8),
            kudu_pb::common::DataType::INT16 => Ok(DataType::Int16),
            kudu_pb::common::DataType::INT32 => Ok(DataType::Int32),
            kudu_pb::common::DataType::INT64 => Ok(DataType::Int64),
            kudu_pb::common::DataType::TIMESTAMP => Ok(DataType::Timestamp),
            kudu_pb::common::DataType::FLOAT => Ok(DataType::Float),
            kudu_pb::common::DataType::DOUBLE => Ok(DataType::Double),
            kudu_pb::common::DataType::BINARY => Ok(DataType::Binary),
            kudu_pb::common::DataType::STRING => Ok(DataType::String),
            _ => Err(Error::VersionMismatch("unknown data type".to_string())),
        }
    }

    #[cfg(any(feature="quickcheck", test))]
    pub fn arbitrary_primary_key<G>(g: &mut G) -> DataType where G: quickcheck::Gen {
        *g.choose(&[
                  DataType::Int8,
                  DataType::Int16,
                  DataType::Int32,
                  DataType::Int64,
                  DataType::Timestamp,
                  DataType::Binary,
                  DataType::String,
        ]).unwrap()
    }
}

#[cfg(any(feature="quickcheck", test))]
impl quickcheck::Arbitrary for DataType {
    fn arbitrary<G>(g: &mut G) -> DataType where G: quickcheck::Gen {
        *g.choose(&[
                  DataType::Bool,
                  DataType::Int8,
                  DataType::Int16,
                  DataType::Int32,
                  DataType::Int64,
                  DataType::Timestamp,
                  DataType::Float,
                  DataType::Double,
                  DataType::Binary,
                  DataType::String,
        ]).unwrap()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum EncodingType {
    Auto,
    Plain,
    Prefix,
    GroupVarint,
    RunLength,
    Dictionary,
    BitShuffle,
}

impl EncodingType {
    fn to_pb(self) -> kudu_pb::common::EncodingType {
        match self {
            EncodingType::Auto => kudu_pb::common::EncodingType::AUTO_ENCODING,
            EncodingType::Plain => kudu_pb::common::EncodingType::PLAIN_ENCODING,
            EncodingType::Prefix => kudu_pb::common::EncodingType::PREFIX_ENCODING,
            EncodingType::GroupVarint => kudu_pb::common::EncodingType::GROUP_VARINT,
            EncodingType::RunLength => kudu_pb::common::EncodingType::RLE,
            EncodingType::Dictionary => kudu_pb::common::EncodingType::DICT_ENCODING,
            EncodingType::BitShuffle => kudu_pb::common::EncodingType::BIT_SHUFFLE,
        }
    }

    fn from_pb(pb: kudu_pb::common::EncodingType) -> Result<EncodingType> {
        match pb {
            kudu_pb::common::EncodingType::AUTO_ENCODING => Ok(EncodingType::Auto),
            kudu_pb::common::EncodingType::PLAIN_ENCODING => Ok(EncodingType::Plain),
            kudu_pb::common::EncodingType::PREFIX_ENCODING => Ok(EncodingType::Prefix),
            kudu_pb::common::EncodingType::GROUP_VARINT => Ok(EncodingType::GroupVarint),
            kudu_pb::common::EncodingType::RLE => Ok(EncodingType::RunLength),
            kudu_pb::common::EncodingType::DICT_ENCODING => Ok(EncodingType::Dictionary),
            kudu_pb::common::EncodingType::BIT_SHUFFLE => Ok(EncodingType::BitShuffle),
            _ => Err(Error::VersionMismatch("unknown encoding type".to_string())),
        }
    }

    #[cfg(any(feature="quickcheck", test))]
    pub fn arbitrary<G>(g: &mut G, data_type: DataType) -> EncodingType where G: quickcheck::Gen {
        match data_type {
            DataType::Bool => *g.choose(&[
                EncodingType::Auto,
                EncodingType::Plain,
                EncodingType::RunLength
            ]).unwrap(),
            DataType::Int8 | DataType::Int16 |
            DataType::Int32 | DataType::Int64 | DataType::Timestamp => *g.choose(&[
                EncodingType::Auto,
                EncodingType::Plain,
                EncodingType::RunLength,
                EncodingType::BitShuffle
            ]).unwrap(),
            DataType::Float | DataType::Double => *g.choose(&[
                EncodingType::Auto,
                EncodingType::Plain,
                EncodingType::BitShuffle
            ]).unwrap(),
            DataType::Binary | DataType::String => *g.choose(&[
                EncodingType::Auto,
                EncodingType::Plain,
                EncodingType::Prefix,
                EncodingType::Dictionary
            ]).unwrap(),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum CompressionType {
    Default,
    None,
    Snappy,
    Lz4,
    Zlib,
}

impl CompressionType {
    fn to_pb(self) -> kudu_pb::common::CompressionType {
        match self {
            CompressionType::Default => kudu_pb::common::CompressionType::DEFAULT_COMPRESSION,
            CompressionType::None => kudu_pb::common::CompressionType::NO_COMPRESSION,
            CompressionType::Snappy => kudu_pb::common::CompressionType::SNAPPY,
            CompressionType::Lz4 => kudu_pb::common::CompressionType::LZ4,
            CompressionType::Zlib => kudu_pb::common::CompressionType::ZLIB,
        }
    }

    fn from_pb(pb: kudu_pb::common::CompressionType) -> Result<CompressionType> {
        match pb {
            kudu_pb::common::CompressionType::DEFAULT_COMPRESSION => Ok(CompressionType::Default),
            kudu_pb::common::CompressionType::NO_COMPRESSION => Ok(CompressionType::None),
            kudu_pb::common::CompressionType::SNAPPY => Ok(CompressionType::Snappy),
            kudu_pb::common::CompressionType::LZ4 => Ok(CompressionType::Lz4),
            kudu_pb::common::CompressionType::ZLIB => Ok(CompressionType::Zlib),
            _ => Err(Error::VersionMismatch("unknown compression type".to_string())),
        }
    }
}

#[cfg(any(feature="quickcheck", test))]
impl quickcheck::Arbitrary for CompressionType {
    fn arbitrary<G>(g: &mut G) -> CompressionType where G: quickcheck::Gen {
        *g.choose(&[
                  CompressionType::Default,
                  CompressionType::None,
                  CompressionType::Snappy,
                  CompressionType::Lz4,
                  CompressionType::Zlib,
        ]).unwrap()
    }
}

/// The role of a Kudu master or tserver in a consensus group.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RaftRole {
    Follower,
    Leader,
    Learner,
    NonParticipant,
    Unknown,
}

impl RaftRole {
    fn from_pb(pb: kudu_pb::consensus_metadata::RaftPeerPB_Role) -> RaftRole {
        match pb {
            kudu_pb::consensus_metadata::RaftPeerPB_Role::UNKNOWN_ROLE => RaftRole::Unknown,
            kudu_pb::consensus_metadata::RaftPeerPB_Role::FOLLOWER => RaftRole::Follower,
            kudu_pb::consensus_metadata::RaftPeerPB_Role::LEADER => RaftRole::Leader,
            kudu_pb::consensus_metadata::RaftPeerPB_Role::LEARNER => RaftRole::Learner,
            kudu_pb::consensus_metadata::RaftPeerPB_Role::NON_PARTICIPANT => RaftRole::NonParticipant,
        }
    }
}

macro_rules! id {
    ($id:ident) => {
        #[derive(Clone, PartialEq, Eq, Hash)]
        pub struct $id {
            id: Uuid,
        }

        impl $id {
            pub fn as_bytes(&self) -> &[u8; 16] {
                self.id.as_bytes()
            }

            fn parse(input: &str) -> Result<$id> {
                Uuid::parse_str(input).map_err(|error| Error::Serialization(format!("{}", error)))
                                      .map(|id| $id { id: id })
            }

            fn parse_bytes(input: &[u8]) -> Result<$id> {
                str::from_utf8(input)
                    .map_err(|error| Error::Serialization(format!("{}", error)))
                    .and_then($id::parse)
            }
        }

        impl fmt::Debug for $id {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                write!(f, "{}", self.id.simple())
            }
        }

        impl fmt::Display for $id {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                write!(f, "{}", self.id.simple())
            }
        }
    }
}

id!(MasterId);
id!(ReplicaId);
id!(TableId);
id!(TabletId);
id!(TabletServerId);
<|endoftext|>


<commit_before>//! Parse and write Iridium Short Burst Data (SBD) messages.
//!
//! # Background
//!
//! Iridium is both a
//! [satellite constellation](https://en.wikipedia.org/wiki/Iridium_satellite_constellation)
//! and a [company](https://en.wikipedia.org/wiki/Iridium_Communications) that provides satellite
//! communications. The Iridium network is used by phones, modems, and other communication devices.
//!
//! One mode of transmitting data over the Iridium network is via Short Burst Data (SBD) messages.
//! These messages carry a payload of some small number of bytes, usually less than one thousand.
//! Messages can be Mobile-Originated (MO), meaning that they are sent *from* an Iridium modem, or
//! Mobile-Terminated (MT), meaning that the are sent *to* an Iridium modem. Mobile-Originated
//! messages are delivered either to an email address via MIME attachment, or directly to a given
//! IP address and port via TCP; this second method is called `DirectIP`.
//!
//! # Usage
//!
//! This is a simple library for reading mobile originated SBD messages from a stream, decoding
//! their headers and data payloads, and writing them back to a stream. This library does not
//! handle mobile terminated messages.
//!
//! MO messages can be read from a byte stream:
//!
//! ```
//! let mut file = std::fs::File::open("data/0-mo.sbd").unwrap();
//! let message = sbd::mo::Message::read_from(file).unwrap();
//! ```
//!
//! To receive MO messages via `DirectIP`, a server is provided.
//! This server will listen for incoming messages forever, storing them in a `Storage`:
//!
//! ```no_run
//! let storage = sbd::storage::FilesystemStorage::open("/var/iridium").unwrap();
//! let mut server = sbd::directip::Server::new("0.0.0.0:10800", storage);
//! server.serve_forever();
//! ```
//!
//! Most of the functionality of this library is exposed by a single executable, named `sbd`.  Use
//! the `sbd` executable to inspect raw sbd files stores on a filesystem, interrogate sbd files on a
//! filesystem, and start that forever-running server to receive Iridium SBD `DirectIP` messages.

#![deny(missing_copy_implementations, missing_debug_implementations, missing_docs, trivial_casts,
        trivial_numeric_casts, unsafe_code, unstable_features, unused_extern_crates,
        unused_import_braces, unused_qualifications)]
#![recursion_limit="128"]

extern crate byteorder;
extern crate chrono;
#[macro_use]
extern crate log;
#[macro_use]
extern crate quick_error;
#[macro_use]
extern crate serde_derive;
extern crate walkdir;

pub mod directip;
pub mod mo;
pub mod storage;

/// Create-specific `Result`.
pub type Result<T> = std::result::Result<T, Error>;

quick_error! {
    /// Crate-specific errors
    #[derive(Debug)]
    pub enum Error {
        /// A wrapper around a `std::io::Error`.
        Io(err: std::io::Error) {
            from()
            cause(err)
            description(err.description())
            display("io error: {}", err)
        }
        /// Invalid protocol revision number.
        InvalidProtocolRevisionNumber(n: u8) {
            description("invalid protocol revision number")
            display("invalid protocol revision number: {}", n)
        }
        /// Invalid information element identifier.
        InvalidInformationElementIdentifier(n: u8) {
            description("invalid information element identifier")
            display("invalid information element identifier: {}", n)
        }
        /// The timestamp is negative, but only positive ones are supported.
        NegativeTimestamp(timestamp: i64) {
            description("only positive timestamps are allowed in mo messages")
            display("negative timestamp: {}", timestamp)
        }
        /// No header on a MO message.
        NoHeader {
            description("no header on a mo message")
        }
        /// No payload on a MO message.
        NoPayload {
            description("no payload on a mo message")
        }
        /// We expected a directory, but this isn't one.
        ///
        /// TODO can this be a PathBuf?
        NotADirectory(s: std::ffi::OsString) {
            description("the os string is not a directory")
            display("this os string is not a directory: {}", s.to_string_lossy())
        }
        /// The overall message length is too long.
        OverallMessageLength(len: usize) {
            description("the overall message length is too long")
            display("the overall message length is too long: {}", len)
        }
        /// The payload is too long.
        PayloadTooLong(len: usize) {
            description("the mo payload is too long")
            display("the payload is too long: {}", len)
        }
        /// Two headers in an MO message.
        TwoHeaders {
            description("two headers in a MO message")
        }
        /// Two payloads in an MO message.
        TwoPayloads {
            description("two payloads in a MO message")
        }
        /// Wrapper around `std::str::Utf8Error`.
        Utf8(err: std::str::Utf8Error) {
            from()
            cause(err)
            description(err.description())
            display("utf8 error: {}", err)
        }
        /// The session status is unknown.
        UnknownSessionStatus(n: u8) {
            description("unknown session status")
            display("uknown session status code: {}", n)
        }
        /// Wrapper around `walkdir::Error`.
        WalkDir(err: walkdir::Error) {
            from()
            cause(err)
            description(err.description())
            display("walkdir error: {}", err)
        }
    }
}
<commit_msg>Reformat the module-level configuration<commit_after>//! Parse and write Iridium Short Burst Data (SBD) messages.
//!
//! # Background
//!
//! Iridium is both a
//! [satellite constellation](https://en.wikipedia.org/wiki/Iridium_satellite_constellation)
//! and a [company](https://en.wikipedia.org/wiki/Iridium_Communications) that provides satellite
//! communications. The Iridium network is used by phones, modems, and other communication devices.
//!
//! One mode of transmitting data over the Iridium network is via Short Burst Data (SBD) messages.
//! These messages carry a payload of some small number of bytes, usually less than one thousand.
//! Messages can be Mobile-Originated (MO), meaning that they are sent *from* an Iridium modem, or
//! Mobile-Terminated (MT), meaning that the are sent *to* an Iridium modem. Mobile-Originated
//! messages are delivered either to an email address via MIME attachment, or directly to a given
//! IP address and port via TCP; this second method is called `DirectIP`.
//!
//! # Usage
//!
//! This is a simple library for reading mobile originated SBD messages from a stream, decoding
//! their headers and data payloads, and writing them back to a stream. This library does not
//! handle mobile terminated messages.
//!
//! MO messages can be read from a byte stream:
//!
//! ```
//! let mut file = std::fs::File::open("data/0-mo.sbd").unwrap();
//! let message = sbd::mo::Message::read_from(file).unwrap();
//! ```
//!
//! To receive MO messages via `DirectIP`, a server is provided.
//! This server will listen for incoming messages forever, storing them in a `Storage`:
//!
//! ```no_run
//! let storage = sbd::storage::FilesystemStorage::open("/var/iridium").unwrap();
//! let mut server = sbd::directip::Server::new("0.0.0.0:10800", storage);
//! server.serve_forever();
//! ```
//!
//! Most of the functionality of this library is exposed by a single executable, named `sbd`.  Use
//! the `sbd` executable to inspect raw sbd files stores on a filesystem, interrogate sbd files on a
//! filesystem, and start that forever-running server to receive Iridium SBD `DirectIP` messages.

#![deny(
    missing_copy_implementations, missing_debug_implementations, missing_docs, trivial_casts,
    trivial_numeric_casts, unsafe_code, unstable_features, unused_extern_crates,
    unused_import_braces, unused_qualifications
)]
#![recursion_limit = "128"]

extern crate byteorder;
extern crate chrono;
#[macro_use]
extern crate log;
#[macro_use]
extern crate quick_error;
#[macro_use]
extern crate serde_derive;
extern crate walkdir;

pub mod directip;
pub mod mo;
pub mod storage;

/// Create-specific `Result`.
pub type Result<T> = std::result::Result<T, Error>;

quick_error! {
    /// Crate-specific errors
    #[derive(Debug)]
    pub enum Error {
        /// A wrapper around a `std::io::Error`.
        Io(err: std::io::Error) {
            from()
            cause(err)
            description(err.description())
            display("io error: {}", err)
        }
        /// Invalid protocol revision number.
        InvalidProtocolRevisionNumber(n: u8) {
            description("invalid protocol revision number")
            display("invalid protocol revision number: {}", n)
        }
        /// Invalid information element identifier.
        InvalidInformationElementIdentifier(n: u8) {
            description("invalid information element identifier")
            display("invalid information element identifier: {}", n)
        }
        /// The timestamp is negative, but only positive ones are supported.
        NegativeTimestamp(timestamp: i64) {
            description("only positive timestamps are allowed in mo messages")
            display("negative timestamp: {}", timestamp)
        }
        /// No header on a MO message.
        NoHeader {
            description("no header on a mo message")
        }
        /// No payload on a MO message.
        NoPayload {
            description("no payload on a mo message")
        }
        /// We expected a directory, but this isn't one.
        ///
        /// TODO can this be a PathBuf?
        NotADirectory(s: std::ffi::OsString) {
            description("the os string is not a directory")
            display("this os string is not a directory: {}", s.to_string_lossy())
        }
        /// The overall message length is too long.
        OverallMessageLength(len: usize) {
            description("the overall message length is too long")
            display("the overall message length is too long: {}", len)
        }
        /// The payload is too long.
        PayloadTooLong(len: usize) {
            description("the mo payload is too long")
            display("the payload is too long: {}", len)
        }
        /// Two headers in an MO message.
        TwoHeaders {
            description("two headers in a MO message")
        }
        /// Two payloads in an MO message.
        TwoPayloads {
            description("two payloads in a MO message")
        }
        /// Wrapper around `std::str::Utf8Error`.
        Utf8(err: std::str::Utf8Error) {
            from()
            cause(err)
            description(err.description())
            display("utf8 error: {}", err)
        }
        /// The session status is unknown.
        UnknownSessionStatus(n: u8) {
            description("unknown session status")
            display("uknown session status code: {}", n)
        }
        /// Wrapper around `walkdir::Error`.
        WalkDir(err: walkdir::Error) {
            from()
            cause(err)
            description(err.description())
            display("walkdir error: {}", err)
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use] extern crate quick_error;
#[macro_use] extern crate string_cache;
extern crate mime;
extern crate hyper;
extern crate url;
extern crate kuchiki;
extern crate image;

mod strategies;
mod util;

use image::GenericImage;
use strategies::Strategy;
use std::io::Read;
use util::AsImageFormat;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        Hyper(error: hyper::Error) { from() }
        Io(error: std::io::Error) { from() }
        Image(error: image::ImageError) { from() }
        Other(msg: String) {
            description(msg)
        }
    }
}

pub struct IconScraper {
    pub document_url: url::Url,
    pub dom: Option<kuchiki::NodeRef>,
    pub icons: Vec<Icon>
}

impl IconScraper {
    pub fn from_url(url: url::Url) -> Self {
        IconScraper {
            document_url: url,
            dom: None,
            icons: vec![]
        }
    }

    pub fn fetch_document(&mut self) -> Result<(), Error> {
        let client = hyper::client::Client::new();
        let mut response = try!(client.get(self.document_url.clone()).send());
        let parser = try!(kuchiki::Html::from_stream(&mut response));
        self.dom = Some(parser.parse());
        Ok(())
    }

    pub fn fetch_icons(&mut self) {
        self.icons = strategies::LinkRelStrategy.get_guesses(&mut self)
            .into_iter()
            .chain(strategies::DefaultFaviconPathStrategy.get_guesses(&mut self).into_iter())
            .filter_map(|mut icon| if icon.fetch_dimensions().is_ok() { Some(icon) } else { None })
            .collect::<Vec<_>>();

        self.icons.sort_by(|a, b| {
            (a.width.unwrap() * a.height.unwrap())
                .cmp(&(b.width.unwrap() * b.height.unwrap()))
        });
    }

    pub fn at_least(self, width: u32, height: u32) -> Option<Icon> {
        self.icons
            .into_iter()
            .skip_while(|icon| icon.width.unwrap() < width || icon.height.unwrap() < height)
            .next()
    }

    pub fn largest(self) -> Option<Icon> {
        if self.icons.len() > 0 {
            Some(self.icons[self.icons.len() - 1])
        }
    }
}

pub struct Icon {
    pub url: url::Url,
    pub raw: Option<Vec<u8>>,
    pub mime_type: Option<mime::Mime>,
    pub width: Option<u32>,
    pub height: Option<u32>
}

impl Icon {
    pub fn from_url(url: url::Url) -> Self {
        Icon {
            url: url,
            raw: None,
            mime_type: None,
            width: None,
            height: None
        }
    }

    pub fn fetch(&mut self) -> Result<(), Error> {
        let client = hyper::client::Client::new();
        let mut response = try!(client.get(self.url.clone()).send());

        let mut bytes: Vec<u8> = vec![];
        try!(response.read_to_end(&mut bytes));
        if !response.status.is_success() {
            return Err(Error::Other(format!("Bad status code: {:?}", response.status)));
        }

        let mime_type: mime::Mime = match response.headers.get::<hyper::header::ContentType>() {
            Some(x) => x.clone().0,
            None => return Err(Error::Other("No Content-Type found.".to_owned()))
        };
        let image_format = match mime_type.as_image_format() {
            Some(x) => x,
            None => return Err(Error::Other(format!("Invalid image type: {:?}", mime_type)))
        };
        let image = try!(image::load_from_memory_with_format(&bytes, image_format));

        self.width = Some(image.width());
        self.height = Some(image.height());
        self.raw = Some(bytes);
        self.mime_type = Some(mime_type);
        Ok(())
    }

    pub fn fetch_dimensions(&mut self) -> Result<(), Error> {
        match (self.width, self.height) {
            (Some(_), Some(_)) => Ok(()),
            _ => self.fetch()
        }
    }
}
<commit_msg>fixup<commit_after>#[macro_use] extern crate quick_error;
#[macro_use] extern crate string_cache;
extern crate mime;
extern crate hyper;
extern crate url;
extern crate kuchiki;
extern crate image;

mod strategies;
mod util;

use image::GenericImage;
use strategies::Strategy;
use std::io::Read;
use util::AsImageFormat;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        Hyper(error: hyper::Error) { from() }
        Io(error: std::io::Error) { from() }
        Image(error: image::ImageError) { from() }
        Other(msg: String) {
            description(msg)
        }
    }
}

pub struct IconScraper {
    pub document_url: url::Url,
    pub dom: Option<kuchiki::NodeRef>,
    pub icons: Vec<Icon>
}

impl IconScraper {
    pub fn from_url(url: url::Url) -> Self {
        IconScraper {
            document_url: url,
            dom: None,
            icons: vec![]
        }
    }

    pub fn fetch_document(&mut self) -> Result<(), Error> {
        let client = hyper::client::Client::new();
        let mut response = try!(client.get(self.document_url.clone()).send());
        let parser = try!(kuchiki::Html::from_stream(&mut response));
        self.dom = Some(parser.parse());
        Ok(())
    }

    pub fn fetch_icons(&mut self) {
        self.icons = strategies::LinkRelStrategy.get_guesses(&mut self)
            .into_iter()
            .chain(strategies::DefaultFaviconPathStrategy.get_guesses(&mut self).into_iter())
            .filter_map(|mut icon| if icon.fetch_dimensions().is_ok() { Some(icon) } else { None })
            .collect::<Vec<_>>();

        self.icons.sort_by(|a, b| {
            (a.width.unwrap() * a.height.unwrap())
                .cmp(&(b.width.unwrap() * b.height.unwrap()))
        });
    }

    pub fn at_least(self, width: u32, height: u32) -> Option<Icon> {
        self.icons
            .into_iter()
            .skip_while(|icon| icon.width.unwrap() < width || icon.height.unwrap() < height)
            .next()
    }

    pub fn largest(self) -> Option<Icon> {
        if self.icons.len() > 0 {
            Some(self.icons[self.icons.len() - 1])
        } else {
            None
        }
    }
}

pub struct Icon {
    pub url: url::Url,
    pub raw: Option<Vec<u8>>,
    pub mime_type: Option<mime::Mime>,
    pub width: Option<u32>,
    pub height: Option<u32>
}

impl Icon {
    pub fn from_url(url: url::Url) -> Self {
        Icon {
            url: url,
            raw: None,
            mime_type: None,
            width: None,
            height: None
        }
    }

    pub fn fetch(&mut self) -> Result<(), Error> {
        let client = hyper::client::Client::new();
        let mut response = try!(client.get(self.url.clone()).send());

        let mut bytes: Vec<u8> = vec![];
        try!(response.read_to_end(&mut bytes));
        if !response.status.is_success() {
            return Err(Error::Other(format!("Bad status code: {:?}", response.status)));
        }

        let mime_type: mime::Mime = match response.headers.get::<hyper::header::ContentType>() {
            Some(x) => x.clone().0,
            None => return Err(Error::Other("No Content-Type found.".to_owned()))
        };
        let image_format = match mime_type.as_image_format() {
            Some(x) => x,
            None => return Err(Error::Other(format!("Invalid image type: {:?}", mime_type)))
        };
        let image = try!(image::load_from_memory_with_format(&bytes, image_format));

        self.width = Some(image.width());
        self.height = Some(image.height());
        self.raw = Some(bytes);
        self.mime_type = Some(mime_type);
        Ok(())
    }

    pub fn fetch_dimensions(&mut self) -> Result<(), Error> {
        match (self.width, self.height) {
            (Some(_), Some(_)) => Ok(()),
            _ => self.fetch()
        }
    }
}
<|endoftext|>