<commit_before>//! This module implements config access.
use std::io;
use std::fs;
use std::env;
use std::path::PathBuf;

use dotenv;
use log;
use java_properties;
use clap::ArgMatches;
use url::Url;
use ini::Ini;

use prelude::*;
use constants::{DEFAULT_URL, VERSION, PROTOCOL_VERSION};

/// Represents the auth information
#[derive(Debug, Clone)]
pub enum Auth {
    Key(String),
    Token(String),
}

/// Represents a DSN
#[derive(Debug, Clone)]
pub struct Dsn {
    pub host: String,
    pub protocol: String,
    pub port: u16,
    pub client_id: String,
    pub secret: String,
    pub project_id: u64,
}

impl Dsn {
    /// Parses a Dsn from a given string.
    fn from_str(dsn: &str) -> Result<Dsn> {
        let url = Url::parse(dsn)?;
        let project_id = if let Some(component_iter) = url.path_segments() {
            let components: Vec<_> = component_iter.collect();
            if components.len() != 1 {
                fail!("invalid dsn: invalid project ID");
            }
            components[0].parse().or(Err("invalid dsn: invalid project id"))?
        } else {
            fail!("invalid dsn: missing project ID");
        };
        if !(url.scheme() == "http" || url.scheme() == "https") {
            fail!(format!("invalid dsn: unknown protocol '{}'", url.scheme()));
        }

        if url.username() == "" {
            fail!("invalid dsn: missing client id");
        }

        Ok(Dsn {
            protocol: url.scheme().into(),
            host: url.host_str().ok_or("invalid dsn: missing host")?.into(),
            port: url.port_or_known_default().unwrap(),
            client_id: url.username().into(),
            secret: url.password().ok_or("invalid dsn: missing secret")?.into(),
            project_id: project_id,
        })
    }

    /// Returns the URL where events should be sent.
    pub fn get_submit_url(&self) -> String {
        format!("{}://{}:{}/api/{}/store/",
                self.protocol,
                self.host,
                self.port,
                self.project_id)
    }

    /// Returns the given auth header (ts is the timestamp of the event)
    pub fn get_auth_header(&self, ts: f64) -> String {
        format!("Sentry \
            sentry_timestamp={}, \
            sentry_client=sentry-cli/{}, \
            sentry_version={}, \
            sentry_key={}, \
            sentry_secret={}",
                ts,
                VERSION,
                PROTOCOL_VERSION,
                self.client_id,
                self.secret)
    }
}

pub fn prepare_environment() {
    dotenv::dotenv().ok();
}

/// Represents the `sentry-cli` config.
#[derive(Clone)]
pub struct Config {
    pub filename: PathBuf,
    pub auth: Option<Auth>,
    pub url: String,
    pub log_level: log::LogLevelFilter,
    pub ini: Ini,
}

impl Config {
    /// Loads the CLI config from the default location and returns it.
    pub fn from_cli_config() -> Result<Config> {
        let (filename, ini) = load_cli_config()?;
        Ok(Config {
            filename: filename,
            auth: get_default_auth(&ini),
            url: get_default_url(&ini),
            log_level: get_default_log_level(&ini)?,
            ini: ini,
        })
    }

    /// Update the environment based on the config
    pub fn configure_environment(&self) {
        if !env::var("http_proxy").is_ok() {
            if let Some(proxy) = self.get_proxy_url() {
                env::set_var("http_proxy", proxy);
            }
        }
        #[cfg(not(windows))]
        {
            use openssl_probe::init_ssl_cert_env_vars;
            init_ssl_cert_env_vars();
        }
    }

    /// Returns the base url (without trailing slashes)
    pub fn get_base_url(&self) -> Result<&str> {
        let base = self.url.trim_right_matches('/');
        if !base.starts_with("http://") && !base.starts_with("https://") {
            fail!("bad sentry url: unknown scheme ({})", base);
        }
        if base.matches('/').count() != 2 {
            fail!("bad sentry url: not on URL root ({})", base);
        }
        Ok(base)
    }

    /// Returns the API URL for a path
    pub fn get_api_endpoint(&self, path: &str) -> Result<String> {
        let base = self.get_base_url()?;
        Ok(format!("{}/api/0/{}", base, path.trim_left_matches('/')))
    }

    /// Indicates whether keepalive support should be enabled.  This
    /// mostly corresponds to an ini config but also has some sensible
    /// default handling.
    pub fn allow_keepalive(&self) -> bool {
        let val = self.ini.get_from(Some("http"), "keepalive");
        match val {
            // keepalive is broken on our dev server.  Since this makes local development
            // quite frustrating we disable keepalive (handle reuse) when we connect to
            // unprotected servers where it does not matter that much.
            None => !self.has_insecure_server(),
            Some(val) => val == "true",
        }
    }

    /// Returns the proxy URL if defined.
    fn get_proxy_url(&self) -> Option<&str> {
        self.ini.get_from(Some("http"), "proxy_url")
    }

    /// Returns the proxy username if defined.
    pub fn get_proxy_username(&self) -> Option<&str> {
        self.ini.get_from(Some("http"), "proxy_username")
    }

    /// Returns the proxy password if defined.
    pub fn get_proxy_password(&self) -> Option<&str> {
        self.ini.get_from(Some("http"), "proxy_password")
    }

    /// Indicates if SSL is enabled or disabled for the server.
    pub fn has_insecure_server(&self) -> bool {
        self.url.starts_with("http://")
    }

    /// Indicates whether SSL verification should be on or off.
    pub fn should_verify_ssl(&self) -> bool {
        let val = self.ini.get_from(Some("http"), "verify_ssl");
        match val {
            None => true,
            Some(val) => val == "true",
        }
    }

    /// Controls the SSL revocation check on windows.  This can be used as a
    /// workaround for misconfigured local SSL proxies.
    pub fn disable_ssl_revocation_check(&self) -> bool {
        let val = self.ini.get_from(Some("http"), "check_ssl_revoke");
        match val {
            None => true,
            Some(val) => val == "true",
        }
    }

    /// Given a match object from clap, this returns the org from it.
    pub fn get_org(&self, matches: &ArgMatches) -> Result<String> {
        Ok(matches.value_of("org")
               .map(|x| x.to_owned())
               .or_else(|| env::var("SENTRY_ORG").ok())
               .or_else(|| self.ini.get_from(Some("defaults"), "org").map(|x| x.to_owned()))
               .ok_or("An organization slug is required (provide with --org)")?)
    }

    /// Given a match object from clap, this returns a tuple in the
    /// form `(org, project)` which can either come from the match
    /// object or some defaults (envvar, ini etc.).
    pub fn get_org_and_project(&self, matches: &ArgMatches) -> Result<(String, String)> {
        let org = self.get_org(matches)?;
        let project = if let Some(project) = matches.value_of("project") {
            project.to_owned()
        } else {
            self.get_project_default()?
        };
        Ok((org, project))
    }

    /// Return the default value for a project.
    pub fn get_project_default(&self) -> Result<String> {
        Ok(env::var("SENTRY_PROJECT").ok()
            .or_else(|| self.ini.get_from(Some("defaults"), "project").map(|x| x.to_owned()))
            .ok_or("A project slug is required")?)
    }

    /// Returns the defaults for org and project.
    pub fn get_org_and_project_defaults(&self) -> (Option<String>, Option<String>) {
        (env::var("SENTRY_ORG")
             .ok()
             .or_else(|| self.ini.get_from(Some("defaults"), "org").map(|x| x.to_owned())),
         env::var("SENTRY_PROJECT")
             .ok()
             .or_else(|| self.ini.get_from(Some("defaults"), "project").map(|x| x.to_owned())))
    }

    /// Returns true if notifications should be displayed
    pub fn show_notifications(&self) -> Result<bool> {
        Ok(self.ini.get_from(Some("ui"), "show_notifications")
            .map(|x| x == "true")
            .unwrap_or(true))
    }

    /// Returns the maximum dsym upload size
    pub fn get_max_dsym_upload_size(&self) -> Result<u64> {
        Ok(self.ini.get_from(Some("dsym"), "max_upload_size")
            .and_then(|x| x.parse().ok())
            .unwrap_or(35 * 1024 * 1024))
    }

    /// Return the DSN
    pub fn get_dsn(&self) -> Result<Dsn> {
        if let Some(ref val) = env::var("SENTRY_DSN").ok() {
            Dsn::from_str(val)
        } else if let Some(val) = self.ini.get_from(Some("auth"), "dsn") {
            Dsn::from_str(val)
        } else {
            fail!("No DSN provided");
        }
    }

    /// Should we nag about updates?
    pub fn disable_update_nagger(&self) -> bool {
        if let Ok(var) = env::var("SENTRY_DISABLE_UPDATE_CHECK") {
            &var == "1" || &var == "true"
        } else {
            if let Some(val) = self.ini.get_from(Some("update"), "disable_check") {
                val == "true"
            } else {
                false
            }
        }
    }
}

fn find_project_config_file() -> Option<PathBuf> {
    env::current_dir().ok().and_then(|mut path| {
        loop {
            path.push(".sentryclirc");
            if path.exists() {
                return Some(path);
            }
            path.set_file_name("sentrycli.ini");
            if path.exists() {
                return Some(path);
            }
            path.pop();
            if !path.pop() {
                return None;
            }
        }
    })
}

fn load_cli_config() -> Result<(PathBuf, Ini)> {
    let mut home_fn = env::home_dir().ok_or("Could not find home dir")?;
    home_fn.push(".sentryclirc");

    let mut rv = match fs::File::open(&home_fn) {
        Ok(mut file) => Ini::read_from(&mut file)?,
        Err(err) => {
            if err.kind() == io::ErrorKind::NotFound {
                Ini::new()
            } else {
                return Err(err).chain_err(
                    || "Failed to load .sentryclirc file from the home folder.");
            }
        }
    };

    let (path, mut rv) = if let Some(project_config_path) = find_project_config_file() {
        let mut f = fs::File::open(&project_config_path)
            .chain_err(|| format!("Failed to load .sentryclirc file from project path ({})",
                project_config_path.display()))?;
        let ini = Ini::read_from(&mut f)?;
        for (section, props) in ini.iter() {
            for (key, value) in props {
                rv.set_to(section.clone(), key.clone(), value.to_owned());
            }
        }
        (project_config_path, rv)
    } else {
        (home_fn, rv)
    };

    if let Ok(prop_path) = env::var("SENTRY_PROPERTIES") {
        match fs::File::open(&prop_path) {
            Ok(f) => {
                let props = match java_properties::read(f) {
                    Ok(props) => props,
                    Err(err) => {
                        return Err(Error::from(format!(
                            "Could not load java style properties file: {}", err)));
                    }
                };
                for (key, value) in props {
                    let mut iter = key.rsplitn(2, '.');
                    if let Some(key) = iter.next() {
                        let section = iter.next();
                        rv.set_to(section, key.to_string(), value);
                    }
                }
            },
            Err(err) => {
                if err.kind() != io::ErrorKind::NotFound {
                    return Err(Error::from(err)).chain_err(
                        || format!("Failed to load file referenced by SENTRY_PROPERTIES ({})",
                                   &prop_path));
                }
            }
        }
    }

    Ok((path, rv))
}

fn get_default_auth(ini: &Ini) -> Option<Auth> {
    if let Some(ref val) = env::var("SENTRY_AUTH_TOKEN").ok() {
        Some(Auth::Token(val.to_owned()))
    } else if let Some(ref val) = env::var("SENTRY_API_KEY").ok() {
        Some(Auth::Key(val.to_owned()))
    } else if let Some(val) = ini.get_from(Some("auth"), "token") {
        Some(Auth::Token(val.to_owned()))
    } else if let Some(val) = ini.get_from(Some("auth"), "api_key") {
        Some(Auth::Key(val.to_owned()))
    } else {
        None
    }
}

fn get_default_url(ini: &Ini) -> String {
    if let Some(ref val) = env::var("SENTRY_URL").ok() {
        val.to_owned()
    } else if let Some(val) = ini.get_from(Some("defaults"), "url") {
        val.to_owned()
    } else {
        DEFAULT_URL.to_owned()
    }
}

fn get_default_log_level(ini: &Ini) -> Result<log::LogLevelFilter> {
    if let Ok(level_str) = env::var("SENTRY_LOG_LEVEL") {
        if let Ok(level) = level_str.parse() {
            return Ok(level);
        }
    }

    if let Some(level_str) = ini.get_from(Some("log"), "level") {
        if let Ok(level) = level_str.parse() {
            return Ok(level);
        }
    }

    Ok(log::LogLevelFilter::Warn)
}
<commit_msg>feat: No longer require the DSN secret (#189)<commit_after>//! This module implements config access.
use std::io;
use std::fs;
use std::env;
use std::path::PathBuf;

use dotenv;
use log;
use java_properties;
use clap::ArgMatches;
use url::Url;
use ini::Ini;

use prelude::*;
use constants::{DEFAULT_URL, VERSION, PROTOCOL_VERSION};

/// Represents the auth information
#[derive(Debug, Clone)]
pub enum Auth {
    Key(String),
    Token(String),
}

/// Represents a DSN
#[derive(Debug, Clone)]
pub struct Dsn {
    pub host: String,
    pub protocol: String,
    pub port: u16,
    pub client_id: String,
    pub secret: Option<String>,
    pub project_id: u64,
}

impl Dsn {
    /// Parses a Dsn from a given string.
    fn from_str(dsn: &str) -> Result<Dsn> {
        let url = Url::parse(dsn)?;
        let project_id = if let Some(component_iter) = url.path_segments() {
            let components: Vec<_> = component_iter.collect();
            if components.len() != 1 {
                fail!("invalid dsn: invalid project ID");
            }
            components[0].parse().or(Err("invalid dsn: invalid project id"))?
        } else {
            fail!("invalid dsn: missing project ID");
        };
        if !(url.scheme() == "http" || url.scheme() == "https") {
            fail!(format!("invalid dsn: unknown protocol '{}'", url.scheme()));
        }

        if url.username() == "" {
            fail!("invalid dsn: missing client id");
        }

        Ok(Dsn {
            protocol: url.scheme().into(),
            host: url.host_str().ok_or("invalid dsn: missing host")?.into(),
            port: url.port_or_known_default().unwrap(),
            client_id: url.username().into(),
            secret: url.password().map(|x| x.to_string()),
            project_id: project_id,
        })
    }

    /// Returns the URL where events should be sent.
    pub fn get_submit_url(&self) -> String {
        format!("{}://{}:{}/api/{}/store/",
                self.protocol,
                self.host,
                self.port,
                self.project_id)
    }

    /// Returns the given auth header (ts is the timestamp of the event)
    pub fn get_auth_header(&self, ts: f64) -> String {
        let mut rv = format!("Sentry \
            sentry_timestamp={}, \
            sentry_client=sentry-cli/{}, \
            sentry_version={}, \
            sentry_key={}",
                ts,
                VERSION,
                PROTOCOL_VERSION,
                self.client_id);
        if let Some(ref secret) = self.secret {
            rv = format!("{}, sentry_secret={}", rv, secret);
        }
        rv
    }
}

pub fn prepare_environment() {
    dotenv::dotenv().ok();
}

/// Represents the `sentry-cli` config.
#[derive(Clone)]
pub struct Config {
    pub filename: PathBuf,
    pub auth: Option<Auth>,
    pub url: String,
    pub log_level: log::LogLevelFilter,
    pub ini: Ini,
}

impl Config {
    /// Loads the CLI config from the default location and returns it.
    pub fn from_cli_config() -> Result<Config> {
        let (filename, ini) = load_cli_config()?;
        Ok(Config {
            filename: filename,
            auth: get_default_auth(&ini),
            url: get_default_url(&ini),
            log_level: get_default_log_level(&ini)?,
            ini: ini,
        })
    }

    /// Update the environment based on the config
    pub fn configure_environment(&self) {
        if !env::var("http_proxy").is_ok() {
            if let Some(proxy) = self.get_proxy_url() {
                env::set_var("http_proxy", proxy);
            }
        }
        #[cfg(not(windows))]
        {
            use openssl_probe::init_ssl_cert_env_vars;
            init_ssl_cert_env_vars();
        }
    }

    /// Returns the base url (without trailing slashes)
    pub fn get_base_url(&self) -> Result<&str> {
        let base = self.url.trim_right_matches('/');
        if !base.starts_with("http://") && !base.starts_with("https://") {
            fail!("bad sentry url: unknown scheme ({})", base);
        }
        if base.matches('/').count() != 2 {
            fail!("bad sentry url: not on URL root ({})", base);
        }
        Ok(base)
    }

    /// Returns the API URL for a path
    pub fn get_api_endpoint(&self, path: &str) -> Result<String> {
        let base = self.get_base_url()?;
        Ok(format!("{}/api/0/{}", base, path.trim_left_matches('/')))
    }

    /// Indicates whether keepalive support should be enabled.  This
    /// mostly corresponds to an ini config but also has some sensible
    /// default handling.
    pub fn allow_keepalive(&self) -> bool {
        let val = self.ini.get_from(Some("http"), "keepalive");
        match val {
            // keepalive is broken on our dev server.  Since this makes local development
            // quite frustrating we disable keepalive (handle reuse) when we connect to
            // unprotected servers where it does not matter that much.
            None => !self.has_insecure_server(),
            Some(val) => val == "true",
        }
    }

    /// Returns the proxy URL if defined.
    fn get_proxy_url(&self) -> Option<&str> {
        self.ini.get_from(Some("http"), "proxy_url")
    }

    /// Returns the proxy username if defined.
    pub fn get_proxy_username(&self) -> Option<&str> {
        self.ini.get_from(Some("http"), "proxy_username")
    }

    /// Returns the proxy password if defined.
    pub fn get_proxy_password(&self) -> Option<&str> {
        self.ini.get_from(Some("http"), "proxy_password")
    }

    /// Indicates if SSL is enabled or disabled for the server.
    pub fn has_insecure_server(&self) -> bool {
        self.url.starts_with("http://")
    }

    /// Indicates whether SSL verification should be on or off.
    pub fn should_verify_ssl(&self) -> bool {
        let val = self.ini.get_from(Some("http"), "verify_ssl");
        match val {
            None => true,
            Some(val) => val == "true",
        }
    }

    /// Controls the SSL revocation check on windows.  This can be used as a
    /// workaround for misconfigured local SSL proxies.
    pub fn disable_ssl_revocation_check(&self) -> bool {
        let val = self.ini.get_from(Some("http"), "check_ssl_revoke");
        match val {
            None => true,
            Some(val) => val == "true",
        }
    }

    /// Given a match object from clap, this returns the org from it.
    pub fn get_org(&self, matches: &ArgMatches) -> Result<String> {
        Ok(matches.value_of("org")
               .map(|x| x.to_owned())
               .or_else(|| env::var("SENTRY_ORG").ok())
               .or_else(|| self.ini.get_from(Some("defaults"), "org").map(|x| x.to_owned()))
               .ok_or("An organization slug is required (provide with --org)")?)
    }

    /// Given a match object from clap, this returns a tuple in the
    /// form `(org, project)` which can either come from the match
    /// object or some defaults (envvar, ini etc.).
    pub fn get_org_and_project(&self, matches: &ArgMatches) -> Result<(String, String)> {
        let org = self.get_org(matches)?;
        let project = if let Some(project) = matches.value_of("project") {
            project.to_owned()
        } else {
            self.get_project_default()?
        };
        Ok((org, project))
    }

    /// Return the default value for a project.
    pub fn get_project_default(&self) -> Result<String> {
        Ok(env::var("SENTRY_PROJECT").ok()
            .or_else(|| self.ini.get_from(Some("defaults"), "project").map(|x| x.to_owned()))
            .ok_or("A project slug is required")?)
    }

    /// Returns the defaults for org and project.
    pub fn get_org_and_project_defaults(&self) -> (Option<String>, Option<String>) {
        (env::var("SENTRY_ORG")
             .ok()
             .or_else(|| self.ini.get_from(Some("defaults"), "org").map(|x| x.to_owned())),
         env::var("SENTRY_PROJECT")
             .ok()
             .or_else(|| self.ini.get_from(Some("defaults"), "project").map(|x| x.to_owned())))
    }

    /// Returns true if notifications should be displayed
    pub fn show_notifications(&self) -> Result<bool> {
        Ok(self.ini.get_from(Some("ui"), "show_notifications")
            .map(|x| x == "true")
            .unwrap_or(true))
    }

    /// Returns the maximum dsym upload size
    pub fn get_max_dsym_upload_size(&self) -> Result<u64> {
        Ok(self.ini.get_from(Some("dsym"), "max_upload_size")
            .and_then(|x| x.parse().ok())
            .unwrap_or(35 * 1024 * 1024))
    }

    /// Return the DSN
    pub fn get_dsn(&self) -> Result<Dsn> {
        if let Some(ref val) = env::var("SENTRY_DSN").ok() {
            Dsn::from_str(val)
        } else if let Some(val) = self.ini.get_from(Some("auth"), "dsn") {
            Dsn::from_str(val)
        } else {
            fail!("No DSN provided");
        }
    }

    /// Should we nag about updates?
    pub fn disable_update_nagger(&self) -> bool {
        if let Ok(var) = env::var("SENTRY_DISABLE_UPDATE_CHECK") {
            &var == "1" || &var == "true"
        } else {
            if let Some(val) = self.ini.get_from(Some("update"), "disable_check") {
                val == "true"
            } else {
                false
            }
        }
    }
}

fn find_project_config_file() -> Option<PathBuf> {
    env::current_dir().ok().and_then(|mut path| {
        loop {
            path.push(".sentryclirc");
            if path.exists() {
                return Some(path);
            }
            path.set_file_name("sentrycli.ini");
            if path.exists() {
                return Some(path);
            }
            path.pop();
            if !path.pop() {
                return None;
            }
        }
    })
}

fn load_cli_config() -> Result<(PathBuf, Ini)> {
    let mut home_fn = env::home_dir().ok_or("Could not find home dir")?;
    home_fn.push(".sentryclirc");

    let mut rv = match fs::File::open(&home_fn) {
        Ok(mut file) => Ini::read_from(&mut file)?,
        Err(err) => {
            if err.kind() == io::ErrorKind::NotFound {
                Ini::new()
            } else {
                return Err(err).chain_err(
                    || "Failed to load .sentryclirc file from the home folder.");
            }
        }
    };

    let (path, mut rv) = if let Some(project_config_path) = find_project_config_file() {
        let mut f = fs::File::open(&project_config_path)
            .chain_err(|| format!("Failed to load .sentryclirc file from project path ({})",
                project_config_path.display()))?;
        let ini = Ini::read_from(&mut f)?;
        for (section, props) in ini.iter() {
            for (key, value) in props {
                rv.set_to(section.clone(), key.clone(), value.to_owned());
            }
        }
        (project_config_path, rv)
    } else {
        (home_fn, rv)
    };

    if let Ok(prop_path) = env::var("SENTRY_PROPERTIES") {
        match fs::File::open(&prop_path) {
            Ok(f) => {
                let props = match java_properties::read(f) {
                    Ok(props) => props,
                    Err(err) => {
                        return Err(Error::from(format!(
                            "Could not load java style properties file: {}", err)));
                    }
                };
                for (key, value) in props {
                    let mut iter = key.rsplitn(2, '.');
                    if let Some(key) = iter.next() {
                        let section = iter.next();
                        rv.set_to(section, key.to_string(), value);
                    }
                }
            },
            Err(err) => {
                if err.kind() != io::ErrorKind::NotFound {
                    return Err(Error::from(err)).chain_err(
                        || format!("Failed to load file referenced by SENTRY_PROPERTIES ({})",
                                   &prop_path));
                }
            }
        }
    }

    Ok((path, rv))
}

fn get_default_auth(ini: &Ini) -> Option<Auth> {
    if let Some(ref val) = env::var("SENTRY_AUTH_TOKEN").ok() {
        Some(Auth::Token(val.to_owned()))
    } else if let Some(ref val) = env::var("SENTRY_API_KEY").ok() {
        Some(Auth::Key(val.to_owned()))
    } else if let Some(val) = ini.get_from(Some("auth"), "token") {
        Some(Auth::Token(val.to_owned()))
    } else if let Some(val) = ini.get_from(Some("auth"), "api_key") {
        Some(Auth::Key(val.to_owned()))
    } else {
        None
    }
}

fn get_default_url(ini: &Ini) -> String {
    if let Some(ref val) = env::var("SENTRY_URL").ok() {
        val.to_owned()
    } else if let Some(val) = ini.get_from(Some("defaults"), "url") {
        val.to_owned()
    } else {
        DEFAULT_URL.to_owned()
    }
}

fn get_default_log_level(ini: &Ini) -> Result<log::LogLevelFilter> {
    if let Ok(level_str) = env::var("SENTRY_LOG_LEVEL") {
        if let Ok(level) = level_str.parse() {
            return Ok(level);
        }
    }

    if let Some(level_str) = ini.get_from(Some("log"), "level") {
        if let Ok(level) = level_str.parse() {
            return Ok(level);
        }
    }

    Ok(log::LogLevelFilter::Warn)
}
<|endoftext|>


<commit_before>//! Common configuration utilities

use std::{
    env,
    fs::OpenOptions,
    io::{self, Read},
    path::{Path, PathBuf},
    str::FromStr,
};

use cfg_if::cfg_if;
use clap::{ArgMatches, ErrorKind as ClapErrorKind};
use directories::ProjectDirs;
use serde::Deserialize;

/// Default configuration file path
pub fn get_default_config_path() -> Option<PathBuf> {
    // config.json in the current working directory ($PWD)
    if let Ok(mut path) = env::current_dir() {
        path.push("config.json");

        if path.exists() {
            return Some(path);
        }
    } else {
        // config.json in the current working directory (relative path)
        let relative_path = PathBuf::from("config.json");
        if relative_path.exists() {
            return Some(relative_path);
        }
    }

    // System standard directories
    if let Some(project_dirs) = ProjectDirs::from("org", "shadowsocks", "shadowsocks-rust") {
        // Linux: $XDG_CONFIG_HOME/shadowsocks-rust/config.json
        //        $HOME/.config/shadowsocks-rust/config.json
        // macOS: $HOME/Library/Application Support/org.shadowsocks.shadowsocks-rust/config.json
        // Windows: {FOLDERID_RoamingAppData}/shadowsocks/shadowsocks-rust/config/config.json

        let mut config_path = project_dirs.config_dir().to_path_buf();
        config_path.push("config.json");

        if config_path.exists() {
            return Some(config_path);
        }
    }

    // UNIX systems, XDG Base Directory
    // https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html
    #[cfg(unix)]
    if let Ok(base_directories) = xdg::BaseDirectories::with_prefix("shadowsocks-rust") {
        // $XDG_CONFIG_HOME/shadowsocks-rust/config.json
        // for dir in $XDG_CONFIG_DIRS; $dir/shadowsocks-rust/config.json
        if let Some(config_path) = base_directories.find_config_file("config.json") {
            return Some(config_path);
        }
    }

    // UNIX global configuration file
    #[cfg(unix)]
    {
        let global_config_path = Path::new("/etc/shadowsocks-rust/config.json");
        if global_config_path.exists() {
            return Some(global_config_path.to_path_buf());
        }
    }

    None
}

/// Error while reading `Config`
#[derive(thiserror::Error, Debug)]
pub enum ConfigError {
    /// Input/Output error
    #[error("{0}")]
    IoError(#[from] io::Error),
    /// JSON parsing error
    #[error("{0}")]
    JsonError(#[from] json5::Error),
    /// Invalid value
    #[error("Invalid value: {0}")]
    InvalidValue(String),
}

/// Configuration Options for shadowsocks service runnables
#[derive(Debug, Clone, Default)]
pub struct Config {
    /// Logger configuration
    #[cfg(feature = "logging")]
    pub log: LogConfig,

    /// Runtime configuration
    pub runtime: RuntimeConfig,
}

impl Config {
    /// Load `Config` from file
    pub fn load_from_file<P: AsRef<Path>>(filename: &P) -> Result<Config, ConfigError> {
        let filename = filename.as_ref();

        let mut reader = OpenOptions::new().read(true).open(filename)?;
        let mut content = String::new();
        reader.read_to_string(&mut content)?;

        Config::load_from_str(&content)
    }

    /// Load `Config` from string
    pub fn load_from_str(s: &str) -> Result<Config, ConfigError> {
        let ssconfig = json5::from_str(s)?;
        Config::load_from_ssconfig(ssconfig)
    }

    fn load_from_ssconfig(ssconfig: SSConfig) -> Result<Config, ConfigError> {
        let mut config = Config::default();

        #[cfg(feature = "logging")]
        if let Some(log) = ssconfig.log {
            let mut nlog = LogConfig::default();
            if let Some(level) = log.level {
                nlog.level = level;
            }

            if let Some(format) = log.format {
                let mut nformat = LogFormatConfig::default();
                if let Some(without_time) = format.without_time {
                    nformat.without_time = without_time;
                }
                nlog.format = nformat;
            }

            if let Some(config_path) = log.config_path {
                nlog.config_path = Some(PathBuf::from(config_path));
            }

            config.log = nlog;
        }

        if let Some(runtime) = ssconfig.runtime {
            let mut nruntime = RuntimeConfig::default();

            #[cfg(feature = "multi-threaded")]
            if let Some(worker_count) = runtime.worker_count {
                nruntime.worker_count = Some(worker_count);
            }

            if let Some(mode) = runtime.mode {
                match mode.parse::<RuntimeMode>() {
                    Ok(m) => nruntime.mode = m,
                    Err(..) => return Err(ConfigError::InvalidValue(mode)),
                }
            }

            config.runtime = nruntime;
        }

        Ok(config)
    }

    /// Set by command line options
    pub fn set_options(&mut self, matches: &ArgMatches) {
        #[cfg(feature = "logging")]
        {
            let debug_level = matches.occurrences_of("VERBOSE");
            if debug_level > 0 {
                self.log.level = debug_level as u32;
            }

            if matches.is_present("LOG_WITHOUT_TIME") {
                self.log.format.without_time = true;
            }

            if let Some(log_config) = matches.value_of("LOG_CONFIG") {
                self.log.config_path = Some(log_config.into());
            }
        }

        #[cfg(feature = "multi-threaded")]
        if matches.is_present("SINGLE_THREADED") {
            self.runtime.mode = RuntimeMode::SingleThread;
        }

        #[cfg(feature = "multi-threaded")]
        match matches.value_of_t::<usize>("WORKER_THREADS") {
            Ok(worker_count) => self.runtime.worker_count = Some(worker_count),
            Err(ref err) if err.kind == ClapErrorKind::ArgumentNotFound => {}
            Err(err) => err.exit(),
        }

        let _ = matches;
    }
}

/// Logger configuration
#[cfg(feature = "logging")]
#[derive(Debug, Clone)]
pub struct LogConfig {
    /// Default logger log level, [0, 3]
    pub level: u32,
    /// Default logger format configuration
    pub format: LogFormatConfig,
    /// Logging configuration file path
    pub config_path: Option<PathBuf>,
}

#[cfg(feature = "logging")]
impl Default for LogConfig {
    fn default() -> LogConfig {
        LogConfig {
            level: 0,
            format: LogFormatConfig::default(),
            config_path: None,
        }
    }
}

/// Logger format configuration
#[cfg(feature = "logging")]
#[derive(Debug, Clone)]
pub struct LogFormatConfig {
    pub without_time: bool,
}

#[cfg(feature = "logging")]
impl Default for LogFormatConfig {
    fn default() -> LogFormatConfig {
        LogFormatConfig { without_time: false }
    }
}

/// Runtime mode (Tokio)
#[derive(Debug, Clone, Copy)]
pub enum RuntimeMode {
    /// Single-Thread Runtime
    SingleThread,
    /// Multi-Thread Runtime
    #[cfg(feature = "multi-threaded")]
    MultiThread,
}

impl Default for RuntimeMode {
    fn default() -> RuntimeMode {
        cfg_if! {
            if #[cfg(feature = "multi-threaded")] {
                RuntimeMode::MultiThread
            } else {
                RuntimeMode::SingleThread
            }
        }
    }
}

/// Parse `RuntimeMode` from string error
#[derive(Debug)]
pub struct RuntimeModeError;

impl FromStr for RuntimeMode {
    type Err = RuntimeModeError;

    fn from_str(s: &str) -> Result<RuntimeMode, Self::Err> {
        match s {
            "single_thread" => Ok(RuntimeMode::SingleThread),
            #[cfg(feature = "multi-threaded")]
            "multi_thread" => Ok(RuntimeMode::MultiThread),
            _ => Err(RuntimeModeError),
        }
    }
}

/// Runtime configuration
#[derive(Debug, Clone, Default)]
pub struct RuntimeConfig {
    /// Multithread runtime worker count, CPU count if not configured
    #[cfg(feature = "multi-threaded")]
    pub worker_count: Option<usize>,
    /// Runtime Mode, single-thread, multi-thread
    pub mode: RuntimeMode,
}

#[derive(Deserialize)]
struct SSConfig {
    #[cfg(feature = "logging")]
    log: Option<SSLogConfig>,
    runtime: Option<SSRuntimeConfig>,
}

#[cfg(feature = "logging")]
#[derive(Deserialize)]
struct SSLogConfig {
    level: Option<u32>,
    format: Option<SSLogFormat>,
    config_path: Option<String>,
}

#[cfg(feature = "logging")]
#[derive(Deserialize)]
struct SSLogFormat {
    without_time: Option<bool>,
}

#[derive(Deserialize)]
struct SSRuntimeConfig {
    #[cfg(feature = "multi-threaded")]
    worker_count: Option<usize>,
    mode: Option<String>,
}
<commit_msg>clap::ErrorKind only used when multi-threaded feature is enabled<commit_after>//! Common configuration utilities

use std::{
    env,
    fs::OpenOptions,
    io::{self, Read},
    path::{Path, PathBuf},
    str::FromStr,
};

use cfg_if::cfg_if;
use clap::ArgMatches;
use directories::ProjectDirs;
use serde::Deserialize;

/// Default configuration file path
pub fn get_default_config_path() -> Option<PathBuf> {
    // config.json in the current working directory ($PWD)
    if let Ok(mut path) = env::current_dir() {
        path.push("config.json");

        if path.exists() {
            return Some(path);
        }
    } else {
        // config.json in the current working directory (relative path)
        let relative_path = PathBuf::from("config.json");
        if relative_path.exists() {
            return Some(relative_path);
        }
    }

    // System standard directories
    if let Some(project_dirs) = ProjectDirs::from("org", "shadowsocks", "shadowsocks-rust") {
        // Linux: $XDG_CONFIG_HOME/shadowsocks-rust/config.json
        //        $HOME/.config/shadowsocks-rust/config.json
        // macOS: $HOME/Library/Application Support/org.shadowsocks.shadowsocks-rust/config.json
        // Windows: {FOLDERID_RoamingAppData}/shadowsocks/shadowsocks-rust/config/config.json

        let mut config_path = project_dirs.config_dir().to_path_buf();
        config_path.push("config.json");

        if config_path.exists() {
            return Some(config_path);
        }
    }

    // UNIX systems, XDG Base Directory
    // https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html
    #[cfg(unix)]
    if let Ok(base_directories) = xdg::BaseDirectories::with_prefix("shadowsocks-rust") {
        // $XDG_CONFIG_HOME/shadowsocks-rust/config.json
        // for dir in $XDG_CONFIG_DIRS; $dir/shadowsocks-rust/config.json
        if let Some(config_path) = base_directories.find_config_file("config.json") {
            return Some(config_path);
        }
    }

    // UNIX global configuration file
    #[cfg(unix)]
    {
        let global_config_path = Path::new("/etc/shadowsocks-rust/config.json");
        if global_config_path.exists() {
            return Some(global_config_path.to_path_buf());
        }
    }

    None
}

/// Error while reading `Config`
#[derive(thiserror::Error, Debug)]
pub enum ConfigError {
    /// Input/Output error
    #[error("{0}")]
    IoError(#[from] io::Error),
    /// JSON parsing error
    #[error("{0}")]
    JsonError(#[from] json5::Error),
    /// Invalid value
    #[error("Invalid value: {0}")]
    InvalidValue(String),
}

/// Configuration Options for shadowsocks service runnables
#[derive(Debug, Clone, Default)]
pub struct Config {
    /// Logger configuration
    #[cfg(feature = "logging")]
    pub log: LogConfig,

    /// Runtime configuration
    pub runtime: RuntimeConfig,
}

impl Config {
    /// Load `Config` from file
    pub fn load_from_file<P: AsRef<Path>>(filename: &P) -> Result<Config, ConfigError> {
        let filename = filename.as_ref();

        let mut reader = OpenOptions::new().read(true).open(filename)?;
        let mut content = String::new();
        reader.read_to_string(&mut content)?;

        Config::load_from_str(&content)
    }

    /// Load `Config` from string
    pub fn load_from_str(s: &str) -> Result<Config, ConfigError> {
        let ssconfig = json5::from_str(s)?;
        Config::load_from_ssconfig(ssconfig)
    }

    fn load_from_ssconfig(ssconfig: SSConfig) -> Result<Config, ConfigError> {
        let mut config = Config::default();

        #[cfg(feature = "logging")]
        if let Some(log) = ssconfig.log {
            let mut nlog = LogConfig::default();
            if let Some(level) = log.level {
                nlog.level = level;
            }

            if let Some(format) = log.format {
                let mut nformat = LogFormatConfig::default();
                if let Some(without_time) = format.without_time {
                    nformat.without_time = without_time;
                }
                nlog.format = nformat;
            }

            if let Some(config_path) = log.config_path {
                nlog.config_path = Some(PathBuf::from(config_path));
            }

            config.log = nlog;
        }

        if let Some(runtime) = ssconfig.runtime {
            let mut nruntime = RuntimeConfig::default();

            #[cfg(feature = "multi-threaded")]
            if let Some(worker_count) = runtime.worker_count {
                nruntime.worker_count = Some(worker_count);
            }

            if let Some(mode) = runtime.mode {
                match mode.parse::<RuntimeMode>() {
                    Ok(m) => nruntime.mode = m,
                    Err(..) => return Err(ConfigError::InvalidValue(mode)),
                }
            }

            config.runtime = nruntime;
        }

        Ok(config)
    }

    /// Set by command line options
    pub fn set_options(&mut self, matches: &ArgMatches) {
        #[cfg(feature = "logging")]
        {
            let debug_level = matches.occurrences_of("VERBOSE");
            if debug_level > 0 {
                self.log.level = debug_level as u32;
            }

            if matches.is_present("LOG_WITHOUT_TIME") {
                self.log.format.without_time = true;
            }

            if let Some(log_config) = matches.value_of("LOG_CONFIG") {
                self.log.config_path = Some(log_config.into());
            }
        }

        #[cfg(feature = "multi-threaded")]
        if matches.is_present("SINGLE_THREADED") {
            self.runtime.mode = RuntimeMode::SingleThread;
        }

        #[cfg(feature = "multi-threaded")]
        match matches.value_of_t::<usize>("WORKER_THREADS") {
            Ok(worker_count) => self.runtime.worker_count = Some(worker_count),
            Err(ref err) if err.kind == clap::ErrorKind::ArgumentNotFound => {}
            Err(err) => err.exit(),
        }

        let _ = matches;
    }
}

/// Logger configuration
#[cfg(feature = "logging")]
#[derive(Debug, Clone)]
pub struct LogConfig {
    /// Default logger log level, [0, 3]
    pub level: u32,
    /// Default logger format configuration
    pub format: LogFormatConfig,
    /// Logging configuration file path
    pub config_path: Option<PathBuf>,
}

#[cfg(feature = "logging")]
impl Default for LogConfig {
    fn default() -> LogConfig {
        LogConfig {
            level: 0,
            format: LogFormatConfig::default(),
            config_path: None,
        }
    }
}

/// Logger format configuration
#[cfg(feature = "logging")]
#[derive(Debug, Clone)]
pub struct LogFormatConfig {
    pub without_time: bool,
}

#[cfg(feature = "logging")]
impl Default for LogFormatConfig {
    fn default() -> LogFormatConfig {
        LogFormatConfig { without_time: false }
    }
}

/// Runtime mode (Tokio)
#[derive(Debug, Clone, Copy)]
pub enum RuntimeMode {
    /// Single-Thread Runtime
    SingleThread,
    /// Multi-Thread Runtime
    #[cfg(feature = "multi-threaded")]
    MultiThread,
}

impl Default for RuntimeMode {
    fn default() -> RuntimeMode {
        cfg_if! {
            if #[cfg(feature = "multi-threaded")] {
                RuntimeMode::MultiThread
            } else {
                RuntimeMode::SingleThread
            }
        }
    }
}

/// Parse `RuntimeMode` from string error
#[derive(Debug)]
pub struct RuntimeModeError;

impl FromStr for RuntimeMode {
    type Err = RuntimeModeError;

    fn from_str(s: &str) -> Result<RuntimeMode, Self::Err> {
        match s {
            "single_thread" => Ok(RuntimeMode::SingleThread),
            #[cfg(feature = "multi-threaded")]
            "multi_thread" => Ok(RuntimeMode::MultiThread),
            _ => Err(RuntimeModeError),
        }
    }
}

/// Runtime configuration
#[derive(Debug, Clone, Default)]
pub struct RuntimeConfig {
    /// Multithread runtime worker count, CPU count if not configured
    #[cfg(feature = "multi-threaded")]
    pub worker_count: Option<usize>,
    /// Runtime Mode, single-thread, multi-thread
    pub mode: RuntimeMode,
}

#[derive(Deserialize)]
struct SSConfig {
    #[cfg(feature = "logging")]
    log: Option<SSLogConfig>,
    runtime: Option<SSRuntimeConfig>,
}

#[cfg(feature = "logging")]
#[derive(Deserialize)]
struct SSLogConfig {
    level: Option<u32>,
    format: Option<SSLogFormat>,
    config_path: Option<String>,
}

#[cfg(feature = "logging")]
#[derive(Deserialize)]
struct SSLogFormat {
    without_time: Option<bool>,
}

#[derive(Deserialize)]
struct SSRuntimeConfig {
    #[cfg(feature = "multi-threaded")]
    worker_count: Option<usize>,
    mode: Option<String>,
}
<|endoftext|>


<commit_before>use regex::Regex;
use std::collections::HashMap;
use std::fs;
use std::io::{Error, Read};
use std::path::Path;
use toml;

#[derive(Debug)]
pub struct Benchmark {
    pub name: String,
    pub cmd: String,
    pub args: Vec<String>,
    pub result_expr: Vec<Regex>,
    pub lower_is_better: bool,
    pub improvement_threshold: f64,
    pub regression_threshold: f64,
}

pub struct Config {
    pub benchmarks: Vec<Benchmark>,
    pub slack_aliases: HashMap<String, String>,
}

pub fn parse_config(cfg: &Path,
                    def_imp_threshold: f64,
                    def_reg_threshold: f64)
                    -> Result<Config, Error> {
    let mut f = try!(fs::File::open(cfg));
    let mut buf = String::new();
    try!(f.read_to_string(&mut buf));

    let value = toml::Parser::new(&buf).parse().unwrap();

    let to_bench = |t: (&String, &toml::Value)| {
        Benchmark {
            name: t.0.clone(),
            cmd: String::from(t.1.lookup("command").unwrap().as_str().unwrap()),
            args: t.1.as_table().unwrap()["args"]
                .as_slice()
                .unwrap()
                .iter()
                .map(|a| String::from(a.as_str().unwrap()))
                .collect(),
            result_expr: t.1
                .lookup("regexs")
                .unwrap()
                .as_slice()
                .unwrap()
                .iter()
                .map(|r| Regex::new(r.as_str().unwrap()).unwrap())
                .collect(),
            lower_is_better: match t.1.lookup("lower_better") {
                None => false,
                Some(v) => v.as_bool().unwrap(),
            },
            improvement_threshold: match t.1.lookup("improvement_threshold") {
                None => def_imp_threshold,
                Some(ref it) => it.as_float().unwrap(),
            },
            regression_threshold: match t.1.lookup("regression_threshold") {
                None => def_reg_threshold,
                Some(ref rt) => rt.as_float().unwrap(),
            },
        }
    };

    // Github <-> Slack username mappings
    let slack_aliases = value.iter()
        .filter(|t| t.0 == "slack-aliases")
        .flat_map(|t| {
            t.1
                .as_table()
                .unwrap()
                .iter()
                .map(|(k, v)| (k.clone(), String::from(v.as_str().unwrap())))
        })
        .collect::<HashMap<_, _>>();

    // Benchmark definitions
    let benchmarks = value.iter().filter(|t| t.0 != "slack-aliases").map(|t| to_bench(t)).collect();

    Ok(Config {
        benchmarks: benchmarks,
        slack_aliases: slack_aliases,
    })
}
<commit_msg>Don't crash if taster config in repo invalid<commit_after>use regex::Regex;
use std::collections::HashMap;
use std::fs;
use std::io::{Error, ErrorKind, Read};
use std::path::Path;
use toml;

#[derive(Debug)]
pub struct Benchmark {
    pub name: String,
    pub cmd: String,
    pub args: Vec<String>,
    pub result_expr: Vec<Regex>,
    pub lower_is_better: bool,
    pub improvement_threshold: f64,
    pub regression_threshold: f64,
}

pub struct Config {
    pub benchmarks: Vec<Benchmark>,
    pub slack_aliases: HashMap<String, String>,
}

pub fn parse_config(cfg: &Path,
                    def_imp_threshold: f64,
                    def_reg_threshold: f64)
                    -> Result<Config, Error> {
    let mut f = try!(fs::File::open(cfg));
    let mut buf = String::new();
    try!(f.read_to_string(&mut buf));

    let value = match toml::Parser::new(&buf).parse() {
        None => return Err(Error::new(ErrorKind::InvalidInput, "failed to parse taster config!")),
        Some(v) => v,
    };

    let to_bench = |t: (&String, &toml::Value)| {
        Benchmark {
            name: t.0.clone(),
            cmd: String::from(t.1.lookup("command").unwrap().as_str().unwrap()),
            args: t.1.as_table().unwrap()["args"]
                .as_slice()
                .unwrap()
                .iter()
                .map(|a| String::from(a.as_str().unwrap()))
                .collect(),
            result_expr: t.1
                .lookup("regexs")
                .unwrap()
                .as_slice()
                .unwrap()
                .iter()
                .map(|r| Regex::new(r.as_str().unwrap()).unwrap())
                .collect(),
            lower_is_better: match t.1.lookup("lower_better") {
                None => false,
                Some(v) => v.as_bool().unwrap(),
            },
            improvement_threshold: match t.1.lookup("improvement_threshold") {
                None => def_imp_threshold,
                Some(ref it) => it.as_float().unwrap(),
            },
            regression_threshold: match t.1.lookup("regression_threshold") {
                None => def_reg_threshold,
                Some(ref rt) => rt.as_float().unwrap(),
            },
        }
    };

    // Github <-> Slack username mappings
    let slack_aliases = value.iter()
        .filter(|t| t.0 == "slack-aliases")
        .flat_map(|t| {
            t.1
                .as_table()
                .unwrap()
                .iter()
                .map(|(k, v)| (k.clone(), String::from(v.as_str().unwrap())))
        })
        .collect::<HashMap<_, _>>();

    // Benchmark definitions
    let benchmarks = value.iter().filter(|t| t.0 != "slack-aliases").map(|t| to_bench(t)).collect();

    Ok(Config {
        benchmarks: benchmarks,
        slack_aliases: slack_aliases,
    })
}
<|endoftext|>


<commit_before>// Format of a relocation entry of a Mach-O file.  Modified from the 4.3BSD
// format.  The modifications from the original format were changing the value
// of the r_symbolnum field for "local" (r_extern == 0) relocation entries.
// This modification is required to support symbols in an arbitrary number of
// sections not just the three sections (text, data and bss) in a 4.3BSD file.
// Also the last 4 bits have had the r_type tag added to them.

// The r_address is not really the address as it's name indicates but an offset.
// In 4.3BSD a.out objects this offset is from the start of the "segment" for
// which relocation entry is for (text or data).  For Mach-O object files it is
// also an offset but from the start of the "section" for which the relocation
// entry is for.  See comments in <mach-o/loader.h> about the r_address feild
// in images for used with the dynamic linker.

// In 4.3BSD a.out objects if r_extern is zero then r_symbolnum is an ordinal
// for the segment the symbol being relocated is in.  These ordinals are the
// symbol types N_TEXT, N_DATA, N_BSS or N_ABS.  In Mach-O object files these
// ordinals refer to the sections in the object file in the order their section
// structures appear in the headers of the object file they are in.  The first
// section has the ordinal 1, the second 2, and so on.  This means that the
// same ordinal in two different object files could refer to two different
// sections.  And further could have still different ordinals when combined
// by the link-editor.  The value R_ABS is used for relocation entries for
// absolute symbols which need no further relocation.
use core::fmt;
use crate::mach;
use scroll::{Pread, Pwrite, IOwrite, SizeWith, IOread};

// TODO: armv7 relocations are scattered, must and r_address with 0x8000_0000 to check if its scattered or not
#[derive(Copy, Clone, Pread, Pwrite, IOwrite, SizeWith, IOread)]
#[repr(C)]
pub struct RelocationInfo {
    /// Offset in the section to what is being relocated
    pub r_address: i32,
    /// Contains all of the relocation info as a bitfield.
    /// r_symbolnum, 24 bits, r_pcrel 1 bit, r_length 2 bits, r_extern 1 bit, r_type 4 bits
    pub r_info: u32,
}

pub const SIZEOF_RELOCATION_INFO: usize = 8;

impl RelocationInfo {
    /// Symbol index if `r_extern` == 1 or section ordinal if `r_extern` == 0. In bits :24
    #[inline]
    pub fn r_symbolnum(self) -> usize {
        (self.r_info & 0x00ff_ffffu32) as usize
    }
    /// Was relocated pc relative already, 1 bit
    #[inline]
    pub fn r_pcrel(self) -> u8 {
        ((self.r_info & 0x0100_0000u32) >> 24) as u8
    }
    /// The length of the relocation, 0=byte, 1=word, 2=long, 3=quad, 2 bits
    #[inline]
    pub fn r_length(self) -> u8 {
        ((self.r_info & 0x0600_0000u32) >> 25) as u8
    }
    /// Does not include value of sym referenced, 1 bit
    #[inline]
    pub fn r_extern(self) -> u8 {
        ((self.r_info & 0x0800_0000) >> 27) as u8
    }
    /// Ff not 0, machine specific relocation type, in bits :4
    #[inline]
    pub fn r_type(self) -> u8 {
        ((self.r_info & 0xf000_0000) >> 28) as u8
    }
    /// If true, this relocation is for a symbol; if false,  or a section ordinal otherwise
    #[inline]
    pub fn is_extern(self) -> bool {
        self.r_extern() == 1
    }
    /// If true, this is a PIC relocation
    #[inline]
    pub fn is_pic(self) -> bool {
        self.r_pcrel() > 0
    }
    /// Returns a string representation of this relocation, given the machine `cputype`
    pub fn to_str(self, cputype: mach::cputype::CpuType) -> &'static str {
        reloc_to_str(self.r_type(), cputype)
    }
}

/// Absolute relocation type for Mach-O files
pub const R_ABS: u8 = 0;

impl fmt::Debug for RelocationInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("RelocationInfo")
            .field("r_address", &format_args!("{:#x}", &self.r_address))
            .field("r_info", &format_args!("{:#x}", &self.r_info))
            .field("r_symbolnum", &format_args!("{:#x}", &self.r_symbolnum()))
            .field("r_pcrel", &(self.r_pcrel()))
            .field("r_length", &self.r_length())
            .field("r_extern", &self.r_extern())
            .field("r_type", &self.r_type())
            .finish()
    }
}

pub type RelocType = u8;

/// Absolute address
pub const X86_64_RELOC_UNSIGNED: RelocType = 0;
/// Signed 32-bit displacement
pub const X86_64_RELOC_SIGNED: RelocType = 1;
/// A CALL/JMP instruction with 32-bit displacement
pub const X86_64_RELOC_BRANCH: RelocType = 2;
/// A MOVQ load of a GOT entry
pub const X86_64_RELOC_GOT_LOAD: RelocType = 3;
/// Other GOT references
pub const X86_64_RELOC_GOT: RelocType = 4;
/// Must be followed by a X86_64_RELOC_UNSIGNED relocation
pub const X86_64_RELOC_SUBTRACTOR: RelocType = 5;
/// for signed 32-bit displacement with a -1 addend
pub const X86_64_RELOC_SIGNED_1: RelocType = 6;
/// for signed 32-bit displacement with a -2 addend
pub const X86_64_RELOC_SIGNED_2: RelocType = 7;
/// for signed 32-bit displacement with a -4 addend
pub const X86_64_RELOC_SIGNED_4: RelocType = 8;
/// for thread local variables
pub const X86_64_RELOC_TLV: RelocType = 9;

// x86 relocations
pub const GENERIC_RELOC_VANILLA: RelocType = 0;
pub const GENERIC_RELOC_PAIR: RelocType = 1;
pub const GENERIC_RELOC_SECTDIFF: RelocType = 2;
pub const GENERIC_RELOC_LOCAL_SECTDIFF: RelocType = 3;
pub const GENERIC_RELOC_PB_LA_P: RelocType = 4;

// arm relocations
pub const ARM_RELOC_VANILLA: RelocType = GENERIC_RELOC_VANILLA;
pub const ARM_RELOC_PAIR: RelocType = GENERIC_RELOC_PAIR;
pub const ARM_RELOC_SECTDIFF: RelocType = GENERIC_RELOC_SECTDIFF;
pub const ARM_RELOC_LOCAL_SECTDIFF: RelocType = 3;
pub const ARM_RELOC_PB_LA_PTR: RelocType = 4;
pub const ARM_RELOC_BR24: RelocType = 5;
pub const ARM_THUMB_RELOC_BR22: RelocType = 6;
/// Obsolete
pub const ARM_THUMB_32BIT_BRANCH: RelocType = 7;
pub const ARM_RELOC_HALF: RelocType = 8;
pub const ARM_RELOC_HALF_SECTDIFF: RelocType = 9;

/// For pointers.
pub const ARM64_RELOC_UNSIGNED: RelocType = 0;
/// Must be followed by an ARM64_RELOC_UNSIGNED
pub const ARM64_RELOC_SUBTRACTOR: RelocType = 1;
/// A B/BL instruction with 26-bit displacement.
pub const ARM64_RELOC_BRANCH26: RelocType = 2;
/// PC-rel distance to page of target.
pub const ARM64_RELOC_PAGE21: RelocType = 3;
/// Offset within page, scaled by r_length.
pub const ARM64_RELOC_PAGEOFF12: RelocType = 4;
/// PC-rel distance to page of GOT slot.
pub const ARM64_RELOC_GOT_LOAD_PAGE21: RelocType = 5;
/// Offset within page of GOT slot, scaled by r_length.
pub const ARM64_RELOC_GOT_LOAD_PAGEOFF12: RelocType = 6;
/// For pointers to GOT slots.
pub const ARM64_RELOC_POINTER_TO_GOT: RelocType = 7;
/// PC-rel distance to page of TLVP slot.
pub const ARM64_RELOC_TLVP_LOAD_PAGE21: RelocType = 8;
/// Offset within page of TLVP slot, scaled by r_length.
pub const ARM64_RELOC_TLVP_LOAD_PAGEOFF12: RelocType = 9;
/// Must be followed by ARM64_RELOC_PAGE21 or ARM64_RELOC_PAGEOFF12.
pub const ARM64_RELOC_ADDEND: RelocType = 10;

pub fn reloc_to_str(reloc: RelocType, cputype: mach::cputype::CpuType) -> &'static str {
    use crate::mach::constants::cputype::*;
    match cputype {
        CPU_TYPE_ARM64 | CPU_TYPE_ARM64_32 => {
            match reloc {
                ARM64_RELOC_UNSIGNED => "ARM64_RELOC_UNSIGNED",
                ARM64_RELOC_SUBTRACTOR => "ARM64_RELOC_SUBTRACTOR",
                ARM64_RELOC_BRANCH26 => "ARM64_RELOC_BRANCH26",
                ARM64_RELOC_PAGE21 => "ARM64_RELOC_PAGE21",
                ARM64_RELOC_PAGEOFF12 => "ARM64_RELOC_PAGEOFF12",
                ARM64_RELOC_GOT_LOAD_PAGE21 => "ARM64_RELOC_GOT_LOAD_PAGE21",
                ARM64_RELOC_GOT_LOAD_PAGEOFF12 => "ARM64_RELOC_GOT_LOAD_PAGEOFF12",
                ARM64_RELOC_POINTER_TO_GOT => "ARM64_RELOC_POINTER_TO_GOT",
                ARM64_RELOC_TLVP_LOAD_PAGE21 => "ARM64_RELOC_TLVP_LOAD_PAGE21",
                ARM64_RELOC_TLVP_LOAD_PAGEOFF12 => "ARM64_RELOC_TLVP_LOAD_PAGEOFF12",
                ARM64_RELOC_ADDEND => "ARM64_RELOC_ADDEND",
                _ => "UNKNOWN",
            }
        },
        CPU_TYPE_X86_64 => {
            match reloc {
                X86_64_RELOC_UNSIGNED => "X86_64_RELOC_UNSIGNED",
                X86_64_RELOC_SIGNED => "X86_64_RELOC_SIGNED",
                X86_64_RELOC_BRANCH => "X86_64_RELOC_BRANCH",
                X86_64_RELOC_GOT_LOAD => "X86_64_RELOC_GOT_LOAD",
                X86_64_RELOC_GOT => "X86_64_RELOC_GOT",
                X86_64_RELOC_SUBTRACTOR => "X86_64_RELOC_SUBTRACTOR",
                X86_64_RELOC_SIGNED_1 => "X86_64_RELOC_SIGNED_1",
                X86_64_RELOC_SIGNED_2 => "X86_64_RELOC_SIGNED_2",
                X86_64_RELOC_SIGNED_4 => "X86_64_RELOC_SIGNED_4",
                X86_64_RELOC_TLV => "X86_64_RELOC_TLV",
                _ => "UNKNOWN",
            }
        },
        CPU_TYPE_ARM => {
            match reloc {
                ARM_RELOC_VANILLA => "ARM_RELOC_VANILLA",
                ARM_RELOC_PAIR => "ARM_RELOC_PAIR",
                ARM_RELOC_SECTDIFF => "ARM_RELOC_SECTDIFF",
                ARM_RELOC_LOCAL_SECTDIFF => "ARM_RELOC_LOCAL_SECTDIFF",
                ARM_RELOC_PB_LA_PTR => "ARM_RELOC_PB_LA_PTR",
                ARM_RELOC_BR24 => "ARM_RELOC_BR24",
                ARM_THUMB_RELOC_BR22 => "ARM_THUMB_RELOC_BR22",
                ARM_THUMB_32BIT_BRANCH => "ARM_THUMB_32BIT_BRANCH",
                ARM_RELOC_HALF => "ARM_RELOC_HALF",
                ARM_RELOC_HALF_SECTDIFF => "ARM_RELOC_HALF_SECTDIFF",
                _ => "UNKNOWN",
            }
        },
        CPU_TYPE_X86 => {
            match reloc {
                GENERIC_RELOC_VANILLA => "GENERIC_RELOC_VANILLA",
                GENERIC_RELOC_PAIR => "GENERIC_RELOC_PAIR",
                GENERIC_RELOC_SECTDIFF => "GENERIC_RELOC_SECTDIFF",
                GENERIC_RELOC_LOCAL_SECTDIFF => "GENERIC_RELOC_LOCAL_SECTDIFF",
                GENERIC_RELOC_PB_LA_P => "GENERIC_RELOC_PB_LA_P",
                _ => "UNKNOWN",
            }
        },
        _ => "BAD_CPUTYPE"
    }
}
<commit_msg>mach: fix generic reloc constants<commit_after>// Format of a relocation entry of a Mach-O file.  Modified from the 4.3BSD
// format.  The modifications from the original format were changing the value
// of the r_symbolnum field for "local" (r_extern == 0) relocation entries.
// This modification is required to support symbols in an arbitrary number of
// sections not just the three sections (text, data and bss) in a 4.3BSD file.
// Also the last 4 bits have had the r_type tag added to them.

// The r_address is not really the address as it's name indicates but an offset.
// In 4.3BSD a.out objects this offset is from the start of the "segment" for
// which relocation entry is for (text or data).  For Mach-O object files it is
// also an offset but from the start of the "section" for which the relocation
// entry is for.  See comments in <mach-o/loader.h> about the r_address feild
// in images for used with the dynamic linker.

// In 4.3BSD a.out objects if r_extern is zero then r_symbolnum is an ordinal
// for the segment the symbol being relocated is in.  These ordinals are the
// symbol types N_TEXT, N_DATA, N_BSS or N_ABS.  In Mach-O object files these
// ordinals refer to the sections in the object file in the order their section
// structures appear in the headers of the object file they are in.  The first
// section has the ordinal 1, the second 2, and so on.  This means that the
// same ordinal in two different object files could refer to two different
// sections.  And further could have still different ordinals when combined
// by the link-editor.  The value R_ABS is used for relocation entries for
// absolute symbols which need no further relocation.
use core::fmt;
use crate::mach;
use scroll::{Pread, Pwrite, IOwrite, SizeWith, IOread};

// TODO: armv7 relocations are scattered, must and r_address with 0x8000_0000 to check if its scattered or not
#[derive(Copy, Clone, Pread, Pwrite, IOwrite, SizeWith, IOread)]
#[repr(C)]
pub struct RelocationInfo {
    /// Offset in the section to what is being relocated
    pub r_address: i32,
    /// Contains all of the relocation info as a bitfield.
    /// r_symbolnum, 24 bits, r_pcrel 1 bit, r_length 2 bits, r_extern 1 bit, r_type 4 bits
    pub r_info: u32,
}

pub const SIZEOF_RELOCATION_INFO: usize = 8;

impl RelocationInfo {
    /// Symbol index if `r_extern` == 1 or section ordinal if `r_extern` == 0. In bits :24
    #[inline]
    pub fn r_symbolnum(self) -> usize {
        (self.r_info & 0x00ff_ffffu32) as usize
    }
    /// Was relocated pc relative already, 1 bit
    #[inline]
    pub fn r_pcrel(self) -> u8 {
        ((self.r_info & 0x0100_0000u32) >> 24) as u8
    }
    /// The length of the relocation, 0=byte, 1=word, 2=long, 3=quad, 2 bits
    #[inline]
    pub fn r_length(self) -> u8 {
        ((self.r_info & 0x0600_0000u32) >> 25) as u8
    }
    /// Does not include value of sym referenced, 1 bit
    #[inline]
    pub fn r_extern(self) -> u8 {
        ((self.r_info & 0x0800_0000) >> 27) as u8
    }
    /// Ff not 0, machine specific relocation type, in bits :4
    #[inline]
    pub fn r_type(self) -> u8 {
        ((self.r_info & 0xf000_0000) >> 28) as u8
    }
    /// If true, this relocation is for a symbol; if false,  or a section ordinal otherwise
    #[inline]
    pub fn is_extern(self) -> bool {
        self.r_extern() == 1
    }
    /// If true, this is a PIC relocation
    #[inline]
    pub fn is_pic(self) -> bool {
        self.r_pcrel() > 0
    }
    /// Returns a string representation of this relocation, given the machine `cputype`
    pub fn to_str(self, cputype: mach::cputype::CpuType) -> &'static str {
        reloc_to_str(self.r_type(), cputype)
    }
}

/// Absolute relocation type for Mach-O files
pub const R_ABS: u8 = 0;

impl fmt::Debug for RelocationInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("RelocationInfo")
            .field("r_address", &format_args!("{:#x}", &self.r_address))
            .field("r_info", &format_args!("{:#x}", &self.r_info))
            .field("r_symbolnum", &format_args!("{:#x}", &self.r_symbolnum()))
            .field("r_pcrel", &(self.r_pcrel()))
            .field("r_length", &self.r_length())
            .field("r_extern", &self.r_extern())
            .field("r_type", &self.r_type())
            .finish()
    }
}

pub type RelocType = u8;

/// Absolute address
pub const X86_64_RELOC_UNSIGNED: RelocType = 0;
/// Signed 32-bit displacement
pub const X86_64_RELOC_SIGNED: RelocType = 1;
/// A CALL/JMP instruction with 32-bit displacement
pub const X86_64_RELOC_BRANCH: RelocType = 2;
/// A MOVQ load of a GOT entry
pub const X86_64_RELOC_GOT_LOAD: RelocType = 3;
/// Other GOT references
pub const X86_64_RELOC_GOT: RelocType = 4;
/// Must be followed by a X86_64_RELOC_UNSIGNED relocation
pub const X86_64_RELOC_SUBTRACTOR: RelocType = 5;
/// for signed 32-bit displacement with a -1 addend
pub const X86_64_RELOC_SIGNED_1: RelocType = 6;
/// for signed 32-bit displacement with a -2 addend
pub const X86_64_RELOC_SIGNED_2: RelocType = 7;
/// for signed 32-bit displacement with a -4 addend
pub const X86_64_RELOC_SIGNED_4: RelocType = 8;
/// for thread local variables
pub const X86_64_RELOC_TLV: RelocType = 9;

// x86 relocations
pub const GENERIC_RELOC_VANILLA: RelocType = 0;
pub const GENERIC_RELOC_PAIR: RelocType = 1;
pub const GENERIC_RELOC_SECTDIFF: RelocType = 2;
pub const GENERIC_RELOC_PB_LA_PTR: RelocType = 3;
pub const GENERIC_RELOC_LOCAL_SECTDIFF: RelocType = 4;
pub const GENERIC_RELOC_TLV: RelocType = 5;

// arm relocations
pub const ARM_RELOC_VANILLA: RelocType = GENERIC_RELOC_VANILLA;
pub const ARM_RELOC_PAIR: RelocType = GENERIC_RELOC_PAIR;
pub const ARM_RELOC_SECTDIFF: RelocType = GENERIC_RELOC_SECTDIFF;
pub const ARM_RELOC_LOCAL_SECTDIFF: RelocType = 3;
pub const ARM_RELOC_PB_LA_PTR: RelocType = 4;
pub const ARM_RELOC_BR24: RelocType = 5;
pub const ARM_THUMB_RELOC_BR22: RelocType = 6;
/// Obsolete
pub const ARM_THUMB_32BIT_BRANCH: RelocType = 7;
pub const ARM_RELOC_HALF: RelocType = 8;
pub const ARM_RELOC_HALF_SECTDIFF: RelocType = 9;

/// For pointers.
pub const ARM64_RELOC_UNSIGNED: RelocType = 0;
/// Must be followed by an ARM64_RELOC_UNSIGNED
pub const ARM64_RELOC_SUBTRACTOR: RelocType = 1;
/// A B/BL instruction with 26-bit displacement.
pub const ARM64_RELOC_BRANCH26: RelocType = 2;
/// PC-rel distance to page of target.
pub const ARM64_RELOC_PAGE21: RelocType = 3;
/// Offset within page, scaled by r_length.
pub const ARM64_RELOC_PAGEOFF12: RelocType = 4;
/// PC-rel distance to page of GOT slot.
pub const ARM64_RELOC_GOT_LOAD_PAGE21: RelocType = 5;
/// Offset within page of GOT slot, scaled by r_length.
pub const ARM64_RELOC_GOT_LOAD_PAGEOFF12: RelocType = 6;
/// For pointers to GOT slots.
pub const ARM64_RELOC_POINTER_TO_GOT: RelocType = 7;
/// PC-rel distance to page of TLVP slot.
pub const ARM64_RELOC_TLVP_LOAD_PAGE21: RelocType = 8;
/// Offset within page of TLVP slot, scaled by r_length.
pub const ARM64_RELOC_TLVP_LOAD_PAGEOFF12: RelocType = 9;
/// Must be followed by ARM64_RELOC_PAGE21 or ARM64_RELOC_PAGEOFF12.
pub const ARM64_RELOC_ADDEND: RelocType = 10;

pub fn reloc_to_str(reloc: RelocType, cputype: mach::cputype::CpuType) -> &'static str {
    use crate::mach::constants::cputype::*;
    match cputype {
        CPU_TYPE_ARM64 | CPU_TYPE_ARM64_32 => {
            match reloc {
                ARM64_RELOC_UNSIGNED => "ARM64_RELOC_UNSIGNED",
                ARM64_RELOC_SUBTRACTOR => "ARM64_RELOC_SUBTRACTOR",
                ARM64_RELOC_BRANCH26 => "ARM64_RELOC_BRANCH26",
                ARM64_RELOC_PAGE21 => "ARM64_RELOC_PAGE21",
                ARM64_RELOC_PAGEOFF12 => "ARM64_RELOC_PAGEOFF12",
                ARM64_RELOC_GOT_LOAD_PAGE21 => "ARM64_RELOC_GOT_LOAD_PAGE21",
                ARM64_RELOC_GOT_LOAD_PAGEOFF12 => "ARM64_RELOC_GOT_LOAD_PAGEOFF12",
                ARM64_RELOC_POINTER_TO_GOT => "ARM64_RELOC_POINTER_TO_GOT",
                ARM64_RELOC_TLVP_LOAD_PAGE21 => "ARM64_RELOC_TLVP_LOAD_PAGE21",
                ARM64_RELOC_TLVP_LOAD_PAGEOFF12 => "ARM64_RELOC_TLVP_LOAD_PAGEOFF12",
                ARM64_RELOC_ADDEND => "ARM64_RELOC_ADDEND",
                _ => "UNKNOWN",
            }
        },
        CPU_TYPE_X86_64 => {
            match reloc {
                X86_64_RELOC_UNSIGNED => "X86_64_RELOC_UNSIGNED",
                X86_64_RELOC_SIGNED => "X86_64_RELOC_SIGNED",
                X86_64_RELOC_BRANCH => "X86_64_RELOC_BRANCH",
                X86_64_RELOC_GOT_LOAD => "X86_64_RELOC_GOT_LOAD",
                X86_64_RELOC_GOT => "X86_64_RELOC_GOT",
                X86_64_RELOC_SUBTRACTOR => "X86_64_RELOC_SUBTRACTOR",
                X86_64_RELOC_SIGNED_1 => "X86_64_RELOC_SIGNED_1",
                X86_64_RELOC_SIGNED_2 => "X86_64_RELOC_SIGNED_2",
                X86_64_RELOC_SIGNED_4 => "X86_64_RELOC_SIGNED_4",
                X86_64_RELOC_TLV => "X86_64_RELOC_TLV",
                _ => "UNKNOWN",
            }
        },
        CPU_TYPE_ARM => {
            match reloc {
                ARM_RELOC_VANILLA => "ARM_RELOC_VANILLA",
                ARM_RELOC_PAIR => "ARM_RELOC_PAIR",
                ARM_RELOC_SECTDIFF => "ARM_RELOC_SECTDIFF",
                ARM_RELOC_LOCAL_SECTDIFF => "ARM_RELOC_LOCAL_SECTDIFF",
                ARM_RELOC_PB_LA_PTR => "ARM_RELOC_PB_LA_PTR",
                ARM_RELOC_BR24 => "ARM_RELOC_BR24",
                ARM_THUMB_RELOC_BR22 => "ARM_THUMB_RELOC_BR22",
                ARM_THUMB_32BIT_BRANCH => "ARM_THUMB_32BIT_BRANCH",
                ARM_RELOC_HALF => "ARM_RELOC_HALF",
                ARM_RELOC_HALF_SECTDIFF => "ARM_RELOC_HALF_SECTDIFF",
                _ => "UNKNOWN",
            }
        },
        CPU_TYPE_X86 => {
            match reloc {
                GENERIC_RELOC_VANILLA => "GENERIC_RELOC_VANILLA",
                GENERIC_RELOC_PAIR => "GENERIC_RELOC_PAIR",
                GENERIC_RELOC_SECTDIFF => "GENERIC_RELOC_SECTDIFF",
                GENERIC_RELOC_PB_LA_PTR => "GENERIC_RELOC_PB_LA_PTR",
                GENERIC_RELOC_LOCAL_SECTDIFF => "GENERIC_RELOC_LOCAL_SECTDIFF",
                GENERIC_RELOC_TLV => "GENERIC_RELOC_TLV",
                _ => "UNKNOWN",
            }
        },
        _ => "BAD_CPUTYPE"
    }
}
<|endoftext|>


<commit_before>use crate::io::{self, BufWriter, IoSlice, Write};
use crate::memchr;

/// Private helper struct for implementing the line-buffered writing logic.
/// This shim temporarily wraps a BufWriter, and uses its internals to
/// implement a line-buffered writer (specifically by using the internal
/// methods like write_to_buf and flush_buf). In this way, a more
/// efficient abstraction can be created than one that only had access to
/// `write` and `flush`, without needlessly duplicating a lot of the
/// implementation details of BufWriter. This also allows existing
/// `BufWriters` to be temporarily given line-buffering logic; this is what
/// enables Stdout to be alternately in line-buffered or block-buffered mode.
#[derive(Debug)]
pub struct LineWriterShim<'a, W: Write> {
    buffer: &'a mut BufWriter<W>,
}

impl<'a, W: Write> LineWriterShim<'a, W> {
    pub fn new(buffer: &'a mut BufWriter<W>) -> Self {
        Self { buffer }
    }

    /// Get a mutable reference to the inner writer (that is, the writer
    /// wrapped by the BufWriter). Be careful with this writer, as writes to
    /// it will bypass the buffer.
    fn inner_mut(&mut self) -> &mut W {
        self.buffer.get_mut()
    }

    /// Get the content currently buffered in self.buffer
    fn buffered(&self) -> &[u8] {
        self.buffer.buffer()
    }

    /// Flush the buffer iff the last byte is a newline (indicating that an
    /// earlier write only succeeded partially, and we want to retry flushing
    /// the buffered line before continuing with a subsequent write)
    fn flush_if_completed_line(&mut self) -> io::Result<()> {
        match self.buffered().last().copied() {
            Some(b'\n') => self.buffer.flush_buf(),
            _ => Ok(()),
        }
    }
}

impl<'a, W: Write> Write for LineWriterShim<'a, W> {
    /// Write some data into this BufReader with line buffering. This means
    /// that, if any newlines are present in the data, the data up to the last
    /// newline is sent directly to the underlying writer, and data after it
    /// is buffered. Returns the number of bytes written.
    ///
    /// This function operates on a "best effort basis"; in keeping with the
    /// convention of `Write::write`, it makes at most one attempt to write
    /// new data to the underlying writer. If that write only reports a partial
    /// success, the remaining data will be buffered.
    ///
    /// Because this function attempts to send completed lines to the underlying
    /// writer, it will also flush the existing buffer if it ends with a
    /// newline, even if the incoming data does not contain any newlines.
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let newline_idx = match memchr::memrchr(b'\n', buf) {
            // If there are no new newlines (that is, if this write is less than
            // one line), just do a regular buffered write (which may flush if
            // we exceed the inner buffer's size)
            None => {
                self.flush_if_completed_line()?;
                return self.buffer.write(buf);
            }
            // Otherwise, arrange for the lines to be written directly to the
            // inner writer.
            Some(newline_idx) => newline_idx + 1,
        };

        // Flush existing content to prepare for our write. We have to do this
        // before attempting to write `buf` in order to maintain consistency;
        // if we add `buf` to the buffer then try to flush it all at once,
        // we're obligated to return Ok(), which would mean suppressing any
        // errors that occur during flush.
        self.buffer.flush_buf()?;

        // This is what we're going to try to write directly to the inner
        // writer. The rest will be buffered, if nothing goes wrong.
        let lines = &buf[..newline_idx];

        // Write `lines` directly to the inner writer. In keeping with the
        // `write` convention, make at most one attempt to add new (unbuffered)
        // data. Because this write doesn't touch the BufWriter state directly,
        // and the buffer is known to be empty, we don't need to worry about
        // self.buffer.panicked here.
        let flushed = self.inner_mut().write(lines)?;

        // If buffer returns Ok(0), propagate that to the caller without
        // doing additional buffering; otherwise we're just guaranteeing
        // an "ErrorKind::WriteZero" later.
        if flushed == 0 {
            return Ok(0);
        }

        // Now that the write has succeeded, buffer the rest (or as much of
        // the rest as possible). If there were any unwritten newlines, we
        // only buffer out to the last unwritten newline that fits in the
        // buffer; this helps prevent flushing partial lines on subsequent
        // calls to LineWriterShim::write.

        // Handle the cases in order of most-common to least-common, under
        // the presumption that most writes succeed in totality, and that most
        // writes are smaller than the buffer.
        // - Is this a partial line (ie, no newlines left in the unwritten tail)
        // - If not, does the data out to the last unwritten newline fit in
        //   the buffer?
        // - If not, scan for the last newline that *does* fit in the buffer
        let tail = if flushed >= newline_idx {
            &buf[flushed..]
        } else if newline_idx - flushed <= self.buffer.capacity() {
            &buf[flushed..newline_idx]
        } else {
            let scan_area = &buf[flushed..];
            let scan_area = &scan_area[..self.buffer.capacity()];
            match memchr::memrchr(b'\n', scan_area) {
                Some(newline_idx) => &scan_area[..newline_idx + 1],
                None => scan_area,
            }
        };

        let buffered = self.buffer.write_to_buf(tail);
        Ok(flushed + buffered)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.buffer.flush()
    }

    /// Write some vectored data into this BufReader with line buffering. This
    /// means that, if any newlines are present in the data, the data up to
    /// and including the buffer containing the last newline is sent directly
    /// to the inner writer, and the data after it is buffered. Returns the
    /// number of bytes written.
    ///
    /// This function operates on a "best effort basis"; in keeping with the
    /// convention of `Write::write`, it makes at most one attempt to write
    /// new data to the underlying writer.
    ///
    /// Because this function attempts to send completed lines to the underlying
    /// writer, it will also flush the existing buffer if it contains any
    /// newlines.
    ///
    /// Because sorting through an array of `IoSlice` can be a bit convoluted,
    /// This method differs from write in the following ways:
    ///
    /// - It attempts to write the full content of all the buffers up to and
    ///   including the one containing the last newline. This means that it
    ///   may attempt to write a partial line, that buffer has data past the
    ///   newline.
    /// - If the write only reports partial success, it does not attempt to
    ///   find the precise location of the written bytes and buffer the rest.
    ///
    /// If the underlying vector doesn't support vectored writing, we instead
    /// simply write the first non-empty buffer with `write`. This way, we
    /// get the benefits of more granular partial-line handling without losing
    /// anything in efficiency
    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
        // If there's no specialized behavior for write_vectored, just use
        // write. This has the benefit of more granular partial-line handling.
        if !self.is_write_vectored() {
            return match bufs.iter().find(|buf| !buf.is_empty()) {
                Some(buf) => self.write(buf),
                None => Ok(0),
            };
        }

        // Find the buffer containing the last newline
        let last_newline_buf_idx = bufs
            .iter()
            .enumerate()
            .rev()
            .find_map(|(i, buf)| memchr::memchr(b'\n', buf).map(|_| i));

        // If there are no new newlines (that is, if this write is less than
        // one line), just do a regular buffered write
        let last_newline_buf_idx = match last_newline_buf_idx {
            // No newlines; just do a normal buffered write
            None => {
                self.flush_if_completed_line()?;
                return self.buffer.write_vectored(bufs);
            }
            Some(i) => i,
        };

        // Flush existing content to prepare for our write
        self.buffer.flush_buf()?;

        // This is what we're going to try to write directly to the inner
        // writer. The rest will be buffered, if nothing goes wrong.
        let (lines, tail) = bufs.split_at(last_newline_buf_idx + 1);

        // Write `lines` directly to the inner writer. In keeping with the
        // `write` convention, make at most one attempt to add new (unbuffered)
        // data. Because this write doesn't touch the BufWriter state directly,
        // and the buffer is known to be empty, we don't need to worry about
        // self.panicked here.
        let flushed = self.inner_mut().write_vectored(lines)?;

        // If inner returns Ok(0), propagate that to the caller without
        // doing additional buffering; otherwise we're just guaranteeing
        // an "ErrorKind::WriteZero" later.
        if flushed == 0 {
            return Ok(0);
        }

        // Don't try to reconstruct the exact amount written; just bail
        // in the event of a partial write
        let lines_len = lines.iter().map(|buf| buf.len()).sum();
        if flushed < lines_len {
            return Ok(flushed);
        }

        // Now that the write has succeeded, buffer the rest (or as much of the
        // rest as possible)
        let buffered: usize = tail
            .iter()
            .filter(|buf| !buf.is_empty())
            .map(|buf| self.buffer.write_to_buf(buf))
            .take_while(|&n| n > 0)
            .sum();

        Ok(flushed + buffered)
    }

    fn is_write_vectored(&self) -> bool {
        self.buffer.is_write_vectored()
    }

    /// Write some data into this BufReader with line buffering. This means
    /// that, if any newlines are present in the data, the data up to the last
    /// newline is sent directly to the underlying writer, and data after it
    /// is buffered.
    ///
    /// Because this function attempts to send completed lines to the underlying
    /// writer, it will also flush the existing buffer if it contains any
    /// newlines, even if the incoming data does not contain any newlines.
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        match memchr::memrchr(b'\n', buf) {
            // If there are no new newlines (that is, if this write is less than
            // one line), just do a regular buffered write (which may flush if
            // we exceed the inner buffer's size)
            None => {
                self.flush_if_completed_line()?;
                self.buffer.write_all(buf)
            }
            Some(newline_idx) => {
                let (lines, tail) = buf.split_at(newline_idx + 1);

                if self.buffered().is_empty() {
                    self.inner_mut().write_all(lines)?;
                } else {
                    // If there is any buffered data, we add the incoming lines
                    // to that buffer before flushing, which saves us at least
                    // one write call. We can't really do this with `write`,
                    // since we can't do this *and* not suppress errors *and*
                    // report a consistent state to the caller in a return
                    // value, but here in write_all it's fine.
                    self.buffer.write_all(lines)?;
                    self.buffer.flush_buf()?;
                }

                self.buffer.write_all(tail)
            }
        }
    }
}
<commit_msg>Fix is_write_vectored in LineWriterShim<commit_after>use crate::io::{self, BufWriter, IoSlice, Write};
use crate::memchr;

/// Private helper struct for implementing the line-buffered writing logic.
/// This shim temporarily wraps a BufWriter, and uses its internals to
/// implement a line-buffered writer (specifically by using the internal
/// methods like write_to_buf and flush_buf). In this way, a more
/// efficient abstraction can be created than one that only had access to
/// `write` and `flush`, without needlessly duplicating a lot of the
/// implementation details of BufWriter. This also allows existing
/// `BufWriters` to be temporarily given line-buffering logic; this is what
/// enables Stdout to be alternately in line-buffered or block-buffered mode.
#[derive(Debug)]
pub struct LineWriterShim<'a, W: Write> {
    buffer: &'a mut BufWriter<W>,
}

impl<'a, W: Write> LineWriterShim<'a, W> {
    pub fn new(buffer: &'a mut BufWriter<W>) -> Self {
        Self { buffer }
    }

    /// Get a reference to the inner writer (that is, the writer
    /// wrapped by the BufWriter).
    fn inner(&self) -> &W {
        self.buffer.get_ref()
    }

    /// Get a mutable reference to the inner writer (that is, the writer
    /// wrapped by the BufWriter). Be careful with this writer, as writes to
    /// it will bypass the buffer.
    fn inner_mut(&mut self) -> &mut W {
        self.buffer.get_mut()
    }

    /// Get the content currently buffered in self.buffer
    fn buffered(&self) -> &[u8] {
        self.buffer.buffer()
    }

    /// Flush the buffer iff the last byte is a newline (indicating that an
    /// earlier write only succeeded partially, and we want to retry flushing
    /// the buffered line before continuing with a subsequent write)
    fn flush_if_completed_line(&mut self) -> io::Result<()> {
        match self.buffered().last().copied() {
            Some(b'\n') => self.buffer.flush_buf(),
            _ => Ok(()),
        }
    }
}

impl<'a, W: Write> Write for LineWriterShim<'a, W> {
    /// Write some data into this BufReader with line buffering. This means
    /// that, if any newlines are present in the data, the data up to the last
    /// newline is sent directly to the underlying writer, and data after it
    /// is buffered. Returns the number of bytes written.
    ///
    /// This function operates on a "best effort basis"; in keeping with the
    /// convention of `Write::write`, it makes at most one attempt to write
    /// new data to the underlying writer. If that write only reports a partial
    /// success, the remaining data will be buffered.
    ///
    /// Because this function attempts to send completed lines to the underlying
    /// writer, it will also flush the existing buffer if it ends with a
    /// newline, even if the incoming data does not contain any newlines.
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let newline_idx = match memchr::memrchr(b'\n', buf) {
            // If there are no new newlines (that is, if this write is less than
            // one line), just do a regular buffered write (which may flush if
            // we exceed the inner buffer's size)
            None => {
                self.flush_if_completed_line()?;
                return self.buffer.write(buf);
            }
            // Otherwise, arrange for the lines to be written directly to the
            // inner writer.
            Some(newline_idx) => newline_idx + 1,
        };

        // Flush existing content to prepare for our write. We have to do this
        // before attempting to write `buf` in order to maintain consistency;
        // if we add `buf` to the buffer then try to flush it all at once,
        // we're obligated to return Ok(), which would mean suppressing any
        // errors that occur during flush.
        self.buffer.flush_buf()?;

        // This is what we're going to try to write directly to the inner
        // writer. The rest will be buffered, if nothing goes wrong.
        let lines = &buf[..newline_idx];

        // Write `lines` directly to the inner writer. In keeping with the
        // `write` convention, make at most one attempt to add new (unbuffered)
        // data. Because this write doesn't touch the BufWriter state directly,
        // and the buffer is known to be empty, we don't need to worry about
        // self.buffer.panicked here.
        let flushed = self.inner_mut().write(lines)?;

        // If buffer returns Ok(0), propagate that to the caller without
        // doing additional buffering; otherwise we're just guaranteeing
        // an "ErrorKind::WriteZero" later.
        if flushed == 0 {
            return Ok(0);
        }

        // Now that the write has succeeded, buffer the rest (or as much of
        // the rest as possible). If there were any unwritten newlines, we
        // only buffer out to the last unwritten newline that fits in the
        // buffer; this helps prevent flushing partial lines on subsequent
        // calls to LineWriterShim::write.

        // Handle the cases in order of most-common to least-common, under
        // the presumption that most writes succeed in totality, and that most
        // writes are smaller than the buffer.
        // - Is this a partial line (ie, no newlines left in the unwritten tail)
        // - If not, does the data out to the last unwritten newline fit in
        //   the buffer?
        // - If not, scan for the last newline that *does* fit in the buffer
        let tail = if flushed >= newline_idx {
            &buf[flushed..]
        } else if newline_idx - flushed <= self.buffer.capacity() {
            &buf[flushed..newline_idx]
        } else {
            let scan_area = &buf[flushed..];
            let scan_area = &scan_area[..self.buffer.capacity()];
            match memchr::memrchr(b'\n', scan_area) {
                Some(newline_idx) => &scan_area[..newline_idx + 1],
                None => scan_area,
            }
        };

        let buffered = self.buffer.write_to_buf(tail);
        Ok(flushed + buffered)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.buffer.flush()
    }

    /// Write some vectored data into this BufReader with line buffering. This
    /// means that, if any newlines are present in the data, the data up to
    /// and including the buffer containing the last newline is sent directly
    /// to the inner writer, and the data after it is buffered. Returns the
    /// number of bytes written.
    ///
    /// This function operates on a "best effort basis"; in keeping with the
    /// convention of `Write::write`, it makes at most one attempt to write
    /// new data to the underlying writer.
    ///
    /// Because this function attempts to send completed lines to the underlying
    /// writer, it will also flush the existing buffer if it contains any
    /// newlines.
    ///
    /// Because sorting through an array of `IoSlice` can be a bit convoluted,
    /// This method differs from write in the following ways:
    ///
    /// - It attempts to write the full content of all the buffers up to and
    ///   including the one containing the last newline. This means that it
    ///   may attempt to write a partial line, that buffer has data past the
    ///   newline.
    /// - If the write only reports partial success, it does not attempt to
    ///   find the precise location of the written bytes and buffer the rest.
    ///
    /// If the underlying vector doesn't support vectored writing, we instead
    /// simply write the first non-empty buffer with `write`. This way, we
    /// get the benefits of more granular partial-line handling without losing
    /// anything in efficiency
    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
        // If there's no specialized behavior for write_vectored, just use
        // write. This has the benefit of more granular partial-line handling.
        if !self.is_write_vectored() {
            return match bufs.iter().find(|buf| !buf.is_empty()) {
                Some(buf) => self.write(buf),
                None => Ok(0),
            };
        }

        // Find the buffer containing the last newline
        let last_newline_buf_idx = bufs
            .iter()
            .enumerate()
            .rev()
            .find_map(|(i, buf)| memchr::memchr(b'\n', buf).map(|_| i));

        // If there are no new newlines (that is, if this write is less than
        // one line), just do a regular buffered write
        let last_newline_buf_idx = match last_newline_buf_idx {
            // No newlines; just do a normal buffered write
            None => {
                self.flush_if_completed_line()?;
                return self.buffer.write_vectored(bufs);
            }
            Some(i) => i,
        };

        // Flush existing content to prepare for our write
        self.buffer.flush_buf()?;

        // This is what we're going to try to write directly to the inner
        // writer. The rest will be buffered, if nothing goes wrong.
        let (lines, tail) = bufs.split_at(last_newline_buf_idx + 1);

        // Write `lines` directly to the inner writer. In keeping with the
        // `write` convention, make at most one attempt to add new (unbuffered)
        // data. Because this write doesn't touch the BufWriter state directly,
        // and the buffer is known to be empty, we don't need to worry about
        // self.panicked here.
        let flushed = self.inner_mut().write_vectored(lines)?;

        // If inner returns Ok(0), propagate that to the caller without
        // doing additional buffering; otherwise we're just guaranteeing
        // an "ErrorKind::WriteZero" later.
        if flushed == 0 {
            return Ok(0);
        }

        // Don't try to reconstruct the exact amount written; just bail
        // in the event of a partial write
        let lines_len = lines.iter().map(|buf| buf.len()).sum();
        if flushed < lines_len {
            return Ok(flushed);
        }

        // Now that the write has succeeded, buffer the rest (or as much of the
        // rest as possible)
        let buffered: usize = tail
            .iter()
            .filter(|buf| !buf.is_empty())
            .map(|buf| self.buffer.write_to_buf(buf))
            .take_while(|&n| n > 0)
            .sum();

        Ok(flushed + buffered)
    }

    fn is_write_vectored(&self) -> bool {
        self.inner().is_write_vectored()
    }

    /// Write some data into this BufReader with line buffering. This means
    /// that, if any newlines are present in the data, the data up to the last
    /// newline is sent directly to the underlying writer, and data after it
    /// is buffered.
    ///
    /// Because this function attempts to send completed lines to the underlying
    /// writer, it will also flush the existing buffer if it contains any
    /// newlines, even if the incoming data does not contain any newlines.
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        match memchr::memrchr(b'\n', buf) {
            // If there are no new newlines (that is, if this write is less than
            // one line), just do a regular buffered write (which may flush if
            // we exceed the inner buffer's size)
            None => {
                self.flush_if_completed_line()?;
                self.buffer.write_all(buf)
            }
            Some(newline_idx) => {
                let (lines, tail) = buf.split_at(newline_idx + 1);

                if self.buffered().is_empty() {
                    self.inner_mut().write_all(lines)?;
                } else {
                    // If there is any buffered data, we add the incoming lines
                    // to that buffer before flushing, which saves us at least
                    // one write call. We can't really do this with `write`,
                    // since we can't do this *and* not suppress errors *and*
                    // report a consistent state to the caller in a return
                    // value, but here in write_all it's fine.
                    self.buffer.write_all(lines)?;
                    self.buffer.flush_buf()?;
                }

                self.buffer.write_all(tail)
            }
        }
    }
}
<|endoftext|>