<commit_before>// Copyright (c) 2016 Anatoly Ikorsky
//
// 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. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

//! ## mysql-async
//! Tokio based asynchronous MySql client library for rust progrumming language.
//!
//! ### Installation
//! Library hosted on [crates.io](https://crates.io/crates/mysql_async/).
//!
//! ```toml
//! [dependencies]
//! mysql = "<desired version>"
//! ```
//!
//! ### Example
//!
//! ```rust
//! extern crate futures;
//! #[macro_use]
//! extern crate mysql_async as my;
//! extern crate tokio_core as tokio;
//! // ...
//!
//! use futures::Future;
//! use my::prelude::*;
//! use tokio::reactor::Core;
//! # use std::env;
//!
//! #[derive(Debug, PartialEq, Eq)]
//! struct Payment {
//!     customer_id: i32,
//!     amount: i32,
//!     account_name: Option<String>,
//! }
//!
//! fn main() {
//!     let mut lp = Core::new().unwrap();
//!
//!     let payments = vec![
//!         Payment { customer_id: 1, amount: 2, account_name: None },
//!         Payment { customer_id: 3, amount: 4, account_name: Some("foo".into()) },
//!         Payment { customer_id: 5, amount: 6, account_name: None },
//!         Payment { customer_id: 7, amount: 8, account_name: None },
//!         Payment { customer_id: 9, amount: 10, account_name: Some("bar".into()) },
//!     ];
//!
//!     # let database_url: String = if let Ok(url) = env::var("DATABASE_URL") {
//!     #     let opts = my::Opts::from_url(&url).expect("DATABASE_URL invalid");
//!     #     if opts.get_db_name().expect("a database name is required").is_empty() {
//!     #         panic!("database name is empty");
//!     #     }
//!     #     url
//!     # } else {
//!     #     "mysql://root:password@127.0.0.1:3307/mysql".into()
//!     # };
//!
//!     let pool = my::Pool::new(database_url, &lp.handle());
//!     let future = pool.get_conn().and_then(|conn| {
//!         // Create temporary table
//!         conn.query(
//!             r"CREATE TEMPORARY TABLE payment (
//!                 customer_id int not null,
//!                 amount int not null,
//!                 account_name text
//!             )"
//!         ).and_then(|result| result.drop_result())
//!     }).and_then(|conn| {
//!         // Save payments
//!         let params = payments.iter().map(|payment| {
//!             params! {
//!                 "customer_id" => payment.customer_id,
//!                 "amount" => payment.amount,
//!                 "account_name" => payment.account_name.clone(),
//!             }
//!         }).collect();
//!
//!         conn.batch_exec(r"INSERT INTO payment (customer_id, amount, account_name)
//!                         VALUES (:customer_id, :amount, :account_name)", params)
//!     }).and_then(|conn| {
//!         // Load payments from database.
//!         conn.prep_exec("SELECT customer_id, amount, account_name FROM payment", ())
//!     }).and_then(|result| {
//!         // Collect payments
//!         result.map(|row| {
//!             let (customer_id, amount, account_name) = my::from_row(row);
//!             Payment {
//!                 customer_id: customer_id,
//!                 amount: amount,
//!                 account_name: account_name,
//!             }
//!         })
//!     }).and_then(|(payments, _ /* conn */)| {
//!         // Destructor of connection will return it to the pool
//!         // But pool should be disconnected explicitly because it is
//!         // asynchronous procedure.
//!         pool.disconnect().map(|_| payments)
//!     });
//!
//!     let loaded_payments = lp.run(future).unwrap();
//!
//!     assert_eq!(loaded_payments, payments);
//! }
//! ```

#![recursion_limit = "1024"]
#![cfg_attr(feature = "nightly", feature(test, const_fn, drop_types_in_const))]

#[cfg(feature = "nightly")]
extern crate test;

#[macro_use]
extern crate bitflags;
extern crate byteorder;
pub extern crate chrono;
pub extern crate either;
#[macro_use]
extern crate error_chain;
extern crate fnv;
#[macro_use]
extern crate futures as lib_futures;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate regex;
extern crate serde;
extern crate serde_json;
extern crate sha1;
pub extern crate time;
extern crate tokio_core as tokio;
extern crate tokio_io;
extern crate twox_hash;
extern crate url;

#[cfg(test)]
extern crate env_logger;

#[macro_use]
pub mod macros;
#[macro_use]
mod value;
mod conn;
/// Mysql constants
pub mod consts;
/// Errors used in this crate
pub mod errors;
mod io;
mod opts;
mod proto;
mod scramble;

#[doc(inline)]
pub use self::conn::Conn;

#[doc(inline)]
pub use self::conn::futures::query_result::{BinaryResult, BinQueryResult, ResultSet,
                                            TextQueryResult, TextResult};

#[doc(inline)]
pub use self::conn::pool::Pool;

#[doc(inline)]
pub use self::conn::stmt::Stmt;

#[doc(inline)]
pub use self::conn::transaction::Transaction;

#[doc(inline)]
pub use self::conn::transaction::futures::query_result::{TransBinQueryResult, TransTextQueryResult};

#[doc(inline)]
pub use self::opts::{Opts, OptsBuilder, WhiteListFsLocalInfileHandler};

#[doc(inline)]
pub use self::proto::{Column, ErrPacket, Row};

#[doc(inline)]
pub use self::value::{from_row, from_row_opt, from_value, from_value_opt, Params, Value,
                      Serialized, Deserialized};

/// Futures used in this crate
pub mod futures {
    #[doc(inline)]
    pub use conn::futures::BatchExec;
    #[doc(inline)]
    pub use conn::futures::Disconnect;
    #[doc(inline)]
    pub use conn::futures::DropExec;
    #[doc(inline)]
    pub use conn::futures::DropQuery;
    #[doc(inline)]
    pub use conn::futures::First;
    #[doc(inline)]
    pub use conn::futures::FirstExec;
    #[doc(inline)]
    pub use conn::futures::NewConn;
    #[doc(inline)]
    pub use conn::futures::Ping;
    #[doc(inline)]
    pub use conn::futures::Prepare;
    #[doc(inline)]
    pub use conn::futures::PrepExec;
    #[doc(inline)]
    pub use conn::futures::Query;
    #[doc(inline)]
    pub use conn::futures::Reset;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::Collect;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::CollectAll;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::DropResult;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::ForEach;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::Map;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::Reduce;
    #[doc(inline)]
    pub use conn::pool::futures::DisconnectPool;
    #[doc(inline)]
    pub use conn::pool::futures::GetConn;
    #[doc(inline)]
    pub use conn::stmt::futures::Batch;
    #[doc(inline)]
    pub use conn::stmt::futures::Execute;
    #[doc(inline)]
    pub use conn::stmt::futures::First as StmtFirst;
    #[doc(inline)]
    pub use conn::transaction::futures::Commit;
    #[doc(inline)]
    pub use conn::transaction::futures::Rollback;
    #[doc(inline)]
    pub use conn::transaction::futures::StartTransaction;
    #[doc(inline)]
    pub use conn::transaction::futures::TransBatchExec;
    #[doc(inline)]
    pub use conn::transaction::futures::TransFirst;
    #[doc(inline)]
    pub use conn::transaction::futures::TransFirstExec;
    #[doc(inline)]
    pub use conn::transaction::futures::TransPrepExec;
    #[doc(inline)]
    pub use conn::transaction::futures::TransQuery;
}

/// Traits used in this crate
pub mod prelude {
    #[doc(inline)]
    pub use conn::futures::query_result::QueryResult;
    #[doc(inline)]
    pub use conn::futures::query_result::ResultKind;
    #[doc(inline)]
    pub use conn::futures::query_result::UnconsumedQueryResult;
    #[doc(inline)]
    pub use opts::LocalInfileHandler;
    #[doc(inline)]
    pub use value::ConvIr;
    #[doc(inline)]
    pub use value::FromRow;
    #[doc(inline)]
    pub use value::FromValue;
    #[doc(inline)]
    pub use value::ToValue;
}

#[cfg(test)]
mod test_misc {
    use std::env;
    use opts;
    lazy_static! {
        pub static ref DATABASE_URL: String = {
            if let Ok(url) = env::var("DATABASE_URL") {
                let opts = opts::Opts::from_url(&url).expect("DATABASE_URL invalid");
                if opts.get_db_name().expect("a database name is required").is_empty() {
                    panic!("database name is empty");
                }
                url
            } else {
                "mysql://root:password@127.0.0.1:3307/mysql".into()
            }
        };
    }
}
<commit_msg>Reexport TransDropQuery and TransDropExec<commit_after>// Copyright (c) 2016 Anatoly Ikorsky
//
// 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. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

//! ## mysql-async
//! Tokio based asynchronous MySql client library for rust progrumming language.
//!
//! ### Installation
//! Library hosted on [crates.io](https://crates.io/crates/mysql_async/).
//!
//! ```toml
//! [dependencies]
//! mysql = "<desired version>"
//! ```
//!
//! ### Example
//!
//! ```rust
//! extern crate futures;
//! #[macro_use]
//! extern crate mysql_async as my;
//! extern crate tokio_core as tokio;
//! // ...
//!
//! use futures::Future;
//! use my::prelude::*;
//! use tokio::reactor::Core;
//! # use std::env;
//!
//! #[derive(Debug, PartialEq, Eq)]
//! struct Payment {
//!     customer_id: i32,
//!     amount: i32,
//!     account_name: Option<String>,
//! }
//!
//! fn main() {
//!     let mut lp = Core::new().unwrap();
//!
//!     let payments = vec![
//!         Payment { customer_id: 1, amount: 2, account_name: None },
//!         Payment { customer_id: 3, amount: 4, account_name: Some("foo".into()) },
//!         Payment { customer_id: 5, amount: 6, account_name: None },
//!         Payment { customer_id: 7, amount: 8, account_name: None },
//!         Payment { customer_id: 9, amount: 10, account_name: Some("bar".into()) },
//!     ];
//!
//!     # let database_url: String = if let Ok(url) = env::var("DATABASE_URL") {
//!     #     let opts = my::Opts::from_url(&url).expect("DATABASE_URL invalid");
//!     #     if opts.get_db_name().expect("a database name is required").is_empty() {
//!     #         panic!("database name is empty");
//!     #     }
//!     #     url
//!     # } else {
//!     #     "mysql://root:password@127.0.0.1:3307/mysql".into()
//!     # };
//!
//!     let pool = my::Pool::new(database_url, &lp.handle());
//!     let future = pool.get_conn().and_then(|conn| {
//!         // Create temporary table
//!         conn.query(
//!             r"CREATE TEMPORARY TABLE payment (
//!                 customer_id int not null,
//!                 amount int not null,
//!                 account_name text
//!             )"
//!         ).and_then(|result| result.drop_result())
//!     }).and_then(|conn| {
//!         // Save payments
//!         let params = payments.iter().map(|payment| {
//!             params! {
//!                 "customer_id" => payment.customer_id,
//!                 "amount" => payment.amount,
//!                 "account_name" => payment.account_name.clone(),
//!             }
//!         }).collect();
//!
//!         conn.batch_exec(r"INSERT INTO payment (customer_id, amount, account_name)
//!                         VALUES (:customer_id, :amount, :account_name)", params)
//!     }).and_then(|conn| {
//!         // Load payments from database.
//!         conn.prep_exec("SELECT customer_id, amount, account_name FROM payment", ())
//!     }).and_then(|result| {
//!         // Collect payments
//!         result.map(|row| {
//!             let (customer_id, amount, account_name) = my::from_row(row);
//!             Payment {
//!                 customer_id: customer_id,
//!                 amount: amount,
//!                 account_name: account_name,
//!             }
//!         })
//!     }).and_then(|(payments, _ /* conn */)| {
//!         // Destructor of connection will return it to the pool
//!         // But pool should be disconnected explicitly because it is
//!         // asynchronous procedure.
//!         pool.disconnect().map(|_| payments)
//!     });
//!
//!     let loaded_payments = lp.run(future).unwrap();
//!
//!     assert_eq!(loaded_payments, payments);
//! }
//! ```

#![recursion_limit = "1024"]
#![cfg_attr(feature = "nightly", feature(test, const_fn, drop_types_in_const))]

#[cfg(feature = "nightly")]
extern crate test;

#[macro_use]
extern crate bitflags;
extern crate byteorder;
pub extern crate chrono;
pub extern crate either;
#[macro_use]
extern crate error_chain;
extern crate fnv;
#[macro_use]
extern crate futures as lib_futures;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
extern crate regex;
extern crate serde;
extern crate serde_json;
extern crate sha1;
pub extern crate time;
extern crate tokio_core as tokio;
extern crate tokio_io;
extern crate twox_hash;
extern crate url;

#[cfg(test)]
extern crate env_logger;

#[macro_use]
pub mod macros;
#[macro_use]
mod value;
mod conn;
/// Mysql constants
pub mod consts;
/// Errors used in this crate
pub mod errors;
mod io;
mod opts;
mod proto;
mod scramble;

#[doc(inline)]
pub use self::conn::Conn;

#[doc(inline)]
pub use self::conn::futures::query_result::{BinaryResult, BinQueryResult, ResultSet,
                                            TextQueryResult, TextResult};

#[doc(inline)]
pub use self::conn::pool::Pool;

#[doc(inline)]
pub use self::conn::stmt::Stmt;

#[doc(inline)]
pub use self::conn::transaction::Transaction;

#[doc(inline)]
pub use self::conn::transaction::futures::query_result::{TransBinQueryResult, TransTextQueryResult};

#[doc(inline)]
pub use self::opts::{Opts, OptsBuilder, WhiteListFsLocalInfileHandler};

#[doc(inline)]
pub use self::proto::{Column, ErrPacket, Row};

#[doc(inline)]
pub use self::value::{from_row, from_row_opt, from_value, from_value_opt, Params, Value,
                      Serialized, Deserialized};

/// Futures used in this crate
pub mod futures {
    #[doc(inline)]
    pub use conn::futures::BatchExec;
    #[doc(inline)]
    pub use conn::futures::Disconnect;
    #[doc(inline)]
    pub use conn::futures::DropExec;
    #[doc(inline)]
    pub use conn::futures::DropQuery;
    #[doc(inline)]
    pub use conn::futures::First;
    #[doc(inline)]
    pub use conn::futures::FirstExec;
    #[doc(inline)]
    pub use conn::futures::NewConn;
    #[doc(inline)]
    pub use conn::futures::Ping;
    #[doc(inline)]
    pub use conn::futures::Prepare;
    #[doc(inline)]
    pub use conn::futures::PrepExec;
    #[doc(inline)]
    pub use conn::futures::Query;
    #[doc(inline)]
    pub use conn::futures::Reset;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::Collect;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::CollectAll;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::DropResult;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::ForEach;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::Map;
    #[doc(inline)]
    pub use conn::futures::query_result::futures::Reduce;
    #[doc(inline)]
    pub use conn::pool::futures::DisconnectPool;
    #[doc(inline)]
    pub use conn::pool::futures::GetConn;
    #[doc(inline)]
    pub use conn::stmt::futures::Batch;
    #[doc(inline)]
    pub use conn::stmt::futures::Execute;
    #[doc(inline)]
    pub use conn::stmt::futures::First as StmtFirst;
    #[doc(inline)]
    pub use conn::transaction::futures::Commit;
    #[doc(inline)]
    pub use conn::transaction::futures::Rollback;
    #[doc(inline)]
    pub use conn::transaction::futures::StartTransaction;
    #[doc(inline)]
    pub use conn::transaction::futures::TransBatchExec;
    #[doc(inline)]
    pub use conn::transaction::futures::TransFirst;
    #[doc(inline)]
    pub use conn::transaction::futures::TransDropQuery;
    #[doc(inline)]
    pub use conn::transaction::futures::TransFirstExec;
    #[doc(inline)]
    pub use conn::transaction::futures::TransDropExec;
    #[doc(inline)]
    pub use conn::transaction::futures::TransPrepExec;
    #[doc(inline)]
    pub use conn::transaction::futures::TransQuery;
}

/// Traits used in this crate
pub mod prelude {
    #[doc(inline)]
    pub use conn::futures::query_result::QueryResult;
    #[doc(inline)]
    pub use conn::futures::query_result::ResultKind;
    #[doc(inline)]
    pub use conn::futures::query_result::UnconsumedQueryResult;
    #[doc(inline)]
    pub use opts::LocalInfileHandler;
    #[doc(inline)]
    pub use value::ConvIr;
    #[doc(inline)]
    pub use value::FromRow;
    #[doc(inline)]
    pub use value::FromValue;
    #[doc(inline)]
    pub use value::ToValue;
}

#[cfg(test)]
mod test_misc {
    use std::env;
    use opts;
    lazy_static! {
        pub static ref DATABASE_URL: String = {
            if let Ok(url) = env::var("DATABASE_URL") {
                let opts = opts::Opts::from_url(&url).expect("DATABASE_URL invalid");
                if opts.get_db_name().expect("a database name is required").is_empty() {
                    panic!("database name is empty");
                }
                url
            } else {
                "mysql://root:password@127.0.0.1:3307/mysql".into()
            }
        };
    }
}
<|endoftext|>


<commit_before>#![feature(collections)]
#![feature(convert)]
mod token;

pub fn compile(sass: String, style: &str) -> String {
    let mut sp = SassParser::new(sass);
    let parsed = sp.parse();
    match style {
        "nested"     => parsed.nested(&sp),
        "compressed" => compressed(parsed, sp),
        "expanded"   => expanded(parsed, sp),
        "compact"    => compact(parsed, sp),
        _            => panic!("Unknown style: {}. Please specify one of nested, compressed, expanded, or compact.", style),
    }
}

fn compressed(sass: SassRuleSet, sp: SassParser) -> String {
    // sass.replace(" ", "").replace("\n", "")
    format!("{:?}", sass)
}

fn expanded(sass: SassRuleSet, sp: SassParser) -> String {
    format!("{:?}", sass)
}

fn compact(sass: SassRuleSet, sp: SassParser) -> String {
    format!("{:?}", sass)
}

#[derive(Debug)]
struct SassRuleSet {
    pub rules: Vec<SassRule>,
}

impl SassRuleSet {
    fn nested(&self, sp: &SassParser) -> String {
        let mut output =  String::from_str("");
        for rule in &self.rules {
            output.push_str((&rule).nested(&sp).as_str());
        }
        output
    }
}

#[derive(Debug)]
struct SassRule {
    pub selectors: Vec<token::Range>,
    pub props_and_values: Vec<PropertyValueSet>,
}

impl SassRule {
    fn nested(&self, sp: &SassParser) -> String {
        let mut output =  String::from_str("");

        for selector in &self.selectors {
            output.push_str(&sp.extract(selector));
        }

        output.push_str(" {");

        for prop_and_val in &self.props_and_values {
            output.push_str(format!("\n  {}: {};", &sp.extract(&prop_and_val.property), &sp.extract(&prop_and_val.value)).as_str());
        }

        output.push_str(" }\n");

        output
    }
}

#[derive(Debug)]
struct PropertyValueSet {
    pub property: token::Range,
    pub value: token::Range,
}

#[derive(Debug)]
struct SassParser {
    pub token: token::Range,
    pub last_token: token::Range,
    pub tokenizer: SassTokenizer,
    sass: String,
}

impl SassParser {
    pub fn new(str: String) -> SassParser {
        let mut tokenizer = SassTokenizer::new(str.clone());
        let initial_token = tokenizer.real_token();
        SassParser {
            token: initial_token.clone(),
            last_token: initial_token.clone(),
            tokenizer: tokenizer,
            sass: str,
        }
    }

    pub fn extract(&self, range: &token::Range) -> &str {
        self.sass.slice_chars(range.start_pos as usize, range.end_pos as usize)
    }

    pub fn bump(&mut self) {
        self.last_token = self.token.clone();
        self.token = self.tokenizer.real_token();
    }

    pub fn parse(&mut self) -> SassRuleSet {
        match self.parse_rules() {
            Ok(sass_rule_set) => sass_rule_set,
            Err(msg) => panic!(msg),
        }
    }

    fn parse_rules(&mut self) -> Result<SassRuleSet, &'static str> {
        let mut rules = vec![];
        while let Some(rule) = try!(self.parse_rule()) {
            rules.push(rule);
        }
        Ok(SassRuleSet { rules: rules })
    }

    fn parse_rule(&mut self) -> Result<Option<SassRule>, &'static str> {
        if self.token.token == token::Eof {
            return Ok(None)
        }

        let mut selectors = vec![];
        while let Some(selector) = try!(self.parse_selector()) {
            selectors.push(selector);
        }
        if selectors.len() == 0 {
            panic!("Empty selector!")
        }

        let mut props_and_values = vec![];
        while let Some(property_value_set) = try!(self.parse_property_value_set()) {
            props_and_values.push(property_value_set);
        }

        // A rule without properties and values isn't *wrong*, per se...
        Ok(Some(SassRule { selectors: selectors, props_and_values: props_and_values }))
    }

    fn parse_selector(&mut self) -> Result<Option<token::Range>, &'static str> {
        if self.token.token == token::Eof {
            return Ok(None)
        }

        if self.token.token == token::Text {
            self.bump();
            return Ok(Some(self.last_token.clone()))
        }

        match &self.token.token {
            &token::OpenDelim(token::Brace) => {
                self.bump();
                Ok(None)
            },
            &token::Comma => {
                self.bump();
                self.parse_selector()
            },
            _ => Err("Unexpected token where we expected selectors"),
        }
    }

    fn parse_property_value_set(&mut self) -> Result<Option<PropertyValueSet>, &'static str> {
        if self.token.token == token::Eof {
            return Ok(None)
        }

        if self.token.token == token::CloseDelim(token::Brace) {
            self.bump();
            return Ok(None)
        }

        let p = self.token.clone();

        match p.token {
            token::Text => {
                self.bump(); // should be colon
                self.bump(); // optional whitespace required for now

                let v = self.token.clone();

                if v.token == token::Text {
                    self.bump(); // semicolon
                    self.bump(); // whitespace
                    return Ok(Some(PropertyValueSet { property: p, value: v }))
                } else {
                    panic!("Expected a value here, instead got who knows what!");
                }
            },
            token::CloseDelim(token::Brace) => {
                self.bump();
                Ok(None)
            },
            _ => Err("Unexpected token where we expected properties and values"),
        }
    }
}

#[derive(Debug)]
struct SassTokenizer {
    pub pos: u32,
    pub last_pos: u32,
    pub curr: Option<char>,
    pub peek_range: token::Range,
    sass: String,
}

impl SassTokenizer {
    pub fn new(str: String) -> SassTokenizer {
        let mut sr = SassTokenizer {
            pos: 0,
            last_pos: 0,
            curr: Some('\n'),
            peek_range: token::Range { start_pos: 0, end_pos: 0, token: token::Eof },
            sass: str,
        };
        sr.bump();
        sr.advance_token();
        sr
    }

    pub fn bump(&mut self) {
        self.last_pos = self.pos;
        let current_pos = self.pos as usize;
        if current_pos < self.sass.len() {
            let ch = char_at(&self.sass, current_pos);

            self.pos = self.pos + (ch.len_utf8() as u32);
            self.curr = Some(ch);
        } else {
            self.curr = None;
        }
    }

    fn advance_token(&mut self) {
        match self.scan_whitespace() {
            Some(whitespace) => {
                self.peek_range = whitespace;
            },
            None => {
                if self.is_eof() {
                    self.peek_range = token::Range { start_pos: self.pos + 1, end_pos: self.pos + 1, token: token::Eof };
                } else {
                    self.peek_range = self.next_token_inner();
                }
            },
        }
    }

    fn next_token(&mut self) -> token::Range {
        let retval = self.peek_range.clone();
        self.advance_token();
        retval
    }

    fn next_token_inner(&mut self) -> token::Range {
        let c = match self.curr {
            Some(c) => c,
            None => return token::Range { start_pos: self.pos + 1, end_pos: self.pos + 1, token: token::Eof },
        };

        if c >= 'a' && c <= 'z' {
            let start = self.last_pos;
            while !self.curr.is_none() && self.curr.unwrap() >= 'a' && self.curr.unwrap() <= 'z' {
                self.bump();
            }
            return token::Range { start_pos: start, end_pos: self.last_pos, token: token::Text }
        }

        match c {
            ';' => { self.bump(); return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::Semi }; },
            ':' => { self.bump(); return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::Colon }; },
            ',' => { self.bump(); return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::Comma }; },
            '{' => { self.bump(); return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::OpenDelim(token::Brace) }; },
            '}' => { self.bump(); return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::CloseDelim(token::Brace) }; },
            _   => { return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::Unknown} },
        }
    }

    fn real_token(&mut self) -> token::Range {
        let mut t = self.next_token();
        loop {
            match t.token {
                token::Whitespace => {
                    t = self.next_token();
                },
                _ => break,
            }
        }
        t
    }

    fn scan_whitespace(&mut self) -> Option<token::Range> {
        match self.curr.unwrap_or('\0') {
            c if is_whitespace(Some(c)) => {
                let start = self.last_pos;
                while is_whitespace(self.curr) { self.bump(); }
                Some(token::Range { start_pos: start, end_pos: self.last_pos, token: token::Whitespace })
            },
            _ => None
        }
    }

    fn is_eof(&self) -> bool {
        self.curr.is_none()
    }
}

pub fn char_at(s: &str, byte: usize) -> char {
    s[byte..].chars().next().unwrap()
}

pub fn is_whitespace(c: Option<char>) -> bool {
    match c.unwrap_or('\x00') { // None can be null for now... it's not whitespace
        ' ' | '\n' | '\t' | '\r' => true,
        _ => false
    }
}
<commit_msg>Implement the rest of the output formats<commit_after>#![feature(collections)]
#![feature(convert)]
mod token;

pub fn compile(sass: String, style: &str) -> String {
    let mut sp = SassParser::new(sass);
    let parsed = sp.parse();
    match style {
        "nested"     => parsed.nested(&sp),
        "compressed" => parsed.compressed(&sp),
        "expanded"   => parsed.expanded(&sp),
        "compact"    => parsed.compact(&sp),
        _            => panic!("Unknown style: {}. Please specify one of nested, compressed, expanded, or compact.", style),
    }
}

#[derive(Debug)]
struct SassRuleSet {
    pub rules: Vec<SassRule>,
}

impl SassRuleSet {
    fn nested(&self, sp: &SassParser) -> String {
        let mut output =  String::from_str("");
        for rule in &self.rules {
            output.push_str((&rule).nested(&sp).as_str());
        }
        output
    }

    fn compressed(&self, sp: &SassParser) -> String {
        let mut output =  String::from_str("");
        for rule in &self.rules {
            output.push_str((&rule).compressed(&sp).as_str());
        }
        output
    }

    fn expanded(&self, sp: &SassParser) -> String {
        let mut output =  String::from_str("");
        for rule in &self.rules {
            output.push_str((&rule).expanded(&sp).as_str());
        }
        output
    }

    fn compact(&self, sp: &SassParser) -> String {
        let mut output =  String::from_str("");
        for rule in &self.rules {
            output.push_str((&rule).compact(&sp).as_str());
        }
        output
    }
}

#[derive(Debug)]
struct SassRule {
    pub selectors: Vec<token::Range>,
    pub props_and_values: Vec<PropertyValueSet>,
}

impl SassRule {
    fn nested(&self, sp: &SassParser) -> String {
        let mut output =  String::from_str("");

        for selector in &self.selectors {
            output.push_str(&sp.extract(selector));
        }

        output.push_str(" {");

        for prop_and_val in &self.props_and_values {
            output.push_str(format!("\n  {}: {};", &sp.extract(&prop_and_val.property), &sp.extract(&prop_and_val.value)).as_str());
        }

        output.push_str(" }\n");

        output
    }

    fn compressed(&self, sp: &SassParser) -> String {
        let mut output =  String::from_str("");

        for selector in &self.selectors {
            output.push_str(&sp.extract(selector));
        }

        output.push_str("{");

        for prop_and_val in &self.props_and_values {
            output.push_str(format!("{}:{}", &sp.extract(&prop_and_val.property), &sp.extract(&prop_and_val.value)).as_str());
        }

        output.push_str("}\n");

        output
    }

    fn expanded(&self, sp: &SassParser) -> String {
        let mut output =  String::from_str("");

        for selector in &self.selectors {
            output.push_str(&sp.extract(selector));
        }

        output.push_str(" {");

        for prop_and_val in &self.props_and_values {
            output.push_str(format!("\n  {}: {};", &sp.extract(&prop_and_val.property), &sp.extract(&prop_and_val.value)).as_str());
        }

        output.push_str("\n}\n");

        output
    }

    fn compact(&self, sp: &SassParser) -> String {
        let mut output =  String::from_str("");

        for selector in &self.selectors {
            output.push_str(&sp.extract(selector));
        }

        output.push_str(" {");

        for prop_and_val in &self.props_and_values {
            output.push_str(format!(" {}: {};", &sp.extract(&prop_and_val.property), &sp.extract(&prop_and_val.value)).as_str());
        }

        output.push_str(" }\n");

        output
    }
}

#[derive(Debug)]
struct PropertyValueSet {
    pub property: token::Range,
    pub value: token::Range,
}

#[derive(Debug)]
struct SassParser {
    pub token: token::Range,
    pub last_token: token::Range,
    pub tokenizer: SassTokenizer,
    sass: String,
}

impl SassParser {
    pub fn new(str: String) -> SassParser {
        let mut tokenizer = SassTokenizer::new(str.clone());
        let initial_token = tokenizer.real_token();
        SassParser {
            token: initial_token.clone(),
            last_token: initial_token.clone(),
            tokenizer: tokenizer,
            sass: str,
        }
    }

    pub fn extract(&self, range: &token::Range) -> &str {
        self.sass.slice_chars(range.start_pos as usize, range.end_pos as usize)
    }

    pub fn bump(&mut self) {
        self.last_token = self.token.clone();
        self.token = self.tokenizer.real_token();
    }

    pub fn parse(&mut self) -> SassRuleSet {
        match self.parse_rules() {
            Ok(sass_rule_set) => sass_rule_set,
            Err(msg) => panic!(msg),
        }
    }

    fn parse_rules(&mut self) -> Result<SassRuleSet, &'static str> {
        let mut rules = vec![];
        while let Some(rule) = try!(self.parse_rule()) {
            rules.push(rule);
        }
        Ok(SassRuleSet { rules: rules })
    }

    fn parse_rule(&mut self) -> Result<Option<SassRule>, &'static str> {
        if self.token.token == token::Eof {
            return Ok(None)
        }

        let mut selectors = vec![];
        while let Some(selector) = try!(self.parse_selector()) {
            selectors.push(selector);
        }
        if selectors.len() == 0 {
            panic!("Empty selector!")
        }

        let mut props_and_values = vec![];
        while let Some(property_value_set) = try!(self.parse_property_value_set()) {
            props_and_values.push(property_value_set);
        }

        // A rule without properties and values isn't *wrong*, per se...
        Ok(Some(SassRule { selectors: selectors, props_and_values: props_and_values }))
    }

    fn parse_selector(&mut self) -> Result<Option<token::Range>, &'static str> {
        if self.token.token == token::Eof {
            return Ok(None)
        }

        if self.token.token == token::Text {
            self.bump();
            return Ok(Some(self.last_token.clone()))
        }

        match &self.token.token {
            &token::OpenDelim(token::Brace) => {
                self.bump();
                Ok(None)
            },
            &token::Comma => {
                self.bump();
                self.parse_selector()
            },
            _ => Err("Unexpected token where we expected selectors"),
        }
    }

    fn parse_property_value_set(&mut self) -> Result<Option<PropertyValueSet>, &'static str> {
        if self.token.token == token::Eof {
            return Ok(None)
        }

        if self.token.token == token::CloseDelim(token::Brace) {
            self.bump();
            return Ok(None)
        }

        let p = self.token.clone();

        match p.token {
            token::Text => {
                self.bump(); // should be colon
                self.bump(); // optional whitespace required for now

                let v = self.token.clone();

                if v.token == token::Text {
                    self.bump(); // semicolon
                    self.bump(); // whitespace
                    return Ok(Some(PropertyValueSet { property: p, value: v }))
                } else {
                    panic!("Expected a value here, instead got who knows what!");
                }
            },
            token::CloseDelim(token::Brace) => {
                self.bump();
                Ok(None)
            },
            _ => Err("Unexpected token where we expected properties and values"),
        }
    }
}

#[derive(Debug)]
struct SassTokenizer {
    pub pos: u32,
    pub last_pos: u32,
    pub curr: Option<char>,
    pub peek_range: token::Range,
    sass: String,
}

impl SassTokenizer {
    pub fn new(str: String) -> SassTokenizer {
        let mut sr = SassTokenizer {
            pos: 0,
            last_pos: 0,
            curr: Some('\n'),
            peek_range: token::Range { start_pos: 0, end_pos: 0, token: token::Eof },
            sass: str,
        };
        sr.bump();
        sr.advance_token();
        sr
    }

    pub fn bump(&mut self) {
        self.last_pos = self.pos;
        let current_pos = self.pos as usize;
        if current_pos < self.sass.len() {
            let ch = char_at(&self.sass, current_pos);

            self.pos = self.pos + (ch.len_utf8() as u32);
            self.curr = Some(ch);
        } else {
            self.curr = None;
        }
    }

    fn advance_token(&mut self) {
        match self.scan_whitespace() {
            Some(whitespace) => {
                self.peek_range = whitespace;
            },
            None => {
                if self.is_eof() {
                    self.peek_range = token::Range { start_pos: self.pos + 1, end_pos: self.pos + 1, token: token::Eof };
                } else {
                    self.peek_range = self.next_token_inner();
                }
            },
        }
    }

    fn next_token(&mut self) -> token::Range {
        let retval = self.peek_range.clone();
        self.advance_token();
        retval
    }

    fn next_token_inner(&mut self) -> token::Range {
        let c = match self.curr {
            Some(c) => c,
            None => return token::Range { start_pos: self.pos + 1, end_pos: self.pos + 1, token: token::Eof },
        };

        if c >= 'a' && c <= 'z' {
            let start = self.last_pos;
            while !self.curr.is_none() && self.curr.unwrap() >= 'a' && self.curr.unwrap() <= 'z' {
                self.bump();
            }
            return token::Range { start_pos: start, end_pos: self.last_pos, token: token::Text }
        }

        match c {
            ';' => { self.bump(); return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::Semi }; },
            ':' => { self.bump(); return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::Colon }; },
            ',' => { self.bump(); return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::Comma }; },
            '{' => { self.bump(); return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::OpenDelim(token::Brace) }; },
            '}' => { self.bump(); return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::CloseDelim(token::Brace) }; },
            _   => { return token::Range { start_pos: self.last_pos, end_pos: self.last_pos, token: token::Unknown} },
        }
    }

    fn real_token(&mut self) -> token::Range {
        let mut t = self.next_token();
        loop {
            match t.token {
                token::Whitespace => {
                    t = self.next_token();
                },
                _ => break,
            }
        }
        t
    }

    fn scan_whitespace(&mut self) -> Option<token::Range> {
        match self.curr.unwrap_or('\0') {
            c if is_whitespace(Some(c)) => {
                let start = self.last_pos;
                while is_whitespace(self.curr) { self.bump(); }
                Some(token::Range { start_pos: start, end_pos: self.last_pos, token: token::Whitespace })
            },
            _ => None
        }
    }

    fn is_eof(&self) -> bool {
        self.curr.is_none()
    }
}

pub fn char_at(s: &str, byte: usize) -> char {
    s[byte..].chars().next().unwrap()
}

pub fn is_whitespace(c: Option<char>) -> bool {
    match c.unwrap_or('\x00') { // None can be null for now... it's not whitespace
        ' ' | '\n' | '\t' | '\r' => true,
        _ => false
    }
}
<|endoftext|>


<commit_before>#[cfg(feature = "use_hyper")]
extern crate hyper;
#[cfg(feature = "use_hyper")]
extern crate url as servo_url;

pub mod server;
pub mod client;
pub mod mock;
pub mod url;
#[cfg(feature = "use_hyper")]
pub mod mockable_hyper;
pub mod mockable_tcp_stream;

pub type Url<'a> = url::Url<'a>;
pub type Mock = mock::Mock;

pub fn mock(method: &str, path: &str) -> Mock {
    Mock::new(method, path)
}

#[cfg(test)]
#[cfg(feature = "mock_hyper")]
mod mock_hyper_tests {
    use hyper::Client;
    use hyper::header::Connection;
    use server;
    use url::Url;
    use std::io::Read;

    #[test]
    fn test_proxying() {
        super::start();

        let client = Client::new();
        let mut res = client.get(Url("http://www.example.com"))
            .header(Connection::close())
            .send()
            .unwrap();

        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();

        assert_eq!(body, "Hello world");
    }
}

#[cfg(test)]
#[cfg(feature = "mock_tcp_stream")]
mod mock_tcp_stream_tests {
    use MockBuilder;
    use super::mock;

    #[test]
    fn test_mock() {
        let mock = mock("GET /");

        assert_eq!("GET /", mock.request_line);
    }
}
<commit_msg>Add mockito::reset()<commit_after>#[cfg(feature = "use_hyper")]
extern crate hyper;
#[cfg(feature = "use_hyper")]
extern crate url as servo_url;

pub mod server;
pub mod client;
pub mod mock;
pub mod url;
#[cfg(feature = "use_hyper")]
pub mod mockable_hyper;
pub mod mockable_tcp_stream;

pub type Url<'a> = url::Url<'a>;
pub type Mock = mock::Mock;

pub fn mock(method: &str, path: &str) -> Mock {
    Mock::new(method, path)
}

pub fn reset() {
    client::reset();
}

#[cfg(test)]
#[cfg(feature = "mock_hyper")]
mod mock_hyper_tests {
    use hyper::Client;
    use hyper::header::Connection;
    use server;
    use url::Url;
    use std::io::Read;

    #[test]
    fn test_proxying() {
        super::start();

        let client = Client::new();
        let mut res = client.get(Url("http://www.example.com"))
            .header(Connection::close())
            .send()
            .unwrap();

        let mut body = String::new();
        res.read_to_string(&mut body).unwrap();

        assert_eq!(body, "Hello world");
    }
}

#[cfg(test)]
#[cfg(feature = "mock_tcp_stream")]
mod mock_tcp_stream_tests {
    use MockBuilder;
    use super::mock;

    #[test]
    fn test_mock() {
        let mock = mock("GET /");

        assert_eq!("GET /", mock.request_line);
    }
}
<|endoftext|>


<commit_before>//! A library providing a generic connection pool.
#![warn(missing_docs)]
#![doc(html_root_url="https://sfackler.github.io/r2d2/doc/v0.5.8")]

#[macro_use]
extern crate log;
extern crate time;
extern crate debug_builders;

use debug_builders::DebugStruct;
use std::collections::VecDeque;
use std::error::Error;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex, MutexGuard, Condvar};
use time::{Duration, SteadyTime};

#[doc(inline)]
pub use config::Config;

use task::ScheduledThreadPool;

pub mod config;
mod task;
mod thunk;

/// A trait which provides connection-specific functionality.
pub trait ManageConnection: Send + Sync + 'static {
    /// The connection type this manager deals with.
    type Connection: Send + 'static;

    /// The error type returned by `Connection`s.
    type Error: 'static;

    /// Attempts to create a new connection.
    fn connect(&self) -> Result<Self::Connection, Self::Error>;

    /// Determines if the connection is still connected to the database.
    ///
    /// A standard implementation would check if a simple query like `SELECT 1`
    /// succeeds.
    fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error>;

    /// *Quickly* determines if the connection is no longer usable.
    ///
    /// This will be called synchronously every time a connection is returned
    /// to the pool, so it should *not* block. If it returns `true`, the
    /// connection will be discarded.
    ///
    /// For example, an implementation might check if the underlying TCP socket
    /// has disconnected. Implementations that do not support this kind of
    /// fast health check may simply return `false`.
    fn has_broken(&self, conn: &mut Self::Connection) -> bool;
}

/// A trait which handles errors reported by the `ManageConnection`.
pub trait HandleError<E>: Send+Sync+'static {
    /// Handles an error.
    fn handle_error(&self, error: E);
}

/// A `HandleError` which does nothing.
#[derive(Copy, Clone, Debug)]
pub struct NoopErrorHandler;

impl<E> HandleError<E> for NoopErrorHandler {
    fn handle_error(&self, _: E) {}
}

/// An `ErrorHandler` which logs at the error level.
#[derive(Copy, Clone, Debug)]
pub struct LoggingErrorHandler;

impl<E> HandleError<E> for LoggingErrorHandler where E: fmt::Debug {
    fn handle_error(&self, error: E) {
        error!("{:?}", error);
    }
}

/// A trait which allows for customization of connections.
pub trait CustomizeConnection<C, E>: Send+Sync+'static {
    /// Called with connections immediately after they are returned from
    /// `ManageConnection::connect`.
    ///
    /// The default implementation simply returns `Ok(())`.
    ///
    /// # Errors
    ///
    /// If this method returns an error, the connection will be discarded.
    #[allow(unused_variables)]
    fn on_acquire(&self, conn: &mut C) -> Result<(), E> {
        Ok(())
    }
}

/// A `CustomizeConnection` which does nothing.
#[derive(Copy, Clone, Debug)]
pub struct NoopConnectionCustomizer;

impl<C, E> CustomizeConnection<C, E> for NoopConnectionCustomizer {}

struct PoolInternals<C> {
    conns: VecDeque<C>,
    num_conns: u32,
}

struct SharedPool<M> where M: ManageConnection {
    config: Config<M::Connection, M::Error>,
    manager: M,
    internals: Mutex<PoolInternals<M::Connection>>,
    cond: Condvar,
    thread_pool: ScheduledThreadPool,
}

fn add_connection<M>(delay: Duration, shared: &Arc<SharedPool<M>>) where M: ManageConnection {
    let new_shared = shared.clone();
    shared.thread_pool.run_after(delay, move || {
        let shared = new_shared;
        let conn = shared.manager.connect().and_then(|mut conn| {
            shared.config.connection_customizer().on_acquire(&mut conn).map(|_| conn)
        });
        match conn {
            Ok(conn) => {
                let mut internals = shared.internals.lock().unwrap();
                internals.conns.push_back(conn);
                internals.num_conns += 1;
                shared.cond.notify_one();
            }
            Err(err) => {
                shared.config.error_handler().handle_error(err);
                add_connection(Duration::seconds(1), &shared);
            },
        }
    });
}

/// A generic connection pool.
pub struct Pool<M> where M: ManageConnection {
    shared: Arc<SharedPool<M>>,
}

impl<M> Drop for Pool<M> where M: ManageConnection {
    fn drop(&mut self) {
        self.shared.thread_pool.clear();
    }
}

impl<M> fmt::Debug for Pool<M> where M: ManageConnection + fmt::Debug {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        DebugStruct::new(fmt, "Pool")
            .field("idle_connections", &self.shared.internals.lock().unwrap().conns.len())
            .field("config", &self.shared.config)
            .field("manager", &self.shared.manager)
            .finish()
    }
}

/// An error returned by `Pool::new` if it fails to initialize connections.
#[derive(Debug)]
pub struct InitializationError;

impl fmt::Display for InitializationError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(self.description())
    }
}

impl Error for InitializationError {
    fn description(&self) -> &str {
        "Unable to initialize connections"
    }
}

/// An error returned by `Pool::get` if it times out without retrieving a connection.
#[derive(Debug)]
pub struct GetTimeout;

impl fmt::Display for GetTimeout {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(self.description())
    }
}

impl Error for GetTimeout {
    fn description(&self) -> &str {
        "Timed out while waiting for a connection"
    }
}

impl<M> Pool<M> where M: ManageConnection {
    /// Creates a new connection pool.
    ///
    /// Returns an `Err` value if `initialization_fail_fast` is set to true in
    /// the configuration and the pool is unable to open all of its
    /// connections.
    pub fn new(config: Config<M::Connection, M::Error>, manager: M)
               -> Result<Pool<M>, InitializationError> {
        let internals = PoolInternals {
            conns: VecDeque::new(),
            num_conns: 0,
        };

        let shared = Arc::new(SharedPool {
            thread_pool: ScheduledThreadPool::new(config.helper_threads() as usize),
            config: config,
            manager: manager,
            internals: Mutex::new(internals),
            cond: Condvar::new(),
        });

        for _ in 0..shared.config.pool_size() {
            add_connection(Duration::zero(), &shared);
        }

        if shared.config.initialization_fail_fast() {
            let end = SteadyTime::now() + shared.config.connection_timeout();
            let mut internals = shared.internals.lock().unwrap();

            while internals.num_conns != shared.config.pool_size() {
                let wait = end - SteadyTime::now();
                if wait <= Duration::zero() {
                    return Err(InitializationError);
                }
                internals = shared.cond.wait_timeout_ms(internals,
                                                        wait.num_milliseconds() as u32)
                    .unwrap().0;
            }
        }

        Ok(Pool {
            shared: shared,
        })
    }

    fn get_inner(&self) -> Result<M::Connection, GetTimeout> {
        let end = SteadyTime::now() + self.shared.config.connection_timeout();
        let mut internals = self.shared.internals.lock().unwrap();

        loop {
            match internals.conns.pop_front() {
                Some(mut conn) => {
                    drop(internals);

                    if self.shared.config.test_on_check_out() {
                        if let Err(e) = self.shared.manager.is_valid(&mut conn) {
                            self.shared.config.error_handler().handle_error(e);
                            internals = self.shared.internals.lock().unwrap();
                            self.handle_broken(&mut internals);
                            continue
                        }
                    }

                    return Ok(conn);
                }
                None => {
                    let now = SteadyTime::now();
                    let mut timeout = (end - now).num_milliseconds();
                    if timeout < 0 {
                        timeout = 0
                    };
                    let (new_internals, no_timeout) =
                        self.shared.cond.wait_timeout_ms(internals, timeout as u32).unwrap();
                    internals = new_internals;

                    if !no_timeout {
                        return Err(GetTimeout);
                    }
                }
            }
        }
    }

    /// Retrieves a connection from the pool.
    ///
    /// Waits for at most `Config::connection_timeout` before returning an
    /// error.
    pub fn get<'a>(&'a self) -> Result<PooledConnection<'a, M>, GetTimeout> {
        Ok(PooledConnection {
            pool: MaybeArcPool::Ref(self),
            conn: Some(try!(self.get_inner())),
        })
    }

    /// Retrieves a connection from a reference counted pool.
    ///
    /// Unlike `Pool::get`, the resulting `PooledConnection` is not bound by
    /// any lifetime, as it will store the `Arc` internally.
    pub fn get_arc(pool: Arc<Self>) -> Result<PooledConnection<'static, M>, GetTimeout> {
        Ok(PooledConnection {
            conn: Some(try!(pool.get_inner())),
            pool: MaybeArcPool::Arc(pool),
        })
    }

    fn handle_broken(&self, internals: &mut MutexGuard<PoolInternals<M::Connection>>) {
        internals.num_conns -= 1;
        add_connection(Duration::zero(), &self.shared);
    }

    fn put_back(&self, mut conn: M::Connection) {
        // This is specified to be fast, but call it before locking anyways
        let broken = self.shared.manager.has_broken(&mut conn);

        let mut internals = self.shared.internals.lock().unwrap();
        if broken {
            self.handle_broken(&mut internals);
        } else {
            internals.conns.push_back(conn);
            self.shared.cond.notify_one();
        }
    }
}

enum MaybeArcPool<'a, M> where M: ManageConnection {
    Ref(&'a Pool<M>),
    Arc(Arc<Pool<M>>),
}

impl<'a, M> Deref for MaybeArcPool<'a, M> where M: ManageConnection {
    type Target = Pool<M>;

    fn deref(&self) -> &Pool<M> {
        match *self {
            MaybeArcPool::Ref(p) => p,
            MaybeArcPool::Arc(ref p) => &**p,
        }
    }
}

/// A smart pointer wrapping a connection.
pub struct PooledConnection<'a, M> where M: ManageConnection {
    pool: MaybeArcPool<'a, M>,
    conn: Option<M::Connection>,
}

impl<'a, M> fmt::Debug for PooledConnection<'a, M>
        where M: ManageConnection + fmt::Debug, M::Connection: fmt::Debug {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        DebugStruct::new(fmt, "PooledConnection")
            .field("pool", &*self.pool)
            .field("connection", self.conn.as_ref().unwrap())
            .finish()
    }
}

impl<'a, M> Drop for PooledConnection<'a, M> where M: ManageConnection {
    fn drop(&mut self) {
        self.pool.put_back(self.conn.take().unwrap());
    }
}

impl<'a, M> Deref for PooledConnection<'a, M> where M: ManageConnection {
    type Target = M::Connection;

    fn deref(&self) -> &M::Connection {
        self.conn.as_ref().unwrap()
    }
}

impl<'a, M> DerefMut for PooledConnection<'a, M> where M: ManageConnection {
    fn deref_mut(&mut self) -> &mut M::Connection {
        self.conn.as_mut().unwrap()
    }
}
<commit_msg>Disallow third party creation of error types<commit_after>//! A library providing a generic connection pool.
#![warn(missing_docs)]
#![doc(html_root_url="https://sfackler.github.io/r2d2/doc/v0.5.8")]

#[macro_use]
extern crate log;
extern crate time;
extern crate debug_builders;

use debug_builders::DebugStruct;
use std::collections::VecDeque;
use std::error::Error;
use std::fmt;
use std::ops::{Deref, DerefMut};
use std::sync::{Arc, Mutex, MutexGuard, Condvar};
use time::{Duration, SteadyTime};

#[doc(inline)]
pub use config::Config;

use task::ScheduledThreadPool;

pub mod config;
mod task;
mod thunk;

/// A trait which provides connection-specific functionality.
pub trait ManageConnection: Send + Sync + 'static {
    /// The connection type this manager deals with.
    type Connection: Send + 'static;

    /// The error type returned by `Connection`s.
    type Error: 'static;

    /// Attempts to create a new connection.
    fn connect(&self) -> Result<Self::Connection, Self::Error>;

    /// Determines if the connection is still connected to the database.
    ///
    /// A standard implementation would check if a simple query like `SELECT 1`
    /// succeeds.
    fn is_valid(&self, conn: &mut Self::Connection) -> Result<(), Self::Error>;

    /// *Quickly* determines if the connection is no longer usable.
    ///
    /// This will be called synchronously every time a connection is returned
    /// to the pool, so it should *not* block. If it returns `true`, the
    /// connection will be discarded.
    ///
    /// For example, an implementation might check if the underlying TCP socket
    /// has disconnected. Implementations that do not support this kind of
    /// fast health check may simply return `false`.
    fn has_broken(&self, conn: &mut Self::Connection) -> bool;
}

/// A trait which handles errors reported by the `ManageConnection`.
pub trait HandleError<E>: Send+Sync+'static {
    /// Handles an error.
    fn handle_error(&self, error: E);
}

/// A `HandleError` which does nothing.
#[derive(Copy, Clone, Debug)]
pub struct NoopErrorHandler;

impl<E> HandleError<E> for NoopErrorHandler {
    fn handle_error(&self, _: E) {}
}

/// An `ErrorHandler` which logs at the error level.
#[derive(Copy, Clone, Debug)]
pub struct LoggingErrorHandler;

impl<E> HandleError<E> for LoggingErrorHandler where E: fmt::Debug {
    fn handle_error(&self, error: E) {
        error!("{:?}", error);
    }
}

/// A trait which allows for customization of connections.
pub trait CustomizeConnection<C, E>: Send+Sync+'static {
    /// Called with connections immediately after they are returned from
    /// `ManageConnection::connect`.
    ///
    /// The default implementation simply returns `Ok(())`.
    ///
    /// # Errors
    ///
    /// If this method returns an error, the connection will be discarded.
    #[allow(unused_variables)]
    fn on_acquire(&self, conn: &mut C) -> Result<(), E> {
        Ok(())
    }
}

/// A `CustomizeConnection` which does nothing.
#[derive(Copy, Clone, Debug)]
pub struct NoopConnectionCustomizer;

impl<C, E> CustomizeConnection<C, E> for NoopConnectionCustomizer {}

struct PoolInternals<C> {
    conns: VecDeque<C>,
    num_conns: u32,
}

struct SharedPool<M> where M: ManageConnection {
    config: Config<M::Connection, M::Error>,
    manager: M,
    internals: Mutex<PoolInternals<M::Connection>>,
    cond: Condvar,
    thread_pool: ScheduledThreadPool,
}

fn add_connection<M>(delay: Duration, shared: &Arc<SharedPool<M>>) where M: ManageConnection {
    let new_shared = shared.clone();
    shared.thread_pool.run_after(delay, move || {
        let shared = new_shared;
        let conn = shared.manager.connect().and_then(|mut conn| {
            shared.config.connection_customizer().on_acquire(&mut conn).map(|_| conn)
        });
        match conn {
            Ok(conn) => {
                let mut internals = shared.internals.lock().unwrap();
                internals.conns.push_back(conn);
                internals.num_conns += 1;
                shared.cond.notify_one();
            }
            Err(err) => {
                shared.config.error_handler().handle_error(err);
                add_connection(Duration::seconds(1), &shared);
            },
        }
    });
}

/// A generic connection pool.
pub struct Pool<M> where M: ManageConnection {
    shared: Arc<SharedPool<M>>,
}

impl<M> Drop for Pool<M> where M: ManageConnection {
    fn drop(&mut self) {
        self.shared.thread_pool.clear();
    }
}

impl<M> fmt::Debug for Pool<M> where M: ManageConnection + fmt::Debug {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        DebugStruct::new(fmt, "Pool")
            .field("idle_connections", &self.shared.internals.lock().unwrap().conns.len())
            .field("config", &self.shared.config)
            .field("manager", &self.shared.manager)
            .finish()
    }
}

/// An error returned by `Pool::new` if it fails to initialize connections.
#[derive(Debug)]
pub struct InitializationError(());

impl fmt::Display for InitializationError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(self.description())
    }
}

impl Error for InitializationError {
    fn description(&self) -> &str {
        "Unable to initialize connections"
    }
}

/// An error returned by `Pool::get` if it times out without retrieving a connection.
#[derive(Debug)]
pub struct GetTimeout(());

impl fmt::Display for GetTimeout {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.write_str(self.description())
    }
}

impl Error for GetTimeout {
    fn description(&self) -> &str {
        "Timed out while waiting for a connection"
    }
}

impl<M> Pool<M> where M: ManageConnection {
    /// Creates a new connection pool.
    ///
    /// Returns an `Err` value if `initialization_fail_fast` is set to true in
    /// the configuration and the pool is unable to open all of its
    /// connections.
    pub fn new(config: Config<M::Connection, M::Error>, manager: M)
               -> Result<Pool<M>, InitializationError> {
        let internals = PoolInternals {
            conns: VecDeque::new(),
            num_conns: 0,
        };

        let shared = Arc::new(SharedPool {
            thread_pool: ScheduledThreadPool::new(config.helper_threads() as usize),
            config: config,
            manager: manager,
            internals: Mutex::new(internals),
            cond: Condvar::new(),
        });

        for _ in 0..shared.config.pool_size() {
            add_connection(Duration::zero(), &shared);
        }

        if shared.config.initialization_fail_fast() {
            let end = SteadyTime::now() + shared.config.connection_timeout();
            let mut internals = shared.internals.lock().unwrap();

            while internals.num_conns != shared.config.pool_size() {
                let wait = end - SteadyTime::now();
                if wait <= Duration::zero() {
                    return Err(InitializationError(()));
                }
                internals = shared.cond.wait_timeout_ms(internals,
                                                        wait.num_milliseconds() as u32)
                    .unwrap().0;
            }
        }

        Ok(Pool {
            shared: shared,
        })
    }

    fn get_inner(&self) -> Result<M::Connection, GetTimeout> {
        let end = SteadyTime::now() + self.shared.config.connection_timeout();
        let mut internals = self.shared.internals.lock().unwrap();

        loop {
            match internals.conns.pop_front() {
                Some(mut conn) => {
                    drop(internals);

                    if self.shared.config.test_on_check_out() {
                        if let Err(e) = self.shared.manager.is_valid(&mut conn) {
                            self.shared.config.error_handler().handle_error(e);
                            internals = self.shared.internals.lock().unwrap();
                            self.handle_broken(&mut internals);
                            continue
                        }
                    }

                    return Ok(conn);
                }
                None => {
                    let now = SteadyTime::now();
                    let mut timeout = (end - now).num_milliseconds();
                    if timeout < 0 {
                        timeout = 0
                    };
                    let (new_internals, no_timeout) =
                        self.shared.cond.wait_timeout_ms(internals, timeout as u32).unwrap();
                    internals = new_internals;

                    if !no_timeout {
                        return Err(GetTimeout(()));
                    }
                }
            }
        }
    }

    /// Retrieves a connection from the pool.
    ///
    /// Waits for at most `Config::connection_timeout` before returning an
    /// error.
    pub fn get<'a>(&'a self) -> Result<PooledConnection<'a, M>, GetTimeout> {
        Ok(PooledConnection {
            pool: MaybeArcPool::Ref(self),
            conn: Some(try!(self.get_inner())),
        })
    }

    /// Retrieves a connection from a reference counted pool.
    ///
    /// Unlike `Pool::get`, the resulting `PooledConnection` is not bound by
    /// any lifetime, as it will store the `Arc` internally.
    pub fn get_arc(pool: Arc<Self>) -> Result<PooledConnection<'static, M>, GetTimeout> {
        Ok(PooledConnection {
            conn: Some(try!(pool.get_inner())),
            pool: MaybeArcPool::Arc(pool),
        })
    }

    fn handle_broken(&self, internals: &mut MutexGuard<PoolInternals<M::Connection>>) {
        internals.num_conns -= 1;
        add_connection(Duration::zero(), &self.shared);
    }

    fn put_back(&self, mut conn: M::Connection) {
        // This is specified to be fast, but call it before locking anyways
        let broken = self.shared.manager.has_broken(&mut conn);

        let mut internals = self.shared.internals.lock().unwrap();
        if broken {
            self.handle_broken(&mut internals);
        } else {
            internals.conns.push_back(conn);
            self.shared.cond.notify_one();
        }
    }
}

enum MaybeArcPool<'a, M> where M: ManageConnection {
    Ref(&'a Pool<M>),
    Arc(Arc<Pool<M>>),
}

impl<'a, M> Deref for MaybeArcPool<'a, M> where M: ManageConnection {
    type Target = Pool<M>;

    fn deref(&self) -> &Pool<M> {
        match *self {
            MaybeArcPool::Ref(p) => p,
            MaybeArcPool::Arc(ref p) => &**p,
        }
    }
}

/// A smart pointer wrapping a connection.
pub struct PooledConnection<'a, M> where M: ManageConnection {
    pool: MaybeArcPool<'a, M>,
    conn: Option<M::Connection>,
}

impl<'a, M> fmt::Debug for PooledConnection<'a, M>
        where M: ManageConnection + fmt::Debug, M::Connection: fmt::Debug {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        DebugStruct::new(fmt, "PooledConnection")
            .field("pool", &*self.pool)
            .field("connection", self.conn.as_ref().unwrap())
            .finish()
    }
}

impl<'a, M> Drop for PooledConnection<'a, M> where M: ManageConnection {
    fn drop(&mut self) {
        self.pool.put_back(self.conn.take().unwrap());
    }
}

impl<'a, M> Deref for PooledConnection<'a, M> where M: ManageConnection {
    type Target = M::Connection;

    fn deref(&self) -> &M::Connection {
        self.conn.as_ref().unwrap()
    }
}

impl<'a, M> DerefMut for PooledConnection<'a, M> where M: ManageConnection {
    fn deref_mut(&mut self) -> &mut M::Connection {
        self.conn.as_mut().unwrap()
    }
}
<|endoftext|>


<commit_before>extern crate log;
extern crate simplelog;

use goblin::pe::PE;
use goblin::Object;
use log::{debug, error, info, trace};
use rayon::prelude::*;
use std::env;
use std::fs;
use std::io::prelude::*;
use zydis;

// TODO: this really shouldn't be public. only for testing.
// TODO: what is the `cfg(doctest)` setting?
//#[cfg(test)]
pub mod rsrc;

pub struct Config {
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 2 {
            return Err("not enough arguments");
        }

        let filename = args[1].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

#[derive(Debug)]
pub enum Error {
    FileAccess,
    FileFormat,
    NotImplemented,
    ValueError,
}

/// Round the given value up to the next multiple of the given base.
///
/// # Panics
///
///   - Base `b` must be at least `2`.
///
/// # Examples
///
/// ```
/// use lancelot::*;
/// assert_eq!(align(0, 2), 0);
/// assert_eq!(align(1, 2), 2);
/// assert_eq!(align(2, 2), 2);
/// assert_eq!(align(3, 2), 4);
/// assert_eq!(align(4, 2), 4);
/// ```
pub fn align(i: usize, b: usize) -> usize {
    if b < 2 {
        panic!("base `b` must be at least: 2");
    }
    let rem = i % b;
    if rem == 0 {
        i
    } else {
        i + (b - rem)
    }
}

pub fn hexdump_ascii(b: u8) -> char {
    if b.is_ascii_graphic() || b == b' ' {
        b as char
    } else {
        '.'
    }
}

pub fn hexdump(buf: &[u8], offset: usize) -> String {
    // 01234567:  00 01 02 03 04 05 06 07  ...............
    // <prefix>   <hex col>                <ascii col>

    let padding = "  ";

    let padding_size = 2;
    let hex_col_size = 3;
    let ascii_col_size = 1;
    let prefix_size = 8 + 1;
    let newline_size = 1;
    let line_size = prefix_size
        + padding_size
        + 16 * hex_col_size
        + padding_size
        + 16 * ascii_col_size
        + newline_size;
    let line_count = align(buf.len(), 0x10) / 0x10;

    let mut ret = String::with_capacity(line_count * line_size);

    let mut line = String::with_capacity(line_size);
    let mut remaining_count = buf.len();
    for line_index in 0..line_count {
        let line_elem_count = 0x10.min(remaining_count);
        let padding_elem_count = 0x10 - line_elem_count;

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        // ^^^^^^^^^
        line.push_str(format!("{:08x}:", offset + 0x10 * line_index).as_str());

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //          ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //            ^^^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push_str(format!("{:02x} ", elem).as_str());
        }
        for _ in 0..padding_elem_count {
            line.push_str("   ");
        }

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                   ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                     ^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push(hexdump_ascii(*elem))
        }
        for _ in 0..padding_elem_count {
            line.push(' ');
        }
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                                    ^
        line.push('\n');

        ret.push_str(line.as_str());
        line.truncate(0x0);
        remaining_count -= line_elem_count;
    }

    ret
}

fn foo(ws: &Workspace) -> Result<(), Error> {
    if let Object::PE(pe) = ws.get_obj().unwrap() {
        info!("module: {}", pe.name.unwrap_or("(unknown)"));

        info!("bitness: {}", if pe.is_64 { "64" } else { "32" });
        info!("image base: 0x{:x}", pe.image_base);
        info!("entry rva: 0x{:x}", pe.entry);

        // like:
        //
        //     sections:
        //       - .text
        //         raw size:     0x18aa00
        //         virtual size: 0x18a9a8
        info!("sections:");
        for section in ws.sections.iter() {
            info!("  - {}", section.name);
            info!("    virtual address: 0x{:x}", section.addr);
            info!("    virtual size: 0x{:x}", section.buf.len());

            info!(
                "\n{}",
                hexdump(&section.buf[..0x1C], pe.image_base + section.addr as usize)
            );

            let decoder =
                zydis::Decoder::new(zydis::MachineMode::Long64, zydis::AddressWidth::_64).unwrap();
            let insns: Vec<_> = section
                .buf
                .par_windows(0x10)
                .map(|ibuf| decoder.decode(ibuf))
                .collect();

            info!("total instructions: {}", insns.len());

            info!(
                "successful disassembles: {}",
                insns.par_iter().filter(|insn| insn.is_ok()).count()
            );

            info!(
                "valid instructions: {}",
                insns
                    .par_iter()
                    .filter(|insn| match insn {
                        Ok(Some(_)) => true,
                        _ => false,
                    })
                    .count()
            );
        }
    }

    Ok(())
}

pub fn read_file(filename: &str) -> Result<Vec<u8>, Error> {
    debug!("read_file: {:?}", filename);

    let mut buf = Vec::new();
    {
        debug!("reading file: {}", filename);
        let mut f = match fs::File::open(filename) {
            Ok(f) => f,
            Err(_) => {
                error!("failed to open file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        let bytes_read = match f.read_to_end(&mut buf) {
            Ok(c) => c,
            Err(_) => {
                error!("failed to read entire file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        debug!("read {} bytes", bytes_read);
        if bytes_read < 0x10 {
            error!("file too small: {}", filename);
            return Err(Error::FileFormat);
        }
    }

    Ok(buf)
}

pub struct Section {
    pub name: String,
    pub addr: u32,
    pub buf: Vec<u8>,
}

pub struct Workspace {
    pub filename: String,
    pub buf: Vec<u8>,
    pub sections: Vec<Section>,
}

impl Workspace {
    /// Parse the given file into its object.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///
    /// # Examples
    ///
    /// ```
    /// use goblin::Object;
    /// use matches::matches;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// assert!(matches!(ws.get_obj().unwrap(), Object::PE(_)));
    /// ```
    ///
    /// you might be tempted to maintain a method `get_pe`,
    /// however, i don't think this is a good idea:
    /// its fragile because the file type may not be PE.
    /// therefore, force clients to be explicit:
    ///
    /// ```
    /// use goblin::Object;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// if let Object::PE(_) = ws.get_obj().unwrap() {
    ///     // everyone is happy!
    /// }
    /// ```
    ///
    /// TODO: demonstrate `Error::FileFormat`.
    pub fn get_obj(&self) -> Result<Object, Error> {
        let obj = match Object::parse(&self.buf) {
            Ok(o) => o,
            Err(e) => {
                error!("failed to parse file: {} error: {:?}", self.filename, e);
                return Err(Error::FileFormat);
            }
        };

        match obj {
            Object::Unknown(_) => {
                error!(
                    "unknown file format, magic: | {:02X} {:02X} | '{}{}' ",
                    self.buf[0],
                    self.buf[1],
                    hexdump_ascii(self.buf[0]),
                    hexdump_ascii(self.buf[1])
                );

                Err(Error::FileFormat)
            }
            _ => Ok(obj),
        }
    }

    pub fn get_pe(&self) -> Result<PE, Error> {
        match self.get_obj()? {
            Object::PE(pe) => {
                return Ok(pe);
            }
            _ => {
                error!("not a PE file");
                return Err(Error::ValueError);
            }
        }
    }

    /// Construct a workspace from the module with the given contents.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///   - `Error::NotImplemented`: when not a PE file.
    ///
    /// # Examples
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let buf = get_buf(Rsrc::K32);
    /// let ws = Workspace::from_buf("kernel32.dll", buf).unwrap();
    /// ```
    ///
    /// TODO: demonstrate ELF file behavior.
    /// TODO: demonstrate MachO file behavior.
    /// TODO: demonstrate unknown file behavior.
    pub fn from_buf(filename: &str, buf: Vec<u8>) -> Result<Workspace, Error> {
        let mut ws = Workspace {
            filename: filename.to_string(),
            buf: buf.clone(),
            sections: vec![],
        };

        match ws.get_obj()? {
            Object::PE(pe) => {
                info!("found PE file");

                // TODO: load PE header, too

                ws.sections
                    .extend(pe.sections.iter().map(|section| -> Section {
                        // TODO: i'm sure this can be abused.
                        // TODO: add tests for weird section names.
                        let name = String::from_utf8_lossy(&section.name[..]).into_owned();

                        // TODO: figure out if we will work with usize, or u64, or what,
                        // then assert usize is ok.
                        // ref: `usize::max_value()`
                        let mut secbuf = vec![0; align(section.virtual_size as usize, 0x200)];

                        {
                            let secsize = section.size_of_raw_data as usize;
                            let rawbuf = &mut secbuf[..secsize];
                            let pstart = section.pointer_to_raw_data as usize;
                            rawbuf.copy_from_slice(&buf[pstart..pstart + secsize]);
                        }

                        Section {
                            name: name,
                            addr: section.virtual_address,
                            buf: secbuf,
                        }
                    }));
            }
            Object::Elf(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Mach(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Archive(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Unknown(_) => {
                return Err(Error::NotImplemented);
            }
        }

        Ok(ws)
    }

    /// Construct a workspace from the module at given a file path.
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let path = get_path(Rsrc::K32);
    /// // This test resource file is mangled. Needs to be fixed before parsing.
    /// // Otherwise, the following would work:
    /// // let ws = Workspace::from_file(&path).unwrap();
    /// ```
    pub fn from_file(filename: &str) -> Result<Workspace, Error> {
        let buf = read_file(filename)?;
        Workspace::from_buf(filename, buf)
    }
}

pub fn run(args: &Config) -> Result<(), Error> {
    debug!("filename: {:?}", args.filename);
    let ws = Workspace::from_file(&args.filename)?;

    if let Object::PE(_) = ws.get_obj()? {
        foo(&ws).expect("failed to foo");
    }

    Ok(())
}
<commit_msg>lib: disassemble instructions during load<commit_after>extern crate log;
extern crate simplelog;

use goblin::pe::PE;
use goblin::Object;
use log::{debug, error, info, trace};
use rayon::prelude::*;
use std::env;
use std::fs;
use std::io::prelude::*;
use zydis;

// TODO: this really shouldn't be public. only for testing.
// TODO: what is the `cfg(doctest)` setting?
//#[cfg(test)]
pub mod rsrc;

pub struct Config {
    pub filename: String,
}

impl Config {
    pub fn from_args(args: env::Args) -> Result<Config, &'static str> {
        let args: Vec<String> = args.collect();

        if args.len() < 2 {
            return Err("not enough arguments");
        }

        let filename = args[1].clone();
        trace!("config: parsed filename: {:?}", filename);

        Ok(Config { filename })
    }
}

pub fn setup_logging(_args: &Config) {
    simplelog::TermLogger::init(simplelog::LevelFilter::Info, simplelog::Config::default())
        .expect("failed to setup logging");
}

#[derive(Debug)]
pub enum Error {
    FileAccess,
    FileFormat,
    NotImplemented,
    ValueError,
}

/// Round the given value up to the next multiple of the given base.
///
/// # Panics
///
///   - Base `b` must be at least `2`.
///
/// # Examples
///
/// ```
/// use lancelot::*;
/// assert_eq!(align(0, 2), 0);
/// assert_eq!(align(1, 2), 2);
/// assert_eq!(align(2, 2), 2);
/// assert_eq!(align(3, 2), 4);
/// assert_eq!(align(4, 2), 4);
/// ```
pub fn align(i: usize, b: usize) -> usize {
    if b < 2 {
        panic!("base `b` must be at least: 2");
    }
    let rem = i % b;
    if rem == 0 {
        i
    } else {
        i + (b - rem)
    }
}

pub fn hexdump_ascii(b: u8) -> char {
    if b.is_ascii_graphic() || b == b' ' {
        b as char
    } else {
        '.'
    }
}

pub fn hexdump(buf: &[u8], offset: usize) -> String {
    // 01234567:  00 01 02 03 04 05 06 07  ...............
    // <prefix>   <hex col>                <ascii col>

    let padding = "  ";

    let padding_size = 2;
    let hex_col_size = 3;
    let ascii_col_size = 1;
    let prefix_size = 8 + 1;
    let newline_size = 1;
    let line_size = prefix_size
        + padding_size
        + 16 * hex_col_size
        + padding_size
        + 16 * ascii_col_size
        + newline_size;
    let line_count = align(buf.len(), 0x10) / 0x10;

    let mut ret = String::with_capacity(line_count * line_size);

    let mut line = String::with_capacity(line_size);
    let mut remaining_count = buf.len();
    for line_index in 0..line_count {
        let line_elem_count = 0x10.min(remaining_count);
        let padding_elem_count = 0x10 - line_elem_count;

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        // ^^^^^^^^^
        line.push_str(format!("{:08x}:", offset + 0x10 * line_index).as_str());

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //          ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //            ^^^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push_str(format!("{:02x} ", elem).as_str());
        }
        for _ in 0..padding_elem_count {
            line.push_str("   ");
        }

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                   ^^
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                     ^
        for elem in &buf[line_index..line_index + line_elem_count] {
            line.push(hexdump_ascii(*elem))
        }
        for _ in 0..padding_elem_count {
            line.push(' ');
        }
        line.push_str(padding);

        // 01234567:  00 01 02 03 04 05 06 07  ...............
        //                                                    ^
        line.push('\n');

        ret.push_str(line.as_str());
        line.truncate(0x0);
        remaining_count -= line_elem_count;
    }

    ret
}

fn foo(ws: &Workspace) -> Result<(), Error> {
    if let Object::PE(pe) = ws.get_obj().unwrap() {
        info!("module: {}", pe.name.unwrap_or("(unknown)"));

        info!("bitness: {}", if pe.is_64 { "64" } else { "32" });
        info!("image base: 0x{:x}", pe.image_base);
        info!("entry rva: 0x{:x}", pe.entry);

        // like:
        //
        //     sections:
        //       - .text
        //         raw size:     0x18aa00
        //         virtual size: 0x18a9a8
        info!("sections:");
        for section in ws.sections.iter() {
            info!("  - {}", section.name);
            info!("    virtual address: 0x{:x}", section.addr);
            info!("    virtual size: 0x{:x}", section.buf.len());

            info!(
                "\n{}",
                hexdump(&section.buf[..0x1C], pe.image_base + section.addr as usize)
            );

            let decoder =
                zydis::Decoder::new(zydis::MachineMode::Long64, zydis::AddressWidth::_64).unwrap();
            let insns: Vec<_> = section
                .buf
                .par_windows(0x10)
                .map(|ibuf| decoder.decode(ibuf))
                .collect();

            info!("total instructions: {}", insns.len());

            info!(
                "successful disassembles: {}",
                insns.par_iter().filter(|insn| insn.is_ok()).count()
            );

            info!(
                "valid instructions: {}",
                insns
                    .par_iter()
                    .filter(|insn| match insn {
                        Ok(Some(_)) => true,
                        _ => false,
                    })
                    .count()
            );
        }
    }

    Ok(())
}

pub fn read_file(filename: &str) -> Result<Vec<u8>, Error> {
    debug!("read_file: {:?}", filename);

    let mut buf = Vec::new();
    {
        debug!("reading file: {}", filename);
        let mut f = match fs::File::open(filename) {
            Ok(f) => f,
            Err(_) => {
                error!("failed to open file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        let bytes_read = match f.read_to_end(&mut buf) {
            Ok(c) => c,
            Err(_) => {
                error!("failed to read entire file: {}", filename);
                return Err(Error::FileAccess);
            }
        };
        debug!("read {} bytes", bytes_read);
        if bytes_read < 0x10 {
            error!("file too small: {}", filename);
            return Err(Error::FileFormat);
        }
    }

    Ok(buf)
}

pub struct Section {
    pub name: String,
    pub addr: u32,
    pub buf: Vec<u8>,
    pub insns: Vec<Option<zydis::ffi::DecodedInstruction>>,
}

pub struct Workspace {
    pub filename: String,
    pub buf: Vec<u8>,
    pub sections: Vec<Section>,
}

impl Workspace {
    /// Parse the given file into its object.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///
    /// # Examples
    ///
    /// ```
    /// use goblin::Object;
    /// use matches::matches;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// assert!(matches!(ws.get_obj().unwrap(), Object::PE(_)));
    /// ```
    ///
    /// you might be tempted to maintain a method `get_pe`,
    /// however, i don't think this is a good idea:
    /// its fragile because the file type may not be PE.
    /// therefore, force clients to be explicit:
    ///
    /// ```
    /// use goblin::Object;
    /// use lancelot::rsrc::*;
    /// let ws = get_workspace(Rsrc::K32);
    /// if let Object::PE(_) = ws.get_obj().unwrap() {
    ///     // everyone is happy!
    /// }
    /// ```
    ///
    /// TODO: demonstrate `Error::FileFormat`.
    pub fn get_obj(&self) -> Result<Object, Error> {
        let obj = match Object::parse(&self.buf) {
            Ok(o) => o,
            Err(e) => {
                error!("failed to parse file: {} error: {:?}", self.filename, e);
                return Err(Error::FileFormat);
            }
        };

        match obj {
            Object::Unknown(_) => {
                error!(
                    "unknown file format, magic: | {:02X} {:02X} | '{}{}' ",
                    self.buf[0],
                    self.buf[1],
                    hexdump_ascii(self.buf[0]),
                    hexdump_ascii(self.buf[1])
                );

                Err(Error::FileFormat)
            }
            _ => Ok(obj),
        }
    }

    pub fn get_pe(&self) -> Result<PE, Error> {
        match self.get_obj()? {
            Object::PE(pe) => {
                return Ok(pe);
            }
            _ => {
                error!("not a PE file");
                return Err(Error::ValueError);
            }
        }
    }

    /// Construct a workspace from the module with the given contents.
    ///
    /// # Errors
    ///   - `Error::FileFormat`: when not able to be parsed by Goblin.
    ///   - `Error::NotImplemented`: when not a PE file.
    ///
    /// # Examples
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let buf = get_buf(Rsrc::K32);
    /// let ws = Workspace::from_buf("kernel32.dll", buf).unwrap();
    /// ```
    ///
    /// TODO: demonstrate ELF file behavior.
    /// TODO: demonstrate MachO file behavior.
    /// TODO: demonstrate unknown file behavior.
    pub fn from_buf(filename: &str, buf: Vec<u8>) -> Result<Workspace, Error> {
        let mut ws = Workspace {
            filename: filename.to_string(),
            buf: buf.clone(),
            sections: vec![],
        };

        match ws.get_obj()? {
            Object::PE(pe) => {
                info!("found PE file");

                let machine;
                let mode;
                if pe.is_64 {
                    machine = zydis::MachineMode::Long64;
                    mode = zydis::AddressWidth::_64;
                } else {
                    // TODO: not sure what `LongCompat32` means, vs `Legacy32`.
                    machine = zydis::MachineMode::LongCompat32;
                    mode = zydis::AddressWidth::_32;
                }
                // TODO: save off decoder into workspace.
                let decoder = zydis::Decoder::new(machine, mode).unwrap();

                // TODO: load PE header, too

                ws.sections
                    .extend(pe.sections.iter().map(|section| -> Section {
                        // TODO: i'm sure this can be abused.
                        // TODO: add tests for weird section names.
                        let name = String::from_utf8_lossy(&section.name[..]).into_owned();

                        // TODO: figure out if we will work with usize, or u64, or what,
                        // then assert usize is ok.
                        // ref: `usize::max_value()`
                        let mut secbuf = vec![0; align(section.virtual_size as usize, 0x200)];

                        {
                            let secsize = section.size_of_raw_data as usize;
                            let rawbuf = &mut secbuf[..secsize];
                            let pstart = section.pointer_to_raw_data as usize;
                            rawbuf.copy_from_slice(&buf[pstart..pstart + secsize]);
                        }

                        let insns: Vec<_> = secbuf
                            .par_windows(0x10)
                            .map(|ibuf| match decoder.decode(ibuf) {
                                Ok(Some(insn)) => Some(insn),
                                _ => None,
                            })
                            .collect();

                        Section {
                            name: name,
                            addr: section.virtual_address,
                            buf: secbuf,
                            insns: insns,
                        }
                    }));
            }
            Object::Elf(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Mach(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Archive(_) => {
                return Err(Error::NotImplemented);
            }
            Object::Unknown(_) => {
                return Err(Error::NotImplemented);
            }
        }

        Ok(ws)
    }

    /// Construct a workspace from the module at given a file path.
    ///
    /// ```
    /// use lancelot::*;
    /// use lancelot::rsrc::*;
    /// let path = get_path(Rsrc::K32);
    /// // This test resource file is mangled. Needs to be fixed before parsing.
    /// // Otherwise, the following would work:
    /// // let ws = Workspace::from_file(&path).unwrap();
    /// ```
    pub fn from_file(filename: &str) -> Result<Workspace, Error> {
        let buf = read_file(filename)?;
        Workspace::from_buf(filename, buf)
    }
}

pub fn run(args: &Config) -> Result<(), Error> {
    debug!("filename: {:?}", args.filename);
    let ws = Workspace::from_file(&args.filename)?;

    if let Object::PE(_) = ws.get_obj()? {
        foo(&ws).expect("failed to foo");
    }

    Ok(())
}
<|endoftext|>