<commit_before>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

use super::objects::{DocType, EsId, Admin};
use chrono;
use regex;
use hyper;
use hyper::status::StatusCode;
use rs_es::error::EsError;
use rs_es;
use rs_es::EsResponse;
use serde;
use serde_json;
use rs_es::operations::search::ScanResult;

use super::objects::{AliasOperations, AliasOperation, AliasParameter};
use rs_es::units::{Duration};
use std::collections::BTreeMap;

// Rubber is an wrapper around elasticsearch API
pub struct Rubber {
    es_client: rs_es::Client,
    // some operation are not implemented in rs_es, we need to use a raw http client
    http_client: hyper::client::Client,
}

/// return the index associated to the given type and dataset
/// this will be an alias over another real index
fn get_main_index(doc_type: &str, dataset: &str) -> String {
    format!("munin_{}_{}", doc_type, dataset)
}

impl Rubber {
    // build a rubber with a connection string (http://host:port/)
    pub fn new(cnx: &str) -> Rubber {
        let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)").unwrap();
        let cap = re.captures(cnx).unwrap();
        let host = cap.name("host").unwrap();
        let port = cap.name("port").unwrap().parse::<u32>().unwrap();
        info!("elastic search host {:?} port {:?}", host, port);

        Rubber {
            es_client: rs_es::Client::new(&host, port),
            http_client: hyper::client::Client::new(),
        }
    }

    pub fn get(&self, path: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a get on {}", path);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .get(&url)
                              .send());
        rs_es::do_req(result)
    }
    fn put(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a put on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .put(&url)
                              .body(body)
                              .send());
        rs_es::do_req(result)
    }
    fn post(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a post on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .post(&url)
                              .body(body)
                              .send());
        rs_es::do_req(result)
    }

    pub fn make_index(&self, doc_type: &str, dataset: &str) -> Result<String, String> {
        let current_time = chrono::UTC::now().format("%Y%m%d_%H%M%S_%f");
        let index_name = format!("munin_{}_{}_{}", doc_type, dataset, current_time);
        info!("creating index {}", index_name);
        self.create_index(&index_name.to_string()).map(|_| index_name)
    }

    fn create_index(&self, name: &String) -> Result<(), String> {
        debug!("creating index");
        // Note: in rs_es it can be done with MappingOperation but for the moment I think
        // storing the mapping in json is more convenient
        let analysis = include_str!("../../../json/settings.json");
        self.put(name, &analysis)
            .map_err(|e| {
                info!("Error while creating new index {}", name);
                e.to_string()
            })
            .and_then(|res| {
                if res.status == StatusCode::Ok {
                    Ok(())
                } else {
                    Err(format!("cannot create index: {:?}", res))
                }
            })
    }

    fn get_last_index(&self, doc_type: &str, dataset: &str) -> Result<Vec<String>, String> {
        debug!("get last index: {base_index}/_aliases",
               base_index = get_main_index(doc_type, dataset));
        self.get(&format!("{base_index}/_aliases",
                          base_index = get_main_index(doc_type, dataset)))
            .map_err(|e| e.to_string())
            .and_then(|res| {
                match res.status {
                    StatusCode::Ok => {
                        let value: serde_json::Value = try!(res.read_response()
                                                               .map_err(|e| e.to_string()));
                        Ok(value.as_object()
                                .and_then(|aliases| Some(aliases.keys().cloned().collect()))
                                .unwrap_or_else(|| {
                                    info!("no previous index to delete for type {} and dataset {}",
                                          doc_type,
                                          dataset);
                                    vec![]
                                }))
                    }
                    StatusCode::NotFound => {
                        info!("impossible to find alias {}, no last index to remove",
                              get_main_index(doc_type, dataset));
                        Ok(vec![])
                    }
                    _ => Err(format!("invalid elasticsearch response: {:?}", res)),
                }

            })
    }

    /// publish the index as the new index for this doc_type and this dataset
    /// move the index alias of the doc_type and the dataset to point to this indexes
    /// and remove the old index
    pub fn publish_index(&mut self,
                         doc_type: &str,
                         dataset: &str,
                         index: String)
                         -> Result<(), String> {
        debug!("publishing index");
        let last_indexes = try!(self.get_last_index(doc_type, dataset));
        let main_index = get_main_index(doc_type, dataset);
        try!(self.alias(&main_index, &vec![index.clone()], &last_indexes));
        try!(self.alias("munin", &vec![main_index.to_string()], &vec![]));
        for i in last_indexes {
            try!(self.delete_index(&i));
        }
        Ok(())
    }

    pub fn is_existing_index(&self, name: &String) -> Result<bool, String> {
        self.get(&name)
            .map_err(|e| e.to_string())
            .map(|res| res.status == StatusCode::Ok)
    }

    /// add a list of new indexes to the alias
    /// remove a list of indexes from the alias
    fn alias(&self, alias: &str, add: &Vec<String>, remove: &Vec<String>) -> Result<(), String> {
        info!("for {}, adding alias {:?}, removing {:?}",
              alias,
              add,
              remove);
        let add_operations = add.iter().map(|x| {
            AliasOperation {
                remove: None,
                add: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let remove_operations = remove.iter().map(|x| {
            AliasOperation {
                add: None,
                remove: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let operations = AliasOperations {
            actions: add_operations.chain(remove_operations).collect(),
        };
        let json = serde_json::to_string(&operations).unwrap();
        self.post("_aliases", &json)
            .map_err(|e| e.to_string())
            .and_then(|res| {
                if res.status == StatusCode::Ok {
                    Ok(())
                } else {
                    error!("failed to change aliases for {}, es response: {:?}",
                           alias,
                           res);
                    Err(format!("failed to post aliases for {}: {:?}", alias, res).to_string())
                }
            })
    }

    pub fn delete_index(&mut self, index: &String) -> Result<(), String> {
        debug!("deleting index {}", &index);
        let res = self.es_client
                      .delete_index(&index)
                      .map(|res| res.acknowledged)
                      .unwrap_or(false);
        if !res {
            Err(format!("Error deleting index {}", &index).into())
        } else {
            Ok(())
        }
    }

    pub fn bulk_index<T, I>(&mut self,
                            index: &String,
                            iter: I)
                            -> Result<usize, rs_es::error::EsError>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        use rs_es::operations::bulk::Action;
        let mut nb = 0;
        let chunk_size = 1000;
        //fold is used for creating the action and optionally set the id of the object
        let mut actions = iter.map(|v| v.es_id()
                                        .into_iter()
                                        .fold(Action::index(v), |action, id| action.with_id(id)));
        loop {
            let chunk = actions.by_ref().take(chunk_size).collect::<Vec<_>>();
            nb += chunk.len();
            try!(self.es_client
                     .bulk(&chunk)
                     .with_index(&index)
                     .with_doc_type(T::doc_type())
                     .send());

            if chunk.len() < chunk_size {
                break;
            }
        }

        Ok(nb)
    }

    /// add all the element of 'iter' into elasticsearch
    ///
    /// To have zero downtime:
    /// first all the elements are added in a temporary index and when all has been indexed
    /// the index is published and the old index is removed
    pub fn index<T, I>(&mut self, doc_type: &str, dataset: &str, iter: I) -> Result<usize, String>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        // TODO better error handling
        let index = try!(self.make_index(doc_type, dataset));
        let nb_elements = try!(self.bulk_index(&index, iter).map_err(|e| e.to_string()));
        try!(self.publish_index(doc_type, dataset, index));
        Ok(nb_elements)
    }

	fn make_admin(&self, value: Option<Box<serde_json::Value>>) -> Option<Admin> {
		value.and_then(|v| serde_json::from_value::<Admin>(*v).ok())
	}

    pub fn get_admins(&mut self, dataset: &str) -> Result<BTreeMap<String, Admin>, rs_es::error::EsError> {

        let mut result: BTreeMap<String, Admin> = BTreeMap::new();
        let index = get_main_index("admin", dataset);
        let mut scan: ScanResult<serde_json::Value> =
        try!(self.es_client
                  .search_query()
                  .with_indexes(&[&index])
                  .with_size(1000)
                  .with_types(&[&"admin"])
                  .scan(&Duration::minutes(1)));
        loop {
        	let page = try!(scan.scroll(&mut self.es_client, &Duration::minutes(1)));
        	if page.hits.hits.len() == 0 {
        		break;
        	}
        	let tmp: Vec<Admin> = page.hits
        	                          .hits
        	                          .into_iter()
        	                          .filter_map(|hit| self.make_admin(hit.source))
        	                          .collect();

        	for ad in tmp {
        		result.insert(ad.id.to_string(), ad);
        	}
        }
        try!(scan.close(&mut self.es_client));
        Ok(result)
	}
}

#[test]
pub fn test_valid_url() {
    Rubber::new("http://localhost:9200");
    Rubber::new("localhost:9200");
}

#[test]
#[should_panic]
pub fn test_invalid_url() {
    Rubber::new("http://bob");
}

#[test]
#[should_panic]
pub fn test_invalid_url_no_port() {
    Rubber::new("localhost");
}
<commit_msg>refactoring<commit_after>// Copyright © 2016, Canal TP and/or its affiliates. All rights reserved.
//
// This file is part of Navitia,
//     the software to build cool stuff with public transport.
//
// Hope you'll enjoy and contribute to this project,
//     powered by Canal TP (www.canaltp.fr).
// Help us simplify mobility and open public transport:
//     a non ending quest to the responsive locomotion way of traveling!
//
// LICENCE: This program is free software; you can redistribute it
// and/or modify it under the terms of the GNU Affero General Public
// License as published by the Free Software Foundation, either
// version 3 of the License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public
// License along with this program. If not, see
// <http://www.gnu.org/licenses/>.
//
// Stay tuned using
// twitter @navitia
// IRC #navitia on freenode
// https://groups.google.com/d/forum/navitia
// www.navitia.io

use super::objects::{DocType, EsId, Admin};
use chrono;
use regex;
use hyper;
use hyper::status::StatusCode;
use rs_es::error::EsError;
use rs_es;
use rs_es::EsResponse;
use serde;
use serde_json;
use rs_es::operations::search::ScanResult;

use super::objects::{AliasOperations, AliasOperation, AliasParameter};
use rs_es::units::{Duration};
use std::collections::BTreeMap;

// Rubber is an wrapper around elasticsearch API
pub struct Rubber {
    es_client: rs_es::Client,
    // some operation are not implemented in rs_es, we need to use a raw http client
    http_client: hyper::client::Client,
}

/// return the index associated to the given type and dataset
/// this will be an alias over another real index
fn get_main_index(doc_type: &str, dataset: &str) -> String {
    format!("munin_{}_{}", doc_type, dataset)
}

impl Rubber {
    // build a rubber with a connection string (http://host:port/)
    pub fn new(cnx: &str) -> Rubber {
        let re = regex::Regex::new(r"(?:https?://)?(?P<host>.+?):(?P<port>\d+)").unwrap();
        let cap = re.captures(cnx).unwrap();
        let host = cap.name("host").unwrap();
        let port = cap.name("port").unwrap().parse::<u32>().unwrap();
        info!("elastic search host {:?} port {:?}", host, port);

        Rubber {
            es_client: rs_es::Client::new(&host, port),
            http_client: hyper::client::Client::new(),
        }
    }

    pub fn get(&self, path: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a get on {}", path);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .get(&url)
                              .send());
        rs_es::do_req(result)
    }
    fn put(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a put on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .put(&url)
                              .body(body)
                              .send());
        rs_es::do_req(result)
    }
    fn post(&self, path: &str, body: &str) -> Result<hyper::client::response::Response, EsError> {
        // Note: a bit duplicate on rs_es because some ES operations are not implemented
        debug!("doing a post on {} with {}", path, body);
        let url = self.es_client.full_url(path);
        let result = try!(self.http_client
                              .post(&url)
                              .body(body)
                              .send());
        rs_es::do_req(result)
    }

    pub fn make_index(&self, doc_type: &str, dataset: &str) -> Result<String, String> {
        let current_time = chrono::UTC::now().format("%Y%m%d_%H%M%S_%f");
        let index_name = format!("munin_{}_{}_{}", doc_type, dataset, current_time);
        info!("creating index {}", index_name);
        self.create_index(&index_name.to_string()).map(|_| index_name)
    }

    fn create_index(&self, name: &String) -> Result<(), String> {
        debug!("creating index");
        // Note: in rs_es it can be done with MappingOperation but for the moment I think
        // storing the mapping in json is more convenient
        let analysis = include_str!("../../../json/settings.json");
        self.put(name, &analysis)
            .map_err(|e| {
                info!("Error while creating new index {}", name);
                e.to_string()
            })
            .and_then(|res| {
                if res.status == StatusCode::Ok {
                    Ok(())
                } else {
                    Err(format!("cannot create index: {:?}", res))
                }
            })
    }

    fn get_last_index(&self, doc_type: &str, dataset: &str) -> Result<Vec<String>, String> {
        debug!("get last index: {base_index}/_aliases",
               base_index = get_main_index(doc_type, dataset));
        self.get(&format!("{base_index}/_aliases",
                          base_index = get_main_index(doc_type, dataset)))
            .map_err(|e| e.to_string())
            .and_then(|res| {
                match res.status {
                    StatusCode::Ok => {
                        let value: serde_json::Value = try!(res.read_response()
                                                               .map_err(|e| e.to_string()));
                        Ok(value.as_object()
                                .and_then(|aliases| Some(aliases.keys().cloned().collect()))
                                .unwrap_or_else(|| {
                                    info!("no previous index to delete for type {} and dataset {}",
                                          doc_type,
                                          dataset);
                                    vec![]
                                }))
                    }
                    StatusCode::NotFound => {
                        info!("impossible to find alias {}, no last index to remove",
                              get_main_index(doc_type, dataset));
                        Ok(vec![])
                    }
                    _ => Err(format!("invalid elasticsearch response: {:?}", res)),
                }

            })
    }

    /// publish the index as the new index for this doc_type and this dataset
    /// move the index alias of the doc_type and the dataset to point to this indexes
    /// and remove the old index
    pub fn publish_index(&mut self,
                         doc_type: &str,
                         dataset: &str,
                         index: String)
                         -> Result<(), String> {
        debug!("publishing index");
        let last_indexes = try!(self.get_last_index(doc_type, dataset));
        let main_index = get_main_index(doc_type, dataset);
        try!(self.alias(&main_index, &vec![index.clone()], &last_indexes));
        try!(self.alias("munin", &vec![main_index.to_string()], &vec![]));
        for i in last_indexes {
            try!(self.delete_index(&i));
        }
        Ok(())
    }

    pub fn is_existing_index(&self, name: &String) -> Result<bool, String> {
        self.get(&name)
            .map_err(|e| e.to_string())
            .map(|res| res.status == StatusCode::Ok)
    }

    /// add a list of new indexes to the alias
    /// remove a list of indexes from the alias
    fn alias(&self, alias: &str, add: &Vec<String>, remove: &Vec<String>) -> Result<(), String> {
        info!("for {}, adding alias {:?}, removing {:?}",
              alias,
              add,
              remove);
        let add_operations = add.iter().map(|x| {
            AliasOperation {
                remove: None,
                add: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let remove_operations = remove.iter().map(|x| {
            AliasOperation {
                add: None,
                remove: Some(AliasParameter {
                    index: x.clone(),
                    alias: alias.to_string(),
                }),
            }
        });
        let operations = AliasOperations {
            actions: add_operations.chain(remove_operations).collect(),
        };
        let json = serde_json::to_string(&operations).unwrap();
        self.post("_aliases", &json)
            .map_err(|e| e.to_string())
            .and_then(|res| {
                if res.status == StatusCode::Ok {
                    Ok(())
                } else {
                    error!("failed to change aliases for {}, es response: {:?}",
                           alias,
                           res);
                    Err(format!("failed to post aliases for {}: {:?}", alias, res).to_string())
                }
            })
    }

    pub fn delete_index(&mut self, index: &String) -> Result<(), String> {
        debug!("deleting index {}", &index);
        let res = self.es_client
                      .delete_index(&index)
                      .map(|res| res.acknowledged)
                      .unwrap_or(false);
        if !res {
            Err(format!("Error deleting index {}", &index).into())
        } else {
            Ok(())
        }
    }

    pub fn bulk_index<T, I>(&mut self,
                            index: &String,
                            iter: I)
                            -> Result<usize, rs_es::error::EsError>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        use rs_es::operations::bulk::Action;
        let mut nb = 0;
        let chunk_size = 1000;
        //fold is used for creating the action and optionally set the id of the object
        let mut actions = iter.map(|v| v.es_id()
                                        .into_iter()
                                        .fold(Action::index(v), |action, id| action.with_id(id)));
        loop {
            let chunk = actions.by_ref().take(chunk_size).collect::<Vec<_>>();
            nb += chunk.len();
            try!(self.es_client
                     .bulk(&chunk)
                     .with_index(&index)
                     .with_doc_type(T::doc_type())
                     .send());

            if chunk.len() < chunk_size {
                break;
            }
        }

        Ok(nb)
    }

    /// add all the element of 'iter' into elasticsearch
    ///
    /// To have zero downtime:
    /// first all the elements are added in a temporary index and when all has been indexed
    /// the index is published and the old index is removed
    pub fn index<T, I>(&mut self, doc_type: &str, dataset: &str, iter: I) -> Result<usize, String>
        where T: serde::Serialize + DocType + EsId,
              I: Iterator<Item = T>
    {
        // TODO better error handling
        let index = try!(self.make_index(doc_type, dataset));
        let nb_elements = try!(self.bulk_index(&index, iter).map_err(|e| e.to_string()));
        try!(self.publish_index(doc_type, dataset, index));
        Ok(nb_elements)
    }

    pub fn get_admins(&mut self, dataset: &str) -> Result<BTreeMap<String, Admin>, rs_es::error::EsError> {

        let mut result: BTreeMap<String, Admin> = BTreeMap::new();
        let index = get_main_index("admin", dataset);
        let mut scan: ScanResult<Admin> =
        try!(self.es_client
                  .search_query()
                  .with_indexes(&[&index])
                  .with_size(1000)
                  .with_types(&[&"admin"])
                  .scan(&Duration::minutes(1)));
        loop {
        	let page = try!(scan.scroll(&mut self.es_client, &Duration::minutes(1)));
        	if page.hits.hits.len() == 0 {
        		break;
        	}
        	for ad in page.hits.hits.into_iter().filter_map(|hit| hit.source) {
        		result.insert(ad.id.to_string(), *ad);
        	}
        }
        try!(scan.close(&mut self.es_client));
        Ok(result)
	}
}

#[test]
pub fn test_valid_url() {
    Rubber::new("http://localhost:9200");
    Rubber::new("localhost:9200");
}

#[test]
#[should_panic]
pub fn test_invalid_url() {
    Rubber::new("http://bob");
}

#[test]
#[should_panic]
pub fn test_invalid_url_no_port() {
    Rubber::new("localhost");
}
<|endoftext|>


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

extern crate clap;
extern crate env_logger;
extern crate habitat_core as hcore;
extern crate habitat_common as common;
extern crate habitat_pkg_export_docker as export_docker;
extern crate handlebars;
extern crate rand;
#[macro_use]
extern crate serde_json;
#[macro_use]
extern crate log;

extern crate failure;
#[macro_use]
extern crate failure_derive;

mod topology;
mod error;
mod manifest;
mod cli;

use std::env;

use hcore::channel;
use hcore::PROGRAM_NAME;
use hcore::url as hurl;
use common::ui::UI;

use export_docker::{Credentials, BuildSpec, Naming, Result};

use manifest::Manifest;

// Synced with the version of the Habitat operator.
pub const VERSION: &'static str = "0.1.0";

fn main() {
    env_logger::init().unwrap();
    let mut ui = UI::default_with_env();
    if let Err(e) = start(&mut ui) {
        let _ = ui.fatal(e);
        std::process::exit(1)
    }
}

fn start(ui: &mut UI) -> Result<()> {
    let m = cli().get_matches();
    debug!("clap cli args: {:?}", m);

    if !m.is_present("NO_DOCKER_IMAGE") {
        gen_docker_img(ui, &m)?;
    }
    let mut manifest = Manifest::new_from_cli_matches(ui, &m)?;
    manifest.generate()
}

fn gen_docker_img(ui: &mut UI, matches: &clap::ArgMatches) -> Result<()> {
    let default_channel = channel::default();
    let default_url = hurl::default_bldr_url();
    let spec = BuildSpec::new_from_cli_matches(&matches, &default_channel, &default_url);
    let naming = Naming::new_from_cli_matches(&matches);

    let docker_image = export_docker::export(ui, spec, &naming)?;
    docker_image.create_report(
        ui,
        env::current_dir()?.join("results"),
    )?;

    if matches.is_present("PUSH_IMAGE") {
        let credentials = Credentials::new(
            naming.registry_type,
            matches.value_of("REGISTRY_USERNAME").unwrap(),
            matches.value_of("REGISTRY_PASSWORD").unwrap(),
        )?;
        docker_image.push(ui, &credentials, naming.registry_url)?;
    }
    if matches.is_present("RM_IMAGE") {
        docker_image.rm(ui)?;
    }

    Ok(())
}

fn cli<'a, 'b>() -> clap::App<'a, 'b> {
    let name: &str = &*PROGRAM_NAME;
    let about = "Creates a Docker image and Kubernetes manifest for a Habitat package. Habitat \
                 operator must be deployed within the Kubernetes cluster before the generated \
                 manifest can be applied to this cluster.";

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

extern crate clap;
extern crate env_logger;
extern crate habitat_core as hcore;
extern crate habitat_common as common;
extern crate habitat_pkg_export_docker as export_docker;
extern crate handlebars;
extern crate rand;
#[macro_use]
extern crate serde_json;
#[macro_use]
extern crate log;

extern crate failure;
#[macro_use]
extern crate failure_derive;

mod topology;
mod error;
mod manifest;
mod cli;

use hcore::PROGRAM_NAME;
use common::ui::UI;

use export_docker::Result;

use manifest::Manifest;

// Synced with the version of the Habitat operator.
pub const VERSION: &'static str = "0.1.0";

fn main() {
    env_logger::init().unwrap();
    let mut ui = UI::default_with_env();
    if let Err(e) = start(&mut ui) {
        let _ = ui.fatal(e);
        std::process::exit(1)
    }
}

fn start(ui: &mut UI) -> Result<()> {
    let m = cli().get_matches();
    debug!("clap cli args: {:?}", m);

    if !m.is_present("NO_DOCKER_IMAGE") {
        export_docker::export_for_cli_matches(ui, &m)?;
    }
    let mut manifest = Manifest::new_from_cli_matches(ui, &m)?;
    manifest.generate()
}

fn cli<'a, 'b>() -> clap::App<'a, 'b> {
    let name: &str = &*PROGRAM_NAME;
    let about = "Creates a Docker image and Kubernetes manifest for a Habitat package. Habitat \
                 operator must be deployed within the Kubernetes cluster before the generated \
                 manifest can be applied to this cluster.";

    cli::Cli::new(name, about).add_all_args().app
}
<|endoftext|>


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

use derive_common::cg;
use proc_macro2::TokenStream;
use syn::{DeriveInput, Ident, Path};
use synstructure::{BindStyle, BindingInfo};

pub fn derive_to_value(
    mut input: DeriveInput,
    trait_path: Path,
    output_type_name: Ident,
    bind_style: BindStyle,
    // Returns whether to apply the field bound for a given item.
    mut field_bound: impl FnMut(&BindingInfo) -> bool,
    // Returns a token stream of the form: trait_path::from_foo(#binding)
    mut call_from: impl FnMut(&BindingInfo) -> TokenStream,
    mut call_to: impl FnMut(&BindingInfo) -> TokenStream,
    // Returns a tokenstream of the form:
    // fn from_function_syntax(foobar) -> Baz {
    //     #first_arg
    // }
    //
    // fn to_function_syntax(foobar) -> Baz {
    //     #second_arg
    // }
    mut trait_impl: impl FnMut(TokenStream, TokenStream) -> TokenStream,
    // if this is provided, the derive for non-generic types will be simplified
    // to this token stream, which should be the body of the impl block.
    non_generic_implementation: impl FnOnce() -> Option<TokenStream>,
) -> TokenStream {
    let mut where_clause = input.generics.where_clause.take();
    cg::propagate_clauses_to_output_type(
        &mut where_clause,
        &input.generics,
        &trait_path,
        &output_type_name,
    );
    let (to_body, from_body) = {
        let params = input.generics.type_params().collect::<Vec<_>>();
        for param in &params {
            cg::add_predicate(
                &mut where_clause,
                parse_quote!(#param: #trait_path),
            );
        }

        let to_body = cg::fmap_match(&input, bind_style, |binding| {
            if field_bound(&binding) {
                let ty = &binding.ast().ty;

                let output_type = cg::map_type_params(
                    ty,
                    &params,
                    &mut |ident| parse_quote!(<#ident as #trait_path>::#output_type_name),
                );

                cg::add_predicate(
                    &mut where_clause,
                    parse_quote!(
                        #ty: #trait_path<#output_type_name = #output_type>
                    ),
                );
            }
            call_to(&binding)
        });
        let from_body = cg::fmap_match(&input, bind_style, |binding| {
            call_from(&binding)
        });

        (to_body, from_body)
    };

    input.generics.where_clause = where_clause;
    let name = &input.ident;
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
    if input.generics.type_params().next().is_none() {
        if let Some(non_generic_implementation) = non_generic_implementation() {
            return quote! {
                impl #impl_generics #trait_path for #name #ty_generics
                #where_clause
                {
                    #non_generic_implementation
                }
            };
        }
    }

    let computed_value_type = cg::fmap_trait_output(
        &input,
        &trait_path,
        &output_type_name,
    );

    let impl_ = trait_impl(from_body, to_body);

    quote! {
        impl #impl_generics #trait_path for #name #ty_generics #where_clause {
            type #output_type_name = #computed_value_type;

            #impl_
        }
    }
}

pub fn derive(input: DeriveInput) -> TokenStream {
    let trait_impl = |from_body, to_body| {
        quote! {
            #[inline]
            fn from_computed_value(computed: &Self::ComputedValue) -> Self {
                match *computed {
                    #from_body
                }
            }

            #[allow(unused_variables)]
            #[inline]
            fn to_computed_value(&self, context: &crate::values::computed::Context) -> Self::ComputedValue {
                match *self {
                    #to_body
                }
            }
       }
    };

    let non_generic_implementation = || {
        Some(quote! {
            type ComputedValue = Self;

            #[inline]
            fn to_computed_value(&self, _: &crate::values::computed::Context) -> Self::ComputedValue {
                std::clone::Clone::clone(self)
            }

            #[inline]
            fn from_computed_value(computed: &Self::ComputedValue) -> Self {
                std::clone::Clone::clone(computed)
            }
        })
    };

    derive_to_value(
        input,
        parse_quote!(crate::values::computed::ToComputedValue),
        parse_quote!(ComputedValue),
        BindStyle::Ref,
        |binding| cg::parse_field_attrs::<ComputedValueAttrs>(&binding.ast()).field_bound,
        |binding| quote!(crate::values::computed::ToComputedValue::from_computed_value(#binding)),
        |binding| quote!(crate::values::computed::ToComputedValue::to_computed_value(#binding, context)),
        trait_impl,
        non_generic_implementation,
    )
}

#[darling(attributes(compute), default)]
#[derive(Default, FromField)]
struct ComputedValueAttrs {
    field_bound: bool,
}
<commit_msg>style: Handle the non-generic path earlier in the derive code.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

use derive_common::cg;
use proc_macro2::TokenStream;
use syn::{DeriveInput, Ident, Path};
use synstructure::{BindStyle, BindingInfo};

pub fn derive_to_value(
    mut input: DeriveInput,
    trait_path: Path,
    output_type_name: Ident,
    bind_style: BindStyle,
    // Returns whether to apply the field bound for a given item.
    mut field_bound: impl FnMut(&BindingInfo) -> bool,
    // Returns a token stream of the form: trait_path::from_foo(#binding)
    mut call_from: impl FnMut(&BindingInfo) -> TokenStream,
    mut call_to: impl FnMut(&BindingInfo) -> TokenStream,
    // Returns a tokenstream of the form:
    // fn from_function_syntax(foobar) -> Baz {
    //     #first_arg
    // }
    //
    // fn to_function_syntax(foobar) -> Baz {
    //     #second_arg
    // }
    mut trait_impl: impl FnMut(TokenStream, TokenStream) -> TokenStream,
    // if this is provided, the derive for non-generic types will be simplified
    // to this token stream, which should be the body of the impl block.
    non_generic_implementation: impl FnOnce() -> Option<TokenStream>,
) -> TokenStream {
    let name = &input.ident;

    if input.generics.type_params().next().is_none() {
        if let Some(non_generic_implementation) = non_generic_implementation() {
            let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
            return quote! {
                impl #impl_generics #trait_path for #name #ty_generics
                #where_clause
                {
                    #non_generic_implementation
                }
            };
        }
    }

    let mut where_clause = input.generics.where_clause.take();
    cg::propagate_clauses_to_output_type(
        &mut where_clause,
        &input.generics,
        &trait_path,
        &output_type_name,
    );
    let (to_body, from_body) = {
        let params = input.generics.type_params().collect::<Vec<_>>();
        for param in &params {
            cg::add_predicate(
                &mut where_clause,
                parse_quote!(#param: #trait_path),
            );
        }

        let to_body = cg::fmap_match(&input, bind_style, |binding| {
            if field_bound(&binding) {
                let ty = &binding.ast().ty;

                let output_type = cg::map_type_params(
                    ty,
                    &params,
                    &mut |ident| parse_quote!(<#ident as #trait_path>::#output_type_name),
                );

                cg::add_predicate(
                    &mut where_clause,
                    parse_quote!(
                        #ty: #trait_path<#output_type_name = #output_type>
                    ),
                );
            }
            call_to(&binding)
        });
        let from_body = cg::fmap_match(&input, bind_style, |binding| {
            call_from(&binding)
        });

        (to_body, from_body)
    };

    input.generics.where_clause = where_clause;
    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
    let computed_value_type = cg::fmap_trait_output(
        &input,
        &trait_path,
        &output_type_name,
    );

    let impl_ = trait_impl(from_body, to_body);

    quote! {
        impl #impl_generics #trait_path for #name #ty_generics #where_clause {
            type #output_type_name = #computed_value_type;

            #impl_
        }
    }
}

pub fn derive(input: DeriveInput) -> TokenStream {
    let trait_impl = |from_body, to_body| {
        quote! {
            #[inline]
            fn from_computed_value(computed: &Self::ComputedValue) -> Self {
                match *computed {
                    #from_body
                }
            }

            #[allow(unused_variables)]
            #[inline]
            fn to_computed_value(&self, context: &crate::values::computed::Context) -> Self::ComputedValue {
                match *self {
                    #to_body
                }
            }
       }
    };

    let non_generic_implementation = || {
        Some(quote! {
            type ComputedValue = Self;

            #[inline]
            fn to_computed_value(&self, _: &crate::values::computed::Context) -> Self::ComputedValue {
                std::clone::Clone::clone(self)
            }

            #[inline]
            fn from_computed_value(computed: &Self::ComputedValue) -> Self {
                std::clone::Clone::clone(computed)
            }
        })
    };

    derive_to_value(
        input,
        parse_quote!(crate::values::computed::ToComputedValue),
        parse_quote!(ComputedValue),
        BindStyle::Ref,
        |binding| cg::parse_field_attrs::<ComputedValueAttrs>(&binding.ast()).field_bound,
        |binding| quote!(crate::values::computed::ToComputedValue::from_computed_value(#binding)),
        |binding| quote!(crate::values::computed::ToComputedValue::to_computed_value(#binding, context)),
        trait_impl,
        non_generic_implementation,
    )
}

#[darling(attributes(compute), default)]
#[derive(Default, FromField)]
struct ComputedValueAttrs {
    field_bound: bool,
}
<|endoftext|>


<commit_before>use mio::net::{TcpListener, TcpStream};
use mio::*;
use std::net;
use std::sync::{Arc, Barrier};
use std::thread::{self, sleep};
use std::time::Duration;

mod util;

use util::{any_local_address, assert_send, assert_sync, init};

#[test]
fn is_send_and_sync() {
    assert_sync::<Poll>();
    assert_send::<Poll>();

    assert_sync::<Registry>();
    assert_send::<Registry>();
}

#[test]
fn run_once_with_nothing() {
    init();

    let mut events = Events::with_capacity(16);
    let mut poll = Poll::new().unwrap();
    poll.poll(&mut events, Some(Duration::from_millis(100)))
        .unwrap();
}

#[test]
fn add_then_drop() {
    init();

    let mut events = Events::with_capacity(16);
    let l = TcpListener::bind(any_local_address()).unwrap();
    let mut poll = Poll::new().unwrap();
    poll.registry()
        .register(&l, Token(1), Interests::READABLE | Interests::WRITABLE)
        .unwrap();
    drop(l);
    poll.poll(&mut events, Some(Duration::from_millis(100)))
        .unwrap();
}

#[test]
fn zero_duration_polls_events() {
    init();

    let mut poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(16);

    let listener = net::TcpListener::bind(any_local_address()).unwrap();
    let addr = listener.local_addr().unwrap();

    let streams: Vec<TcpStream> = (0..3)
        .map(|n| {
            let stream = TcpStream::connect(addr).unwrap();
            poll.registry()
                .register(&stream, Token(n), Interests::WRITABLE)
                .unwrap();
            stream
        })
        .collect();

    // Ensure the TcpStreams have some time to connection and for the events to
    // show up.
    sleep(Duration::from_millis(10));

    // Even when passing a zero duration timeout we still want do the system
    // call.
    poll.poll(&mut events, Some(Duration::from_nanos(0)))
        .unwrap();
    assert!(!events.is_empty());

    // Both need to live until here.
    drop(streams);
    drop(listener);
}

#[test]
fn test_poll_closes_fd() {
    init();

    for _ in 0..2000 {
        let mut poll = Poll::new().unwrap();
        let mut events = Events::with_capacity(4);

        poll.poll(&mut events, Some(Duration::from_millis(0)))
            .unwrap();

        drop(poll);
    }
}

#[test]
fn test_drop_cancels_interest_and_shuts_down() {
    init();

    use mio::net::TcpStream;
    use std::io;
    use std::io::Read;
    use std::net::TcpListener;
    use std::thread;

    let l = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = l.local_addr().unwrap();

    let t = thread::spawn(move || {
        let mut s = l.incoming().next().unwrap().unwrap();
        s.set_read_timeout(Some(Duration::from_secs(5)))
            .expect("set_read_timeout");
        let r = s.read(&mut [0; 16]);
        match r {
            Ok(_) => (),
            Err(e) => {
                if e.kind() != io::ErrorKind::UnexpectedEof {
                    panic!(e);
                }
            }
        }
    });

    let mut poll = Poll::new().unwrap();
    let mut s = TcpStream::connect(addr).unwrap();

    poll.registry()
        .register(&s, Token(1), Interests::READABLE | Interests::WRITABLE)
        .unwrap();
    let mut events = Events::with_capacity(16);
    'outer: loop {
        poll.poll(&mut events, None).unwrap();
        for event in &events {
            if event.token() == Token(1) {
                // connected
                break 'outer;
            }
        }
    }

    let mut b = [0; 1024];
    match s.read(&mut b) {
        Ok(_) => panic!("unexpected ok"),
        Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => (),
        Err(e) => panic!("unexpected error: {:?}", e),
    }

    drop(s);
    t.join().unwrap();
}

#[test]
fn test_registry_behind_arc() {
    // `Registry` should work behind an `Arc`, being `Sync` and `Send`.
    init();

    let mut poll = Poll::new().unwrap();
    let registry = Arc::new(poll.registry().try_clone().unwrap());
    let mut events = Events::with_capacity(128);

    let listener = TcpListener::bind(any_local_address()).unwrap();
    let addr = listener.local_addr().unwrap();
    let barrier = Arc::new(Barrier::new(3));

    let registry2 = Arc::clone(&registry);
    let registry3 = Arc::clone(&registry);
    let barrier2 = Arc::clone(&barrier);
    let barrier3 = Arc::clone(&barrier);

    let handle1 = thread::spawn(move || {
        registry2
            .register(&listener, Token(0), Interests::READABLE)
            .unwrap();
        barrier2.wait();
    });
    let handle2 = thread::spawn(move || {
        let stream = TcpStream::connect(addr).unwrap();
        registry3
            .register(&stream, Token(1), Interests::READABLE | Interests::WRITABLE)
            .unwrap();
        barrier3.wait();
    });

    poll.poll(&mut events, Some(Duration::from_millis(1000)))
        .unwrap();
    assert!(events.iter().count() >= 1);

    // Let the threads return.
    barrier.wait();

    handle1.join().unwrap();
    handle2.join().unwrap();
}

// On kqueue platforms registering twice (not *re*registering) works.
#[test]
#[cfg(any(target_os = "linux", target_os = "windows"))]
pub fn test_double_register() {
    init();
    let poll = Poll::new().unwrap();

    let l = TcpListener::bind("127.0.0.1:0".parse().unwrap()).unwrap();

    poll.registry()
        .register(&l, Token(0), Interests::READABLE)
        .unwrap();

    assert!(poll
        .registry()
        .register(&l, Token(1), Interests::READABLE)
        .is_err());
}
<commit_msg>Add tests for the registration process<commit_after>use mio::net::{TcpListener, TcpStream};
use mio::{event, Events, Interests, Poll, Registry, Token};

use std::net;
use std::sync::{Arc, Barrier, Mutex};
use std::thread::{self, sleep};
use std::time::Duration;
use std::{fmt, io};

mod util;

use util::{any_local_address, assert_send, assert_sync, init};

#[test]
fn is_send_and_sync() {
    assert_sync::<Poll>();
    assert_send::<Poll>();

    assert_sync::<Registry>();
    assert_send::<Registry>();
}

#[test]
fn run_once_with_nothing() {
    init();

    let mut events = Events::with_capacity(16);
    let mut poll = Poll::new().unwrap();
    poll.poll(&mut events, Some(Duration::from_millis(100)))
        .unwrap();
}

#[test]
fn add_then_drop() {
    init();

    let mut events = Events::with_capacity(16);
    let l = TcpListener::bind(any_local_address()).unwrap();
    let mut poll = Poll::new().unwrap();
    poll.registry()
        .register(&l, Token(1), Interests::READABLE | Interests::WRITABLE)
        .unwrap();
    drop(l);
    poll.poll(&mut events, Some(Duration::from_millis(100)))
        .unwrap();
}

#[test]
fn zero_duration_polls_events() {
    init();

    let mut poll = Poll::new().unwrap();
    let mut events = Events::with_capacity(16);

    let listener = net::TcpListener::bind(any_local_address()).unwrap();
    let addr = listener.local_addr().unwrap();

    let streams: Vec<TcpStream> = (0..3)
        .map(|n| {
            let stream = TcpStream::connect(addr).unwrap();
            poll.registry()
                .register(&stream, Token(n), Interests::WRITABLE)
                .unwrap();
            stream
        })
        .collect();

    // Ensure the TcpStreams have some time to connection and for the events to
    // show up.
    sleep(Duration::from_millis(10));

    // Even when passing a zero duration timeout we still want do the system
    // call.
    poll.poll(&mut events, Some(Duration::from_nanos(0)))
        .unwrap();
    assert!(!events.is_empty());

    // Both need to live until here.
    drop(streams);
    drop(listener);
}

#[test]
fn test_poll_closes_fd() {
    init();

    for _ in 0..2000 {
        let mut poll = Poll::new().unwrap();
        let mut events = Events::with_capacity(4);

        poll.poll(&mut events, Some(Duration::from_millis(0)))
            .unwrap();

        drop(poll);
    }
}

#[test]
fn test_drop_cancels_interest_and_shuts_down() {
    init();

    use mio::net::TcpStream;
    use std::io;
    use std::io::Read;
    use std::net::TcpListener;
    use std::thread;

    let l = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = l.local_addr().unwrap();

    let t = thread::spawn(move || {
        let mut s = l.incoming().next().unwrap().unwrap();
        s.set_read_timeout(Some(Duration::from_secs(5)))
            .expect("set_read_timeout");
        let r = s.read(&mut [0; 16]);
        match r {
            Ok(_) => (),
            Err(e) => {
                if e.kind() != io::ErrorKind::UnexpectedEof {
                    panic!(e);
                }
            }
        }
    });

    let mut poll = Poll::new().unwrap();
    let mut s = TcpStream::connect(addr).unwrap();

    poll.registry()
        .register(&s, Token(1), Interests::READABLE | Interests::WRITABLE)
        .unwrap();
    let mut events = Events::with_capacity(16);
    'outer: loop {
        poll.poll(&mut events, None).unwrap();
        for event in &events {
            if event.token() == Token(1) {
                // connected
                break 'outer;
            }
        }
    }

    let mut b = [0; 1024];
    match s.read(&mut b) {
        Ok(_) => panic!("unexpected ok"),
        Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => (),
        Err(e) => panic!("unexpected error: {:?}", e),
    }

    drop(s);
    t.join().unwrap();
}

#[test]
fn test_registry_behind_arc() {
    // `Registry` should work behind an `Arc`, being `Sync` and `Send`.
    init();

    let mut poll = Poll::new().unwrap();
    let registry = Arc::new(poll.registry().try_clone().unwrap());
    let mut events = Events::with_capacity(128);

    let listener = TcpListener::bind(any_local_address()).unwrap();
    let addr = listener.local_addr().unwrap();
    let barrier = Arc::new(Barrier::new(3));

    let registry2 = Arc::clone(&registry);
    let registry3 = Arc::clone(&registry);
    let barrier2 = Arc::clone(&barrier);
    let barrier3 = Arc::clone(&barrier);

    let handle1 = thread::spawn(move || {
        registry2
            .register(&listener, Token(0), Interests::READABLE)
            .unwrap();
        barrier2.wait();
    });
    let handle2 = thread::spawn(move || {
        let stream = TcpStream::connect(addr).unwrap();
        registry3
            .register(&stream, Token(1), Interests::READABLE | Interests::WRITABLE)
            .unwrap();
        barrier3.wait();
    });

    poll.poll(&mut events, Some(Duration::from_millis(1000)))
        .unwrap();
    assert!(events.iter().count() >= 1);

    // Let the threads return.
    barrier.wait();

    handle1.join().unwrap();
    handle2.join().unwrap();
}

// On kqueue platforms registering twice (not *re*registering) works.
#[test]
#[cfg(any(target_os = "linux", target_os = "windows"))]
pub fn test_double_register() {
    init();
    let poll = Poll::new().unwrap();

    let l = TcpListener::bind("127.0.0.1:0".parse().unwrap()).unwrap();

    poll.registry()
        .register(&l, Token(0), Interests::READABLE)
        .unwrap();

    assert!(poll
        .registry()
        .register(&l, Token(1), Interests::READABLE)
        .is_err());
}

struct TestEventSource(Mutex<TestEventSourceData>);

struct TestEventSourceData {
    registrations: Vec<(Token, Interests)>,
    reregistrations: Vec<(Token, Interests)>,
    deregister_count: usize,
}

impl TestEventSource {
    fn new() -> TestEventSource {
        TestEventSource(Mutex::new(TestEventSourceData {
            registrations: Vec::new(),
            reregistrations: Vec::new(),
            deregister_count: 0,
        }))
    }
}

impl event::Source for TestEventSource {
    fn register(&self, _registry: &Registry, token: Token, interests: Interests) -> io::Result<()> {
        let mut inner = self.0.lock().unwrap();
        inner.registrations.push((token, interests));
        Ok(())
    }

    fn reregister(
        &self,
        _registry: &Registry,
        token: Token,
        interests: Interests,
    ) -> io::Result<()> {
        let mut inner = self.0.lock().unwrap();
        inner.reregistrations.push((token, interests));
        Ok(())
    }

    fn deregister(&self, _registry: &Registry) -> io::Result<()> {
        let mut inner = self.0.lock().unwrap();
        inner.deregister_count += 1;
        Ok(())
    }
}

#[test]
fn poll_registration() {
    init();
    let poll = Poll::new().unwrap();
    let registry = poll.registry();

    let source = TestEventSource::new();
    let token = Token(0);
    let interests = Interests::READABLE;
    registry.register(&source, token, interests).unwrap();
    {
        let source = source.0.lock().unwrap();
        assert_eq!(source.registrations.len(), 1);
        assert_eq!(source.registrations.get(0), Some(&(token, interests)));
        assert!(source.reregistrations.is_empty());
        assert_eq!(source.deregister_count, 0);
    }

    let re_token = Token(0);
    let re_interests = Interests::READABLE;
    registry
        .reregister(&source, re_token, re_interests)
        .unwrap();
    {
        let source = source.0.lock().unwrap();
        assert_eq!(source.registrations.len(), 1);
        assert_eq!(source.reregistrations.len(), 1);
        assert_eq!(
            source.reregistrations.get(0),
            Some(&(re_token, re_interests))
        );
        assert_eq!(source.deregister_count, 0);
    }

    registry.deregister(&source).unwrap();
    {
        let source = source.0.lock().unwrap();
        assert_eq!(source.registrations.len(), 1);
        assert_eq!(source.reregistrations.len(), 1);
        assert_eq!(source.deregister_count, 1);
    }
}

struct ErroneousTestEventSource;

impl event::Source for ErroneousTestEventSource {
    fn register(
        &self,
        _registry: &Registry,
        _token: Token,
        _interests: Interests,
    ) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "register"))
    }

    fn reregister(
        &self,
        _registry: &Registry,
        _token: Token,
        _interests: Interests,
    ) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "reregister"))
    }

    fn deregister(&self, _registry: &Registry) -> io::Result<()> {
        Err(io::Error::new(io::ErrorKind::Other, "deregister"))
    }
}

#[test]
fn poll_erroneous_registration() {
    init();
    let poll = Poll::new().unwrap();
    let registry = poll.registry();

    let source = ErroneousTestEventSource;
    let token = Token(0);
    let interests = Interests::READABLE;
    assert_error(registry.register(&source, token, interests), "register");
    assert_error(registry.reregister(&source, token, interests), "reregister");
    assert_error(registry.deregister(&source), "deregister");
}

/// Assert that `result` is an error and the formatted error (via
/// `fmt::Display`) equals `expected_msg`.
pub fn assert_error<T, E: fmt::Display>(result: Result<T, E>, expected_msg: &str) {
    match result {
        Ok(_) => panic!("unexpected OK result"),
        Err(err) => assert!(
            err.to_string().contains(expected_msg),
            "wanted: {}, got: {}",
            err,
            expected_msg
        ),
    }
}
<|endoftext|>


<commit_before>use std::result;
use std::io;
use byteorder;

pub type Result<T> = result::Result<T, Error>;

#[derive(Debug)]
pub enum Error {
    IncorrectPacketFormat,
    InvalidTopicPath,
    UnsupportedProtocolName,
    UnsupportedProtocolVersion,
    UnsupportedQualityOfService,
    UnsupportedPacketType,
    UnsupportedConnectReturnCode,
    PayloadSizeIncorrect,
    PayloadTooLong,
    PayloadRequired,
    TopicNameMustNotContainWildcard,
    MalformedRemainingLength,
    UnexpectedEOF,
    Io(io::Error)
}

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

impl From<byteorder::Error> for Error {
    fn from(err: byteorder::Error) -> Error {
        match err {
            byteorder::Error::UnexpectedEOF => Error::UnexpectedEOF,
            byteorder::Error::Io(err) => Error::Io(err)
        }
    }
}
<commit_msg>mqtt3: Rename UnexpectedEof<commit_after>use std::result;
use std::io;
use byteorder;

pub type Result<T> = result::Result<T, Error>;

#[derive(Debug)]
pub enum Error {
    IncorrectPacketFormat,
    InvalidTopicPath,
    UnsupportedProtocolName,
    UnsupportedProtocolVersion,
    UnsupportedQualityOfService,
    UnsupportedPacketType,
    UnsupportedConnectReturnCode,
    PayloadSizeIncorrect,
    PayloadTooLong,
    PayloadRequired,
    TopicNameMustNotContainWildcard,
    MalformedRemainingLength,
    UnexpectedEof,
    Io(io::Error)
}

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

impl From<byteorder::Error> for Error {
    fn from(err: byteorder::Error) -> Error {
        match err {
            byteorder::Error::UnexpectedEOF => Error::UnexpectedEof,
            byteorder::Error::Io(err) => Error::Io(err)
        }
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "i128", feature(i128_type, i128))]
#![cfg_attr(
    feature = "cargo-clippy",
    allow(cast_lossless, string_lit_as_bytes)
)]
#![allow(non_snake_case)]

extern crate itoa;

macro_rules! test {
    (
        $(
            $(#[$attr:meta])*
            $name:ident($value:expr, $expected:expr)
        ),*
    ) => {
        $(
            $(#[$attr])*
            #[test]
            fn $name() {
                #[cfg(feature = "std")]
                {
                    let mut buf = [b'\0'; 40];
                    let len = itoa::write(&mut buf[..], $value).unwrap();
                    assert_eq!(&buf[0..len], $expected.as_bytes());
                }

                let mut s = String::new();
                itoa::fmt(&mut s, $value).unwrap();
                assert_eq!(s, $expected);
            }
        )*
    }
}

test!{
    test_u64_0(0u64, "0"),
    test_u64_half(<u32>::max_value() as u64, "4294967295"),
    test_u64_max(<u64>::max_value(), "18446744073709551615"),
    test_i64_min(<i64>::min_value(), "-9223372036854775808"),

    test_i16_0(0i16, "0"),
    test_i16_min(<i16>::min_value(), "-32768"),

    #[cfg(feature = "i128")]
    test_u128_0(0u128, "0"),
    #[cfg(feature = "i128")]
    test_u128_max(<u128>::max_value(), "340282366920938463463374607431768211455"),
    #[cfg(feature = "i128")]
    test_i128_min(<i128>::min_value(), "-170141183460469231731687303715884105728")
}
<commit_msg>No longer need i128 feature gate which has been stabilized<commit_after>#![cfg_attr(
    feature = "cargo-clippy",
    allow(cast_lossless, string_lit_as_bytes)
)]
#![allow(non_snake_case)]

extern crate itoa;

macro_rules! test {
    (
        $(
            $(#[$attr:meta])*
            $name:ident($value:expr, $expected:expr)
        ),*
    ) => {
        $(
            $(#[$attr])*
            #[test]
            fn $name() {
                #[cfg(feature = "std")]
                {
                    let mut buf = [b'\0'; 40];
                    let len = itoa::write(&mut buf[..], $value).unwrap();
                    assert_eq!(&buf[0..len], $expected.as_bytes());
                }

                let mut s = String::new();
                itoa::fmt(&mut s, $value).unwrap();
                assert_eq!(s, $expected);
            }
        )*
    }
}

test!{
    test_u64_0(0u64, "0"),
    test_u64_half(<u32>::max_value() as u64, "4294967295"),
    test_u64_max(<u64>::max_value(), "18446744073709551615"),
    test_i64_min(<i64>::min_value(), "-9223372036854775808"),

    test_i16_0(0i16, "0"),
    test_i16_min(<i16>::min_value(), "-32768"),

    #[cfg(feature = "i128")]
    test_u128_0(0u128, "0"),
    #[cfg(feature = "i128")]
    test_u128_max(<u128>::max_value(), "340282366920938463463374607431768211455"),
    #[cfg(feature = "i128")]
    test_i128_min(<i128>::min_value(), "-170141183460469231731687303715884105728")
}
<|endoftext|>


<commit_before>extern crate petgraph;

use petgraph::{
    GraphMap,
};

use petgraph::algo::{
    dijkstra,
};

#[test]
fn simple() {
    //let root = TypedArena::<Node<_>>::new();
    let mut gr = GraphMap::new();
    //let node = |&: name: &'static str| Ptr(root.alloc(Node(name.to_string())));
    let a = gr.add_node(("A"));
    let b = gr.add_node(("B"));
    let c = gr.add_node(("C"));
    let d = gr.add_node(("D"));
    let e = gr.add_node(("E"));
    let f = gr.add_node(("F"));
    gr.add_edge(a, b, 7);
    gr.add_edge(a, c, 9);
    gr.add_edge(a, d, 14);
    gr.add_edge(b, c, 10);
    gr.add_edge(c, d, 2);
    gr.add_edge(d, e, 9);
    gr.add_edge(b, f, 15);
    gr.add_edge(c, f, 11);

    assert!(gr.add_edge(e, f, 5));

    // duplicate edges
    assert!(!gr.add_edge(f, b, 15));
    assert!(!gr.add_edge(f, e, 6));
    println!("{:?}", gr);

    assert_eq!(gr.node_count(), 6);
    assert_eq!(gr.edge_count(), 9);

    // check updated edge weight
    assert_eq!(gr.edge_weight(e, f), Some(&6));
    let scores = dijkstra(&gr, a, None, |gr, n| gr.edges(n).map(|(n, &e)| (n, e)));
    let mut scores: Vec<_> = scores.into_iter().collect();
    scores.sort();
    assert_eq!(scores,
       vec![("A", 0), ("B", 7), ("C", 9), ("D", 11), ("E", 20), ("F", 20)]);
}

#[test]
fn remov()
{
    let mut g = GraphMap::new();
    g.add_node(1);
    g.add_node(2);
    g.add_edge(1, 2, -1);

    assert_eq!(g.edge_weight(1, 2), Some(&-1));
    assert_eq!(g.edge_weight(2, 1), Some(&-1));
    assert_eq!(g.neighbors(1).count(), 1);

    let noexist = g.remove_edge(2, 3);
    assert_eq!(noexist, None);

    let exist = g.remove_edge(2, 1);
    assert_eq!(exist, Some(-1));
    assert_eq!(g.edge_count(), 0);
    assert_eq!(g.edge_weight(1, 2), None);
    assert_eq!(g.edge_weight(2, 1), None);
    assert_eq!(g.neighbors(1).count(), 0);
}

#[test]
fn dfs() {
    let mut gr = GraphMap::new();
    let h = gr.add_node("H");
    let i = gr.add_node("I");
    let j = gr.add_node("J");
    let k = gr.add_node("K");
    // Z is disconnected.
    let _ = gr.add_node("Z");
    gr.add_edge(h, i, 1.);
    gr.add_edge(h, j, 3.);
    gr.add_edge(i, j, 1.);
    gr.add_edge(i, k, 2.);

    /*
    assert_eq!(DfsIter::new(&gr, h).count(), 4);
    assert_eq!(DfsIter::new(&gr, i).count(), 4);
    assert_eq!(DfsIter::new(&gr, z).count(), 1);
    */
}

<commit_msg>Use Dfs in graphmap test to make sure Dfs + GraphMap is tested<commit_after>extern crate petgraph;

use petgraph::{
    GraphMap,
    Dfs,
};

use petgraph::algo::{
    dijkstra,
};

#[test]
fn simple() {
    //let root = TypedArena::<Node<_>>::new();
    let mut gr = GraphMap::new();
    //let node = |&: name: &'static str| Ptr(root.alloc(Node(name.to_string())));
    let a = gr.add_node(("A"));
    let b = gr.add_node(("B"));
    let c = gr.add_node(("C"));
    let d = gr.add_node(("D"));
    let e = gr.add_node(("E"));
    let f = gr.add_node(("F"));
    gr.add_edge(a, b, 7);
    gr.add_edge(a, c, 9);
    gr.add_edge(a, d, 14);
    gr.add_edge(b, c, 10);
    gr.add_edge(c, d, 2);
    gr.add_edge(d, e, 9);
    gr.add_edge(b, f, 15);
    gr.add_edge(c, f, 11);

    assert!(gr.add_edge(e, f, 5));

    // duplicate edges
    assert!(!gr.add_edge(f, b, 15));
    assert!(!gr.add_edge(f, e, 6));
    println!("{:?}", gr);

    assert_eq!(gr.node_count(), 6);
    assert_eq!(gr.edge_count(), 9);

    // check updated edge weight
    assert_eq!(gr.edge_weight(e, f), Some(&6));
    let scores = dijkstra(&gr, a, None, |gr, n| gr.edges(n).map(|(n, &e)| (n, e)));
    let mut scores: Vec<_> = scores.into_iter().collect();
    scores.sort();
    assert_eq!(scores,
       vec![("A", 0), ("B", 7), ("C", 9), ("D", 11), ("E", 20), ("F", 20)]);
}

#[test]
fn remov()
{
    let mut g = GraphMap::new();
    g.add_node(1);
    g.add_node(2);
    g.add_edge(1, 2, -1);

    assert_eq!(g.edge_weight(1, 2), Some(&-1));
    assert_eq!(g.edge_weight(2, 1), Some(&-1));
    assert_eq!(g.neighbors(1).count(), 1);

    let noexist = g.remove_edge(2, 3);
    assert_eq!(noexist, None);

    let exist = g.remove_edge(2, 1);
    assert_eq!(exist, Some(-1));
    assert_eq!(g.edge_count(), 0);
    assert_eq!(g.edge_weight(1, 2), None);
    assert_eq!(g.edge_weight(2, 1), None);
    assert_eq!(g.neighbors(1).count(), 0);
}

#[test]
fn dfs() {
    let mut gr = GraphMap::new();
    let h = gr.add_node("H");
    let i = gr.add_node("I");
    let j = gr.add_node("J");
    let k = gr.add_node("K");
    // Z is disconnected.
    let z = gr.add_node("Z");
    gr.add_edge(h, i, 1.);
    gr.add_edge(h, j, 3.);
    gr.add_edge(i, j, 1.);
    gr.add_edge(i, k, 2.);

    println!("{:?}", gr);

    {
        let mut cnt = 0;
        let mut dfs = Dfs::new(&gr, h);
        while let Some(_) = dfs.next(&gr) { cnt += 1; }
        assert_eq!(cnt, 4);
    }
    {
        let mut cnt = 0;
        let mut dfs = Dfs::new(&gr, z);
        while let Some(_) = dfs.next(&gr) { cnt += 1; }
        assert_eq!(cnt, 1);
    }

    /*
    assert_eq!(DfsIter::new(&gr, h).count(), 4);
    assert_eq!(DfsIter::new(&gr, i).count(), 4);
    assert_eq!(DfsIter::new(&gr, z).count(), 1);
    */
}

<|endoftext|>


<commit_before>//! Common types for implementing federation authorization.

use headers::{authorization::Credentials, HeaderValue};

use ruma_common::{OwnedServerName, OwnedServerSigningKeyId};
use tracing::debug;
use yap::{IntoTokens, TokenLocation, Tokens};

/// Typed representation of an `Authorization` header of scheme `X-Matrix`, as defined in the
/// [Matrix Server-Server API][spec]. Includes an implementation of
/// [`headers::authorization::Credentials`] for automatically handling the encoding and decoding
/// when using a web framework that supports typed headers.
///
/// [spec]: https://spec.matrix.org/v1.3/server-server-api/#request-authentication
#[non_exhaustive]
pub struct XMatrix {
    /// The server name of the sending server.
    pub origin: OwnedServerName,
    /// The server name of the receiving sender. For compatibility with older servers, recipients
    /// should accept requests without this parameter, but MUST always send it. If this property is
    /// included, but the value does not match the receiving server's name, the receiving server
    /// must deny the request with an HTTP status code 401 Unauthorized.
    pub destination: Option<OwnedServerName>,
    /// The ID - including the algorithm name - of the sending server's key that was used to sign
    /// the request.
    pub key: OwnedServerSigningKeyId,
    /// The signature of the JSON.
    pub sig: String,
}

impl XMatrix {
    /// Construct a new X-Matrix Authorization header.
    pub fn new(
        origin: OwnedServerName,
        destination: Option<OwnedServerName>,
        key: OwnedServerSigningKeyId,
        sig: String,
    ) -> Self {
        Self { origin, destination, key, sig }
    }
}

fn parse_token<'a>(tokens: &mut impl Tokens<Item = &'a u8>) -> Option<Vec<u8>> {
    tokens.optional(|t| {
        let token: Vec<u8> = t.tokens_while(|c| is_tchar(**c)).copied().collect();
        if !token.is_empty() {
            Some(token)
        } else {
            debug!("Returning early because of empty token at {}", t.location().offset());
            None
        }
    })
}

fn parse_token_with_colons<'a>(tokens: &mut impl Tokens<Item = &'a u8>) -> Option<Vec<u8>> {
    tokens.optional(|t| {
        let token: Vec<u8> = t.tokens_while(|c| is_tchar(**c) || **c == b':').copied().collect();
        if !token.is_empty() {
            Some(token)
        } else {
            debug!("Returning early because of empty token at {}", t.location().offset());
            None
        }
    })
}

fn parse_quoted<'a>(tokens: &mut impl Tokens<Item = &'a u8>) -> Option<Vec<u8>> {
    tokens.optional(|t| {
        if !(t.token(&b'"')) {
            return None;
        }
        let mut buffer = Vec::new();
        loop {
            match t.next()? {
                // quoted pair
                b'\\' => {
                    let escaped = t.next().filter(|c| {
                        if is_quoted_pair(**c) {
                            true
                        } else {
                            debug!(
                                "Encountered an illegal character {} at location {}",
                                **c as char,
                                t.location().offset()
                            );
                            false
                        }
                    })?;
                    buffer.push(*escaped)
                }
                // end of quote
                b'"' => break,
                // regular character
                c if is_qdtext(*c) => buffer.push(*c),
                // Invalid character
                c => {
                    debug!(
                        "Encountered an illegal character {} at location {}",
                        *c as char,
                        t.location().offset()
                    );
                    return None;
                }
            }
        }
        Some(buffer)
    })
}

fn parse_xmatrix_field<'a>(tokens: &mut impl Tokens<Item = &'a u8>) -> Option<(String, Vec<u8>)> {
    tokens.optional(|t| {
        let name = parse_token(t).and_then(|name| {
            let name = std::str::from_utf8(&name).ok()?.to_ascii_lowercase();
            match name.as_str() {
                "origin" | "destination" | "key" | "sig" => Some(name),
                name => {
                    debug!(
                        "Encountered an invalid field name {} at location {}",
                        name,
                        t.location().offset()
                    );
                    None
                }
            }
        })?;

        if !t.token(&b'=') {
            return None;
        }

        let value = parse_quoted(t).or_else(|| parse_token_with_colons(t))?;

        Some((name, value))
    })
}

fn parse_xmatrix<'a>(tokens: &mut impl Tokens<Item = &'a u8>) -> Option<XMatrix> {
    tokens.optional(|t| {
        if !t.tokens(b"X-Matrix ".into_tokens()) {
            debug!("Failed to parse X-Matrix credentials, didn't start with 'X-Matrix '");
            return None;
        }
        let mut origin = None;
        let mut destination = None;
        let mut key = None;
        let mut sig = None;

        for (name, value) in t.sep_by(|t| parse_xmatrix_field(t), |t| t.token(&b',')) {
            match name.as_str() {
                "origin" => {
                    if origin.is_some() {
                        debug!("Field origin duplicated in X-Matrix Authorization header");
                    }
                    origin = Some(std::str::from_utf8(&value).ok()?.try_into().ok()?);
                }
                "destination" => {
                    if destination.is_some() {
                        debug!("Field destination duplicated in X-Matrix Authorization header");
                    }
                    destination = Some(std::str::from_utf8(&value).ok()?.try_into().ok()?);
                }
                "key" => {
                    if key.is_some() {
                        debug!("Field key duplicated in X-Matrix Authorization header");
                    }
                    key = Some(std::str::from_utf8(&value).ok()?.try_into().ok()?);
                }
                "sig" => {
                    if sig.is_some() {
                        debug!("Field sig duplicated in X-Matrix Authorization header");
                    }
                    sig = Some(std::str::from_utf8(&value).ok()?.to_owned());
                }
                name => {
                    debug!("Unknown field {} found in X-Matrix Authorization header", name);
                }
            }
        }

        Some(XMatrix { origin: origin?, destination, key: key?, sig: sig? })
    })
}

fn is_alpha(c: u8) -> bool {
    (0x41..=0x5A).contains(&c) || (0x61..=0x7A).contains(&c)
}

fn is_digit(c: u8) -> bool {
    (0x30..=0x39).contains(&c)
}

fn is_tchar(c: u8) -> bool {
    const TOKEN_CHARS: [u8; 15] =
        [b'!', b'#', b'$', b'%', b'&', b'\'', b'*', b'+', b'-', b'.', b'^', b'_', b'`', b'|', b'~'];
    is_alpha(c) || is_digit(c) || TOKEN_CHARS.contains(&c)
}

fn is_qdtext(c: u8) -> bool {
    c == b'\t'
        || c == b' '
        || c == 0x21
        || (0x23..=0x5B).contains(&c)
        || (0x5D..=0x7E).contains(&c)
        || is_obs_text(c)
}

fn is_obs_text(c: u8) -> bool {
    c >= 0x80 // The spec does contain an upper limit of 0xFF here, but that's enforced by the type
}

fn is_vchar(c: u8) -> bool {
    (0x21..=0x7E).contains(&c)
}

fn is_quoted_pair(c: u8) -> bool {
    c == b'\t' || c == b' ' || is_vchar(c) || is_obs_text(c)
}

impl Credentials for XMatrix {
    const SCHEME: &'static str = "X-Matrix";

    fn decode(value: &HeaderValue) -> Option<Self> {
        let value: Vec<u8> = value.as_bytes().to_vec();
        parse_xmatrix(&mut value.into_tokens())
    }

    fn encode(&self) -> HeaderValue {
        if let Some(destination) = &self.destination {
            format!(
                "X-Matrix origin=\"{}\",destination=\"{}\",key=\"{}\",sig=\"{}\"",
                self.origin, destination, self.key, self.sig
            )
        } else {
            format!("X-Matrix origin=\"{}\",key=\"{}\",sig=\"{}\"", self.origin, self.key, self.sig)
        }
        .try_into()
        .expect("header format is static")
    }
}

#[cfg(test)]
mod tests {
    use headers::{authorization::Credentials, HeaderValue};
    use ruma_common::OwnedServerName;

    use super::XMatrix;

    #[test]
    fn xmatrix_auth_pre_1_3() {
        let header = HeaderValue::from_static(
            "X-Matrix origin=\"origin.hs.example.com\",key=\"ed25519:key1\",sig=\"ABCDEF...\"",
        );
        let origin = "origin.hs.example.com".try_into().unwrap();
        let key = "ed25519:key1".try_into().unwrap();
        let sig = "ABCDEF...".to_string();
        let credentials: XMatrix = Credentials::decode(&header).unwrap();
        assert_eq!(credentials.origin, origin);
        assert_eq!(credentials.destination, None);
        assert_eq!(credentials.key, key);
        assert_eq!(credentials.sig, sig);

        let credentials = XMatrix::new(origin, None, key, sig);

        assert_eq!(credentials.encode(), header);
    }

    #[test]
    fn xmatrix_auth_1_3() {
        let header = HeaderValue::from_static("X-Matrix origin=\"origin.hs.example.com\",destination=\"destination.hs.example.com\",key=\"ed25519:key1\",sig=\"ABCDEF...\"");
        let origin: OwnedServerName = "origin.hs.example.com".try_into().unwrap();
        let destination: OwnedServerName = "destination.hs.example.com".try_into().unwrap();
        let key = "ed25519:key1".try_into().unwrap();
        let sig = "ABCDEF...".to_string();
        let credentials: XMatrix = Credentials::decode(&header).unwrap();
        assert_eq!(credentials.origin, origin);
        assert_eq!(credentials.destination, Some(destination.clone()));
        assert_eq!(credentials.key, key);
        assert_eq!(credentials.sig, sig);

        let credentials = XMatrix::new(origin, Some(destination), key, sig);

        assert_eq!(credentials.encode(), header);
    }
}
<commit_msg>server-util: Fix clippy lints<commit_after>//! Common types for implementing federation authorization.

use headers::{authorization::Credentials, HeaderValue};

use ruma_common::{OwnedServerName, OwnedServerSigningKeyId};
use tracing::debug;
use yap::{IntoTokens, TokenLocation, Tokens};

/// Typed representation of an `Authorization` header of scheme `X-Matrix`, as defined in the
/// [Matrix Server-Server API][spec]. Includes an implementation of
/// [`headers::authorization::Credentials`] for automatically handling the encoding and decoding
/// when using a web framework that supports typed headers.
///
/// [spec]: https://spec.matrix.org/v1.3/server-server-api/#request-authentication
#[non_exhaustive]
pub struct XMatrix {
    /// The server name of the sending server.
    pub origin: OwnedServerName,
    /// The server name of the receiving sender. For compatibility with older servers, recipients
    /// should accept requests without this parameter, but MUST always send it. If this property is
    /// included, but the value does not match the receiving server's name, the receiving server
    /// must deny the request with an HTTP status code 401 Unauthorized.
    pub destination: Option<OwnedServerName>,
    /// The ID - including the algorithm name - of the sending server's key that was used to sign
    /// the request.
    pub key: OwnedServerSigningKeyId,
    /// The signature of the JSON.
    pub sig: String,
}

impl XMatrix {
    /// Construct a new X-Matrix Authorization header.
    pub fn new(
        origin: OwnedServerName,
        destination: Option<OwnedServerName>,
        key: OwnedServerSigningKeyId,
        sig: String,
    ) -> Self {
        Self { origin, destination, key, sig }
    }
}

fn parse_token<'a>(tokens: &mut impl Tokens<Item = &'a u8>) -> Option<Vec<u8>> {
    tokens.optional(|t| {
        let token: Vec<u8> = t.tokens_while(|c| is_tchar(**c)).copied().collect();
        if !token.is_empty() {
            Some(token)
        } else {
            debug!("Returning early because of empty token at {}", t.location().offset());
            None
        }
    })
}

fn parse_token_with_colons<'a>(tokens: &mut impl Tokens<Item = &'a u8>) -> Option<Vec<u8>> {
    tokens.optional(|t| {
        let token: Vec<u8> = t.tokens_while(|c| is_tchar(**c) || **c == b':').copied().collect();
        if !token.is_empty() {
            Some(token)
        } else {
            debug!("Returning early because of empty token at {}", t.location().offset());
            None
        }
    })
}

fn parse_quoted<'a>(tokens: &mut impl Tokens<Item = &'a u8>) -> Option<Vec<u8>> {
    tokens.optional(|t| {
        if !(t.token(&b'"')) {
            return None;
        }
        let mut buffer = Vec::new();
        loop {
            match t.next()? {
                // quoted pair
                b'\\' => {
                    let escaped = t.next().filter(|c| {
                        if is_quoted_pair(**c) {
                            true
                        } else {
                            debug!(
                                "Encountered an illegal character {} at location {}",
                                **c as char,
                                t.location().offset()
                            );
                            false
                        }
                    })?;
                    buffer.push(*escaped);
                }
                // end of quote
                b'"' => break,
                // regular character
                c if is_qdtext(*c) => buffer.push(*c),
                // Invalid character
                c => {
                    debug!(
                        "Encountered an illegal character {} at location {}",
                        *c as char,
                        t.location().offset()
                    );
                    return None;
                }
            }
        }
        Some(buffer)
    })
}

fn parse_xmatrix_field<'a>(tokens: &mut impl Tokens<Item = &'a u8>) -> Option<(String, Vec<u8>)> {
    tokens.optional(|t| {
        let name = parse_token(t).and_then(|name| {
            let name = std::str::from_utf8(&name).ok()?.to_ascii_lowercase();
            match name.as_str() {
                "origin" | "destination" | "key" | "sig" => Some(name),
                name => {
                    debug!(
                        "Encountered an invalid field name {} at location {}",
                        name,
                        t.location().offset()
                    );
                    None
                }
            }
        })?;

        if !t.token(&b'=') {
            return None;
        }

        let value = parse_quoted(t).or_else(|| parse_token_with_colons(t))?;

        Some((name, value))
    })
}

fn parse_xmatrix<'a>(tokens: &mut impl Tokens<Item = &'a u8>) -> Option<XMatrix> {
    tokens.optional(|t| {
        if !t.tokens(b"X-Matrix ".into_tokens()) {
            debug!("Failed to parse X-Matrix credentials, didn't start with 'X-Matrix '");
            return None;
        }
        let mut origin = None;
        let mut destination = None;
        let mut key = None;
        let mut sig = None;

        for (name, value) in t.sep_by(|t| parse_xmatrix_field(t), |t| t.token(&b',')) {
            match name.as_str() {
                "origin" => {
                    if origin.is_some() {
                        debug!("Field origin duplicated in X-Matrix Authorization header");
                    }
                    origin = Some(std::str::from_utf8(&value).ok()?.try_into().ok()?);
                }
                "destination" => {
                    if destination.is_some() {
                        debug!("Field destination duplicated in X-Matrix Authorization header");
                    }
                    destination = Some(std::str::from_utf8(&value).ok()?.try_into().ok()?);
                }
                "key" => {
                    if key.is_some() {
                        debug!("Field key duplicated in X-Matrix Authorization header");
                    }
                    key = Some(std::str::from_utf8(&value).ok()?.try_into().ok()?);
                }
                "sig" => {
                    if sig.is_some() {
                        debug!("Field sig duplicated in X-Matrix Authorization header");
                    }
                    sig = Some(std::str::from_utf8(&value).ok()?.to_owned());
                }
                name => {
                    debug!("Unknown field {} found in X-Matrix Authorization header", name);
                }
            }
        }

        Some(XMatrix { origin: origin?, destination, key: key?, sig: sig? })
    })
}

fn is_alpha(c: u8) -> bool {
    (0x41..=0x5A).contains(&c) || (0x61..=0x7A).contains(&c)
}

fn is_digit(c: u8) -> bool {
    (0x30..=0x39).contains(&c)
}

fn is_tchar(c: u8) -> bool {
    const TOKEN_CHARS: [u8; 15] =
        [b'!', b'#', b'$', b'%', b'&', b'\'', b'*', b'+', b'-', b'.', b'^', b'_', b'`', b'|', b'~'];
    is_alpha(c) || is_digit(c) || TOKEN_CHARS.contains(&c)
}

fn is_qdtext(c: u8) -> bool {
    c == b'\t'
        || c == b' '
        || c == 0x21
        || (0x23..=0x5B).contains(&c)
        || (0x5D..=0x7E).contains(&c)
        || is_obs_text(c)
}

fn is_obs_text(c: u8) -> bool {
    c >= 0x80 // The spec does contain an upper limit of 0xFF here, but that's enforced by the type
}

fn is_vchar(c: u8) -> bool {
    (0x21..=0x7E).contains(&c)
}

fn is_quoted_pair(c: u8) -> bool {
    c == b'\t' || c == b' ' || is_vchar(c) || is_obs_text(c)
}

impl Credentials for XMatrix {
    const SCHEME: &'static str = "X-Matrix";

    fn decode(value: &HeaderValue) -> Option<Self> {
        let value: Vec<u8> = value.as_bytes().to_vec();
        parse_xmatrix(&mut value.into_tokens())
    }

    fn encode(&self) -> HeaderValue {
        if let Some(destination) = &self.destination {
            format!(
                "X-Matrix origin=\"{}\",destination=\"{}\",key=\"{}\",sig=\"{}\"",
                self.origin, destination, self.key, self.sig
            )
        } else {
            format!("X-Matrix origin=\"{}\",key=\"{}\",sig=\"{}\"", self.origin, self.key, self.sig)
        }
        .try_into()
        .expect("header format is static")
    }
}

#[cfg(test)]
mod tests {
    use headers::{authorization::Credentials, HeaderValue};
    use ruma_common::OwnedServerName;

    use super::XMatrix;

    #[test]
    fn xmatrix_auth_pre_1_3() {
        let header = HeaderValue::from_static(
            "X-Matrix origin=\"origin.hs.example.com\",key=\"ed25519:key1\",sig=\"ABCDEF...\"",
        );
        let origin = "origin.hs.example.com".try_into().unwrap();
        let key = "ed25519:key1".try_into().unwrap();
        let sig = "ABCDEF...".to_owned();
        let credentials: XMatrix = Credentials::decode(&header).unwrap();
        assert_eq!(credentials.origin, origin);
        assert_eq!(credentials.destination, None);
        assert_eq!(credentials.key, key);
        assert_eq!(credentials.sig, sig);

        let credentials = XMatrix::new(origin, None, key, sig);

        assert_eq!(credentials.encode(), header);
    }

    #[test]
    fn xmatrix_auth_1_3() {
        let header = HeaderValue::from_static("X-Matrix origin=\"origin.hs.example.com\",destination=\"destination.hs.example.com\",key=\"ed25519:key1\",sig=\"ABCDEF...\"");
        let origin: OwnedServerName = "origin.hs.example.com".try_into().unwrap();
        let destination: OwnedServerName = "destination.hs.example.com".try_into().unwrap();
        let key = "ed25519:key1".try_into().unwrap();
        let sig = "ABCDEF...".to_owned();
        let credentials: XMatrix = Credentials::decode(&header).unwrap();
        assert_eq!(credentials.origin, origin);
        assert_eq!(credentials.destination, Some(destination.clone()));
        assert_eq!(credentials.key, key);
        assert_eq!(credentials.sig, sig);

        let credentials = XMatrix::new(origin, Some(destination), key, sig);

        assert_eq!(credentials.encode(), header);
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2015-2017 the original author or authors.
 *
 * 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.
 */

use common::run_java;

mod common;

#[test]
fn basic() {
    assert!(!run_java("org.cloudfoundry.jvmkill.ThreadExhaustion", "", &[], &["Resource exhaustion event: the JVM was unable to create a thread.",
        "ERROR: PoolStats action failed: cannot determine memory usage statistics since the JVM is unable to create a thread"]));
}

#[test]
fn print_memory_usage_0() {
    assert!(!run_java("org.cloudfoundry.jvmkill.ThreadExhaustion", "=printMemoryUsage=0", &[], &["jvmkill killing current process"]));
}

#[test]
fn time_10_count_2() {
    assert!(!run_java("org.cloudfoundry.jvmkill.ThreadExhaustion", "=time=10,count=2,printHeapHistogram=1,heapHistogramMaxEntries=10,printMemoryUsage=0",
                      &["Lorg/cloudfoundry/jvmkill/Sleeper;"], &["ResourceExhausted! (1/2)"]));
}

#[test]
fn parallel_time_10_count_2() {
    assert!(!run_java("org.cloudfoundry.jvmkill.ParallelThreadExhaustion", "=time=10,count=2", &[],
                      &["ResourceExhausted! (1/2)", "jvmkill killing current process"]));
}
<commit_msg>Improve test stability<commit_after>/*
 * Copyright (c) 2015-2017 the original author or authors.
 *
 * 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.
 */

use common::run_java;

mod common;

#[test]
fn basic() {
    assert!(!run_java("org.cloudfoundry.jvmkill.ThreadExhaustion", "", &[], &["Resource exhaustion event: the JVM was unable to create a thread.",
        "ERROR: PoolStats action failed: cannot determine memory usage statistics since the JVM is unable to create a thread"]));
}

#[test]
fn print_memory_usage_0() {
    assert!(!run_java("org.cloudfoundry.jvmkill.ThreadExhaustion", "=printMemoryUsage=0", &[], &["jvmkill killing current process"]));
}

#[test]
fn time_10_count_2() {
    assert!(!run_java("org.cloudfoundry.jvmkill.ThreadExhaustion", "=time=10,count=2,printHeapHistogram=1,heapHistogramMaxEntries=10,printMemoryUsage=0",
                      &["Ljava/lang/Class;"], &["ResourceExhausted! (1/2)"]));
}

#[test]
fn parallel_time_10_count_2() {
    assert!(!run_java("org.cloudfoundry.jvmkill.ParallelThreadExhaustion", "=time=10,count=2", &[],
                      &["ResourceExhausted! (1/2)", "jvmkill killing current process"]));
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use std::{fs::File, io::prelude::*};
use vm_genesis::{default_config, encode_genesis_transaction, GENESIS_KEYPAIR};

const CONFIG_LOCATION: &str = "genesis/vm_config.toml";
const GENESIS_LOCATION: &str = "genesis/genesis.blob";

use proto_conv::IntoProtoBytes;

fn main() {
    println!(
        "Creating genesis binary blob at {} from configuration file {}",
        GENESIS_LOCATION, CONFIG_LOCATION
    );
    let config = default_config();
    config.save_config(CONFIG_LOCATION);

    // Generate a genesis blob used for vm tests.
    let genesis_txn = encode_genesis_transaction(&GENESIS_KEYPAIR.0, GENESIS_KEYPAIR.1.clone());
    let mut file = File::create(GENESIS_LOCATION).unwrap();
    file.write_all(&genesis_txn.into_proto_bytes().unwrap())
        .unwrap();
}
<commit_msg>[language] test that flags genesis blob changes<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use std::{fs::File, io::prelude::*};
use vm_genesis::{default_config, encode_genesis_transaction, GENESIS_KEYPAIR};

const CONFIG_LOCATION: &str = "genesis/vm_config.toml";
const GENESIS_LOCATION: &str = "genesis/genesis.blob";

use proto_conv::IntoProtoBytes;

/// Generate the genesis blob used by the Libra blockchain
fn generate_genesis_blob() -> Vec<u8> {
    encode_genesis_transaction(&GENESIS_KEYPAIR.0, GENESIS_KEYPAIR.1.clone())
        .into_proto_bytes()
        .expect("Generating genesis block failed")
}

fn main() {
    println!(
        "Creating genesis binary blob at {} from configuration file {}",
        GENESIS_LOCATION, CONFIG_LOCATION
    );
    let config = default_config();
    config.save_config(CONFIG_LOCATION);

    let mut file = File::create(GENESIS_LOCATION).unwrap();
    file.write_all(&generate_genesis_blob()).unwrap();
}

// A test that fails if the generated genesis blob is different from the one on  disk. Intended
// to catch commits that
// - accidentally change the genesis block
// - change it without remembering to update the on-disk copy
// - cause generation of the genesis block to fail
#[test]
fn genesis_blob_unchanged() {
    let mut genesis_file = File::open(GENESIS_LOCATION).unwrap();
    let mut old_genesis_bytes = vec![];
    genesis_file.read_to_end(&mut old_genesis_bytes).unwrap();
    assert!(old_genesis_bytes == generate_genesis_blob(),
            format!("The freshly generated genesis file is different from the one on disk at {}. Did you forget to regenerate the genesis file via `cargo run` inside libra/language/vm/vm_genesis?", GENESIS_LOCATION));
}
<|endoftext|>