<commit_before>// Type inference

use core::{Match, Src, Scope, Session};
use nameres::resolve_path_with_str;
use core::Namespace;
use core;
use ast;
use scopes;
use matchers;
use core::SearchType::ExactMatch;
use util::txt_matches;

fn find_start_of_function_body(src: &str) -> usize {
    // TODO: this should ignore anything inside parens so as to skip the arg list
    src.find('{').unwrap()
}

// Removes the body of the statement (anything in the braces {...}), leaving just
// the header
// TODO: this should skip parens (e.g. function arguments)
pub fn generate_skeleton_for_parsing(src: &str) -> String {
    let mut s = String::new();
    let n = src.find('{').unwrap();
    s.push_str(&src[..n+1]);
    s.push_str("};");
    s
}

pub fn first_param_is_self(blob: &str) -> bool {
    // skip generic arg
    // consider 'pub fn map<U, F: FnOnce(T) -> U>(self, f: F)'
    // we have to match the '>'
    match blob.find('(') {
        None => false,
        Some(probable_param_start) => {
            let skip_generic = match blob.find('<') {
                None => 0,
                Some(generic_start) if generic_start < probable_param_start => {
                    let mut level = 0;
                    let mut prev = ' ';
                    let mut skip_generic = 0;
                    for (i, c) in blob[generic_start..].char_indices() {
                        match c {
                            '<' => level += 1,
                            '>' if prev == '-' => (),
                            '>' => level -= 1,
                            _ => (),
                        }
                        prev = c;
                        if level == 0 {
                            skip_generic = i;
                            break;
                        }
                    }
                    skip_generic
                },
                Some(..) => 0,
            };
            while let Some(start) = blob[skip_generic..].find('(') {
                let end = scopes::find_closing_paren(blob, start + 1);
                let is_self = txt_matches(ExactMatch, "self", &blob[(start + 1)..end]);
                debug!("searching fn args: |{}| {}",
                       &blob[(start + 1)..end],
                       is_self);
                return is_self;
            }
            false
        }
    }
}

#[test]
fn generates_skeleton_for_mod() {
    let src = "mod foo { blah };";
    let out = generate_skeleton_for_parsing(src);
    assert_eq!("mod foo {};", out);
}

fn get_type_of_self_arg(m: &Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    debug!("get_type_of_self_arg {:?}", m);
    scopes::find_impl_start(msrc, m.point, 0).and_then(|start| {
        let decl = generate_skeleton_for_parsing(&msrc.from(start));
        debug!("get_type_of_self_arg impl skeleton |{}|", decl);

        if decl.starts_with("impl") {
            let implres = ast::parse_impl(decl);
            debug!("get_type_of_self_arg implres |{:?}|", implres);
            resolve_path_with_str(&implres.name_path.expect("failed parsing impl name"),
                                  &m.filepath, start,
                                  ExactMatch, Namespace::Type,
                                  session).nth(0).map(core::Ty::Match)
        } else {
            // // must be a trait
            ast::parse_trait(decl).name.and_then(|name| {
                Some(core::Ty::Match(Match {
                           matchstr: name,
                           filepath: m.filepath.clone(),
                           point: start,
                           local: m.local,
                           mtype: core::MatchType::Trait,
                           contextstr: matchers::first_line(&msrc[start..]),
                           generic_args: Vec::new(),
                           generic_types: Vec::new(),
                           docs: String::new(),
                }))
            })
        }
    })
}

fn get_type_of_fnarg(m: &Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    if m.matchstr == "self" {
        return get_type_of_self_arg(m, msrc, session);
    }

    let stmtstart = scopes::find_stmt_start(msrc, m.point).unwrap();
    let block = msrc.from(stmtstart);
    if let Some((start, end)) = block.iter_stmts().next() {
        let blob = &msrc[(stmtstart+start)..(stmtstart+end)];
        // wrap in "impl blah { }" so that methods get parsed correctly too
        let mut s = String::new();
        s.push_str("impl blah {");
        let impl_header_len = s.len();
        s.push_str(&blob[..(find_start_of_function_body(blob)+1)]);
        s.push_str("}}");
        let argpos = m.point - (stmtstart+start) + impl_header_len;
        return ast::parse_fn_arg_type(s, argpos, Scope::from_match(m), session);
    }
    None
}

fn get_type_of_let_expr(m: &Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    // ASSUMPTION: this is being called on a let decl
    let point = scopes::find_stmt_start(msrc, m.point).unwrap();
    let src = msrc.from(point);

    if let Some((start, end)) = src.iter_stmts().next() {
        let blob = &src[start..end];
        debug!("get_type_of_let_expr calling parse_let |{}|", blob);

        let pos = m.point - point - start;
        let scope = Scope{ filepath: m.filepath.clone(), point: m.point };
        ast::get_let_type(blob.to_owned(), pos, scope, session)
    } else {
        None
    }
}

fn get_type_of_let_block_expr(m: &Match, msrc: Src, session: &Session, prefix: &str) -> Option<core::Ty> {
    // ASSUMPTION: this is being called on an if let or while let decl
    let stmtstart = scopes::find_stmt_start(msrc, m.point).unwrap();
    let stmt = msrc.from(stmtstart);
    let point = stmt.find(prefix).unwrap();
    let src = core::new_source(generate_skeleton_for_parsing(&stmt[point..]));

    if let Some((start, end)) = src.as_src().iter_stmts().next() {
        let blob = &src[start..end];
        debug!("get_type_of_let_block_expr calling get_let_type |{}|", blob);

        let pos = m.point - stmtstart - point - start;
        let scope = Scope{ filepath: m.filepath.clone(), point: m.point };
        ast::get_let_type(blob.to_owned(), pos, scope, session)
    } else {
        None
    }
}

fn get_type_of_for_expr(m: &Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    let stmtstart = scopes::find_stmt_start(msrc, m.point).unwrap();
    let stmt = msrc.from(stmtstart);
    let forpos = stmt.find("for ").unwrap();
    let inpos = stmt.find(" in ").unwrap();
    // XXX: this need not be the correct brace, see generate_skeleton_for_parsing
    let bracepos = stmt.find('{').unwrap();
    let mut src = stmt[..forpos].to_owned();
    src.push_str("if let Some(");
    src.push_str(&stmt[forpos+4..inpos]);
    src.push_str(") = ");
    let iter_stmt = &stmt[inpos+4..bracepos];

    // TODO: Remove these lines when iter()/iter_mut() method lookup on
    //       built in types is properly supported
    let mut iter_stmt_trimmed = iter_stmt.replace(".iter()", ".into_iter()");
    iter_stmt_trimmed = iter_stmt_trimmed.replace(".iter_mut()", ".into_iter()");

    src.push_str(&iter_stmt_trimmed);
    src = src.trim_right().to_owned();
    src.push_str(".into_iter().next() { }}");

    let src = core::new_source(src);

    if let Some((start, end)) = src.as_src().iter_stmts().next() {
        let blob = &src[start..end];
        debug!("get_type_of_for_expr: |{}| {} {} {} {}", blob, m.point, stmtstart, forpos, start);

        let pos = m.point + 8 - stmtstart - forpos - start;
        let scope = Scope{ filepath: m.filepath.clone(), point: m.point + 8 };

        ast::get_let_type(blob.to_owned(), pos, scope, session)
    } else {
        None
    }
}

pub fn get_struct_field_type(fieldname: &str, structmatch: &Match, session: &Session) -> Option<core::Ty> {
    assert!(structmatch.mtype == core::MatchType::Struct);

    let src = session.load_file(&structmatch.filepath);

    let opoint = scopes::find_stmt_start(src.as_src(), structmatch.point);
    let structsrc = scopes::end_of_next_scope(&src[opoint.unwrap()..]);

    let fields = ast::parse_struct_fields(structsrc.to_owned(), Scope::from_match(structmatch));
    for (field, _, ty) in fields.into_iter() {
        if fieldname == field {
            return ty;
        }
    }
    None
}

pub fn get_tuplestruct_field_type(fieldnum: usize, structmatch: &Match, session: &Session) -> Option<core::Ty> {
    let src = session.load_file(&structmatch.filepath);

    let structsrc = if let core::MatchType::EnumVariant = structmatch.mtype {
        // decorate the enum variant src to make it look like a tuple struct
        let to = src[structmatch.point..].find('(')
            .map(|n| scopes::find_closing_paren(&src, structmatch.point + n+1))
            .unwrap();
        "struct ".to_owned() + &src[structmatch.point..(to+1)] + ";"
    } else {
        assert!(structmatch.mtype == core::MatchType::Struct);
        let opoint = scopes::find_stmt_start(src.as_src(), structmatch.point);
        (*get_first_stmt(src.as_src().from(opoint.unwrap()))).to_owned()
    };

    debug!("get_tuplestruct_field_type structsrc=|{}|", structsrc);

    let fields = ast::parse_struct_fields(structsrc, Scope::from_match(structmatch));

    for (i, (_, _, ty)) in fields.into_iter().enumerate() {
        if i == fieldnum {
            return ty;
        }
    }
    None
}

pub fn get_first_stmt(src: Src) -> Src {
    match src.iter_stmts().next() {
        Some((from, to)) => src.from_to(from, to),
        None => src
    }
}

pub fn get_type_of_match(m: Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    debug!("get_type_of match {:?} ", m);

    match m.mtype {
        core::MatchType::Let => get_type_of_let_expr(&m, msrc, session),
        core::MatchType::IfLet => get_type_of_let_block_expr(&m, msrc, session, "if let"),
        core::MatchType::WhileLet => get_type_of_let_block_expr(&m, msrc, session, "while let"),
        core::MatchType::For => get_type_of_for_expr(&m, msrc, session),
        core::MatchType::FnArg => get_type_of_fnarg(&m, msrc, session),
        core::MatchType::MatchArm => get_type_from_match_arm(&m, msrc, session),
        core::MatchType::Struct |
        core::MatchType::Enum |
        core::MatchType::Function |
        core::MatchType::Module => Some(core::Ty::Match(m)),
        _ => { debug!("!!! WARNING !!! Can't get type of {:?}", m.mtype); None }
    }
}

macro_rules! otry {
    ($e:expr) => (match $e { Some(e) => e, None => return None })
}

pub fn get_type_from_match_arm(m: &Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    // We construct a faux match stmt and then parse it. This is because the
    // match stmt may be incomplete (half written) in the real code

    // skip to end of match arm pattern so we can search backwards
    let arm = otry!(msrc[m.point..].find("=>")) + m.point;
    let scopestart = scopes::scope_start(msrc, arm);

    let stmtstart = otry!(scopes::find_stmt_start(msrc, scopestart-1));
    debug!("PHIL preblock is {} {}", stmtstart, scopestart);
    let preblock = &msrc[stmtstart..scopestart];
    let matchstart = otry!(preblock.rfind("match ")) + stmtstart;

    let lhs_start = scopes::get_start_of_pattern(&msrc, arm);
    let lhs = &msrc[lhs_start..arm];
    // construct faux match statement and recreate point
    let mut fauxmatchstmt = msrc[matchstart..scopestart].to_owned();
    let faux_prefix_size = fauxmatchstmt.len();
    fauxmatchstmt = fauxmatchstmt + lhs + " => () };";
    let faux_point = faux_prefix_size + (m.point - lhs_start);

    debug!("fauxmatchstmt for parsing is pt:{} src:|{}|", faux_point, fauxmatchstmt);

    ast::get_match_arm_type(fauxmatchstmt, faux_point,
                            // scope is used to locate expression, so send
                            // it the start of the match expr
                            Scope {
                                filepath: m.filepath.clone(),
                                point: matchstart,
                            }, session)
}

pub fn get_function_declaration(fnmatch: &Match, session: &Session) -> String {
    let src = session.load_file(&fnmatch.filepath);
    let start = scopes::find_stmt_start(src.as_src(), fnmatch.point).unwrap();
    let def_end: &[_] = &['{', ';'];
    let end = src[start..].find(def_end).unwrap();
    src[start..end+start].to_owned()
}

pub fn get_return_type_of_function(fnmatch: &Match, contextm: &Match, session: &Session) -> Option<core::Ty> {
    let src = session.load_file(&fnmatch.filepath);
    let point = scopes::find_stmt_start(src.as_src(), fnmatch.point).unwrap();
    let mut out = src[point..].find(|c| {c == '{' || c == ';'}).and_then(|n| {
        // wrap in "impl blah { }" so that methods get parsed correctly too
        let mut decl = String::new();
        decl.push_str("impl blah {");
        decl.push_str(&src[point..(point+n+1)]);
        if decl.ends_with(';') {
            decl.pop();
            decl.push_str("{}}");
        }
        else {
            decl.push_str("}}");
        }
        debug!("get_return_type_of_function: passing in |{}|", decl);
        ast::parse_fn_output(decl, Scope::from_match(fnmatch))
    });

    // Convert output arg of type Self to the correct type
    if let Some(core::Ty::PathSearch(ref mut path, _)) = out {
        if let Some(ref mut path_seg) = path.clone().segments.get(0) {
            if "Self" == &path_seg.name {
                return get_type_of_self_arg(fnmatch, src.as_src(), session);
            }
        }
    }

    // Convert a generic output arg to the correct type
    if let Some(core::Ty::PathSearch(ref mut path, _)) = out {
        if let Some(ref mut path_seg) = path.clone().segments.get(0) {
            if path.segments.len() == 1 && path_seg.types.is_empty() {
                for type_name in &fnmatch.generic_args {
                    if type_name == &path_seg.name {
                        return Some(core::Ty::Match(contextm.clone()));
                    }
                }
            }
        }
    };
    out
}
<commit_msg>remove not needed mut and clone from get_return_type_of_function<commit_after>// Type inference

use core::{Match, Src, Scope, Session};
use nameres::resolve_path_with_str;
use core::Namespace;
use core;
use ast;
use scopes;
use matchers;
use core::SearchType::ExactMatch;
use util::txt_matches;

fn find_start_of_function_body(src: &str) -> usize {
    // TODO: this should ignore anything inside parens so as to skip the arg list
    src.find('{').unwrap()
}

// Removes the body of the statement (anything in the braces {...}), leaving just
// the header
// TODO: this should skip parens (e.g. function arguments)
pub fn generate_skeleton_for_parsing(src: &str) -> String {
    let mut s = String::new();
    let n = src.find('{').unwrap();
    s.push_str(&src[..n+1]);
    s.push_str("};");
    s
}

pub fn first_param_is_self(blob: &str) -> bool {
    // skip generic arg
    // consider 'pub fn map<U, F: FnOnce(T) -> U>(self, f: F)'
    // we have to match the '>'
    match blob.find('(') {
        None => false,
        Some(probable_param_start) => {
            let skip_generic = match blob.find('<') {
                None => 0,
                Some(generic_start) if generic_start < probable_param_start => {
                    let mut level = 0;
                    let mut prev = ' ';
                    let mut skip_generic = 0;
                    for (i, c) in blob[generic_start..].char_indices() {
                        match c {
                            '<' => level += 1,
                            '>' if prev == '-' => (),
                            '>' => level -= 1,
                            _ => (),
                        }
                        prev = c;
                        if level == 0 {
                            skip_generic = i;
                            break;
                        }
                    }
                    skip_generic
                },
                Some(..) => 0,
            };
            while let Some(start) = blob[skip_generic..].find('(') {
                let end = scopes::find_closing_paren(blob, start + 1);
                let is_self = txt_matches(ExactMatch, "self", &blob[(start + 1)..end]);
                debug!("searching fn args: |{}| {}",
                       &blob[(start + 1)..end],
                       is_self);
                return is_self;
            }
            false
        }
    }
}

#[test]
fn generates_skeleton_for_mod() {
    let src = "mod foo { blah };";
    let out = generate_skeleton_for_parsing(src);
    assert_eq!("mod foo {};", out);
}

fn get_type_of_self_arg(m: &Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    debug!("get_type_of_self_arg {:?}", m);
    scopes::find_impl_start(msrc, m.point, 0).and_then(|start| {
        let decl = generate_skeleton_for_parsing(&msrc.from(start));
        debug!("get_type_of_self_arg impl skeleton |{}|", decl);

        if decl.starts_with("impl") {
            let implres = ast::parse_impl(decl);
            debug!("get_type_of_self_arg implres |{:?}|", implres);
            resolve_path_with_str(&implres.name_path.expect("failed parsing impl name"),
                                  &m.filepath, start,
                                  ExactMatch, Namespace::Type,
                                  session).nth(0).map(core::Ty::Match)
        } else {
            // // must be a trait
            ast::parse_trait(decl).name.and_then(|name| {
                Some(core::Ty::Match(Match {
                           matchstr: name,
                           filepath: m.filepath.clone(),
                           point: start,
                           local: m.local,
                           mtype: core::MatchType::Trait,
                           contextstr: matchers::first_line(&msrc[start..]),
                           generic_args: Vec::new(),
                           generic_types: Vec::new(),
                           docs: String::new(),
                }))
            })
        }
    })
}

fn get_type_of_fnarg(m: &Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    if m.matchstr == "self" {
        return get_type_of_self_arg(m, msrc, session);
    }

    let stmtstart = scopes::find_stmt_start(msrc, m.point).unwrap();
    let block = msrc.from(stmtstart);
    if let Some((start, end)) = block.iter_stmts().next() {
        let blob = &msrc[(stmtstart+start)..(stmtstart+end)];
        // wrap in "impl blah { }" so that methods get parsed correctly too
        let mut s = String::new();
        s.push_str("impl blah {");
        let impl_header_len = s.len();
        s.push_str(&blob[..(find_start_of_function_body(blob)+1)]);
        s.push_str("}}");
        let argpos = m.point - (stmtstart+start) + impl_header_len;
        return ast::parse_fn_arg_type(s, argpos, Scope::from_match(m), session);
    }
    None
}

fn get_type_of_let_expr(m: &Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    // ASSUMPTION: this is being called on a let decl
    let point = scopes::find_stmt_start(msrc, m.point).unwrap();
    let src = msrc.from(point);

    if let Some((start, end)) = src.iter_stmts().next() {
        let blob = &src[start..end];
        debug!("get_type_of_let_expr calling parse_let |{}|", blob);

        let pos = m.point - point - start;
        let scope = Scope{ filepath: m.filepath.clone(), point: m.point };
        ast::get_let_type(blob.to_owned(), pos, scope, session)
    } else {
        None
    }
}

fn get_type_of_let_block_expr(m: &Match, msrc: Src, session: &Session, prefix: &str) -> Option<core::Ty> {
    // ASSUMPTION: this is being called on an if let or while let decl
    let stmtstart = scopes::find_stmt_start(msrc, m.point).unwrap();
    let stmt = msrc.from(stmtstart);
    let point = stmt.find(prefix).unwrap();
    let src = core::new_source(generate_skeleton_for_parsing(&stmt[point..]));

    if let Some((start, end)) = src.as_src().iter_stmts().next() {
        let blob = &src[start..end];
        debug!("get_type_of_let_block_expr calling get_let_type |{}|", blob);

        let pos = m.point - stmtstart - point - start;
        let scope = Scope{ filepath: m.filepath.clone(), point: m.point };
        ast::get_let_type(blob.to_owned(), pos, scope, session)
    } else {
        None
    }
}

fn get_type_of_for_expr(m: &Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    let stmtstart = scopes::find_stmt_start(msrc, m.point).unwrap();
    let stmt = msrc.from(stmtstart);
    let forpos = stmt.find("for ").unwrap();
    let inpos = stmt.find(" in ").unwrap();
    // XXX: this need not be the correct brace, see generate_skeleton_for_parsing
    let bracepos = stmt.find('{').unwrap();
    let mut src = stmt[..forpos].to_owned();
    src.push_str("if let Some(");
    src.push_str(&stmt[forpos+4..inpos]);
    src.push_str(") = ");
    let iter_stmt = &stmt[inpos+4..bracepos];

    // TODO: Remove these lines when iter()/iter_mut() method lookup on
    //       built in types is properly supported
    let mut iter_stmt_trimmed = iter_stmt.replace(".iter()", ".into_iter()");
    iter_stmt_trimmed = iter_stmt_trimmed.replace(".iter_mut()", ".into_iter()");

    src.push_str(&iter_stmt_trimmed);
    src = src.trim_right().to_owned();
    src.push_str(".into_iter().next() { }}");

    let src = core::new_source(src);

    if let Some((start, end)) = src.as_src().iter_stmts().next() {
        let blob = &src[start..end];
        debug!("get_type_of_for_expr: |{}| {} {} {} {}", blob, m.point, stmtstart, forpos, start);

        let pos = m.point + 8 - stmtstart - forpos - start;
        let scope = Scope{ filepath: m.filepath.clone(), point: m.point + 8 };

        ast::get_let_type(blob.to_owned(), pos, scope, session)
    } else {
        None
    }
}

pub fn get_struct_field_type(fieldname: &str, structmatch: &Match, session: &Session) -> Option<core::Ty> {
    assert!(structmatch.mtype == core::MatchType::Struct);

    let src = session.load_file(&structmatch.filepath);

    let opoint = scopes::find_stmt_start(src.as_src(), structmatch.point);
    let structsrc = scopes::end_of_next_scope(&src[opoint.unwrap()..]);

    let fields = ast::parse_struct_fields(structsrc.to_owned(), Scope::from_match(structmatch));
    for (field, _, ty) in fields.into_iter() {
        if fieldname == field {
            return ty;
        }
    }
    None
}

pub fn get_tuplestruct_field_type(fieldnum: usize, structmatch: &Match, session: &Session) -> Option<core::Ty> {
    let src = session.load_file(&structmatch.filepath);

    let structsrc = if let core::MatchType::EnumVariant = structmatch.mtype {
        // decorate the enum variant src to make it look like a tuple struct
        let to = src[structmatch.point..].find('(')
            .map(|n| scopes::find_closing_paren(&src, structmatch.point + n+1))
            .unwrap();
        "struct ".to_owned() + &src[structmatch.point..(to+1)] + ";"
    } else {
        assert!(structmatch.mtype == core::MatchType::Struct);
        let opoint = scopes::find_stmt_start(src.as_src(), structmatch.point);
        (*get_first_stmt(src.as_src().from(opoint.unwrap()))).to_owned()
    };

    debug!("get_tuplestruct_field_type structsrc=|{}|", structsrc);

    let fields = ast::parse_struct_fields(structsrc, Scope::from_match(structmatch));

    for (i, (_, _, ty)) in fields.into_iter().enumerate() {
        if i == fieldnum {
            return ty;
        }
    }
    None
}

pub fn get_first_stmt(src: Src) -> Src {
    match src.iter_stmts().next() {
        Some((from, to)) => src.from_to(from, to),
        None => src
    }
}

pub fn get_type_of_match(m: Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    debug!("get_type_of match {:?} ", m);

    match m.mtype {
        core::MatchType::Let => get_type_of_let_expr(&m, msrc, session),
        core::MatchType::IfLet => get_type_of_let_block_expr(&m, msrc, session, "if let"),
        core::MatchType::WhileLet => get_type_of_let_block_expr(&m, msrc, session, "while let"),
        core::MatchType::For => get_type_of_for_expr(&m, msrc, session),
        core::MatchType::FnArg => get_type_of_fnarg(&m, msrc, session),
        core::MatchType::MatchArm => get_type_from_match_arm(&m, msrc, session),
        core::MatchType::Struct |
        core::MatchType::Enum |
        core::MatchType::Function |
        core::MatchType::Module => Some(core::Ty::Match(m)),
        _ => { debug!("!!! WARNING !!! Can't get type of {:?}", m.mtype); None }
    }
}

macro_rules! otry {
    ($e:expr) => (match $e { Some(e) => e, None => return None })
}

pub fn get_type_from_match_arm(m: &Match, msrc: Src, session: &Session) -> Option<core::Ty> {
    // We construct a faux match stmt and then parse it. This is because the
    // match stmt may be incomplete (half written) in the real code

    // skip to end of match arm pattern so we can search backwards
    let arm = otry!(msrc[m.point..].find("=>")) + m.point;
    let scopestart = scopes::scope_start(msrc, arm);

    let stmtstart = otry!(scopes::find_stmt_start(msrc, scopestart-1));
    debug!("PHIL preblock is {} {}", stmtstart, scopestart);
    let preblock = &msrc[stmtstart..scopestart];
    let matchstart = otry!(preblock.rfind("match ")) + stmtstart;

    let lhs_start = scopes::get_start_of_pattern(&msrc, arm);
    let lhs = &msrc[lhs_start..arm];
    // construct faux match statement and recreate point
    let mut fauxmatchstmt = msrc[matchstart..scopestart].to_owned();
    let faux_prefix_size = fauxmatchstmt.len();
    fauxmatchstmt = fauxmatchstmt + lhs + " => () };";
    let faux_point = faux_prefix_size + (m.point - lhs_start);

    debug!("fauxmatchstmt for parsing is pt:{} src:|{}|", faux_point, fauxmatchstmt);

    ast::get_match_arm_type(fauxmatchstmt, faux_point,
                            // scope is used to locate expression, so send
                            // it the start of the match expr
                            Scope {
                                filepath: m.filepath.clone(),
                                point: matchstart,
                            }, session)
}

pub fn get_function_declaration(fnmatch: &Match, session: &Session) -> String {
    let src = session.load_file(&fnmatch.filepath);
    let start = scopes::find_stmt_start(src.as_src(), fnmatch.point).unwrap();
    let def_end: &[_] = &['{', ';'];
    let end = src[start..].find(def_end).unwrap();
    src[start..end+start].to_owned()
}

pub fn get_return_type_of_function(fnmatch: &Match, contextm: &Match, session: &Session) -> Option<core::Ty> {
    let src = session.load_file(&fnmatch.filepath);
    let point = scopes::find_stmt_start(src.as_src(), fnmatch.point).unwrap();
    let out = src[point..].find(|c| {c == '{' || c == ';'}).and_then(|n| {
        // wrap in "impl blah { }" so that methods get parsed correctly too
        let mut decl = String::new();
        decl.push_str("impl blah {");
        decl.push_str(&src[point..(point+n+1)]);
        if decl.ends_with(';') {
            decl.pop();
            decl.push_str("{}}");
        }
        else {
            decl.push_str("}}");
        }
        debug!("get_return_type_of_function: passing in |{}|", decl);
        ast::parse_fn_output(decl, Scope::from_match(fnmatch))
    });

    // Convert output arg of type Self to the correct type
    if let Some(core::Ty::PathSearch(ref path, _)) = out {
        if let Some(ref path_seg) = path.segments.get(0) {
            if "Self" == path_seg.name {
                return get_type_of_self_arg(fnmatch, src.as_src(), session);
            }
        }
    }

    // Convert a generic output arg to the correct type
    if let Some(core::Ty::PathSearch(ref path, _)) = out {
        if let Some(ref path_seg) = path.segments.get(0) {
            if path.segments.len() == 1 && path_seg.types.is_empty() {
                for type_name in &fnmatch.generic_args {
                    if type_name == &path_seg.name {
                        return Some(core::Ty::Match(contextm.clone()));
                    }
                }
            }
        }
    };
    out
}
<|endoftext|>


<commit_before>// Copyright 2014 The Rustastic SMTP Developers
//
// 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.

//! The `server` module contains things needed to build an SMTP server, but useless for
//! an SMTP client.

use std::io::net::tcp::{TcpListener, TcpAcceptor, TcpStream};
use std::io::{Listener, Acceptor, IoError, Reader, Writer, InvalidInput};
use super::common::stream::{SmtpStream};
use std::sync::Arc;
use std::ascii::OwnedAsciiExt;
use super::common::transaction::SmtpTransaction;
use super::common::mailbox::Mailbox;

mod handler;

/// Hooks into different places of the SMTP server to allow its customization.
pub trait SmtpServerEventHandler {
    /// Called after getting a recipient mailbox.
    ///
    /// If `Ok(())` is returned, a 250 response is sent. If `Err(())` is returned, a 550 response
    /// is sent and the recipient is discarded.
    #[allow(unused_variable)]
    fn handle_rcpt(&mut self, transaction: &SmtpTransaction, mailbox: &Mailbox) -> Result<(), ()> {
        Ok(())
    }

    /// Called after getting the entire message for this transaction.
    ///
    /// At this point, the client has not yet been notified of success or failure. If `Ok(())` is
    /// returned, a success message will be sent, else a `554 Transaction failed` message will be
    /// be sent and the transaction will be discarded.
    #[allow(unused_variable)]
    fn handle_transaction(&mut self, transaction: &SmtpTransaction) -> Result<(), ()> {
        Ok(())
    }
}

/// Represents the configuration of an SMTP server.
pub struct SmtpServerConfig {
    /// Maximum number of recipients per SMTP transaction.
    pub max_recipients: uint,
    /// Port on which to listen for incoming messages.
    pub port: u16,
    /// If `true`, debug messages will be printed to the console during transactions.
    pub debug: bool,
    /// The IP on which to `bind (2)` the `TcpListener`.
    pub ip: &'static str,
    /// The domain name used to identify the SMTP server.
    pub domain: &'static str,
    /// The maximum message size, including headers and ending sequence.
    pub max_message_size: uint,
    /// The maximum line size, including `<CRLF>`. At least 1000 per RFC 5321.
    pub max_line_size: uint,
    //pub timeout: uint, // at least 5 minutes
    //pub max_clients: uint, // maximum clients to handle at any given time
    //pub max_pending_clients: uint, // maximum clients to put on hold while handling other clients
}

/// Represents an SMTP server which handles client transactions with any kind of stream.
///
/// This is useful for testing purposes as we can test the server from a plain text file. It
/// should not be used for other purposes directly. Use `SmtpServer` instead.
pub struct SmtpServer<S: 'static+Writer+Reader, A: Acceptor<S>, E: 'static+SmtpServerEventHandler> {
    // Underlying acceptor that allows accepting client connections to handle them.
    acceptor: A,
    config: Arc<SmtpServerConfig>,
    event_handler: E,
    handlers: Arc<Vec<handler::SmtpHandler<S, E>>>
}

/// Represents an error during creation of an SMTP server.
#[deriving(Show)]
pub enum SmtpServerError {
    /// The system call `bind` failed.
    BindFailed(IoError),
    /// The system call `listen` failed.
    ListenFailed(IoError),
}

#[test]
fn test_smtp_server_error() {
    // fail!();
}

impl<E: SmtpServerEventHandler+Clone+Send> SmtpServer<TcpStream, TcpAcceptor, E> {
    /// Creates a new SMTP server that listens on `0.0.0.0:2525`.
    pub fn new(config: SmtpServerConfig, event_handler: E) -> Result<SmtpServer<TcpStream, TcpAcceptor, E>, SmtpServerError> {
        // TODO: Add config checks to force limits to be spec compliant.
        match TcpListener::bind(config.ip, config.port) {
            Ok(listener) => {
                if config.debug {
                    println!("rsmtp: info: binding on ip {}", config.ip);
                }
                match listener.listen() {
                    Ok(acceptor) => {
                        if config.debug {
                            println!("rsmtp: info: listening on port {}", config.port);
                        }
                        Ok(SmtpServer::new_from_acceptor(acceptor, config, event_handler))
                    },
                    Err(err) => Err(ListenFailed(err))
                }
            },
            Err(err) => Err(BindFailed(err))
        }
    }
}

impl<S: Writer+Reader+Send, A: Acceptor<S>, E: SmtpServerEventHandler+Clone+Send> SmtpServer<S, A, E> {
    /// Creates a new SMTP server from an `Acceptor` implementor. Useful for testing.
    pub fn new_from_acceptor(acceptor: A, config: SmtpServerConfig, event_handler: E) -> SmtpServer<S, A, E> {
        SmtpServer {
            acceptor: acceptor,
            config: Arc::new(config),
            event_handler: event_handler,
            handlers: Arc::new(handler::get_handlers::<S, E>())
        }
    }

    /// Run the SMTP server.
    pub fn run(&mut self) {
        for mut stream_res in self.acceptor.incoming() {
            match stream_res {
                Ok(stream) => {
                    let handlers = self.handlers.clone();
                    let config = self.config.clone();
                    let mut event_handler = self.event_handler.clone();
                    spawn(proc() {
                        let mut stream = SmtpStream::new(stream, config.max_message_size, config.max_line_size);
                        // WAIT FOR: https://github.com/rust-lang/rust/issues/15802
                        //stream.stream.set_deadline(local_config.timeout);
                        let mut transaction = SmtpTransaction::new();

                        // Send the opening welcome message.
                        stream.write_line(format!("220 {}", config.domain).as_slice()).unwrap();

                        // Debug arrival of this client.
                        if config.debug {
                            println!("rsmtp: omsg: 220 {}", config.domain);
                        }

                        // Forever, looooop over command lines and handle them.
                        'main_loop: loop {
                            match stream.read_line() {
                                Ok(bytes) => {
                                    let line = String::from_utf8_lossy(bytes.as_slice()).into_string();

                                    if config.debug {
                                        println!("rsmtp: imsg: '{}'", line);
                                    }

                                    // Check if the line is a valid command. If so, do what needs to be done.
                                    for h in handlers.deref().iter() {
                                        // Don't check lines shorter than required. This also avoids getting an
                                        // out of bounds error below.
                                        if line.len() < h.command_start.len() {
                                            continue;
                                        }
                                        let line_start = line.as_slice().slice_to(h.command_start.len())
                                            .into_string().into_ascii_upper();
                                        // Check that the begining of the command matches an existing SMTP
                                        // command. This could be something like "HELO " or "RCPT TO:".
                                        if line_start.as_slice().starts_with(h.command_start.as_slice()) {
                                            if h.allowed_states.contains(&transaction.state) {
                                                let rest = line.as_slice().slice_from(h.command_start.len());
                                                // We're good to go!
                                                (h.callback)(
                                                    &mut stream,
                                                    &mut transaction,
                                                    config.deref(),
                                                    &mut event_handler,
                                                    rest
                                                ).unwrap();
                                                continue 'main_loop;
                                            } else {
                                                // Bad sequence of commands.
                                                stream.write_line("503 Bad sequence of commands").unwrap();
                                                // Debug to console.
                                                if config.debug {
                                                    println!("rsmtp: omsg: 503 Bad sequence of commands");
                                                }
                                                continue 'main_loop;
                                            }
                                        }
                                    }
                                },
                                Err(err) => {
                                    // If the line was too long, notify the client.
                                    match err.kind {
                                        InvalidInput => {
                                            stream.write_line("500 Command line too long, max is 512 bytes").unwrap();
                                            // Debug to console.
                                            if config.debug {
                                                println!("rsmtp: omsg: 500 Command line too long, max is 512 bytes");
                                            }
                                            continue 'main_loop;
                                        },
                                        _ => {
                                            // If we get here, the error is unexpected. What to do with it?
                                            fail!(err);
                                        }
                                    }
                                }
                            }
                            // No valid command was given.
                            stream.write_line("500 Command unrecognized").unwrap();
                            // Debug to console.
                            if config.debug {
                                println!("rsmtp: omsg: 500 Command unrecognized");
                            }
                        }
                    });
                },
                // Ignore accept error. Is this right? If you think not, please open an issue on Github.
                _ => {}
            }
        }
    }
}

#[test]
fn test_smtp_server_new() {
    // fail!();
}

#[test]
fn test_smtp_server_new_from_acceptor() {
    // fail!();
}

#[test]
fn test_smtp_server_handlers() {
    // fail!();
}

#[test]
fn test_smtp_server_run() {
    // fail!();
}
<commit_msg>ideas for new event handler methods<commit_after>// Copyright 2014 The Rustastic SMTP Developers
//
// 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.

//! The `server` module contains things needed to build an SMTP server, but useless for
//! an SMTP client.

use std::io::net::tcp::{TcpListener, TcpAcceptor, TcpStream};
use std::io::{Listener, Acceptor, IoError, Reader, Writer, InvalidInput};
use super::common::stream::{SmtpStream};
use std::sync::Arc;
use std::ascii::OwnedAsciiExt;
use super::common::transaction::SmtpTransaction;
use super::common::mailbox::Mailbox;

mod handler;

/// Hooks into different places of the SMTP server to allow its customization.
///
/// The implementor of this trait you pass to your server is cloned for each
/// new client, which means that you can safely make it have its own fields.
pub trait SmtpServerEventHandler {
    /// Called after getting a HELO command.
    ///
    /// This could be used to check if the sender comes from a banned server,
    /// to log the server information or anything else you desire.
    ///
    /// If `Err(())` is returned, the connection is aborted.
    fn handle_connection(&mut self, server_domain: &str, client_ip: &IpAddr) -> Result<(), ()> {
        Ok(())
    }

    /// Called after getting a MAIL command with a sender address.
    ///
    /// The sender address is either `Some(Mailbox)` or `None`. If it is `None`,
    /// it means that the reverse-path (as described in RFC 5321) was null,
    /// which can happen when an email server sends a delivery failure
    /// notification.
    ///
    /// If `Ok(())` is returned, a 250 response is sent. If `Err(())` is returned, a 550 response
    /// is sent and the sender is discarded.
    fn handle_sender_address(&mut self, mailbox: Option<Mailbox>) -> Result<(), ()> {
        Ok(())
    }

    /// Called after getting a RCPT command.
    ///
    /// If `Ok(())` is returned, a 250 response is sent. If `Err(())` is returned, a 550 response
    /// is sent and the recipient is discarded.
    #[allow(unused_variable)]
    fn handle_receiver_address(&mut self, mailbox: Mailbox) -> Result<(), ()> {
        Ok(())
    }

    /// Called when we know the first body part is coming.
    ///
    /// This could be used to initiate a connection to an HTTP API if that's
    /// where you want to send the body.
    ///
    /// If `Err(())` is returned, the connection is aborted.
    fn handle_body_start(&mut self) -> Result<(), ()> {
        Ok(())
    }

    /// Called after getting a part of the body.
    ///
    /// This can happen in several cases:
    ///   * when reading a line of input after a DATA command
    ///   * when getting a chunck of input after a BDAT command
    ///
    /// This can be used to parse the body on the fly or push it to an HTTP
    /// API or whatever you wish to do.
    ///
    /// If `Err(())` is returned, the connection is aborted.
    fn handle_body_part(&mut self, part: &[u8]) -> Result<(), ()> {
        Ok(())
    }

    /// Called after getting the last body part.
    ///
    /// If you are sending body parts to an HTTP API, this method could be used
    /// to close the HTTP client.
    ///
    /// If `Err(())` is returned, the connection is aborted.
    fn handle_body_end(&mut self) -> Result<(), ()> {
        Ok(())
    }
}

/// Represents the configuration of an SMTP server.
pub struct SmtpServerConfig {
    /// Maximum number of recipients per SMTP transaction.
    pub max_recipients: uint,
    /// Port on which to listen for incoming messages.
    pub port: u16,
    /// If `true`, debug messages will be printed to the console during transactions.
    pub debug: bool,
    /// The IP on which to `bind (2)` the `TcpListener`.
    pub ip: &'static str,
    /// The domain name used to identify the SMTP server.
    pub domain: &'static str,
    /// The maximum message size, including headers and ending sequence.
    pub max_message_size: uint,
    /// The maximum line size, including `<CRLF>`. At least 1000 per RFC 5321.
    pub max_line_size: uint,
    //pub timeout: uint, // at least 5 minutes
    //pub max_clients: uint, // maximum clients to handle at any given time
    //pub max_pending_clients: uint, // maximum clients to put on hold while handling other clients
}

/// Represents an SMTP server which handles client transactions with any kind of stream.
///
/// This is useful for testing purposes as we can test the server from a plain text file. It
/// should not be used for other purposes directly. Use `SmtpServer` instead.
pub struct SmtpServer<S: 'static+Writer+Reader, A: Acceptor<S>, E: 'static+SmtpServerEventHandler> {
    // Underlying acceptor that allows accepting client connections to handle them.
    acceptor: A,
    config: Arc<SmtpServerConfig>,
    event_handler: E,
    handlers: Arc<Vec<handler::SmtpHandler<S, E>>>
}

/// Represents an error during creation of an SMTP server.
#[deriving(Show)]
pub enum SmtpServerError {
    /// The system call `bind` failed.
    BindFailed(IoError),
    /// The system call `listen` failed.
    ListenFailed(IoError),
}

#[test]
fn test_smtp_server_error() {
    // fail!();
}

impl<E: SmtpServerEventHandler+Clone+Send> SmtpServer<TcpStream, TcpAcceptor, E> {
    /// Creates a new SMTP server that listens on `0.0.0.0:2525`.
    pub fn new(config: SmtpServerConfig, event_handler: E) -> Result<SmtpServer<TcpStream, TcpAcceptor, E>, SmtpServerError> {
        // TODO: Add config checks to force limits to be spec compliant.
        match TcpListener::bind(config.ip, config.port) {
            Ok(listener) => {
                if config.debug {
                    println!("rsmtp: info: binding on ip {}", config.ip);
                }
                match listener.listen() {
                    Ok(acceptor) => {
                        if config.debug {
                            println!("rsmtp: info: listening on port {}", config.port);
                        }
                        Ok(SmtpServer::new_from_acceptor(acceptor, config, event_handler))
                    },
                    Err(err) => Err(ListenFailed(err))
                }
            },
            Err(err) => Err(BindFailed(err))
        }
    }
}

impl<S: Writer+Reader+Send, A: Acceptor<S>, E: SmtpServerEventHandler+Clone+Send> SmtpServer<S, A, E> {
    /// Creates a new SMTP server from an `Acceptor` implementor. Useful for testing.
    pub fn new_from_acceptor(acceptor: A, config: SmtpServerConfig, event_handler: E) -> SmtpServer<S, A, E> {
        SmtpServer {
            acceptor: acceptor,
            config: Arc::new(config),
            event_handler: event_handler,
            handlers: Arc::new(handler::get_handlers::<S, E>())
        }
    }

    /// Run the SMTP server.
    pub fn run(&mut self) {
        for mut stream_res in self.acceptor.incoming() {
            match stream_res {
                Ok(stream) => {
                    let handlers = self.handlers.clone();
                    let config = self.config.clone();
                    let mut event_handler = self.event_handler.clone();
                    spawn(proc() {
                        let mut stream = SmtpStream::new(stream, config.max_message_size, config.max_line_size);
                        // WAIT FOR: https://github.com/rust-lang/rust/issues/15802
                        //stream.stream.set_deadline(local_config.timeout);
                        let mut transaction = SmtpTransaction::new();

                        // Send the opening welcome message.
                        stream.write_line(format!("220 {}", config.domain).as_slice()).unwrap();

                        // Debug arrival of this client.
                        if config.debug {
                            println!("rsmtp: omsg: 220 {}", config.domain);
                        }

                        // Forever, looooop over command lines and handle them.
                        'main_loop: loop {
                            match stream.read_line() {
                                Ok(bytes) => {
                                    let line = String::from_utf8_lossy(bytes.as_slice()).into_string();

                                    if config.debug {
                                        println!("rsmtp: imsg: '{}'", line);
                                    }

                                    // Check if the line is a valid command. If so, do what needs to be done.
                                    for h in handlers.deref().iter() {
                                        // Don't check lines shorter than required. This also avoids getting an
                                        // out of bounds error below.
                                        if line.len() < h.command_start.len() {
                                            continue;
                                        }
                                        let line_start = line.as_slice().slice_to(h.command_start.len())
                                            .into_string().into_ascii_upper();
                                        // Check that the begining of the command matches an existing SMTP
                                        // command. This could be something like "HELO " or "RCPT TO:".
                                        if line_start.as_slice().starts_with(h.command_start.as_slice()) {
                                            if h.allowed_states.contains(&transaction.state) {
                                                let rest = line.as_slice().slice_from(h.command_start.len());
                                                // We're good to go!
                                                (h.callback)(
                                                    &mut stream,
                                                    &mut transaction,
                                                    config.deref(),
                                                    &mut event_handler,
                                                    rest
                                                ).unwrap();
                                                continue 'main_loop;
                                            } else {
                                                // Bad sequence of commands.
                                                stream.write_line("503 Bad sequence of commands").unwrap();
                                                // Debug to console.
                                                if config.debug {
                                                    println!("rsmtp: omsg: 503 Bad sequence of commands");
                                                }
                                                continue 'main_loop;
                                            }
                                        }
                                    }
                                },
                                Err(err) => {
                                    // If the line was too long, notify the client.
                                    match err.kind {
                                        InvalidInput => {
                                            stream.write_line("500 Command line too long, max is 512 bytes").unwrap();
                                            // Debug to console.
                                            if config.debug {
                                                println!("rsmtp: omsg: 500 Command line too long, max is 512 bytes");
                                            }
                                            continue 'main_loop;
                                        },
                                        _ => {
                                            // If we get here, the error is unexpected. What to do with it?
                                            fail!(err);
                                        }
                                    }
                                }
                            }
                            // No valid command was given.
                            stream.write_line("500 Command unrecognized").unwrap();
                            // Debug to console.
                            if config.debug {
                                println!("rsmtp: omsg: 500 Command unrecognized");
                            }
                        }
                    });
                },
                // Ignore accept error. Is this right? If you think not, please open an issue on Github.
                _ => {}
            }
        }
    }
}

#[test]
fn test_smtp_server_new() {
    // fail!();
}

#[test]
fn test_smtp_server_new_from_acceptor() {
    // fail!();
}

#[test]
fn test_smtp_server_handlers() {
    // fail!();
}

#[test]
fn test_smtp_server_run() {
    // fail!();
}
<|endoftext|>


<commit_before>//! HTTP Server
use std::fmt;
use std::io::{ErrorKind, BufWriter, Write};
use std::marker::PhantomData;
use std::net::{SocketAddr, ToSocketAddrs};
use std::path::Path;
use std::thread::{self, JoinHandle};

use num_cpus;

pub use self::request::Request;
pub use self::response::Response;

pub use net::{Fresh, Streaming};

use HttpError::HttpIoError;
use {HttpResult};
use buffer::BufReader;
use header::{Headers, Connection, Expect};
use header::ConnectionOption::{Close, KeepAlive};
use method::Method;
use net::{NetworkListener, NetworkStream, HttpListener};
use status::StatusCode;
use uri::RequestUri;
use version::HttpVersion::{Http10, Http11};

use self::listener::ListenerPool;

pub mod request;
pub mod response;

mod listener;

/// A server can listen on a TCP socket.
///
/// Once listening, it will create a `Request`/`Response` pair for each
/// incoming connection, and hand them to the provided handler.
#[derive(Debug)]
pub struct Server<'a, H: Handler, L = HttpListener> {
    handler: H,
    ssl: Option<(&'a Path, &'a Path)>,
    _marker: PhantomData<L>
}

macro_rules! try_option(
    ($e:expr) => {{
        match $e {
            Some(v) => v,
            None => return None
        }
    }}
);

impl<'a, H: Handler, L: NetworkListener> Server<'a, H, L> {
    /// Creates a new server with the provided handler.
    pub fn new(handler: H) -> Server<'a, H, L> {
        Server {
            handler: handler,
            ssl: None,
            _marker: PhantomData
        }
    }
}

impl<'a, H: Handler + 'static> Server<'a, H, HttpListener> {
    /// Creates a new server that will handle `HttpStream`s.
    pub fn http(handler: H) -> Server<'a, H, HttpListener> {
        Server::new(handler)
    }
    /// Creates a new server that will handler `HttpStreams`s using a TLS connection.
    pub fn https(handler: H, cert: &'a Path, key: &'a Path) -> Server<'a, H, HttpListener> {
        Server {
            handler: handler,
            ssl: Some((cert, key)),
            _marker: PhantomData
        }
    }
}

impl<'a, H: Handler + 'static> Server<'a, H, HttpListener> {
    /// Binds to a socket, and starts handling connections using a task pool.
    pub fn listen_threads<T: ToSocketAddrs>(self, addr: T, threads: usize) -> HttpResult<Listening> {
        let listener = try!(match self.ssl {
            Some((cert, key)) => HttpListener::https(addr, cert, key),
            None => HttpListener::http(addr)
        });
        self.with_listener(listener, threads)
    }

    /// Binds to a socket and starts handling connections.
    pub fn listen<T: ToSocketAddrs>(self, addr: T) -> HttpResult<Listening> {
        self.listen_threads(addr, num_cpus::get() * 5 / 4)
    }
}
impl<
'a,
H: Handler + 'static,
L: NetworkListener<Stream=S> + Send + 'static,
S: NetworkStream + Clone + Send> Server<'a, H, L> {
    /// Creates a new server that will handle `HttpStream`s.
    pub fn with_listener(self, mut listener: L, threads: usize) -> HttpResult<Listening> {
        let socket = try!(listener.local_addr());
        let handler = self.handler;

        debug!("threads = {:?}", threads);
        let pool = ListenerPool::new(listener.clone());
        let work = move |mut stream| handle_connection(&mut stream, &handler);

        let guard = thread::spawn(move || pool.accept(work, threads));

        Ok(Listening {
            _guard: Some(guard),
            socket: socket,
        })
    }
}

fn handle_connection<'h, S, H>(mut stream: &mut S, handler: &'h H)
where S: NetworkStream + Clone, H: Handler {
    debug!("Incoming stream");
    let addr = match stream.peer_addr() {
        Ok(addr) => addr,
        Err(e) => {
            error!("Peer Name error: {:?}", e);
            return;
        }
    };

    // FIXME: Use Type ascription
    let stream_clone: &mut NetworkStream = &mut stream.clone();
    let mut rdr = BufReader::new(stream_clone);
    let mut wrt = BufWriter::new(stream);

    let mut keep_alive = true;
    while keep_alive {
        let req = match Request::new(&mut rdr, addr) {
            Ok(req) => req,
            Err(HttpIoError(ref e)) if e.kind() == ErrorKind::ConnectionAborted => {
                trace!("tcp closed, cancelling keep-alive loop");
                break;
            }
            Err(HttpIoError(e)) => {
                debug!("ioerror in keepalive loop = {:?}", e);
                break;
            }
            Err(e) => {
                //TODO: send a 400 response
                error!("request error = {:?}", e);
                break;
            }
        };

        if req.version == Http11 && req.headers.get() == Some(&Expect::Continue) {
            let status = handler.check_continue((&req.method, &req.uri, &req.headers));
            match write!(&mut wrt, "{} {}\r\n\r\n", Http11, status) {
                Ok(..) => (),
                Err(e) => {
                    error!("error writing 100-continue: {:?}", e);
                    break;
                }
            }

            if status != StatusCode::Continue {
                debug!("non-100 status ({}) for Expect 100 request", status);
                break;
            }
        }

        keep_alive = match (req.version, req.headers.get::<Connection>()) {
            (Http10, Some(conn)) if !conn.contains(&KeepAlive) => false,
            (Http11, Some(conn)) if conn.contains(&Close)  => false,
            _ => true
        };
        let mut res = Response::new(&mut wrt);
        res.version = req.version;
        handler.handle(req, res);
        debug!("keep_alive = {:?}", keep_alive);
    }
}

/// A listening server, which can later be closed.
pub struct Listening {
    #[cfg(feature = "nightly")]
    _guard: Option<JoinHandle<()>>,
    #[cfg(not(feature = "nightly"))]
    _guard: Option<JoinHandle>,
    /// The socket addresses that the server is bound to.
    pub socket: SocketAddr,
}

impl fmt::Debug for Listening {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Listening {{ socket: {:?} }}", self.socket)
    }
}

impl Drop for Listening {
    fn drop(&mut self) {
        let _ = self._guard.take().map(|g| g.join());
    }
}

impl Listening {
    /// Stop the server from listening to its socket address.
    pub fn close(&mut self) -> HttpResult<()> {
        let _ = self._guard.take();
        debug!("closing server");
        //try!(self.acceptor.close());
        Ok(())
    }
}

/// A handler that can handle incoming requests for a server.
pub trait Handler: Sync + Send {
    /// Receives a `Request`/`Response` pair, and should perform some action on them.
    ///
    /// This could reading from the request, and writing to the response.
    fn handle<'a, 'k>(&'a self, Request<'a, 'k>, Response<'a, Fresh>);

    /// Called when a Request includes a `Expect: 100-continue` header.
    ///
    /// By default, this will always immediately response with a `StatusCode::Continue`,
    /// but can be overridden with custom behavior.
    fn check_continue(&self, _: (&Method, &RequestUri, &Headers)) -> StatusCode {
        StatusCode::Continue
    }
}

impl<F> Handler for F where F: Fn(Request, Response<Fresh>), F: Sync + Send {
    fn handle<'a, 'k>(&'a self, req: Request<'a, 'k>, res: Response<'a, Fresh>) {
        self(req, res)
    }
}

#[cfg(test)]
mod tests {
    use header::Headers;
    use method::Method;
    use mock::MockStream;
    use status::StatusCode;
    use uri::RequestUri;

    use super::{Request, Response, Fresh, Handler, handle_connection};

    #[test]
    fn test_check_continue_default() {
        let mut mock = MockStream::with_input(b"\
            POST /upload HTTP/1.1\r\n\
            Host: example.domain\r\n\
            Expect: 100-continue\r\n\
            Content-Length: 10\r\n\
            \r\n\
            1234567890\
        ");

        fn handle(_: Request, res: Response<Fresh>) {
            res.start().unwrap().end().unwrap();
        }

        handle_connection(&mut mock, &handle);
        let cont = b"HTTP/1.1 100 Continue\r\n\r\n";
        assert_eq!(&mock.write[..cont.len()], cont);
        let res = b"HTTP/1.1 200 OK\r\n";
        assert_eq!(&mock.write[cont.len()..cont.len() + res.len()], res);
    }

    #[test]
    fn test_check_continue_reject() {
        struct Reject;
        impl Handler for Reject {
            fn handle<'a, 'k>(&'a self, _: Request<'a, 'k>, res: Response<'a, Fresh>) {
                res.start().unwrap().end().unwrap();
            }

            fn check_continue(&self, _: (&Method, &RequestUri, &Headers)) -> StatusCode {
                StatusCode::ExpectationFailed
            }
        }

        let mut mock = MockStream::with_input(b"\
            POST /upload HTTP/1.1\r\n\
            Host: example.domain\r\n\
            Expect: 100-continue\r\n\
            Content-Length: 10\r\n\
            \r\n\
            1234567890\
        ");

        handle_connection(&mut mock, &Reject);
        assert_eq!(mock.write, &b"HTTP/1.1 417 Expectation Failed\r\n\r\n"[..]);
    }
}
<commit_msg>fix(server): JoinHandle type parameter<commit_after>//! HTTP Server
use std::fmt;
use std::io::{ErrorKind, BufWriter, Write};
use std::marker::PhantomData;
use std::net::{SocketAddr, ToSocketAddrs};
use std::path::Path;
use std::thread::{self, JoinHandle};

use num_cpus;

pub use self::request::Request;
pub use self::response::Response;

pub use net::{Fresh, Streaming};

use HttpError::HttpIoError;
use {HttpResult};
use buffer::BufReader;
use header::{Headers, Connection, Expect};
use header::ConnectionOption::{Close, KeepAlive};
use method::Method;
use net::{NetworkListener, NetworkStream, HttpListener};
use status::StatusCode;
use uri::RequestUri;
use version::HttpVersion::{Http10, Http11};

use self::listener::ListenerPool;

pub mod request;
pub mod response;

mod listener;

/// A server can listen on a TCP socket.
///
/// Once listening, it will create a `Request`/`Response` pair for each
/// incoming connection, and hand them to the provided handler.
#[derive(Debug)]
pub struct Server<'a, H: Handler, L = HttpListener> {
    handler: H,
    ssl: Option<(&'a Path, &'a Path)>,
    _marker: PhantomData<L>
}

macro_rules! try_option(
    ($e:expr) => {{
        match $e {
            Some(v) => v,
            None => return None
        }
    }}
);

impl<'a, H: Handler, L: NetworkListener> Server<'a, H, L> {
    /// Creates a new server with the provided handler.
    pub fn new(handler: H) -> Server<'a, H, L> {
        Server {
            handler: handler,
            ssl: None,
            _marker: PhantomData
        }
    }
}

impl<'a, H: Handler + 'static> Server<'a, H, HttpListener> {
    /// Creates a new server that will handle `HttpStream`s.
    pub fn http(handler: H) -> Server<'a, H, HttpListener> {
        Server::new(handler)
    }
    /// Creates a new server that will handler `HttpStreams`s using a TLS connection.
    pub fn https(handler: H, cert: &'a Path, key: &'a Path) -> Server<'a, H, HttpListener> {
        Server {
            handler: handler,
            ssl: Some((cert, key)),
            _marker: PhantomData
        }
    }
}

impl<'a, H: Handler + 'static> Server<'a, H, HttpListener> {
    /// Binds to a socket, and starts handling connections using a task pool.
    pub fn listen_threads<T: ToSocketAddrs>(self, addr: T, threads: usize) -> HttpResult<Listening> {
        let listener = try!(match self.ssl {
            Some((cert, key)) => HttpListener::https(addr, cert, key),
            None => HttpListener::http(addr)
        });
        self.with_listener(listener, threads)
    }

    /// Binds to a socket and starts handling connections.
    pub fn listen<T: ToSocketAddrs>(self, addr: T) -> HttpResult<Listening> {
        self.listen_threads(addr, num_cpus::get() * 5 / 4)
    }
}
impl<
'a,
H: Handler + 'static,
L: NetworkListener<Stream=S> + Send + 'static,
S: NetworkStream + Clone + Send> Server<'a, H, L> {
    /// Creates a new server that will handle `HttpStream`s.
    pub fn with_listener(self, mut listener: L, threads: usize) -> HttpResult<Listening> {
        let socket = try!(listener.local_addr());
        let handler = self.handler;

        debug!("threads = {:?}", threads);
        let pool = ListenerPool::new(listener.clone());
        let work = move |mut stream| handle_connection(&mut stream, &handler);

        let guard = thread::spawn(move || pool.accept(work, threads));

        Ok(Listening {
            _guard: Some(guard),
            socket: socket,
        })
    }
}

fn handle_connection<'h, S, H>(mut stream: &mut S, handler: &'h H)
where S: NetworkStream + Clone, H: Handler {
    debug!("Incoming stream");
    let addr = match stream.peer_addr() {
        Ok(addr) => addr,
        Err(e) => {
            error!("Peer Name error: {:?}", e);
            return;
        }
    };

    // FIXME: Use Type ascription
    let stream_clone: &mut NetworkStream = &mut stream.clone();
    let mut rdr = BufReader::new(stream_clone);
    let mut wrt = BufWriter::new(stream);

    let mut keep_alive = true;
    while keep_alive {
        let req = match Request::new(&mut rdr, addr) {
            Ok(req) => req,
            Err(HttpIoError(ref e)) if e.kind() == ErrorKind::ConnectionAborted => {
                trace!("tcp closed, cancelling keep-alive loop");
                break;
            }
            Err(HttpIoError(e)) => {
                debug!("ioerror in keepalive loop = {:?}", e);
                break;
            }
            Err(e) => {
                //TODO: send a 400 response
                error!("request error = {:?}", e);
                break;
            }
        };

        if req.version == Http11 && req.headers.get() == Some(&Expect::Continue) {
            let status = handler.check_continue((&req.method, &req.uri, &req.headers));
            match write!(&mut wrt, "{} {}\r\n\r\n", Http11, status) {
                Ok(..) => (),
                Err(e) => {
                    error!("error writing 100-continue: {:?}", e);
                    break;
                }
            }

            if status != StatusCode::Continue {
                debug!("non-100 status ({}) for Expect 100 request", status);
                break;
            }
        }

        keep_alive = match (req.version, req.headers.get::<Connection>()) {
            (Http10, Some(conn)) if !conn.contains(&KeepAlive) => false,
            (Http11, Some(conn)) if conn.contains(&Close)  => false,
            _ => true
        };
        let mut res = Response::new(&mut wrt);
        res.version = req.version;
        handler.handle(req, res);
        debug!("keep_alive = {:?}", keep_alive);
    }
}

/// A listening server, which can later be closed.
pub struct Listening {
    _guard: Option<JoinHandle<()>>,
    /// The socket addresses that the server is bound to.
    pub socket: SocketAddr,
}

impl fmt::Debug for Listening {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Listening {{ socket: {:?} }}", self.socket)
    }
}

impl Drop for Listening {
    fn drop(&mut self) {
        let _ = self._guard.take().map(|g| g.join());
    }
}

impl Listening {
    /// Stop the server from listening to its socket address.
    pub fn close(&mut self) -> HttpResult<()> {
        let _ = self._guard.take();
        debug!("closing server");
        //try!(self.acceptor.close());
        Ok(())
    }
}

/// A handler that can handle incoming requests for a server.
pub trait Handler: Sync + Send {
    /// Receives a `Request`/`Response` pair, and should perform some action on them.
    ///
    /// This could reading from the request, and writing to the response.
    fn handle<'a, 'k>(&'a self, Request<'a, 'k>, Response<'a, Fresh>);

    /// Called when a Request includes a `Expect: 100-continue` header.
    ///
    /// By default, this will always immediately response with a `StatusCode::Continue`,
    /// but can be overridden with custom behavior.
    fn check_continue(&self, _: (&Method, &RequestUri, &Headers)) -> StatusCode {
        StatusCode::Continue
    }
}

impl<F> Handler for F where F: Fn(Request, Response<Fresh>), F: Sync + Send {
    fn handle<'a, 'k>(&'a self, req: Request<'a, 'k>, res: Response<'a, Fresh>) {
        self(req, res)
    }
}

#[cfg(test)]
mod tests {
    use header::Headers;
    use method::Method;
    use mock::MockStream;
    use status::StatusCode;
    use uri::RequestUri;

    use super::{Request, Response, Fresh, Handler, handle_connection};

    #[test]
    fn test_check_continue_default() {
        let mut mock = MockStream::with_input(b"\
            POST /upload HTTP/1.1\r\n\
            Host: example.domain\r\n\
            Expect: 100-continue\r\n\
            Content-Length: 10\r\n\
            \r\n\
            1234567890\
        ");

        fn handle(_: Request, res: Response<Fresh>) {
            res.start().unwrap().end().unwrap();
        }

        handle_connection(&mut mock, &handle);
        let cont = b"HTTP/1.1 100 Continue\r\n\r\n";
        assert_eq!(&mock.write[..cont.len()], cont);
        let res = b"HTTP/1.1 200 OK\r\n";
        assert_eq!(&mock.write[cont.len()..cont.len() + res.len()], res);
    }

    #[test]
    fn test_check_continue_reject() {
        struct Reject;
        impl Handler for Reject {
            fn handle<'a, 'k>(&'a self, _: Request<'a, 'k>, res: Response<'a, Fresh>) {
                res.start().unwrap().end().unwrap();
            }

            fn check_continue(&self, _: (&Method, &RequestUri, &Headers)) -> StatusCode {
                StatusCode::ExpectationFailed
            }
        }

        let mut mock = MockStream::with_input(b"\
            POST /upload HTTP/1.1\r\n\
            Host: example.domain\r\n\
            Expect: 100-continue\r\n\
            Content-Length: 10\r\n\
            \r\n\
            1234567890\
        ");

        handle_connection(&mut mock, &Reject);
        assert_eq!(mock.write, &b"HTTP/1.1 417 Expectation Failed\r\n\r\n"[..]);
    }
}
<|endoftext|>


<commit_before>use std::io::{self, Read};
use std::iter::Iterator;
use std::path::{Component, Path, PathBuf};
use std::slice;

use flate2::read::GzDecoder;
use tar;
use time::Timespec;

use backend::Backend;
use collections::{CollectionsStatus, SignatureFile};

pub struct BackupFiles {
    chains: Vec<Chain>
}

pub struct Snapshot {
    pub time: Timespec
}

pub struct File {
    pub name: String,
    pub last_modified: Timespec
}

/// Iterator over a list of backup snapshots.
pub type Snapshots<'a> = slice::Iter<'a, Snapshot>;

enum DiffType {
    Signature,
    Snapshot,
    Deleted
}

/// Store separately informations about the signatures and informations about the paths in the
/// signatures. This allows to reuse informations between snapshots and avoid duplicating them.
struct Chain {
    timestamps: Vec<Timespec>,
    files: Vec<PathSnapshots>
}

struct PathSnapshots {
    // the directory or file path
    path: PathBuf,
    // all the snapshots for this path
    snapshots: Vec<PathSnapshot>
}

struct PathSnapshot {
    // info are None if the snapshot has deleted this path
    info: Option<PathInfo>,
    // the index of the snapshot in the chain
    index: u8
}

/// Informations about a path inside a snapshot.
struct PathInfo {
    mtime: Timespec
}

impl BackupFiles {
    pub fn new<B: Backend>(backend: &B) -> io::Result<BackupFiles> {
        let collection = {
            let filenames = try!(backend.get_file_names());
            CollectionsStatus::from_filenames(&filenames)
        };
        let mut chains: Vec<Chain> = Vec::new();
        let coll_chains = collection.signature_chains();
        for coll_chain in coll_chains {
            // translate collections::SignatureChain into a Chain
            let mut chain = Chain{ timestamps: Vec::new(), files: Vec::new() };
            // add to the chain the full signature and all the incremental signatures
            // if an error occurs in the full signature exit
            let file = try!(backend.open_file(coll_chain.fullsig.file_name.as_ref()));
            try!(add_sigfile_to_chain(&mut chain, file, &coll_chain.fullsig));
            for inc in coll_chain.inclist.iter() {
                // TODO: if an error occurs here, do not exit with an error, instead
                // break the iteration and store the error inside the chain
                let file = try!(backend.open_file(inc.file_name.as_ref()));
                try!(add_sigfile_to_chain(&mut chain, file, &inc));
            }
            chains.push(chain);
        }
        Ok(BackupFiles{ chains: chains })
    }

    pub fn snapshots(&self) -> Snapshots {
        unimplemented!()
    }
}

fn add_sigfile_to_chain<R: Read>(chain: &mut Chain,
                                 file: R,
                                 sigfile: &SignatureFile) -> io::Result<()>
{
    let result = {
        let snapshot_id = chain.files.len() as u8;
        if sigfile.compressed {
            let gz_decoder = try!(GzDecoder::new(file));
            add_sigtar_to_snapshots(&mut chain.files, tar::Archive::new(gz_decoder), snapshot_id)
        }
        else {
            add_sigtar_to_snapshots(&mut chain.files, tar::Archive::new(file), snapshot_id)
        }
    };
    if result.is_ok() {
        // add to the list of signatures only if everything is ok
        // we do not need to cleanup the chain if someting went wrong, because if the list of
        // signatures is not updated, the change is not observable
        chain.timestamps.push(sigfile.time);
    }
    result
}

fn add_sigtar_to_snapshots<R: Read>(snapshots: &mut Vec<PathSnapshots>,
                                    mut tar: tar::Archive<R>,
                                    snapshot_id: u8) -> io::Result<()>
{
    let mut new_files: Vec<PathSnapshots> = Vec::new();
    {
        let mut old_snapshots = snapshots.iter_mut();
        for tarfile in try!(tar.files_mut()) {
            // we can ignore paths with errors
            // the only problem here is that we miss some change in the chain, but it is better
            // than abort the whole signature
            let tarfile = unwrap_or_continue!(tarfile);
            let header = tarfile.header();
            let path = unwrap_or_continue!(header.path());
            let (difftype, path) = unwrap_opt_or_continue!(parse_snapshot_path(&path));
            let info = match difftype {
                DiffType::Signature | DiffType::Snapshot => {
                    let time = Timespec::new(header.mtime().unwrap_or(0 as u64) as i64, 0);
                    Some(PathInfo{ mtime: time })
                }
                _ => None
            };
            let new_snapshot = PathSnapshot{ info: info, index: snapshot_id };
            // find the current path in the old snapshots
            // note: they are ordered
            while let Some(path_snapshots) = old_snapshots.next() {
                if path_snapshots.path.as_path() < path {
                    continue;
                }
                if path_snapshots.path.as_path() == path {
                    // this path is already present in old snapshots: update them
                    path_snapshots.snapshots.push(new_snapshot);
                }
                else {
                    // the path is not present in the old snapshots: add to new list
                    new_files.push(PathSnapshots{
                        path: path.to_path_buf(),
                        snapshots: vec![new_snapshot]
                    });
                }
                break;
            }
        }
    }
    // merge the new files with old snapthots
    if !new_files.is_empty() {
        // TODO: Performance hurt here: we have two sorted arrays to merge,
        // better to use this algorithm: http://stackoverflow.com/a/4553321/1667955
        snapshots.extend(new_files.into_iter());
        snapshots.sort_by(|a, b| a.path.cmp(&b.path));
    }
    Ok(())
}

fn parse_snapshot_path(path: &Path) -> Option<(DiffType, &Path)> {
    // split the path in (first directory, the remaining path)
    // the first is the type, the remaining is the real path
    let mut pcomps = path.components();
    let pfirst = try_opt!(pcomps.next());
    if let Component::Normal(strfirst) = pfirst {
        let difftype = match strfirst.to_str() {
            Some("signature") => DiffType::Signature,
            Some("snapshot")  => DiffType::Snapshot,
            Some("deleted")   => DiffType::Deleted,
            _                 => { return None; }
        };
        let realpath = pcomps.as_path();
        Some((difftype, realpath))
    }
    else {
        None
    }
}

<commit_msg>First implementation of signatures module<commit_after>use std::io::{self, Read};
use std::iter::Iterator;
use std::path::{Component, Path, PathBuf};
use std::slice;

use flate2::read::GzDecoder;
use tar;
use time::Timespec;

use backend::Backend;
use collections::{CollectionsStatus, SignatureFile};

pub struct BackupFiles {
    chains: Vec<Chain>
}

pub struct Snapshot<'a> {
    index: u8,
    chain: &'a Chain
}

pub struct File<'a> {
    pub name: &'a str,
    pub last_modified: Timespec
}

/// Iterator over a list of backup snapshots.
pub struct Snapshots<'a> {
    chain_iter: slice::Iter<'a, Chain>,
    chain: Option<&'a Chain>,
    snapshot_id: u8
}

pub struct SnapshotFiles<'a> {
    index: u8,
    iter: slice::Iter<'a, PathSnapshots>
}

enum DiffType {
    Signature,
    Snapshot,
    Deleted
}

/// Store separately informations about the signatures and informations about the paths in the
/// signatures. This allows to reuse informations between snapshots and avoid duplicating them.
struct Chain {
    timestamps: Vec<Timespec>,
    files: Vec<PathSnapshots>
}

struct PathSnapshots {
    // the directory or file path
    path: PathBuf,
    // all the snapshots for this path
    snapshots: Vec<PathSnapshot>
}

struct PathSnapshot {
    // info are None if the snapshot has deleted this path
    info: Option<PathInfo>,
    // the index of the snapshot in the chain
    index: u8
}

/// Informations about a path inside a snapshot.
struct PathInfo {
    mtime: Timespec
}

impl BackupFiles {
    pub fn new<B: Backend>(backend: &B) -> io::Result<BackupFiles> {
        let collection = {
            let filenames = try!(backend.get_file_names());
            CollectionsStatus::from_filenames(&filenames)
        };
        let mut chains: Vec<Chain> = Vec::new();
        let coll_chains = collection.signature_chains();
        for coll_chain in coll_chains {
            // translate collections::SignatureChain into a Chain
            let mut chain = Chain{ timestamps: Vec::new(), files: Vec::new() };
            // add to the chain the full signature and all the incremental signatures
            // if an error occurs in the full signature exit
            let file = try!(backend.open_file(coll_chain.fullsig.file_name.as_ref()));
            try!(add_sigfile_to_chain(&mut chain, file, &coll_chain.fullsig));
            for inc in coll_chain.inclist.iter() {
                // TODO: if an error occurs here, do not exit with an error, instead
                // break the iteration and store the error inside the chain
                let file = try!(backend.open_file(inc.file_name.as_ref()));
                try!(add_sigfile_to_chain(&mut chain, file, &inc));
            }
            chains.push(chain);
        }
        Ok(BackupFiles{ chains: chains })
    }

    pub fn snapshots(&self) -> Snapshots {
        let mut iter = self.chains.iter();
        let first_chain = iter.next();
        Snapshots{ chain_iter: iter, chain: first_chain, snapshot_id: 0 as u8 }
    }
}

impl<'a> Iterator for Snapshots<'a> {
    type Item = Snapshot<'a>;

    fn next(&mut self) -> Option<Snapshot<'a>> {
        loop {
            if let Some(chain) = self.chain {
                if let Some(_) = chain.timestamps.get(self.snapshot_id as usize) {
                    self.snapshot_id += 1;
                    return Some(Snapshot{ index: self.snapshot_id, chain: chain })
                }
                else {
                    // this chain is completed
                    // go to next chain
                    self.chain = self.chain_iter.next();
                }
            }
            else {
                // no other chains are present
                return None
            }
        }
    }
}

impl<'a> Snapshot<'a> {
    pub fn time(&self) -> Timespec {
        self.chain.timestamps[self.index as usize]
    }

    pub fn files(&self) -> SnapshotFiles<'a> {
        SnapshotFiles{ index: self.index, iter: self.chain.files.iter() }
    }
}

impl<'a> Iterator for SnapshotFiles<'a> {
    type Item = File<'a>;

    fn next(&mut self) -> Option<File<'a>> {
        while let Some(path_snapshots) = self.iter.next() {
            if let Some(s) = path_snapshots.snapshots.iter().rev().find(|s| s.index <= self.index) {
                // now we have a path info present in this snaptshot
                // if it is not deleted return it
                if let Some(ref info) = s.info {
                    return Some(File{ name: path_snapshots.path.to_str().unwrap(), last_modified: info.mtime })
                }
            }
        }
        None
    }
}

fn add_sigfile_to_chain<R: Read>(chain: &mut Chain,
                                 file: R,
                                 sigfile: &SignatureFile) -> io::Result<()>
{
    let result = {
        let snapshot_id = chain.files.len() as u8;
        if sigfile.compressed {
            let gz_decoder = try!(GzDecoder::new(file));
            add_sigtar_to_snapshots(&mut chain.files, tar::Archive::new(gz_decoder), snapshot_id)
        }
        else {
            add_sigtar_to_snapshots(&mut chain.files, tar::Archive::new(file), snapshot_id)
        }
    };
    if result.is_ok() {
        // add to the list of signatures only if everything is ok
        // we do not need to cleanup the chain if someting went wrong, because if the list of
        // signatures is not updated, the change is not observable
        chain.timestamps.push(sigfile.time);
    }
    result
}

fn add_sigtar_to_snapshots<R: Read>(snapshots: &mut Vec<PathSnapshots>,
                                    mut tar: tar::Archive<R>,
                                    snapshot_id: u8) -> io::Result<()>
{
    let mut new_files: Vec<PathSnapshots> = Vec::new();
    {
        let mut old_snapshots = snapshots.iter_mut();
        for tarfile in try!(tar.files_mut()) {
            // we can ignore paths with errors
            // the only problem here is that we miss some change in the chain, but it is better
            // than abort the whole signature
            let tarfile = unwrap_or_continue!(tarfile);
            let header = tarfile.header();
            let path = unwrap_or_continue!(header.path());
            let (difftype, path) = unwrap_opt_or_continue!(parse_snapshot_path(&path));
            let info = match difftype {
                DiffType::Signature | DiffType::Snapshot => {
                    let time = Timespec::new(header.mtime().unwrap_or(0 as u64) as i64, 0);
                    Some(PathInfo{ mtime: time })
                }
                _ => None
            };
            let new_snapshot = PathSnapshot{ info: info, index: snapshot_id };
            // find the current path in the old snapshots
            // note: they are ordered
            while let Some(path_snapshots) = old_snapshots.next() {
                if path_snapshots.path.as_path() < path {
                    continue;
                }
                if path_snapshots.path.as_path() == path {
                    // this path is already present in old snapshots: update them
                    path_snapshots.snapshots.push(new_snapshot);
                }
                else {
                    // the path is not present in the old snapshots: add to new list
                    new_files.push(PathSnapshots{
                        path: path.to_path_buf(),
                        snapshots: vec![new_snapshot]
                    });
                }
                break;
            }
        }
    }
    // merge the new files with old snapthots
    if !new_files.is_empty() {
        // TODO: Performance hurt here: we have two sorted arrays to merge,
        // better to use this algorithm: http://stackoverflow.com/a/4553321/1667955
        snapshots.extend(new_files.into_iter());
        snapshots.sort_by(|a, b| a.path.cmp(&b.path));
    }
    Ok(())
}

fn parse_snapshot_path(path: &Path) -> Option<(DiffType, &Path)> {
    // split the path in (first directory, the remaining path)
    // the first is the type, the remaining is the real path
    let mut pcomps = path.components();
    let pfirst = try_opt!(pcomps.next());
    if let Component::Normal(strfirst) = pfirst {
        let difftype = match strfirst.to_str() {
            Some("signature") => DiffType::Signature,
            Some("snapshot")  => DiffType::Snapshot,
            Some("deleted")   => DiffType::Deleted,
            _                 => { return None; }
        };
        let realpath = pcomps.as_path();
        Some((difftype, realpath))
    }
    else {
        None
    }
}

<|endoftext|>


<commit_before>pub mod encoder;
pub mod decoder;

#[cfg(test)]
mod tests {
    use super::{decoder, encoder};

    #[test]
    fn test_end_of_frame() {
        use std::io::{Read, Write};

        let mut enc = encoder::Encoder::new(Vec::new(), 1).unwrap();
        enc.write_all(b"foo").unwrap();
        let mut compressed = enc.finish().unwrap();

        // Add footer/whatever to underlying storage.
        compressed.push(0);

        // Drain zstd stream until end-of-frame.
        let mut dec = decoder::Decoder::new(&compressed[..]).unwrap();
        let mut buf = Vec::new();
        dec.read_to_end(&mut buf).unwrap();
        assert_eq!(&buf, b"foo");
    }

    #[test]
    fn test_flush() {
        use std::io::Write;

        let buf = Vec::new();
        let mut z = encoder::Encoder::new(buf, 19).unwrap();

        z.write_all(b"hello").unwrap();

        z.flush().unwrap(); // Might corrupt stream
        let buf = z.finish().unwrap();

        let s = ::decode_all(&buf[..]).unwrap();
        let s = ::std::str::from_utf8(&s).unwrap();
        assert_eq!(s, "hello");
    }
}
<commit_msg>Add invalid and incomplete frame tests<commit_after>pub mod encoder;
pub mod decoder;

#[cfg(test)]
mod tests {
    use super::{decoder, encoder};

    #[test]
    fn test_end_of_frame() {
        use std::io::{Read, Write};

        let mut enc = encoder::Encoder::new(Vec::new(), 1).unwrap();
        enc.write_all(b"foo").unwrap();
        let mut compressed = enc.finish().unwrap();

        // Add footer/whatever to underlying storage.
        compressed.push(0);

        // Drain zstd stream until end-of-frame.
        let mut dec = decoder::Decoder::new(&compressed[..]).unwrap();
        let mut buf = Vec::new();
        dec.read_to_end(&mut buf).unwrap();
        assert_eq!(&buf, b"foo");
    }

    #[test]
    fn test_flush() {
        use std::io::Write;

        let buf = Vec::new();
        let mut z = encoder::Encoder::new(buf, 19).unwrap();

        z.write_all(b"hello").unwrap();

        z.flush().unwrap(); // Might corrupt stream
        let buf = z.finish().unwrap();

        let s = ::decode_all(&buf[..]).unwrap();
        let s = ::std::str::from_utf8(&s).unwrap();
        assert_eq!(s, "hello");
    }

    #[test]
    fn test_invalid_frame() {
        use std::io::Read;

        // I really hope this data is invalid...
        let data = &[1u8, 2u8, 3u8, 4u8, 5u8];
        let mut dec = decoder::Decoder::new(&data[..]).unwrap();
        assert!(dec.read_to_end(&mut Vec::new()).is_err());
    }

    #[test]
    fn test_incomplete_frame() {
        use std::io::{Read, Write};

        let mut enc = encoder::Encoder::new(Vec::new(), 1).unwrap();
        enc.write_all(b"This is a regular string").unwrap();
        let mut compressed = enc.finish().unwrap();

        let half_size = compressed.len() - 2;
        compressed.truncate(half_size);

        let mut dec = decoder::Decoder::new(&compressed[..]).unwrap();
        assert!(dec.read_to_end(&mut Vec::new()).is_err());
    }
}
<|endoftext|>