<commit_before>#![warn(clippy::all, rust_2018_idioms)]

use cargo_registry::{env_optional, metrics::LogEncoder, util::errors::AppResult, App, Env};
use std::{env, fs::File, process::Command, sync::Arc, time::Duration};

use conduit_hyper::Service;
use futures_util::future::FutureExt;
use prometheus::Encoder;
use reqwest::blocking::Client;
use std::io::Write;
use tokio::io::AsyncWriteExt;
use tokio::signal::unix::{signal, SignalKind};
use tracing::Level;
use tracing_subscriber::{filter, prelude::*};

const CORE_THREADS: usize = 4;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let _sentry = cargo_registry::sentry::init();

    // Initialize logging

    let log_filter = env::var("RUST_LOG")
        .unwrap_or_default()
        .parse::<filter::Targets>()
        .expect("Invalid RUST_LOG value");

    let sentry_filter = filter::Targets::new().with_default(Level::INFO);

    tracing_subscriber::registry()
        .with(tracing_subscriber::fmt::layer().with_filter(log_filter))
        .with(sentry::integrations::tracing::layer().with_filter(sentry_filter))
        .init();

    let config = cargo_registry::config::Server::default();
    let env = config.env();
    let client = Client::new();
    let app = Arc::new(App::new(config, Some(client)));

    // Start the background thread periodically persisting download counts to the database.
    downloads_counter_thread(app.clone());

    // Start the background thread periodically logging instance metrics.
    log_instance_metrics_thread(app.clone());

    let handler = cargo_registry::build_handler(app.clone());

    let heroku = dotenv::var("HEROKU").is_ok();
    let fastboot = dotenv::var("USE_FASTBOOT").is_ok();
    let dev_docker = dotenv::var("DEV_DOCKER").is_ok();

    let ip = if dev_docker {
        [0, 0, 0, 0]
    } else {
        [127, 0, 0, 1]
    };
    let port = match (heroku, env_optional("PORT")) {
        (false, Some(port)) => port,
        _ => 8888,
    };

    let threads = dotenv::var("SERVER_THREADS")
        .map(|s| s.parse().expect("SERVER_THREADS was not a valid number"))
        .unwrap_or_else(|_| match env {
            Env::Development => 5,
            // A large default because this can be easily changed via env and in production we
            // want the logging middleware to accurately record the start time.
            _ => 500,
        });

    println!("Booting with a hyper based server");

    let rt = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .worker_threads(CORE_THREADS)
        .max_blocking_threads(threads as usize)
        .build()
        .unwrap();

    let handler = Arc::new(conduit_hyper::BlockingHandler::new(handler));
    let make_service =
        hyper::service::make_service_fn(move |socket: &hyper::server::conn::AddrStream| {
            let addr = socket.remote_addr();
            let handler = handler.clone();
            async move { Service::from_blocking(handler, addr) }
        });

    let addr = (ip, port).into();
    #[allow(clippy::async_yields_async)]
    let server = rt.block_on(async { hyper::Server::bind(&addr).serve(make_service) });

    let mut sig_int = rt.block_on(async { signal(SignalKind::interrupt()) })?;
    let mut sig_term = rt.block_on(async { signal(SignalKind::terminate()) })?;

    // When the user configures PORT=0 the operative system will allocate a random unused port.
    // This fetches that random port and uses it to display the "listening on port" message later.
    let addr = server.local_addr();

    let server = server.with_graceful_shutdown(async move {
        // Wait for either signal
        futures_util::select! {
            _ = sig_int.recv().fuse() => {},
            _ = sig_term.recv().fuse() => {},
        };
        let mut stdout = tokio::io::stdout();
        stdout.write_all(b"Starting graceful shutdown\n").await.ok();
    });

    let server = rt.spawn(async { server.await.unwrap() });

    // Do not change this line! Removing the line or changing its contents in any way will break
    // the test suite :)
    println!("Listening at http://{addr}");

    // Creating this file tells heroku to tell nginx that the application is ready
    // to receive traffic.
    if heroku {
        let path = if fastboot {
            "/tmp/backend-initialized"
        } else {
            "/tmp/app-initialized"
        };
        println!("Writing to {path}");
        File::create(path).unwrap();

        // Launch nginx via the Heroku nginx buildpack
        // `wait()` is never called on the child process, but it should be okay to leave a zombie
        // process around on shutdown when Heroku is tearing down the entire container anyway.
        Command::new("./script/start-web.sh")
            .spawn()
            .expect("Couldn't spawn nginx");
    }

    // Block the main thread until the server has shutdown
    rt.block_on(async { server.await.unwrap() });

    println!("Persisting remaining downloads counters");
    match app.downloads_counter.persist_all_shards(&app) {
        Ok(stats) => stats.log(),
        Err(err) => println!("downloads_counter error: {err}"),
    }

    println!("Server has gracefully shutdown!");
    Ok(())
}

fn downloads_counter_thread(app: Arc<App>) {
    let interval = Duration::from_millis(
        (app.config.downloads_persist_interval_ms / app.downloads_counter.shards_count()) as u64,
    );

    std::thread::spawn(move || loop {
        std::thread::sleep(interval);

        match app.downloads_counter.persist_next_shard(&app) {
            Ok(stats) => stats.log(),
            Err(err) => println!("downloads_counter error: {err}"),
        }
    });
}

fn log_instance_metrics_thread(app: Arc<App>) {
    // Only run the thread if the configuration is provided
    let interval = if let Some(secs) = app.config.instance_metrics_log_every_seconds {
        Duration::from_secs(secs)
    } else {
        return;
    };

    std::thread::spawn(move || loop {
        if let Err(err) = log_instance_metrics_inner(&app) {
            eprintln!("log_instance_metrics error: {err}");
        }
        std::thread::sleep(interval);
    });
}

fn log_instance_metrics_inner(app: &App) -> AppResult<()> {
    let families = app.instance_metrics.gather(app)?;

    let mut stdout = std::io::stdout();
    LogEncoder::new().encode(&families, &mut stdout)?;
    stdout.flush()?;

    Ok(())
}
<commit_msg>Simplify spawning the http server in production (#5160)<commit_after>#![warn(clippy::all, rust_2018_idioms)]

use cargo_registry::{env_optional, metrics::LogEncoder, util::errors::AppResult, App, Env};
use std::{env, fs::File, process::Command, sync::Arc, time::Duration};

use conduit_hyper::Service;
use futures_util::future::FutureExt;
use prometheus::Encoder;
use reqwest::blocking::Client;
use std::io::{self, Write};
use tokio::io::AsyncWriteExt;
use tokio::signal::unix::{signal, SignalKind};
use tracing::Level;
use tracing_subscriber::{filter, prelude::*};

const CORE_THREADS: usize = 4;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let _sentry = cargo_registry::sentry::init();

    // Initialize logging

    let log_filter = env::var("RUST_LOG")
        .unwrap_or_default()
        .parse::<filter::Targets>()
        .expect("Invalid RUST_LOG value");

    let sentry_filter = filter::Targets::new().with_default(Level::INFO);

    tracing_subscriber::registry()
        .with(tracing_subscriber::fmt::layer().with_filter(log_filter))
        .with(sentry::integrations::tracing::layer().with_filter(sentry_filter))
        .init();

    let config = cargo_registry::config::Server::default();
    let env = config.env();
    let client = Client::new();
    let app = Arc::new(App::new(config, Some(client)));

    // Start the background thread periodically persisting download counts to the database.
    downloads_counter_thread(app.clone());

    // Start the background thread periodically logging instance metrics.
    log_instance_metrics_thread(app.clone());

    let handler = cargo_registry::build_handler(app.clone());

    let heroku = dotenv::var("HEROKU").is_ok();
    let fastboot = dotenv::var("USE_FASTBOOT").is_ok();
    let dev_docker = dotenv::var("DEV_DOCKER").is_ok();

    let ip = if dev_docker {
        [0, 0, 0, 0]
    } else {
        [127, 0, 0, 1]
    };
    let port = match (heroku, env_optional("PORT")) {
        (false, Some(port)) => port,
        _ => 8888,
    };

    let threads = dotenv::var("SERVER_THREADS")
        .map(|s| s.parse().expect("SERVER_THREADS was not a valid number"))
        .unwrap_or_else(|_| match env {
            Env::Development => 5,
            // A large default because this can be easily changed via env and in production we
            // want the logging middleware to accurately record the start time.
            _ => 500,
        });

    let rt = tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .worker_threads(CORE_THREADS)
        .max_blocking_threads(threads as usize)
        .build()
        .unwrap();

    let handler = Arc::new(conduit_hyper::BlockingHandler::new(handler));
    let make_service =
        hyper::service::make_service_fn(move |socket: &hyper::server::conn::AddrStream| {
            let addr = socket.remote_addr();
            let handler = handler.clone();
            async move { Service::from_blocking(handler, addr) }
        });

    let (addr, server) = rt.block_on(async {
        let server = hyper::Server::bind(&(ip, port).into()).serve(make_service);

        // When the user configures PORT=0 the operating system will allocate a random unused port.
        // This fetches that random port and uses it to display the the correct url later.
        let addr = server.local_addr();

        let mut sig_int = signal(SignalKind::interrupt())?;
        let mut sig_term = signal(SignalKind::terminate())?;
        let server = server.with_graceful_shutdown(async move {
            // Wait for either signal
            tokio::select! {
                _ = sig_int.recv().fuse() => {},
                _ = sig_term.recv().fuse() => {},
            };
            tokio::io::stdout()
                .write_all(b"Starting graceful shutdown\n")
                .await
                .ok();
        });

        Ok::<_, io::Error>((addr, server))
    })?;

    // Do not change this line! Removing the line or changing its contents in any way will break
    // the test suite :)
    println!("Listening at http://{addr}");

    // Creating this file tells heroku to tell nginx that the application is ready
    // to receive traffic.
    if heroku {
        let path = if fastboot {
            "/tmp/backend-initialized"
        } else {
            "/tmp/app-initialized"
        };
        println!("Writing to {path}");
        File::create(path).unwrap();

        // Launch nginx via the Heroku nginx buildpack
        // `wait()` is never called on the child process, but it should be okay to leave a zombie
        // process around on shutdown when Heroku is tearing down the entire container anyway.
        Command::new("./script/start-web.sh")
            .spawn()
            .expect("Couldn't spawn nginx");
    }

    // Block the main thread until the server has shutdown
    rt.block_on(server)?;

    println!("Persisting remaining downloads counters");
    match app.downloads_counter.persist_all_shards(&app) {
        Ok(stats) => stats.log(),
        Err(err) => println!("downloads_counter error: {err}"),
    }

    println!("Server has gracefully shutdown!");
    Ok(())
}

fn downloads_counter_thread(app: Arc<App>) {
    let interval = Duration::from_millis(
        (app.config.downloads_persist_interval_ms / app.downloads_counter.shards_count()) as u64,
    );

    std::thread::spawn(move || loop {
        std::thread::sleep(interval);

        match app.downloads_counter.persist_next_shard(&app) {
            Ok(stats) => stats.log(),
            Err(err) => println!("downloads_counter error: {err}"),
        }
    });
}

fn log_instance_metrics_thread(app: Arc<App>) {
    // Only run the thread if the configuration is provided
    let interval = if let Some(secs) = app.config.instance_metrics_log_every_seconds {
        Duration::from_secs(secs)
    } else {
        return;
    };

    std::thread::spawn(move || loop {
        if let Err(err) = log_instance_metrics_inner(&app) {
            eprintln!("log_instance_metrics error: {err}");
        }
        std::thread::sleep(interval);
    });
}

fn log_instance_metrics_inner(app: &App) -> AppResult<()> {
    let families = app.instance_metrics.gather(app)?;

    let mut stdout = std::io::stdout();
    LogEncoder::new().encode(&families, &mut stdout)?;
    stdout.flush()?;

    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2017 int08h LLC
//
// 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.

//!
//! Roughtime server
//!
//! # Configuration
//! The `roughenough` server is configured via a config file:
//!
//! ```yaml
//! interface: 127.0.0.1
//! port: 8686
//! seed: f61075c988feb9cb700a4a6a3291bfbc9cab11b9c9eca8c802468eb38a43d7d3
//! batch_size: 64
//! ```
//!
//! Where:
//!
//!   * **interface** - IP address or interface name for listening to client requests
//!   * **port** - UDP port to listen to requests
//!   * **seed** - A 32-byte hexadecimal value used as the seed to generate the
//!                server's long-term key pair. **This is a secret value**, treat it
//!                with care.
//!   * **batch_size** - The number of requests to process in one batch. All nonces
//!                      in a batch are used to build a Merkle tree, the root of which
//!                      is signed.
//!
//! # Running the Server
//!
//! ```bash
//! $ cargo run --release --bin server /path/to/config.file
//! ```

extern crate byteorder;
extern crate ring;
extern crate roughenough;
extern crate time;
extern crate untrusted;
extern crate ctrlc;
extern crate yaml_rust;
#[macro_use]
extern crate log;
extern crate simple_logger;
extern crate mio;
extern crate mio_extras;
extern crate hex;

use std::env;
use std::process;
use std::fs::File;
use std::io::{Read, ErrorKind};
use std::time::Duration;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::thread;

use mio::{Poll, Token, Ready, PollOpt, Events};
use mio::net::UdpSocket;
use mio_extras::timer::Timer;

use byteorder::{LittleEndian, WriteBytesExt};

use roughenough::{RtMessage, Tag, Error};
use roughenough::{VERSION, CERTIFICATE_CONTEXT, MIN_REQUEST_LENGTH, SIGNED_RESPONSE_CONTEXT};
use roughenough::sign::Signer;
use roughenough::merkle::*;

use ring::rand;
use ring::rand::SecureRandom;

use yaml_rust::YamlLoader;

const MESSAGE: Token = Token(0);
const STATUS: Token = Token(1);

pub static NUM_RESPONSES: AtomicUsize = AtomicUsize::new(0);

fn create_ephemeral_key() -> Signer {
    let rng = rand::SystemRandom::new();
    let mut seed = [0u8; 32];
    rng.fill(&mut seed).unwrap();

    Signer::new(&seed)
}

fn make_dele_bytes(ephemeral_key: &Signer) -> Result<Vec<u8>, Error> {
    let zeros = [0u8; 8];
    let max = [0xff; 8];

    let mut dele_msg = RtMessage::new(3);
    dele_msg.add_field(Tag::PUBK, ephemeral_key.public_key_bytes())?;
    dele_msg.add_field(Tag::MINT, &zeros)?;
    dele_msg.add_field(Tag::MAXT, &max)?;

    dele_msg.encode()
}

fn make_key_and_cert(seed: &[u8]) -> (Signer, Vec<u8>) {
    let mut long_term_key = Signer::new(seed);
    let ephemeral_key = create_ephemeral_key();

    info!("Long-term public key: {}", hex::encode(long_term_key.public_key_bytes()));
    info!("Ephemeral public key: {}", hex::encode(ephemeral_key.public_key_bytes()));

    // Make DELE and sign it with long-term key
    let dele_bytes = make_dele_bytes(&ephemeral_key).unwrap();
    let dele_signature = {
        long_term_key.update(CERTIFICATE_CONTEXT.as_bytes());
        long_term_key.update(&dele_bytes);
        long_term_key.sign()
    };

    // Create CERT
    let cert_bytes = {
        let mut cert_msg = RtMessage::new(2);
        cert_msg.add_field(Tag::SIG, &dele_signature).unwrap();
        cert_msg.add_field(Tag::DELE, &dele_bytes).unwrap();

        cert_msg.encode().unwrap()
    };

    (ephemeral_key, cert_bytes)
}

struct SRep {
    raw_bytes: Vec<u8>,
    signature: Vec<u8>
}

fn make_srep(ephemeral_key: &mut Signer, root: &[u8]) -> SRep {
    //   create SREP
    //   sign SREP
    //   create response:
    //    - SIG
    //    - PATH (always 0)
    //    - SREP
    //    - CERT (pre-created)
    //    - INDX (always 0)

    let mut radi = [0; 4];
    let mut midp = [0; 8];

    // one second (in microseconds)
    (&mut radi as &mut [u8]).write_u32::<LittleEndian>(1_000_000).unwrap();

    // current epoch time in microseconds
    let now = {
        let tv = time::get_time();
        let secs = (tv.sec as u64) * 1_000_000;
        let nsecs = (tv.nsec as u64) / 1_000;

        secs + nsecs
    };
    (&mut midp as &mut [u8]).write_u64::<LittleEndian>(now).unwrap();

    // Signed response SREP
    let srep_bytes = {
        let mut srep_msg = RtMessage::new(3);
        srep_msg.add_field(Tag::RADI, &radi).unwrap();
        srep_msg.add_field(Tag::MIDP, &midp).unwrap();
        srep_msg.add_field(Tag::ROOT, root).unwrap();

        srep_msg.encode().unwrap()
    };

    // signature on SREP
    let srep_signature = {
        ephemeral_key.update(SIGNED_RESPONSE_CONTEXT.as_bytes());
        ephemeral_key.update(&srep_bytes);
        ephemeral_key.sign()
    };

    SRep {
        raw_bytes: srep_bytes,
        signature: srep_signature
    }
}

fn make_response(srep: &SRep, cert_bytes: &[u8], path: &[u8], idx: u32) -> RtMessage {

    let mut index = [0; 4];
    (&mut index as &mut [u8]).write_u32::<LittleEndian>(idx).unwrap();

    let mut response = RtMessage::new(5);
    response.add_field(Tag::SIG, &srep.signature).unwrap();
    response.add_field(Tag::PATH, &path).unwrap();
    response.add_field(Tag::SREP, &srep.raw_bytes).unwrap();
    response.add_field(Tag::CERT, cert_bytes).unwrap();
    response.add_field(Tag::INDX, &index).unwrap();

    response
}

// extract the client's nonce from its request
fn nonce_from_request(buf: &[u8], num_bytes: usize) -> Result<&[u8], Error> {
    if num_bytes < MIN_REQUEST_LENGTH as usize {
        return Err(Error::RequestTooShort);
    }

    let tag_count = &buf[..4];
    let expected_nonc = &buf[8..12];
    let expected_pad = &buf[12..16];

    let tag_count_is_2 = tag_count == [0x02, 0x00, 0x00, 0x00];
    let tag1_is_nonc = expected_nonc == Tag::NONC.wire_value();
    let tag2_is_pad = expected_pad == Tag::PAD.wire_value();

    if tag_count_is_2 && tag1_is_nonc && tag2_is_pad {
        Ok(&buf[0x10..0x50])
    } else {
        Err(Error::InvalidRequest)
    }
}

fn load_config(config_file: &str) -> (SocketAddr, Vec<u8>, u8) {
    let mut infile = File::open(config_file)
        .expect("failed to open config file");

    let mut contents = String::new();
    infile.read_to_string(&mut contents)
        .expect("could not read config file");

    let cfg = YamlLoader::load_from_str(&contents)
        .expect("could not parse config file");

    if cfg.len() != 1 {
        panic!("empty or malformed config file");
    }

    let mut port: u16 = 0;
    let mut iface: String = "unknown".to_string();
    let mut seed: String = "".to_string();
    let mut batch_size: u8 = 1;

    for (key, value) in cfg[0].as_hash().unwrap() {
        match key.as_str().unwrap() {
            "port" => port = value.as_i64().unwrap() as u16,
            "interface" => iface = value.as_str().unwrap().to_string(),
            "seed" => seed = value.as_str().unwrap().to_string(),
            "batch_size" => batch_size = value.as_i64().unwrap() as u8,
            _ => warn!("ignoring unknown config key '{}'", key.as_str().unwrap())
        }
    }

    let addr = format!("{}:{}", iface, port);
    let sock_addr: SocketAddr = addr.parse()
        .expect(&format!("could not create socket address from {}", addr));

    let binseed = hex::decode(seed)
        .expect("seed value invalid; 'seed' should be 32 byte hex value");

    (sock_addr, binseed, batch_size)
}

fn polling_loop(addr: &SocketAddr, mut ephemeral_key: &mut Signer, cert_bytes: &[u8], batch_size: u8) {
    let keep_running = Arc::new(AtomicBool::new(true));
    let kr = keep_running.clone();

    ctrlc::set_handler(move || kr.store(false, Ordering::Release))
        .expect("failed setting Ctrl-C handler");

    let socket = UdpSocket::bind(addr).expect("failed to bind to socket");
    let status_duration = Duration::from_secs(6);
    let poll_duration = Some(Duration::from_millis(100));

    let mut timer: Timer<()> = Timer::default();
    timer.set_timeout(status_duration, ());

    let mut buf = [0u8; 65_536];
    let mut events = Events::with_capacity(32);
    let mut num_bad_requests = 0u64;

    let poll = Poll::new().unwrap();
    poll.register(&socket, MESSAGE, Ready::readable(), PollOpt::edge()).unwrap();
    poll.register(&timer, STATUS, Ready::readable(), PollOpt::edge()).unwrap();

    let mut merkle = MerkleTree::new();
    let mut requests = Vec::with_capacity(batch_size as usize);

    macro_rules! check_ctrlc {
        () => {
            if !keep_running.load(Ordering::Acquire) {
                warn!("Ctrl-C caught, exiting...");
                return;
            }
        }
    }

    loop {
        check_ctrlc!();

        poll.poll(&mut events, poll_duration).expect("poll failed");

        for event in events.iter() {

            match event.token() {
                MESSAGE => {

                    let mut done = false;

                    'process_batch: loop {
                        check_ctrlc!();

                        merkle.reset();
                        requests.clear();

                        let resp_start = NUM_RESPONSES.load(Ordering::SeqCst);

                        for i in 0..batch_size {
                            match socket.recv_from(&mut buf) {
                                Ok((num_bytes, src_addr)) => {
                                    if let Ok(nonce) = nonce_from_request(&buf, num_bytes) {
                                        requests.push((Vec::from(nonce), src_addr));
                                        merkle.push_leaf(nonce);
                                    } else {
                                        num_bad_requests += 1;
                                        info!("Invalid request ({} bytes) from {} (#{} in batch, resp #{})", num_bytes, src_addr, i, resp_start + i as usize);
                                    }
                                },
                                Err(e) => match e.kind() {
                                    ErrorKind::WouldBlock => {
                                        done = true;
                                        break;
                                    },
                                    _ => panic!("recv_from failed with {:?}", e)
                                }
                            };
                        }

                        if requests.is_empty() {
                            break 'process_batch
                        }

                        let root = merkle.compute_root();
                        let srep = make_srep(&mut ephemeral_key, &root);

                        for (i, &(ref nonce, ref src_addr)) in requests.iter().enumerate() {
                            let paths = merkle.get_paths(i);

                            let resp = make_response(&srep, cert_bytes, &paths, i as u32);
                            let resp_bytes = resp.encode().unwrap();

                            let bytes_sent = socket.send_to(&resp_bytes, &src_addr).expect("send_to failed");
                            let num_responses = NUM_RESPONSES.fetch_add(1, Ordering::SeqCst);

                            info!("Responded {} bytes to {} for '{}..' (#{} in batch, resp #{})", bytes_sent, src_addr, hex::encode(&nonce[0..4]), i, num_responses);
                        }
                        if done {
                            break 'process_batch
                        }
                    }

                }

                STATUS => {
                    info!("responses {}, invalid requests {}", NUM_RESPONSES.load(Ordering::SeqCst), num_bad_requests);
                    timer.set_timeout(status_duration, ());
                }

                _ => unreachable!()
            }
        }
    }
}

pub fn main() {
    use log::Level;

    simple_logger::init_with_level(Level::Info).unwrap();

    info!("Roughenough server v{} starting", VERSION);

    let mut args = env::args();
    if args.len() != 2 {
        error!("Usage: server /path/to/config.file");
        process::exit(1);
    }

    let (addr, key_seed, batch_size) = load_config(&args.nth(1).unwrap());
    let (mut ephemeral_key, cert_bytes) = make_key_and_cert(&key_seed);

    info!("Server listening on {}", addr);

    if env::var("BENCH").is_ok()  {
        log::set_max_level(log::LevelFilter::Warn);

        thread::spawn(|| {
            loop {
                let old = time::get_time().sec;
                let old_reqs = NUM_RESPONSES.load(Ordering::SeqCst);

                thread::sleep(Duration::from_secs(1));

                let new = time::get_time().sec;
                let new_reqs = NUM_RESPONSES.load(Ordering::SeqCst);

                warn!("Processing at {:?} reqs/sec", (new_reqs - old_reqs) / (new - old) as usize);
            }
        });
    }

    polling_loop(&addr, &mut ephemeral_key, &cert_bytes, batch_size);

    info!("Done.");
    process::exit(0);
}
<commit_msg>Use an Arc to share response counter<commit_after>// Copyright 2017 int08h LLC
//
// 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.

//!
//! Roughtime server
//!
//! # Configuration
//! The `roughenough` server is configured via a config file:
//!
//! ```yaml
//! interface: 127.0.0.1
//! port: 8686
//! seed: f61075c988feb9cb700a4a6a3291bfbc9cab11b9c9eca8c802468eb38a43d7d3
//! batch_size: 64
//! ```
//!
//! Where:
//!
//!   * **interface** - IP address or interface name for listening to client requests
//!   * **port** - UDP port to listen to requests
//!   * **seed** - A 32-byte hexadecimal value used as the seed to generate the
//!                server's long-term key pair. **This is a secret value**, treat it
//!                with care.
//!   * **batch_size** - The number of requests to process in one batch. All nonces
//!                      in a batch are used to build a Merkle tree, the root of which
//!                      is signed.
//!
//! # Running the Server
//!
//! ```bash
//! $ cargo run --release --bin server /path/to/config.file
//! ```

extern crate byteorder;
extern crate ring;
extern crate roughenough;
extern crate time;
extern crate untrusted;
extern crate ctrlc;
extern crate yaml_rust;
#[macro_use]
extern crate log;
extern crate simple_logger;
extern crate mio;
extern crate mio_extras;
extern crate hex;

use std::env;
use std::process;
use std::fs::File;
use std::io::{Read, ErrorKind};
use std::time::Duration;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::thread;

use mio::{Poll, Token, Ready, PollOpt, Events};
use mio::net::UdpSocket;
use mio_extras::timer::Timer;

use byteorder::{LittleEndian, WriteBytesExt};

use roughenough::{RtMessage, Tag, Error};
use roughenough::{VERSION, CERTIFICATE_CONTEXT, MIN_REQUEST_LENGTH, SIGNED_RESPONSE_CONTEXT};
use roughenough::sign::Signer;
use roughenough::merkle::*;

use ring::rand;
use ring::rand::SecureRandom;

use yaml_rust::YamlLoader;

const MESSAGE: Token = Token(0);
const STATUS: Token = Token(1);

fn create_ephemeral_key() -> Signer {
    let rng = rand::SystemRandom::new();
    let mut seed = [0u8; 32];
    rng.fill(&mut seed).unwrap();

    Signer::new(&seed)
}

fn make_dele_bytes(ephemeral_key: &Signer) -> Result<Vec<u8>, Error> {
    let zeros = [0u8; 8];
    let max = [0xff; 8];

    let mut dele_msg = RtMessage::new(3);
    dele_msg.add_field(Tag::PUBK, ephemeral_key.public_key_bytes())?;
    dele_msg.add_field(Tag::MINT, &zeros)?;
    dele_msg.add_field(Tag::MAXT, &max)?;

    dele_msg.encode()
}

fn make_key_and_cert(seed: &[u8]) -> (Signer, Vec<u8>) {
    let mut long_term_key = Signer::new(seed);
    let ephemeral_key = create_ephemeral_key();

    info!("Long-term public key: {}", hex::encode(long_term_key.public_key_bytes()));
    info!("Ephemeral public key: {}", hex::encode(ephemeral_key.public_key_bytes()));

    // Make DELE and sign it with long-term key
    let dele_bytes = make_dele_bytes(&ephemeral_key).unwrap();
    let dele_signature = {
        long_term_key.update(CERTIFICATE_CONTEXT.as_bytes());
        long_term_key.update(&dele_bytes);
        long_term_key.sign()
    };

    // Create CERT
    let cert_bytes = {
        let mut cert_msg = RtMessage::new(2);
        cert_msg.add_field(Tag::SIG, &dele_signature).unwrap();
        cert_msg.add_field(Tag::DELE, &dele_bytes).unwrap();

        cert_msg.encode().unwrap()
    };

    (ephemeral_key, cert_bytes)
}

struct SRep {
    raw_bytes: Vec<u8>,
    signature: Vec<u8>
}

fn make_srep(ephemeral_key: &mut Signer, root: &[u8]) -> SRep {
    //   create SREP
    //   sign SREP
    //   create response:
    //    - SIG
    //    - PATH (always 0)
    //    - SREP
    //    - CERT (pre-created)
    //    - INDX (always 0)

    let mut radi = [0; 4];
    let mut midp = [0; 8];

    // one second (in microseconds)
    (&mut radi as &mut [u8]).write_u32::<LittleEndian>(1_000_000).unwrap();

    // current epoch time in microseconds
    let now = {
        let tv = time::get_time();
        let secs = (tv.sec as u64) * 1_000_000;
        let nsecs = (tv.nsec as u64) / 1_000;

        secs + nsecs
    };
    (&mut midp as &mut [u8]).write_u64::<LittleEndian>(now).unwrap();

    // Signed response SREP
    let srep_bytes = {
        let mut srep_msg = RtMessage::new(3);
        srep_msg.add_field(Tag::RADI, &radi).unwrap();
        srep_msg.add_field(Tag::MIDP, &midp).unwrap();
        srep_msg.add_field(Tag::ROOT, root).unwrap();

        srep_msg.encode().unwrap()
    };

    // signature on SREP
    let srep_signature = {
        ephemeral_key.update(SIGNED_RESPONSE_CONTEXT.as_bytes());
        ephemeral_key.update(&srep_bytes);
        ephemeral_key.sign()
    };

    SRep {
        raw_bytes: srep_bytes,
        signature: srep_signature
    }
}

fn make_response(srep: &SRep, cert_bytes: &[u8], path: &[u8], idx: u32) -> RtMessage {

    let mut index = [0; 4];
    (&mut index as &mut [u8]).write_u32::<LittleEndian>(idx).unwrap();

    let mut response = RtMessage::new(5);
    response.add_field(Tag::SIG, &srep.signature).unwrap();
    response.add_field(Tag::PATH, &path).unwrap();
    response.add_field(Tag::SREP, &srep.raw_bytes).unwrap();
    response.add_field(Tag::CERT, cert_bytes).unwrap();
    response.add_field(Tag::INDX, &index).unwrap();

    response
}

// extract the client's nonce from its request
fn nonce_from_request(buf: &[u8], num_bytes: usize) -> Result<&[u8], Error> {
    if num_bytes < MIN_REQUEST_LENGTH as usize {
        return Err(Error::RequestTooShort);
    }

    let tag_count = &buf[..4];
    let expected_nonc = &buf[8..12];
    let expected_pad = &buf[12..16];

    let tag_count_is_2 = tag_count == [0x02, 0x00, 0x00, 0x00];
    let tag1_is_nonc = expected_nonc == Tag::NONC.wire_value();
    let tag2_is_pad = expected_pad == Tag::PAD.wire_value();

    if tag_count_is_2 && tag1_is_nonc && tag2_is_pad {
        Ok(&buf[0x10..0x50])
    } else {
        Err(Error::InvalidRequest)
    }
}

fn load_config(config_file: &str) -> (SocketAddr, Vec<u8>, u8) {
    let mut infile = File::open(config_file)
        .expect("failed to open config file");

    let mut contents = String::new();
    infile.read_to_string(&mut contents)
        .expect("could not read config file");

    let cfg = YamlLoader::load_from_str(&contents)
        .expect("could not parse config file");

    if cfg.len() != 1 {
        panic!("empty or malformed config file");
    }

    let mut port: u16 = 0;
    let mut iface: String = "unknown".to_string();
    let mut seed: String = "".to_string();
    let mut batch_size: u8 = 1;

    for (key, value) in cfg[0].as_hash().unwrap() {
        match key.as_str().unwrap() {
            "port" => port = value.as_i64().unwrap() as u16,
            "interface" => iface = value.as_str().unwrap().to_string(),
            "seed" => seed = value.as_str().unwrap().to_string(),
            "batch_size" => batch_size = value.as_i64().unwrap() as u8,
            _ => warn!("ignoring unknown config key '{}'", key.as_str().unwrap())
        }
    }

    let addr = format!("{}:{}", iface, port);
    let sock_addr: SocketAddr = addr.parse()
        .expect(&format!("could not create socket address from {}", addr));

    let binseed = hex::decode(seed)
        .expect("seed value invalid; 'seed' should be 32 byte hex value");

    (sock_addr, binseed, batch_size)
}

fn polling_loop(addr: &SocketAddr, mut ephemeral_key: &mut Signer, cert_bytes: &[u8], batch_size: u8, response_counter: Arc<AtomicUsize>) {
    let keep_running = Arc::new(AtomicBool::new(true));
    let kr = keep_running.clone();

    ctrlc::set_handler(move || kr.store(false, Ordering::Release))
        .expect("failed setting Ctrl-C handler");

    let socket = UdpSocket::bind(addr).expect("failed to bind to socket");
    let status_duration = Duration::from_secs(6);
    let poll_duration = Some(Duration::from_millis(100));

    let mut timer: Timer<()> = Timer::default();
    timer.set_timeout(status_duration, ());

    let mut buf = [0u8; 65_536];
    let mut events = Events::with_capacity(32);
    let mut num_bad_requests = 0u64;

    let poll = Poll::new().unwrap();
    poll.register(&socket, MESSAGE, Ready::readable(), PollOpt::edge()).unwrap();
    poll.register(&timer, STATUS, Ready::readable(), PollOpt::edge()).unwrap();

    let mut merkle = MerkleTree::new();
    let mut requests = Vec::with_capacity(batch_size as usize);

    macro_rules! check_ctrlc {
        () => {
            if !keep_running.load(Ordering::Acquire) {
                warn!("Ctrl-C caught, exiting...");
                return;
            }
        }
    }

    loop {
        check_ctrlc!();

        poll.poll(&mut events, poll_duration).expect("poll failed");

        for event in events.iter() {

            match event.token() {
                MESSAGE => {

                    let mut done = false;

                    'process_batch: loop {
                        check_ctrlc!();

                        merkle.reset();
                        requests.clear();

                        let resp_start = response_counter.load(Ordering::SeqCst);

                        for i in 0..batch_size {
                            match socket.recv_from(&mut buf) {
                                Ok((num_bytes, src_addr)) => {
                                    if let Ok(nonce) = nonce_from_request(&buf, num_bytes) {
                                        requests.push((Vec::from(nonce), src_addr));
                                        merkle.push_leaf(nonce);
                                    } else {
                                        num_bad_requests += 1;
                                        info!("Invalid request ({} bytes) from {} (#{} in batch, resp #{})", num_bytes, src_addr, i, resp_start + i as usize);
                                    }
                                },
                                Err(e) => match e.kind() {
                                    ErrorKind::WouldBlock => {
                                        done = true;
                                        break;
                                    },
                                    _ => panic!("recv_from failed with {:?}", e)
                                }
                            };
                        }

                        if requests.is_empty() {
                            break 'process_batch
                        }

                        let root = merkle.compute_root();
                        let srep = make_srep(&mut ephemeral_key, &root);

                        for (i, &(ref nonce, ref src_addr)) in requests.iter().enumerate() {
                            let paths = merkle.get_paths(i);

                            let resp = make_response(&srep, cert_bytes, &paths, i as u32);
                            let resp_bytes = resp.encode().unwrap();

                            let bytes_sent = socket.send_to(&resp_bytes, &src_addr).expect("send_to failed");
                            let num_responses = response_counter.fetch_add(1, Ordering::SeqCst);

                            info!("Responded {} bytes to {} for '{}..' (#{} in batch, resp #{})", bytes_sent, src_addr, hex::encode(&nonce[0..4]), i, num_responses);
                        }
                        if done {
                            break 'process_batch
                        }
                    }

                }

                STATUS => {
                    info!("responses {}, invalid requests {}", response_counter.load(Ordering::SeqCst), num_bad_requests);
                    timer.set_timeout(status_duration, ());
                }

                _ => unreachable!()
            }
        }
    }
}

pub fn main() {
    use log::Level;

    simple_logger::init_with_level(Level::Info).unwrap();

    info!("Roughenough server v{} starting", VERSION);

    let mut args = env::args();
    if args.len() != 2 {
        error!("Usage: server /path/to/config.file");
        process::exit(1);
    }

    let (addr, key_seed, batch_size) = load_config(&args.nth(1).unwrap());
    let (mut ephemeral_key, cert_bytes) = make_key_and_cert(&key_seed);

    info!("Server listening on {}", addr);

    let response_counter = Arc::new(AtomicUsize::new(0));

    if env::var("BENCH").is_ok()  {
        log::set_max_level(log::LevelFilter::Warn);
        let response_counter = response_counter.clone();

        thread::spawn(move || {
            loop {
                let old = time::get_time().sec;
                let old_reqs = response_counter.load(Ordering::SeqCst);

                thread::sleep(Duration::from_secs(1));

                let new = time::get_time().sec;
                let new_reqs = response_counter.load(Ordering::SeqCst);

                warn!("Processing at {:?} reqs/sec", (new_reqs - old_reqs) / (new - old) as usize);
            }
        });
    }

    polling_loop(&addr, &mut ephemeral_key, &cert_bytes, batch_size, response_counter.clone());

    info!("Done.");
    process::exit(0);
}
<|endoftext|>


<commit_before>#![warn(clippy::all, rust_2018_idioms)]

use cargo_registry::{boot, App, Env};
use std::{
    fs::File,
    sync::{mpsc::channel, Arc, Mutex},
    thread,
    time::Duration,
};

use civet::Server as CivetServer;
use conduit_hyper::Service;
use futures_util::future::FutureExt;
use reqwest::blocking::Client;

#[allow(clippy::large_enum_variant)]
enum Server {
    Civet(CivetServer),
    Hyper(tokio::runtime::Runtime, tokio::task::JoinHandle<()>),
}

use Server::*;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    env_logger::init();

    let config = cargo_registry::Config::default();
    let client = Client::new();

    let app = App::new(config.clone(), Some(client));
    let app = cargo_registry::build_handler(Arc::new(app));

    // On every server restart, ensure the categories available in the database match
    // the information in *src/categories.toml*.
    let categories_toml = include_str!("../boot/categories.toml");
    boot::categories::sync(categories_toml).unwrap();

    let heroku = dotenv::var("HEROKU").is_ok();
    let fastboot = dotenv::var("USE_FASTBOOT").is_ok();

    let port = if heroku {
        8888
    } else {
        dotenv::var("PORT")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(8888)
    };
    let threads = dotenv::var("SERVER_THREADS")
        .map(|s| s.parse().expect("SERVER_THREADS was not a valid number"))
        .unwrap_or_else(|_| {
            if config.env == Env::Development {
                5
            } else {
                50
            }
        });

    let server = if dotenv::var("USE_HYPER").is_ok() {
        use tokio::io::AsyncWriteExt;
        use tokio::signal::unix::{signal, SignalKind};

        println!("Booting with a hyper based server");

        let mut rt = tokio::runtime::Builder::new()
            .threaded_scheduler()
            .enable_all()
            .core_threads(4)
            .max_threads(threads as usize)
            .build()
            .unwrap();

        let handler = Arc::new(conduit_hyper::BlockingHandler::new(app));
        let make_service =
            hyper::service::make_service_fn(move |socket: &hyper::server::conn::AddrStream| {
                let addr = socket.remote_addr();
                let handler = handler.clone();
                async move { Service::from_blocking(handler, addr) }
            });

        let addr = ([127, 0, 0, 1], port).into();
        let server = rt.block_on(async { hyper::Server::bind(&addr).serve(make_service) });

        let mut sig_int = rt.block_on(async { signal(SignalKind::interrupt()) })?;
        let mut sig_term = rt.block_on(async { signal(SignalKind::terminate()) })?;

        let server = server.with_graceful_shutdown(async move {
            // Wait for either signal
            futures_util::select! {
                _ = sig_int.recv().fuse() => (),
                _ = sig_term.recv().fuse() => (),
            };
            let mut stdout = tokio::io::stdout();
            stdout.write_all(b"Starting graceful shutdown\n").await.ok();
        });

        let server = rt.spawn(async { server.await.unwrap() });
        Hyper(rt, server)
    } else {
        println!("Booting with a civet based server");
        let mut cfg = civet::Config::new();
        cfg.port(port).threads(threads).keep_alive(true);
        Civet(CivetServer::start(cfg, app).unwrap())
    };

    println!("listening on port {}", port);

    // Give tokio a chance to spawn the first worker thread
    thread::sleep(Duration::from_millis(10));

    // Creating this file tells heroku to tell nginx that the application is ready
    // to receive traffic.
    if heroku {
        let path = if fastboot {
            "/tmp/backend-initialized"
        } else {
            "/tmp/app-initialized"
        };
        println!("Writing to {}", path);
        File::create(path).unwrap();
    }

    // Block the main thread until the server has shutdown
    match server {
        Hyper(mut rt, server) => {
            rt.block_on(async { server.await.unwrap() });
        }
        Civet(server) => {
            let (tx, rx) = channel::<()>();
            ctrlc_handler(move || tx.send(()).unwrap_or(()));
            rx.recv().unwrap();
            drop(server);
        }
    }

    println!("Server has gracefully shutdown!");
    Ok(())
}

fn ctrlc_handler<F>(f: F)
where
    F: FnOnce() + Send + 'static,
{
    let call_once = Mutex::new(Some(f));

    ctrlc::set_handler(move || {
        if let Some(f) = call_once.lock().unwrap().take() {
            println!("Starting graceful shutdown");
            f();
        } else {
            println!("Already sent signal to start graceful shutdown");
        }
    })
    .unwrap();
}
<commit_msg>Account for core_threads when setting up background threads<commit_after>#![warn(clippy::all, rust_2018_idioms)]

use cargo_registry::{boot, App, Env};
use std::{
    fs::File,
    sync::{mpsc::channel, Arc, Mutex},
    thread,
    time::Duration,
};

use civet::Server as CivetServer;
use conduit_hyper::Service;
use futures_util::future::FutureExt;
use reqwest::blocking::Client;

const CORE_THREADS: usize = 4;

#[allow(clippy::large_enum_variant)]
enum Server {
    Civet(CivetServer),
    Hyper(tokio::runtime::Runtime, tokio::task::JoinHandle<()>),
}

use Server::*;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize logging
    env_logger::init();

    let config = cargo_registry::Config::default();
    let client = Client::new();

    let app = App::new(config.clone(), Some(client));
    let app = cargo_registry::build_handler(Arc::new(app));

    // On every server restart, ensure the categories available in the database match
    // the information in *src/categories.toml*.
    let categories_toml = include_str!("../boot/categories.toml");
    boot::categories::sync(categories_toml).unwrap();

    let heroku = dotenv::var("HEROKU").is_ok();
    let fastboot = dotenv::var("USE_FASTBOOT").is_ok();

    let port = if heroku {
        8888
    } else {
        dotenv::var("PORT")
            .ok()
            .and_then(|s| s.parse().ok())
            .unwrap_or(8888)
    };
    let threads = dotenv::var("SERVER_THREADS")
        .map(|s| s.parse().expect("SERVER_THREADS was not a valid number"))
        .unwrap_or_else(|_| {
            if config.env == Env::Development {
                5
            } else {
                50
            }
        });

    let server = if dotenv::var("USE_HYPER").is_ok() {
        use tokio::io::AsyncWriteExt;
        use tokio::signal::unix::{signal, SignalKind};

        println!("Booting with a hyper based server");

        let mut rt = tokio::runtime::Builder::new()
            .threaded_scheduler()
            .enable_all()
            .core_threads(CORE_THREADS)
            .max_threads(threads as usize + CORE_THREADS)
            .build()
            .unwrap();

        let handler = Arc::new(conduit_hyper::BlockingHandler::new(app));
        let make_service =
            hyper::service::make_service_fn(move |socket: &hyper::server::conn::AddrStream| {
                let addr = socket.remote_addr();
                let handler = handler.clone();
                async move { Service::from_blocking(handler, addr) }
            });

        let addr = ([127, 0, 0, 1], port).into();
        let server = rt.block_on(async { hyper::Server::bind(&addr).serve(make_service) });

        let mut sig_int = rt.block_on(async { signal(SignalKind::interrupt()) })?;
        let mut sig_term = rt.block_on(async { signal(SignalKind::terminate()) })?;

        let server = server.with_graceful_shutdown(async move {
            // Wait for either signal
            futures_util::select! {
                _ = sig_int.recv().fuse() => (),
                _ = sig_term.recv().fuse() => (),
            };
            let mut stdout = tokio::io::stdout();
            stdout.write_all(b"Starting graceful shutdown\n").await.ok();
        });

        let server = rt.spawn(async { server.await.unwrap() });
        Hyper(rt, server)
    } else {
        println!("Booting with a civet based server");
        let mut cfg = civet::Config::new();
        cfg.port(port).threads(threads).keep_alive(true);
        Civet(CivetServer::start(cfg, app).unwrap())
    };

    println!("listening on port {}", port);

    // Give tokio a chance to spawn the first worker thread
    thread::sleep(Duration::from_millis(10));

    // Creating this file tells heroku to tell nginx that the application is ready
    // to receive traffic.
    if heroku {
        let path = if fastboot {
            "/tmp/backend-initialized"
        } else {
            "/tmp/app-initialized"
        };
        println!("Writing to {}", path);
        File::create(path).unwrap();
    }

    // Block the main thread until the server has shutdown
    match server {
        Hyper(mut rt, server) => {
            rt.block_on(async { server.await.unwrap() });
        }
        Civet(server) => {
            let (tx, rx) = channel::<()>();
            ctrlc_handler(move || tx.send(()).unwrap_or(()));
            rx.recv().unwrap();
            drop(server);
        }
    }

    println!("Server has gracefully shutdown!");
    Ok(())
}

fn ctrlc_handler<F>(f: F)
where
    F: FnOnce() + Send + 'static,
{
    let call_once = Mutex::new(Some(f));

    ctrlc::set_handler(move || {
        if let Some(f) = call_once.lock().unwrap().take() {
            println!("Starting graceful shutdown");
            f();
        } else {
            println!("Already sent signal to start graceful shutdown");
        }
    })
    .unwrap();
}
<|endoftext|>


<commit_before>//!
//! The API for communicating with devices.
//!
//! This API is provided as Traits to be implemented:
//!
//! - by the low-level layers of the FoxBox, including the adapters;
//! - by test suites and tools that need to simulate connected devices.
//!
//! In turn, this API is used to implement:
//!
//! - the public-facing REST and WebSocket API;
//! - the rules API (ThinkerBell).
//!
//!

use devices::*;
use requests::*;
use values::Value;

/// An error produced by one of the APIs in this module.
pub enum Error {
    /// There is no such node connected to the Foxbox, even indirectly.
    NoSuchNode(NodeId),

    /// There is no such service connected to the Foxbox, even indirectly.
    NoSuchService(ServiceId),

    /// Attempting to set a value with the wrong type
    TypeError,
}

/// An event during watching.
pub enum WatchEvent {
    /// A new value is available.
    Value{from: ServiceId, value: Value},

    /// The set of devices being watched has changed, typically either
    /// because a tag was edited or because a device was connected or
    /// disconnected. Value `SetChanged(n)` means that the set now
    /// holds `n` input services.
    SetChanged(usize),
}

/// The public API.
pub trait API {
    /// Get a list of nodes matching some conditions
    ///
    /// # REST API
    ///
    /// `GET /api/v1/node/list`
    fn get_nodes(&NodeRequest) -> Vec<Node>;

    /// Add a tag to an existing node.
    ///
    /// Tags can be used to locate nodes.
    ///
    /// # REST API
    ///
    /// `PUT /api/v1/node/tag/$NodeId`
    fn put_node_tag(&NodeId, String) -> Result<(), Error>;

    /// Add a tag to an existing node.
    ///
    /// Tags can be used to locate nodes.
    ///
    /// # REST API
    ///
    /// `DELETE /api/v1/node/tag/$NodeId`
    fn delete_node_tag(&NodeId, String) -> Result<(), Error>;

    /// Get a list of inputs matching some conditions
    ///
    /// # REST API
    ///
    /// `GET /api/v1/service/list`
    fn get_input_services(&InputRequest) -> Vec<Service<Input>>;
    fn get_output_services(&OutputRequest) -> Vec<Service<Output>>;

    /// Add a tag to an existing service.
    ///
    /// Tags can be used to locate service.
    ///
    /// # REST API
    ///
    /// `PUT /api/v1/service/tag/$ServiceId`
    fn put_service_tag(&ServiceId, String) -> Result<(), Error>;

    /// Add a tag to an existing service.
    ///
    /// Tags can be used to locate services.
    ///
    /// # REST API
    ///
    /// `DELETE /api/v1/service/tag/$ServiceId`
    fn delete_service_tag(&ServiceId, String) -> Result<(), Error>;

    /// Read one value from an input enpoint
    ///
    /// # REST API
    ///
    /// GET /api/v1/service/value/$ServiceId
    fn get_service_value(&Service<Input>) -> Result<Value, Error>;

    /// Send one value to an output enpoint
    ///
    /// # REST API
    ///
    /// `PUT /api/v1/service/value/$ServiceId`
    fn put_service_value(&Service<Output>, Value) -> Result<(), Error>;

    /// Watch for any change
    ///
    /// # WebSocket API
    ///
    /// `/api/v1/service/watch/$ServiceId`
    fn register_service_watch<F>(WatchOptions, cb: F)
                                 -> Result<Self::WatchGuard, Error>
        where F: Fn(WatchEvent) + Send;

    /// A value that causes a disconnection once it is dropped.
    type WatchGuard;
}

/// Options for watching changes in one or more services.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct WatchOptions {
    /// The set of inputs to watch. Note that the actual inputs in the
    /// set may change over time.
    source: InputRequest,
}

impl WatchOptions {
    pub fn new() -> Self {
        WatchOptions {
            source: InputRequest::new(),
        }
    }

    /// Restrict to input services in a given set. Note that if all
    /// the inputs do not have the same `ServiceKind`, the call to
    /// `register_service_watch` will fail.
    ///
    /// Also note that the actual input services that are part of the
    /// set may change with time, for instance if devices are added
    /// ore removed.  as it changes.
    pub fn with_inputs(self, req: InputRequest) -> Self {
        WatchOptions {
            source: self.source.and(req),
            ..self
        }
    }
}
<commit_msg>[fxbox/taxonomy] Refined WatchEvent on inputs being added/removed<commit_after>//!
//! The API for communicating with devices.
//!
//! This API is provided as Traits to be implemented:
//!
//! - by the low-level layers of the FoxBox, including the adapters;
//! - by test suites and tools that need to simulate connected devices.
//!
//! In turn, this API is used to implement:
//!
//! - the public-facing REST and WebSocket API;
//! - the rules API (ThinkerBell).
//!
//!

use devices::*;
use requests::*;
use values::Value;

/// An error produced by one of the APIs in this module.
pub enum Error {
    /// There is no such node connected to the Foxbox, even indirectly.
    NoSuchNode(NodeId),

    /// There is no such service connected to the Foxbox, even indirectly.
    NoSuchService(ServiceId),

    /// Attempting to set a value with the wrong type
    TypeError,
}

/// An event during watching.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub enum WatchEvent {
    /// A new value is available.
    Value {
        /// The service that sent the value.
        from: ServiceId,

        /// The actual value.
        value: Value
    },

    /// The set of devices being watched has changed, typically either
    /// because a tag was edited or because a device was
    /// removed. Payload is the id of the device that was removed.
    InputRemoved(ServiceId),

    /// The set of devices being watched has changed, typically either
    /// because a tag was edited or because a device was
    /// added. Payload is the id of the device that was added.
    InputAdded(ServiceId),
}

/// The public API.
pub trait API {
    /// Get a list of nodes matching some conditions
    ///
    /// # REST API
    ///
    /// `GET /api/v1/node/list`
    fn get_nodes(&NodeRequest) -> Vec<Node>;

    /// Add a tag to an existing node.
    ///
    /// Tags can be used to locate nodes.
    ///
    /// # REST API
    ///
    /// `PUT /api/v1/node/tag/$NodeId`
    fn put_node_tag(&NodeId, String) -> Result<(), Error>;

    /// Add a tag to an existing node.
    ///
    /// Tags can be used to locate nodes.
    ///
    /// # REST API
    ///
    /// `DELETE /api/v1/node/tag/$NodeId`
    fn delete_node_tag(&NodeId, String) -> Result<(), Error>;

    /// Get a list of inputs matching some conditions
    ///
    /// # REST API
    ///
    /// `GET /api/v1/service/list`
    fn get_input_services(&InputRequest) -> Vec<Service<Input>>;
    fn get_output_services(&OutputRequest) -> Vec<Service<Output>>;

    /// Add a tag to an existing service.
    ///
    /// Tags can be used to locate service.
    ///
    /// # REST API
    ///
    /// `PUT /api/v1/service/tag/$ServiceId`
    fn put_service_tag(&ServiceId, String) -> Result<(), Error>;

    /// Add a tag to an existing service.
    ///
    /// Tags can be used to locate services.
    ///
    /// # REST API
    ///
    /// `DELETE /api/v1/service/tag/$ServiceId`
    fn delete_service_tag(&ServiceId, String) -> Result<(), Error>;

    /// Read one value from an input enpoint
    ///
    /// # REST API
    ///
    /// GET /api/v1/service/value/$ServiceId
    fn get_service_value(&Service<Input>) -> Result<Value, Error>;

    /// Send one value to an output enpoint
    ///
    /// # REST API
    ///
    /// `PUT /api/v1/service/value/$ServiceId`
    fn put_service_value(&Service<Output>, Value) -> Result<(), Error>;

    /// Watch for any change
    ///
    /// # WebSocket API
    ///
    /// `/api/v1/service/watch/$ServiceId`
    fn register_service_watch<F>(WatchOptions, cb: F)
                                 -> Result<Self::WatchGuard, Error>
        where F: Fn(WatchEvent) + Send;

    /// A value that causes a disconnection once it is dropped.
    type WatchGuard;
}

/// Options for watching changes in one or more services.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct WatchOptions {
    /// The set of inputs to watch. Note that the actual inputs in the
    /// set may change over time.
    source: InputRequest,
}

impl WatchOptions {
    pub fn new() -> Self {
        WatchOptions {
            source: InputRequest::new(),
        }
    }

    /// Restrict to input services in a given set. Note that if all
    /// the inputs do not have the same `ServiceKind`, the call to
    /// `register_service_watch` will fail.
    ///
    /// Also note that the actual input services that are part of the
    /// set may change with time, for instance if devices are added
    /// ore removed.  as it changes.
    pub fn with_inputs(self, req: InputRequest) -> Self {
        WatchOptions {
            source: self.source.and(req),
            ..self
        }
    }
}
<|endoftext|>


<commit_before>use libc::c_long;

use std::ffi::CString;
use std::str;

use nix::errno::Errno::*;
use nix::Error::Sys;
use nix::mqueue::{mq_open, mq_close, mq_send, mq_receive, mq_getattr, mq_setattr, mq_unlink, mq_set_nonblock, mq_remove_nonblock};
use nix::mqueue::{MqAttr, MQ_OFlag};
use nix::sys::stat::Mode;

#[test]
fn test_mq_send_and_receive() {
    const MSG_SIZE: c_long =  32;
    let attr =  MqAttr::new(0, 10, MSG_SIZE, 0);
    let mq_name= &CString::new(b"/a_nix_test_queue".as_ref()).unwrap();

    let mqd0 = mq_open(mq_name, MQ_OFlag::O_CREAT | MQ_OFlag::O_WRONLY,
                       Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH,
                       Some(&attr)).unwrap();
    let msg_to_send = "msg_1";
    mq_send(mqd0, msg_to_send.as_bytes(), 1).unwrap();

    let mqd1 = mq_open(mq_name, MQ_OFlag::O_CREAT | MQ_OFlag::O_RDONLY,
                       Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH,
                       Some(&attr)).unwrap();
    let mut buf = [0u8; 32];
    let mut prio = 0u32;
    let len = mq_receive(mqd1, &mut buf, &mut prio).unwrap();
    assert!(prio == 1);

    mq_close(mqd1).unwrap();
    mq_close(mqd0).unwrap();
    assert_eq!(msg_to_send, str::from_utf8(&buf[0..len]).unwrap());
}

#[test]
fn test_mq_getattr() {
    const MSG_SIZE: c_long =  32;
    let initial_attr =  MqAttr::new(0, 10, MSG_SIZE, 0);
    let mq_name = &CString::new("/attr_test_get_attr".as_bytes().as_ref()).unwrap();
    let mqd = mq_open(mq_name, MQ_OFlag::O_CREAT | MQ_OFlag::O_WRONLY, Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH, Some(&initial_attr)).unwrap();
    let read_attr = mq_getattr(mqd);
    assert!(read_attr.unwrap() == initial_attr);
    mq_close(mqd).unwrap();
}

// FIXME: Fix failures for mips in QEMU
#[test]
#[cfg_attr(any(target_arch = "mips", target_arch = "mips64"), ignore)]
fn test_mq_setattr() {
    const MSG_SIZE: c_long =  32;
    let initial_attr =  MqAttr::new(0, 10, MSG_SIZE, 0);
    let mq_name = &CString::new("/attr_test_get_attr".as_bytes().as_ref()).unwrap();
    let mqd = mq_open(mq_name, MQ_OFlag::O_CREAT | MQ_OFlag::O_WRONLY, Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH, Some(&initial_attr)).unwrap();

    let new_attr =  MqAttr::new(0, 20, MSG_SIZE * 2, 100);
    let old_attr = mq_setattr(mqd, &new_attr);
    assert!(old_attr.unwrap() == initial_attr);

    let new_attr_get = mq_getattr(mqd);
    // The following tests make sense. No changes here because according to the Linux man page only
    // O_NONBLOCK can be set (see tests below)
    assert!(new_attr_get.unwrap() != new_attr);

    let new_attr_non_blocking =  MqAttr::new(MQ_OFlag::O_NONBLOCK.bits() as c_long, 10, MSG_SIZE, 0);
    mq_setattr(mqd, &new_attr_non_blocking).unwrap();
    let new_attr_get = mq_getattr(mqd);

    // now the O_NONBLOCK flag has been set
    assert!(new_attr_get.unwrap() != initial_attr);
    assert!(new_attr_get.unwrap() == new_attr_non_blocking);
    mq_close(mqd).unwrap();
}

// FIXME: Fix failures for mips in QEMU
#[test]
#[cfg_attr(any(target_arch = "mips", target_arch = "mips64"), ignore)]
fn test_mq_set_nonblocking() {
    const MSG_SIZE: c_long =  32;
    let initial_attr =  MqAttr::new(0, 10, MSG_SIZE, 0);
    let mq_name = &CString::new("/attr_test_get_attr".as_bytes().as_ref()).unwrap();
    let mqd = mq_open(mq_name, MQ_OFlag::O_CREAT | MQ_OFlag::O_WRONLY, Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH, Some(&initial_attr)).unwrap();
    mq_set_nonblock(mqd).unwrap();
    let new_attr = mq_getattr(mqd);
    assert!(new_attr.unwrap().flags() == MQ_OFlag::O_NONBLOCK.bits() as c_long);
    mq_remove_nonblock(mqd).unwrap();
    let new_attr = mq_getattr(mqd);
    assert!(new_attr.unwrap().flags() == 0);
    mq_close(mqd).unwrap();
}

#[test]
fn test_mq_unlink() {
    const MSG_SIZE: c_long =  32;
    let initial_attr =  MqAttr::new(0, 10, MSG_SIZE, 0);
    let mq_name_opened = &CString::new("/mq_unlink_test".as_bytes().as_ref()).unwrap();
    let mq_name_not_opened = &CString::new("/mq_unlink_test".as_bytes().as_ref()).unwrap();
    let mqd = mq_open(mq_name_opened, MQ_OFlag::O_CREAT | MQ_OFlag::O_WRONLY, Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH, Some(&initial_attr)).unwrap();

    let res_unlink = mq_unlink(mq_name_opened);
    assert!(res_unlink == Ok(()) );

    let res_unlink_not_opened = mq_unlink(mq_name_not_opened);
    assert!(res_unlink_not_opened == Err(Sys(ENOENT)) );

    mq_close(mqd).unwrap();
    let res_unlink_after_close = mq_unlink(mq_name_opened);
    assert!(res_unlink_after_close == Err(Sys(ENOENT)) );

}
<commit_msg>Use byte strings instead of String.as_bytes()<commit_after>use libc::c_long;

use std::ffi::CString;
use std::str;

use nix::errno::Errno::*;
use nix::Error::Sys;
use nix::mqueue::{mq_open, mq_close, mq_send, mq_receive, mq_getattr, mq_setattr, mq_unlink, mq_set_nonblock, mq_remove_nonblock};
use nix::mqueue::{MqAttr, MQ_OFlag};
use nix::sys::stat::Mode;

#[test]
fn test_mq_send_and_receive() {
    const MSG_SIZE: c_long =  32;
    let attr =  MqAttr::new(0, 10, MSG_SIZE, 0);
    let mq_name= &CString::new(b"/a_nix_test_queue".as_ref()).unwrap();

    let mqd0 = mq_open(mq_name, MQ_OFlag::O_CREAT | MQ_OFlag::O_WRONLY,
                       Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH,
                       Some(&attr)).unwrap();
    let msg_to_send = "msg_1";
    mq_send(mqd0, msg_to_send.as_bytes(), 1).unwrap();

    let mqd1 = mq_open(mq_name, MQ_OFlag::O_CREAT | MQ_OFlag::O_RDONLY,
                       Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH,
                       Some(&attr)).unwrap();
    let mut buf = [0u8; 32];
    let mut prio = 0u32;
    let len = mq_receive(mqd1, &mut buf, &mut prio).unwrap();
    assert!(prio == 1);

    mq_close(mqd1).unwrap();
    mq_close(mqd0).unwrap();
    assert_eq!(msg_to_send, str::from_utf8(&buf[0..len]).unwrap());
}

#[test]
fn test_mq_getattr() {
    const MSG_SIZE: c_long =  32;
    let initial_attr =  MqAttr::new(0, 10, MSG_SIZE, 0);
    let mq_name = &CString::new(b"/attr_test_get_attr".as_ref()).unwrap();
    let mqd = mq_open(mq_name, MQ_OFlag::O_CREAT | MQ_OFlag::O_WRONLY, Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH, Some(&initial_attr)).unwrap();
    let read_attr = mq_getattr(mqd);
    assert!(read_attr.unwrap() == initial_attr);
    mq_close(mqd).unwrap();
}

// FIXME: Fix failures for mips in QEMU
#[test]
#[cfg_attr(any(target_arch = "mips", target_arch = "mips64"), ignore)]
fn test_mq_setattr() {
    const MSG_SIZE: c_long =  32;
    let initial_attr =  MqAttr::new(0, 10, MSG_SIZE, 0);
    let mq_name = &CString::new(b"/attr_test_get_attr".as_ref()).unwrap();
    let mqd = mq_open(mq_name, MQ_OFlag::O_CREAT | MQ_OFlag::O_WRONLY, Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH, Some(&initial_attr)).unwrap();

    let new_attr =  MqAttr::new(0, 20, MSG_SIZE * 2, 100);
    let old_attr = mq_setattr(mqd, &new_attr);
    assert!(old_attr.unwrap() == initial_attr);

    let new_attr_get = mq_getattr(mqd);
    // The following tests make sense. No changes here because according to the Linux man page only
    // O_NONBLOCK can be set (see tests below)
    assert!(new_attr_get.unwrap() != new_attr);

    let new_attr_non_blocking =  MqAttr::new(MQ_OFlag::O_NONBLOCK.bits() as c_long, 10, MSG_SIZE, 0);
    mq_setattr(mqd, &new_attr_non_blocking).unwrap();
    let new_attr_get = mq_getattr(mqd);

    // now the O_NONBLOCK flag has been set
    assert!(new_attr_get.unwrap() != initial_attr);
    assert!(new_attr_get.unwrap() == new_attr_non_blocking);
    mq_close(mqd).unwrap();
}

// FIXME: Fix failures for mips in QEMU
#[test]
#[cfg_attr(any(target_arch = "mips", target_arch = "mips64"), ignore)]
fn test_mq_set_nonblocking() {
    const MSG_SIZE: c_long =  32;
    let initial_attr =  MqAttr::new(0, 10, MSG_SIZE, 0);
    let mq_name = &CString::new(b"/attr_test_get_attr".as_ref()).unwrap();
    let mqd = mq_open(mq_name, MQ_OFlag::O_CREAT | MQ_OFlag::O_WRONLY, Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH, Some(&initial_attr)).unwrap();
    mq_set_nonblock(mqd).unwrap();
    let new_attr = mq_getattr(mqd);
    assert!(new_attr.unwrap().flags() == MQ_OFlag::O_NONBLOCK.bits() as c_long);
    mq_remove_nonblock(mqd).unwrap();
    let new_attr = mq_getattr(mqd);
    assert!(new_attr.unwrap().flags() == 0);
    mq_close(mqd).unwrap();
}

#[test]
fn test_mq_unlink() {
    const MSG_SIZE: c_long =  32;
    let initial_attr =  MqAttr::new(0, 10, MSG_SIZE, 0);
    let mq_name_opened = &CString::new(b"/mq_unlink_test".as_ref()).unwrap();
    let mq_name_not_opened = &CString::new(b"/mq_unlink_test".as_ref()).unwrap();
    let mqd = mq_open(mq_name_opened, MQ_OFlag::O_CREAT | MQ_OFlag::O_WRONLY, Mode::S_IWUSR | Mode::S_IRUSR | Mode::S_IRGRP | Mode::S_IROTH, Some(&initial_attr)).unwrap();

    let res_unlink = mq_unlink(mq_name_opened);
    assert!(res_unlink == Ok(()) );

    let res_unlink_not_opened = mq_unlink(mq_name_not_opened);
    assert!(res_unlink_not_opened == Err(Sys(ENOENT)) );

    mq_close(mqd).unwrap();
    let res_unlink_after_close = mq_unlink(mq_name_opened);
    assert!(res_unlink_after_close == Err(Sys(ENOENT)) );

}
<|endoftext|>


<commit_before>use std::default::Default;

use r2d2;

#[test]
fn default_ok() {
    let config: r2d2::Config = Default::default();
    assert!(config.validate().is_ok());
}

#[test]
fn test_zero_pool_size_err() {
    let config = r2d2::Config {
        pool_size: 0,
        ..Default::default()
    };
    assert_eq!(Err(r2d2::ZeroPoolSize), config.validate());
}

#[test]
fn test_zero_helper_tasks_err() {
    let config = r2d2::Config {
        helper_tasks: 0,
        ..Default::default()
    };
    assert_eq!(Err(r2d2::ZeroHelperTasks), config.validate());
}
<commit_msg>Fix tests<commit_after>use std::default::Default;

use r2d2;

#[test]
fn default_ok() {
    let config: r2d2::Config = Default::default();
    assert!(config.validate().is_ok());
}

#[test]
fn test_zero_pool_size_err() {
    let config = r2d2::Config {
        pool_size: 0,
        ..Default::default()
    };
    assert_eq!(Err(r2d2::ConfigError::ZeroPoolSize), config.validate());
}

#[test]
fn test_zero_helper_tasks_err() {
    let config = r2d2::Config {
        helper_tasks: 0,
        ..Default::default()
    };
    assert_eq!(Err(r2d2::ConfigError::ZeroHelperTasks), config.validate());
}
<|endoftext|>


<commit_before>use std::net::TcpStream;
use std::collections::HashMap;
use std::io;
use std::io::{Read,Write};
use std::ops::Deref;
use libc;

use unix_socket::UnixStream;
use rustc_serialize::hex::ToHex;
use dbus_serialize::types::{Value,BasicValue};
use dbus_serialize::decoder::DBusDecoder;

use message;
use message::{Message,HeaderFieldName,MessageBuf};
use demarshal::{demarshal,DemarshalError};
use marshal::Marshal;

trait StreamSocket : Read + Write { }

impl StreamSocket for TcpStream { }
impl StreamSocket for UnixStream { }

enum Socket {
    Tcp(TcpStream),
    Uds(UnixStream)
}

pub struct Connection {
    sock: Socket,
    next_serial: u32,
    queue: Vec<Message>,
}

#[derive(Debug)]
pub enum Error {
    Disconnected,
    IOError(io::Error),
    DemarshalError(DemarshalError),
    BadData,
    AuthFailed,
}

impl From<io::Error> for Error {
    fn from(x: io::Error) -> Self {
        Error::IOError(x)
    }
}

impl From<DemarshalError> for Error {
    fn from(x: DemarshalError) -> Self {
        Error::DemarshalError(x)
    }
}

fn read_exactly(sock: &mut StreamSocket, buf: &mut Vec<u8>, len: usize) -> Result<(),Error> {
    buf.truncate(0);
    buf.reserve(len);
    if try!(sock.take(len as u64).read_to_end(buf)) != len {
        return Err(Error::Disconnected);
    }
    Ok(())
}

fn read_line(sock: &mut StreamSocket) -> Result<String,Error> {
    let mut line = "".to_string();
    let mut last = '\0';

    loop {
        let mut buf = vec![0];
        match sock.read(&mut buf) {
            Ok(x) if x > 0 => (),
            _ => return Err(Error::Disconnected)
        };
        let chr = buf[0] as char;
        line.push(chr);
        if chr == '\n' && last == '\r' {
            break;
        }
        last = chr;
    }
    Ok(line)
}

impl Connection {
    fn get_sock(&mut self) -> &mut StreamSocket {
        match self.sock {
            Socket::Tcp(ref mut x) => x,
            Socket::Uds(ref mut x) => x
        }
    }

    fn auth_anonymous(&mut self) -> Result<(),Error> {
        let sock = self.get_sock();

        // Authenticate to the daemon
        let buf = vec![0];
        try!(sock.write_all(&buf));

        try!(sock.write_all(b"AUTH ANONYMOUS 6c69626462757320312e382e3132\r\n"));

        // Read response
        let resp = try!(read_line(sock));
        if !resp.starts_with("OK ") {
            return Err(Error::AuthFailed);
        }

        // Ready for action
        try!(sock.write_all(b"BEGIN\r\n"));
        Ok(())
    }

    fn auth_external(&mut self) -> Result<(),Error> {
        let sock = self.get_sock();

        // Authenticate to the daemon
        let buf = vec![0];
        try!(sock.write_all(&buf));

        let uid = unsafe {
            libc::funcs::posix88::unistd::getuid()
        };
        let uid_str = uid.to_string();
        let uid_hex = uid_str.into_bytes().to_hex();
        let cmd = "AUTH EXTERNAL ".to_string() + &uid_hex + "\r\n";
        try!(sock.write_all(&cmd.into_bytes()));

        // Read response
        let resp = try!(read_line(sock));
        if !resp.starts_with("OK ") {
            return Err(Error::AuthFailed);
        }

        // Ready for action
        try!(sock.write_all(b"BEGIN\r\n"));
        Ok(())
    }

    fn say_hello(&mut self) -> Result<(),Error> {
        let mut msg = message::create_method_call("org.freedesktop.DBus",
                                                  "/org/freedesktop/DBus",
                                                  "org.freedesktop.DBus",
                                                  "Hello");
        try!(self.call_sync(&mut msg));
        // XXX: validate Hello reply
        Ok(())
    }

    pub fn connect_uds(addr: &str) -> Result<Connection,Error> {
        let sock = try!(UnixStream::connect(addr));
        let mut conn = Connection {
            sock: Socket::Uds(sock),
            queue: Vec::new(),
            next_serial: 1
        };

        try!(conn.auth_external());
        try!(conn.say_hello());
        Ok(conn)
    }

    pub fn connect_tcp(addr: &str) -> Result<Connection,Error> {
        let sock = try!(TcpStream::connect(addr));
        let mut conn = Connection {
            sock: Socket::Tcp(sock),
            queue: Vec::new(),
            next_serial: 1
        };

        try!(conn.auth_anonymous());
        try!(conn.say_hello());
        Ok(conn)
    }

    pub fn send(&mut self, mbuf: &mut MessageBuf) -> Result<u32, Error> {
        let mut msg = &mut mbuf.0;
        // A minimum header with no body is 16 bytes
        let mut len = msg.len() as u32;
        if len < 16 {
            return Err(Error::BadData);
        }

        // Get the current length from the message, which only include the length of the header.
        // That field should actually be the length of only the body, so update that now
        let old_len = message::get_length(msg);
        len -= old_len;
        let mut buf = Vec::new();
        len.dbus_encode(&mut buf);
        // Update the message with a correct serial number, as well
        let this_serial = self.next_serial;
        self.next_serial += 1;
        this_serial.dbus_encode(&mut buf);
        message::set_length(msg, &buf);

        let sock = self.get_sock();
        try!(sock.write_all(msg));
        Ok(this_serial)
    }

    pub fn call_sync(&mut self, mbuf: &mut MessageBuf) -> Result<Vec<Value>,Error> {
        // XXX: assert that this is a method call with reply
        let serial = try!(self.send(mbuf));
        // We need a local queue so that read_msg doesn't just give us
        // the same one over and over
        let mut queue = Vec::new();
        loop {
            let mut msg = try!(self.read_msg());
            match msg.headers.remove(&(HeaderFieldName::ReplySerial as u8)) {
                Some(Value::Variant(x)) => {
                    let obj : Value = *x.object;
                    let reply_serial : u32 = DBusDecoder::decode(obj).unwrap();
                    if reply_serial == serial {
                        // Move our queued messages into the Connection's queue
                        for _ in 0..queue.len() {
                            self.queue.push(queue.remove(0));
                        }
                        return Ok(msg.body);
                    }
                }
                _ => ()
            };
            queue.push(msg);
        }
    }

    pub fn read_msg(&mut self) -> Result<Message,Error> {
        match self.queue.get(0) {
            Some(_) => return Ok(self.queue.remove(0)),
            _ => ()
        };
        let mut buf = Vec::new();
        let sock = self.get_sock();

        // Read and demarshal the fixed portion of the header
        try!(read_exactly(sock, &mut buf, 12));
        let mut offset = 0;
        let mut sig = "(yyyyuu)".to_string();
        let header = match try!(demarshal(&mut buf, &mut offset, &mut sig)) {
            Value::Struct(x) => x,
            x => panic!("Demarshal didn't return what we asked for: {:?}", x)
        };

        let mut v = header.objects;
        let mut msg : Message = Default::default();
        let endian : u8 = DBusDecoder::decode(v.remove(0)).unwrap();
        if endian == 'B' as u8 {
            msg.big_endian = true;
        }
        msg.message_type = message::MessageType(DBusDecoder::decode(v.remove(0)).unwrap());
        msg.flags = DBusDecoder::decode::<u8>(v.remove(0)).unwrap();
        msg.version = DBusDecoder::decode::<u8>(v.remove(0)).unwrap();
        let body_len = DBusDecoder::decode::<u32>(v.remove(0)).unwrap();
        msg.serial = DBusDecoder::decode::<u32>(v.remove(0)).unwrap();

        // Read array length
        try!(read_exactly(sock, &mut buf, 4));
        // demarshal consumes the buf, so save a copy for when we demarshal the entire array
        let mut buf_copy = buf.clone();
        offset = 12;
        sig = "u".to_string();
        let data = demarshal(&mut buf, &mut offset, &mut sig).ok().unwrap();
        let arr_len = DBusDecoder::decode::<u32>(data).unwrap() as usize;

        // Make buf_copy big enough for the entire array, and fill it
        buf_copy.reserve(arr_len);
        if try!(sock.take(arr_len as u64).read_to_end(&mut buf_copy)) != arr_len {
            return Err(Error::Disconnected);
        };

        offset = 12;
        sig = "a(yv)".to_string();
        let header_fields = match try!(demarshal(&mut buf_copy, &mut offset, &mut sig)) {
            Value::Array(x) => x,
            x => panic!("Demarshal didn't return what we asked for: {:?}", x)
        };

        msg.headers = HashMap::new();
        for i in header_fields.objects {
            let mut st = match i {
                Value::Struct(x) => x,
                x => panic!("Demarshal didn't return what we asked for: {:?}", x)
            };
            let val = st.objects.remove(1);
            let code = DBusDecoder::decode::<u8>(st.objects.remove(0)).unwrap();
            msg.headers.insert(code, val);
        }

        // Read the padding, if any
        let trailing_pad = 8 - (offset % 8);
        if trailing_pad % 8 != 0 {
            try!(read_exactly(sock, &mut buf, trailing_pad));
        }

        // Finally, read the entire body
        if body_len > 0 {
            let v = match msg.headers.get(&(HeaderFieldName::Signature as u8)) {
                Some(&Value::Variant(ref x)) => x,
                _ => return Err(Error::DemarshalError(DemarshalError::BadSignature))
            };

            let sigval = match v.object.deref() {
                &Value::BasicValue(BasicValue::Signature(ref x)) => x,
                _ => return Err(Error::DemarshalError(DemarshalError::BadSignature))
            };

            let mut body = Vec::new();
            try!(read_exactly(sock, &mut body, body_len as usize));

            let mut sig = "(".to_string() + &sigval.0 + ")";
            offset = 0;
            let objs = match try!(demarshal(&mut body, &mut offset, &mut sig)) {
                Value::Struct(x) => x.objects,
                x => panic!("Didn't get a struct: {:?}", x)
            };
            for x in objs {
                msg.body.push(x);
            }
        }

        Ok(msg)
    }
}

#[test]
fn test_connect () {
    let mut conn = Connection::connect_uds("/var/run/dbus/system_bus_socket").unwrap();
    let mut msg = message::create_method_call("org.freedesktop.DBus", "/org/freedesktop/DBus",
                                          "org.freedesktop.DBus", "ListNames");
    let resp = conn.call_sync(&mut msg).unwrap();
    println!("ListNames: {:?}", resp);
}

#[cfg(tcp)]
#[test]
fn test_tcp () {
    let mut conn = Connection::connect_tcp("localhost:12345").unwrap();
    let mut msg = message::create_method_call("org.freedesktop.DBus", "/org/freedesktop/DBus",
                                          "org.freedesktop.DBus", "ListNames");
    conn.send(&mut msg).ok();
    let msg = conn.read_msg().unwrap();
    println!("{:?}", msg.body);
    //loop {
    //    conn.read_msg().unwrap();
    //}
}
<commit_msg>connection.rs: validate reply from Hello message<commit_after>use std::net::TcpStream;
use std::collections::HashMap;
use std::io;
use std::io::{Read,Write};
use std::ops::Deref;
use libc;

use unix_socket::UnixStream;
use rustc_serialize::hex::ToHex;
use dbus_serialize::types::{Value,BasicValue};
use dbus_serialize::decoder::DBusDecoder;

use message;
use message::{Message,HeaderFieldName,MessageBuf};
use demarshal::{demarshal,DemarshalError};
use marshal::Marshal;

trait StreamSocket : Read + Write { }

impl StreamSocket for TcpStream { }
impl StreamSocket for UnixStream { }

enum Socket {
    Tcp(TcpStream),
    Uds(UnixStream)
}

pub struct Connection {
    sock: Socket,
    next_serial: u32,
    queue: Vec<Message>,
}

#[derive(Debug)]
pub enum Error {
    Disconnected,
    IOError(io::Error),
    DemarshalError(DemarshalError),
    BadData,
    AuthFailed,
}

impl From<io::Error> for Error {
    fn from(x: io::Error) -> Self {
        Error::IOError(x)
    }
}

impl From<DemarshalError> for Error {
    fn from(x: DemarshalError) -> Self {
        Error::DemarshalError(x)
    }
}

fn read_exactly(sock: &mut StreamSocket, buf: &mut Vec<u8>, len: usize) -> Result<(),Error> {
    buf.truncate(0);
    buf.reserve(len);
    if try!(sock.take(len as u64).read_to_end(buf)) != len {
        return Err(Error::Disconnected);
    }
    Ok(())
}

fn read_line(sock: &mut StreamSocket) -> Result<String,Error> {
    let mut line = "".to_string();
    let mut last = '\0';

    loop {
        let mut buf = vec![0];
        match sock.read(&mut buf) {
            Ok(x) if x > 0 => (),
            _ => return Err(Error::Disconnected)
        };
        let chr = buf[0] as char;
        line.push(chr);
        if chr == '\n' && last == '\r' {
            break;
        }
        last = chr;
    }
    Ok(line)
}

impl Connection {
    fn get_sock(&mut self) -> &mut StreamSocket {
        match self.sock {
            Socket::Tcp(ref mut x) => x,
            Socket::Uds(ref mut x) => x
        }
    }

    fn auth_anonymous(&mut self) -> Result<(),Error> {
        let sock = self.get_sock();

        // Authenticate to the daemon
        let buf = vec![0];
        try!(sock.write_all(&buf));

        try!(sock.write_all(b"AUTH ANONYMOUS 6c69626462757320312e382e3132\r\n"));

        // Read response
        let resp = try!(read_line(sock));
        if !resp.starts_with("OK ") {
            return Err(Error::AuthFailed);
        }

        // Ready for action
        try!(sock.write_all(b"BEGIN\r\n"));
        Ok(())
    }

    fn auth_external(&mut self) -> Result<(),Error> {
        let sock = self.get_sock();

        // Authenticate to the daemon
        let buf = vec![0];
        try!(sock.write_all(&buf));

        let uid = unsafe {
            libc::funcs::posix88::unistd::getuid()
        };
        let uid_str = uid.to_string();
        let uid_hex = uid_str.into_bytes().to_hex();
        let cmd = "AUTH EXTERNAL ".to_string() + &uid_hex + "\r\n";
        try!(sock.write_all(&cmd.into_bytes()));

        // Read response
        let resp = try!(read_line(sock));
        if !resp.starts_with("OK ") {
            return Err(Error::AuthFailed);
        }

        // Ready for action
        try!(sock.write_all(b"BEGIN\r\n"));
        Ok(())
    }

    fn say_hello(&mut self) -> Result<(String),Error> {
        let mut msg = message::create_method_call("org.freedesktop.DBus",
                                                  "/org/freedesktop/DBus",
                                                  "org.freedesktop.DBus",
                                                  "Hello");
        match try!(self.call_sync(&mut msg)).get(0) {
            Some(&Value::BasicValue(BasicValue::String(ref x))) => Ok(x.to_string()),
            _ => Err(Error::BadData)
        }
    }

    pub fn connect_uds(addr: &str) -> Result<Connection,Error> {
        let sock = try!(UnixStream::connect(addr));
        let mut conn = Connection {
            sock: Socket::Uds(sock),
            queue: Vec::new(),
            next_serial: 1
        };

        try!(conn.auth_external());
        try!(conn.say_hello());
        Ok(conn)
    }

    pub fn connect_tcp(addr: &str) -> Result<Connection,Error> {
        let sock = try!(TcpStream::connect(addr));
        let mut conn = Connection {
            sock: Socket::Tcp(sock),
            queue: Vec::new(),
            next_serial: 1
        };

        try!(conn.auth_anonymous());
        try!(conn.say_hello());
        Ok(conn)
    }

    pub fn send(&mut self, mbuf: &mut MessageBuf) -> Result<u32, Error> {
        let mut msg = &mut mbuf.0;
        // A minimum header with no body is 16 bytes
        let mut len = msg.len() as u32;
        if len < 16 {
            return Err(Error::BadData);
        }

        // Get the current length from the message, which only include the length of the header.
        // That field should actually be the length of only the body, so update that now
        let old_len = message::get_length(msg);
        len -= old_len;
        let mut buf = Vec::new();
        len.dbus_encode(&mut buf);
        // Update the message with a correct serial number, as well
        let this_serial = self.next_serial;
        self.next_serial += 1;
        this_serial.dbus_encode(&mut buf);
        message::set_length(msg, &buf);

        let sock = self.get_sock();
        try!(sock.write_all(msg));
        Ok(this_serial)
    }

    pub fn call_sync(&mut self, mbuf: &mut MessageBuf) -> Result<Vec<Value>,Error> {
        // XXX: assert that this is a method call with reply
        let serial = try!(self.send(mbuf));
        // We need a local queue so that read_msg doesn't just give us
        // the same one over and over
        let mut queue = Vec::new();
        loop {
            let mut msg = try!(self.read_msg());
            match msg.headers.remove(&(HeaderFieldName::ReplySerial as u8)) {
                Some(Value::Variant(x)) => {
                    let obj : Value = *x.object;
                    let reply_serial : u32 = DBusDecoder::decode(obj).unwrap();
                    if reply_serial == serial {
                        // Move our queued messages into the Connection's queue
                        for _ in 0..queue.len() {
                            self.queue.push(queue.remove(0));
                        }
                        return Ok(msg.body);
                    }
                }
                _ => ()
            };
            queue.push(msg);
        }
    }

    pub fn read_msg(&mut self) -> Result<Message,Error> {
        match self.queue.get(0) {
            Some(_) => return Ok(self.queue.remove(0)),
            _ => ()
        };
        let mut buf = Vec::new();
        let sock = self.get_sock();

        // Read and demarshal the fixed portion of the header
        try!(read_exactly(sock, &mut buf, 12));
        let mut offset = 0;
        let mut sig = "(yyyyuu)".to_string();
        let header = match try!(demarshal(&mut buf, &mut offset, &mut sig)) {
            Value::Struct(x) => x,
            x => panic!("Demarshal didn't return what we asked for: {:?}", x)
        };

        let mut v = header.objects;
        let mut msg : Message = Default::default();
        let endian : u8 = DBusDecoder::decode(v.remove(0)).unwrap();
        if endian == 'B' as u8 {
            msg.big_endian = true;
        }
        msg.message_type = message::MessageType(DBusDecoder::decode(v.remove(0)).unwrap());
        msg.flags = DBusDecoder::decode::<u8>(v.remove(0)).unwrap();
        msg.version = DBusDecoder::decode::<u8>(v.remove(0)).unwrap();
        let body_len = DBusDecoder::decode::<u32>(v.remove(0)).unwrap();
        msg.serial = DBusDecoder::decode::<u32>(v.remove(0)).unwrap();

        // Read array length
        try!(read_exactly(sock, &mut buf, 4));
        // demarshal consumes the buf, so save a copy for when we demarshal the entire array
        let mut buf_copy = buf.clone();
        offset = 12;
        sig = "u".to_string();
        let data = demarshal(&mut buf, &mut offset, &mut sig).ok().unwrap();
        let arr_len = DBusDecoder::decode::<u32>(data).unwrap() as usize;

        // Make buf_copy big enough for the entire array, and fill it
        buf_copy.reserve(arr_len);
        if try!(sock.take(arr_len as u64).read_to_end(&mut buf_copy)) != arr_len {
            return Err(Error::Disconnected);
        };

        offset = 12;
        sig = "a(yv)".to_string();
        let header_fields = match try!(demarshal(&mut buf_copy, &mut offset, &mut sig)) {
            Value::Array(x) => x,
            x => panic!("Demarshal didn't return what we asked for: {:?}", x)
        };

        msg.headers = HashMap::new();
        for i in header_fields.objects {
            let mut st = match i {
                Value::Struct(x) => x,
                x => panic!("Demarshal didn't return what we asked for: {:?}", x)
            };
            let val = st.objects.remove(1);
            let code = DBusDecoder::decode::<u8>(st.objects.remove(0)).unwrap();
            msg.headers.insert(code, val);
        }

        // Read the padding, if any
        let trailing_pad = 8 - (offset % 8);
        if trailing_pad % 8 != 0 {
            try!(read_exactly(sock, &mut buf, trailing_pad));
        }

        // Finally, read the entire body
        if body_len > 0 {
            let v = match msg.headers.get(&(HeaderFieldName::Signature as u8)) {
                Some(&Value::Variant(ref x)) => x,
                _ => return Err(Error::DemarshalError(DemarshalError::BadSignature))
            };

            let sigval = match v.object.deref() {
                &Value::BasicValue(BasicValue::Signature(ref x)) => x,
                _ => return Err(Error::DemarshalError(DemarshalError::BadSignature))
            };

            let mut body = Vec::new();
            try!(read_exactly(sock, &mut body, body_len as usize));

            let mut sig = "(".to_string() + &sigval.0 + ")";
            offset = 0;
            let objs = match try!(demarshal(&mut body, &mut offset, &mut sig)) {
                Value::Struct(x) => x.objects,
                x => panic!("Didn't get a struct: {:?}", x)
            };
            for x in objs {
                msg.body.push(x);
            }
        }

        Ok(msg)
    }
}

#[test]
fn test_connect () {
    let mut conn = Connection::connect_uds("/var/run/dbus/system_bus_socket").unwrap();
    let mut msg = message::create_method_call("org.freedesktop.DBus", "/org/freedesktop/DBus",
                                          "org.freedesktop.DBus", "ListNames");
    let resp = conn.call_sync(&mut msg).unwrap();
    println!("ListNames: {:?}", resp);
}

#[cfg(tcp)]
#[test]
fn test_tcp () {
    let mut conn = Connection::connect_tcp("localhost:12345").unwrap();
    let mut msg = message::create_method_call("org.freedesktop.DBus", "/org/freedesktop/DBus",
                                          "org.freedesktop.DBus", "ListNames");
    conn.send(&mut msg).ok();
    let msg = conn.read_msg().unwrap();
    println!("{:?}", msg.body);
    //loop {
    //    conn.read_msg().unwrap();
    //}
}
<|endoftext|>