<commit_before>extern crate futures;
extern crate grpc;
#[macro_use]
extern crate log;
extern crate log_ndc_env_logger;

mod test_misc;

use std::sync::Arc;

use futures::stream::StreamExt;

use grpc::rt::*;
use grpc::*;

use futures::executor;
use futures::{future, TryFutureExt, TryStreamExt};
use std::thread;
use test_misc::*;

/// Single method server on random port
fn new_server<H>(service: &str, method: &str, handler: H) -> Server
where
    H: MethodHandler<String, String> + 'static + Sync + Send,
    H: GrpcStreamingFlavor,
{
    let mut methods = Vec::new();
    methods.push(ServerMethod::new(
        string_string_method(
            &format!("{}{}", service, method),
            <H as GrpcStreamingFlavor>::streaming(),
        ),
        handler,
    ));
    let mut server = ServerBuilder::new_plain();
    server.http.set_port(0);
    server.add_service(ServerServiceDefinition::new(service, methods));
    server.build().expect("server")
}

/// Single unary method server
fn new_server_unary<H>(service: &str, method: &str, handler: H) -> Server
where
    H: Fn(
            ServerHandlerContext,
            ServerRequestSingle<String>,
            ServerResponseUnarySink<String>,
        ) -> grpc::Result<()>
        + Sync
        + Send
        + 'static,
{
    new_server(service, method, MethodHandlerUnary::new(handler))
}

/// Single server streaming method server
fn new_server_server_streaming<H>(service: &str, method: &str, handler: H) -> Server
where
    H: Fn(
            ServerHandlerContext,
            ServerRequestSingle<String>,
            ServerResponseSink<String>,
        ) -> grpc::Result<()>
        + Sync
        + Send
        + 'static,
{
    new_server(service, method, MethodHandlerServerStreaming::new(handler))
}

/// Single server streaming method server
fn new_server_client_streaming<H>(service: &str, method: &str, handler: H) -> Server
where
    H: Fn(
            ServerHandlerContext,
            ServerRequest<String>,
            ServerResponseUnarySink<String>,
        ) -> grpc::Result<()>
        + Sync
        + Send
        + 'static,
{
    new_server(service, method, MethodHandlerClientStreaming::new(handler))
}

/// Tester for unary methods
struct TesterUnary {
    name: String,
    client: Client,
    _server: Server,
}

impl TesterUnary {
    fn new<H>(handler: H) -> TesterUnary
    where
        H: Fn(
                ServerHandlerContext,
                ServerRequestSingle<String>,
                ServerResponseUnarySink<String>,
            ) -> grpc::Result<()>
            + Sync
            + Send
            + 'static,
    {
        let server = new_server_unary("/text", "/Unary", handler);
        let port = server.local_addr().port().expect("port");
        TesterUnary {
            name: "/text/Unary".to_owned(),
            _server: server,
            client: ClientBuilder::new(BIND_HOST, port).build().unwrap(),
        }
    }

    fn call(&self, param: &str) -> GrpcFuture<String> {
        self.client
            .call_unary(
                RequestOptions::new(),
                param.to_owned(),
                string_string_method(&self.name, GrpcStreaming::Unary),
            )
            .drop_metadata()
    }

    fn call_expect_error<F: FnOnce(&Error) -> bool>(&self, param: &str, expect: F) {
        match executor::block_on(self.call(param)) {
            Ok(r) => panic!("expecting error, got: {:?}", r),
            Err(e) => assert!(expect(&e), "wrong error: {:?}", e),
        }
    }

    fn call_expect_grpc_error<F: FnOnce(&str) -> bool>(&self, param: &str, expect: F) {
        self.call_expect_error(param, |e| match e {
            &Error::GrpcMessage(GrpcMessageError {
                ref grpc_message, ..
            }) if expect(&grpc_message) => true,
            _ => false,
        });
    }

    fn call_expect_grpc_error_contain(&self, param: &str, expect: &str) {
        self.call_expect_grpc_error(param, |m| m.find(expect).is_some());
    }
}

/// Tester for server streaming methods
struct TesterServerStreaming {
    name: String,
    client: Client,
    _server: Server,
}

impl TesterServerStreaming {
    fn new<H>(handler: H) -> Self
    where
        H: Fn(
                ServerHandlerContext,
                ServerRequestSingle<String>,
                ServerResponseSink<String>,
            ) -> grpc::Result<()>
            + Sync
            + Send
            + 'static,
    {
        let server = new_server_server_streaming("/test", "/ServerStreaming", handler);
        let port = server.local_addr().port().expect("port");
        TesterServerStreaming {
            name: "/test/ServerStreaming".to_owned(),
            _server: server,
            client: ClientBuilder::new(BIND_HOST, port).build().unwrap(),
        }
    }

    fn call(&self, param: &str) -> GrpcStream<String> {
        self.client
            .call_server_streaming(
                RequestOptions::new(),
                param.to_owned(),
                string_string_method(&self.name, GrpcStreaming::ServerStreaming),
            )
            .drop_metadata()
    }
}

struct TesterClientStreaming {
    name: String,
    client: Client,
    _server: Server,
}

impl TesterClientStreaming {
    fn new<H>(handler: H) -> Self
    where
        H: Fn(
                ServerHandlerContext,
                ServerRequest<String>,
                ServerResponseUnarySink<String>,
            ) -> grpc::Result<()>
            + Sync
            + Send
            + 'static,
    {
        let server = new_server_client_streaming("/test", "/ClientStreaming", handler);
        let port = server.local_addr().port().expect("port");
        TesterClientStreaming {
            name: "/test/ClientStreaming".to_owned(),
            _server: server,
            client: ClientBuilder::new(BIND_HOST, port).build().unwrap(),
        }
    }

    fn call(&self) -> (ClientRequestSink<String>, SingleResponse<String>) {
        executor::block_on(self.client.call_client_streaming(
            RequestOptions::new(),
            string_string_method(&self.name, GrpcStreaming::ClientStreaming),
        ))
        .unwrap()
    }
}

#[test]
fn unary() {
    init_logger();

    let tester = TesterUnary::new(|_m, req, resp| resp.finish(req.message));

    assert_eq!("aa", executor::block_on(tester.call("aa")).unwrap());
}

#[test]
fn error_in_handler() {
    init_logger();

    let tester = TesterUnary::new(|_m, _req, _resp| Err(grpc::Error::Other("my error")));

    tester.call_expect_grpc_error_contain("aa", "grpc server handler did not close the sender");
}

// TODO
//#[test]
fn _panic_in_handler() {
    init_logger();

    let tester = TesterUnary::new(|_m, _req, _resp| panic!("icnap"));

    tester.call_expect_grpc_error("aa", |m| {
        m.find("Panic").is_some() && m.find("icnap").is_some()
    });
}

#[test]
fn server_streaming() {
    init_logger();

    let test_sync = Arc::new(TestSync::new());
    let test_sync_server = test_sync.clone();

    let tester = TesterServerStreaming::new(move |_m, req, mut resp| {
        let sync = test_sync_server.clone();
        thread::spawn(move || {
            for i in 0..3 {
                sync.take(i * 2 + 1);
                resp.send_data(format!("{}{}", req.message, i)).unwrap();
            }
            resp.send_trailers(Metadata::new()).unwrap();
        });
        Ok(())
    });

    let mut rs = tester.call("x");

    test_sync.take(0);
    assert_eq!("x0", executor::block_on(rs.next()).unwrap().unwrap());
    test_sync.take(2);
    assert_eq!("x1", executor::block_on(rs.next()).unwrap().unwrap());
    test_sync.take(4);
    assert_eq!("x2", executor::block_on(rs.next()).unwrap().unwrap());
    assert!(executor::block_on(rs.next()).is_none());
}

#[test]
fn client_streaming() {
    init_logger();

    let tester = TesterClientStreaming::new(move |m, req, resp| {
        let request_stream = req.into_stream();
        m.ctx.loop_remote().spawn(
            request_stream
                .try_fold(String::new(), |mut s, message| {
                    s.push_str(&message);
                    future::ok::<_, Error>(s)
                })
                .map_ok(|r| {
                    resp.finish(r).unwrap();
                })
                .map_err(|_| ()),
        );
        Ok(())
    });

    let (mut tx, result) = tester.call();
    tx.send_data("aa".to_owned()).expect("aa");
    tx.send_data("bb".to_owned()).expect("bb");
    tx.send_data("cc".to_owned()).expect("cc");
    tx.finish().expect("finish");

    assert_eq!(
        "aabbcc",
        executor::block_on(result.drop_metadata()).unwrap()
    );
}
<commit_msg>Use a little of async/await<commit_after>extern crate futures;
extern crate grpc;
#[macro_use]
extern crate log;
extern crate log_ndc_env_logger;

mod test_misc;

use std::sync::Arc;

use futures::stream::StreamExt;

use grpc::rt::*;
use grpc::*;

use futures::executor;
use futures::{future, TryFutureExt, TryStreamExt};
use std::thread;
use test_misc::*;

/// Single method server on random port
fn new_server<H>(service: &str, method: &str, handler: H) -> Server
where
    H: MethodHandler<String, String> + 'static + Sync + Send,
    H: GrpcStreamingFlavor,
{
    let mut methods = Vec::new();
    methods.push(ServerMethod::new(
        string_string_method(
            &format!("{}{}", service, method),
            <H as GrpcStreamingFlavor>::streaming(),
        ),
        handler,
    ));
    let mut server = ServerBuilder::new_plain();
    server.http.set_port(0);
    server.add_service(ServerServiceDefinition::new(service, methods));
    server.build().expect("server")
}

/// Single unary method server
fn new_server_unary<H>(service: &str, method: &str, handler: H) -> Server
where
    H: Fn(
            ServerHandlerContext,
            ServerRequestSingle<String>,
            ServerResponseUnarySink<String>,
        ) -> grpc::Result<()>
        + Sync
        + Send
        + 'static,
{
    new_server(service, method, MethodHandlerUnary::new(handler))
}

/// Single server streaming method server
fn new_server_server_streaming<H>(service: &str, method: &str, handler: H) -> Server
where
    H: Fn(
            ServerHandlerContext,
            ServerRequestSingle<String>,
            ServerResponseSink<String>,
        ) -> grpc::Result<()>
        + Sync
        + Send
        + 'static,
{
    new_server(service, method, MethodHandlerServerStreaming::new(handler))
}

/// Single server streaming method server
fn new_server_client_streaming<H>(service: &str, method: &str, handler: H) -> Server
where
    H: Fn(
            ServerHandlerContext,
            ServerRequest<String>,
            ServerResponseUnarySink<String>,
        ) -> grpc::Result<()>
        + Sync
        + Send
        + 'static,
{
    new_server(service, method, MethodHandlerClientStreaming::new(handler))
}

/// Tester for unary methods
struct TesterUnary {
    name: String,
    client: Client,
    _server: Server,
}

impl TesterUnary {
    fn new<H>(handler: H) -> TesterUnary
    where
        H: Fn(
                ServerHandlerContext,
                ServerRequestSingle<String>,
                ServerResponseUnarySink<String>,
            ) -> grpc::Result<()>
            + Sync
            + Send
            + 'static,
    {
        let server = new_server_unary("/text", "/Unary", handler);
        let port = server.local_addr().port().expect("port");
        TesterUnary {
            name: "/text/Unary".to_owned(),
            _server: server,
            client: ClientBuilder::new(BIND_HOST, port).build().unwrap(),
        }
    }

    fn call(&self, param: &str) -> GrpcFuture<String> {
        self.client
            .call_unary(
                RequestOptions::new(),
                param.to_owned(),
                string_string_method(&self.name, GrpcStreaming::Unary),
            )
            .drop_metadata()
    }

    fn call_expect_error<F: FnOnce(&Error) -> bool>(&self, param: &str, expect: F) {
        match executor::block_on(self.call(param)) {
            Ok(r) => panic!("expecting error, got: {:?}", r),
            Err(e) => assert!(expect(&e), "wrong error: {:?}", e),
        }
    }

    fn call_expect_grpc_error<F: FnOnce(&str) -> bool>(&self, param: &str, expect: F) {
        self.call_expect_error(param, |e| match e {
            &Error::GrpcMessage(GrpcMessageError {
                ref grpc_message, ..
            }) if expect(&grpc_message) => true,
            _ => false,
        });
    }

    fn call_expect_grpc_error_contain(&self, param: &str, expect: &str) {
        self.call_expect_grpc_error(param, |m| m.find(expect).is_some());
    }
}

/// Tester for server streaming methods
struct TesterServerStreaming {
    name: String,
    client: Client,
    _server: Server,
}

impl TesterServerStreaming {
    fn new<H>(handler: H) -> Self
    where
        H: Fn(
                ServerHandlerContext,
                ServerRequestSingle<String>,
                ServerResponseSink<String>,
            ) -> grpc::Result<()>
            + Sync
            + Send
            + 'static,
    {
        let server = new_server_server_streaming("/test", "/ServerStreaming", handler);
        let port = server.local_addr().port().expect("port");
        TesterServerStreaming {
            name: "/test/ServerStreaming".to_owned(),
            _server: server,
            client: ClientBuilder::new(BIND_HOST, port).build().unwrap(),
        }
    }

    fn call(&self, param: &str) -> GrpcStream<String> {
        self.client
            .call_server_streaming(
                RequestOptions::new(),
                param.to_owned(),
                string_string_method(&self.name, GrpcStreaming::ServerStreaming),
            )
            .drop_metadata()
    }
}

struct TesterClientStreaming {
    name: String,
    client: Client,
    _server: Server,
}

impl TesterClientStreaming {
    fn new<H>(handler: H) -> Self
    where
        H: Fn(
                ServerHandlerContext,
                ServerRequest<String>,
                ServerResponseUnarySink<String>,
            ) -> grpc::Result<()>
            + Sync
            + Send
            + 'static,
    {
        let server = new_server_client_streaming("/test", "/ClientStreaming", handler);
        let port = server.local_addr().port().expect("port");
        TesterClientStreaming {
            name: "/test/ClientStreaming".to_owned(),
            _server: server,
            client: ClientBuilder::new(BIND_HOST, port).build().unwrap(),
        }
    }

    fn call(&self) -> (ClientRequestSink<String>, SingleResponse<String>) {
        executor::block_on(self.client.call_client_streaming(
            RequestOptions::new(),
            string_string_method(&self.name, GrpcStreaming::ClientStreaming),
        ))
        .unwrap()
    }
}

#[test]
fn unary() {
    init_logger();

    let tester = TesterUnary::new(|_m, req, resp| resp.finish(req.message));

    assert_eq!("aa", executor::block_on(tester.call("aa")).unwrap());
}

#[test]
fn error_in_handler() {
    init_logger();

    let tester = TesterUnary::new(|_m, _req, _resp| Err(grpc::Error::Other("my error")));

    tester.call_expect_grpc_error_contain("aa", "grpc server handler did not close the sender");
}

// TODO
//#[test]
fn _panic_in_handler() {
    init_logger();

    let tester = TesterUnary::new(|_m, _req, _resp| panic!("icnap"));

    tester.call_expect_grpc_error("aa", |m| {
        m.find("Panic").is_some() && m.find("icnap").is_some()
    });
}

#[test]
fn server_streaming() {
    init_logger();

    let test_sync = Arc::new(TestSync::new());
    let test_sync_server = test_sync.clone();

    let tester = TesterServerStreaming::new(move |_m, req, mut resp| {
        let sync = test_sync_server.clone();
        thread::spawn(move || {
            for i in 0..3 {
                sync.take(i * 2 + 1);
                resp.send_data(format!("{}{}", req.message, i)).unwrap();
            }
            resp.send_trailers(Metadata::new()).unwrap();
        });
        Ok(())
    });

    let mut rs = tester.call("x");

    test_sync.take(0);

    executor::block_on(async {
        assert_eq!("x0", rs.next().await.unwrap().unwrap());
        test_sync.take(2);
        assert_eq!("x1", rs.next().await.unwrap().unwrap());
        test_sync.take(4);
        assert_eq!("x2", rs.next().await.unwrap().unwrap());
        assert!(rs.next().await.is_none());
    });
}

#[test]
fn client_streaming() {
    init_logger();

    let tester = TesterClientStreaming::new(move |m, req, resp| {
        let request_stream = req.into_stream();
        m.ctx.loop_remote().spawn(
            request_stream
                .try_fold(String::new(), |mut s, message| {
                    s.push_str(&message);
                    future::ok::<_, Error>(s)
                })
                .map_ok(|r| {
                    resp.finish(r).unwrap();
                })
                .map_err(|_| ()),
        );
        Ok(())
    });

    let (mut tx, result) = tester.call();
    tx.send_data("aa".to_owned()).expect("aa");
    tx.send_data("bb".to_owned()).expect("bb");
    tx.send_data("cc".to_owned()).expect("cc");
    tx.finish().expect("finish");

    assert_eq!(
        "aabbcc",
        executor::block_on(result.drop_metadata()).unwrap()
    );
}
<|endoftext|>


<commit_before>extern crate crypto;

use std::fs::File;
use std::io::{BufReader, BufRead};

use crypto::aessafe::{AesSafe128Decryptor, AesSafe128Encryptor};
use crypto::symmetriccipher::{BlockDecryptor, BlockEncryptor};

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

fn vec_u8_from_hex_string(s: &String) -> Vec<u8> {
    s.as_bytes().chunks(2).map(|a| hex_to_u8(a[0]) * 16 + hex_to_u8(a[1])).collect()
}

fn u8_slice_inc(v: &mut [u8]) {
    let mut inc = 1;

    for i in (0..v.len()).rev() {
        if v[i] == 0xff && inc == 1 {
            v[i] = 0;
            inc = 1;
        } else {
            v[i] += inc;
            inc = 0;
        }
    }
}

fn ctr_decrypt(key: &Vec<u8>, msg: &Vec<u8>) -> Vec<u8> {
    let mut v = Vec::with_capacity(msg.len());
    let (iv, msg) = msg.split_at(16);
    // TODO: find better way to copy arrays
    let mut prev = &mut [0; 16];

    for i in 0..16 {
        prev[i] = iv[i];
    }

    let mut m_i = &mut [0; 16];
    let decryptor = AesSafe128Encryptor::new(&key[..]);

    for c_i in msg.chunks(16) {
        decryptor.encrypt_block(prev, &mut m_i[..]);

        for i in 0..c_i.len() {
            v.push(c_i[i] ^ m_i[i]);
        }

        u8_slice_inc(prev);
    }

    v
}

fn main() {
    let input_file = File::open("in.txt").unwrap();
    let input = BufReader::new(input_file);
    let input: Vec<_> = input.lines().filter_map(|res| res.ok())
        .filter(|s| s.len() > 0)
        .map(|s| vec_u8_from_hex_string(&s))
        .collect();
    let input: Vec<_> = input.chunks(2).collect();

    for v in input {
        println!("key: {:?}", v[0]);
        println!("cpt: {:?}", v[1]);

        let msg = ctr_decrypt(&v[0], &v[1]);
        println!("msg: {:?}", msg);

        let pt = String::from_utf8(msg).unwrap();
        println!("pt: {:?}", pt);

        println!("");
    }
}

<commit_msg>ctr cleanup<commit_after>extern crate crypto;

use std::fs::File;
use std::io::{BufReader, BufRead};

use crypto::aessafe::AesSafe128Encryptor;
use crypto::symmetriccipher::BlockEncryptor;

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

fn vec_u8_from_hex_string(s: &String) -> Vec<u8> {
    s.as_bytes().chunks(2).map(|a| hex_to_u8(a[0]) * 16 + hex_to_u8(a[1])).collect()
}

fn u8_slice_inc(v: &mut [u8]) {
    let mut inc = 1;

    for i in (0..v.len()).rev() {
        if v[i] == 0xff && inc == 1 {
            v[i] = 0;
            inc = 1;
        } else {
            v[i] += inc;
            inc = 0;
        }
    }
}

fn ctr_decrypt(key: &Vec<u8>, msg: &Vec<u8>) -> Vec<u8> {
    let mut v = Vec::with_capacity(msg.len());
    let (iv, msg) = msg.split_at(16);
    // TODO: find better way to copy arrays
    let mut prev = &mut [0; 16];

    for i in 0..16 {
        prev[i] = iv[i];
    }

    let mut m_i = &mut [0; 16];
    let encryptor = AesSafe128Encryptor::new(&key[..]);

    for c_i in msg.chunks(16) {
        encryptor.encrypt_block(prev, &mut m_i[..]);

        for i in 0..c_i.len() {
            v.push(c_i[i] ^ m_i[i]);
        }

        u8_slice_inc(prev);
    }

    v
}

fn main() {
    let input_file = File::open("in.txt").unwrap();
    let input = BufReader::new(input_file);
    let input: Vec<_> = input.lines().filter_map(|res| res.ok())
        .filter(|s| s.len() > 0)
        .map(|s| vec_u8_from_hex_string(&s))
        .collect();
    let input: Vec<_> = input.chunks(2).collect();

    for v in input {
        println!("key: {:?}", v[0]);
        println!("cpt: {:?}", v[1]);

        let msg = ctr_decrypt(&v[0], &v[1]);
        println!("msg: {:?}", msg);

        let pt = String::from_utf8(msg).unwrap();
        println!("pt: {:?}", pt);

        println!("");
    }
}

<|endoftext|>


<commit_before>use std::io;

use legion_prof::analyze::print_statistics;
use legion_prof::serialize::deserialize;
use legion_prof::state::{State, Timestamp};
use legion_prof::trace_viewer::emit_trace;
use legion_prof::visualize::emit_interactive_visualization;

fn main() -> io::Result<()> {
    let matches = clap::App::new("Integrity Checker")
        .about("Legion Prof: application profiler")
        .arg(
            clap::Arg::with_name("filenames")
                .help("input Legion Prof log filenames")
                .required(true)
                .multiple(true),
        )
        .arg(
            clap::Arg::with_name("output")
                .short("o")
                .long("output")
                .takes_value(true)
                .default_value("legion_prof")
                .help("output directory pathname"),
        )
        .arg(
            clap::Arg::with_name("start-trim")
                .long("start-trim")
                .takes_value(true)
                .help("start time in micro-seconds to trim the profile"),
        )
        .arg(
            clap::Arg::with_name("stop-trim")
                .long("stop-trim")
                .takes_value(true)
                .help("stop time in micro-seconds to trim the profile"),
        )
        .arg(
            clap::Arg::with_name("force")
                .short("f")
                .long("force")
                .help("overwrite output directory if it exists"),
        )
        .arg(
            clap::Arg::with_name("statistics")
                .short("s")
                .long("statistics")
                .help("print statistics"),
        )
        .arg(
            clap::Arg::with_name("trace")
                .short("t")
                .long("trace-viewer")
                .help("emit JSON for Google Trace Viewer"),
        )
        .get_matches();

    let filenames = matches.values_of_os("filenames").unwrap();
    let output = matches.value_of_os("output").unwrap();
    let force = matches.is_present("force");
    let statistics = matches.is_present("stats");
    let trace = matches.is_present("trace");
    let start_trim = matches
        .value_of("start-trim")
        .map(|x| Timestamp(x.parse::<u64>().unwrap()));
    let stop_trim = matches
        .value_of("stop-trim")
        .map(|x| Timestamp(x.parse::<u64>().unwrap()));

    let mut state = State::default();
    for filename in filenames {
        println!("Reading log file {:?}...", filename);
        let records = deserialize(filename)?;
        println!("Matched {} objects", records.len());
        state.process_records(&records);
    }

    state.trim_time_range(start_trim, stop_trim);
    state.sort_time_range();
    if statistics {
        print_statistics(&mut state);
    }
    if trace {
        emit_trace(&mut state, output, force)?;
    }
    if !statistics && !trace {
        state.assign_colors();
        emit_interactive_visualization(&mut state, output, force, true, true, true)?;
    }

    Ok(())
}
<commit_msg>legion_prof_rs: Measure trim in correct units.<commit_after>use std::io;

use legion_prof::analyze::print_statistics;
use legion_prof::serialize::deserialize;
use legion_prof::state::{State, Timestamp};
use legion_prof::trace_viewer::emit_trace;
use legion_prof::visualize::emit_interactive_visualization;

fn main() -> io::Result<()> {
    let matches = clap::App::new("Integrity Checker")
        .about("Legion Prof: application profiler")
        .arg(
            clap::Arg::with_name("filenames")
                .help("input Legion Prof log filenames")
                .required(true)
                .multiple(true),
        )
        .arg(
            clap::Arg::with_name("output")
                .short("o")
                .long("output")
                .takes_value(true)
                .default_value("legion_prof")
                .help("output directory pathname"),
        )
        .arg(
            clap::Arg::with_name("start-trim")
                .long("start-trim")
                .takes_value(true)
                .help("start time in micro-seconds to trim the profile"),
        )
        .arg(
            clap::Arg::with_name("stop-trim")
                .long("stop-trim")
                .takes_value(true)
                .help("stop time in micro-seconds to trim the profile"),
        )
        .arg(
            clap::Arg::with_name("force")
                .short("f")
                .long("force")
                .help("overwrite output directory if it exists"),
        )
        .arg(
            clap::Arg::with_name("statistics")
                .short("s")
                .long("statistics")
                .help("print statistics"),
        )
        .arg(
            clap::Arg::with_name("trace")
                .short("t")
                .long("trace-viewer")
                .help("emit JSON for Google Trace Viewer"),
        )
        .get_matches();

    let filenames = matches.values_of_os("filenames").unwrap();
    let output = matches.value_of_os("output").unwrap();
    let force = matches.is_present("force");
    let statistics = matches.is_present("stats");
    let trace = matches.is_present("trace");
    let start_trim = matches
        .value_of("start-trim")
        .map(|x| Timestamp(x.parse::<u64>().unwrap() * 1000));
    let stop_trim = matches
        .value_of("stop-trim")
        .map(|x| Timestamp(x.parse::<u64>().unwrap() * 1000));

    let mut state = State::default();
    for filename in filenames {
        println!("Reading log file {:?}...", filename);
        let records = deserialize(filename)?;
        println!("Matched {} objects", records.len());
        state.process_records(&records);
    }

    state.trim_time_range(start_trim, stop_trim);
    state.sort_time_range();
    if statistics {
        print_statistics(&mut state);
    }
    if trace {
        emit_trace(&mut state, output, force)?;
    }
    if !statistics && !trace {
        state.assign_colors();
        emit_interactive_visualization(&mut state, output, force, true, true, true)?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! A Job is an abstraction of async operation in service worker lifecycle propagation.
//! Each Job is uniquely identified by its scope_url, and is keyed accordingly under
//! the script thread. The script thread contains a JobQueue, which stores all scheduled Jobs
//! by multiple service worker clients in a Vec.

use dom::bindings::cell::DOMRefCell;
use dom::bindings::error::Error;
use dom::bindings::js::JS;
use dom::bindings::refcounted::{Trusted, TrustedPromise};
use dom::bindings::reflector::DomObject;
use dom::client::Client;
use dom::globalscope::GlobalScope;
use dom::promise::Promise;
use dom::serviceworkerregistration::ServiceWorkerRegistration;
use dom::urlhelper::UrlHelper;
use js::jsapi::JSAutoCompartment;
use script_thread::ScriptThread;
use servo_url::ServoUrl;
use std::cmp::PartialEq;
use std::collections::HashMap;
use std::rc::Rc;
use task::Task;
use task_source::TaskSource;
use task_source::dom_manipulation::DOMManipulationTaskSource;

#[derive(Clone, Copy, Debug, JSTraceable, PartialEq)]
pub enum JobType {
    Register,
    Unregister,
    Update
}

#[derive(Clone)]
pub enum SettleType {
    Resolve(Trusted<ServiceWorkerRegistration>),
    Reject(Error)
}

#[must_root]
#[derive(JSTraceable)]
pub struct Job {
    pub job_type: JobType,
    pub scope_url: ServoUrl,
    pub script_url: ServoUrl,
    pub promise: Rc<Promise>,
    pub equivalent_jobs: Vec<Job>,
    // client can be a window client, worker client so `Client` will be an enum in future
    pub client: JS<Client>,
    pub referrer: ServoUrl
}

impl Job {
    #[allow(unrooted_must_root)]
    // https://w3c.github.io/ServiceWorker/#create-job-algorithm
    pub fn create_job(job_type: JobType,
                      scope_url: ServoUrl,
                      script_url: ServoUrl,
                      promise: Rc<Promise>,
                      client: &Client) -> Job {
        Job {
            job_type: job_type,
            scope_url: scope_url,
            script_url: script_url,
            promise: promise,
            equivalent_jobs: vec![],
            client: JS::from_ref(client),
            referrer: client.creation_url()
        }
    }
    #[allow(unrooted_must_root)]
    pub fn append_equivalent_job(&mut self, job: Job) {
        self.equivalent_jobs.push(job);
    }
}

impl PartialEq for Job {
    // Equality criteria as described in https://w3c.github.io/ServiceWorker/#dfn-job-equivalent
    fn eq(&self, other: &Self) -> bool {
        let same_job = self.job_type == other.job_type;
        if same_job {
            match self.job_type {
                JobType::Register | JobType::Update => {
                    self.scope_url == other.scope_url && self.script_url == other.script_url
                },
                JobType::Unregister => self.scope_url == other.scope_url
            }
        } else {
            false
        }
    }
}

#[must_root]
#[derive(JSTraceable)]
pub struct JobQueue(pub DOMRefCell<HashMap<ServoUrl, Vec<Job>>>);

impl JobQueue {
    pub fn new() -> JobQueue {
        JobQueue(DOMRefCell::new(HashMap::new()))
    }
    #[allow(unrooted_must_root)]
    // https://w3c.github.io/ServiceWorker/#schedule-job-algorithm
    pub fn schedule_job(&self, job: Job, script_thread: &ScriptThread) {
        debug!("scheduling {:?} job", job.job_type);
        let mut queue_ref = self.0.borrow_mut();
        let job_queue = queue_ref.entry(job.scope_url.clone()).or_insert(vec![]);
        // Step 1
        if job_queue.is_empty() {
            let scope_url = job.scope_url.clone();
            job_queue.push(job);
            let _ = script_thread.schedule_job_queue(scope_url);
            debug!("queued task to run newly-queued job");
        } else {
            // Step 2
            let mut last_job = job_queue.pop().unwrap();
            if job == last_job && !last_job.promise.is_settled() {
                last_job.append_equivalent_job(job);
                job_queue.push(last_job);
                debug!("appended equivalent job");
            } else {
                // restore the popped last_job
                job_queue.push(last_job);
                // and push this new job to job queue
                job_queue.push(job);
                debug!("pushed onto job queue job");
            }
        }
    }

    #[allow(unrooted_must_root)]
    // https://w3c.github.io/ServiceWorker/#run-job-algorithm
    pub fn run_job(&self, scope_url: ServoUrl, script_thread: &ScriptThread) {
        debug!("running a job");
        let url = {
            let queue_ref = self.0.borrow();
            let front_job = {
                let job_vec = queue_ref.get(&scope_url);
                job_vec.unwrap().first().unwrap()
            };
            let front_scope_url = front_job.scope_url.clone();
            match front_job.job_type {
                JobType::Register => self.run_register(front_job, scope_url, script_thread),
                JobType::Update => self.update(front_job, script_thread),
                JobType::Unregister => unreachable!(),
            };
            front_scope_url
        };
        self.finish_job(url, script_thread);
    }

    #[allow(unrooted_must_root)]
    // https://w3c.github.io/ServiceWorker/#register-algorithm
    fn run_register(&self, job: &Job, scope_url: ServoUrl, script_thread: &ScriptThread) {
        debug!("running register job");
        // Step 1-3
        if !UrlHelper::is_origin_trustworthy(&job.script_url) {
            // Step 1.1
            reject_job_promise(job,
                               Error::Type("Invalid script ServoURL".to_owned()),
                               script_thread.dom_manipulation_task_source());
            // Step 1.2 (see run_job)
            return;
        } else if job.script_url.origin() != job.referrer.origin() || job.scope_url.origin() != job.referrer.origin() {
            // Step 2.1/3.1
            reject_job_promise(job,
                               Error::Security,
                               script_thread.dom_manipulation_task_source());
            // Step 2.2/3.2 (see run_job)
            return;
        }

        // Step 4-5
        if let Some(reg) = script_thread.handle_get_registration(&job.scope_url) {
            // Step 5.1
            if reg.get_uninstalling() {
                reg.set_uninstalling(false);
            }
            // Step 5.3
            if let Some(ref newest_worker) = reg.get_newest_worker() {
                if (&*newest_worker).get_script_url() == job.script_url {
                    // Step 5.3.1
                    resolve_job_promise(job, &*reg, script_thread.dom_manipulation_task_source());
                    // Step 5.3.2 (see run_job)
                    return;
                }
            }
        } else {
            // Step 6.1
            let global = &*job.client.global();
            let pipeline = global.pipeline_id();
            let new_reg = ServiceWorkerRegistration::new(&*global, &job.script_url, scope_url);
            script_thread.handle_serviceworker_registration(&job.scope_url, &*new_reg, pipeline);
        }
        // Step 7
        self.update(job, script_thread)
    }

    #[allow(unrooted_must_root)]
    // https://w3c.github.io/ServiceWorker/#finish-job-algorithm
    pub fn finish_job(&self, scope_url: ServoUrl, script_thread: &ScriptThread) {
        debug!("finishing previous job");
        let run_job = if let Some(job_vec) = (*self.0.borrow_mut()).get_mut(&scope_url) {
            assert_eq!(job_vec.first().as_ref().unwrap().scope_url, scope_url);
            let _  = job_vec.remove(0);
            !job_vec.is_empty()
        } else {
            warn!("non-existent job vector for Servourl: {:?}", scope_url);
            false
        };

        if run_job {
            debug!("further jobs in queue after finishing");
            self.run_job(scope_url, script_thread);
        }
    }

    // https://w3c.github.io/ServiceWorker/#update-algorithm
    fn update(&self, job: &Job, script_thread: &ScriptThread) {
        debug!("running update job");
        // Step 1
        let reg = match script_thread.handle_get_registration(&job.scope_url) {
            Some(reg) => reg,
            None => {
                let err_type = Error::Type("No registration to update".to_owned());
                // Step 2.1
                reject_job_promise(job, err_type, script_thread.dom_manipulation_task_source());
                // Step 2.2 (see run_job)
                return;
            }
        };
        // Step 2
        if reg.get_uninstalling() {
            let err_type = Error::Type("Update called on an uninstalling registration".to_owned());
            // Step 2.1
            reject_job_promise(job, err_type, script_thread.dom_manipulation_task_source());
            // Step 2.2 (see run_job)
            return;
        }
        // Step 3
        let newest_worker = reg.get_newest_worker();
        let newest_worker_url = newest_worker.as_ref().map(|w| w.get_script_url());
        // Step 4
        if newest_worker_url.as_ref() == Some(&job.script_url)  && job.job_type == JobType::Update {
            let err_type = Error::Type("Invalid script ServoURL".to_owned());
            // Step 4.1
            reject_job_promise(job, err_type, script_thread.dom_manipulation_task_source());
            // Step 4.2 (see run_job)
            return;
        }
        // Step 8
        if let Some(newest_worker) = newest_worker {
            job.client.set_controller(&*newest_worker);
            // Step 8.1
            resolve_job_promise(job, &*reg, script_thread.dom_manipulation_task_source());
            // Step 8.2 present in run_job
        }
        // TODO Step 9 (create new service worker)
    }
}

struct AsyncPromiseSettle {
    global: Trusted<GlobalScope>,
    promise: TrustedPromise,
    settle_type: SettleType,
}

impl Task for AsyncPromiseSettle {
    #[allow(unrooted_must_root)]
    fn run(self: Box<Self>) {
        let global = self.global.root();
        let settle_type = self.settle_type.clone();
        let promise = self.promise.root();
        settle_job_promise(&*global, &*promise, settle_type)
    }
}

impl AsyncPromiseSettle {
    #[allow(unrooted_must_root)]
    fn new(promise: Rc<Promise>, settle_type: SettleType) -> AsyncPromiseSettle {
        AsyncPromiseSettle {
            global: Trusted::new(&*promise.global()),
            promise: TrustedPromise::new(promise),
            settle_type: settle_type,
        }
    }
}

fn settle_job_promise(global: &GlobalScope, promise: &Promise, settle: SettleType) {
    let _ac = JSAutoCompartment::new(global.get_cx(), promise.reflector().get_jsobject().get());
    match settle {
        SettleType::Resolve(reg) => promise.resolve_native(global.get_cx(), &*reg.root()),
        SettleType::Reject(err) => promise.reject_error(global.get_cx(), err),
    };
}

fn queue_settle_promise_for_job(job: &Job, settle: SettleType, task_source: &DOMManipulationTaskSource) {
    let task = box AsyncPromiseSettle::new(job.promise.clone(), settle);
    let global = job.client.global();
    let _ = task_source.queue(task, &*global);

}

// https://w3c.github.io/ServiceWorker/#reject-job-promise-algorithm
// https://w3c.github.io/ServiceWorker/#resolve-job-promise-algorithm
fn queue_settle_promise(job: &Job, settle: SettleType, task_source: &DOMManipulationTaskSource) {
    // Step 1
    queue_settle_promise_for_job(job, settle.clone(), task_source);
    // Step 2
    for job in &job.equivalent_jobs {
        queue_settle_promise_for_job(job, settle.clone(), task_source);
    }
}

fn reject_job_promise(job: &Job, err: Error, task_source: &DOMManipulationTaskSource) {
    queue_settle_promise(job, SettleType::Reject(err), task_source)
}

fn resolve_job_promise(job: &Job, reg: &ServiceWorkerRegistration, task_source: &DOMManipulationTaskSource) {
    queue_settle_promise(job, SettleType::Resolve(Trusted::new(reg)), task_source)
}
<commit_msg>Use task! to settle promises for jobs<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! A Job is an abstraction of async operation in service worker lifecycle propagation.
//! Each Job is uniquely identified by its scope_url, and is keyed accordingly under
//! the script thread. The script thread contains a JobQueue, which stores all scheduled Jobs
//! by multiple service worker clients in a Vec.

use dom::bindings::cell::DOMRefCell;
use dom::bindings::error::Error;
use dom::bindings::js::JS;
use dom::bindings::refcounted::{Trusted, TrustedPromise};
use dom::bindings::reflector::DomObject;
use dom::client::Client;
use dom::globalscope::GlobalScope;
use dom::promise::Promise;
use dom::serviceworkerregistration::ServiceWorkerRegistration;
use dom::urlhelper::UrlHelper;
use js::jsapi::JSAutoCompartment;
use script_thread::ScriptThread;
use servo_url::ServoUrl;
use std::cmp::PartialEq;
use std::collections::HashMap;
use std::rc::Rc;
use task_source::TaskSource;
use task_source::dom_manipulation::DOMManipulationTaskSource;

#[derive(Clone, Copy, Debug, JSTraceable, PartialEq)]
pub enum JobType {
    Register,
    Unregister,
    Update
}

#[derive(Clone)]
pub enum SettleType {
    Resolve(Trusted<ServiceWorkerRegistration>),
    Reject(Error)
}

#[must_root]
#[derive(JSTraceable)]
pub struct Job {
    pub job_type: JobType,
    pub scope_url: ServoUrl,
    pub script_url: ServoUrl,
    pub promise: Rc<Promise>,
    pub equivalent_jobs: Vec<Job>,
    // client can be a window client, worker client so `Client` will be an enum in future
    pub client: JS<Client>,
    pub referrer: ServoUrl
}

impl Job {
    #[allow(unrooted_must_root)]
    // https://w3c.github.io/ServiceWorker/#create-job-algorithm
    pub fn create_job(job_type: JobType,
                      scope_url: ServoUrl,
                      script_url: ServoUrl,
                      promise: Rc<Promise>,
                      client: &Client) -> Job {
        Job {
            job_type: job_type,
            scope_url: scope_url,
            script_url: script_url,
            promise: promise,
            equivalent_jobs: vec![],
            client: JS::from_ref(client),
            referrer: client.creation_url()
        }
    }
    #[allow(unrooted_must_root)]
    pub fn append_equivalent_job(&mut self, job: Job) {
        self.equivalent_jobs.push(job);
    }
}

impl PartialEq for Job {
    // Equality criteria as described in https://w3c.github.io/ServiceWorker/#dfn-job-equivalent
    fn eq(&self, other: &Self) -> bool {
        let same_job = self.job_type == other.job_type;
        if same_job {
            match self.job_type {
                JobType::Register | JobType::Update => {
                    self.scope_url == other.scope_url && self.script_url == other.script_url
                },
                JobType::Unregister => self.scope_url == other.scope_url
            }
        } else {
            false
        }
    }
}

#[must_root]
#[derive(JSTraceable)]
pub struct JobQueue(pub DOMRefCell<HashMap<ServoUrl, Vec<Job>>>);

impl JobQueue {
    pub fn new() -> JobQueue {
        JobQueue(DOMRefCell::new(HashMap::new()))
    }
    #[allow(unrooted_must_root)]
    // https://w3c.github.io/ServiceWorker/#schedule-job-algorithm
    pub fn schedule_job(&self, job: Job, script_thread: &ScriptThread) {
        debug!("scheduling {:?} job", job.job_type);
        let mut queue_ref = self.0.borrow_mut();
        let job_queue = queue_ref.entry(job.scope_url.clone()).or_insert(vec![]);
        // Step 1
        if job_queue.is_empty() {
            let scope_url = job.scope_url.clone();
            job_queue.push(job);
            let _ = script_thread.schedule_job_queue(scope_url);
            debug!("queued task to run newly-queued job");
        } else {
            // Step 2
            let mut last_job = job_queue.pop().unwrap();
            if job == last_job && !last_job.promise.is_settled() {
                last_job.append_equivalent_job(job);
                job_queue.push(last_job);
                debug!("appended equivalent job");
            } else {
                // restore the popped last_job
                job_queue.push(last_job);
                // and push this new job to job queue
                job_queue.push(job);
                debug!("pushed onto job queue job");
            }
        }
    }

    #[allow(unrooted_must_root)]
    // https://w3c.github.io/ServiceWorker/#run-job-algorithm
    pub fn run_job(&self, scope_url: ServoUrl, script_thread: &ScriptThread) {
        debug!("running a job");
        let url = {
            let queue_ref = self.0.borrow();
            let front_job = {
                let job_vec = queue_ref.get(&scope_url);
                job_vec.unwrap().first().unwrap()
            };
            let front_scope_url = front_job.scope_url.clone();
            match front_job.job_type {
                JobType::Register => self.run_register(front_job, scope_url, script_thread),
                JobType::Update => self.update(front_job, script_thread),
                JobType::Unregister => unreachable!(),
            };
            front_scope_url
        };
        self.finish_job(url, script_thread);
    }

    #[allow(unrooted_must_root)]
    // https://w3c.github.io/ServiceWorker/#register-algorithm
    fn run_register(&self, job: &Job, scope_url: ServoUrl, script_thread: &ScriptThread) {
        debug!("running register job");
        // Step 1-3
        if !UrlHelper::is_origin_trustworthy(&job.script_url) {
            // Step 1.1
            reject_job_promise(job,
                               Error::Type("Invalid script ServoURL".to_owned()),
                               script_thread.dom_manipulation_task_source());
            // Step 1.2 (see run_job)
            return;
        } else if job.script_url.origin() != job.referrer.origin() || job.scope_url.origin() != job.referrer.origin() {
            // Step 2.1/3.1
            reject_job_promise(job,
                               Error::Security,
                               script_thread.dom_manipulation_task_source());
            // Step 2.2/3.2 (see run_job)
            return;
        }

        // Step 4-5
        if let Some(reg) = script_thread.handle_get_registration(&job.scope_url) {
            // Step 5.1
            if reg.get_uninstalling() {
                reg.set_uninstalling(false);
            }
            // Step 5.3
            if let Some(ref newest_worker) = reg.get_newest_worker() {
                if (&*newest_worker).get_script_url() == job.script_url {
                    // Step 5.3.1
                    resolve_job_promise(job, &*reg, script_thread.dom_manipulation_task_source());
                    // Step 5.3.2 (see run_job)
                    return;
                }
            }
        } else {
            // Step 6.1
            let global = &*job.client.global();
            let pipeline = global.pipeline_id();
            let new_reg = ServiceWorkerRegistration::new(&*global, &job.script_url, scope_url);
            script_thread.handle_serviceworker_registration(&job.scope_url, &*new_reg, pipeline);
        }
        // Step 7
        self.update(job, script_thread)
    }

    #[allow(unrooted_must_root)]
    // https://w3c.github.io/ServiceWorker/#finish-job-algorithm
    pub fn finish_job(&self, scope_url: ServoUrl, script_thread: &ScriptThread) {
        debug!("finishing previous job");
        let run_job = if let Some(job_vec) = (*self.0.borrow_mut()).get_mut(&scope_url) {
            assert_eq!(job_vec.first().as_ref().unwrap().scope_url, scope_url);
            let _  = job_vec.remove(0);
            !job_vec.is_empty()
        } else {
            warn!("non-existent job vector for Servourl: {:?}", scope_url);
            false
        };

        if run_job {
            debug!("further jobs in queue after finishing");
            self.run_job(scope_url, script_thread);
        }
    }

    // https://w3c.github.io/ServiceWorker/#update-algorithm
    fn update(&self, job: &Job, script_thread: &ScriptThread) {
        debug!("running update job");
        // Step 1
        let reg = match script_thread.handle_get_registration(&job.scope_url) {
            Some(reg) => reg,
            None => {
                let err_type = Error::Type("No registration to update".to_owned());
                // Step 2.1
                reject_job_promise(job, err_type, script_thread.dom_manipulation_task_source());
                // Step 2.2 (see run_job)
                return;
            }
        };
        // Step 2
        if reg.get_uninstalling() {
            let err_type = Error::Type("Update called on an uninstalling registration".to_owned());
            // Step 2.1
            reject_job_promise(job, err_type, script_thread.dom_manipulation_task_source());
            // Step 2.2 (see run_job)
            return;
        }
        // Step 3
        let newest_worker = reg.get_newest_worker();
        let newest_worker_url = newest_worker.as_ref().map(|w| w.get_script_url());
        // Step 4
        if newest_worker_url.as_ref() == Some(&job.script_url)  && job.job_type == JobType::Update {
            let err_type = Error::Type("Invalid script ServoURL".to_owned());
            // Step 4.1
            reject_job_promise(job, err_type, script_thread.dom_manipulation_task_source());
            // Step 4.2 (see run_job)
            return;
        }
        // Step 8
        if let Some(newest_worker) = newest_worker {
            job.client.set_controller(&*newest_worker);
            // Step 8.1
            resolve_job_promise(job, &*reg, script_thread.dom_manipulation_task_source());
            // Step 8.2 present in run_job
        }
        // TODO Step 9 (create new service worker)
    }
}

fn settle_job_promise(global: &GlobalScope, promise: &Promise, settle: SettleType) {
    let _ac = JSAutoCompartment::new(global.get_cx(), promise.reflector().get_jsobject().get());
    match settle {
        SettleType::Resolve(reg) => promise.resolve_native(global.get_cx(), &*reg.root()),
        SettleType::Reject(err) => promise.reject_error(global.get_cx(), err),
    };
}

#[allow(unrooted_must_root)]
fn queue_settle_promise_for_job(job: &Job, settle: SettleType, task_source: &DOMManipulationTaskSource) {
    let global = job.client.global();
    let promise = TrustedPromise::new(job.promise.clone());
    // FIXME(nox): Why are errors silenced here?
    let _ = task_source.queue(
        box task!(settle_promise_for_job: move || {
            let promise = promise.root();
            settle_job_promise(&promise.global(), &promise, settle)
        }),
        &*global,
    );
}

// https://w3c.github.io/ServiceWorker/#reject-job-promise-algorithm
// https://w3c.github.io/ServiceWorker/#resolve-job-promise-algorithm
fn queue_settle_promise(job: &Job, settle: SettleType, task_source: &DOMManipulationTaskSource) {
    // Step 1
    queue_settle_promise_for_job(job, settle.clone(), task_source);
    // Step 2
    for job in &job.equivalent_jobs {
        queue_settle_promise_for_job(job, settle.clone(), task_source);
    }
}

fn reject_job_promise(job: &Job, err: Error, task_source: &DOMManipulationTaskSource) {
    queue_settle_promise(job, SettleType::Reject(err), task_source)
}

fn resolve_job_promise(job: &Job, reg: &ServiceWorkerRegistration, task_source: &DOMManipulationTaskSource) {
    queue_settle_promise(job, SettleType::Resolve(Trusted::new(reg)), task_source)
}
<|endoftext|>


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

//! CtlGateway, short for Control Gateway, is a TCP based client and server connection for
//! sending and receiving command and control requests to a running Supervisor. Operational tasks
//! such as starting, stopping, loading, and unloading services are exposed through the
//! CtlGateway.
//!
//! The [`ctl_gateway.client`] and [`ctl_gateway.server`] speak a streaming, multiplexed, binary
//! protocol defined in [`protocol.codec`].

pub mod server;

use std::{borrow::Cow,
          fmt,
          fs::{self,
               File},
          io::{self,
               Write},
          path::Path};

use regex::Regex;
use termcolor::{Buffer,
                Color,
                ColorSpec,
                WriteColor};

use crate::{api_client::DisplayProgress,
            common::ui::{ColorPrinter,
                         UIWriter},
            hcore::{self,
                    output},
            protocol};
use futures::prelude::*;

use crate::error::{Error,
                   Result};

lazy_static! {
    /// Shamelessly stolen from https://github.com/chalk/ansi-regex/blob/master/index.js
    static ref STRIP_ANSI_CODES: Regex = Regex::new(
        r"[\x1b\x9b][\[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?[0-9A-PRZcf-nqry=><]")
        .unwrap();
}

/// Time to wait in milliseconds for a client connection to timeout.
pub const REQ_TIMEOUT: u64 = 10_000;
static LOGKEY: &'static str = "AG";

/// The control gateway secret should only be readable by the
/// Supervisor process
#[cfg(not(windows))]
pub const CTL_SECRET_PERMISSIONS: u32 = 0o600;

/// Used by modules outside of the CtlGateway for seamlessly replying to transactional messages.
/// This type is used in functions which can be called by the CtlGateway such as
/// [`Manager::service_load`] and [`Manager::service_unload`].
///
/// A _bare request_ can be instantiated with `CtlRequest::default()` and used as a no-op stub when
/// calling functions which typically respond to the client. This is useful when the caller isn't
/// a client but needs to perform the same operation. For example, loading an initial service for
/// the Supervisor has no need to reply back to a networked client, so passing a bare request will
/// no-op any calls to send messages back to the client.
#[derive(Clone, Default)]
pub struct CtlRequest {
    /// The sending side of the CtlGateway's server. Replies are sent across this channel and then
    /// eventually over the network back to the client.
    tx: Option<server::CtlSender>,
    /// Transaction for the given request.
    transaction: Option<protocol::codec::SrvTxn>,
}

impl CtlRequest {
    /// Create a new CtlRequest from an optional [`server.CtlSender`] and
    /// [`protocol.codec.SrvTxn`].
    pub fn new(
        tx: Option<server::CtlSender>,
        transaction: Option<protocol::codec::SrvTxn>,
    ) -> Self {
        CtlRequest { tx, transaction }
    }

    /// Reply to the transaction with the given message but indicate to the receiver that this is
    /// not the final message for the transaction.
    pub fn reply_partial<T>(&mut self, msg: T)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        self.send_msg(msg, false);
    }

    /// Reply to the transaction with the given message and indicate to the receiver that this is
    /// the final message for the transaction.
    pub fn reply_complete<T>(&mut self, msg: T)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        self.send_msg(msg, true);
    }

    /// Returns true if the request is transactional and false if not.
    pub fn transactional(&self) -> bool { self.transaction.is_some() && self.tx.is_some() }

    fn send_msg<T>(&mut self, msg: T, complete: bool)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        if !self.transactional() {
            warn!(
                "Attempted to reply to a non-transactional message with {:?}",
                msg
            );
            return;
        }
        let mut wire: protocol::codec::SrvMessage = msg.into();
        wire.reply_for(self.transaction.unwrap(), complete);
        self.tx.as_ref().unwrap().start_send(wire).ok(); // ignore Err return
    }
}

impl UIWriter for CtlRequest {
    type ProgressBar = NetProgressBar;

    fn out(&mut self) -> &mut dyn ColorPrinter { self }

    fn err(&mut self) -> &mut dyn ColorPrinter { self }

    fn is_out_a_terminal(&self) -> bool { true }

    fn is_err_a_terminal(&self) -> bool { true }

    fn progress(&self) -> Option<Self::ProgressBar> {
        if self.is_out_a_terminal() {
            Some(Self::ProgressBar::new(self.clone()))
        } else {
            None
        }
    }
}

impl ColorPrinter for CtlRequest {
    fn print(&mut self, buf: &[u8], color: Option<Color>, bold: bool) -> io::Result<()> {
        // The protocol reply is destined for the client, so (for now,
        // at least), we'll apply colored output.
        //
        // `line` will also have a newline character at the end, FYI.
        let mut msg = protocol::ctl::ConsoleLine::default();
        msg.line = String::from_utf8_lossy(buf).to_string();
        msg.color = color_to_string(color);
        msg.bold = bold;
        self.reply_partial(msg);

        // Down here, however, we're doing double-duty by *also*
        // sending the output to the Supervisor's output stream. In
        // _this_ case, we want to honor whatever global logging flags
        // have been set; in particular, if we're emitting
        // JSON-formatted logs or if we've been instructed to not
        // output ANSI color codes, we need to strip those codes out.
        //
        // TODO (CM): Obviously, this feels hacky to the extreme. It'd
        // be nice to find a cleaner way to model this, but the fact
        // that CtlRequest is sending output to two destinations with
        // different formatting requirements complicates things a bit.
        //
        // TODO (MW): For `han sup run` scenarios, it would be nice to
        // support color on older versions of Windows. We could refactor
        // outputln! to support a cross platform ColorWriter similar to
        // what we use in common::UI.
        let mut ansi_buffer = Buffer::ansi();
        ansi_buffer.set_color(ColorSpec::new().set_bold(bold).set_fg(color))?;
        ansi_buffer.write_all(buf)?;
        ansi_buffer.flush()?;
        let line = String::from_utf8_lossy(ansi_buffer.as_slice()).to_string();

        let maybe_stripped = if output::is_json() || !output::is_color() {
            STRIP_ANSI_CODES.replace_all(&line, "")
        } else {
            Cow::Owned(line)
        };
        // TODO (CM): Consider pulling this newline trimming up into
        // the macro
        outputln!("{}", maybe_stripped.trim_right_matches('\n'));

        Ok(())
    }
}

fn color_to_string(color: Option<Color>) -> Option<String> {
    match color {
        Some(c) => Some(format!("{:?}", c).to_string()),
        None => None,
    }
}

/// A wrapper around a [`protocol.ctl.NetProgress`] and [`CtlRequest`]. This type implements
/// traits for writing it's progress to the console.
pub struct NetProgressBar {
    inner: protocol::ctl::NetProgress,
    req: CtlRequest,
}

impl NetProgressBar {
    /// Create a new progress bar.
    pub fn new(req: CtlRequest) -> Self {
        NetProgressBar {
            inner: protocol::ctl::NetProgress::default(),
            req,
        }
    }
}

impl DisplayProgress for NetProgressBar {
    fn size(&mut self, size: u64) { self.inner.total = size; }

    fn finish(&mut self) {}
}

impl io::Write for NetProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.inner.position += buf.len() as u64;
        self.req.reply_partial(self.inner.clone());
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> { Ok(()) }
}

/// First attempts to read the secret key used to authenticate with the `CtlGateway` from disk
/// and, if not found, will generate a new key and write it to disk.
pub fn readgen_secret_key<T>(sup_root: T) -> Result<String>
where
    T: AsRef<Path>,
{
    let mut out = String::new();
    fs::create_dir_all(&sup_root)
        .map_err(|e| sup_error!(Error::CtlSecretIo(sup_root.as_ref().to_path_buf(), e)))?;
    if protocol::read_secret_key(&sup_root, &mut out)
        .ok()
        .unwrap_or(false)
    {
        Ok(out)
    } else {
        let secret_key_path = protocol::secret_key_path(sup_root);
        {
            let mut f = File::create(&secret_key_path)?;
            protocol::generate_secret_key(&mut out);
            f.write_all(out.as_bytes())?;
            f.sync_all()?;
        }
        set_permissions(&secret_key_path)?;
        Ok(out)
    }
}

#[cfg(not(windows))]
fn set_permissions<T: AsRef<Path>>(path: T) -> hcore::error::Result<()> {
    use crate::hcore::util::posix_perm;

    posix_perm::set_permissions(path.as_ref(), CTL_SECRET_PERMISSIONS)
}

#[cfg(windows)]
fn set_permissions<T: AsRef<Path>>(path: T) -> hcore::error::Result<()> {
    use hcore::util::win_perm;

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

//! CtlGateway, short for Control Gateway, is a TCP based client and server connection for
//! sending and receiving command and control requests to a running Supervisor. Operational tasks
//! such as starting, stopping, loading, and unloading services are exposed through the
//! CtlGateway.
//!
//! The [`ctl_gateway.client`] and [`ctl_gateway.server`] speak a streaming, multiplexed, binary
//! protocol defined in [`protocol.codec`].

pub mod server;

use std::{borrow::Cow,
          fmt,
          fs::{self,
               File},
          io::{self,
               Write},
          path::Path};

use regex::Regex;
use termcolor::{Buffer,
                Color,
                ColorSpec,
                WriteColor};

use crate::{api_client::DisplayProgress,
            common::ui::{ColorPrinter,
                         UIWriter},
            hcore::{self,
                    output},
            protocol};
use futures::prelude::*;

use crate::error::{Error,
                   Result};

lazy_static! {
    /// Shamelessly stolen from https://github.com/chalk/ansi-regex/blob/master/index.js
    static ref STRIP_ANSI_CODES: Regex = Regex::new(
        r"[\x1b\x9b][\[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?[0-9A-PRZcf-nqry=><]")
        .unwrap();
}

/// Time to wait in milliseconds for a client connection to timeout.
pub const REQ_TIMEOUT: u64 = 10_000;
static LOGKEY: &'static str = "AG";

/// The control gateway secret should only be readable by the
/// Supervisor process
#[cfg(not(windows))]
pub const CTL_SECRET_PERMISSIONS: u32 = 0o600;

/// Used by modules outside of the CtlGateway for seamlessly replying to transactional messages.
/// This type is used in functions which can be called by the CtlGateway such as
/// [`Manager::service_load`] and [`Manager::service_unload`].
///
/// A _bare request_ can be instantiated with `CtlRequest::default()` and used as a no-op stub when
/// calling functions which typically respond to the client. This is useful when the caller isn't
/// a client but needs to perform the same operation. For example, loading an initial service for
/// the Supervisor has no need to reply back to a networked client, so passing a bare request will
/// no-op any calls to send messages back to the client.
#[derive(Clone, Default)]
pub struct CtlRequest {
    /// The sending side of the CtlGateway's server. Replies are sent across this channel and then
    /// eventually over the network back to the client.
    tx: Option<server::CtlSender>,
    /// Transaction for the given request.
    transaction: Option<protocol::codec::SrvTxn>,
}

impl CtlRequest {
    /// Create a new CtlRequest from an optional [`server.CtlSender`] and
    /// [`protocol.codec.SrvTxn`].
    pub fn new(
        tx: Option<server::CtlSender>,
        transaction: Option<protocol::codec::SrvTxn>,
    ) -> Self {
        CtlRequest { tx, transaction }
    }

    /// Reply to the transaction with the given message but indicate to the receiver that this is
    /// not the final message for the transaction.
    pub fn reply_partial<T>(&mut self, msg: T)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        self.send_msg(msg, false);
    }

    /// Reply to the transaction with the given message and indicate to the receiver that this is
    /// the final message for the transaction.
    pub fn reply_complete<T>(&mut self, msg: T)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        self.send_msg(msg, true);
    }

    /// Returns true if the request is transactional and false if not.
    pub fn transactional(&self) -> bool { self.transaction.is_some() && self.tx.is_some() }

    fn send_msg<T>(&mut self, msg: T, complete: bool)
    where
        T: Into<protocol::codec::SrvMessage> + fmt::Debug,
    {
        if !self.transactional() {
            warn!(
                "Attempted to reply to a non-transactional message with {:?}",
                msg
            );
            return;
        }
        let mut wire: protocol::codec::SrvMessage = msg.into();
        wire.reply_for(self.transaction.unwrap(), complete);
        self.tx.as_ref().unwrap().start_send(wire).ok(); // ignore Err return
    }
}

impl UIWriter for CtlRequest {
    type ProgressBar = NetProgressBar;

    fn out(&mut self) -> &mut dyn ColorPrinter { self }

    fn err(&mut self) -> &mut dyn ColorPrinter { self }

    fn is_out_a_terminal(&self) -> bool { true }

    fn is_err_a_terminal(&self) -> bool { true }

    fn progress(&self) -> Option<Self::ProgressBar> {
        if self.is_out_a_terminal() {
            Some(Self::ProgressBar::new(self.clone()))
        } else {
            None
        }
    }
}

impl ColorPrinter for CtlRequest {
    fn print(&mut self, buf: &[u8], color: Option<Color>, bold: bool) -> io::Result<()> {
        // The protocol reply is destined for the client, so (for now,
        // at least), we'll apply colored output.
        //
        // `line` will also have a newline character at the end, FYI.
        let mut msg = protocol::ctl::ConsoleLine::default();
        msg.line = String::from_utf8_lossy(buf).to_string();
        msg.color = color_to_string(color);
        msg.bold = bold;
        self.reply_partial(msg);

        // Down here, however, we're doing double-duty by *also*
        // sending the output to the Supervisor's output stream. In
        // _this_ case, we want to honor whatever global logging flags
        // have been set; in particular, if we're emitting
        // JSON-formatted logs or if we've been instructed to not
        // output ANSI color codes, we need to strip those codes out.
        //
        // TODO (CM): Obviously, this feels hacky to the extreme. It'd
        // be nice to find a cleaner way to model this, but the fact
        // that CtlRequest is sending output to two destinations with
        // different formatting requirements complicates things a bit.
        //
        // TODO (MW): For `hab sup run` scenarios, it would be nice to
        // support color on older versions of Windows. We could refactor
        // outputln! to support a cross platform ColorWriter similar to
        // what we use in common::UI.
        let mut ansi_buffer = Buffer::ansi();
        ansi_buffer.set_color(ColorSpec::new().set_bold(bold).set_fg(color))?;
        ansi_buffer.write_all(buf)?;
        ansi_buffer.flush()?;
        let line = String::from_utf8_lossy(ansi_buffer.as_slice()).to_string();

        let maybe_stripped = if output::is_json() || !output::is_color() {
            STRIP_ANSI_CODES.replace_all(&line, "")
        } else {
            Cow::Owned(line)
        };
        // TODO (CM): Consider pulling this newline trimming up into
        // the macro
        outputln!("{}", maybe_stripped.trim_right_matches('\n'));

        Ok(())
    }
}

fn color_to_string(color: Option<Color>) -> Option<String> {
    match color {
        Some(c) => Some(format!("{:?}", c).to_string()),
        None => None,
    }
}

/// A wrapper around a [`protocol.ctl.NetProgress`] and [`CtlRequest`]. This type implements
/// traits for writing it's progress to the console.
pub struct NetProgressBar {
    inner: protocol::ctl::NetProgress,
    req: CtlRequest,
}

impl NetProgressBar {
    /// Create a new progress bar.
    pub fn new(req: CtlRequest) -> Self {
        NetProgressBar {
            inner: protocol::ctl::NetProgress::default(),
            req,
        }
    }
}

impl DisplayProgress for NetProgressBar {
    fn size(&mut self, size: u64) { self.inner.total = size; }

    fn finish(&mut self) {}
}

impl io::Write for NetProgressBar {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.inner.position += buf.len() as u64;
        self.req.reply_partial(self.inner.clone());
        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> { Ok(()) }
}

/// First attempts to read the secret key used to authenticate with the `CtlGateway` from disk
/// and, if not found, will generate a new key and write it to disk.
pub fn readgen_secret_key<T>(sup_root: T) -> Result<String>
where
    T: AsRef<Path>,
{
    let mut out = String::new();
    fs::create_dir_all(&sup_root)
        .map_err(|e| sup_error!(Error::CtlSecretIo(sup_root.as_ref().to_path_buf(), e)))?;
    if protocol::read_secret_key(&sup_root, &mut out)
        .ok()
        .unwrap_or(false)
    {
        Ok(out)
    } else {
        let secret_key_path = protocol::secret_key_path(sup_root);
        {
            let mut f = File::create(&secret_key_path)?;
            protocol::generate_secret_key(&mut out);
            f.write_all(out.as_bytes())?;
            f.sync_all()?;
        }
        set_permissions(&secret_key_path)?;
        Ok(out)
    }
}

#[cfg(not(windows))]
fn set_permissions<T: AsRef<Path>>(path: T) -> hcore::error::Result<()> {
    use crate::hcore::util::posix_perm;

    posix_perm::set_permissions(path.as_ref(), CTL_SECRET_PERMISSIONS)
}

#[cfg(windows)]
fn set_permissions<T: AsRef<Path>>(path: T) -> hcore::error::Result<()> {
    use hcore::util::win_perm;

    win_perm::harden_path(path.as_ref())
}
<|endoftext|>


<commit_before>use std::cmp;

pub trait ZipLongestIteratorExt: Iterator + Sized {
    /// Creates an iterator which iterates over both this and the specified
    /// iterators simultaneously, yielding pairs of two optional elements.
    /// When both iterators return None, all further invocations of next() will
    /// return None.
    ///
    /// # Example
    ///
    /// ```rust
    /// let a = [0i];
    /// let b = [1i, 2i];
    /// let mut it = a.iter().zip(b.iter());
    /// let (x0, x1, x2) = (0i, 1i, 2i);
    /// assert_eq!(it.next().unwrap(), (Some(&x0), Some(&x1)));
    /// assert_eq!(it.next().unwrap(), (None, Some(&x2)));
    /// assert!(it.next().is_none());
    /// ```
    #[inline]
    fn zip_longest<U: Iterator>(self, other: U) -> ZipLongest<Self, U> {
        ZipLongest{a: self, b: other}
    }
}

/// An iterator which iterates two other iterators simultaneously
#[derive(Clone)]
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct ZipLongest<T, U> {
    a: T,
    b: U
}

impl<A, B, T: Iterator<Item = A>, U: Iterator<Item = B>> Iterator for ZipLongest<T, U> {
    type Item = EitherOrBoth<A, B>;

    #[inline]
    fn next(&mut self) -> Option<EitherOrBoth<A, B>> {
        match (self.a.next(), self.b.next()) {
            (None, None) => None,
            (Some(a), None) => Some(EitherOrBoth::Left(a)),
            (None, Some(b)) => Some(EitherOrBoth::Right(b)),
            (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let (a_lower, a_upper) = self.a.size_hint();
        let (b_lower, b_upper) = self.b.size_hint();

        let lower = cmp::max(a_lower, b_lower);

        let upper = match (a_upper, b_upper) {
            (Some(x), Some(y)) => Some(cmp::max(x,y)),
            _ => None
        };

        (lower, upper)
    }
}

impl<T, U> DoubleEndedIterator for ZipLongest<T, U>
where T: DoubleEndedIterator + ExactSizeIterator, U: DoubleEndedIterator + ExactSizeIterator {
    #[inline]
    fn next_back(&mut self) -> Option<<Self as Iterator>::Item> {
        use std::cmp::Ordering::{Equal, Greater, Less};
        match self.a.len().cmp(&self.b.len()) {
            Equal => match (self.a.next_back(), self.b.next_back()) {
                (None, None) => None,
                (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
                // XXX these can only happen if .len() is inconsistent with .next_back()
                (Some(a), None) => Some(EitherOrBoth::Left(a)),
                (None, Some(b)) => Some(EitherOrBoth::Right(b)),
            },
            Greater => self.a.next_back().map(EitherOrBoth::Left),
            Less => self.b.next_back().map(EitherOrBoth::Right),
        }
    }
}

impl<T: ExactSizeIterator, U: ExactSizeIterator> ExactSizeIterator for ZipLongest<T, U> {}

impl<I> ZipLongestIteratorExt for I where I: Iterator {}

/// A value yielded by `ZipLongest`.
/// Contains one or two values,
/// depending on which of the input iterators are exhausted.
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum EitherOrBoth<A, B> {
    /// Neither input iterator is exhausted yet, yielding two values.
    Both(A, B),
    /// The parameter iterator of `.zip_longest()` is exhausted,
    /// only yielding a value from the `self` iterator.
    Left(A),
    /// The `self` iterator of `.zip_longest()` is exhausted,
    /// only yielding a value from the parameter iterator.
    Right(B),
}

#[test]
fn test_iterator_size_hint() {
    use std::usize;

    let c = 0i32..;
    let v: &[_] = &[0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9];
    let v2 = &[10i32, 11, 12];
    let vi = v.iter();
    assert_eq!(c.zip_longest(vi.clone()).size_hint(), (usize::MAX, None));
    assert_eq!(vi.zip_longest(v2.iter()).size_hint(), (10, Some(10)));
}

#[test]
fn test_double_ended() {
    let xs = [1i32, 2, 3, 4, 5, 6];
    let ys = [1i32, 2, 3, 7];
    let a = xs.iter().map(|&x| x);
    let b = ys.iter().map(|&x| x);
    let mut it = a.zip_longest(b);
    assert_eq!(it.next(), Some(EitherOrBoth::Both(1, 1)));
    assert_eq!(it.next(), Some(EitherOrBoth::Both(2, 2)));
    assert_eq!(it.next_back(), Some(EitherOrBoth::Left(6)));
    assert_eq!(it.next_back(), Some(EitherOrBoth::Left(5)));
    assert_eq!(it.next_back(), Some(EitherOrBoth::Both(4, 7)));
    assert_eq!(it.next(), Some(EitherOrBoth::Both(3, 3)));
    assert_eq!(it.next(), None);
}
<commit_msg>Make tests pass on Rust 1.18 and older<commit_after>use std::cmp;

pub trait ZipLongestIteratorExt: Iterator + Sized {
    /// Creates an iterator which iterates over both this and the specified
    /// iterators simultaneously, yielding pairs of two optional elements.
    /// When both iterators return None, all further invocations of next() will
    /// return None.
    ///
    /// # Example
    ///
    /// ```rust
    /// let a = [0i];
    /// let b = [1i, 2i];
    /// let mut it = a.iter().zip(b.iter());
    /// let (x0, x1, x2) = (0i, 1i, 2i);
    /// assert_eq!(it.next().unwrap(), (Some(&x0), Some(&x1)));
    /// assert_eq!(it.next().unwrap(), (None, Some(&x2)));
    /// assert!(it.next().is_none());
    /// ```
    #[inline]
    fn zip_longest<U: Iterator>(self, other: U) -> ZipLongest<Self, U> {
        ZipLongest{a: self, b: other}
    }
}

/// An iterator which iterates two other iterators simultaneously
#[derive(Clone)]
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
pub struct ZipLongest<T, U> {
    a: T,
    b: U
}

impl<A, B, T: Iterator<Item = A>, U: Iterator<Item = B>> Iterator for ZipLongest<T, U> {
    type Item = EitherOrBoth<A, B>;

    #[inline]
    fn next(&mut self) -> Option<EitherOrBoth<A, B>> {
        match (self.a.next(), self.b.next()) {
            (None, None) => None,
            (Some(a), None) => Some(EitherOrBoth::Left(a)),
            (None, Some(b)) => Some(EitherOrBoth::Right(b)),
            (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let (a_lower, a_upper) = self.a.size_hint();
        let (b_lower, b_upper) = self.b.size_hint();

        let lower = cmp::max(a_lower, b_lower);

        let upper = match (a_upper, b_upper) {
            (Some(x), Some(y)) => Some(cmp::max(x,y)),
            _ => None
        };

        (lower, upper)
    }
}

impl<T, U> DoubleEndedIterator for ZipLongest<T, U>
where T: DoubleEndedIterator + ExactSizeIterator, U: DoubleEndedIterator + ExactSizeIterator {
    #[inline]
    fn next_back(&mut self) -> Option<<Self as Iterator>::Item> {
        use std::cmp::Ordering::{Equal, Greater, Less};
        match self.a.len().cmp(&self.b.len()) {
            Equal => match (self.a.next_back(), self.b.next_back()) {
                (None, None) => None,
                (Some(a), Some(b)) => Some(EitherOrBoth::Both(a, b)),
                // XXX these can only happen if .len() is inconsistent with .next_back()
                (Some(a), None) => Some(EitherOrBoth::Left(a)),
                (None, Some(b)) => Some(EitherOrBoth::Right(b)),
            },
            Greater => self.a.next_back().map(EitherOrBoth::Left),
            Less => self.b.next_back().map(EitherOrBoth::Right),
        }
    }
}

impl<T: ExactSizeIterator, U: ExactSizeIterator> ExactSizeIterator for ZipLongest<T, U> {}

impl<I> ZipLongestIteratorExt for I where I: Iterator {}

/// A value yielded by `ZipLongest`.
/// Contains one or two values,
/// depending on which of the input iterators are exhausted.
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum EitherOrBoth<A, B> {
    /// Neither input iterator is exhausted yet, yielding two values.
    Both(A, B),
    /// The parameter iterator of `.zip_longest()` is exhausted,
    /// only yielding a value from the `self` iterator.
    Left(A),
    /// The `self` iterator of `.zip_longest()` is exhausted,
    /// only yielding a value from the parameter iterator.
    Right(B),
}

#[test]
fn test_iterator_size_hint() {
    use std::usize;

    let c = std::iter::repeat(0);
    let v: &[_] = &[0i32, 1, 2, 3, 4, 5, 6, 7, 8, 9];
    let v2 = &[10i32, 11, 12];
    let vi = v.iter();
    assert_eq!(c.zip_longest(vi.clone()).size_hint(), (usize::MAX, None));
    assert_eq!(vi.zip_longest(v2.iter()).size_hint(), (10, Some(10)));
}

#[test]
fn test_double_ended() {
    let xs = [1i32, 2, 3, 4, 5, 6];
    let ys = [1i32, 2, 3, 7];
    let a = xs.iter().map(|&x| x);
    let b = ys.iter().map(|&x| x);
    let mut it = a.zip_longest(b);
    assert_eq!(it.next(), Some(EitherOrBoth::Both(1, 1)));
    assert_eq!(it.next(), Some(EitherOrBoth::Both(2, 2)));
    assert_eq!(it.next_back(), Some(EitherOrBoth::Left(6)));
    assert_eq!(it.next_back(), Some(EitherOrBoth::Left(5)));
    assert_eq!(it.next_back(), Some(EitherOrBoth::Both(4, 7)));
    assert_eq!(it.next(), Some(EitherOrBoth::Both(3, 3)));
    assert_eq!(it.next(), None);
}
<|endoftext|>


<commit_before>use std::io::{Error, ErrorKind, Result};
use std::ffi::CString;
use libc;
use message::Message;
use std::cmp::Ordering;
use std::ptr;
use std::boxed::Box;
use std::rc::Rc;
use std::ops::Deref;

enum CLcm {}

/// An LCM instance that hnadles publishing and subscribing,
/// as well as encoding and decoding messages.
pub struct Lcm {
    lcm: *mut CLcm,
    subscriptions: Vec<Rc<LcmSubscription>>,
}

type CLcmHandler = extern "C" fn(*const CLcmRecvBuf, *const libc::c_char, *const libc::c_void);

#[derive(Eq, PartialEq, Hash)]
enum CLcmSubscription {}
pub struct LcmSubscription {
    subscription: *mut CLcmSubscription,
    handler: Box<FnMut(*const CLcmRecvBuf)>,
}

#[repr(C)]
#[derive(Debug)]
struct CLcmRecvBuf {
    data: *const libc::c_void,
    data_size: libc::uint32_t,
    recv_utime: libc::int64_t,
    lcm: *const CLcm,
}

#[link(name = "lcm")]
extern "C" {
    fn lcm_create(provider: *const libc::c_char) -> *mut CLcm;

    fn lcm_destroy(lcm: *mut CLcm);

    fn lcm_get_fileno(lcm: *mut CLcm) -> libc::c_int;

    fn lcm_subscribe(lcm: *mut CLcm,
                     channel: *const libc::c_char,
                     handler: CLcmHandler,
                     userdata: *const libc::c_void)
                     -> *mut CLcmSubscription;

    fn lcm_unsubscribe(lcm: *mut CLcm, handler: *const CLcmSubscription) -> libc::c_int;

    fn lcm_publish(lcm: *mut CLcm,
                   channel: *const libc::c_char,
                   data: *const libc::c_void,
                   datalen: libc::c_uint)
                   -> libc::c_int;

    fn lcm_handle(lcm: *mut CLcm) -> libc::c_int;

    fn lcm_handle_timeout(lcm: *mut CLcm, timeout_millis: libc::c_int) -> libc::c_int;

    fn lcm_subscription_set_queue_capacity(handler: *const CLcmSubscription,
                                           num_messages: libc::c_int)
                                           -> libc::c_int;
}

impl Lcm {
    /// Creates a new `Lcm` instance.
    ///
    /// ```
    /// use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// ```
    pub fn new() -> Result<Lcm> {
        let lcm = unsafe { lcm_create(ptr::null()) };
        match lcm.is_null() {
            true => Err(Error::new(ErrorKind::Other, "Failed to initialize LCM.")),
            false => {
                Ok(Lcm {
                    lcm: lcm,
                    subscriptions: Vec::new(),
                })
            }
        }
    }

    pub fn get_fileno(&self) -> libc::c_int {
        unsafe { lcm_get_fileno(self.lcm) }
    }

    /// Subscribes a callback to a particular topic.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// use std::sync::mpsc::channel;
    ///
    /// let mut lcm = Lcm::new().unwrap();
    /// let (tx, rx) = channel::<exlcm::Example>();
    /// lcm.subscribe("POSITION", |pos| { tx.send(pos).unwrap(); }
    /// ```
    pub fn subscribe<M, F>(&mut self, channel: &str, mut callback: Box<F>) -> Rc<LcmSubscription>
        where M: Message + Default,
              F: FnMut(M) + 'static
    {
        let channel = CString::new(channel).unwrap();

        let handler = Box::new(move |rbuf: *const CLcmRecvBuf| {
            let buf = unsafe {
                let ref rbuf = *rbuf;
                let data = rbuf.data as *mut u8;
                let len = rbuf.data_size as usize;
                Vec::from_raw_parts(data, len, len)
            };
            match M::decode_with_hash(&mut buf.as_slice()) {
                Ok(msg) => callback(msg),
                Err(_) => {}
            }
        });

        let mut subscription = Rc::new(LcmSubscription {
            subscription: ptr::null_mut(),
            handler: handler,
        });

        let user_data = (subscription.deref() as *const LcmSubscription) as *const libc::c_void;

        let c_subscription = unsafe {
            lcm_subscribe(self.lcm,
                          channel.as_ptr(),
                          Lcm::handler_callback::<M>,
                          user_data)
        };

        Rc::get_mut(&mut subscription).unwrap().subscription = c_subscription;
        self.subscriptions.push(subscription.clone());

        subscription
    }

    /// Unsubscribes a message handler.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("POSITION", handler_function);
    /// // ...
    /// lcm.unsubscribe(handler);
    /// ```
    pub fn unsubscribe(&mut self, handler: LcmSubscription) -> Result<()> {
        let result = unsafe { lcm_unsubscribe(self.lcm, handler.subscription) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM: Failed to unsubscribe")),
        }
    }

    /// Publishes a message on the specified channel.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    ///
    /// let mut my_data = exlcm::Example::new();
    /// my_data.timestamp = 0;
    /// my_data.position[0] = 1.0;
    /// my_data.position[1] = 2.0;
    /// my_data.position[2] = 3.0;
    ///
    /// lcm.publish("POSITION", &my_data).unwrap();
    /// ```
    pub fn publish<M>(&mut self, channel: &str, message: &M) -> Result<()> where M: Message + Sized {
        let channel = CString::new(channel).unwrap();
        let buffer = message.encode_with_hash()?;
        let datalen = buffer.len() as libc::c_uint;
        let result = unsafe {
            lcm_publish(self.lcm,
                        channel.as_ptr(),
                        buffer.as_ptr() as *mut libc::c_void,
                        datalen)
        };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    ///     lcm.handle().unwrap();
    /// }
    /// ```
    pub fn handle(&mut self) -> Result<()> {
        let result = unsafe { lcm_handle(self.lcm) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message, up to a time limit.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    ///     lcm.handle_timeout(1000).unwrap();
    /// }
    /// ```
    pub fn handle_timeout(&mut self, timeout_millis: i64) -> Result<()> {
        let timeout = timeout_millis as libc::c_int;
        let result = unsafe { lcm_handle_timeout(self.lcm, timeout) };
        match result.cmp(&0) {
            Ordering::Less => Err(Error::new(ErrorKind::Other, "LCM Error")),
            Ordering::Equal => Err(Error::new(ErrorKind::Other, "LCM Timeout")),
            Ordering::Greater => Ok(()),
        }
    }

    /// Adjusts the maximum number of received messages that can be queued up for a subscription.
    /// The default is `30`.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("POSITION", handler_function);
    /// lcm.subscription_set_queue_capacity(handler, 30);
    /// ```
    pub fn subscription_set_queue_capacity(handler: Rc<LcmSubscription>, num_messages: usize) {
        let handler = handler.subscription;
        let num_messages = num_messages as libc::c_int;
        unsafe { lcm_subscription_set_queue_capacity(handler, num_messages) };
    }

    extern "C" fn handler_callback<M>(rbuf: *const CLcmRecvBuf,
                                      _: *const libc::c_char,
                                      user_data: *const libc::c_void)
        where M: Message
    {

        let sub = user_data as *mut LcmSubscription;
        let sub = unsafe { &mut *sub };
        (sub.handler)(rbuf);
    }
}

impl Drop for Lcm {
    fn drop(&mut self) {
        unsafe { lcm_destroy(self.lcm) };
    }
}
<commit_msg>Use receive buffer as a slice instead of a Vec<commit_after>use std::io::{Error, ErrorKind, Result};
use std::ffi::CString;
use libc;
use message::Message;
use std::cmp::Ordering;
use std::ptr;
use std::boxed::Box;
use std::rc::Rc;
use std::ops::Deref;
use std::slice;

enum CLcm {}

/// An LCM instance that hnadles publishing and subscribing,
/// as well as encoding and decoding messages.
pub struct Lcm {
    lcm: *mut CLcm,
    subscriptions: Vec<Rc<LcmSubscription>>,
}

type CLcmHandler = extern "C" fn(*const CLcmRecvBuf, *const libc::c_char, *const libc::c_void);

#[derive(Eq, PartialEq, Hash)]
enum CLcmSubscription {}
pub struct LcmSubscription {
    subscription: *mut CLcmSubscription,
    handler: Box<FnMut(*const CLcmRecvBuf)>,
}

#[repr(C)]
#[derive(Debug)]
struct CLcmRecvBuf {
    data: *const libc::c_void,
    data_size: libc::uint32_t,
    recv_utime: libc::int64_t,
    lcm: *const CLcm,
}

#[link(name = "lcm")]
extern "C" {
    fn lcm_create(provider: *const libc::c_char) -> *mut CLcm;

    fn lcm_destroy(lcm: *mut CLcm);

    fn lcm_get_fileno(lcm: *mut CLcm) -> libc::c_int;

    fn lcm_subscribe(lcm: *mut CLcm,
                     channel: *const libc::c_char,
                     handler: CLcmHandler,
                     userdata: *const libc::c_void)
                     -> *mut CLcmSubscription;

    fn lcm_unsubscribe(lcm: *mut CLcm, handler: *const CLcmSubscription) -> libc::c_int;

    fn lcm_publish(lcm: *mut CLcm,
                   channel: *const libc::c_char,
                   data: *const libc::c_void,
                   datalen: libc::c_uint)
                   -> libc::c_int;

    fn lcm_handle(lcm: *mut CLcm) -> libc::c_int;

    fn lcm_handle_timeout(lcm: *mut CLcm, timeout_millis: libc::c_int) -> libc::c_int;

    fn lcm_subscription_set_queue_capacity(handler: *const CLcmSubscription,
                                           num_messages: libc::c_int)
                                           -> libc::c_int;
}

impl Lcm {
    /// Creates a new `Lcm` instance.
    ///
    /// ```
    /// use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// ```
    pub fn new() -> Result<Lcm> {
        let lcm = unsafe { lcm_create(ptr::null()) };
        match lcm.is_null() {
            true => Err(Error::new(ErrorKind::Other, "Failed to initialize LCM.")),
            false => {
                Ok(Lcm {
                    lcm: lcm,
                    subscriptions: Vec::new(),
                })
            }
        }
    }

    pub fn get_fileno(&self) -> libc::c_int {
        unsafe { lcm_get_fileno(self.lcm) }
    }

    /// Subscribes a callback to a particular topic.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// use std::sync::mpsc::channel;
    ///
    /// let mut lcm = Lcm::new().unwrap();
    /// let (tx, rx) = channel::<exlcm::Example>();
    /// lcm.subscribe("POSITION", |pos| { tx.send(pos).unwrap(); }
    /// ```
    pub fn subscribe<M, F>(&mut self, channel: &str, mut callback: Box<F>) -> Rc<LcmSubscription>
        where M: Message + Default,
              F: FnMut(M) + 'static
    {
        let channel = CString::new(channel).unwrap();

        let handler = Box::new(move |rbuf: *const CLcmRecvBuf| {
            let mut buf = unsafe {
                let ref rbuf = *rbuf;
                let data = rbuf.data as *mut u8;
                let len = rbuf.data_size as usize;
                slice::from_raw_parts(data, len)
            };
            match M::decode_with_hash(&mut buf) {
                Ok(msg) => callback(msg),
                Err(_) => {}
            }
        });

        let mut subscription = Rc::new(LcmSubscription {
            subscription: ptr::null_mut(),
            handler: handler,
        });

        let user_data = (subscription.deref() as *const LcmSubscription) as *const libc::c_void;

        let c_subscription = unsafe {
            lcm_subscribe(self.lcm,
                          channel.as_ptr(),
                          Lcm::handler_callback::<M>,
                          user_data)
        };

        Rc::get_mut(&mut subscription).unwrap().subscription = c_subscription;
        self.subscriptions.push(subscription.clone());

        subscription
    }

    /// Unsubscribes a message handler.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("POSITION", handler_function);
    /// // ...
    /// lcm.unsubscribe(handler);
    /// ```
    pub fn unsubscribe(&mut self, handler: LcmSubscription) -> Result<()> {
        let result = unsafe { lcm_unsubscribe(self.lcm, handler.subscription) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM: Failed to unsubscribe")),
        }
    }

    /// Publishes a message on the specified channel.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    ///
    /// let mut my_data = exlcm::Example::new();
    /// my_data.timestamp = 0;
    /// my_data.position[0] = 1.0;
    /// my_data.position[1] = 2.0;
    /// my_data.position[2] = 3.0;
    ///
    /// lcm.publish("POSITION", &my_data).unwrap();
    /// ```
    pub fn publish<M>(&mut self, channel: &str, message: &M) -> Result<()> where M: Message + Sized {
        let channel = CString::new(channel).unwrap();
        let buffer = message.encode_with_hash()?;
        let datalen = buffer.len() as libc::c_uint;
        let result = unsafe {
            lcm_publish(self.lcm,
                        channel.as_ptr(),
                        buffer.as_ptr() as *mut libc::c_void,
                        datalen)
        };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    ///     lcm.handle().unwrap();
    /// }
    /// ```
    pub fn handle(&mut self) -> Result<()> {
        let result = unsafe { lcm_handle(self.lcm) };
        match result {
            0 => Ok(()),
            _ => Err(Error::new(ErrorKind::Other, "LCM Error")),
        }
    }

    /// Waits for and dispatches the next incoming message, up to a time limit.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// lcm.subscribe("POSITION", handler_function);
    /// loop {
    ///     lcm.handle_timeout(1000).unwrap();
    /// }
    /// ```
    pub fn handle_timeout(&mut self, timeout_millis: i64) -> Result<()> {
        let timeout = timeout_millis as libc::c_int;
        let result = unsafe { lcm_handle_timeout(self.lcm, timeout) };
        match result.cmp(&0) {
            Ordering::Less => Err(Error::new(ErrorKind::Other, "LCM Error")),
            Ordering::Equal => Err(Error::new(ErrorKind::Other, "LCM Timeout")),
            Ordering::Greater => Ok(()),
        }
    }

    /// Adjusts the maximum number of received messages that can be queued up for a subscription.
    /// The default is `30`.
    ///
    /// ```
    /// # use lcm::Lcm;
    /// let mut lcm = Lcm::new().unwrap();
    /// let handler = lcm.subscribe("POSITION", handler_function);
    /// lcm.subscription_set_queue_capacity(handler, 30);
    /// ```
    pub fn subscription_set_queue_capacity(handler: Rc<LcmSubscription>, num_messages: usize) {
        let handler = handler.subscription;
        let num_messages = num_messages as libc::c_int;
        unsafe { lcm_subscription_set_queue_capacity(handler, num_messages) };
    }

    extern "C" fn handler_callback<M>(rbuf: *const CLcmRecvBuf,
                                      _: *const libc::c_char,
                                      user_data: *const libc::c_void)
        where M: Message
    {

        let sub = user_data as *mut LcmSubscription;
        let sub = unsafe { &mut *sub };
        (sub.handler)(rbuf);
    }
}

impl Drop for Lcm {
    fn drop(&mut self) {
        unsafe { lcm_destroy(self.lcm) };
    }
}
<|endoftext|>