<commit_before>#![crate_id = "rustful_macros#0.1-pre"]

#![comment = "Macros for rustful"]
#![license = "MIT"]
#![crate_type = "dylib"]

#![doc(html_root_url = "http://ogeon.github.io/rustful/doc/")]

#![feature(macro_rules, macro_registrar, managed_boxes, quote, phase)]

//!This crate provides some helpful macros for rustful, including `router!` and `routes!`.
//!
//!#`router!`
//!The `router!` macro generates a `Router` containing the provided handlers and routing tree.
//!This can be useful to lower the risk of typing errors, among other things.
//!
//!##Example 1
//!
//!```rust
//!#![feature(phase)]
//!#[phase(syntax)]
//!extern crate rustful_macros;
//!
//!extern crate rustful;
//!
//!...
//!
//!
//!let router = router!{
//!	"/about" => Get: about_us,
//!	"/user/:user" => Get: show_user,
//!	"/product/:name" => Get: show_product,
//!	"/*" => Get: show_error,
//!	"/" => Get: show_welcome
//!};
//!```
//!
//!##Example 2
//!
//!```rust
//!#![feature(phase)]
//!#[phase(syntax)]
//!extern crate rustful_macros;
//!
//!extern crate rustful;
//!
//!...
//!
//!
//!let router = router!{
//!	"/" => Get: show_home,
//!	"home" => Get: show_home,
//!	"user/:username" => {Get: show_user, Post: save_user},
//!	"product" => {
//!		Get: show_all_products,
//!
//!		"json" => Get: send_all_product_data
//!		":id" => {
//!			Get: show_product,
//!			Post | Delete: edit_product,
//!
//!			"json" => Get: send_product_data
//!		}
//!	}
//!};
//!```
//!
//!#`routes!`
//!
//!The `routes!` macro generates a vector of routes, which can be used to create a `Router`.
//!It has the same syntax as `routes!`.
//!
//!```rust
//!#![feature(phase)]
//!#[phase(syntax)]
//!extern crate rustful_macros;
//!
//!extern crate rustful;
//!
//!...
//!
//!
//!let routes = routes!{
//!	"/" => Get: show_home,
//!	"home" => Get: show_home,
//!	"user/:username" => {Get: show_user, Post: save_user},
//!	"product" => {
//!		Get: show_all_products,
//!
//!		"json" => Get: send_all_product_data
//!		":id" => {
//!			Get: show_product,
//!			Post | Delete: edit_product,
//!
//!			"json" => Get: send_product_data
//!		}
//!	}
//!};
//!
//!let router = Router::from_routes(routes);
//!```

extern crate syntax;

use std::path::BytesContainer;

use syntax::{ast, codemap};
use syntax::ext::base::{
	SyntaxExtension, ExtCtxt, MacResult, MacExpr,
	NormalTT, BasicMacroExpander
};
use syntax::ext::build::AstBuilder;
use syntax::ext::quote::rt::ExtParseUtils;
use syntax::parse;
use syntax::parse::token;
use syntax::parse::parser;
use syntax::parse::parser::Parser;
//use syntax::print::pprust;

#[macro_registrar]
#[doc(hidden)]
pub fn macro_registrar(register: |ast::Name, SyntaxExtension|) {
	let expander = box BasicMacroExpander{expander: expand_router, span: None};
	register(token::intern("router"), NormalTT(expander, None));

	let expander = box BasicMacroExpander{expander: expand_routes, span: None};
	register(token::intern("routes"), NormalTT(expander, None));
}

fn expand_router(cx: &mut ExtCtxt, sp: codemap::Span, tts: &[ast::TokenTree]) -> Box<MacResult> {
	let router_ident = cx.ident_of("router");
	let insert_method = cx.ident_of("insert_item");

	let mut calls: Vec<@ast::Stmt> = vec!(
		cx.stmt_let(sp, true, router_ident, quote_expr!(&cx, ::rustful::Router::new()))
	);

	for (path, method, handler) in parse_routes(cx, tts).move_iter() {
		let path_expr = cx.parse_expr(format!("\"{}\"", path));
		let method_expr = cx.expr_path(method);
		let handler_expr = cx.expr_path(handler);
		calls.push(cx.stmt_expr(
			cx.expr_method_call(sp, cx.expr_ident(sp, router_ident), insert_method, vec!(method_expr, path_expr, handler_expr))
		));
	}

	let block = cx.expr_block(cx.block(sp, calls, Some(cx.expr_ident(sp, router_ident))));

	MacExpr::new(block)
}

fn expand_routes(cx: &mut ExtCtxt, sp: codemap::Span, tts: &[ast::TokenTree]) -> Box<MacResult> {
	let routes = parse_routes(cx, tts).move_iter().map(|(path, method, handler)| {
		let path_expr = cx.parse_expr(format!("\"{}\"", path));
		let method_expr = cx.expr_path(method);
		let handler_expr = cx.expr_path(handler);
		mk_tup(sp, vec!(method_expr, path_expr, handler_expr))
	}).collect();

	MacExpr::new(cx.expr_vec(sp, routes))
}

fn parse_routes(cx: &mut ExtCtxt, tts: &[ast::TokenTree]) -> Vec<(String, ast::Path, ast::Path)> {

	let mut parser = parse::new_parser_from_tts(
		cx.parse_sess(), cx.cfg(), Vec::from_slice(tts)
	);

	parse_subroutes("", cx, &mut parser)
}

fn parse_subroutes(base: &str, cx: &mut ExtCtxt, parser: &mut Parser) -> Vec<(String, ast::Path, ast::Path)> {
	let mut routes = Vec::new();

	while !parser.eat(&token::EOF) {
		match parser.parse_optional_str() {
			Some((ref s, _)) => {
				if !parser.eat(&token::FAT_ARROW) {
					parser.expect(&token::FAT_ARROW);
					break;
				}

				let mut new_base = base.to_string();
				match s.container_as_str() {
					Some(s) => {
						new_base.push_str(s.trim_chars('/'));
						new_base.push_str("/");
					},
					None => cx.span_err(parser.span, "invalid path")
				}

				if parser.eat(&token::EOF) {
					cx.span_err(parser.span, "unexpected end of routing tree");
				}

				if parser.eat(&token::LBRACE) {
					let subroutes = parse_subroutes(new_base.as_slice(), cx, parser);
					routes.push_all(subroutes.as_slice());

					if parser.eat(&token::RBRACE) {
						if !parser.eat(&token::COMMA) {
							break;
						}
					} else {
						parser.expect_one_of([token::COMMA, token::RBRACE], []);
					}
				} else {
					for (method, handler) in parse_handler(parser).move_iter() {
						routes.push((new_base.clone(), method, handler))
					}

					if !parser.eat(&token::COMMA) {
						break;
					}
				}
			},
			None => {
				for (method, handler) in parse_handler(parser).move_iter() {
					routes.push((base.to_string(), method, handler))
				}

				if !parser.eat(&token::COMMA) {
					break;
				}
			}
		}
	}

	routes
}

fn parse_handler(parser: &mut Parser) -> Vec<(ast::Path, ast::Path)> {
	let mut methods = Vec::new();

	loop {
		methods.push(parser.parse_path(parser::NoTypesAllowed).path);

		if parser.eat(&token::COLON) {
			break;
		}

		if !parser.eat(&token::BINOP(token::OR)) {
			parser.expect_one_of([token::COLON, token::BINOP(token::OR)], []);
		}
	}

	let handler = parser.parse_path(parser::NoTypesAllowed).path;

	methods.move_iter().map(|m| (m, handler.clone())).collect()
}

fn mk_tup(sp: codemap::Span, content: Vec<@ast::Expr>) -> @ast::Expr {
	dummy_expr(sp, ast::ExprTup(content))
}

fn dummy_expr(sp: codemap::Span, e: ast::Expr_) -> @ast::Expr {
	@ast::Expr {
		id: ast::DUMMY_NODE_ID,
		node: e,
		span: sp,
	}
}

/**
A macro for assigning content types.

It takes a main type, a sub type and a parameter list. Instead of this:

```
response.headers.content_type = Some(MediaType {
	type_: String::from_str("text"),
	subtype: String::from_str("html"),
	parameters: vec!((String::from_str("charset"), String::from_str("UTF-8")))
});
```

it can be written like this:

```
response.headers.content_type = content_type!("text", "html", "charset": "UTF-8");
```

The `"charset": "UTF-8"` part defines the parameter list for the content type.
It may contain more than one parameter, or be omitted:

```
response.headers.content_type = content_type!("application", "octet-stream", "type": "image/gif", "padding": "4");
```

```
response.headers.content_type = content_type!("image", "png");
```
**/
#[macro_export]
macro_rules! content_type(
	($main_type:expr, $sub_type:expr) => ({
		Some(::http::headers::content_type::MediaType {
			type_: String::from_str($main_type),
			subtype: String::from_str($sub_type),
			parameters: Vec::new()
		})
	});

	($main_type:expr, $sub_type:expr, $($param:expr: $value:expr),+) => ({
		Some(::http::headers::content_type::MediaType {
			type_: String::from_str($main_type),
			subtype: String::from_str($sub_type),
			parameters: vec!( $( (String::from_str($param), String::from_str($value)) ),+ )
		})
	});
)<commit_msg>router! and routes! macros takes expressions<commit_after>#![crate_id = "rustful_macros#0.1-pre"]

#![comment = "Macros for rustful"]
#![license = "MIT"]
#![crate_type = "dylib"]

#![doc(html_root_url = "http://ogeon.github.io/rustful/doc/")]

#![feature(macro_rules, macro_registrar, managed_boxes, quote, phase)]

//!This crate provides some helpful macros for rustful, including `router!` and `routes!`.
//!
//!#`router!`
//!The `router!` macro generates a `Router` containing the provided handlers and routing tree.
//!This can be useful to lower the risk of typing errors, among other things.
//!
//!##Example 1
//!
//!```rust
//!#![feature(phase)]
//!#[phase(syntax)]
//!extern crate rustful_macros;
//!
//!extern crate rustful;
//!
//!...
//!
//!
//!let router = router!{
//!	"/about" => Get: about_us,
//!	"/user/:user" => Get: show_user,
//!	"/product/:name" => Get: show_product,
//!	"/*" => Get: show_error,
//!	"/" => Get: show_welcome
//!};
//!```
//!
//!##Example 2
//!
//!```rust
//!#![feature(phase)]
//!#[phase(syntax)]
//!extern crate rustful_macros;
//!
//!extern crate rustful;
//!
//!...
//!
//!
//!let router = router!{
//!	"/" => Get: show_home,
//!	"home" => Get: show_home,
//!	"user/:username" => {Get: show_user, Post: save_user},
//!	"product" => {
//!		Get: show_all_products,
//!
//!		"json" => Get: send_all_product_data
//!		":id" => {
//!			Get: show_product,
//!			Post | Delete: edit_product,
//!
//!			"json" => Get: send_product_data
//!		}
//!	}
//!};
//!```
//!
//!#`routes!`
//!
//!The `routes!` macro generates a vector of routes, which can be used to create a `Router`.
//!It has the same syntax as `routes!`.
//!
//!```rust
//!#![feature(phase)]
//!#[phase(syntax)]
//!extern crate rustful_macros;
//!
//!extern crate rustful;
//!
//!...
//!
//!
//!let routes = routes!{
//!	"/" => Get: show_home,
//!	"home" => Get: show_home,
//!	"user/:username" => {Get: show_user, Post: save_user},
//!	"product" => {
//!		Get: show_all_products,
//!
//!		"json" => Get: send_all_product_data
//!		":id" => {
//!			Get: show_product,
//!			Post | Delete: edit_product,
//!
//!			"json" => Get: send_product_data
//!		}
//!	}
//!};
//!
//!let router = Router::from_routes(routes);
//!```

extern crate syntax;

use std::path::BytesContainer;

use syntax::{ast, codemap};
use syntax::ext::base::{
	SyntaxExtension, ExtCtxt, MacResult, MacExpr,
	NormalTT, BasicMacroExpander
};
use syntax::ext::build::AstBuilder;
use syntax::ext::quote::rt::ExtParseUtils;
use syntax::parse;
use syntax::parse::token;
use syntax::parse::parser;
use syntax::parse::parser::Parser;
//use syntax::print::pprust;

#[macro_registrar]
#[doc(hidden)]
pub fn macro_registrar(register: |ast::Name, SyntaxExtension|) {
	let expander = box BasicMacroExpander{expander: expand_router, span: None};
	register(token::intern("router"), NormalTT(expander, None));

	let expander = box BasicMacroExpander{expander: expand_routes, span: None};
	register(token::intern("routes"), NormalTT(expander, None));
}

fn expand_router(cx: &mut ExtCtxt, sp: codemap::Span, tts: &[ast::TokenTree]) -> Box<MacResult> {
	let router_ident = cx.ident_of("router");
	let insert_method = cx.ident_of("insert_item");

	let mut calls: Vec<@ast::Stmt> = vec!(
		cx.stmt_let(sp, true, router_ident, quote_expr!(&cx, ::rustful::Router::new()))
	);

	for (path, method, handler) in parse_routes(cx, tts).move_iter() {
		let path_expr = cx.parse_expr(format!("\"{}\"", path));
		let method_expr = cx.expr_path(method);
		calls.push(cx.stmt_expr(
			cx.expr_method_call(sp, cx.expr_ident(sp, router_ident), insert_method, vec!(method_expr, path_expr, handler))
		));
	}

	let block = cx.expr_block(cx.block(sp, calls, Some(cx.expr_ident(sp, router_ident))));

	MacExpr::new(block)
}

fn expand_routes(cx: &mut ExtCtxt, sp: codemap::Span, tts: &[ast::TokenTree]) -> Box<MacResult> {
	let routes = parse_routes(cx, tts).move_iter().map(|(path, method, handler)| {
		let path_expr = cx.parse_expr(format!("\"{}\"", path));
		let method_expr = cx.expr_path(method);
		mk_tup(sp, vec!(method_expr, path_expr, handler))
	}).collect();

	MacExpr::new(cx.expr_vec(sp, routes))
}

fn parse_routes(cx: &mut ExtCtxt, tts: &[ast::TokenTree]) -> Vec<(String, ast::Path, @ast::Expr)> {

	let mut parser = parse::new_parser_from_tts(
		cx.parse_sess(), cx.cfg(), Vec::from_slice(tts)
	);

	parse_subroutes("", cx, &mut parser)
}

fn parse_subroutes(base: &str, cx: &mut ExtCtxt, parser: &mut Parser) -> Vec<(String, ast::Path, @ast::Expr)> {
	let mut routes = Vec::new();

	while !parser.eat(&token::EOF) {
		match parser.parse_optional_str() {
			Some((ref s, _)) => {
				if !parser.eat(&token::FAT_ARROW) {
					parser.expect(&token::FAT_ARROW);
					break;
				}

				let mut new_base = base.to_string();
				match s.container_as_str() {
					Some(s) => {
						new_base.push_str(s.trim_chars('/'));
						new_base.push_str("/");
					},
					None => cx.span_err(parser.span, "invalid path")
				}

				if parser.eat(&token::EOF) {
					cx.span_err(parser.span, "unexpected end of routing tree");
				}

				if parser.eat(&token::LBRACE) {
					let subroutes = parse_subroutes(new_base.as_slice(), cx, parser);
					routes.push_all(subroutes.as_slice());

					if parser.eat(&token::RBRACE) {
						if !parser.eat(&token::COMMA) {
							break;
						}
					} else {
						parser.expect_one_of([token::COMMA, token::RBRACE], []);
					}
				} else {
					for (method, handler) in parse_handler(parser).move_iter() {
						routes.push((new_base.clone(), method, handler))
					}

					if !parser.eat(&token::COMMA) {
						break;
					}
				}
			},
			None => {
				for (method, handler) in parse_handler(parser).move_iter() {
					routes.push((base.to_string(), method, handler))
				}

				if !parser.eat(&token::COMMA) {
					break;
				}
			}
		}
	}

	routes
}

fn parse_handler(parser: &mut Parser) -> Vec<(ast::Path, @ast::Expr)> {
	let mut methods = Vec::new();

	loop {
		methods.push(parser.parse_path(parser::NoTypesAllowed).path);

		if parser.eat(&token::COLON) {
			break;
		}

		if !parser.eat(&token::BINOP(token::OR)) {
			parser.expect_one_of([token::COLON, token::BINOP(token::OR)], []);
		}
	}

	let handler = parser.parse_expr();

	methods.move_iter().map(|m| (m, handler.clone())).collect()
}

fn mk_tup(sp: codemap::Span, content: Vec<@ast::Expr>) -> @ast::Expr {
	dummy_expr(sp, ast::ExprTup(content))
}

fn dummy_expr(sp: codemap::Span, e: ast::Expr_) -> @ast::Expr {
	@ast::Expr {
		id: ast::DUMMY_NODE_ID,
		node: e,
		span: sp,
	}
}

/**
A macro for assigning content types.

It takes a main type, a sub type and a parameter list. Instead of this:

```
response.headers.content_type = Some(MediaType {
	type_: String::from_str("text"),
	subtype: String::from_str("html"),
	parameters: vec!((String::from_str("charset"), String::from_str("UTF-8")))
});
```

it can be written like this:

```
response.headers.content_type = content_type!("text", "html", "charset": "UTF-8");
```

The `"charset": "UTF-8"` part defines the parameter list for the content type.
It may contain more than one parameter, or be omitted:

```
response.headers.content_type = content_type!("application", "octet-stream", "type": "image/gif", "padding": "4");
```

```
response.headers.content_type = content_type!("image", "png");
```
**/
#[macro_export]
macro_rules! content_type(
	($main_type:expr, $sub_type:expr) => ({
		Some(::http::headers::content_type::MediaType {
			type_: String::from_str($main_type),
			subtype: String::from_str($sub_type),
			parameters: Vec::new()
		})
	});

	($main_type:expr, $sub_type:expr, $($param:expr: $value:expr),+) => ({
		Some(::http::headers::content_type::MediaType {
			type_: String::from_str($main_type),
			subtype: String::from_str($sub_type),
			parameters: vec!( $( (String::from_str($param), String::from_str($value)) ),+ )
		})
	});
)<|endoftext|>


<commit_before>use syntax::abi::Abi;
use syntax::ast;
use syntax::codemap::{DUMMY_SP, Span, respan};
use syntax::ptr::P;

use attr::AttrBuilder;
use block::BlockBuilder;
use generics::GenericsBuilder;
use ident::ToIdent;
use invoke::{Invoke, Identity};
use lifetime::IntoLifetime;

//////////////////////////////////////////////////////////////////////////////

pub struct MethodBuilder<F=Identity> {
    callback: F,
    span: Span,
    attrs: Vec<ast::Attribute>,
    abi: Abi,
    generics: ast::Generics,
    unsafety: ast::Unsafety,
    id: ast::Ident,
    vis: ast::Visibility,
}

impl MethodBuilder {
    pub fn new<I: ToIdent>(id: I) -> Self {
        MethodBuilder::new_with_callback(id, Identity)
    }
}

impl<F> MethodBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    pub fn new_with_callback<I>(id: I, callback: F) -> Self
        where I: ToIdent,
    {
        MethodBuilder {
            callback: callback,
            span: DUMMY_SP,
            attrs: vec![],
            abi: Abi::Rust,
            generics: GenericsBuilder::new().build(),
            unsafety: ast::Unsafety::Normal,
            id: id.to_ident(),
            vis: ast::Visibility::Inherited,
        }
    }

    pub fn span(mut self, span: Span) -> Self {
        self.span = span;
        self
    }

    pub fn with_attr(mut self, attr: ast::Attribute) -> Self {
        self.attrs.push(attr);
        self
    }

    pub fn attr(self) -> AttrBuilder<Self> {
        AttrBuilder::new_with_callback(self)
    }

    pub fn unsafe_(mut self) -> Self {
        self.unsafety = ast::Unsafety::Normal;
        self
    }

    pub fn abi(mut self, abi: Abi) -> Self {
        self.abi = abi;
        self
    }

    pub fn vis(mut self, vis: ast::Visibility) -> Self {
        self.vis = vis;
        self
    }

    pub fn with_generics(mut self, generics: ast::Generics) -> Self {
        self.generics = generics;
        self
    }

    pub fn generics(self) -> GenericsBuilder<Self> {
        GenericsBuilder::new_with_callback(self)
    }

    pub fn self_(self) -> SelfBuilder<Self> {
        SelfBuilder::new_with_callback(self)
    }
}

impl<F> Invoke<ast::Attribute> for MethodBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    type Result = Self;

    fn invoke(self, attr: ast::Attribute) -> Self {
        self.with_attr(attr)
    }
}

impl<F> Invoke<ast::Generics> for MethodBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    type Result = Self;

    fn invoke(mut self, generics: ast::Generics) -> Self {
        self.with_generics(generics)
    }
}

impl<F> Invoke<ast::ExplicitSelf> for MethodBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    type Result = MethodSelfBuilder<F>;

    fn invoke(self, explicit_self: ast::ExplicitSelf) -> MethodSelfBuilder<F> {
        MethodSelfBuilder {
            builder: self,
            explicit_self: explicit_self,
        }
    }
}

//////////////////////////////////////////////////////////////////////////////

pub struct MethodSelfBuilder<F> {
    builder: MethodBuilder<F>,
    explicit_self: ast::ExplicitSelf,
}

impl<F> Invoke<P<ast::FnDecl>> for MethodSelfBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    type Result = BlockBuilder<MethodSelfFnDeclBuilder<F>>;

    fn invoke(self, fn_decl: P<ast::FnDecl>) -> BlockBuilder<MethodSelfFnDeclBuilder<F>> {
        BlockBuilder::new_with_callback(MethodSelfFnDeclBuilder {
            builder: self.builder,
            explicit_self: self.explicit_self,
            fn_decl: fn_decl,
        })
    }
}

//////////////////////////////////////////////////////////////////////////////

pub struct MethodSelfFnDeclBuilder<F> {
    builder: MethodBuilder<F>,
    explicit_self: ast::ExplicitSelf,
    fn_decl: P<ast::FnDecl>,
}

impl<F> Invoke<P<ast::Block>> for MethodSelfFnDeclBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    type Result = F::Result;

    fn invoke(self, block: P<ast::Block>) -> F::Result {
        let method_sig = ast::MethodSig {
            unsafety: self.builder.unsafety,
            abi: self.builder.abi,
            decl: self.fn_decl,
            generics: self.builder.generics,
            explicit_self: self.explicit_self,
        };

        self.builder.callback.invoke(P(ast::ImplItem {
            id: ast::DUMMY_NODE_ID,
            ident: self.builder.id,
            vis: self.builder.vis,
            attrs: self.builder.attrs,
            node: ast::MethodImplItem(method_sig, block),
            span: self.builder.span,
        }))
    }
}

//////////////////////////////////////////////////////////////////////////////

pub struct SelfBuilder<F> {
    callback: F,
    span: Span,
}

impl<F> SelfBuilder<F>
    where F: Invoke<ast::ExplicitSelf>,
{
    pub fn new_with_callback(callback: F) -> Self {
        SelfBuilder {
            callback: callback,
            span: DUMMY_SP,
        }
    }

    pub fn build(self, self_: ast::ExplicitSelf) -> F::Result {
        self.callback.invoke(self_)
    }

    pub fn span(mut self, span: Span) -> Self {
        self.span = span;
        self
    }

    pub fn build_self_(self, self_: ast::ExplicitSelf_) -> F::Result {
        let self_ = respan(self.span, self_);
        self.build(self_)
    }

    pub fn static_(self) -> F::Result {
        self.build_self_(ast::ExplicitSelf_::SelfStatic)
    }

    pub fn value(self) -> F::Result {
        self.build_self_(ast::ExplicitSelf_::SelfValue("self".to_ident()))
    }

    pub fn ref_(self) -> F::Result {
        self.build_self_(ast::ExplicitSelf_::SelfRegion(
            None,
            ast::Mutability::MutImmutable,
            "self".to_ident(),
        ))
    }

    pub fn ref_lifetime<L>(self, lifetime: L) -> F::Result
        where L: IntoLifetime,
    {
        self.build_self_(ast::ExplicitSelf_::SelfRegion(
            Some(lifetime.into_lifetime()),
            ast::Mutability::MutImmutable,
            "self".to_ident(),
        ))
    }

    pub fn ref_mut(self) -> F::Result {
        self.build_self_(ast::ExplicitSelf_::SelfRegion(
            None,
            ast::Mutability::MutMutable,
            "self".to_ident(),
        ))
    }

    pub fn ref_mut_lifetime<L>(self, lifetime: L) -> F::Result
        where L: IntoLifetime,
    {
        self.build_self_(ast::ExplicitSelf_::SelfRegion(
            Some(lifetime.into_lifetime()),
            ast::Mutability::MutMutable,
            "self".to_ident(),
        ))
    }

    /*
    pub fn ty(self) -> TyBuilder<F::Result> {
        TyBuilder::new_with_callback(self)
    }
    */
}

impl<F> Invoke<P<ast::Ty>> for SelfBuilder<F>
    where F: Invoke<ast::ExplicitSelf>,
{
    type Result = F::Result;

    fn invoke(self, ty: P<ast::Ty>) -> F::Result {
        self.build_self_(ast::ExplicitSelf_::SelfExplicit(ty, "self".to_ident()))
    }
}
<commit_msg>Silence a warning<commit_after>use syntax::abi::Abi;
use syntax::ast;
use syntax::codemap::{DUMMY_SP, Span, respan};
use syntax::ptr::P;

use attr::AttrBuilder;
use block::BlockBuilder;
use generics::GenericsBuilder;
use ident::ToIdent;
use invoke::{Invoke, Identity};
use lifetime::IntoLifetime;

//////////////////////////////////////////////////////////////////////////////

pub struct MethodBuilder<F=Identity> {
    callback: F,
    span: Span,
    attrs: Vec<ast::Attribute>,
    abi: Abi,
    generics: ast::Generics,
    unsafety: ast::Unsafety,
    id: ast::Ident,
    vis: ast::Visibility,
}

impl MethodBuilder {
    pub fn new<I: ToIdent>(id: I) -> Self {
        MethodBuilder::new_with_callback(id, Identity)
    }
}

impl<F> MethodBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    pub fn new_with_callback<I>(id: I, callback: F) -> Self
        where I: ToIdent,
    {
        MethodBuilder {
            callback: callback,
            span: DUMMY_SP,
            attrs: vec![],
            abi: Abi::Rust,
            generics: GenericsBuilder::new().build(),
            unsafety: ast::Unsafety::Normal,
            id: id.to_ident(),
            vis: ast::Visibility::Inherited,
        }
    }

    pub fn span(mut self, span: Span) -> Self {
        self.span = span;
        self
    }

    pub fn with_attr(mut self, attr: ast::Attribute) -> Self {
        self.attrs.push(attr);
        self
    }

    pub fn attr(self) -> AttrBuilder<Self> {
        AttrBuilder::new_with_callback(self)
    }

    pub fn unsafe_(mut self) -> Self {
        self.unsafety = ast::Unsafety::Normal;
        self
    }

    pub fn abi(mut self, abi: Abi) -> Self {
        self.abi = abi;
        self
    }

    pub fn vis(mut self, vis: ast::Visibility) -> Self {
        self.vis = vis;
        self
    }

    pub fn with_generics(mut self, generics: ast::Generics) -> Self {
        self.generics = generics;
        self
    }

    pub fn generics(self) -> GenericsBuilder<Self> {
        GenericsBuilder::new_with_callback(self)
    }

    pub fn self_(self) -> SelfBuilder<Self> {
        SelfBuilder::new_with_callback(self)
    }
}

impl<F> Invoke<ast::Attribute> for MethodBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    type Result = Self;

    fn invoke(self, attr: ast::Attribute) -> Self {
        self.with_attr(attr)
    }
}

impl<F> Invoke<ast::Generics> for MethodBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    type Result = Self;

    fn invoke(self, generics: ast::Generics) -> Self {
        self.with_generics(generics)
    }
}

impl<F> Invoke<ast::ExplicitSelf> for MethodBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    type Result = MethodSelfBuilder<F>;

    fn invoke(self, explicit_self: ast::ExplicitSelf) -> MethodSelfBuilder<F> {
        MethodSelfBuilder {
            builder: self,
            explicit_self: explicit_self,
        }
    }
}

//////////////////////////////////////////////////////////////////////////////

pub struct MethodSelfBuilder<F> {
    builder: MethodBuilder<F>,
    explicit_self: ast::ExplicitSelf,
}

impl<F> Invoke<P<ast::FnDecl>> for MethodSelfBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    type Result = BlockBuilder<MethodSelfFnDeclBuilder<F>>;

    fn invoke(self, fn_decl: P<ast::FnDecl>) -> BlockBuilder<MethodSelfFnDeclBuilder<F>> {
        BlockBuilder::new_with_callback(MethodSelfFnDeclBuilder {
            builder: self.builder,
            explicit_self: self.explicit_self,
            fn_decl: fn_decl,
        })
    }
}

//////////////////////////////////////////////////////////////////////////////

pub struct MethodSelfFnDeclBuilder<F> {
    builder: MethodBuilder<F>,
    explicit_self: ast::ExplicitSelf,
    fn_decl: P<ast::FnDecl>,
}

impl<F> Invoke<P<ast::Block>> for MethodSelfFnDeclBuilder<F>
    where F: Invoke<P<ast::ImplItem>>,
{
    type Result = F::Result;

    fn invoke(self, block: P<ast::Block>) -> F::Result {
        let method_sig = ast::MethodSig {
            unsafety: self.builder.unsafety,
            abi: self.builder.abi,
            decl: self.fn_decl,
            generics: self.builder.generics,
            explicit_self: self.explicit_self,
        };

        self.builder.callback.invoke(P(ast::ImplItem {
            id: ast::DUMMY_NODE_ID,
            ident: self.builder.id,
            vis: self.builder.vis,
            attrs: self.builder.attrs,
            node: ast::MethodImplItem(method_sig, block),
            span: self.builder.span,
        }))
    }
}

//////////////////////////////////////////////////////////////////////////////

pub struct SelfBuilder<F> {
    callback: F,
    span: Span,
}

impl<F> SelfBuilder<F>
    where F: Invoke<ast::ExplicitSelf>,
{
    pub fn new_with_callback(callback: F) -> Self {
        SelfBuilder {
            callback: callback,
            span: DUMMY_SP,
        }
    }

    pub fn build(self, self_: ast::ExplicitSelf) -> F::Result {
        self.callback.invoke(self_)
    }

    pub fn span(mut self, span: Span) -> Self {
        self.span = span;
        self
    }

    pub fn build_self_(self, self_: ast::ExplicitSelf_) -> F::Result {
        let self_ = respan(self.span, self_);
        self.build(self_)
    }

    pub fn static_(self) -> F::Result {
        self.build_self_(ast::ExplicitSelf_::SelfStatic)
    }

    pub fn value(self) -> F::Result {
        self.build_self_(ast::ExplicitSelf_::SelfValue("self".to_ident()))
    }

    pub fn ref_(self) -> F::Result {
        self.build_self_(ast::ExplicitSelf_::SelfRegion(
            None,
            ast::Mutability::MutImmutable,
            "self".to_ident(),
        ))
    }

    pub fn ref_lifetime<L>(self, lifetime: L) -> F::Result
        where L: IntoLifetime,
    {
        self.build_self_(ast::ExplicitSelf_::SelfRegion(
            Some(lifetime.into_lifetime()),
            ast::Mutability::MutImmutable,
            "self".to_ident(),
        ))
    }

    pub fn ref_mut(self) -> F::Result {
        self.build_self_(ast::ExplicitSelf_::SelfRegion(
            None,
            ast::Mutability::MutMutable,
            "self".to_ident(),
        ))
    }

    pub fn ref_mut_lifetime<L>(self, lifetime: L) -> F::Result
        where L: IntoLifetime,
    {
        self.build_self_(ast::ExplicitSelf_::SelfRegion(
            Some(lifetime.into_lifetime()),
            ast::Mutability::MutMutable,
            "self".to_ident(),
        ))
    }

    /*
    pub fn ty(self) -> TyBuilder<F::Result> {
        TyBuilder::new_with_callback(self)
    }
    */
}

impl<F> Invoke<P<ast::Ty>> for SelfBuilder<F>
    where F: Invoke<ast::ExplicitSelf>,
{
    type Result = F::Result;

    fn invoke(self, ty: P<ast::Ty>) -> F::Result {
        self.build_self_(ast::ExplicitSelf_::SelfExplicit(ty, "self".to_ident()))
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Various utility functions used throughout rustbuild.
//!
//! Simple things like testing the various filesystem operations here and there,
//! not a lot of interesting happenings here unfortunately.

use std::env;
use std::ffi::OsString;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::Instant;

use filetime::FileTime;

/// Returns the `name` as the filename of a static library for `target`.
pub fn staticlib(name: &str, target: &str) -> String {
    if target.contains("windows") {
        format!("{}.lib", name)
    } else {
        format!("lib{}.a", name)
    }
}

/// Returns the last-modified time for `path`, or zero if it doesn't exist.
pub fn mtime(path: &Path) -> FileTime {
    fs::metadata(path).map(|f| {
        FileTime::from_last_modification_time(&f)
    }).unwrap_or(FileTime::zero())
}

/// Copies a file from `src` to `dst`, attempting to use hard links and then
/// falling back to an actually filesystem copy if necessary.
pub fn copy(src: &Path, dst: &Path) {
    let res = fs::hard_link(src, dst);
    let res = res.or_else(|_| fs::copy(src, dst).map(|_| ()));
    if let Err(e) = res {
        panic!("failed to copy `{}` to `{}`: {}", src.display(),
               dst.display(), e)
    }
}

/// Copies the `src` directory recursively to `dst`. Both are assumed to exist
/// when this function is called.
pub fn cp_r(src: &Path, dst: &Path) {
    for f in t!(fs::read_dir(src)) {
        let f = t!(f);
        let path = f.path();
        let name = path.file_name().unwrap();
        let dst = dst.join(name);
        if t!(f.file_type()).is_dir() {
            t!(fs::create_dir_all(&dst));
            cp_r(&path, &dst);
        } else {
            let _ = fs::remove_file(&dst);
            copy(&path, &dst);
        }
    }
}

/// Copies the `src` directory recursively to `dst`. Both are assumed to exist
/// when this function is called. Unwanted files or directories can be skipped
/// by returning `false` from the filter function.
pub fn cp_filtered<F: Fn(&Path) -> bool>(src: &Path, dst: &Path, filter: &F) {
    // Inner function does the actual work
    fn recurse<F: Fn(&Path) -> bool>(src: &Path, dst: &Path, relative: &Path, filter: &F) {
        for f in t!(fs::read_dir(src)) {
            let f = t!(f);
            let path = f.path();
            let name = path.file_name().unwrap();
            let dst = dst.join(name);
            let relative = relative.join(name);
            // Only copy file or directory if the filter function returns true
            if filter(&relative) {
                if t!(f.file_type()).is_dir() {
                    let _ = fs::remove_dir_all(&dst);
                    t!(fs::create_dir(&dst));
                    recurse(&path, &dst, &relative, filter);
                } else {
                    let _ = fs::remove_file(&dst);
                    copy(&path, &dst);
                }
            }
        }
    }
    // Immediately recurse with an empty relative path
    recurse(src, dst, Path::new(""), filter)
}

/// Given an executable called `name`, return the filename for the
/// executable for a particular target.
pub fn exe(name: &str, target: &str) -> String {
    if target.contains("windows") {
        format!("{}.exe", name)
    } else {
        name.to_string()
    }
}

/// Returns whether the file name given looks like a dynamic library.
pub fn is_dylib(name: &str) -> bool {
    name.ends_with(".dylib") || name.ends_with(".so") || name.ends_with(".dll")
}

/// Returns the corresponding relative library directory that the compiler's
/// dylibs will be found in.
pub fn libdir(target: &str) -> &'static str {
    if target.contains("windows") {"bin"} else {"lib"}
}

/// Adds a list of lookup paths to `cmd`'s dynamic library lookup path.
pub fn add_lib_path(path: Vec<PathBuf>, cmd: &mut Command) {
    let mut list = dylib_path();
    for path in path {
        list.insert(0, path);
    }
    cmd.env(dylib_path_var(), t!(env::join_paths(list)));
}

/// Returns whether `dst` is up to date given that the file or files in `src`
/// are used to generate it.
///
/// Uses last-modified time checks to verify this.
pub fn up_to_date(src: &Path, dst: &Path) -> bool {
    let threshold = mtime(dst);
    let meta = match fs::metadata(src) {
        Ok(meta) => meta,
        Err(e) => panic!("source {:?} failed to get metadata: {}", src, e),
    };
    if meta.is_dir() {
        dir_up_to_date(src, &threshold)
    } else {
        FileTime::from_last_modification_time(&meta) <= threshold
    }
}

fn dir_up_to_date(src: &Path, threshold: &FileTime) -> bool {
    t!(fs::read_dir(src)).map(|e| t!(e)).all(|e| {
        let meta = t!(e.metadata());
        if meta.is_dir() {
            dir_up_to_date(&e.path(), threshold)
        } else {
            FileTime::from_last_modification_time(&meta) < *threshold
        }
    })
}

/// Returns the environment variable which the dynamic library lookup path
/// resides in for this platform.
pub fn dylib_path_var() -> &'static str {
    if cfg!(target_os = "windows") {
        "PATH"
    } else if cfg!(target_os = "macos") {
        "DYLD_LIBRARY_PATH"
    } else {
        "LD_LIBRARY_PATH"
    }
}

/// Parses the `dylib_path_var()` environment variable, returning a list of
/// paths that are members of this lookup path.
pub fn dylib_path() -> Vec<PathBuf> {
    env::split_paths(&env::var_os(dylib_path_var()).unwrap_or(OsString::new()))
        .collect()
}

/// `push` all components to `buf`. On windows, append `.exe` to the last component.
pub fn push_exe_path(mut buf: PathBuf, components: &[&str]) -> PathBuf {
    let (&file, components) = components.split_last().expect("at least one component required");
    let mut file = file.to_owned();

    if cfg!(windows) {
        file.push_str(".exe");
    }

    for c in components {
        buf.push(c);
    }

    buf.push(file);

    buf
}

pub struct TimeIt(Instant);

/// Returns an RAII structure that prints out how long it took to drop.
pub fn timeit() -> TimeIt {
    TimeIt(Instant::now())
}

impl Drop for TimeIt {
    fn drop(&mut self) {
        let time = self.0.elapsed();
        println!("\tfinished in {}.{:03}",
                 time.as_secs(),
                 time.subsec_nanos() / 1_000_000);
    }
}
<commit_msg>Auto merge of #38405 - alexcrichton:rustbuild-fixes, r=japaric<commit_after>// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Various utility functions used throughout rustbuild.
//!
//! Simple things like testing the various filesystem operations here and there,
//! not a lot of interesting happenings here unfortunately.

use std::env;
use std::ffi::OsString;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::time::Instant;

use filetime::FileTime;

/// Returns the `name` as the filename of a static library for `target`.
pub fn staticlib(name: &str, target: &str) -> String {
    if target.contains("windows") {
        format!("{}.lib", name)
    } else {
        format!("lib{}.a", name)
    }
}

/// Returns the last-modified time for `path`, or zero if it doesn't exist.
pub fn mtime(path: &Path) -> FileTime {
    fs::metadata(path).map(|f| {
        FileTime::from_last_modification_time(&f)
    }).unwrap_or(FileTime::zero())
}

/// Copies a file from `src` to `dst`, attempting to use hard links and then
/// falling back to an actually filesystem copy if necessary.
pub fn copy(src: &Path, dst: &Path) {
    // A call to `hard_link` will fail if `dst` exists, so remove it if it
    // already exists so we can try to help `hard_link` succeed.
    let _ = fs::remove_file(&dst);

    // Attempt to "easy copy" by creating a hard link (symlinks don't work on
    // windows), but if that fails just fall back to a slow `copy` operation.
    let res = fs::hard_link(src, dst);
    let res = res.or_else(|_| fs::copy(src, dst).map(|_| ()));
    if let Err(e) = res {
        panic!("failed to copy `{}` to `{}`: {}", src.display(),
               dst.display(), e)
    }
}

/// Copies the `src` directory recursively to `dst`. Both are assumed to exist
/// when this function is called.
pub fn cp_r(src: &Path, dst: &Path) {
    for f in t!(fs::read_dir(src)) {
        let f = t!(f);
        let path = f.path();
        let name = path.file_name().unwrap();
        let dst = dst.join(name);
        if t!(f.file_type()).is_dir() {
            t!(fs::create_dir_all(&dst));
            cp_r(&path, &dst);
        } else {
            let _ = fs::remove_file(&dst);
            copy(&path, &dst);
        }
    }
}

/// Copies the `src` directory recursively to `dst`. Both are assumed to exist
/// when this function is called. Unwanted files or directories can be skipped
/// by returning `false` from the filter function.
pub fn cp_filtered<F: Fn(&Path) -> bool>(src: &Path, dst: &Path, filter: &F) {
    // Inner function does the actual work
    fn recurse<F: Fn(&Path) -> bool>(src: &Path, dst: &Path, relative: &Path, filter: &F) {
        for f in t!(fs::read_dir(src)) {
            let f = t!(f);
            let path = f.path();
            let name = path.file_name().unwrap();
            let dst = dst.join(name);
            let relative = relative.join(name);
            // Only copy file or directory if the filter function returns true
            if filter(&relative) {
                if t!(f.file_type()).is_dir() {
                    let _ = fs::remove_dir_all(&dst);
                    t!(fs::create_dir(&dst));
                    recurse(&path, &dst, &relative, filter);
                } else {
                    let _ = fs::remove_file(&dst);
                    copy(&path, &dst);
                }
            }
        }
    }
    // Immediately recurse with an empty relative path
    recurse(src, dst, Path::new(""), filter)
}

/// Given an executable called `name`, return the filename for the
/// executable for a particular target.
pub fn exe(name: &str, target: &str) -> String {
    if target.contains("windows") {
        format!("{}.exe", name)
    } else {
        name.to_string()
    }
}

/// Returns whether the file name given looks like a dynamic library.
pub fn is_dylib(name: &str) -> bool {
    name.ends_with(".dylib") || name.ends_with(".so") || name.ends_with(".dll")
}

/// Returns the corresponding relative library directory that the compiler's
/// dylibs will be found in.
pub fn libdir(target: &str) -> &'static str {
    if target.contains("windows") {"bin"} else {"lib"}
}

/// Adds a list of lookup paths to `cmd`'s dynamic library lookup path.
pub fn add_lib_path(path: Vec<PathBuf>, cmd: &mut Command) {
    let mut list = dylib_path();
    for path in path {
        list.insert(0, path);
    }
    cmd.env(dylib_path_var(), t!(env::join_paths(list)));
}

/// Returns whether `dst` is up to date given that the file or files in `src`
/// are used to generate it.
///
/// Uses last-modified time checks to verify this.
pub fn up_to_date(src: &Path, dst: &Path) -> bool {
    let threshold = mtime(dst);
    let meta = match fs::metadata(src) {
        Ok(meta) => meta,
        Err(e) => panic!("source {:?} failed to get metadata: {}", src, e),
    };
    if meta.is_dir() {
        dir_up_to_date(src, &threshold)
    } else {
        FileTime::from_last_modification_time(&meta) <= threshold
    }
}

fn dir_up_to_date(src: &Path, threshold: &FileTime) -> bool {
    t!(fs::read_dir(src)).map(|e| t!(e)).all(|e| {
        let meta = t!(e.metadata());
        if meta.is_dir() {
            dir_up_to_date(&e.path(), threshold)
        } else {
            FileTime::from_last_modification_time(&meta) < *threshold
        }
    })
}

/// Returns the environment variable which the dynamic library lookup path
/// resides in for this platform.
pub fn dylib_path_var() -> &'static str {
    if cfg!(target_os = "windows") {
        "PATH"
    } else if cfg!(target_os = "macos") {
        "DYLD_LIBRARY_PATH"
    } else {
        "LD_LIBRARY_PATH"
    }
}

/// Parses the `dylib_path_var()` environment variable, returning a list of
/// paths that are members of this lookup path.
pub fn dylib_path() -> Vec<PathBuf> {
    env::split_paths(&env::var_os(dylib_path_var()).unwrap_or(OsString::new()))
        .collect()
}

/// `push` all components to `buf`. On windows, append `.exe` to the last component.
pub fn push_exe_path(mut buf: PathBuf, components: &[&str]) -> PathBuf {
    let (&file, components) = components.split_last().expect("at least one component required");
    let mut file = file.to_owned();

    if cfg!(windows) {
        file.push_str(".exe");
    }

    for c in components {
        buf.push(c);
    }

    buf.push(file);

    buf
}

pub struct TimeIt(Instant);

/// Returns an RAII structure that prints out how long it took to drop.
pub fn timeit() -> TimeIt {
    TimeIt(Instant::now())
}

impl Drop for TimeIt {
    fn drop(&mut self) {
        let time = self.0.elapsed();
        println!("\tfinished in {}.{:03}",
                 time.as_secs(),
                 time.subsec_nanos() / 1_000_000);
    }
}
<|endoftext|>


<commit_before>extern crate assert_cli;

#[cfg(test)]
mod integration {
    use assert_cli;

    #[test]
    fn mhost_wo_args() {
        mhost()
            .fails()
            .and()
            .stderr().contains("error: The following required arguments were not provided:")
            .unwrap()
    }

    #[test]
    fn mhost_simple_lookup() {
        mhost()
            .with_args(&["www.example.com"])
            .succeeds()
            .and()
            .stdout().contains("* IPv4: 93.184.216.34")
            .and()
            .stdout().contains("* IPv6: 2606:2800:220:1:248:1893:25c8:1946")
            .unwrap()
    }

    #[test]
    fn mhost_simple_ipv4_reverse_lookup() {
        mhost()
            .with_args(&["8.8.8.8"])
            .succeeds()
            .and()
            .stdout().contains("* PTR: google-public-dns-a.google.com.")
            .unwrap()
    }

    #[test]
    fn mhost_simple_ipv6_reverse_lookup() {
        mhost()
            .with_args(&["2001:4860:4860::8888"])
            .succeeds()
            .and()
            .stdout().contains("* PTR: google-public-dns-a.google.com.")
            .unwrap()
    }

    fn mhost() -> assert_cli::Assert {
        assert_cli::Assert::command(&["cargo", "run", "--release", "--features=bin", "--"])
    }
}
<commit_msg>Add RUST_BACKTRACE to integration tests<commit_after>extern crate assert_cli;

#[cfg(test)]
mod integration {
    use assert_cli;

    #[test]
    fn mhost_wo_args() {
        mhost()
            .fails()
            .and()
            .stderr().contains("error: The following required arguments were not provided:")
            .unwrap()
    }

    #[test]
    fn mhost_simple_lookup() {
        mhost()
            .with_args(&["www.example.com"])
            .succeeds()
            .and()
            .stdout().contains("* IPv4: 93.184.216.34")
            .and()
            .stdout().contains("* IPv6: 2606:2800:220:1:248:1893:25c8:1946")
            .unwrap()
    }

    #[test]
    fn mhost_simple_ipv4_reverse_lookup() {
        mhost()
            .with_args(&["8.8.8.8"])
            .succeeds()
            .and()
            .stdout().contains("* PTR: google-public-dns-a.google.com.")
            .unwrap()
    }

    #[test]
    fn mhost_simple_ipv6_reverse_lookup() {
        mhost()
            .with_args(&["2001:4860:4860::8888"])
            .succeeds()
            .and()
            .stdout().contains("* PTR: google-public-dns-a.google.com.")
            .unwrap()
    }

    fn mhost() -> assert_cli::Assert {
        ::std::env::set_var("RUST_BACKTRACE", "1");
        assert_cli::Assert::command(&["cargo", "run", "--release", "--features=bin", "--"])
    }
}
<|endoftext|>


<commit_before>use crate::convert::Convert;
use std::default::Default;
use std::hash::{Hasher};
use arrayref::*;

use const_random::const_random;

///Const random provides randomzied keys with no runtime cost.
const DEFAULT_KEYS: [u64;2] = [const_random!(u64), const_random!(u64)];
///Just a simple bit pattern.
const PAD : u128 = 0xF0E1_D2C3_B4A5_9687_7869_5A4B_3C2D_1E0F;

/// A `Hasher` for hashing an arbitrary stream of bytes.
///
/// Instances of [AHasher] represent state that is updated while hashing data.
///
/// Each method updates the internal state based on the new data provided. Once
/// all of the data has been provided, the resulting hash can be obtained by calling
/// `finish()`
///
/// [Clone] is also provided in case you wish to calculate hashes for two different items that
/// start with the same data.
///
#[derive(Debug, Clone)]
pub struct AHasher {
    buffer: [u64; 2],
}

/// Provides a [Hasher] is typically used (e.g. by [HashMap]) to create
/// [AHasher]s for each key such that they are hashed independently of one
/// another, since [AHasher]s contain state.
///
/// Constructs a new [AHasher] with compile time generated constants keys.
/// So the key will be the same from one instance to another,
/// but different from build to the next. So if it is possible for a potential
/// attacker to have access to your compiled binary it would be better
/// to specify keys generated at runtime.
///
/// # Examples
///
/// ```
/// use ahash::AHasher;
/// use std::hash::Hasher;
///
/// let mut hasher_1 = AHasher::default();
/// let mut hasher_2 = AHasher::default();
///
/// hasher_1.write_u32(8128);
/// hasher_2.write_u32(8128);
///
/// assert_eq!(hasher_1.finish(), hasher_2.finish());
/// ```
/// [Hasher]: std::hash::Hasher
/// [HashMap]: std::collections::HashMap
impl Default for AHasher {
    #[inline]
    fn default() -> AHasher {
        AHasher { buffer: DEFAULT_KEYS }
    }
}

impl AHasher {
    /// Creates a new hasher keyed to the provided keys.
    /// # Example
    ///
    /// ```
    /// use std::hash::Hasher;
    /// use ahash::AHasher;
    ///
    /// let mut hasher = AHasher::new_with_keys(123, 456);
    ///
    /// hasher.write_u32(1989);
    /// hasher.write_u8(11);
    /// hasher.write_u8(9);
    /// hasher.write(b"Huh?");
    ///
    /// println!("Hash is {:x}!", hasher.finish());
    /// ```
    pub fn new_with_keys(key0: u64, key1: u64) -> AHasher {
        AHasher { buffer: [key0, key1] }
    }
}

/// Provides methods to hash all of the primitive types.
impl Hasher for AHasher {
    #[inline]
    fn write_u8(&mut self, i: u8) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u16(&mut self, i: u16) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u32(&mut self, i: u32) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u128(&mut self, i: u128) {
        self.buffer = aeshash(self.buffer.convert(),i).convert();
    }

    #[inline]
    fn write_usize(&mut self, i: usize) {
        self.write_u64(i as u64);
    }

    #[inline]
    fn write_u64(&mut self, i: u64) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write(&mut self, input: &[u8]) {
        let mut data = input;
        let length = data.len() as u64;
        //This will be scrambled by the first AES round in any branch.
        self.buffer[1] ^= length;
        //A 'binary search' on sizes reduces the number of comparisons.
        if data.len() >= 8 {
            if data.len() >= 16 {
                if data.len() > 128 {
                    let mut par_block: u128 = self.buffer.convert();
                    par_block ^= PAD;
                    while data.len() > 128 {
                        let (b1, rest) = data.split_at(16);
                        let b1: u128 = (*as_array!(b1, 16)).convert();
                        par_block = aeshash(par_block, b1);
                        data = rest;
                        let (b2, rest) = data.split_at(16);
                        let b2: u128 = (*as_array!(b2, 16)).convert();
                        self.buffer = aeshash(self.buffer.convert(), b2).convert();
                        data = rest;
                    }
                    self.buffer = aeshash(par_block, self.buffer.convert()).convert();
                }
                while data.len() > 32 {
                    //len 33-128
                    let (block, rest) = data.split_at(16);
                    let block: u128 = (*as_array!(block, 16)).convert();
                    self.buffer = aeshash(self.buffer.convert(),block).convert();
                    data = rest;
                }
                //len 16-32
                let block = (*array_ref!(data, 0, 16)).convert();
                self.buffer = aeshash(self.buffer.convert(),block).convert();
                let block = (*array_ref!(data, data.len()-16, 16)).convert();
                self.buffer = aeshash(self.buffer.convert(),block).convert();
            } else {
                //len 8-15
                let block: [u64; 2] = [(*array_ref!(data, 0, 8)).convert(),
                    (*array_ref!(data, data.len()-8, 8)).convert()];
                self.buffer = aeshash(self.buffer.convert(),block.convert()).convert();
            }
        } else {
            if data.len() >= 2 {
                if data.len() >= 4 {
                    //len 4-7
                    let block: [u32; 2] = [(*array_ref!(data, 0, 4)).convert(),
                        (*array_ref!(data, data.len()-4, 4)).convert()];
                    let block: [u64;2] = [block[1] as u64, block[0] as u64];
                    self.buffer = aeshash(self.buffer.convert(),block.convert()).convert()
                } else {
                    //len 2-3
                    let block: [u16; 2] = [(*array_ref!(data, 0, 2)).convert(),
                        (*array_ref!(data, data.len()-2, 2)).convert()];
                    let block: u32 = block.convert();
                    self.buffer = aeshash(self.buffer.convert(), block as u128).convert();
                }
            } else {
                if data.len() > 0 {
                    //len 1
                    self.buffer = aeshash(self.buffer.convert(), data[0] as u128).convert();
                }
            }
        }
    }
    #[inline]
    fn finish(&self) -> u64 {
        let result: [u64; 2] = aeshash(aeshash(self.buffer.convert(), PAD), PAD).convert();
        result[0]//.wrapping_add(result[1])
    }
}

#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes"))]
#[inline(always)]
fn aeshash(value: u128, xor: u128) -> u128 {
    use std::mem::transmute;
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    unsafe {
        let value = transmute(value);
        transmute(_mm_aesdec_si128(value, transmute(xor)))
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::hash::{BuildHasherDefault};
    use crate::convert::Convert;
    use crate::aes_hash::*;

    #[test]
    fn test_builder() {
        let mut map = HashMap::<u32, u64, BuildHasherDefault<AHasher>>::default();
        map.insert(1, 3);
    }

    #[test]
    fn test_default() {
        let hasher_a = AHasher::default();
        assert_ne!(0, hasher_a.buffer[0]);
        assert_ne!(0, hasher_a.buffer[1]);
        assert_ne!(hasher_a.buffer[0], hasher_a.buffer[1]);
        let hasher_b = AHasher::default();
        assert_eq!(hasher_a.buffer[0], hasher_b.buffer[0]);
        assert_eq!(hasher_a.buffer[1], hasher_b.buffer[1]);
    }

    #[test]
    fn test_hash() {
        let mut result: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value: [u64; 2] = [1 << 32, 0xFEDCBA9876543210];
        result = aeshash(value.convert(), result.convert()).convert();
        result = aeshash(result.convert(), result.convert()).convert();
        let mut result2: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value2: [u64; 2] = [1, 0xFEDCBA9876543210];
        result2 = aeshash(value2.convert(), result2.convert()).convert();
        result2 = aeshash(result2.convert(), result.convert()).convert();
        let result: [u8; 16] = result.convert();
        let result2: [u8; 16] = result2.convert();
        assert_ne!(hex::encode(result), hex::encode(result2));
    }

    #[test]
    fn test_conversion() {
        let input: &[u8] = "dddddddd".as_bytes();
        let bytes: u64 = as_array!(input, 8).convert();
        assert_eq!(bytes, 0x6464646464646464);
    }
}
<commit_msg>Minor tweak to improve perf at 8 and 16 bytes.<commit_after>use crate::convert::Convert;
use std::default::Default;
use std::hash::{Hasher};
use arrayref::*;

use const_random::const_random;

///Const random provides randomzied keys with no runtime cost.
const DEFAULT_KEYS: [u64;2] = [const_random!(u64), const_random!(u64)];
///Just a simple bit pattern.
const PAD : u128 = 0xF0E1_D2C3_B4A5_9687_7869_5A4B_3C2D_1E0F;

/// A `Hasher` for hashing an arbitrary stream of bytes.
///
/// Instances of [AHasher] represent state that is updated while hashing data.
///
/// Each method updates the internal state based on the new data provided. Once
/// all of the data has been provided, the resulting hash can be obtained by calling
/// `finish()`
///
/// [Clone] is also provided in case you wish to calculate hashes for two different items that
/// start with the same data.
///
#[derive(Debug, Clone)]
pub struct AHasher {
    buffer: [u64; 2],
}

/// Provides a [Hasher] is typically used (e.g. by [HashMap]) to create
/// [AHasher]s for each key such that they are hashed independently of one
/// another, since [AHasher]s contain state.
///
/// Constructs a new [AHasher] with compile time generated constants keys.
/// So the key will be the same from one instance to another,
/// but different from build to the next. So if it is possible for a potential
/// attacker to have access to your compiled binary it would be better
/// to specify keys generated at runtime.
///
/// # Examples
///
/// ```
/// use ahash::AHasher;
/// use std::hash::Hasher;
///
/// let mut hasher_1 = AHasher::default();
/// let mut hasher_2 = AHasher::default();
///
/// hasher_1.write_u32(8128);
/// hasher_2.write_u32(8128);
///
/// assert_eq!(hasher_1.finish(), hasher_2.finish());
/// ```
/// [Hasher]: std::hash::Hasher
/// [HashMap]: std::collections::HashMap
impl Default for AHasher {
    #[inline]
    fn default() -> AHasher {
        AHasher { buffer: DEFAULT_KEYS }
    }
}

impl AHasher {
    /// Creates a new hasher keyed to the provided keys.
    /// # Example
    ///
    /// ```
    /// use std::hash::Hasher;
    /// use ahash::AHasher;
    ///
    /// let mut hasher = AHasher::new_with_keys(123, 456);
    ///
    /// hasher.write_u32(1989);
    /// hasher.write_u8(11);
    /// hasher.write_u8(9);
    /// hasher.write(b"Huh?");
    ///
    /// println!("Hash is {:x}!", hasher.finish());
    /// ```
    pub fn new_with_keys(key0: u64, key1: u64) -> AHasher {
        AHasher { buffer: [key0, key1] }
    }
}

/// Provides methods to hash all of the primitive types.
impl Hasher for AHasher {
    #[inline]
    fn write_u8(&mut self, i: u8) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u16(&mut self, i: u16) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u32(&mut self, i: u32) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write_u128(&mut self, i: u128) {
        self.buffer = aeshash(self.buffer.convert(),i).convert();
    }

    #[inline]
    fn write_usize(&mut self, i: usize) {
        self.write_u64(i as u64);
    }

    #[inline]
    fn write_u64(&mut self, i: u64) {
        self.write_u128(i as u128);
    }

    #[inline]
    fn write(&mut self, input: &[u8]) {
        let mut data = input;
        let length = data.len() as u64;
        //This will be scrambled by the first AES round in any branch.
        self.buffer[1] ^= length;
        //A 'binary search' on sizes reduces the number of comparisons.
        if data.len() > 8 {
            if data.len() > 16 {
                if data.len() > 128 {
                    let mut par_block: u128 = self.buffer.convert();
                    par_block ^= PAD;
                    while data.len() > 128 {
                        let (b1, rest) = data.split_at(16);
                        let b1: u128 = (*as_array!(b1, 16)).convert();
                        par_block = aeshash(par_block, b1);
                        data = rest;
                        let (b2, rest) = data.split_at(16);
                        let b2: u128 = (*as_array!(b2, 16)).convert();
                        self.buffer = aeshash(self.buffer.convert(), b2).convert();
                        data = rest;
                    }
                    self.buffer = aeshash(par_block, self.buffer.convert()).convert();
                }
                while data.len() > 32 {
                    //len 33-128
                    let (block, rest) = data.split_at(16);
                    let block: u128 = (*as_array!(block, 16)).convert();
                    self.buffer = aeshash(self.buffer.convert(),block).convert();
                    data = rest;
                }
                //len 17-32
                let block = (*array_ref!(data, 0, 16)).convert();
                self.buffer = aeshash(self.buffer.convert(),block).convert();
                let block = (*array_ref!(data, data.len()-16, 16)).convert();
                self.buffer = aeshash(self.buffer.convert(),block).convert();
            } else {
                //len 9-16
                let block: [u64; 2] = [(*array_ref!(data, 0, 8)).convert(),
                    (*array_ref!(data, data.len()-8, 8)).convert()];
                self.buffer = aeshash(self.buffer.convert(),block.convert()).convert();
            }
        } else {
            if data.len() >= 2 {
                if data.len() >= 4 {
                    //len 4-8
                    let block: [u32; 2] = [(*array_ref!(data, 0, 4)).convert(),
                        (*array_ref!(data, data.len()-4, 4)).convert()];
                    let block: [u64;2] = [block[1] as u64, block[0] as u64];
                    self.buffer = aeshash(self.buffer.convert(),block.convert()).convert()
                } else {
                    //len 2-3
                    let block: [u16; 2] = [(*array_ref!(data, 0, 2)).convert(),
                        (*array_ref!(data, data.len()-2, 2)).convert()];
                    let block: u32 = block.convert();
                    self.buffer = aeshash(self.buffer.convert(), block as u128).convert();
                }
            } else {
                if data.len() > 0 {
                    //len 1
                    self.buffer = aeshash(self.buffer.convert(), data[0] as u128).convert();
                }
            }
        }
    }
    #[inline]
    fn finish(&self) -> u64 {
        let result: [u64; 2] = aeshash(aeshash(self.buffer.convert(), PAD), PAD).convert();
        result[0]//.wrapping_add(result[1])
    }
}

#[cfg(all(any(target_arch = "x86", target_arch = "x86_64"), target_feature = "aes"))]
#[inline(always)]
fn aeshash(value: u128, xor: u128) -> u128 {
    use std::mem::transmute;
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    unsafe {
        let value = transmute(value);
        transmute(_mm_aesdec_si128(value, transmute(xor)))
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::hash::{BuildHasherDefault};
    use crate::convert::Convert;
    use crate::aes_hash::*;

    #[test]
    fn test_builder() {
        let mut map = HashMap::<u32, u64, BuildHasherDefault<AHasher>>::default();
        map.insert(1, 3);
    }

    #[test]
    fn test_default() {
        let hasher_a = AHasher::default();
        assert_ne!(0, hasher_a.buffer[0]);
        assert_ne!(0, hasher_a.buffer[1]);
        assert_ne!(hasher_a.buffer[0], hasher_a.buffer[1]);
        let hasher_b = AHasher::default();
        assert_eq!(hasher_a.buffer[0], hasher_b.buffer[0]);
        assert_eq!(hasher_a.buffer[1], hasher_b.buffer[1]);
    }

    #[test]
    fn test_hash() {
        let mut result: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value: [u64; 2] = [1 << 32, 0xFEDCBA9876543210];
        result = aeshash(value.convert(), result.convert()).convert();
        result = aeshash(result.convert(), result.convert()).convert();
        let mut result2: [u64; 2] = [0x6c62272e07bb0142, 0x62b821756295c58d];
        let value2: [u64; 2] = [1, 0xFEDCBA9876543210];
        result2 = aeshash(value2.convert(), result2.convert()).convert();
        result2 = aeshash(result2.convert(), result.convert()).convert();
        let result: [u8; 16] = result.convert();
        let result2: [u8; 16] = result2.convert();
        assert_ne!(hex::encode(result), hex::encode(result2));
    }

    #[test]
    fn test_conversion() {
        let input: &[u8] = "dddddddd".as_bytes();
        let bytes: u64 = as_array!(input, 8).convert();
        assert_eq!(bytes, 0x6464646464646464);
    }
}
<|endoftext|>


<commit_before>//! Network related functionality for `KafkaClient`.
//!
//! This module is crate private and not exposed to the public except
//! through re-exports of individual items from within
//! `kafka::client`.

use std::collections::HashMap;
use std::fmt;
use std::io::{Read, Write};
use std::mem;
use std::net::{TcpStream, Shutdown};
use std::time::{Instant, Duration};

#[cfg(feature = "security")]
use openssl::ssl::{SslContext, SslStream};

use error::Result;

// --------------------------------------------------------------------

/// Security relevant configuration options for `KafkaClient`.
// This will be expanded in the future. See #51.
#[cfg(feature = "security")]
#[derive(Debug)]
pub struct SecurityConfig(SslContext);

#[cfg(feature = "security")]
impl SecurityConfig {
    /// In the future this will also support a kerbos via #51.
    pub fn new(ssl: SslContext) -> SecurityConfig {
        SecurityConfig(ssl)
    }
}

// --------------------------------------------------------------------

struct Pooled<T> {
    last_checkout: Instant,
    item: T,
}

impl<T> Pooled<T> {
    fn new(last_checkout: Instant, item: T) -> Self {
        Pooled { last_checkout: last_checkout, item: item }
    }
}

impl<T: fmt::Debug> fmt::Debug for Pooled<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Pooled {{ last_checkout: {:?}, item: {:?} }}",
               self.last_checkout, self.item)
    }
}

#[derive(Debug)]
pub struct Config {
    rw_timeout: Option<Duration>,
    idle_timeout: Duration,
    #[cfg(feature = "security")]
    security_config: Option<SecurityConfig>,
}

impl Config {
    #[cfg(not(feature = "security"))]
    fn new_conn(&self, host: &str) -> Result<KafkaConnection> {
        // XXX ideally we'd print some uniq ID of the connection here
        // to differentiate between two connections to the same host
        debug!("Establishing connection to '{}'", host);
        KafkaConnection::new(host, self.timeout)
    }

    #[cfg(feature = "security")]
    fn new_conn(&self, host: &str) -> Result<KafkaConnection> {
        // XXX ideally we'd print some uniq ID of the connection here
        // to differentiate between two connections to the same host
        debug!("Establishing connection to '{}'", host);
        KafkaConnection::new(host,
                             self.rw_timeout,
                             self.security_config.as_ref().map(|c| &c.0))
    }
}

#[derive(Debug)]
pub struct Connections {
    conns: HashMap<String, Pooled<KafkaConnection>>,
    config: Config,
}

impl Connections {
    #[cfg(not(feature = "security"))]
    pub fn new(rw_timeout: Option<Duration>, idle_timeout: Duration)
               -> Connections
    {
        ConnectionPool {
            conns: HashMap::new(),
            rw_timeout: rw_timeout,
            idle_timeout: idle_timeout,
        }
    }

    #[cfg(feature = "security")]
    pub fn new(rw_timeout: Option<Duration>, idle_timeout: Duration)
               -> Connections
    {
        Self::new_with_security(rw_timeout, idle_timeout, None)
    }

    #[cfg(feature = "security")]
    pub fn new_with_security(rw_timeout: Option<Duration>,
                             idle_timeout: Duration,
                             security: Option<SecurityConfig>)
                             -> Connections
    {
        Connections {
            conns: HashMap::new(),
            config: Config {
                rw_timeout: rw_timeout,
                idle_timeout: idle_timeout,
                security_config: security,
            },
        }
    }

    pub fn set_idle_timeout(&mut self, idle_timeout: Duration) {
        self.config.idle_timeout = idle_timeout;
    }

    pub fn idle_timeout(&self) -> Duration {
        self.config.idle_timeout
    }

    pub fn get_conn<'a>(&'a mut self, host: &str, now: Instant)
                        -> Result<&'a mut KafkaConnection>
    {
        if let Some(conn) = self.conns.get_mut(host) {
            if now.duration_since(conn.last_checkout) >= self.config.idle_timeout {
                debug!("Idle timeout reached for connection to '{}'", host);
                let new_conn = try!(self.config.new_conn(host));
                let _ = conn.item.shutdown();
                conn.item = new_conn;
            }
            conn.last_checkout = now;
            let kconn: &mut KafkaConnection = &mut conn.item;
            // ~ decouple the lifetimes to make the borrowck happy;
            // this is safe since we're immediatelly returning the
            // reference and the rest of the code in this method is
            // not affected
            return Ok(unsafe { mem::transmute(kconn) });
        }
        self.conns.insert(host.to_owned(),
                          Pooled::new(now, try!(self.config.new_conn(host))));
        Ok(&mut self.conns.get_mut(host).unwrap().item)
    }

    pub fn get_conn_any(&mut self, now: Instant) -> Option<&mut KafkaConnection> {
        for (host, conn) in &mut self.conns {
            if now.duration_since(conn.last_checkout) >= self.config.idle_timeout {
                debug!("Idle timeout ({:?}) reached for connection to '{}'",
                       self.config.idle_timeout, host);
                let new_conn = match self.config.new_conn(host.as_str()) {
                    Ok(new_conn) => {
                        let _ = conn.item.shutdown();
                        new_conn
                    }
                    Err(e) => {
                        debug!("Failed to establish connection to {}: {:?}",
                               host, e);
                        continue;
                    }
                };
                conn.item = new_conn;
            }
            conn.last_checkout = now;
            let kconn: &mut KafkaConnection = &mut conn.item;
            return Some(kconn);
        }
        None
    }
}

// --------------------------------------------------------------------

#[cfg(not(feature = "security"))]
type KafkaStream = TcpStream;

#[cfg(feature = "security")]
use self::openssled::KafkaStream;

#[cfg(feature = "security")]
mod openssled {
    use std::io::{self, Read, Write};
    use std::net::{TcpStream, Shutdown};
    use std::time::Duration;

    use openssl::ssl::SslStream;

    pub enum KafkaStream {
        Plain(TcpStream),
        Ssl(SslStream<TcpStream>)
    }

    impl KafkaStream {
        fn get_ref(&self) -> &TcpStream {
            match *self {
                KafkaStream::Plain(ref s) => s,
                KafkaStream::Ssl(ref s) => s.get_ref()
            }
        }

        pub fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
            self.get_ref().set_read_timeout(dur)
        }

        pub fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
            self.get_ref().set_write_timeout(dur)
        }

        pub fn shutdown(&mut self, how: Shutdown) -> io::Result<()> {
            self.get_ref().shutdown(how)
        }
    }

    impl Read for KafkaStream {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.read(buf),
                KafkaStream::Ssl(ref mut s) => s.read(buf),
            }
        }
    }

    impl Write for KafkaStream {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.write(buf),
                KafkaStream::Ssl(ref mut s) => s.write(buf),
            }
        }
        fn flush(&mut self) -> io::Result<()> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.flush(),
                KafkaStream::Ssl(ref mut s) => s.flush(),
            }
        }
    }
}

/// A TCP stream to a remote Kafka broker.
pub struct KafkaConnection {
    host: String,
    stream: KafkaStream,
}

impl fmt::Debug for KafkaConnection {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "KafkaConnection {{ host: \"{}\" }}", self.host)
    }
}

impl KafkaConnection {

    pub fn send(&mut self, msg: &[u8]) -> Result<usize> {
        let r = self.stream.write(&msg[..]).map_err(From::from);
        debug!("Sent {} bytes to '{}' => {:?}", msg.len(), self.host, r);
        r
    }

    pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
        let r = (&mut self.stream).read_exact(buf).map_err(From::from);
        debug!("Read {} bytes from '{}' => {:?}", buf.len(), self.host, r);
        r
    }

    pub fn read_exact_alloc(&mut self, size: u64) -> Result<Vec<u8>> {
        let size: usize = size as usize;
        let mut buffer: Vec<u8> = Vec::with_capacity(size);
        // this is safe actually: we are setting the len to the
        // buffers capacity and either fully populate it in the
        // following call to `read_exact` or discard the vector (in
        // the error case)
        unsafe { buffer.set_len(size) };
        try!(self.read_exact(buffer.as_mut_slice()));
        Ok(buffer)
    }

    fn shutdown(&mut self) -> Result<()> {
        let r = self.stream.shutdown(Shutdown::Both);
        debug!("Shut down connection to '{}' => {:?}", self.host, r);
        r.map_err(From::from)
    }

    fn from_stream(stream: KafkaStream, host: &str, rw_timeout: Option<Duration>)
                   -> Result<KafkaConnection>
    {
        try!(stream.set_read_timeout(rw_timeout));
        try!(stream.set_write_timeout(rw_timeout));
        Ok(KafkaConnection { host: host.to_owned(), stream: stream })
    }

    #[cfg(not(feature = "security"))]
    fn new(host: &str, rw_timeout: Option<Duration>) -> Result<KafkaConnection> {
        Ok(KafkaConnection::from_stream(try!(TcpStream::connect(host)), host, rw_timeout))
    }

    #[cfg(feature = "security")]
    fn new(host: &str, rw_timeout: Option<Duration>, security: Option<&SslContext>)
           -> Result<KafkaConnection>
    {
        let stream = try!(TcpStream::connect(host));
        let stream = match security {
            Some(ctx) => KafkaStream::Ssl(try!(SslStream::connect(ctx, stream))),
            None => KafkaStream::Plain(stream),
        };
        KafkaConnection::from_stream(stream, host, rw_timeout)
    }
}
<commit_msg>Traceable debug log of kafka::client::network<commit_after>//! Network related functionality for `KafkaClient`.
//!
//! This module is crate private and not exposed to the public except
//! through re-exports of individual items from within
//! `kafka::client`.

use std::collections::HashMap;
use std::fmt;
use std::io::{Read, Write};
use std::mem;
use std::net::{TcpStream, Shutdown};
use std::time::{Instant, Duration};

#[cfg(feature = "security")]
use openssl::ssl::{SslContext, SslStream};

use error::Result;

// --------------------------------------------------------------------

/// Security relevant configuration options for `KafkaClient`.
// This will be expanded in the future. See #51.
#[cfg(feature = "security")]
#[derive(Debug)]
pub struct SecurityConfig(SslContext);

#[cfg(feature = "security")]
impl SecurityConfig {
    /// In the future this will also support a kerbos via #51.
    pub fn new(ssl: SslContext) -> SecurityConfig {
        SecurityConfig(ssl)
    }
}

// --------------------------------------------------------------------

struct Pooled<T> {
    last_checkout: Instant,
    item: T,
}

impl<T> Pooled<T> {
    fn new(last_checkout: Instant, item: T) -> Self {
        Pooled { last_checkout: last_checkout, item: item }
    }
}

impl<T: fmt::Debug> fmt::Debug for Pooled<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Pooled {{ last_checkout: {:?}, item: {:?} }}",
               self.last_checkout, self.item)
    }
}

#[derive(Debug)]
pub struct Config {
    rw_timeout: Option<Duration>,
    idle_timeout: Duration,
    #[cfg(feature = "security")]
    security_config: Option<SecurityConfig>,
}

impl Config {
    #[cfg(not(feature = "security"))]
    fn new_conn(&self, id: u32, host: &str) -> Result<KafkaConnection> {
        KafkaConnection::new(id, host, self.rw_timeout)
            .map(|c| {
                debug!("Established: {:?}", c);
                c
            })
    }

    #[cfg(feature = "security")]
    fn new_conn(&self, id: u32, host: &str) -> Result<KafkaConnection> {
        KafkaConnection::new(id,
                             host,
                             self.rw_timeout,
                             self.security_config.as_ref().map(|c| &c.0))
            .map(|c| {
                debug!("Established: {:?}", c);
                c
            })
    }
}

#[derive(Debug)]
pub struct State {
    num_conns: u32,
}

impl State {
    fn new() -> State {
        State { num_conns: 0 }
    }

    fn next_conn_id(&mut self) -> u32 {
        let c = self.num_conns;
        self.num_conns = self.num_conns.wrapping_add(1);
        c
    }
}

#[derive(Debug)]
pub struct Connections {
    conns: HashMap<String, Pooled<KafkaConnection>>,
    state: State,
    config: Config,
}

impl Connections {
    #[cfg(not(feature = "security"))]
    pub fn new(rw_timeout: Option<Duration>, idle_timeout: Duration) -> Connections {
        Connections {
            conns: HashMap::new(),
            state: State::new(),
            config: Config {
                rw_timeout: rw_timeout,
                idle_timeout: idle_timeout,
            },
        }
    }

    #[cfg(feature = "security")]
    pub fn new(rw_timeout: Option<Duration>, idle_timeout: Duration) -> Connections {
        Self::new_with_security(rw_timeout, idle_timeout, None)
    }

    #[cfg(feature = "security")]
    pub fn new_with_security(rw_timeout: Option<Duration>,
                             idle_timeout: Duration,
                             security: Option<SecurityConfig>)
                             -> Connections
    {
        Connections {
            conns: HashMap::new(),
            state: State::new(),
            config: Config {
                rw_timeout: rw_timeout,
                idle_timeout: idle_timeout,
                security_config: security,
            },
        }
    }

    pub fn set_idle_timeout(&mut self, idle_timeout: Duration) {
        self.config.idle_timeout = idle_timeout;
    }

    pub fn idle_timeout(&self) -> Duration {
        self.config.idle_timeout
    }

    pub fn get_conn<'a>(&'a mut self, host: &str, now: Instant)
                        -> Result<&'a mut KafkaConnection>
    {
        if let Some(conn) = self.conns.get_mut(host) {
            if now.duration_since(conn.last_checkout) >= self.config.idle_timeout {
                debug!("Idle timeout reached: {:?}", conn.item);
                let new_conn = try!(self.config.new_conn(self.state.next_conn_id(), host));
                let _ = conn.item.shutdown();
                conn.item = new_conn;
            }
            conn.last_checkout = now;
            let kconn: &mut KafkaConnection = &mut conn.item;
            // ~ decouple the lifetimes to make the borrowck happy;
            // this is safe since we're immediatelly returning the
            // reference and the rest of the code in this method is
            // not affected
            return Ok(unsafe { mem::transmute(kconn) });
        }
        let cid = self.state.next_conn_id();
        self.conns.insert(host.to_owned(),
                          Pooled::new(now, try!(self.config.new_conn(cid, host))));
        Ok(&mut self.conns.get_mut(host).unwrap().item)
    }

    pub fn get_conn_any(&mut self, now: Instant) -> Option<&mut KafkaConnection> {
        for (host, conn) in &mut self.conns {
            if now.duration_since(conn.last_checkout) >= self.config.idle_timeout {
                debug!("Idle timeout reached: {:?}", conn.item);
                let new_conn_id = self.state.next_conn_id();
                let new_conn = match self.config.new_conn(new_conn_id, host.as_str()) {
                    Ok(new_conn) => {
                        let _ = conn.item.shutdown();
                        new_conn
                    }
                    Err(e) => {
                        debug!("Failed to establish connection to {}: {:?}", host, e);
                        continue;
                    }
                };
                conn.item = new_conn;
            }
            conn.last_checkout = now;
            let kconn: &mut KafkaConnection = &mut conn.item;
            return Some(kconn);
        }
        None
    }
}

// --------------------------------------------------------------------

trait IsSecured {
    fn is_secured(&self) -> bool;
}

#[cfg(not(feature = "security"))]
type KafkaStream = TcpStream;

#[cfg(not(feature = "security"))]
impl IsSecured for KafkaStream {
    fn is_secured(&self) -> bool {
        false
    }
}

#[cfg(feature = "security")]
use self::openssled::KafkaStream;

#[cfg(feature = "security")]
mod openssled {
    use std::io::{self, Read, Write};
    use std::net::{TcpStream, Shutdown};
    use std::time::Duration;

    use openssl::ssl::SslStream;

    use super::IsSecured;

    pub enum KafkaStream {
        Plain(TcpStream),
        Ssl(SslStream<TcpStream>)
    }

    impl IsSecured for KafkaStream {
        fn is_secured(&self) -> bool {
            match *self {
                KafkaStream::Ssl(_) => true,
                _ => false,
            }
        }
    }

    impl KafkaStream {
        fn get_ref(&self) -> &TcpStream {
            match *self {
                KafkaStream::Plain(ref s) => s,
                KafkaStream::Ssl(ref s) => s.get_ref()
            }
        }

        pub fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
            self.get_ref().set_read_timeout(dur)
        }

        pub fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
            self.get_ref().set_write_timeout(dur)
        }

        pub fn shutdown(&mut self, how: Shutdown) -> io::Result<()> {
            self.get_ref().shutdown(how)
        }
    }

    impl Read for KafkaStream {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.read(buf),
                KafkaStream::Ssl(ref mut s) => s.read(buf),
            }
        }
    }

    impl Write for KafkaStream {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.write(buf),
                KafkaStream::Ssl(ref mut s) => s.write(buf),
            }
        }
        fn flush(&mut self) -> io::Result<()> {
            match *self {
                KafkaStream::Plain(ref mut s) => s.flush(),
                KafkaStream::Ssl(ref mut s) => s.flush(),
            }
        }
    }
}

/// A TCP stream to a remote Kafka broker.
pub struct KafkaConnection {
    // a surrogate identifier to distinguish between
    // connections to the same host in debug messages
    id: u32,
    // "host:port"
    host: String,
    // the (wrapped) tcp stream
    stream: KafkaStream,
}

impl fmt::Debug for KafkaConnection {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "KafkaConnection {{ id: {}, secured: {}, host: \"{}\" }}",
               self.id, self.stream.is_secured(), self.host)
    }
}

impl KafkaConnection {

    pub fn send(&mut self, msg: &[u8]) -> Result<usize> {
        let r = self.stream.write(&msg[..]).map_err(From::from);
        debug!("Sent {} bytes to: {:?} => {:?}", msg.len(), self, r);
        r
    }

    pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<()> {
        let r = (&mut self.stream).read_exact(buf).map_err(From::from);
        debug!("Read {} bytes from: {:?} => {:?}", buf.len(), self, r);
        r
    }

    pub fn read_exact_alloc(&mut self, size: u64) -> Result<Vec<u8>> {
        let size: usize = size as usize;
        let mut buffer: Vec<u8> = Vec::with_capacity(size);
        // this is safe actually: we are setting the len to the
        // buffers capacity and either fully populate it in the
        // following call to `read_exact` or discard the vector (in
        // the error case)
        unsafe { buffer.set_len(size) };
        try!(self.read_exact(buffer.as_mut_slice()));
        Ok(buffer)
    }

    fn shutdown(&mut self) -> Result<()> {
        let r = self.stream.shutdown(Shutdown::Both);
        debug!("Shut down: {:?} => {:?}", self, r);
        r.map_err(From::from)
    }

    fn from_stream(stream: KafkaStream, id: u32, host: &str, rw_timeout: Option<Duration>)
                   -> Result<KafkaConnection>
    {
        try!(stream.set_read_timeout(rw_timeout));
        try!(stream.set_write_timeout(rw_timeout));
        Ok(KafkaConnection { id: id, host: host.to_owned(), stream: stream })
    }

    #[cfg(not(feature = "security"))]
    fn new(id: u32, host: &str, rw_timeout: Option<Duration>) -> Result<KafkaConnection> {
        KafkaConnection::from_stream(try!(TcpStream::connect(host)), id, host, rw_timeout)
    }

    #[cfg(feature = "security")]
    fn new(id: u32, host: &str, rw_timeout: Option<Duration>, security: Option<&SslContext>)
           -> Result<KafkaConnection>
    {
        let stream = try!(TcpStream::connect(host));
        let stream = match security {
            Some(ctx) => KafkaStream::Ssl(try!(SslStream::connect(ctx, stream))),
            None => KafkaStream::Plain(stream),
        };
        KafkaConnection::from_stream(stream, id, host, rw_timeout)
    }
}
<|endoftext|>