<commit_before>#![allow(dead_code)]

use ast::*;
use visit::{Visitor};
use std::cell::RefCell;
use ptr::*;

pub type OptionalBinding = Option<B<Binding>>;
/*impl<T> Visitor<T> for CodeGenVisitor where T: std::fmt::Display {
    fn visit(&self, t: T) {
        println!("{}", t);
    }
}*/

/*struct CodeGenVisitor;
struct ExpVisitor{
    result : i32
}

impl<'a> Visitor<&'a NumExpression> for ExpVisitor {
    fn visit(&mut self, n: &'a NumExpression) {
        self.result = n.value;
    }
}

impl<'a> Visitor<&'a AddExpression> for ExpVisitor {
    fn visit(&mut self, add: &'a AddExpression) {
        let b1 = *(add.e1);
        b1.accept(self);
        let r1 = self.result;
        add.e2.accept(self);
        let r2 = self.result;
        self.result = r1 + r2;
    }
}*/

struct ExpEvaluator{
    value : i32
}

impl ExpEvaluator{
    fn get_value(&mut self, expr : & Expr)->i32{
        match expr{
            &Expr::NumExpr(value) => value,
            _ => 1
        }
    }
}

impl<'a> Visitor<'a> for ExpEvaluator{
    fn visit_expr(&mut self, expr: &'a Expr){
        match expr{
            &Expr::AddExpr(ref left, ref right) => {
                let v1 = self.get_value(left);
                let v2 = self.get_value(right);
                self.value = v1 + v2;
            },
            _ => {}
        }
    }
}

pub struct TypeChecker{
    //block_stack : Vec<RefCell<&'a  Block>>,
    pub sym_tab : Vec<(String, OptionalBinding)>,
    //decl_cnt : u32,
    //decl_cnt_stack : Vec<u32>,
    pub ty : TType
}

impl TypeChecker{
    pub fn new()->Self{
        TypeChecker {sym_tab : Vec::new(), ty : TType::TNil}
    }

    fn get_type_for(&self){//}->&TType{
        //self.block_stack
    }
}

impl<'a> Visitor<'a> for TypeChecker{
    fn visit_expr(&mut self, expr: &'a Expr){
        match expr{
            //FIXME remove NilExpr; this is only for unit testing
            &Expr::NilExpr => self.ty = TType::TString,
            &Expr::NumExpr(_) => self.ty = TType::TInt32,
            &Expr::StringExpr(_) => self.ty = TType::TString,
            &Expr::IdentExpr(ref id) =>{
                //search in the symtab for id's existence and get the type
                let mut found = false;
                for &(ref _id, ref _binding) in &self.sym_tab{ //iterator returns a ref to tuple while iterating; so &(_,_) has to be used
                    if *_id == *id{
                        found = true;
                        self.ty = match **_binding.as_ref().unwrap(){
                            Binding::TypeBinding(ref ty) |
                            Binding::VarBinding(ref ty) |
                            Binding::FuncBinding(ref ty) => ty.clone()
                        };
                        break;
                    }
                }
                if !found{
                    panic!("{} not found", id);
                }
            },
            &Expr::AddExpr(ref left, ref right) => {
                self.visit_expr(left);
                let left_ty = self.ty.clone();
                if left_ty != TType::TInt32{
                    panic!("Expected left operand of int type");
                }
                self.visit_expr(right);
                if self.ty != TType::TInt32{
                    panic!("Expected right operand of int type")
                }
            },
            &Expr::SeqExpr(ref opt_expr_list) => {
                for b_expr in opt_expr_list.as_ref().unwrap() {
                    self.visit_expr(&*b_expr);
                }
            },
            &Expr::LetExpr(ref decls, ref opt_expr) => {
                self.sym_tab.push(("<marker>".to_string(), None));

                for dec in decls{ //decls is a &
                    self.visit_decl(dec);
                }

                match opt_expr {
                    &Some(ref b_expr) => {
                        self.visit_expr(&*expr);
                    },
                    _ => {}
                }
                //pop till marker and then pop marker
                while self.sym_tab.last().unwrap().0 != "<marker>".to_string(){
                    self.sym_tab.pop();
                }
                self.sym_tab.pop();
            },

            _ => {}
        }
    }

    fn visit_decl(&mut self, decl : &'a Decl){
        match decl{
            &Decl::VarDec(ref id, ref ty, ref expr) => {
                self.visit_expr(expr);
                if *ty != self.ty{
                    panic!("Types mismatch");
                }
                self.sym_tab.push((id.clone(), Some(B(Binding::VarBinding(self.ty.clone())))));
            },
            &Decl::FunDec(ref id, ref params, ref ret_type, ref body) => {
                self.visit_expr(&body);
                if self.ty != *ret_type{
                    panic!("Return type doesn't match with the type of the last expression.");
                }
                self.sym_tab.push((id.clone(), Some(B(Binding::FuncBinding(ret_type.clone())))));
            },
            &Decl::TypeDec(ref id, ref ty) => {
                self.sym_tab.push((id.clone(), Some(B(Binding::TypeBinding(self.ty.clone())))));
            }
        }
    }
}

struct PrettyPrintVisitor;

impl<'a> Visitor<'a> for PrettyPrintVisitor{
    fn visit_expr(&mut self, expr:&'a Expr){
        match expr{
            &Expr::LetExpr(ref v, _) => {
                println!("(let");
                for d in v{
                    println!("\t(");

                }
                println!(")");
            },
            &Expr::AddExpr(ref left, ref right) => {
                self.visit_expr(left);
                println!(" Plus ");
                self.visit_expr(right);
            },
            &Expr::NumExpr(value) => {
                println!(" Num({}) ", value);
            },
            &Expr::IdentExpr(ref value) => {
                println!(" Ident({}) ", value);
            },
            _ => {}
        }
    }

    fn visit_block(&mut self, block: &'a Block){
        for s in &block.statements{
            self.visit_stmt(&*s);
        }
    }

    fn visit_stmt(&mut self, stmt : &'a Stmt){
        match stmt{
            &Stmt::VarDeclStmt(ref local) => {
                println!("(var {} type {} init ", local.ident, local.ty);
                self.visit_expr(&*local.expr);
                println!(")");
            },
            &Stmt::ExprStmt(ref expr) => {
                self.visit_expr(expr);
            }
        }
    }
}

//FIXME remove this struct
struct SymbolTableBuilder<'a>{
    block_stack : Vec<RefCell<&'a  Block>>
}

impl<'a> SymbolTableBuilder<'a>{
    fn new()->Self{
        SymbolTableBuilder {block_stack : Vec::new()}
    }
}

impl<'a> Visitor<'a> for SymbolTableBuilder<'a>{
    fn visit_block(&mut self, block : &'a Block){
        self.block_stack.push(RefCell::new(block));
        for s in &block.statements{
            self.visit_stmt(s);
        }
        self.block_stack.pop();
    }

    fn visit_stmt(&mut self, stmt : &'a Stmt){
        match stmt{
            &Stmt::VarDeclStmt(ref local) => {
                    //FIXME deduce the correct type
                    let mut block = self.block_stack.last_mut().unwrap().borrow_mut();
                    block.sym_tab.borrow_mut().insert(local.ident.clone(), TType::TInt32);
                },
            &Stmt::ExprStmt(ref expr) => {
                let b =  &**expr; //*expr is deref B which is Box<T>; **expr is deref Box<T> which is T; &**expr is therefore &T
                match b {
                    &Expr::BlockExpr(ref block) => {
                        self.visit_block(block);
                    },
                    _ => {}
                }
            }
        }
    }
}

#[test]
fn test_pp_visit_add_expr(){
    let mut p = PrettyPrintVisitor;
    p.visit_expr(&Expr::AddExpr(B(Expr::NumExpr(1)), B(Expr::NumExpr(2))));
}

#[test]
fn test_pp_visit_block(){
    let mut p = PrettyPrintVisitor;
    let mut b = Block::new();
    // let l = Local::new("a".to_string(), LuaType::LNil, B(Expr::NumExpr(1)));
    // b.statements.push(B(Stmt::VarDeclStmt(l)));
    p.visit_block(&b);
}

//#[test]
fn test_pp_visit_add(){
    let mut p = PrettyPrintVisitor;
    let mut b = Block::new();
    p.visit_block(&b);
}

#[test]
fn test_ty_set_for_num() {
    let mut v = TypeChecker::new();
    v.visit_expr(&Expr::NumExpr(23));
    assert_eq!(TType::TInt32, v.ty);
}

#[test]
fn test_ty_set_for_int_id() {
    let mut v = TypeChecker::new();
    v.sym_tab.push(("a".to_string(), Some(B(Binding::VarBinding(TType::TInt32)))));
    v.visit_expr(&Expr::IdentExpr("a".to_string()));
    assert_eq!(TType::TInt32, v.ty);
}

#[test]
fn test_type_match_int_for_var_dec() {
    let mut v = TypeChecker::new();
    v.visit_decl(&Decl::VarDec("a".to_string(), TType::TInt32, B(Expr::NumExpr(4))));
    assert_eq!(TType::TInt32, v.ty);
    assert_eq!(v.sym_tab.len(), 1);
    assert_eq!(v.sym_tab[0].0, "a".to_string());
    //assert_eq!(v.sym_tab[0].1, TType::TInt32);
}

#[test]
fn test_type_match_string_for_var_dec() {
    let mut v = TypeChecker::new();
    v.visit_decl(&Decl::VarDec("a".to_string(), TType::TString, B(Expr::NilExpr)));
    assert_eq!(TType::TString, v.ty);
    assert_eq!(v.sym_tab.len(), 1);
    assert_eq!(v.sym_tab[0].0, "a".to_string());
    //assert_eq!(v.sym_tab[0].1, TType::TString);
}

#[test]
#[should_panic]
fn test_type_check_for_var_dec_type_mismatch() {
    let mut v = TypeChecker::new();
    v.visit_decl(&Decl::VarDec("a".to_string(), TType::TInt32, B(Expr::NilExpr)));
    assert_eq!(TType::TInt32, v.ty);
}

#[test]
fn test_correct_types_for_add_expr() {
    let mut v = TypeChecker::new();
    v.visit_expr(&Expr::AddExpr(B(Expr::NumExpr(4)), B(Expr::NumExpr(4))));
    assert_eq!(v.ty, TType::TInt32);
}

#[test]
#[should_panic(expected="Expected left operand of int type")]
fn test_left_type_invalid_for_add_expr() {
    let mut v = TypeChecker::new();
    v.visit_expr(&Expr::AddExpr(B(Expr::NilExpr), B(Expr::NumExpr(4))));
}

#[test]
#[should_panic(expected="Expected right operand of int type")]
fn test_right_type_invalid_for_add_expr() {
    let mut v = TypeChecker::new();
    v.visit_expr(&Expr::AddExpr(B(Expr::NumExpr(4)), B(Expr::NilExpr)));
}

#[test]
fn test_var_hiding() {
    //let mut v =
}

#[test]
fn test_func_decl_correct_return_type() {
    let mut v = TypeChecker::new();
    v.visit_decl(&Decl::FunDec(String::from("foo"), None, TType::TInt32, B(Expr::NumExpr(4))));
}

#[test]
#[should_panic(expected="Return type doesn't match with the type of the last expression.")]
fn test_func_decl_incorrect_return_type() {
    let mut v = TypeChecker::new();
    v.visit_decl(&Decl::FunDec(String::from("foo"), None, TType::TString, B(Expr::NumExpr(4))));
}
<commit_msg>Macroize symbol table storage<commit_after>#![allow(dead_code)]

use ast::*;
use visit::{Visitor};
use std::cell::RefCell;
use ptr::*;

pub type OptionalBinding = Option<B<Binding>>;
/*impl<T> Visitor<T> for CodeGenVisitor where T: std::fmt::Display {
    fn visit(&self, t: T) {
        println!("{}", t);
    }
}*/

/*struct CodeGenVisitor;
struct ExpVisitor{
    result : i32
}

impl<'a> Visitor<&'a NumExpression> for ExpVisitor {
    fn visit(&mut self, n: &'a NumExpression) {
        self.result = n.value;
    }
}

impl<'a> Visitor<&'a AddExpression> for ExpVisitor {
    fn visit(&mut self, add: &'a AddExpression) {
        let b1 = *(add.e1);
        b1.accept(self);
        let r1 = self.result;
        add.e2.accept(self);
        let r2 = self.result;
        self.result = r1 + r2;
    }
}*/

struct ExpEvaluator{
    value : i32
}

impl ExpEvaluator{
    fn get_value(&mut self, expr : & Expr)->i32{
        match expr{
            &Expr::NumExpr(value) => value,
            _ => 1
        }
    }
}

impl<'a> Visitor<'a> for ExpEvaluator{
    fn visit_expr(&mut self, expr: &'a Expr){
        match expr{
            &Expr::AddExpr(ref left, ref right) => {
                let v1 = self.get_value(left);
                let v2 = self.get_value(right);
                self.value = v1 + v2;
            },
            _ => {}
        }
    }
}

pub struct TypeChecker{
    //block_stack : Vec<RefCell<&'a  Block>>,
    pub sym_tab : Vec<(String, OptionalBinding)>,
    //decl_cnt : u32,
    //decl_cnt_stack : Vec<u32>,
    pub ty : TType
}

impl TypeChecker{
    pub fn new()->Self{
        TypeChecker {sym_tab : Vec::new(), ty : TType::TNil}
    }

    fn get_type_for(&self){//}->&TType{
        //self.block_stack
    }
}

impl<'a> Visitor<'a> for TypeChecker{
    fn visit_expr(&mut self, expr: &'a Expr){
        match expr{
            //FIXME remove NilExpr; this is only for unit testing
            &Expr::NilExpr => self.ty = TType::TString,
            &Expr::NumExpr(_) => self.ty = TType::TInt32,
            &Expr::StringExpr(_) => self.ty = TType::TString,
            &Expr::IdentExpr(ref id) =>{
                //search in the symtab for id's existence and get the type
                let mut found = false;
                for &(ref _id, ref _binding) in &self.sym_tab{ //iterator returns a ref to tuple while iterating; so &(_,_) has to be used
                    if *_id == *id{
                        found = true;
                        self.ty = match **_binding.as_ref().unwrap(){
                            Binding::TypeBinding(ref ty) |
                            Binding::VarBinding(ref ty) |
                            Binding::FuncBinding(ref ty) => ty.clone()
                        };
                        break;
                    }
                }
                if !found{
                    panic!("{} not found", id);
                }
            },
            &Expr::AddExpr(ref left, ref right) => {
                self.visit_expr(left);
                let left_ty = self.ty.clone();
                if left_ty != TType::TInt32{
                    panic!("Expected left operand of int type");
                }
                self.visit_expr(right);
                if self.ty != TType::TInt32{
                    panic!("Expected right operand of int type")
                }
            },
            &Expr::SeqExpr(ref opt_expr_list) => {
                for b_expr in opt_expr_list.as_ref().unwrap() {
                    self.visit_expr(&*b_expr);
                }
            },
            &Expr::LetExpr(ref decls, ref opt_expr) => {
                self.sym_tab.push(("<marker>".to_string(), None));

                for dec in decls{ //decls is a &
                    self.visit_decl(dec);
                }

                match opt_expr {
                    &Some(ref b_expr) => {
                        self.visit_expr(&*expr);
                    },
                    _ => {}
                }
                //pop till marker and then pop marker
                while self.sym_tab.last().unwrap().0 != "<marker>".to_string(){
                    self.sym_tab.pop();
                }
                self.sym_tab.pop();
            },

            _ => {}
        }
    }

    fn visit_decl(&mut self, decl : &'a Decl){
        macro_rules! store_into_sym_tab {
            ($self_ : ident, $i : ident, $p : path) => {
                $self_.sym_tab.push(($i.clone(), Some(B($p($self_.ty.clone())))));
            }
        }
        match decl{
            &Decl::VarDec(ref id, ref ty, ref expr) => {
                self.visit_expr(expr);
                if *ty != self.ty{
                    panic!("Types mismatch");
                }
                store_into_sym_tab!(self, id, Binding::VarBinding);
            },
            &Decl::FunDec(ref id, ref params, ref ret_type, ref body) => {
                self.visit_expr(&body);
                if self.ty != *ret_type{
                    panic!("Return type doesn't match with the type of the last expression.");
                }
                store_into_sym_tab!(self, id, Binding::FuncBinding);
            },
            &Decl::TypeDec(ref id, ref ty) => {
                store_into_sym_tab!(self, id, Binding::TypeBinding);
            }
        }
    }
}

struct PrettyPrintVisitor;

impl<'a> Visitor<'a> for PrettyPrintVisitor{
    fn visit_expr(&mut self, expr:&'a Expr){
        match expr{
            &Expr::LetExpr(ref v, _) => {
                println!("(let");
                for d in v{
                    println!("\t(");

                }
                println!(")");
            },
            &Expr::AddExpr(ref left, ref right) => {
                self.visit_expr(left);
                println!(" Plus ");
                self.visit_expr(right);
            },
            &Expr::NumExpr(value) => {
                println!(" Num({}) ", value);
            },
            &Expr::IdentExpr(ref value) => {
                println!(" Ident({}) ", value);
            },
            _ => {}
        }
    }

    fn visit_block(&mut self, block: &'a Block){
        for s in &block.statements{
            self.visit_stmt(&*s);
        }
    }

    fn visit_stmt(&mut self, stmt : &'a Stmt){
        match stmt{
            &Stmt::VarDeclStmt(ref local) => {
                println!("(var {} type {} init ", local.ident, local.ty);
                self.visit_expr(&*local.expr);
                println!(")");
            },
            &Stmt::ExprStmt(ref expr) => {
                self.visit_expr(expr);
            }
        }
    }
}

//FIXME remove this struct
struct SymbolTableBuilder<'a>{
    block_stack : Vec<RefCell<&'a  Block>>
}

impl<'a> SymbolTableBuilder<'a>{
    fn new()->Self{
        SymbolTableBuilder {block_stack : Vec::new()}
    }
}

impl<'a> Visitor<'a> for SymbolTableBuilder<'a>{
    fn visit_block(&mut self, block : &'a Block){
        self.block_stack.push(RefCell::new(block));
        for s in &block.statements{
            self.visit_stmt(s);
        }
        self.block_stack.pop();
    }

    fn visit_stmt(&mut self, stmt : &'a Stmt){
        match stmt{
            &Stmt::VarDeclStmt(ref local) => {
                    //FIXME deduce the correct type
                    let mut block = self.block_stack.last_mut().unwrap().borrow_mut();
                    block.sym_tab.borrow_mut().insert(local.ident.clone(), TType::TInt32);
                },
            &Stmt::ExprStmt(ref expr) => {
                let b =  &**expr; //*expr is deref B which is Box<T>; **expr is deref Box<T> which is T; &**expr is therefore &T
                match b {
                    &Expr::BlockExpr(ref block) => {
                        self.visit_block(block);
                    },
                    _ => {}
                }
            }
        }
    }
}

#[test]
fn test_pp_visit_add_expr(){
    let mut p = PrettyPrintVisitor;
    p.visit_expr(&Expr::AddExpr(B(Expr::NumExpr(1)), B(Expr::NumExpr(2))));
}

#[test]
fn test_pp_visit_block(){
    let mut p = PrettyPrintVisitor;
    let mut b = Block::new();
    // let l = Local::new("a".to_string(), LuaType::LNil, B(Expr::NumExpr(1)));
    // b.statements.push(B(Stmt::VarDeclStmt(l)));
    p.visit_block(&b);
}

//#[test]
fn test_pp_visit_add(){
    let mut p = PrettyPrintVisitor;
    let mut b = Block::new();
    p.visit_block(&b);
}

#[test]
fn test_ty_set_for_num() {
    let mut v = TypeChecker::new();
    v.visit_expr(&Expr::NumExpr(23));
    assert_eq!(TType::TInt32, v.ty);
}

#[test]
fn test_ty_set_for_int_id() {
    let mut v = TypeChecker::new();
    v.sym_tab.push(("a".to_string(), Some(B(Binding::VarBinding(TType::TInt32)))));
    v.visit_expr(&Expr::IdentExpr("a".to_string()));
    assert_eq!(TType::TInt32, v.ty);
}

#[test]
fn test_type_match_int_for_var_dec() {
    let mut v = TypeChecker::new();
    v.visit_decl(&Decl::VarDec("a".to_string(), TType::TInt32, B(Expr::NumExpr(4))));
    assert_eq!(TType::TInt32, v.ty);
    assert_eq!(v.sym_tab.len(), 1);
    assert_eq!(v.sym_tab[0].0, "a".to_string());
    //assert_eq!(v.sym_tab[0].1, TType::TInt32);
}

#[test]
fn test_type_match_string_for_var_dec() {
    let mut v = TypeChecker::new();
    v.visit_decl(&Decl::VarDec("a".to_string(), TType::TString, B(Expr::NilExpr)));
    assert_eq!(TType::TString, v.ty);
    assert_eq!(v.sym_tab.len(), 1);
    assert_eq!(v.sym_tab[0].0, "a".to_string());
    //assert_eq!(v.sym_tab[0].1, TType::TString);
}

#[test]
#[should_panic]
fn test_type_check_for_var_dec_type_mismatch() {
    let mut v = TypeChecker::new();
    v.visit_decl(&Decl::VarDec("a".to_string(), TType::TInt32, B(Expr::NilExpr)));
    assert_eq!(TType::TInt32, v.ty);
}

#[test]
fn test_correct_types_for_add_expr() {
    let mut v = TypeChecker::new();
    v.visit_expr(&Expr::AddExpr(B(Expr::NumExpr(4)), B(Expr::NumExpr(4))));
    assert_eq!(v.ty, TType::TInt32);
}

#[test]
#[should_panic(expected="Expected left operand of int type")]
fn test_left_type_invalid_for_add_expr() {
    let mut v = TypeChecker::new();
    v.visit_expr(&Expr::AddExpr(B(Expr::NilExpr), B(Expr::NumExpr(4))));
}

#[test]
#[should_panic(expected="Expected right operand of int type")]
fn test_right_type_invalid_for_add_expr() {
    let mut v = TypeChecker::new();
    v.visit_expr(&Expr::AddExpr(B(Expr::NumExpr(4)), B(Expr::NilExpr)));
}

#[test]
fn test_var_hiding() {
    //let mut v =
}

#[test]
fn test_func_decl_correct_return_type() {
    let mut v = TypeChecker::new();
    v.visit_decl(&Decl::FunDec(String::from("foo"), None, TType::TInt32, B(Expr::NumExpr(4))));
}

#[test]
#[should_panic(expected="Return type doesn't match with the type of the last expression.")]
fn test_func_decl_incorrect_return_type() {
    let mut v = TypeChecker::new();
    v.visit_decl(&Decl::FunDec(String::from("foo"), None, TType::TString, B(Expr::NumExpr(4))));
}
<|endoftext|>


<commit_before>use std::cmp::Ordering;

use super::*;

pub(crate) fn prefix_encode(prefix: &[u8], buf: &[u8]) -> Vec<u8> {
    assert!(
        prefix <= buf,
        "prefix {:?} must be lexicographically <= to the encoded buf {:?}",
        prefix,
        buf
    );
    let limit = std::cmp::min(std::u8::MAX as usize, buf.len());
    let mut prefix_len = 0usize;
    for (i, c) in prefix.iter().take(limit).enumerate() {
        if buf[i] == *c {
            prefix_len += 1;
        } else {
            break;
        }
    }
    let mut ret = Vec::with_capacity(1 + buf.len() - prefix_len);
    unsafe {
        ret.set_len(1 + buf.len() - prefix_len);
    }
    ret[1..].copy_from_slice(&buf[prefix_len..]);
    ret[0] = prefix_len as u8;
    ret
}

pub(crate) fn prefix_decode(prefix: &[u8], buf: &[u8]) -> Vec<u8> {
    assert!(buf.len() >= 1);
    let prefix_len = buf[0] as usize;
    let mut ret = Vec::with_capacity(prefix_len + buf.len() - 1);
    unsafe {
        ret.set_len(prefix_len + buf.len() - 1);
    }
    ret[0..prefix_len].copy_from_slice(&prefix[0..prefix_len]);
    ret[prefix_len..].copy_from_slice(&buf[1..]);
    ret
}

// NB: the correctness of this function depends on
// the invariant that the prefix is ALWAYS lexicographically
// Less than or Equal to the keys that have been encoded
// using it. Otherwise this comparison would make no sense.
pub(crate) fn prefix_cmp(a: &[u8], b: &[u8]) -> Ordering {
    if a.is_empty() && b.is_empty() {
        return Ordering::Equal;
    } else if a.is_empty() && !b.is_empty() {
        return Ordering::Less;
    } else if !a.is_empty() && b.is_empty() {
        return Ordering::Greater;
    }

    if a[0] > b[0] {
        return Ordering::Less;
    } else if a[0] < b[0] {
        return Ordering::Greater;
    } else {
        return a[1..].cmp(&b[1..]);
    }
}

#[test]
fn test_prefix() {
    let prefix = b"cat";
    assert_eq!(
        prefix_encode(prefix, prefix),
        vec![prefix.len() as u8]
    );
    assert_eq!(prefix_encode(prefix, b"catt"), vec![3, b't']);
    assert_eq!(prefix_encode(prefix, b"ca"), vec![2]);
    assert_eq!(prefix_encode(prefix, b"cab"), vec![2, b'b']);
    assert_eq!(
        prefix_encode(prefix, b"cvar"),
        vec![1, b'v', b'a', b'r']
    );
    assert_eq!(
        prefix_encode(prefix, b"zig"),
        vec![0, b'z', b'i', b'g']
    );

    let prefix = b"";
    assert_eq!(
        prefix_encode(prefix, prefix),
        vec![prefix.len() as u8]
    );
    assert_eq!(prefix_encode(prefix, b"ca"), vec![0, b'c', b'a']);
    assert_eq!(
        prefix_encode(prefix, b"cat"),
        vec![0, b'c', b'a', b't']
    );
    assert_eq!(
        prefix_encode(prefix, b"cab"),
        vec![0, b'c', b'a', b'b']
    );
    assert_eq!(
        prefix_encode(prefix, b"cvar"),
        vec![0, b'c', b'v', b'a', b'r']
    );
    assert_eq!(
        prefix_encode(prefix, b"zig"),
        vec![0, b'z', b'i', b'g']
    );

    let rtt =
        vec![b"" as &[u8], b"\x00cat", b"\x00", b"oyunwytounw\x00"];
    for item in rtt {
        assert_eq!(
            prefix_decode(prefix, &*prefix_encode(prefix, item)),
            item.to_vec()
        );
    }
}

#[test]
fn test_prefix_cmp() {
    assert_eq!(prefix_cmp(&[], &[]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[0], &[]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[], &[0]), Ordering::Less);

    assert_eq!(prefix_cmp(&[3], &[4]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[4, 3], &[3, 4]), Ordering::Less);

    assert_eq!(prefix_cmp(&[1], &[1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 3], &[1, 1]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 3]), Ordering::Less);
}
<commit_msg>Fix prefix test that failed the correcness invariant.<commit_after>use std::cmp::Ordering;

use super::*;

pub(crate) fn prefix_encode(prefix: &[u8], buf: &[u8]) -> Vec<u8> {
    assert!(
        prefix <= buf,
        "prefix {:?} must be lexicographically <= to the encoded buf {:?}",
        prefix,
        buf
    );
    let limit = std::cmp::min(std::u8::MAX as usize, buf.len());
    let mut prefix_len = 0usize;
    for (i, c) in prefix.iter().take(limit).enumerate() {
        if buf[i] == *c {
            prefix_len += 1;
        } else {
            break;
        }
    }
    let mut ret = Vec::with_capacity(1 + buf.len() - prefix_len);
    unsafe {
        ret.set_len(1 + buf.len() - prefix_len);
    }
    ret[1..].copy_from_slice(&buf[prefix_len..]);
    ret[0] = prefix_len as u8;
    ret
}

pub(crate) fn prefix_decode(prefix: &[u8], buf: &[u8]) -> Vec<u8> {
    assert!(buf.len() >= 1);
    let prefix_len = buf[0] as usize;
    let mut ret = Vec::with_capacity(prefix_len + buf.len() - 1);
    unsafe {
        ret.set_len(prefix_len + buf.len() - 1);
    }
    ret[0..prefix_len].copy_from_slice(&prefix[0..prefix_len]);
    ret[prefix_len..].copy_from_slice(&buf[1..]);
    ret
}

// NB: the correctness of this function depends on
// the invariant that the prefix is ALWAYS lexicographically
// Less than or Equal to the keys that have been encoded
// using it. Otherwise this comparison would make no sense.
pub(crate) fn prefix_cmp(a: &[u8], b: &[u8]) -> Ordering {
    if a.is_empty() && b.is_empty() {
        return Ordering::Equal;
    } else if a.is_empty() && !b.is_empty() {
        return Ordering::Less;
    } else if !a.is_empty() && b.is_empty() {
        return Ordering::Greater;
    }

    if a[0] > b[0] {
        return Ordering::Less;
    } else if a[0] < b[0] {
        return Ordering::Greater;
    } else {
        return a[1..].cmp(&b[1..]);
    }
}

#[test]
fn test_prefix() {
    let prefix = b"cat";
    assert_eq!(
        prefix_encode(prefix, prefix),
        vec![prefix.len() as u8]
    );
    assert_eq!(prefix_encode(prefix, b"catt"), vec![3, b't']);
    assert_eq!(prefix_encode(prefix, b"cb"), vec![1, b'b']);
    assert_eq!(prefix_encode(prefix, b"caz"), vec![2, b'z']);
    assert_eq!(
        prefix_encode(prefix, b"cvar"),
        vec![1, b'v', b'a', b'r']
    );
    assert_eq!(
        prefix_encode(prefix, b"zig"),
        vec![0, b'z', b'i', b'g']
    );

    let prefix = b"";
    assert_eq!(
        prefix_encode(prefix, prefix),
        vec![prefix.len() as u8]
    );
    assert_eq!(prefix_encode(prefix, b"ca"), vec![0, b'c', b'a']);
    assert_eq!(
        prefix_encode(prefix, b"cat"),
        vec![0, b'c', b'a', b't']
    );
    assert_eq!(
        prefix_encode(prefix, b"cab"),
        vec![0, b'c', b'a', b'b']
    );
    assert_eq!(
        prefix_encode(prefix, b"cvar"),
        vec![0, b'c', b'v', b'a', b'r']
    );
    assert_eq!(
        prefix_encode(prefix, b"zig"),
        vec![0, b'z', b'i', b'g']
    );

    let rtt =
        vec![b"" as &[u8], b"\x00cat", b"\x00", b"oyunwytounw\x00"];
    for item in rtt {
        assert_eq!(
            prefix_decode(prefix, &*prefix_encode(prefix, item)),
            item.to_vec()
        );
    }
}

#[test]
fn test_prefix_cmp() {
    assert_eq!(prefix_cmp(&[], &[]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[0], &[]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[], &[0]), Ordering::Less);

    assert_eq!(prefix_cmp(&[3], &[4]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[4, 3], &[3, 4]), Ordering::Less);

    assert_eq!(prefix_cmp(&[1], &[1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 1]), Ordering::Equal);
    assert_eq!(prefix_cmp(&[1, 3], &[1, 1]), Ordering::Greater);
    assert_eq!(prefix_cmp(&[1, 1], &[1, 3]), Ordering::Less);
}
<|endoftext|>


<commit_before>use dns::question::Question;
use dns::answer::Answer;

pub struct Message {
    pub id: u16, // 2 bytes
    pub query_response: u16, // 1 bit
    pub operation_code: u16, // 4 bits
    pub authoritative_answer: u16, // 1 bit
    pub truncation_flag: u16, // 1 bit
    pub recursion_desired: u16, // 1 bit
    pub recursion_available: u16, // 1 bit
    pub unused: u16, // 3 bits
    pub error_code: u16, // 4 bits
    pub question_count: u16, // 2 bytes
    pub answer_count: u16, // 2 bytes
    pub ns_count: u16, // 2 bytes
    pub ar_count: u16, // 2 bytes

    pub questions: Vec<Question>,
    pub answers: Vec<Answer>,
}

impl Message {
    pub fn unpack(buffer: &[u8]) -> Message {
        let id: u16 = (buffer[0] as u16) << 8 | buffer[1] as u16;
        let body: u16 = (buffer[2] as u16) << 8 | buffer[3] as u16;
        let question_count: u16 = (buffer[4] as u16) << 8 | buffer[5] as u16;
        let answer_count: u16 = (buffer[6] as u16) << 8 | buffer[7] as u16;
        let ns_count: u16 = (buffer[8] as u16) << 8 | buffer[9] as u16;
        let ar_count: u16 = (buffer[10] as u16) << 8 | buffer[11] as u16;

        let mut offset: usize = 12;

        let mut questions = Vec::with_capacity(question_count as usize);

        for _ in 0..question_count {
            match Question::unpack(buffer, offset) {
                (question, updated_offset) => {
                    questions.push(question);
                    offset = updated_offset;
                }
            }
        }

        Message {
            id: id,
            query_response: (body & (1 << 15)) >> 15,
            operation_code: (body & (15 << 11)) >> 11,
            authoritative_answer: (body & (1 >> 10)) >> 10,
            truncation_flag: (body & (1 >> 9)) >> 9,
            recursion_desired: (body & (1 >> 8)) >> 8,
            recursion_available: (body & (1 >> 7)) >> 7,
            unused: (body & (7 << 4)) >> 4,
            error_code: (body & (15 << 0)) >> 0,
            question_count: question_count,
            answer_count: answer_count,
            ns_count: ns_count,
            ar_count: ar_count,
            questions: questions,
            answers: Vec::new(), // TODO parse answers
        }
    }

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

        buffer.push((self.id >> 8) as u8);
        buffer.push(self.id as u8);

        let mut body: u16 = 0;

        body = body | self.query_response << 15;
        body = body | self.operation_code << 11;
        body = body | self.authoritative_answer << 10;
        body = body | self.truncation_flag << 9;
        body = body | self.recursion_desired << 8;
        body = body | self.recursion_available << 7;
        body = body | self.unused << 4;
        body = body | self.error_code;

        buffer.push((body >> 8) as u8);
        buffer.push(body as u8);

        buffer.push((self.question_count >> 8) as u8);
        buffer.push(self.question_count as u8);

        buffer.push((self.answer_count >> 8) as u8);
        buffer.push(self.answer_count as u8);

        buffer.push((self.ns_count >> 8) as u8);
        buffer.push(self.ns_count as u8);

        buffer.push((self.ar_count >> 8) as u8);
        buffer.push(self.ar_count as u8);

        for question in &self.questions {
            for part in &question.name {
                buffer.push(part.len() as u8);
                let bytes = part.to_owned().into_bytes();
                for byte in &bytes {
                    buffer.push(*byte);
                }
            }

            buffer.push(0 as u8);

            buffer.push((question.rrtype >> 8) as u8);
            buffer.push(question.rrtype as u8);

            buffer.push((question.class >> 8) as u8);
            buffer.push(question.class as u8);
        }

        for answer in &self.answers {
            for part in &answer.name {
                buffer.push(part.len() as u8);
                let bytes = part.to_owned().into_bytes();
                for byte in &bytes {
                    buffer.push(*byte);
                }
            }

            buffer.push(0 as u8);

            buffer.push((answer.rrtype >> 8) as u8);
            buffer.push(answer.rrtype as u8);

            buffer.push((answer.class >> 8) as u8);
            buffer.push(answer.class as u8);

            buffer.push(((answer.ttl & (256 << 24)) >> 24) as u8);
            buffer.push(((answer.ttl & (256 << 16)) >> 16) as u8);
            buffer.push(((answer.ttl & (256 << 8)) >> 8) as u8);
            buffer.push(((answer.ttl & (256 << 0)) >> 0) as u8);

            buffer.push((answer.length >> 8) as u8);
            buffer.push(answer.length as u8);

            for byte in &answer.data {
                buffer.push(*byte);
            }
        }

        buffer
    }
}
<commit_msg>Use iter() in some for loops<commit_after>use dns::question::Question;
use dns::answer::Answer;

pub struct Message {
    pub id: u16, // 2 bytes
    pub query_response: u16, // 1 bit
    pub operation_code: u16, // 4 bits
    pub authoritative_answer: u16, // 1 bit
    pub truncation_flag: u16, // 1 bit
    pub recursion_desired: u16, // 1 bit
    pub recursion_available: u16, // 1 bit
    pub unused: u16, // 3 bits
    pub error_code: u16, // 4 bits
    pub question_count: u16, // 2 bytes
    pub answer_count: u16, // 2 bytes
    pub ns_count: u16, // 2 bytes
    pub ar_count: u16, // 2 bytes

    pub questions: Vec<Question>,
    pub answers: Vec<Answer>,
}

impl Message {
    pub fn unpack(buffer: &[u8]) -> Message {
        let id: u16 = (buffer[0] as u16) << 8 | buffer[1] as u16;
        let body: u16 = (buffer[2] as u16) << 8 | buffer[3] as u16;
        let question_count: u16 = (buffer[4] as u16) << 8 | buffer[5] as u16;
        let answer_count: u16 = (buffer[6] as u16) << 8 | buffer[7] as u16;
        let ns_count: u16 = (buffer[8] as u16) << 8 | buffer[9] as u16;
        let ar_count: u16 = (buffer[10] as u16) << 8 | buffer[11] as u16;

        let mut offset: usize = 12;

        let mut questions = Vec::with_capacity(question_count as usize);

        for _ in 0..question_count {
            match Question::unpack(buffer, offset) {
                (question, updated_offset) => {
                    questions.push(question);
                    offset = updated_offset;
                }
            }
        }

        Message {
            id: id,
            query_response: (body & (1 << 15)) >> 15,
            operation_code: (body & (15 << 11)) >> 11,
            authoritative_answer: (body & (1 >> 10)) >> 10,
            truncation_flag: (body & (1 >> 9)) >> 9,
            recursion_desired: (body & (1 >> 8)) >> 8,
            recursion_available: (body & (1 >> 7)) >> 7,
            unused: (body & (7 << 4)) >> 4,
            error_code: (body & (15 << 0)) >> 0,
            question_count: question_count,
            answer_count: answer_count,
            ns_count: ns_count,
            ar_count: ar_count,
            questions: questions,
            answers: Vec::new(), // TODO parse answers
        }
    }

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

        buffer.push((self.id >> 8) as u8);
        buffer.push(self.id as u8);

        let mut body: u16 = 0;

        body = body | self.query_response << 15;
        body = body | self.operation_code << 11;
        body = body | self.authoritative_answer << 10;
        body = body | self.truncation_flag << 9;
        body = body | self.recursion_desired << 8;
        body = body | self.recursion_available << 7;
        body = body | self.unused << 4;
        body = body | self.error_code;

        buffer.push((body >> 8) as u8);
        buffer.push(body as u8);

        buffer.push((self.question_count >> 8) as u8);
        buffer.push(self.question_count as u8);

        buffer.push((self.answer_count >> 8) as u8);
        buffer.push(self.answer_count as u8);

        buffer.push((self.ns_count >> 8) as u8);
        buffer.push(self.ns_count as u8);

        buffer.push((self.ar_count >> 8) as u8);
        buffer.push(self.ar_count as u8);

        for question in self.questions.iter() {
            for part in question.name.iter() {
                buffer.push(part.len() as u8);

                for byte in part.to_owned().into_bytes().iter() {
                    buffer.push(*byte);
                }
            }

            buffer.push(0 as u8);

            buffer.push((question.rrtype >> 8) as u8);
            buffer.push(question.rrtype as u8);

            buffer.push((question.class >> 8) as u8);
            buffer.push(question.class as u8);
        }

        for answer in self.answers.iter() {
            for part in answer.name.iter() {
                buffer.push(part.len() as u8);

                for byte in part.to_owned().into_bytes().iter() {
                    buffer.push(*byte);
                }
            }

            buffer.push(0 as u8);

            buffer.push((answer.rrtype >> 8) as u8);
            buffer.push(answer.rrtype as u8);

            buffer.push((answer.class >> 8) as u8);
            buffer.push(answer.class as u8);

            buffer.push(((answer.ttl & (256 << 24)) >> 24) as u8);
            buffer.push(((answer.ttl & (256 << 16)) >> 16) as u8);
            buffer.push(((answer.ttl & (256 << 8)) >> 8) as u8);
            buffer.push(((answer.ttl & (256 << 0)) >> 0) as u8);

            buffer.push((answer.length >> 8) as u8);
            buffer.push(answer.length as u8);

            for byte in &answer.data {
                buffer.push(*byte);
            }
        }

        buffer
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use url::Url;

use ffi;
use ffi::opt;
use ffi::easy::Easy;
use http::Response;
use http::body::{Body,ToBody};
use {ProgressCb,ErrCode};

static DEFAULT_TIMEOUT_MS: uint = 30_000;

pub struct Handle {
    easy: Easy,
}

impl Handle {
    pub fn new() -> Handle {
        Handle { easy: Easy::new() }
            .timeout(DEFAULT_TIMEOUT_MS)
            .connect_timeout(DEFAULT_TIMEOUT_MS)
    }

    pub fn timeout(mut self, ms: uint) -> Handle {
        self.easy.setopt(opt::TIMEOUT_MS, ms).unwrap();
        self
    }

    pub fn connect_timeout(mut self, ms: uint) -> Handle {
        self.easy.setopt(opt::CONNECTTIMEOUT_MS, ms).unwrap();
        self
    }

    pub fn proxy<U: ToUrl>(mut self, proxy: U) -> Handle {
        proxy.with_url_str(|s| {
            self.easy.setopt(opt::PROXY, s).unwrap();
        });

        self
    }

    pub fn get<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Get).uri(uri)
    }

    pub fn head<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Head).uri(uri)
    }

    pub fn post<'a, 'b, U: ToUrl, B: ToBody<'b>>(&'a mut self, uri: U, body: B) -> Request<'a, 'b> {
        Request::new(self, Post).uri(uri).body(body)
    }

    pub fn put<'a, 'b, U: ToUrl, B: ToBody<'b>>(&'a mut self, uri: U, body: B) -> Request<'a, 'b> {
        Request::new(self, Put).uri(uri).body(body)
    }

    pub fn delete<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Delete).uri(uri)
    }
}

pub enum Method {
    Options,
    Get,
    Head,
    Post,
    Put,
    Delete,
    Trace,
    Connect
}

pub struct Request<'a, 'b> {
    err: Option<ErrCode>,
    handle: &'a mut Handle,
    method: Method,
    headers: HashMap<String, Vec<String>>,
    body: Option<Body<'b>>,
    body_type: Option<BodyType>,
    content_type: bool, // whether or not the content type was set
    expect_continue: bool, // whether to expect a 100 continue from the server
    progress: Option<ProgressCb<'b>>,
    follow: bool,
}

enum BodyType {
    Fixed(uint),
    Chunked,
}

impl<'a, 'b> Request<'a, 'b> {
    fn new<'a, 'b>(handle: &'a mut Handle, method: Method) -> Request<'a, 'b> {
        Request {
            err: None,
            handle: handle,
            method: method,
            headers: HashMap::new(),
            body: None,
            body_type: None,
            content_type: false,
            expect_continue: false,
            progress: None,
            follow: false,
        }
    }

    pub fn uri<U: ToUrl>(mut self, uri: U) -> Request<'a, 'b> {
        uri.with_url_str(|s| {
            match self.handle.easy.setopt(opt::URL, s) {
                Ok(_) => {}
                Err(e) => self.err = Some(e)
            }
        });

        self
    }

    pub fn body<B: ToBody<'b>>(mut self, body: B) -> Request<'a, 'b> {
        self.body = Some(body.to_body());
        self
    }

    pub fn content_type(mut self, ty: &str) -> Request<'a, 'b> {
        if !self.content_type {
            self.content_type = true;
            append_header(&mut self.headers, "Content-Type", ty);
        }

        self
    }

    pub fn content_length(mut self, len: uint) -> Request<'a, 'b> {
        self.body_type = Some(Fixed(len));
        self
    }

    pub fn chunked(mut self) -> Request<'a, 'b> {
        self.body_type = Some(Chunked);
        self
    }

    pub fn expect_continue(mut self) -> Request<'a, 'b> {
        self.expect_continue = true;
        self
    }

    pub fn header(mut self, name: &str, val: &str) -> Request<'a, 'b> {
        append_header(&mut self.headers, name, val);
        self
    }

    pub fn get_header<'a>(&'a self, name: &str) -> Option<&'a [String]> {
        self.headers.find_equiv(&name).map(|a| a.as_slice())
    }

    pub fn headers<'c, 'd, I: Iterator<(&'c str, &'d str)>>(mut self, mut hdrs: I) -> Request<'a, 'b> {
        for (name, val) in hdrs {
            append_header(&mut self.headers, name, val);
        }

        self
    }

    pub fn progress(mut self, cb: ProgressCb<'b>) -> Request<'a, 'b> {
        self.progress = Some(cb);
        self
    }

    pub fn follow_redirects(mut self, follow: bool) -> Request<'a, 'b> {
        self.follow = follow;
        self
    }

    pub fn exec(self) -> Result<Response, ErrCode> {
        // Deconstruct the struct
        let Request {
            err,
            handle,
            method,
            mut headers,
            mut body,
            body_type,
            content_type,
            expect_continue,
            progress,
            follow,
            ..
        } = self;

        if follow {
            try!(handle.easy.setopt(opt::FOLLOWLOCATION, 1i));
        }

        match err {
            Some(e) => return Err(e),
            None => {}
        }

        // Clear custom headers set from the previous request
        try!(handle.easy.setopt(opt::HTTPHEADER, 0u));

        match method {
            Get => try!(handle.easy.setopt(opt::HTTPGET, 1i)),
            Head => try!(handle.easy.setopt(opt::NOBODY, 1i)),
            Post => try!(handle.easy.setopt(opt::POST, 1i)),
            Put => try!(handle.easy.setopt(opt::UPLOAD, 1i)),
            Delete => {
                if body.is_some() {
                    try!(handle.easy.setopt(opt::UPLOAD, 1i));
                }

                try!(handle.easy.setopt(opt::CUSTOMREQUEST, "DELETE"));
            }
            _ => unimplemented!()
        }

        match body.as_ref() {
            None => {}
            Some(body) => {
                debug!("handling body");

                let body_type = body_type.unwrap_or(match body.get_size() {
                    Some(len) => Fixed(len),
                    None => Chunked,
                });

                match body_type {
                    Fixed(len) => {
                        match method {
                            Post => try!(handle.easy.setopt(opt::POSTFIELDSIZE, len)),
                            Put | Delete  => try!(handle.easy.setopt(opt::INFILESIZE, len)),
                            _ => {}
                        }
                        append_header(&mut headers, "Content-Length",
                                      len.to_string().as_slice());
                    }
                    Chunked => {
                        append_header(&mut headers, "Transfer-Encoding",
                                      "chunked");
                    }

                }

                if !content_type {
                    append_header(&mut headers, "Content-Type", "application/octet-stream");
                }

                if !expect_continue {
                    append_header(&mut headers, "Expect", "");
                }
            }
        }

        let mut ffi_headers = ffi::List::new();

        if !headers.is_empty() {
            let mut buf = Vec::new();

            for (k, v) in headers.iter() {
                buf.push_all(k.as_bytes());
                buf.push_all(b": ");

                for v in v.iter() {
                    buf.push_all(v.as_bytes());
                    buf.push(0);
                    ffi_headers.push_bytes(buf.as_slice());
                    buf.truncate(k.len() + 2);
                }

                buf.truncate(0);
            }

            try!(handle.easy.setopt(opt::HTTPHEADER, &ffi_headers));
        }

        handle.easy.perform(body.as_mut(), progress)
    }
}

fn append_header(map: &mut HashMap<String, Vec<String>>, key: &str, val: &str) {
    map.find_or_insert(key.to_string(), Vec::new()).push(val.to_string());
}

pub trait ToUrl{
    fn with_url_str(self, f: |&str|);
}

impl<'a> ToUrl for &'a str {
    fn with_url_str(self, f: |&str|) {
        f(self);
    }
}

impl<'a> ToUrl for &'a Url {
    fn with_url_str(self, f: |&str|) {
        self.to_string().with_url_str(f);
    }
}

impl ToUrl for String {
    fn with_url_str(self, f: |&str|) {
        self.as_slice().with_url_str(f);
    }
}

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

    #[test]
    fn get_header() {
        let mut h = Handle::new();
        let r = h.get("/foo").header("foo", "bar");
        assert_eq!(r.get_header("foo"), Some(&["bar".to_string()]));
    }
}
<commit_msg>Update to Rust master<commit_after>use std::collections::HashMap;
use url::Url;

use ffi;
use ffi::opt;
use ffi::easy::Easy;
use http::Response;
use http::body::{Body,ToBody};
use {ProgressCb,ErrCode};

static DEFAULT_TIMEOUT_MS: uint = 30_000;

pub struct Handle {
    easy: Easy,
}

impl Handle {
    pub fn new() -> Handle {
        Handle { easy: Easy::new() }
            .timeout(DEFAULT_TIMEOUT_MS)
            .connect_timeout(DEFAULT_TIMEOUT_MS)
    }

    pub fn timeout(mut self, ms: uint) -> Handle {
        self.easy.setopt(opt::TIMEOUT_MS, ms).unwrap();
        self
    }

    pub fn connect_timeout(mut self, ms: uint) -> Handle {
        self.easy.setopt(opt::CONNECTTIMEOUT_MS, ms).unwrap();
        self
    }

    pub fn proxy<U: ToUrl>(mut self, proxy: U) -> Handle {
        proxy.with_url_str(|s| {
            self.easy.setopt(opt::PROXY, s).unwrap();
        });

        self
    }

    pub fn get<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Get).uri(uri)
    }

    pub fn head<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Head).uri(uri)
    }

    pub fn post<'a, 'b, U: ToUrl, B: ToBody<'b>>(&'a mut self, uri: U, body: B) -> Request<'a, 'b> {
        Request::new(self, Post).uri(uri).body(body)
    }

    pub fn put<'a, 'b, U: ToUrl, B: ToBody<'b>>(&'a mut self, uri: U, body: B) -> Request<'a, 'b> {
        Request::new(self, Put).uri(uri).body(body)
    }

    pub fn delete<'a, 'b, U: ToUrl>(&'a mut self, uri: U) -> Request<'a, 'b> {
        Request::new(self, Delete).uri(uri)
    }
}

pub enum Method {
    Options,
    Get,
    Head,
    Post,
    Put,
    Delete,
    Trace,
    Connect
}

pub struct Request<'a, 'b> {
    err: Option<ErrCode>,
    handle: &'a mut Handle,
    method: Method,
    headers: HashMap<String, Vec<String>>,
    body: Option<Body<'b>>,
    body_type: Option<BodyType>,
    content_type: bool, // whether or not the content type was set
    expect_continue: bool, // whether to expect a 100 continue from the server
    progress: Option<ProgressCb<'b>>,
    follow: bool,
}

enum BodyType {
    Fixed(uint),
    Chunked,
}

impl<'a, 'b> Request<'a, 'b> {
    fn new<'a, 'b>(handle: &'a mut Handle, method: Method) -> Request<'a, 'b> {
        Request {
            err: None,
            handle: handle,
            method: method,
            headers: HashMap::new(),
            body: None,
            body_type: None,
            content_type: false,
            expect_continue: false,
            progress: None,
            follow: false,
        }
    }

    pub fn uri<U: ToUrl>(mut self, uri: U) -> Request<'a, 'b> {
        uri.with_url_str(|s| {
            match self.handle.easy.setopt(opt::URL, s) {
                Ok(_) => {}
                Err(e) => self.err = Some(e)
            }
        });

        self
    }

    pub fn body<B: ToBody<'b>>(mut self, body: B) -> Request<'a, 'b> {
        self.body = Some(body.to_body());
        self
    }

    pub fn content_type(mut self, ty: &str) -> Request<'a, 'b> {
        if !self.content_type {
            self.content_type = true;
            append_header(&mut self.headers, "Content-Type", ty);
        }

        self
    }

    pub fn content_length(mut self, len: uint) -> Request<'a, 'b> {
        self.body_type = Some(Fixed(len));
        self
    }

    pub fn chunked(mut self) -> Request<'a, 'b> {
        self.body_type = Some(Chunked);
        self
    }

    pub fn expect_continue(mut self) -> Request<'a, 'b> {
        self.expect_continue = true;
        self
    }

    pub fn header(mut self, name: &str, val: &str) -> Request<'a, 'b> {
        append_header(&mut self.headers, name, val);
        self
    }

    pub fn get_header<'a>(&'a self, name: &str) -> Option<&'a [String]> {
        self.headers.find_equiv(&name).map(|a| a.as_slice())
    }

    pub fn headers<'c, 'd, I: Iterator<(&'c str, &'d str)>>(mut self, mut hdrs: I) -> Request<'a, 'b> {
        for (name, val) in hdrs {
            append_header(&mut self.headers, name, val);
        }

        self
    }

    pub fn progress(mut self, cb: ProgressCb<'b>) -> Request<'a, 'b> {
        self.progress = Some(cb);
        self
    }

    pub fn follow_redirects(mut self, follow: bool) -> Request<'a, 'b> {
        self.follow = follow;
        self
    }

    pub fn exec(self) -> Result<Response, ErrCode> {
        // Deconstruct the struct
        let Request {
            err,
            handle,
            method,
            mut headers,
            mut body,
            body_type,
            content_type,
            expect_continue,
            progress,
            follow,
            ..
        } = self;

        if follow {
            try!(handle.easy.setopt(opt::FOLLOWLOCATION, 1i));
        }

        match err {
            Some(e) => return Err(e),
            None => {}
        }

        // Clear custom headers set from the previous request
        try!(handle.easy.setopt(opt::HTTPHEADER, 0u));

        match method {
            Get => try!(handle.easy.setopt(opt::HTTPGET, 1i)),
            Head => try!(handle.easy.setopt(opt::NOBODY, 1i)),
            Post => try!(handle.easy.setopt(opt::POST, 1i)),
            Put => try!(handle.easy.setopt(opt::UPLOAD, 1i)),
            Delete => {
                if body.is_some() {
                    try!(handle.easy.setopt(opt::UPLOAD, 1i));
                }

                try!(handle.easy.setopt(opt::CUSTOMREQUEST, "DELETE"));
            }
            _ => unimplemented!()
        }

        match body.as_ref() {
            None => {}
            Some(body) => {
                debug!("handling body");

                let body_type = body_type.unwrap_or(match body.get_size() {
                    Some(len) => Fixed(len),
                    None => Chunked,
                });

                match body_type {
                    Fixed(len) => {
                        match method {
                            Post => try!(handle.easy.setopt(opt::POSTFIELDSIZE, len)),
                            Put | Delete  => try!(handle.easy.setopt(opt::INFILESIZE, len)),
                            _ => {}
                        }
                        append_header(&mut headers, "Content-Length",
                                      len.to_string().as_slice());
                    }
                    Chunked => {
                        append_header(&mut headers, "Transfer-Encoding",
                                      "chunked");
                    }

                }

                if !content_type {
                    append_header(&mut headers, "Content-Type", "application/octet-stream");
                }

                if !expect_continue {
                    append_header(&mut headers, "Expect", "");
                }
            }
        }

        let mut ffi_headers = ffi::List::new();

        if !headers.is_empty() {
            let mut buf = Vec::new();

            for (k, v) in headers.iter() {
                buf.push_all(k.as_bytes());
                buf.push_all(b": ");

                for v in v.iter() {
                    buf.push_all(v.as_bytes());
                    buf.push(0);
                    ffi_headers.push_bytes(buf.as_slice());
                    buf.truncate(k.len() + 2);
                }

                buf.truncate(0);
            }

            try!(handle.easy.setopt(opt::HTTPHEADER, &ffi_headers));
        }

        handle.easy.perform(body.as_mut(), progress)
    }
}

fn append_header(map: &mut HashMap<String, Vec<String>>, key: &str, val: &str) {
    map.find_or_insert(key.to_string(), Vec::new()).push(val.to_string());
}

pub trait ToUrl{
    fn with_url_str(self, f: |&str|);
}

impl<'a> ToUrl for &'a str {
    fn with_url_str(self, f: |&str|) {
        f(self);
    }
}

impl<'a> ToUrl for &'a Url {
    fn with_url_str(self, f: |&str|) {
        self.to_string().with_url_str(f);
    }
}

impl ToUrl for String {
    fn with_url_str(self, f: |&str|) {
        self.as_slice().with_url_str(f);
    }
}

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

    #[test]
    fn get_header() {
        let mut h = Handle::new();
        let r = h.get("/foo").header("foo", "bar");
        assert_eq!(r.get_header("foo"), Some(["bar".to_string()].as_slice()));
    }
}
<|endoftext|>


<commit_before>extern crate openssl;

use std::io;
use std::fmt;
use std::error;
use self::openssl::pkcs12;
use self::openssl::error::ErrorStack;
use self::openssl::ssl::{self, SslMethod, SslConnectorBuilder, SslConnector, SslAcceptorBuilder,
                         SslAcceptor, MidHandshakeSslStream, SslContextBuilder};
use self::openssl::x509::X509;

use Protocol;

fn supported_protocols(protocols: &[Protocol], ctx: &mut SslContextBuilder) {
    // This constant is only defined on OpenSSL 1.0.2 and above, so manually do it.
    let ssl_op_no_ssl_mask = ssl::SSL_OP_NO_SSLV2 | ssl::SSL_OP_NO_SSLV3 | ssl::SSL_OP_NO_TLSV1 |
        ssl::SSL_OP_NO_TLSV1_1 |
        ssl::SSL_OP_NO_TLSV1_2;

    ctx.clear_options(ssl_op_no_ssl_mask);
    let mut options = ssl_op_no_ssl_mask;
    for protocol in protocols {
        let op = match *protocol {
            Protocol::Sslv3 => ssl::SSL_OP_NO_SSLV3,
            Protocol::Tlsv10 => ssl::SSL_OP_NO_TLSV1,
            Protocol::Tlsv11 => ssl::SSL_OP_NO_TLSV1_1,
            Protocol::Tlsv12 => ssl::SSL_OP_NO_TLSV1_2,
            Protocol::__NonExhaustive => unreachable!(),
        };
        options &= !op;
    }
    ctx.set_options(options);
}

pub struct Error(ssl::Error);

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

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

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

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

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

impl From<ErrorStack> for Error {
    fn from(err: ErrorStack) -> Error {
        ssl::Error::Ssl(err).into()
    }
}

pub struct Pkcs12(pkcs12::ParsedPkcs12);

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let pkcs12 = try!(pkcs12::Pkcs12::from_der(buf));
        let parsed = try!(pkcs12.parse(pass));
        Ok(Pkcs12(parsed))
    }
}

pub struct Certificate(X509);

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

pub struct MidHandshakeTlsStream<S>(MidHandshakeSslStream<S>);

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

impl<S> MidHandshakeTlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

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

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

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

impl<S> From<ssl::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: ssl::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            ssl::HandshakeError::SetupFailure(e) => {
                HandshakeError::Failure(Error(ssl::Error::Ssl(e)))
            }
            ssl::HandshakeError::Failure(e) => HandshakeError::Failure(Error(e.into_error())),
            ssl::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

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

pub struct TlsConnectorBuilder(SslConnectorBuilder);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        // FIXME clear chain certs to clean up if called multiple times
        try!(self.0.set_certificate(&pkcs12.0.cert));
        try!(self.0.set_private_key(&pkcs12.0.pkey));
        try!(self.0.check_private_key());
        for cert in pkcs12.0.chain {
            try!(self.0.add_extra_chain_cert(cert));
        }
        Ok(())
    }

    pub fn add_root_certificate(&mut self, cert: Certificate) -> Result<(), Error> {
        try!(self.0.cert_store_mut().add_cert(cert.0));
        Ok(())
    }

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

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

#[derive(Clone)]
pub struct TlsConnector(SslConnector);

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        let builder = try!(SslConnectorBuilder::new(SslMethod::tls()));
        Ok(TlsConnectorBuilder(builder))
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.connect(domain, stream));
        Ok(TlsStream(s))
    }

    pub fn connect_no_domain<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.danger_connect_without_providing_domain_for_certificate_verification_and_server_name_indication(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslConnectorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslConnectorBuilder`.
    fn builder(&self) -> &SslConnectorBuilder;

    /// Returns a mutable reference to the inner `SslConnectorBuilder`.
    fn builder_mut(&mut self) -> &mut SslConnectorBuilder;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn from_openssl(builder: SslConnectorBuilder) -> ::TlsConnectorBuilder {
        ::TlsConnectorBuilder(TlsConnectorBuilder(builder))
    }

    fn builder(&self) -> &SslConnectorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslConnectorBuilder {
        &mut (self.0).0
    }
}

pub struct TlsAcceptorBuilder(SslAcceptorBuilder);

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

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

#[derive(Clone)]
pub struct TlsAcceptor(SslAcceptor);

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        let builder = try!(SslAcceptorBuilder::mozilla_intermediate(
            SslMethod::tls(),
            &pkcs12.0.pkey,
            &pkcs12.0.cert,
            &pkcs12.0.chain,
        ));
        Ok(TlsAcceptorBuilder(builder))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.accept(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsAcceptorBuilder`.
pub trait TlsAcceptorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslAcceptorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslAcceptorBuilder`.
    fn builder(&self) -> &SslAcceptorBuilder;

    /// Returns a mutable reference to the inner `SslAcceptorBuilder`.
    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder;
}

impl TlsAcceptorBuilderExt for ::TlsAcceptorBuilder {
    fn from_openssl(builder: SslAcceptorBuilder) -> ::TlsAcceptorBuilder {
        ::TlsAcceptorBuilder(TlsAcceptorBuilder(builder))
    }

    fn builder(&self) -> &SslAcceptorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder {
        &mut (self.0).0
    }
}

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

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

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(self.0.ssl().pending())
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        match self.0.shutdown() {
            Ok(_) |
            Err(ssl::Error::ZeroReturn) => Ok(()),
            Err(ssl::Error::Stream(e)) |
            Err(ssl::Error::WantRead(e)) |
            Err(ssl::Error::WantWrite(e)) => Err(e),
            Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
        }
    }

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

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

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

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

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

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

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

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

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

/// OpenSSL-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying OpenSSL error for inspection.
    fn openssl_error(&self) -> &ssl::Error;
}

impl ErrorExt for ::Error {
    fn openssl_error(&self) -> &ssl::Error {
        &(self.0).0
    }
}
<commit_msg>More fixes<commit_after>extern crate openssl;

use std::io;
use std::fmt;
use std::error;
use self::openssl::pkcs12;
use self::openssl::error::ErrorStack;
use self::openssl::ssl::{self, SslMethod, SslConnectorBuilder, SslConnector, SslAcceptorBuilder,
                         SslAcceptor, MidHandshakeSslStream, SslContextBuilder};
use self::openssl::x509::X509;

use Protocol;

fn supported_protocols(protocols: &[Protocol], ctx: &mut SslContextBuilder) {
    // This constant is only defined on OpenSSL 1.0.2 and above, so manually do it.
    let ssl_op_no_ssl_mask = ssl::SSL_OP_NO_SSLV2 | ssl::SSL_OP_NO_SSLV3 | ssl::SSL_OP_NO_TLSV1 |
        ssl::SSL_OP_NO_TLSV1_1 |
        ssl::SSL_OP_NO_TLSV1_2;

    ctx.clear_options(ssl_op_no_ssl_mask);
    let mut options = ssl_op_no_ssl_mask;
    for protocol in protocols {
        let op = match *protocol {
            Protocol::Sslv3 => ssl::SSL_OP_NO_SSLV3,
            Protocol::Tlsv10 => ssl::SSL_OP_NO_TLSV1,
            Protocol::Tlsv11 => ssl::SSL_OP_NO_TLSV1_1,
            Protocol::Tlsv12 => ssl::SSL_OP_NO_TLSV1_2,
            Protocol::__NonExhaustive => unreachable!(),
        };
        options &= !op;
    }
    ctx.set_options(options);
}

pub struct Error(ssl::Error);

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

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

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

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

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

impl From<ErrorStack> for Error {
    fn from(err: ErrorStack) -> Error {
        ssl::Error::Ssl(err).into()
    }
}

pub struct Pkcs12(pkcs12::ParsedPkcs12);

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let pkcs12 = try!(pkcs12::Pkcs12::from_der(buf));
        let parsed = try!(pkcs12.parse(pass));
        Ok(Pkcs12(parsed))
    }
}

pub struct Certificate(X509);

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

pub struct MidHandshakeTlsStream<S>(MidHandshakeSslStream<S>);

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

impl<S> MidHandshakeTlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

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

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

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

impl<S> From<ssl::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: ssl::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            ssl::HandshakeError::SetupFailure(e) => {
                HandshakeError::Failure(Error(ssl::Error::Ssl(e)))
            }
            ssl::HandshakeError::Failure(e) => HandshakeError::Failure(Error(e.into_error())),
            ssl::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

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

pub struct TlsConnectorBuilder(SslConnectorBuilder);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        // FIXME clear chain certs to clean up if called multiple times
        try!(self.0.set_certificate(&pkcs12.0.cert));
        try!(self.0.set_private_key(&pkcs12.0.pkey));
        try!(self.0.check_private_key());
        for cert in pkcs12.0.chain {
            try!(self.0.add_extra_chain_cert(cert));
        }
        Ok(())
    }

    pub fn add_root_certificate(&mut self, cert: Certificate) -> Result<(), Error> {
        try!(self.0.cert_store_mut().add_cert(cert.0));
        Ok(())
    }

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

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

#[derive(Clone)]
pub struct TlsConnector(SslConnector);

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        let builder = try!(SslConnectorBuilder::new(SslMethod::tls()));
        Ok(TlsConnectorBuilder(builder))
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.connect(domain, stream));
        Ok(TlsStream(s))
    }

    pub fn connect_no_domain<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.danger_connect_without_providing_domain_for_certificate_verification_and_server_name_indication(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslConnectorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslConnectorBuilder`.
    fn builder(&self) -> &SslConnectorBuilder;

    /// Returns a mutable reference to the inner `SslConnectorBuilder`.
    fn builder_mut(&mut self) -> &mut SslConnectorBuilder;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn from_openssl(builder: SslConnectorBuilder) -> ::TlsConnectorBuilder {
        ::TlsConnectorBuilder(TlsConnectorBuilder(builder))
    }

    fn builder(&self) -> &SslConnectorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslConnectorBuilder {
        &mut (self.0).0
    }
}

pub struct TlsAcceptorBuilder(SslAcceptorBuilder);

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

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

#[derive(Clone)]
pub struct TlsAcceptor(SslAcceptor);

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        let builder = try!(SslAcceptorBuilder::mozilla_intermediate(
            SslMethod::tls(),
            &pkcs12.0.pkey,
            &pkcs12.0.cert,
            &pkcs12.0.chain,
        ));
        Ok(TlsAcceptorBuilder(builder))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.accept(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsAcceptorBuilder`.
pub trait TlsAcceptorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslAcceptorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslAcceptorBuilder`.
    fn builder(&self) -> &SslAcceptorBuilder;

    /// Returns a mutable reference to the inner `SslAcceptorBuilder`.
    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder;
}

impl TlsAcceptorBuilderExt for ::TlsAcceptorBuilder {
    fn from_openssl(builder: SslAcceptorBuilder) -> ::TlsAcceptorBuilder {
        ::TlsAcceptorBuilder(TlsAcceptorBuilder(builder))
    }

    fn builder(&self) -> &SslAcceptorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder {
        &mut (self.0).0
    }
}

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

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

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(self.0.ssl().pending())
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        match self.0.shutdown() {
            Ok(_) |
            Err(ssl::Error::ZeroReturn) => Ok(()),
            Err(ssl::Error::Stream(e)) |
            Err(ssl::Error::WantRead(e)) |
            Err(ssl::Error::WantWrite(e)) => Err(e),
            Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
        }
    }

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

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

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

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

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

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

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

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

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

/// OpenSSL-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying OpenSSL error for inspection.
    fn openssl_error(&self) -> &ssl::Error;
}

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


<commit_before>extern crate openssl;

use std::io;
use std::fmt;
use std::error;
use self::openssl::pkcs12;
use self::openssl::error::ErrorStack;
use self::openssl::ssl::{self, SslMethod, SslConnectorBuilder, SslConnector, SslAcceptorBuilder,
                         SslAcceptor, MidHandshakeSslStream, SslContextBuilder};
use self::openssl::x509::X509;

use Protocol;

fn supported_protocols(protocols: &[Protocol], ctx: &mut SslContextBuilder) {
    // This constant is only defined on OpenSSL 1.0.2 and above, so manually do it.
    let ssl_op_no_ssl_mask = ssl::SSL_OP_NO_SSLV2 | ssl::SSL_OP_NO_SSLV3 | ssl::SSL_OP_NO_TLSV1 |
        ssl::SSL_OP_NO_TLSV1_1 |
        ssl::SSL_OP_NO_TLSV1_2;

    ctx.clear_options(ssl_op_no_ssl_mask);
    let mut options = ssl_op_no_ssl_mask;
    for protocol in protocols {
        let op = match *protocol {
            Protocol::Sslv3 => ssl::SSL_OP_NO_SSLV3,
            Protocol::Tlsv10 => ssl::SSL_OP_NO_TLSV1,
            Protocol::Tlsv11 => ssl::SSL_OP_NO_TLSV1_1,
            Protocol::Tlsv12 => ssl::SSL_OP_NO_TLSV1_2,
            Protocol::__NonExhaustive => unreachable!(),
        };
        options &= !op;
    }
    ctx.set_options(options);
}

pub struct Error(ssl::Error);

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

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

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

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

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

impl From<ErrorStack> for Error {
    fn from(err: ErrorStack) -> Error {
        ssl::Error::Ssl(err).into()
    }
}

pub struct Pkcs12(pkcs12::ParsedPkcs12);

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let pkcs12 = try!(pkcs12::Pkcs12::from_der(buf));
        let parsed = try!(pkcs12.parse(pass));
        Ok(Pkcs12(parsed))
    }
}

pub struct Certificate(X509);

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

pub struct MidHandshakeTlsStream<S>(MidHandshakeSslStream<S>);

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

impl<S> MidHandshakeTlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

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

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

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

impl<S> From<ssl::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: ssl::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            ssl::HandshakeError::SetupFailure(e) => {
                HandshakeError::Failure(Error(ssl::Error::Ssl(e)))
            }
            ssl::HandshakeError::Failure(e) => HandshakeError::Failure(Error(e.into_error())),
            ssl::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

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

pub struct TlsConnectorBuilder(SslConnectorBuilder);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        let ctx = self.0.builder_mut();
        // FIXME clear chain certs to clean up if called multiple times
        try!(ctx.set_certificate(&pkcs12.0.cert));
        try!(ctx.set_private_key(&pkcs12.0.pkey));
        try!(ctx.check_private_key());
        for cert in pkcs12.0.chain {
            try!(ctx.add_extra_chain_cert(cert));
        }
        Ok(())
    }

    pub fn add_root_certificate(&mut self, cert: Certificate) -> Result<(), Error> {
        try!(self.0.builder_mut().cert_store_mut().add_cert(cert.0));
        Ok(())
    }

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

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

#[derive(Clone)]
pub struct TlsConnector(SslConnector);

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        let builder = try!(SslConnectorBuilder::new(SslMethod::tls()));
        Ok(TlsConnectorBuilder(builder))
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.connect(domain, stream));
        Ok(TlsStream(s))
    }

    pub fn connect_no_domain<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.danger_connect_without_providing_domain_for_certificate_verification_and_server_name_indication(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslConnectorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslConnectorBuilder`.
    fn builder(&self) -> &SslConnectorBuilder;

    /// Returns a mutable reference to the inner `SslConnectorBuilder`.
    fn builder_mut(&mut self) -> &mut SslConnectorBuilder;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn from_openssl(builder: SslConnectorBuilder) -> ::TlsConnectorBuilder {
        ::TlsConnectorBuilder(TlsConnectorBuilder(builder))
    }

    fn builder(&self) -> &SslConnectorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslConnectorBuilder {
        &mut (self.0).0
    }
}

pub struct TlsAcceptorBuilder(SslAcceptorBuilder);

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

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

#[derive(Clone)]
pub struct TlsAcceptor(SslAcceptor);

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        let builder = try!(SslAcceptorBuilder::mozilla_intermediate(
            SslMethod::tls(),
            &pkcs12.0.pkey,
            &pkcs12.0.cert,
            &pkcs12.0.chain,
        ));
        Ok(TlsAcceptorBuilder(builder))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.accept(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsAcceptorBuilder`.
pub trait TlsAcceptorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslAcceptorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslAcceptorBuilder`.
    fn builder(&self) -> &SslAcceptorBuilder;

    /// Returns a mutable reference to the inner `SslAcceptorBuilder`.
    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder;
}

impl TlsAcceptorBuilderExt for ::TlsAcceptorBuilder {
    fn from_openssl(builder: SslAcceptorBuilder) -> ::TlsAcceptorBuilder {
        ::TlsAcceptorBuilder(TlsAcceptorBuilder(builder))
    }

    fn builder(&self) -> &SslAcceptorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder {
        &mut (self.0).0
    }
}

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

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

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(self.0.ssl().pending())
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        match self.0.shutdown() {
            Ok(_) |
            Err(ssl::Error::ZeroReturn) => Ok(()),
            Err(ssl::Error::Stream(e)) |
            Err(ssl::Error::WantRead(e)) |
            Err(ssl::Error::WantWrite(e)) => Err(e),
            Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
        }
    }

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

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

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

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

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

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

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

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

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

/// OpenSSL-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying OpenSSL error for inspection.
    fn openssl_error(&self) -> &ssl::Error;
}

impl ErrorExt for ::Error {
    fn openssl_error(&self) -> &ssl::Error {
        &(self.0).0
    }
}
<commit_msg>Add from_pem() for Certificate<commit_after>extern crate openssl;

use std::io;
use std::fmt;
use std::error;
use self::openssl::pkcs12;
use self::openssl::error::ErrorStack;
use self::openssl::ssl::{self, SslMethod, SslConnectorBuilder, SslConnector, SslAcceptorBuilder,
                         SslAcceptor, MidHandshakeSslStream, SslContextBuilder};
use self::openssl::x509::X509;

use Protocol;

fn supported_protocols(protocols: &[Protocol], ctx: &mut SslContextBuilder) {
    // This constant is only defined on OpenSSL 1.0.2 and above, so manually do it.
    let ssl_op_no_ssl_mask = ssl::SSL_OP_NO_SSLV2 | ssl::SSL_OP_NO_SSLV3 | ssl::SSL_OP_NO_TLSV1 |
        ssl::SSL_OP_NO_TLSV1_1 |
        ssl::SSL_OP_NO_TLSV1_2;

    ctx.clear_options(ssl_op_no_ssl_mask);
    let mut options = ssl_op_no_ssl_mask;
    for protocol in protocols {
        let op = match *protocol {
            Protocol::Sslv3 => ssl::SSL_OP_NO_SSLV3,
            Protocol::Tlsv10 => ssl::SSL_OP_NO_TLSV1,
            Protocol::Tlsv11 => ssl::SSL_OP_NO_TLSV1_1,
            Protocol::Tlsv12 => ssl::SSL_OP_NO_TLSV1_2,
            Protocol::__NonExhaustive => unreachable!(),
        };
        options &= !op;
    }
    ctx.set_options(options);
}

pub struct Error(ssl::Error);

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

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

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

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

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

impl From<ErrorStack> for Error {
    fn from(err: ErrorStack) -> Error {
        ssl::Error::Ssl(err).into()
    }
}

pub struct Pkcs12(pkcs12::ParsedPkcs12);

impl Pkcs12 {
    pub fn from_der(buf: &[u8], pass: &str) -> Result<Pkcs12, Error> {
        let pkcs12 = try!(pkcs12::Pkcs12::from_der(buf));
        let parsed = try!(pkcs12.parse(pass));
        Ok(Pkcs12(parsed))
    }
}

pub struct Certificate(X509);

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

pub struct MidHandshakeTlsStream<S>(MidHandshakeSslStream<S>);

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

impl<S> MidHandshakeTlsStream<S> {
    pub fn get_ref(&self) -> &S {
        self.0.get_ref()
    }

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

impl<S> MidHandshakeTlsStream<S>
where
    S: io::Read + io::Write,
{
    pub fn handshake(self) -> Result<TlsStream<S>, HandshakeError<S>> {
        match self.0.handshake() {
            Ok(s) => Ok(TlsStream(s)),
            Err(e) => Err(e.into()),
        }
    }
}

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

impl<S> From<ssl::HandshakeError<S>> for HandshakeError<S> {
    fn from(e: ssl::HandshakeError<S>) -> HandshakeError<S> {
        match e {
            ssl::HandshakeError::SetupFailure(e) => {
                HandshakeError::Failure(Error(ssl::Error::Ssl(e)))
            }
            ssl::HandshakeError::Failure(e) => HandshakeError::Failure(Error(e.into_error())),
            ssl::HandshakeError::Interrupted(s) => {
                HandshakeError::Interrupted(MidHandshakeTlsStream(s))
            }
        }
    }
}

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

pub struct TlsConnectorBuilder(SslConnectorBuilder);

impl TlsConnectorBuilder {
    pub fn identity(&mut self, pkcs12: Pkcs12) -> Result<(), Error> {
        let ctx = self.0.builder_mut();
        // FIXME clear chain certs to clean up if called multiple times
        try!(ctx.set_certificate(&pkcs12.0.cert));
        try!(ctx.set_private_key(&pkcs12.0.pkey));
        try!(ctx.check_private_key());
        for cert in pkcs12.0.chain {
            try!(ctx.add_extra_chain_cert(cert));
        }
        Ok(())
    }

    pub fn add_root_certificate(&mut self, cert: Certificate) -> Result<(), Error> {
        try!(self.0.builder_mut().cert_store_mut().add_cert(cert.0));
        Ok(())
    }

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

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

#[derive(Clone)]
pub struct TlsConnector(SslConnector);

impl TlsConnector {
    pub fn builder() -> Result<TlsConnectorBuilder, Error> {
        let builder = try!(SslConnectorBuilder::new(SslMethod::tls()));
        Ok(TlsConnectorBuilder(builder))
    }

    pub fn connect<S>(&self, domain: &str, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.connect(domain, stream));
        Ok(TlsStream(s))
    }

    pub fn connect_no_domain<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.danger_connect_without_providing_domain_for_certificate_verification_and_server_name_indication(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsConnectorBuilder`.
pub trait TlsConnectorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslConnectorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslConnectorBuilder`.
    fn builder(&self) -> &SslConnectorBuilder;

    /// Returns a mutable reference to the inner `SslConnectorBuilder`.
    fn builder_mut(&mut self) -> &mut SslConnectorBuilder;
}

impl TlsConnectorBuilderExt for ::TlsConnectorBuilder {
    fn from_openssl(builder: SslConnectorBuilder) -> ::TlsConnectorBuilder {
        ::TlsConnectorBuilder(TlsConnectorBuilder(builder))
    }

    fn builder(&self) -> &SslConnectorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslConnectorBuilder {
        &mut (self.0).0
    }
}

pub struct TlsAcceptorBuilder(SslAcceptorBuilder);

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

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

#[derive(Clone)]
pub struct TlsAcceptor(SslAcceptor);

impl TlsAcceptor {
    pub fn builder(pkcs12: Pkcs12) -> Result<TlsAcceptorBuilder, Error> {
        let builder = try!(SslAcceptorBuilder::mozilla_intermediate(
            SslMethod::tls(),
            &pkcs12.0.pkey,
            &pkcs12.0.cert,
            &pkcs12.0.chain,
        ));
        Ok(TlsAcceptorBuilder(builder))
    }

    pub fn accept<S>(&self, stream: S) -> Result<TlsStream<S>, HandshakeError<S>>
    where
        S: io::Read + io::Write,
    {
        let s = try!(self.0.accept(stream));
        Ok(TlsStream(s))
    }
}

/// OpenSSL-specific extensions to `TlsAcceptorBuilder`.
pub trait TlsAcceptorBuilderExt {
    /// Initialize `TlsAcceptorBuilderExt` from an `SslAcceptorBuilder`.
    fn from_openssl(builder: SslAcceptorBuilder) -> Self;

    /// Returns a shared reference to the inner `SslAcceptorBuilder`.
    fn builder(&self) -> &SslAcceptorBuilder;

    /// Returns a mutable reference to the inner `SslAcceptorBuilder`.
    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder;
}

impl TlsAcceptorBuilderExt for ::TlsAcceptorBuilder {
    fn from_openssl(builder: SslAcceptorBuilder) -> ::TlsAcceptorBuilder {
        ::TlsAcceptorBuilder(TlsAcceptorBuilder(builder))
    }

    fn builder(&self) -> &SslAcceptorBuilder {
        &(self.0).0
    }

    fn builder_mut(&mut self) -> &mut SslAcceptorBuilder {
        &mut (self.0).0
    }
}

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

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

impl<S: io::Read + io::Write> TlsStream<S> {
    pub fn buffered_read_size(&self) -> Result<usize, Error> {
        Ok(self.0.ssl().pending())
    }

    pub fn shutdown(&mut self) -> io::Result<()> {
        match self.0.shutdown() {
            Ok(_) |
            Err(ssl::Error::ZeroReturn) => Ok(()),
            Err(ssl::Error::Stream(e)) |
            Err(ssl::Error::WantRead(e)) |
            Err(ssl::Error::WantWrite(e)) => Err(e),
            Err(e) => return Err(io::Error::new(io::ErrorKind::Other, e)),
        }
    }

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

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

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

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

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

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

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

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

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

/// OpenSSL-specific extensions to `Error`
pub trait ErrorExt {
    /// Extract the underlying OpenSSL error for inspection.
    fn openssl_error(&self) -> &ssl::Error;
}

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