<commit_before>//! This module handles auto-magic editing actions applied together with users
//! edits. For example, if the user typed
//!
//! ```text
//!     foo
//!         .bar()
//!         .baz()
//!     |   // <- cursor is here
//! ```
//!
//! and types `.` next, we want to indent the dot.
//!
//! Language server executes such typing assists synchronously. That is, they
//! block user's typing and should be pretty fast for this reason!

mod on_enter;

use ra_db::{FilePosition, SourceDatabase};
use ra_fmt::leading_indent;
use ra_ide_db::{source_change::SourceFileEdit, RootDatabase};
use ra_syntax::{
    algo::find_node_at_offset,
    ast::{self, AstToken},
    AstNode, SourceFile, TextRange, TextSize,
};

use ra_text_edit::TextEdit;

use crate::SourceChange;

pub(crate) use on_enter::on_enter;

pub(crate) const TRIGGER_CHARS: &str = ".=>";

// Feature: On Typing Assists
//
// Some features trigger on typing certain characters:
//
// - typing `let =` tries to smartly add `;` if `=` is followed by an existing expression
// - Enter inside comments automatically inserts `///`
// - typing `.` in a chain method call auto-indents
pub(crate) fn on_char_typed(
    db: &RootDatabase,
    position: FilePosition,
    char_typed: char,
) -> Option<SourceChange> {
    assert!(TRIGGER_CHARS.contains(char_typed));
    let file = &db.parse(position.file_id).tree();
    assert_eq!(file.syntax().text().char_at(position.offset), Some(char_typed));
    let edit = on_char_typed_inner(file, position.offset, char_typed)?;
    Some(SourceFileEdit { file_id: position.file_id, edit }.into())
}

fn on_char_typed_inner(file: &SourceFile, offset: TextSize, char_typed: char) -> Option<TextEdit> {
    assert!(TRIGGER_CHARS.contains(char_typed));
    match char_typed {
        '.' => on_dot_typed(file, offset),
        '=' => on_eq_typed(file, offset),
        '>' => on_arrow_typed(file, offset),
        _ => unreachable!(),
    }
}

/// Returns an edit which should be applied after `=` was typed. Primarily,
/// this works when adding `let =`.
// FIXME: use a snippet completion instead of this hack here.
fn on_eq_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
    assert_eq!(file.syntax().text().char_at(offset), Some('='));
    let let_stmt: ast::LetStmt = find_node_at_offset(file.syntax(), offset)?;
    if let_stmt.semicolon_token().is_some() {
        return None;
    }
    if let Some(expr) = let_stmt.initializer() {
        let expr_range = expr.syntax().text_range();
        if expr_range.contains(offset) && offset != expr_range.start() {
            return None;
        }
        if file.syntax().text().slice(offset..expr_range.start()).contains_char('\n') {
            return None;
        }
    } else {
        return None;
    }
    let offset = let_stmt.syntax().text_range().end();
    Some(TextEdit::insert(offset, ";".to_string()))
}

/// Returns an edit which should be applied when a dot ('.') is typed on a blank line, indenting the line appropriately.
fn on_dot_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
    assert_eq!(file.syntax().text().char_at(offset), Some('.'));
    let whitespace =
        file.syntax().token_at_offset(offset).left_biased().and_then(ast::Whitespace::cast)?;

    let current_indent = {
        let text = whitespace.text();
        let newline = text.rfind('\n')?;
        &text[newline + 1..]
    };
    let current_indent_len = TextSize::of(current_indent);

    // Make sure dot is a part of call chain
    let field_expr = ast::FieldExpr::cast(whitespace.syntax().parent())?;
    let prev_indent = leading_indent(field_expr.syntax())?;
    let target_indent = format!("    {}", prev_indent);
    let target_indent_len = TextSize::of(&target_indent);
    if current_indent_len == target_indent_len {
        return None;
    }

    Some(TextEdit::replace(TextRange::new(offset - current_indent_len, offset), target_indent))
}

/// Adds a space after an arrow when `fn foo() { ... }` is turned into `fn foo() -> { ... }`
fn on_arrow_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
    let file_text = file.syntax().text();
    assert_eq!(file_text.char_at(offset), Some('>'));
    let after_arrow = offset + TextSize::of('>');
    if file_text.char_at(after_arrow) != Some('{') {
        return None;
    }
    if find_node_at_offset::<ast::RetType>(file.syntax(), offset).is_none() {
        return None;
    }

    Some(TextEdit::insert(after_arrow, " ".to_string()))
}

#[cfg(test)]
mod tests {
    use test_utils::{assert_eq_text, extract_offset};

    use super::*;

    fn do_type_char(char_typed: char, before: &str) -> Option<String> {
        let (offset, before) = extract_offset(before);
        let edit = TextEdit::insert(offset, char_typed.to_string());
        let mut before = before.to_string();
        edit.apply(&mut before);
        let parse = SourceFile::parse(&before);
        on_char_typed_inner(&parse.tree(), offset, char_typed).map(|it| {
            it.apply(&mut before);
            before.to_string()
        })
    }

    fn type_char(char_typed: char, before: &str, after: &str) {
        let actual = do_type_char(char_typed, before)
            .unwrap_or_else(|| panic!("typing `{}` did nothing", char_typed));

        assert_eq_text!(after, &actual);
    }

    fn type_char_noop(char_typed: char, before: &str) {
        let file_change = do_type_char(char_typed, before);
        assert!(file_change.is_none())
    }

    #[test]
    fn test_on_eq_typed() {
        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        // }
        // ");
        type_char(
            '=',
            r"
fn foo() {
    let foo <|> 1 + 1
}
",
            r"
fn foo() {
    let foo = 1 + 1;
}
",
        );
        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        //     let bar = 1;
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        //     let bar = 1;
        // }
        // ");
    }

    #[test]
    fn indents_new_chain_call() {
        type_char(
            '.',
            r"
            fn main() {
                xs.foo()
                <|>
            }
            ",
            r"
            fn main() {
                xs.foo()
                    .
            }
            ",
        );
        type_char_noop(
            '.',
            r"
            fn main() {
                xs.foo()
                    <|>
            }
            ",
        )
    }

    #[test]
    fn indents_new_chain_call_with_semi() {
        type_char(
            '.',
            r"
            fn main() {
                xs.foo()
                <|>;
            }
            ",
            r"
            fn main() {
                xs.foo()
                    .;
            }
            ",
        );
        type_char_noop(
            '.',
            r"
            fn main() {
                xs.foo()
                    <|>;
            }
            ",
        )
    }

    #[test]
    fn indents_continued_chain_call() {
        type_char(
            '.',
            r"
            fn main() {
                xs.foo()
                    .first()
                <|>
            }
            ",
            r"
            fn main() {
                xs.foo()
                    .first()
                    .
            }
            ",
        );
        type_char_noop(
            '.',
            r"
            fn main() {
                xs.foo()
                    .first()
                    <|>
            }
            ",
        );
    }

    #[test]
    fn indents_middle_of_chain_call() {
        type_char(
            '.',
            r"
            fn source_impl() {
                let var = enum_defvariant_list().unwrap()
                <|>
                    .nth(92)
                    .unwrap();
            }
            ",
            r"
            fn source_impl() {
                let var = enum_defvariant_list().unwrap()
                    .
                    .nth(92)
                    .unwrap();
            }
            ",
        );
        type_char_noop(
            '.',
            r"
            fn source_impl() {
                let var = enum_defvariant_list().unwrap()
                    <|>
                    .nth(92)
                    .unwrap();
            }
            ",
        );
    }

    #[test]
    fn dont_indent_freestanding_dot() {
        type_char_noop(
            '.',
            r"
            fn main() {
                <|>
            }
            ",
        );
        type_char_noop(
            '.',
            r"
            fn main() {
            <|>
            }
            ",
        );
    }

    #[test]
    fn adds_space_after_return_type() {
        type_char('>', "fn foo() -<|>{ 92 }", "fn foo() -> { 92 }")
    }
}
<commit_msg>Indent chain `.` even if there's more stuff afterwards<commit_after>//! This module handles auto-magic editing actions applied together with users
//! edits. For example, if the user typed
//!
//! ```text
//!     foo
//!         .bar()
//!         .baz()
//!     |   // <- cursor is here
//! ```
//!
//! and types `.` next, we want to indent the dot.
//!
//! Language server executes such typing assists synchronously. That is, they
//! block user's typing and should be pretty fast for this reason!

mod on_enter;

use ra_db::{FilePosition, SourceDatabase};
use ra_fmt::leading_indent;
use ra_ide_db::{source_change::SourceFileEdit, RootDatabase};
use ra_syntax::{
    algo::find_node_at_offset,
    ast::{self, AstToken},
    AstNode, SourceFile,
    SyntaxKind::{FIELD_EXPR, METHOD_CALL_EXPR},
    TextRange, TextSize,
};

use ra_text_edit::TextEdit;

use crate::SourceChange;

pub(crate) use on_enter::on_enter;

pub(crate) const TRIGGER_CHARS: &str = ".=>";

// Feature: On Typing Assists
//
// Some features trigger on typing certain characters:
//
// - typing `let =` tries to smartly add `;` if `=` is followed by an existing expression
// - Enter inside comments automatically inserts `///`
// - typing `.` in a chain method call auto-indents
pub(crate) fn on_char_typed(
    db: &RootDatabase,
    position: FilePosition,
    char_typed: char,
) -> Option<SourceChange> {
    assert!(TRIGGER_CHARS.contains(char_typed));
    let file = &db.parse(position.file_id).tree();
    assert_eq!(file.syntax().text().char_at(position.offset), Some(char_typed));
    let edit = on_char_typed_inner(file, position.offset, char_typed)?;
    Some(SourceFileEdit { file_id: position.file_id, edit }.into())
}

fn on_char_typed_inner(file: &SourceFile, offset: TextSize, char_typed: char) -> Option<TextEdit> {
    assert!(TRIGGER_CHARS.contains(char_typed));
    match char_typed {
        '.' => on_dot_typed(file, offset),
        '=' => on_eq_typed(file, offset),
        '>' => on_arrow_typed(file, offset),
        _ => unreachable!(),
    }
}

/// Returns an edit which should be applied after `=` was typed. Primarily,
/// this works when adding `let =`.
// FIXME: use a snippet completion instead of this hack here.
fn on_eq_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
    assert_eq!(file.syntax().text().char_at(offset), Some('='));
    let let_stmt: ast::LetStmt = find_node_at_offset(file.syntax(), offset)?;
    if let_stmt.semicolon_token().is_some() {
        return None;
    }
    if let Some(expr) = let_stmt.initializer() {
        let expr_range = expr.syntax().text_range();
        if expr_range.contains(offset) && offset != expr_range.start() {
            return None;
        }
        if file.syntax().text().slice(offset..expr_range.start()).contains_char('\n') {
            return None;
        }
    } else {
        return None;
    }
    let offset = let_stmt.syntax().text_range().end();
    Some(TextEdit::insert(offset, ";".to_string()))
}

/// Returns an edit which should be applied when a dot ('.') is typed on a blank line, indenting the line appropriately.
fn on_dot_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
    assert_eq!(file.syntax().text().char_at(offset), Some('.'));
    let whitespace =
        file.syntax().token_at_offset(offset).left_biased().and_then(ast::Whitespace::cast)?;

    let current_indent = {
        let text = whitespace.text();
        let newline = text.rfind('\n')?;
        &text[newline + 1..]
    };
    let current_indent_len = TextSize::of(current_indent);

    let parent = whitespace.syntax().parent();
    // Make sure dot is a part of call chain
    if !matches!(parent.kind(), FIELD_EXPR | METHOD_CALL_EXPR) {
        return None;
    }
    let prev_indent = leading_indent(&parent)?;
    let target_indent = format!("    {}", prev_indent);
    let target_indent_len = TextSize::of(&target_indent);
    if current_indent_len == target_indent_len {
        return None;
    }

    Some(TextEdit::replace(TextRange::new(offset - current_indent_len, offset), target_indent))
}

/// Adds a space after an arrow when `fn foo() { ... }` is turned into `fn foo() -> { ... }`
fn on_arrow_typed(file: &SourceFile, offset: TextSize) -> Option<TextEdit> {
    let file_text = file.syntax().text();
    assert_eq!(file_text.char_at(offset), Some('>'));
    let after_arrow = offset + TextSize::of('>');
    if file_text.char_at(after_arrow) != Some('{') {
        return None;
    }
    if find_node_at_offset::<ast::RetType>(file.syntax(), offset).is_none() {
        return None;
    }

    Some(TextEdit::insert(after_arrow, " ".to_string()))
}

#[cfg(test)]
mod tests {
    use test_utils::{assert_eq_text, extract_offset};

    use super::*;

    fn do_type_char(char_typed: char, before: &str) -> Option<String> {
        let (offset, before) = extract_offset(before);
        let edit = TextEdit::insert(offset, char_typed.to_string());
        let mut before = before.to_string();
        edit.apply(&mut before);
        let parse = SourceFile::parse(&before);
        on_char_typed_inner(&parse.tree(), offset, char_typed).map(|it| {
            it.apply(&mut before);
            before.to_string()
        })
    }

    fn type_char(char_typed: char, ra_fixture_before: &str, ra_fixture_after: &str) {
        let actual = do_type_char(char_typed, ra_fixture_before)
            .unwrap_or_else(|| panic!("typing `{}` did nothing", char_typed));

        assert_eq_text!(ra_fixture_after, &actual);
    }

    fn type_char_noop(char_typed: char, before: &str) {
        let file_change = do_type_char(char_typed, before);
        assert!(file_change.is_none())
    }

    #[test]
    fn test_on_eq_typed() {
        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        // }
        // ");
        type_char(
            '=',
            r"
fn foo() {
    let foo <|> 1 + 1
}
",
            r"
fn foo() {
    let foo = 1 + 1;
}
",
        );
        //     do_check(r"
        // fn foo() {
        //     let foo =<|>
        //     let bar = 1;
        // }
        // ", r"
        // fn foo() {
        //     let foo =;
        //     let bar = 1;
        // }
        // ");
    }

    #[test]
    fn indents_new_chain_call() {
        type_char(
            '.',
            r"
            fn main() {
                xs.foo()
                <|>
            }
            ",
            r"
            fn main() {
                xs.foo()
                    .
            }
            ",
        );
        type_char_noop(
            '.',
            r"
            fn main() {
                xs.foo()
                    <|>
            }
            ",
        )
    }

    #[test]
    fn indents_new_chain_call_with_semi() {
        type_char(
            '.',
            r"
            fn main() {
                xs.foo()
                <|>;
            }
            ",
            r"
            fn main() {
                xs.foo()
                    .;
            }
            ",
        );
        type_char_noop(
            '.',
            r"
            fn main() {
                xs.foo()
                    <|>;
            }
            ",
        )
    }

    #[test]
    fn indents_new_chain_call_with_let() {
        type_char(
            '.',
            r#"
fn main() {
    let _ = foo
    <|>
    bar()
}
"#,
            r#"
fn main() {
    let _ = foo
        .
    bar()
}
"#,
        );
    }

    #[test]
    fn indents_continued_chain_call() {
        type_char(
            '.',
            r"
            fn main() {
                xs.foo()
                    .first()
                <|>
            }
            ",
            r"
            fn main() {
                xs.foo()
                    .first()
                    .
            }
            ",
        );
        type_char_noop(
            '.',
            r"
            fn main() {
                xs.foo()
                    .first()
                    <|>
            }
            ",
        );
    }

    #[test]
    fn indents_middle_of_chain_call() {
        type_char(
            '.',
            r"
            fn source_impl() {
                let var = enum_defvariant_list().unwrap()
                <|>
                    .nth(92)
                    .unwrap();
            }
            ",
            r"
            fn source_impl() {
                let var = enum_defvariant_list().unwrap()
                    .
                    .nth(92)
                    .unwrap();
            }
            ",
        );
        type_char_noop(
            '.',
            r"
            fn source_impl() {
                let var = enum_defvariant_list().unwrap()
                    <|>
                    .nth(92)
                    .unwrap();
            }
            ",
        );
    }

    #[test]
    fn dont_indent_freestanding_dot() {
        type_char_noop(
            '.',
            r"
            fn main() {
                <|>
            }
            ",
        );
        type_char_noop(
            '.',
            r"
            fn main() {
            <|>
            }
            ",
        );
    }

    #[test]
    fn adds_space_after_return_type() {
        type_char('>', "fn foo() -<|>{ 92 }", "fn foo() -> { 92 }")
    }
}
<|endoftext|>


<commit_before>use std::borrow::ToOwned;
use std::env;
use std::error::Error;
use std::fmt;
use std::fs;
use std::io::{Read,Write};
use std::io;
use std::path;
use std::process;

use clap::{Arg,App,SubCommand,AppSettings};
use atomicwrites::{AtomicFile,AllowOverwrite};

use utils;
use utils::CustomPathExt;

#[inline]
fn get_pwd() -> path::PathBuf {
    env::current_dir().ok().expect("Failed to get CWD")
}

#[inline]
fn get_envvar(key: &str) -> Option<String> {
    match env::var(key) {
        Ok(x) => Some(x),
        Err(env::VarError::NotPresent) => None,
        Err(env::VarError::NotUnicode(_)) => panic!(format!("{} is not unicode.", key)),
    }
}

fn build_index(outfile: &path::Path, dir: &path::Path) -> MainResult<()> {
    if !dir.is_dir() {
        return Err(MainError::new("MATES_DIR must be a directory.").into());
    };

    let af = AtomicFile::new(&outfile, AllowOverwrite);
    let mut errors = false;

    try!(af.write(|outf| {
        for entry in try!(fs::read_dir(dir)) {
            let entry = match entry {
                Ok(x) => x,
                Err(e) => {
                    println!("Error while listing directory: {}", e);
                    errors = true;
                    continue;
                }
            };

            let pathbuf = entry.path();

            if pathbuf.str_extension().unwrap_or("") != "vcf" || !pathbuf.is_file() {
                continue;
            };

            let contact = match utils::Contact::from_file(&pathbuf) {
                Ok(x) => x,
                Err(e) => {
                    println!("Error while reading {}: {}", pathbuf.display(), e);
                    errors = true;
                    continue
                }
            };

            match utils::index_item_from_contact(&contact) {
                Ok(index_string) => {
                    try!(outf.write_all(index_string.as_bytes()));
                },
                Err(e) => {
                    println!("Error while indexing {}: {}", pathbuf.display(), e);
                    errors = true;
                    continue
                }
            };
        };
        Ok(())
    }));

    if errors {
        Err(MainError::new("Several errors happened while generating the index.").into())
    } else {
        Ok(())
    }
}

pub fn cli_main() {
    match cli_main_raw() {
        Err(e) => {
            writeln!(&mut io::stderr(), "{}", e).unwrap();
            process::exit(1);
        },
        _ => ()
    };
}

pub fn cli_main_raw() -> MainResult<()> {
    let matches = App::new("mates")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Markus Unterwaditzer")
        .about("A simple commandline addressbook")
        .setting(AppSettings::SubcommandRequired)
        .subcommand(SubCommand::with_name("index")
                    .about("Rewrite/create the index"))
        .subcommand(SubCommand::with_name("mutt-query")
                    .about("Search for contact, output is usable for mutt's query_command.")
                    .arg(Arg::with_name("query").index(1)))
        .subcommand(SubCommand::with_name("file-query")
                    .about("Search for contact, return just the filename.")
                    .arg(Arg::with_name("query").index(1)))
        .subcommand(SubCommand::with_name("email-query")
                    .about("Search for contact, return \"name <email>\".")
                    .arg(Arg::with_name("query").index(1)))
        .subcommand(SubCommand::with_name("add")
                    .about("Take mail from stdin, add sender to contacts. Print filename."))
        .subcommand(SubCommand::with_name("edit")
                    .about(
                        "Open contact (given by filepath or search-string) in $MATES_EDITOR. If
                        the file is cleared, the contact is removed. As a further convenience it
                        also clears stdin, which is necessary for editors and most interactive
                        programs to not act weird when piped to."
                    )
                    .arg(Arg::with_name("file-or-query").index(1)))
        .get_matches();

    let command = matches.subcommand_name().unwrap();

    let config = match Configuration::new() {
        Ok(x) => x,
        Err(e) => {
            return Err(MainError::new(format!("Error while reading configuration: {}", e)).into());
        }
    };

    let submatches = matches.subcommand_matches(command).expect("Internal error.");

    match command {
        "index" => {
            println!("Rebuilding index file \"{}\"...", config.index_path.display());
            try!(build_index(&config.index_path, &config.vdir_path));
        },
        "mutt-query" => {
            let query = submatches.value_of("query").unwrap_or("");
            try!(mutt_query(&config, &query[..]));
        },
        "file-query" => {
            let query = submatches.value_of("query").unwrap_or("");
            try!(file_query(&config, &query[..]));
        },
        "email-query" => {
            let query = submatches.value_of("query").unwrap_or("");
            try!(email_query(&config, &query[..]));
        },
        "add" => {
            let stdin = io::stdin();
            let mut email = String::new();
            try!(stdin.lock().read_to_string(&mut email));
            let contact = try!(utils::add_contact_from_email(
                &config.vdir_path,
                &email[..]
            ));
            println!("{}", contact.path.display());

            let mut index_fp = try!(fs::OpenOptions::new()
                                    .append(true)
                                    .write(true)
                                    .open(&config.index_path));

            let index_entry = try!(utils::index_item_from_contact(&contact));
            try!(index_fp.write_all(index_entry.as_bytes()));
        },
        "edit" => {
            let query = submatches.value_of("file-or-query").unwrap_or("");
            try!(edit_contact(&config, &query[..]));
        },
        _ => {
            return Err(MainError::new(format!("Invalid command: {}", command)).into());
        }
    };
    Ok(())
}

fn edit_contact(config: &Configuration, query: &str) -> MainResult<()> {
    let results = if get_pwd().join(query).is_file() {
        vec![path::PathBuf::from(query)]
    } else {
        try!(utils::file_query(config, query)).into_iter().collect()
    };

    if results.len() < 1 {
        return Err(MainError::new("No such contact.").into());
    } else if results.len() > 1 {
        return Err(MainError::new("Ambiguous query.").into());
    }

    let fpath = &results[0];
    let mut process = try!(process::Command::new("sh")
        .arg("-c")
        // clear stdin, http://unix.stackexchange.com/a/77593
        .arg("$0 \"$1\" < $2")
        .arg(&config.editor_cmd[..])
        .arg(fpath.as_os_str())
        .arg("/dev/tty")
        .stdin(process::Stdio::inherit())
        .stdout(process::Stdio::inherit())
        .stderr(process::Stdio::inherit())
        .spawn());

    try!(utils::handle_process(&mut process));

    let fcontent = {
        let mut fcontent = String::new();
        let mut file = try!(fs::File::open(fpath));
        try!(file.read_to_string(&mut fcontent));
        fcontent
    };

    if (&fcontent[..]).trim().len() == 0 {
        try!(fs::remove_file(fpath));
        return Err(MainError::new("Contact emptied, file removed.").into());
    };

    Ok(())
}

fn mutt_query<'a>(config: &Configuration, query: &str) -> MainResult<()> {
    println!("");  // For some reason mutt requires an empty line
    for item in try!(utils::index_query(config, query)) {
        if item.email.len() > 0 && item.name.len() > 0 {
            println!("{}\t{}", item.email, item.name);
        };
    };
    Ok(())
}

fn file_query<'a>(config: &Configuration, query: &str) -> MainResult<()> {
    for path in try!(utils::file_query(config, query)).iter() {
        println!("{}", path.display());
    };
    Ok(())
}

fn email_query<'a>(config: &Configuration, query: &str) -> MainResult<()> {
    for item in try!(utils::index_query(config, query)) {
        if item.name.len() > 0 && item.email.len() > 0 {
            println!("{} <{}>", item.name, item.email);
        };
    };
    Ok(())
}

pub struct Configuration {
    pub index_path: path::PathBuf,
    pub vdir_path: path::PathBuf,
    pub editor_cmd: String,
    pub grep_cmd: String
}

impl Configuration {
    pub fn new() -> Result<Configuration, String> {
        Ok(Configuration {
            index_path: match get_envvar("MATES_INDEX") {
                Some(x) => path::PathBuf::from(&x),
                None => match get_envvar("HOME") {
                    Some(home) => get_pwd().join(&home).join(".mates_index"),
                    None => return Err("Unable to determine user's home directory.".to_owned())
                }
            },
            vdir_path: match get_envvar("MATES_DIR") {
                Some(x) => path::PathBuf::from(&x),
                None => return Err("MATES_DIR must be set to your vdir path (directory of vcf-files).".to_owned())
            },
            editor_cmd: match get_envvar("MATES_EDITOR") {
                Some(x) => x,
                None => match get_envvar("EDITOR") {
                    Some(x) => x,
                    None => return Err("MATES_EDITOR or EDITOR must be set.".to_owned())
                }
            },
            grep_cmd: match get_envvar("MATES_GREP") {
                Some(x) => x,
                None => "grep -i".to_owned()
            }
        })
    }
}

#[derive(PartialEq, Eq, Debug)]
pub struct MainError {
    desc: String,
}

pub type MainResult<T> = Result<T, Box<Error>>;

impl Error for MainError {
    fn description(&self) -> &str {
        &self.desc[..]
    }

    fn cause(&self) -> Option<&Error> {
        None
    }
}

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

impl MainError {
    pub fn new<T: Into<String>>(desc: T) -> Self {
        MainError {
            desc: desc.into(),
        }
    }
}
<commit_msg>Fix formatting<commit_after>use std::borrow::ToOwned;
use std::env;
use std::error::Error;
use std::fmt;
use std::fs;
use std::io::{Read,Write};
use std::io;
use std::path;
use std::process;

use clap::{Arg,App,SubCommand,AppSettings};
use atomicwrites::{AtomicFile,AllowOverwrite};

use utils;
use utils::CustomPathExt;

#[inline]
fn get_pwd() -> path::PathBuf {
    env::current_dir().ok().expect("Failed to get CWD")
}

#[inline]
fn get_envvar(key: &str) -> Option<String> {
    match env::var(key) {
        Ok(x) => Some(x),
        Err(env::VarError::NotPresent) => None,
        Err(env::VarError::NotUnicode(_)) => panic!(format!("{} is not unicode.", key)),
    }
}

fn build_index(outfile: &path::Path, dir: &path::Path) -> MainResult<()> {
    if !dir.is_dir() {
        return Err(MainError::new("MATES_DIR must be a directory.").into());
    };

    let af = AtomicFile::new(&outfile, AllowOverwrite);
    let mut errors = false;

    try!(af.write(|outf| {
        for entry in try!(fs::read_dir(dir)) {
            let entry = match entry {
                Ok(x) => x,
                Err(e) => {
                    println!("Error while listing directory: {}", e);
                    errors = true;
                    continue;
                }
            };

            let pathbuf = entry.path();

            if pathbuf.str_extension().unwrap_or("") != "vcf" || !pathbuf.is_file() {
                continue;
            };

            let contact = match utils::Contact::from_file(&pathbuf) {
                Ok(x) => x,
                Err(e) => {
                    println!("Error while reading {}: {}", pathbuf.display(), e);
                    errors = true;
                    continue
                }
            };

            match utils::index_item_from_contact(&contact) {
                Ok(index_string) => {
                    try!(outf.write_all(index_string.as_bytes()));
                },
                Err(e) => {
                    println!("Error while indexing {}: {}", pathbuf.display(), e);
                    errors = true;
                    continue
                }
            };
        };
        Ok(())
    }));

    if errors {
        Err(MainError::new("Several errors happened while generating the index.").into())
    } else {
        Ok(())
    }
}

pub fn cli_main() {
    match cli_main_raw() {
        Err(e) => {
            writeln!(&mut io::stderr(), "{}", e).unwrap();
            process::exit(1);
        },
        _ => ()
    };
}

pub fn cli_main_raw() -> MainResult<()> {
    let matches = App::new("mates")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Markus Unterwaditzer")
        .about("A simple commandline addressbook")
        .setting(AppSettings::SubcommandRequired)
        .subcommand(SubCommand::with_name("index")
                    .about("Rewrite/create the index"))
        .subcommand(SubCommand::with_name("mutt-query")
                    .about("Search for contact, output is usable for mutt's query_command.")
                    .arg(Arg::with_name("query").index(1)))
        .subcommand(SubCommand::with_name("file-query")
                    .about("Search for contact, return just the filename.")
                    .arg(Arg::with_name("query").index(1)))
        .subcommand(SubCommand::with_name("email-query")
                    .about("Search for contact, return \"name <email>\".")
                    .arg(Arg::with_name("query").index(1)))
        .subcommand(SubCommand::with_name("add")
                    .about("Take mail from stdin, add sender to contacts. Print filename."))
        .subcommand(SubCommand::with_name("edit")
                    .about(
                        "Open contact (given by filepath or search-string) in $MATES_EDITOR. If \
                        the file is cleared, the contact is removed. As a further convenience it \
                        also clears stdin, which is necessary for editors and most interactive \
                        programs to not act weird when piped to."
                    )
                    .arg(Arg::with_name("file-or-query").index(1)))
        .get_matches();

    let command = matches.subcommand_name().unwrap();

    let config = match Configuration::new() {
        Ok(x) => x,
        Err(e) => {
            return Err(MainError::new(format!("Error while reading configuration: {}", e)).into());
        }
    };

    let submatches = matches.subcommand_matches(command).expect("Internal error.");

    match command {
        "index" => {
            println!("Rebuilding index file \"{}\"...", config.index_path.display());
            try!(build_index(&config.index_path, &config.vdir_path));
        },
        "mutt-query" => {
            let query = submatches.value_of("query").unwrap_or("");
            try!(mutt_query(&config, &query[..]));
        },
        "file-query" => {
            let query = submatches.value_of("query").unwrap_or("");
            try!(file_query(&config, &query[..]));
        },
        "email-query" => {
            let query = submatches.value_of("query").unwrap_or("");
            try!(email_query(&config, &query[..]));
        },
        "add" => {
            let stdin = io::stdin();
            let mut email = String::new();
            try!(stdin.lock().read_to_string(&mut email));
            let contact = try!(utils::add_contact_from_email(
                &config.vdir_path,
                &email[..]
            ));
            println!("{}", contact.path.display());

            let mut index_fp = try!(fs::OpenOptions::new()
                                    .append(true)
                                    .write(true)
                                    .open(&config.index_path));

            let index_entry = try!(utils::index_item_from_contact(&contact));
            try!(index_fp.write_all(index_entry.as_bytes()));
        },
        "edit" => {
            let query = submatches.value_of("file-or-query").unwrap_or("");
            try!(edit_contact(&config, &query[..]));
        },
        _ => {
            return Err(MainError::new(format!("Invalid command: {}", command)).into());
        }
    };
    Ok(())
}

fn edit_contact(config: &Configuration, query: &str) -> MainResult<()> {
    let results = if get_pwd().join(query).is_file() {
        vec![path::PathBuf::from(query)]
    } else {
        try!(utils::file_query(config, query)).into_iter().collect()
    };

    if results.len() < 1 {
        return Err(MainError::new("No such contact.").into());
    } else if results.len() > 1 {
        return Err(MainError::new("Ambiguous query.").into());
    }

    let fpath = &results[0];
    let mut process = try!(process::Command::new("sh")
        .arg("-c")
        // clear stdin, http://unix.stackexchange.com/a/77593
        .arg("$0 \"$1\" < $2")
        .arg(&config.editor_cmd[..])
        .arg(fpath.as_os_str())
        .arg("/dev/tty")
        .stdin(process::Stdio::inherit())
        .stdout(process::Stdio::inherit())
        .stderr(process::Stdio::inherit())
        .spawn());

    try!(utils::handle_process(&mut process));

    let fcontent = {
        let mut fcontent = String::new();
        let mut file = try!(fs::File::open(fpath));
        try!(file.read_to_string(&mut fcontent));
        fcontent
    };

    if (&fcontent[..]).trim().len() == 0 {
        try!(fs::remove_file(fpath));
        return Err(MainError::new("Contact emptied, file removed.").into());
    };

    Ok(())
}

fn mutt_query<'a>(config: &Configuration, query: &str) -> MainResult<()> {
    println!("");  // For some reason mutt requires an empty line
    for item in try!(utils::index_query(config, query)) {
        if item.email.len() > 0 && item.name.len() > 0 {
            println!("{}\t{}", item.email, item.name);
        };
    };
    Ok(())
}

fn file_query<'a>(config: &Configuration, query: &str) -> MainResult<()> {
    for path in try!(utils::file_query(config, query)).iter() {
        println!("{}", path.display());
    };
    Ok(())
}

fn email_query<'a>(config: &Configuration, query: &str) -> MainResult<()> {
    for item in try!(utils::index_query(config, query)) {
        if item.name.len() > 0 && item.email.len() > 0 {
            println!("{} <{}>", item.name, item.email);
        };
    };
    Ok(())
}

pub struct Configuration {
    pub index_path: path::PathBuf,
    pub vdir_path: path::PathBuf,
    pub editor_cmd: String,
    pub grep_cmd: String
}

impl Configuration {
    pub fn new() -> Result<Configuration, String> {
        Ok(Configuration {
            index_path: match get_envvar("MATES_INDEX") {
                Some(x) => path::PathBuf::from(&x),
                None => match get_envvar("HOME") {
                    Some(home) => get_pwd().join(&home).join(".mates_index"),
                    None => return Err("Unable to determine user's home directory.".to_owned())
                }
            },
            vdir_path: match get_envvar("MATES_DIR") {
                Some(x) => path::PathBuf::from(&x),
                None => return Err("MATES_DIR must be set to your vdir path (directory of vcf-files).".to_owned())
            },
            editor_cmd: match get_envvar("MATES_EDITOR") {
                Some(x) => x,
                None => match get_envvar("EDITOR") {
                    Some(x) => x,
                    None => return Err("MATES_EDITOR or EDITOR must be set.".to_owned())
                }
            },
            grep_cmd: match get_envvar("MATES_GREP") {
                Some(x) => x,
                None => "grep -i".to_owned()
            }
        })
    }
}

#[derive(PartialEq, Eq, Debug)]
pub struct MainError {
    desc: String,
}

pub type MainResult<T> = Result<T, Box<Error>>;

impl Error for MainError {
    fn description(&self) -> &str {
        &self.desc[..]
    }

    fn cause(&self) -> Option<&Error> {
        None
    }
}

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

impl MainError {
    pub fn new<T: Into<String>>(desc: T) -> Self {
        MainError {
            desc: desc.into(),
        }
    }
}
<|endoftext|>


<commit_before>extern crate rocket;

use rocket::request::{FromForm, FormItems};

pub struct WhereStructResult {
    sql: String, values: Vec<String>,
}

// My Results
type QueryOpResult = Result<&'static str, String>;
type WhereResult = Result<WhereStructResult, String>;

// TODO: Move there later
#[derive(Clone)]
pub struct Params {
    pub k: String, pub v: String,
}
// implement QueryString
pub struct QueryString {
    pub query: Vec<Params>
}

impl<'f> FromForm<'f> for QueryString {
    type Error = ();

    fn from_form(items: &mut FormItems<'f>, strict: bool) -> Result<QueryString, ()> {
        let mut qs: QueryString = QueryString{ query: Vec::new() };
        for (k, v) in items {
            qs.query.push(Params{k: k.to_string(), v: v.url_decode().unwrap()});
        }
        Ok(qs)
    }
}

pub fn query_operator(qop: &str) -> QueryOpResult {
    match &qop[..] {
        "$eq" => Ok("="),
        "$ne" => Ok("!="),
        "$gt" => Ok(">"),
        "$lt" => Ok("<"),
        "$gte" => Ok(">="),
        "$lte" => Ok("<="),
        "$in" => Ok("IN"),
        "$nin" => Ok("NOT IN"),
        "$notnull" => Ok("IS NOT NULL"),
        "$null" => Ok("IS NULL"),
         _ => Err(format!("operator {} not found", qop)),
    }
}

fn collect_params(v: &Vec<Params>, collect_values: bool) -> Vec<String> {
    let mut pid: i32 = 0;
    v.iter()
        .map(|param| {
            pid = pid + 1;
            let p = param.clone();
            let result = match p.v.contains(".") {
                true => {
                    let value_splited = p.v.split(".").collect::<Vec<&str>>();
                    let operator = match query_operator(value_splited[0]) {
                        Ok(op) => op,
                        _ => "operator not found"
                    };
                    (operator, value_splited[1])
                },
                _ => ("=", p.v.as_str())
            };
            if collect_values {
                format!("{}", result.1)
            } else {
                format!("{}{}${}", p.k, result.0, pid)
            }
        })
        .collect()
}

// Get QueryString and return chunck WHERE statement with values
pub fn ext_where(qs: QueryString) -> WhereResult {
    let where_syntax: String;
    let ref iter_params = qs.query;
    let where_values = collect_params(iter_params, true);
    let where_fields = collect_params(iter_params, false);

    if where_fields.len() > 1 {
        where_syntax = format!("WHERE {}", where_fields.join(" AND "));
    } else {
        where_syntax = format!("WHERE {}", where_fields.join(" "));
    }

    Ok(WhereStructResult {
        sql: where_syntax,
        values: where_values
    })
}

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

    #[test]
    fn get_opetators() {
        use std::collections::HashMap;

        let mut test_cases = HashMap::new();
        test_cases.insert("$eq", "=");
        test_cases.insert("$ne", "!=");
        test_cases.insert("$gt", ">");
        test_cases.insert("$gte", ">=");
        test_cases.insert("$lt", "<");
        test_cases.insert("$lte", "<=");
        test_cases.insert("$in", "IN");
        test_cases.insert("$nin", "NOT IN");
        test_cases.insert("$notnull", "IS NOT NULL");
        test_cases.insert("$null", "IS NULL");

        for (op, expected) in &test_cases {
            match query_operator(&op.to_string()) {
                Ok(x) => assert_eq!(expected.to_string(), x),
                Err(e) => assert_eq!("", e)
            }
        }
    }

    #[test]
    fn get_opetators_error() {
        let invalid = "!tn".to_string();
        match query_operator(&invalid) {
            Ok(_) => println!("no result"),
            Err(e) => assert_eq!(format!("operator {} not found", invalid), e)
        }
    }

    #[test]
    fn where_by_request_one_querystring() {
        let mut hm = Vec::new();
        hm.push(Params{k: "user_id".to_string(), v: "5".to_string()});

        let qs = QueryString{ query: hm };
        let result = ext_where(qs)
            .unwrap();
        assert_eq!(result.sql, "WHERE user_id=$1".to_string());
        assert_eq!(result.values, vec!["5".to_string()])
    }

    #[test]
    fn where_by_request_n_querystring() {
        let mut hm = Vec::new();
        hm.push(Params{k: "user_id".to_string(), v: "5".to_string()});
        hm.push(Params{k: "name".to_string(), v: "goga".to_string()});

        let qs = QueryString{ query: hm };
        let result = ext_where(qs)
            .unwrap();
        assert_eq!(result.sql, "WHERE user_id=$1 AND name=$2".to_string());
        assert_eq!(result.values, vec!["5".to_string(), "goga".to_string()])
    }

    #[test]
    fn where_by_request_with_query_operators(){
        let mut hm = Vec::new();
        hm.push(Params{k: "user_id".to_string(), v: "$ne.5".to_string()});
        hm.push(Params{k: "name".to_string(), v: "goga".to_string()});

        let qs = QueryString{ query: hm };
        let result = ext_where(qs)
            .unwrap();
        assert_eq!(result.sql, "WHERE user_id!=$1 AND name=$2".to_string());
        assert_eq!(result.values, vec!["5".to_string(), "goga".to_string()])
    }

    #[test]
    fn where_by_request_with_query_invalid_operators(){
        let mut hm = Vec::new();
        hm.push(Params{k: "user_id".to_string(), v: "$nt.5".to_string()});
        hm.push(Params{k: "name".to_string(), v: "goga".to_string()});

        let qs = QueryString{ query: hm };
        match ext_where(qs) {
            Err(e) => assert_eq!(e, "operator $nt not found"),
            Ok(_) => println!("no result")
        }
    }
}
<commit_msg>Remove clone from reference<commit_after>extern crate rocket;

use rocket::request::{FromForm, FormItems};

pub struct WhereStructResult {
    sql: String, values: Vec<String>,
}

// My Results
type QueryOpResult = Result<&'static str, String>;
type WhereResult = Result<WhereStructResult, String>;

// TODO: Move there later
pub struct Params {
    pub k: String, pub v: String,
}
// implement QueryString
pub struct QueryString {
    pub query: Vec<Params>
}

impl<'f> FromForm<'f> for QueryString {
    type Error = ();

    fn from_form(items: &mut FormItems<'f>, strict: bool) -> Result<QueryString, ()> {
        let mut qs: QueryString = QueryString{ query: Vec::new() };
        for (k, v) in items {
            qs.query.push(Params{k: k.to_string(), v: v.url_decode().unwrap()});
        }
        Ok(qs)
    }
}

pub fn query_operator(qop: &str) -> QueryOpResult {
    match &qop[..] {
        "$eq" => Ok("="),
        "$ne" => Ok("!="),
        "$gt" => Ok(">"),
        "$lt" => Ok("<"),
        "$gte" => Ok(">="),
        "$lte" => Ok("<="),
        "$in" => Ok("IN"),
        "$nin" => Ok("NOT IN"),
        "$notnull" => Ok("IS NOT NULL"),
        "$null" => Ok("IS NULL"),
         _ => Err(format!("operator {} not found", qop)),
    }
}

fn collect_params(v: &Vec<Params>, collect_values: bool) -> Vec<String> {
    let mut pid: i32 = 0;
    v.iter()
        .map(|param| {
            pid = pid + 1;
            let result = match param.v.contains(".") {
                true => {
                    let value_splited = param.v.split(".").collect::<Vec<&str>>();
                    let operator = match query_operator(value_splited[0]) {
                        Ok(op) => op,
                        _ => "operator not found"
                    };
                    (operator, value_splited[1])
                },
                _ => ("=", param.v.as_str())
            };
            if collect_values {
                format!("{}", result.1)
            } else {
                format!("{}{}${}", param.k, result.0, pid)
            }
        })
        .collect()
}

// Get QueryString and return chunck WHERE statement with values
pub fn ext_where(qs: QueryString) -> WhereResult {
    let where_syntax: String;
    let ref iter_params = qs.query;
    let where_values = collect_params(iter_params, true);
    let where_fields = collect_params(iter_params, false);

    if where_fields.len() > 1 {
        where_syntax = format!("WHERE {}", where_fields.join(" AND "));
    } else {
        where_syntax = format!("WHERE {}", where_fields.join(" "));
    }

    Ok(WhereStructResult {
        sql: where_syntax,
        values: where_values
    })
}

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

    #[test]
    fn get_opetators() {
        use std::collections::HashMap;

        let mut test_cases = HashMap::new();
        test_cases.insert("$eq", "=");
        test_cases.insert("$ne", "!=");
        test_cases.insert("$gt", ">");
        test_cases.insert("$gte", ">=");
        test_cases.insert("$lt", "<");
        test_cases.insert("$lte", "<=");
        test_cases.insert("$in", "IN");
        test_cases.insert("$nin", "NOT IN");
        test_cases.insert("$notnull", "IS NOT NULL");
        test_cases.insert("$null", "IS NULL");

        for (op, expected) in &test_cases {
            match query_operator(&op.to_string()) {
                Ok(x) => assert_eq!(expected.to_string(), x),
                Err(e) => assert_eq!("", e)
            }
        }
    }

    #[test]
    fn get_opetators_error() {
        let invalid = "!tn".to_string();
        match query_operator(&invalid) {
            Ok(_) => println!("no result"),
            Err(e) => assert_eq!(format!("operator {} not found", invalid), e)
        }
    }

    #[test]
    fn where_by_request_one_querystring() {
        let mut hm = Vec::new();
        hm.push(Params{k: "user_id".to_string(), v: "5".to_string()});

        let qs = QueryString{ query: hm };
        let result = ext_where(qs)
            .unwrap();
        assert_eq!(result.sql, "WHERE user_id=$1".to_string());
        assert_eq!(result.values, vec!["5".to_string()])
    }

    #[test]
    fn where_by_request_n_querystring() {
        let mut hm = Vec::new();
        hm.push(Params{k: "user_id".to_string(), v: "5".to_string()});
        hm.push(Params{k: "name".to_string(), v: "goga".to_string()});

        let qs = QueryString{ query: hm };
        let result = ext_where(qs)
            .unwrap();
        assert_eq!(result.sql, "WHERE user_id=$1 AND name=$2".to_string());
        assert_eq!(result.values, vec!["5".to_string(), "goga".to_string()])
    }

    #[test]
    fn where_by_request_with_query_operators(){
        let mut hm = Vec::new();
        hm.push(Params{k: "user_id".to_string(), v: "$ne.5".to_string()});
        hm.push(Params{k: "name".to_string(), v: "goga".to_string()});

        let qs = QueryString{ query: hm };
        let result = ext_where(qs)
            .unwrap();
        assert_eq!(result.sql, "WHERE user_id!=$1 AND name=$2".to_string());
        assert_eq!(result.values, vec!["5".to_string(), "goga".to_string()])
    }

    #[test]
    fn where_by_request_with_query_invalid_operators(){
        let mut hm = Vec::new();
        hm.push(Params{k: "user_id".to_string(), v: "$nt.5".to_string()});
        hm.push(Params{k: "name".to_string(), v: "goga".to_string()});

        let qs = QueryString{ query: hm };
        match ext_where(qs) {
            Err(e) => assert_eq!(e, "operator $nt not found"),
            Ok(_) => println!("no result")
        }
    }
}
<|endoftext|>


<commit_before>//! This module defines a series of convenience modifiers for changing
//! Responses.
//!
//! Modifiers can be used to edit `Response`s through the owning method `set`
//! or the mutating `set_mut`, both of which are defined in the `Set` trait.
//!
//! For Iron, the `Modifier` interface offers extensible and ergonomic response
//! creation while avoiding the introduction of many highly specific `Response`
//! constructors.
//!
//! The simplest case of a modifier is probably the one used to change the
//! return status code:
//!
//! ```
//! # use iron::prelude::*;
//! # use iron::status;
//! let r = Response::with(status::NotFound);
//! assert_eq!(r.status.unwrap().to_u16(), 404);
//! ```
//!
//! You can also pass in a tuple of modifiers, they will all be applied. Here's
//! an example of a modifier 2-tuple that will change the status code and the
//! body message:
//!
//! ```
//! # use iron::prelude::*;
//! # use iron::status;
//! Response::with((status::ImATeapot, "I am a tea pot!"));
//! ```
//!
//! There is also a `Redirect` modifier:
//!
//! ```
//! # use iron::prelude::*;
//! # use iron::status;
//! # use iron::modifiers;
//! # use iron::Url;
//! let url = Url::parse("http://doc.rust-lang.org").unwrap();
//! Response::with((status::Found, modifiers::Redirect(url)));
//! ```
//!
//! The modifiers are applied depending on their type. Currently the easiest
//! way to see how different types are used as modifiers, take a look at [the
//! source code](https://github.com/iron/iron/blob/master/src/modifiers.rs).
//!
//! For more information about the modifier system, see
//! [rust-modifier](https://github.com/reem/rust-modifier).

use std::fs::File;
use std::io;
use std::path::{Path, PathBuf};

use modifier::Modifier;

use hyper::mime::Mime;

use {status, headers, Request, Response, Set, Url};

use mime_types;
use response::{WriteBody, BodyReader};

lazy_static! {
    static ref MIME_TYPES: mime_types::Types = mime_types::Types::new().unwrap();
}

impl Modifier<Response> for Mime {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentType(self))
    }
}

impl Modifier<Response> for Box<WriteBody> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.body = Some(self);
    }
}

impl <R: io::Read + Send + 'static> Modifier<Response> for BodyReader<R> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.body = Some(Box::new(self));
    }
}

impl Modifier<Response> for String {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.into_bytes().modify(res);
    }
}

impl Modifier<Response> for Vec<u8> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentLength(self.len() as u64));
        res.body = Some(Box::new(self));
    }
}

impl<'a> Modifier<Response> for &'a str {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_owned().modify(res);
    }
}

impl<'a> Modifier<Response> for &'a [u8] {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_vec().modify(res);
    }
}

impl Modifier<Response> for File {
    fn modify(self, res: &mut Response) {
        // Set the content type based on the file extension if a path is available.
        if let Ok(metadata) = self.metadata() {
            res.headers.set(headers::ContentLength(metadata.len()));
        }

        res.body = Some(Box::new(self));
    }
}

impl<'a> Modifier<Response> for &'a Path {
    /// Set the body to the contents of the File at this path.
    ///
    /// ## Panics
    ///
    /// Panics if there is no file at the passed-in Path.
    fn modify(self, res: &mut Response) {
        File::open(self)
            .expect(&format!("No such file: {}", self.display()))
            .modify(res);

        let mime_str = MIME_TYPES.mime_for_path(self);
        let _ = mime_str.parse().map(|mime: Mime| res.set_mut(mime));
    }
}

impl Modifier<Response> for PathBuf {
    /// Set the body to the contents of the File at this path.
    ///
    /// ## Panics
    ///
    /// Panics if there is no file at the passed-in Path.
    fn modify(self, res: &mut Response) {
        File::open(&self)
            .expect(&format!("No such file: {}", self.display()))
            .modify(res);
    }
}

impl Modifier<Response> for status::Status {
    fn modify(self, res: &mut Response) {
        res.status = Some(self);
    }
}

/// A modifier for changing headers on requests and responses.
pub struct Header<H: headers::Header + headers::HeaderFormat>(pub H);

impl<H> Modifier<Response> for Header<H>
where H: headers::Header + headers::HeaderFormat {
    fn modify(self, res: &mut Response) {
        res.headers.set(self.0);
    }
}

impl<'a, 'b, H> Modifier<Request<'a, 'b>> for Header<H>
where H: headers::Header + headers::HeaderFormat {
    fn modify(self, res: &mut Request) {
        res.headers.set(self.0);
    }
}

/// A modifier for creating redirect responses.
pub struct Redirect(pub Url);

impl Modifier<Response> for Redirect {
    fn modify(self, res: &mut Response) {
        let Redirect(url) = self;
        res.headers.set(headers::Location(url.to_string()));
    }
}

/// A modifier for creating redirect responses.
pub struct RedirectRaw(String);

impl Modifier<Response> for RedirectRaw {
    fn modify(self, res: &mut Response) {
        let RedirectRaw(path) = self;
        res.headers.set(headers::Location(path));
    }
}
<commit_msg>make RedirectRaw's field public to allow construction<commit_after>//! This module defines a series of convenience modifiers for changing
//! Responses.
//!
//! Modifiers can be used to edit `Response`s through the owning method `set`
//! or the mutating `set_mut`, both of which are defined in the `Set` trait.
//!
//! For Iron, the `Modifier` interface offers extensible and ergonomic response
//! creation while avoiding the introduction of many highly specific `Response`
//! constructors.
//!
//! The simplest case of a modifier is probably the one used to change the
//! return status code:
//!
//! ```
//! # use iron::prelude::*;
//! # use iron::status;
//! let r = Response::with(status::NotFound);
//! assert_eq!(r.status.unwrap().to_u16(), 404);
//! ```
//!
//! You can also pass in a tuple of modifiers, they will all be applied. Here's
//! an example of a modifier 2-tuple that will change the status code and the
//! body message:
//!
//! ```
//! # use iron::prelude::*;
//! # use iron::status;
//! Response::with((status::ImATeapot, "I am a tea pot!"));
//! ```
//!
//! There is also a `Redirect` modifier:
//!
//! ```
//! # use iron::prelude::*;
//! # use iron::status;
//! # use iron::modifiers;
//! # use iron::Url;
//! let url = Url::parse("http://doc.rust-lang.org").unwrap();
//! Response::with((status::Found, modifiers::Redirect(url)));
//! ```
//!
//! The modifiers are applied depending on their type. Currently the easiest
//! way to see how different types are used as modifiers, take a look at [the
//! source code](https://github.com/iron/iron/blob/master/src/modifiers.rs).
//!
//! For more information about the modifier system, see
//! [rust-modifier](https://github.com/reem/rust-modifier).

use std::fs::File;
use std::io;
use std::path::{Path, PathBuf};

use modifier::Modifier;

use hyper::mime::Mime;

use {status, headers, Request, Response, Set, Url};

use mime_types;
use response::{WriteBody, BodyReader};

lazy_static! {
    static ref MIME_TYPES: mime_types::Types = mime_types::Types::new().unwrap();
}

impl Modifier<Response> for Mime {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentType(self))
    }
}

impl Modifier<Response> for Box<WriteBody> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.body = Some(self);
    }
}

impl <R: io::Read + Send + 'static> Modifier<Response> for BodyReader<R> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.body = Some(Box::new(self));
    }
}

impl Modifier<Response> for String {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.into_bytes().modify(res);
    }
}

impl Modifier<Response> for Vec<u8> {
    #[inline]
    fn modify(self, res: &mut Response) {
        res.headers.set(headers::ContentLength(self.len() as u64));
        res.body = Some(Box::new(self));
    }
}

impl<'a> Modifier<Response> for &'a str {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_owned().modify(res);
    }
}

impl<'a> Modifier<Response> for &'a [u8] {
    #[inline]
    fn modify(self, res: &mut Response) {
        self.to_vec().modify(res);
    }
}

impl Modifier<Response> for File {
    fn modify(self, res: &mut Response) {
        // Set the content type based on the file extension if a path is available.
        if let Ok(metadata) = self.metadata() {
            res.headers.set(headers::ContentLength(metadata.len()));
        }

        res.body = Some(Box::new(self));
    }
}

impl<'a> Modifier<Response> for &'a Path {
    /// Set the body to the contents of the File at this path.
    ///
    /// ## Panics
    ///
    /// Panics if there is no file at the passed-in Path.
    fn modify(self, res: &mut Response) {
        File::open(self)
            .expect(&format!("No such file: {}", self.display()))
            .modify(res);

        let mime_str = MIME_TYPES.mime_for_path(self);
        let _ = mime_str.parse().map(|mime: Mime| res.set_mut(mime));
    }
}

impl Modifier<Response> for PathBuf {
    /// Set the body to the contents of the File at this path.
    ///
    /// ## Panics
    ///
    /// Panics if there is no file at the passed-in Path.
    fn modify(self, res: &mut Response) {
        File::open(&self)
            .expect(&format!("No such file: {}", self.display()))
            .modify(res);
    }
}

impl Modifier<Response> for status::Status {
    fn modify(self, res: &mut Response) {
        res.status = Some(self);
    }
}

/// A modifier for changing headers on requests and responses.
pub struct Header<H: headers::Header + headers::HeaderFormat>(pub H);

impl<H> Modifier<Response> for Header<H>
where H: headers::Header + headers::HeaderFormat {
    fn modify(self, res: &mut Response) {
        res.headers.set(self.0);
    }
}

impl<'a, 'b, H> Modifier<Request<'a, 'b>> for Header<H>
where H: headers::Header + headers::HeaderFormat {
    fn modify(self, res: &mut Request) {
        res.headers.set(self.0);
    }
}

/// A modifier for creating redirect responses.
pub struct Redirect(pub Url);

impl Modifier<Response> for Redirect {
    fn modify(self, res: &mut Response) {
        let Redirect(url) = self;
        res.headers.set(headers::Location(url.to_string()));
    }
}

/// A modifier for creating redirect responses.
pub struct RedirectRaw(pub String);

impl Modifier<Response> for RedirectRaw {
    fn modify(self, res: &mut Response) {
        let RedirectRaw(path) = self;
        res.headers.set(headers::Location(path));
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::net::SocketAddr;
use std::pin::Pin;
use std::str;
use std::task::{Context, Poll};

use futures::prelude::*;
use futures::Future;

use hyper::Client;

use tokio::net::UdpSocket;
use tokio::time::timeout;

use hyper::body::Bytes;

use crate::aio::Gateway;
use crate::common::{messages, parsing, SearchOptions};
use crate::errors::SearchError;

const MAX_RESPONSE_SIZE: usize = 1500;

/// Search for a gateway with the provided options
pub async fn search_gateway(options: SearchOptions) -> Result<Gateway, SearchError> {
    // Create socket for future calls
    let socket = UdpSocket::bind(&options.bind_addr).await?;

    let search_future = SearchFuture::search(socket, options.broadcast_address).await?;
    // Create future and issue request
    match options.timeout {
        Some(t) => timeout(t, search_future).await?,
        None => search_future.await,
    }
}

pub struct SearchFuture {
    socket: UdpSocket,
    pending: HashMap<SocketAddr, RequestState>,
}

enum RequestState {
    Connecting(String, Pin<Box<dyn Future<Output = Result<Bytes, SearchError>> + Send>>),
    Done(String),
    Error,
}

impl SearchFuture {
    // Create a new search
    async fn search(mut socket: UdpSocket, addr: SocketAddr) -> Result<SearchFuture, SearchError> {
        debug!(
            "sending broadcast request to: {} on interface: {:?}",
            addr,
            socket.local_addr()
        );
        socket.send_to(messages::SEARCH_REQUEST.as_bytes(), &addr).await?;
        Ok(SearchFuture {
            socket,
            pending: HashMap::new(),
        })
    }

    // Handle a UDP response message
    fn handle_broadcast_resp(from: SocketAddr, data: &[u8]) -> Result<(SocketAddr, String), SearchError> {
        debug!("handling broadcast response from: {}", from);

        // Convert response to text
        let text = str::from_utf8(&data).map_err(|e| SearchError::from(e))?;

        // Parse socket address and path
        let (addr, path) = parsing::parse_search_result(text)?;

        Ok((SocketAddr::V4(addr), path))
    }

    // Issue a control URL request over HTTP using the provided
    fn request_control_url(
        addr: SocketAddr,
        path: String,
    ) -> Result<Pin<Box<dyn Future<Output = Result<Bytes, SearchError>> + Send>>, SearchError> {
        let client = Client::new();

        let uri = match format!("http://{}{}", addr, path).parse() {
            Ok(uri) => uri,
            Err(err) => return Err(SearchError::from(err)),
        };

        debug!("requesting control url from: {}", uri);

        Ok(Box::pin(
            client
                .get(uri)
                .and_then(|resp| hyper::body::to_bytes(resp.into_body()))
                .map_err(|e| SearchError::from(e)),
        ))
    }

    // Process a control response to extract the control URL
    fn handle_control_resp(addr: SocketAddr, resp: Bytes) -> Result<String, SearchError> {
        debug!("handling control response from: {}", addr);

        // Create a cursor over the response data
        let c = std::io::Cursor::new(&resp);

        // Parse control URL out of body
        let url = parsing::parse_control_url(c)?;

        Ok(url)
    }
}

impl Future for SearchFuture {
    type Output = Result<Gateway, SearchError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<Gateway, SearchError>> {
        // Poll for (and handle) incoming messages
        let mut buff = [0u8; MAX_RESPONSE_SIZE];
        let resp = self.socket.poll_recv_from(cx, &mut buff);
        if let Poll::Ready(Ok((n, from))) = resp {
            // Try handle response messages
            if let Ok((addr, path)) = Self::handle_broadcast_resp(from, &buff[0..n]) {
                if !self.pending.contains_key(&addr) {
                    debug!("received broadcast response from: {}", from);

                    // Issue control request
                    match Self::request_control_url(addr, path.clone()) {
                        // Store pending requests
                        Ok(f) => {
                            self.pending.insert(addr, RequestState::Connecting(path, f));
                        }
                        Err(e) => return Poll::Ready(Err(e)),
                    }
                } else {
                    debug!("received duplicate broadcast response from: {}, dropping", from);
                }
            }
        }
        if let Poll::Ready(Err(err)) = resp {
            return Poll::Ready(Err(err.into()));
        }

        // Poll on any outstanding control requests
        for (addr, state) in &mut self.pending {
            // Poll if we're in the connecting state
            let (root_url, resp) = {
                let (root_url, c) = match state {
                    RequestState::Connecting(root_url, c) => (root_url, c),
                    _ => continue,
                };

                match c.as_mut().poll(cx)? {
                    Poll::Ready(resp) => (root_url.clone(), resp),
                    _ => continue,
                }
            };

            // Handle any responses
            if let Ok(control_url) = Self::handle_control_resp(*addr, resp) {
                debug!("received control url from: {} (url: {})", addr, control_url);
                *state = RequestState::Done(url.clone());

                match addr {
                    SocketAddr::V4(a) => {
                        let g = Gateway {
                            addr: *a,
                            root_url,
                            control_url,
                        };
                        return Poll::Ready(Ok(g));
                    }
                    _ => warn!("unsupported IPv6 gateway response from addr: {}", addr),
                }
            } else {
                *state = RequestState::Error;
            }
        }

        Poll::Pending
    }
}
<commit_msg>Rename variable to be more explicit<commit_after>use std::collections::HashMap;
use std::net::SocketAddr;
use std::pin::Pin;
use std::str;
use std::task::{Context, Poll};

use futures::prelude::*;
use futures::Future;

use hyper::Client;

use tokio::net::UdpSocket;
use tokio::time::timeout;

use hyper::body::Bytes;

use crate::aio::Gateway;
use crate::common::{messages, parsing, SearchOptions};
use crate::errors::SearchError;

const MAX_RESPONSE_SIZE: usize = 1500;

/// Search for a gateway with the provided options
pub async fn search_gateway(options: SearchOptions) -> Result<Gateway, SearchError> {
    // Create socket for future calls
    let socket = UdpSocket::bind(&options.bind_addr).await?;

    let search_future = SearchFuture::search(socket, options.broadcast_address).await?;
    // Create future and issue request
    match options.timeout {
        Some(t) => timeout(t, search_future).await?,
        None => search_future.await,
    }
}

pub struct SearchFuture {
    socket: UdpSocket,
    pending: HashMap<SocketAddr, RequestState>,
}

enum RequestState {
    Connecting(String, Pin<Box<dyn Future<Output = Result<Bytes, SearchError>> + Send>>),
    Done(String),
    Error,
}

impl SearchFuture {
    // Create a new search
    async fn search(mut socket: UdpSocket, addr: SocketAddr) -> Result<SearchFuture, SearchError> {
        debug!(
            "sending broadcast request to: {} on interface: {:?}",
            addr,
            socket.local_addr()
        );
        socket.send_to(messages::SEARCH_REQUEST.as_bytes(), &addr).await?;
        Ok(SearchFuture {
            socket,
            pending: HashMap::new(),
        })
    }

    // Handle a UDP response message
    fn handle_broadcast_resp(from: SocketAddr, data: &[u8]) -> Result<(SocketAddr, String), SearchError> {
        debug!("handling broadcast response from: {}", from);

        // Convert response to text
        let text = str::from_utf8(&data).map_err(|e| SearchError::from(e))?;

        // Parse socket address and path
        let (addr, root_url) = parsing::parse_search_result(text)?;

        Ok((SocketAddr::V4(addr), root_url))
    }

    // Issue a control URL request over HTTP using the provided
    fn request_control_url(
        addr: SocketAddr,
        path: String,
    ) -> Result<Pin<Box<dyn Future<Output = Result<Bytes, SearchError>> + Send>>, SearchError> {
        let client = Client::new();

        let uri = match format!("http://{}{}", addr, path).parse() {
            Ok(uri) => uri,
            Err(err) => return Err(SearchError::from(err)),
        };

        debug!("requesting control url from: {}", uri);

        Ok(Box::pin(
            client
                .get(uri)
                .and_then(|resp| hyper::body::to_bytes(resp.into_body()))
                .map_err(|e| SearchError::from(e)),
        ))
    }

    // Process a control response to extract the control URL
    fn handle_control_resp(addr: SocketAddr, resp: Bytes) -> Result<String, SearchError> {
        debug!("handling control response from: {}", addr);

        // Create a cursor over the response data
        let c = std::io::Cursor::new(&resp);

        // Parse control URL out of body
        let url = parsing::parse_control_url(c)?;

        Ok(url)
    }
}

impl Future for SearchFuture {
    type Output = Result<Gateway, SearchError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<Gateway, SearchError>> {
        // Poll for (and handle) incoming messages
        let mut buff = [0u8; MAX_RESPONSE_SIZE];
        let resp = self.socket.poll_recv_from(cx, &mut buff);
        if let Poll::Ready(Ok((n, from))) = resp {
            // Try handle response messages
            if let Ok((addr, root_url)) = Self::handle_broadcast_resp(from, &buff[0..n]) {
                if !self.pending.contains_key(&addr) {
                    debug!("received broadcast response from: {}", from);

                    // Issue control request
                    match Self::request_control_url(addr, root_url.clone()) {
                        // Store pending requests
                        Ok(f) => {
                            self.pending.insert(addr, RequestState::Connecting(root_url, f));
                        }
                        Err(e) => return Poll::Ready(Err(e)),
                    }
                } else {
                    debug!("received duplicate broadcast response from: {}, dropping", from);
                }
            }
        }
        if let Poll::Ready(Err(err)) = resp {
            return Poll::Ready(Err(err.into()));
        }

        // Poll on any outstanding control requests
        for (addr, state) in &mut self.pending {
            // Poll if we're in the connecting state
            let (root_url, resp) = {
                let (root_url, c) = match state {
                    RequestState::Connecting(root_url, c) => (root_url, c),
                    _ => continue,
                };

                match c.as_mut().poll(cx)? {
                    Poll::Ready(resp) => (root_url.clone(), resp),
                    _ => continue,
                }
            };

            // Handle any responses
            if let Ok(control_url) = Self::handle_control_resp(*addr, resp) {
                debug!("received control url from: {} (url: {})", addr, control_url);
                *state = RequestState::Done(url.clone());

                match addr {
                    SocketAddr::V4(a) => {
                        let g = Gateway {
                            addr: *a,
                            root_url,
                            control_url,
                        };
                        return Poll::Ready(Ok(g));
                    }
                    _ => warn!("unsupported IPv6 gateway response from addr: {}", addr),
                }
            } else {
                *state = RequestState::Error;
            }
        }

        Poll::Pending
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::path::Path;
use std::time::Duration;

use hab_core::util::sys;
use hyper::client::Client as HyperClient;
use hyper::client::RequestBuilder;
use hyper::client::pool::{Config, Pool};
use hyper::header::UserAgent;
use hyper::http::h1::Http11Protocol;
use hyper::net::HttpsConnector;
use hyper_openssl::OpensslClient;
use openssl::ssl::{SslConnectorBuilder, SslConnector, SslMethod, SslOption, SSL_OP_NO_SSLV2,
                   SSL_OP_NO_SSLV3, SSL_OP_NO_COMPRESSION};
use url::Url;

use error::Result;
use net::ProxyHttpsConnector;
use proxy::{ProxyInfo, proxy_unless_domain_exempted};
use ssl;

// Read and write TCP socket timeout for Hyper/HTTP client calls.
const CLIENT_SOCKET_RW_TIMEOUT: u64 = 30;

header! { (ProxyAuthorization, "Proxy-Authorization") => [String] }

/// A generic wrapper around a Hyper HTTP client intended for API-like usage.
///
/// When an `ApiClient` is created, it has a constant URL base which is assumed to be some API
/// endpoint. This allows the underlying Hyper client to load and use any relevant HTTP proxy
/// support and to provide reasonable User-Agent HTTP headers, etc.
#[derive(Debug)]
pub struct ApiClient {
    /// The base URL for the client.
    endpoint: Url,
    /// An instance of a `hyper::Client` which is configured with an SSL context and optionally
    /// using an HTTP proxy.
    inner: HyperClient,
    /// Proxy information, if a proxy is being used.
    proxy: Option<ProxyInfo>,
    /// The URL scheme of the endpoint.
    target_scheme: String,
    /// The `User-Agent` header string to use for HTTP calls.
    user_agent_header: UserAgent,
}

impl ApiClient {
    /// Creates and returns a new `ApiClient` instance.
    ///
    /// # Errors
    ///
    /// * If the underlying Hyper client cannot be created
    /// * If a suitable SSL context cannot be established
    /// * If an HTTP proxy cannot be correctly setup
    /// * If a `User-Agent` HTTP header string cannot be constructed
    pub fn new(endpoint: &Url,
               product: &str,
               version: &str,
               fs_root_path: Option<&Path>)
               -> Result<Self> {
        Ok(ApiClient {
               endpoint: endpoint.clone(),
               inner: try!(new_hyper_client(Some(endpoint), fs_root_path)),
               proxy: try!(proxy_unless_domain_exempted(Some(endpoint))),
               target_scheme: endpoint.scheme().to_string(),
               user_agent_header: try!(user_agent(product, version)),
           })
    }

    /// Builds an HTTP GET request for a given path.
    pub fn get(&self, path: &str) -> RequestBuilder {
        self.get_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP GET request for a given path with the ability to customize the target URL.
    pub fn get_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("GET {} with {:?}", &url, &self);
        self.add_headers(self.inner.get(url))
    }

    /// Builds an HTTP HEAD request for a given path.
    pub fn head(&self, path: &str) -> RequestBuilder {
        self.head_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP HEAD request for a given path with the ability to customize the target URL.
    pub fn head_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("HEAD {} with {:?}", &url, &self);
        self.add_headers(self.inner.head(url))
    }

    /// Builds an HTTP PATCH request for a given path.
    pub fn patch(&self, path: &str) -> RequestBuilder {
        self.patch_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP PATCH request for a given path with the ability to customize the target URL.
    pub fn patch_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("PATH {} with {:?}", &url, &self);
        self.add_headers(self.inner.patch(url))
    }

    /// Builds an HTTP POST request for a given path.
    pub fn post(&self, path: &str) -> RequestBuilder {
        self.post_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP POST request for a given path with the ability to customize the target URL.
    pub fn post_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("POST {} with {:?}", &url, &self);
        self.add_headers(self.inner.post(url))
    }

    /// Builds an HTTP PUT request for a given path.
    pub fn put(&self, path: &str) -> RequestBuilder {
        self.put_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP PUT request for a given path with the ability to customize the target URL.
    pub fn put_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("PUT {} with {:?}", &url, &self);
        self.add_headers(self.inner.put(url))
    }

    /// Builds an HTTP DELETE request for a given path.
    pub fn delete(&self, path: &str) -> RequestBuilder {
        self.delete_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP DELETE request for a given path with the ability to customize the target URL.
    pub fn delete_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("DELETE {} with {:?}", &url, &self);
        self.add_headers(self.inner.delete(url))
    }

    fn add_headers<'a>(&'a self, rb: RequestBuilder<'a>) -> RequestBuilder {
        let mut rb = rb.header(self.user_agent_header.clone());
        // If the target URL is an `"http"` scheme and we're using a proxy server, then add the
        // proxy authorization header if appropriate. Note that for `"https"` targets, the proxy
        // server will be operating in TCP tunneling mode and will be authenticated on connection to
        // the proxy server which is why we should not add an additional header in this latter
        // case.
        if self.target_scheme == "http" {
            if let Some(ref info) = self.proxy {
                if let Some(header_value) = info.authorization_header_value() {
                    rb = rb.header(ProxyAuthorization(header_value));
                }
            }
        }
        rb
    }

    fn url_for(&self, path: &str) -> Url {
        let mut url = self.endpoint.clone();
        if path.is_empty() {
            return url;
        }
        if url.path().ends_with("/") {
            url.set_path(&format!("{}{}", self.endpoint.path(), path));
        } else {
            url.set_path(&format!("{}/{}", self.endpoint.path(), path));
        }
        url
    }
}

/// Builds a new hyper HTTP client with appropriate SSL configuration and HTTP/HTTPS proxy support.
///
/// ## Linux Platforms
///
/// We need a set of root certificates when connected to SSL/TLS web endpoints and this usually
/// boiled down to using an all-in-one certificate file (such as a `cacert.pem` file) or a directory
/// of files which are certificates. The strategy to location or use a reasonable set of
/// certificates is as follows:
///
/// 1. If the `SSL_CERT_FILE` environment variable is set, then use its value for the certificates.
///    Internally this is triggering default OpenSSL behavior for this environment variable.
/// 2. If the `SSL_CERT_DIR` environment variable is set, then use its value for the directory
///    containing certificates. Like the `SSL_CERT_FILE` case above, this triggers default OpenSSL
///    behavior for this environment variable.
/// 3. If the `core/cacerts` Habitat package is installed locally, then use the latest release's
///    `cacert.pem` file.
/// 4. If none of the conditions above are met, then a `cacert.pem` will be written in an SSL cache
///    directory (by default `/hab/cache/ssl` for a root user and `$HOME/.hab/cache/ssl` for a
///    non-root user) and this will be used. The contents of this file will be inlined in this
///    crate at build time as a fallback insurance policy, meaning that if the a program using this
///    code is operating in a minimal environment which may not contain system certificates, it can
///    still operate. Once a `core/cacerts` Habitat package is present, the logic would fall back
///    to preferring the package version to the cached/inline file version.
///
/// ## Mac Platforms
///
/// The Mac platform uses a Security Framework to store and find root certificates and the hyper
/// library will default to using this on the Mac. Therefore the behavior on the Mac remains
/// unchanged and will use the system's certificates.
///
fn new_hyper_client(for_domain: Option<&Url>, fs_root_path: Option<&Path>) -> Result<HyperClient> {
    let connector = try!(ssl_connector(fs_root_path));
    let ssl_client = OpensslClient::from(connector);
    let timeout = Some(Duration::from_secs(CLIENT_SOCKET_RW_TIMEOUT));

    match try!(proxy_unless_domain_exempted(for_domain)) {
        Some(proxy) => {
            debug!("Using proxy {}:{}...", proxy.host(), proxy.port());
            let connector = try!(ProxyHttpsConnector::new(proxy, ssl_client));
            let pool = Pool::with_connector(Config::default(), connector);
            let mut client = HyperClient::with_protocol(Http11Protocol::with_connector(pool));
            client.set_read_timeout(timeout);
            client.set_write_timeout(timeout);
            Ok(client)
        }
        None => {
            let connector = HttpsConnector::new(ssl_client);
            let pool = Pool::with_connector(Config::default(), connector);
            let mut client = HyperClient::with_protocol(Http11Protocol::with_connector(pool));
            client.set_read_timeout(timeout);
            client.set_write_timeout(timeout);
            Ok(client)
        }
    }
}

/// Returns an HTTP User-Agent string type for use by Hyper when making HTTP requests.
///
/// The general form for Habitat-related clients are of the following form:
///
/// ```text
/// <PRODUCT>/<VERSION> (<TARGET>; <KERNEL_RELEASE>)
/// ```
///
/// where:
///
/// * `<PRODUCT>`: is the provided product name
/// * `<VERSION>`: is the provided version string which may also include a release number
/// * `<TARGET>`: is the machine architecture and the kernel separated by a dash in lower case
/// * `<KERNEL_RELEASE>`: is the kernel release string from `uname`
///
/// For example:
///
/// ```text
/// hab/0.6.0/20160606153031 (x86_64-darwin; 14.5.0)
/// ```
///
/// # Errors
///
/// * If system information cannot be obtained via `uname`
fn user_agent(product: &str, version: &str) -> Result<UserAgent> {
    let uname = try!(sys::uname());
    let ua = format!("{}/{} ({}-{}; {})",
                     product.trim(),
                     version.trim(),
                     uname.machine.trim().to_lowercase(),
                     uname.sys_name.trim().to_lowercase(),
                     uname.release.trim().to_lowercase());
    debug!("User-Agent: {}", &ua);
    Ok(UserAgent(ua))
}

fn ssl_connector(fs_root_path: Option<&Path>) -> Result<SslConnector> {
    let mut conn = try!(SslConnectorBuilder::new(SslMethod::tls()));
    let mut options = SslOption::empty();
    options.toggle(SSL_OP_NO_SSLV2);
    options.toggle(SSL_OP_NO_SSLV3);
    options.toggle(SSL_OP_NO_COMPRESSION);
    try!(ssl::set_ca(conn.builder_mut(), fs_root_path));
    conn.builder_mut().set_options(options);
    try!(conn.builder_mut()
             .set_cipher_list("ALL!EXPORT!EXPORT40!EXPORT56!aNULL!LOW!RC4@STRENGTH"));
    Ok(conn.build())
}
<commit_msg>Add `--channel` flag to `svc start` and `svc load` commands<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::path::Path;
use std::time::Duration;

use hab_core::util::sys;
use hyper::client::{Client as HyperClient, IntoUrl, RequestBuilder};
use hyper::client::pool::{Config, Pool};
use hyper::header::UserAgent;
use hyper::http::h1::Http11Protocol;
use hyper::net::HttpsConnector;
use hyper_openssl::OpensslClient;
use openssl::ssl::{SslConnectorBuilder, SslConnector, SslMethod, SslOption, SSL_OP_NO_SSLV2,
                   SSL_OP_NO_SSLV3, SSL_OP_NO_COMPRESSION};
use url::Url;

use error::Result;
use net::ProxyHttpsConnector;
use proxy::{ProxyInfo, proxy_unless_domain_exempted};
use ssl;

// Read and write TCP socket timeout for Hyper/HTTP client calls.
const CLIENT_SOCKET_RW_TIMEOUT: u64 = 30;

header! { (ProxyAuthorization, "Proxy-Authorization") => [String] }

/// A generic wrapper around a Hyper HTTP client intended for API-like usage.
///
/// When an `ApiClient` is created, it has a constant URL base which is assumed to be some API
/// endpoint. This allows the underlying Hyper client to load and use any relevant HTTP proxy
/// support and to provide reasonable User-Agent HTTP headers, etc.
#[derive(Debug)]
pub struct ApiClient {
    /// The base URL for the client.
    endpoint: Url,
    /// An instance of a `hyper::Client` which is configured with an SSL context and optionally
    /// using an HTTP proxy.
    inner: HyperClient,
    /// Proxy information, if a proxy is being used.
    proxy: Option<ProxyInfo>,
    /// The URL scheme of the endpoint.
    target_scheme: String,
    /// The `User-Agent` header string to use for HTTP calls.
    user_agent_header: UserAgent,
}

impl ApiClient {
    /// Creates and returns a new `ApiClient` instance.
    ///
    /// # Errors
    ///
    /// * If the underlying Hyper client cannot be created
    /// * If a suitable SSL context cannot be established
    /// * If an HTTP proxy cannot be correctly setup
    /// * If a `User-Agent` HTTP header string cannot be constructed
    pub fn new<T>(endpoint: T,
                  product: &str,
                  version: &str,
                  fs_root_path: Option<&Path>)
                  -> Result<Self>
        where T: IntoUrl
    {
        let endpoint = endpoint.into_url()?;
        Ok(ApiClient {
               inner: new_hyper_client(&endpoint, fs_root_path)?,
               proxy: proxy_unless_domain_exempted(Some(&endpoint))?,
               target_scheme: endpoint.scheme().to_string(),
               endpoint: endpoint,
               user_agent_header: user_agent(product, version)?,
           })
    }

    /// Builds an HTTP GET request for a given path.
    pub fn get(&self, path: &str) -> RequestBuilder {
        self.get_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP GET request for a given path with the ability to customize the target URL.
    pub fn get_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("GET {} with {:?}", &url, &self);
        self.add_headers(self.inner.get(url))
    }

    /// Builds an HTTP HEAD request for a given path.
    pub fn head(&self, path: &str) -> RequestBuilder {
        self.head_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP HEAD request for a given path with the ability to customize the target URL.
    pub fn head_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("HEAD {} with {:?}", &url, &self);
        self.add_headers(self.inner.head(url))
    }

    /// Builds an HTTP PATCH request for a given path.
    pub fn patch(&self, path: &str) -> RequestBuilder {
        self.patch_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP PATCH request for a given path with the ability to customize the target URL.
    pub fn patch_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("PATH {} with {:?}", &url, &self);
        self.add_headers(self.inner.patch(url))
    }

    /// Builds an HTTP POST request for a given path.
    pub fn post(&self, path: &str) -> RequestBuilder {
        self.post_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP POST request for a given path with the ability to customize the target URL.
    pub fn post_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("POST {} with {:?}", &url, &self);
        self.add_headers(self.inner.post(url))
    }

    /// Builds an HTTP PUT request for a given path.
    pub fn put(&self, path: &str) -> RequestBuilder {
        self.put_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP PUT request for a given path with the ability to customize the target URL.
    pub fn put_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("PUT {} with {:?}", &url, &self);
        self.add_headers(self.inner.put(url))
    }

    /// Builds an HTTP DELETE request for a given path.
    pub fn delete(&self, path: &str) -> RequestBuilder {
        self.delete_with_custom_url(path, |_| {})
    }

    /// Builds an HTTP DELETE request for a given path with the ability to customize the target URL.
    pub fn delete_with_custom_url<F>(&self, path: &str, mut customize_url: F) -> RequestBuilder
        where F: FnMut(&mut Url)
    {
        let mut url = self.url_for(path);
        customize_url(&mut url);
        debug!("DELETE {} with {:?}", &url, &self);
        self.add_headers(self.inner.delete(url))
    }

    fn add_headers<'a>(&'a self, rb: RequestBuilder<'a>) -> RequestBuilder {
        let mut rb = rb.header(self.user_agent_header.clone());
        // If the target URL is an `"http"` scheme and we're using a proxy server, then add the
        // proxy authorization header if appropriate. Note that for `"https"` targets, the proxy
        // server will be operating in TCP tunneling mode and will be authenticated on connection to
        // the proxy server which is why we should not add an additional header in this latter
        // case.
        if self.target_scheme == "http" {
            if let Some(ref info) = self.proxy {
                if let Some(header_value) = info.authorization_header_value() {
                    rb = rb.header(ProxyAuthorization(header_value));
                }
            }
        }
        rb
    }

    fn url_for(&self, path: &str) -> Url {
        let mut url = self.endpoint.clone();
        if path.is_empty() {
            return url;
        }
        if url.path().ends_with("/") {
            url.set_path(&format!("{}{}", self.endpoint.path(), path));
        } else {
            url.set_path(&format!("{}/{}", self.endpoint.path(), path));
        }
        url
    }
}

/// Builds a new hyper HTTP client with appropriate SSL configuration and HTTP/HTTPS proxy support.
///
/// ## Linux Platforms
///
/// We need a set of root certificates when connected to SSL/TLS web endpoints and this usually
/// boiled down to using an all-in-one certificate file (such as a `cacert.pem` file) or a directory
/// of files which are certificates. The strategy to location or use a reasonable set of
/// certificates is as follows:
///
/// 1. If the `SSL_CERT_FILE` environment variable is set, then use its value for the certificates.
///    Internally this is triggering default OpenSSL behavior for this environment variable.
/// 2. If the `SSL_CERT_DIR` environment variable is set, then use its value for the directory
///    containing certificates. Like the `SSL_CERT_FILE` case above, this triggers default OpenSSL
///    behavior for this environment variable.
/// 3. If the `core/cacerts` Habitat package is installed locally, then use the latest release's
///    `cacert.pem` file.
/// 4. If none of the conditions above are met, then a `cacert.pem` will be written in an SSL cache
///    directory (by default `/hab/cache/ssl` for a root user and `$HOME/.hab/cache/ssl` for a
///    non-root user) and this will be used. The contents of this file will be inlined in this
///    crate at build time as a fallback insurance policy, meaning that if the a program using this
///    code is operating in a minimal environment which may not contain system certificates, it can
///    still operate. Once a `core/cacerts` Habitat package is present, the logic would fall back
///    to preferring the package version to the cached/inline file version.
///
/// ## Mac Platforms
///
/// The Mac platform uses a Security Framework to store and find root certificates and the hyper
/// library will default to using this on the Mac. Therefore the behavior on the Mac remains
/// unchanged and will use the system's certificates.
///
fn new_hyper_client(url: &Url, fs_root_path: Option<&Path>) -> Result<HyperClient> {
    let connector = try!(ssl_connector(fs_root_path));
    let ssl_client = OpensslClient::from(connector);
    let timeout = Some(Duration::from_secs(CLIENT_SOCKET_RW_TIMEOUT));

    match proxy_unless_domain_exempted(Some(url))? {
        Some(proxy) => {
            debug!("Using proxy {}:{}...", proxy.host(), proxy.port());
            let connector = try!(ProxyHttpsConnector::new(proxy, ssl_client));
            let pool = Pool::with_connector(Config::default(), connector);
            let mut client = HyperClient::with_protocol(Http11Protocol::with_connector(pool));
            client.set_read_timeout(timeout);
            client.set_write_timeout(timeout);
            Ok(client)
        }
        None => {
            let connector = HttpsConnector::new(ssl_client);
            let pool = Pool::with_connector(Config::default(), connector);
            let mut client = HyperClient::with_protocol(Http11Protocol::with_connector(pool));
            client.set_read_timeout(timeout);
            client.set_write_timeout(timeout);
            Ok(client)
        }
    }
}

/// Returns an HTTP User-Agent string type for use by Hyper when making HTTP requests.
///
/// The general form for Habitat-related clients are of the following form:
///
/// ```text
/// <PRODUCT>/<VERSION> (<TARGET>; <KERNEL_RELEASE>)
/// ```
///
/// where:
///
/// * `<PRODUCT>`: is the provided product name
/// * `<VERSION>`: is the provided version string which may also include a release number
/// * `<TARGET>`: is the machine architecture and the kernel separated by a dash in lower case
/// * `<KERNEL_RELEASE>`: is the kernel release string from `uname`
///
/// For example:
///
/// ```text
/// hab/0.6.0/20160606153031 (x86_64-darwin; 14.5.0)
/// ```
///
/// # Errors
///
/// * If system information cannot be obtained via `uname`
fn user_agent(product: &str, version: &str) -> Result<UserAgent> {
    let uname = try!(sys::uname());
    let ua = format!("{}/{} ({}-{}; {})",
                     product.trim(),
                     version.trim(),
                     uname.machine.trim().to_lowercase(),
                     uname.sys_name.trim().to_lowercase(),
                     uname.release.trim().to_lowercase());
    debug!("User-Agent: {}", &ua);
    Ok(UserAgent(ua))
}

fn ssl_connector(fs_root_path: Option<&Path>) -> Result<SslConnector> {
    let mut conn = try!(SslConnectorBuilder::new(SslMethod::tls()));
    let mut options = SslOption::empty();
    options.toggle(SSL_OP_NO_SSLV2);
    options.toggle(SSL_OP_NO_SSLV3);
    options.toggle(SSL_OP_NO_COMPRESSION);
    try!(ssl::set_ca(conn.builder_mut(), fs_root_path));
    conn.builder_mut().set_options(options);
    try!(conn.builder_mut()
             .set_cipher_list("ALL!EXPORT!EXPORT40!EXPORT56!aNULL!LOW!RC4@STRENGTH"));
    Ok(conn.build())
}
<|endoftext|>