<commit_before>#![feature(test)]

#[macro_use] extern crate adapton ;
extern crate test ;

use std::rc::Rc;
use adapton::adapton_sigs::* ;
use adapton::engine::* ;
use adapton::macros::* ;

/// Direction Graph where each node has exactly two outgoing edges
#[derive(Debug,PartialEq,Eq,Hash,Clone)]
pub struct BiDiGraph {
  id:usize,
  lsucc:Artic<BiDiGraph>,
  rsucc:Artic<BiDiGraph>,
}

pub fn count_forever (g:BiDiGraph) -> usize {
  let cl = eageric!(count_forever, g:force(&g.lsucc));
  let cr = eageric!(count_forever, g:force(&g.rsucc));
  1 + cl.1 + cr.1
}

pub fn count_visit_cases (g:BiDiGraph) -> usize {
  fn visit (g:BiDiGraph,
            visit_flag:bool,
            rec:&(Fn(BiDiGraph)-> usize)) -> usize {
    if visit_flag { 0 } else {
      let gl = force(&g.lsucc);
      let gr = force(&g.rsucc);
      match (g.id == gl.id, g.id == gr.id, gl.id == gr.id) {
        (false, false, false) => 1 + rec(gl) + rec(gr),
        (true,  false, false) => 1 + 0       + rec(gr),
        (false, true,  false) => 1 + rec(gl) + 0,
        (false, false, true)  => 1 + rec(gl) + 0,
        (true,  true,  true)  => 1 + 0       + 0,
        (true,  true,  false) => unreachable!(),
        (true,  false, true)  => unreachable!(),
        (false, true,  true)  => unreachable!(),
      }
    }
  };
  let t = cothunk(prog_pt!(stringify!(visit)),
                  Rc::new(Box::new(visit)), g);
  force(&t)
}

pub fn count_visit (g:BiDiGraph) -> usize {
  fn visit (g:BiDiGraph,
            visit_flag:bool,
            rec:&(Fn(BiDiGraph)-> usize)) -> usize {
    if visit_flag { 0 } else {
      let gl = force(&g.lsucc);
      let gr = force(&g.rsucc);
      1 + rec(gl) + rec(gr)
    }
  };
  let t = cothunk(prog_pt!(stringify!(visit)),
                  Rc::new(Box::new(visit)), g);
  force(&t)
}

pub fn count_visit_trip (g:BiDiGraph) -> usize {
  fn visit (g:BiDiGraph,
            visit_flag:bool, st:Trip,
            rec:&(Fn(BiDiGraph,Trip)-> (usize,Trip))) -> (usize,Trip) {
    if visit_flag { (0,st) } else {
      let gl      = force(&g.lsucc);
      let gr      = force(&g.rsucc);
      let (cl,st) = rec(gl,st);
      let (cr,st) = rec(gr,st);
      (1 + cr + cl, st)
    }
  };
  let t = cothunk2(prog_pt!(stringify!(visit)),
                   Rc::new(Box::new(visit)), g);
  force(&t)
}

pub fn bomb (_x:usize) -> BiDiGraph { panic!("BOMB!") }

// graph_RGB:
//
//      R
//    /  \
//   Y -- G
//
// Creates a directed K3 graph shown above, with nodes {R,G,B}.
// Each node has two successors: The other two nodes in the set.
// We show undirected edges above, due to limitations of ASCII.
//
pub fn graph_rgb () -> BiDiGraph {
  let nr = name_of_str("r");
  let ng = name_of_str("g");
  let ny = name_of_str("y");

  let ag : Artic<BiDiGraph> = thunkic!(ng.clone() =>> bomb, _x:0);
  let ay : Artic<BiDiGraph> = thunkic!(ny.clone() =>> bomb, _x:0);

  let r  = BiDiGraph{id:0, lsucc:ay.clone(), rsucc:ag.clone()};
  let ar = cell(nr, r.clone());

  let g  = BiDiGraph{id:1, lsucc:ar.clone(), rsucc:ay};
  let ag = cell(ng, g);

  let y  = BiDiGraph{id:2, lsucc:ag, rsucc:ar};
  let _  = cell(ny, y);

  return r
}

#[test]
pub fn count_graph() {

}
<commit_msg>minor; formatting<commit_after>#![feature(test)]

#[macro_use] extern crate adapton ;
extern crate test ;

use std::rc::Rc;
use adapton::adapton_sigs::* ;
use adapton::engine::* ;
use adapton::macros::* ;

/// Direction Graph where each node has exactly two outgoing edges
#[derive(Debug,PartialEq,Eq,Hash,Clone)]
pub struct BiDiGraph {
  id:usize,
  lsucc:Artic<BiDiGraph>,
  rsucc:Artic<BiDiGraph>,
}

pub fn count_forever (g:BiDiGraph) -> usize {
  let cl = eageric!(count_forever, g:force(&g.lsucc));
  let cr = eageric!(count_forever, g:force(&g.rsucc));
  1 + cl.1 + cr.1
}

pub fn count_visit_cases (g:BiDiGraph) -> usize {
  fn visit (g:BiDiGraph,
            visit_flag:bool,
            rec:&(Fn(BiDiGraph)-> usize)) -> usize {
    if visit_flag { 0 } else {
      let gl = force(&g.lsucc);
      let gr = force(&g.rsucc);
      match (g.id == gl.id, g.id == gr.id, gl.id == gr.id) {
        (false, false, false) => 1 + rec(gl) + rec(gr),
        (true,  false, false) => 1 + 0       + rec(gr),
        (false, true,  false) => 1 + rec(gl) + 0,
        (false, false, true)  => 1 + rec(gl) + 0,
        (true,  true,  true)  => 1 + 0       + 0,
        (true,  true,  false) => unreachable!(),
        (true,  false, true)  => unreachable!(),
        (false, true,  true)  => unreachable!(),
      }
    }
  };
  let t = cothunk(prog_pt!(stringify!(visit)),
                  Rc::new(Box::new(visit)), g);
  force(&t)
}

pub fn count_visit (g:BiDiGraph) -> usize {
  fn visit (g:BiDiGraph,
            visit_flag:bool,
            rec:&(Fn(BiDiGraph)-> usize)) -> usize {
    if visit_flag { 0 } else {
      let gl = force(&g.lsucc);
      let gr = force(&g.rsucc);
      1 + rec(gl) + rec(gr)
    }
  };
  let t = cothunk(prog_pt!(stringify!(visit)),
                  Rc::new(Box::new(visit)), g);
  force(&t)
}

pub fn count_visit_trip (g:BiDiGraph) -> usize {
  fn visit (g:BiDiGraph,
            visit_flag:bool, st:Trip,
            rec:&(Fn(BiDiGraph,Trip)-> (usize,Trip))) -> (usize,Trip) {
    if visit_flag {
      (0, st)
    } else {
      let gl      = force(&g.lsucc);
      let gr      = force(&g.rsucc);
      let (cl,st) = rec(gl,st);
      let (cr,st) = rec(gr,st);
      (1 + cr + cl, st)
    }
  };
  let t = cothunk2(prog_pt!(stringify!(visit)),
                   Rc::new(Box::new(visit)), g);
  force(&t)
}

pub fn bomb (_x:usize) -> BiDiGraph { panic!("BOMB!") }

// graph_RGB:
//
//      R
//    /  \
//   Y -- G
//
// Creates a directed K3 graph shown above, with nodes {R,G,B}.
// Each node has two successors: The other two nodes in the set.
// We show undirected edges above, due to limitations of ASCII.
//
pub fn graph_rgb () -> BiDiGraph {
  let nr = name_of_str("r");
  let ng = name_of_str("g");
  let ny = name_of_str("y");
  let ag : Artic<BiDiGraph> = thunkic!(ng.clone() =>> bomb, _x:0);
  let ay : Artic<BiDiGraph> = thunkic!(ny.clone() =>> bomb, _x:0);
  let r  = BiDiGraph{id:0, lsucc:ay.clone(), rsucc:ag.clone()};
  let ar = cell(nr, r.clone());
  let g  = BiDiGraph{id:1, lsucc:ar.clone(), rsucc:ay};
  let ag = cell(ng, g);
  let y  = BiDiGraph{id:2, lsucc:ag, rsucc:ar};
  let _  = cell(ny, y);
  return r
}

#[test]
pub fn count_graph() {

}
<|endoftext|>


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

// This defines a base target-configuration for native UEFI systems. The UEFI specification has
// quite detailed sections on the ABI of all the supported target architectures. In almost all
// cases it simply follows what Microsoft Windows does. Hence, whenever in doubt, see the MSDN
// documentation.
// UEFI uses COFF/PE32+ format for binaries. All binaries must be statically linked. No dynamic
// linker is supported. As native to COFF, binaries are position-dependent, but will be relocated
// by the loader if the pre-chosen memory location is already in use.
// UEFI forbids running code on anything but the boot-CPU. Not interrupts are allowed other than
// the timer-interrupt. Device-drivers are required to use polling-based models. Furthermore, all
// code runs in the same environment, no process separation is supported.

use spec::{LinkArgs, LinkerFlavor, LldFlavor, PanicStrategy, TargetOptions};
use std::default::Default;

pub fn opts() -> TargetOptions {
    let mut pre_link_args = LinkArgs::new();

    pre_link_args.insert(LinkerFlavor::Lld(LldFlavor::Link), vec![
            // Suppress the verbose logo and authorship debugging output, which would needlessly
            // clog any log files.
            "/NOLOGO".to_string(),

            // UEFI is fully compatible to non-executable data pages. Tell the compiler that
            // non-code sections can be marked as non-executable, including stack pages.
            "/NXCOMPAT".to_string(),

            // There is no runtime for UEFI targets, prevent them from being linked. UEFI targets
            // must be freestanding.
            "/nodefaultlib".to_string(),

            // Non-standard subsystems have no default entry-point in PE+ files. We have to define
            // one. "efi_main" seems to be a common choice amongst other implementations and the
            // spec.
            "/entry:efi_main".to_string(),

            // COFF images have a "Subsystem" field in their header, which defines what kind of
            // program it is. UEFI has 3 fields reserved, which are EFI_APPLICATION,
            // EFI_BOOT_SERVICE_DRIVER, and EFI_RUNTIME_DRIVER. We default to EFI_APPLICATION,
            // which is very likely the most common option. Individual projects can override this
            // with custom linker flags.
            // The subsystem-type only has minor effects on the application. It defines the memory
            // regions the application is loaded into (runtime-drivers need to be put into
            // reserved areas), as well as whether a return from the entry-point is treated as
            // exit (default for applications).
            "/subsystem:efi_application".to_string(),
        ]);

    TargetOptions {
        dynamic_linking: false,
        executables: true,
        disable_redzone: true,
        exe_suffix: ".efi".to_string(),
        allows_weak_linkage: false,
        panic_strategy: PanicStrategy::Abort,
        singlethread: true,
        emit_debug_gdb_scripts: false,

        linker: Some("lld-link".to_string()),
        lld_flavor: LldFlavor::Link,
        pre_link_args,

        .. Default::default()
    }
}
<commit_msg>Rollup merge of #56913 - tirr-c:uefi-stack-probes, r=alexcrichton<commit_after>// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// This defines a base target-configuration for native UEFI systems. The UEFI specification has
// quite detailed sections on the ABI of all the supported target architectures. In almost all
// cases it simply follows what Microsoft Windows does. Hence, whenever in doubt, see the MSDN
// documentation.
// UEFI uses COFF/PE32+ format for binaries. All binaries must be statically linked. No dynamic
// linker is supported. As native to COFF, binaries are position-dependent, but will be relocated
// by the loader if the pre-chosen memory location is already in use.
// UEFI forbids running code on anything but the boot-CPU. Not interrupts are allowed other than
// the timer-interrupt. Device-drivers are required to use polling-based models. Furthermore, all
// code runs in the same environment, no process separation is supported.

use spec::{LinkArgs, LinkerFlavor, LldFlavor, PanicStrategy, TargetOptions};
use std::default::Default;

pub fn opts() -> TargetOptions {
    let mut pre_link_args = LinkArgs::new();

    pre_link_args.insert(LinkerFlavor::Lld(LldFlavor::Link), vec![
            // Suppress the verbose logo and authorship debugging output, which would needlessly
            // clog any log files.
            "/NOLOGO".to_string(),

            // UEFI is fully compatible to non-executable data pages. Tell the compiler that
            // non-code sections can be marked as non-executable, including stack pages.
            "/NXCOMPAT".to_string(),

            // There is no runtime for UEFI targets, prevent them from being linked. UEFI targets
            // must be freestanding.
            "/nodefaultlib".to_string(),

            // Non-standard subsystems have no default entry-point in PE+ files. We have to define
            // one. "efi_main" seems to be a common choice amongst other implementations and the
            // spec.
            "/entry:efi_main".to_string(),

            // COFF images have a "Subsystem" field in their header, which defines what kind of
            // program it is. UEFI has 3 fields reserved, which are EFI_APPLICATION,
            // EFI_BOOT_SERVICE_DRIVER, and EFI_RUNTIME_DRIVER. We default to EFI_APPLICATION,
            // which is very likely the most common option. Individual projects can override this
            // with custom linker flags.
            // The subsystem-type only has minor effects on the application. It defines the memory
            // regions the application is loaded into (runtime-drivers need to be put into
            // reserved areas), as well as whether a return from the entry-point is treated as
            // exit (default for applications).
            "/subsystem:efi_application".to_string(),
        ]);

    TargetOptions {
        dynamic_linking: false,
        executables: true,
        disable_redzone: true,
        exe_suffix: ".efi".to_string(),
        allows_weak_linkage: false,
        panic_strategy: PanicStrategy::Abort,
        stack_probes: true,
        singlethread: true,
        emit_debug_gdb_scripts: false,

        linker: Some("lld-link".to_string()),
        lld_flavor: LldFlavor::Link,
        pre_link_args,

        .. Default::default()
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate askama;

use askama::Template;

#[derive(Template)]
#[template(path = "simple.html")]
struct VariablesTemplate<'a> {
    strvar: &'a str,
    num: i64,
    i18n: String,
}

#[test]
fn test_variables() {
    let s = VariablesTemplate {
        strvar: "foo",
        num: 42,
        i18n: "Iñtërnâtiônàlizætiøn".to_string(),
    };
    assert_eq!(s.render().unwrap(), "\nhello world, foo\n\
                                     with number: 42\n\
                                     Iñtërnâtiônàlizætiøn is important\n\
                                     in vars too: Iñtërnâtiônàlizætiøn");
}

#[derive(Template)]
#[template(path = "if.html")]
struct IfTemplate {
    cond: bool,
}

#[test]
fn test_if() {
    let s = IfTemplate { cond: true };
    assert_eq!(s.render().unwrap(), "true");
}

#[derive(Template)]
#[template(path = "else.html")]
struct ElseTemplate {
    cond: bool,
}

#[test]
fn test_else() {
    let s = ElseTemplate { cond: false };
    assert_eq!(s.render().unwrap(), "false");
}

#[derive(Template)]
#[template(path = "else-if.html")]
struct ElseIfTemplate {
    cond: bool,
    check: bool,
}

#[test]
fn test_else_if() {
    let s = ElseIfTemplate { cond: false, check: true };
    assert_eq!(s.render().unwrap(), "checked");
}

#[derive(Template)]
#[template(path = "literals.html")]
struct LiteralsTemplate {}

#[test]
fn test_literals() {
    let s = LiteralsTemplate {};
    assert_eq!(s.render().unwrap(), "a");
}

struct Holder {
    a: usize,
}

#[derive(Template)]
#[template(path = "attr.html")]
struct AttrTemplate {
    inner: Holder,
}

#[test]
fn test_attr() {
    let t = AttrTemplate { inner: Holder { a: 5 } };
    assert_eq!(t.render().unwrap(), "5");
}

#[derive(Template)]
#[template(path = "tuple-attr.html")]
struct TupleAttrTemplate<'a> {
    tuple: (&'a str, &'a str),
}

#[test]
fn test_tuple_attr() {
    let t = TupleAttrTemplate { tuple: ("foo", "bar") };
    assert_eq!(t.render().unwrap(), "foobar");
}

struct NestedHolder {
    holder: Holder,
}

#[derive(Template)]
#[template(path = "nested-attr.html")]
struct NestedAttrTemplate {
    inner: NestedHolder,
}

#[test]
fn test_nested_attr() {
    let t = NestedAttrTemplate {
        inner: NestedHolder { holder: Holder { a: 5 } }
    };
    assert_eq!(t.render().unwrap(), "5");
}

#[derive(Template)]
#[template(path = "option.html")]
struct OptionTemplate<'a> {
    var: Option<&'a str>,
}

#[test]
fn test_option() {
    let some = OptionTemplate { var: Some("foo") };
    assert_eq!(some.render().unwrap(), "some: foo");
    let none = OptionTemplate { var: None };
    assert_eq!(none.render().unwrap(), "none");
}

#[derive(Template)]
#[template(path = "generics.html")]
struct GenericsTemplate<T: std::fmt::Display, U = u8>
    where U: std::fmt::Display {
    t: T,
    u: U,
}

#[test]
fn test_generics() {
    let t = GenericsTemplate { t: "a", u: 42 };
    assert_eq!(t.render().unwrap(), "a42");
}

#[derive(Template)]
#[template(path = "composition.html")]
struct CompositionTemplate {
    foo: IfTemplate,
}

#[test]
fn test_composition() {
    let t = CompositionTemplate { foo: IfTemplate { cond: true } };
    assert_eq!(t.render().unwrap(), "composed: true");
}
<commit_msg>Add test for path expressions (see #56)<commit_after>#[macro_use]
extern crate askama;

use askama::Template;

#[derive(Template)]
#[template(path = "simple.html")]
struct VariablesTemplate<'a> {
    strvar: &'a str,
    num: i64,
    i18n: String,
}

#[test]
fn test_variables() {
    let s = VariablesTemplate {
        strvar: "foo",
        num: 42,
        i18n: "Iñtërnâtiônàlizætiøn".to_string(),
    };
    assert_eq!(s.render().unwrap(), "\nhello world, foo\n\
                                     with number: 42\n\
                                     Iñtërnâtiônàlizætiøn is important\n\
                                     in vars too: Iñtërnâtiônàlizætiøn");
}

#[derive(Template)]
#[template(path = "if.html")]
struct IfTemplate {
    cond: bool,
}

#[test]
fn test_if() {
    let s = IfTemplate { cond: true };
    assert_eq!(s.render().unwrap(), "true");
}

#[derive(Template)]
#[template(path = "else.html")]
struct ElseTemplate {
    cond: bool,
}

#[test]
fn test_else() {
    let s = ElseTemplate { cond: false };
    assert_eq!(s.render().unwrap(), "false");
}

#[derive(Template)]
#[template(path = "else-if.html")]
struct ElseIfTemplate {
    cond: bool,
    check: bool,
}

#[test]
fn test_else_if() {
    let s = ElseIfTemplate { cond: false, check: true };
    assert_eq!(s.render().unwrap(), "checked");
}

#[derive(Template)]
#[template(path = "literals.html")]
struct LiteralsTemplate {}

#[test]
fn test_literals() {
    let s = LiteralsTemplate {};
    assert_eq!(s.render().unwrap(), "a");
}

struct Holder {
    a: usize,
}

#[derive(Template)]
#[template(path = "attr.html")]
struct AttrTemplate {
    inner: Holder,
}

#[test]
fn test_attr() {
    let t = AttrTemplate { inner: Holder { a: 5 } };
    assert_eq!(t.render().unwrap(), "5");
}

#[derive(Template)]
#[template(path = "tuple-attr.html")]
struct TupleAttrTemplate<'a> {
    tuple: (&'a str, &'a str),
}

#[test]
fn test_tuple_attr() {
    let t = TupleAttrTemplate { tuple: ("foo", "bar") };
    assert_eq!(t.render().unwrap(), "foobar");
}

struct NestedHolder {
    holder: Holder,
}

#[derive(Template)]
#[template(path = "nested-attr.html")]
struct NestedAttrTemplate {
    inner: NestedHolder,
}

#[test]
fn test_nested_attr() {
    let t = NestedAttrTemplate {
        inner: NestedHolder { holder: Holder { a: 5 } }
    };
    assert_eq!(t.render().unwrap(), "5");
}

#[derive(Template)]
#[template(path = "option.html")]
struct OptionTemplate<'a> {
    var: Option<&'a str>,
}

#[test]
fn test_option() {
    let some = OptionTemplate { var: Some("foo") };
    assert_eq!(some.render().unwrap(), "some: foo");
    let none = OptionTemplate { var: None };
    assert_eq!(none.render().unwrap(), "none");
}

#[derive(Template)]
#[template(path = "generics.html")]
struct GenericsTemplate<T: std::fmt::Display, U = u8>
    where U: std::fmt::Display {
    t: T,
    u: U,
}

#[test]
fn test_generics() {
    let t = GenericsTemplate { t: "a", u: 42 };
    assert_eq!(t.render().unwrap(), "a42");
}

#[derive(Template)]
#[template(path = "composition.html")]
struct CompositionTemplate {
    foo: IfTemplate,
}

#[test]
fn test_composition() {
    let t = CompositionTemplate { foo: IfTemplate { cond: true } };
    assert_eq!(t.render().unwrap(), "composed: true");
}

#[derive(PartialEq, Eq)]
enum Alphabet {
    Alpha,
}

#[derive(Template)]
#[template(source = "{% if x == Alphabet::Alpha %}true{% endif %}", ext = "txt")]
struct PathCompareTemplate {
    x: Alphabet,
}

#[test]
fn test_path_compare() {
    let t = PathCompareTemplate { x: Alphabet::Alpha };
    assert_eq!(t.render().unwrap(), "true");
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::fmt::Display;
use std::fmt::Formatter;
use std::fmt::Error as FmtError;
use std::clone::Clone;

#[derive(Clone, Copy, Debug)]
pub enum StoreErrorKind {
    FileError,
    IdLocked,
    IdNotFound,
    OutOfMemory,
    FileNotFound,
    FileNotCreated,
    // maybe more
}

fn store_error_type_as_str(e: &StoreErrorKind) -> &'static str {
    match e {
        &StoreErrorKind::FileError   => "File Error",
        &StoreErrorKind::IdLocked    => "ID locked",
        &StoreErrorKind::IdNotFound  => "ID not found",
        &StoreErrorKind::OutOfMemory => "Out of Memory",
        &StoreErrorKind::FileNotFound => "File corresponding to ID not found",
        &StoreErrorKind::FileNotCreated => "File corresponding to ID could not be created",
    }
}

impl Display for StoreErrorKind {

    fn fmt(&self, fmt: &mut Formatter) -> Result<(), FmtError> {
        try!(write!(fmt, "{}", store_error_type_as_str(self)));
        Ok(())
    }

}

#[derive(Debug)]
pub struct StoreError {
    err_type: StoreErrorKind,
    cause: Option<Box<Error>>,
}

impl StoreError {

    pub fn new(errtype: StoreErrorKind, cause: Option<Box<Error>>)
        -> StoreError
    {
        StoreError {
            err_type: errtype,
            cause: cause,
        }
    }

    pub fn err_type(&self) -> StoreErrorKind {
        self.err_type.clone()
    }

}

impl Display for StoreError {

    fn fmt(&self, fmt: &mut Formatter) -> Result<(), FmtError> {
        try!(write!(fmt, "[{}]", store_error_type_as_str(&self.err_type.clone())));
        Ok(())
    }

}

impl Error for StoreError {

    fn description(&self) -> &str {
        store_error_type_as_str(&self.err_type.clone())
    }

    fn cause(&self) -> Option<&Error> {
        self.cause.as_ref().map(|e| &**e)
    }

}

<commit_msg>error.rs: Add documentation<commit_after>use std::error::Error;
use std::fmt::Display;
use std::fmt::Formatter;
use std::fmt::Error as FmtError;
use std::clone::Clone;

/**
 * Kind of store error
 */
#[derive(Clone, Copy, Debug)]
pub enum StoreErrorKind {
    FileError,
    IdLocked,
    IdNotFound,
    OutOfMemory,
    FileNotFound,
    FileNotCreated,
    // maybe more
}

fn store_error_type_as_str(e: &StoreErrorKind) -> &'static str {
    match e {
        &StoreErrorKind::FileError   => "File Error",
        &StoreErrorKind::IdLocked    => "ID locked",
        &StoreErrorKind::IdNotFound  => "ID not found",
        &StoreErrorKind::OutOfMemory => "Out of Memory",
        &StoreErrorKind::FileNotFound => "File corresponding to ID not found",
        &StoreErrorKind::FileNotCreated => "File corresponding to ID could not be created",
    }
}

impl Display for StoreErrorKind {

    fn fmt(&self, fmt: &mut Formatter) -> Result<(), FmtError> {
        try!(write!(fmt, "{}", store_error_type_as_str(self)));
        Ok(())
    }

}

/**
 * Store error type
 */
#[derive(Debug)]
pub struct StoreError {
    err_type: StoreErrorKind,
    cause: Option<Box<Error>>,
}

impl StoreError {

    /**
     * Build a new StoreError from an StoreErrorKind, optionally with cause
     */
    pub fn new(errtype: StoreErrorKind, cause: Option<Box<Error>>)
        -> StoreError
    {
        StoreError {
            err_type: errtype,
            cause: cause,
        }
    }

    /**
     * Get the error type of this StoreError
     */
    pub fn err_type(&self) -> StoreErrorKind {
        self.err_type.clone()
    }

}

impl Display for StoreError {

    fn fmt(&self, fmt: &mut Formatter) -> Result<(), FmtError> {
        try!(write!(fmt, "[{}]", store_error_type_as_str(&self.err_type.clone())));
        Ok(())
    }

}

impl Error for StoreError {

    fn description(&self) -> &str {
        store_error_type_as_str(&self.err_type.clone())
    }

    fn cause(&self) -> Option<&Error> {
        self.cause.as_ref().map(|e| &**e)
    }

}

<|endoftext|>


<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
use super::model;
use regex;
use rs_es;
use rs_es::query::Query as rs_q;
use rs_es::operations::search::SearchResult;
use rs_es::units as rs_u;
use mimir;
use serde_json;
use serde;
use mimir::objects::{MimirObject, Admin, Addr, Stop};

fn build_rs_client(cnx: &String) -> rs_es::Client {
    let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)").unwrap();
    let cap = re.captures(&cnx).unwrap();
    let host = cap.name("host").unwrap().as_str();
    let port = cap.name("port").unwrap().as_str().parse::<u32>().unwrap();

    rs_es::Client::new(&host, port)
}

/// takes a ES json blob and build a Place from it
/// it uses the _type field of ES to know which type of the Place enum to fill
pub fn make_place(doc_type: String, value: Option<Box<serde_json::Value>>) -> Option<mimir::Place> {
    value.and_then(|v| {
        fn convert<T: serde::Deserialize>(v: serde_json::Value,
                                          f: fn(T) -> mimir::Place)
                                          -> Option<mimir::Place> {
            serde_json::from_value::<T>(v)
                .map_err(|err| warn!("Impossible to load ES result: {}", err))
                .ok()
                .map(f)
        }
        match doc_type.as_ref() {
            "addr" => convert(*v, mimir::Place::Addr),
            "street" => convert(*v, mimir::Place::Street),
            "admin" => convert(*v, mimir::Place::Admin),
            "poi" => convert(*v, mimir::Place::Poi),
            "stop" => convert(*v, mimir::Place::Stop),
            _ => {
                warn!("unknown ES return value, _type field = {}", doc_type);
                None
            }
        }
    })
}

#[derive(Debug, Eq, PartialEq)]
enum MatchType {
    Prefix,
    Fuzzy,
}

fn build_query(q: &str,
               match_type: MatchType,
               coord: &Option<model::Coord>,
               shape: Option<Vec<rs_es::units::Location>>)
               -> rs_es::query::Query {
    use rs_es::query::functions::Function;
    // we order the type of object we want
    // Note: the addresses are boosted more because even if we don't want them first
    // because they are more severely filtered
    let boost_addr = rs_q::build_term("_type", Addr::doc_type()).with_boost(5000).build();
    let boost_admin = rs_q::build_term("_type", Admin::doc_type()).with_boost(3000).build();
    let boost_stop = rs_q::build_term("_type", Stop::doc_type()).with_boost(2000).build();

    let main_match_type = match match_type {
        MatchType::Prefix => "label.prefix",
        MatchType::Fuzzy => "label.ngram",
    };

    let boost_main_match_query = rs_q::build_multi_match(vec![main_match_type.to_string()],
                                                         q.to_string())
        .with_boost(100)
        .build();
    let boost_zipcode_match_query = rs_q::build_multi_match(vec!["zip_codes.prefix".to_string()],
                                                            q.to_string())
        .with_boost(500)
        .build();

    let mut should_query = vec![boost_addr,
                                boost_admin,
                                boost_stop,
                                boost_main_match_query,
                                boost_zipcode_match_query];

    // for fuzzy search we also search by the prefix index (with a greater boost than ngram)
    // to have better results
    if match_type == MatchType::Fuzzy {
        should_query.push(rs_q::build_match("label.prefix", q.to_string())
            .with_boost(1000)
            .build());
    }

    if let &Some(ref c) = coord {
        // if we have coordinate, we boost we result near this coordinate
        let boost_on_proximity =
            rs_q::build_function_score()
                .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
                .with_boost(500)
                .with_function(Function::build_decay("coord",
                                           rs_u::Location::LatLon(c.lat, c.lon),
                                           rs_u::Distance::new(50f64,
                                                               rs_u::DistanceUnit::Kilometer))
                                   .build_exp())
                .build();
        should_query.push(boost_on_proximity);
    } else {
        // if we don't have coords, we take the field `weight` into account
        let boost_on_weight = rs_q::build_function_score()
            .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
            .with_boost(300)
            .with_query(rs_q::build_match_all().build())
            .with_function(Function::build_field_value_factor("weight")
                .with_factor(1)
                .with_modifier(rs_es::query::functions::Modifier::Log1p)
                .build())
            .build();
        should_query.push(boost_on_weight);
    }

    let sub_query = rs_q::build_bool()
        .with_should(should_query)
        .build();

    use rs_es::query::{CombinationMinimumShouldMatch, MinimumShouldMatch};
    let min_match = |min, threshold| {
        CombinationMinimumShouldMatch::new(MinimumShouldMatch::from(min),
                                           MinimumShouldMatch::from(threshold))
    };

    let minimum_should_match = match match_type {
        MatchType::Prefix => {
            // when we search [0, 2] fields, we need to match 100% of them
            // when we search [3, 4] fields, we need to match 90% of them
            // when we search [5, +] fields, we only need to match 70% on them
            // this way if the input has more fields than the documents
            // (additional information like the country, the region, ...)
            // they can be ignored
            MinimumShouldMatch::from(vec![min_match(2, 90f64), min_match(4, 70f64)])
        }
        // for fuzzy search we lower our expectation and we accept 50% of token match
        MatchType::Fuzzy => MinimumShouldMatch::from(50f64),
    };

    let mut must = vec![rs_q::build_multi_match(vec![main_match_type.to_string(),
                                                     "zip_codes.prefix".to_string()],
                                                q.to_string())
                            .with_minimum_should_match(minimum_should_match)
                            .build()];

    if let Some(s) = shape {
        must.push(rs_q::build_geo_polygon("coord", s).build());
    }

    // filter to handle house number
    // we either want:
    // * to exactly match the document house_number
    // * or that the document has no house_number
    let filter = rs_q::build_bool()
        .with_should(vec![rs_q::build_bool()
                              .with_must_not(rs_q::build_exists("house_number").build())
                              .build(),
                          rs_q::build_match("house_number", q.to_string()).build()])
        .with_must(must)
        .build();

    rs_q::build_bool()
        .with_must(vec![sub_query])
        .with_filter(filter)
        .build()
}

fn is_existing_index(client: &mut rs_es::Client, index: &String) -> bool {
    !index.is_empty() && client.open_index(&index).is_ok()
}

fn make_indexes<'a>(all_data: bool,
                    pt_dataset_index: &'a Option<String>,
                    client: &mut rs_es::Client)
                    -> Vec<&'a str> {
	let mut result: Vec<&str> = vec![];
    if all_data {
        result.push("munin");
    } else {
        if is_existing_index(client, &"munin_geo_data".to_string()) {
            result.push("munin_geo_data");
        }
        match *pt_dataset_index {
            Some(ref dataset) => {
                if is_existing_index(client, dataset) {
                    result.push(dataset);
                }
            },
            _ => {}
        };
    }
	return result;
}

fn query(q: &str,
         pt_dataset: &Option<&str>,
         all_data: bool,
         cnx: &str,
         match_type: MatchType,
         offset: u64,
         limit: u64,
         coord: &Option<model::Coord>,
         shape: Option<Vec<rs_es::units::Location>>)
         -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    let query = build_query(q, match_type, coord, shape);

    let mut client = build_rs_client(&cnx.to_string());

    let pt_dataset_index = pt_dataset.map(|d| format!("munin_stop_{}", d));
    let indexes = make_indexes(all_data, &pt_dataset_index, &mut client);

    let result: SearchResult<serde_json::Value> = try!(client.search_query()
        .with_indexes(&indexes)
        .with_query(&query)
        .with_from(offset)
        .with_size(limit)
        .send());

    debug!("{} documents found", result.hits.total);

    // for the moment rs-es does not handle enum Document,
    // so we need to convert the ES glob to a Place
    Ok(result.hits
        .hits
        .into_iter()
        .filter_map(|hit| make_place(hit.doc_type, hit.source))
        .collect())
}

pub fn autocomplete(q: &str,
                    pt_dataset: &Option<&str>,
                    all_data: bool,
                    offset: u64,
                    limit: u64,
                    coord: Option<model::Coord>,
                    cnx: &str,
                    shape: Option<Vec<(f64, f64)>>)
                    -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    fn make_shape(shape: &Option<Vec<(f64, f64)>>) -> Option<Vec<rs_es::units::Location>> {
        shape.as_ref().map(|v| v.iter().map(|&l| l.into()).collect())
    }

    // First we try a prety exact match on the prefix.
    // If there are no results then we do a new fuzzy search (matching ngrams)
    let results = try!(query(&q,
                             &pt_dataset,
                             all_data,
                             cnx,
                             MatchType::Prefix,
                             offset,
                             limit,
                             &coord,
                             make_shape(&shape)));
    if results.is_empty() {
        query(&q,
              &pt_dataset,
              all_data,
              cnx,
              MatchType::Fuzzy,
              offset,
              limit,
              &coord,
              make_shape(&shape))
    } else {
        Ok(results)
    }
}
<commit_msg>fmt passed<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io
use super::model;
use regex;
use rs_es;
use rs_es::query::Query as rs_q;
use rs_es::operations::search::SearchResult;
use rs_es::units as rs_u;
use mimir;
use serde_json;
use serde;
use mimir::objects::{MimirObject, Admin, Addr, Stop};

fn build_rs_client(cnx: &String) -> rs_es::Client {
    let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)").unwrap();
    let cap = re.captures(&cnx).unwrap();
    let host = cap.name("host").unwrap().as_str();
    let port = cap.name("port").unwrap().as_str().parse::<u32>().unwrap();

    rs_es::Client::new(&host, port)
}

/// takes a ES json blob and build a Place from it
/// it uses the _type field of ES to know which type of the Place enum to fill
pub fn make_place(doc_type: String, value: Option<Box<serde_json::Value>>) -> Option<mimir::Place> {
    value.and_then(|v| {
        fn convert<T: serde::Deserialize>(v: serde_json::Value,
                                          f: fn(T) -> mimir::Place)
                                          -> Option<mimir::Place> {
            serde_json::from_value::<T>(v)
                .map_err(|err| warn!("Impossible to load ES result: {}", err))
                .ok()
                .map(f)
        }
        match doc_type.as_ref() {
            "addr" => convert(*v, mimir::Place::Addr),
            "street" => convert(*v, mimir::Place::Street),
            "admin" => convert(*v, mimir::Place::Admin),
            "poi" => convert(*v, mimir::Place::Poi),
            "stop" => convert(*v, mimir::Place::Stop),
            _ => {
                warn!("unknown ES return value, _type field = {}", doc_type);
                None
            }
        }
    })
}

#[derive(Debug, Eq, PartialEq)]
enum MatchType {
    Prefix,
    Fuzzy,
}

fn build_query(q: &str,
               match_type: MatchType,
               coord: &Option<model::Coord>,
               shape: Option<Vec<rs_es::units::Location>>)
               -> rs_es::query::Query {
    use rs_es::query::functions::Function;
    // we order the type of object we want
    // Note: the addresses are boosted more because even if we don't want them first
    // because they are more severely filtered
    let boost_addr = rs_q::build_term("_type", Addr::doc_type()).with_boost(5000).build();
    let boost_admin = rs_q::build_term("_type", Admin::doc_type()).with_boost(3000).build();
    let boost_stop = rs_q::build_term("_type", Stop::doc_type()).with_boost(2000).build();

    let main_match_type = match match_type {
        MatchType::Prefix => "label.prefix",
        MatchType::Fuzzy => "label.ngram",
    };

    let boost_main_match_query = rs_q::build_multi_match(vec![main_match_type.to_string()],
                                                         q.to_string())
        .with_boost(100)
        .build();
    let boost_zipcode_match_query = rs_q::build_multi_match(vec!["zip_codes.prefix".to_string()],
                                                            q.to_string())
        .with_boost(500)
        .build();

    let mut should_query = vec![boost_addr,
                                boost_admin,
                                boost_stop,
                                boost_main_match_query,
                                boost_zipcode_match_query];

    // for fuzzy search we also search by the prefix index (with a greater boost than ngram)
    // to have better results
    if match_type == MatchType::Fuzzy {
        should_query.push(rs_q::build_match("label.prefix", q.to_string())
            .with_boost(1000)
            .build());
    }

    if let &Some(ref c) = coord {
        // if we have coordinate, we boost we result near this coordinate
        let boost_on_proximity =
            rs_q::build_function_score()
                .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
                .with_boost(500)
                .with_function(Function::build_decay("coord",
                                           rs_u::Location::LatLon(c.lat, c.lon),
                                           rs_u::Distance::new(50f64,
                                                               rs_u::DistanceUnit::Kilometer))
                                   .build_exp())
                .build();
        should_query.push(boost_on_proximity);
    } else {
        // if we don't have coords, we take the field `weight` into account
        let boost_on_weight = rs_q::build_function_score()
            .with_boost_mode(rs_es::query::compound::BoostMode::Multiply)
            .with_boost(300)
            .with_query(rs_q::build_match_all().build())
            .with_function(Function::build_field_value_factor("weight")
                .with_factor(1)
                .with_modifier(rs_es::query::functions::Modifier::Log1p)
                .build())
            .build();
        should_query.push(boost_on_weight);
    }

    let sub_query = rs_q::build_bool()
        .with_should(should_query)
        .build();

    use rs_es::query::{CombinationMinimumShouldMatch, MinimumShouldMatch};
    let min_match = |min, threshold| {
        CombinationMinimumShouldMatch::new(MinimumShouldMatch::from(min),
                                           MinimumShouldMatch::from(threshold))
    };

    let minimum_should_match = match match_type {
        MatchType::Prefix => {
            // when we search [0, 2] fields, we need to match 100% of them
            // when we search [3, 4] fields, we need to match 90% of them
            // when we search [5, +] fields, we only need to match 70% on them
            // this way if the input has more fields than the documents
            // (additional information like the country, the region, ...)
            // they can be ignored
            MinimumShouldMatch::from(vec![min_match(2, 90f64), min_match(4, 70f64)])
        }
        // for fuzzy search we lower our expectation and we accept 50% of token match
        MatchType::Fuzzy => MinimumShouldMatch::from(50f64),
    };

    let mut must = vec![rs_q::build_multi_match(vec![main_match_type.to_string(),
                                                     "zip_codes.prefix".to_string()],
                                                q.to_string())
                            .with_minimum_should_match(minimum_should_match)
                            .build()];

    if let Some(s) = shape {
        must.push(rs_q::build_geo_polygon("coord", s).build());
    }

    // filter to handle house number
    // we either want:
    // * to exactly match the document house_number
    // * or that the document has no house_number
    let filter = rs_q::build_bool()
        .with_should(vec![rs_q::build_bool()
                              .with_must_not(rs_q::build_exists("house_number").build())
                              .build(),
                          rs_q::build_match("house_number", q.to_string()).build()])
        .with_must(must)
        .build();

    rs_q::build_bool()
        .with_must(vec![sub_query])
        .with_filter(filter)
        .build()
}

fn is_existing_index(client: &mut rs_es::Client, index: &String) -> bool {
    !index.is_empty() && client.open_index(&index).is_ok()
}

fn make_indexes<'a>(all_data: bool,
                    pt_dataset_index: &'a Option<String>,
                    client: &mut rs_es::Client)
                    -> Vec<&'a str> {
    let mut result: Vec<&str> = vec![];
    if all_data {
        result.push("munin");
    } else {
        if is_existing_index(client, &"munin_geo_data".to_string()) {
            result.push("munin_geo_data");
        }
        match *pt_dataset_index {
            Some(ref dataset) => {
                if is_existing_index(client, dataset) {
                    result.push(dataset);
                }
            }
            _ => {}
        };
    }
    return result;
}

fn query(q: &str,
         pt_dataset: &Option<&str>,
         all_data: bool,
         cnx: &str,
         match_type: MatchType,
         offset: u64,
         limit: u64,
         coord: &Option<model::Coord>,
         shape: Option<Vec<rs_es::units::Location>>)
         -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    let query = build_query(q, match_type, coord, shape);

    let mut client = build_rs_client(&cnx.to_string());

    let pt_dataset_index = pt_dataset.map(|d| format!("munin_stop_{}", d));
    let indexes = make_indexes(all_data, &pt_dataset_index, &mut client);

    let result: SearchResult<serde_json::Value> = try!(client.search_query()
        .with_indexes(&indexes)
        .with_query(&query)
        .with_from(offset)
        .with_size(limit)
        .send());

    debug!("{} documents found", result.hits.total);

    // for the moment rs-es does not handle enum Document,
    // so we need to convert the ES glob to a Place
    Ok(result.hits
        .hits
        .into_iter()
        .filter_map(|hit| make_place(hit.doc_type, hit.source))
        .collect())
}

pub fn autocomplete(q: &str,
                    pt_dataset: &Option<&str>,
                    all_data: bool,
                    offset: u64,
                    limit: u64,
                    coord: Option<model::Coord>,
                    cnx: &str,
                    shape: Option<Vec<(f64, f64)>>)
                    -> Result<Vec<mimir::Place>, rs_es::error::EsError> {
    fn make_shape(shape: &Option<Vec<(f64, f64)>>) -> Option<Vec<rs_es::units::Location>> {
        shape.as_ref().map(|v| v.iter().map(|&l| l.into()).collect())
    }

    // First we try a prety exact match on the prefix.
    // If there are no results then we do a new fuzzy search (matching ngrams)
    let results = try!(query(&q,
                             &pt_dataset,
                             all_data,
                             cnx,
                             MatchType::Prefix,
                             offset,
                             limit,
                             &coord,
                             make_shape(&shape)));
    if results.is_empty() {
        query(&q,
              &pt_dataset,
              all_data,
              cnx,
              MatchType::Fuzzy,
              offset,
              limit,
              &coord,
              make_shape(&shape))
    } else {
        Ok(results)
    }
}
<|endoftext|>


<commit_before>use std::path::{Path, PathBuf};

use log::{debug, info, trace, warn};
use failure::{Error};

use flirt;
use flirt::pat;
use flirt::sig;
use super::super::super::util;
use super::super::super::arch::{RVA};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

#[derive(Clone, Debug)]
pub struct FlirtConfig {
    pub pat_dir: PathBuf,
    pub sig_dir: PathBuf,
}

impl Default for FlirtConfig {
    fn default() -> FlirtConfig {
        FlirtConfig {
            // TODO: per OS, pick a better default
            pat_dir: PathBuf::from("~/.lancelot/sig/flirt/pat/"),
            sig_dir: PathBuf::from("~/.lancelot/sig/flirt/sig/"),
        }
    }
}

pub struct FlirtAnalyzer {
    sigs: flirt::FlirtSignatureSet,
}

impl FlirtAnalyzer {
    /// parse the given file system path into a collection of FLIRT signatures.
    fn load_pat_file(path: &Path) -> Result<Vec<flirt::FlirtSignature>, Error> {
        let buf = util::read_file(path.to_str().unwrap())?;  // danger
        let s = String::from_utf8(buf)?;
        pat::parse(&s)
    }

    fn load_sig_file(path: &Path) -> Result<Vec<flirt::FlirtSignature>, Error> {
        let buf = util::read_file(path.to_str().unwrap())?;  // danger
        sig::parse(&buf)
    }

    /// remove sigantures from the given collection that don't match basic criteria:
    ///   - the signature must specify a name (or its no use to us)
    ///   - functions must be at least 8 bytes long
    ///   - shorter functions must not have too many wildcards
    fn filter_flirt_signatures(sigs: Vec<flirt::FlirtSignature>) -> Vec<flirt::FlirtSignature> {
        sigs.into_iter()
            .filter(|sig| {
                if let None = sig.get_name() {
                    // must have a name that we can apply.
                    return false;
                }

                let wc_count = sig.byte_sig.0.iter().take(sig.size_of_function as usize).filter(|b| match b {
                    flirt::SigElement::Wildcard => true,
                    flirt::SigElement::Byte(_) => false,
                }).count();

                if sig.size_of_function < 0x8 {
                    // lancelot specific: don't use signatures for functions less than 0x8 bytes.
                    // this just seems too short to be unique and specific.
                    trace!("sig too short: {} {:?}", sig.size_of_function, sig.get_name());
                    return false;
                } else if sig.size_of_function < 0x10 && wc_count > 0 {
                    // lancelot specific: don't allow many wildcards for short functions.
                    trace!("sig too many wildcards: {}/{} {:?}", wc_count, sig.size_of_function, sig.get_name());
                    return false;
                } else if sig.size_of_function < 0x18 && wc_count > 4 {
                    trace!("sig too many wildcards: {}/{} {:?}", wc_count, sig.size_of_function, sig.get_name());
                    return false;
                } else if sig.size_of_function < 0x20 && wc_count > 0x10 {
                    trace!("sig too many wildcards: {}/{} {:?}", wc_count, sig.size_of_function, sig.get_name());
                    return false;
                } else {
                    return true;
                }
            })
            .collect()
    }

    /// parse the given directory into a set of FLIRT signatures.
    ///
    /// applies some filtering on the signatures for better accuracy:
    ///   - the signature must specify a name (or its no use to us)
    ///   - functions must be at least 8 bytes long
    ///   - shorter functions must not have too many wildcards
    ///
    /// the directory structure should look like:
    ///   - flat directory of .pat or .sig files
    fn load_flirt_directory(path: &Path) -> Result<Vec<flirt::FlirtSignature>, Error> {
        let mut sigs = vec![];

        for entry in std::fs::read_dir(path)? {
            let entry = entry?;
            let path = entry.path();

            if path.is_dir() {
                continue
            }

            let ext = path.extension().and_then(|ext| ext.to_str()).unwrap_or("");

            if ext == "pat" {
                debug!("FLIRT analyzer: loading .pat file: {:?}", path);
                sigs.extend(FlirtAnalyzer::load_pat_file(&path).unwrap()); //danger
            } else if ext == ".sig" {
                debug!("FLIRT analyzer: loading .sig file: {:?}", path);
                sigs.extend(FlirtAnalyzer::load_sig_file(&path).unwrap()); //danger
            } else {
                debug!("FLIRT analyzer: skipping file: {:?}", path);
            }
        }
        debug!("loaded {} total FLIRT signatures", sigs.len());

        let sigs = FlirtAnalyzer::filter_flirt_signatures(sigs);
        info!("filtered to {} usable FLIRT signatures", sigs.len());

        Ok(sigs)
    }

    pub fn new(config: FlirtConfig) -> FlirtAnalyzer {
        // TODO: add startup signatures to detect runtime/signature set

        let mut sigs = vec![];

        for path in [config.pat_dir, config.sig_dir].iter() {
            sigs.extend(
                if path.exists() {
                    match FlirtAnalyzer::load_flirt_directory(&path) {
                        Ok(sigs) => sigs,
                        Err(_) => vec![],
                    }
                } else {
                    warn!("FLIRT directory does not exist: {:?}", path);
                    vec![]
                }
            );
        }

        FlirtAnalyzer{
            sigs: flirt::FlirtSignatureSet::with_signatures(sigs)
        }
    }
}

impl Analyzer for FlirtAnalyzer {
    fn get_name(&self) -> String {
        "FLIRT function signature analyzer".to_string()
    }

    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let mut buf = [0u8; 0xFF];
        let functions: Vec<RVA> = ws.get_functions().cloned().collect();

        for &fva in functions.iter() {
            if let Ok(buf) = ws.read_bytes_into(fva, &mut buf[..]) {
                let matches = self.sigs.r#match(buf);

                // no matches
                if matches.len() == 0 {
                    continue;
                }

                if matches.len() > 1 {
                    // more than one match.
                    //
                    // the only time this is acceptable is if we've loaded multiple signature sets
                    //  and these match the same functions.
                    // in this case, all the names should be the same.
                    // so, lets ensure their names don't conflict.
                    //
                    // implementation: the first name should be the same as all the other names.
                    let name1 = matches[0].get_name();
                    if let Some(m2) = matches[1..]
                        .iter()
                        .find(|m| m.get_name() != name1) {

                        debug!("ambiguous FLIRT signature match: {}: {:?} and {:?}",
                               fva, name1, m2.get_name());
                        continue;
                    }
                };

                let match_ = matches[0];

                // TODO: should not apply the same symbol name to more than one location?
                // TODO: apply reference names

                // can unwrap name cause its guaranteed to have a name due to filter above.
                let name = match_.get_name().unwrap();
                debug!("FLIRT signature match: {} {}", fva, name);
                ws.make_symbol(fva, name).unwrap();  // danger
                continue;
            }
        }

        ws.analyze()?;

        Ok(())
    }
}
<commit_msg>analysis: flirt: fix expected extension for sig files<commit_after>use std::path::{Path, PathBuf};

use log::{debug, info, trace, warn};
use failure::{Error};

use flirt;
use flirt::pat;
use flirt::sig;
use super::super::super::util;
use super::super::super::arch::{RVA};
use super::super::super::workspace::Workspace;
use super::super::{Analyzer};

#[derive(Clone, Debug)]
pub struct FlirtConfig {
    pub pat_dir: PathBuf,
    pub sig_dir: PathBuf,
}

impl Default for FlirtConfig {
    fn default() -> FlirtConfig {
        FlirtConfig {
            // TODO: per OS, pick a better default
            pat_dir: PathBuf::from("~/.lancelot/sig/flirt/pat/"),
            sig_dir: PathBuf::from("~/.lancelot/sig/flirt/sig/"),
        }
    }
}

pub struct FlirtAnalyzer {
    sigs: flirt::FlirtSignatureSet,
}

impl FlirtAnalyzer {
    /// parse the given file system path into a collection of FLIRT signatures.
    fn load_pat_file(path: &Path) -> Result<Vec<flirt::FlirtSignature>, Error> {
        let buf = util::read_file(path.to_str().unwrap())?;  // danger
        let s = String::from_utf8(buf)?;
        pat::parse(&s)
    }

    fn load_sig_file(path: &Path) -> Result<Vec<flirt::FlirtSignature>, Error> {
        let buf = util::read_file(path.to_str().unwrap())?;  // danger
        sig::parse(&buf)
    }

    /// remove sigantures from the given collection that don't match basic criteria:
    ///   - the signature must specify a name (or its no use to us)
    ///   - functions must be at least 8 bytes long
    ///   - shorter functions must not have too many wildcards
    fn filter_flirt_signatures(sigs: Vec<flirt::FlirtSignature>) -> Vec<flirt::FlirtSignature> {
        sigs.into_iter()
            .filter(|sig| {
                if let None = sig.get_name() {
                    // must have a name that we can apply.
                    return false;
                }

                let wc_count = sig.byte_sig.0.iter().take(sig.size_of_function as usize).filter(|b| match b {
                    flirt::SigElement::Wildcard => true,
                    flirt::SigElement::Byte(_) => false,
                }).count();

                if sig.size_of_function < 0x8 {
                    // lancelot specific: don't use signatures for functions less than 0x8 bytes.
                    // this just seems too short to be unique and specific.
                    trace!("sig too short: {} {:?}", sig.size_of_function, sig.get_name());
                    return false;
                } else if sig.size_of_function < 0x10 && wc_count > 0 {
                    // lancelot specific: don't allow many wildcards for short functions.
                    trace!("sig too many wildcards: {}/{} {:?}", wc_count, sig.size_of_function, sig.get_name());
                    return false;
                } else if sig.size_of_function < 0x18 && wc_count > 4 {
                    trace!("sig too many wildcards: {}/{} {:?}", wc_count, sig.size_of_function, sig.get_name());
                    return false;
                } else if sig.size_of_function < 0x20 && wc_count > 0x10 {
                    trace!("sig too many wildcards: {}/{} {:?}", wc_count, sig.size_of_function, sig.get_name());
                    return false;
                } else {
                    return true;
                }
            })
            .collect()
    }

    /// parse the given directory into a set of FLIRT signatures.
    ///
    /// applies some filtering on the signatures for better accuracy:
    ///   - the signature must specify a name (or its no use to us)
    ///   - functions must be at least 8 bytes long
    ///   - shorter functions must not have too many wildcards
    ///
    /// the directory structure should look like:
    ///   - flat directory of .pat or .sig files
    fn load_flirt_directory(path: &Path) -> Result<Vec<flirt::FlirtSignature>, Error> {
        let mut sigs = vec![];

        for entry in std::fs::read_dir(path)? {
            let entry = entry?;
            let path = entry.path();

            if path.is_dir() {
                continue
            }

            let ext = path.extension().and_then(|ext| ext.to_str()).unwrap_or("");

            if ext == "pat" {
                debug!("FLIRT analyzer: loading .pat file: {:?}", path);
                sigs.extend(FlirtAnalyzer::load_pat_file(&path).unwrap()); //danger
            } else if ext == "sig" {
                debug!("FLIRT analyzer: loading .sig file: {:?}", path);
                sigs.extend(FlirtAnalyzer::load_sig_file(&path).unwrap()); //danger
            } else {
                debug!("FLIRT analyzer: skipping file: {:?}", path);
            }
        }
        debug!("loaded {} total FLIRT signatures", sigs.len());

        let sigs = FlirtAnalyzer::filter_flirt_signatures(sigs);
        info!("filtered to {} usable FLIRT signatures", sigs.len());

        Ok(sigs)
    }

    pub fn new(config: FlirtConfig) -> FlirtAnalyzer {
        // TODO: add startup signatures to detect runtime/signature set

        let mut sigs = vec![];

        for path in [config.pat_dir, config.sig_dir].iter() {
            sigs.extend(
                if path.exists() {
                    match FlirtAnalyzer::load_flirt_directory(&path) {
                        Ok(sigs) => sigs,
                        Err(_) => vec![],
                    }
                } else {
                    warn!("FLIRT directory does not exist: {:?}", path);
                    vec![]
                }
            );
        }

        FlirtAnalyzer{
            sigs: flirt::FlirtSignatureSet::with_signatures(sigs)
        }
    }
}

impl Analyzer for FlirtAnalyzer {
    fn get_name(&self) -> String {
        "FLIRT function signature analyzer".to_string()
    }

    fn analyze(&self, ws: &mut Workspace) -> Result<(), Error> {
        let mut buf = [0u8; 0xFF];
        let functions: Vec<RVA> = ws.get_functions().cloned().collect();

        for &fva in functions.iter() {
            if let Ok(buf) = ws.read_bytes_into(fva, &mut buf[..]) {
                let matches = self.sigs.r#match(buf);

                // no matches
                if matches.len() == 0 {
                    continue;
                }

                if matches.len() > 1 {
                    // more than one match.
                    //
                    // the only time this is acceptable is if we've loaded multiple signature sets
                    //  and these match the same functions.
                    // in this case, all the names should be the same.
                    // so, lets ensure their names don't conflict.
                    //
                    // implementation: the first name should be the same as all the other names.
                    let name1 = matches[0].get_name();
                    if let Some(m2) = matches[1..]
                        .iter()
                        .find(|m| m.get_name() != name1) {

                        debug!("ambiguous FLIRT signature match: {}: {:?} and {:?}",
                               fva, name1, m2.get_name());
                        continue;
                    }
                };

                let match_ = matches[0];

                // TODO: should not apply the same symbol name to more than one location?
                // TODO: apply reference names

                // can unwrap name cause its guaranteed to have a name due to filter above.
                let name = match_.get_name().unwrap();
                debug!("FLIRT signature match: {} {}", fva, name);
                ws.make_symbol(fva, name).unwrap();  // danger
                continue;
            }
        }

        ws.analyze()?;

        Ok(())
    }
}
<|endoftext|>


<commit_before>extern crate difference;
extern crate quickcheck;

use difference::{Changeset, Difference};
use quickcheck::{TestResult, quickcheck, QuickCheck};
use std::fmt;

const DEBUG: bool = false;

struct Check<'a> {
    old: &'a str,
    new: &'a str,
    changeset: Changeset,
}

impl<'a> fmt::Display for Check<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Changeset::new({:?}, {:?}, {:?}) -> [",
            self.old,
            self.new,
            self.changeset.split
        )?;

        let mut iter = self.changeset.diffs.iter();
        if let Some(d) = iter.next() {
            write!(f, "{:?}", d)?;
        }
        for d in iter {
            write!(f, " {:?}", d)?;
        }
        write!(f, "]")
    }
}

fn check_changeset(old: &str, new: &str, split: &str) -> TestResult {
    Check::new(old, new, split).check()
}

impl<'a> Check<'a> {
    fn new(old: &'a str, new: &'a str, split: &'a str) -> Check<'a> {
        Check {
            old: old,
            new: new,
            changeset: Changeset::new(old, new, split),
        }
    }

    fn check(&self) -> TestResult {
        let split = &self.changeset.split;

        let mut old = self.old.trim_left_matches(split);
        let mut new = self.new.trim_left_matches(split);

        macro_rules! expect {
            ($ident:ident, $pattern:expr, $diff:expr, $split:expr) => {
                if !$ident.starts_with($pattern) {
                    return TestResult::error(format!("`{:?}` does not match `{}` at {:?} for {}",
                                                      $diff,
                                                      stringify!($ident),
                                                      $ident,
                                                      self));
                }
                $ident = &$ident[$pattern.len()..].trim_left_matches($split);
            }
        }

        for d in &self.changeset.diffs {
            if DEBUG {
                println!("assert `{:?}` (old: {:?}, new: {:?})", d, old, new);
            }

            match *d {
                Difference::Same(ref x) => {
                    expect!(old, x, d, split);
                    expect!(new, x, d, split);
                }
                Difference::Add(ref x) => {
                    expect!(new, x, d, split);
                }
                Difference::Rem(ref x) => {
                    expect!(old, x, d, split);
                }
            }
        }
        if !old.is_empty() {
            return TestResult::error(format!(
                "expected end of string in `old` at {:?} for {}",
                old,
                self
            ));
        }
        if !new.is_empty() {
            return TestResult::error(format!(
                "expected end of string in `new` at {:?} for {}",
                new,
                self
            ));
        }
        TestResult::passed()
    }
}

#[test]
fn simple() {
    quickcheck(check_changeset("a", "a a", " "));
}

#[test]
#[ignore]
fn issue_19() {
    // this should work but it doesn't
    // https://github.com/johannhof/difference.rs/issues/19
    quickcheck(check_changeset("a b : g", "b a : b b : g g", " "));
}

#[test]
#[ignore]
fn fuzzy() {
    fn prop(old: Vec<usize>, new: Vec<usize>, words: Vec<char>) -> TestResult {
        if words.is_empty() {
            return TestResult::discard();
        }

        fn map_to_words(input: &[usize], words: &[char]) -> String {
            input.iter().enumerate().fold(
                String::new(),
                |mut acc, (i, x)| {
                    if i > 0 {
                        acc.push(' ');
                    }
                    acc.push(words[x % words.len()]);
                    acc
                },
            )
        }
        let old = map_to_words(&old, &words);
        let new = map_to_words(&new, &words);

        check_changeset(&old, &new, " ")
    }

    QuickCheck::new()
        .tests(100) // max successful tests
        .max_tests(10000) // max attempts
        .quickcheck(prop as fn(Vec<usize>, Vec<usize>, Vec<char>) -> TestResult);
}
<commit_msg>Rewrite quickcheck test.<commit_after>extern crate difference;
extern crate quickcheck;

use difference::{Changeset, Difference};
use quickcheck::{TestResult, quickcheck, QuickCheck};
use std::fmt;

const DEBUG: bool = false;

struct Check<'a> {
    old: &'a str,
    new: &'a str,
    changeset: Changeset,
}

impl<'a> fmt::Display for Check<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Changeset::new({:?}, {:?}, {:?}) -> [",
            self.old,
            self.new,
            self.changeset.split
        )?;

        let mut iter = self.changeset.diffs.iter();
        if let Some(d) = iter.next() {
            write!(f, "{:?}", d)?;
        }
        for d in iter {
            write!(f, " {:?}", d)?;
        }
        write!(f, "]")
    }
}

fn check_changeset(old: &str, new: &str, split: &str) -> TestResult {
    Check::new(old, new, split).check()
}

impl<'a> Check<'a> {
    fn new(old: &'a str, new: &'a str, split: &'a str) -> Check<'a> {
        Check {
            old: old,
            new: new,
            changeset: Changeset::new(old, new, split),
        }
    }

    fn check(&self) -> TestResult {
        let split = &self.changeset.split;

        let mut old: Vec<&str> = Vec::new();
        let mut new: Vec<&str> = Vec::new();

        for d in &self.changeset.diffs {
            if DEBUG {
                println!("assert `{:?}` (old: {:?}, new: {:?})", d, old, new);
            }

            match *d {
                Difference::Same(ref x) => {
                    old.push(x);
                    new.push(x);
                }
                Difference::Add(ref x) => {
                    new.push(x);
                }
                Difference::Rem(ref x) => {
                    old.push(x);
                }
            }
        }
        let got_old = old.join(split);
        let got_new = new.join(split);
        if got_old != self.old {
            return TestResult::error(format!("Diff output implies old=`{:?}`, not `{:?}` in {}",
                        got_old, self.old, self,
                ));
        }
        if got_new != self.new {
            return TestResult::error(format!("Diff output implies new=`{:?}`, not `{:?}` in {}",
                        got_new, self.new, self,
                ));
        }

        TestResult::passed()
    }
}

#[test]
fn simple() {
    quickcheck(check_changeset("a", "a a", " "));
}

#[test]
#[ignore]
fn issue_19() {
    // this should work but it doesn't
    // https://github.com/johannhof/difference.rs/issues/19
    quickcheck(check_changeset("a b : g", "b a : b b : g g", " "));
}

#[test]
#[ignore]
fn fuzzy() {
    fn prop(old: Vec<usize>, new: Vec<usize>, words: Vec<char>) -> TestResult {
        if words.is_empty() {
            return TestResult::discard();
        }

        fn map_to_words(input: &[usize], words: &[char]) -> String {
            input.iter().enumerate().fold(
                String::new(),
                |mut acc, (i, x)| {
                    if i > 0 {
                        acc.push(' ');
                    }
                    acc.push(words[x % words.len()]);
                    acc
                },
            )
        }
        let old = map_to_words(&old, &words);
        let new = map_to_words(&new, &words);

        check_changeset(&old, &new, " ")
    }

    QuickCheck::new()
        .tests(100) // max successful tests
        .max_tests(10000) // max attempts
        .quickcheck(prop as fn(Vec<usize>, Vec<usize>, Vec<char>) -> TestResult);
}
<|endoftext|>


<commit_before>use std::io::{self, IoSlice, IoSliceMut, Read, Write};
use std::net::{self, Shutdown, SocketAddr};
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
use std::sync::mpsc::channel;
use std::sync::{Arc, Barrier};
use std::thread;

use mio::net::TcpStream;
use mio::{Interests, Token};

mod util;

use util::{
    any_local_address, any_local_ipv6_address, assert_send, assert_sync, assert_would_block,
    expect_events, expect_no_events, init, init_with_poll, ExpectEvent,
};

const DATA1: &[u8] = b"Hello world!";
const DATA2: &[u8] = b"Hello mars!";
// TODO: replace with `DATA1.len()` once `const_slice_len` is stable.
const DATA1_LEN: usize = 12;
const DATA2_LEN: usize = 11;

const ID1: Token = Token(0);
const ID2: Token = Token(1);

#[test]
fn is_send_and_sync() {
    assert_send::<TcpStream>();
    assert_sync::<TcpStream>();
}

#[test]
fn tcp_stream_ipv4() {
    smoke_test_tcp_stream(any_local_address());
}

#[test]
fn tcp_stream_ipv6() {
    smoke_test_tcp_stream(any_local_ipv6_address());
}

fn smoke_test_tcp_stream(addr: SocketAddr) {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, addr) = echo_listener(addr, 1);
    let mut stream = TcpStream::connect(addr).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::WRITABLE.add(Interests::READABLE))
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::WRITABLE)],
    );

    let mut buf = [0; 16];
    assert_would_block(stream.peek(&mut buf));
    assert_would_block(stream.read(&mut buf));

    // NOTE: the call to `peer_addr` must happen after we received a writable
    // event as the stream might not yet be connected.
    assert_eq!(stream.peer_addr().unwrap(), addr);
    assert!(stream.local_addr().unwrap().ip().is_loopback());

    let n = stream.write(&DATA1).expect("unable to write to stream");
    assert_eq!(n, DATA1.len());

    stream.flush().unwrap();

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    let n = stream.peek(&mut buf).expect("unable to peek from stream");
    assert_eq!(n, DATA1.len());
    assert_eq!(&buf[..n], DATA1);

    let n = stream.read(&mut buf).expect("unable to read from stream");
    assert_eq!(n, DATA1.len());
    assert_eq!(&buf[..n], DATA1);

    assert!(stream.take_error().unwrap().is_none());

    let bufs = [IoSlice::new(&DATA1), IoSlice::new(&DATA2)];
    let n = stream
        .write_vectored(&bufs)
        .expect("unable to write vectored to stream");
    assert_eq!(n, DATA1.len() + DATA2.len());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    let mut buf1 = [1; DATA1_LEN];
    let mut buf2 = [2; DATA2_LEN + 1];
    let mut bufs = [IoSliceMut::new(&mut buf1), IoSliceMut::new(&mut buf2)];
    let n = stream
        .read_vectored(&mut bufs)
        .expect("unable to read vectored from stream");
    assert_eq!(n, DATA1.len() + DATA2.len());
    assert_eq!(&buf1, DATA1);
    assert_eq!(&buf2[..DATA2.len()], DATA2);
    assert_eq!(buf2[DATA2.len()], 2); // Last byte should be unchanged.

    // Close the connection to allow the listener to shutdown.
    drop(stream);

    thread_handle.join().expect("unable to join thread");
}

#[test]
fn try_clone() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let mut stream1 = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream1, ID1, Interests::WRITABLE)
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::WRITABLE)],
    );

    let n = stream1.write(DATA1).unwrap();
    assert_eq!(n, DATA1.len());

    let mut stream2 = stream1.try_clone().unwrap();

    // When using `try_clone` the `TcpStream` needs to be deregistered!
    poll.registry().deregister(&stream1).unwrap();
    drop(stream1);

    poll.registry()
        .register(&stream2, ID2, Interests::READABLE)
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID2, Interests::READABLE)],
    );

    let mut buf = [0; 20];
    let n = stream2.read(&mut buf).unwrap();
    assert_eq!(n, DATA1.len());
    assert_eq!(&buf[..n], DATA1);

    drop(stream2);
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn ttl() {
    init();

    let barrier = Arc::new(Barrier::new(2));
    let (thread_handle, address) = start_listener(1, Some(barrier.clone()));

    let stream = TcpStream::connect(address).unwrap();

    const TTL: u32 = 10;
    stream.set_ttl(TTL).unwrap();
    assert_eq!(stream.ttl().unwrap(), TTL);
    assert!(stream.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn nodelay() {
    init();

    let barrier = Arc::new(Barrier::new(2));
    let (thread_handle, address) = start_listener(1, Some(barrier.clone()));

    let stream = TcpStream::connect(address).unwrap();

    const NO_DELAY: bool = true;
    stream.set_nodelay(NO_DELAY).unwrap();
    assert_eq!(stream.nodelay().unwrap(), NO_DELAY);
    assert!(stream.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn shutdown_read() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let mut stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::WRITABLE.add(Interests::READABLE))
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::WRITABLE)],
    );

    let n = stream.write(DATA2).unwrap();
    assert_eq!(n, DATA2.len());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    stream.shutdown(Shutdown::Read).unwrap();

    let mut buf = [0; 20];
    let n = stream.read(&mut buf).unwrap();
    assert_eq!(n, 0);

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

#[test]
#[ignore = "This test is flaky, it doesn't always receive an event after shutting down the write side"]
fn shutdown_write() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let mut stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::WRITABLE.add(Interests::READABLE))
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::WRITABLE)],
    );

    let n = stream.write(DATA1).unwrap();
    assert_eq!(n, DATA1.len());

    stream.shutdown(Shutdown::Write).unwrap();

    let err = stream.write(DATA2).unwrap_err();
    assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);

    // FIXME: we don't always receive the following event.
    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    // Read should be ok.
    let mut buf = [0; 20];
    let n = stream.read(&mut buf).unwrap();
    assert_eq!(n, DATA1.len());
    assert_eq!(&buf[..n], DATA1);

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn shutdown_both() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let mut stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::WRITABLE.add(Interests::READABLE))
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::WRITABLE)],
    );

    let n = stream.write(DATA1).unwrap();
    assert_eq!(n, DATA1.len());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    stream.shutdown(Shutdown::Both).unwrap();

    let mut buf = [0; 20];
    let n = stream.read(&mut buf).unwrap();
    assert_eq!(n, 0);

    let err = stream.write(DATA2).unwrap_err();
    assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

#[test]
#[cfg(unix)]
fn raw_fd() {
    init();

    let (thread_handle, address) = start_listener(1, None);

    let stream = TcpStream::connect(address).unwrap();
    let address = stream.local_addr().unwrap();

    let raw_fd1 = stream.as_raw_fd();
    let raw_fd2 = stream.into_raw_fd();
    assert_eq!(raw_fd1, raw_fd2);

    let stream = unsafe { TcpStream::from_raw_fd(raw_fd2) };
    assert_eq!(stream.as_raw_fd(), raw_fd1);
    assert_eq!(stream.local_addr().unwrap(), address);

    thread_handle.join().expect("unable to join thread");
}

#[test]
fn registering() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::READABLE)
        .expect("unable to register TCP stream");

    expect_no_events(&mut poll, &mut events);

    // NOTE: more tests are done in the smoke tests above.

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn reregistering() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::READABLE)
        .expect("unable to register TCP stream");

    poll.registry()
        .reregister(&stream, ID2, Interests::WRITABLE)
        .expect("unable to reregister TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID2, Interests::WRITABLE)],
    );

    assert_eq!(stream.peer_addr().unwrap(), address);

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn deregistering() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::READABLE)
        .expect("unable to register TCP stream");

    poll.registry()
        .deregister(&stream)
        .expect("unable to deregister TCP stream");

    expect_no_events(&mut poll, &mut events);

    // We do expect to be connected.
    assert_eq!(stream.peer_addr().unwrap(), address);

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

/// Start a listener that accepts `n_connections` connections on the returned
/// address. It echos back any data it reads from the connection before
/// accepting another one.
fn echo_listener(addr: SocketAddr, n_connections: usize) -> (thread::JoinHandle<()>, SocketAddr) {
    let (sender, receiver) = channel();
    let thread_handle = thread::spawn(move || {
        let listener = net::TcpListener::bind(addr).unwrap();
        let local_address = listener.local_addr().unwrap();
        sender.send(local_address).unwrap();

        let mut buf = [0; 128];
        for _ in 0..n_connections {
            let (mut stream, _) = listener.accept().unwrap();

            loop {
                let n = stream.read(&mut buf).expect("error reading");
                if n == 0 {
                    break;
                }
                let written = stream.write(&buf[..n]).expect("error writing");
                assert_eq!(written, n, "short write");
            }
        }
    });
    (thread_handle, receiver.recv().unwrap())
}

/// Start a listener that accepts `n_connections` connections on the returned
/// address. If a barrier is provided it will wait on it before closing the
/// connection.
fn start_listener(
    n_connections: usize,
    barrier: Option<Arc<Barrier>>,
) -> (thread::JoinHandle<()>, SocketAddr) {
    let (sender, receiver) = channel();
    let thread_handle = thread::spawn(move || {
        let listener = net::TcpListener::bind(any_local_address()).unwrap();
        let local_address = listener.local_addr().unwrap();
        sender.send(local_address).unwrap();

        for _ in 0..n_connections {
            let (stream, _) = listener.accept().unwrap();
            if let Some(ref barrier) = barrier {
                barrier.wait();
            }
            drop(stream);
        }
    });
    (thread_handle, receiver.recv().unwrap())
}
<commit_msg>Loosen TcpStream shutdown reading side test<commit_after>use std::io::{self, IoSlice, IoSliceMut, Read, Write};
use std::net::{self, Shutdown, SocketAddr};
#[cfg(unix)]
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
use std::sync::mpsc::channel;
use std::sync::{Arc, Barrier};
use std::thread;

use mio::net::TcpStream;
use mio::{Interests, Token};

mod util;

use util::{
    any_local_address, any_local_ipv6_address, assert_send, assert_sync, assert_would_block,
    expect_events, expect_no_events, init, init_with_poll, ExpectEvent,
};

const DATA1: &[u8] = b"Hello world!";
const DATA2: &[u8] = b"Hello mars!";
// TODO: replace with `DATA1.len()` once `const_slice_len` is stable.
const DATA1_LEN: usize = 12;
const DATA2_LEN: usize = 11;

const ID1: Token = Token(0);
const ID2: Token = Token(1);

#[test]
fn is_send_and_sync() {
    assert_send::<TcpStream>();
    assert_sync::<TcpStream>();
}

#[test]
fn tcp_stream_ipv4() {
    smoke_test_tcp_stream(any_local_address());
}

#[test]
fn tcp_stream_ipv6() {
    smoke_test_tcp_stream(any_local_ipv6_address());
}

fn smoke_test_tcp_stream(addr: SocketAddr) {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, addr) = echo_listener(addr, 1);
    let mut stream = TcpStream::connect(addr).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::WRITABLE.add(Interests::READABLE))
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::WRITABLE)],
    );

    let mut buf = [0; 16];
    assert_would_block(stream.peek(&mut buf));
    assert_would_block(stream.read(&mut buf));

    // NOTE: the call to `peer_addr` must happen after we received a writable
    // event as the stream might not yet be connected.
    assert_eq!(stream.peer_addr().unwrap(), addr);
    assert!(stream.local_addr().unwrap().ip().is_loopback());

    let n = stream.write(&DATA1).expect("unable to write to stream");
    assert_eq!(n, DATA1.len());

    stream.flush().unwrap();

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    let n = stream.peek(&mut buf).expect("unable to peek from stream");
    assert_eq!(n, DATA1.len());
    assert_eq!(&buf[..n], DATA1);

    let n = stream.read(&mut buf).expect("unable to read from stream");
    assert_eq!(n, DATA1.len());
    assert_eq!(&buf[..n], DATA1);

    assert!(stream.take_error().unwrap().is_none());

    let bufs = [IoSlice::new(&DATA1), IoSlice::new(&DATA2)];
    let n = stream
        .write_vectored(&bufs)
        .expect("unable to write vectored to stream");
    assert_eq!(n, DATA1.len() + DATA2.len());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    let mut buf1 = [1; DATA1_LEN];
    let mut buf2 = [2; DATA2_LEN + 1];
    let mut bufs = [IoSliceMut::new(&mut buf1), IoSliceMut::new(&mut buf2)];
    let n = stream
        .read_vectored(&mut bufs)
        .expect("unable to read vectored from stream");
    assert_eq!(n, DATA1.len() + DATA2.len());
    assert_eq!(&buf1, DATA1);
    assert_eq!(&buf2[..DATA2.len()], DATA2);
    assert_eq!(buf2[DATA2.len()], 2); // Last byte should be unchanged.

    // Close the connection to allow the listener to shutdown.
    drop(stream);

    thread_handle.join().expect("unable to join thread");
}

#[test]
fn try_clone() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let mut stream1 = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream1, ID1, Interests::WRITABLE)
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::WRITABLE)],
    );

    let n = stream1.write(DATA1).unwrap();
    assert_eq!(n, DATA1.len());

    let mut stream2 = stream1.try_clone().unwrap();

    // When using `try_clone` the `TcpStream` needs to be deregistered!
    poll.registry().deregister(&stream1).unwrap();
    drop(stream1);

    poll.registry()
        .register(&stream2, ID2, Interests::READABLE)
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID2, Interests::READABLE)],
    );

    let mut buf = [0; 20];
    let n = stream2.read(&mut buf).unwrap();
    assert_eq!(n, DATA1.len());
    assert_eq!(&buf[..n], DATA1);

    drop(stream2);
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn ttl() {
    init();

    let barrier = Arc::new(Barrier::new(2));
    let (thread_handle, address) = start_listener(1, Some(barrier.clone()));

    let stream = TcpStream::connect(address).unwrap();

    const TTL: u32 = 10;
    stream.set_ttl(TTL).unwrap();
    assert_eq!(stream.ttl().unwrap(), TTL);
    assert!(stream.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn nodelay() {
    init();

    let barrier = Arc::new(Barrier::new(2));
    let (thread_handle, address) = start_listener(1, Some(barrier.clone()));

    let stream = TcpStream::connect(address).unwrap();

    const NO_DELAY: bool = true;
    stream.set_nodelay(NO_DELAY).unwrap();
    assert_eq!(stream.nodelay().unwrap(), NO_DELAY);
    assert!(stream.take_error().unwrap().is_none());

    barrier.wait();
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn shutdown_read() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let mut stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::WRITABLE.add(Interests::READABLE))
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::WRITABLE)],
    );

    let n = stream.write(DATA2).unwrap();
    assert_eq!(n, DATA2.len());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    stream.shutdown(Shutdown::Read).unwrap();

    // Shutting down the reading side is different on each platform. For example
    // on Linux based systems we can still read.
    #[cfg(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "ios",
        target_os = "macos",
        target_os = "netbsd",
        target_os = "openbsd"
    ))]
    {
        let mut buf = [0; 20];
        let n = stream.read(&mut buf).unwrap();
        assert_eq!(n, 0);
    }

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

#[test]
#[ignore = "This test is flaky, it doesn't always receive an event after shutting down the write side"]
fn shutdown_write() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let mut stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::WRITABLE.add(Interests::READABLE))
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::WRITABLE)],
    );

    let n = stream.write(DATA1).unwrap();
    assert_eq!(n, DATA1.len());

    stream.shutdown(Shutdown::Write).unwrap();

    let err = stream.write(DATA2).unwrap_err();
    assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);

    // FIXME: we don't always receive the following event.
    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    // Read should be ok.
    let mut buf = [0; 20];
    let n = stream.read(&mut buf).unwrap();
    assert_eq!(n, DATA1.len());
    assert_eq!(&buf[..n], DATA1);

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn shutdown_both() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let mut stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::WRITABLE.add(Interests::READABLE))
        .expect("unable to register TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::WRITABLE)],
    );

    let n = stream.write(DATA1).unwrap();
    assert_eq!(n, DATA1.len());

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID1, Interests::READABLE)],
    );

    stream.shutdown(Shutdown::Both).unwrap();

    // Shutting down the reading side is different on each platform. For example
    // on Linux based systems we can still read.
    #[cfg(any(
        target_os = "dragonfly",
        target_os = "freebsd",
        target_os = "ios",
        target_os = "macos",
        target_os = "netbsd",
        target_os = "openbsd"
    ))]
    {
        let mut buf = [0; 20];
        let n = stream.read(&mut buf).unwrap();
        assert_eq!(n, 0);
    }

    let err = stream.write(DATA2).unwrap_err();
    assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

#[test]
#[cfg(unix)]
fn raw_fd() {
    init();

    let (thread_handle, address) = start_listener(1, None);

    let stream = TcpStream::connect(address).unwrap();
    let address = stream.local_addr().unwrap();

    let raw_fd1 = stream.as_raw_fd();
    let raw_fd2 = stream.into_raw_fd();
    assert_eq!(raw_fd1, raw_fd2);

    let stream = unsafe { TcpStream::from_raw_fd(raw_fd2) };
    assert_eq!(stream.as_raw_fd(), raw_fd1);
    assert_eq!(stream.local_addr().unwrap(), address);

    thread_handle.join().expect("unable to join thread");
}

#[test]
fn registering() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::READABLE)
        .expect("unable to register TCP stream");

    expect_no_events(&mut poll, &mut events);

    // NOTE: more tests are done in the smoke tests above.

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn reregistering() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::READABLE)
        .expect("unable to register TCP stream");

    poll.registry()
        .reregister(&stream, ID2, Interests::WRITABLE)
        .expect("unable to reregister TCP stream");

    expect_events(
        &mut poll,
        &mut events,
        vec![ExpectEvent::new(ID2, Interests::WRITABLE)],
    );

    assert_eq!(stream.peer_addr().unwrap(), address);

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

#[test]
fn deregistering() {
    let (mut poll, mut events) = init_with_poll();

    let (thread_handle, address) = echo_listener(any_local_address(), 1);

    let stream = TcpStream::connect(address).unwrap();

    poll.registry()
        .register(&stream, ID1, Interests::READABLE)
        .expect("unable to register TCP stream");

    poll.registry()
        .deregister(&stream)
        .expect("unable to deregister TCP stream");

    expect_no_events(&mut poll, &mut events);

    // We do expect to be connected.
    assert_eq!(stream.peer_addr().unwrap(), address);

    drop(stream);
    thread_handle.join().expect("unable to join thread");
}

/// Start a listener that accepts `n_connections` connections on the returned
/// address. It echos back any data it reads from the connection before
/// accepting another one.
fn echo_listener(addr: SocketAddr, n_connections: usize) -> (thread::JoinHandle<()>, SocketAddr) {
    let (sender, receiver) = channel();
    let thread_handle = thread::spawn(move || {
        let listener = net::TcpListener::bind(addr).unwrap();
        let local_address = listener.local_addr().unwrap();
        sender.send(local_address).unwrap();

        let mut buf = [0; 128];
        for _ in 0..n_connections {
            let (mut stream, _) = listener.accept().unwrap();

            loop {
                let n = stream.read(&mut buf).expect("error reading");
                if n == 0 {
                    break;
                }
                let written = stream.write(&buf[..n]).expect("error writing");
                assert_eq!(written, n, "short write");
            }
        }
    });
    (thread_handle, receiver.recv().unwrap())
}

/// Start a listener that accepts `n_connections` connections on the returned
/// address. If a barrier is provided it will wait on it before closing the
/// connection.
fn start_listener(
    n_connections: usize,
    barrier: Option<Arc<Barrier>>,
) -> (thread::JoinHandle<()>, SocketAddr) {
    let (sender, receiver) = channel();
    let thread_handle = thread::spawn(move || {
        let listener = net::TcpListener::bind(any_local_address()).unwrap();
        let local_address = listener.local_addr().unwrap();
        sender.send(local_address).unwrap();

        for _ in 0..n_connections {
            let (stream, _) = listener.accept().unwrap();
            if let Some(ref barrier) = barrier {
                barrier.wait();
            }
            drop(stream);
        }
    });
    (thread_handle, receiver.recv().unwrap())
}
<|endoftext|>