<commit_before>/// Easy to use Let's Encrypt client to issue and renew TLS certs

extern crate acme_client;
extern crate clap;
extern crate env_logger;

use std::io::{self, Write};
use acme_client::Directory;
use acme_client::error::Result;
use clap::{Arg, App, SubCommand, ArgMatches};

fn main() {
    let matches = App::new(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .usage("letsencrypt-rs sign -D example.org -P /var/www -k domain.key -o domain.crt\
                \n    letsencrypt-rs revoke -K user_or_domain.key -C signed.crt")
        .subcommand(SubCommand::with_name("sign")
            .about("Signs a certificate")
            .display_order(1)
            .arg(Arg::with_name("USER_KEY_PATH")
                .help("User private key path to use it in account registration.")
                .long("user-key")
                .short("U")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN_KEY_PATH")
                .help("Domain private key path to use it in CSR generation.")
                .short("K")
                .long("domain-key")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN")
                .help("Domain name to obtain certificate. You can use more than one domain name.")
                .short("D")
                .long("domain")
                .multiple(true)
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("PUBLIC_DIR")
                .help("Directory to save ACME simple http challenge. This option is required.")
                .short("P")
                .long("public-dir")
                .takes_value(true))
            .arg(Arg::with_name("EMAIL")
                .help("Contact email address (optional).")
                .short("E")
                .long("email")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN_CSR")
                .help("Path to domain certificate signing request.")
                .short("C")
                .long("domain-csr")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_USER_KEY")
                .help("Path to save private user key.")
                .long("save-user-key")
                .short("u")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_DOMAIN_KEY")
                .help("Path to save domain private key.")
                .short("k")
                .long("save-domain-key")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_DOMAIN_CSR")
                .help("Path to save domain certificate signing request.")
                .long("save-csr")
                .short("S")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_SIGNED_CERTIFICATE")
                .help("Path to save signed certificate. Default is STDOUT.")
                .short("o")
                .long("save-crt")
                .takes_value(true))
            .arg(Arg::with_name("CHAIN")
                .help("Chains the signed certificate with Let's Encrypt Authority X3 \
                       (IdenTrust cross-signed) intermediate certificate.")
                .short("c")
                .long("chain")
                .takes_value(false))
            .arg(Arg::with_name("DNS_CHALLENGE")
                 .help("Use DNS challenge instead of HTTP. This option requires user \
                        to generate a TXT record for domain")
                 .short("d")
                 .long("dns")
                 .takes_value(false)))
        .subcommand(SubCommand::with_name("revoke")
            .about("Revokes a signed certificate")
            .display_order(2)
            .arg(Arg::with_name("USER_KEY")
                .help("User or domain private key path.")
                .long("user-key")
                .short("K")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("SIGNED_CRT")
                .help("Path to signed domain certificate to revoke.")
                .long("signed-crt")
                .short("C")
                .required(true)
                .takes_value(true)))
        .arg(Arg::with_name("verbose")
             .help("Show verbose output")
             .short("v")
             .multiple(true))
        .get_matches();

    init_logger(matches.occurrences_of("verbose"));

    let res = if let Some(matches) = matches.subcommand_matches("sign") {
        sign_certificate(matches)
    } else if let Some(matches) = matches.subcommand_matches("revoke") {
        revoke_certificate(matches)
    } else {
        println!("{}", matches.usage());
        Ok(())
    };

    if let Err(e) = res {
        writeln!(io::stderr(), "{}", e).expect("Failed to write stderr");
        ::std::process::exit(1);
    }
}

fn sign_certificate(matches: &ArgMatches) -> Result<()> {
    let directory = Directory::lets_encrypt()?;

    let mut account_registration = directory.account_registration();

    if let Some(email) = matches.value_of("EMAIL") {
        account_registration = account_registration.email(email);
    }

    if let Some(user_key_path) = matches.value_of("USER_KEY_PATH") {
        account_registration = account_registration.pkey_from_file(user_key_path)?;
    }

    let account = account_registration.register()?;
    let domains: Vec<_> = matches.values_of("DOMAIN")
        .ok_or("You need to provide at least one domain name")?
        .collect();
    for domain in &domains {
        let authorization = account.authorization(domain)?;
        if !matches.is_present("DNS_CHALLENGE") {
            let challenge = authorization.get_http_challenge().ok_or("HTTP challenge not found")?;
            challenge.save_key_authorization(matches.value_of("PUBLIC_DIR")
                                                 .ok_or("--public-dir not defined. \
                                                            You need to define a public \
                                                            directory to use http challenge \
                                                            verification")?)?;
            challenge.validate()?;
        } else {
            let challenge = authorization.get_dns_challenge().ok_or("DNS challenge not found")?;
            println!("Please create a TXT record for _acme-challenge.{}: {}\n\
                      Press enter to continue",
                     domain,
                     challenge.signature()?);
            io::stdin().read_line(&mut String::new()).unwrap();
            challenge.validate()?;
        }
    }

    let mut certificate_signer = account.certificate_signer(&domains);

    if let Some(domain_key_path) = matches.value_of("DOMAIN_KEY_PATH") {
        if let Some(csr_path) = matches.value_of("DOMAIN_CSR") {
            certificate_signer = certificate_signer.csr_from_file(domain_key_path, csr_path)?;
        } else {
            certificate_signer = certificate_signer.pkey_from_file(domain_key_path)?;
        }
    }

    let certificate = certificate_signer.sign_certificate()?;
    let signed_certificate_path = matches.value_of("SAVE_SIGNED_CERTIFICATE")
        .ok_or("You need to save signed certificate")?;
    if matches.is_present("CHAIN") {
        certificate.save_signed_certificate_and_chain(None, signed_certificate_path)?;
    } else {
        certificate.save_signed_certificate(signed_certificate_path)?;
    }

    if let Some(path) = matches.value_of("SAVE_DOMAIN_KEY") {
        certificate.save_private_key(path)?;
    }
    if let Some(path) = matches.value_of("SAVE_DOMAIN_CSR") {
        certificate.save_csr(path)?;
    }
    if let Some(path) = matches.value_of("SAVE_USER_KEY") {
        account.save_private_key(path)?;
    }

    Ok(())
}

fn revoke_certificate(matches: &ArgMatches) -> Result<()> {
    let directory = Directory::lets_encrypt()?;
    let account = directory.account_registration()
        .pkey_from_file(matches.value_of("USER_KEY")
                            .ok_or("You need to provide user \
                                   or domain private key used \
                                   to sign certificate.")?)?
        .register()?;
    account.revoke_certificate_from_file(matches.value_of("SIGNED_CRT")
                                             .ok_or("You need to provide \
                                                    a signed certificate to \
                                                    revoke.")?)?;
    Ok(())
}

fn init_logger(level: u64) {
    let level = match level {
        0 => "",
        1 => "acme_client=info",
        _ => "acme_client=debug",
    };
    let mut builder = env_logger::LogBuilder::new();
    builder.parse(&::std::env::var("RUST_LOG").unwrap_or(level.to_owned()));
    let _ = builder.init();
}
<commit_msg>eprintln for stderr<commit_after>/// Easy to use Let's Encrypt client to issue and renew TLS certs

extern crate acme_client;
extern crate clap;
extern crate env_logger;

use std::io::{self, Write};
use acme_client::Directory;
use acme_client::error::Result;
use clap::{Arg, App, SubCommand, ArgMatches};

fn main() {
    let matches = App::new(env!("CARGO_PKG_NAME"))
        .version(env!("CARGO_PKG_VERSION"))
        .about(env!("CARGO_PKG_DESCRIPTION"))
        .usage("letsencrypt-rs sign -D example.org -P /var/www -k domain.key -o domain.crt\
                \n    letsencrypt-rs revoke -K user_or_domain.key -C signed.crt")
        .subcommand(SubCommand::with_name("sign")
            .about("Signs a certificate")
            .display_order(1)
            .arg(Arg::with_name("USER_KEY_PATH")
                .help("User private key path to use it in account registration.")
                .long("user-key")
                .short("U")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN_KEY_PATH")
                .help("Domain private key path to use it in CSR generation.")
                .short("K")
                .long("domain-key")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN")
                .help("Domain name to obtain certificate. You can use more than one domain name.")
                .short("D")
                .long("domain")
                .multiple(true)
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("PUBLIC_DIR")
                .help("Directory to save ACME simple http challenge. This option is required.")
                .short("P")
                .long("public-dir")
                .takes_value(true))
            .arg(Arg::with_name("EMAIL")
                .help("Contact email address (optional).")
                .short("E")
                .long("email")
                .takes_value(true))
            .arg(Arg::with_name("DOMAIN_CSR")
                .help("Path to domain certificate signing request.")
                .short("C")
                .long("domain-csr")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_USER_KEY")
                .help("Path to save private user key.")
                .long("save-user-key")
                .short("u")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_DOMAIN_KEY")
                .help("Path to save domain private key.")
                .short("k")
                .long("save-domain-key")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_DOMAIN_CSR")
                .help("Path to save domain certificate signing request.")
                .long("save-csr")
                .short("S")
                .takes_value(true))
            .arg(Arg::with_name("SAVE_SIGNED_CERTIFICATE")
                .help("Path to save signed certificate. Default is STDOUT.")
                .short("o")
                .long("save-crt")
                .takes_value(true))
            .arg(Arg::with_name("CHAIN")
                .help("Chains the signed certificate with Let's Encrypt Authority X3 \
                       (IdenTrust cross-signed) intermediate certificate.")
                .short("c")
                .long("chain")
                .takes_value(false))
            .arg(Arg::with_name("DNS_CHALLENGE")
                 .help("Use DNS challenge instead of HTTP. This option requires user \
                        to generate a TXT record for domain")
                 .short("d")
                 .long("dns")
                 .takes_value(false)))
        .subcommand(SubCommand::with_name("revoke")
            .about("Revokes a signed certificate")
            .display_order(2)
            .arg(Arg::with_name("USER_KEY")
                .help("User or domain private key path.")
                .long("user-key")
                .short("K")
                .required(true)
                .takes_value(true))
            .arg(Arg::with_name("SIGNED_CRT")
                .help("Path to signed domain certificate to revoke.")
                .long("signed-crt")
                .short("C")
                .required(true)
                .takes_value(true)))
        .arg(Arg::with_name("verbose")
             .help("Show verbose output")
             .short("v")
             .multiple(true))
        .get_matches();

    init_logger(matches.occurrences_of("verbose"));

    let res = if let Some(matches) = matches.subcommand_matches("sign") {
        sign_certificate(matches)
    } else if let Some(matches) = matches.subcommand_matches("revoke") {
        revoke_certificate(matches)
    } else {
        println!("{}", matches.usage());
        Ok(())
    };

    if let Err(e) = res {
        eprintln!("{}", e);
        ::std::process::exit(1);
    }
}

fn sign_certificate(matches: &ArgMatches) -> Result<()> {
    let directory = Directory::lets_encrypt()?;

    let mut account_registration = directory.account_registration();

    if let Some(email) = matches.value_of("EMAIL") {
        account_registration = account_registration.email(email);
    }

    if let Some(user_key_path) = matches.value_of("USER_KEY_PATH") {
        account_registration = account_registration.pkey_from_file(user_key_path)?;
    }

    let account = account_registration.register()?;
    let domains: Vec<_> = matches.values_of("DOMAIN")
        .ok_or("You need to provide at least one domain name")?
        .collect();
    for domain in &domains {
        let authorization = account.authorization(domain)?;
        if !matches.is_present("DNS_CHALLENGE") {
            let challenge = authorization.get_http_challenge().ok_or("HTTP challenge not found")?;
            challenge.save_key_authorization(matches.value_of("PUBLIC_DIR")
                                                 .ok_or("--public-dir not defined. \
                                                            You need to define a public \
                                                            directory to use http challenge \
                                                            verification")?)?;
            challenge.validate()?;
        } else {
            let challenge = authorization.get_dns_challenge().ok_or("DNS challenge not found")?;
            println!("Please create a TXT record for _acme-challenge.{}: {}\n\
                      Press enter to continue",
                     domain,
                     challenge.signature()?);
            io::stdin().read_line(&mut String::new()).unwrap();
            challenge.validate()?;
        }
    }

    let mut certificate_signer = account.certificate_signer(&domains);

    if let Some(domain_key_path) = matches.value_of("DOMAIN_KEY_PATH") {
        if let Some(csr_path) = matches.value_of("DOMAIN_CSR") {
            certificate_signer = certificate_signer.csr_from_file(domain_key_path, csr_path)?;
        } else {
            certificate_signer = certificate_signer.pkey_from_file(domain_key_path)?;
        }
    }

    let certificate = certificate_signer.sign_certificate()?;
    let signed_certificate_path = matches.value_of("SAVE_SIGNED_CERTIFICATE")
        .ok_or("You need to save signed certificate")?;
    if matches.is_present("CHAIN") {
        certificate.save_signed_certificate_and_chain(None, signed_certificate_path)?;
    } else {
        certificate.save_signed_certificate(signed_certificate_path)?;
    }

    if let Some(path) = matches.value_of("SAVE_DOMAIN_KEY") {
        certificate.save_private_key(path)?;
    }
    if let Some(path) = matches.value_of("SAVE_DOMAIN_CSR") {
        certificate.save_csr(path)?;
    }
    if let Some(path) = matches.value_of("SAVE_USER_KEY") {
        account.save_private_key(path)?;
    }

    Ok(())
}

fn revoke_certificate(matches: &ArgMatches) -> Result<()> {
    let directory = Directory::lets_encrypt()?;
    let account = directory.account_registration()
        .pkey_from_file(matches.value_of("USER_KEY")
                            .ok_or("You need to provide user \
                                   or domain private key used \
                                   to sign certificate.")?)?
        .register()?;
    account.revoke_certificate_from_file(matches.value_of("SIGNED_CRT")
                                             .ok_or("You need to provide \
                                                    a signed certificate to \
                                                    revoke.")?)?;
    Ok(())
}

fn init_logger(level: u64) {
    let level = match level {
        0 => "",
        1 => "acme_client=info",
        _ => "acme_client=debug",
    };
    let mut builder = env_logger::LogBuilder::new();
    builder.parse(&::std::env::var("RUST_LOG").unwrap_or(level.to_owned()));
    let _ = builder.init();
}
<|endoftext|>


<commit_before>// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::fmt;
use std::sync::{Arc, Mutex};

pub use devices::virtio::balloon::device::BalloonStats;
use devices::virtio::balloon::Error as BalloonError;
pub use devices::virtio::BALLOON_DEV_ID;
use devices::virtio::{Balloon, BalloonConfig};

use serde::{Deserialize, Serialize};

type MutexBalloon = Arc<Mutex<Balloon>>;

/// Errors associated with the operations allowed on the balloon.
#[derive(Debug)]
pub enum BalloonConfigError {
    /// The user made a request on an inexistent balloon device.
    DeviceNotFound,
    /// Device not activated yet.
    DeviceNotActive,
    /// The user tried to enable/disable the statistics after boot.
    InvalidStatsUpdate,
    /// Amount of pages requested is too large.
    TooManyPagesRequested,
    /// The user polled the statistics of a balloon device that
    /// does not have the statistics enabled.
    StatsNotFound,
    /// Failed to create a balloon device.
    CreateFailure(devices::virtio::balloon::Error),
    /// Failed to update the configuration of the ballon device.
    UpdateFailure(std::io::Error),
}

impl fmt::Display for BalloonConfigError {
    fn fmt(&self, f: &mut fmt::Formatter) -> std::fmt::Result {
        use self::BalloonConfigError::*;
        match self {
            DeviceNotFound => write!(f, "No balloon device found."),
            DeviceNotActive => write!(f, "Device is inactive, check balloon driver is enabled."),
            InvalidStatsUpdate => write!(f, "Cannot enable/disable the statistics after boot."),
            TooManyPagesRequested => write!(f, "Amount of pages requested is too large."),
            StatsNotFound => write!(f, "Statistics for the balloon device are not enabled"),
            CreateFailure(e) => write!(f, "Error creating the balloon device: {:?}", e),
            UpdateFailure(e) => write!(
                f,
                "Error updating the balloon device configuration: {:?}",
                e
            ),
        }
    }
}

impl From<BalloonError> for BalloonConfigError {
    fn from(error: BalloonError) -> Self {
        match error {
            BalloonError::DeviceNotFound => Self::DeviceNotFound,
            BalloonError::DeviceNotActive => Self::DeviceNotActive,
            BalloonError::InterruptError(io_error) => Self::UpdateFailure(io_error),
            BalloonError::StatisticsStateChange => Self::InvalidStatsUpdate,
            BalloonError::StatisticsDisabled => Self::StatsNotFound,
            BalloonError::TooManyPagesRequested => Self::TooManyPagesRequested,
            e => Self::CreateFailure(e),
        }
    }
}

type Result<T> = std::result::Result<T, BalloonConfigError>;

/// This struct represents the strongly typed equivalent of the json body
/// from balloon related requests.
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct BalloonDeviceConfig {
    /// Target balloon size in MB.
    pub amount_mb: u32,
    /// Option to deflate the balloon in case the guest is out of memory.
    pub deflate_on_oom: bool,
    /// Interval in seconds between refreshing statistics.
    #[serde(default)]
    pub stats_polling_interval_s: u16,
}

impl From<BalloonConfig> for BalloonDeviceConfig {
    fn from(state: BalloonConfig) -> Self {
        BalloonDeviceConfig {
            amount_mb: state.amount_mb,
            deflate_on_oom: state.deflate_on_oom,
            stats_polling_interval_s: state.stats_polling_interval_s,
        }
    }
}

/// The data fed into a balloon update request. Currently, only the number
/// of pages and the stats polling interval can be updated.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct BalloonUpdateConfig {
    /// Target balloon size in MB.
    pub amount_mb: u32,
}

/// The data fed into a balloon statistics interval update request.
/// Note that the state of the statistics cannot be changed from ON to OFF
/// or vice versa after boot, only the interval of polling can be changed
/// if the statistics were activated in the device configuration.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct BalloonUpdateStatsConfig {
    /// Interval in seconds between refreshing statistics.
    pub stats_polling_interval_s: u16,
}

/// A builder for `Balloon` devices from 'BalloonDeviceConfig'.
pub struct BalloonBuilder {
    inner: Option<MutexBalloon>,
}

#[cfg(not(test))]
impl Default for BalloonBuilder {
    fn default() -> BalloonBuilder {
        BalloonBuilder { inner: None }
    }
}

impl BalloonBuilder {
    /// Creates an empty Balloon Store.
    pub fn new() -> Self {
        Self { inner: None }
    }

    /// Inserts a Balloon device in the store.
    /// If an entry already exists, it will overwrite it.
    pub fn set(&mut self, cfg: BalloonDeviceConfig) -> Result<()> {
        self.inner = Some(Arc::new(Mutex::new(
            Balloon::new(
                cfg.amount_mb,
                cfg.deflate_on_oom,
                cfg.stats_polling_interval_s,
                // `restored` flag is false because this code path
                // is never called by snapshot restore functionality.
                false,
            )
            .map_err(BalloonConfigError::CreateFailure)?,
        )));

        Ok(())
    }

    /// Provides a reference to the Balloon if present.
    pub fn get(&self) -> Option<&MutexBalloon> {
        self.inner.as_ref()
    }

    /// Returns the same structure that was used to configure the device.
    pub fn get_config(&self) -> Result<BalloonDeviceConfig> {
        self.get()
            .ok_or(BalloonConfigError::DeviceNotFound)
            .map(|balloon_mutex| balloon_mutex.lock().expect("Poisoned lock").config())
            .map(BalloonDeviceConfig::from)
    }
}

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

    pub(crate) fn default_config() -> BalloonDeviceConfig {
        BalloonDeviceConfig {
            amount_mb: 0,
            deflate_on_oom: false,
            stats_polling_interval_s: 0,
        }
    }

    impl Default for BalloonBuilder {
        fn default() -> BalloonBuilder {
            let mut balloon = BalloonBuilder::new();
            assert!(balloon.set(BalloonDeviceConfig::default()).is_ok());
            balloon
        }
    }

    #[test]
    fn test_balloon_create() {
        let default_balloon_config = default_config();
        let balloon_config = BalloonDeviceConfig {
            amount_mb: 0,
            deflate_on_oom: false,
            stats_polling_interval_s: 0,
        };
        assert_eq!(default_balloon_config, balloon_config);
        let mut builder = BalloonBuilder::new();
        assert!(builder.get().is_none());

        builder.set(balloon_config).unwrap();
        assert_eq!(builder.get().unwrap().lock().unwrap().num_pages(), 0);
        assert_eq!(builder.get_config().unwrap(), default_balloon_config);

        let _update_config = BalloonUpdateConfig { amount_mb: 5 };
        let _stats_update_config = BalloonUpdateStatsConfig {
            stats_polling_interval_s: 5,
        };
    }

    #[test]
    fn test_from_balloon_state() {
        let expected_balloon_config = BalloonDeviceConfig {
            amount_mb: 5,
            deflate_on_oom: false,
            stats_polling_interval_s: 3,
        };

        let actual_balloon_config = BalloonDeviceConfig::from(BalloonConfig {
            amount_mb: 5,
            deflate_on_oom: false,
            stats_polling_interval_s: 3,
        });

        assert_eq!(expected_balloon_config, actual_balloon_config);
    }

    #[test]
    fn test_error_messages() {
        use super::BalloonConfigError::*;
        use std::io;
        let err = CreateFailure(devices::virtio::balloon::Error::EventFd(
            io::Error::from_raw_os_error(0),
        ));
        let _ = format!("{}{:?}", err, err);

        let err = UpdateFailure(io::Error::from_raw_os_error(0));
        let _ = format!("{}{:?}", err, err);

        let err = DeviceNotFound;
        let _ = format!("{}{:?}", err, err);

        let err = InvalidStatsUpdate;
        let _ = format!("{}{:?}", err, err);

        let err = TooManyPagesRequested;
        let _ = format!("{}{:?}", err, err);

        let err = StatsNotFound;
        let _ = format!("{}{:?}", err, err);
    }
}
<commit_msg>balloon: change wording of error message<commit_after>// Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::fmt;
use std::sync::{Arc, Mutex};

pub use devices::virtio::balloon::device::BalloonStats;
use devices::virtio::balloon::Error as BalloonError;
pub use devices::virtio::BALLOON_DEV_ID;
use devices::virtio::{Balloon, BalloonConfig};

use serde::{Deserialize, Serialize};

type MutexBalloon = Arc<Mutex<Balloon>>;

/// Errors associated with the operations allowed on the balloon.
#[derive(Debug)]
pub enum BalloonConfigError {
    /// The user made a request on an inexistent balloon device.
    DeviceNotFound,
    /// Device not activated yet.
    DeviceNotActive,
    /// The user tried to enable/disable the statistics after boot.
    InvalidStatsUpdate,
    /// Amount of pages requested is too large.
    TooManyPagesRequested,
    /// The user polled the statistics of a balloon device that
    /// does not have the statistics enabled.
    StatsNotFound,
    /// Failed to create a balloon device.
    CreateFailure(devices::virtio::balloon::Error),
    /// Failed to update the configuration of the ballon device.
    UpdateFailure(std::io::Error),
}

impl fmt::Display for BalloonConfigError {
    fn fmt(&self, f: &mut fmt::Formatter) -> std::fmt::Result {
        use self::BalloonConfigError::*;
        match self {
            DeviceNotFound => write!(f, "No balloon device found."),
            DeviceNotActive => write!(
                f,
                "Device is inactive, check if balloon driver is enabled in guest kernel."
            ),
            InvalidStatsUpdate => write!(f, "Cannot enable/disable the statistics after boot."),
            TooManyPagesRequested => write!(f, "Amount of pages requested is too large."),
            StatsNotFound => write!(f, "Statistics for the balloon device are not enabled"),
            CreateFailure(e) => write!(f, "Error creating the balloon device: {:?}", e),
            UpdateFailure(e) => write!(
                f,
                "Error updating the balloon device configuration: {:?}",
                e
            ),
        }
    }
}

impl From<BalloonError> for BalloonConfigError {
    fn from(error: BalloonError) -> Self {
        match error {
            BalloonError::DeviceNotFound => Self::DeviceNotFound,
            BalloonError::DeviceNotActive => Self::DeviceNotActive,
            BalloonError::InterruptError(io_error) => Self::UpdateFailure(io_error),
            BalloonError::StatisticsStateChange => Self::InvalidStatsUpdate,
            BalloonError::StatisticsDisabled => Self::StatsNotFound,
            BalloonError::TooManyPagesRequested => Self::TooManyPagesRequested,
            e => Self::CreateFailure(e),
        }
    }
}

type Result<T> = std::result::Result<T, BalloonConfigError>;

/// This struct represents the strongly typed equivalent of the json body
/// from balloon related requests.
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct BalloonDeviceConfig {
    /// Target balloon size in MB.
    pub amount_mb: u32,
    /// Option to deflate the balloon in case the guest is out of memory.
    pub deflate_on_oom: bool,
    /// Interval in seconds between refreshing statistics.
    #[serde(default)]
    pub stats_polling_interval_s: u16,
}

impl From<BalloonConfig> for BalloonDeviceConfig {
    fn from(state: BalloonConfig) -> Self {
        BalloonDeviceConfig {
            amount_mb: state.amount_mb,
            deflate_on_oom: state.deflate_on_oom,
            stats_polling_interval_s: state.stats_polling_interval_s,
        }
    }
}

/// The data fed into a balloon update request. Currently, only the number
/// of pages and the stats polling interval can be updated.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct BalloonUpdateConfig {
    /// Target balloon size in MB.
    pub amount_mb: u32,
}

/// The data fed into a balloon statistics interval update request.
/// Note that the state of the statistics cannot be changed from ON to OFF
/// or vice versa after boot, only the interval of polling can be changed
/// if the statistics were activated in the device configuration.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct BalloonUpdateStatsConfig {
    /// Interval in seconds between refreshing statistics.
    pub stats_polling_interval_s: u16,
}

/// A builder for `Balloon` devices from 'BalloonDeviceConfig'.
pub struct BalloonBuilder {
    inner: Option<MutexBalloon>,
}

#[cfg(not(test))]
impl Default for BalloonBuilder {
    fn default() -> BalloonBuilder {
        BalloonBuilder { inner: None }
    }
}

impl BalloonBuilder {
    /// Creates an empty Balloon Store.
    pub fn new() -> Self {
        Self { inner: None }
    }

    /// Inserts a Balloon device in the store.
    /// If an entry already exists, it will overwrite it.
    pub fn set(&mut self, cfg: BalloonDeviceConfig) -> Result<()> {
        self.inner = Some(Arc::new(Mutex::new(
            Balloon::new(
                cfg.amount_mb,
                cfg.deflate_on_oom,
                cfg.stats_polling_interval_s,
                // `restored` flag is false because this code path
                // is never called by snapshot restore functionality.
                false,
            )
            .map_err(BalloonConfigError::CreateFailure)?,
        )));

        Ok(())
    }

    /// Provides a reference to the Balloon if present.
    pub fn get(&self) -> Option<&MutexBalloon> {
        self.inner.as_ref()
    }

    /// Returns the same structure that was used to configure the device.
    pub fn get_config(&self) -> Result<BalloonDeviceConfig> {
        self.get()
            .ok_or(BalloonConfigError::DeviceNotFound)
            .map(|balloon_mutex| balloon_mutex.lock().expect("Poisoned lock").config())
            .map(BalloonDeviceConfig::from)
    }
}

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

    pub(crate) fn default_config() -> BalloonDeviceConfig {
        BalloonDeviceConfig {
            amount_mb: 0,
            deflate_on_oom: false,
            stats_polling_interval_s: 0,
        }
    }

    impl Default for BalloonBuilder {
        fn default() -> BalloonBuilder {
            let mut balloon = BalloonBuilder::new();
            assert!(balloon.set(BalloonDeviceConfig::default()).is_ok());
            balloon
        }
    }

    #[test]
    fn test_balloon_create() {
        let default_balloon_config = default_config();
        let balloon_config = BalloonDeviceConfig {
            amount_mb: 0,
            deflate_on_oom: false,
            stats_polling_interval_s: 0,
        };
        assert_eq!(default_balloon_config, balloon_config);
        let mut builder = BalloonBuilder::new();
        assert!(builder.get().is_none());

        builder.set(balloon_config).unwrap();
        assert_eq!(builder.get().unwrap().lock().unwrap().num_pages(), 0);
        assert_eq!(builder.get_config().unwrap(), default_balloon_config);

        let _update_config = BalloonUpdateConfig { amount_mb: 5 };
        let _stats_update_config = BalloonUpdateStatsConfig {
            stats_polling_interval_s: 5,
        };
    }

    #[test]
    fn test_from_balloon_state() {
        let expected_balloon_config = BalloonDeviceConfig {
            amount_mb: 5,
            deflate_on_oom: false,
            stats_polling_interval_s: 3,
        };

        let actual_balloon_config = BalloonDeviceConfig::from(BalloonConfig {
            amount_mb: 5,
            deflate_on_oom: false,
            stats_polling_interval_s: 3,
        });

        assert_eq!(expected_balloon_config, actual_balloon_config);
    }

    #[test]
    fn test_error_messages() {
        use super::BalloonConfigError::*;
        use std::io;
        let err = CreateFailure(devices::virtio::balloon::Error::EventFd(
            io::Error::from_raw_os_error(0),
        ));
        let _ = format!("{}{:?}", err, err);

        let err = UpdateFailure(io::Error::from_raw_os_error(0));
        let _ = format!("{}{:?}", err, err);

        let err = DeviceNotFound;
        let _ = format!("{}{:?}", err, err);

        let err = InvalidStatsUpdate;
        let _ = format!("{}{:?}", err, err);

        let err = TooManyPagesRequested;
        let _ = format!("{}{:?}", err, err);

        let err = StatsNotFound;
        let _ = format!("{}{:?}", err, err);
    }
}
<|endoftext|>


<commit_before>/* kernel::sgash.rs */
#[allow(unused_imports)];
use core::*;
use core::str::*;
use core::option::{Some, Option, None};
use core::iter::Iterator;
use kernel::*;
use kernel::cstr::cstr;
use super::super::platform::*;
use kernel::memory::Allocator;

static PROMPT: &'static str = &"sgash> ";
static PROMPT_COLOR: u32 = 0xFFAF00;
static mut count: uint = 0;

static mut buffer: cstr = cstr {
    p: 0 as *mut u8,
    p_cstr_i: 0,
    size: 0,
};

static mut history: cstr = cstr {
    p: 0 as *mut u8,
    p_cstr_i: 0,
    size: 0,
};

/* Thanks to https://github.com/jvns/puddle/blob/master/src/stdio.rs */
pub fn putnum(x: uint, max: uint) {
    let mut i = max;
    if i == 0 {
        i = 1000000;
    }
    while(i > 0) {
        // Get the offset from each decimal point up to a certain value
        putchar(((((x / i) % 10) as u8) + ('0' as u8)) as char);
        i /= 10;
    }
    putchar(' ');
}

pub fn putkeycode(x: u8) {
    putnum(x as uint, 100);
}

pub fn putchar(key: char) {
    unsafe {
        /*
        * We need to include a blank asm call to prevent rustc
        * from optimizing this part out
        */
        asm!("");
        io::write_char(key, io::UART0);
    }
}

pub unsafe fn showstr(msg: &str) {
    putstr(msg);
    drawstr(msg);
}

pub fn putstr(msg: &str) {
    for c in slice::iter(as_bytes(msg)) {
        putchar(*c as char);
    }
}

pub unsafe fn drawstr(msg: &str) {
    for c in slice::iter(as_bytes(msg)) {
        drawchar(*c as char);
    }
}

unsafe fn drawchar(x: char) {
    io::restore();
    if ((x as u8) as uint) == 10 || x == '\n' {
        io::CURSOR_Y += io::CURSOR_HEIGHT;
        io::CURSOR_X = 0u32;
    } else {
        if (((x as u8) as uint) >= 32 && ((x as u8) as uint) <= 125) {
            io::draw_char(x);
            io::CURSOR_X += io::CURSOR_WIDTH;
        }
    }
    io::backup();
    io::draw_cursor();
}

unsafe fn backspace() {
    io::restore();
    if (io::CURSOR_X >= io::CURSOR_WIDTH) {
        io::CURSOR_X -= io::CURSOR_WIDTH;
        io::draw_char(' ');
    }
    io::backup();
    io::draw_cursor();
}

pub unsafe fn parsekey(x: char) {
    let x = x as u8;
    // Set this to false to learn the keycodes of various keys!
    // Key codes are printed backwards because life is hard

    match x {
        13 =>  {
            parse_buffer();
            prompt();
            buffer.reset();
        }
        127 =>  {
            putchar('');
            putchar(' ');
            putchar('');
            buffer.delete_char();
            backspace();
        }
        _ =>  {
            buffer.add_u8(x);
            if io::CURSOR_X < io::SCREEN_WIDTH-io::CURSOR_WIDTH {
                putchar(x as char);
                drawchar(x as char);
            }
        }
    }
}

pub unsafe fn parse_buffer() {
    showstr(&"\n");
    buffer.map(putchar);
    buffer.map(drawchar);
    showstr(&"\n");
    if buffer.streq(&"history") {
        history.map(drawchar);
        history.map(putchar);
    }
    else {
        let nhistory = history.join(buffer);
        history = nhistory;
        heap.free(history.p);
        history.add_char('\n');
    }
    let (command, args) = buffer.split(' ');
    if command.streq(&"echo") {
        args.map(drawchar);
        args.map(putchar);
        showstr(&"\n");
    }
}

pub unsafe fn screen() {
    putstr(&"\n                                                               ");
    putstr(&"\n                                                               ");
    putstr(&"\n                       7=..~$=..:7                             ");
    putstr(&"\n                  +$: =$$$+$$$?$$$+ ,7?                        ");
    putstr(&"\n                  $$$$$$$$$$$$$$$$$$Z$$                        ");
    putstr(&"\n              7$$$$$$$$$$$$. .Z$$$$$Z$$$$$$                    ");
    putstr(&"\n           ~..7$$Z$$$$$7+7$+.?Z7=7$$Z$$Z$$$..:                 ");
    putstr(&"\n          ~$$$$$$$$7:     :ZZZ,     :7ZZZZ$$$$=                ");
    putstr(&"\n           Z$$$$$?                    .+ZZZZ$$                 ");
    putstr(&"\n       +$ZZ$$$Z7                         7ZZZ$Z$$I.            ");
    putstr(&"\n        $$$$ZZZZZZZZZZZZZZZZZZZZZZZZI,    ,ZZZ$$Z              ");
    putstr(&"\n      :+$$$$ZZZZZZZZZZZZZZZZZZZZZZZZZZZ=    $ZZ$$+~,           ");
    putstr(&"\n     ?$Z$$$$ZZZZZZZZZZZZZZZZZZZZZZZZZZZZI   7ZZZ$ZZI           ");
    putstr(&"\n      =Z$$+7Z$$7ZZZZZZZZ$$$$$$$ZZZZZZZZZZ  ~Z$?$ZZ?            ");
    putstr(&"\n    :$Z$Z...$Z  $ZZZZZZZ~       ~ZZZZZZZZ,.ZZ...Z$Z$~          ");
    putstr(&"\n    7ZZZZZI$ZZ  $ZZZZZZZ~       =ZZZZZZZ7..ZZ$?$ZZZZ$          ");
    putstr(&"\n      ZZZZ$:    $ZZZZZZZZZZZZZZZZZZZZZZ=     ~$ZZZ$:           ");
    putstr(&"\n    7Z$ZZ$,     $ZZZZZZZZZZZZZZZZZZZZ7         ZZZ$Z$          ");
    putstr(&"\n   =ZZZZZZ,     $ZZZZZZZZZZZZZZZZZZZZZZ,       ZZZ$ZZ+         ");
    putstr(&"\n     ,ZZZZ,     $ZZZZZZZ:     =ZZZZZZZZZ     ZZZZZ$:           ");
    putstr(&"\n    =$ZZZZ+     ZZZZZZZZ~       ZZZZZZZZ~   =ZZZZZZZI          ");
    putstr(&"\n    $ZZ$ZZZ$$Z$$ZZZZZZZZZ$$$$   IZZZZZZZZZ$ZZZZZZZZZ$          ");
    putstr(&"\n      :ZZZZZZZZZZZZZZZZZZZZZZ   ~ZZZZZZZZZZZZZZZZZ~            ");
    putstr(&"\n     ,Z$$ZZZZZZZZZZZZZZZZZZZZ    ZZZZZZZZZZZZZZZZZZ~           ");
    putstr(&"\n     =$ZZZZZZZZZZZZZZZZZZZZZZ     $ZZZZZZZZZZZZZZZ$+           ");
    putstr(&"\n        IZZZZZ:.                        . ,ZZZZZ$              ");
    putstr(&"\n       ~$ZZZZZZZZZZZ                 ZZZZ$ZZZZZZZ+             ");
    putstr(&"\n           Z$ZZZ. ,Z~               =Z:.,ZZZ$Z                 ");
    putstr(&"\n          ,ZZZZZ..~Z$.             .7Z:..ZZZZZ:                ");
    putstr(&"\n          ~7+:$ZZZZZZZZI=:.   .,=IZZZZZZZ$Z:=7=                ");
    putstr(&"\n              $$ZZZZZZZZZZZZZZZZZZZZZZ$ZZZZ                    ");
    putstr(&"\n              ==..$ZZZ$ZZZZZZZZZZZ$ZZZZ .~+                    ");
    putstr(&"\n                  I$?.?ZZZ$ZZZ$ZZZI =$7                        ");
    putstr(&"\n                       $7..I$7..I$,                            ");
    putstr(&"\n");
    putstr(&"\n _                     _     _                         _  ");
    putstr(&"\n| |                   (_)   | |                       | | ");
    putstr(&"\n| | ____ ___  ____     _____| |_____  ____ ____  _____| | ");
    putstr(&"\n| |/ ___) _ \\|  _ \\   |  _   _) ___ |/ ___)  _ \\| ___ | | ");
    putstr(&"\n| | |  | |_| | | | |  | |  \\ \\| ____| |   | | | | ____| | ");
    putstr(&"\n|_|_|  \\____/|_| |_|  |_|   \\_\\_____)_|   |_| |_|_____)__)\n\n");
}

pub unsafe fn init() {
    buffer = cstr::new();
    history = cstr::news(512);
    screen();
    prompt();
}

pub unsafe fn prompt() {
    putstr(PROMPT);
    let prev_c = super::super::io::FG_COLOR;
    super::super::io::set_fg(PROMPT_COLOR);
    drawstr(PROMPT);
    super::super::io::set_fg(prev_c);
}

<commit_msg>only push to history if command is nonempty<commit_after>/* kernel::sgash.rs */
#[allow(unused_imports)];
use core::*;
use core::str::*;
use core::option::{Some, Option, None};
use core::iter::Iterator;
use kernel::*;
use kernel::cstr::cstr;
use super::super::platform::*;
use kernel::memory::Allocator;

static PROMPT: &'static str = &"sgash> ";
static PROMPT_COLOR: u32 = 0xFFAF00;
static mut count: uint = 0;

static mut buffer: cstr = cstr {
    p: 0 as *mut u8,
    p_cstr_i: 0,
    size: 0,
};

static mut history: cstr = cstr {
    p: 0 as *mut u8,
    p_cstr_i: 0,
    size: 0,
};

/* Thanks to https://github.com/jvns/puddle/blob/master/src/stdio.rs */
pub fn putnum(x: uint, max: uint) {
    let mut i = max;
    if i == 0 {
        i = 1000000;
    }
    while(i > 0) {
        // Get the offset from each decimal point up to a certain value
        putchar(((((x / i) % 10) as u8) + ('0' as u8)) as char);
        i /= 10;
    }
    putchar(' ');
}

pub fn putkeycode(x: u8) {
    putnum(x as uint, 100);
}

pub fn putchar(key: char) {
    unsafe {
        /*
        * We need to include a blank asm call to prevent rustc
        * from optimizing this part out
        */
        asm!("");
        io::write_char(key, io::UART0);
    }
}

pub unsafe fn showstr(msg: &str) {
    putstr(msg);
    drawstr(msg);
}

pub fn putstr(msg: &str) {
    for c in slice::iter(as_bytes(msg)) {
        putchar(*c as char);
    }
}

pub unsafe fn drawstr(msg: &str) {
    for c in slice::iter(as_bytes(msg)) {
        drawchar(*c as char);
    }
}

unsafe fn drawchar(x: char) {
    io::restore();
    if ((x as u8) as uint) == 10 || x == '\n' {
        io::CURSOR_Y += io::CURSOR_HEIGHT;
        io::CURSOR_X = 0u32;
    } else {
        if (((x as u8) as uint) >= 32 && ((x as u8) as uint) <= 125) {
            io::draw_char(x);
            io::CURSOR_X += io::CURSOR_WIDTH;
        }
    }
    io::backup();
    io::draw_cursor();
}

unsafe fn backspace() {
    io::restore();
    if (io::CURSOR_X >= io::CURSOR_WIDTH) {
        io::CURSOR_X -= io::CURSOR_WIDTH;
        io::draw_char(' ');
    }
    io::backup();
    io::draw_cursor();
}

pub unsafe fn parsekey(x: char) {
    let x = x as u8;
    // Set this to false to learn the keycodes of various keys!
    // Key codes are printed backwards because life is hard

    match x {
        13 =>  {
            parse_buffer();
            prompt();
            buffer.reset();
        }
        127 =>  {
            putchar('');
            putchar(' ');
            putchar('');
            buffer.delete_char();
            backspace();
        }
        _ =>  {
            buffer.add_u8(x);
            if io::CURSOR_X < io::SCREEN_WIDTH-io::CURSOR_WIDTH {
                putchar(x as char);
                drawchar(x as char);
            }
        }
    }
}

pub unsafe fn parse_buffer() {
    showstr(&"\n");
    buffer.map(putchar);
    buffer.map(drawchar);
    showstr(&"\n");
    if buffer.streq(&"history") {
        history.map(drawchar);
        history.map(putchar);
    }
    else if buffer.len() > 0 {
        let nhistory = history.join(buffer);
        history = nhistory;
        heap.free(history.p);
        history.add_char('\n');
    }
    let (command, args) = buffer.split(' ');
    if command.streq(&"echo") {
        args.map(drawchar);
        args.map(putchar);
        showstr(&"\n");
    }
}

pub unsafe fn screen() {
    putstr(&"\n                                                               ");
    putstr(&"\n                                                               ");
    putstr(&"\n                       7=..~$=..:7                             ");
    putstr(&"\n                  +$: =$$$+$$$?$$$+ ,7?                        ");
    putstr(&"\n                  $$$$$$$$$$$$$$$$$$Z$$                        ");
    putstr(&"\n              7$$$$$$$$$$$$. .Z$$$$$Z$$$$$$                    ");
    putstr(&"\n           ~..7$$Z$$$$$7+7$+.?Z7=7$$Z$$Z$$$..:                 ");
    putstr(&"\n          ~$$$$$$$$7:     :ZZZ,     :7ZZZZ$$$$=                ");
    putstr(&"\n           Z$$$$$?                    .+ZZZZ$$                 ");
    putstr(&"\n       +$ZZ$$$Z7                         7ZZZ$Z$$I.            ");
    putstr(&"\n        $$$$ZZZZZZZZZZZZZZZZZZZZZZZZI,    ,ZZZ$$Z              ");
    putstr(&"\n      :+$$$$ZZZZZZZZZZZZZZZZZZZZZZZZZZZ=    $ZZ$$+~,           ");
    putstr(&"\n     ?$Z$$$$ZZZZZZZZZZZZZZZZZZZZZZZZZZZZI   7ZZZ$ZZI           ");
    putstr(&"\n      =Z$$+7Z$$7ZZZZZZZZ$$$$$$$ZZZZZZZZZZ  ~Z$?$ZZ?            ");
    putstr(&"\n    :$Z$Z...$Z  $ZZZZZZZ~       ~ZZZZZZZZ,.ZZ...Z$Z$~          ");
    putstr(&"\n    7ZZZZZI$ZZ  $ZZZZZZZ~       =ZZZZZZZ7..ZZ$?$ZZZZ$          ");
    putstr(&"\n      ZZZZ$:    $ZZZZZZZZZZZZZZZZZZZZZZ=     ~$ZZZ$:           ");
    putstr(&"\n    7Z$ZZ$,     $ZZZZZZZZZZZZZZZZZZZZ7         ZZZ$Z$          ");
    putstr(&"\n   =ZZZZZZ,     $ZZZZZZZZZZZZZZZZZZZZZZ,       ZZZ$ZZ+         ");
    putstr(&"\n     ,ZZZZ,     $ZZZZZZZ:     =ZZZZZZZZZ     ZZZZZ$:           ");
    putstr(&"\n    =$ZZZZ+     ZZZZZZZZ~       ZZZZZZZZ~   =ZZZZZZZI          ");
    putstr(&"\n    $ZZ$ZZZ$$Z$$ZZZZZZZZZ$$$$   IZZZZZZZZZ$ZZZZZZZZZ$          ");
    putstr(&"\n      :ZZZZZZZZZZZZZZZZZZZZZZ   ~ZZZZZZZZZZZZZZZZZ~            ");
    putstr(&"\n     ,Z$$ZZZZZZZZZZZZZZZZZZZZ    ZZZZZZZZZZZZZZZZZZ~           ");
    putstr(&"\n     =$ZZZZZZZZZZZZZZZZZZZZZZ     $ZZZZZZZZZZZZZZZ$+           ");
    putstr(&"\n        IZZZZZ:.                        . ,ZZZZZ$              ");
    putstr(&"\n       ~$ZZZZZZZZZZZ                 ZZZZ$ZZZZZZZ+             ");
    putstr(&"\n           Z$ZZZ. ,Z~               =Z:.,ZZZ$Z                 ");
    putstr(&"\n          ,ZZZZZ..~Z$.             .7Z:..ZZZZZ:                ");
    putstr(&"\n          ~7+:$ZZZZZZZZI=:.   .,=IZZZZZZZ$Z:=7=                ");
    putstr(&"\n              $$ZZZZZZZZZZZZZZZZZZZZZZ$ZZZZ                    ");
    putstr(&"\n              ==..$ZZZ$ZZZZZZZZZZZ$ZZZZ .~+                    ");
    putstr(&"\n                  I$?.?ZZZ$ZZZ$ZZZI =$7                        ");
    putstr(&"\n                       $7..I$7..I$,                            ");
    putstr(&"\n");
    putstr(&"\n _                     _     _                         _  ");
    putstr(&"\n| |                   (_)   | |                       | | ");
    putstr(&"\n| | ____ ___  ____     _____| |_____  ____ ____  _____| | ");
    putstr(&"\n| |/ ___) _ \\|  _ \\   |  _   _) ___ |/ ___)  _ \\| ___ | | ");
    putstr(&"\n| | |  | |_| | | | |  | |  \\ \\| ____| |   | | | | ____| | ");
    putstr(&"\n|_|_|  \\____/|_| |_|  |_|   \\_\\_____)_|   |_| |_|_____)__)\n\n");
}

pub unsafe fn init() {
    buffer = cstr::new();
    history = cstr::news(512);
    screen();
    prompt();
}

pub unsafe fn prompt() {
    putstr(PROMPT);
    let prev_c = super::super::io::FG_COLOR;
    super::super::io::set_fg(PROMPT_COLOR);
    drawstr(PROMPT);
    super::super::io::set_fg(prev_c);
}

<|endoftext|>


<commit_before>#![deny(warnings)]
#![feature(const_fn)]

extern crate orbclient;
extern crate syscall;

use orbclient::event;

use std::{env, str};
use std::error::Error;
use std::fs::File;
use std::io::{self, Read, Write};
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
use std::process::{Command, Stdio};

use console::Console;
use getpty::getpty;

mod console;
mod getpty;

fn main() {
    let shell = env::args().nth(1).unwrap_or("sh".to_string());

    let (master_fd, tty_path) = getpty();

    let slave_stdin = File::open(&tty_path).unwrap();
    let slave_stdout = File::create(&tty_path).unwrap();
    let slave_stderr = File::create(&tty_path).unwrap();

    let width = 640;
    let height = 480;

    env::set_var("COLUMNS", format!("{}", width / 8));
    env::set_var("LINES", format!("{}", height / 16));
    env::set_var("TTY", format!("{}", tty_path.display()));

    match unsafe {
        Command::new(&shell)
            .stdin(Stdio::from_raw_fd(slave_stdin.into_raw_fd()))
            .stdout(Stdio::from_raw_fd(slave_stdout.into_raw_fd()))
            .stderr(Stdio::from_raw_fd(slave_stderr.into_raw_fd()))
            .spawn()
    } {
        Ok(_process) => {
            let mut console = Console::new(width, height);

            let mut event_file = File::open("event:").expect("terminal: failed to open event file");

            let window_fd = console.window.as_raw_fd();
            syscall::fevent(window_fd, syscall::flag::EVENT_READ).expect("terminal: failed to fevent console window");

            let mut master = unsafe { File::from_raw_fd(master_fd) };
            syscall::fevent(master_fd, syscall::flag::EVENT_READ).expect("terminal: failed to fevent master PTY");

            let mut handle_event = |event_id: usize, event_count: usize| -> bool {
                if event_id == window_fd {
                    for event in console.window.events() {
                        if event.code == event::EVENT_QUIT {
                            println!("window quit");
                            return false;
                        }

                        if let Some(line) = console.event(event) {
                            if let Err(err) = master.write(&line.as_bytes()) {
                                let term_stderr = io::stderr();
                                let mut term_stderr = term_stderr.lock();

                                let _ = term_stderr.write(b"failed to write stdin: ");
                                let _ = term_stderr.write(err.description().as_bytes());
                                let _ = term_stderr.write(b"\n");
                                return false;
                            }
                        }
                    }
                } else if event_id == master_fd {
                    let mut packet = [0; 4096];
                    let count = master.read(&mut packet).expect("terminal: failed to read master PTY");
                    if count == 0 {
                        if event_count == 0 {
                            return false;
                        }
                    } else {
                        if packet[0] & 1 == 1 {
                            console.inner.redraw = true;
                        }
                        console.write(&packet[1..count])
                    }
                } else {
                    println!("Unknown event {}", event_id);
                }

                true
            };

            handle_event(window_fd, 0);
            handle_event(master_fd, 0);

            'events: loop {
                let mut sys_event = syscall::Event::default();
                event_file.read(&mut sys_event).expect("terminal: failed to read event file");
                if ! handle_event(sys_event.id, sys_event.data) {
                    break 'events;
                }
                println!("handled");
            }
        },
        Err(err) => {
            let term_stderr = io::stderr();
            let mut term_stderr = term_stderr.lock();
            let _ = term_stderr.write(b"failed to execute '");
            let _ = term_stderr.write(shell.as_bytes());
            let _ = term_stderr.write(b"': ");
            let _ = term_stderr.write(err.description().as_bytes());
            let _ = term_stderr.write(b"\n");
        }
    }
}
<commit_msg>Correctly wait on shell<commit_after>#![deny(warnings)]
#![feature(const_fn)]

extern crate orbclient;
extern crate syscall;

use orbclient::event;

use std::{env, str};
use std::error::Error;
use std::fs::File;
use std::io::{self, Read, Write};
use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
use std::process::{Command, Stdio};

use console::Console;
use getpty::getpty;

mod console;
mod getpty;

fn main() {
    let shell = env::args().nth(1).unwrap_or("sh".to_string());

    let (master_fd, tty_path) = getpty();

    let slave_stdin = File::open(&tty_path).unwrap();
    let slave_stdout = File::create(&tty_path).unwrap();
    let slave_stderr = File::create(&tty_path).unwrap();

    let width = 640;
    let height = 480;

    env::set_var("COLUMNS", format!("{}", width / 8));
    env::set_var("LINES", format!("{}", height / 16));
    env::set_var("TTY", format!("{}", tty_path.display()));

    match unsafe {
        Command::new(&shell)
            .stdin(Stdio::from_raw_fd(slave_stdin.into_raw_fd()))
            .stdout(Stdio::from_raw_fd(slave_stdout.into_raw_fd()))
            .stderr(Stdio::from_raw_fd(slave_stderr.into_raw_fd()))
            .spawn()
    } {
        Ok(mut process) => {
            {
                let mut console = Console::new(width, height);

                let mut event_file = File::open("event:").expect("terminal: failed to open event file");

                let window_fd = console.window.as_raw_fd();
                syscall::fevent(window_fd, syscall::flag::EVENT_READ).expect("terminal: failed to fevent console window");

                let mut master = unsafe { File::from_raw_fd(master_fd) };
                syscall::fevent(master_fd, syscall::flag::EVENT_READ).expect("terminal: failed to fevent master PTY");

                let mut handle_event = |event_id: usize, event_count: usize| -> bool {
                    if event_id == window_fd {
                        for event in console.window.events() {
                            if event.code == event::EVENT_QUIT {
                                return false;
                            }

                            if let Some(line) = console.event(event) {
                                if let Err(err) = master.write(&line.as_bytes()) {
                                    let term_stderr = io::stderr();
                                    let mut term_stderr = term_stderr.lock();

                                    let _ = term_stderr.write(b"failed to write stdin: ");
                                    let _ = term_stderr.write(err.description().as_bytes());
                                    let _ = term_stderr.write(b"\n");
                                    return false;
                                }
                            }
                        }
                    } else if event_id == master_fd {
                        let mut packet = [0; 4096];
                        let count = master.read(&mut packet).expect("terminal: failed to read master PTY");
                        if count == 0 {
                            if event_count == 0 {
                                return false;
                            }
                        } else {
                            if packet[0] & 1 == 1 {
                                console.inner.redraw = true;
                            }
                            console.write(&packet[1..count])
                        }
                    } else {
                        println!("Unknown event {}", event_id);
                    }

                    true
                };

                handle_event(window_fd, 0);
                handle_event(master_fd, 0);

                'events: loop {
                    let mut sys_event = syscall::Event::default();
                    event_file.read(&mut sys_event).expect("terminal: failed to read event file");
                    if ! handle_event(sys_event.id, sys_event.data) {
                        break 'events;
                    }
                }
            }

            process.wait().expect("terminal: failed to wait on shell");
        },
        Err(err) => {
            let term_stderr = io::stderr();
            let mut term_stderr = term_stderr.lock();
            let _ = term_stderr.write(b"failed to execute '");
            let _ = term_stderr.write(shell.as_bytes());
            let _ = term_stderr.write(b"': ");
            let _ = term_stderr.write(err.description().as_bytes());
            let _ = term_stderr.write(b"\n");
        }
    }
}
<|endoftext|>


<commit_before>use luminance::{Depth32F, Dim2, Flat, RGBA32F, Uniformable};
use luminance_gl::gl33::{Framebuffer, Pipe, Pipeline, RenderCommand, ShadingCommand, Texture};

use camera::{Camera, Freefly};
use extra::shaders::default::{DEFAULT_3D_INST, DEFAULT_3D_PROJ, DEFAULT_3D_VIEW, DefaultProgram3D};
use id::Id;
use object::Object;
use projection::Projectable;
use renderer::Renderer;
use scene::Scene;
use shader::Program;
use transform::Transformable;

pub type Texture2D<A> = Texture<Flat, Dim2, A>;

/// Simple renderer that takes a camera, a set of model and applies a shader on them. This renderer
/// outputs a single color map along with the depth map.
pub struct Simple<'a> {
  program: DefaultProgram3D<'a>,
  framebuffer: Framebuffer<Flat, Dim2, Texture2D<RGBA32F>, Texture2D<Depth32F>>
}

impl<'a> Simple<'a> {
  pub fn new_from(w: u32, h: u32, scene: &'a mut Scene<'a>) -> Self {
    Simple {
      program: DefaultProgram3D::new_from(scene).unwrap(),
      framebuffer: Framebuffer::new((w, h), 0).unwrap()
    }
  }

  pub fn renderer(&'a mut self, camera: &'a Camera<Freefly>, objects: &'a [Object<'a>]) -> impl Renderer<'a, Output=(&'a Texture2D<RGBA32F>, &'a Texture2D<Depth32F>)> {
    SimpleRenderer {
      program: &self.program,
      framebuffer: &self.framebuffer,
      camera: camera,
      objects: objects
    }
  }
}

struct SimpleRenderer<'a> {
  program: &'a Id<'a, Program>,
  framebuffer: &'a Framebuffer<Flat, Dim2, Texture2D<RGBA32F>, Texture2D<Depth32F>>,
  camera: &'a Camera<Freefly>,
  objects: &'a [Object<'a>]
}

impl<'a> Renderer<'a> for SimpleRenderer<'a> {
  type Output = (&'a Texture2D<RGBA32F>, &'a Texture2D<Depth32F>);

  fn render(&'a mut self, scene: &'a mut Scene<'a>) -> Self::Output {
    let program = scene.get_by_id(self.program).unwrap();

    // reify objects
    let objects: Vec<_> = self.objects.iter().map(|object| {
      (object, scene.get_by_id(&object.model).unwrap())
    }).collect();

    let tessellations = objects.iter().flat_map(|&(ref object, ref model)| {
      model.parts.iter().map(move |part| {
        Pipe::new(move |program| {
                    program.update(&DEFAULT_3D_INST, *object.transform().as_ref());
                  },
                  &part.tess)
      })
    }).collect();

    Pipeline::new(self.framebuffer, [0., 0., 0., 0.], &[], &[], vec![
      Pipe::new(|program| {
                  // update the camera
                  program.update(&DEFAULT_3D_PROJ, *self.camera.projection().as_ref());
                  program.update(&DEFAULT_3D_VIEW, *self.camera.transform().as_ref());
                },
                ShadingCommand::new(&program, vec![
                  Pipe::new(|_| {}, RenderCommand::new(None, true, tessellations, 1, None))
                ]))
    ]).run();

    (&self.framebuffer.color_slot, &self.framebuffer.depth_slot)
  }
}
<commit_msg>Changed lifetimes in SimpleRenderer.<commit_after>use luminance::{Depth32F, Dim2, Flat, RGBA32F};
use luminance_gl::gl33::{Framebuffer, Pipe, Pipeline, RenderCommand, ShadingCommand, Texture};

use camera::{Camera, Freefly};
use extra::shaders::default::{DEFAULT_3D_INST, DEFAULT_3D_PROJ, DEFAULT_3D_VIEW, DefaultProgram3D};
use id::Id;
use object::Object;
use projection::Projectable;
use renderer::Renderer;
use scene::Scene;
use shader::Program;
use transform::Transformable;

pub type Texture2D<A> = Texture<Flat, Dim2, A>;

/// Simple renderer that takes a camera, a set of model and applies a shader on them. This renderer
/// outputs a single color map along with the depth map.
pub struct Simple<'a> {
  program: DefaultProgram3D<'a>,
  framebuffer: Framebuffer<Flat, Dim2, Texture2D<RGBA32F>, Texture2D<Depth32F>>
}

impl<'a> Simple<'a> {
  pub fn new_from(w: u32, h: u32, scene: &mut Scene<'a>) -> Self {
    Simple {
      program: DefaultProgram3D::new_from(scene).unwrap(),
      framebuffer: Framebuffer::new((w, h), 0).unwrap()
    }
  }

  pub fn renderer(&'a mut self, camera: &'a Camera<Freefly>, objects: &'a [Object<'a>]) -> impl Renderer<'a, Output=(&'a Texture2D<RGBA32F>, &'a Texture2D<Depth32F>)> {
    SimpleRenderer {
      program: &self.program,
      framebuffer: &self.framebuffer,
      camera: camera,
      objects: objects
    }
  }
}

struct SimpleRenderer<'a> {
  program: &'a Id<'a, Program>,
  framebuffer: &'a Framebuffer<Flat, Dim2, Texture2D<RGBA32F>, Texture2D<Depth32F>>,
  camera: &'a Camera<Freefly>,
  objects: &'a [Object<'a>]
}

impl<'a> Renderer<'a> for SimpleRenderer<'a> {
  type Output = (&'a Texture2D<RGBA32F>, &'a Texture2D<Depth32F>);

  fn render(&'a mut self, scene: &'a mut Scene<'a>) -> Self::Output {
    let program = scene.get_by_id(self.program).unwrap();

    // reify objects
    let objects: Vec<_> = self.objects.iter().map(|object| {
      (object, scene.get_by_id(&object.model).unwrap())
    }).collect();

    let tessellations = objects.iter().flat_map(|&(ref object, ref model)| {
      model.parts.iter().map(move |part| {
        Pipe::new(move |program| {
                    program.update(&DEFAULT_3D_INST, *object.transform().as_ref());
                  },
                  &part.tess)
      })
    }).collect();

    Pipeline::new(self.framebuffer, [0., 0., 0., 0.], &[], &[], vec![
      Pipe::new(|program| {
                  // update the camera
                  program.update(&DEFAULT_3D_PROJ, *self.camera.projection().as_ref());
                  program.update(&DEFAULT_3D_VIEW, *self.camera.transform().as_ref());
                },
                ShadingCommand::new(&program, vec![
                  Pipe::new(|_| {}, RenderCommand::new(None, true, tessellations, 1, None))
                ]))
    ]).run();

    (&self.framebuffer.color_slot, &self.framebuffer.depth_slot)
  }
}
<|endoftext|>


<commit_before>//! Defines the `Source` Trait.

use serde_json;

use std::path::{Path, PathBuf};
use std::fs::{self, File};
use std::io::{Read, Write};
use std::rc::Rc;
use std::cell::RefCell;
use std::error::Error;
use std::fmt;

use r2api::api_trait::R2Api;
use r2pipe::r2::R2;
use r2api::structs::{FunctionInfo, LFlagInfo, LOpInfo, LRegInfo, LSectionInfo, LStringInfo, LSymbolInfo,
LImportInfo, LExportInfo, LRelocInfo, LEntryInfo};

#[derive(Debug)]
pub enum SourceErr {
    SrcErr(&'static str),
    OtherErr(Box<Error>),
}

impl fmt::Display for SourceErr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &SourceErr::SrcErr(s) => write!(f, "{}", s),
            &SourceErr::OtherErr(ref e) => write!(f, "{}", e),
        }
    }
}

impl<T: 'static + Error> From<T> for SourceErr {
    fn from(e: T) -> SourceErr {
        SourceErr::OtherErr(Box::new(e))
    }
}

// TODO: Split this up/compose this from more basic traits to avoid reimplementation
// as currently this is a re-implementation of r2api
pub trait Source {
    fn functions(&self) -> Result<Vec<FunctionInfo>, SourceErr>;
    fn instructions_at(&self, u64) -> Result<Vec<LOpInfo>, SourceErr>;
    fn register_profile(&self) -> Result<LRegInfo, SourceErr>;
    fn flags(&self) -> Result<Vec<LFlagInfo>, SourceErr>;
    fn sections(&self) -> Result<Vec<LSectionInfo>, SourceErr>;
    fn symbols(&self) -> Result<Vec<LSymbolInfo>, SourceErr> { unimplemented!() }
    fn imports(&self) -> Result<Vec<LImportInfo>, SourceErr> { unimplemented!() }
    fn exports(&self) -> Result<Vec<LExportInfo>, SourceErr> { unimplemented!() }
    fn relocs(&self) -> Result<Vec<LRelocInfo>, SourceErr> { unimplemented!() }
    fn libraries(&self) -> Result<Vec<String>, SourceErr> { unimplemented!() }
    fn entrypoint(&self) -> Result<Vec<LEntryInfo>, SourceErr> { unimplemented!() }
    fn disassemble_n_bytes(&self, n: u64, at: u64) -> Result<Vec<LOpInfo>, SourceErr> { unimplemented!() }
    fn disassemble_n_insts(&self, n: u64, at: u64) -> Result<Vec<LOpInfo>, SourceErr> { unimplemented!() }

    fn send(&self, _: &str) -> Result<(), SourceErr> { Ok(()) }

    // Non essential / functions with default implementation.
    fn function_at(&self, address: u64) -> Result<FunctionInfo, SourceErr> {
        for f in self.functions()? {
            match f.offset {
                Some(off) if address == off => return Ok(f),
                _ => {}
            }
        }
        Err(SourceErr::SrcErr("No known function offset"))
    }

    fn function_named(&self, fn_name: &str) -> Result<FunctionInfo, SourceErr> {
        for f in self.functions()? {
            match f.name {
                Some(ref name) if name == fn_name => return Ok(f.clone()),
                _ => {}
            }
        }
        Err(SourceErr::SrcErr("No known function"))
    }

    fn instructions_at_fn(&self, fn_name: &str) -> Result<Vec<LOpInfo>, SourceErr> {
        let fn_ = self.function_named(fn_name)?;
        if let Some(offset) = fn_.offset {
            self.instructions_at(offset)
        } else {
            Err(SourceErr::SrcErr("No known function"))
        }
    }

    fn flag_at(&self, address: u64) -> Result<LFlagInfo, SourceErr> {
        for flag in self.flags()? {
            if flag.offset == address {
                return Ok(flag);
            }
        }
        Err(SourceErr::SrcErr("No known flag at offset"))
    }

    fn section_of(&self, address: u64) -> Result<LSectionInfo, SourceErr> {
        for s in self.sections()? {
            let addr = s.vaddr.expect("Invalid section");
            let size = s.size.expect("Invalid section size");
            if address >= addr && address < addr + size {
                return Ok(s);
            }
        }
        Err(SourceErr::SrcErr("No known section at addr"))
    }
}

// Cause R2Api requires borrowing mutably, while `Source` takes self which
// is immutable.
// The only problem with this is that r2pipe is not thread safe, therefore
// using `r2` from multiple threads will cause results to be inconsistent.
// However, this should not be a problem as `Rc` prevents concurrent accesses from multiple
// threads.
//
// TODO: For threading, this should be replaced by `RwLock<T>` or `Mutex<T>`
// But it is still up for discussion if this should be done at the r2pipe level rather than here.
// RwLock allows for multiple concurrent readers and one writer, however this will not improve
// the performance when compared to Mutex as every function takes `self` mutably (as it has to
// communicate with r2 which amounts to writing out to process pipe). Therefore, it'd make
// sense to have some sort of cached information so that concurrent reads can occur. This should
// be invalidated whenever some information is exported back to radare or some analysis is run on
// r2.
pub type WrappedR2Api<R> = Rc<RefCell<R>>;

// Implementation of `Source` trait for R2.
impl<R: R2Api> Source for WrappedR2Api<R> {
    fn functions(&self) -> Result<Vec<FunctionInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.fn_list()?)
    }

    fn instructions_at(&self, address: u64) -> Result<Vec<LOpInfo>, SourceErr> {
        if let Ok(fn_info) = self.try_borrow_mut()?.function(&format!("{}", address)) {
            fn_info.ops.ok_or(SourceErr::SrcErr("No Instructions found"))
        } else {
            Err(SourceErr::SrcErr("No Instructions found"))
        }
    }

    fn register_profile(&self) -> Result<LRegInfo, SourceErr> {
        Ok(self.try_borrow_mut()?.reg_info()?)
    }

    fn flags(&self) -> Result<Vec<LFlagInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.flag_info()?)
    }

    fn sections(&self) -> Result<Vec<LSectionInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.sections()?)
    }

    fn symbols(&self) -> Result<Vec<LSymbolInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.symbols()?)
    }

    fn imports(&self) -> Result<Vec<LImportInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.imports()?)
    }

    fn exports(&self) -> Result<Vec<LExportInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.exports()?)
    }

    fn relocs(&self) -> Result<Vec<LRelocInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.relocs()?)
    }

    fn libraries(&self) -> Result<Vec<String>, SourceErr> {
        Ok(self.try_borrow_mut()?.libraries()?)
    }

    fn entrypoint(&self) -> Result<Vec<LEntryInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.entrypoint()?)
    }

    fn disassemble_n_bytes(&self, n: u64, at: u64) -> Result<Vec<LOpInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.disassemble_n_bytes(n, Some(at))?)
    }

    fn disassemble_n_insts(&self, n: u64, at: u64) -> Result<Vec<LOpInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.disassemble_n_insts(n, Some(at))?)
    }

    fn send(&self, s: &str) -> Result<(), SourceErr> {
        unimplemented!()
    }
}

#[derive(Clone, Debug)]
/// File source is used to load information from json files.
/// The files in the directory must be of the form "<base_name>_<suffix>.json". This means that the
/// files must have specific names.
///
/// To construct a FileSource painlessly (and save it for future use), use the `from::From` trait
/// implemented on `R2`.
pub struct FileSource {
    /// Directory where files should be loaded from.
    pub dir: String,
    /// Base name for the sample/example. suffixes will be appended based on what information is
    /// requested for. Hence, files in the directory must have specific names.
    pub base_name: String,
}

impl FileSource {
    fn read_file(&self, suffix: &str) -> String {
        let mut path = PathBuf::from(&self.dir);
        path.push(&format!("{}_{}.json", self.base_name, suffix));
        let mut f = File::open(path).expect("Failed to open file");
        let mut json_str = String::new();
        let _ = f.read_to_string(&mut json_str).expect("Failed to read file");
        json_str
    }

    fn write_file(&mut self, suffix: &str, data: &str) {
        let mut path = PathBuf::from(&self.dir);
        path.push(&format!("{}_{}.json", self.base_name, suffix));
        let mut f = File::create(path).expect("Failed to open file");
        f.write_all(data.to_string()
                        .as_bytes()).expect("Failed to read file");
    }
}

mod suffix {
    pub const FUNCTION_INFO: &'static str = "fn_info";
    pub const INSTRUCTIONS: &'static str = "insts";
    pub const REGISTER: &'static str = "register_profile";
    pub const FLAG: &'static str = "flags";
    pub const SECTION: &'static str = "sections";
    pub const STRING: &'static str = "strings";
}

impl FileSource {
    pub fn open(f: Option<&str>) -> FileSource {
        let path = Path::new(f.unwrap());
        let dir = path.parent().unwrap().to_str().unwrap();
        let base_name = path.file_name().unwrap().to_str().unwrap();
        FileSource {
            dir: dir.to_owned(),
            base_name: base_name.to_owned(),
        }
    }

    pub fn strings(&mut self) -> Vec<LStringInfo> {
        serde_json::from_str(&self.read_file(suffix::STRING)).expect("Failed to decode json")
    }
}

impl Source for FileSource {
    fn functions(&self) -> Result<Vec<FunctionInfo>, SourceErr> {
        Ok(serde_json::from_str(&self.read_file(suffix::FUNCTION_INFO))?)
    }

    fn instructions_at(&self, address: u64) -> Result<Vec<LOpInfo>, SourceErr> {
        let suffix = format!("{}_{:#X}", suffix::INSTRUCTIONS, address);
        Ok(serde_json::from_str(&self.read_file(&suffix))?)
    }

    fn register_profile(&self) -> Result<LRegInfo, SourceErr> {
        Ok(serde_json::from_str(&self.read_file(suffix::REGISTER))?)
    }

    fn flags(&self) -> Result<Vec<LFlagInfo>, SourceErr> {
        Ok(serde_json::from_str(&self.read_file(suffix::FLAG))?)
    }

    fn sections(&self) -> Result<Vec<LSectionInfo>, SourceErr> {
        Ok(serde_json::from_str(&self.read_file(suffix::SECTION))?)
    }
}

impl<R: R2Api> From<WrappedR2Api<R>> for FileSource {
    fn from(mut r2: WrappedR2Api<R>) -> FileSource {
        let bin_info = r2.borrow_mut().bin_info().expect("Failed to load bin_info");
        let fname = bin_info.core.unwrap().file.unwrap();
        let fname = Path::new(&fname).file_stem().unwrap();
        let mut dir = PathBuf::from(".");
        dir.push(&fname);
        fs::create_dir_all(&dir).expect("Failed to create directory");
        let mut fsource = FileSource {
            dir: dir.to_str().unwrap().to_owned(),
            base_name: fname.to_str().unwrap().to_owned(),
        };

        {
            let fns = r2.functions().expect("Failed to load function info");
            {
                let json_str = serde_json::to_string(&fns).expect("Failed to encode to json");
                fsource.write_file(suffix::FUNCTION_INFO, &json_str);
            }

            for f in fns {
                let result = r2.instructions_at(f.offset.unwrap()).expect("Failed to load instructions");
                let json_str = serde_json::to_string(&result).expect("Failed to encode to json");
                let suffix = format!("{}_{:#X}", suffix::INSTRUCTIONS, f.offset.unwrap());
                fsource.write_file(&suffix, &json_str)
            }

            {
                let reg = r2.register_profile().expect("Failed to load register info");
                let json_str = serde_json::to_string(&reg).expect("Failed to encode to json");
                fsource.write_file(suffix::REGISTER, &json_str);
            }

            {
                let flags = r2.flags().expect("Failed to load flag info");
                let json_str = serde_json::to_string(&flags).expect("Failed to encode to json");
                fsource.write_file(suffix::FLAG, &json_str);
            }

            {
                let sections = r2.sections().expect("Failed to load section map");
                let json_str = serde_json::to_string(&sections).expect("Failed to encode to json");
                fsource.write_file(suffix::SECTION, &json_str);
            }

            {
                let strings = r2.borrow_mut().strings(false).expect("Unable to load String info from r2");
                let json_str = serde_json::to_string(&strings).expect("Failed to encode to json");
                fsource.write_file(suffix::STRING, &json_str);
            }
        }

        fsource
    }
}
<commit_msg>Exposes raw in source<commit_after>//! Defines the `Source` Trait.

use serde_json;

use std::path::{Path, PathBuf};
use std::fs::{self, File};
use std::io::{Read, Write};
use std::rc::Rc;
use std::cell::RefCell;
use std::error::Error;
use std::fmt;

use r2api::api_trait::R2Api;
use r2pipe::r2::R2;
use r2api::structs::{FunctionInfo, LFlagInfo, LOpInfo, LRegInfo, LSectionInfo, LStringInfo, LSymbolInfo,
LImportInfo, LExportInfo, LRelocInfo, LEntryInfo};

#[derive(Debug)]
pub enum SourceErr {
    SrcErr(&'static str),
    OtherErr(Box<Error>),
}

impl fmt::Display for SourceErr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &SourceErr::SrcErr(s) => write!(f, "{}", s),
            &SourceErr::OtherErr(ref e) => write!(f, "{}", e),
        }
    }
}

impl<T: 'static + Error> From<T> for SourceErr {
    fn from(e: T) -> SourceErr {
        SourceErr::OtherErr(Box::new(e))
    }
}

// TODO: Split this up/compose this from more basic traits to avoid reimplementation
// as currently this is a re-implementation of r2api
pub trait Source {
    fn functions(&self) -> Result<Vec<FunctionInfo>, SourceErr>;
    fn instructions_at(&self, u64) -> Result<Vec<LOpInfo>, SourceErr>;
    fn register_profile(&self) -> Result<LRegInfo, SourceErr>;
    fn flags(&self) -> Result<Vec<LFlagInfo>, SourceErr>;
    fn sections(&self) -> Result<Vec<LSectionInfo>, SourceErr>;
    fn symbols(&self) -> Result<Vec<LSymbolInfo>, SourceErr> { unimplemented!() }
    fn imports(&self) -> Result<Vec<LImportInfo>, SourceErr> { unimplemented!() }
    fn exports(&self) -> Result<Vec<LExportInfo>, SourceErr> { unimplemented!() }
    fn relocs(&self) -> Result<Vec<LRelocInfo>, SourceErr> { unimplemented!() }
    fn libraries(&self) -> Result<Vec<String>, SourceErr> { unimplemented!() }
    fn entrypoint(&self) -> Result<Vec<LEntryInfo>, SourceErr> { unimplemented!() }
    fn disassemble_n_bytes(&self, n: u64, at: u64) -> Result<Vec<LOpInfo>, SourceErr> { unimplemented!() }
    fn disassemble_n_insts(&self, n: u64, at: u64) -> Result<Vec<LOpInfo>, SourceErr> { unimplemented!() }
    fn raw(&self, cmd: String) -> Result<String, SourceErr> { unimplemented!() }

    fn send(&self, _: &str) -> Result<(), SourceErr> { Ok(()) }

    // Non essential / functions with default implementation.
    fn function_at(&self, address: u64) -> Result<FunctionInfo, SourceErr> {
        for f in self.functions()? {
            match f.offset {
                Some(off) if address == off => return Ok(f),
                _ => {}
            }
        }
        Err(SourceErr::SrcErr("No known function offset"))
    }

    fn function_named(&self, fn_name: &str) -> Result<FunctionInfo, SourceErr> {
        for f in self.functions()? {
            match f.name {
                Some(ref name) if name == fn_name => return Ok(f.clone()),
                _ => {}
            }
        }
        Err(SourceErr::SrcErr("No known function"))
    }

    fn instructions_at_fn(&self, fn_name: &str) -> Result<Vec<LOpInfo>, SourceErr> {
        let fn_ = self.function_named(fn_name)?;
        if let Some(offset) = fn_.offset {
            self.instructions_at(offset)
        } else {
            Err(SourceErr::SrcErr("No known function"))
        }
    }

    fn flag_at(&self, address: u64) -> Result<LFlagInfo, SourceErr> {
        for flag in self.flags()? {
            if flag.offset == address {
                return Ok(flag);
            }
        }
        Err(SourceErr::SrcErr("No known flag at offset"))
    }

    fn section_of(&self, address: u64) -> Result<LSectionInfo, SourceErr> {
        for s in self.sections()? {
            let addr = s.vaddr.expect("Invalid section");
            let size = s.size.expect("Invalid section size");
            if address >= addr && address < addr + size {
                return Ok(s);
            }
        }
        Err(SourceErr::SrcErr("No known section at addr"))
    }
}

// Cause R2Api requires borrowing mutably, while `Source` takes self which
// is immutable.
// The only problem with this is that r2pipe is not thread safe, therefore
// using `r2` from multiple threads will cause results to be inconsistent.
// However, this should not be a problem as `Rc` prevents concurrent accesses from multiple
// threads.
//
// TODO: For threading, this should be replaced by `RwLock<T>` or `Mutex<T>`
// But it is still up for discussion if this should be done at the r2pipe level rather than here.
// RwLock allows for multiple concurrent readers and one writer, however this will not improve
// the performance when compared to Mutex as every function takes `self` mutably (as it has to
// communicate with r2 which amounts to writing out to process pipe). Therefore, it'd make
// sense to have some sort of cached information so that concurrent reads can occur. This should
// be invalidated whenever some information is exported back to radare or some analysis is run on
// r2.
pub type WrappedR2Api<R> = Rc<RefCell<R>>;

// Implementation of `Source` trait for R2.
impl<R: R2Api> Source for WrappedR2Api<R> {
    fn functions(&self) -> Result<Vec<FunctionInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.fn_list()?)
    }

    fn instructions_at(&self, address: u64) -> Result<Vec<LOpInfo>, SourceErr> {
        if let Ok(fn_info) = self.try_borrow_mut()?.function(&format!("{}", address)) {
            fn_info.ops.ok_or(SourceErr::SrcErr("No Instructions found"))
        } else {
            Err(SourceErr::SrcErr("No Instructions found"))
        }
    }

    fn register_profile(&self) -> Result<LRegInfo, SourceErr> {
        Ok(self.try_borrow_mut()?.reg_info()?)
    }

    fn flags(&self) -> Result<Vec<LFlagInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.flag_info()?)
    }

    fn sections(&self) -> Result<Vec<LSectionInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.sections()?)
    }

    fn symbols(&self) -> Result<Vec<LSymbolInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.symbols()?)
    }

    fn imports(&self) -> Result<Vec<LImportInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.imports()?)
    }

    fn exports(&self) -> Result<Vec<LExportInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.exports()?)
    }

    fn relocs(&self) -> Result<Vec<LRelocInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.relocs()?)
    }

    fn libraries(&self) -> Result<Vec<String>, SourceErr> {
        Ok(self.try_borrow_mut()?.libraries()?)
    }

    fn entrypoint(&self) -> Result<Vec<LEntryInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.entrypoint()?)
    }

    fn disassemble_n_bytes(&self, n: u64, at: u64) -> Result<Vec<LOpInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.disassemble_n_bytes(n, Some(at))?)
    }

    fn disassemble_n_insts(&self, n: u64, at: u64) -> Result<Vec<LOpInfo>, SourceErr> {
        Ok(self.try_borrow_mut()?.disassemble_n_insts(n, Some(at))?)
    }

    fn raw(&self, cmd: String) -> Result<String, SourceErr> {
        Ok(self.try_borrow_mut()?.raw(cmd))
    }

    fn send(&self, s: &str) -> Result<(), SourceErr> {
        unimplemented!()
    }
}

#[derive(Clone, Debug)]
/// File source is used to load information from json files.
/// The files in the directory must be of the form "<base_name>_<suffix>.json". This means that the
/// files must have specific names.
///
/// To construct a FileSource painlessly (and save it for future use), use the `from::From` trait
/// implemented on `R2`.
pub struct FileSource {
    /// Directory where files should be loaded from.
    pub dir: String,
    /// Base name for the sample/example. suffixes will be appended based on what information is
    /// requested for. Hence, files in the directory must have specific names.
    pub base_name: String,
}

impl FileSource {
    fn read_file(&self, suffix: &str) -> String {
        let mut path = PathBuf::from(&self.dir);
        path.push(&format!("{}_{}.json", self.base_name, suffix));
        let mut f = File::open(path).expect("Failed to open file");
        let mut json_str = String::new();
        let _ = f.read_to_string(&mut json_str).expect("Failed to read file");
        json_str
    }

    fn write_file(&mut self, suffix: &str, data: &str) {
        let mut path = PathBuf::from(&self.dir);
        path.push(&format!("{}_{}.json", self.base_name, suffix));
        let mut f = File::create(path).expect("Failed to open file");
        f.write_all(data.to_string()
                        .as_bytes()).expect("Failed to read file");
    }
}

mod suffix {
    pub const FUNCTION_INFO: &'static str = "fn_info";
    pub const INSTRUCTIONS: &'static str = "insts";
    pub const REGISTER: &'static str = "register_profile";
    pub const FLAG: &'static str = "flags";
    pub const SECTION: &'static str = "sections";
    pub const STRING: &'static str = "strings";
}

impl FileSource {
    pub fn open(f: Option<&str>) -> FileSource {
        let path = Path::new(f.unwrap());
        let dir = path.parent().unwrap().to_str().unwrap();
        let base_name = path.file_name().unwrap().to_str().unwrap();
        FileSource {
            dir: dir.to_owned(),
            base_name: base_name.to_owned(),
        }
    }

    pub fn strings(&mut self) -> Vec<LStringInfo> {
        serde_json::from_str(&self.read_file(suffix::STRING)).expect("Failed to decode json")
    }
}

impl Source for FileSource {
    fn functions(&self) -> Result<Vec<FunctionInfo>, SourceErr> {
        Ok(serde_json::from_str(&self.read_file(suffix::FUNCTION_INFO))?)
    }

    fn instructions_at(&self, address: u64) -> Result<Vec<LOpInfo>, SourceErr> {
        let suffix = format!("{}_{:#X}", suffix::INSTRUCTIONS, address);
        Ok(serde_json::from_str(&self.read_file(&suffix))?)
    }

    fn register_profile(&self) -> Result<LRegInfo, SourceErr> {
        Ok(serde_json::from_str(&self.read_file(suffix::REGISTER))?)
    }

    fn flags(&self) -> Result<Vec<LFlagInfo>, SourceErr> {
        Ok(serde_json::from_str(&self.read_file(suffix::FLAG))?)
    }

    fn sections(&self) -> Result<Vec<LSectionInfo>, SourceErr> {
        Ok(serde_json::from_str(&self.read_file(suffix::SECTION))?)
    }
}

impl<R: R2Api> From<WrappedR2Api<R>> for FileSource {
    fn from(mut r2: WrappedR2Api<R>) -> FileSource {
        let bin_info = r2.borrow_mut().bin_info().expect("Failed to load bin_info");
        let fname = bin_info.core.unwrap().file.unwrap();
        let fname = Path::new(&fname).file_stem().unwrap();
        let mut dir = PathBuf::from(".");
        dir.push(&fname);
        fs::create_dir_all(&dir).expect("Failed to create directory");
        let mut fsource = FileSource {
            dir: dir.to_str().unwrap().to_owned(),
            base_name: fname.to_str().unwrap().to_owned(),
        };

        {
            let fns = r2.functions().expect("Failed to load function info");
            {
                let json_str = serde_json::to_string(&fns).expect("Failed to encode to json");
                fsource.write_file(suffix::FUNCTION_INFO, &json_str);
            }

            for f in fns {
                let result = r2.instructions_at(f.offset.unwrap()).expect("Failed to load instructions");
                let json_str = serde_json::to_string(&result).expect("Failed to encode to json");
                let suffix = format!("{}_{:#X}", suffix::INSTRUCTIONS, f.offset.unwrap());
                fsource.write_file(&suffix, &json_str)
            }

            {
                let reg = r2.register_profile().expect("Failed to load register info");
                let json_str = serde_json::to_string(&reg).expect("Failed to encode to json");
                fsource.write_file(suffix::REGISTER, &json_str);
            }

            {
                let flags = r2.flags().expect("Failed to load flag info");
                let json_str = serde_json::to_string(&flags).expect("Failed to encode to json");
                fsource.write_file(suffix::FLAG, &json_str);
            }

            {
                let sections = r2.sections().expect("Failed to load section map");
                let json_str = serde_json::to_string(&sections).expect("Failed to encode to json");
                fsource.write_file(suffix::SECTION, &json_str);
            }

            {
                let strings = r2.borrow_mut().strings(false).expect("Unable to load String info from r2");
                let json_str = serde_json::to_string(&strings).expect("Failed to encode to json");
                fsource.write_file(suffix::STRING, &json_str);
            }
        }

        fsource
    }
}
<|endoftext|>


<commit_before>#[cfg(feature = "tokio-runtime")]
use tokio::net::{TcpListener, TcpStream};

#[cfg(feature = "async-std-runtime")]
use async_std::net::{TcpListener, TcpStream};

use super::make_framed;
use super::AcceptStopHandle;
use crate::async_rt;
use crate::codec::FramedIo;
use crate::endpoint::{Endpoint, Host, Port};
use crate::task_handle::TaskHandle;
use crate::ZmqResult;

use futures::{select, FutureExt};

pub(crate) async fn connect(host: &Host, port: Port) -> ZmqResult<(FramedIo, Endpoint)> {
    let raw_socket = TcpStream::connect((host.to_string().as_str(), port)).await?;
    raw_socket.set_nodelay(true)?;
    let peer_addr = raw_socket.peer_addr()?;

    Ok((make_framed(raw_socket), Endpoint::from_tcp_addr(peer_addr)))
}

pub(crate) async fn begin_accept<T>(
    mut host: Host,
    port: Port,
    cback: impl Fn(ZmqResult<(FramedIo, Endpoint)>) -> T + Send + 'static,
) -> ZmqResult<(Endpoint, AcceptStopHandle)>
where
    T: std::future::Future<Output = ()> + Send + 'static,
{
    let listener = TcpListener::bind((host.to_string().as_str(), port)).await?;
    let resolved_addr = listener.local_addr()?;
    let (stop_channel, stop_callback) = futures::channel::oneshot::channel::<()>();
    let task_handle = async_rt::task::spawn(async move {
        let mut stop_callback = stop_callback.fuse();
        loop {
            select! {
                incoming = listener.accept().fuse() => {
                    let maybe_accepted: Result<_, _> = incoming.and_then(|(raw_socket, remote_addr)|{
                        raw_socket.set_nodelay(true).map(|_| {
                            (raw_socket, remote_addr)
                        })
                    }).map(|(raw_socket, remote_addr)| {
                        (make_framed(raw_socket), Endpoint::from_tcp_addr(remote_addr))
                    }).map_err(|err| err.into());
                    async_rt::task::spawn(cback(maybe_accepted));
                },
                _ = stop_callback => {
                    break
                }
            }
        }
        Ok(())
    });
    debug_assert_ne!(resolved_addr.port(), 0);
    let port = resolved_addr.port();
    let resolved_host: Host = resolved_addr.ip().into();
    if let Host::Ipv4(ip) = host {
        debug_assert_eq!(ip, resolved_addr.ip());
        host = resolved_host;
    } else if let Host::Ipv6(ip) = host {
        debug_assert_eq!(ip, resolved_addr.ip());
        host = resolved_host;
    }
    Ok((
        Endpoint::Tcp(host, port),
        AcceptStopHandle(TaskHandle::new(stop_channel, task_handle)),
    ))
}
<commit_msg>Disable tcp set_nodelay on windows.<commit_after>#[cfg(feature = "tokio-runtime")]
use tokio::net::{TcpListener, TcpStream};

#[cfg(feature = "async-std-runtime")]
use async_std::net::{TcpListener, TcpStream};

use super::make_framed;
use super::AcceptStopHandle;
use crate::async_rt;
use crate::codec::FramedIo;
use crate::endpoint::{Endpoint, Host, Port};
use crate::task_handle::TaskHandle;
use crate::ZmqResult;

use futures::{select, FutureExt};

pub(crate) async fn connect(host: &Host, port: Port) -> ZmqResult<(FramedIo, Endpoint)> {
    let raw_socket = TcpStream::connect((host.to_string().as_str(), port)).await?;
    // For some reason set_nodelay doesn't work on windows. See
    // https://github.com/zeromq/zmq.rs/issues/148 for details
    #[cfg(not(windows))]
    raw_socket.set_nodelay(true)?;
    let peer_addr = raw_socket.peer_addr()?;

    Ok((make_framed(raw_socket), Endpoint::from_tcp_addr(peer_addr)))
}

pub(crate) async fn begin_accept<T>(
    mut host: Host,
    port: Port,
    cback: impl Fn(ZmqResult<(FramedIo, Endpoint)>) -> T + Send + 'static,
) -> ZmqResult<(Endpoint, AcceptStopHandle)>
where
    T: std::future::Future<Output = ()> + Send + 'static,
{
    let listener = TcpListener::bind((host.to_string().as_str(), port)).await?;
    let resolved_addr = listener.local_addr()?;
    let (stop_channel, stop_callback) = futures::channel::oneshot::channel::<()>();
    let task_handle = async_rt::task::spawn(async move {
        let mut stop_callback = stop_callback.fuse();
        loop {
            select! {
                incoming = listener.accept().fuse() => {
                    let maybe_accepted: Result<_, _> = incoming.and_then(|(raw_socket, remote_addr)|{
                        raw_socket.set_nodelay(true).map(|_| {
                            (raw_socket, remote_addr)
                        })
                    }).map(|(raw_socket, remote_addr)| {
                        (make_framed(raw_socket), Endpoint::from_tcp_addr(remote_addr))
                    }).map_err(|err| err.into());
                    async_rt::task::spawn(cback(maybe_accepted));
                },
                _ = stop_callback => {
                    break
                }
            }
        }
        Ok(())
    });
    debug_assert_ne!(resolved_addr.port(), 0);
    let port = resolved_addr.port();
    let resolved_host: Host = resolved_addr.ip().into();
    if let Host::Ipv4(ip) = host {
        debug_assert_eq!(ip, resolved_addr.ip());
        host = resolved_host;
    } else if let Host::Ipv6(ip) = host {
        debug_assert_eq!(ip, resolved_addr.ip());
        host = resolved_host;
    }
    Ok((
        Endpoint::Tcp(host, port),
        AcceptStopHandle(TaskHandle::new(stop_channel, task_handle)),
    ))
}
<|endoftext|>