<commit_before>#![feature(type_ascription)]

#[macro_use]
extern crate rustless;
extern crate jsonway;
extern crate iron;
extern crate hyper;
extern crate valico;
extern crate env_logger;
extern crate clap;
extern crate serde;
extern crate time;
extern crate rand;
extern crate log;
extern crate colored;

extern crate exonum;
extern crate blockchain_explorer;
extern crate cryptocurrency;

use std::net::SocketAddr;
use std::path::Path;
use std::thread;
use std::default::Default;
use std::env;

use clap::{Arg, App, SubCommand};
use rustless::json::ToJson;
use rustless::{Application, Api, Nesting, Versioning, Response};
use rustless::batteries::cookie::{Cookie, CookieExt, CookieJar};
use rustless::batteries::swagger;
use valico::json_dsl;
use hyper::status::StatusCode;
use rand::{Rng, thread_rng};
use log::{LogRecord, LogLevel};
use env_logger::LogBuilder;
use colored::*;

use exonum::node::{Node, Configuration, TxSender, NodeChannel};
use exonum::storage::{Database, MemoryDB, LevelDB, LevelDBOptions};
use exonum::storage::{Result as StorageResult, Error as StorageError};
use exonum::crypto::{gen_keypair, PublicKey, SecretKey, HexValue};
use exonum::messages::Message;
use exonum::config::ConfigFile;
use exonum::node::config::GenesisConfig;

use cryptocurrency::{CurrencyBlockchain, CurrencyTx, TxIssue, TxTransfer, TxCreateWallet};
use cryptocurrency::api::CurrencyApi;

pub type CurrencyTxSender<B> = TxSender<B, NodeChannel<B>>;

fn save_keypair_in_cookies(storage: &mut CookieJar,
                           public_key: &PublicKey,
                           secret_key: &SecretKey) {
    let p = storage.permanent();
    let e = p.encrypted();

    e.add(Cookie::new("public_key".to_string(), public_key.to_hex()));
    e.add(Cookie::new("secret_key".to_string(), secret_key.to_hex()));
}

fn load_hex_value_from_cookie<'a>(storage: &'a CookieJar, key: &str) -> StorageResult<Vec<u8>> {
    if let Some(cookie) = storage.find(key) {
        if let Ok(value) = HexValue::from_hex(cookie.value) {
            return Ok(value);
        }
    }
    Err(StorageError::new(format!("Unable to find value with given key {}", key)))
}

fn load_keypair_from_cookies(storage: &CookieJar) -> StorageResult<(PublicKey, SecretKey)> {
    let p = storage.permanent();
    let e = p.encrypted();

    let public_key = PublicKey::from_slice(load_hex_value_from_cookie(&e, "public_key")?.as_ref());
    let secret_key = SecretKey::from_slice(load_hex_value_from_cookie(&e, "secret_key")?.as_ref());

    let public_key = public_key.ok_or(StorageError::new("Unable to read public key"))?;
    let secret_key = secret_key.ok_or(StorageError::new("Unable to read secret key"))?;
    Ok((public_key, secret_key))
}

fn blockchain_explorer_api<D: Database>(api: &mut Api,
                                        b: CurrencyBlockchain<D>,
                                        cfg: Configuration) {
    blockchain_explorer::make_api::<CurrencyBlockchain<D>, CurrencyTx>(api, b, cfg);
}

fn cryptocurrency_api<D: Database>(api: &mut Api,
                                   blockchain: CurrencyBlockchain<D>,
                                   channel: CurrencyTxSender<CurrencyBlockchain<D>>,
                                   cfg: Configuration) {
    api.namespace("wallets", move |api| {
        let ch = channel.clone();
        api.post("create", move |endpoint| {
            endpoint.summary("Create a new wallet for user with given name");
            endpoint.params(|params| {
                params.req_typed("name", json_dsl::string());
            });

            endpoint.handle(move |client, params| {
                let name = params.find("name").unwrap().as_str().unwrap();
                // TODO make secure
                let (public_key, secret_key) = gen_keypair();
                {
                    let mut cookies = client.request.cookies();
                    save_keypair_in_cookies(&mut cookies, &public_key, &secret_key);
                }
                let tx = TxCreateWallet::new(&public_key, &name, &secret_key);

                let tx_hash = tx.hash().to_hex();
                ch.send(CurrencyTx::CreateWallet(tx));
                let json = &jsonway::object(|json| json.set("tx_hash", tx_hash)).unwrap();
                client.json(json)
            })
        });

        let ch = channel.clone();
        api.post("issue", move |endpoint| {
            endpoint.params(|params| {
                params.req_typed("amount", json_dsl::i64());
            });

            endpoint.handle(move |client, params| {
                let (public_key, secret_key) = {
                    let r = {
                        let cookies = client.request.cookies();
                        load_keypair_from_cookies(&cookies)
                    };
                    match r {
                        Ok((p, s)) => (p, s),
                        Err(e) => return client.error(e),
                    }
                };

                let amount = params.find("amount").unwrap().as_i64().unwrap();
                let seed = thread_rng().gen::<u64>();
                let tx = TxIssue::new(&public_key, amount, seed, &secret_key);

                let tx_hash = tx.hash().to_hex();
                ch.send(CurrencyTx::Issue(tx));
                let json = &jsonway::object(|json| json.set("tx_hash", tx_hash)).unwrap();
                client.json(json)
            })
        });

        let ch = channel.clone();
        api.post("transfer", move |endpoint| {
            endpoint.params(|params| {
                params.req_typed("amount", json_dsl::i64());
                params.req_typed("to", json_dsl::string());
            });

            endpoint.handle(move |client, params| {
                let (from_key, secret_key) = {
                    let r = {
                        let cookies = client.request.cookies();
                        load_keypair_from_cookies(&cookies)
                    };
                    match r {
                        Ok((p, s)) => (p, s),
                        Err(e) => return client.error(e),
                    }
                };
                // TODO remove unwrap
                let amount = params.find("amount").unwrap().as_i64().unwrap();
                let to = params.find("to").unwrap().as_str().unwrap();
                let seed = thread_rng().gen::<u64>();
                match PublicKey::from_hex(to) {
                    Ok(to_key) => {
                        let tx = TxTransfer::new(&from_key, &to_key, amount, seed, &secret_key);

                        let tx_hash = tx.hash().to_hex();
                        ch.send(CurrencyTx::Transfer(tx));
                        let json = &jsonway::object(|json| json.set("tx_hash", tx_hash)).unwrap();
                        client.json(json)
                    }
                    Err(e) => return client.error(e),
                }
            })
        });

        let b = blockchain.clone();
        let c = cfg.clone();
        api.get("info", move |endpoint| {
            endpoint.handle(move |client, _| {
                let (public_key, _) = {
                    let r = {
                        let cookies = client.request.cookies();
                        load_keypair_from_cookies(&cookies)
                    };
                    match r {
                        Ok((p, s)) => (p, s),
                        Err(e) => return client.error(e),
                    }
                };
                let currency_api = CurrencyApi::<D>::new(b.clone(), c.clone());
                match currency_api.wallet_info(&public_key) {
                    Ok(Some(info)) => client.json(&info.to_json()),
                    _ => client.error(StorageError::new("Unable to get wallet info")),
                }
            })
        })
    });
}

fn run_node<D: Database>(blockchain: CurrencyBlockchain<D>,
                         node_cfg: Configuration,
                         port: Option<u16>) {
    if let Some(port) = port {
        let mut node = Node::new(blockchain.clone(), node_cfg.clone());
        let channel = node.channel();

        let api_thread = thread::spawn(move || {
            let channel = channel.clone();
            let blockchain = blockchain.clone();

            let api = Api::build(move |api| {
                // Specify API version
                api.version("v1", Versioning::Path);
                api.prefix("api");

                api.error_formatter(|err, _media| {
                    if let Some(e) = err.downcast::<StorageError>() {
                        let body = format!("An internal error occured: {}", e);
                        Some(Response::from(StatusCode::InternalServerError, Box::new(body)))
                    } else {
                        None
                    }
                });

                blockchain_explorer_api(api, blockchain.clone(), node_cfg.clone());
                cryptocurrency_api(api, blockchain.clone(), channel.clone(), node_cfg);
                api.mount(swagger::create_api("docs"));
            });

            let listen_address: SocketAddr = format!("127.0.0.1:{}", port).parse().unwrap();
            println!("Cryptocurrency node server started on {}", listen_address);

            let mut app = Application::new(api);

            swagger::enable(&mut app,
                            swagger::Spec {
                                info: swagger::Info {
                                    title: "Cryptocurrency API".to_string(),
                                    description: Some("Simple API to demonstration".to_string()),
                                    contact: Some(swagger::Contact {
                                        name: "Aleksey Sidorov".to_string(),
                                        url: Some("aleksei.sidorov@xdev.re".to_string()),
                                        ..Default::default()
                                    }),
                                    license: Some(swagger::License {
                                        name: "Demo".to_string(),
                                        url: "http://exonum.com".to_string(),
                                    }),
                                    ..Default::default()
                                },
                                ..Default::default()
                            });

            let mut chain = iron::Chain::new(app);
            let api_key = b"abacabasecretcoinblabla23nx8Hasojd8";
            let cookie = ::rustless::batteries::cookie::new(api_key);
            chain.link(cookie);
            iron::Iron::new(chain).http(listen_address).unwrap();
        });

        node.run().unwrap();
        api_thread.join().unwrap();
    } else {
        Node::new(blockchain, node_cfg).run().unwrap();
    }
}

fn main() {
    let format = |record: &LogRecord| {
        let now = time::now_utc();
        let level = match record.level() {
            LogLevel::Error => "ERROR".red(),
            LogLevel::Warn => "WARN".yellow(),
            LogLevel::Info => "INFO".green(),
            LogLevel::Debug => "DEBUG".cyan(),
            LogLevel::Trace => "TRACE".white(),
        };
        format!("{} - [ {} ] - {}",
                now.asctime().to_string().bold(),
                level,
                record.args())
    };

    let mut builder = LogBuilder::new();
    builder.format(format);

    if env::var("RUST_LOG").is_ok() {
        builder.parse(&env::var("RUST_LOG").unwrap());
    }

    builder.init().unwrap();

    let app = App::new("Simple cryptocurrency demo program")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Aleksey S. <aleksei.sidorov@xdev.re>")
        .about("Demo cryptocurrency validator node")
        .arg(Arg::with_name("CONFIG")
            .short("c")
            .long("config")
            .value_name("CONFIG_PATH")
            .help("Sets a node config file")
            .required(true)
            .takes_value(true))
        .subcommand(SubCommand::with_name("generate")
            .about("Generates default configuration file")
            .version(env!("CARGO_PKG_VERSION"))
            .author("Aleksey S. <aleksei.sidorov@xdev.re>")
            .arg(Arg::with_name("START_PORT")
                .short("p")
                .long("port")
                .value_name("START_PORT")
                .help("Port for first validator")
                .takes_value(true))
            .arg(Arg::with_name("COUNT")
                .help("Validators count")
                .required(true)
                .index(1)))
        .subcommand(SubCommand::with_name("run")
            .about("Run demo node with the given validator id")
            .version(env!("CARGO_PKG_VERSION"))
            .author("Aleksey S. <aleksei.sidorov@xdev.re>")
            .arg(Arg::with_name("LEVELDB_PATH")
                .short("d")
                .long("leveldb-path")
                .value_name("LEVELDB_PATH")
                .help("Use leveldb database with the given path")
                .takes_value(true))
            .arg(Arg::with_name("HTTP_PORT")
                .short("p")
                .long("port")
                .value_name("HTTP_PORT")
                .help("Run http server on given port")
                .takes_value(true))
            .arg(Arg::with_name("PEERS")
                .long("known-peers")
                .value_name("PEERS")
                .help("Comma separated list of known validator ids")
                .takes_value(true))
            .arg(Arg::with_name("VALIDATOR")
                .help("Sets a validator id")
                .required(true)
                .index(1)));

    let matches = app.get_matches();
    let path = Path::new(matches.value_of("CONFIG").unwrap());
    match matches.subcommand() {
        ("generate", Some(matches)) => {
            let count: u8 = matches.value_of("COUNT").unwrap().parse().unwrap();
            let port: Option<u16> = matches.value_of("START_PORT").map(|x| x.parse().unwrap());
            let cfg = GenesisConfig::gen(count, port);
            ConfigFile::save(&cfg, &path).unwrap();
            println!("The configuration was successfully written to file {:?}",
                     path);
        }
        ("run", Some(matches)) => {
            let cfg: GenesisConfig = ConfigFile::load(path).unwrap();
            let idx: usize = matches.value_of("VALIDATOR").unwrap().parse().unwrap();
            let port: Option<u16> = matches.value_of("HTTP_PORT").map(|x| x.parse().unwrap());
            let peers = match matches.value_of("PEERS") {
                Some(string) => {
                    string.split(" ")
                        .map(|x| -> usize { x.parse().unwrap() })
                        .map(|x| cfg.validators[x].address)
                        .collect()
                }
                None => {
                    cfg.validators
                        .iter()
                        .map(|v| v.address)
                        .collect()
                }
            };
            let node_cfg = cfg.gen_node_configuration(idx, peers);
            match matches.value_of("LEVELDB_PATH") {
                Some(ref db_path) => {
                    println!("Using levedb storage with path: {}", db_path);
                    let mut options = LevelDBOptions::new();
                    options.create_if_missing = true;
                    let leveldb = LevelDB::new(&Path::new(db_path), options).unwrap();

                    let blockchain = CurrencyBlockchain { db: leveldb };
                    run_node(blockchain, node_cfg, port);
                }
                None => {
                    println!("Using memorydb storage");

                    let blockchain = CurrencyBlockchain { db: MemoryDB::new() };
                    run_node(blockchain, node_cfg, port);
                }
            };
        }
        _ => {
            unreachable!("Wrong subcommand");
        }
    }
}
<commit_msg>Disable encryption in cookies<commit_after>#![feature(type_ascription)]

#[macro_use]
extern crate rustless;
extern crate jsonway;
extern crate iron;
extern crate hyper;
extern crate valico;
extern crate env_logger;
extern crate clap;
extern crate serde;
extern crate time;
extern crate rand;
extern crate log;
extern crate colored;

extern crate exonum;
extern crate blockchain_explorer;
extern crate cryptocurrency;

use std::net::SocketAddr;
use std::path::Path;
use std::thread;
use std::default::Default;
use std::env;

use clap::{Arg, App, SubCommand};
use rustless::json::ToJson;
use rustless::{Application, Api, Nesting, Versioning, Response};
use rustless::batteries::cookie::{Cookie, CookieExt, CookieJar};
use rustless::batteries::swagger;
use valico::json_dsl;
use hyper::status::StatusCode;
use rand::{Rng, thread_rng};
use log::{LogRecord, LogLevel};
use env_logger::LogBuilder;
use colored::*;

use exonum::node::{Node, Configuration, TxSender, NodeChannel};
use exonum::storage::{Database, MemoryDB, LevelDB, LevelDBOptions};
use exonum::storage::{Result as StorageResult, Error as StorageError};
use exonum::crypto::{gen_keypair, PublicKey, SecretKey, HexValue};
use exonum::messages::Message;
use exonum::config::ConfigFile;
use exonum::node::config::GenesisConfig;

use cryptocurrency::{CurrencyBlockchain, CurrencyTx, TxIssue, TxTransfer, TxCreateWallet};
use cryptocurrency::api::CurrencyApi;

pub type CurrencyTxSender<B> = TxSender<B, NodeChannel<B>>;

fn save_keypair_in_cookies(storage: &mut CookieJar,
                           public_key: &PublicKey,
                           secret_key: &SecretKey) {
    let p = storage.permanent();

    p.add(Cookie::new("public_key".to_string(), public_key.to_hex()));
    p.add(Cookie::new("secret_key".to_string(), secret_key.to_hex()));
}

fn load_hex_value_from_cookie<'a>(storage: &'a CookieJar, key: &str) -> StorageResult<Vec<u8>> {
    if let Some(cookie) = storage.find(key) {
        if let Ok(value) = HexValue::from_hex(cookie.value) {
            return Ok(value);
        }
    }
    Err(StorageError::new(format!("Unable to find value with given key {}", key)))
}

fn load_keypair_from_cookies(storage: &CookieJar) -> StorageResult<(PublicKey, SecretKey)> {
    let p = storage.permanent();

    let public_key = PublicKey::from_slice(load_hex_value_from_cookie(&p, "public_key")?.as_ref());
    let secret_key = SecretKey::from_slice(load_hex_value_from_cookie(&p, "secret_key")?.as_ref());

    let public_key = public_key.ok_or(StorageError::new("Unable to read public key"))?;
    let secret_key = secret_key.ok_or(StorageError::new("Unable to read secret key"))?;
    Ok((public_key, secret_key))
}

fn blockchain_explorer_api<D: Database>(api: &mut Api,
                                        b: CurrencyBlockchain<D>,
                                        cfg: Configuration) {
    blockchain_explorer::make_api::<CurrencyBlockchain<D>, CurrencyTx>(api, b, cfg);
}

fn cryptocurrency_api<D: Database>(api: &mut Api,
                                   blockchain: CurrencyBlockchain<D>,
                                   channel: CurrencyTxSender<CurrencyBlockchain<D>>,
                                   cfg: Configuration) {
    api.namespace("wallets", move |api| {
        let ch = channel.clone();
        api.post("create", move |endpoint| {
            endpoint.summary("Create a new wallet for user with given name");
            endpoint.params(|params| {
                params.req_typed("name", json_dsl::string());
            });

            endpoint.handle(move |client, params| {
                let name = params.find("name").unwrap().as_str().unwrap();
                // TODO make secure
                let (public_key, secret_key) = gen_keypair();
                {
                    let mut cookies = client.request.cookies();
                    save_keypair_in_cookies(&mut cookies, &public_key, &secret_key);
                }
                let tx = TxCreateWallet::new(&public_key, &name, &secret_key);

                let tx_hash = tx.hash().to_hex();
                ch.send(CurrencyTx::CreateWallet(tx));
                let json = &jsonway::object(|json| json.set("tx_hash", tx_hash)).unwrap();
                client.json(json)
            })
        });

        let ch = channel.clone();
        api.post("issue", move |endpoint| {
            endpoint.params(|params| {
                params.req_typed("amount", json_dsl::i64());
            });

            endpoint.handle(move |client, params| {
                let (public_key, secret_key) = {
                    let r = {
                        let cookies = client.request.cookies();
                        load_keypair_from_cookies(&cookies)
                    };
                    match r {
                        Ok((p, s)) => (p, s),
                        Err(e) => return client.error(e),
                    }
                };

                let amount = params.find("amount").unwrap().as_i64().unwrap();
                let seed = thread_rng().gen::<u64>();
                let tx = TxIssue::new(&public_key, amount, seed, &secret_key);

                let tx_hash = tx.hash().to_hex();
                ch.send(CurrencyTx::Issue(tx));
                let json = &jsonway::object(|json| json.set("tx_hash", tx_hash)).unwrap();
                client.json(json)
            })
        });

        let ch = channel.clone();
        api.post("transfer", move |endpoint| {
            endpoint.params(|params| {
                params.req_typed("amount", json_dsl::i64());
                params.req_typed("to", json_dsl::string());
            });

            endpoint.handle(move |client, params| {
                let (from_key, secret_key) = {
                    let r = {
                        let cookies = client.request.cookies();
                        load_keypair_from_cookies(&cookies)
                    };
                    match r {
                        Ok((p, s)) => (p, s),
                        Err(e) => return client.error(e),
                    }
                };
                // TODO remove unwrap
                let amount = params.find("amount").unwrap().as_i64().unwrap();
                let to = params.find("to").unwrap().as_str().unwrap();
                let seed = thread_rng().gen::<u64>();
                match PublicKey::from_hex(to) {
                    Ok(to_key) => {
                        let tx = TxTransfer::new(&from_key, &to_key, amount, seed, &secret_key);

                        let tx_hash = tx.hash().to_hex();
                        ch.send(CurrencyTx::Transfer(tx));
                        let json = &jsonway::object(|json| json.set("tx_hash", tx_hash)).unwrap();
                        client.json(json)
                    }
                    Err(e) => return client.error(e),
                }
            })
        });

        let b = blockchain.clone();
        let c = cfg.clone();
        api.get("info", move |endpoint| {
            endpoint.handle(move |client, _| {
                let (public_key, _) = {
                    let r = {
                        let cookies = client.request.cookies();
                        load_keypair_from_cookies(&cookies)
                    };
                    match r {
                        Ok((p, s)) => (p, s),
                        Err(e) => return client.error(e),
                    }
                };
                let currency_api = CurrencyApi::<D>::new(b.clone(), c.clone());
                match currency_api.wallet_info(&public_key) {
                    Ok(Some(info)) => client.json(&info.to_json()),
                    _ => client.error(StorageError::new("Unable to get wallet info")),
                }
            })
        })
    });
}

fn run_node<D: Database>(blockchain: CurrencyBlockchain<D>,
                         node_cfg: Configuration,
                         port: Option<u16>) {
    if let Some(port) = port {
        let mut node = Node::new(blockchain.clone(), node_cfg.clone());
        let channel = node.channel();

        let api_thread = thread::spawn(move || {
            let channel = channel.clone();
            let blockchain = blockchain.clone();

            let api = Api::build(move |api| {
                // Specify API version
                api.version("v1", Versioning::Path);
                api.prefix("api");

                api.error_formatter(|err, _media| {
                    if let Some(e) = err.downcast::<StorageError>() {
                        let body = format!("An internal error occured: {}", e);
                        Some(Response::from(StatusCode::InternalServerError, Box::new(body)))
                    } else {
                        None
                    }
                });

                blockchain_explorer_api(api, blockchain.clone(), node_cfg.clone());
                cryptocurrency_api(api, blockchain.clone(), channel.clone(), node_cfg);
                api.mount(swagger::create_api("docs"));
            });

            let listen_address: SocketAddr = format!("127.0.0.1:{}", port).parse().unwrap();
            println!("Cryptocurrency node server started on {}", listen_address);

            let mut app = Application::new(api);

            swagger::enable(&mut app,
                            swagger::Spec {
                                info: swagger::Info {
                                    title: "Cryptocurrency API".to_string(),
                                    description: Some("Simple API to demonstration".to_string()),
                                    contact: Some(swagger::Contact {
                                        name: "Aleksey Sidorov".to_string(),
                                        url: Some("aleksei.sidorov@xdev.re".to_string()),
                                        ..Default::default()
                                    }),
                                    license: Some(swagger::License {
                                        name: "Demo".to_string(),
                                        url: "http://exonum.com".to_string(),
                                    }),
                                    ..Default::default()
                                },
                                ..Default::default()
                            });

            let mut chain = iron::Chain::new(app);
            let api_key = b"abacabasecretcoinblabla23nx8Hasojd8";
            let cookie = ::rustless::batteries::cookie::new(api_key);
            chain.link(cookie);
            iron::Iron::new(chain).http(listen_address).unwrap();
        });

        node.run().unwrap();
        api_thread.join().unwrap();
    } else {
        Node::new(blockchain, node_cfg).run().unwrap();
    }
}

fn main() {
    let format = |record: &LogRecord| {
        let now = time::now_utc();
        let level = match record.level() {
            LogLevel::Error => "ERROR".red(),
            LogLevel::Warn => "WARN".yellow(),
            LogLevel::Info => "INFO".green(),
            LogLevel::Debug => "DEBUG".cyan(),
            LogLevel::Trace => "TRACE".white(),
        };
        format!("{} - [ {} ] - {}",
                now.asctime().to_string().bold(),
                level,
                record.args())
    };

    let mut builder = LogBuilder::new();
    builder.format(format);

    if env::var("RUST_LOG").is_ok() {
        builder.parse(&env::var("RUST_LOG").unwrap());
    }

    builder.init().unwrap();

    let app = App::new("Simple cryptocurrency demo program")
        .version(env!("CARGO_PKG_VERSION"))
        .author("Aleksey S. <aleksei.sidorov@xdev.re>")
        .about("Demo cryptocurrency validator node")
        .arg(Arg::with_name("CONFIG")
            .short("c")
            .long("config")
            .value_name("CONFIG_PATH")
            .help("Sets a node config file")
            .required(true)
            .takes_value(true))
        .subcommand(SubCommand::with_name("generate")
            .about("Generates default configuration file")
            .version(env!("CARGO_PKG_VERSION"))
            .author("Aleksey S. <aleksei.sidorov@xdev.re>")
            .arg(Arg::with_name("START_PORT")
                .short("p")
                .long("port")
                .value_name("START_PORT")
                .help("Port for first validator")
                .takes_value(true))
            .arg(Arg::with_name("COUNT")
                .help("Validators count")
                .required(true)
                .index(1)))
        .subcommand(SubCommand::with_name("run")
            .about("Run demo node with the given validator id")
            .version(env!("CARGO_PKG_VERSION"))
            .author("Aleksey S. <aleksei.sidorov@xdev.re>")
            .arg(Arg::with_name("LEVELDB_PATH")
                .short("d")
                .long("leveldb-path")
                .value_name("LEVELDB_PATH")
                .help("Use leveldb database with the given path")
                .takes_value(true))
            .arg(Arg::with_name("HTTP_PORT")
                .short("p")
                .long("port")
                .value_name("HTTP_PORT")
                .help("Run http server on given port")
                .takes_value(true))
            .arg(Arg::with_name("PEERS")
                .long("known-peers")
                .value_name("PEERS")
                .help("Comma separated list of known validator ids")
                .takes_value(true))
            .arg(Arg::with_name("VALIDATOR")
                .help("Sets a validator id")
                .required(true)
                .index(1)));

    let matches = app.get_matches();
    let path = Path::new(matches.value_of("CONFIG").unwrap());
    match matches.subcommand() {
        ("generate", Some(matches)) => {
            let count: u8 = matches.value_of("COUNT").unwrap().parse().unwrap();
            let port: Option<u16> = matches.value_of("START_PORT").map(|x| x.parse().unwrap());
            let cfg = GenesisConfig::gen(count, port);
            ConfigFile::save(&cfg, &path).unwrap();
            println!("The configuration was successfully written to file {:?}",
                     path);
        }
        ("run", Some(matches)) => {
            let cfg: GenesisConfig = ConfigFile::load(path).unwrap();
            let idx: usize = matches.value_of("VALIDATOR").unwrap().parse().unwrap();
            let port: Option<u16> = matches.value_of("HTTP_PORT").map(|x| x.parse().unwrap());
            let peers = match matches.value_of("PEERS") {
                Some(string) => {
                    string.split(" ")
                        .map(|x| -> usize { x.parse().unwrap() })
                        .map(|x| cfg.validators[x].address)
                        .collect()
                }
                None => {
                    cfg.validators
                        .iter()
                        .map(|v| v.address)
                        .collect()
                }
            };
            let node_cfg = cfg.gen_node_configuration(idx, peers);
            match matches.value_of("LEVELDB_PATH") {
                Some(ref db_path) => {
                    println!("Using levedb storage with path: {}", db_path);
                    let mut options = LevelDBOptions::new();
                    options.create_if_missing = true;
                    let leveldb = LevelDB::new(&Path::new(db_path), options).unwrap();

                    let blockchain = CurrencyBlockchain { db: leveldb };
                    run_node(blockchain, node_cfg, port);
                }
                None => {
                    println!("Using memorydb storage");

                    let blockchain = CurrencyBlockchain { db: MemoryDB::new() };
                    run_node(blockchain, node_cfg, port);
                }
            };
        }
        _ => {
            unreachable!("Wrong subcommand");
        }
    }
}
<|endoftext|>


<commit_before>use std::{
    net::{Ipv4Addr, SocketAddrV4, TcpListener, TcpStream},
    path::PathBuf,
};

use clap::Parser;
use dicom_core::{dicom_value, DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_encoding::transfer_syntax::TransferSyntaxIndex;
use dicom_object::{FileMetaTableBuilder, InMemDicomObject, StandardDataDictionary};
use dicom_transfer_syntax_registry::TransferSyntaxRegistry;
use dicom_ul::{pdu::PDataValueType, Pdu};
use tracing::{debug, error, info, warn, Level};

use crate::transfer::ABSTRACT_SYNTAXES;

mod transfer;

/// DICOM C-STORE SCP
#[derive(Debug, Parser)]
struct App {
    /// verbose mode
    #[clap(short = 'v', long = "verbose")]
    verbose: bool,
    /// the calling Application Entity title
    #[structopt(long = "calling-ae-title", default_value = "STORE-SCP")]
    calling_ae_title: String,
    /// enforce max pdu length
    #[clap(short = 's', long = "strict")]
    strict: bool,
    /// Only accept native/uncompressed transfer syntaxes
    #[clap(long)]
    uncompressed_only: bool,
    /// max pdu length
    #[clap(short = 'm', long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,
    /// output directory for incoming objects
    #[clap(short = 'o', default_value = ".")]
    out_dir: PathBuf,
    /// Which port to listen on
    #[clap(short, default_value = "11111")]
    port: u16,
}

fn run(scu_stream: TcpStream, args: &App) -> Result<(), Box<dyn std::error::Error>> {
    let App {
        verbose,
        calling_ae_title,
        strict,
        uncompressed_only,
        max_pdu_length,
        out_dir,
        port: _,
    } = args;
    let verbose = *verbose;

    let mut buffer: Vec<u8> = Vec::with_capacity(*max_pdu_length as usize);
    let mut instance_buffer: Vec<u8> = Vec::with_capacity(1024 * 1024);
    let mut msgid = 1;
    let mut sop_class_uid = "".to_string();
    let mut sop_instance_uid = "".to_string();

    let mut options = dicom_ul::association::ServerAssociationOptions::new()
        .accept_any()
        .ae_title(calling_ae_title)
        .strict(*strict);

    if *uncompressed_only {
        options = options
            .with_transfer_syntax("1.2.840.10008.1.2")
            .with_transfer_syntax("1.2.840.10008.1.2.1");
    } else {
        for ts in TransferSyntaxRegistry.iter() {
            if !ts.unsupported() {
                options = options.with_transfer_syntax(ts.uid());
            }
        }
    };

    for uid in ABSTRACT_SYNTAXES {
        options = options.with_abstract_syntax(*uid);
    }

    let mut association = options.establish(scu_stream)?;

    info!("New association from {}", association.client_ae_title());
    debug!(
        "> Presentation contexts: {:?}",
        association.presentation_contexts()
    );

    loop {
        match association.receive() {
            Ok(mut pdu) => {
                if verbose {
                    debug!("scu ----> scp: {}", pdu.short_description());
                }
                match pdu {
                    Pdu::PData { ref mut data } => {
                        if data[0].value_type == PDataValueType::Data && !data[0].is_last {
                            instance_buffer.append(&mut data[0].data);
                        } else if data[0].value_type == PDataValueType::Command && data[0].is_last {
                            // commands are always in implict VR LE
                            let ts =
                                dicom_transfer_syntax_registry::entries::IMPLICIT_VR_LITTLE_ENDIAN
                                    .erased();
                            let data_value = &data[0];
                            let v = &data_value.data;

                            let obj = InMemDicomObject::read_dataset_with_ts(v.as_slice(), &ts)?;
                            msgid = obj.element(tags::MESSAGE_ID)?.to_int()?;
                            sop_class_uid = obj
                                .element(tags::AFFECTED_SOP_CLASS_UID)?
                                .to_str()?
                                .to_string();
                            sop_instance_uid = obj
                                .element(tags::AFFECTED_SOP_INSTANCE_UID)?
                                .to_str()?
                                .to_string();
                            instance_buffer.clear();
                        } else if data[0].value_type == PDataValueType::Data && data[0].is_last {
                            instance_buffer.append(&mut data[0].data);

                            let presentation_context = association
                                .presentation_contexts()
                                .iter()
                                .filter(|pc| pc.id == data[0].presentation_context_id)
                                .next()
                                .unwrap();
                            let ts = &presentation_context.transfer_syntax;

                            let obj = InMemDicomObject::read_dataset_with_ts(
                                instance_buffer.as_slice(),
                                TransferSyntaxRegistry.get(ts).unwrap(),
                            )?;
                            let file_meta = FileMetaTableBuilder::new()
                                .media_storage_sop_class_uid(
                                    obj.element_by_name("SOPClassUID")?.to_str()?,
                                )
                                .media_storage_sop_instance_uid(
                                    obj.element_by_name("SOPInstanceUID")?.to_str()?,
                                )
                                .transfer_syntax(ts)
                                .build()?;
                            let file_obj = obj.with_exact_meta(file_meta);

                            // write the files to the current directory with their SOPInstanceUID as filenames
                            let mut file_path = out_dir.clone();
                            file_path
                                .push(sop_instance_uid.trim_end_matches('\0').to_string() + ".dcm");
                            file_obj.write_to_file(&file_path)?;
                            info!("Stored {}", file_path.display());

                            // send C-STORE-RSP object
                            // commands are always in implict VR LE
                            let ts =
                                dicom_transfer_syntax_registry::entries::IMPLICIT_VR_LITTLE_ENDIAN
                                    .erased();

                            let obj =
                                create_cstore_response(msgid, &sop_class_uid, &sop_instance_uid);

                            let mut obj_data = Vec::new();

                            obj.write_dataset_with_ts(&mut obj_data, &ts)?;

                            let pdu_response = Pdu::PData {
                                data: vec![dicom_ul::pdu::PDataValue {
                                    presentation_context_id: data[0].presentation_context_id,
                                    value_type: PDataValueType::Command,
                                    is_last: true,
                                    data: obj_data,
                                }],
                            };
                            association.send(&pdu_response)?;
                        }
                    }
                    Pdu::ReleaseRQ => {
                        buffer.clear();
                        association.send(&Pdu::ReleaseRP)?;
                        info!(
                            "Released association with {}",
                            association.client_ae_title()
                        );
                    }
                    _ => {}
                }
            }
            Err(err @ dicom_ul::association::server::Error::Receive { .. }) => {
                debug!("{}", err);
                break;
            }
            Err(err) => {
                warn!("Unexpected error: {}", err);
                break;
            }
        }
    }
    info!("Association with {} dropped", association.client_ae_title());
    Ok(())
}

fn create_cstore_response(
    message_id: u16,
    sop_class_uid: &str,
    sop_instance_uid: &str,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + sop_class_uid.len() as i32
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + sop_instance_uid.len() as i32,
        ),
    ));

    // service
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, sop_class_uid),
    ));
    // command
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x8001]),
    ));
    // message ID being responded to
    obj.put(DataElement::new(
        tags::MESSAGE_ID_BEING_RESPONDED_TO,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));
    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0101]),
    ));
    // status https://dicom.nema.org/dicom/2013/output/chtml/part07/chapter_C.html
    obj.put(DataElement::new(
        tags::STATUS,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));
    // SOPInstanceUID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_INSTANCE_UID,
        VR::UI,
        dicom_value!(Str, sop_instance_uid),
    ));

    obj
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if args.verbose {
                Level::DEBUG
            } else {
                Level::INFO
            })
            .finish(),
    )
    .unwrap_or_else(|e| {
        eprintln!("Could not set up global logger: {}", e);
    });

    std::fs::create_dir_all(&args.out_dir).unwrap_or_else(|e| {
        error!("Could not create output directory: {}", e);
        std::process::exit(-2);
    });

    let listen_addr = SocketAddrV4::new(Ipv4Addr::from(0), args.port);
    let listener = TcpListener::bind(&listen_addr)?;
    info!(
        "{} listening on: tcp://{}",
        &args.calling_ae_title, listen_addr
    );

    for stream in listener.incoming() {
        match stream {
            Ok(scu_stream) => {
                if let Err(e) = run(scu_stream, &args) {
                    error!("{}", e);
                }
            }
            Err(e) => {
                error!("{}", e);
            }
        }
    }

    Ok(())
}
<commit_msg>[storescp] Increase resilience<commit_after>use std::{
    net::{Ipv4Addr, SocketAddrV4, TcpListener, TcpStream},
    path::PathBuf,
};

use clap::Parser;
use dicom_core::{dicom_value, DataElement, PrimitiveValue, VR};
use dicom_dictionary_std::tags;
use dicom_encoding::transfer_syntax::TransferSyntaxIndex;
use dicom_object::{FileMetaTableBuilder, InMemDicomObject, StandardDataDictionary};
use dicom_transfer_syntax_registry::TransferSyntaxRegistry;
use dicom_ul::{pdu::PDataValueType, Pdu};
use tracing::{debug, error, info, warn, Level};

use crate::transfer::ABSTRACT_SYNTAXES;

mod transfer;

/// DICOM C-STORE SCP
#[derive(Debug, Parser)]
struct App {
    /// verbose mode
    #[clap(short = 'v', long = "verbose")]
    verbose: bool,
    /// the calling Application Entity title
    #[structopt(long = "calling-ae-title", default_value = "STORE-SCP")]
    calling_ae_title: String,
    /// enforce max pdu length
    #[clap(short = 's', long = "strict")]
    strict: bool,
    /// Only accept native/uncompressed transfer syntaxes
    #[clap(long)]
    uncompressed_only: bool,
    /// max pdu length
    #[clap(short = 'm', long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,
    /// output directory for incoming objects
    #[clap(short = 'o', default_value = ".")]
    out_dir: PathBuf,
    /// Which port to listen on
    #[clap(short, default_value = "11111")]
    port: u16,
}

fn run(scu_stream: TcpStream, args: &App) -> Result<(), Box<dyn std::error::Error>> {
    let App {
        verbose,
        calling_ae_title,
        strict,
        uncompressed_only,
        max_pdu_length,
        out_dir,
        port: _,
    } = args;
    let verbose = *verbose;

    let mut buffer: Vec<u8> = Vec::with_capacity(*max_pdu_length as usize);
    let mut instance_buffer: Vec<u8> = Vec::with_capacity(1024 * 1024);
    let mut msgid = 1;
    let mut sop_class_uid = "".to_string();
    let mut sop_instance_uid = "".to_string();

    let mut options = dicom_ul::association::ServerAssociationOptions::new()
        .accept_any()
        .ae_title(calling_ae_title)
        .strict(*strict);

    if *uncompressed_only {
        options = options
            .with_transfer_syntax("1.2.840.10008.1.2")
            .with_transfer_syntax("1.2.840.10008.1.2.1");
    } else {
        for ts in TransferSyntaxRegistry.iter() {
            if !ts.unsupported() {
                options = options.with_transfer_syntax(ts.uid());
            }
        }
    };

    for uid in ABSTRACT_SYNTAXES {
        options = options.with_abstract_syntax(*uid);
    }

    let mut association = options.establish(scu_stream)?;

    info!("New association from {}", association.client_ae_title());
    debug!(
        "> Presentation contexts: {:?}",
        association.presentation_contexts()
    );

    loop {
        match association.receive() {
            Ok(mut pdu) => {
                if verbose {
                    debug!("scu ----> scp: {}", pdu.short_description());
                }
                match pdu {
                    Pdu::PData { ref mut data } => {
                        if data.is_empty() {
                            debug!("Ignoring empty PData PDU");
                            continue;
                        }

                        if data[0].value_type == PDataValueType::Data && !data[0].is_last {
                            instance_buffer.append(&mut data[0].data);
                        } else if data[0].value_type == PDataValueType::Command && data[0].is_last {
                            // commands are always in implict VR LE
                            let ts =
                                dicom_transfer_syntax_registry::entries::IMPLICIT_VR_LITTLE_ENDIAN
                                    .erased();
                            let data_value = &data[0];
                            let v = &data_value.data;

                            let obj = InMemDicomObject::read_dataset_with_ts(v.as_slice(), &ts)?;
                            msgid = obj.element(tags::MESSAGE_ID)?.to_int()?;
                            sop_class_uid = obj
                                .element(tags::AFFECTED_SOP_CLASS_UID)?
                                .to_str()?
                                .to_string();
                            sop_instance_uid = obj
                                .element(tags::AFFECTED_SOP_INSTANCE_UID)?
                                .to_str()?
                                .to_string();
                            instance_buffer.clear();
                        } else if data[0].value_type == PDataValueType::Data && data[0].is_last {
                            instance_buffer.append(&mut data[0].data);

                            let presentation_context = association
                                .presentation_contexts()
                                .iter()
                                .find(|pc| pc.id == data[0].presentation_context_id)
                                .ok_or("missing presentation context")?;
                            let ts = &presentation_context.transfer_syntax;

                            let obj = InMemDicomObject::read_dataset_with_ts(
                                instance_buffer.as_slice(),
                                TransferSyntaxRegistry.get(ts).unwrap(),
                            )?;
                            let file_meta = FileMetaTableBuilder::new()
                                .media_storage_sop_class_uid(
                                    obj.element_by_name("SOPClassUID")?.to_str()?,
                                )
                                .media_storage_sop_instance_uid(
                                    obj.element_by_name("SOPInstanceUID")?.to_str()?,
                                )
                                .transfer_syntax(ts)
                                .build()?;
                            let file_obj = obj.with_exact_meta(file_meta);

                            // write the files to the current directory with their SOPInstanceUID as filenames
                            let mut file_path = out_dir.clone();
                            file_path
                                .push(sop_instance_uid.trim_end_matches('\0').to_string() + ".dcm");
                            file_obj.write_to_file(&file_path)?;
                            info!("Stored {}", file_path.display());

                            // send C-STORE-RSP object
                            // commands are always in implict VR LE
                            let ts =
                                dicom_transfer_syntax_registry::entries::IMPLICIT_VR_LITTLE_ENDIAN
                                    .erased();

                            let obj =
                                create_cstore_response(msgid, &sop_class_uid, &sop_instance_uid);

                            let mut obj_data = Vec::new();

                            obj.write_dataset_with_ts(&mut obj_data, &ts)?;

                            let pdu_response = Pdu::PData {
                                data: vec![dicom_ul::pdu::PDataValue {
                                    presentation_context_id: data[0].presentation_context_id,
                                    value_type: PDataValueType::Command,
                                    is_last: true,
                                    data: obj_data,
                                }],
                            };
                            association.send(&pdu_response)?;
                        }
                    }
                    Pdu::ReleaseRQ => {
                        buffer.clear();
                        association.send(&Pdu::ReleaseRP)?;
                        info!(
                            "Released association with {}",
                            association.client_ae_title()
                        );
                    }
                    _ => {}
                }
            }
            Err(err @ dicom_ul::association::server::Error::Receive { .. }) => {
                debug!("{}", err);
                break;
            }
            Err(err) => {
                warn!("Unexpected error: {}", err);
                break;
            }
        }
    }
    info!("Association with {} dropped", association.client_ae_title());
    Ok(())
}

fn create_cstore_response(
    message_id: u16,
    sop_class_uid: &str,
    sop_instance_uid: &str,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::new_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            8 + sop_class_uid.len() as i32
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + 8
                + 2
                + sop_instance_uid.len() as i32,
        ),
    ));

    // service
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, sop_class_uid),
    ));
    // command
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x8001]),
    ));
    // message ID being responded to
    obj.put(DataElement::new(
        tags::MESSAGE_ID_BEING_RESPONDED_TO,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));
    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0101]),
    ));
    // status https://dicom.nema.org/dicom/2013/output/chtml/part07/chapter_C.html
    obj.put(DataElement::new(
        tags::STATUS,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));
    // SOPInstanceUID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_INSTANCE_UID,
        VR::UI,
        dicom_value!(Str, sop_instance_uid),
    ));

    obj
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = App::from_args();

    tracing::subscriber::set_global_default(
        tracing_subscriber::FmtSubscriber::builder()
            .with_max_level(if args.verbose {
                Level::DEBUG
            } else {
                Level::INFO
            })
            .finish(),
    )
    .unwrap_or_else(|e| {
        eprintln!("Could not set up global logger: {}", e);
    });

    std::fs::create_dir_all(&args.out_dir).unwrap_or_else(|e| {
        error!("Could not create output directory: {}", e);
        std::process::exit(-2);
    });

    let listen_addr = SocketAddrV4::new(Ipv4Addr::from(0), args.port);
    let listener = TcpListener::bind(&listen_addr)?;
    info!(
        "{} listening on: tcp://{}",
        &args.calling_ae_title, listen_addr
    );

    for stream in listener.incoming() {
        match stream {
            Ok(scu_stream) => {
                if let Err(e) = run(scu_stream, &args) {
                    error!("{}", e);
                }
            }
            Err(e) => {
                error!("{}", e);
            }
        }
    }

    Ok(())
}
<|endoftext|>


<commit_before>use dicom::core::smallvec;
use dicom::{core::dicom_value, dictionary_std::tags};
use dicom::{
    core::{DataElement, PrimitiveValue, VR},
    encoding::transfer_syntax,
    object::{mem::InMemDicomObject, open_file, StandardDataDictionary},
    transfer_syntax::TransferSyntaxRegistry,
};
use dicom_ul::pdu::Pdu;
use dicom_ul::{
    association::ClientAssociationOptions,
    pdu::{PDataValue, PDataValueType},
};
use smallvec::smallvec;
use std::{io::Write, path::PathBuf};
use structopt::StructOpt;
use transfer_syntax::TransferSyntaxIndex;

/// DICOM C-STORE SCU
#[derive(Debug, StructOpt)]
struct App {
    /// socket address to STORE SCP (example: "127.0.0.1:104")
    addr: String,
    /// the DICOM file to store
    file: PathBuf,
    /// verbose mode
    #[structopt(short = "v", long = "verbose")]
    verbose: bool,
    /// the C-STORE message ID
    #[structopt(short = "m", long = "message-id", default_value = "1")]
    message_id: u16,
    /// the calling AE title
    #[structopt(long = "calling-ae-title", default_value = "STORE-SCU")]
    calling_ae_title: String,
    /// the called AE title
    #[structopt(long = "called-ae-title", default_value = "ANY-SCP")]
    called_ae_title: String,
    /// the maximum PDU length
    #[structopt(long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let App {
        addr,
        file,
        verbose,
        message_id,
        calling_ae_title,
        called_ae_title,
        max_pdu_length,
    } = App::from_args();

    if verbose {
        println!("Opening file '{}'...", file.display());
    }

    let dicom_file = open_file(file)?;

    let meta = dicom_file.meta();

    let storage_sop_class_uid = &meta.media_storage_sop_class_uid;
    let storage_sop_instance_uid = &meta.media_storage_sop_instance_uid;
    let transfer_syntax = &meta.transfer_syntax;

    if verbose {
        println!("Establishing association with '{}'...", &addr);
    }

    let mut scu = ClientAssociationOptions::new()
        .with_abstract_syntax(storage_sop_class_uid.to_string())
        .with_transfer_syntax(transfer_syntax.to_string())
        .calling_ae_title(calling_ae_title)
        .called_ae_title(called_ae_title)
        .max_pdu_length(max_pdu_length)
        .establish(addr)?;

    if verbose {
        println!("Association established");
    }
    let ts = TransferSyntaxRegistry
        .get(&transfer_syntax)
        .expect("Poorly negotiated transfer syntax");

    if verbose {
        println!("Transfer Syntax: {}", ts.name());
    }

    let cmd = store_req_command(
        &storage_sop_class_uid,
        &storage_sop_instance_uid,
        message_id,
    );

    let mut cmd_data = Vec::with_capacity(128);
    cmd.write_dataset_with_ts(
        &mut cmd_data,
        &dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
    )?;

    // !!! ensure that the negotiated transfer syntax is used

    let mut object_data = Vec::with_capacity(1024);
    dicom_file.write_dataset_with_ts(&mut object_data, &ts)?;

    let nbytes = cmd_data.len() + object_data.len();

    if verbose {
        println!("Sending payload (~ {} Kb)...", nbytes / 1024);
    }

    let pc_selected = scu
        .presentation_contexts()
        .first()
        .ok_or("No presentation context accepted")?
        .clone();

    if nbytes < max_pdu_length as usize - 100 {
        let pdu = Pdu::PData {
            data: vec![
                PDataValue {
                    presentation_context_id: pc_selected.id,
                    value_type: PDataValueType::Command,
                    is_last: true,
                    data: cmd_data,
                },
                PDataValue {
                    presentation_context_id: pc_selected.id,
                    value_type: PDataValueType::Data,
                    is_last: true,
                    data: object_data,
                },
            ],
        };

        scu.send(&pdu)?;
    } else {
        let pdu = Pdu::PData {
            data: vec![PDataValue {
                presentation_context_id: pc_selected.id,
                value_type: PDataValueType::Command,
                is_last: true,
                data: cmd_data,
            }],
        };

        scu.send(&pdu)?;

        scu.send_pdata(pc_selected.id).write_all(&object_data)?;
    }

    if verbose {
        println!("Awaiting response...");
    }

    let rsp_pdu = scu.receive()?;

    match rsp_pdu {
        Pdu::PData { data } => {
            let data_value = &data[0];

            let cmd_obj = InMemDicomObject::read_dataset_with_ts(
                &data_value.data[..],
                &dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
            )?;
            if verbose {
                println!("Response: {:?}", cmd_obj);
            }
            let status = cmd_obj.element(tags::STATUS)?.to_int::<u16>()?;
            if status == 0 {
                println!("Sucessfully stored instance `{}`", storage_sop_instance_uid);
            } else {
                println!(
                    "Failed to store instance '{}' (status code {})",
                    storage_sop_instance_uid, status
                );
            }

            scu.release()?;
        }

        pdu @ Pdu::Unknown { .. }
        | pdu @ Pdu::AssociationRQ { .. }
        | pdu @ Pdu::AssociationAC { .. }
        | pdu @ Pdu::AssociationRJ { .. }
        | pdu @ Pdu::ReleaseRQ
        | pdu @ Pdu::ReleaseRP
        | pdu @ Pdu::AbortRQ { .. } => {
            eprintln!("Unexpected SCP response: {:?}", pdu);
            std::process::exit(-2);
        }
    }

    Ok(())
}

fn store_req_command(
    storage_sop_class_uid: &str,
    storage_sop_instance_uid: &str,
    message_id: u16,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::create_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(0),
    ));

    // SOP Class UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, storage_sop_class_uid),
    ));

    // command field
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x0001]),
    ));

    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));

    //priority
    obj.put(DataElement::new(
        tags::PRIORITY,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));

    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));

    // affected SOP Instance UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_INSTANCE_UID,
        VR::UI,
        dicom_value!(Str, storage_sop_instance_uid),
    ));

    obj
}
<commit_msg>[storescu] carefully examine transfer syntaxes<commit_after>use dicom::core::smallvec;
use dicom::{core::dicom_value, dictionary_std::tags};
use dicom::{
    core::{DataElement, PrimitiveValue, VR},
    encoding::transfer_syntax,
    object::{mem::InMemDicomObject, open_file, StandardDataDictionary},
    transfer_syntax::TransferSyntaxRegistry,
};
use dicom_ul::pdu::Pdu;
use dicom_ul::{
    association::ClientAssociationOptions,
    pdu::{PDataValue, PDataValueType},
};
use smallvec::smallvec;
use std::{io::Write, path::PathBuf};
use structopt::StructOpt;
use transfer_syntax::TransferSyntaxIndex;

/// DICOM C-STORE SCU
#[derive(Debug, StructOpt)]
struct App {
    /// socket address to STORE SCP (example: "127.0.0.1:104")
    addr: String,
    /// the DICOM file to store
    file: PathBuf,
    /// verbose mode
    #[structopt(short = "v", long = "verbose")]
    verbose: bool,
    /// the C-STORE message ID
    #[structopt(short = "m", long = "message-id", default_value = "1")]
    message_id: u16,
    /// the calling AE title
    #[structopt(long = "calling-ae-title", default_value = "STORE-SCU")]
    calling_ae_title: String,
    /// the called AE title
    #[structopt(long = "called-ae-title", default_value = "ANY-SCP")]
    called_ae_title: String,
    /// the maximum PDU length
    #[structopt(long = "max-pdu-length", default_value = "16384")]
    max_pdu_length: u32,
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let App {
        addr,
        file,
        verbose,
        message_id,
        calling_ae_title,
        called_ae_title,
        max_pdu_length,
    } = App::from_args();

    if verbose {
        println!("Opening file '{}'...", file.display());
    }

    let dicom_file = open_file(file)?;

    let meta = dicom_file.meta();

    let storage_sop_class_uid = &meta.media_storage_sop_class_uid;
    let storage_sop_instance_uid = &meta.media_storage_sop_instance_uid;
    let transfer_syntax = &meta.transfer_syntax.trim_end_matches("\0");

    let file_ts_desc = TransferSyntaxRegistry
        .get(transfer_syntax)
        .ok_or("Unsupported file transfer syntax")?;

    if verbose {
        println!("Establishing association with '{}'...", &addr);
    }

    let mut scu = ClientAssociationOptions::new()
        .with_abstract_syntax(storage_sop_class_uid.to_string())
        .with_transfer_syntax(transfer_syntax.to_string())
        .calling_ae_title(calling_ae_title)
        .called_ae_title(called_ae_title)
        .max_pdu_length(max_pdu_length)
        .establish(addr)?;

    if verbose {
        println!("Association established");
    }

    // TODO(#106) transfer syntax conversion is currently not supported
    let pc_selected = if let Some(pc) = scu.presentation_contexts().iter().find(|pc| {
        // Check support for this transfer syntax.
        // If it is the same as the file, we're good.
        // Otherwise, uncompressed data set encoding
        // and native pixel data is required on both ends.
        let ts = dbg!(&pc.transfer_syntax);
        ts == transfer_syntax
            || TransferSyntaxRegistry
                .get(&pc.transfer_syntax)
                .filter(|ts| file_ts_desc.is_codec_free() && ts.is_codec_free())
                .map(|_| true)
                .unwrap_or(false)
    }) {
        pc.clone()
    } else {
        eprintln!("Could not choose a transfer syntax");
        let _ = scu.abort();
        std::process::exit(-2);
    };

    let ts = TransferSyntaxRegistry
        .get(&pc_selected.transfer_syntax)
        .ok_or_else(|| "Poorly negotiated transfer syntax")?;

    if verbose {
        println!("Transfer Syntax: {}", ts.name());
    }

    let cmd = store_req_command(
        &storage_sop_class_uid,
        &storage_sop_instance_uid,
        message_id,
    );

    let mut cmd_data = Vec::with_capacity(128);
    cmd.write_dataset_with_ts(
        &mut cmd_data,
        &dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
    )?;

    let mut object_data = Vec::with_capacity(1024);
    dicom_file.write_dataset_with_ts(&mut object_data, &ts)?;

    let nbytes = cmd_data.len() + object_data.len();

    if verbose {
        println!("Sending payload (~ {} Kb)...", nbytes / 1024);
    }

    if nbytes < max_pdu_length as usize - 100 {
        let pdu = Pdu::PData {
            data: vec![
                PDataValue {
                    presentation_context_id: pc_selected.id,
                    value_type: PDataValueType::Command,
                    is_last: true,
                    data: cmd_data,
                },
                PDataValue {
                    presentation_context_id: pc_selected.id,
                    value_type: PDataValueType::Data,
                    is_last: true,
                    data: object_data,
                },
            ],
        };

        scu.send(&pdu)?;
    } else {
        let pdu = Pdu::PData {
            data: vec![PDataValue {
                presentation_context_id: pc_selected.id,
                value_type: PDataValueType::Command,
                is_last: true,
                data: cmd_data,
            }],
        };

        scu.send(&pdu)?;

        scu.send_pdata(pc_selected.id).write_all(&object_data)?;
    }

    if verbose {
        println!("Awaiting response...");
    }

    let rsp_pdu = scu.receive()?;

    match rsp_pdu {
        Pdu::PData { data } => {
            let data_value = &data[0];

            let cmd_obj = InMemDicomObject::read_dataset_with_ts(
                &data_value.data[..],
                &dicom::transfer_syntax::entries::IMPLICIT_VR_LITTLE_ENDIAN.erased(),
            )?;
            if verbose {
                println!("Response: {:?}", cmd_obj);
            }
            let status = cmd_obj.element(tags::STATUS)?.to_int::<u16>()?;
            if status == 0 {
                println!("Sucessfully stored instance `{}`", storage_sop_instance_uid);
            } else {
                println!(
                    "Failed to store instance '{}' (status code {})",
                    storage_sop_instance_uid, status
                );
            }

            scu.release()?;
        }

        pdu @ Pdu::Unknown { .. }
        | pdu @ Pdu::AssociationRQ { .. }
        | pdu @ Pdu::AssociationAC { .. }
        | pdu @ Pdu::AssociationRJ { .. }
        | pdu @ Pdu::ReleaseRQ
        | pdu @ Pdu::ReleaseRP
        | pdu @ Pdu::AbortRQ { .. } => {
            eprintln!("Unexpected SCP response: {:?}", pdu);
            let _ = scu.abort();
            std::process::exit(-2);
        }
    }

    Ok(())
}

fn store_req_command(
    storage_sop_class_uid: &str,
    storage_sop_instance_uid: &str,
    message_id: u16,
) -> InMemDicomObject<StandardDataDictionary> {
    let mut obj = InMemDicomObject::create_empty();

    // group length
    obj.put(DataElement::new(
        tags::COMMAND_GROUP_LENGTH,
        VR::UL,
        PrimitiveValue::from(
            12 + 8
                + even_len(storage_sop_class_uid.len())
                + 10
                + 10
                + 10
                + 10
                + 8
                + even_len(storage_sop_instance_uid.len()),
        ),
    ));

    // SOP Class UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_CLASS_UID,
        VR::UI,
        dicom_value!(Str, storage_sop_class_uid),
    ));

    // command field
    obj.put(DataElement::new(
        tags::COMMAND_FIELD,
        VR::US,
        dicom_value!(U16, [0x0001]),
    ));

    // message ID
    obj.put(DataElement::new(
        tags::MESSAGE_ID,
        VR::US,
        dicom_value!(U16, [message_id]),
    ));

    //priority
    obj.put(DataElement::new(
        tags::PRIORITY,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));

    // data set type
    obj.put(DataElement::new(
        tags::COMMAND_DATA_SET_TYPE,
        VR::US,
        dicom_value!(U16, [0x0000]),
    ));

    // affected SOP Instance UID
    obj.put(DataElement::new(
        tags::AFFECTED_SOP_INSTANCE_UID,
        VR::UI,
        dicom_value!(Str, storage_sop_instance_uid),
    ));

    obj
}

fn even_len(l: usize) -> u32 {
    ((l + 1) & !1) as u32
}

#[cfg(test)]
mod tests {
    use super::even_len;
    #[test]
    fn test_even_len() {
        assert_eq!(even_len(0), 0);
        assert_eq!(even_len(1), 2);
        assert_eq!(even_len(2), 2);
        assert_eq!(even_len(3), 4);
        assert_eq!(even_len(4), 4);
        assert_eq!(even_len(5), 6);
    }
}
<|endoftext|>


<commit_before>#![deny(missing_docs)]
#![warn(rust_2018_idioms)]
#![doc(html_root_url = "https://docs.rs/amq-protocol-tcp/4.2.0/")]

//! # AMQP URI TCP connection handling
//!
//! amq-protocol-tcp is a library aiming at providing tools to help
//! connecting to an AMQP URI

use amq_protocol_uri::{AMQPScheme, AMQPUri};
use mio::{Events, Poll, PollOpt, Ready, Token};
use tcp_stream::HandshakeError;

use std::io;

/// Re-export TcpStream
pub use tcp_stream::{Identity, TcpStream};

/// Trait providing a method to connect to a TcpStream
pub trait AMQPUriTcpExt {
    /// connect to a TcpStream
    fn connect<S, F: FnOnce(TcpStream, AMQPUri, Option<(Poll, Token)>) -> S>(
        self,
        f: F,
    ) -> io::Result<S>
    where
        Self: Sized,
    {
        self.connect_full(f, None, None)
    }
    /// connect to a TcpStream, registering it to the given Poll with the given Token to handle the
    /// handshake process. You should reregister it afterwards to better fit your needs
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri, Option<(Poll, Token)>) -> S>(
        self,
        f: F,
        poll: Option<(Poll, Token)>,
        identity: Option<Identity<'_, '_>>,
    ) -> io::Result<S>;
}

impl AMQPUriTcpExt for AMQPUri {
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri, Option<(Poll, Token)>) -> S>(
        self,
        f: F,
        poll: Option<(Poll, Token)>,
        identity: Option<Identity<'_, '_>>,
    ) -> io::Result<S> {
        let stream =
            TcpStream::connect(format!("{}:{}", self.authority.host, self.authority.port))?;

        if let Some((poll, token)) = poll.as_ref() {
            poll.register(&stream, *token, Ready::all(), PollOpt::edge())?;
        }

        match self.scheme {
            AMQPScheme::AMQP => Ok((stream, poll)),
            AMQPScheme::AMQPS => connect_amqps(stream, &self.authority.host, poll, identity),
        }
        .map(|(s, poll)| f(s, self, poll))
    }
}

fn connect_amqps(
    stream: TcpStream,
    host: &str,
    poll: Option<(Poll, Token)>,
    identity: Option<Identity<'_, '_>>,
) -> io::Result<(TcpStream, Option<(Poll, Token)>)> {
    let mut events = Events::with_capacity(1024);
    let mut res = stream.into_tls(host, identity);

    while let Err(error) = res {
        if let Some((poll, _)) = poll.as_ref() {
            poll.poll(&mut events, None)?;
        }
        match error {
            HandshakeError::Failure(io_err) => return Err(io_err),
            HandshakeError::WouldBlock(mid) => res = mid.handshake(),
        };
    }

    Ok((res.unwrap(), poll))
}

impl AMQPUriTcpExt for &str {
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri, Option<(Poll, Token)>) -> S>(
        self,
        f: F,
        poll: Option<(Poll, Token)>,
        identity: Option<Identity<'_, '_>>,
    ) -> io::Result<S> {
        self.parse::<AMQPUri>()
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?
            .connect_full(f, poll, identity)
    }
}
<commit_msg>tcp: fix deadlock on tls error<commit_after>#![deny(missing_docs)]
#![warn(rust_2018_idioms)]
#![doc(html_root_url = "https://docs.rs/amq-protocol-tcp/4.2.0/")]

//! # AMQP URI TCP connection handling
//!
//! amq-protocol-tcp is a library aiming at providing tools to help
//! connecting to an AMQP URI

use amq_protocol_uri::{AMQPScheme, AMQPUri};
use mio::{Events, Poll, PollOpt, Ready, Token};
use tcp_stream::HandshakeError;

use std::io;

/// Re-export TcpStream
pub use tcp_stream::{Identity, TcpStream};

/// Trait providing a method to connect to a TcpStream
pub trait AMQPUriTcpExt {
    /// connect to a TcpStream
    fn connect<S, F: FnOnce(TcpStream, AMQPUri, Option<(Poll, Token)>) -> S>(
        self,
        f: F,
    ) -> io::Result<S>
    where
        Self: Sized,
    {
        self.connect_full(f, None, None)
    }
    /// connect to a TcpStream, registering it to the given Poll with the given Token to handle the
    /// handshake process. You should reregister it afterwards to better fit your needs
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri, Option<(Poll, Token)>) -> S>(
        self,
        f: F,
        poll: Option<(Poll, Token)>,
        identity: Option<Identity<'_, '_>>,
    ) -> io::Result<S>;
}

impl AMQPUriTcpExt for AMQPUri {
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri, Option<(Poll, Token)>) -> S>(
        self,
        f: F,
        poll: Option<(Poll, Token)>,
        identity: Option<Identity<'_, '_>>,
    ) -> io::Result<S> {
        let stream =
            TcpStream::connect(format!("{}:{}", self.authority.host, self.authority.port))?;

        if let Some((poll, token)) = poll.as_ref() {
            poll.register(&stream, *token, Ready::all(), PollOpt::edge())?;
        }

        match self.scheme {
            AMQPScheme::AMQP => Ok((stream, poll)),
            AMQPScheme::AMQPS => connect_amqps(stream, &self.authority.host, poll, identity),
        }
        .map(|(s, poll)| f(s, self, poll))
    }
}

fn connect_amqps(
    stream: TcpStream,
    host: &str,
    poll: Option<(Poll, Token)>,
    identity: Option<Identity<'_, '_>>,
) -> io::Result<(TcpStream, Option<(Poll, Token)>)> {
    let mut events = Events::with_capacity(1024);
    let mut res = stream.into_tls(host, identity);

    while let Err(error) = res {
        match error {
            HandshakeError::Failure(io_err) => return Err(io_err),
            HandshakeError::WouldBlock(mid) => {
                if let Some((poll, _)) = poll.as_ref() {
                    poll.poll(&mut events, None)?;
                }
                res = mid.handshake()
            }
        };
    }

    Ok((res.unwrap(), poll))
}

impl AMQPUriTcpExt for &str {
    fn connect_full<S, F: FnOnce(TcpStream, AMQPUri, Option<(Poll, Token)>) -> S>(
        self,
        f: F,
        poll: Option<(Poll, Token)>,
        identity: Option<Identity<'_, '_>>,
    ) -> io::Result<S> {
        self.parse::<AMQPUri>()
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?
            .connect_full(f, poll, identity)
    }
}
<|endoftext|>


<commit_before>use json_patch::merge;
use serde_json::Value;

/// The MMDS is the Microvm Metadata Service represented as an untyped json.
#[derive(Clone)]
pub struct MMDS {
    data_store: Value,
    is_initialized: bool,
}

impl Default for MMDS {
    fn default() -> Self {
        MMDS {
            data_store: Value::default(),
            is_initialized: false,
        }
    }
}

impl MMDS {
    /// This method is needed to provide the correct status code for API request.
    /// When the MMDS structure is initialized for the first time via the API, the
    /// status code should be 201 (Created) and when the structure is updated, the
    /// status code should be 204 (Updated).
    pub fn is_initialized(&self) -> bool {
        return self.is_initialized;
    }

    pub fn put_data(&mut self, data: Value) {
        self.data_store = data;
        self.is_initialized = true;
    }

    pub fn patch_data(&mut self, patch_data: Value) {
        merge(&mut self.data_store, &patch_data);
    }

    pub fn get_data_str(&self) -> String {
        if self.data_store.is_null() {
            return String::from("{}");
        }
        return self.data_store.to_string();
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use serde_json;

    #[test]
    fn test_mmds() {
        let mut mmds = MMDS::default();
        assert_eq!(mmds.is_initialized(), false);

        let mut mmds_json = "{\"meta-data\":{\"iam\":\"dummy\"},\"user-data\":\"1522850095\"}";

        mmds.put_data(serde_json::from_str(mmds_json).unwrap());
        assert_eq!(mmds.is_initialized(), true);

        assert_eq!(mmds.get_data_str(), mmds_json);

        // update the user-data field add test that patch works as expected
        let patch_json = "{\"user-data\":\"10\"}";
        mmds.patch_data(serde_json::from_str(patch_json).unwrap());
        mmds_json = "{\"meta-data\":{\"iam\":\"dummy\"},\"user-data\":\"10\"}";
        assert_eq!(mmds.get_data_str(), mmds_json);
    }
}
<commit_msg>MMDS: get data by path<commit_after>use json_patch::merge;
use serde_json::Value;

/// The MMDS is the Microvm Metadata Service represented as an untyped json.
#[derive(Clone)]
pub struct MMDS {
    data_store: Value,
    is_initialized: bool,
}

#[derive(Debug, PartialEq)]
pub enum Error {
    NotFound,
    UnsupportedValueType,
}

impl Default for MMDS {
    fn default() -> Self {
        MMDS {
            data_store: Value::default(),
            is_initialized: false,
        }
    }
}

impl MMDS {
    /// This method is needed to provide the correct status code for API request.
    /// When the MMDS structure is initialized for the first time via the API, the
    /// status code should be 201 (Created) and when the structure is updated, the
    /// status code should be 204 (Updated).
    pub fn is_initialized(&self) -> bool {
        return self.is_initialized;
    }

    pub fn put_data(&mut self, data: Value) {
        // TODO: we should add a data validator and only accept Strings, arrays & dictionaries
        // https://github.com/aws/PRIVATE-firecracker/issues/401
        self.data_store = data;
        self.is_initialized = true;
    }

    pub fn patch_data(&mut self, patch_data: Value) {
        merge(&mut self.data_store, &patch_data);
    }

    pub fn get_data_str(&self) -> String {
        if self.data_store.is_null() {
            return String::from("{}");
        }
        return self.data_store.to_string();
    }

    // Helper function for getting all the keys from Value::Object.
    // Returns a Vec<String> with all keys. If the key corresponds to a
    // dictionary, a "/" is appended to the key name.
    // If the `dict` is Value::Null, Error::NotFound is thrown.
    // If the `dict` is not a dictionary, a Vec with the value corresponding to
    // the key is returned.
    fn get_keys(dict: &Value) -> Result<Vec<String>, Error> {
        if dict.is_null() {
            return Err(Error::NotFound);
        }

        let mut ret = Vec::new();
        match dict.as_object() {
            Some(map) => {
                // When the object is a map, push all the keys in the Vec.
                for key in map.keys() {
                    let mut key = key.clone();
                    if dict[&key].is_object() {
                        key.push_str("/");
                    }

                    ret.push(key);
                }
                return Ok(ret);
            }
            None => {
                // When the object is not a map, return the value.
                match dict.as_str() {
                    Some(val) => {
                        ret.push(val.to_string());
                        return Ok(ret);
                    }
                    None => return Err(Error::UnsupportedValueType),
                };
            }
        };
    }

    // Helper function for converting a Value to String by following the IMDS specs.
    // The only supported Value is String. Throws UnsupportedValueType when the type of the
    // Value is not String.
    fn get_value_as_string(val: &Value) -> Result<String, Error> {
        if val.is_null() {
            return Err(Error::NotFound);
        }

        if val.is_object() {
            return Ok(String::new());
        }

        match val.as_str() {
            Some(value) => Ok(value.to_string()),
            None => Err(Error::UnsupportedValueType),
        }
    }

    /// This function replicates the behavior of the Instance Metadata Service
    /// https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-metadata.html
    /// When the path ends with / there are two cases:
    /// 1. For a (key, value) pair where the value is a dictionary, it will return all the keys
    /// in the dictionary.
    /// 2. For a (key, value) pair where the value is a simple type (bool, string, number),
    /// it will return the value.
    ///
    /// When the path does not end with / there are also two cases to cover:
    /// 1. The value corresponding to that path is a dictionary, the function returns
    /// an empty string.
    /// 2. The value corresponding to that path is a simple type, the function returns the value.
    ///
    /// When the path is not found, a NotFound error is returned.
    pub fn get_value(&self, path: String) -> Result<Vec<String>, Error> {
        // The pointer function splits the input by "/". With a trailing "/", pointer does not
        // know how to get the object.
        let value = match path.ends_with('/') {
            true => self.data_store.pointer(&path.as_str()[..(path.len() - 1)]),
            false => self.data_store.pointer(path.as_str()),
        };

        match value {
            Some(val) => {
                // If path ends with /, return all keys in Value::Object
                if path.ends_with("/") {
                    MMDS::get_keys(val)
                } else {
                    match MMDS::get_value_as_string(val) {
                        Ok(value) => Ok(vec![value]),
                        Err(e) => Err(e),
                    }
                }
            }
            None => return Err(Error::NotFound),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json;

    #[test]
    fn test_mmds() {
        let mut mmds = MMDS::default();
        assert_eq!(mmds.is_initialized(), false);

        let mut mmds_json = "{\"meta-data\":{\"iam\":\"dummy\"},\"user-data\":\"1522850095\"}";

        mmds.put_data(serde_json::from_str(mmds_json).unwrap());
        assert_eq!(mmds.is_initialized(), true);

        assert_eq!(mmds.get_data_str(), mmds_json);

        // update the user-data field add test that patch works as expected
        let patch_json = "{\"user-data\":\"10\"}";
        mmds.patch_data(serde_json::from_str(patch_json).unwrap());
        mmds_json = "{\"meta-data\":{\"iam\":\"dummy\"},\"user-data\":\"10\"}";
        assert_eq!(mmds.get_data_str(), mmds_json);
    }

    #[test]
    fn test_get_value() {
        let mut mmds = MMDS::default();
        let data = r#"{
            "name": {
                "first": "John",
                "second": "Doe"
            },
            "age": "43",
            "phones": {
                "home": {
                    "RO": "+40 1234567",
                    "UK": "+44 1234567"
                },
                "mobile": "+44 2345678"
            }
        }"#;

        let data_store: Value = serde_json::from_str(data).unwrap();
        mmds.put_data(data_store);

        // Test invalid path.
        match mmds.get_value("/invalid_path".to_string()) {
            Ok(_) => assert!(false),
            Err(e) => assert_eq!(e, Error::NotFound),
        };
        match mmds.get_value("/invalid_path/".to_string()) {
            Ok(_) => assert!(false),
            Err(e) => assert_eq!(e, Error::NotFound),
        };

        // Test path ends with /; Value is a dictionary.
        match mmds.get_value("/phones/".to_string()) {
            Ok(ret) => assert_eq!(ret, vec!["home/", "mobile"]),
            Err(e) => assert!(false),
        };
        match mmds.get_value("/phones/home/".to_string()) {
            Ok(ret) => assert_eq!(ret, vec!["RO", "UK"]),
            Err(_) => assert!(false),
        };

        // Test path ends with /; Value is a String.
        match mmds.get_value("/phones/mobile/".to_string()) {
            Ok(ret) => assert_eq!(ret, vec!["+44 2345678"]),
            Err(_) => assert!(false),
        };

        // Test path does NOT end with /; Value is a dictionary.
        match mmds.get_value("/phones".to_string()) {
            Ok(ret) => assert_eq!(ret, vec![""]),
            Err(_) => assert!(false),
        };

        // Test path does NOT end with /; Value is a String.
        match mmds.get_value("/phones/mobile".to_string()) {
            Ok(ret) => assert_eq!(ret, vec!["+44 2345678"]),
            Err(_) => assert!(false),
        };
    }

    #[test]
    fn test_get_element_from_array() {
        let mut mmds = MMDS::default();
        let data = r#"{
            "phones": [
                "+40 1234567",
                "+44 1234567"
            ]
        }"#;

        let data_store: Value = serde_json::from_str(data).unwrap();
        mmds.put_data(data_store);

        // Test path does NOT end with /; Value is a String.
        match mmds.get_value("/phones/0".to_string()) {
            Ok(ret) => assert_eq!(ret, vec!["+40 1234567"]),
            Err(_) => assert!(false),
        };
    }

    #[test]
    fn test_invalid_types() {
        let mut mmds = MMDS::default();
        let data = r#"{
            "name": {
                "first": "John",
                "second": "Doe"
            },
            "age": 43
        }"#;

        let data_store: Value = serde_json::from_str(data).unwrap();
        // TODO: This should fail; we should only accept String types
        mmds.put_data(data_store);

        match mmds.get_value("/age".to_string()) {
            Ok(_) => assert!(false),
            Err(e) => assert_eq!(e, Error::UnsupportedValueType),
        };

        match mmds.get_value("/age/".to_string()) {
            Ok(_) => assert!(false),
            Err(e) => assert_eq!(e, Error::UnsupportedValueType),
        };
    }
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![license = "MIT/ASL2"]
#![crate_id = "github.com/DaGenix/rust-crypto#rust-crypto-util:0.1"]

extern crate getopts;
extern crate rust_crypto = "rust-crypto";

use std::io;
use std::os;

use getopts::{optopt, optflag, getopts, Matches};

use rust_crypto::scrypt;

fn print_usage() {
    println!("Usage: rust-crypto-util <algorithm> [options]");
    println!("-h, --help\tUsage");
    println!("");
    println!("Algorithms:");
    println!(" * Scrypt (scrypt)");
    println!("");
    println!("Scrypt options:");
    println!("--logn\t\tThe Log N parameter");
    println!("-r\t\tThe R parameter");
    println!("-p\t\tThe P parameter");
    println!("--dklen\t\tThe DkLen parameter");
    println!("--rawsalt\tThe salt parameter is supplied on STDIN");
    println!("--rawpassword\tThe password parameter is supplied on STDIN");
    println!("--rawoutput\tThe resulting value should be output directly to STDOUT");
}

fn run_scrypt(matches: &Matches) {
    if !matches.opt_present("logn") || !matches.opt_present("r") || !matches.opt_present("p") ||
       !matches.opt_present("dklen") {
        print_usage();
        return;
    }
    let logn = from_str::<u8>(matches.opt_str("logn").unwrap().as_slice()).unwrap();
    let r = from_str::<u32>(matches.opt_str("r").unwrap().as_slice()).unwrap();
    let p = from_str::<u32>(matches.opt_str("p").unwrap().as_slice()).unwrap();
    let dklen = from_str::<uint>(matches.opt_str("dklen").unwrap().as_slice()).unwrap();

    if !matches.opt_present("rawsalt") || !matches.opt_present("rawpassword") ||
       !matches.opt_present("rawoutput") {
        println!("Required options missing.");
        return;
    }

    let salt_len = io::stdio::stdin_raw().read_be_u32().unwrap();
    let salt = io::stdio::stdin_raw().read_exact(salt_len as uint).unwrap();
    let pass_len = io::stdio::stdin_raw().read_be_u32().unwrap();
    let pass = io::stdio::stdin_raw().read_exact(pass_len as uint).unwrap();

    let params = scrypt::ScryptParams::new(logn, r, p);
    let mut output = Vec::from_elem(dklen, 0u8);
    scrypt::scrypt(pass.as_slice(), salt.as_slice(), &params, output.as_mut_slice());

    match io::stdout().write(output.as_slice()) {
        Ok(_) => { },
        Err(_) => fail!("Error writing result")
    }
}

fn main() {
    let args: Vec<String> = os::args().iter().map(|x| x.to_string()).collect();

    let opts = vec![
        // General parameters:
        optflag("h", "help", "Print help"),

        // Scrypt parameters:
        optopt("", "logn", "Log-N parameter for Scrypt", ""),
        optopt("r", "", "R parameter for Scrypt", ""),
        optopt("p", "", "P parameter for Scrypt", ""),
        optopt("", "dklen", "Length of the derived key", ""),
        optflag("", "rawsalt", "Use a raw salt value"),
        optflag("", "rawpassword", "Use a raw password value"),
        optflag("", "rawoutput", "Use raw output mode"),
    ];

    let matches = match getopts(args.tail(), opts.as_slice()) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_err_msg()) }
    };

    if matches.opt_present("h") || matches.opt_present("help") {
        print_usage();
        return;
    }

    if matches.free.is_empty() {
        print_usage();
        return;
    }
    let algorithm_name = matches.free.get(0).as_slice();

    match algorithm_name {
        "scrypt" => run_scrypt(&matches),
        _ => {
            print_usage();
            return;
        }
    }
}
<commit_msg>Fix deprecation warning from using .to_err_msg()<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#![license = "MIT/ASL2"]
#![crate_id = "github.com/DaGenix/rust-crypto#rust-crypto-util:0.1"]

extern crate getopts;
extern crate rust_crypto = "rust-crypto";

use std::io;
use std::os;

use getopts::{optopt, optflag, getopts, Matches};

use rust_crypto::scrypt;

fn print_usage() {
    println!("Usage: rust-crypto-util <algorithm> [options]");
    println!("-h, --help\tUsage");
    println!("");
    println!("Algorithms:");
    println!(" * Scrypt (scrypt)");
    println!("");
    println!("Scrypt options:");
    println!("--logn\t\tThe Log N parameter");
    println!("-r\t\tThe R parameter");
    println!("-p\t\tThe P parameter");
    println!("--dklen\t\tThe DkLen parameter");
    println!("--rawsalt\tThe salt parameter is supplied on STDIN");
    println!("--rawpassword\tThe password parameter is supplied on STDIN");
    println!("--rawoutput\tThe resulting value should be output directly to STDOUT");
}

fn run_scrypt(matches: &Matches) {
    if !matches.opt_present("logn") || !matches.opt_present("r") || !matches.opt_present("p") ||
       !matches.opt_present("dklen") {
        print_usage();
        return;
    }
    let logn = from_str::<u8>(matches.opt_str("logn").unwrap().as_slice()).unwrap();
    let r = from_str::<u32>(matches.opt_str("r").unwrap().as_slice()).unwrap();
    let p = from_str::<u32>(matches.opt_str("p").unwrap().as_slice()).unwrap();
    let dklen = from_str::<uint>(matches.opt_str("dklen").unwrap().as_slice()).unwrap();

    if !matches.opt_present("rawsalt") || !matches.opt_present("rawpassword") ||
       !matches.opt_present("rawoutput") {
        println!("Required options missing.");
        return;
    }

    let salt_len = io::stdio::stdin_raw().read_be_u32().unwrap();
    let salt = io::stdio::stdin_raw().read_exact(salt_len as uint).unwrap();
    let pass_len = io::stdio::stdin_raw().read_be_u32().unwrap();
    let pass = io::stdio::stdin_raw().read_exact(pass_len as uint).unwrap();

    let params = scrypt::ScryptParams::new(logn, r, p);
    let mut output = Vec::from_elem(dklen, 0u8);
    scrypt::scrypt(pass.as_slice(), salt.as_slice(), &params, output.as_mut_slice());

    match io::stdout().write(output.as_slice()) {
        Ok(_) => { },
        Err(_) => fail!("Error writing result")
    }
}

fn main() {
    let args: Vec<String> = os::args().iter().map(|x| x.to_string()).collect();

    let opts = vec![
        // General parameters:
        optflag("h", "help", "Print help"),

        // Scrypt parameters:
        optopt("", "logn", "Log-N parameter for Scrypt", ""),
        optopt("r", "", "R parameter for Scrypt", ""),
        optopt("p", "", "P parameter for Scrypt", ""),
        optopt("", "dklen", "Length of the derived key", ""),
        optflag("", "rawsalt", "Use a raw salt value"),
        optflag("", "rawpassword", "Use a raw password value"),
        optflag("", "rawoutput", "Use raw output mode"),
    ];

    let matches = match getopts(args.tail(), opts.as_slice()) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_str()) }
    };

    if matches.opt_present("h") || matches.opt_present("help") {
        print_usage();
        return;
    }

    if matches.free.is_empty() {
        print_usage();
        return;
    }
    let algorithm_name = matches.free.get(0).as_slice();

    match algorithm_name {
        "scrypt" => run_scrypt(&matches),
        _ => {
            print_usage();
            return;
        }
    }
}
<|endoftext|>


<commit_before>// FIXME: This is a huge hack to find the static version of the library
// instead of the shared. It looks in a very specific place that only has
// relevance to servo.

#[cfg(target_os = "linux")]
#[link_args = "../harfbuzz/src/.libs/libharfbuzz.a -lpango-1.0 -lglib-2.0"]
#[no_link]
extern mod m { }

#[cfg(target_os = "macos")]
#[link_args = "../harfbuzz/src/.libs/libharfbuzz.a -lstdc++"]
#[no_link]
extern mod m { }
<commit_msg>Fixed linking problems<commit_after>// FIXME: This is a huge hack to find the static version of the library
// instead of the shared. It looks in a very specific place that only has
// relevance to servo.

#[cfg(target_os = "linux")]
#[link_args = "../harfbuzz/src/.libs/libharfbuzz.a -lpango-1.0 -lglib-2.0"]
#[link_args = "-lstdc++"]
#[no_link]
extern mod m { }

#[cfg(target_os = "macos")]
#[link_args = "../harfbuzz/src/.libs/libharfbuzz.a -lstdc++"]
#[no_link]
extern mod m { }
<|endoftext|>