<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Alex Lyon <arcterus@mail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) RFILE refsize rfilename fsize tsize

#[macro_use]
extern crate uucore;

use clap::{App, Arg};
use std::fs::{metadata, OpenOptions};
use std::io::ErrorKind;
use std::path::Path;

#[derive(Eq, PartialEq)]
enum TruncateMode {
    Absolute,
    Reference,
    Extend,
    Reduce,
    AtMost,
    AtLeast,
    RoundDown,
    RoundUp,
}

static ABOUT: &str = "Shrink or extend the size of each file to the specified size.";
static VERSION: &str = env!("CARGO_PKG_VERSION");

pub mod options {
    pub static IO_BLOCKS: &str = "io-blocks";
    pub static NO_CREATE: &str = "no-create";
    pub static REFERENCE: &str = "reference";
    pub static SIZE: &str = "size";
}

static ARG_FILES: &str = "files";

fn get_usage() -> String {
    format!("{0} [OPTION]... [FILE]...", executable!())
}

fn get_long_usage() -> String {
    String::from(
        "
    SIZE is an integer with an optional prefix and optional unit.
    The available units (K, M, G, T, P, E, Z, and Y) use the following format:
        'KB' =>           1000 (kilobytes)
        'K'  =>           1024 (kibibytes)
        'MB' =>      1000*1000 (megabytes)
        'M'  =>      1024*1024 (mebibytes)
        'GB' => 1000*1000*1000 (gigabytes)
        'G'  => 1024*1024*1024 (gibibytes)
    SIZE may also be prefixed by one of the following to adjust the size of each
    file based on its current size:
        '+'  => extend by
        '-'  => reduce by
        '<'  => at most
        '>'  => at least
        '/'  => round down to multiple of
        '%'  => round up to multiple of",
    )
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();
    let long_usage = get_long_usage();

    let matches = App::new(executable!())
        .version(VERSION)
        .about(ABOUT)
        .usage(&usage[..])
        .after_help(&long_usage[..])
        .arg(
            Arg::with_name(options::IO_BLOCKS)
            .short("o")
            .long(options::IO_BLOCKS)
            .help("treat SIZE as the number of I/O blocks of the file rather than bytes (NOT IMPLEMENTED)")
        )
        .arg(
            Arg::with_name(options::NO_CREATE)
            .short("c")
            .long(options::NO_CREATE)
            .help("do not create files that do not exist")
        )
        .arg(
            Arg::with_name(options::REFERENCE)
            .short("r")
            .long(options::REFERENCE)
            .help("base the size of each file on the size of RFILE")
            .value_name("RFILE")
        )
        .arg(
            Arg::with_name(options::SIZE)
            .short("s")
            .long("size")
            .help("set or adjust the size of each file according to SIZE, which is in bytes unless --io-blocks is specified")
            .value_name("SIZE")
        )
        .arg(Arg::with_name(ARG_FILES).multiple(true).takes_value(true).min_values(1))
        .get_matches_from(args);

    let files: Vec<String> = matches
        .values_of(ARG_FILES)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    if files.is_empty() {
        show_error!("Missing an argument");
        return 1;
    } else {
        let io_blocks = matches.is_present(options::IO_BLOCKS);
        let no_create = matches.is_present(options::NO_CREATE);
        let reference = matches.value_of(options::REFERENCE).map(String::from);
        let size = matches.value_of(options::SIZE).map(String::from);
        if reference.is_none() && size.is_none() {
            crash!(1, "you must specify either --reference or --size");
        } else {
            truncate(no_create, io_blocks, reference, size, files);
        }
    }

    0
}

fn truncate(
    no_create: bool,
    _: bool,
    reference: Option<String>,
    size: Option<String>,
    filenames: Vec<String>,
) {
    let (modsize, mode) = match size {
        Some(size_string) => {
            // Trim any whitespace.
            let size_string = size_string.trim();

            // Get the modifier character from the size string, if any. For
            // example, if the argument is "+123", then the modifier is '+'.
            let c = size_string.chars().next().unwrap();

            let mode = match c {
                '+' => TruncateMode::Extend,
                '-' => TruncateMode::Reduce,
                '<' => TruncateMode::AtMost,
                '>' => TruncateMode::AtLeast,
                '/' => TruncateMode::RoundDown,
                '%' => TruncateMode::RoundUp,
                _ => TruncateMode::Absolute, /* assume that the size is just a number */
            };

            // If there was a modifier character, strip it.
            let size_string = match mode {
                TruncateMode::Absolute => size_string,
                _ => &size_string[1..],
            };
            let num_bytes = match parse_size(size_string) {
                Ok(b) => b,
                Err(_) => crash!(1, "Invalid number: ‘{}’", size_string),
            };
            (num_bytes, mode)
        }
        None => (0, TruncateMode::Reference),
    };

    let refsize = match reference {
        Some(ref rfilename) => {
            match mode {
                // Only Some modes work with a reference
                TruncateMode::Reference => (), //No --size was given
                TruncateMode::Extend => (),
                TruncateMode::Reduce => (),
                _ => crash!(1, "you must specify a relative ‘--size’ with ‘--reference’"),
            };
            match metadata(rfilename) {
                Ok(meta) => meta.len(),
                Err(f) => match f.kind() {
                    ErrorKind::NotFound => {
                        crash!(1, "cannot stat '{}': No such file or directory", rfilename)
                    }
                    _ => crash!(1, "{}", f.to_string()),
                },
            }
        }
        None => 0,
    };
    for filename in &filenames {
        let path = Path::new(filename);
        match OpenOptions::new()
            .read(true)
            .write(true)
            .create(!no_create)
            .open(path)
        {
            Ok(file) => {
                let fsize = match reference {
                    Some(_) => refsize,
                    None => match metadata(filename) {
                        Ok(meta) => meta.len(),
                        Err(f) => {
                            show_warning!("{}", f.to_string());
                            continue;
                        }
                    },
                };
                let tsize: u64 = match mode {
                    TruncateMode::Absolute => modsize,
                    TruncateMode::Reference => fsize,
                    TruncateMode::Extend => fsize + modsize,
                    TruncateMode::Reduce => fsize - modsize,
                    TruncateMode::AtMost => fsize.min(modsize),
                    TruncateMode::AtLeast => fsize.max(modsize),
                    TruncateMode::RoundDown => fsize - fsize % modsize,
                    TruncateMode::RoundUp => fsize + fsize % modsize,
                };
                match file.set_len(tsize) {
                    Ok(_) => {}
                    Err(f) => crash!(1, "{}", f.to_string()),
                };
            }
            Err(f) => crash!(1, "{}", f.to_string()),
        }
    }
}

/// Parse a size string into a number of bytes.
///
/// A size string comprises an integer and an optional unit. The unit
/// may be K, M, G, T, P, E, Z, or Y (powers of 1024) or KB, MB,
/// etc. (powers of 1000).
///
/// # Errors
///
/// This function returns an error if the string does not begin with a
/// numeral, or if the unit is not one of the supported units described
/// in the preceding section.
///
/// # Examples
///
/// ```rust,ignore
/// assert_eq!(parse_size("123").unwrap(), 123);
/// assert_eq!(parse_size("123K").unwrap(), 123 * 1024);
/// assert_eq!(parse_size("123KB").unwrap(), 123 * 1000);
/// ```
fn parse_size(size: &str) -> Result<u64, ()> {
    // Get the numeric part of the size argument. For example, if the
    // argument is "123K", then the numeric part is "123".
    let numeric_string: String = size.chars().take_while(|c| c.is_digit(10)).collect();
    let number: u64 = match numeric_string.parse() {
        Ok(n) => n,
        Err(_) => return Err(()),
    };

    // Get the alphabetic units part of the size argument and compute
    // the factor it represents. For example, if the argument is "123K",
    // then the unit part is "K" and the factor is 1024. This may be the
    // empty string, in which case, the factor is 1.
    let n = numeric_string.len();
    let (base, exponent): (u64, u32) = match &size[n..] {
        "" => (1, 0),
        "K" | "k" => (1024, 1),
        "M" | "m" => (1024, 2),
        "G" | "g" => (1024, 3),
        "T" | "t" => (1024, 4),
        "P" | "p" => (1024, 5),
        "E" | "e" => (1024, 6),
        "Z" | "z" => (1024, 7),
        "Y" | "y" => (1024, 8),
        "KB" | "kB" => (1000, 1),
        "MB" | "mB" => (1000, 2),
        "GB" | "gB" => (1000, 3),
        "TB" | "tB" => (1000, 4),
        "PB" | "pB" => (1000, 5),
        "EB" | "eB" => (1000, 6),
        "ZB" | "zB" => (1000, 7),
        "YB" | "yB" => (1000, 8),
        _ => return Err(()),
    };
    let factor = base.pow(exponent);
    Ok(number * factor)
}

#[cfg(test)]
mod tests {
    use crate::parse_size;

    #[test]
    fn test_parse_size_zero() {
        assert_eq!(parse_size("0").unwrap(), 0);
        assert_eq!(parse_size("0K").unwrap(), 0);
        assert_eq!(parse_size("0KB").unwrap(), 0);
    }

    #[test]
    fn test_parse_size_without_factor() {
        assert_eq!(parse_size("123").unwrap(), 123);
    }

    #[test]
    fn test_parse_size_kilobytes() {
        assert_eq!(parse_size("123K").unwrap(), 123 * 1024);
        assert_eq!(parse_size("123KB").unwrap(), 123 * 1000);
    }

    #[test]
    fn test_parse_size_megabytes() {
        assert_eq!(parse_size("123").unwrap(), 123);
        assert_eq!(parse_size("123M").unwrap(), 123 * 1024 * 1024);
        assert_eq!(parse_size("123MB").unwrap(), 123 * 1000 * 1000);
    }
}
<commit_msg>truncate: remove read permissions from OpenOptions<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Alex Lyon <arcterus@mail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

// spell-checker:ignore (ToDO) RFILE refsize rfilename fsize tsize

#[macro_use]
extern crate uucore;

use clap::{App, Arg};
use std::fs::{metadata, OpenOptions};
use std::io::ErrorKind;
use std::path::Path;

#[derive(Eq, PartialEq)]
enum TruncateMode {
    Absolute,
    Reference,
    Extend,
    Reduce,
    AtMost,
    AtLeast,
    RoundDown,
    RoundUp,
}

static ABOUT: &str = "Shrink or extend the size of each file to the specified size.";
static VERSION: &str = env!("CARGO_PKG_VERSION");

pub mod options {
    pub static IO_BLOCKS: &str = "io-blocks";
    pub static NO_CREATE: &str = "no-create";
    pub static REFERENCE: &str = "reference";
    pub static SIZE: &str = "size";
}

static ARG_FILES: &str = "files";

fn get_usage() -> String {
    format!("{0} [OPTION]... [FILE]...", executable!())
}

fn get_long_usage() -> String {
    String::from(
        "
    SIZE is an integer with an optional prefix and optional unit.
    The available units (K, M, G, T, P, E, Z, and Y) use the following format:
        'KB' =>           1000 (kilobytes)
        'K'  =>           1024 (kibibytes)
        'MB' =>      1000*1000 (megabytes)
        'M'  =>      1024*1024 (mebibytes)
        'GB' => 1000*1000*1000 (gigabytes)
        'G'  => 1024*1024*1024 (gibibytes)
    SIZE may also be prefixed by one of the following to adjust the size of each
    file based on its current size:
        '+'  => extend by
        '-'  => reduce by
        '<'  => at most
        '>'  => at least
        '/'  => round down to multiple of
        '%'  => round up to multiple of",
    )
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = get_usage();
    let long_usage = get_long_usage();

    let matches = App::new(executable!())
        .version(VERSION)
        .about(ABOUT)
        .usage(&usage[..])
        .after_help(&long_usage[..])
        .arg(
            Arg::with_name(options::IO_BLOCKS)
            .short("o")
            .long(options::IO_BLOCKS)
            .help("treat SIZE as the number of I/O blocks of the file rather than bytes (NOT IMPLEMENTED)")
        )
        .arg(
            Arg::with_name(options::NO_CREATE)
            .short("c")
            .long(options::NO_CREATE)
            .help("do not create files that do not exist")
        )
        .arg(
            Arg::with_name(options::REFERENCE)
            .short("r")
            .long(options::REFERENCE)
            .help("base the size of each file on the size of RFILE")
            .value_name("RFILE")
        )
        .arg(
            Arg::with_name(options::SIZE)
            .short("s")
            .long("size")
            .help("set or adjust the size of each file according to SIZE, which is in bytes unless --io-blocks is specified")
            .value_name("SIZE")
        )
        .arg(Arg::with_name(ARG_FILES).multiple(true).takes_value(true).min_values(1))
        .get_matches_from(args);

    let files: Vec<String> = matches
        .values_of(ARG_FILES)
        .map(|v| v.map(ToString::to_string).collect())
        .unwrap_or_default();

    if files.is_empty() {
        show_error!("Missing an argument");
        return 1;
    } else {
        let io_blocks = matches.is_present(options::IO_BLOCKS);
        let no_create = matches.is_present(options::NO_CREATE);
        let reference = matches.value_of(options::REFERENCE).map(String::from);
        let size = matches.value_of(options::SIZE).map(String::from);
        if reference.is_none() && size.is_none() {
            crash!(1, "you must specify either --reference or --size");
        } else {
            truncate(no_create, io_blocks, reference, size, files);
        }
    }

    0
}

fn truncate(
    no_create: bool,
    _: bool,
    reference: Option<String>,
    size: Option<String>,
    filenames: Vec<String>,
) {
    let (modsize, mode) = match size {
        Some(size_string) => {
            // Trim any whitespace.
            let size_string = size_string.trim();

            // Get the modifier character from the size string, if any. For
            // example, if the argument is "+123", then the modifier is '+'.
            let c = size_string.chars().next().unwrap();

            let mode = match c {
                '+' => TruncateMode::Extend,
                '-' => TruncateMode::Reduce,
                '<' => TruncateMode::AtMost,
                '>' => TruncateMode::AtLeast,
                '/' => TruncateMode::RoundDown,
                '%' => TruncateMode::RoundUp,
                _ => TruncateMode::Absolute, /* assume that the size is just a number */
            };

            // If there was a modifier character, strip it.
            let size_string = match mode {
                TruncateMode::Absolute => size_string,
                _ => &size_string[1..],
            };
            let num_bytes = match parse_size(size_string) {
                Ok(b) => b,
                Err(_) => crash!(1, "Invalid number: ‘{}’", size_string),
            };
            (num_bytes, mode)
        }
        None => (0, TruncateMode::Reference),
    };

    let refsize = match reference {
        Some(ref rfilename) => {
            match mode {
                // Only Some modes work with a reference
                TruncateMode::Reference => (), //No --size was given
                TruncateMode::Extend => (),
                TruncateMode::Reduce => (),
                _ => crash!(1, "you must specify a relative ‘--size’ with ‘--reference’"),
            };
            match metadata(rfilename) {
                Ok(meta) => meta.len(),
                Err(f) => match f.kind() {
                    ErrorKind::NotFound => {
                        crash!(1, "cannot stat '{}': No such file or directory", rfilename)
                    }
                    _ => crash!(1, "{}", f.to_string()),
                },
            }
        }
        None => 0,
    };
    for filename in &filenames {
        let path = Path::new(filename);
        match OpenOptions::new().write(true).create(!no_create).open(path) {
            Ok(file) => {
                let fsize = match reference {
                    Some(_) => refsize,
                    None => match metadata(filename) {
                        Ok(meta) => meta.len(),
                        Err(f) => {
                            show_warning!("{}", f.to_string());
                            continue;
                        }
                    },
                };
                let tsize: u64 = match mode {
                    TruncateMode::Absolute => modsize,
                    TruncateMode::Reference => fsize,
                    TruncateMode::Extend => fsize + modsize,
                    TruncateMode::Reduce => fsize - modsize,
                    TruncateMode::AtMost => fsize.min(modsize),
                    TruncateMode::AtLeast => fsize.max(modsize),
                    TruncateMode::RoundDown => fsize - fsize % modsize,
                    TruncateMode::RoundUp => fsize + fsize % modsize,
                };
                match file.set_len(tsize) {
                    Ok(_) => {}
                    Err(f) => crash!(1, "{}", f.to_string()),
                };
            }
            Err(f) => crash!(1, "{}", f.to_string()),
        }
    }
}

/// Parse a size string into a number of bytes.
///
/// A size string comprises an integer and an optional unit. The unit
/// may be K, M, G, T, P, E, Z, or Y (powers of 1024) or KB, MB,
/// etc. (powers of 1000).
///
/// # Errors
///
/// This function returns an error if the string does not begin with a
/// numeral, or if the unit is not one of the supported units described
/// in the preceding section.
///
/// # Examples
///
/// ```rust,ignore
/// assert_eq!(parse_size("123").unwrap(), 123);
/// assert_eq!(parse_size("123K").unwrap(), 123 * 1024);
/// assert_eq!(parse_size("123KB").unwrap(), 123 * 1000);
/// ```
fn parse_size(size: &str) -> Result<u64, ()> {
    // Get the numeric part of the size argument. For example, if the
    // argument is "123K", then the numeric part is "123".
    let numeric_string: String = size.chars().take_while(|c| c.is_digit(10)).collect();
    let number: u64 = match numeric_string.parse() {
        Ok(n) => n,
        Err(_) => return Err(()),
    };

    // Get the alphabetic units part of the size argument and compute
    // the factor it represents. For example, if the argument is "123K",
    // then the unit part is "K" and the factor is 1024. This may be the
    // empty string, in which case, the factor is 1.
    let n = numeric_string.len();
    let (base, exponent): (u64, u32) = match &size[n..] {
        "" => (1, 0),
        "K" | "k" => (1024, 1),
        "M" | "m" => (1024, 2),
        "G" | "g" => (1024, 3),
        "T" | "t" => (1024, 4),
        "P" | "p" => (1024, 5),
        "E" | "e" => (1024, 6),
        "Z" | "z" => (1024, 7),
        "Y" | "y" => (1024, 8),
        "KB" | "kB" => (1000, 1),
        "MB" | "mB" => (1000, 2),
        "GB" | "gB" => (1000, 3),
        "TB" | "tB" => (1000, 4),
        "PB" | "pB" => (1000, 5),
        "EB" | "eB" => (1000, 6),
        "ZB" | "zB" => (1000, 7),
        "YB" | "yB" => (1000, 8),
        _ => return Err(()),
    };
    let factor = base.pow(exponent);
    Ok(number * factor)
}

#[cfg(test)]
mod tests {
    use crate::parse_size;

    #[test]
    fn test_parse_size_zero() {
        assert_eq!(parse_size("0").unwrap(), 0);
        assert_eq!(parse_size("0K").unwrap(), 0);
        assert_eq!(parse_size("0KB").unwrap(), 0);
    }

    #[test]
    fn test_parse_size_without_factor() {
        assert_eq!(parse_size("123").unwrap(), 123);
    }

    #[test]
    fn test_parse_size_kilobytes() {
        assert_eq!(parse_size("123K").unwrap(), 123 * 1024);
        assert_eq!(parse_size("123KB").unwrap(), 123 * 1000);
    }

    #[test]
    fn test_parse_size_megabytes() {
        assert_eq!(parse_size("123").unwrap(), 123);
        assert_eq!(parse_size("123M").unwrap(), 123 * 1024 * 1024);
        assert_eq!(parse_size("123MB").unwrap(), 123 * 1000 * 1000);
    }
}
<|endoftext|>


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

//! Client for connecting and communicating with a server listener which speaks SrvProtocol.
//!
//! Currently all functions will block and wait for a response from the remote. This is likely to
//! change to a futures based implementation as our needs increase.
//!
//! # RPC Call Example
//!
//! ```ignore
//! let conn = SrvClient::connect(&listen_addr, secret_key).wait()?;
//! let msg = protocols::ctl::ServiceGetDefaultCfg::new();
//! conn.call(msg).for_each(|reply| {
//!     match reply.message_id() {
//!         "ServiceCfg" => {
//!             let m = reply.parse::<protocols::types::ServiceCfg>().unwrap();
//!             println!("{}", m.get_default());
//!         }
//!         "NetErr" => {
//!             let m = reply.parse::<protocols::net::NetErr>().unwrap();
//!             println!("{}", m);
//!         }
//!         _ => (),
//!     }
//!     Ok(())
//! })
//! ```

#[macro_use]
extern crate futures;
extern crate habitat_sup_protocol as protocol;
#[macro_use]
extern crate log;
extern crate prost;
extern crate tokio;
extern crate tokio_io;

use std::error;
use std::fmt;
use std::io;
use std::net::SocketAddr;
use std::path::PathBuf;

use futures::prelude::*;
use futures::sink;
use protocol::codec::*;
use protocol::net::NetErr;
use tokio::net::TcpStream;
use tokio_io::AsyncRead;

pub type SrvSend = sink::Send<SrvStream>;

/// Error types returned by a [`SrvClient`].
#[derive(Debug)]
pub enum SrvClientError {
    /// The remote server unexpectedly closed the connection.
    ConnectionClosed,
    /// Unable to locate a secret key on disk.
    CtlSecretNotFound(PathBuf),
    /// Decoding a message from the remote failed.
    Decode(prost::DecodeError),
    /// An Os level IO error occurred.
    Io(io::Error),
    /// An RPC call to the remote was received but failed.
    NetErr(NetErr),
}

impl error::Error for SrvClientError {
    fn description(&self) -> &str {
        match *self {
            SrvClientError::ConnectionClosed => "Connection closed",
            SrvClientError::CtlSecretNotFound(_) => "Ctl secret key not found",
            SrvClientError::Decode(ref err) => err.description(),
            SrvClientError::Io(ref err) => err.description(),
            SrvClientError::NetErr(ref err) => err.description(),
        }
    }
}

impl fmt::Display for SrvClientError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let content = match *self {
            SrvClientError::ConnectionClosed => format!("Connection closed"),
            SrvClientError::CtlSecretNotFound(ref path) => format!(
                "No Supervisor CtlGateway secret set in `cli.toml` or found at {}. Run \
                 `hab setup` or run the Supervisor for the first time before attempting to \
                 command the Supervisor.",
                path.display()
            ),
            SrvClientError::Decode(ref err) => format!("{}", err),
            SrvClientError::Io(ref err) => format!(
                "Unable to contact the Supervisor.\n\n\
                If the Supervisor you are contacting is local, this probably means it is not running. You can run a Supervisor in the foreground with:\n\n\
                   hab sup run\n\n\
                Or try restarting the Supervisor through your oerating systems init process.\n\n\
                Original error is:\n\n\
                   {}",
                err
            ),
            SrvClientError::NetErr(ref err) => format!("{}", err),
        };
        write!(f, "{}", content)
    }
}

impl From<NetErr> for SrvClientError {
    fn from(err: NetErr) -> Self {
        SrvClientError::NetErr(err)
    }
}

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

impl From<prost::DecodeError> for SrvClientError {
    fn from(err: prost::DecodeError) -> Self {
        SrvClientError::Decode(err)
    }
}

/// Client for connecting and communicating with a server listener which speaks SrvProtocol.
///
/// See module doc for usage.
pub struct SrvClient {
    /// Sending and receiving framed socket pair for connecting and communicating to a remote
    /// server.
    socket: SrvStream,
    /// Transaction ID counter.
    ///
    /// Useful for determining the next transaction ID in sequence embed in
    /// the next transactional message sent.
    current_txn: SrvTxn,
}

impl SrvClient {
    /// Connect to the given remote server and authenticate with the given secret_key.
    pub fn connect<T>(
        addr: &SocketAddr,
        secret_key: T,
    ) -> Box<Future<Item = SrvClient, Error = SrvClientError> + 'static>
    where
        T: ToString,
    {
        let secret_key = secret_key.to_string();
        let conn = TcpStream::connect(addr)
            .map_err(SrvClientError::from)
            .and_then(move |socket| {
                let client = Self::new(socket, None);
                let mut request = protocol::ctl::Handshake::default();
                request.secret_key = Some(secret_key);
                client
                    .call(request)
                    .into_future()
                    .map_err(|(err, _)| err)
                    .and_then(move |(m, io)| {
                        m.map_or_else(
                            || Err(SrvClientError::ConnectionClosed),
                            move |m| {
                                m.try_ok()
                                    .map_err(SrvClientError::from)
                                    .and_then(|()| Ok(io.into_inner()))
                            },
                        )
                    })
            });
        Box::new(conn)
    }

    pub fn read_secret_key() -> Result<String, SrvClientError> {
        let mut buf = String::new();
        protocol::read_secret_key(protocol::sup_root(None::<String>, None::<String>), &mut buf)
            .map_err(SrvClientError::from)?;
        Ok(buf)
    }

    fn new(socket: TcpStream, current_txn: Option<SrvTxn>) -> Self {
        SrvClient {
            socket: socket.framed(SrvCodec::new()),
            current_txn: current_txn.unwrap_or_default(),
        }
    }

    /// Send a transactional request to the connected server. The returned `SrvReply` is a Stream
    /// containing one or more `SrvMessage` responses for the given request.
    pub fn call<T>(mut self, request: T) -> SrvReply
    where
        T: Into<SrvMessage> + fmt::Debug,
    {
        self.current_txn.increment();
        let mut msg: SrvMessage = request.into();
        msg.set_transaction(self.current_txn);
        trace!("Sending SrvMessage -> {:?}", msg);
        SrvReply::new(self.socket.send(msg), self.current_txn)
    }

    /// Send a non-transactional request to the connected server.
    pub fn cast<T>(self, request: T) -> SrvSend
    where
        T: Into<SrvMessage> + fmt::Debug,
    {
        let message: SrvMessage = request.into();
        self.socket.send(message)
    }
}

/// A `Future` that will resolve into a stream of one or more `SrvMessage` replies.
#[must_use = "futures do nothing unless polled"]
pub struct SrvReply {
    io: sink::Send<SrvStream>,
    state: SrvReplyState,
    txn_id: SrvTxn,
}

impl SrvReply {
    fn new(io: sink::Send<SrvStream>, txn_id: SrvTxn) -> Self {
        SrvReply {
            io: io,
            state: SrvReplyState::Sending,
            txn_id: txn_id,
        }
    }

    /// Consume the `SrvReply` and return the contained `SrvClient`.
    pub fn into_inner(self) -> SrvClient {
        match self.state {
            SrvReplyState::Receiving(io, true) => {
                SrvClient::new(io.into_inner(), Some(self.txn_id))
            }
            _ => panic!("into_inner called before complete"),
        }
    }
}

impl Stream for SrvReply {
    type Item = SrvMessage;
    type Error = SrvClientError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        loop {
            match self.state {
                SrvReplyState::Sending => {
                    let io = try_ready!(self.io.poll());
                    self.state = SrvReplyState::Receiving(io, false);
                    continue;
                }
                SrvReplyState::Receiving(_, true) => return Ok(Async::Ready(None)),
                SrvReplyState::Receiving(ref mut io, ref mut complete) => {
                    match try_ready!(io.poll()) {
                        Some(msg) => {
                            *complete = msg.is_complete();
                            return Ok(Async::Ready(Some(msg)));
                        }
                        None => return Err(SrvClientError::ConnectionClosed),
                    }
                }
            }
        }
    }
}

enum SrvReplyState {
    /// Request is sending.
    Sending,
    /// Request is sent and awaiting message(s). Receiving is complete when the bool at `self.1`
    /// is true.
    Receiving(SrvStream, bool),
}
<commit_msg>Fix typos per PR feedback<commit_after>// Copyright (c) 2018 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Client for connecting and communicating with a server listener which speaks SrvProtocol.
//!
//! Currently all functions will block and wait for a response from the remote. This is likely to
//! change to a futures based implementation as our needs increase.
//!
//! # RPC Call Example
//!
//! ```ignore
//! let conn = SrvClient::connect(&listen_addr, secret_key).wait()?;
//! let msg = protocols::ctl::ServiceGetDefaultCfg::new();
//! conn.call(msg).for_each(|reply| {
//!     match reply.message_id() {
//!         "ServiceCfg" => {
//!             let m = reply.parse::<protocols::types::ServiceCfg>().unwrap();
//!             println!("{}", m.get_default());
//!         }
//!         "NetErr" => {
//!             let m = reply.parse::<protocols::net::NetErr>().unwrap();
//!             println!("{}", m);
//!         }
//!         _ => (),
//!     }
//!     Ok(())
//! })
//! ```

#[macro_use]
extern crate futures;
extern crate habitat_sup_protocol as protocol;
#[macro_use]
extern crate log;
extern crate prost;
extern crate tokio;
extern crate tokio_io;

use std::error;
use std::fmt;
use std::io;
use std::net::SocketAddr;
use std::path::PathBuf;

use futures::prelude::*;
use futures::sink;
use protocol::codec::*;
use protocol::net::NetErr;
use tokio::net::TcpStream;
use tokio_io::AsyncRead;

pub type SrvSend = sink::Send<SrvStream>;

/// Error types returned by a [`SrvClient`].
#[derive(Debug)]
pub enum SrvClientError {
    /// The remote server unexpectedly closed the connection.
    ConnectionClosed,
    /// Unable to locate a secret key on disk.
    CtlSecretNotFound(PathBuf),
    /// Decoding a message from the remote failed.
    Decode(prost::DecodeError),
    /// An Os level IO error occurred.
    Io(io::Error),
    /// An RPC call to the remote was received but failed.
    NetErr(NetErr),
}

impl error::Error for SrvClientError {
    fn description(&self) -> &str {
        match *self {
            SrvClientError::ConnectionClosed => "Connection closed",
            SrvClientError::CtlSecretNotFound(_) => "Ctl secret key not found",
            SrvClientError::Decode(ref err) => err.description(),
            SrvClientError::Io(ref err) => err.description(),
            SrvClientError::NetErr(ref err) => err.description(),
        }
    }
}

impl fmt::Display for SrvClientError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let content = match *self {
            SrvClientError::ConnectionClosed => format!("Connection closed"),
            SrvClientError::CtlSecretNotFound(ref path) => format!(
                "No Supervisor CtlGateway secret set in `cli.toml` or found at {}. Run \
                 `hab setup` or run the Supervisor for the first time before attempting to \
                 command the Supervisor.",
                path.display()
            ),
            SrvClientError::Decode(ref err) => format!("{}", err),
            SrvClientError::Io(ref err) => format!(
                "Unable to contact the Supervisor.\n\n\
                If the Supervisor you are contacting is local, this probably means it is not running. You can run a Supervisor in the foreground with:\n\n\
                   hab sup run\n\n\
                Or try restarting the Supervisor through your operating system's init process or Windows service.\n\n\
                Original error is:\n\n\
                   {}",
                err
            ),
            SrvClientError::NetErr(ref err) => format!("{}", err),
        };
        write!(f, "{}", content)
    }
}

impl From<NetErr> for SrvClientError {
    fn from(err: NetErr) -> Self {
        SrvClientError::NetErr(err)
    }
}

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

impl From<prost::DecodeError> for SrvClientError {
    fn from(err: prost::DecodeError) -> Self {
        SrvClientError::Decode(err)
    }
}

/// Client for connecting and communicating with a server listener which speaks SrvProtocol.
///
/// See module doc for usage.
pub struct SrvClient {
    /// Sending and receiving framed socket pair for connecting and communicating to a remote
    /// server.
    socket: SrvStream,
    /// Transaction ID counter.
    ///
    /// Useful for determining the next transaction ID in sequence embed in
    /// the next transactional message sent.
    current_txn: SrvTxn,
}

impl SrvClient {
    /// Connect to the given remote server and authenticate with the given secret_key.
    pub fn connect<T>(
        addr: &SocketAddr,
        secret_key: T,
    ) -> Box<Future<Item = SrvClient, Error = SrvClientError> + 'static>
    where
        T: ToString,
    {
        let secret_key = secret_key.to_string();
        let conn = TcpStream::connect(addr)
            .map_err(SrvClientError::from)
            .and_then(move |socket| {
                let client = Self::new(socket, None);
                let mut request = protocol::ctl::Handshake::default();
                request.secret_key = Some(secret_key);
                client
                    .call(request)
                    .into_future()
                    .map_err(|(err, _)| err)
                    .and_then(move |(m, io)| {
                        m.map_or_else(
                            || Err(SrvClientError::ConnectionClosed),
                            move |m| {
                                m.try_ok()
                                    .map_err(SrvClientError::from)
                                    .and_then(|()| Ok(io.into_inner()))
                            },
                        )
                    })
            });
        Box::new(conn)
    }

    pub fn read_secret_key() -> Result<String, SrvClientError> {
        let mut buf = String::new();
        protocol::read_secret_key(protocol::sup_root(None::<String>, None::<String>), &mut buf)
            .map_err(SrvClientError::from)?;
        Ok(buf)
    }

    fn new(socket: TcpStream, current_txn: Option<SrvTxn>) -> Self {
        SrvClient {
            socket: socket.framed(SrvCodec::new()),
            current_txn: current_txn.unwrap_or_default(),
        }
    }

    /// Send a transactional request to the connected server. The returned `SrvReply` is a Stream
    /// containing one or more `SrvMessage` responses for the given request.
    pub fn call<T>(mut self, request: T) -> SrvReply
    where
        T: Into<SrvMessage> + fmt::Debug,
    {
        self.current_txn.increment();
        let mut msg: SrvMessage = request.into();
        msg.set_transaction(self.current_txn);
        trace!("Sending SrvMessage -> {:?}", msg);
        SrvReply::new(self.socket.send(msg), self.current_txn)
    }

    /// Send a non-transactional request to the connected server.
    pub fn cast<T>(self, request: T) -> SrvSend
    where
        T: Into<SrvMessage> + fmt::Debug,
    {
        let message: SrvMessage = request.into();
        self.socket.send(message)
    }
}

/// A `Future` that will resolve into a stream of one or more `SrvMessage` replies.
#[must_use = "futures do nothing unless polled"]
pub struct SrvReply {
    io: sink::Send<SrvStream>,
    state: SrvReplyState,
    txn_id: SrvTxn,
}

impl SrvReply {
    fn new(io: sink::Send<SrvStream>, txn_id: SrvTxn) -> Self {
        SrvReply {
            io: io,
            state: SrvReplyState::Sending,
            txn_id: txn_id,
        }
    }

    /// Consume the `SrvReply` and return the contained `SrvClient`.
    pub fn into_inner(self) -> SrvClient {
        match self.state {
            SrvReplyState::Receiving(io, true) => {
                SrvClient::new(io.into_inner(), Some(self.txn_id))
            }
            _ => panic!("into_inner called before complete"),
        }
    }
}

impl Stream for SrvReply {
    type Item = SrvMessage;
    type Error = SrvClientError;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        loop {
            match self.state {
                SrvReplyState::Sending => {
                    let io = try_ready!(self.io.poll());
                    self.state = SrvReplyState::Receiving(io, false);
                    continue;
                }
                SrvReplyState::Receiving(_, true) => return Ok(Async::Ready(None)),
                SrvReplyState::Receiving(ref mut io, ref mut complete) => {
                    match try_ready!(io.poll()) {
                        Some(msg) => {
                            *complete = msg.is_complete();
                            return Ok(Async::Ready(Some(msg)));
                        }
                        None => return Err(SrvClientError::ConnectionClosed),
                    }
                }
            }
        }
    }
}

enum SrvReplyState {
    /// Request is sending.
    Sending,
    /// Request is sent and awaiting message(s). Receiving is complete when the bool at `self.1`
    /// is true.
    Receiving(SrvStream, bool),
}
<|endoftext|>


<commit_before>// Copyright 2018 pdb Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! Definitions for PE headers contained in PDBs.

// PDBs contain PE section headers in one or two streams. `pdb::pe` is responsible for parsing them.

use std::fmt;

use crate::common::*;

/// A PE `IMAGE_SECTION_HEADER`, as described in [the Microsoft documentation](https://msdn.microsoft.com/en-us/library/windows/desktop/ms680341(v=vs.85).aspx).
#[derive(Copy, Clone, PartialEq, Eq, Default)]
pub struct ImageSectionHeader {
    /// An 8-byte, null-padded UTF-8 string. There is no terminating null character if the string is
    /// exactly eight characters long. For longer names, this member contains a forward slash (`/`)
    /// followed by an ASCII representation of a decimal number that is an offset into the string
    /// table. Executable images do not use a string table and do not support section names longer
    /// than eight characters.
    pub name: [u8; 8],

    /// The file address.
    pub physical_address: u32,

    /// The address of the first byte of the section when loaded into memory, relative to the image
    /// base. For object files, this is the address of the first byte before relocation is applied.
    pub virtual_address: u32,

    /// The size of the initialized data on disk, in bytes. This value must be a multiple of the
    /// `FileAlignment` member of the `IMAGE_OPTIONAL_HEADER` structure. If this value is less than
    /// the `VirtualSize` member, the remainder of the section is filled with zeroes. If the section
    /// contains only uninitialized data, the member is zero.
    pub size_of_raw_data: u32,

    /// A file pointer to the first page within the COFF file. This value must be a multiple of the
    /// `FileAlignment` member of the `IMAGE_OPTIONAL_HEADER` structure. If a section contains only
    /// uninitialized data, set this member is zero.
    pub pointer_to_raw_data: u32,

    /// A file pointer to the beginning of the relocation entries for the section. If there are no
    /// relocations, this value is zero.
    pub pointer_to_relocations: u32,

    /// A file pointer to the beginning of the line-number entries for the section. If there are no
    /// COFF line numbers, this value is zero.
    pub pointer_to_line_numbers: u32,

    /// The number of relocation entries for the section. This value is zero for executable images.
    pub number_of_relocations: u16,

    /// The number of line-number entries for the section.
    pub number_of_line_numbers: u16,

    /// The characteristics of the image.
    pub characteristics: u32,
}

impl ImageSectionHeader {
    pub(crate) fn parse(parse_buffer: &mut ParseBuffer<'_>) -> Result<Self> {
        let name_bytes = parse_buffer.take(8)?;

        Ok(ImageSectionHeader {
            name: [
                name_bytes[0],
                name_bytes[1],
                name_bytes[2],
                name_bytes[3],
                name_bytes[4],
                name_bytes[5],
                name_bytes[6],
                name_bytes[7],
            ],
            physical_address: parse_buffer.parse_u32()?,
            virtual_address: parse_buffer.parse_u32()?,
            size_of_raw_data: parse_buffer.parse_u32()?,
            pointer_to_raw_data: parse_buffer.parse_u32()?,
            pointer_to_relocations: parse_buffer.parse_u32()?,
            pointer_to_line_numbers: parse_buffer.parse_u32()?,
            number_of_relocations: parse_buffer.parse_u16()?,
            number_of_line_numbers: parse_buffer.parse_u16()?,
            characteristics: parse_buffer.parse_u32()?,
        })
    }

    /// Returns the name of the section.
    pub fn name(&self) -> &str {
        let end = self
            .name
            .iter()
            .position(|ch| *ch == 0)
            .unwrap_or(self.name.len());

        // The spec guarantees that the name is a proper UTF-8 string.
        // TODO: Look up long names from the string table.
        std::str::from_utf8(&self.name[0..end]).unwrap_or("")
    }
}

impl fmt::Debug for ImageSectionHeader {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ImageSectionHeader")
            .field("name()", &self.name().to_string())
            .field(
                "physical_address",
                &format_args!("{:#x}", self.physical_address),
            )
            .field(
                "virtual_address",
                &format_args!("{:#x}", self.virtual_address),
            )
            .field("size_of_raw_data", &self.size_of_raw_data)
            .field(
                "pointer_to_raw_data",
                &format_args!("{:#x}", self.pointer_to_raw_data),
            )
            .field(
                "pointer_to_relocations",
                &format_args!("{:#x}", self.pointer_to_relocations),
            )
            .field(
                "pointer_to_line_numbers",
                &format_args!("{:#x}", self.pointer_to_line_numbers),
            )
            .field("number_of_relocations", &self.number_of_relocations)
            .field("number_of_line_numbers", &self.number_of_line_numbers)
            .field(
                "characteristics",
                &format_args!("{:#x}", self.characteristics),
            )
            .finish()
    }
}

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

    #[test]
    fn image_section_header() {
        let bytes: Vec<u8> = vec![
            0x2E, 0x64, 0x61, 0x74, 0x61, 0x00, 0x00, 0x00, 0x48, 0x35, 0x09, 0x00, 0x00, 0xD0,
            0x1E, 0x00, 0x00, 0xFE, 0x00, 0x00, 0x00, 0xA2, 0x1E, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0xC8,
        ];

        let mut parse_buffer = ParseBuffer::from(bytes.as_slice());

        let ish = ImageSectionHeader::parse(&mut parse_buffer).expect("parse");
        assert_eq!(&ish.name, b".data\0\0\0");
        assert_eq!(ish.name(), ".data");
        assert_eq!(ish.physical_address, 0x93548);
        assert_eq!(ish.virtual_address, 0x001e_d000);
        assert_eq!(ish.size_of_raw_data, 0xfe00);
        assert_eq!(ish.pointer_to_raw_data, 0x001e_a200);
        assert_eq!(ish.pointer_to_relocations, 0);
        assert_eq!(ish.pointer_to_line_numbers, 0);
        assert_eq!(ish.number_of_relocations, 0);
        assert_eq!(ish.number_of_line_numbers, 0);
        assert_eq!(ish.characteristics, 0xc800_0040);
    }
}
<commit_msg>Rename ImageSectionHeader::physical_address to virtual_size<commit_after>// Copyright 2018 pdb Developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.

//! Definitions for PE headers contained in PDBs.

// PDBs contain PE section headers in one or two streams. `pdb::pe` is responsible for parsing them.

use std::fmt;

use crate::common::*;

/// A PE `IMAGE_SECTION_HEADER`, as described in [the Microsoft documentation](https://msdn.microsoft.com/en-us/library/windows/desktop/ms680341(v=vs.85).aspx).
#[derive(Copy, Clone, PartialEq, Eq, Default)]
pub struct ImageSectionHeader {
    /// An 8-byte, null-padded UTF-8 string. There is no terminating null character if the string is
    /// exactly eight characters long. For longer names, this member contains a forward slash (`/`)
    /// followed by an ASCII representation of a decimal number that is an offset into the string
    /// table. Executable images do not use a string table and do not support section names longer
    /// than eight characters.
    pub name: [u8; 8],

    /// The total size of the section when loaded into memory, in bytes. If this value is greater
    /// than the [`size_of_raw_data`](Self::size_of_raw_data) member, the section is filled with
    /// zeroes. This field is valid only for executable images and should be set to `0` for object
    /// files.
    ///
    /// In object files, this field would be replaced with the physical file address. Such headers
    /// are never embedded in PDBs.
    pub virtual_size: u32,

    /// The address of the first byte of the section when loaded into memory, relative to the image
    /// base. For object files, this is the address of the first byte before relocation is applied.
    pub virtual_address: u32,

    /// The size of the initialized data on disk, in bytes. This value must be a multiple of the
    /// `FileAlignment` member of the [`IMAGE_OPTIONAL_HEADER`] structure. If this value is less than
    /// the [`virtual_size`](Self::virtual_size) member, the remainder of the section is filled with
    /// zeroes. If the section contains only uninitialized data, the member is zero.
    ///
    /// [`IMAGE_OPTIONAL_HEADER`]: https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-image_optional_header32
    pub size_of_raw_data: u32,

    /// A file pointer to the first page within the COFF file. This value must be a multiple of the
    /// `FileAlignment` member of the [`IMAGE_OPTIONAL_HEADER`] structure. If a section contains only
    /// uninitialized data, set this member is zero.
    ///
    /// [`IMAGE_OPTIONAL_HEADER`]: https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-image_optional_header32
    pub pointer_to_raw_data: u32,

    /// A file pointer to the beginning of the relocation entries for the section. If there are no
    /// relocations, this value is zero.
    pub pointer_to_relocations: u32,

    /// A file pointer to the beginning of the line-number entries for the section. If there are no
    /// COFF line numbers, this value is zero.
    pub pointer_to_line_numbers: u32,

    /// The number of relocation entries for the section. This value is zero for executable images.
    pub number_of_relocations: u16,

    /// The number of line-number entries for the section.
    pub number_of_line_numbers: u16,

    /// The characteristics of the image.
    pub characteristics: u32,
}

impl ImageSectionHeader {
    pub(crate) fn parse(parse_buffer: &mut ParseBuffer<'_>) -> Result<Self> {
        let name_bytes = parse_buffer.take(8)?;

        Ok(ImageSectionHeader {
            name: [
                name_bytes[0],
                name_bytes[1],
                name_bytes[2],
                name_bytes[3],
                name_bytes[4],
                name_bytes[5],
                name_bytes[6],
                name_bytes[7],
            ],
            virtual_size: parse_buffer.parse_u32()?,
            virtual_address: parse_buffer.parse_u32()?,
            size_of_raw_data: parse_buffer.parse_u32()?,
            pointer_to_raw_data: parse_buffer.parse_u32()?,
            pointer_to_relocations: parse_buffer.parse_u32()?,
            pointer_to_line_numbers: parse_buffer.parse_u32()?,
            number_of_relocations: parse_buffer.parse_u16()?,
            number_of_line_numbers: parse_buffer.parse_u16()?,
            characteristics: parse_buffer.parse_u32()?,
        })
    }

    /// Returns the name of the section.
    pub fn name(&self) -> &str {
        let end = self
            .name
            .iter()
            .position(|ch| *ch == 0)
            .unwrap_or(self.name.len());

        // The spec guarantees that the name is a proper UTF-8 string.
        // TODO: Look up long names from the string table.
        std::str::from_utf8(&self.name[0..end]).unwrap_or("")
    }
}

impl fmt::Debug for ImageSectionHeader {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ImageSectionHeader")
            .field("name()", &self.name().to_string())
            .field("virtual_size", &format_args!("{:#x}", self.virtual_size))
            .field(
                "virtual_address",
                &format_args!("{:#x}", self.virtual_address),
            )
            .field("size_of_raw_data", &self.size_of_raw_data)
            .field(
                "pointer_to_raw_data",
                &format_args!("{:#x}", self.pointer_to_raw_data),
            )
            .field(
                "pointer_to_relocations",
                &format_args!("{:#x}", self.pointer_to_relocations),
            )
            .field(
                "pointer_to_line_numbers",
                &format_args!("{:#x}", self.pointer_to_line_numbers),
            )
            .field("number_of_relocations", &self.number_of_relocations)
            .field("number_of_line_numbers", &self.number_of_line_numbers)
            .field(
                "characteristics",
                &format_args!("{:#x}", self.characteristics),
            )
            .finish()
    }
}

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

    #[test]
    fn image_section_header() {
        let bytes: Vec<u8> = vec![
            0x2E, 0x64, 0x61, 0x74, 0x61, 0x00, 0x00, 0x00, 0x48, 0x35, 0x09, 0x00, 0x00, 0xD0,
            0x1E, 0x00, 0x00, 0xFE, 0x00, 0x00, 0x00, 0xA2, 0x1E, 0x00, 0x00, 0x00, 0x00, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0xC8,
        ];

        let mut parse_buffer = ParseBuffer::from(bytes.as_slice());

        let ish = ImageSectionHeader::parse(&mut parse_buffer).expect("parse");
        assert_eq!(&ish.name, b".data\0\0\0");
        assert_eq!(ish.name(), ".data");
        assert_eq!(ish.virtual_size, 0x93548);
        assert_eq!(ish.virtual_address, 0x001e_d000);
        assert_eq!(ish.size_of_raw_data, 0xfe00);
        assert_eq!(ish.pointer_to_raw_data, 0x001e_a200);
        assert_eq!(ish.pointer_to_relocations, 0);
        assert_eq!(ish.pointer_to_line_numbers, 0);
        assert_eq!(ish.number_of_relocations, 0);
        assert_eq!(ish.number_of_line_numbers, 0);
        assert_eq!(ish.characteristics, 0xc800_0040);
    }
}
<|endoftext|>


<commit_before>//! # generator run time support
//!
//! generator run time context management
//!
use std::mem;
use std::ptr;
use std::any::Any;
use std::intrinsics::type_name;

use stack::Stack;
use reg_context::RegContext;

/// each thread has it's own generator context stack
thread_local!(static ROOT_CONTEXT: Box<Context> = {
    let mut root = Box::new(Context::new(0));
    let p = &mut *root as *mut _;
    root.parent = p; // init top to current
    root
});

/// yield panic error types
#[allow(dead_code)]
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Error {
    /// Cancel panic
    Cancel,
    /// Type mismatch panic
    TypeErr,
    /// Stack overflow panic
    StackErr,
    /// Wrong Context panic
    ContextErr,
}

/// generator context
pub struct Context {
    /// generator regs context
    pub regs: RegContext,
    /// generator execution stack
    pub stack: Stack,
    /// passed in para for send
    pub para: *mut Any,
    /// this is just a buffer for the return value
    pub ret: *mut Any,
    /// track generator ref, yield will -1, send will +1
    pub _ref: u32,
    /// propagate panic
    pub err: Option<Box<Any + Send>>,
    /// context local storage
    pub local_data: *mut u8,

    /// child context
    child: *mut Context,
    /// parent context
    pub parent: *mut Context,
}

impl Context {
    /// return a default generator context
    pub fn new(size: usize) -> Context {
        Context {
            regs: RegContext::empty(),
            stack: Stack::new(size),
            para: unsafe { mem::uninitialized() },
            ret: unsafe { mem::uninitialized() },
            _ref: 1, // none zero means it's not running
            err: None,
            child: ptr::null_mut(),
            parent: ptr::null_mut(),
            local_data: ptr::null_mut(),
        }
    }

    /// judge it's generator context
    #[inline]
    pub fn is_generator(&self) -> bool {
        self.parent != self as *const _ as *mut _
    }

    /// get current generator send para
    #[inline]
    pub fn get_para(&self) -> Option
        where A: Any
    {
        fn error() -> ! {
            let t = unsafe { type_name::() };
            error!("get yield type error detected, expected type: {}", t);
            panic!(Error::TypeErr)
        }

        let para = unsafe { &mut *self.para };
        match para.downcast_mut::<Option>() {
            Some(v) => v.take(),
            None => error::(),
        }
    }

    /// set current generator return value
    #[inline]
    pub fn set_ret<T>(&mut self, v: T)
        where T: Any
    {
        fn error<T>() -> ! {
            let t = unsafe { type_name::<T>() };
            error!("yield type error detected, expected type: {}", t);
            panic!(Error::TypeErr)
        }

        let ret = unsafe { &mut *self.ret };
        match ret.downcast_mut::<Option<T>>() {
            Some(r) => *r = Some(v),
            None => error::<T>(),
        }
    }
}

/// Coroutine managing environment
pub struct ContextStack {
    root: *mut Context,
}

impl ContextStack {
    #[inline]
    pub fn current() -> ContextStack {
        let root = ROOT_CONTEXT.with(|r| &**r as *const _ as *mut Context);
        ContextStack { root: root }
    }

    /// get the top context
    #[inline]
    pub fn top(&self) -> &'static mut Context {
        let root = unsafe { &mut *self.root };
        unsafe { &mut *root.parent }
    }

    /// get the coroutine context
    #[inline]
    pub fn co_ctx(&self) -> &'static mut Context {
        let root = unsafe { &mut *self.root };

        // search from top
        let mut ctx = unsafe { &mut *root.parent };
        while ctx as *const _ != root as *const _ {
            if !ctx.local_data.is_null() {
                return ctx;
            }
            ctx = unsafe { &mut *ctx.parent };
        }
        // not find any coroutine
        unreachable!("there is no coroutine found!");
    }

    /// push the context to the thread context list
    #[inline]
    pub fn push_context(&self, ctx: *mut Context) {
        let root = unsafe { &mut *self.root };
        let ctx = unsafe { &mut *ctx };
        let top = unsafe { &mut *root.parent };
        let new_top = ctx.parent;

        // link top and new ctx
        top.child = ctx;
        ctx.parent = top;

        // save the new top
        root.parent = new_top;
    }

    /// pop the context from the thread context list and return it's parent context
    #[inline]
    pub fn pop_context(&self, ctx: *mut Context) -> &'static mut Context {
        let root = unsafe { &mut *self.root };
        let ctx = unsafe { &mut *ctx };
        let parent = unsafe { &mut *ctx.parent };

        // save the old top in ctx's parent
        ctx.parent = root.parent;
        // unlink ctx and it's parent
        parent.child = ptr::null_mut();

        // save the new top
        root.parent = parent;

        parent
    }
}

/// check the current context if it's generator
#[inline]
pub fn is_generator() -> bool {
    let env = ContextStack::current();
    let root = unsafe { &mut *env.root };
    !root.child.is_null()
}

/// get the current context local data
/// only coroutine support local data
#[inline]
pub fn get_local_data() -> *mut u8 {
    let env = ContextStack::current();
    let ctx = env.co_ctx();
    ctx.local_data
}

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

    #[test]
    fn test_is_context() {
        // this is the root context
        assert!(!is_generator());
    }
}
<commit_msg>fix get_local_data bug<commit_after>//! # generator run time support
//!
//! generator run time context management
//!
use std::mem;
use std::ptr;
use std::any::Any;
use std::intrinsics::type_name;

use stack::Stack;
use reg_context::RegContext;

/// each thread has it's own generator context stack
thread_local!(static ROOT_CONTEXT: Box<Context> = {
    let mut root = Box::new(Context::new(0));
    let p = &mut *root as *mut _;
    root.parent = p; // init top to current
    root
});

/// yield panic error types
#[allow(dead_code)]
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub enum Error {
    /// Cancel panic
    Cancel,
    /// Type mismatch panic
    TypeErr,
    /// Stack overflow panic
    StackErr,
    /// Wrong Context panic
    ContextErr,
}

/// generator context
pub struct Context {
    /// generator regs context
    pub regs: RegContext,
    /// generator execution stack
    pub stack: Stack,
    /// passed in para for send
    pub para: *mut Any,
    /// this is just a buffer for the return value
    pub ret: *mut Any,
    /// track generator ref, yield will -1, send will +1
    pub _ref: u32,
    /// propagate panic
    pub err: Option<Box<Any + Send>>,
    /// context local storage
    pub local_data: *mut u8,

    /// child context
    child: *mut Context,
    /// parent context
    pub parent: *mut Context,
}

impl Context {
    /// return a default generator context
    pub fn new(size: usize) -> Context {
        Context {
            regs: RegContext::empty(),
            stack: Stack::new(size),
            para: unsafe { mem::uninitialized() },
            ret: unsafe { mem::uninitialized() },
            _ref: 1, // none zero means it's not running
            err: None,
            child: ptr::null_mut(),
            parent: ptr::null_mut(),
            local_data: ptr::null_mut(),
        }
    }

    /// judge it's generator context
    #[inline]
    pub fn is_generator(&self) -> bool {
        self.parent != self as *const _ as *mut _
    }

    /// get current generator send para
    #[inline]
    pub fn get_para(&self) -> Option
        where A: Any
    {
        fn error() -> ! {
            let t = unsafe { type_name::() };
            error!("get yield type error detected, expected type: {}", t);
            panic!(Error::TypeErr)
        }

        let para = unsafe { &mut *self.para };
        match para.downcast_mut::<Option>() {
            Some(v) => v.take(),
            None => error::(),
        }
    }

    /// set current generator return value
    #[inline]
    pub fn set_ret<T>(&mut self, v: T)
        where T: Any
    {
        fn error<T>() -> ! {
            let t = unsafe { type_name::<T>() };
            error!("yield type error detected, expected type: {}", t);
            panic!(Error::TypeErr)
        }

        let ret = unsafe { &mut *self.ret };
        match ret.downcast_mut::<Option<T>>() {
            Some(r) => *r = Some(v),
            None => error::<T>(),
        }
    }
}

/// Coroutine managing environment
pub struct ContextStack {
    root: *mut Context,
}

impl ContextStack {
    #[inline]
    pub fn current() -> ContextStack {
        let root = ROOT_CONTEXT.with(|r| &**r as *const _ as *mut Context);
        ContextStack { root: root }
    }

    /// get the top context
    #[inline]
    pub fn top(&self) -> &'static mut Context {
        let root = unsafe { &mut *self.root };
        unsafe { &mut *root.parent }
    }

    /// get the coroutine context
    #[inline]
    pub fn co_ctx(&self) -> &'static mut Context {
        let root = unsafe { &mut *self.root };

        // search from top
        let mut ctx = unsafe { &mut *root.parent };
        while ctx as *const _ != root as *const _ {
            if !ctx.local_data.is_null() {
                return ctx;
            }
            ctx = unsafe { &mut *ctx.parent };
        }
        // not find any coroutine
        unreachable!("there is no coroutine found!");
    }

    /// push the context to the thread context list
    #[inline]
    pub fn push_context(&self, ctx: *mut Context) {
        let root = unsafe { &mut *self.root };
        let ctx = unsafe { &mut *ctx };
        let top = unsafe { &mut *root.parent };
        let new_top = ctx.parent;

        // link top and new ctx
        top.child = ctx;
        ctx.parent = top;

        // save the new top
        root.parent = new_top;
    }

    /// pop the context from the thread context list and return it's parent context
    #[inline]
    pub fn pop_context(&self, ctx: *mut Context) -> &'static mut Context {
        let root = unsafe { &mut *self.root };
        let ctx = unsafe { &mut *ctx };
        let parent = unsafe { &mut *ctx.parent };

        // save the old top in ctx's parent
        ctx.parent = root.parent;
        // unlink ctx and it's parent
        parent.child = ptr::null_mut();

        // save the new top
        root.parent = parent;

        parent
    }
}

/// check the current context if it's generator
#[inline]
pub fn is_generator() -> bool {
    let env = ContextStack::current();
    let root = unsafe { &mut *env.root };
    !root.child.is_null()
}

/// get the current context local data
/// only coroutine support local data
#[inline]
pub fn get_local_data() -> *mut u8 {
    let env = ContextStack::current();
    let root = unsafe {&mut *env.root};

    // search from top
    let mut ctx = unsafe { &mut *root.parent };
    while ctx as *const _ != root as *const _ {
        if !ctx.local_data.is_null() {
            return ctx.local_data;
        }
        ctx = unsafe { &mut *ctx.parent };
    }

    ptr::null_mut()
}

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

    #[test]
    fn test_is_context() {
        // this is the root context
        assert!(!is_generator());
    }
}
<|endoftext|>


<commit_before>#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum SuffixType {
    Small,
    Large,
    LeftMostSmall,
}

// Generates a suffix array and sorts them using the "induced sorting" method
// (Thanks to the python implementation in http://zork.net/~st/jottings/sais.html)
pub fn suffix_array(mut input: Vec<u8>) -> Vec<usize> {
    input.push(0);
    let mut type_map = vec![SuffixType::Small; input.len()];
    let mut bucket_sizes = [0; 256];
    let mut bucket_heads = vec![0; 256];
    let mut bucket_tails = vec![0; 256];

    for i in (0..input.len() - 1).rev() {
        bucket_sizes[input[i] as usize] += 1;

        if input[i] >= input[i + 1] {
            if type_map[i + 1] == SuffixType::Small {
                type_map[i + 1] = SuffixType::LeftMostSmall;
            }

            type_map[i] = SuffixType::Large;
        }
    }

    input.pop();

    let mut idx = 1;
    for (i, size) in bucket_sizes.iter().enumerate() {
        bucket_heads[i] = idx;
        idx += *size;
        bucket_tails[i] = idx - 1;
    }

    let mut approx_sa = {
        let mut vec = vec![-1i32; input.len() + 1];
        let mut bucket_tails = bucket_tails.clone();
        for i in 0..input.len() {
            if type_map[i] != SuffixType::LeftMostSmall {
                continue
            }

            let bucket_idx = input[i] as usize;
            vec[bucket_tails[bucket_idx]] = i as i32;
            bucket_tails[bucket_idx] -= 1;
        }

        vec[0] = input.len() as i32;    // null byte
        vec
    };

    fn induced_sort_large(input: &[u8], approx_sa: &mut [i32],
                          mut bucket_heads: Vec<usize>, type_map: &[SuffixType]) {
        for i in 0..approx_sa.len() {
            let j = approx_sa[i] - 1;
            if j < 0 || type_map[j as usize] != SuffixType::Large {
                continue
            }

            let bucket_idx = input[j as usize] as usize;
            approx_sa[bucket_heads[bucket_idx]] = j;
            bucket_heads[bucket_idx] += 1;
        }
    }

    fn induced_sort_small(input: &[u8], approx_sa: &mut [i32],
                          mut bucket_tails: Vec<usize>, type_map: &[SuffixType]) {
        for i in (0..approx_sa.len()).rev() {
            let j = approx_sa[i] - 1;
            if j < 0 || type_map[j as usize] == SuffixType::Large {
                continue
            }

            let bucket_idx = input[j as usize] as usize;
            approx_sa[bucket_tails[bucket_idx]] = j;
            bucket_tails[bucket_idx] -= 1;
        }
    }

    induced_sort_large(&input, &mut approx_sa, bucket_heads.clone(), &type_map);
    induced_sort_small(&input, &mut approx_sa, bucket_tails.clone(), &type_map);

    fn is_equal_lms(input: &[u8], type_map: &[SuffixType], j: usize, k: usize) -> bool {
        if j == input.len() || k == input.len() {
            return false    // null byte
        }

        for i in 0..input.len() {
            let first_lms = type_map[i + j] == SuffixType::LeftMostSmall;
            let second_lms = type_map[i + k] == SuffixType::LeftMostSmall;
            if i > 0 && first_lms && second_lms {
                return true
            } else if (first_lms != second_lms) || (input[i + j] != input[i + k]) {
                return false
            }
        }

        false
    }

    let mut lms_bytes = vec![-1i32; input.len() + 1];
    let mut byte = 0usize;
    lms_bytes[approx_sa[0] as usize] = 0;
    let mut last_idx = approx_sa[0];

    for i in 1..approx_sa.len() {
        let idx = approx_sa[i];
        if type_map[idx as usize] != SuffixType::LeftMostSmall {
            continue
        }

        if !is_equal_lms(&input, &type_map, last_idx as usize, idx as usize) {
            byte += 1;
        }

        last_idx = idx;
        lms_bytes[idx as usize] = byte as i32;
    }

    drop(approx_sa);

    let mut summary_string = Vec::with_capacity(lms_bytes.len());
    let mut summary_index = Vec::with_capacity(lms_bytes.len());
    for (i, b) in lms_bytes.drain(..).enumerate() {
        if b == -1 {
            continue
        }

        summary_index.push(i);
        summary_string.push(b as u8);
    }

    let mut sa = {
        let mut bucket_tails = bucket_tails.clone();
        let summary_sa = if byte == summary_string.len() - 1 {
            let mut sa = vec![0usize; summary_string.len() + 1];
            sa[0] = summary_string.len();
            for i in 0..summary_string.len() {
                let idx = summary_string[i];
                sa[idx as usize + 1] = i;
            }

            sa
        } else {
            suffix_array(summary_string)
        };

        let mut suffix_idx = vec![-1i32; input.len() + 1];
        for i in (2..summary_sa.len()).rev() {
            let idx = summary_index[summary_sa[i] as usize];
            let bucket_idx = input[idx] as usize;
            suffix_idx[bucket_tails[bucket_idx]] = idx as i32;
            bucket_tails[bucket_idx] -= 1;
        }

        suffix_idx[0] = input.len() as i32;

        suffix_idx
    };

    induced_sort_large(&input, &mut sa, bucket_heads, &type_map);
    induced_sort_small(&input, &mut sa, bucket_tails, &type_map);

    sa.into_iter().map(|i| i as usize).collect()
}
<commit_msg>Fix some overflow/indexing errors<commit_after>use std::usize;

// FIXME: prefer this over Option<usize> (won't require another byte of memory)
const _MARKER: usize = usize::MAX;

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum SuffixType {
    Small,
    Large,
    LeftMostSmall,
}

// Generates a suffix array and sorts them using the "induced sorting" method
// (Thanks to the python implementation in http://zork.net/~st/jottings/sais.html)
pub fn suffix_array(input: Vec<u8>) -> Vec<usize> {
    let mut type_map = vec![SuffixType::Small; input.len() + 1];
    let mut bucket_sizes = vec![0; 256];
    let mut bucket_heads = vec![0; 256];
    let mut bucket_tails = vec![0; 256];

    type_map[input.len()] = SuffixType::LeftMostSmall;      // assume that this is a null byte
    type_map[input.len() - 1] = SuffixType::Large;          // should be L-type
    bucket_sizes[input[input.len() - 1] as usize] = 1;      // don't forget the last value!

    // group the bytes into S-type or L-type (also mark LMS types)
    for i in (0..input.len() - 1).rev() {
        bucket_sizes[input[i] as usize] += 1;

        if input[i] > input[i + 1] ||
           (input[i] == input[i + 1] && type_map[i + 1] == SuffixType::Large) {
            if type_map[i + 1] == SuffixType::Small {
                type_map[i + 1] = SuffixType::LeftMostSmall;
            }

            type_map[i] = SuffixType::Large;
        }
    }

    let mut idx = 1;
    // fill the bucket heads and tails
    for (i, size) in bucket_sizes.iter().enumerate() {
        bucket_heads[i] = idx;
        idx += *size;
        bucket_tails[i] = idx - 1;
    }

    let mut approx_sa = {       // build the first (approximate) SA
        let mut vec = vec![-1i32; input.len() + 1];
        let mut bucket_tails = bucket_tails.clone();
        for i in 0..input.len() {
            if type_map[i] != SuffixType::LeftMostSmall {
                continue        // ignore the LMS types for now
            }

            let bucket_idx = input[i] as usize;
            vec[bucket_tails[bucket_idx]] = i as i32;
            bucket_tails[bucket_idx] -= 1;
        }

        vec[0] = input.len() as i32;       // null byte
        vec
    };

    fn induced_sort_large(input: &[u8], approx_sa: &mut [i32],
                          mut bucket_heads: Vec<usize>, type_map: &[SuffixType]) {
        for i in 0..approx_sa.len() {
            let j = approx_sa[i] - 1;
            if j < 0 || type_map[j as usize] != SuffixType::Large {
                continue
            }

            let bucket_idx = input[j as usize] as usize;
            approx_sa[bucket_heads[bucket_idx]] = j;
            bucket_heads[bucket_idx] += 1;
        }
    }

    fn induced_sort_small(input: &[u8], approx_sa: &mut [i32],
                          mut bucket_tails: Vec<usize>, type_map: &[SuffixType]) {
        for i in (0..approx_sa.len()).rev() {
            let j = approx_sa[i] - 1;
            if j < 0 || type_map[j as usize] == SuffixType::Large {
                continue
            }

            let bucket_idx = input[j as usize] as usize;
            approx_sa[bucket_tails[bucket_idx]] = j;
            bucket_tails[bucket_idx] -= 1;
        }
    }

    induced_sort_large(&input, &mut approx_sa, bucket_heads.clone(), &type_map);
    induced_sort_small(&input, &mut approx_sa, bucket_tails.clone(), &type_map);

    fn is_equal_lms(input: &[u8], type_map: &[SuffixType], j: usize, k: usize) -> bool {
        if j == input.len() || k == input.len() {
            return false    // null byte
        }

        for i in 0..(input.len() + 1) {
            let first_lms = type_map[i + j] == SuffixType::LeftMostSmall;
            let second_lms = type_map[i + k] == SuffixType::LeftMostSmall;
            if first_lms && second_lms && i > 0 {
                return true
            } else if (first_lms != second_lms) || (input[i + j] != input[i + k]) {
                return false
            }
        }

        false
    }

    let mut lms_bytes = vec![input.len(); input.len() + 1];
    let mut byte = 0;
    lms_bytes[approx_sa[0] as usize] = 0;
    let mut last_idx = approx_sa[0] as usize;

    for i in 1..approx_sa.len() {
        let idx = if approx_sa[i] == -1 { input.len() } else { approx_sa[i] as usize };
        if type_map[idx] != SuffixType::LeftMostSmall {
            continue
        }

        if !is_equal_lms(&input, &type_map, last_idx, idx) {
            byte += 1;
        }

        last_idx = idx;
        lms_bytes[idx] = byte;
    }

    drop(approx_sa);

    let mut summary_string = Vec::with_capacity(lms_bytes.len());
    let mut summary_index = Vec::with_capacity(lms_bytes.len());
    for (i, b) in lms_bytes.drain(..).enumerate() {
        if b == input.len() {
            continue
        }

        summary_index.push(i);
        summary_string.push(b as u8);
    }

    let mut sa = {
        let mut bucket_tails = bucket_tails.clone();
        let summary_sa = if byte == summary_string.len() - 1 {
            let mut sa = vec![0usize; summary_string.len() + 1];
            sa[0] = summary_string.len();
            for i in 0..summary_string.len() {
                let idx = summary_string[i];
                sa[idx as usize + 1] = i;
            }

            sa
        } else {
            suffix_array(summary_string)
        };

        let mut suffix_idx = vec![-1i32; input.len() + 1];
        for i in (2..summary_sa.len()).rev() {
            let idx = summary_index[summary_sa[i] as usize];
            let bucket_idx = *input.get(idx).unwrap_or(&0) as usize;
            suffix_idx[bucket_tails[bucket_idx]] = idx as i32;
            bucket_tails[bucket_idx] -= 1;
        }

        suffix_idx[0] = input.len() as i32;

        suffix_idx
    };

    induced_sort_large(&input, &mut sa, bucket_heads, &type_map);
    induced_sort_small(&input, &mut sa, bucket_tails, &type_map);

    sa.into_iter().map(|i| i as usize).collect()
}
<|endoftext|>


<commit_before>#![warn(bad_style,
        unused, unused_extern_crates, unused_import_braces,
        unused_qualifications, unused_results, unused_typecasts)]

#![feature(core)]

extern crate curl;
extern crate "oauth-client" as oauth;
extern crate rand;

use std::borrow::IntoCow;
use std::collections::HashMap;
use std::str;
use std::string::CowString;
use curl::http;
use curl::http::handle::Method;
use oauth::Token;
use rand::Rng;

mod api {
    pub const REQUEST_TOKEN: &'static str = "http://term.ie/oauth/example/request_token.php";
    pub const ACCESS_TOKEN: &'static str = "http://term.ie/oauth/example/access_token.php";
    pub const ECHO: &'static str = "http://term.ie/oauth/example/echo_api.php";
}

fn split_query<'a>(query: &'a str) -> HashMap<CowString<'a>, CowString<'a>> {
    let mut param = HashMap::new();
    for q in query.split('&') {
        let mut s = q.splitn(2, '=');
        let k = s.next().unwrap();
        let v = s.next().unwrap();
        let _ = param.insert(k.into_cow(), v.into_cow());
    }
    param
}

fn get_request_token(consumer: &Token) -> Token<'static> {
    let header = oauth::authorization_header(Method::Get, api::REQUEST_TOKEN, consumer, None, None);
    let resp = http::handle()
        .get(api::REQUEST_TOKEN)
        .header("Authorization", header.as_slice())
        .exec()
        .unwrap();
    let resp = str::from_utf8(resp.get_body())
        .unwrap()
        .to_string();
    println!("get_request_token response: {:?}", resp);
    let param = split_query(resp.as_slice());
    Token::new(param.get("oauth_token").unwrap().to_string(),
               param.get("oauth_token_secret").unwrap().to_string())
}

fn get_access_token(consumer: &Token, request: &Token) -> Token<'static> {
    let header = oauth::authorization_header(Method::Get, api::ACCESS_TOKEN, consumer, Some(request), None);
    let resp = http::handle()
        .get(api::ACCESS_TOKEN)
        .header("Authorization", header.as_slice())
        .exec()
        .unwrap();
    let resp = str::from_utf8(resp.get_body())
        .unwrap()
        .to_string();
    println!("get_access_token response: {:?}", resp);
    let param = split_query(resp.as_slice());
    Token::new(param.get("oauth_token").unwrap().to_string(),
               param.get("oauth_token_secret").unwrap().to_string())
}

fn echo(consumer: &Token, access: &Token) {
    let header = oauth::authorization_header(Method::Post, api::ECHO, consumer, Some(access), None);

    let mut rng = rand::thread_rng();
    let req_body = rng.gen_ascii_chars().take(100).collect::<String>();
    let resp = http::handle()
        .post(api::ECHO, req_body.as_slice())
        .header("Authorization", header.as_slice())
        .exec()
        .unwrap();
    let resp = str::from_utf8(resp.get_body()).unwrap();
    println!("echo response: {:?}", resp);
    let resp_body = resp.as_slice();
    assert_eq!("", resp_body);
}

fn main() {
    let consumer = Token::new("key", "secret");
    println!("consumer: {:?}", consumer);

    let request = get_request_token(&consumer);
    println!("request: {:?}", request);

    let access = get_access_token(&consumer, &request);
    println!("access: {:?}", access);

    echo(&consumer, &access);

    println!("OK");
}

<commit_msg>update to rustc 1.0.0-nightly (e2fa53e59 2015-03-20) (built 2015-03-21)<commit_after>#![warn(bad_style,
        unused, unused_extern_crates, unused_import_braces,
        unused_qualifications, unused_results, unused_typecasts)]

#![feature(core)]

extern crate curl;
extern crate "oauth-client" as oauth;
extern crate rand;

use std::borrow::{Cow, IntoCow};
use std::collections::HashMap;
use std::str;
use curl::http;
use curl::http::handle::Method;
use oauth::Token;
use rand::Rng;

mod api {
    pub const REQUEST_TOKEN: &'static str = "http://term.ie/oauth/example/request_token.php";
    pub const ACCESS_TOKEN: &'static str = "http://term.ie/oauth/example/access_token.php";
    pub const ECHO: &'static str = "http://term.ie/oauth/example/echo_api.php";
}

fn split_query<'a>(query: &'a str) -> HashMap<Cow<'a, str>, Cow<'a, str>> {
    let mut param = HashMap::new();
    for q in query.split('&') {
        let mut s = q.splitn(2, '=');
        let k = s.next().unwrap();
        let v = s.next().unwrap();
        let _ = param.insert(k.into_cow(), v.into_cow());
    }
    param
}

fn get_request_token(consumer: &Token) -> Token<'static> {
    let header = oauth::authorization_header(Method::Get, api::REQUEST_TOKEN, consumer, None, None);
    let resp = http::handle()
        .get(api::REQUEST_TOKEN)
        .header("Authorization", header.as_slice())
        .exec()
        .unwrap();
    let resp = str::from_utf8(resp.get_body())
        .unwrap()
        .to_string();
    println!("get_request_token response: {:?}", resp);
    let param = split_query(resp.as_slice());
    Token::new(param.get("oauth_token").unwrap().to_string(),
               param.get("oauth_token_secret").unwrap().to_string())
}

fn get_access_token(consumer: &Token, request: &Token) -> Token<'static> {
    let header = oauth::authorization_header(Method::Get, api::ACCESS_TOKEN, consumer, Some(request), None);
    let resp = http::handle()
        .get(api::ACCESS_TOKEN)
        .header("Authorization", header.as_slice())
        .exec()
        .unwrap();
    let resp = str::from_utf8(resp.get_body())
        .unwrap()
        .to_string();
    println!("get_access_token response: {:?}", resp);
    let param = split_query(resp.as_slice());
    Token::new(param.get("oauth_token").unwrap().to_string(),
               param.get("oauth_token_secret").unwrap().to_string())
}

fn echo(consumer: &Token, access: &Token) {
    let header = oauth::authorization_header(Method::Post, api::ECHO, consumer, Some(access), None);

    let mut rng = rand::thread_rng();
    let req_body = rng.gen_ascii_chars().take(100).collect::<String>();
    let resp = http::handle()
        .post(api::ECHO, req_body.as_slice())
        .header("Authorization", header.as_slice())
        .exec()
        .unwrap();
    let resp = str::from_utf8(resp.get_body()).unwrap();
    println!("echo response: {:?}", resp);
    let resp_body = resp.as_slice();
    assert_eq!("", resp_body);
}

fn main() {
    let consumer = Token::new("key", "secret");
    println!("consumer: {:?}", consumer);

    let request = get_request_token(&consumer);
    println!("request: {:?}", request);

    let access = get_access_token(&consumer, &request);
    println!("access: {:?}", access);

    echo(&consumer, &access);

    println!("OK");
}

<|endoftext|>