<commit_before>#![crate_name = "stdbuf"]
#![allow(unstable)]

/*
* This file is part of the uutils coreutils package.
*
* (c) Dorota Kapturkiewicz <dokaptur@gmail.com>
*
* For the full copyright and license information, please view the LICENSE
* file that was distributed with this source code.
*/

extern crate getopts;
extern crate libc;
use getopts::{optopt, optflag, getopts, usage, Matches, OptGroup};
use std::io::process::{Command, StdioContainer};
use std::io::fs::PathExtensions;
use std::iter::range_inclusive;
use std::num::Int;
use std::os;

#[path = "../common/util.rs"]
#[macro_use]
mod util;

static NAME: &'static str = "stdbuf";
static VERSION: &'static str = "1.0.0";
static LIBSTDBUF: &'static str = "libstdbuf";

enum BufferType {
    Default,
    Line,
    Size(u64)
}

struct ProgramOptions {
    stdin: BufferType,
    stdout: BufferType,
    stderr: BufferType,
}

enum ErrMsg {
    Retry,
    Fatal
}

enum OkMsg {
    Buffering,
    Help,
    Version
}

#[cfg(target_os = "linux")]
fn preload_strings() -> (&'static str, &'static str) {
    ("LD_PRELOAD", ".so")
}

#[cfg(target_os = "macos")]
fn preload_strings() -> (&'static str, &'static str) {
    ("DYLD_INSERT_LIBRARIES", ".dylib")
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn preload_strings() -> (&'static str, &'static str) {
    crash!(1, "Command not supported for this operating system!")
}

fn print_version() {
    println!("{} version {}", NAME, VERSION);
}

fn print_usage(opts: &[OptGroup]) {
    let brief =
        "Usage: stdbuf OPTION... COMMAND\n \
        Run COMMAND, with modified buffering operations for its standard streams\n \
        Mandatory arguments to long options are mandatory for short options too.";
    let explanation =
        "If MODE is 'L' the corresponding stream will be line buffered.\n \
        This option is invalid with standard input.\n\n \
        If MODE is '0' the corresponding stream will be unbuffered.\n\n \
        Otherwise MODE is a number which may be followed by one of the following:\n\n \
        KB 1000, K 1024, MB 1000*1000, M 1024*1024, and so on for G, T, P, E, Z, Y.\n \
        In this case the corresponding stream will be fully buffered with the buffer size set to MODE bytes.\n\n \
        NOTE: If COMMAND adjusts the buffering of its standard streams ('tee' does for e.g.) then that will override \
        corresponding settings changed by 'stdbuf'.\n \
        Also some filters (like 'dd' and 'cat' etc.) don't use streams for I/O, \
        and are thus unaffected by 'stdbuf' settings.\n";
    println!("{}\n{}", getopts::usage(brief, opts), explanation);
}

fn parse_size(size: &str) -> Option<u64> {
    let ext = size.trim_left_matches(|&: c: char| c.is_digit(10));
    let num = size.trim_right_matches(|&: c: char| c.is_alphabetic());
    let mut recovered = num.to_string();
    recovered.push_str(ext);
    if recovered.as_slice() != size {
        return None;
    }
    let buf_size: u64 = match num.parse() {
        Some(m) => m,
        None => return None,
    };
    let (power, base): (usize, u64) = match ext {
        "" => (0, 0),
        "KB" => (1, 1024),
        "K" => (1, 1000),
        "MB" => (2, 1024),
        "M" => (2, 1000),
        "GB" => (3, 1024),
        "G" => (3, 1000),
        "TB" => (4, 1024),
        "T" => (4, 1000),
        "PB" => (5, 1024),
        "P" => (5, 1000),
        "EB" => (6, 1024),
        "E" => (6, 1000),
        "ZB" => (7, 1024),
        "Z" => (7, 1000),
        "YB" => (8, 1024),
        "Y" => (8, 1000),
        _ => return None,
    };
    Some(buf_size * base.pow(power))
}

fn check_option(matches: &Matches, name: &str, modified: &mut bool) -> Option<BufferType> {
    match matches.opt_str(name) {
        Some(value) => {
            *modified = true;
            match value.as_slice() {
                "L" => {
                    if name == "input" {
                        show_info!("stdbuf: line buffering stdin is meaningless");
                        None
                    } else {
                        Some(BufferType::Line)
                    }
                },
                x => {
                    let size = match parse_size(x) {
                        Some(m) => m,
                        None => { show_error!("Invalid mode {}", x); return None }
                    };
                    Some(BufferType::Size(size))
                },
            }
        },
        None => Some(BufferType::Default),
    }
}

fn parse_options(args: &[String], options: &mut ProgramOptions, optgrps: &[OptGroup]) -> Result<OkMsg, ErrMsg> {
    let matches = match getopts(args, optgrps) {
        Ok(m) => m,
        Err(_) => return Err(ErrMsg::Retry)
    };
    if matches.opt_present("help") {
        return Ok(OkMsg::Help);
    }
    if matches.opt_present("version") {
        return Ok(OkMsg::Version);
    }
    let mut modified = false;
    options.stdin = try!(check_option(&matches, "input", &mut modified).ok_or(ErrMsg::Fatal));
    options.stdout = try!(check_option(&matches, "output", &mut modified).ok_or(ErrMsg::Fatal));
    options.stderr = try!(check_option(&matches, "error", &mut modified).ok_or(ErrMsg::Fatal));

    if matches.free.len() != 1 {
        return Err(ErrMsg::Retry);
    }
    if !modified {
        show_error!("stdbuf: you must specify a buffering mode option");
        return Err(ErrMsg::Fatal);
    }
    Ok(OkMsg::Buffering)
}

fn set_command_env(command: &mut Command, buffer_name: &str, buffer_type: BufferType) {
    match buffer_type {
        BufferType::Size(m) => { command.env(buffer_name, m.to_string()); },
        BufferType::Line => { command.env(buffer_name, "L"); },
        BufferType::Default => {},
    }
}

fn get_preload_env() -> (String, String) {
    let (preload, extension) = preload_strings();
    let mut libstdbuf = LIBSTDBUF.to_string();
    libstdbuf.push_str(extension);
    // First search for library in directory of executable.
    let mut path = match os::self_exe_path() {
        Some(exe_path) => exe_path,
        None => crash!(1, "Impossible to fetch the path of this executable.")
    };
    path.push(libstdbuf.as_slice());
    if path.exists() {
        match path.as_str() {
            Some(s) => { return (preload.to_string(), s.to_string()); },
            None => crash!(1, "Error while converting path.")
        };
    }
    // We assume library is in LD_LIBRARY_PATH/ DYLD_LIBRARY_PATH.
    (preload.to_string(), libstdbuf)
}

pub fn uumain(args: Vec<String>) -> isize {
    let optgrps = [
        optopt("i", "input", "adjust standard input stream buffering", "MODE"),
        optopt("o", "output", "adjust standard output stream buffering", "MODE"),
        optopt("e", "error", "adjust standard error stream buffering", "MODE"),
        optflag("", "help", "display this help and exit"),
        optflag("", "version", "output version information and exit"),
    ];
    let mut options = ProgramOptions {stdin: BufferType::Default, stdout: BufferType::Default, stderr: BufferType::Default};
    let mut command_idx = -1;
    for i in range_inclusive(1, args.len()) {
        match parse_options(&args[1 .. i], &mut options, &optgrps) {
            Ok(OkMsg::Buffering) => {
                command_idx = i - 1;
                break;
            },
            Ok(OkMsg::Help) => {
                print_usage(&optgrps);
                return 0;
            },
            Ok(OkMsg::Version) => {
                print_version();
                return 0;
            },
            Err(ErrMsg::Fatal) => break,
            Err(ErrMsg::Retry) => continue,
        }
    };
    if command_idx == -1 {
        crash!(125, "Invalid options\nTry 'stdbuf --help' for more information.");
    }
    let ref command_name = args[command_idx];
    let mut command = Command::new(command_name);
    let (preload_env, libstdbuf) = get_preload_env();
    command.args(&args[command_idx + 1 ..]).env(preload_env.as_slice(), libstdbuf.as_slice());
    command.stdin(StdioContainer::InheritFd(0)).stdout(StdioContainer::InheritFd(1)).stderr(StdioContainer::InheritFd(2));
    set_command_env(&mut command, "_STDBUF_I", options.stdin);
    set_command_env(&mut command, "_STDBUF_O", options.stdout);
    set_command_env(&mut command, "_STDBUF_E", options.stderr);
    if let Err(e) = command.spawn() {
        crash!(1, "failed to execute process: {}", e);
    }
    0
}
<commit_msg>stdbuf - fix issue 512<commit_after>#![crate_name = "stdbuf"]
#![allow(unstable)]

/*
* This file is part of the uutils coreutils package.
*
* (c) Dorota Kapturkiewicz <dokaptur@gmail.com>
*
* For the full copyright and license information, please view the LICENSE
* file that was distributed with this source code.
*/

extern crate getopts;
extern crate libc;
use getopts::{optopt, optflag, getopts, usage, Matches, OptGroup};
use std::io::process::{Command, StdioContainer, ProcessExit};
use std::io::fs::PathExtensions;
use std::iter::range_inclusive;
use std::num::Int;
use std::os;

#[path = "../common/util.rs"]
#[macro_use]
mod util;

static NAME: &'static str = "stdbuf";
static VERSION: &'static str = "1.0.0";
static LIBSTDBUF: &'static str = "libstdbuf";

enum BufferType {
    Default,
    Line,
    Size(u64)
}

struct ProgramOptions {
    stdin: BufferType,
    stdout: BufferType,
    stderr: BufferType,
}

enum ErrMsg {
    Retry,
    Fatal
}

enum OkMsg {
    Buffering,
    Help,
    Version
}

#[cfg(target_os = "linux")]
fn preload_strings() -> (&'static str, &'static str) {
    ("LD_PRELOAD", ".so")
}

#[cfg(target_os = "macos")]
fn preload_strings() -> (&'static str, &'static str) {
    ("DYLD_INSERT_LIBRARIES", ".dylib")
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
fn preload_strings() -> (&'static str, &'static str) {
    crash!(1, "Command not supported for this operating system!")
}

fn print_version() {
    println!("{} version {}", NAME, VERSION);
}

fn print_usage(opts: &[OptGroup]) {
    let brief =
        "Usage: stdbuf OPTION... COMMAND\n \
        Run COMMAND, with modified buffering operations for its standard streams\n \
        Mandatory arguments to long options are mandatory for short options too.";
    let explanation =
        "If MODE is 'L' the corresponding stream will be line buffered.\n \
        This option is invalid with standard input.\n\n \
        If MODE is '0' the corresponding stream will be unbuffered.\n\n \
        Otherwise MODE is a number which may be followed by one of the following:\n\n \
        KB 1000, K 1024, MB 1000*1000, M 1024*1024, and so on for G, T, P, E, Z, Y.\n \
        In this case the corresponding stream will be fully buffered with the buffer size set to MODE bytes.\n\n \
        NOTE: If COMMAND adjusts the buffering of its standard streams ('tee' does for e.g.) then that will override \
        corresponding settings changed by 'stdbuf'.\n \
        Also some filters (like 'dd' and 'cat' etc.) don't use streams for I/O, \
        and are thus unaffected by 'stdbuf' settings.\n";
    println!("{}\n{}", getopts::usage(brief, opts), explanation);
}

fn parse_size(size: &str) -> Option<u64> {
    let ext = size.trim_left_matches(|&: c: char| c.is_digit(10));
    let num = size.trim_right_matches(|&: c: char| c.is_alphabetic());
    let mut recovered = num.to_string();
    recovered.push_str(ext);
    if recovered.as_slice() != size {
        return None;
    }
    let buf_size: u64 = match num.parse() {
        Some(m) => m,
        None => return None,
    };
    let (power, base): (usize, u64) = match ext {
        "" => (0, 0),
        "KB" => (1, 1024),
        "K" => (1, 1000),
        "MB" => (2, 1024),
        "M" => (2, 1000),
        "GB" => (3, 1024),
        "G" => (3, 1000),
        "TB" => (4, 1024),
        "T" => (4, 1000),
        "PB" => (5, 1024),
        "P" => (5, 1000),
        "EB" => (6, 1024),
        "E" => (6, 1000),
        "ZB" => (7, 1024),
        "Z" => (7, 1000),
        "YB" => (8, 1024),
        "Y" => (8, 1000),
        _ => return None,
    };
    Some(buf_size * base.pow(power))
}

fn check_option(matches: &Matches, name: &str, modified: &mut bool) -> Option<BufferType> {
    match matches.opt_str(name) {
        Some(value) => {
            *modified = true;
            match value.as_slice() {
                "L" => {
                    if name == "input" {
                        show_info!("line buffering stdin is meaningless");
                        None
                    } else {
                        Some(BufferType::Line)
                    }
                },
                x => {
                    let size = match parse_size(x) {
                        Some(m) => m,
                        None => { show_error!("Invalid mode {}", x); return None }
                    };
                    Some(BufferType::Size(size))
                },
            }
        },
        None => Some(BufferType::Default),
    }
}

fn parse_options(args: &[String], options: &mut ProgramOptions, optgrps: &[OptGroup]) -> Result<OkMsg, ErrMsg> {
    let matches = match getopts(args, optgrps) {
        Ok(m) => m,
        Err(_) => return Err(ErrMsg::Retry)
    };
    if matches.opt_present("help") {
        return Ok(OkMsg::Help);
    }
    if matches.opt_present("version") {
        return Ok(OkMsg::Version);
    }
    let mut modified = false;
    options.stdin = try!(check_option(&matches, "input", &mut modified).ok_or(ErrMsg::Fatal));
    options.stdout = try!(check_option(&matches, "output", &mut modified).ok_or(ErrMsg::Fatal));
    options.stderr = try!(check_option(&matches, "error", &mut modified).ok_or(ErrMsg::Fatal));

    if matches.free.len() != 1 {
        return Err(ErrMsg::Retry);
    }
    if !modified {
        show_error!("you must specify a buffering mode option");
        return Err(ErrMsg::Fatal);
    }
    Ok(OkMsg::Buffering)
}

fn set_command_env(command: &mut Command, buffer_name: &str, buffer_type: BufferType) {
    match buffer_type {
        BufferType::Size(m) => { command.env(buffer_name, m.to_string()); },
        BufferType::Line => { command.env(buffer_name, "L"); },
        BufferType::Default => {},
    }
}

fn get_preload_env() -> (String, String) {
    let (preload, extension) = preload_strings();
    let mut libstdbuf = LIBSTDBUF.to_string();
    libstdbuf.push_str(extension);
    // First search for library in directory of executable.
    let mut path = match os::self_exe_path() {
        Some(exe_path) => exe_path,
        None => crash!(1, "Impossible to fetch the path of this executable.")
    };
    path.push(libstdbuf.as_slice());
    if path.exists() {
        match path.as_str() {
            Some(s) => { return (preload.to_string(), s.to_string()); },
            None => crash!(1, "Error while converting path.")
        };
    }
    // We assume library is in LD_LIBRARY_PATH/ DYLD_LIBRARY_PATH.
    (preload.to_string(), libstdbuf)
}

pub fn uumain(args: Vec<String>) -> isize {
    let optgrps = [
        optopt("i", "input", "adjust standard input stream buffering", "MODE"),
        optopt("o", "output", "adjust standard output stream buffering", "MODE"),
        optopt("e", "error", "adjust standard error stream buffering", "MODE"),
        optflag("", "help", "display this help and exit"),
        optflag("", "version", "output version information and exit"),
    ];
    let mut options = ProgramOptions {stdin: BufferType::Default, stdout: BufferType::Default, stderr: BufferType::Default};
    let mut command_idx = -1;
    for i in range_inclusive(1, args.len()) {
        match parse_options(&args[1 .. i], &mut options, &optgrps) {
            Ok(OkMsg::Buffering) => {
                command_idx = i - 1;
                break;
            },
            Ok(OkMsg::Help) => {
                print_usage(&optgrps);
                return 0;
            },
            Ok(OkMsg::Version) => {
                print_version();
                return 0;
            },
            Err(ErrMsg::Fatal) => break,
            Err(ErrMsg::Retry) => continue,
        }
    };
    if command_idx == -1 {
        crash!(125, "Invalid options\nTry 'stdbuf --help' for more information.");
    }
    let ref command_name = args[command_idx];
    let mut command = Command::new(command_name);
    let (preload_env, libstdbuf) = get_preload_env();
    command.args(&args[command_idx + 1 ..]).env(preload_env.as_slice(), libstdbuf.as_slice());
    command.stdin(StdioContainer::InheritFd(0)).stdout(StdioContainer::InheritFd(1)).stderr(StdioContainer::InheritFd(2));
    set_command_env(&mut command, "_STDBUF_I", options.stdin);
    set_command_env(&mut command, "_STDBUF_O", options.stdout);
    set_command_env(&mut command, "_STDBUF_E", options.stderr);
    let mut process = match command.spawn() {
        Ok(p) => p,
        Err(e) => crash!(1, "failed to execute process: {}", e)
    };
    match process.wait() {
        Ok(status) => {
            match status {
                ProcessExit::ExitStatus(i) => return i,
                ProcessExit::ExitSignal(i) => crash!(1, "process killed by signal {}", i),
            }
        },
        Err(e) => crash!(1, "{}", e)
    };
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::str::FromStr;

use hyper::Uri;

use client::{self, ApiClient};

static DATASTORE_ROOT: &str = "https://datastore.googleapis.com/v1";

pub struct DatastoreService {}
pub type Hub<'a> = client::Hub<'a, DatastoreService>;

pub type ValueMap = HashMap<String, Value>;

#[derive(Serialize, Default, Debug)]
pub struct BeginTransactionRequest {}

#[derive(Deserialize, Default, Debug)]
pub struct BeginTransactionResponse {
    pub transaction: String,
}

#[derive(Serialize, Default, Debug)]
pub struct CommitRequest {
    pub transaction: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub mode: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub mutations: Option<Vec<Mutation>>,
}

#[derive(Deserialize, Default, Debug)]
pub struct CommitResponse {
    #[serde(rename="mutationResults")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mutation_results: Option<Vec<MutationResult>>,

    #[serde(rename="indexUpdates")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub index_updates: Option<i32>,
}

#[derive(Serialize, Default, Debug)]
pub struct Mutation {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub insert: Option<Entity>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub upsert: Option<Entity>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub update: Option<Entity>,

    #[serde(rename="baseVersion")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub base_version: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub delete: Option<Key>,
}

#[derive(Deserialize, Default, Debug)]
pub struct MutationResult {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub version: Option<String>,

    #[serde(rename="conflictDetected")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub conflict_detected: Option<bool>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub key: Option<Key>,
}

#[derive(Serialize, Default, Debug)]
pub struct LookupRequest {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub keys: Option<Vec<Key>>,

    #[serde(rename="readOptions")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub read_options: Option<ReadOptions>,
}

#[derive(Deserialize, Default, Debug)]
pub struct LookupResponse {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub found: Option<Vec<EntityResult>>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub missing: Option<Vec<EntityResult>>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub deferred: Option<Vec<Key>>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct ReadOptions {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub transaction: Option<String>,

    #[serde(rename="readConsistency")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub read_consistency: Option<String>,
}

#[derive(Deserialize, Default, Debug)]
pub struct EntityResult {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cursor: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub version: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub entity: Option<Entity>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct Entity {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub properties: Option<HashMap<String, Value>>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub key: Option<Key>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct Key {
    pub path: Vec<PathElement>,

    #[serde(rename="partitionId")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub partition_id: Option<PartitionId>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct PathElement {
    pub kind: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct PartitionId {
    #[serde(rename="projectId")]
    pub project_id: String,

    #[serde(rename="namespaceId")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub namespace_id: Option<String>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct Value {
    #[serde(rename="entityValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub entity_value: Option<Entity>,

    #[serde(rename="timestampValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub timestamp_value: Option<String>,

    #[serde(rename="stringValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub string_value: Option<String>,

    #[serde(rename="doubleValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub double_value: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub meaning: Option<i32>,

    #[serde(rename="excludeFromIndexes")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub exclude_from_indexes: Option<bool>,

    #[serde(rename="blobValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub blob_value: Option<String>,

    #[serde(rename="keyValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub key_value: Option<Key>,

    #[serde(rename="booleanValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub boolean_value: Option<bool>,

    #[serde(rename="integerValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub integer_value: Option<String>,

    #[serde(rename="nullValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub null_value: Option<String>,

    //#[serde(rename="geoPointValue")]
    //#[serde(skip_serializing_if = "Option::is_none")]
    //pub geo_point_value: Option<LatLng>,

    //#[serde(rename="arrayValue")]
    //#[serde(skip_serializing_if = "Option::is_none")]
    //pub array_value: Option<ArrayValue>,
}

impl<'a> Hub<'a> {
    pub fn insert_entity_auto_id(&self,
                                 kind: &str,
                                 ns: &str,
                                 props: ValueMap)
                                 -> client::Result<String> {
        let key = self.mk_key(kind, Some(ns), None, None);
        let res = self.insert(key, props)?;
        let mut results = res.mutation_results.expect("mutations to be valid");

        let mut key = results.remove(0).key.expect("key to be valid");
        Ok(key.path.remove(0).id.expect("id to be valid"))
    }

    pub fn insert_entity_by_name(&self,
                                 kind: &str,
                                 ns: &str,
                                 name: &str,
                                 props: ValueMap)
                                 -> client::Result<()> {
        let key = self.mk_key(kind, Some(ns), Some(name), None);
        self.insert(key, props)?;
        Ok(())
    }

    fn insert(&self, key: Key, props: ValueMap) -> client::Result<CommitResponse> {
        let entity = Entity {
            key: Some(key),
            properties: Some(props),
        };

        let insert = Mutation {
            insert: Some(entity),
            ..Default::default()
        };

        let req = CommitRequest {
            mutations: Some(vec![insert]),
            ..Default::default()
        };

        let txn_id = self.begin_transaction()?;
        self.commit(&txn_id, req)
    }

    pub fn lookup_by_id(&self, kind: &str, ns: &str, id: &str) -> client::Result<Option<ValueMap>> {
        let key = self.mk_key(kind, Some(ns), None, Some(id));
        self.lookup_one(key)
    }

    pub fn lookup_by_name(&self,
                          kind: &str,
                          ns: &str,
                          name: &str)
                          -> client::Result<Option<ValueMap>> {
        let key = self.mk_key(kind, Some(ns), Some(name), None);
        self.lookup_one(key)
    }

    fn lookup_one(&self, key: Key) -> client::Result<Option<ValueMap>> {
        let txn_id = self.begin_transaction()?;

        let read_options = ReadOptions {
            transaction: Some(txn_id.clone()),
            ..Default::default()
        };

        let req = LookupRequest {
            keys: Some(vec![key]),
            read_options: Some(read_options),
        };

        let uri = self.mk_uri("lookup");
        let res = self.post::<_, LookupResponse>(&uri, req, &[])?;

        self.commit(&txn_id, CommitRequest::default())?;

        Ok(res.found
               .and_then(|mut f| if f.len() != 1 {
                             None
                         } else {
                             f.remove(0).entity
                         })
               .and_then(|e| e.properties))
    }

    pub fn update_by_id(&self,
                        kind: &str,
                        ns: &str,
                        id: &str,
                        props: ValueMap)
                        -> client::Result<()> {
        let key = self.mk_key(kind, Some(ns), None, Some(id));

        let entity = Entity {
            key: Some(key),
            properties: Some(props),
        };

        let update = Mutation {
            update: Some(entity),
            ..Default::default()
        };

        let req = CommitRequest {
            mutations: Some(vec![update]),
            ..Default::default()
        };

        let txn_id = self.begin_transaction()?;
        self.commit(&txn_id, req)?;

        Ok(())
    }

    pub fn delete_by_id(&self, kind: &str, ns: &str, id: &str) -> client::Result<()> {
        let key = self.mk_key(kind, Some(ns), None, Some(id));

        let delete = Mutation {
            delete: Some(key),
            ..Default::default()
        };

        let req = CommitRequest {
            mutations: Some(vec![delete]),
            ..Default::default()
        };

        let txn_id = self.begin_transaction()?;
        self.commit(&txn_id, req)?;

        Ok(())
    }

    fn begin_transaction(&self) -> client::Result<String> {
        let uri = self.mk_uri("beginTransaction");
        let req = BeginTransactionRequest::default();
        self.post::<_, BeginTransactionResponse>(&uri, req, &[])
            .map(|r| r.transaction)
    }

    fn commit(&self, txn_id: &str, mut req: CommitRequest) -> client::Result<CommitResponse> {
        req.transaction = txn_id.to_string();
        let uri = self.mk_uri("commit");
        self.post(&uri, req, &[])
    }

    fn mk_uri(&self, action: &str) -> Uri {
        let path = format!("{}/projects/{}:{}",
                           DATASTORE_ROOT,
                           self.project_id(),
                           action);
        Uri::from_str(&path).expect("uri to be valid")
    }

    fn mk_key(&self, kind: &str, ns: Option<&str>, name: Option<&str>, id: Option<&str>) -> Key {
        let partition_id = PartitionId {
            project_id: self.project_id().to_string(),
            namespace_id: ns.map(|ns| ns.to_string()),
        };

        let path = PathElement {
            kind: kind.to_string(),
            name: name.map(|name| name.to_string()),
            id: id.map(|id| id.to_string()),
        };

        Key {
            path: vec![path],
            partition_id: Some(partition_id),
        }
    }
}
<commit_msg>Don't use a transaction for single-key lookups<commit_after>use std::collections::HashMap;
use std::str::FromStr;

use hyper::Uri;

use client::{self, ApiClient};

static DATASTORE_ROOT: &str = "https://datastore.googleapis.com/v1";

pub struct DatastoreService {}
pub type Hub<'a> = client::Hub<'a, DatastoreService>;

pub type ValueMap = HashMap<String, Value>;

#[derive(Serialize, Default, Debug)]
pub struct BeginTransactionRequest {}

#[derive(Deserialize, Default, Debug)]
pub struct BeginTransactionResponse {
    pub transaction: String,
}

#[derive(Serialize, Default, Debug)]
pub struct CommitRequest {
    pub transaction: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub mode: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub mutations: Option<Vec<Mutation>>,
}

#[derive(Deserialize, Default, Debug)]
pub struct CommitResponse {
    #[serde(rename="mutationResults")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mutation_results: Option<Vec<MutationResult>>,

    #[serde(rename="indexUpdates")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub index_updates: Option<i32>,
}

#[derive(Serialize, Default, Debug)]
pub struct Mutation {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub insert: Option<Entity>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub upsert: Option<Entity>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub update: Option<Entity>,

    #[serde(rename="baseVersion")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub base_version: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub delete: Option<Key>,
}

#[derive(Deserialize, Default, Debug)]
pub struct MutationResult {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub version: Option<String>,

    #[serde(rename="conflictDetected")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub conflict_detected: Option<bool>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub key: Option<Key>,
}

#[derive(Serialize, Default, Debug)]
pub struct LookupRequest {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub keys: Option<Vec<Key>>,

    #[serde(rename="readOptions")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub read_options: Option<ReadOptions>,
}

#[derive(Deserialize, Default, Debug)]
pub struct LookupResponse {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub found: Option<Vec<EntityResult>>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub missing: Option<Vec<EntityResult>>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub deferred: Option<Vec<Key>>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct ReadOptions {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub transaction: Option<String>,

    #[serde(rename="readConsistency")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub read_consistency: Option<String>,
}

#[derive(Deserialize, Default, Debug)]
pub struct EntityResult {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cursor: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub version: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub entity: Option<Entity>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct Entity {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub properties: Option<HashMap<String, Value>>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub key: Option<Key>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct Key {
    pub path: Vec<PathElement>,

    #[serde(rename="partitionId")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub partition_id: Option<PartitionId>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct PathElement {
    pub kind: String,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub id: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct PartitionId {
    #[serde(rename="projectId")]
    pub project_id: String,

    #[serde(rename="namespaceId")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub namespace_id: Option<String>,
}

#[derive(Serialize, Deserialize, Default, Debug)]
pub struct Value {
    #[serde(rename="entityValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub entity_value: Option<Entity>,

    #[serde(rename="timestampValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub timestamp_value: Option<String>,

    #[serde(rename="stringValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub string_value: Option<String>,

    #[serde(rename="doubleValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub double_value: Option<f64>,

    #[serde(skip_serializing_if = "Option::is_none")]
    pub meaning: Option<i32>,

    #[serde(rename="excludeFromIndexes")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub exclude_from_indexes: Option<bool>,

    #[serde(rename="blobValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub blob_value: Option<String>,

    #[serde(rename="keyValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub key_value: Option<Key>,

    #[serde(rename="booleanValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub boolean_value: Option<bool>,

    #[serde(rename="integerValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub integer_value: Option<String>,

    #[serde(rename="nullValue")]
    #[serde(skip_serializing_if = "Option::is_none")]
    pub null_value: Option<String>,

    //#[serde(rename="geoPointValue")]
    //#[serde(skip_serializing_if = "Option::is_none")]
    //pub geo_point_value: Option<LatLng>,

    //#[serde(rename="arrayValue")]
    //#[serde(skip_serializing_if = "Option::is_none")]
    //pub array_value: Option<ArrayValue>,
}

impl<'a> Hub<'a> {
    pub fn insert_entity_auto_id(&self,
                                 kind: &str,
                                 ns: &str,
                                 props: ValueMap)
                                 -> client::Result<String> {
        let key = self.mk_key(kind, Some(ns), None, None);
        let res = self.insert(key, props)?;
        let mut results = res.mutation_results.expect("mutations to be valid");

        let mut key = results.remove(0).key.expect("key to be valid");
        Ok(key.path.remove(0).id.expect("id to be valid"))
    }

    pub fn insert_entity_by_name(&self,
                                 kind: &str,
                                 ns: &str,
                                 name: &str,
                                 props: ValueMap)
                                 -> client::Result<()> {
        let key = self.mk_key(kind, Some(ns), Some(name), None);
        self.insert(key, props)?;
        Ok(())
    }

    fn insert(&self, key: Key, props: ValueMap) -> client::Result<CommitResponse> {
        let entity = Entity {
            key: Some(key),
            properties: Some(props),
        };

        let insert = Mutation {
            insert: Some(entity),
            ..Default::default()
        };

        let req = CommitRequest {
            mutations: Some(vec![insert]),
            ..Default::default()
        };

        let txn_id = self.begin_transaction()?;
        self.commit(&txn_id, req)
    }

    pub fn lookup_by_id(&self, kind: &str, ns: &str, id: &str) -> client::Result<Option<ValueMap>> {
        let key = self.mk_key(kind, Some(ns), None, Some(id));
        self.lookup_one(key)
    }

    pub fn lookup_by_name(&self,
                          kind: &str,
                          ns: &str,
                          name: &str)
                          -> client::Result<Option<ValueMap>> {
        let key = self.mk_key(kind, Some(ns), Some(name), None);
        self.lookup_one(key)
    }

    fn lookup_one(&self, key: Key) -> client::Result<Option<ValueMap>> {
        let req = LookupRequest {
            keys: Some(vec![key]),
            read_options: None,
        };

        let uri = self.mk_uri("lookup");
        let res = self.post::<_, LookupResponse>(&uri, req, &[])?;

        Ok(res.found
               .and_then(|mut f| if f.len() != 1 {
                             None
                         } else {
                             f.remove(0).entity
                         })
               .and_then(|e| e.properties))
    }

    pub fn update_by_id(&self,
                        kind: &str,
                        ns: &str,
                        id: &str,
                        props: ValueMap)
                        -> client::Result<()> {
        let key = self.mk_key(kind, Some(ns), None, Some(id));

        let entity = Entity {
            key: Some(key),
            properties: Some(props),
        };

        let update = Mutation {
            update: Some(entity),
            ..Default::default()
        };

        let req = CommitRequest {
            mutations: Some(vec![update]),
            ..Default::default()
        };

        let txn_id = self.begin_transaction()?;
        self.commit(&txn_id, req)?;

        Ok(())
    }

    pub fn delete_by_id(&self, kind: &str, ns: &str, id: &str) -> client::Result<()> {
        let key = self.mk_key(kind, Some(ns), None, Some(id));

        let delete = Mutation {
            delete: Some(key),
            ..Default::default()
        };

        let req = CommitRequest {
            mutations: Some(vec![delete]),
            ..Default::default()
        };

        let txn_id = self.begin_transaction()?;
        self.commit(&txn_id, req)?;

        Ok(())
    }

    fn begin_transaction(&self) -> client::Result<String> {
        let uri = self.mk_uri("beginTransaction");
        let req = BeginTransactionRequest::default();
        self.post::<_, BeginTransactionResponse>(&uri, req, &[])
            .map(|r| r.transaction)
    }

    fn commit(&self, txn_id: &str, mut req: CommitRequest) -> client::Result<CommitResponse> {
        req.transaction = txn_id.to_string();
        let uri = self.mk_uri("commit");
        self.post(&uri, req, &[])
    }

    fn mk_uri(&self, action: &str) -> Uri {
        let path = format!("{}/projects/{}:{}",
                           DATASTORE_ROOT,
                           self.project_id(),
                           action);
        Uri::from_str(&path).expect("uri to be valid")
    }

    fn mk_key(&self, kind: &str, ns: Option<&str>, name: Option<&str>, id: Option<&str>) -> Key {
        let partition_id = PartitionId {
            project_id: self.project_id().to_string(),
            namespace_id: ns.map(|ns| ns.to_string()),
        };

        let path = PathElement {
            kind: kind.to_string(),
            name: name.map(|name| name.to_string()),
            id: id.map(|id| id.to_string()),
        };

        Key {
            path: vec![path],
            partition_id: Some(partition_id),
        }
    }
}
<|endoftext|>


<commit_before>use rustc_serialize::{Decodable, Encodable};
use std::vec::Vec;
use swarm::socket::{SwarmResult, SwarmSocket};
use swarm::SwarmMsg;
use agent::{SwarmAddr, SwarmAgent};
use Location;
use std::io::net::ip::{SocketAddr, ToSocketAddr};
use bincode::DecoderReader;
use bincode::EncoderWriter;
use std::io::{BufReader};
use std::io::MemWriter;
use std::fmt::Show;

#[derive(Clone, Eq, PartialEq, Show, RustcDecodable, RustcEncodable)]
pub enum IronSwarmRPC<Loc> {
    HRTBT(SwarmAgent<Loc>),
    HRTBTACK(Vec<SwarmAgent<Loc>>),
    JOIN(SwarmAgent<Loc>),
    INFO(Loc, SwarmMsg<Loc>),
    BROADCAST(SwarmMsg<Loc>),
}

pub struct SwarmNetwork<Loc> {
    socket: SwarmSocket,
    local_agent: SwarmAgent<Loc>,
    neighbors: Vec<SwarmAgent<Loc>>
}

const MAX_NEIGHBORS: u8 = 3;

impl<Loc> SwarmNetwork<Loc> {
    fn new<A: ToSocketAddr>(loc: Loc, address: A) -> SwarmNetwork<Loc> {
        let mut socket = SwarmSocket::new(address);
        let addr = socket.socket_name();
        let agent = SwarmAgent::new(loc, addr);

        SwarmNetwork {
            socket: socket,
            local_agent: agent,
            neighbors: Vec::new()
        }
    }

    fn address(&mut self) -> SocketAddr {
        self.socket.socket_name()
    }

    fn update_location(&mut self, location: Loc) {
        self.local_agent.update_location(location);
    }
}

// Send RPC's.
impl<Loc: Encodable + Clone> SwarmNetwork<Loc> {
    fn send_heartbeat<A: ToSocketAddr>(&mut self,
                                       hrtbt: SwarmAgent<Loc>,
                                       dest: A) -> SwarmResult<()> {
        let rpc = IronSwarmRPC::HRTBT(hrtbt);
        self.socket.send_packet(rpc, dest)
    }

    fn send_heartbeat_ack<A: ToSocketAddr>(&mut self,
                                           neighbors: Vec<SwarmAgent<Loc>>,
                                           dest: A) -> SwarmResult<()> {
        let rpc = IronSwarmRPC::HRTBTACK(neighbors);
        self.socket.send_packet(rpc, dest)
    }

    fn send_info<A: ToSocketAddr>(&mut self,
                                  loc: Loc,
                                  msg: SwarmMsg<Loc>,
                                  dest: A) -> SwarmResult<()> {
        let rpc = IronSwarmRPC::INFO(loc, msg);
        self.socket.send_packet(rpc, dest)
    }

    fn send_broadcast<A: ToSocketAddr>(&mut self,
                                       msg: SwarmMsg<Loc>,
                                       dest: A) -> SwarmResult<()> {
        let rpc = IronSwarmRPC::BROADCAST(msg);
        self.socket.send_packet(rpc, dest)
    }

    fn send_join<A: ToSocketAddr>(&mut self,
                                  agn: SwarmAgent<Loc>,
                                  dest: A) -> SwarmResult<()> {
        let rpc = IronSwarmRPC::JOIN(agn);
        self.socket.send_packet(rpc, dest)
    }

    fn join<A: ToSocketAddr>(&mut self, address: A) {
        let agn = self.local_agent.clone();
        self.send_join(agn, address);
    }

    fn heartbeat(&mut self) -> SwarmResult<()> {
        let addresses: Vec<SwarmAddr> = self.neighbors.iter().
            map(|n| n.address().clone()).collect();

        let agn = self.local_agent.clone();
        for dest in addresses.iter() {
            try!(self.send_heartbeat(agn.clone(), dest));
        }
        Ok(())
    }
}

// React to received RPC's.
impl<Loc: Show + Location + Encodable + Decodable + PartialEq + Clone> SwarmNetwork<Loc> {
    fn next_msg(&mut self) -> SwarmResult<IronSwarmRPC<Loc>> {
        self.socket.recv_msg()
    }

    fn dispatch_rpc(&mut self) -> SwarmResult<()> {
        let rpc = try!(self.next_msg());

        match rpc {
            IronSwarmRPC::HRTBT(agn) => {
                self.respond_to_heartbeat(agn);
            }
            IronSwarmRPC::HRTBTACK(ack_vec) => {
                let new_neighbors: Vec<SwarmAgent<Loc>> = {
                    let niter = self.neighbors.iter();
                    ack_vec.
                        into_iter().
                        filter(|n| {
                            *n != self.local_agent || niter.all(|old| *old != *n)
                        }).collect()
                };
                self.neighbors.push_all(new_neighbors.as_slice());
            }
            IronSwarmRPC::JOIN(join_agn) => {
                self.route_join_request(join_agn)
            }
            IronSwarmRPC::INFO(loc, msg) => {}
            IronSwarmRPC::BROADCAST(msg) => {}
        }

        Ok(())
    }

    fn find_closest_neighbor(&self, loc: &Loc) -> Option<SwarmAgent<Loc>> {
        self.neighbors.
            iter().
            cloned().
            min_by(|a| a.location().distance(loc))
    }

    fn respond_to_heartbeat(&mut self, agn: SwarmAgent<Loc>) {
        if self.neighbors.len() >= MAX_NEIGHBORS as usize {
            return
        }

        {
            let dest = agn.address();
            let neighbors = self.neighbors.clone();

            self.send_heartbeat_ack(neighbors, dest);
        }

        if self.neighbors.iter().all(|n| *n != agn) {
            self.neighbors.push(agn);
        }
    }

    fn route_join_request(&mut self, join_agn: SwarmAgent<Loc>) {
        let closest_agent = self.find_closest_neighbor(join_agn.location());

        match closest_agent {
            Some(send_agn) => {
                let send_dist = send_agn.location().distance(join_agn.location());
                let self_dist = self.local_agent.location().distance(join_agn.location());

                if send_dist < self_dist ||
                    self.neighbors.len() >= MAX_NEIGHBORS as usize {
                        self.send_join(join_agn, send_agn.address());
                    } else {
                        self.neighbors.push(join_agn);
                    }
            }
            // TODO: Already have too many neighbors
            None => {
                self.neighbors.push(join_agn);
            }
        }
    }
}

#[cfg(test)]
mod test {
    extern crate bincode;
    use std::io::net::ip::{SocketAddr, Ipv4Addr};
    use std::io::test::next_test_port;
    use agent::{SwarmAgent};
    use artifact::SwarmArtifact;
    use swarm::socket::SwarmResult;
    use swarm::SwarmMsg;
    use Location;
    use swarm::network::IronSwarmRPC;
    use super::SwarmNetwork;
    use bincode::{decode, encode};

    fn construct_artifact() -> SwarmArtifact<isize> {
        let loc = 9is;

        SwarmArtifact::new(loc)
    }

    fn construct_agent() -> SwarmAgent<isize> {
        SwarmAgent::new(9, local_socket())
    }

    fn construct_swarm_msg() -> SwarmMsg<isize> {
        SwarmMsg::new_artifact_msg(construct_agent(), construct_artifact())
    }

    fn local_socket() -> SocketAddr {
        SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: next_test_port() }
    }

    fn bincode_rpc_tester(rpc: IronSwarmRPC<isize>) {
        let orig_rpc = rpc.clone();
        let limit = bincode::SizeLimit::Infinite;
        let encoded = bincode::encode(&rpc, limit).ok().unwrap();
        let dec_rpc: IronSwarmRPC<isize> =
            bincode::decode(encoded.as_slice()).ok().unwrap();
        assert_eq!(orig_rpc, dec_rpc);
    }

    fn send_broadcast_tester(from_nework: &mut SwarmNetwork<isize>,
                           to_network: &mut SwarmNetwork<isize>) -> SwarmResult<()> {
        let msg = construct_swarm_msg();
        let exp_rpc = IronSwarmRPC::BROADCAST(msg.clone());

        try!(from_nework.send_broadcast(msg, to_network.address()));
        let recv_rpc = try!(to_network.next_msg());

        assert_eq!(exp_rpc, recv_rpc);
        Ok(())
    }

    fn next_msg_rpc_tester(from_network: &mut SwarmNetwork<isize>,
                           to_network: &mut SwarmNetwork<isize>,
                           rpc: IronSwarmRPC<isize>) -> SwarmResult<()> {
        let orig_rpc = rpc.clone();
        let net_sock = to_network.address();
        try!(from_network.socket.send_packet(rpc, net_sock));

        let dec_rpc = try!(to_network.next_msg());
        assert_eq!(orig_rpc, dec_rpc);
        Ok(())
    }

    #[test]
    fn next_msg_test() {
        let mut ack_vec = Vec::new();
        ack_vec.push(construct_agent());
        let mut from_network = SwarmNetwork::new(0is, local_socket());
        let mut to_network = SwarmNetwork::new(1is, local_socket());

        let res = next_msg_rpc_tester(&mut from_network, &mut to_network,
                                      IronSwarmRPC::HRTBT(construct_agent()));
        assert!(res.is_ok());
        let res = next_msg_rpc_tester(&mut from_network, &mut to_network,
                                      IronSwarmRPC::HRTBTACK(ack_vec));
        assert!(res.is_ok());
        let res = next_msg_rpc_tester(&mut from_network, &mut to_network,
                                      IronSwarmRPC::JOIN(construct_agent()));
        assert!(res.is_ok());
        let res = next_msg_rpc_tester(&mut from_network, &mut to_network,
                                      IronSwarmRPC::INFO(10, construct_swarm_msg()));
        assert!(res.is_ok());
        let res = next_msg_rpc_tester(&mut from_network, &mut to_network,
                                      IronSwarmRPC::BROADCAST(construct_swarm_msg()));
        assert!(res.is_ok());
    }

    #[test]
    fn bincode_test() {
        let mut ack_vec = Vec::new();
        ack_vec.push(construct_agent());

        bincode_rpc_tester(IronSwarmRPC::HRTBT(construct_agent()));
        bincode_rpc_tester(IronSwarmRPC::HRTBTACK(ack_vec));
        bincode_rpc_tester(IronSwarmRPC::JOIN(construct_agent()));
        bincode_rpc_tester(IronSwarmRPC::INFO(10, construct_swarm_msg()));
        bincode_rpc_tester(IronSwarmRPC::BROADCAST(construct_swarm_msg()));
    }

    #[test]
    fn send_rpc_test() {
        let mut network_to = SwarmNetwork::new(0is, local_socket());
        let mut network_from = SwarmNetwork::new(10is, local_socket());

        let res = send_broadcast_tester(&mut network_from, &mut network_to);
        assert!(res.is_ok());
    }

    #[test]
    fn join_self_closest_test() {
        let mut network1 = SwarmNetwork::new(0is, local_socket());
        let mut network2 = SwarmNetwork::new(10is, local_socket());
        let mut joining = SwarmNetwork::new(2is, local_socket());
        network1.neighbors.push(network2.local_agent.clone());

        joining.join(network1.address());
        network1.dispatch_rpc();

        assert_eq!(network1.neighbors.len(), 2);
    }

    #[test]
    fn join_other_closest_test() {
        let mut network1 = SwarmNetwork::new(0is, local_socket());
        let mut network2 = SwarmNetwork::new(10is, local_socket());
        let mut joining = SwarmNetwork::new(9is, local_socket());
        network1.neighbors.push(network2.local_agent.clone());

        joining.join(network1.address());
        network1.dispatch_rpc();
        network2.dispatch_rpc();

        assert_eq!(network1.neighbors.len(), 1);
        assert_eq!(network2.neighbors.len(), 1);
    }

    #[test]
    fn join_self_closest_over_max_test() {
        let mut network1 = SwarmNetwork::new(10is, local_socket());
        let mut network2 = SwarmNetwork::new(1is, local_socket());
        let mut network3 = SwarmNetwork::new(2is, local_socket());
        let mut network4 = SwarmNetwork::new(3is, local_socket());
        let mut joining = SwarmNetwork::new(20is, local_socket());
        network1.neighbors.push(network2.local_agent.clone());
        network1.neighbors.push(network3.local_agent.clone());
        network1.neighbors.push(network4.local_agent.clone());

        joining.join(network1.address());
        network1.dispatch_rpc();
        network4.dispatch_rpc();

        assert_eq!(network1.neighbors.len(), 3);
        assert_eq!(network4.neighbors.len(), 1);
    }

    #[test]
    fn htbt_and_ack_add_neighbor_test() {
        let mut network1 = SwarmNetwork::new(0is, local_socket());
        let mut network2 = SwarmNetwork::new(10is, local_socket());
        network1.neighbors.push(network2.local_agent.clone());

        // Neighbor of ACK'd swarm agent
        let mut network3 = SwarmNetwork::new(9is, local_socket());
        network2.neighbors.push(network3.local_agent.clone());

        // New agent that is not in neighbor list.
        let mut network4 = SwarmNetwork::new(11is, local_socket());
        network4.neighbors.push(network1.local_agent.clone());

        network1.heartbeat();
        network2.dispatch_rpc();
        network1.dispatch_rpc();

        network4.heartbeat();
        network1.dispatch_rpc();

        assert_eq!(network1.neighbors.len(), 3);
    }

    #[test]
    fn hrtbt_over_max_no_ack_test() {
        let mut network1 = SwarmNetwork::new(10is, local_socket());
        let mut network2 = SwarmNetwork::new(1is, local_socket());
        let mut network3 = SwarmNetwork::new(2is, local_socket());
        let mut network4 = SwarmNetwork::new(3is, local_socket());
        let mut network5 = SwarmNetwork::new(20is, local_socket());

        network1.neighbors.push(network2.local_agent.clone());
        network1.neighbors.push(network3.local_agent.clone());
        network1.neighbors.push(network4.local_agent.clone());

        network5.neighbors.push(network1.local_agent.clone());

        network5.heartbeat();
        network1.dispatch_rpc();
        let res = network5.dispatch_rpc();
        assert!(res.is_err());

        assert_eq!(network1.neighbors.len(), 3);
        assert_eq!(network2.neighbors.len(), 0);
        assert_eq!(network3.neighbors.len(), 0);
        assert_eq!(network4.neighbors.len(), 0);
        assert_eq!(network5.neighbors.len(), 1);
    }
}
<commit_msg>Filter neighbors properly on HRTB_ACK reception.<commit_after>use rustc_serialize::{Decodable, Encodable};
use std::vec::Vec;
use swarm::socket::{SwarmResult, SwarmSocket};
use swarm::SwarmMsg;
use agent::{SwarmAddr, SwarmAgent};
use Location;
use std::io::net::ip::{SocketAddr, ToSocketAddr};
use bincode::DecoderReader;
use bincode::EncoderWriter;
use std::io::{BufReader};
use std::io::MemWriter;
use std::fmt::Show;

#[derive(Clone, Eq, PartialEq, Show, RustcDecodable, RustcEncodable)]
pub enum IronSwarmRPC<Loc> {
    HRTBT(SwarmAgent<Loc>),
    HRTBTACK(Vec<SwarmAgent<Loc>>),
    JOIN(SwarmAgent<Loc>),
    INFO(Loc, SwarmMsg<Loc>),
    BROADCAST(SwarmMsg<Loc>),
}

pub struct SwarmNetwork<Loc> {
    socket: SwarmSocket,
    local_agent: SwarmAgent<Loc>,
    neighbors: Vec<SwarmAgent<Loc>>
}

const MAX_NEIGHBORS: u8 = 3;

impl<Loc> SwarmNetwork<Loc> {
    fn new<A: ToSocketAddr>(loc: Loc, address: A) -> SwarmNetwork<Loc> {
        let mut socket = SwarmSocket::new(address);
        let addr = socket.socket_name();
        let agent = SwarmAgent::new(loc, addr);

        SwarmNetwork {
            socket: socket,
            local_agent: agent,
            neighbors: Vec::new()
        }
    }

    fn address(&mut self) -> SocketAddr {
        self.socket.socket_name()
    }

    fn update_location(&mut self, location: Loc) {
        self.local_agent.update_location(location);
    }
}

// Send RPC's.
impl<Loc: Encodable + Clone> SwarmNetwork<Loc> {
    fn send_heartbeat<A: ToSocketAddr>(&mut self,
                                       hrtbt: SwarmAgent<Loc>,
                                       dest: A) -> SwarmResult<()> {
        let rpc = IronSwarmRPC::HRTBT(hrtbt);
        self.socket.send_packet(rpc, dest)
    }

    fn send_heartbeat_ack<A: ToSocketAddr>(&mut self,
                                           neighbors: Vec<SwarmAgent<Loc>>,
                                           dest: A) -> SwarmResult<()> {
        let rpc = IronSwarmRPC::HRTBTACK(neighbors);
        self.socket.send_packet(rpc, dest)
    }

    fn send_info<A: ToSocketAddr>(&mut self,
                                  loc: Loc,
                                  msg: SwarmMsg<Loc>,
                                  dest: A) -> SwarmResult<()> {
        let rpc = IronSwarmRPC::INFO(loc, msg);
        self.socket.send_packet(rpc, dest)
    }

    fn send_broadcast<A: ToSocketAddr>(&mut self,
                                       msg: SwarmMsg<Loc>,
                                       dest: A) -> SwarmResult<()> {
        let rpc = IronSwarmRPC::BROADCAST(msg);
        self.socket.send_packet(rpc, dest)
    }

    fn send_join<A: ToSocketAddr>(&mut self,
                                  agn: SwarmAgent<Loc>,
                                  dest: A) -> SwarmResult<()> {
        let rpc = IronSwarmRPC::JOIN(agn);
        self.socket.send_packet(rpc, dest)
    }

    fn join<A: ToSocketAddr>(&mut self, address: A) {
        let agn = self.local_agent.clone();
        self.send_join(agn, address);
    }

    fn heartbeat(&mut self) -> SwarmResult<()> {
        let addresses: Vec<SwarmAddr> = self.neighbors.iter().
            map(|n| n.address().clone()).collect();

        let agn = self.local_agent.clone();
        for dest in addresses.iter() {
            try!(self.send_heartbeat(agn.clone(), dest));
        }
        Ok(())
    }
}

// React to received RPC's.
impl<Loc: Show + Location + Encodable + Decodable + PartialEq + Clone> SwarmNetwork<Loc> {
    fn next_msg(&mut self) -> SwarmResult<IronSwarmRPC<Loc>> {
        self.socket.recv_msg()
    }

    fn dispatch_rpc(&mut self) -> SwarmResult<()> {
        let rpc = try!(self.next_msg());

        match rpc {
            IronSwarmRPC::HRTBT(agn) => {
                self.respond_to_heartbeat(agn);
            }
            IronSwarmRPC::HRTBTACK(ack_vec) => {
                let new_neighbors: Vec<SwarmAgent<Loc>> = {
                    let niter = self.neighbors.iter();
                    ack_vec.
                        into_iter().
                        filter(|n| {
                            *n != self.local_agent && niter.all(|old| *old != *n)
                        }).collect()
                };
                self.neighbors.push_all(new_neighbors.as_slice());
            }
            IronSwarmRPC::JOIN(join_agn) => {
                self.route_join_request(join_agn)
            }
            IronSwarmRPC::INFO(loc, msg) => {}
            IronSwarmRPC::BROADCAST(msg) => {}
        }

        Ok(())
    }

    fn find_closest_neighbor(&self, loc: &Loc) -> Option<SwarmAgent<Loc>> {
        self.neighbors.
            iter().
            cloned().
            min_by(|a| a.location().distance(loc))
    }

    fn respond_to_heartbeat(&mut self, agn: SwarmAgent<Loc>) {
        if self.neighbors.len() >= MAX_NEIGHBORS as usize {
            return
        }

        {
            let dest = agn.address();
            let mut neighbors = self.neighbors.clone();
            // Add the sending agent to the list of neighbors.
            neighbors.push(self.local_agent.clone());
            let res = self.send_heartbeat_ack(neighbors, dest);

            if res.is_err() {
                panic!("Err: {:?}", res.err().unwrap());
            }
        }

        if self.neighbors.iter().all(|n| *n != agn) {
            self.neighbors.push(agn);
        }
    }

    fn route_join_request(&mut self, join_agn: SwarmAgent<Loc>) {
        let closest_agent = self.find_closest_neighbor(join_agn.location());

        match closest_agent {
            Some(send_agn) => {
                let send_dist = send_agn.location().distance(join_agn.location());
                let self_dist = self.local_agent.location().distance(join_agn.location());

                if send_dist < self_dist ||
                    self.neighbors.len() >= MAX_NEIGHBORS as usize {
                        self.send_join(join_agn, send_agn.address());
                    } else {
                        self.neighbors.push(join_agn);
                    }
            }
            // TODO: Already have too many neighbors
            None => {
                self.neighbors.push(join_agn);
            }
        }
    }
}

#[cfg(test)]
mod test {
    extern crate bincode;
    use std::io::net::ip::{SocketAddr, Ipv4Addr};
    use std::io::test::next_test_port;
    use agent::{SwarmAgent};
    use artifact::SwarmArtifact;
    use swarm::socket::SwarmResult;
    use swarm::SwarmMsg;
    use Location;
    use swarm::network::IronSwarmRPC;
    use super::SwarmNetwork;
    use bincode::{decode, encode};

    fn construct_artifact() -> SwarmArtifact<isize> {
        let loc = 9is;

        SwarmArtifact::new(loc)
    }

    fn construct_agent() -> SwarmAgent<isize> {
        SwarmAgent::new(9, local_socket())
    }

    fn construct_swarm_msg() -> SwarmMsg<isize> {
        SwarmMsg::new_artifact_msg(construct_agent(), construct_artifact())
    }

    fn local_socket() -> SocketAddr {
        SocketAddr { ip: Ipv4Addr(127, 0, 0, 1), port: next_test_port() }
    }

    fn bincode_rpc_tester(rpc: IronSwarmRPC<isize>) {
        let orig_rpc = rpc.clone();
        let limit = bincode::SizeLimit::Infinite;
        let encoded = bincode::encode(&rpc, limit).ok().unwrap();
        let dec_rpc: IronSwarmRPC<isize> =
            bincode::decode(encoded.as_slice()).ok().unwrap();
        assert_eq!(orig_rpc, dec_rpc);
    }

    fn send_broadcast_tester(from_nework: &mut SwarmNetwork<isize>,
                           to_network: &mut SwarmNetwork<isize>) -> SwarmResult<()> {
        let msg = construct_swarm_msg();
        let exp_rpc = IronSwarmRPC::BROADCAST(msg.clone());

        try!(from_nework.send_broadcast(msg, to_network.address()));
        let recv_rpc = try!(to_network.next_msg());

        assert_eq!(exp_rpc, recv_rpc);
        Ok(())
    }

    fn next_msg_rpc_tester(from_network: &mut SwarmNetwork<isize>,
                           to_network: &mut SwarmNetwork<isize>,
                           rpc: IronSwarmRPC<isize>) -> SwarmResult<()> {
        let orig_rpc = rpc.clone();
        let net_sock = to_network.address();
        try!(from_network.socket.send_packet(rpc, net_sock));

        let dec_rpc = try!(to_network.next_msg());
        assert_eq!(orig_rpc, dec_rpc);
        Ok(())
    }

    #[test]
    fn next_msg_test() {
        let mut ack_vec = Vec::new();
        ack_vec.push(construct_agent());
        let mut from_network = SwarmNetwork::new(0is, local_socket());
        let mut to_network = SwarmNetwork::new(1is, local_socket());

        let res = next_msg_rpc_tester(&mut from_network, &mut to_network,
                                      IronSwarmRPC::HRTBT(construct_agent()));
        assert!(res.is_ok());
        let res = next_msg_rpc_tester(&mut from_network, &mut to_network,
                                      IronSwarmRPC::HRTBTACK(ack_vec));
        assert!(res.is_ok());
        let res = next_msg_rpc_tester(&mut from_network, &mut to_network,
                                      IronSwarmRPC::JOIN(construct_agent()));
        assert!(res.is_ok());
        let res = next_msg_rpc_tester(&mut from_network, &mut to_network,
                                      IronSwarmRPC::INFO(10, construct_swarm_msg()));
        assert!(res.is_ok());
        let res = next_msg_rpc_tester(&mut from_network, &mut to_network,
                                      IronSwarmRPC::BROADCAST(construct_swarm_msg()));
        assert!(res.is_ok());
    }

    #[test]
    fn bincode_test() {
        let mut ack_vec = Vec::new();
        ack_vec.push(construct_agent());

        bincode_rpc_tester(IronSwarmRPC::HRTBT(construct_agent()));
        bincode_rpc_tester(IronSwarmRPC::HRTBTACK(ack_vec));
        bincode_rpc_tester(IronSwarmRPC::JOIN(construct_agent()));
        bincode_rpc_tester(IronSwarmRPC::INFO(10, construct_swarm_msg()));
        bincode_rpc_tester(IronSwarmRPC::BROADCAST(construct_swarm_msg()));
    }

    #[test]
    fn send_rpc_test() {
        let mut network_to = SwarmNetwork::new(0is, local_socket());
        let mut network_from = SwarmNetwork::new(10is, local_socket());

        let res = send_broadcast_tester(&mut network_from, &mut network_to);
        assert!(res.is_ok());
    }

    #[test]
    fn join_self_closest_test() {
        let mut network1 = SwarmNetwork::new(0is, local_socket());
        let mut network2 = SwarmNetwork::new(10is, local_socket());
        let mut joining = SwarmNetwork::new(2is, local_socket());
        network1.neighbors.push(network2.local_agent.clone());

        joining.join(network1.address());
        network1.dispatch_rpc();

        assert_eq!(network1.neighbors.len(), 2);
    }

    #[test]
    fn join_other_closest_test() {
        let mut network1 = SwarmNetwork::new(0is, local_socket());
        let mut network2 = SwarmNetwork::new(10is, local_socket());
        let mut joining = SwarmNetwork::new(9is, local_socket());
        network1.neighbors.push(network2.local_agent.clone());

        joining.join(network1.address());
        network1.dispatch_rpc();
        network2.dispatch_rpc();

        assert_eq!(network1.neighbors.len(), 1);
        assert_eq!(network2.neighbors.len(), 1);
    }

    #[test]
    fn join_self_closest_over_max_test() {
        let mut network1 = SwarmNetwork::new(10is, local_socket());
        let mut network2 = SwarmNetwork::new(1is, local_socket());
        let mut network3 = SwarmNetwork::new(2is, local_socket());
        let mut network4 = SwarmNetwork::new(3is, local_socket());
        let mut joining = SwarmNetwork::new(20is, local_socket());
        network1.neighbors.push(network2.local_agent.clone());
        network1.neighbors.push(network3.local_agent.clone());
        network1.neighbors.push(network4.local_agent.clone());

        joining.join(network1.address());
        network1.dispatch_rpc();
        network4.dispatch_rpc();

        assert_eq!(network1.neighbors.len(), 3);
        assert_eq!(network4.neighbors.len(), 1);
    }

    #[test]
    fn htbt_and_ack_add_neighbor_test() {
        let mut network1 = SwarmNetwork::new(0is, local_socket());
        let mut network2 = SwarmNetwork::new(10is, local_socket());
        network1.neighbors.push(network2.local_agent.clone());

        // Neighbor of ACK'd swarm agent
        let mut network3 = SwarmNetwork::new(9is, local_socket());
        network2.neighbors.push(network3.local_agent.clone());

        // New agent that is not in neighbor list.
        let mut network4 = SwarmNetwork::new(11is, local_socket());
        network4.neighbors.push(network1.local_agent.clone());
        network4.neighbors.push(network2.local_agent.clone());

        network1.heartbeat();
        let mut res = network2.dispatch_rpc();
        assert!(res.is_ok());
        res = network1.dispatch_rpc();
        assert!(res.is_ok());

        network4.heartbeat();
        res = network1.dispatch_rpc();
        assert!(res.is_ok());
        res = network2.dispatch_rpc();
        assert!(res.is_ok());
        res = network4.dispatch_rpc();
        assert!(res.is_ok());

        assert_eq!(network1.neighbors.len(), 3);
        assert_eq!(network2.neighbors.len(), 3);
    }

    #[test]
    fn hrtbt_ack_filter_neighbors_test() {
        let mut network1 = SwarmNetwork::new(0is, local_socket());
        let mut network2 = SwarmNetwork::new(10is, local_socket());
        let mut network3 = SwarmNetwork::new(9is, local_socket());

        network1.neighbors.push(network2.local_agent.clone());
        network1.neighbors.push(network3.local_agent.clone());
        network2.neighbors.push(network3.local_agent.clone());

        network1.heartbeat();
        let mut res = network2.dispatch_rpc();
        assert!(res.is_ok());
        res = network1.dispatch_rpc();
        assert!(res.is_ok());

        assert_eq!(network1.neighbors.len(), 2);
    }

    #[test]
    fn hrtbt_over_max_no_ack_test() {
        let mut network1 = SwarmNetwork::new(10is, local_socket());
        let mut network2 = SwarmNetwork::new(1is, local_socket());
        let mut network3 = SwarmNetwork::new(2is, local_socket());
        let mut network4 = SwarmNetwork::new(3is, local_socket());
        let mut network5 = SwarmNetwork::new(20is, local_socket());

        network1.neighbors.push(network2.local_agent.clone());
        network1.neighbors.push(network3.local_agent.clone());
        network1.neighbors.push(network4.local_agent.clone());

        network5.neighbors.push(network1.local_agent.clone());

        network5.heartbeat();
        network1.dispatch_rpc();
        let res = network5.dispatch_rpc();
        assert!(res.is_err());

        assert_eq!(network1.neighbors.len(), 3);
        assert_eq!(network2.neighbors.len(), 0);
        assert_eq!(network3.neighbors.len(), 0);
        assert_eq!(network4.neighbors.len(), 0);
        assert_eq!(network5.neighbors.len(), 1);
    }
}
<|endoftext|>


<commit_before>use core::mem;
use core::ops::Range;
use alloc::string::String;
use alloc::vec::Vec;

use super::data::{Map, Stat, TimeSpec};
use super::flag::*;
use super::number::*;
use super::validate::*;

// Copied from std
pub struct EscapeDefault {
    range: Range<usize>,
    data: [u8; 4],
}

pub fn escape_default(c: u8) -> EscapeDefault {
    let (data, len) = match c {
        b'\t' => ([b'\\', b't', 0, 0], 2),
        b'\r' => ([b'\\', b'r', 0, 0], 2),
        b'\n' => ([b'\\', b'n', 0, 0], 2),
        b'\\' => ([b'\\', b'\\', 0, 0], 2),
        b'\'' => ([b'\\', b'\'', 0, 0], 2),
        b'"' => ([b'\\', b'"', 0, 0], 2),
        b'\x20' ... b'\x7e' => ([c, 0, 0, 0], 1),
        _ => ([b'\\', b'x', hexify(c >> 4), hexify(c & 0xf)], 4),
    };

    return EscapeDefault { range: (0.. len), data: data };

    fn hexify(b: u8) -> u8 {
        match b {
            0 ... 9 => b'0' + b,
            _ => b'a' + b - 10,
        }
    }
}

impl Iterator for EscapeDefault {
    type Item = u8;
    fn next(&mut self) -> Option<u8> { self.range.next().map(|i| self.data[i]) }
    fn size_hint(&self) -> (usize, Option<usize>) { self.range.size_hint() }
}

struct ByteStr<'a>(&'a[u8]);

impl<'a> ::core::fmt::Debug for ByteStr<'a> {
    fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
        write!(f, "\"")?;
        for i in self.0 {
            for ch in escape_default(*i) {
                write!(f, "{}", ch as char)?;
            }
        }
        write!(f, "\"")?;
        Ok(())
    }
}

pub fn format_call(a: usize, b: usize, c: usize, d: usize, e: usize, f: usize) -> String {
    match a {
        SYS_OPEN => format!(
            "open({:?}, {:#X})",
            validate_slice(b as *const u8, c).map(ByteStr),
            d
        ),
        SYS_CHMOD => format!(
            "chmod({:?}, {:#o})",
            validate_slice(b as *const u8, c).map(ByteStr),
            d
        ),
        SYS_RMDIR => format!(
            "rmdir({:?})",
            validate_slice(b as *const u8, c).map(ByteStr)
        ),
        SYS_UNLINK => format!(
            "unlink({:?})",
            validate_slice(b as *const u8, c).map(ByteStr)
        ),
        SYS_CLOSE => format!(
            "close({})", b
        ),
        SYS_DUP => format!(
            "dup({}, {:?})",
            b,
            validate_slice(c as *const u8, d).map(ByteStr)
        ),
        SYS_DUP2 => format!(
            "dup2({}, {}, {:?})",
            b,
            c,
            validate_slice(d as *const u8, e).map(ByteStr)
        ),
        SYS_READ => format!(
            "read({}, {:#X}, {})",
            b,
            c,
            d
        ),
        SYS_WRITE => format!(
            "write({}, {:#X}, {})",
            b,
            c,
            d
        ),
        SYS_LSEEK => format!(
            "lseek({}, {}, {} ({}))",
            b,
            c as isize,
            match d {
                SEEK_SET => "SEEK_SET",
                SEEK_CUR => "SEEK_CUR",
                SEEK_END => "SEEK_END",
                _ => "UNKNOWN"
            },
            d
        ),
        SYS_FCNTL => format!(
            "fcntl({}, {} ({}), {:#X})",
            b,
            match c {
                F_DUPFD => "F_DUPFD",
                F_GETFD => "F_GETFD",
                F_SETFD => "F_SETFD",
                F_SETFL => "F_SETFL",
                F_GETFL => "F_GETFL",
                _ => "UNKNOWN"
            },
            c,
            d
        ),
        SYS_FMAP => format!(
            "fmap({}, {:?})",
            b,
            validate_slice(
                c as *const Map,
                d/mem::size_of::<Map>()
            ),
        ),
        SYS_FUNMAP => format!(
            "funmap({:#X})",
            b
        ),
        SYS_FPATH => format!(
            "fpath({}, {:#X}, {})",
            b,
            c,
            d
        ),
        SYS_FSTAT => format!(
            "fstat({}, {:?})",
            b,
            validate_slice(
                c as *const Stat,
                d/mem::size_of::<Stat>()
            ),
        ),
        SYS_FSTATVFS => format!(
            "fstatvfs({}, {:#X}, {})",
            b,
            c,
            d
        ),
        SYS_FSYNC => format!(
            "fsync({})",
            b
        ),
        SYS_FTRUNCATE => format!(
            "ftruncate({}, {})",
            b,
            c
        ),

        SYS_BRK => format!(
            "brk({:#X})",
            b
        ),
        SYS_CHDIR => format!(
            "chdir({:?})",
            validate_slice(b as *const u8, c).map(ByteStr)
        ),
        SYS_CLOCK_GETTIME => format!(
            "clock_gettime({}, {:?})",
            b,
            validate_slice_mut(c as *mut TimeSpec, 1)
        ),
        SYS_CLONE => format!(
            "clone({})",
            b
        ),
        SYS_EXIT => format!(
            "exit({})",
            b
        ),
        //TODO: Cleanup, do not allocate
        SYS_FEXEC => format!(
            "fexec({}, {:?}, {:?})",
            b,
            validate_slice(
                c as *const [usize; 2],
                d
            ).map(|slice| {
                slice.iter().map(|a|
                    validate_slice(a[0] as *const u8, a[1]).ok()
                    .and_then(|s| ::core::str::from_utf8(s).ok())
                ).collect::<Vec<Option<&str>>>()
            }),
            validate_slice(
                e as *const [usize; 2],
                f
            ).map(|slice| {
                slice.iter().map(|a|
                    validate_slice(a[0] as *const u8, a[1]).ok()
                    .and_then(|s| ::core::str::from_utf8(s).ok())
                ).collect::<Vec<Option<&str>>>()
            })
        ),
        SYS_FUTEX => format!(
            "futex({:#X} [{:?}], {}, {}, {}, {})",
            b,
            validate_slice_mut(b as *mut i32, 1).map(|uaddr| &mut uaddr[0]),
            c,
            d,
            e,
            f
        ),
        SYS_GETCWD => format!(
            "getcwd({:#X}, {})",
            b,
            c
        ),
        SYS_GETEGID => format!("getegid()"),
        SYS_GETENS => format!("getens()"),
        SYS_GETEUID => format!("geteuid()"),
        SYS_GETGID => format!("getgid()"),
        SYS_GETNS => format!("getns()"),
        SYS_GETPID => format!("getpid()"),
        SYS_GETUID => format!("getuid()"),
        SYS_IOPL => format!(
            "iopl({})",
            b
        ),
        SYS_KILL => format!(
            "kill({}, {})",
            b,
            c
        ),
        SYS_SIGRETURN => format!("sigreturn()"),
        SYS_SIGACTION => format!(
            "sigaction({}, {:#X}, {:#X}, {:#X})",
            b,
            c,
            d,
            e
        ),
        SYS_SIGPROCMASK => format!(
            "sigprocmask({}, {:?}, {:?})",
            b,
            validate_slice(c as *const [u64; 2], 1),
            validate_slice(d as *const [u64; 2], 1)
        ),
        SYS_MKNS => format!(
            "mkns({:?})",
            validate_slice(b as *const [usize; 2], c)
        ),
        SYS_NANOSLEEP => format!(
            "nanosleep({:?}, ({}, {}))",
            validate_slice(b as *const TimeSpec, 1),
            c,
            d
        ),
        SYS_PHYSALLOC => format!(
            "physalloc({})",
            b
        ),
        SYS_PHYSFREE => format!(
            "physfree({:#X}, {})",
            b,
            c
        ),
        SYS_PHYSMAP => format!(
            "physmap({:#X}, {}, {:#X})",
            b,
            c,
            d
        ),
        SYS_PHYSUNMAP => format!(
            "physunmap({:#X})",
            b
        ),
        SYS_VIRTTOPHYS => format!(
            "virttophys({:#X})",
            b
        ),
        SYS_PIPE2 => format!(
            "pipe2({:?}, {})",
            validate_slice_mut(b as *mut usize, 2),
            c
        ),
        SYS_SETREGID => format!(
            "setregid({}, {})",
            b,
            c
        ),
        SYS_SETRENS => format!(
            "setrens({}, {})",
            b,
            c
        ),
        SYS_SETREUID => format!(
            "setreuid({}, {})",
            b,
            c
        ),
        SYS_UMASK => format!(
            "umask({:#o}",
            b
        ),
        SYS_WAITPID => format!(
            "waitpid({}, {:#X}, {})",
            b,
            c,
            d
        ),
        SYS_YIELD => format!("yield()"),
        _ => format!(
            "UNKNOWN{} {:#X}({:#X}, {:#X}, {:#X}, {:#X}, {:#X})",
            a, a,
            b,
            c,
            d,
            e,
            f
        )
    }
}
<commit_msg>Add debugging for mprotect<commit_after>use core::mem;
use core::ops::Range;
use alloc::string::String;
use alloc::vec::Vec;

use super::data::{Map, Stat, TimeSpec};
use super::flag::*;
use super::number::*;
use super::validate::*;

// Copied from std
pub struct EscapeDefault {
    range: Range<usize>,
    data: [u8; 4],
}

pub fn escape_default(c: u8) -> EscapeDefault {
    let (data, len) = match c {
        b'\t' => ([b'\\', b't', 0, 0], 2),
        b'\r' => ([b'\\', b'r', 0, 0], 2),
        b'\n' => ([b'\\', b'n', 0, 0], 2),
        b'\\' => ([b'\\', b'\\', 0, 0], 2),
        b'\'' => ([b'\\', b'\'', 0, 0], 2),
        b'"' => ([b'\\', b'"', 0, 0], 2),
        b'\x20' ... b'\x7e' => ([c, 0, 0, 0], 1),
        _ => ([b'\\', b'x', hexify(c >> 4), hexify(c & 0xf)], 4),
    };

    return EscapeDefault { range: (0.. len), data: data };

    fn hexify(b: u8) -> u8 {
        match b {
            0 ... 9 => b'0' + b,
            _ => b'a' + b - 10,
        }
    }
}

impl Iterator for EscapeDefault {
    type Item = u8;
    fn next(&mut self) -> Option<u8> { self.range.next().map(|i| self.data[i]) }
    fn size_hint(&self) -> (usize, Option<usize>) { self.range.size_hint() }
}

struct ByteStr<'a>(&'a[u8]);

impl<'a> ::core::fmt::Debug for ByteStr<'a> {
    fn fmt(&self, f: &mut ::core::fmt::Formatter) -> ::core::fmt::Result {
        write!(f, "\"")?;
        for i in self.0 {
            for ch in escape_default(*i) {
                write!(f, "{}", ch as char)?;
            }
        }
        write!(f, "\"")?;
        Ok(())
    }
}

pub fn format_call(a: usize, b: usize, c: usize, d: usize, e: usize, f: usize) -> String {
    match a {
        SYS_OPEN => format!(
            "open({:?}, {:#X})",
            validate_slice(b as *const u8, c).map(ByteStr),
            d
        ),
        SYS_CHMOD => format!(
            "chmod({:?}, {:#o})",
            validate_slice(b as *const u8, c).map(ByteStr),
            d
        ),
        SYS_RMDIR => format!(
            "rmdir({:?})",
            validate_slice(b as *const u8, c).map(ByteStr)
        ),
        SYS_UNLINK => format!(
            "unlink({:?})",
            validate_slice(b as *const u8, c).map(ByteStr)
        ),
        SYS_CLOSE => format!(
            "close({})", b
        ),
        SYS_DUP => format!(
            "dup({}, {:?})",
            b,
            validate_slice(c as *const u8, d).map(ByteStr)
        ),
        SYS_DUP2 => format!(
            "dup2({}, {}, {:?})",
            b,
            c,
            validate_slice(d as *const u8, e).map(ByteStr)
        ),
        SYS_READ => format!(
            "read({}, {:#X}, {})",
            b,
            c,
            d
        ),
        SYS_WRITE => format!(
            "write({}, {:#X}, {})",
            b,
            c,
            d
        ),
        SYS_LSEEK => format!(
            "lseek({}, {}, {} ({}))",
            b,
            c as isize,
            match d {
                SEEK_SET => "SEEK_SET",
                SEEK_CUR => "SEEK_CUR",
                SEEK_END => "SEEK_END",
                _ => "UNKNOWN"
            },
            d
        ),
        SYS_FCNTL => format!(
            "fcntl({}, {} ({}), {:#X})",
            b,
            match c {
                F_DUPFD => "F_DUPFD",
                F_GETFD => "F_GETFD",
                F_SETFD => "F_SETFD",
                F_SETFL => "F_SETFL",
                F_GETFL => "F_GETFL",
                _ => "UNKNOWN"
            },
            c,
            d
        ),
        SYS_FMAP => format!(
            "fmap({}, {:?})",
            b,
            validate_slice(
                c as *const Map,
                d/mem::size_of::<Map>()
            ),
        ),
        SYS_FUNMAP => format!(
            "funmap({:#X})",
            b
        ),
        SYS_FPATH => format!(
            "fpath({}, {:#X}, {})",
            b,
            c,
            d
        ),
        SYS_FSTAT => format!(
            "fstat({}, {:?})",
            b,
            validate_slice(
                c as *const Stat,
                d/mem::size_of::<Stat>()
            ),
        ),
        SYS_FSTATVFS => format!(
            "fstatvfs({}, {:#X}, {})",
            b,
            c,
            d
        ),
        SYS_FSYNC => format!(
            "fsync({})",
            b
        ),
        SYS_FTRUNCATE => format!(
            "ftruncate({}, {})",
            b,
            c
        ),

        SYS_BRK => format!(
            "brk({:#X})",
            b
        ),
        SYS_CHDIR => format!(
            "chdir({:?})",
            validate_slice(b as *const u8, c).map(ByteStr)
        ),
        SYS_CLOCK_GETTIME => format!(
            "clock_gettime({}, {:?})",
            b,
            validate_slice_mut(c as *mut TimeSpec, 1)
        ),
        SYS_CLONE => format!(
            "clone({})",
            b
        ),
        SYS_EXIT => format!(
            "exit({})",
            b
        ),
        //TODO: Cleanup, do not allocate
        SYS_FEXEC => format!(
            "fexec({}, {:?}, {:?})",
            b,
            validate_slice(
                c as *const [usize; 2],
                d
            ).map(|slice| {
                slice.iter().map(|a|
                    validate_slice(a[0] as *const u8, a[1]).ok()
                    .and_then(|s| ::core::str::from_utf8(s).ok())
                ).collect::<Vec<Option<&str>>>()
            }),
            validate_slice(
                e as *const [usize; 2],
                f
            ).map(|slice| {
                slice.iter().map(|a|
                    validate_slice(a[0] as *const u8, a[1]).ok()
                    .and_then(|s| ::core::str::from_utf8(s).ok())
                ).collect::<Vec<Option<&str>>>()
            })
        ),
        SYS_FUTEX => format!(
            "futex({:#X} [{:?}], {}, {}, {}, {})",
            b,
            validate_slice_mut(b as *mut i32, 1).map(|uaddr| &mut uaddr[0]),
            c,
            d,
            e,
            f
        ),
        SYS_GETCWD => format!(
            "getcwd({:#X}, {})",
            b,
            c
        ),
        SYS_GETEGID => format!("getegid()"),
        SYS_GETENS => format!("getens()"),
        SYS_GETEUID => format!("geteuid()"),
        SYS_GETGID => format!("getgid()"),
        SYS_GETNS => format!("getns()"),
        SYS_GETPID => format!("getpid()"),
        SYS_GETUID => format!("getuid()"),
        SYS_IOPL => format!(
            "iopl({})",
            b
        ),
        SYS_KILL => format!(
            "kill({}, {})",
            b,
            c
        ),
        SYS_SIGRETURN => format!("sigreturn()"),
        SYS_SIGACTION => format!(
            "sigaction({}, {:#X}, {:#X}, {:#X})",
            b,
            c,
            d,
            e
        ),
        SYS_SIGPROCMASK => format!(
            "sigprocmask({}, {:?}, {:?})",
            b,
            validate_slice(c as *const [u64; 2], 1),
            validate_slice(d as *const [u64; 2], 1)
        ),
        SYS_MKNS => format!(
            "mkns({:?})",
            validate_slice(b as *const [usize; 2], c)
        ),
        SYS_MPROTECT => format!(
            "mprotect({:#X}, {}, {:#X})",
            b,
            c,
            d
        ),
        SYS_NANOSLEEP => format!(
            "nanosleep({:?}, ({}, {}))",
            validate_slice(b as *const TimeSpec, 1),
            c,
            d
        ),
        SYS_PHYSALLOC => format!(
            "physalloc({})",
            b
        ),
        SYS_PHYSFREE => format!(
            "physfree({:#X}, {})",
            b,
            c
        ),
        SYS_PHYSMAP => format!(
            "physmap({:#X}, {}, {:#X})",
            b,
            c,
            d
        ),
        SYS_PHYSUNMAP => format!(
            "physunmap({:#X})",
            b
        ),
        SYS_VIRTTOPHYS => format!(
            "virttophys({:#X})",
            b
        ),
        SYS_PIPE2 => format!(
            "pipe2({:?}, {})",
            validate_slice_mut(b as *mut usize, 2),
            c
        ),
        SYS_SETREGID => format!(
            "setregid({}, {})",
            b,
            c
        ),
        SYS_SETRENS => format!(
            "setrens({}, {})",
            b,
            c
        ),
        SYS_SETREUID => format!(
            "setreuid({}, {})",
            b,
            c
        ),
        SYS_UMASK => format!(
            "umask({:#o}",
            b
        ),
        SYS_WAITPID => format!(
            "waitpid({}, {:#X}, {})",
            b,
            c,
            d
        ),
        SYS_YIELD => format!("yield()"),
        _ => format!(
            "UNKNOWN{} {:#X}({:#X}, {:#X}, {:#X}, {:#X}, {:#X})",
            a, a,
            b,
            c,
            d,
            e,
            f
        )
    }
}
<|endoftext|>


<commit_before>#![deny(warnings)]
#![feature(const_fn)]

extern crate orbclient;

use orbclient::event;

use std::env;
use std::io::{Read, Write, self};
use std::process::{Command, Stdio};
use std::sync::{Arc, Mutex};
use std::thread;
use std::error::Error;

use console::Console;

mod console;

fn main() {
    let shell = env::args().nth(1).unwrap_or("sh".to_string());

    let width = 640;
    let height = 480;

    env::set_var("COLUMNS", format!("{}", width/8));
    env::set_var("LINES", format!("{}", height/16));
    match Command::new(&shell).stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::piped()).spawn() {
        Ok(process) => {
            let output_mutex = Arc::new(Mutex::new(Vec::new()));

            {
                let mut stdout = process.stdout.unwrap();
                let stdout_output_mutex = output_mutex.clone();
                thread::spawn(move || {
                    let term_stderr = io::stderr();
                    let mut term_stderr = term_stderr.lock();
                    'stdout: loop {
                        let mut buf = [0; 4096];
                        match stdout.read(&mut buf) {
                            Ok(0) => break 'stdout,
                            Ok(count) => {
                                if let Ok(mut stdout_output) = stdout_output_mutex.lock() {
                                    stdout_output.extend_from_slice(&buf[..count]);
                                } else {
                                    let _ = term_stderr.write(b"failed to lock stdout output mutex.\n");
                                    break 'stdout;
                                }
                            },
                            Err(err) => {
                                let _ = term_stderr.write(b"failed to read stdout: ");
                                let _ = term_stderr.write(err.description().as_bytes());
                                let _ = term_stderr.write(b"\n");
                                break 'stdout;
                            }
                        }
                    }
                });
            }

            {
                let mut stderr = process.stderr.unwrap();
                let stderr_output_mutex = output_mutex.clone();
                thread::spawn(move || {
                    let mut term_stderr = io::stderr();
                    'stderr: loop {
                        let mut buf = [0; 4096];
                        match stderr.read(&mut buf) {
                            Ok(0) => break 'stderr,
                            Ok(count) => {
                                match stderr_output_mutex.lock() {
                                    Ok(mut stderr_output) => stderr_output.extend_from_slice(&buf[..count]),
                                    Err(_) => {
                                        let _ = term_stderr.write(b"failed to lock stderr output mutex.\n");
                                        break 'stderr;
                                    }
                                }
                            },
                            Err(err) => {
                                let _ = term_stderr.write(b"failed to read stderr: ");
                                let _ = term_stderr.write(err.description().as_bytes());
                                let _ = term_stderr.write(b"\n");
                                break 'stderr;
                            }
                        }
                    }
                });
            }

            let mut stdin = process.stdin.unwrap();
            let mut console = Console::new(width, height);
            'events: loop {
                match output_mutex.lock() {
                    Ok(mut output) => {
                        if !output.is_empty() {
                            console.write(&output);
                            output.clear();
                        }
                    },
                    Err(_) => {
                        let term_stderr = io::stderr();
                        let mut term_stderr = term_stderr.lock();
                        let _ = term_stderr.write(b"failed to lock stdout mutex.\n");
                        break 'events;
                    }
                }

                for event in console.window.events() {
                    if event.code == event::EVENT_QUIT {
                        break 'events;
                    }

                    if let Some(line) = console.event(event) {
                        if let Err(err) = stdin.write(&line.as_bytes()) {
                            let term_stderr = io::stderr();
                            let mut term_stderr = term_stderr.lock();

                            let _ = term_stderr.write(b"failed to write stdin: ");
                            let _ = term_stderr.write(err.description().as_bytes());
                            let _ = term_stderr.write(b"\n");
                            break 'events;
                        }
                    }
                }

                thread::sleep_ms(30);
            }
        },
        Err(err) => {
            let term_stderr = io::stderr();
            let mut term_stderr = term_stderr.lock();
            let _ = term_stderr.write(b"failed to execute '");
            let _ = term_stderr.write(shell.as_bytes());
            let _ = term_stderr.write(b"': ");
            let _ = term_stderr.write(err.description().as_bytes());
            let _ = term_stderr.write(b"\n");
        }
    }
}
<commit_msg>Make terminal sleep for a shorter time<commit_after>#![deny(warnings)]
#![feature(const_fn)]

extern crate orbclient;

use orbclient::event;

use std::env;
use std::io::{Read, Write, self};
use std::process::{Command, Stdio};
use std::sync::{Arc, Mutex};
use std::thread;
use std::error::Error;

use console::Console;

mod console;

fn main() {
    let shell = env::args().nth(1).unwrap_or("sh".to_string());

    let width = 640;
    let height = 480;

    env::set_var("COLUMNS", format!("{}", width/8));
    env::set_var("LINES", format!("{}", height/16));
    match Command::new(&shell).stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::piped()).spawn() {
        Ok(process) => {
            let output_mutex = Arc::new(Mutex::new(Vec::new()));

            {
                let mut stdout = process.stdout.unwrap();
                let stdout_output_mutex = output_mutex.clone();
                thread::spawn(move || {
                    let term_stderr = io::stderr();
                    let mut term_stderr = term_stderr.lock();
                    'stdout: loop {
                        let mut buf = [0; 4096];
                        match stdout.read(&mut buf) {
                            Ok(0) => break 'stdout,
                            Ok(count) => {
                                if let Ok(mut stdout_output) = stdout_output_mutex.lock() {
                                    stdout_output.extend_from_slice(&buf[..count]);
                                } else {
                                    let _ = term_stderr.write(b"failed to lock stdout output mutex.\n");
                                    break 'stdout;
                                }
                            },
                            Err(err) => {
                                let _ = term_stderr.write(b"failed to read stdout: ");
                                let _ = term_stderr.write(err.description().as_bytes());
                                let _ = term_stderr.write(b"\n");
                                break 'stdout;
                            }
                        }
                    }
                });
            }

            {
                let mut stderr = process.stderr.unwrap();
                let stderr_output_mutex = output_mutex.clone();
                thread::spawn(move || {
                    let mut term_stderr = io::stderr();
                    'stderr: loop {
                        let mut buf = [0; 4096];
                        match stderr.read(&mut buf) {
                            Ok(0) => break 'stderr,
                            Ok(count) => {
                                match stderr_output_mutex.lock() {
                                    Ok(mut stderr_output) => stderr_output.extend_from_slice(&buf[..count]),
                                    Err(_) => {
                                        let _ = term_stderr.write(b"failed to lock stderr output mutex.\n");
                                        break 'stderr;
                                    }
                                }
                            },
                            Err(err) => {
                                let _ = term_stderr.write(b"failed to read stderr: ");
                                let _ = term_stderr.write(err.description().as_bytes());
                                let _ = term_stderr.write(b"\n");
                                break 'stderr;
                            }
                        }
                    }
                });
            }

            let mut stdin = process.stdin.unwrap();
            let mut console = Console::new(width, height);
            'events: loop {
                match output_mutex.lock() {
                    Ok(mut output) => {
                        if !output.is_empty() {
                            console.write(&output);
                            output.clear();
                        }
                    },
                    Err(_) => {
                        let term_stderr = io::stderr();
                        let mut term_stderr = term_stderr.lock();
                        let _ = term_stderr.write(b"failed to lock stdout mutex.\n");
                        break 'events;
                    }
                }

                for event in console.window.events() {
                    if event.code == event::EVENT_QUIT {
                        break 'events;
                    }

                    if let Some(line) = console.event(event) {
                        if let Err(err) = stdin.write(&line.as_bytes()) {
                            let term_stderr = io::stderr();
                            let mut term_stderr = term_stderr.lock();

                            let _ = term_stderr.write(b"failed to write stdin: ");
                            let _ = term_stderr.write(err.description().as_bytes());
                            let _ = term_stderr.write(b"\n");
                            break 'events;
                        }
                    }
                }

                thread::sleep_ms(1);
            }
        },
        Err(err) => {
            let term_stderr = io::stderr();
            let mut term_stderr = term_stderr.lock();
            let _ = term_stderr.write(b"failed to execute '");
            let _ = term_stderr.write(shell.as_bytes());
            let _ = term_stderr.write(b"': ");
            let _ = term_stderr.write(err.description().as_bytes());
            let _ = term_stderr.write(b"\n");
        }
    }
}
<|endoftext|>