<commit_before>extern crate base64;
extern crate byteorder;
extern crate hyper;
extern crate rayon;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;

use byteorder::{BigEndian, WriteBytesExt};

use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

use std::io::{Read, Write};

#[derive(Debug, Clone)]
pub struct Log {
    pub description: String,
    pub url: String,
    is_google: bool,
}

#[derive(Debug, Deserialize)]
pub struct SignedCertificateTimestamp {
    sct_version: u8,
    id: String,
    timestamp: u64,
    extensions: String,
    signature: String,
}

impl SignedCertificateTimestamp {
    pub fn to_raw_bytes(&self) -> Vec<u8> {
        let mut b = Vec::new();
        b.write_u8(self.sct_version).unwrap();

        let log_id = base64::decode(&self.id).unwrap();
        b.write(&log_id).unwrap();

        b.write_u64::<BigEndian>(self.timestamp).unwrap();

        let extensions = base64::decode(&self.extensions).unwrap();
        assert!(extensions.len() <= 65535);
        b.write_u16::<BigEndian>(extensions.len() as u16)
            .unwrap();
        b.write(&extensions).unwrap();

        let signature = base64::decode(&self.signature).unwrap();
        b.write(&signature).unwrap();

        return b;
    }
}

#[derive(Deserialize)]
struct LogsResponseLogs {
    description: String,
    url: String,
    operated_by: Vec<u32>,
    disqualified_at: Option<u64>,
}

#[derive(Deserialize)]
struct LogsResponseOperators {
    name: String,
    id: u32,
}

#[derive(Deserialize)]
struct LogsResponse {
    logs: Vec<LogsResponseLogs>,
    operators: Vec<LogsResponseOperators>,
}

const LOG_LIST_URL: &'static str = "https://www.gstatic.com/ct/log_list/log_list.json";

pub fn fetch_trusted_ct_logs(http_client: &hyper::Client) -> Vec<Log> {
    let response = http_client.get(LOG_LIST_URL).send().unwrap();
    // Limit the response to 10MB at most, to be resillient to DoS.
    let logs_response: LogsResponse = serde_json::from_reader(response.take(10 * 1024 * 1024))
        .unwrap();

    let google_id = logs_response
        .operators
        .iter()
        .find(|o| o.name == "Google")
        .map(|o| o.id);

    return logs_response
               .logs
               .into_iter()
               .filter(|log| log.disqualified_at.is_none())
               .map(|log| {
                        Log {
                            url: log.url,
                            description: log.description,
                            is_google: log.operated_by.contains(&google_id.unwrap()),
                        }
                    })
               .collect();
}

fn submit_to_log(http_client: &hyper::Client,
                 url: &str,
                 payload: &[u8])
                 -> Option<SignedCertificateTimestamp> {
    let mut url = "https://".to_string() + url;
    if !url.ends_with("/") {
        url += "/";
    }
    url += "ct/v1/add-chain";
    let response = http_client
        .post(&url)
        .body(hyper::client::Body::BufBody(payload, payload.len()))
        .header(hyper::header::ContentType::json())
        .send();
    let response = match response {
        Ok(r) => r,
        // TODO: maybe not all of these should be silently ignored.
        Err(_) => return None,
    };

    // 400, 403, and probably some others generally indicate a log doesn't accept certs from this
    // root, or that the log isn't accepting new submissions.
    if response.status.is_client_error() {
        return None;
    }

    // Limt the response to 10MB (well above what would ever be needed) to be resilient to DoS in
    // the face of a dumb or malicious log.
    return Some(serde_json::from_reader(response.take(10 * 1024 * 1024)).unwrap());
}

#[derive(Serialize)]
struct AddChainRequest {
    chain: Vec<String>,
}

pub fn submit_cert_to_logs<'a>(http_client: &hyper::Client,
                           logs: &'a [Log],
                           cert: Vec<Vec<u8>>)
                           -> Vec<(&'a Log, SignedCertificateTimestamp)> {
    let payload = serde_json::to_vec(&AddChainRequest {
                                          chain: cert.iter().map(|r| base64::encode(r)).collect(),
                                      })
            .unwrap();

    return logs.par_iter()
               .filter_map(|ref log| {
                               let sct = submit_to_log(http_client, &log.url, &payload);
                               return sct.map(|s| (log.clone(), s));
                           })
               .collect();
}
<commit_msg>cargo fmt<commit_after>extern crate base64;
extern crate byteorder;
extern crate hyper;
extern crate rayon;
#[macro_use]
extern crate serde_derive;
extern crate serde_json;

use byteorder::{BigEndian, WriteBytesExt};

use rayon::iter::{IntoParallelRefIterator, ParallelIterator};

use std::io::{Read, Write};

#[derive(Debug, Clone)]
pub struct Log {
    pub description: String,
    pub url: String,
    is_google: bool,
}

#[derive(Debug, Deserialize)]
pub struct SignedCertificateTimestamp {
    sct_version: u8,
    id: String,
    timestamp: u64,
    extensions: String,
    signature: String,
}

impl SignedCertificateTimestamp {
    pub fn to_raw_bytes(&self) -> Vec<u8> {
        let mut b = Vec::new();
        b.write_u8(self.sct_version).unwrap();

        let log_id = base64::decode(&self.id).unwrap();
        b.write(&log_id).unwrap();

        b.write_u64::<BigEndian>(self.timestamp).unwrap();

        let extensions = base64::decode(&self.extensions).unwrap();
        assert!(extensions.len() <= 65535);
        b.write_u16::<BigEndian>(extensions.len() as u16)
            .unwrap();
        b.write(&extensions).unwrap();

        let signature = base64::decode(&self.signature).unwrap();
        b.write(&signature).unwrap();

        return b;
    }
}

#[derive(Deserialize)]
struct LogsResponseLogs {
    description: String,
    url: String,
    operated_by: Vec<u32>,
    disqualified_at: Option<u64>,
}

#[derive(Deserialize)]
struct LogsResponseOperators {
    name: String,
    id: u32,
}

#[derive(Deserialize)]
struct LogsResponse {
    logs: Vec<LogsResponseLogs>,
    operators: Vec<LogsResponseOperators>,
}

const LOG_LIST_URL: &'static str = "https://www.gstatic.com/ct/log_list/log_list.json";

pub fn fetch_trusted_ct_logs(http_client: &hyper::Client) -> Vec<Log> {
    let response = http_client.get(LOG_LIST_URL).send().unwrap();
    // Limit the response to 10MB at most, to be resillient to DoS.
    let logs_response: LogsResponse = serde_json::from_reader(response.take(10 * 1024 * 1024))
        .unwrap();

    let google_id = logs_response
        .operators
        .iter()
        .find(|o| o.name == "Google")
        .map(|o| o.id);

    return logs_response
               .logs
               .into_iter()
               .filter(|log| log.disqualified_at.is_none())
               .map(|log| {
                        Log {
                            url: log.url,
                            description: log.description,
                            is_google: log.operated_by.contains(&google_id.unwrap()),
                        }
                    })
               .collect();
}

fn submit_to_log(http_client: &hyper::Client,
                 url: &str,
                 payload: &[u8])
                 -> Option<SignedCertificateTimestamp> {
    let mut url = "https://".to_string() + url;
    if !url.ends_with("/") {
        url += "/";
    }
    url += "ct/v1/add-chain";
    let response = http_client
        .post(&url)
        .body(hyper::client::Body::BufBody(payload, payload.len()))
        .header(hyper::header::ContentType::json())
        .send();
    let response = match response {
        Ok(r) => r,
        // TODO: maybe not all of these should be silently ignored.
        Err(_) => return None,
    };

    // 400, 403, and probably some others generally indicate a log doesn't accept certs from this
    // root, or that the log isn't accepting new submissions.
    if response.status.is_client_error() {
        return None;
    }

    // Limt the response to 10MB (well above what would ever be needed) to be resilient to DoS in
    // the face of a dumb or malicious log.
    return Some(serde_json::from_reader(response.take(10 * 1024 * 1024)).unwrap());
}

#[derive(Serialize)]
struct AddChainRequest {
    chain: Vec<String>,
}

pub fn submit_cert_to_logs<'a>(http_client: &hyper::Client,
                               logs: &'a [Log],
                               cert: Vec<Vec<u8>>)
                               -> Vec<(&'a Log, SignedCertificateTimestamp)> {
    let payload = serde_json::to_vec(&AddChainRequest {
                                          chain: cert.iter().map(|r| base64::encode(r)).collect(),
                                      })
            .unwrap();

    return logs.par_iter()
               .filter_map(|ref log| {
                               let sct = submit_to_log(http_client, &log.url, &payload);
                               return sct.map(|s| (log.clone(), s));
                           })
               .collect();
}
<|endoftext|>


<commit_before>// Copyright 2017 The Exonum Team
//
// 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.

//! The time oracle service for Exonum.

#![deny(missing_debug_implementations, missing_docs)]

extern crate serde;
extern crate serde_json;
#[macro_use]
extern crate exonum;
extern crate router;
extern crate bodyparser;
extern crate iron;

use iron::prelude::*;
use iron::Handler;
use router::Router;

use std::time::SystemTime;
use std::collections::HashMap;

use exonum::blockchain::{Blockchain, Service, ServiceContext, Schema, Transaction, ApiContext};
use exonum::messages::{RawTransaction, Message};
use exonum::encoding::serialize::json::reexport::Value;
use exonum::storage::{Fork, Snapshot, MapIndex, Entry};
use exonum::crypto::PublicKey;
use exonum::encoding;
use exonum::helpers::fabric::{ServiceFactory, Context};
use exonum::api::Api;

/// Time service id.
const SERVICE_ID: u16 = 4;
/// `TxTime` transaction id.
const TX_TIME_ID: u16 = 1;
/// Time service name.
const SERVICE_NAME: &str = "exonum_time";

encoding_struct! {
    /// Time information.
    struct Time {
        const SIZE = 12;

        /// Field that stores `SystemTime`.
        field time:     SystemTime  [00 => 12]
    }
}

/// `Exonum-time` service database schema.
#[derive(Debug)]
pub struct TimeSchema<T> {
    view: T,
}

impl<T: AsRef<Snapshot>> TimeSchema<T> {
    /// Constructs schema for the given `snapshot`.
    pub fn new(view: T) -> Self {
        TimeSchema { view }
    }

    /// Returns the table that stores `Time` struct for every validator.
    pub fn validators_time(&self) -> MapIndex<&Snapshot, PublicKey, Time> {
        MapIndex::new(
            format!("{}.validators_time", SERVICE_NAME),
            self.view.as_ref(),
        )
    }

    /// Returns stored `Time`.
    pub fn time(&self) -> Entry<&Snapshot, Time> {
        Entry::new(format!("{}.time", SERVICE_NAME), self.view.as_ref())
    }
}

impl<'a> TimeSchema<&'a mut Fork> {
    /// Mutable reference to the ['validators_time'][1] index.
    ///
    /// [1]: struct.TimeSchema.html#method.validators_time
    pub fn validators_time_mut(&mut self) -> MapIndex<&mut Fork, PublicKey, Time> {
        MapIndex::new(format!("{}.validators_time", SERVICE_NAME), self.view)
    }

    /// Mutable reference to the ['time'][1] index.
    ///
    /// [1]: struct.TimeSchema.html#method.time
    pub fn time_mut(&mut self) -> Entry<&mut Fork, Time> {
        Entry::new(format!("{}.time", SERVICE_NAME), self.view)
    }
}

message! {
    /// Transaction that is sent by the validator after the commit of the block.
    struct TxTime {
        const TYPE = SERVICE_ID;
        const ID = TX_TIME_ID;
        const SIZE = 44;
        /// Validator's time.
        field time:     SystemTime  [00 => 12]
        /// Validator's public key.
        field pub_key:  &PublicKey  [12 => 44]
    }
}

impl Transaction for TxTime {
    fn verify(&self) -> bool {
        self.verify_signature(self.pub_key())
    }

    fn execute(&self, view: &mut Fork) {
        let validator_keys = Schema::new(&view).actual_configuration().validator_keys;

        // The transaction must be signed by the validator.
        let signed = validator_keys.iter().any(|&validator| {
            validator.service_key == *self.pub_key()
        });
        if !signed {
            return;
        }

        let mut schema = TimeSchema::new(view);
        match schema.validators_time().get(self.pub_key()) {
            // The validator time in the storage should be less than in the transaction.
            Some(ref storage_time) if storage_time.time() >= self.time() => {
                return;
            }
            // Write the time for the validator.
            _ => {
                schema.validators_time_mut().put(
                    self.pub_key(),
                    Time::new(self.time()),
                )
            }
        }

        // Find all known times for the validators.
        let mut validators_time: Vec<SystemTime>;
        {
            let idx = schema.validators_time();
            validators_time = idx.iter()
                .filter_map(|pair| {
                    validator_keys
                        .iter()
                        .find(|validator| validator.service_key == pair.0)
                        .map(|_| pair.1.time())
                })
                .collect();
        }

        // The largest number of Byzantine nodes.
        let max_byzantine_nodes = (validator_keys.len() - 1) / 3;
        if validators_time.len() <= 2 * max_byzantine_nodes {
            return;
        }
        // Ordering time from highest to lowest.
        validators_time.sort_by(|a, b| b.cmp(a));

        match schema.time().get() {
            // Selected time should be longer than the time in the storage.
            Some(ref current_time)
                if current_time.time() >= validators_time[max_byzantine_nodes] => {
                return;
            }
            _ => {
                // Change the time in the storage.
                schema.time_mut().set(Time::new(
                    validators_time[max_byzantine_nodes],
                ));
            }
        }
    }
}

/// Implements the node API.
#[derive(Clone)]
struct TimeApi {
    blockchain: Blockchain,
}

/// Shortcut to get data from storage.
impl TimeApi {
    /// Endpoint for getting value of the time that is saved in storage.
    fn get_current_time(&self, _: &mut Request) -> IronResult<Response> {
        let view = self.blockchain.snapshot();
        let schema = TimeSchema::new(&view);
        let current_time = schema.time().get();
        self.ok_response(&serde_json::to_value(current_time).unwrap())
    }

    /// Endpoint for getting time values for all validators.
    fn get_validators_time(&self, _: &mut Request) -> IronResult<Response> {
        let view = self.blockchain.snapshot();
        let schema = TimeSchema::new(&view);
        let idx = schema.validators_time();

        let validators_time: HashMap<_, _> = idx.iter().collect();

        if validators_time.is_empty() {
            self.not_found_response(&serde_json::to_value("Validators time database is empty")
                .unwrap())
        } else {
            self.ok_response(&serde_json::to_value(validators_time).unwrap())
        }
    }

    fn wire_private(&self, router: &mut Router) {
        let self_ = self.clone();
        let get_validators_time = move |req: &mut Request| self_.get_validators_time(req);
        router.get(
            "/validators_time",
            get_validators_time,
            "get_validators_time",
        );
    }
}

impl Api for TimeApi {
    fn wire(&self, router: &mut Router) {
        let self_ = self.clone();
        let get_current_time = move |req: &mut Request| self_.get_current_time(req);
        router.get("/current_time", get_current_time, "get_current_time");
    }
}

/// A helper trait that provides the node with a current time.
pub trait TimeProvider: Send + Sync + ::std::fmt::Debug {
    /// Returns the current time.
    fn current_time(&self) -> SystemTime;
}

#[derive(Debug)]
struct SystemTimeProvider;

impl TimeProvider for SystemTimeProvider {
    fn current_time(&self) -> SystemTime {
        SystemTime::now()
    }
}

/// Define the service.
#[derive(Debug)]
pub struct TimeService {
    /// Current time.
    time: Box<TimeProvider>,
}

impl Default for TimeService {
    fn default() -> TimeService {
        TimeService { time: Box::new(SystemTimeProvider) as Box<TimeProvider> }
    }
}

impl TimeService {
    /// Create a new `TimeService`.
    pub fn new() -> TimeService {
        TimeService::default()
    }

    /// Create a new `TimeService` with time provider `T`.
    pub fn with_provider<T: Into<Box<TimeProvider>>>(time_provider: T) -> TimeService {
        TimeService { time: time_provider.into() }
    }
}

impl Service for TimeService {
    fn service_name(&self) -> &'static str {
        SERVICE_NAME
    }

    fn service_id(&self) -> u16 {
        SERVICE_ID
    }

    fn tx_from_raw(&self, raw: RawTransaction) -> Result<Box<Transaction>, encoding::Error> {
        match raw.message_type() {
            TX_TIME_ID => Ok(Box::new(TxTime::from_raw(raw)?)),
            _ => {
                let error =
                    encoding::Error::IncorrectMessageType { message_type: raw.message_type() };
                Err(error)
            }
        }
    }

    fn initialize(&self, _fork: &mut Fork) -> Value {
        Value::Null
    }

    /// Creates transaction after commit of the block.
    fn handle_commit(&self, context: &ServiceContext) {
        // The transaction must be created by the validator.
        if context.validator_id().is_none() {
            return;
        }
        let (pub_key, sec_key) = (*context.public_key(), context.secret_key().clone());
        context
            .transaction_sender()
            .send(Box::new(
                TxTime::new(self.time.current_time(), &pub_key, &sec_key),
            ))
            .unwrap();
    }

    fn private_api_handler(&self, ctx: &ApiContext) -> Option<Box<Handler>> {
        let mut router = Router::new();
        let api = TimeApi { blockchain: ctx.blockchain().clone() };
        api.wire_private(&mut router);
        Some(Box::new(router))
    }

    fn public_api_handler(&self, ctx: &ApiContext) -> Option<Box<Handler>> {
        let mut router = Router::new();
        let api = TimeApi { blockchain: ctx.blockchain().clone() };
        api.wire(&mut router);
        Some(Box::new(router))
    }
}

/// A time service creator for the `NodeBuilder`.
#[derive(Debug)]
pub struct TimeServiceFactory;

impl ServiceFactory for TimeServiceFactory {
    fn make_service(&mut self, _: &Context) -> Box<Service> {
        Box::new(TimeService::new())
    }
}
<commit_msg>Change json format<commit_after>// Copyright 2017 The Exonum Team
//
// 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.

//! The time oracle service for Exonum.

#![deny(missing_debug_implementations, missing_docs)]

extern crate serde;
extern crate serde_json;
#[macro_use]
extern crate exonum;
extern crate router;
extern crate bodyparser;
extern crate iron;

use iron::prelude::*;
use iron::Handler;
use router::Router;

use std::time::SystemTime;

use exonum::blockchain::{Blockchain, Service, ServiceContext, Schema, Transaction, ApiContext};
use exonum::messages::{RawTransaction, Message};
use exonum::encoding::serialize::json::reexport::Value;
use exonum::storage::{Fork, Snapshot, MapIndex, Entry};
use exonum::crypto::PublicKey;
use exonum::encoding;
use exonum::helpers::fabric::{ServiceFactory, Context};
use exonum::api::Api;

/// Time service id.
const SERVICE_ID: u16 = 4;
/// `TxTime` transaction id.
const TX_TIME_ID: u16 = 1;
/// Time service name.
const SERVICE_NAME: &str = "exonum_time";

encoding_struct! {
    /// Time information.
    struct Time {
        const SIZE = 12;

        /// Field that stores `SystemTime`.
        field time:     SystemTime  [00 => 12]
    }
}

/// `Exonum-time` service database schema.
#[derive(Debug)]
pub struct TimeSchema<T> {
    view: T,
}

impl<T: AsRef<Snapshot>> TimeSchema<T> {
    /// Constructs schema for the given `snapshot`.
    pub fn new(view: T) -> Self {
        TimeSchema { view }
    }

    /// Returns the table that stores `Time` struct for every validator.
    pub fn validators_time(&self) -> MapIndex<&Snapshot, PublicKey, Time> {
        MapIndex::new(
            format!("{}.validators_time", SERVICE_NAME),
            self.view.as_ref(),
        )
    }

    /// Returns stored `Time`.
    pub fn time(&self) -> Entry<&Snapshot, Time> {
        Entry::new(format!("{}.time", SERVICE_NAME), self.view.as_ref())
    }
}

impl<'a> TimeSchema<&'a mut Fork> {
    /// Mutable reference to the ['validators_time'][1] index.
    ///
    /// [1]: struct.TimeSchema.html#method.validators_time
    pub fn validators_time_mut(&mut self) -> MapIndex<&mut Fork, PublicKey, Time> {
        MapIndex::new(format!("{}.validators_time", SERVICE_NAME), self.view)
    }

    /// Mutable reference to the ['time'][1] index.
    ///
    /// [1]: struct.TimeSchema.html#method.time
    pub fn time_mut(&mut self) -> Entry<&mut Fork, Time> {
        Entry::new(format!("{}.time", SERVICE_NAME), self.view)
    }
}

message! {
    /// Transaction that is sent by the validator after the commit of the block.
    struct TxTime {
        const TYPE = SERVICE_ID;
        const ID = TX_TIME_ID;
        const SIZE = 44;
        /// Validator's time.
        field time:     SystemTime  [00 => 12]
        /// Validator's public key.
        field pub_key:  &PublicKey  [12 => 44]
    }
}

impl Transaction for TxTime {
    fn verify(&self) -> bool {
        self.verify_signature(self.pub_key())
    }

    fn execute(&self, view: &mut Fork) {
        let validator_keys = Schema::new(&view).actual_configuration().validator_keys;

        // The transaction must be signed by the validator.
        let signed = validator_keys.iter().any(|&validator| {
            validator.service_key == *self.pub_key()
        });
        if !signed {
            return;
        }

        let mut schema = TimeSchema::new(view);
        match schema.validators_time().get(self.pub_key()) {
            // The validator time in the storage should be less than in the transaction.
            Some(ref storage_time) if storage_time.time() >= self.time() => {
                return;
            }
            // Write the time for the validator.
            _ => {
                schema.validators_time_mut().put(
                    self.pub_key(),
                    Time::new(self.time()),
                )
            }
        }

        // Find all known times for the validators.
        let mut validators_time: Vec<SystemTime>;
        {
            let idx = schema.validators_time();
            validators_time = idx.iter()
                .filter_map(|pair| {
                    validator_keys
                        .iter()
                        .find(|validator| validator.service_key == pair.0)
                        .map(|_| pair.1.time())
                })
                .collect();
        }

        // The largest number of Byzantine nodes.
        let max_byzantine_nodes = (validator_keys.len() - 1) / 3;
        if validators_time.len() <= 2 * max_byzantine_nodes {
            return;
        }
        // Ordering time from highest to lowest.
        validators_time.sort_by(|a, b| b.cmp(a));

        match schema.time().get() {
            // Selected time should be longer than the time in the storage.
            Some(ref current_time)
                if current_time.time() >= validators_time[max_byzantine_nodes] => {
                return;
            }
            _ => {
                // Change the time in the storage.
                schema.time_mut().set(Time::new(
                    validators_time[max_byzantine_nodes],
                ));
            }
        }
    }
}

/// Implements the node API.
#[derive(Clone)]
struct TimeApi {
    blockchain: Blockchain,
}

encoding_struct! {
    /// Structure for saving validator's public key and last known local time.
    struct ValidatorTime {
        const SIZE = 44;

        /// Validator's public key.
        field public_key:     &PublicKey      [00 => 32]
        /// Validator's time.
        field time:           SystemTime      [32 => 44]
    }
}

/// Shortcut to get data from storage.
impl TimeApi {
    /// Endpoint for getting value of the time that is saved in storage.
    fn get_current_time(&self, _: &mut Request) -> IronResult<Response> {
        let view = self.blockchain.snapshot();
        let schema = TimeSchema::new(&view);
        let current_time = schema.time().get();
        self.ok_response(&serde_json::to_value(current_time).unwrap())
    }

    /// Endpoint for getting time values for all validators.
    fn get_validators_time(&self, _: &mut Request) -> IronResult<Response> {
        let view = self.blockchain.snapshot();
        let schema = TimeSchema::new(&view);
        let idx = schema.validators_time();

        let validators_time: Vec<_> = idx.iter()
            .map(|(public_key, time)| {
                ValidatorTime::new(&public_key, time.time())
            })
            .collect();

        self.ok_response(&serde_json::to_value(validators_time).unwrap())
    }

    fn wire_private(&self, router: &mut Router) {
        let self_ = self.clone();
        let get_validators_time = move |req: &mut Request| self_.get_validators_time(req);
        router.get(
            "/validators_time",
            get_validators_time,
            "get_validators_time",
        );
    }
}

impl Api for TimeApi {
    fn wire(&self, router: &mut Router) {
        let self_ = self.clone();
        let get_current_time = move |req: &mut Request| self_.get_current_time(req);
        router.get("/current_time", get_current_time, "get_current_time");
    }
}

/// A helper trait that provides the node with a current time.
pub trait TimeProvider: Send + Sync + ::std::fmt::Debug {
    /// Returns the current time.
    fn current_time(&self) -> SystemTime;
}

#[derive(Debug)]
struct SystemTimeProvider;

impl TimeProvider for SystemTimeProvider {
    fn current_time(&self) -> SystemTime {
        SystemTime::now()
    }
}

/// Define the service.
#[derive(Debug)]
pub struct TimeService {
    /// Current time.
    time: Box<TimeProvider>,
}

impl Default for TimeService {
    fn default() -> TimeService {
        TimeService { time: Box::new(SystemTimeProvider) as Box<TimeProvider> }
    }
}

impl TimeService {
    /// Create a new `TimeService`.
    pub fn new() -> TimeService {
        TimeService::default()
    }

    /// Create a new `TimeService` with time provider `T`.
    pub fn with_provider<T: Into<Box<TimeProvider>>>(time_provider: T) -> TimeService {
        TimeService { time: time_provider.into() }
    }
}

impl Service for TimeService {
    fn service_name(&self) -> &'static str {
        SERVICE_NAME
    }

    fn service_id(&self) -> u16 {
        SERVICE_ID
    }

    fn tx_from_raw(&self, raw: RawTransaction) -> Result<Box<Transaction>, encoding::Error> {
        match raw.message_type() {
            TX_TIME_ID => Ok(Box::new(TxTime::from_raw(raw)?)),
            _ => {
                let error =
                    encoding::Error::IncorrectMessageType { message_type: raw.message_type() };
                Err(error)
            }
        }
    }

    fn initialize(&self, _fork: &mut Fork) -> Value {
        Value::Null
    }

    /// Creates transaction after commit of the block.
    fn handle_commit(&self, context: &ServiceContext) {
        // The transaction must be created by the validator.
        if context.validator_id().is_none() {
            return;
        }
        let (pub_key, sec_key) = (*context.public_key(), context.secret_key().clone());
        context
            .transaction_sender()
            .send(Box::new(
                TxTime::new(self.time.current_time(), &pub_key, &sec_key),
            ))
            .unwrap();
    }

    fn private_api_handler(&self, ctx: &ApiContext) -> Option<Box<Handler>> {
        let mut router = Router::new();
        let api = TimeApi { blockchain: ctx.blockchain().clone() };
        api.wire_private(&mut router);
        Some(Box::new(router))
    }

    fn public_api_handler(&self, ctx: &ApiContext) -> Option<Box<Handler>> {
        let mut router = Router::new();
        let api = TimeApi { blockchain: ctx.blockchain().clone() };
        api.wire(&mut router);
        Some(Box::new(router))
    }
}

/// A time service creator for the `NodeBuilder`.
#[derive(Debug)]
pub struct TimeServiceFactory;

impl ServiceFactory for TimeServiceFactory {
    fn make_service(&mut self, _: &Context) -> Box<Service> {
        Box::new(TimeService::new())
    }
}
<|endoftext|>


<commit_before>//! corepack is a no_std support for messagepack in serde.
//
// 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 https://mozilla.org/MPL/2.0/.

#![cfg_attr(feature = "alloc", feature(alloc))]
#![allow(overflowing_literals)]

// testing requires std to be available
#![cfg_attr(all(not(feature = "std"), not(test)), no_std)]
#[cfg(all(not(feature = "std"), not(test)))]
extern crate core as std;
extern crate serde;
extern crate byteorder;
#[cfg(test)]
#[macro_use]
extern crate serde_derive;

#[cfg(feature = "alloc")]
#[macro_use]
extern crate alloc;

#[cfg(feature = "alloc")]
use alloc::Vec;

pub use ser::Serializer;
pub use de::Deserializer;

pub mod error;
pub mod read;

mod defs;
mod seq_serializer;
mod map_serializer;
mod variant_deserializer;
mod ext_deserializer;
mod seq_deserializer;

mod ser;
mod de;

/// Parse V out of a stream of bytes.
pub fn from_iter<I, V>(mut iter: I) -> Result<V, error::Error>
    where I: Iterator<Item = u8>,
          V: serde::de::DeserializeOwned
{
    let mut de = Deserializer::new(read::CopyRead::new(|buf: &mut [u8]| {
        for i in 0..buf.len() {
            if let Some(byte) = iter.next() {
                buf[i] = byte;
            } else {
                return Err(error::Error::EndOfStream);
            }
        }

        Ok(())
    }));

    V::deserialize(&mut de)
}

/// Parse V out of a slice of bytes.
pub fn from_bytes<'a, V>(bytes: &'a [u8]) -> Result<V, error::Error>
    where V: serde::Deserialize<'a>
{
    let mut position: usize = 0;

    let mut de = Deserializer::new(read::BorrowRead::new(|len: usize| if position + len >
                                                                         bytes.len() {
        Err(error::Error::EndOfStream)
    } else {
        let result = &bytes[position..position + len];

        position += len;

        Ok(result)
    }));

    V::deserialize(&mut de)
}

/// Serialize V into a byte buffer.
pub fn to_bytes<V>(value: V) -> Result<Vec<u8>, error::Error>
    where V: serde::Serialize
{
    let mut bytes = vec![];

    {
        let mut ser = Serializer::new(|buf| {
            bytes.extend_from_slice(buf);
            Ok(())
        });

        try!(value.serialize(&mut ser));
    }

    Ok(bytes)
}

#[cfg(test)]
mod test {
    use serde::Serialize;
    use serde::de::DeserializeOwned;
    use std::fmt::Debug;

    #[derive(PartialEq, Eq, Debug, Serialize, Deserialize)]
    enum T {
        A(usize),
        B,
        C(i8, i8),
        D { a: isize, b: String },
    }

    fn test_through<T>(expected: T)
        where T: Serialize + DeserializeOwned + PartialEq + Debug
    {
        let x = ::to_bytes(&expected).expect("Failed to serialize");

        let actual = ::from_bytes(&x).expect("Failed to deserialize");

        assert_eq!(expected, actual);
    }

    #[test]
    fn test_str() {
        test_through(format!("Hello World!"))
    }

    #[test]
    fn test_enum() {
        test_through(T::B)
    }

    #[test]
    fn test_enum_newtype() {
        test_through(T::A(42))
    }

    #[test]
    fn test_enum_tuple() {
        test_through(T::C(-3, 22))
    }

    #[test]
    fn test_enum_struct() {
        test_through(T::D {
            a: 9001,
            b: "Hello world!".into(),
        })
    }

    #[test]
    fn test_option() {
        test_through(Some(7))
    }

    #[test]
    fn test_unit_option() {
        test_through(Some(()))
    }

    #[test]
    fn test_char() {
        test_through('b')
    }
}
<commit_msg>test(lib): Improve tests<commit_after>//! corepack is a no_std support for messagepack in serde.
//
// 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 https://mozilla.org/MPL/2.0/.

#![cfg_attr(feature = "alloc", feature(alloc))]
#![allow(overflowing_literals)]

// testing requires std to be available
#![cfg_attr(all(not(feature = "std"), not(test)), no_std)]
#[cfg(all(not(feature = "std"), not(test)))]
extern crate core as std;
extern crate serde;
extern crate byteorder;
#[cfg(test)]
#[macro_use]
extern crate serde_derive;

#[cfg(feature = "alloc")]
#[macro_use]
extern crate alloc;

#[cfg(feature = "alloc")]
use alloc::Vec;

pub use ser::Serializer;
pub use de::Deserializer;

pub mod error;
pub mod read;

mod defs;
mod seq_serializer;
mod map_serializer;
mod variant_deserializer;
mod ext_deserializer;
mod seq_deserializer;

mod ser;
mod de;

/// Parse V out of a stream of bytes.
pub fn from_iter<I, V>(mut iter: I) -> Result<V, error::Error>
    where I: Iterator<Item = u8>,
          V: serde::de::DeserializeOwned
{
    let mut de = Deserializer::new(read::CopyRead::new(|buf: &mut [u8]| {
        for i in 0..buf.len() {
            if let Some(byte) = iter.next() {
                buf[i] = byte;
            } else {
                return Err(error::Error::EndOfStream);
            }
        }

        Ok(())
    }));

    V::deserialize(&mut de)
}

/// Parse V out of a slice of bytes.
pub fn from_bytes<'a, V>(bytes: &'a [u8]) -> Result<V, error::Error>
    where V: serde::Deserialize<'a>
{
    let mut position: usize = 0;

    let mut de = Deserializer::new(read::BorrowRead::new(|len: usize| if position + len >
                                                                         bytes.len() {
        Err(error::Error::EndOfStream)
    } else {
        let result = &bytes[position..position + len];

        position += len;

        Ok(result)
    }));

    V::deserialize(&mut de)
}

/// Serialize V into a byte buffer.
pub fn to_bytes<V>(value: V) -> Result<Vec<u8>, error::Error>
    where V: serde::Serialize
{
    let mut bytes = vec![];

    {
        let mut ser = Serializer::new(|buf| {
            bytes.extend_from_slice(buf);
            Ok(())
        });

        try!(value.serialize(&mut ser));
    }

    Ok(bytes)
}

#[cfg(test)]
mod test {
    use serde::Serialize;
    use serde::de::DeserializeOwned;
    use std::fmt::Debug;
    use std::ffi::CString;

    #[derive(PartialEq, Eq, Debug, Serialize, Deserialize)]
    enum T {
        A(usize),
        B,
        C(i8, i8),
        D { a: isize, b: String },
    }

    fn test_through<T>(item: T, expected: &[u8])
        where T: Serialize + DeserializeOwned + PartialEq + Debug
    {
        let actual = ::to_bytes(&item).expect("Failed to serialize");

        assert_eq!(expected, &*actual);

        let deserialized_item = ::from_bytes(&actual).expect("Failed to deserialize");

        assert_eq!(item, deserialized_item);
    }

    #[test]
    fn test_str() {
        test_through(format!("Hello World!"),
            &[0xac, 0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x20, 0x57, 0x6f, 0x72, 0x6c, 0x64, 0x21]);
    }

    #[test]
    fn test_enum() {
        test_through(T::B, &[0x92, 0x01, 0xc0])
    }

    #[test]
    fn test_enum_newtype() {
        test_through(T::A(42), &[0x92, 0x00, 0x2a])
    }

    #[test]
    fn test_enum_tuple() {
        test_through(T::C(-3, 22), &[0x92, 0x02, 0x92, 0xfd, 0x16])
    }

    #[test]
    fn test_enum_struct() {
        test_through(T::D {
            a: 9001,
            b: "Hello world!".into(),
        }, &[0x92, // array with two elements
                0x03, // 3 (variant index)
                0x82, // map with two entries
                    0xa1, 0x61, // entry one, fixstr length one: 'a'
                        0xd1, 0x23, 0x29, // i16: 9001
                    0xa1, 0x62, // entry two, fixstr length one: 'b'
                        // fixstr, length 12: Hello world!
                        0xac, 0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x20, 0x77, 0x6f, 0x72, 0x6c, 0x64, 0x21])
    }

    #[test]
    fn test_option() {
        test_through(Some(7), &[0x92, 0xc3, 0x07])
    }

    #[test]
    fn test_option_none() {
        test_through::<Option<usize>>(None, &[0x91, 0xc2])
    }

    #[test]
    fn test_unit_option() {
        test_through(Some(()), &[0x92, 0xc3, 0xc0])
    }

    #[test]
    fn test_char() {
        test_through('b', &[0xa1, 0x62])
    }

    #[test]
    fn test_false() {
        test_through(false, &[0xc2])
    }

    #[test]
    fn test_byte_array() {
        test_through(CString::new("hello").unwrap(), &[0xc4, 0x05, 0x68, 0x65, 0x6c, 0x6c, 0x6f])
    }

    #[test]
    fn test_float() {
        test_through(4.5, &[0xcb, 0x40, 0x12, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])
    }

    #[test]
    fn test_float32() {
        test_through(3.2f32, &[0xca, 0x40, 0x4c, 0xcc, 0xcd])
    }
}
<|endoftext|>


<commit_before>//! MySQL Proxy Server

#[macro_use]
extern crate log;
extern crate env_logger;
#[macro_use]
extern crate futures;
#[macro_use]
extern crate tokio_core;
extern crate futures_cpupool;
extern crate byteorder;

use std::cell::RefCell;
use std::rc::Rc;
use std::env;
use std::io::{self, Read, Write, Error, ErrorKind};
use std::net::{SocketAddr, Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
use std::net::Shutdown;
use std::str;
use std::time::Duration;

use futures::{Future, Poll, Async};
use futures::stream::Stream;
use futures_cpupool::CpuPool;
use tokio_core::{Loop, LoopHandle, TcpStream};
use tokio_core::io::{read_exact, write_all, Window};

use byteorder::*;

#[derive(Copy,Clone)]
pub enum PacketType {
    ComSleep = 0x00,
    ComQuit = 0x01,
    ComInitDb = 0x02,
    ComQuery = 0x03,
    ComFieldList = 0x04,
    ComCreateDb = 0x05,
    ComDropDb = 0x06,
    ComRefresh = 0x07,
    ComShutdown = 0x08,
    ComStatistics = 0x09,
    ComProcessInfo = 0x0a,
    ComConnect = 0x0b,
    ComProcessKill= 0x0c,
    ComDebug = 0x0d,
    ComPing = 0x0e,
    ComTime = 0x0f,
    ComDelayedInsert = 0x10,
    ComChangeUser = 0x11,
    ComBinlogDump = 0x12,
    ComTableDump = 0x13,
    ComConnectOut = 0x14,
    ComRegisterSlave = 0x15,
    ComStmtPrepare = 0x16,
    ComStmtExecute = 0x17,
    ComStmtSendLongData = 0x18,
    ComStmtClose = 0x19,
    ComStmtReset = 0x1a,
    ComDaemon= 0x1d,
    ComBinlogDumpGtid = 0x1e,
    ComResetConnection = 0x1f,
}

pub struct Packet {
    pub bytes: Vec<u8>
}

impl Packet {

    pub fn error_packet(code: u16, state: [u8; 5], msg: String) -> Self {

        // start building payload
        let mut payload: Vec<u8> = Vec::with_capacity(9 + msg.len());
        payload.push(0xff);  // packet type
        payload.write_u16::<LittleEndian>(code).unwrap(); // error code
        payload.extend_from_slice("#".as_bytes()); // sql_state_marker
        payload.extend_from_slice(&state); // SQL STATE
        payload.extend_from_slice(msg.as_bytes());

        // create header with length and sequence id
        let mut header: Vec<u8> = Vec::with_capacity(4 + 9 + msg.len());
        header.write_u32::<LittleEndian>(payload.len() as u32).unwrap();
        header.pop(); // we need 3 byte length, so discard last byte
        header.push(1); // sequence_id

        // combine the vectors
        header.extend_from_slice(&payload);

        // now move the vector into the packet
        Packet { bytes: header }
    }

    pub fn sequence_id(&self) -> u8 {
        self.bytes[3]
    }

    //TODO: should return Result<PacketType, ?>
    pub fn packet_type(&self) -> PacketType {
        match self.bytes[4] {
            0x00 => PacketType::ComSleep,
            0x01 => PacketType::ComQuit,
            0x02 => PacketType::ComInitDb,
            0x03 => PacketType::ComQuery,
            0x04 => PacketType::ComFieldList,
            0x05 => PacketType::ComCreateDb,
            0x06 => PacketType::ComDropDb,
            0x07 => PacketType::ComRefresh,
            0x08 => PacketType::ComShutdown,
            0x09 => PacketType::ComStatistics,
            0x0a => PacketType::ComProcessInfo,
            0x0b => PacketType::ComConnect,
            0x0c => PacketType::ComProcessKill,
            0x0d => PacketType::ComDebug,
            0x0e => PacketType::ComPing,
            0x0f => PacketType::ComTime,
            0x10 => PacketType::ComDelayedInsert,
            0x11 => PacketType::ComChangeUser,
            0x12 => PacketType::ComBinlogDump,
            0x13 => PacketType::ComTableDump,
            0x14 => PacketType::ComConnectOut,
            0x15 => PacketType::ComRegisterSlave,
            0x16 => PacketType::ComStmtPrepare,
            0x17 => PacketType::ComStmtExecute,
            0x18 => PacketType::ComStmtSendLongData,
            0x19 => PacketType::ComStmtClose,
            0x1a => PacketType::ComStmtReset,
            0x1d => PacketType::ComDaemon,
            0x1e => PacketType::ComBinlogDumpGtid,
            0x1f => PacketType::ComResetConnection,
            _ => panic!("Unsupported packet type")
        }
    }

}

pub enum Action {
    Forward,                // forward the packet unmodified
    Mutate(Packet),         // mutate the packet
    Respond(Vec<Packet>)    // handle directly and optionally return some packets
}

pub trait PacketHandler {
    fn handle_request(&self, p: &Packet) -> Action;
    fn handle_response(&self, p: &Packet) -> Action;
}

fn parse_packet_length(header: &[u8]) -> usize {
    (((header[2] as u32) << 16) |
    ((header[1] as u32) << 8) |
    header[0] as u32) as usize
}

pub struct Client {
    pub pool: CpuPool,
    pub handle: LoopHandle,
}

impl Client {

    pub fn serve<H: PacketHandler + 'static>(self, conn: TcpStream,
                                             handler: H,
                                        ) -> Box<Future<Item = (u64, u64), Error = Error>> {

        let addr = Ipv4Addr::new(127, 0, 0, 1);
        let port = 3306;
        let addr = SocketAddr::V4(SocketAddrV4::new(addr, port));

        let handle = self.handle.clone();
        let connected = handle.tcp_connect(&addr);

        let pair = connected.and_then(move |server| Ok((conn,server)));

        Box::new(pair.and_then(move |(c1, c2)| {
            let c1 = Rc::new(c1);
            let c2 = Rc::new(c2);
            Pipe::new(c1, c2, handler)
        }))
    }
}

#[derive(Debug)]
enum Direction {
    Client,
    Server
}

struct ConnReader {
    direction: Direction,
    stream: Rc<TcpStream>,
    read_buf: Vec<u8>,
    read_pos: usize,
    read_amt: u64,
}

struct ConnWriter {
    direction: Direction,
    stream: Rc<TcpStream>,
    write_buf: Vec<u8>,
    write_pos: usize,
    write_amt: u64,
}

impl ConnReader {

    fn new(stream: Rc<TcpStream>, direction: Direction) -> Self {
        ConnReader {
            stream: stream,
            direction: direction,
            read_buf: vec![0u8; 4096],
            read_pos: 0,
            read_amt: 0,
        }
    }

    /// Read from the socket until the status is NotReady
    fn read(&mut self) -> Poll<(), io::Error> {
        loop {
            try_ready!(self.stream.poll_read());
            //TODO: ensure capacity first
            let n = try_nb!((&*self.stream).read(&mut self.read_buf[self.read_pos..]));
            if n == 0 {
                return Err(Error::new(ErrorKind::Other, "connection closed"));
            }
            self.read_amt += n as u64;
            self.read_pos += n;
        }
    }

    fn next(&mut self) -> Option<Packet> {
        // do we have a header
        if self.read_pos > 3 {
            let l = parse_packet_length(&self.read_buf);
            // do we have the whole packet?
            let s = 4 + l;
            if self.read_pos >= s {
                let mut temp : Vec<u8> = Vec::with_capacity(s);
                temp.extend_from_slice(&self.read_buf[0..s]);
                let p = Packet { bytes: temp };

                // shift data down
                let mut j = 0;
                for i in s .. self.read_pos {
                    self.read_buf[j] = self.read_buf[i];
                    j += 1;
                }
                self.read_pos -= s;

                Some(p)
            } else {
                None
            }
        } else {
            None
        }
    }
}

impl ConnWriter {

    fn new(stream: Rc<TcpStream>, direction: Direction) -> Self {
        ConnWriter{
            stream: stream,
            direction: direction,
            write_buf: vec![0u8; 4096],
            write_pos: 0,
            write_amt: 0
        }
    }

    /// Write a packet to the write buffer
    fn push(&mut self, p: &Packet) {
        for i in 0 .. p.bytes.len() {
            self.write_buf[self.write_pos + i] = p.bytes[i];
        }
        self.write_pos += p.bytes.len();
    }

    /// Writes the contents of the write buffer to the socket
    fn write(&mut self) -> Poll<(), io::Error> {
        while self.write_pos > 0 {
            try_ready!(self.stream.poll_write());
            let s = try!((&*self.stream).write(&self.write_buf[0..self.write_pos]));

            let mut j = 0;
            for i in s .. self.write_pos {
                self.write_buf[j] = self.write_buf[i];
                j += 1;
            }
            self.write_pos -= s;

        }
        return Ok(Async::Ready(()));
    }
}

struct Pipe<H: PacketHandler + 'static> {
    client_reader: ConnReader,
    client_writer: ConnWriter,
    server_reader: ConnReader,
    server_writer: ConnWriter,
    handler: H,
}

impl<H> Pipe<H> where H: PacketHandler + 'static {
    fn new(client: Rc<TcpStream>,
           server: Rc<TcpStream>,
           handler: H
           ) -> Pipe<H> {

        Pipe {
            client_reader: ConnReader::new(client.clone(), Direction::Client),
            client_writer: ConnWriter::new(client, Direction::Client),
            server_reader: ConnReader::new(server.clone(), Direction::Server),
            server_writer: ConnWriter::new(server, Direction::Server),
            handler: handler,
        }
    }
}

impl<H> Future for Pipe<H> where H: PacketHandler + 'static {
    type Item = (u64, u64);
    type Error = Error;

    fn poll(&mut self) -> Poll<(u64, u64), io::Error> {
        loop {
            let client_read = self.client_reader.read();

            // if the client connection has closed, close the server connection too
            match &client_read {
                &Err(ref e) => { self.server_writer.stream.shutdown(Shutdown::Write).unwrap(); },
                _ => {}
            }

            // process buffered requests
            while let Some(request) = self.client_reader.next() {
                match self.handler.handle_request(&request) {
                    Action::Forward => self.server_writer.push(&request),
                    Action::Mutate(ref p2) => self.server_writer.push(p2),
                    Action::Respond(ref v) => {
                        for p in v {
                            self.client_writer.push(&p);
                        }
                    }
                };
            }

            // try writing to server
            let server_write = self.server_writer.write();

            // try reading from server
            let server_read = self.server_reader.read();

            // if the server connection has closed, close the client connection too
            match &server_read {
                &Err(ref e) => { self.client_writer.stream.shutdown(Shutdown::Write).unwrap(); },
                _ => {}
            }

            // process buffered responses
            while let Some(response) = self.server_reader.next() {
                match self.handler.handle_response(&response) {
                    Action::Forward => self.client_writer.push(&response),
                    Action::Mutate(ref p2) => self.client_writer.push(p2),
                    Action::Respond(ref v) => {
                        for p in v {
                            self.server_writer.push(&p);
                        }
                    }
                };
            }

            // try writing to client
            let client_write = self.client_writer.write();

            try_ready!(client_read);
            try_ready!(client_write);
            try_ready!(server_read);
            try_ready!(server_write);
        }

    }

}

#[allow(dead_code)]
pub fn print_packet_chars(buf: &[u8]) {
    print!("[");
    for i in 0..buf.len() {
        print!("{} ", buf[i] as char);
    }
    println!("]");
}

#[allow(dead_code)]
pub fn print_packet_bytes(buf: &[u8]) {
    print!("[");
    for i in 0..buf.len() {
        if i%8==0 { println!(""); }
        print!("{:#04x} ",buf[i]);
    }
    println!("]");
}

<commit_msg>perform all writes at end of poll method<commit_after>//! MySQL Proxy Server

#[macro_use]
extern crate log;
extern crate env_logger;
#[macro_use]
extern crate futures;
#[macro_use]
extern crate tokio_core;
extern crate futures_cpupool;
extern crate byteorder;

use std::cell::RefCell;
use std::rc::Rc;
use std::env;
use std::io::{self, Read, Write, Error, ErrorKind};
use std::net::{SocketAddr, Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
use std::net::Shutdown;
use std::str;
use std::time::Duration;

use futures::{Future, Poll, Async};
use futures::stream::Stream;
use futures_cpupool::CpuPool;
use tokio_core::{Loop, LoopHandle, TcpStream};
use tokio_core::io::{read_exact, write_all, Window};

use byteorder::*;

#[derive(Copy,Clone)]
pub enum PacketType {
    ComSleep = 0x00,
    ComQuit = 0x01,
    ComInitDb = 0x02,
    ComQuery = 0x03,
    ComFieldList = 0x04,
    ComCreateDb = 0x05,
    ComDropDb = 0x06,
    ComRefresh = 0x07,
    ComShutdown = 0x08,
    ComStatistics = 0x09,
    ComProcessInfo = 0x0a,
    ComConnect = 0x0b,
    ComProcessKill= 0x0c,
    ComDebug = 0x0d,
    ComPing = 0x0e,
    ComTime = 0x0f,
    ComDelayedInsert = 0x10,
    ComChangeUser = 0x11,
    ComBinlogDump = 0x12,
    ComTableDump = 0x13,
    ComConnectOut = 0x14,
    ComRegisterSlave = 0x15,
    ComStmtPrepare = 0x16,
    ComStmtExecute = 0x17,
    ComStmtSendLongData = 0x18,
    ComStmtClose = 0x19,
    ComStmtReset = 0x1a,
    ComDaemon= 0x1d,
    ComBinlogDumpGtid = 0x1e,
    ComResetConnection = 0x1f,
}

pub struct Packet {
    pub bytes: Vec<u8>
}

impl Packet {

    pub fn error_packet(code: u16, state: [u8; 5], msg: String) -> Self {

        // start building payload
        let mut payload: Vec<u8> = Vec::with_capacity(9 + msg.len());
        payload.push(0xff);  // packet type
        payload.write_u16::<LittleEndian>(code).unwrap(); // error code
        payload.extend_from_slice("#".as_bytes()); // sql_state_marker
        payload.extend_from_slice(&state); // SQL STATE
        payload.extend_from_slice(msg.as_bytes());

        // create header with length and sequence id
        let mut header: Vec<u8> = Vec::with_capacity(4 + 9 + msg.len());
        header.write_u32::<LittleEndian>(payload.len() as u32).unwrap();
        header.pop(); // we need 3 byte length, so discard last byte
        header.push(1); // sequence_id

        // combine the vectors
        header.extend_from_slice(&payload);

        // now move the vector into the packet
        Packet { bytes: header }
    }

    pub fn sequence_id(&self) -> u8 {
        self.bytes[3]
    }

    //TODO: should return Result<PacketType, ?>
    pub fn packet_type(&self) -> PacketType {
        match self.bytes[4] {
            0x00 => PacketType::ComSleep,
            0x01 => PacketType::ComQuit,
            0x02 => PacketType::ComInitDb,
            0x03 => PacketType::ComQuery,
            0x04 => PacketType::ComFieldList,
            0x05 => PacketType::ComCreateDb,
            0x06 => PacketType::ComDropDb,
            0x07 => PacketType::ComRefresh,
            0x08 => PacketType::ComShutdown,
            0x09 => PacketType::ComStatistics,
            0x0a => PacketType::ComProcessInfo,
            0x0b => PacketType::ComConnect,
            0x0c => PacketType::ComProcessKill,
            0x0d => PacketType::ComDebug,
            0x0e => PacketType::ComPing,
            0x0f => PacketType::ComTime,
            0x10 => PacketType::ComDelayedInsert,
            0x11 => PacketType::ComChangeUser,
            0x12 => PacketType::ComBinlogDump,
            0x13 => PacketType::ComTableDump,
            0x14 => PacketType::ComConnectOut,
            0x15 => PacketType::ComRegisterSlave,
            0x16 => PacketType::ComStmtPrepare,
            0x17 => PacketType::ComStmtExecute,
            0x18 => PacketType::ComStmtSendLongData,
            0x19 => PacketType::ComStmtClose,
            0x1a => PacketType::ComStmtReset,
            0x1d => PacketType::ComDaemon,
            0x1e => PacketType::ComBinlogDumpGtid,
            0x1f => PacketType::ComResetConnection,
            _ => panic!("Unsupported packet type")
        }
    }

}

pub enum Action {
    Forward,                // forward the packet unmodified
    Mutate(Packet),         // mutate the packet
    Respond(Vec<Packet>)    // handle directly and optionally return some packets
}

pub trait PacketHandler {
    fn handle_request(&self, p: &Packet) -> Action;
    fn handle_response(&self, p: &Packet) -> Action;
}

fn parse_packet_length(header: &[u8]) -> usize {
    (((header[2] as u32) << 16) |
    ((header[1] as u32) << 8) |
    header[0] as u32) as usize
}

pub struct Client {
    pub pool: CpuPool,
    pub handle: LoopHandle,
}

impl Client {

    pub fn serve<H: PacketHandler + 'static>(self, conn: TcpStream,
                                             handler: H,
                                        ) -> Box<Future<Item = (u64, u64), Error = Error>> {

        let addr = Ipv4Addr::new(127, 0, 0, 1);
        let port = 3306;
        let addr = SocketAddr::V4(SocketAddrV4::new(addr, port));

        let handle = self.handle.clone();
        let connected = handle.tcp_connect(&addr);

        let pair = connected.and_then(move |server| Ok((conn,server)));

        Box::new(pair.and_then(move |(c1, c2)| {
            let c1 = Rc::new(c1);
            let c2 = Rc::new(c2);
            Pipe::new(c1, c2, handler)
        }))
    }
}

#[derive(Debug)]
enum Direction {
    Client,
    Server
}

struct ConnReader {
    direction: Direction,
    stream: Rc<TcpStream>,
    read_buf: Vec<u8>,
    read_pos: usize,
    read_amt: u64,
}

struct ConnWriter {
    direction: Direction,
    stream: Rc<TcpStream>,
    write_buf: Vec<u8>,
    write_pos: usize,
    write_amt: u64,
}

impl ConnReader {

    fn new(stream: Rc<TcpStream>, direction: Direction) -> Self {
        ConnReader {
            stream: stream,
            direction: direction,
            read_buf: vec![0u8; 4096],
            read_pos: 0,
            read_amt: 0,
        }
    }

    /// Read from the socket until the status is NotReady
    fn read(&mut self) -> Poll<(), io::Error> {
        loop {
            try_ready!(self.stream.poll_read());
            //TODO: ensure capacity first
            let n = try_nb!((&*self.stream).read(&mut self.read_buf[self.read_pos..]));
            if n == 0 {
                return Err(Error::new(ErrorKind::Other, "connection closed"));
            }
            self.read_amt += n as u64;
            self.read_pos += n;
        }
    }

    fn next(&mut self) -> Option<Packet> {
        // do we have a header
        if self.read_pos > 3 {
            let l = parse_packet_length(&self.read_buf);
            // do we have the whole packet?
            let s = 4 + l;
            if self.read_pos >= s {
                let mut temp : Vec<u8> = Vec::with_capacity(s);
                temp.extend_from_slice(&self.read_buf[0..s]);
                let p = Packet { bytes: temp };

                // shift data down
                let mut j = 0;
                for i in s .. self.read_pos {
                    self.read_buf[j] = self.read_buf[i];
                    j += 1;
                }
                self.read_pos -= s;

                Some(p)
            } else {
                None
            }
        } else {
            None
        }
    }
}

impl ConnWriter {

    fn new(stream: Rc<TcpStream>, direction: Direction) -> Self {
        ConnWriter{
            stream: stream,
            direction: direction,
            write_buf: vec![0u8; 4096],
            write_pos: 0,
            write_amt: 0
        }
    }

    /// Write a packet to the write buffer
    fn push(&mut self, p: &Packet) {
        for i in 0 .. p.bytes.len() {
            self.write_buf[self.write_pos + i] = p.bytes[i];
        }
        self.write_pos += p.bytes.len();
    }

    /// Writes the contents of the write buffer to the socket
    fn write(&mut self) -> Poll<(), io::Error> {
        while self.write_pos > 0 {
            try_ready!(self.stream.poll_write());
            let s = try!((&*self.stream).write(&self.write_buf[0..self.write_pos]));

            let mut j = 0;
            for i in s .. self.write_pos {
                self.write_buf[j] = self.write_buf[i];
                j += 1;
            }
            self.write_pos -= s;

        }
        return Ok(Async::Ready(()));
    }
}

struct Pipe<H: PacketHandler + 'static> {
    client_reader: ConnReader,
    client_writer: ConnWriter,
    server_reader: ConnReader,
    server_writer: ConnWriter,
    handler: H,
}

impl<H> Pipe<H> where H: PacketHandler + 'static {
    fn new(client: Rc<TcpStream>,
           server: Rc<TcpStream>,
           handler: H
           ) -> Pipe<H> {

        Pipe {
            client_reader: ConnReader::new(client.clone(), Direction::Client),
            client_writer: ConnWriter::new(client, Direction::Client),
            server_reader: ConnReader::new(server.clone(), Direction::Server),
            server_writer: ConnWriter::new(server, Direction::Server),
            handler: handler,
        }
    }
}

impl<H> Future for Pipe<H> where H: PacketHandler + 'static {
    type Item = (u64, u64);
    type Error = Error;

    fn poll(&mut self) -> Poll<(u64, u64), io::Error> {
        loop {
            let client_read = self.client_reader.read();

            // if the client connection has closed, close the server connection too
            match &client_read {
                &Err(ref e) => { self.server_writer.stream.shutdown(Shutdown::Write).unwrap(); },
                _ => {}
            }

            // process buffered requests
            while let Some(request) = self.client_reader.next() {
                match self.handler.handle_request(&request) {
                    Action::Forward => self.server_writer.push(&request),
                    Action::Mutate(ref p2) => self.server_writer.push(p2),
                    Action::Respond(ref v) => {
                        for p in v {
                            self.client_writer.push(&p);
                        }
                    }
                };
            }

            // try reading from server
            let server_read = self.server_reader.read();

            // if the server connection has closed, close the client connection too
            match &server_read {
                &Err(ref e) => { self.client_writer.stream.shutdown(Shutdown::Write).unwrap(); },
                _ => {}
            }

            // process buffered responses
            while let Some(response) = self.server_reader.next() {
                match self.handler.handle_response(&response) {
                    Action::Forward => self.client_writer.push(&response),
                    Action::Mutate(ref p2) => self.client_writer.push(p2),
                    Action::Respond(ref v) => {
                        for p in v {
                            self.server_writer.push(&p);
                        }
                    }
                };
            }

            // perform all of the writes at the end, since the request handlers may have
            // queued packets in either, or both directions

            // try writing to client
            let client_write = self.client_writer.write();

            // try writing to server
            let server_write = self.server_writer.write();

            try_ready!(client_read);
            try_ready!(client_write);
            try_ready!(server_read);
            try_ready!(server_write);
        }

    }

}

#[allow(dead_code)]
pub fn print_packet_chars(buf: &[u8]) {
    print!("[");
    for i in 0..buf.len() {
        print!("{} ", buf[i] as char);
    }
    println!("]");
}

#[allow(dead_code)]
pub fn print_packet_bytes(buf: &[u8]) {
    print!("[");
    for i in 0..buf.len() {
        if i%8==0 { println!(""); }
        print!("{:#04x} ",buf[i]);
    }
    println!("]");
}

<|endoftext|>


<commit_before>/*
 * Copyright (c) 2016 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

//! Parse config files.
//!
//! # Usage
//!
//! Call the `parse` function on the input.

/*
 * TODO: create a new method instead of using words and find/unwrap.
 * TODO: add the attribute #[special_command] in the Command derive macro.
 * TODO: auto-include files.
 * TODO: support set = without spaces around =.
 * TODO: Add array type.
 */

#![warn(missing_docs)]

pub mod error;
pub mod key;
#[macro_use]
mod macros;
#[doc(hidden)]
pub mod position;
pub mod settings;
mod string;

use std::collections::HashMap;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::marker::PhantomData;
use std::path::{Path, PathBuf};

use error::{Error, Result};
use error::ErrorType::{MissingArgument, NoCommand, Parse, UnknownCommand};
use key::{Key, parse_keys};
use position::Pos;
use string::{StrExt, check_ident, maybe_word, word, words};

use Command::*;
use Value::*;

/// Trait to specify the completion values for a type.
pub trait CompletionValues {
    /// Get the completion values for the type.
    fn completion_values() -> Vec<String>;
}

impl CompletionValues for bool {
    fn completion_values() -> Vec<String> {
        vec!["true".to_string(), "false".to_string()]
    }
}

impl CompletionValues for i64 {
    fn completion_values() -> Vec<String> {
        vec![]
    }
}

impl CompletionValues for String {
    fn completion_values() -> Vec<String> {
        vec![]
    }
}

/// The `EnumFromStr` trait is used to specify how to construct an enum value from a string.
pub trait EnumFromStr
    where Self: Sized
{
    /// Create the enum value from the `variant` string and an `argument` string.
    fn create(variant: &str, argument: &str) -> std::result::Result<Self, String>;

    /// Check wether the enum variant has an argument.
    fn has_argument(variant: &str) -> std::result::Result<bool, String>;
}

/// Tre `EnumMetaData` trait is used to get associated meta-data for the enum variants.
/// The meta-data is specified using the following attributes:
/// #[completion(hidden)]
/// #[special_command]
/// #[help(Command help)]
pub trait EnumMetaData {
    /// Get the metadata associated with the enum.
    fn get_metadata() -> HashMap<String, MetaData>;
}

/// Command/setting meta-data coming from the attributes.
/// See `EnumMetaData` to see the list of supported attributes.
#[derive(Debug)]
pub struct MetaData {
    /// Whether this command/setting should be shown in the completion or not.
    pub completion_hidden: bool,
    /// The help text associated with this command/setting.
    pub help_text: String,
    /// Whether this is a special command or not.
    /// This is not applicable to settings.
    pub is_special_command: bool,
}

/// Trait specifying the value completions for settings.
pub trait SettingCompletion {
    /// Get the value completions of all the setting.
    fn get_value_completions() -> HashMap<String, Vec<String>>;
}

/// The `Command` enum represents a command from a config file.
#[derive(Debug, PartialEq)]
pub enum Command<T> {
    /// A command from the application library.
    App(String),
    /// A custom command.
    Custom(T),
    /// A map command creates a new key mapping.
    Map {
        /// The action that will be executed when the `keys` are pressed.
        action: String,
        /// The key shortcut to trigger the action.
        keys: Vec<Key>,
        /// The mode in which this mapping is available.
        mode: String,
    },
    /// A set command sets a value to an option.
    Set(String, Value),
    /// An unmap command removes a key mapping.
    Unmap {
        /// The key shortcut to remove.
        keys: Vec<Key>,
        /// The mode in which this mapping is available.
        mode: String,
    },
}

/// The parsing configuration.
#[derive(Default)]
pub struct Config {
    /// The application library commands.
    pub application_commands: Vec<String>,
    /// The available mapping modes for the map command.
    pub mapping_modes: Vec<String>,
}

/// The config parser.
pub struct Parser<T> {
    column: usize,
    config: Config,
    include_path: PathBuf,
    line: usize,
    _phantom: PhantomData<T>,
}

impl<T: EnumFromStr> Parser<T> {
    /// Create a new parser without config.
    pub fn new() -> Self {
        Parser {
            column: 1,
            config: Config::default(),
            include_path: Path::new("./").to_path_buf(),
            line: 1,
            _phantom: PhantomData,
        }
    }

    /// Create a new parser with config.
    pub fn new_with_config(config: Config) -> Self {
        Parser {
            column: 1,
            config: config,
            include_path: Path::new("./").to_path_buf(),
            line: 1,
            _phantom: PhantomData,
        }
    }

    /// Check that we reached the end of the line.
    fn check_eol(&self, line: &str, index: usize) -> Result<()> {
        if line.len() > index {
            let rest = &line[index..];
            if let Some(word) = maybe_word(rest) {
                let index = rest.find(word).unwrap(); // NOTE: the line contains the word, hence unwrap.
                return Err(Box::new(Error::new(
                    Parse,
                    rest.to_string(),
                    "<end of line>".to_string(),
                    Pos::new(self.line, self.column + index),
                )));
            }
        }
        Ok(())
    }

    /// Parse a custom command or return an error if it does not exist.
    fn custom_command(&self, line: &str, word: &str, start_index: usize, index: usize) -> Result<Command<T>> {
        let args =
            if line.len() > start_index {
                line[start_index..].trim()
            }
            else if let Ok(true) = T::has_argument(word) {
                return Err(Box::new(self.missing_args(start_index)));
            }
            else {
                ""
            };
        if let Ok(command) = T::create(word, args) {
            Ok(Custom(command))
        }
        else if self.config.application_commands.contains(&word.to_string()) {
            Ok(App(word.to_string()))
        }
        else {
            Err(Box::new(Error::new(
                UnknownCommand,
                word.to_string(),
                "command or comment".to_string(),
                Pos::new(self.line, index + 1)
            )))
        }
    }

    /// Get the rest of the line, starting at `column`, returning an error if the column is greater
    /// than the line's length.
    fn get_rest<'a>(&self, line: &'a str, column: usize) -> Result<&'a str> {
        if line.len() > column {
            Ok(&line[column..])
        }
        else {
            Err(Box::new(self.missing_args(column)))
        }
    }

    /// Parse a line.
    fn line(&mut self, line: &str) -> Result<Vec<Command<T>>> {
        if let Some(word) = maybe_word(line) {
            // NOTE: the word is in the line, hence unwrap.
            let index = line.find(word).unwrap();
            let start_index = index + word.len() + 1;
            self.column = start_index + 1;

            let (start3, end3) = word.rsplit_at(3);
            let (start5, end5) = word.rsplit_at(5);
            if word.starts_with('#') {
                return Ok(vec![]);
            }

            if word == "include" {
                let rest = self.get_rest(line, start_index)?;
                self.include_command(rest)
            }
            else {
                let command =
                    if word == "set" {
                        let rest = self.get_rest(line, start_index)?;
                        self.set_command(rest)
                    }
                    else if end3 == "map" && self.config.mapping_modes.contains(&start3.to_string()) {
                        let rest = self.get_rest(line, start_index)?;
                        self.map_command(rest, start3)
                    }
                    else if end5 == "unmap" && self.config.mapping_modes.contains(&start5.to_string()) {
                        let rest = self.get_rest(line, start_index)?;
                        self.unmap_command(rest, start5)
                    }
                    else {
                        self.custom_command(line, word, start_index, index)
                    };
                command.map(|command| vec![command])
            }
        }
        else {
            Ok(vec![])
        }
    }

    /// Parse an include command.
    fn include_command(&mut self, line: &str) -> Result<Vec<Command<T>>> {
        let word = word(line);
        // NOTE: the line contains the word, hence unwrap.
        let index = line.find(word).unwrap();
        let after_index = index + word.len() + 1;
        self.column += after_index;
        self.check_eol(line, after_index)?;
        let path = Path::new(&self.include_path).join(word);
        let file = File::open(path)?;
        let buf_reader = BufReader::new(file);
        let commands = self.parse(buf_reader)?;
        Ok(commands)
    }

    /// Parse a map command.
    fn map_command(&self, line: &str, mode: &str) -> Result<Command<T>> {
        let word = word(line);

        // NOTE: the line contains the word, hence unwrap.
        let index = line.find(word).unwrap();

        let rest = &line[index + word.len() ..].trim();
        if !rest.is_empty() {
            Ok(Map {
                action: rest.to_string(),
                keys: parse_keys(word, self.line, self.column + index)?,
                mode: mode.to_string(),
            })
        }
        else {
            Err(Box::new(Error::new(
                Parse,
                "<end of line>".to_string(),
                "mapping action".to_string(),
                Pos::new(self.line, self.column + line.len())
            )))
        }
    }

    /// Get an missing arguments error.
    fn missing_args(&self, column: usize) -> Error {
        Error::new(
            MissingArgument,
            "<end of line>".to_string(),
            "command arguments".to_string(),
            Pos::new(self.line, column)
        )
    }

    /// Parse settings.
    pub fn parse<R: BufRead>(&mut self, input: R) -> Result<Vec<Command<T>>> {
        let mut commands = vec![];
        for (line_num, input_line) in input.lines().enumerate() {
            self.line = line_num + 1;
            let mut new_commands = self.line(&input_line?)?;
            commands.append(&mut new_commands);
        }
        Ok(commands)
    }

    /// Parse a single line of settings.
    pub fn parse_line(&mut self, line: &str) -> Result<Command<T>> {
        let mut commands = self.parse(line.as_bytes())?;
        match commands.pop() {
            Some(command) => Ok(command),
            None => Err(Box::new(Error::new(
                        NoCommand,
                        "comment or <end of line>".to_string(),
                        "command".to_string(),
                        Pos::new(self.line, self.column + line.len())
                    )))
        }
    }

    /// Parse a set command.
    fn set_command(&mut self, line: &str) -> Result<Command<T>> {
        if let Some(words) = words(line, 2) {
            // NOTE: the line contains the word, hence unwrap.
            let index = line.find(words[0]).unwrap();
            let identifier = check_ident(words[0].to_string(), &Pos::new(self.line, self.column + index))?;

            let operator = words[1];
            // NOTE: the operator is in the line, hence unwrap.
            let operator_index = line.find(operator).unwrap();
            if operator == "=" {
                let rest = &line[operator_index + 1..];
                self.column += operator_index + 1;
                Ok(Set(identifier.to_string(), self.value(rest)?))
            }
            else {
                Err(Box::new(Error::new(
                    Parse,
                    operator.to_string(),
                    "=".to_string(),
                    Pos::new(self.line, self.column + operator_index)
                )))
            }
        }
        else {
            Err(Box::new(Error::new(
                Parse,
                "<end of line>".to_string(),
                "=".to_string(),
                Pos::new(self.line, self.column + line.len()),
            )))
        }
    }

    /// Set the directory where the include command will look for files to include.
    pub fn set_include_path<P: AsRef<Path>>(&mut self, directory: P) {
        self.include_path = directory.as_ref().to_path_buf();
    }

    /// Parse an unmap command.
    fn unmap_command(&mut self, line: &str, mode: &str) -> Result<Command<T>> {
        let word = word(line);

        // NOTE: the line contains the word, hence unwrap.
        let index = line.find(word).unwrap();

        let after_index = index + word.len() + 1;
        self.column += after_index;
        self.check_eol(line, after_index)?;
        Ok(Unmap {
            keys: parse_keys(word, self.line, self.column + index)?,
            mode: mode.to_string(),
        })
    }

    /// Parse a value.
    fn value(&self, input: &str) -> Result<Value> {
        let string: String = input.chars().take_while(|&character| character != '#').collect();
        let string = string.trim();
        match string {
            "" => Err(Box::new(Error::new(
                      Parse,
                      "<end of line>".to_string(),
                      "value".to_string(),
                      Pos::new(self.line, self.column + string.len())
                  ))),
            "true" => Ok(Bool(true)),
            "false" => Ok(Bool(false)),
            _ => {
                if string.chars().all(|character| character.is_digit(10)) {
                    // NOTE: the string only contains digit, hence unwrap.
                    Ok(Int(string.parse().unwrap()))
                }
                else if string.chars().all(|character| character.is_digit(10) || character == '.') {
                    // NOTE: the string only contains digit or dot, hence unwrap.
                    Ok(Float(string.parse().unwrap()))
                }
                else {
                    Ok(Str(input.trim().to_string()))
                }
            },
        }
    }
}

/// The `Value` enum represents a value along with its type.
#[derive(Debug, PartialEq)]
pub enum Value {
    /// Boolean value.
    Bool(bool),
    /// Floating-point value.
    Float(f64),
    /// Integer value.
    Int(i64),
    /// String value.
    Str(String),
}

impl Value {
    /// Get a string representation of the value.
    pub fn to_type(&self) -> &str {
        match *self {
            Bool(_) => "bool",
            Float(_) => "float",
            Int(_) => "int",
            Str(_) => "string",
        }
    }
}
<commit_msg>Dummy commit.<commit_after>/*
 * Copyright (c) 2016 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

//! Parse config files.
//!
//! # Usage
//!
//! Call the `parse` function on the input.

/*
 * TODO: create a new method instead of using words and find/unwrap.
 * TODO: add the attribute #[special_command] in the Command derive macro.
 * TODO: auto-include files.
 * TODO: support set = without spaces around =.
 * TODO: Add array type.
 */

#![warn(missing_docs)]

pub mod error;
pub mod key;
#[macro_use]
mod macros;
#[doc(hidden)]
pub mod position;
pub mod settings;
mod string;

use std::collections::HashMap;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::marker::PhantomData;
use std::path::{Path, PathBuf};

use error::{Error, Result};
use error::ErrorType::{MissingArgument, NoCommand, Parse, UnknownCommand};
use key::{Key, parse_keys};
use position::Pos;
use string::{StrExt, check_ident, maybe_word, word, words};

use Command::*;
use Value::*;

/// Trait to specify the completion values for a type.
pub trait CompletionValues {
    /// Get the completion values for the type.
    fn completion_values() -> Vec<String>;
}

impl CompletionValues for bool {
    fn completion_values() -> Vec<String> {
        vec!["true".to_string(), "false".to_string()]
    }
}

impl CompletionValues for i64 {
    fn completion_values() -> Vec<String> {
        vec![]
    }
}

impl CompletionValues for String {
    fn completion_values() -> Vec<String> {
        vec![]
    }
}

/// The `EnumFromStr` trait is used to specify how to construct an enum value from a string.
pub trait EnumFromStr
    where Self: Sized
{
    /// Create the enum value from the `variant` string and an `argument` string.
    fn create(variant: &str, argument: &str) -> std::result::Result<Self, String>;

    /// Check wether the enum variant has an argument.
    fn has_argument(variant: &str) -> std::result::Result<bool, String>;
}

/// Tre `EnumMetaData` trait is used to get associated meta-data for the enum variants.
/// The meta-data is specified using the following attributes:
/// #[completion(hidden)]
/// #[special_command]
/// #[help(Command help)]
pub trait EnumMetaData {
    /// Get the metadata associated with the enum.
    fn get_metadata() -> HashMap<String, MetaData>;
}

/// Command/setting meta-data coming from the attributes.
/// See `EnumMetaData` to see the list of supported attributes.
#[derive(Debug)]
pub struct MetaData {
    /// Whether this command/setting should be shown in the completion or not.
    pub completion_hidden: bool,
    /// The help text associated with this command/setting.
    pub help_text: String,
    /// Whether this is a special command or not.
    /// This is not applicable to settings.
    pub is_special_command: bool,
}

/// Trait specifying the value completions for settings.
pub trait SettingCompletion {
    /// Get the value completions of all the setting.
    fn get_value_completions() -> HashMap<String, Vec<String>>;
}

/// The `Command` enum represents a command from a config file.
#[derive(Debug, PartialEq)]
pub enum Command<T> {
    /// A command from the application library.
    App(String),
    /// A custom command.
    Custom(T),
    /// A map command creates a new key mapping.
    Map {
        /// The action that will be executed when the `keys` are pressed.
        action: String,
        /// The key shortcut to trigger the action.
        keys: Vec<Key>,
        /// The mode in which this mapping is available.
        mode: String,
    },
    /// A set command sets a value to an option.
    Set(String, Value),
    /// An unmap command removes a key mapping.
    Unmap {
        /// The key shortcut to remove.
        keys: Vec<Key>,
        /// The mode in which this mapping is available.
        mode: String,
    },
}

/// The parsing configuration.
#[derive(Default)]
pub struct Config {
    /// The application library commands.
    pub application_commands: Vec<String>,
    /// The available mapping modes for the map command.
    pub mapping_modes: Vec<String>,
}

/// The config parser.
pub struct Parser<T> {
    column: usize,
    config: Config,
    include_path: PathBuf,
    line: usize,
    _phantom: PhantomData<T>,
}

impl<T: EnumFromStr> Parser<T> {
    /// Create a new parser without config.
    pub fn new() -> Self {
        Parser {
            column: 1,
            config: Config::default(),
            include_path: Path::new("./").to_path_buf(),
            line: 1,
            _phantom: PhantomData,
        }
    }

    /// Create a new parser with config.
    pub fn new_with_config(config: Config) -> Self {
        Parser {
            column: 1,
            config: config,
            include_path: Path::new("./").to_path_buf(),
            line: 1,
            _phantom: PhantomData,
        }
    }

    /// Check that we reached the end of the line.
    fn check_eol(&self, line: &str, index: usize) -> Result<()> {
        if line.len() > index {
            let rest = &line[index..];
            if let Some(word) = maybe_word(rest) {
                let index = rest.find(word).unwrap(); // NOTE: the line contains the word, hence unwrap.
                return Err(Box::new(Error::new(
                    Parse,
                    rest.to_string(),
                    "<end of line>".to_string(),
                    Pos::new(self.line, self.column + index),
                )));
            }
        }
        Ok(())
    }

    /// Parse a custom command or return an error if it does not exist.
    fn custom_command(&self, line: &str, word: &str, start_index: usize, index: usize) -> Result<Command<T>> {
        let args =
            if line.len() > start_index {
                line[start_index..].trim()
            }
            else if let Ok(true) = T::has_argument(word) {
                return Err(Box::new(self.missing_args(start_index)));
            }
            else {
                ""
            };
        if let Ok(command) = T::create(word, args) {
            Ok(Custom(command))
        }
        else if self.config.application_commands.contains(&word.to_string()) {
            Ok(App(word.to_string()))
        }
        else {
            Err(Box::new(Error::new(
                UnknownCommand,
                word.to_string(),
                "command or comment".to_string(),
                Pos::new(self.line, index + 1)
            )))
        }
    }

    /// Get the rest of the line, starting at `column`, returning an error if the column is greater
    /// than the line's length.
    fn get_rest<'a>(&self, line: &'a str, column: usize) -> Result<&'a str> {
        if line.len() > column {
            Ok(&line[column..])
        }
        else {
            Err(Box::new(self.missing_args(column)))
        }
    }

    /// Parse a line.
    fn line(&mut self, line: &str) -> Result<Vec<Command<T>>> {
        if let Some(word) = maybe_word(line) {
            // NOTE: the word is in the line, hence unwrap.
            let index = line.find(word).unwrap();
            let start_index = index + word.len() + 1;
            self.column = start_index + 1;

            let (start3, end3) = word.rsplit_at(3);
            let (start5, end5) = word.rsplit_at(5);
            if word.starts_with('#') {
                return Ok(vec![]);
            }

            if word == "include" {
                let rest = self.get_rest(line, start_index)?;
                self.include_command(rest)
            }
            else {
                let command =
                    if word == "set" {
                        let rest = self.get_rest(line, start_index)?;
                        self.set_command(rest)
                    }
                    else if end3 == "map" && self.config.mapping_modes.contains(&start3.to_string()) {
                        let rest = self.get_rest(line, start_index)?;
                        self.map_command(rest, start3)
                    }
                    else if end5 == "unmap" && self.config.mapping_modes.contains(&start5.to_string()) {
                        let rest = self.get_rest(line, start_index)?;
                        self.unmap_command(rest, start5)
                    }
                    else {
                        self.custom_command(line, word, start_index, index)
                    };
                command.map(|command| vec![command])
            }
        }
        else {
            Ok(vec![])
        }
    }

    /// Parse an include command.
    fn include_command(&mut self, line: &str) -> Result<Vec<Command<T>>> {
        let word = word(line);
        // NOTE: the line contains the word, hence unwrap.
        let index = line.find(word).unwrap();
        let after_index = index + word.len() + 1;
        self.column += after_index;
        self.check_eol(line, after_index)?;
        let path = Path::new(&self.include_path).join(word);
        let file = File::open(path)?;
        let buf_reader = BufReader::new(file);
        let commands = self.parse(buf_reader)?;
        Ok(commands)
    }

    /// Parse a map command.
    fn map_command(&self, line: &str, mode: &str) -> Result<Command<T>> {
        let word = word(line);

        // NOTE: the line contains the word, hence unwrap.
        let index = line.find(word).unwrap();

        let rest = &line[index + word.len() ..].trim();
        if !rest.is_empty() {
            Ok(Map {
                action: rest.to_string(),
                keys: parse_keys(word, self.line, self.column + index)?,
                mode: mode.to_string(),
            })
        }
        else {
            Err(Box::new(Error::new(
                Parse,
                "<end of line>".to_string(),
                "mapping action".to_string(),
                Pos::new(self.line, self.column + line.len())
            )))
        }
    }

    /// Get an missing arguments error.
    fn missing_args(&self, column: usize) -> Error {
        Error::new(
            MissingArgument,
            "<end of line>".to_string(),
            "command arguments".to_string(),
            Pos::new(self.line, column)
        )
    }

    /// Parse settings.
    pub fn parse<R: BufRead>(&mut self, input: R) -> Result<Vec<Command<T>>> {
        let mut commands = vec![];
        for (line_num, input_line) in input.lines().enumerate() {
            self.line = line_num + 1;
            let mut new_commands = self.line(&input_line?)?;
            commands.append(&mut new_commands);
        }
        Ok(commands)
    }

    /// Parse a single line of settings.
    pub fn parse_line(&mut self, line: &str) -> Result<Command<T>> {
        let mut commands = self.parse(line.as_bytes())?;
        match commands.pop() {
            Some(command) => Ok(command),
            None => Err(Box::new(Error::new(
                        NoCommand,
                        "comment or <end of line>".to_string(),
                        "command".to_string(),
                        Pos::new(self.line, self.column + line.len())
                    )))
        }
    }

    /// Parse a set command.
    fn set_command(&mut self, line: &str) -> Result<Command<T>> {
        if let Some(words) = words(line, 2) {
            // NOTE: the line contains the word, hence unwrap.
            let index = line.find(words[0]).unwrap();
            let identifier = check_ident(words[0].to_string(), &Pos::new(self.line, self.column + index))?;

            let operator = words[1];
            // NOTE: the operator is in the line, hence unwrap.
            let operator_index = line.find(operator).unwrap();
            if operator == "=" {
                let rest = &line[operator_index + 1..];
                self.column += operator_index + 1;
                Ok(Set(identifier.to_string(), self.value(rest)?))
            }
            else {
                Err(Box::new(Error::new(
                    Parse,
                    operator.to_string(),
                    "=".to_string(),
                    Pos::new(self.line, self.column + operator_index)
                )))
            }
        }
        else {
            Err(Box::new(Error::new(
                Parse,
                "<end of line>".to_string(),
                "=".to_string(),
                Pos::new(self.line, self.column + line.len()),
            )))
        }
    }

    /// Set the directory where the include command will look for files to include.
    pub fn set_include_path<P: AsRef<Path>>(&mut self, directory: P) {
        self.include_path = directory.as_ref().to_path_buf();
    }

    /// Parse an unmap command.
    fn unmap_command(&mut self, line: &str, mode: &str) -> Result<Command<T>> {
        let word = word(line);

        // NOTE: the line contains the word, hence unwrap.
        let index = line.find(word).unwrap();

        let after_index = index + word.len() + 1;
        self.column += after_index;
        self.check_eol(line, after_index)?;
        Ok(Unmap {
            keys: parse_keys(word, self.line, self.column + index)?,
            mode: mode.to_string(),
        })
    }

    /// Parse a value.
    fn value(&self, input: &str) -> Result<Value> {
        let string: String = input.chars().take_while(|&character| character != '#').collect();
        let string = string.trim();
        match string {
            "" => Err(Box::new(Error::new(
                      Parse,
                      "<end of line>".to_string(),
                      "value".to_string(),
                      Pos::new(self.line, self.column + string.len())
                  ))),
            "true" => Ok(Bool(true)),
            "false" => Ok(Bool(false)),
            _ => {
                if string.chars().all(|character| character.is_digit(10)) {
                    // NOTE: the string only contains digit, hence unwrap.
                    Ok(Int(string.parse().unwrap()))
                }
                else if string.chars().all(|character| character.is_digit(10) || character == '.') {
                    // NOTE: the string only contains digit or dot, hence unwrap.
                    Ok(Float(string.parse().unwrap()))
                }
                else {
                    Ok(Str(input.trim().to_string()))
                }
            },
        }
    }
}

/// The `Value` enum represents a value along with its type.
#[derive(Debug, PartialEq)]
pub enum Value {
    /// Boolean value.
    Bool(bool),
    /// Floating-point value.
    Float(f64),
    /// Integer value.
    Int(i64),
    /// String value.
    Str(String),
}

impl Value {
    /// Get a string representation of the value.
    pub fn to_type(&self) -> &str {
        match *self {
            Bool(_) => "bool",
            Float(_) => "float",
            Int(_) => "int",
            Str(_) => "string",
        }
    }
}
<|endoftext|>


<commit_before>#![feature(lang_items)]
#![feature(const_fn)]
#![feature(unique)]
#![feature(alloc)]
#![feature(asm)]
#![feature(core_intrinsics)]
#![feature(naked_functions)]
#![feature(collections)]
#![feature(drop_types_in_const)]
#![feature(box_syntax)]

#![no_std]

#[macro_use]
extern crate x86;
extern crate rlibc;
extern crate spin;
extern crate multiboot2;
extern crate alloc_opsys;
extern crate alloc;

#[macro_use]
extern crate collections;

#[macro_use]
extern crate bitflags;

#[macro_use]
extern crate once;

#[macro_use]
mod vga_buffer;
mod memory;
mod interrupts;
mod drivers;
mod io;
mod schedule;
mod kernel;

// Main entry point, need no_mangle so we can call from assembly
// Extern to abide with C calling convention
#[no_mangle]
pub extern "C" fn kernel_main(multiboot_info_address: usize) {
    // Initialise the hardware
    init_cpu();
    memory::init(multiboot_info_address);

    // Setup the heap allocator
    alloc_opsys::init();

    // Setup the kernel
    kernel::init();

    // Initialize interrupts
    interrupts::init();

    vga_buffer::clear_screen();
    kprintln!("OpSys v{}", "0.0.1");

    loop {
    }
}

fn init_cpu() {
    // Configure the CPU flags as required
    enable_nxe_bit(); // Allow the NO_EXECUTABLE flag
    enable_write_protect_bit(); // Ensure non-writable by default
}

fn enable_nxe_bit() {
    use x86::msr::{IA32_EFER, rdmsr, wrmsr};

    let nxe_bit = 1 << 11;
    unsafe {
        let efer = rdmsr(IA32_EFER);
        wrmsr(IA32_EFER, efer | nxe_bit);
    }
}

fn enable_write_protect_bit() {
    use x86::controlregs::{cr0, cr0_write};

    let wp_bit = 1 << 16;
    unsafe { cr0_write(cr0() | wp_bit) };
}

// For stack-unwinding, not supported currently
#[lang = "eh_personality"]
extern "C" fn eh_personality() {}

// For panic!
#[no_mangle]
#[lang = "panic_fmt"]
extern "C" fn panic_fmt(fmt: core::fmt::Arguments, file: &str, line: u32) -> ! {
    kprintln!("\n\nPANIC in {} at line {}:", file, line);
    kprintln!("    {}", fmt);

    // Hang here.
    loop {
    }
}

#[allow(non_snake_case)]
#[no_mangle]
pub extern "C" fn _Unwind_Resume() -> ! {
    loop {}
}
<commit_msg>Export panic_fmt to fix compiler warning<commit_after>#![feature(lang_items)]
#![feature(const_fn)]
#![feature(unique)]
#![feature(alloc)]
#![feature(asm)]
#![feature(core_intrinsics)]
#![feature(naked_functions)]
#![feature(collections)]
#![feature(drop_types_in_const)]
#![feature(box_syntax)]

#![no_std]

#[macro_use]
extern crate x86;
extern crate rlibc;
extern crate spin;
extern crate multiboot2;
extern crate alloc_opsys;
extern crate alloc;

#[macro_use]
extern crate collections;

#[macro_use]
extern crate bitflags;

#[macro_use]
extern crate once;

#[macro_use]
mod vga_buffer;
mod memory;
mod interrupts;
mod drivers;
mod io;
mod schedule;
mod kernel;

// Main entry point, need no_mangle so we can call from assembly
// Extern to abide with C calling convention
#[no_mangle]
pub extern "C" fn kernel_main(multiboot_info_address: usize) {
    // Initialise the hardware
    init_cpu();
    memory::init(multiboot_info_address);

    // Setup the heap allocator
    alloc_opsys::init();

    // Setup the kernel
    kernel::init();

    // Initialize interrupts
    interrupts::init();

    vga_buffer::clear_screen();
    kprintln!("OpSys v{}", "0.0.1");

    loop {
    }
}

fn init_cpu() {
    // Configure the CPU flags as required
    enable_nxe_bit(); // Allow the NO_EXECUTABLE flag
    enable_write_protect_bit(); // Ensure non-writable by default
}

fn enable_nxe_bit() {
    use x86::msr::{IA32_EFER, rdmsr, wrmsr};

    let nxe_bit = 1 << 11;
    unsafe {
        let efer = rdmsr(IA32_EFER);
        wrmsr(IA32_EFER, efer | nxe_bit);
    }
}

fn enable_write_protect_bit() {
    use x86::controlregs::{cr0, cr0_write};

    let wp_bit = 1 << 16;
    unsafe { cr0_write(cr0() | wp_bit) };
}

// For stack-unwinding, not supported currently
#[lang = "eh_personality"]
extern "C" fn eh_personality() {}

// For panic!
#[no_mangle]
#[lang = "panic_fmt"]
pub extern "C" fn panic_fmt(fmt: core::fmt::Arguments, file: &str, line: u32) -> ! {
    kprintln!("\n\nPANIC in {} at line {}:", file, line);
    kprintln!("    {}", fmt);

    // Hang here.
    loop {
    }
}

#[allow(non_snake_case)]
#[no_mangle]
pub extern "C" fn _Unwind_Resume() -> ! {
    loop {}
}
<|endoftext|>