<commit_before> #![feature(buf_stream)]

#[macro_use]
extern crate log;
extern crate url;
extern crate capnp;

pub mod yak_capnp;

use std::net::TcpStream;
use std::io::{self,BufStream,Write};
use capnp::serialize_packed;
use capnp::{MessageBuilder, MessageReader, MallocMessageBuilder, ReaderOptions};

use url::{Url,SchemeType,UrlParser};

use yak_capnp::{client_request,client_response};

fn yak_url_scheme(_scheme: &str) -> SchemeType {
  SchemeType::Relative(0)
}

pub struct Client {
  connection: BufStream<TcpStream>,
  space: String,
}

#[derive(Debug)]
pub enum YakError {
  UrlParseError(url::ParseError),
  InvalidUrl(Url),
  IoError(io::Error),
  CapnpError(capnp::Error),
  CapnpNotInSchema(capnp::NotInSchema),
  ProtocolError
}

#[derive(PartialEq,Eq,PartialOrd,Ord,Debug)]
pub struct Datum {
  pub content: Vec<u8>
}

pub struct Request {
  space: String,
  operation: Operation,
}
pub enum Operation {
  Truncate ,
  Read { key: Vec<u8> },
  Write { key: Vec<u8>, value: Vec<u8> },
}

impl Request {
  fn encode<B: MessageBuilder>(&self, message: &mut B) {
    match &self.operation {
      &Operation::Truncate => Self::encode_truncate(message, &self.space),
      &Operation::Read { key: ref key } => Self::encode_read(message, &self.space, &key),
      &Operation::Write { key: ref key, value: ref val } => Self::encode_write(message, &self.space, &key, &val),
    }
  }

  fn truncate(space: &str) -> Request {
    Request { space: space.to_string(), operation: Operation::Truncate }
  }

  fn read(space: &str, key: &[u8]) -> Request {
    Request { space: space.to_string(), operation: Operation::Read { key: key.to_vec() } }
  }

  fn write(space: &str, key: &[u8], value: &[u8]) -> Request {
    Request { space: space.to_string(), operation: Operation::Write { key: key.to_owned(), value: value.to_owned() } }
  }

  fn encode_truncate<B: MessageBuilder>(message: &mut B, space: &str) {
    let mut rec = message.init_root::<client_request::Builder>();
    let mut trunc = rec.init_truncate();
    trunc.set_space(space)
  }

  fn encode_write<B: MessageBuilder>(message: &mut B, space: &str, key: &[u8], val: &[u8]) {
    let mut rec = message.init_root::<client_request::Builder>();
    let mut req = rec.init_write();
    req.set_space(space);
    req.set_key(key);
    req.set_value(val);
  }

  fn encode_read<B: MessageBuilder>(message: &mut B, space: &str, key: &[u8]) {
      let mut rec = message.init_root::<client_request::Builder>();
      let mut req = rec.init_read();
      req.set_space(space);
      req.set_key(key)
  }

}

impl Client {
  pub fn connect(loc: &str) -> Result<Client, YakError> {
    let mut p = UrlParser::new();
    p.scheme_type_mapper(yak_url_scheme);
    let url = try!(p.parse(loc));
    debug!("yak:url: {:?}", url);
    let (addr, space) = match (url.domain(), url.port(), url.serialize_path()) {
      (Some(host), Some(port), Some(path)) => ((host, port), path),
      _ => return Err(YakError::InvalidUrl(url.clone()))
    };

    let sock = try!(TcpStream::connect(addr));
    debug!("connected:{:?}", sock);
    Ok(Client { connection: BufStream::new(sock), space: space })
  }

  fn expect_ok<R: MessageReader>(message: &R) -> Result<(), YakError> {
    let msg = try!(message.get_root::<client_response::Reader>());
    debug!("Got response");
    match try!(msg.which()) {
      client_response::Ok(()) => Ok(()),
      other => Err(YakError::ProtocolError)
    }
  }

  fn expect_datum_list<R: MessageReader>(message_reader: &R) -> Result<Vec<Datum>, YakError> {
    let msg = try!(message_reader.get_root::<client_response::Reader>());
    match try!(msg.which()) {
      client_response::OkData(d) => {
        debug!("Got response Data: ");
        let mut data = Vec::with_capacity(try!(d).len() as usize);
        for it in try!(d).iter() {
          let val : Vec<u8> = try!(it.get_value()).iter().map(|v|v.clone()).collect();
          data.push(Datum { content: val });
        }

        Ok(data)
      },
      other => Err(YakError::ProtocolError)
    }
  }

  fn send(&mut self, req: Request) -> Result<(), YakError> {
    let mut message = MallocMessageBuilder::new_default();
    req.encode(&mut message);
    try!(serialize_packed::write_message(&mut self.connection, &mut message));
    try!(self.connection.flush());
    Ok(())
  }

  pub fn truncate(&mut self) -> Result<(), YakError> {
    let req = Request::truncate(&self.space);
    try!(self.send(req));
    debug!("Waiting for response");

    let message_reader = try!(serialize_packed::read_message(&mut self.connection, ReaderOptions::new()));
    Self::expect_ok(&message_reader)
  }

  pub fn write(&mut self, key: &[u8], val: &[u8]) -> Result<(), YakError> {
    let req = Request::write(&self.space, key, val);
    try!(self.send(req));
    debug!("Waiting for response");

    let message_reader = try!(serialize_packed::read_message(&mut self.connection, ReaderOptions::new()));
    Self::expect_ok(&message_reader)
  }

  pub fn read(&mut self, key: &[u8]) -> Result<Vec<Datum>, YakError> {
    let req = Request::read(&self.space, key);
    try!(self.send(req));
    debug!("Waiting for response");

    let message_reader = try!(serialize_packed::read_message(&mut self.connection, ReaderOptions::new()));
    Self::expect_datum_list(&message_reader)
  }
}

impl From<url::ParseError> for YakError {
  fn from(err: url::ParseError) -> YakError {
    YakError::UrlParseError(err)
  }
}

impl From<io::Error> for YakError {
  fn from(err: io::Error) -> YakError {
    YakError::IoError(err)
  }
}

impl From<capnp::Error> for YakError {
  fn from(err: capnp::Error) -> YakError {
    YakError::CapnpError(err)
  }
}

impl From<capnp::NotInSchema> for YakError {
  fn from(err: capnp::NotInSchema) -> YakError {
    YakError::CapnpNotInSchema(err)
  }
}
<commit_msg>Factor out response reading code.<commit_after> #![feature(buf_stream)]

#[macro_use]
extern crate log;
extern crate url;
extern crate capnp;

pub mod yak_capnp;

use std::net::TcpStream;
use std::io::{self,BufStream,Write};
use capnp::serialize_packed;
use capnp::{MessageBuilder, MessageReader, MallocMessageBuilder, ReaderOptions};

use url::{Url,SchemeType,UrlParser};

use yak_capnp::{client_request,client_response};

fn yak_url_scheme(_scheme: &str) -> SchemeType {
  SchemeType::Relative(0)
}

pub struct Client {
  connection: BufStream<TcpStream>,
  space: String,
}

#[derive(Debug)]
pub enum YakError {
  UrlParseError(url::ParseError),
  InvalidUrl(Url),
  IoError(io::Error),
  CapnpError(capnp::Error),
  CapnpNotInSchema(capnp::NotInSchema),
  ProtocolError
}

#[derive(PartialEq,Eq,PartialOrd,Ord,Debug, Clone)]
pub struct Datum {
  pub content: Vec<u8>
}

pub struct Request {
  space: String,
  operation: Operation,
}
pub enum Operation {
  Truncate ,
  Read { key: Vec<u8> },
  Write { key: Vec<u8>, value: Vec<u8> },
}

impl Request {
  fn encode<B: MessageBuilder>(&self, message: &mut B) {
    match &self.operation {
      &Operation::Truncate => Self::encode_truncate(message, &self.space),
      &Operation::Read { key: ref key } => Self::encode_read(message, &self.space, &key),
      &Operation::Write { key: ref key, value: ref val } => Self::encode_write(message, &self.space, &key, &val),
    }
  }

  fn truncate(space: &str) -> Request {
    Request { space: space.to_string(), operation: Operation::Truncate }
  }

  fn read(space: &str, key: &[u8]) -> Request {
    Request { space: space.to_string(), operation: Operation::Read { key: key.to_vec() } }
  }

  fn write(space: &str, key: &[u8], value: &[u8]) -> Request {
    Request { space: space.to_string(), operation: Operation::Write { key: key.to_owned(), value: value.to_owned() } }
  }

  fn encode_truncate<B: MessageBuilder>(message: &mut B, space: &str) {
    let mut rec = message.init_root::<client_request::Builder>();
    let mut trunc = rec.init_truncate();
    trunc.set_space(space)
  }

  fn encode_write<B: MessageBuilder>(message: &mut B, space: &str, key: &[u8], val: &[u8]) {
    let mut rec = message.init_root::<client_request::Builder>();
    let mut req = rec.init_write();
    req.set_space(space);
    req.set_key(key);
    req.set_value(val);
  }

  fn encode_read<B: MessageBuilder>(message: &mut B, space: &str, key: &[u8]) {
      let mut rec = message.init_root::<client_request::Builder>();
      let mut req = rec.init_read();
      req.set_space(space);
      req.set_key(key)
  }
}

pub enum Response {
  Okay,
  OkayData(Vec<Datum>)
}

impl Response {
  fn decode<R: MessageReader>(message: &R) -> Result<Response, YakError> {
    let msg = try!(message.get_root::<client_response::Reader>());
    match try!(msg.which()) {
      client_response::Ok(()) => Ok(Response::Okay),
      client_response::OkData(d) => {
        debug!("Got response Data: ");
        let mut data = Vec::with_capacity(try!(d).len() as usize);
        for it in try!(d).iter() {
          let val : Vec<u8> = try!(it.get_value()).iter().map(|v|v.clone()).collect();
          data.push(Datum { content: val });
        }
        Ok(Response::OkayData(data))
      }
    }
  }

  pub fn expect_ok(&self) -> Result<(), YakError> {
    match self {
      &Response::Okay => Ok(()),
      &Response::OkayData(_) => Err(YakError::ProtocolError)
    }
  }

  pub fn expect_datum_list(&self) -> Result<Vec<Datum>, YakError> {
    match self {
      &Response::OkayData(ref result) => Ok(result.clone()),
      &Response::Okay => Err(YakError::ProtocolError)
    }
  }
}

impl Client {
  pub fn connect(loc: &str) -> Result<Client, YakError> {
    let mut p = UrlParser::new();
    p.scheme_type_mapper(yak_url_scheme);
    let url = try!(p.parse(loc));
    debug!("yak:url: {:?}", url);
    let (addr, space) = match (url.domain(), url.port(), url.serialize_path()) {
      (Some(host), Some(port), Some(path)) => ((host, port), path),
      _ => return Err(YakError::InvalidUrl(url.clone()))
    };

    let sock = try!(TcpStream::connect(addr));
    debug!("connected:{:?}", sock);
    Ok(Client { connection: BufStream::new(sock), space: space })
  }

  fn expect_ok<R: MessageReader>(message: &R) -> Result<(), YakError> {
    let msg = try!(message.get_root::<client_response::Reader>());
    debug!("Got response");
    match try!(msg.which()) {
      client_response::Ok(()) => Ok(()),
      other => Err(YakError::ProtocolError)
    }
  }

  fn expect_datum_list<R: MessageReader>(message_reader: &R) -> Result<Vec<Datum>, YakError> {
    let msg = try!(message_reader.get_root::<client_response::Reader>());
    match try!(msg.which()) {
      client_response::OkData(d) => {
        debug!("Got response Data: ");
        let mut data = Vec::with_capacity(try!(d).len() as usize);
        for it in try!(d).iter() {
          let val : Vec<u8> = try!(it.get_value()).iter().map(|v|v.clone()).collect();
          data.push(Datum { content: val });
        }

        Ok(data)
      },
      other => Err(YakError::ProtocolError)
    }
  }

  fn send(&mut self, req: Request) -> Result<(), YakError> {
    let mut message = MallocMessageBuilder::new_default();
    req.encode(&mut message);
    try!(serialize_packed::write_message(&mut self.connection, &mut message));
    try!(self.connection.flush());
    Ok(())
  }

  fn read_response(&mut self) -> Result<Response,YakError> {
    let message_reader =
      try!(serialize_packed::read_message(
	&mut self.connection, ReaderOptions::new()));
    let resp = try!(Response::decode(&message_reader));
    Ok(resp)
  }

  pub fn truncate(&mut self) -> Result<(), YakError> {
    let req = Request::truncate(&self.space);
    try!(self.send(req));
    debug!("Waiting for response");
    try!(self.read_response()).expect_ok()
  }

  pub fn write(&mut self, key: &[u8], val: &[u8]) -> Result<(), YakError> {
    let req = Request::write(&self.space, key, val);
    try!(self.send(req));
    debug!("Waiting for response");
    try!(self.read_response()).expect_ok()
  }

  pub fn read(&mut self, key: &[u8]) -> Result<Vec<Datum>, YakError> {
    let req = Request::read(&self.space, key);
    try!(self.send(req));
    debug!("Waiting for response");

    try!(self.read_response()).expect_datum_list()
  }
}

impl From<url::ParseError> for YakError {
  fn from(err: url::ParseError) -> YakError {
    YakError::UrlParseError(err)
  }
}

impl From<io::Error> for YakError {
  fn from(err: io::Error) -> YakError {
    YakError::IoError(err)
  }
}

impl From<capnp::Error> for YakError {
  fn from(err: capnp::Error) -> YakError {
    YakError::CapnpError(err)
  }
}

impl From<capnp::NotInSchema> for YakError {
  fn from(err: capnp::NotInSchema) -> YakError {
    YakError::CapnpNotInSchema(err)
  }
}
<|endoftext|>


<commit_before>//! Conversion methods for JSON values.

use rustwlc::{Geometry, Point, Size};
use std::collections::BTreeMap;

use rlua;
use rustc_serialize::json::{Json, ToJson};

/// Converts a Json map into a Lua value
pub fn json_to_lua<'lua>(lua: &'lua rlua::Lua, json: Json)
                         -> Result<rlua::Value<'lua>, rlua::Error> {
    use rlua::Value;
    match json {
        Json::String(val)  => Ok(Value::String(lua.create_string(&val))),
        Json::Boolean(val) => Ok(Value::Boolean(val)),
        Json::F64(val)     => Ok(Value::Number(val)),
        Json::I64(val)     => Ok(Value::Number(val as f64)),
        Json::U64(val)     => Ok(Value::Number(val as f64)),
        Json::Null         => Ok(Value::Nil),
        Json::Array(vals)  => {
            let mut new_vals = Vec::with_capacity(vals.len());
            for (idx, val) in vals.into_iter().enumerate() {
                new_vals.push((idx, json_to_lua(lua, val)?));
            }
            lua.create_table_from(new_vals)
                .map(|table| Value::Table(table))
        },
        Json::Object(vals) => {
            let mut new_vals = BTreeMap::new();
            for (k, v) in vals {
                new_vals.insert(k, json_to_lua(lua, v)?);
            }
            lua.create_table_from(new_vals)
                .map(|table| Value::Table(table))
        }
    }
}

/// Converts an `rlua::Value` to a `Json`.
pub fn lua_to_json<'lua>(lua_value: rlua::Value<'lua>)
                         -> Result<Json, rlua::Error> {
    use rlua::Value::*;
    use rlua::Error;
    match lua_value {
        Nil => Ok(Json::Null),
        String(val) => Ok(Json::String(val.to_str().unwrap().into())),
        Number(val) => Ok(Json::F64(val)),
        Integer(val) => Ok(Json::I64(val as _)),
        Boolean(val) => Ok(Json::Boolean(val)),
        Table(arr) => lua_table_to_json(arr),
        _ => Err(Error::RuntimeError(
            format!("Did not expect {:#?}", lua_value)))
    }
}

/// Convert an AnyLuaValue to a Json array using numerical indicies.
///
/// # Result
/// This function returns an Err if the Lua object has a non-String key.
fn lua_table_to_json<'lua>(table: rlua::Table<'lua>)
                               -> Result<Json, rlua::Error> {
    use rlua::Value;
    use rlua::Error;
    // Check if every key is a number
    let mut counter = 0.0; // Account for first index?

    for entry in table.clone().pairs::<Value, Value>() {
        let (key, _ )= entry?;
        match key {
            Value::Number(num) => {
                counter += num;
            }
            Value::String(_) => {
                break;
            }
            // Non-string keys are not allowed
            _ => return Err(Error::FromLuaConversionError {
                from: "Lua table",
                to: "JSON object",
                message: Some(format!("Could not convert {:#?}", table))
            })
        }
    }

    let len = table.len().unwrap();
    // Gauss' trick
    let desired_sum = (len * (len + 1)) / 2;
    if counter != desired_sum as f64 {
        return lua_object_to_json(table)
    }

    let mut json_arr: Vec<Json> = Vec::with_capacity(len as _);

    for entry in table.pairs::<Value, Value>() {
        let (_, val) = entry?;
        let lua_val = lua_to_json(val)?;
        json_arr.push(lua_val);
    }
    Ok(Json::Array(json_arr))
}

/// Converts an AnyLuaValue object to a Json object.
///
/// Will return an Err if the Lua object uses non-String keys.
fn lua_object_to_json<'lua>(table: rlua::Table<'lua>)
                          -> Result<Json, rlua::Error> {
    use rlua::{Value, Error};
    let mut json_obj: BTreeMap<String, Json> = BTreeMap::new();

    for entry in table.clone().pairs::<Value, Value>() {
        let (key, val) = entry?;
        match key {
            Value::String(text) => {
                let text = text.to_str().unwrap().into();
                json_obj.insert(text, lua_to_json(val)?);
            },
            Value::Number(ix) => {
                json_obj.insert(ix.to_string(), lua_to_json(val)?);
            }
            val => {
                return Err(Error::RuntimeError(
                    format!("Did not expect {:#?} as key in {:#?}",
                            val, table)))
            }
        }
    }
    Ok(Json::Object(json_obj))
}

pub fn size_to_json(size: Size) -> Json {
    let mut map = BTreeMap::new();
    map.insert("w".into(), size.w.to_json());
    map.insert("h".into(), size.h.to_json());
    map.to_json()
}

pub fn point_to_json(point: Point) -> Json {
    let mut map = BTreeMap::new();
    map.insert("x".into(), point.x.to_json());
    map.insert("y".into(), point.y.to_json());
    map.to_json()
}

pub fn geometry_to_json(geometry: Geometry) -> Json {
    let mut map = BTreeMap::new();
    let origin = point_to_json(geometry.origin);
    let size = size_to_json(geometry.size);
    map.insert("origin".into(), origin);
    map.insert("size".into(), size);
    map.to_json()
}
<commit_msg>Turned unwraps into error propogation<commit_after>//! Conversion methods for JSON values.

use rustwlc::{Geometry, Point, Size};
use std::collections::BTreeMap;

use rlua;
use rustc_serialize::json::{Json, ToJson};

/// Converts a Json map into a Lua value
pub fn json_to_lua<'lua>(lua: &'lua rlua::Lua, json: Json)
                         -> Result<rlua::Value<'lua>, rlua::Error> {
    use rlua::Value;
    match json {
        Json::String(val)  => Ok(Value::String(lua.create_string(&val))),
        Json::Boolean(val) => Ok(Value::Boolean(val)),
        Json::F64(val)     => Ok(Value::Number(val)),
        Json::I64(val)     => Ok(Value::Number(val as f64)),
        Json::U64(val)     => Ok(Value::Number(val as f64)),
        Json::Null         => Ok(Value::Nil),
        Json::Array(vals)  => {
            let mut new_vals = Vec::with_capacity(vals.len());
            for (idx, val) in vals.into_iter().enumerate() {
                new_vals.push((idx, json_to_lua(lua, val)?));
            }
            lua.create_table_from(new_vals)
                .map(|table| Value::Table(table))
        },
        Json::Object(vals) => {
            let mut new_vals = BTreeMap::new();
            for (k, v) in vals {
                new_vals.insert(k, json_to_lua(lua, v)?);
            }
            lua.create_table_from(new_vals)
                .map(|table| Value::Table(table))
        }
    }
}

/// Converts an `rlua::Value` to a `Json`.
pub fn lua_to_json<'lua>(lua_value: rlua::Value<'lua>)
                         -> Result<Json, rlua::Error> {
    use rlua::Value::*;
    use rlua::Error;
    match lua_value {
        Nil => Ok(Json::Null),
        String(val) =>
            val.to_str()
            .map_err(|err| Error::RuntimeError(format!("Bad string: {:#?}",
                                                       err)))
                .map(|string| Json::String(string.into())),
        Number(val) => Ok(Json::F64(val)),
        Integer(val) => Ok(Json::I64(val as _)),
        Boolean(val) => Ok(Json::Boolean(val)),
        Table(arr) => lua_table_to_json(arr),
        _ => Err(Error::RuntimeError(
            format!("Did not expect {:#?}", lua_value)))
    }
}

/// Convert an AnyLuaValue to a Json array using numerical indicies.
///
/// # Result
/// This function returns an Err if the Lua object has a non-String key.
fn lua_table_to_json<'lua>(table: rlua::Table<'lua>)
                               -> Result<Json, rlua::Error> {
    use rlua::Value;
    use rlua::Error;
    // Check if every key is a number
    let mut counter = 0.0; // Account for first index?

    for entry in table.clone().pairs::<Value, Value>() {
        let (key, _ )= entry?;
        match key {
            Value::Number(num) => {
                counter += num;
            }
            Value::String(_) => {
                break;
            }
            // Non-string keys are not allowed
            _ => return Err(Error::FromLuaConversionError {
                from: "Lua table",
                to: "JSON object",
                message: Some(format!("Could not convert {:#?}", table))
            })
        }
    }

    let len = table.len()?;
    // Gauss' trick
    let desired_sum = (len * (len + 1)) / 2;
    if counter != desired_sum as f64 {
        return lua_object_to_json(table)
    }

    let mut json_arr: Vec<Json> = Vec::with_capacity(len as _);

    for entry in table.pairs::<Value, Value>() {
        let (_, val) = entry?;
        let lua_val = lua_to_json(val)?;
        json_arr.push(lua_val);
    }
    Ok(Json::Array(json_arr))
}

/// Converts an AnyLuaValue object to a Json object.
///
/// Will return an Err if the Lua object uses non-String keys.
fn lua_object_to_json<'lua>(table: rlua::Table<'lua>)
                          -> Result<Json, rlua::Error> {
    use rlua::{Value, Error};
    let mut json_obj: BTreeMap<String, Json> = BTreeMap::new();

    for entry in table.clone().pairs::<Value, Value>() {
        let (key, val) = entry?;
        match key {
            Value::String(text) => {
                let text = text.to_str()?.into();
                json_obj.insert(text, lua_to_json(val)?);
            },
            Value::Number(ix) => {
                json_obj.insert(ix.to_string(), lua_to_json(val)?);
            }
            val => {
                return Err(Error::RuntimeError(
                    format!("Did not expect {:#?} as key in {:#?}",
                            val, table)))
            }
        }
    }
    Ok(Json::Object(json_obj))
}

pub fn size_to_json(size: Size) -> Json {
    let mut map = BTreeMap::new();
    map.insert("w".into(), size.w.to_json());
    map.insert("h".into(), size.h.to_json());
    map.to_json()
}

pub fn point_to_json(point: Point) -> Json {
    let mut map = BTreeMap::new();
    map.insert("x".into(), point.x.to_json());
    map.insert("y".into(), point.y.to_json());
    map.to_json()
}

pub fn geometry_to_json(geometry: Geometry) -> Json {
    let mut map = BTreeMap::new();
    let origin = point_to_json(geometry.origin);
    let size = size_to_json(geometry.size);
    map.insert("origin".into(), origin);
    map.insert("size".into(), size);
    map.to_json()
}
<|endoftext|>


<commit_before>#![feature(io, path)]

extern crate icepick;

use icepick::search::Search;
use std::old_io::{File, BufferedReader};

#[allow(dead_code)]
fn main() {
    let lines = read_lines("benches/30000.txt".to_string());

    let mut search = Search::blank(&lines, None, 20);

    search = search.append_to_search("t").append_to_search("o").append_to_search("a").append_to_search("w").append_to_search("c").backspace().backspace().backspace().append_to_search("w").backspace().append_to_search("a");

    println!("\n{}", search.selection().unwrap_or("None".to_string()));
}

fn read_lines(fname: String) -> Vec<String> {
    let path = Path::new(fname);
    let mut file = BufferedReader::new(File::open(&path));
    let mut result = Vec::new();
    loop {
        match file.read_line() {
            Ok(line) => result.push(line),
            Err(_) => break,
        }
    }
    result
}
<commit_msg>Use new io in bench_main<commit_after>#![feature(io, path)]

extern crate icepick;

use icepick::search::Search;
use std::io::prelude::*;
use std::io::BufReader;
use std::fs::{File, OpenOptions};
use std::path::Path;

#[allow(dead_code)]
fn main() {
    let lines = read_lines("benches/30000.txt");

    let mut search = Search::blank(&lines, None, 20);

    search = search.append_to_search("t").append_to_search("o").append_to_search("a").append_to_search("w").append_to_search("c").backspace().backspace().backspace().append_to_search("w").backspace().append_to_search("a");

    println!("\n{}", search.selection().unwrap_or("None".to_string()));
}

fn read_lines(fname: &str) -> Vec<String> {
    let path = Path::new(fname);
    let mut file = BufReader::new(OpenOptions::new().read(true).write(true).open(&path).unwrap());
    let mut result = Vec::new();
    loop {
        let mut buf = String::new();
        match file.read_line(&mut buf) {
            Ok(_) => result.push(buf),
            Err(_) => break,
        }
    }
    result
}
<|endoftext|>


<commit_before>use {ffi, Return, Result, Raii, Handle, Statement};
use super::types::FixedSizedType;

/// Allows types to be used with `Statement::bind_parameter`
pub unsafe trait InputParameter {
    fn c_data_type(&self) -> ffi::SqlCDataType;
    fn column_size(&self) -> ffi::SQLULEN;
    fn decimal_digits(&self) -> ffi::SQLSMALLINT;
    fn value_ptr(&self) -> ffi::SQLPOINTER;
    fn indicator(&self) -> ffi::SQLLEN;
}

impl<'a, 'b, S> Statement<'a, 'b, S> {
    /// Binds a parameter to a parameter marker in an SQL statement.
    ///
    /// # Result
    /// This method will destroy the statement and create a new one which may not outlive the bound
    /// parameter. This is to ensure that the statement will not derefernce an invalid pointer
    /// during execution.
    ///
    /// # Arguments
    /// * `parameter_index` - Index of the marker to bind to the parameter. Starting at `1`
    /// * `value` - Reference to bind to the marker
    ///
    /// # Example
    /// ```
    /// # use odbc::*;
    /// # fn do_odbc_stuff() -> std::result::Result<(), Box<std::error::Error>> {
    ///     let env = Environment::new()?.set_odbc_version_3()?;
    ///     let conn = DataSource::with_parent(&env)?.connect("TestDataSource", "", "")?;
    ///     let stmt = Statement::with_parent(&conn)?;
    ///     let param = 1968;
    ///     let stmt = stmt.bind_parameter(1, &param)?;
    ///     let sql_text = "SELECT TITLE FROM MOVIES WHERE YEAR = ?";
    ///     if let Data(mut stmt) = stmt.exec_direct(sql_text)? {
    ///         // ...
    ///     }
    /// #   Ok(())
    /// # }
    /// ```
    pub fn bind_parameter<'c, T>(mut self,
                                 parameter_index: u16,
                                 value: &'c T)
                                 -> Result<Statement<'a, 'c, S>>
        where T: InputParameter,
              T: ?Sized,
              'b: 'c
    {
        self.raii.bind_input_parameter(parameter_index, value).into_result(&self)?;
        Ok(self)
    }
}

impl Raii<ffi::Stmt> {
    fn bind_input_parameter<T>(&mut self, parameter_index: u16, value: &T) -> Return<()>
        where T: InputParameter,
              T: ?Sized
    {
        match unsafe {
                  ffi::SQLBindParameter(
                self.handle(),
                parameter_index,
                ffi::SQL_PARAM_INPUT,
                value.c_data_type(),
                ffi::SQL_UNKNOWN_TYPE,
                value.column_size(),
                value.decimal_digits(),
                value.value_ptr(),
                0, // buffer length
                &value.indicator() as * const ffi::SQLLEN as * mut ffi::SQLLEN// str len or ind ptr
            )
              } {
            ffi::SQL_SUCCESS => Return::Success(()),
            ffi::SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(()),
            ffi::SQL_ERROR => Return::Error,
            r => panic!("Unexpected return from SQLBindParameter: {:?}", r),
        }
    }
}

unsafe impl InputParameter for str {
    fn c_data_type(&self) -> ffi::SqlCDataType {
        ffi::SQL_C_CHAR
    }

    fn column_size(&self) -> ffi::SQLULEN {
        self.as_bytes().len() as ffi::SQLULEN
    }

    fn decimal_digits(&self) -> ffi::SQLSMALLINT {
        0
    }

    fn value_ptr(&self) -> ffi::SQLPOINTER {
        self.as_bytes().as_ptr() as ffi::SQLPOINTER
    }

    fn indicator(&self) -> ffi::SQLLEN {
        self.as_bytes().len() as ffi::SQLLEN
    }
}

unsafe impl InputParameter for String {
    fn c_data_type(&self) -> ffi::SqlCDataType {
        ffi::SQL_C_CHAR
    }

    fn column_size(&self) -> ffi::SQLULEN {
        self.as_bytes().len() as ffi::SQLULEN
    }

    fn decimal_digits(&self) -> ffi::SQLSMALLINT {
        0
    }

    fn value_ptr(&self) -> ffi::SQLPOINTER {
        self.as_bytes().as_ptr() as ffi::SQLPOINTER
    }

    fn indicator(&self) -> ffi::SQLLEN {
        self.as_bytes().len() as ffi::SQLLEN
    }
}

unsafe impl<T> InputParameter for T
    where T: FixedSizedType
{
    fn c_data_type(&self) -> ffi::SqlCDataType {
        T::c_data_type()
    }

    fn column_size(&self) -> ffi::SQLULEN {
        use std::mem::size_of;
        size_of::<Self>() as ffi::SQLULEN
    }

    fn decimal_digits(&self) -> ffi::SQLSMALLINT {
        0
    }

    fn value_ptr(&self) -> ffi::SQLPOINTER {
        self as *const Self as ffi::SQLPOINTER
    }

    fn indicator(&self) -> ffi::SQLLEN {
        0
    }
}

<commit_msg>fix indentation in code sample<commit_after>use {ffi, Return, Result, Raii, Handle, Statement};
use super::types::FixedSizedType;

/// Allows types to be used with `Statement::bind_parameter`
pub unsafe trait InputParameter {
    fn c_data_type(&self) -> ffi::SqlCDataType;
    fn column_size(&self) -> ffi::SQLULEN;
    fn decimal_digits(&self) -> ffi::SQLSMALLINT;
    fn value_ptr(&self) -> ffi::SQLPOINTER;
    fn indicator(&self) -> ffi::SQLLEN;
}

impl<'a, 'b, S> Statement<'a, 'b, S> {
    /// Binds a parameter to a parameter marker in an SQL statement.
    ///
    /// # Result
    /// This method will destroy the statement and create a new one which may not outlive the bound
    /// parameter. This is to ensure that the statement will not derefernce an invalid pointer
    /// during execution.
    ///
    /// # Arguments
    /// * `parameter_index` - Index of the marker to bind to the parameter. Starting at `1`
    /// * `value` - Reference to bind to the marker
    ///
    /// # Example
    /// ```
    /// # use odbc::*;
    /// # fn do_odbc_stuff() -> std::result::Result<(), Box<std::error::Error>> {
    /// let env = Environment::new()?.set_odbc_version_3()?;
    /// let conn = DataSource::with_parent(&env)?.connect("TestDataSource", "", "")?;
    /// let stmt = Statement::with_parent(&conn)?;
    /// let param = 1968;
    /// let stmt = stmt.bind_parameter(1, &param)?;
    /// let sql_text = "SELECT TITLE FROM MOVIES WHERE YEAR = ?";
    /// if let Data(mut stmt) = stmt.exec_direct(sql_text)? {
    ///     // ...
    /// }
    /// #   Ok(())
    /// # }
    /// ```
    pub fn bind_parameter<'c, T>(mut self,
                                 parameter_index: u16,
                                 value: &'c T)
                                 -> Result<Statement<'a, 'c, S>>
        where T: InputParameter,
              T: ?Sized,
              'b: 'c
    {
        self.raii.bind_input_parameter(parameter_index, value).into_result(&self)?;
        Ok(self)
    }
}

impl Raii<ffi::Stmt> {
    fn bind_input_parameter<T>(&mut self, parameter_index: u16, value: &T) -> Return<()>
        where T: InputParameter,
              T: ?Sized
    {
        match unsafe {
                  ffi::SQLBindParameter(
                self.handle(),
                parameter_index,
                ffi::SQL_PARAM_INPUT,
                value.c_data_type(),
                ffi::SQL_UNKNOWN_TYPE,
                value.column_size(),
                value.decimal_digits(),
                value.value_ptr(),
                0, // buffer length
                &value.indicator() as * const ffi::SQLLEN as * mut ffi::SQLLEN// str len or ind ptr
            )
              } {
            ffi::SQL_SUCCESS => Return::Success(()),
            ffi::SQL_SUCCESS_WITH_INFO => Return::SuccessWithInfo(()),
            ffi::SQL_ERROR => Return::Error,
            r => panic!("Unexpected return from SQLBindParameter: {:?}", r),
        }
    }
}

unsafe impl InputParameter for str {
    fn c_data_type(&self) -> ffi::SqlCDataType {
        ffi::SQL_C_CHAR
    }

    fn column_size(&self) -> ffi::SQLULEN {
        self.as_bytes().len() as ffi::SQLULEN
    }

    fn decimal_digits(&self) -> ffi::SQLSMALLINT {
        0
    }

    fn value_ptr(&self) -> ffi::SQLPOINTER {
        self.as_bytes().as_ptr() as ffi::SQLPOINTER
    }

    fn indicator(&self) -> ffi::SQLLEN {
        self.as_bytes().len() as ffi::SQLLEN
    }
}

unsafe impl InputParameter for String {
    fn c_data_type(&self) -> ffi::SqlCDataType {
        ffi::SQL_C_CHAR
    }

    fn column_size(&self) -> ffi::SQLULEN {
        self.as_bytes().len() as ffi::SQLULEN
    }

    fn decimal_digits(&self) -> ffi::SQLSMALLINT {
        0
    }

    fn value_ptr(&self) -> ffi::SQLPOINTER {
        self.as_bytes().as_ptr() as ffi::SQLPOINTER
    }

    fn indicator(&self) -> ffi::SQLLEN {
        self.as_bytes().len() as ffi::SQLLEN
    }
}

unsafe impl<T> InputParameter for T
    where T: FixedSizedType
{
    fn c_data_type(&self) -> ffi::SqlCDataType {
        T::c_data_type()
    }

    fn column_size(&self) -> ffi::SQLULEN {
        use std::mem::size_of;
        size_of::<Self>() as ffi::SQLULEN
    }

    fn decimal_digits(&self) -> ffi::SQLSMALLINT {
        0
    }

    fn value_ptr(&self) -> ffi::SQLPOINTER {
        self as *const Self as ffi::SQLPOINTER
    }

    fn indicator(&self) -> ffi::SQLLEN {
        0
    }
}

<|endoftext|>


<commit_before>// Copyright (c) 2014 Robert Clipsham <robert@octarineparrot.com>
//
// 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.

extern crate libc;

use std::collections::{RingBuf, Deque};
use std::cmp;
use std::io::{IoResult, IoError};
use std::mem;
use std::option::{Option, Some};
use std::sync::Arc;

use bindings::bpf;
use packet::ethernet::{EthernetHeader, MutableEthernetHeader};
use datalink::{DataLinkChannelType, Layer2, Layer3};
use internal;
use util::NetworkInterface;

// NOTE buffer must be word aligned.
pub fn datalink_channel(network_interface: &NetworkInterface,
                        write_buffer_size: uint,
                        read_buffer_size: uint,
                        channel_type: DataLinkChannelType)
    -> IoResult<(DataLinkSenderImpl, DataLinkReceiverImpl)> {
    #[cfg(target_os = "freebsd")]
    fn get_fd() -> libc::c_int {
        unsafe {
            libc::open("/dev/bpf".to_c_str().as_ptr(), libc::O_RDWR, 0)
        }
    }

    #[cfg(target_os = "macos")]
    fn get_fd() -> libc::c_int {
        // FIXME This is an arbitrary number of attempts
        for i in range(0, 1_000i) {
            let fd = unsafe {
                let file_name = format!("/dev/bpf{}", i);
                libc::open(file_name.to_c_str().as_ptr(), libc::O_RDWR, 0)
            };
            if fd != -1 {
                return fd;
            }
        }
        return -1;
    }

    #[cfg(target_os = "freebsd")]
    fn set_feedback(fd: libc::c_int) -> Result<(), IoError> {
        let one: libc::c_uint = 1;
        if unsafe { bpf::ioctl(fd, bpf::BIOCFEEDBACK, &one) } == -1 {
            let err = IoError::last_error();
            unsafe { libc::close(fd); }
            return Err(err);
        }
        Ok(())
    }

    #[cfg(target_os = "macos")]
    fn set_feedback(_fd: libc::c_int) -> Result<(), IoError> {
        Ok(())
    }

    match channel_type {
        Layer2 => (),
        Layer3(_) => unimplemented!(),
    }

    let fd = get_fd();
    if fd == -1 {
        return Err(IoError::last_error());
    }
    let mut iface: bpf::ifreq = unsafe { mem::zeroed() };
    let mut i = 0;
    for c in network_interface.name.as_slice().bytes() {
        iface.ifr_name[i] = c as i8;
        i += 1;
    }

    let buflen = read_buffer_size as libc::c_uint;
    // NOTE Buffer length must be set before binding to an interface
    //      otherwise this will return Invalid Argument
    if unsafe { bpf::ioctl(fd, bpf::BIOCSBLEN, &buflen) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // Set the interface to use
    if unsafe { bpf::ioctl(fd, bpf::BIOCSETIF, &iface) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // Return from read as soon as packets are available - don't wait to fill the buffer
    let one: libc::c_uint = 1;
    if unsafe { bpf::ioctl(fd, bpf::BIOCIMMEDIATE, &one) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    let mut header_size = 0;

    // Get the device type
    let mut dlt: libc::c_uint = 0;
    if unsafe { bpf::ioctl(fd, bpf::BIOCGDLT, &mut dlt) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // The loopback IoError::last_error()device does weird things
    // FIXME This should really just be another L2 packet header type
    if dlt == bpf::DLT_NULL {
        header_size = 4;

        // Allow packets to be read back after they are written
        match set_feedback(fd) {
            Err(e) => return Err(e),
            _ => ()
        }

    } else {
        // Don't fill in source MAC
        if unsafe { bpf::ioctl(fd, bpf::BIOCSHDRCMPLT, &one) } == -1 {
            let err = IoError::last_error();
            unsafe { libc::close(fd); }
            return Err(err);
        }
    }

    let fd = Arc::new(internal::FileDesc { fd: fd });
    let sender = DataLinkSenderImpl {
        fd: fd.clone(),
        write_buffer: Vec::from_elem(write_buffer_size, 0u8),
        header_size: header_size,
    };
    let receiver = DataLinkReceiverImpl {
        fd: fd,
        read_buffer: Vec::from_elem(write_buffer_size, 0u8),
        header_size: header_size,
    };

    Ok((sender, receiver))
}

pub struct DataLinkSenderImpl {
    fd: Arc<internal::FileDesc>,
    write_buffer: Vec<u8>,
    header_size: uint,
}

impl DataLinkSenderImpl {
    pub fn build_and_send(&mut self, num_packets: uint, packet_size: uint,
                          func: |MutableEthernetHeader| -> ()) -> Option<IoResult<()>> {
        let len = num_packets * (packet_size + self.header_size);
        if len >= self.write_buffer.len() {
            None
        } else {
            let min = cmp::min(self.write_buffer.len(), len);
            for chunk in self.write_buffer.mut_slice_to(min)
                                          .mut_chunks(packet_size + self.header_size) {
                // If we're sending on the loopback device, the first 4 bytes must be set to
                // AF_INET
                if self.header_size == 4 {
                    unsafe {
                        *(chunk.as_mut_ptr() as *mut u32) = libc::AF_INET as u32;
                    }
                }
                {
                    let eh = MutableEthernetHeader::new(chunk.mut_slice_from(self.header_size));
                    func(eh);
                }
                match unsafe { libc::write(self.fd.fd,
                                           chunk.as_ptr() as *const libc::c_void,
                                           chunk.len() as libc::size_t) } {
                    len if len == -1 => return Some(Err(IoError::last_error())),
                    _ => ()
                }
            }
            Some(Ok(()))
        }
    }
}

pub struct DataLinkReceiverImpl {
    fd: Arc<internal::FileDesc>,
    read_buffer: Vec<u8>,
    header_size: uint,
}

impl DataLinkReceiverImpl {
    pub fn iter<'a>(&'a mut self) -> DataLinkChannelIteratorImpl<'a> {
        let buflen = self.read_buffer.len();
        DataLinkChannelIteratorImpl {
            pc: self,
            // Enough room for minimally sized packets without reallocating
            packets: RingBuf::with_capacity(buflen / 64)
        }
    }
}

pub struct DataLinkChannelIteratorImpl<'a> {
    pc: &'a mut DataLinkReceiverImpl,
    packets: RingBuf<(uint, uint)>,
}

impl<'a> DataLinkChannelIteratorImpl<'a> {
    pub fn next<'c>(&'c mut self) -> IoResult<EthernetHeader<'c>> {
        if self.packets.is_empty() {
            let buflen = match unsafe {
                libc::read(self.pc.fd.fd,
                           self.pc.read_buffer.as_ptr() as *mut libc::c_void,
                           self.pc.read_buffer.len() as libc::size_t)
            } {
                len if len > 0 => len,
                _ => return Err(IoError::last_error())
            };
            let mut ptr = self.pc.read_buffer.as_mut_ptr();
            let end = unsafe { self.pc.read_buffer.as_ptr().offset(buflen as int) };
            while (ptr as *const u8) < end {
                unsafe {
                    let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                    let start = ptr as int +
                                (*packet).bh_hdrlen as int -
                                self.pc.read_buffer.as_ptr() as int;
                    self.packets.push((start as uint + self.pc.header_size,
                                      (*packet).bh_caplen as uint - self.pc.header_size));
                    let offset = (*packet).bh_hdrlen as int + (*packet).bh_caplen as int;
                    ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                }
            }
        }
        let (start, len) = self.packets.pop_front().unwrap();
        Ok(EthernetHeader::new(self.pc.read_buffer.slice(start, start + len)))
    }
}

<commit_msg>Fix incorrect buffer sizes with bpf<commit_after>// Copyright (c) 2014 Robert Clipsham <robert@octarineparrot.com>
//
// 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.

extern crate libc;

use std::collections::{RingBuf, Deque};
use std::cmp;
use std::io::{IoResult, IoError};
use std::mem;
use std::option::{Option, Some};
use std::sync::Arc;

use bindings::bpf;
use packet::ethernet::{EthernetHeader, MutableEthernetHeader};
use datalink::{DataLinkChannelType, Layer2, Layer3};
use internal;
use util::NetworkInterface;

// NOTE buffer must be word aligned.
pub fn datalink_channel(network_interface: &NetworkInterface,
                        write_buffer_size: uint,
                        read_buffer_size: uint,
                        channel_type: DataLinkChannelType)
    -> IoResult<(DataLinkSenderImpl, DataLinkReceiverImpl)> {
    #[cfg(target_os = "freebsd")]
    fn get_fd() -> libc::c_int {
        unsafe {
            libc::open("/dev/bpf".to_c_str().as_ptr(), libc::O_RDWR, 0)
        }
    }

    #[cfg(target_os = "macos")]
    fn get_fd() -> libc::c_int {
        // FIXME This is an arbitrary number of attempts
        for i in range(0, 1_000i) {
            let fd = unsafe {
                let file_name = format!("/dev/bpf{}", i);
                libc::open(file_name.to_c_str().as_ptr(), libc::O_RDWR, 0)
            };
            if fd != -1 {
                return fd;
            }
        }
        return -1;
    }

    #[cfg(target_os = "freebsd")]
    fn set_feedback(fd: libc::c_int) -> Result<(), IoError> {
        let one: libc::c_uint = 1;
        if unsafe { bpf::ioctl(fd, bpf::BIOCFEEDBACK, &one) } == -1 {
            let err = IoError::last_error();
            unsafe { libc::close(fd); }
            return Err(err);
        }
        Ok(())
    }

    #[cfg(target_os = "macos")]
    fn set_feedback(_fd: libc::c_int) -> Result<(), IoError> {
        Ok(())
    }

    match channel_type {
        Layer2 => (),
        Layer3(_) => unimplemented!(),
    }

    let fd = get_fd();
    if fd == -1 {
        return Err(IoError::last_error());
    }
    let mut iface: bpf::ifreq = unsafe { mem::zeroed() };
    let mut i = 0;
    for c in network_interface.name.as_slice().bytes() {
        iface.ifr_name[i] = c as i8;
        i += 1;
    }

    let buflen = read_buffer_size as libc::c_uint;
    // NOTE Buffer length must be set before binding to an interface
    //      otherwise this will return Invalid Argument
    if unsafe { bpf::ioctl(fd, bpf::BIOCSBLEN, &buflen) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // Set the interface to use
    if unsafe { bpf::ioctl(fd, bpf::BIOCSETIF, &iface) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // Return from read as soon as packets are available - don't wait to fill the buffer
    let one: libc::c_uint = 1;
    if unsafe { bpf::ioctl(fd, bpf::BIOCIMMEDIATE, &one) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    let mut header_size = 0;

    // Get the device type
    let mut dlt: libc::c_uint = 0;
    if unsafe { bpf::ioctl(fd, bpf::BIOCGDLT, &mut dlt) } == -1 {
        let err = IoError::last_error();
        unsafe { libc::close(fd); }
        return Err(err);
    }

    // The loopback IoError::last_error()device does weird things
    // FIXME This should really just be another L2 packet header type
    if dlt == bpf::DLT_NULL {
        header_size = 4;

        // Allow packets to be read back after they are written
        match set_feedback(fd) {
            Err(e) => return Err(e),
            _ => ()
        }

    } else {
        // Don't fill in source MAC
        if unsafe { bpf::ioctl(fd, bpf::BIOCSHDRCMPLT, &one) } == -1 {
            let err = IoError::last_error();
            unsafe { libc::close(fd); }
            return Err(err);
        }
    }

    let fd = Arc::new(internal::FileDesc { fd: fd });
    let sender = DataLinkSenderImpl {
        fd: fd.clone(),
        write_buffer: Vec::from_elem(write_buffer_size, 0u8),
        header_size: header_size,
    };
    let receiver = DataLinkReceiverImpl {
        fd: fd,
        read_buffer: Vec::from_elem(read_buffer_size, 0u8),
        header_size: header_size,
    };

    Ok((sender, receiver))
}

pub struct DataLinkSenderImpl {
    fd: Arc<internal::FileDesc>,
    write_buffer: Vec<u8>,
    header_size: uint,
}

impl DataLinkSenderImpl {
    pub fn build_and_send(&mut self, num_packets: uint, packet_size: uint,
                          func: |MutableEthernetHeader| -> ()) -> Option<IoResult<()>> {
        let len = num_packets * (packet_size + self.header_size);
        if len >= self.write_buffer.len() {
            None
        } else {
            let min = cmp::min(self.write_buffer.len(), len);
            for chunk in self.write_buffer.mut_slice_to(min)
                                          .mut_chunks(packet_size + self.header_size) {
                // If we're sending on the loopback device, the first 4 bytes must be set to
                // AF_INET
                if self.header_size == 4 {
                    unsafe {
                        *(chunk.as_mut_ptr() as *mut u32) = libc::AF_INET as u32;
                    }
                }
                {
                    let eh = MutableEthernetHeader::new(chunk.mut_slice_from(self.header_size));
                    func(eh);
                }
                match unsafe { libc::write(self.fd.fd,
                                           chunk.as_ptr() as *const libc::c_void,
                                           chunk.len() as libc::size_t) } {
                    len if len == -1 => return Some(Err(IoError::last_error())),
                    _ => ()
                }
            }
            Some(Ok(()))
        }
    }
}

pub struct DataLinkReceiverImpl {
    fd: Arc<internal::FileDesc>,
    read_buffer: Vec<u8>,
    header_size: uint,
}

impl DataLinkReceiverImpl {
    pub fn iter<'a>(&'a mut self) -> DataLinkChannelIteratorImpl<'a> {
        let buflen = self.read_buffer.len();
        DataLinkChannelIteratorImpl {
            pc: self,
            // Enough room for minimally sized packets without reallocating
            packets: RingBuf::with_capacity(buflen / 64)
        }
    }
}

pub struct DataLinkChannelIteratorImpl<'a> {
    pc: &'a mut DataLinkReceiverImpl,
    packets: RingBuf<(uint, uint)>,
}

impl<'a> DataLinkChannelIteratorImpl<'a> {
    pub fn next<'c>(&'c mut self) -> IoResult<EthernetHeader<'c>> {
        if self.packets.is_empty() {
            let buflen = match unsafe {
                libc::read(self.pc.fd.fd,
                           self.pc.read_buffer.as_ptr() as *mut libc::c_void,
                           self.pc.read_buffer.len() as libc::size_t)
            } {
                len if len > 0 => len,
                _ => return Err(IoError::last_error())
            };
            let mut ptr = self.pc.read_buffer.as_mut_ptr();
            let end = unsafe { self.pc.read_buffer.as_ptr().offset(buflen as int) };
            while (ptr as *const u8) < end {
                unsafe {
                    let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                    let start = ptr as int +
                                (*packet).bh_hdrlen as int -
                                self.pc.read_buffer.as_ptr() as int;
                    self.packets.push((start as uint + self.pc.header_size,
                                      (*packet).bh_caplen as uint - self.pc.header_size));
                    let offset = (*packet).bh_hdrlen as int + (*packet).bh_caplen as int;
                    ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                }
            }
        }
        let (start, len) = self.packets.pop_front().unwrap();
        Ok(EthernetHeader::new(self.pc.read_buffer.slice(start, start + len)))
    }
}

<|endoftext|>


<commit_before>use std::error::Error;
use std::fmt::{Debug, Display, Formatter};
use std::fmt;
use std::result::Result;
use std::path::{Path, PathBuf};
use std::convert::From;
use std::convert::Into;

use regex::Regex;

#[derive(Debug)]
#[derive(Clone)]
#[derive(PartialEq)]
#[derive(Eq)]
// #[derive(Display)]
pub enum FileIDType {
    NONE,
    UUID,
}

#[derive(Clone)]
pub struct FileID {
    id: Option<String>,
    id_type: FileIDType,
}

impl FileID {

    pub fn new(id_type: FileIDType, id: String) -> FileID {
        FileID {
            id: Some(id),
            id_type: id_type,
        }
    }

    pub fn is_valid(&self) -> bool {
        self.id.is_some()
    }

    pub fn get_type(&self) -> FileIDType {
        self.id_type.clone()
    }

}

impl Debug for FileID {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileID[{:?}]: {:?}",
               self.id_type,
               self.id);
        Ok(())
    }

}

impl Display for FileID {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileID[{:?}]: {:?}",
               self.id_type,
               self.id);
        Ok(())
    }

}

impl Into<String> for FileID {

    fn into(self) -> String {
        if let Some(id) = self.id {
            id.clone()
        } else {
            String::from("INVALID")
        }
    }

}

impl From<String> for FileID {

    fn from(s: String) -> FileID {
        FileID::from(&s)
    }

}

impl<'a> From<&'a String> for FileID {

    fn from(string: &'a String) -> FileID {

        let regex = Regex::new(r"([:alnum:]*)-([:upper:]*)-([A-Za-z0-9-_]*)\.(.*)").unwrap();
        let s = string.split("/").last().unwrap_or("");

        debug!("Regex build: {:?}", regex);
        debug!("Matching string: '{}'", s);
        regex.captures(s).and_then(|capts| {
            // first one is the whole string, index 1-N are the matches.
            if capts.len() != 5 {
                debug!("Matches, but not expected number of groups");
                return None;
            }
            debug!("Matches: {}", capts.len());

            let modname     = capts.at(1).unwrap();
            let hashname    = capts.at(2).unwrap();
            let mut hash    = capts.at(3).unwrap();

            debug!("Destructure FilePath to ID:");
            debug!("                  FilePath: {:?}", s);
            debug!("               Module Name: {:?}", modname);
            debug!("                 Hash Name: {:?}", hashname);
            debug!("                      Hash: {:?}", hash);

            let idtype = select_id_type_from_str(hashname);
            match idtype {
                FileIDType::NONE => hash = "INVALID",
                _ => {},
            }

            Some(FileID::new(idtype, String::from(hash)))
        }).unwrap_or({
            debug!("Did not match");
            FileID {
                id_type: FileIDType::NONE,
                id: None,
            }
        })
    }

}

impl From<PathBuf> for FileID {

    fn from(s: PathBuf) -> FileID {
        unimplemented!()
    }

}

impl<'a> From<&'a PathBuf> for FileID {

    fn from(s: &'a PathBuf) -> FileID {
        unimplemented!()
    }

}

pub struct FileIDError {
    summary: String,
    descrip: String,
}

impl FileIDError {

    pub fn new(s: String, d: String) -> FileIDError {
        FileIDError {
            summary: s,
            descrip: d,
        }
    }

}

impl<'a> Error for FileIDError {

    fn description(&self) -> &str {
        &self.summary[..]
    }

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

}

impl<'a> Debug for FileIDError {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileIDError: '{}'\n{}", self.summary, self.descrip);
        Ok(())
    }

}

impl<'a> Display for FileIDError {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileIDError: '{}'", self.summary);
        Ok(())
    }

}

fn select_id_type_from_str(s: &str) -> FileIDType {
    match s {
        "UUID"  => FileIDType::UUID,
        _       => FileIDType::NONE,
    }
}

pub type FileIDResult = Result<FileID, FileIDError>;

#[cfg(test)]
mod test {

    use super::{FileID, FileIDType};

    #[test]
    fn file_id_from_string() {
        setup_logger();

        let s1 = String::from("/home/user/testmodule-UUID-some-id.imag");
        let s2 = String::from("/home/user/testmodule-UUID-some-id.extension.imag");
        let s3 = String::from("/home/user/testmodule-NOHASH-some-id.imag");

        let id1 = FileID::from(s1);
        let id2 = FileID::from(s2);
        let id3 = FileID::from(s3);

        println!("Id 1 : {:?}", id1);
        println!("Id 2 : {:?}", id2);
        println!("Id 3 : {:?}", id3);

        assert_eq!(FileIDType::UUID, id1.get_type());
        assert_eq!(FileIDType::UUID, id2.get_type());
        assert_eq!(FileIDType::NONE, id3.get_type());

        let f1 : String = id1.into();
        let f2 : String = id2.into();
        let f3 : String = id3.into();

        assert_eq!(String::from("some-id"), f1);
        assert_eq!(String::from("some-id"), f2);
        assert_eq!(String::from("INVALID"), f3);
    }

    fn setup_logger() {
        extern crate log;
        use log::{LogLevelFilter, set_logger};
        use runtime::ImagLogger;

        log::set_logger(|max_log_lvl| {
            let lvl = LogLevelFilter::Debug;
            max_log_lvl.set(lvl);
            Box::new(ImagLogger::new(lvl.to_log_level().unwrap()))
        });
        debug!("Init logger for test");
    }

}

<commit_msg>Impl Into<String> for FileIDType<commit_after>use std::error::Error;
use std::fmt::{Debug, Display, Formatter};
use std::fmt;
use std::result::Result;
use std::path::{Path, PathBuf};
use std::convert::From;
use std::convert::Into;

use regex::Regex;

#[derive(Debug)]
#[derive(Clone)]
#[derive(PartialEq)]
#[derive(Eq)]
// #[derive(Display)]
pub enum FileIDType {
    NONE,
    UUID,
}

impl Into<String> for FileIDType {

    fn into(self) -> String {
        let s = match self {
            FileIDType::UUID => "UUID",
            FileIDType::NONE => "",
        };

        String::from(s)
    }

}

#[derive(Clone)]
pub struct FileID {
    id: Option<String>,
    id_type: FileIDType,
}

impl FileID {

    pub fn new(id_type: FileIDType, id: String) -> FileID {
        FileID {
            id: Some(id),
            id_type: id_type,
        }
    }

    pub fn is_valid(&self) -> bool {
        self.id.is_some()
    }

    pub fn get_type(&self) -> FileIDType {
        self.id_type.clone()
    }

}

impl Debug for FileID {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileID[{:?}]: {:?}",
               self.id_type,
               self.id);
        Ok(())
    }

}

impl Display for FileID {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileID[{:?}]: {:?}",
               self.id_type,
               self.id);
        Ok(())
    }

}

impl Into<String> for FileID {

    fn into(self) -> String {
        if let Some(id) = self.id {
            id.clone()
        } else {
            String::from("INVALID")
        }
    }

}

impl From<String> for FileID {

    fn from(s: String) -> FileID {
        FileID::from(&s)
    }

}

impl<'a> From<&'a String> for FileID {

    fn from(string: &'a String) -> FileID {

        let regex = Regex::new(r"([:alnum:]*)-([:upper:]*)-([A-Za-z0-9-_]*)\.(.*)").unwrap();
        let s = string.split("/").last().unwrap_or("");

        debug!("Regex build: {:?}", regex);
        debug!("Matching string: '{}'", s);
        regex.captures(s).and_then(|capts| {
            // first one is the whole string, index 1-N are the matches.
            if capts.len() != 5 {
                debug!("Matches, but not expected number of groups");
                return None;
            }
            debug!("Matches: {}", capts.len());

            let modname     = capts.at(1).unwrap();
            let hashname    = capts.at(2).unwrap();
            let mut hash    = capts.at(3).unwrap();

            debug!("Destructure FilePath to ID:");
            debug!("                  FilePath: {:?}", s);
            debug!("               Module Name: {:?}", modname);
            debug!("                 Hash Name: {:?}", hashname);
            debug!("                      Hash: {:?}", hash);

            let idtype = select_id_type_from_str(hashname);
            match idtype {
                FileIDType::NONE => hash = "INVALID",
                _ => {},
            }

            Some(FileID::new(idtype, String::from(hash)))
        }).unwrap_or({
            debug!("Did not match");
            FileID {
                id_type: FileIDType::NONE,
                id: None,
            }
        })
    }

}

impl From<PathBuf> for FileID {

    fn from(s: PathBuf) -> FileID {
        unimplemented!()
    }

}

impl<'a> From<&'a PathBuf> for FileID {

    fn from(s: &'a PathBuf) -> FileID {
        unimplemented!()
    }

}

pub struct FileIDError {
    summary: String,
    descrip: String,
}

impl FileIDError {

    pub fn new(s: String, d: String) -> FileIDError {
        FileIDError {
            summary: s,
            descrip: d,
        }
    }

}

impl<'a> Error for FileIDError {

    fn description(&self) -> &str {
        &self.summary[..]
    }

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

}

impl<'a> Debug for FileIDError {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileIDError: '{}'\n{}", self.summary, self.descrip);
        Ok(())
    }

}

impl<'a> Display for FileIDError {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "FileIDError: '{}'", self.summary);
        Ok(())
    }

}

fn select_id_type_from_str(s: &str) -> FileIDType {
    match s {
        "UUID"  => FileIDType::UUID,
        _       => FileIDType::NONE,
    }
}

pub type FileIDResult = Result<FileID, FileIDError>;

#[cfg(test)]
mod test {

    use super::{FileID, FileIDType};

    #[test]
    fn file_id_from_string() {
        setup_logger();

        let s1 = String::from("/home/user/testmodule-UUID-some-id.imag");
        let s2 = String::from("/home/user/testmodule-UUID-some-id.extension.imag");
        let s3 = String::from("/home/user/testmodule-NOHASH-some-id.imag");

        let id1 = FileID::from(s1);
        let id2 = FileID::from(s2);
        let id3 = FileID::from(s3);

        println!("Id 1 : {:?}", id1);
        println!("Id 2 : {:?}", id2);
        println!("Id 3 : {:?}", id3);

        assert_eq!(FileIDType::UUID, id1.get_type());
        assert_eq!(FileIDType::UUID, id2.get_type());
        assert_eq!(FileIDType::NONE, id3.get_type());

        let f1 : String = id1.into();
        let f2 : String = id2.into();
        let f3 : String = id3.into();

        assert_eq!(String::from("some-id"), f1);
        assert_eq!(String::from("some-id"), f2);
        assert_eq!(String::from("INVALID"), f3);
    }

    fn setup_logger() {
        extern crate log;
        use log::{LogLevelFilter, set_logger};
        use runtime::ImagLogger;

        log::set_logger(|max_log_lvl| {
            let lvl = LogLevelFilter::Debug;
            max_log_lvl.set(lvl);
            Box::new(ImagLogger::new(lvl.to_log_level().unwrap()))
        });
        debug!("Init logger for test");
    }

}

<|endoftext|>


<commit_before>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// 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.

//! Support for sending and receiving data link layer packets using the /dev/bpf device

extern crate libc;

use bindings::bpf;
use datalink::{self, NetworkInterface};
use datalink::{DataLinkChannelIterator, DataLinkReceiver, DataLinkSender};
use datalink::Channel::Ethernet;
use internal;
use sockets;

use std::collections::VecDeque;
use std::ffi::CString;
use std::io;
use std::iter::repeat;
use std::mem;
use std::ptr;
use std::sync::Arc;
use std::time::Duration;

static ETHERNET_HEADER_SIZE: usize = 14;

/// BPF-specific configuration
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096
    pub read_buffer_size: usize,

    /// The read timeout. Defaults to None.
    pub read_timeout: Option<Duration>,

    /// The write timeout. Defaults to None.
    pub write_timeout: Option<Duration>,

    /// The number of /dev/bpf* file descriptors to attempt before failing.
    ///
    /// This setting is only used on OS X - FreeBSD uses a single /dev/bpf rather than creating a
    /// new descriptor each time one is opened.
    ///
    /// Defaults to: 1000
    pub bpf_fd_attempts: usize,
}

impl<'a> From<&'a datalink::Config> for Config {
    fn from(config: &datalink::Config) -> Config {
        Config {
            write_buffer_size: config.write_buffer_size,
            read_buffer_size: config.read_buffer_size,
            bpf_fd_attempts: config.bpf_fd_attempts,
            read_timeout: config.read_timeout,
            write_timeout: config.write_timeout,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
            bpf_fd_attempts: 1000,
            read_timeout: None,
            write_timeout: None,
        }
    }
}

/// Create a datalink channel using the /dev/bpf device
// NOTE buffer must be word aligned.
#[inline]
pub fn channel(network_interface: &NetworkInterface,
               config: Config)
    -> io::Result<datalink::Channel> {
    #[cfg(target_os = "freebsd")]
    fn get_fd(_attempts: usize) -> libc::c_int {
        unsafe {
            libc::open(CString::new(&b"/dev/bpf"[..]).unwrap().as_ptr(),
                       libc::O_RDWR,
                       0)
        }
    }

    #[cfg(target_os = "macos")]
    fn get_fd(attempts: usize) -> libc::c_int {
        for i in 0..attempts {
            let fd = unsafe {
                let file_name = format!("/dev/bpf{}", i);
                libc::open(CString::new(file_name.as_bytes()).unwrap().as_ptr(),
                           libc::O_RDWR,
                           0)
            };
            if fd != -1 {
                return fd;
            }
        }

        -1
    }

    #[cfg(target_os = "freebsd")]
    fn set_feedback(fd: libc::c_int) -> io::Result<()> {
        let one: libc::c_uint = 1;
        if unsafe { bpf::ioctl(fd, bpf::BIOCFEEDBACK, &one) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
        Ok(())
    }

    #[cfg(target_os = "macos")]
    fn set_feedback(_fd: libc::c_int) -> io::Result<()> {
        Ok(())
    }

    let fd = get_fd(config.bpf_fd_attempts);
    if fd == -1 {
        return Err(io::Error::last_os_error());
    }
    let mut iface: bpf::ifreq = unsafe { mem::zeroed() };
    for (i, c) in network_interface.name.bytes().enumerate() {
        iface.ifr_name[i] = c as i8;
    }

    let buflen = config.read_buffer_size as libc::c_uint;
    // NOTE Buffer length must be set before binding to an interface
    //      otherwise this will return Invalid Argument
    if unsafe { bpf::ioctl(fd, bpf::BIOCSBLEN, &buflen) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Set the interface to use
    if unsafe { bpf::ioctl(fd, bpf::BIOCSETIF, &iface) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Return from read as soon as packets are available - don't wait to fill the
    // buffer
    let one: libc::c_uint = 1;
    if unsafe { bpf::ioctl(fd, bpf::BIOCIMMEDIATE, &one) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Get the device type
    let mut dlt: libc::c_uint = 0;
    if unsafe { bpf::ioctl(fd, bpf::BIOCGDLT, &mut dlt) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    let mut loopback = false;
    let mut allocated_read_buffer_size = config.read_buffer_size;
    // The loopback device does weird things
    // FIXME This should really just be another L2 packet header type
    if dlt == bpf::DLT_NULL {
        loopback = true;
        // So we can guaranatee that we can have a header before the packet.
        // Loopback packets arrive without the header.
        allocated_read_buffer_size += ETHERNET_HEADER_SIZE;

        // Allow packets to be read back after they are written
        if let Err(e) = set_feedback(fd) {
            return Err(e);
        }
    } else {
        // Don't fill in source MAC
        if unsafe { bpf::ioctl(fd, bpf::BIOCSHDRCMPLT, &one) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
    }

    // Enable nonblocking
    if unsafe { libc::fcntl(fd, libc::F_SETFL, libc::O_NONBLOCK) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            sockets::close(fd);
        }
        return Err(err);
    }

    let fd = Arc::new(internal::FileDesc { fd: fd });
    let mut sender = Box::new(DataLinkSenderImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        write_buffer: repeat(0u8).take(config.write_buffer_size).collect(),
        loopback: loopback,
        timeout: config.write_timeout.map(|to| internal::duration_to_timespec(to)),
    });
    unsafe {
        libc::FD_ZERO(&mut sender.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut sender.fd_set as *mut libc::fd_set);
    }
    let mut receiver = Box::new(DataLinkReceiverImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        read_buffer: repeat(0u8).take(allocated_read_buffer_size).collect(),
        loopback: loopback,
        timeout: config.read_timeout.map(|to| internal::duration_to_timespec(to)),
    });
    unsafe {
        libc::FD_ZERO(&mut receiver.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut receiver.fd_set as *mut libc::fd_set);
    }

    Ok(Ethernet(sender, receiver))
}

struct DataLinkSenderImpl {
    fd: Arc<internal::FileDesc>,
    fd_set: libc::fd_set,
    write_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
}

impl DataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(&mut [u8]))
        -> Option<io::Result<()>> {
        let len = num_packets * packet_size;
        if len >= self.write_buffer.len() {
            None
        } else {
            // If we're sending on the loopback device, discard the ethernet header.
            // The OS will prepend the packet with 4 bytes set to AF_INET.
            let offset = if self.loopback { ETHERNET_HEADER_SIZE } else { 0 };
            for chunk in self.write_buffer[..len].chunks_mut(packet_size) {
                func(chunk);
                let ret = unsafe {
                    libc::pselect(self.fd.fd + 1,
                                  ptr::null_mut(),
                                  &mut self.fd_set as *mut libc::fd_set,
                                  ptr::null_mut(),
                                  self.timeout
                                      .as_ref()
                                      .map(|to| to as *const libc::timespec)
                                      .unwrap_or(ptr::null()),
                                  ptr::null())
                };
                if ret == -1 {
                    // Error occured!
                    return Some(Err(io::Error::last_os_error()));
                } else if ret == 0 {
                    return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
                } else {
                    match unsafe {
                        libc::write(self.fd.fd,
                                    chunk.as_ptr().offset(offset as isize) as *const libc::c_void,
                                    (chunk.len() - offset) as libc::size_t)
                    } {
                        len if len == -1 => return Some(Err(io::Error::last_os_error())),
                        _ => (),
                    }
                }
            }
            Some(Ok(()))
        }
    }

    #[inline]
    fn send_to(&mut self,
               packet: &[u8],
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        // If we're sending on the loopback device, discard the ethernet header.
        // The OS will prepend the packet with 4 bytes set to AF_INET.
        let offset = if self.loopback { ETHERNET_HEADER_SIZE } else { 0 };
        let ret = unsafe {
            libc::pselect(self.fd.fd + 1,
                          ptr::null_mut(),
                          &mut self.fd_set as *mut libc::fd_set,
                          ptr::null_mut(),
                          self.timeout
                              .as_ref()
                              .map(|to| to as *const libc::timespec)
                              .unwrap_or(ptr::null()),
                          ptr::null())
        };
        if ret == -1 {
            // Error occured!
            return Some(Err(io::Error::last_os_error()));
        } else if ret == 0 {
            return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
        } else {
            match unsafe {
                libc::write(self.fd.fd,
                            packet.as_ptr().offset(offset as isize) as *const libc::c_void,
                            (packet.len() - offset) as libc::size_t)
            } {
                len if len == -1 => Some(Err(io::Error::last_os_error())),
                _ => Some(Ok(())),
            }
        }
    }
}

struct DataLinkReceiverImpl {
    fd: Arc<internal::FileDesc>,
    fd_set: libc::fd_set,
    read_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
}

impl DataLinkReceiver for DataLinkReceiverImpl {
    fn iter<'a>(&'a mut self) -> Box<DataLinkChannelIterator + 'a> {
        let buflen = self.read_buffer.len();
        Box::new(DataLinkChannelIteratorImpl {
            pc: self,
            // Enough room for minimally sized packets without reallocating
            packets: VecDeque::with_capacity(buflen / 64),
        })
    }
}

struct DataLinkChannelIteratorImpl<'a> {
    pc: &'a mut DataLinkReceiverImpl,
    packets: VecDeque<(usize, usize)>,
}

impl<'a> DataLinkChannelIterator<'a> for DataLinkChannelIteratorImpl<'a> {
    fn next(&mut self) -> io::Result<&[u8]> {
        // Loopback packets arrive with a 4 byte header instead of normal ethernet header.
        // Discard that header and replace with zeroed out ethernet header.
        let (header_size, buffer_offset) = if self.pc.loopback {
            (4, ETHERNET_HEADER_SIZE)
        } else {
            (0, 0)
        };
        if self.packets.is_empty() {
            let buffer = &mut self.pc.read_buffer[buffer_offset..];
            let ret = unsafe {
                libc::pselect(self.pc.fd.fd + 1,
                              &mut self.pc.fd_set as *mut libc::fd_set,
                              ptr::null_mut(),
                              ptr::null_mut(),
                              self.pc
                                  .timeout
                                  .as_ref()
                                  .map(|to| to as *const libc::timespec)
                                  .unwrap_or(ptr::null()),
                              ptr::null())
            };
            if ret == -1 {
                return Err(io::Error::last_os_error());
            } else if ret == 0 {
                return Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out"));
            } else {
                let buflen = match unsafe {
                    libc::read(self.pc.fd.fd,
                               buffer.as_ptr() as *mut libc::c_void,
                               buffer.len() as libc::size_t)
                } {
                    len if len > 0 => len,
                    _ => return Err(io::Error::last_os_error()),
                };
                let mut ptr = buffer.as_mut_ptr();
                let end = unsafe { buffer.as_ptr().offset(buflen as isize) };
                while (ptr as *const u8) < end {
                    unsafe {
                        let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                        let start = ptr as isize + (*packet).bh_hdrlen as isize -
                                    buffer.as_ptr() as isize;
                        self.packets.push_back((start as usize + header_size,
                                                (*packet).bh_caplen as usize - header_size));
                        let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                        ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                    }
                }
            }
        }
        let (start, mut len) = self.packets.pop_front().unwrap();
        len += buffer_offset;
        // Zero out part that will become fake ethernet header if on loopback.
        for i in (&mut self.pc.read_buffer[start..start + buffer_offset]).iter_mut() {
            *i = 0;
        }
        Ok(&self.pc.read_buffer[start..start + len])
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    #[path = "unix_interfaces.rs"]
    mod interfaces;
    interfaces::interfaces()
}
<commit_msg>Get rid of DataLinkChannelIterator in bpf<commit_after>// Copyright (c) 2014-2016 Robert Clipsham <robert@octarineparrot.com>
//
// 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.

//! Support for sending and receiving data link layer packets using the /dev/bpf device

extern crate libc;

use bindings::bpf;
use datalink::{self, NetworkInterface};
use datalink::{DataLinkReceiver, DataLinkSender};
use datalink::Channel::Ethernet;
use internal;
use sockets;

use std::collections::VecDeque;
use std::ffi::CString;
use std::io;
use std::iter::repeat;
use std::mem;
use std::ptr;
use std::sync::Arc;
use std::time::Duration;

static ETHERNET_HEADER_SIZE: usize = 14;

/// BPF-specific configuration
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct Config {
    /// The size of buffer to use when writing packets. Defaults to 4096
    pub write_buffer_size: usize,

    /// The size of buffer to use when reading packets. Defaults to 4096
    pub read_buffer_size: usize,

    /// The read timeout. Defaults to None.
    pub read_timeout: Option<Duration>,

    /// The write timeout. Defaults to None.
    pub write_timeout: Option<Duration>,

    /// The number of /dev/bpf* file descriptors to attempt before failing.
    ///
    /// This setting is only used on OS X - FreeBSD uses a single /dev/bpf rather than creating a
    /// new descriptor each time one is opened.
    ///
    /// Defaults to: 1000
    pub bpf_fd_attempts: usize,
}

impl<'a> From<&'a datalink::Config> for Config {
    fn from(config: &datalink::Config) -> Config {
        Config {
            write_buffer_size: config.write_buffer_size,
            read_buffer_size: config.read_buffer_size,
            bpf_fd_attempts: config.bpf_fd_attempts,
            read_timeout: config.read_timeout,
            write_timeout: config.write_timeout,
        }
    }
}

impl Default for Config {
    fn default() -> Config {
        Config {
            write_buffer_size: 4096,
            read_buffer_size: 4096,
            bpf_fd_attempts: 1000,
            read_timeout: None,
            write_timeout: None,
        }
    }
}

/// Create a datalink channel using the /dev/bpf device
// NOTE buffer must be word aligned.
#[inline]
pub fn channel(network_interface: &NetworkInterface,
               config: Config)
    -> io::Result<datalink::Channel> {
    #[cfg(target_os = "freebsd")]
    fn get_fd(_attempts: usize) -> libc::c_int {
        unsafe {
            libc::open(CString::new(&b"/dev/bpf"[..]).unwrap().as_ptr(),
                       libc::O_RDWR,
                       0)
        }
    }

    #[cfg(target_os = "macos")]
    fn get_fd(attempts: usize) -> libc::c_int {
        for i in 0..attempts {
            let fd = unsafe {
                let file_name = format!("/dev/bpf{}", i);
                libc::open(CString::new(file_name.as_bytes()).unwrap().as_ptr(),
                           libc::O_RDWR,
                           0)
            };
            if fd != -1 {
                return fd;
            }
        }

        -1
    }

    #[cfg(target_os = "freebsd")]
    fn set_feedback(fd: libc::c_int) -> io::Result<()> {
        let one: libc::c_uint = 1;
        if unsafe { bpf::ioctl(fd, bpf::BIOCFEEDBACK, &one) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
        Ok(())
    }

    #[cfg(target_os = "macos")]
    fn set_feedback(_fd: libc::c_int) -> io::Result<()> {
        Ok(())
    }

    let fd = get_fd(config.bpf_fd_attempts);
    if fd == -1 {
        return Err(io::Error::last_os_error());
    }
    let mut iface: bpf::ifreq = unsafe { mem::zeroed() };
    for (i, c) in network_interface.name.bytes().enumerate() {
        iface.ifr_name[i] = c as i8;
    }

    let buflen = config.read_buffer_size as libc::c_uint;
    // NOTE Buffer length must be set before binding to an interface
    //      otherwise this will return Invalid Argument
    if unsafe { bpf::ioctl(fd, bpf::BIOCSBLEN, &buflen) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Set the interface to use
    if unsafe { bpf::ioctl(fd, bpf::BIOCSETIF, &iface) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Return from read as soon as packets are available - don't wait to fill the
    // buffer
    let one: libc::c_uint = 1;
    if unsafe { bpf::ioctl(fd, bpf::BIOCIMMEDIATE, &one) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    // Get the device type
    let mut dlt: libc::c_uint = 0;
    if unsafe { bpf::ioctl(fd, bpf::BIOCGDLT, &mut dlt) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            libc::close(fd);
        }
        return Err(err);
    }

    let mut loopback = false;
    let mut allocated_read_buffer_size = config.read_buffer_size;
    // The loopback device does weird things
    // FIXME This should really just be another L2 packet header type
    if dlt == bpf::DLT_NULL {
        loopback = true;
        // So we can guaranatee that we can have a header before the packet.
        // Loopback packets arrive without the header.
        allocated_read_buffer_size += ETHERNET_HEADER_SIZE;

        // Allow packets to be read back after they are written
        if let Err(e) = set_feedback(fd) {
            return Err(e);
        }
    } else {
        // Don't fill in source MAC
        if unsafe { bpf::ioctl(fd, bpf::BIOCSHDRCMPLT, &one) } == -1 {
            let err = io::Error::last_os_error();
            unsafe {
                libc::close(fd);
            }
            return Err(err);
        }
    }

    // Enable nonblocking
    if unsafe { libc::fcntl(fd, libc::F_SETFL, libc::O_NONBLOCK) } == -1 {
        let err = io::Error::last_os_error();
        unsafe {
            sockets::close(fd);
        }
        return Err(err);
    }

    let fd = Arc::new(internal::FileDesc { fd: fd });
    let mut sender = Box::new(DataLinkSenderImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        write_buffer: repeat(0u8).take(config.write_buffer_size).collect(),
        loopback: loopback,
        timeout: config.write_timeout.map(|to| internal::duration_to_timespec(to)),
    });
    unsafe {
        libc::FD_ZERO(&mut sender.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut sender.fd_set as *mut libc::fd_set);
    }
    let mut receiver = Box::new(DataLinkReceiverImpl {
        fd: fd.clone(),
        fd_set: unsafe { mem::zeroed() },
        read_buffer: repeat(0u8).take(allocated_read_buffer_size).collect(),
        loopback: loopback,
        timeout: config.read_timeout.map(|to| internal::duration_to_timespec(to)),
        // Enough room for minimally sized packets without reallocating
        packets: VecDeque::with_capacity(allocated_read_buffer_size / 64),
    });
    unsafe {
        libc::FD_ZERO(&mut receiver.fd_set as *mut libc::fd_set);
        libc::FD_SET(fd.fd, &mut receiver.fd_set as *mut libc::fd_set);
    }

    Ok(Ethernet(sender, receiver))
}

struct DataLinkSenderImpl {
    fd: Arc<internal::FileDesc>,
    fd_set: libc::fd_set,
    write_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
}

impl DataLinkSender for DataLinkSenderImpl {
    #[inline]
    fn build_and_send(&mut self,
                      num_packets: usize,
                      packet_size: usize,
                      func: &mut FnMut(&mut [u8]))
        -> Option<io::Result<()>> {
        let len = num_packets * packet_size;
        if len >= self.write_buffer.len() {
            None
        } else {
            // If we're sending on the loopback device, discard the ethernet header.
            // The OS will prepend the packet with 4 bytes set to AF_INET.
            let offset = if self.loopback { ETHERNET_HEADER_SIZE } else { 0 };
            for chunk in self.write_buffer[..len].chunks_mut(packet_size) {
                func(chunk);
                let ret = unsafe {
                    libc::pselect(self.fd.fd + 1,
                                  ptr::null_mut(),
                                  &mut self.fd_set as *mut libc::fd_set,
                                  ptr::null_mut(),
                                  self.timeout
                                      .as_ref()
                                      .map(|to| to as *const libc::timespec)
                                      .unwrap_or(ptr::null()),
                                  ptr::null())
                };
                if ret == -1 {
                    // Error occured!
                    return Some(Err(io::Error::last_os_error()));
                } else if ret == 0 {
                    return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
                } else {
                    match unsafe {
                        libc::write(self.fd.fd,
                                    chunk.as_ptr().offset(offset as isize) as *const libc::c_void,
                                    (chunk.len() - offset) as libc::size_t)
                    } {
                        len if len == -1 => return Some(Err(io::Error::last_os_error())),
                        _ => (),
                    }
                }
            }
            Some(Ok(()))
        }
    }

    #[inline]
    fn send_to(&mut self,
               packet: &[u8],
               _dst: Option<NetworkInterface>)
        -> Option<io::Result<()>> {
        // If we're sending on the loopback device, discard the ethernet header.
        // The OS will prepend the packet with 4 bytes set to AF_INET.
        let offset = if self.loopback { ETHERNET_HEADER_SIZE } else { 0 };
        let ret = unsafe {
            libc::pselect(self.fd.fd + 1,
                          ptr::null_mut(),
                          &mut self.fd_set as *mut libc::fd_set,
                          ptr::null_mut(),
                          self.timeout
                              .as_ref()
                              .map(|to| to as *const libc::timespec)
                              .unwrap_or(ptr::null()),
                          ptr::null())
        };
        if ret == -1 {
            // Error occured!
            return Some(Err(io::Error::last_os_error()));
        } else if ret == 0 {
            return Some(Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out")));
        } else {
            match unsafe {
                libc::write(self.fd.fd,
                            packet.as_ptr().offset(offset as isize) as *const libc::c_void,
                            (packet.len() - offset) as libc::size_t)
            } {
                len if len == -1 => Some(Err(io::Error::last_os_error())),
                _ => Some(Ok(())),
            }
        }
    }
}

struct DataLinkReceiverImpl {
    fd: Arc<internal::FileDesc>,
    fd_set: libc::fd_set,
    read_buffer: Vec<u8>,
    loopback: bool,
    timeout: Option<libc::timespec>,
    packets: VecDeque<(usize, usize)>,
}

impl DataLinkReceiver for DataLinkReceiverImpl {
    fn next(&mut self) -> io::Result<&[u8]> {
        // Loopback packets arrive with a 4 byte header instead of normal ethernet header.
        // Discard that header and replace with zeroed out ethernet header.
        let (header_size, buffer_offset) = if self.loopback {
            (4, ETHERNET_HEADER_SIZE)
        } else {
            (0, 0)
        };
        if self.packets.is_empty() {
            let buffer = &mut self.read_buffer[buffer_offset..];
            let ret = unsafe {
                libc::pselect(self.fd.fd + 1,
                              &mut self.fd_set as *mut libc::fd_set,
                              ptr::null_mut(),
                              ptr::null_mut(),
                              self
                                  .timeout
                                  .as_ref()
                                  .map(|to| to as *const libc::timespec)
                                  .unwrap_or(ptr::null()),
                              ptr::null())
            };
            if ret == -1 {
                return Err(io::Error::last_os_error());
            } else if ret == 0 {
                return Err(io::Error::new(io::ErrorKind::TimedOut, "Timed out"));
            } else {
                let buflen = match unsafe {
                    libc::read(self.fd.fd,
                               buffer.as_ptr() as *mut libc::c_void,
                               buffer.len() as libc::size_t)
                } {
                    len if len > 0 => len,
                    _ => return Err(io::Error::last_os_error()),
                };
                let mut ptr = buffer.as_mut_ptr();
                let end = unsafe { buffer.as_ptr().offset(buflen as isize) };
                while (ptr as *const u8) < end {
                    unsafe {
                        let packet: *const bpf::bpf_hdr = mem::transmute(ptr);
                        let start = ptr as isize + (*packet).bh_hdrlen as isize -
                                    buffer.as_ptr() as isize;
                        self.packets.push_back((start as usize + header_size,
                                                (*packet).bh_caplen as usize - header_size));
                        let offset = (*packet).bh_hdrlen as isize + (*packet).bh_caplen as isize;
                        ptr = ptr.offset(bpf::BPF_WORDALIGN(offset));
                    }
                }
            }
        }
        let (start, mut len) = self.packets.pop_front().unwrap();
        len += buffer_offset;
        // Zero out part that will become fake ethernet header if on loopback.
        for i in (&mut self.read_buffer[start..start + buffer_offset]).iter_mut() {
            *i = 0;
        }
        Ok(&self.read_buffer[start..start + len])
    }
}

/// Get a list of available network interfaces for the current machine.
pub fn interfaces() -> Vec<NetworkInterface> {
    #[path = "unix_interfaces.rs"]
    mod interfaces;
    interfaces::interfaces()
}
<|endoftext|>