<commit_before>use std::collections::HashMap;
use std::env;
use std::fs::{self, File};
use std::iter::repeat;
use std::path::PathBuf;
use std::time::Duration;

use curl::easy::Easy;
use git2;
use registry::{Registry, NewCrate, NewCrateDependency};
use term::color::BLACK;

use url::percent_encoding::{percent_encode, QUERY_ENCODE_SET};

use core::source::Source;
use core::{Package, SourceId, Workspace};
use core::dependency::Kind;
use core::manifest::ManifestMetadata;
use ops;
use sources::{RegistrySource};
use util::config;
use util::paths;
use util::{CargoResult, human, ChainError, ToUrl};
use util::config::{Config, ConfigValue, Location};
use util::important_paths::find_root_manifest_for_wd;

pub struct RegistryConfig {
    pub index: Option<String>,
    pub token: Option<String>,
}

pub struct PublishOpts<'cfg> {
    pub config: &'cfg Config,
    pub token: Option<String>,
    pub index: Option<String>,
    pub verify: bool,
    pub allow_dirty: bool,
    pub jobs: Option<u32>,
    pub dry_run: bool,
}

pub fn publish(ws: &Workspace, opts: &PublishOpts) -> CargoResult<()> {
    let pkg = ws.current()?;

    if !pkg.publish() {
        bail!("some crates cannot be published.\n\
               `{}` is marked as unpublishable", pkg.name());
    }

    let (mut registry, reg_id) = registry(opts.config,
                                               opts.token.clone(),
                                               opts.index.clone())?;
    verify_dependencies(&pkg, &reg_id)?;

    // Prepare a tarball, with a non-surpressable warning if metadata
    // is missing since this is being put online.
    let tarball = ops::package(ws, &ops::PackageOpts {
        config: opts.config,
        verify: opts.verify,
        list: false,
        check_metadata: true,
        allow_dirty: opts.allow_dirty,
        jobs: opts.jobs,
    })?.unwrap();

    // Upload said tarball to the specified destination
    opts.config.shell().status("Uploading", pkg.package_id().to_string())?;
    transmit(opts.config, &pkg, tarball.file(), &mut registry, opts.dry_run)?;

    Ok(())
}

fn verify_dependencies(pkg: &Package, registry_src: &SourceId)
                       -> CargoResult<()> {
    for dep in pkg.dependencies().iter() {
        if dep.source_id().is_path() {
            if !dep.specified_req() {
                bail!("all path dependencies must have a version specified \
                       when publishing.\ndependency `{}` does not specify \
                       a version", dep.name())
            }
        } else if dep.source_id() != registry_src {
            bail!("all dependencies must come from the same source.\n\
                   dependency `{}` comes from {} instead",
                  dep.name(), dep.source_id())
        }
    }
    Ok(())
}

fn transmit(config: &Config,
            pkg: &Package,
            tarball: &File,
            registry: &mut Registry,
            dry_run: bool) -> CargoResult<()> {
    let deps = pkg.dependencies().iter().map(|dep| {
        NewCrateDependency {
            optional: dep.is_optional(),
            default_features: dep.uses_default_features(),
            name: dep.name().to_string(),
            features: dep.features().to_vec(),
            version_req: dep.version_req().to_string(),
            target: dep.platform().map(|s| s.to_string()),
            kind: match dep.kind() {
                Kind::Normal => "normal",
                Kind::Build => "build",
                Kind::Development => "dev",
            }.to_string(),
        }
    }).collect::<Vec<NewCrateDependency>>();
    let manifest = pkg.manifest();
    let ManifestMetadata {
        ref authors, ref description, ref homepage, ref documentation,
        ref keywords, ref readme, ref repository, ref license, ref license_file,
    } = *manifest.metadata();
    let readme = match *readme {
        Some(ref readme) => Some(paths::read(&pkg.root().join(readme))?),
        None => None,
    };
    match *license_file {
        Some(ref file) => {
            if fs::metadata(&pkg.root().join(file)).is_err() {
                bail!("the license file `{}` does not exist", file)
            }
        }
        None => {}
    }

    // Do not upload if performing a dry run
    if dry_run {
        config.shell().warn("aborting upload due to dry run")?;
        return Ok(());
    }

    registry.publish(&NewCrate {
        name: pkg.name().to_string(),
        vers: pkg.version().to_string(),
        deps: deps,
        features: pkg.summary().features().clone(),
        authors: authors.clone(),
        description: description.clone(),
        homepage: homepage.clone(),
        documentation: documentation.clone(),
        keywords: keywords.clone(),
        readme: readme,
        repository: repository.clone(),
        license: license.clone(),
        license_file: license_file.clone(),
    }, tarball).map_err(|e| {
        human(e.to_string())
    })
}

pub fn registry_configuration(config: &Config) -> CargoResult<RegistryConfig> {
    let index = config.get_string("registry.index")?.map(|p| p.val);
    let token = config.get_string("registry.token")?.map(|p| p.val);
    Ok(RegistryConfig { index: index, token: token })
}

pub fn registry(config: &Config,
                token: Option<String>,
                index: Option<String>) -> CargoResult<(Registry, SourceId)> {
    // Parse all configuration options
    let RegistryConfig {
        token: token_config,
        index: _index_config,
    } = registry_configuration(config)?;
    let token = token.or(token_config);
    let sid = match index {
        Some(index) => SourceId::for_registry(&index.to_url()?),
        None => SourceId::crates_io(config)?,
    };
    let api_host = {
        let mut src = RegistrySource::remote(&sid, config);
        src.update().chain_error(|| {
            human(format!("failed to update {}", sid))
        })?;
        (src.config()?).unwrap().api
    };
    let handle = http_handle(config)?;
    Ok((Registry::new_handle(api_host, token, handle), sid))
}

/// Create a new HTTP handle with appropriate global configuration for cargo.
pub fn http_handle(config: &Config) -> CargoResult<Easy> {
    if !config.network_allowed() {
        bail!("attempting to make an HTTP request, but --frozen was \
               specified")
    }

    // The timeout option for libcurl by default times out the entire transfer,
    // but we probably don't want this. Instead we only set timeouts for the
    // connect phase as well as a "low speed" timeout so if we don't receive
    // many bytes in a large-ish period of time then we time out.
    let mut handle = Easy::new();
    handle.connect_timeout(Duration::new(30, 0))?;
    handle.low_speed_limit(10 /* bytes per second */)?;
    handle.low_speed_time(Duration::new(30, 0))?;
    if let Some(proxy) = http_proxy(config)? {
        handle.proxy(&proxy)?;
    }
    if let Some(cainfo) = config.get_path("http.cainfo")? {
        handle.cainfo(&cainfo.val)?;
    }
    if let Some(timeout) = http_timeout(config)? {
        handle.connect_timeout(Duration::new(timeout as u64, 0))?;
        handle.low_speed_time(Duration::new(timeout as u64, 0))?;
    }
    Ok(handle)
}

/// Find an explicit HTTP proxy if one is available.
///
/// Favor cargo's `http.proxy`, then git's `http.proxy`. Proxies specified
/// via environment variables are picked up by libcurl.
fn http_proxy(config: &Config) -> CargoResult<Option<String>> {
    match config.get_string("http.proxy")? {
        Some(s) => return Ok(Some(s.val)),
        None => {}
    }
    match git2::Config::open_default() {
        Ok(cfg) => {
            match cfg.get_str("http.proxy") {
                Ok(s) => return Ok(Some(s.to_string())),
                Err(..) => {}
            }
        }
        Err(..) => {}
    }
    Ok(None)
}

/// Determine if an http proxy exists.
///
/// Checks the following for existence, in order:
///
/// * cargo's `http.proxy`
/// * git's `http.proxy`
/// * http_proxy env var
/// * HTTP_PROXY env var
/// * https_proxy env var
/// * HTTPS_PROXY env var
pub fn http_proxy_exists(config: &Config) -> CargoResult<bool> {
    if http_proxy(config)?.is_some() {
        Ok(true)
    } else {
        Ok(["http_proxy", "HTTP_PROXY",
           "https_proxy", "HTTPS_PROXY"].iter().any(|v| env::var(v).is_ok()))
    }
}

pub fn http_timeout(config: &Config) -> CargoResult<Option<i64>> {
    match config.get_i64("http.timeout")? {
        Some(s) => return Ok(Some(s.val)),
        None => {}
    }
    Ok(env::var("HTTP_TIMEOUT").ok().and_then(|s| s.parse().ok()))
}

pub fn registry_login(config: &Config, token: String) -> CargoResult<()> {
    let RegistryConfig { index, token: _ } = registry_configuration(config)?;
    let mut map = HashMap::new();
    let p = config.cwd().to_path_buf();
    match index {
        Some(index) => {
            map.insert("index".to_string(), ConfigValue::String(index, p.clone()));
        }
        None => {}
    }
    map.insert("token".to_string(), ConfigValue::String(token, p));

    config::set_config(config, Location::Global, "registry",
                       ConfigValue::Table(map, PathBuf::from(".")))
}

pub struct OwnersOptions {
    pub krate: Option<String>,
    pub token: Option<String>,
    pub index: Option<String>,
    pub to_add: Option<Vec<String>>,
    pub to_remove: Option<Vec<String>>,
    pub list: bool,
}

pub fn modify_owners(config: &Config, opts: &OwnersOptions) -> CargoResult<()> {
    let name = match opts.krate {
        Some(ref name) => name.clone(),
        None => {
            let manifest_path = find_root_manifest_for_wd(None, config.cwd())?;
            let pkg = Package::for_path(&manifest_path, config)?;
            pkg.name().to_string()
        }
    };

    let (mut registry, _) = registry(config, opts.token.clone(),
                                          opts.index.clone())?;

    match opts.to_add {
        Some(ref v) => {
            let v = v.iter().map(|s| &s[..]).collect::<Vec<_>>();
            config.shell().status("Owner", format!("adding {:?} to crate {}",
                                                        v, name))?;
            registry.add_owners(&name, &v).map_err(|e| {
                human(format!("failed to add owners to crate {}: {}", name, e))
            })?;
        }
        None => {}
    }

    match opts.to_remove {
        Some(ref v) => {
            let v = v.iter().map(|s| &s[..]).collect::<Vec<_>>();
            config.shell().status("Owner", format!("removing {:?} from crate {}",
                                                        v, name))?;
            registry.remove_owners(&name, &v).map_err(|e| {
                human(format!("failed to remove owners from crate {}: {}", name, e))
            })?;
        }
        None => {}
    }

    if opts.list {
        let owners = registry.list_owners(&name).map_err(|e| {
            human(format!("failed to list owners of crate {}: {}", name, e))
        })?;
        for owner in owners.iter() {
            print!("{}", owner.login);
            match (owner.name.as_ref(), owner.email.as_ref()) {
                (Some(name), Some(email)) => println!(" ({} <{}>)", name, email),
                (Some(s), None) |
                (None, Some(s)) => println!(" ({})", s),
                (None, None) => println!(""),
            }
        }
    }

    Ok(())
}

pub fn yank(config: &Config,
            krate: Option<String>,
            version: Option<String>,
            token: Option<String>,
            index: Option<String>,
            undo: bool) -> CargoResult<()> {
    let name = match krate {
        Some(name) => name,
        None => {
            let manifest_path = find_root_manifest_for_wd(None, config.cwd())?;
            let pkg = Package::for_path(&manifest_path, config)?;
            pkg.name().to_string()
        }
    };
    let version = match version {
        Some(v) => v,
        None => bail!("a version must be specified to yank")
    };

    let (mut registry, _) = registry(config, token, index)?;

    if undo {
        config.shell().status("Unyank", format!("{}:{}", name, version))?;
        registry.unyank(&name, &version).map_err(|e| {
            human(format!("failed to undo a yank: {}", e))
        })?;
    } else {
        config.shell().status("Yank", format!("{}:{}", name, version))?;
        registry.yank(&name, &version).map_err(|e| {
            human(format!("failed to yank: {}", e))
        })?;
    }

    Ok(())
}

pub fn search(query: &str,
              config: &Config,
              index: Option<String>,
              limit: u8) -> CargoResult<()> {
    fn truncate_with_ellipsis(s: &str, max_length: usize) -> String {
        if s.len() < max_length {
            s.to_string()
        } else {
            format!("{}…", &s[..max_length - 1])
        }
    }

    let (mut registry, _) = registry(config, None, index)?;
    let (crates, total_crates) = registry.search(query, limit).map_err(|e| {
        human(format!("failed to retrieve search results from the registry: {}", e))
    })?;

    let list_items = crates.iter()
        .map(|krate| (
            format!("{} ({})", krate.name, krate.max_version),
            krate.description.as_ref().map(|desc|
                truncate_with_ellipsis(&desc.replace("\n", " "), 128))
        ))
        .collect::<Vec<_>>();
    let description_margin = list_items.iter()
        .map(|&(ref left, _)| left.len() + 4)
        .max()
        .unwrap_or(0);

    for (name, description) in list_items.into_iter() {
        let line = match description {
            Some(desc) => {
                let space = repeat(' ').take(description_margin - name.len())
                                       .collect::<String>();
                name + &space + &desc
            }
            None => name
        };
        config.shell().say(line, BLACK)?;
    }

    let search_max_limit = 100;
    if total_crates > limit as u32 && limit < search_max_limit {
        config.shell().say(
            format!("... and {} crates more (use --limit N to see more)",
                    total_crates - limit as u32),
            BLACK)
        ?;
    } else if total_crates > limit as u32 && limit >= search_max_limit {
        config.shell().say(
            format!(
                "... and {} crates more (go to http://crates.io/search?q={} to see more)",
                total_crates - limit as u32,
                percent_encode(query.as_bytes(), QUERY_ENCODE_SET)
            ),
            BLACK)
        ?;
    }

    Ok(())
}
<commit_msg>Auto merge of #3485 - sfackler:useragent, r=alexcrichton<commit_after>use std::collections::HashMap;
use std::env;
use std::fs::{self, File};
use std::iter::repeat;
use std::path::PathBuf;
use std::time::Duration;

use curl::easy::Easy;
use git2;
use registry::{Registry, NewCrate, NewCrateDependency};
use term::color::BLACK;

use url::percent_encoding::{percent_encode, QUERY_ENCODE_SET};

use version;
use core::source::Source;
use core::{Package, SourceId, Workspace};
use core::dependency::Kind;
use core::manifest::ManifestMetadata;
use ops;
use sources::{RegistrySource};
use util::config;
use util::paths;
use util::{CargoResult, human, ChainError, ToUrl};
use util::config::{Config, ConfigValue, Location};
use util::important_paths::find_root_manifest_for_wd;

pub struct RegistryConfig {
    pub index: Option<String>,
    pub token: Option<String>,
}

pub struct PublishOpts<'cfg> {
    pub config: &'cfg Config,
    pub token: Option<String>,
    pub index: Option<String>,
    pub verify: bool,
    pub allow_dirty: bool,
    pub jobs: Option<u32>,
    pub dry_run: bool,
}

pub fn publish(ws: &Workspace, opts: &PublishOpts) -> CargoResult<()> {
    let pkg = ws.current()?;

    if !pkg.publish() {
        bail!("some crates cannot be published.\n\
               `{}` is marked as unpublishable", pkg.name());
    }

    let (mut registry, reg_id) = registry(opts.config,
                                               opts.token.clone(),
                                               opts.index.clone())?;
    verify_dependencies(&pkg, &reg_id)?;

    // Prepare a tarball, with a non-surpressable warning if metadata
    // is missing since this is being put online.
    let tarball = ops::package(ws, &ops::PackageOpts {
        config: opts.config,
        verify: opts.verify,
        list: false,
        check_metadata: true,
        allow_dirty: opts.allow_dirty,
        jobs: opts.jobs,
    })?.unwrap();

    // Upload said tarball to the specified destination
    opts.config.shell().status("Uploading", pkg.package_id().to_string())?;
    transmit(opts.config, &pkg, tarball.file(), &mut registry, opts.dry_run)?;

    Ok(())
}

fn verify_dependencies(pkg: &Package, registry_src: &SourceId)
                       -> CargoResult<()> {
    for dep in pkg.dependencies().iter() {
        if dep.source_id().is_path() {
            if !dep.specified_req() {
                bail!("all path dependencies must have a version specified \
                       when publishing.\ndependency `{}` does not specify \
                       a version", dep.name())
            }
        } else if dep.source_id() != registry_src {
            bail!("all dependencies must come from the same source.\n\
                   dependency `{}` comes from {} instead",
                  dep.name(), dep.source_id())
        }
    }
    Ok(())
}

fn transmit(config: &Config,
            pkg: &Package,
            tarball: &File,
            registry: &mut Registry,
            dry_run: bool) -> CargoResult<()> {
    let deps = pkg.dependencies().iter().map(|dep| {
        NewCrateDependency {
            optional: dep.is_optional(),
            default_features: dep.uses_default_features(),
            name: dep.name().to_string(),
            features: dep.features().to_vec(),
            version_req: dep.version_req().to_string(),
            target: dep.platform().map(|s| s.to_string()),
            kind: match dep.kind() {
                Kind::Normal => "normal",
                Kind::Build => "build",
                Kind::Development => "dev",
            }.to_string(),
        }
    }).collect::<Vec<NewCrateDependency>>();
    let manifest = pkg.manifest();
    let ManifestMetadata {
        ref authors, ref description, ref homepage, ref documentation,
        ref keywords, ref readme, ref repository, ref license, ref license_file,
    } = *manifest.metadata();
    let readme = match *readme {
        Some(ref readme) => Some(paths::read(&pkg.root().join(readme))?),
        None => None,
    };
    match *license_file {
        Some(ref file) => {
            if fs::metadata(&pkg.root().join(file)).is_err() {
                bail!("the license file `{}` does not exist", file)
            }
        }
        None => {}
    }

    // Do not upload if performing a dry run
    if dry_run {
        config.shell().warn("aborting upload due to dry run")?;
        return Ok(());
    }

    registry.publish(&NewCrate {
        name: pkg.name().to_string(),
        vers: pkg.version().to_string(),
        deps: deps,
        features: pkg.summary().features().clone(),
        authors: authors.clone(),
        description: description.clone(),
        homepage: homepage.clone(),
        documentation: documentation.clone(),
        keywords: keywords.clone(),
        readme: readme,
        repository: repository.clone(),
        license: license.clone(),
        license_file: license_file.clone(),
    }, tarball).map_err(|e| {
        human(e.to_string())
    })
}

pub fn registry_configuration(config: &Config) -> CargoResult<RegistryConfig> {
    let index = config.get_string("registry.index")?.map(|p| p.val);
    let token = config.get_string("registry.token")?.map(|p| p.val);
    Ok(RegistryConfig { index: index, token: token })
}

pub fn registry(config: &Config,
                token: Option<String>,
                index: Option<String>) -> CargoResult<(Registry, SourceId)> {
    // Parse all configuration options
    let RegistryConfig {
        token: token_config,
        index: _index_config,
    } = registry_configuration(config)?;
    let token = token.or(token_config);
    let sid = match index {
        Some(index) => SourceId::for_registry(&index.to_url()?),
        None => SourceId::crates_io(config)?,
    };
    let api_host = {
        let mut src = RegistrySource::remote(&sid, config);
        src.update().chain_error(|| {
            human(format!("failed to update {}", sid))
        })?;
        (src.config()?).unwrap().api
    };
    let handle = http_handle(config)?;
    Ok((Registry::new_handle(api_host, token, handle), sid))
}

/// Create a new HTTP handle with appropriate global configuration for cargo.
pub fn http_handle(config: &Config) -> CargoResult<Easy> {
    if !config.network_allowed() {
        bail!("attempting to make an HTTP request, but --frozen was \
               specified")
    }

    // The timeout option for libcurl by default times out the entire transfer,
    // but we probably don't want this. Instead we only set timeouts for the
    // connect phase as well as a "low speed" timeout so if we don't receive
    // many bytes in a large-ish period of time then we time out.
    let mut handle = Easy::new();
    handle.connect_timeout(Duration::new(30, 0))?;
    handle.low_speed_limit(10 /* bytes per second */)?;
    handle.low_speed_time(Duration::new(30, 0))?;
    handle.useragent(&version())?;
    if let Some(proxy) = http_proxy(config)? {
        handle.proxy(&proxy)?;
    }
    if let Some(cainfo) = config.get_path("http.cainfo")? {
        handle.cainfo(&cainfo.val)?;
    }
    if let Some(timeout) = http_timeout(config)? {
        handle.connect_timeout(Duration::new(timeout as u64, 0))?;
        handle.low_speed_time(Duration::new(timeout as u64, 0))?;
    }
    Ok(handle)
}

/// Find an explicit HTTP proxy if one is available.
///
/// Favor cargo's `http.proxy`, then git's `http.proxy`. Proxies specified
/// via environment variables are picked up by libcurl.
fn http_proxy(config: &Config) -> CargoResult<Option<String>> {
    match config.get_string("http.proxy")? {
        Some(s) => return Ok(Some(s.val)),
        None => {}
    }
    match git2::Config::open_default() {
        Ok(cfg) => {
            match cfg.get_str("http.proxy") {
                Ok(s) => return Ok(Some(s.to_string())),
                Err(..) => {}
            }
        }
        Err(..) => {}
    }
    Ok(None)
}

/// Determine if an http proxy exists.
///
/// Checks the following for existence, in order:
///
/// * cargo's `http.proxy`
/// * git's `http.proxy`
/// * http_proxy env var
/// * HTTP_PROXY env var
/// * https_proxy env var
/// * HTTPS_PROXY env var
pub fn http_proxy_exists(config: &Config) -> CargoResult<bool> {
    if http_proxy(config)?.is_some() {
        Ok(true)
    } else {
        Ok(["http_proxy", "HTTP_PROXY",
           "https_proxy", "HTTPS_PROXY"].iter().any(|v| env::var(v).is_ok()))
    }
}

pub fn http_timeout(config: &Config) -> CargoResult<Option<i64>> {
    match config.get_i64("http.timeout")? {
        Some(s) => return Ok(Some(s.val)),
        None => {}
    }
    Ok(env::var("HTTP_TIMEOUT").ok().and_then(|s| s.parse().ok()))
}

pub fn registry_login(config: &Config, token: String) -> CargoResult<()> {
    let RegistryConfig { index, token: _ } = registry_configuration(config)?;
    let mut map = HashMap::new();
    let p = config.cwd().to_path_buf();
    match index {
        Some(index) => {
            map.insert("index".to_string(), ConfigValue::String(index, p.clone()));
        }
        None => {}
    }
    map.insert("token".to_string(), ConfigValue::String(token, p));

    config::set_config(config, Location::Global, "registry",
                       ConfigValue::Table(map, PathBuf::from(".")))
}

pub struct OwnersOptions {
    pub krate: Option<String>,
    pub token: Option<String>,
    pub index: Option<String>,
    pub to_add: Option<Vec<String>>,
    pub to_remove: Option<Vec<String>>,
    pub list: bool,
}

pub fn modify_owners(config: &Config, opts: &OwnersOptions) -> CargoResult<()> {
    let name = match opts.krate {
        Some(ref name) => name.clone(),
        None => {
            let manifest_path = find_root_manifest_for_wd(None, config.cwd())?;
            let pkg = Package::for_path(&manifest_path, config)?;
            pkg.name().to_string()
        }
    };

    let (mut registry, _) = registry(config, opts.token.clone(),
                                          opts.index.clone())?;

    match opts.to_add {
        Some(ref v) => {
            let v = v.iter().map(|s| &s[..]).collect::<Vec<_>>();
            config.shell().status("Owner", format!("adding {:?} to crate {}",
                                                        v, name))?;
            registry.add_owners(&name, &v).map_err(|e| {
                human(format!("failed to add owners to crate {}: {}", name, e))
            })?;
        }
        None => {}
    }

    match opts.to_remove {
        Some(ref v) => {
            let v = v.iter().map(|s| &s[..]).collect::<Vec<_>>();
            config.shell().status("Owner", format!("removing {:?} from crate {}",
                                                        v, name))?;
            registry.remove_owners(&name, &v).map_err(|e| {
                human(format!("failed to remove owners from crate {}: {}", name, e))
            })?;
        }
        None => {}
    }

    if opts.list {
        let owners = registry.list_owners(&name).map_err(|e| {
            human(format!("failed to list owners of crate {}: {}", name, e))
        })?;
        for owner in owners.iter() {
            print!("{}", owner.login);
            match (owner.name.as_ref(), owner.email.as_ref()) {
                (Some(name), Some(email)) => println!(" ({} <{}>)", name, email),
                (Some(s), None) |
                (None, Some(s)) => println!(" ({})", s),
                (None, None) => println!(""),
            }
        }
    }

    Ok(())
}

pub fn yank(config: &Config,
            krate: Option<String>,
            version: Option<String>,
            token: Option<String>,
            index: Option<String>,
            undo: bool) -> CargoResult<()> {
    let name = match krate {
        Some(name) => name,
        None => {
            let manifest_path = find_root_manifest_for_wd(None, config.cwd())?;
            let pkg = Package::for_path(&manifest_path, config)?;
            pkg.name().to_string()
        }
    };
    let version = match version {
        Some(v) => v,
        None => bail!("a version must be specified to yank")
    };

    let (mut registry, _) = registry(config, token, index)?;

    if undo {
        config.shell().status("Unyank", format!("{}:{}", name, version))?;
        registry.unyank(&name, &version).map_err(|e| {
            human(format!("failed to undo a yank: {}", e))
        })?;
    } else {
        config.shell().status("Yank", format!("{}:{}", name, version))?;
        registry.yank(&name, &version).map_err(|e| {
            human(format!("failed to yank: {}", e))
        })?;
    }

    Ok(())
}

pub fn search(query: &str,
              config: &Config,
              index: Option<String>,
              limit: u8) -> CargoResult<()> {
    fn truncate_with_ellipsis(s: &str, max_length: usize) -> String {
        if s.len() < max_length {
            s.to_string()
        } else {
            format!("{}…", &s[..max_length - 1])
        }
    }

    let (mut registry, _) = registry(config, None, index)?;
    let (crates, total_crates) = registry.search(query, limit).map_err(|e| {
        human(format!("failed to retrieve search results from the registry: {}", e))
    })?;

    let list_items = crates.iter()
        .map(|krate| (
            format!("{} ({})", krate.name, krate.max_version),
            krate.description.as_ref().map(|desc|
                truncate_with_ellipsis(&desc.replace("\n", " "), 128))
        ))
        .collect::<Vec<_>>();
    let description_margin = list_items.iter()
        .map(|&(ref left, _)| left.len() + 4)
        .max()
        .unwrap_or(0);

    for (name, description) in list_items.into_iter() {
        let line = match description {
            Some(desc) => {
                let space = repeat(' ').take(description_margin - name.len())
                                       .collect::<String>();
                name + &space + &desc
            }
            None => name
        };
        config.shell().say(line, BLACK)?;
    }

    let search_max_limit = 100;
    if total_crates > limit as u32 && limit < search_max_limit {
        config.shell().say(
            format!("... and {} crates more (use --limit N to see more)",
                    total_crates - limit as u32),
            BLACK)
        ?;
    } else if total_crates > limit as u32 && limit >= search_max_limit {
        config.shell().say(
            format!(
                "... and {} crates more (go to http://crates.io/search?q={} to see more)",
                total_crates - limit as u32,
                percent_encode(query.as_bytes(), QUERY_ENCODE_SET)
            ),
            BLACK)
        ?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>use itertools::Itertools;
use std::iter::successors;
use syntax::{
    algo::neighbor,
    ast::{self, AstNode},
    Direction,
};

use crate::{AssistContext, AssistId, AssistKind, Assists, TextRange};

// Assist: merge_match_arms
//
// Merges the current match arm with the following if their bodies are identical.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         $0Action::Move(..) => foo(),
//         Action::Stop => foo(),
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move(..) | Action::Stop => foo(),
//     }
// }
// ```
pub(crate) fn merge_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let current_arm = ctx.find_node_at_offset::<ast::MatchArm>()?;
    // Don't try to handle arms with guards for now - can add support for this later
    if current_arm.guard().is_some() {
        return None;
    }
    let current_expr = current_arm.expr()?;
    let current_text_range = current_arm.syntax().text_range();
    let current_arm_types = get_arm_types(&ctx, &current_arm);

    // We check if the following match arms match this one. We could, but don't,
    // compare to the previous match arm as well.
    let arms_to_merge = successors(Some(current_arm), |it| neighbor(it, Direction::Next))
        .take_while(|arm| match arm.expr() {
            Some(expr) if arm.guard().is_none() => {
                let same_text = expr.syntax().text() == current_expr.syntax().text();
                if !same_text {
                    return false;
                }

                return are_same_types(&current_arm_types, arm, ctx);
            }
            _ => false,
        })
        .collect::<Vec<_>>();

    if arms_to_merge.len() <= 1 {
        return None;
    }

    acc.add(
        AssistId("merge_match_arms", AssistKind::RefactorRewrite),
        "Merge match arms",
        current_text_range,
        |edit| {
            let pats = if arms_to_merge.iter().any(contains_placeholder) {
                "_".into()
            } else {
                arms_to_merge
                    .iter()
                    .filter_map(ast::MatchArm::pat)
                    .map(|x| x.syntax().to_string())
                    .collect::<Vec<String>>()
                    .join(" | ")
            };

            let arm = format!("{} => {},", pats, current_expr.syntax().text());

            if let [first, .., last] = &*arms_to_merge {
                let start = first.syntax().text_range().start();
                let end = last.syntax().text_range().end();

                edit.replace(TextRange::new(start, end), arm);
            }
        },
    )
}

fn contains_placeholder(a: &ast::MatchArm) -> bool {
    matches!(a.pat(), Some(ast::Pat::WildcardPat(..)))
}

fn are_same_types(
    current_arm_types: &Vec<Option<hir::TypeInfo>>,
    arm: &ast::MatchArm,
    ctx: &AssistContext,
) -> bool {
    let arm_types = get_arm_types(&ctx, &arm);
    for i in 0..arm_types.len() {
        let other_arm_type = &arm_types[i];
        let current_arm_type = &current_arm_types[i];
        if let (Some(other_arm_type), Some(current_arm_type)) = (other_arm_type, current_arm_type) {
            return &other_arm_type.original == &current_arm_type.original;
        }
    }

    return true;
}

fn get_arm_types(ctx: &AssistContext, arm: &ast::MatchArm) -> Vec<Option<hir::TypeInfo>> {
    match arm.pat() {
        Some(ast::Pat::TupleStructPat(tp)) => {
            tp.fields().into_iter().map(|field| ctx.sema.type_of_pat(&field)).collect_vec()
        }
        _ => Vec::new(),
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn merge_match_arms_single_patterns() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32$0 }
        X::B => { 1i32 }
        X::C => { 2i32 }
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B => { 1i32 },
        X::C => { 2i32 }
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_multiple_patterns() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B => {$0 1i32 },
        X::C | X::D => { 1i32 },
        X::E => { 2i32 },
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B | X::C | X::D => { 1i32 },
        X::E => { 2i32 },
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_placeholder_pattern() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32 },
        X::B => { 2i$032 },
        _ => { 2i32 }
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32 },
        _ => { 2i32 },
    }
}
"#,
        );
    }

    #[test]
    fn merges_all_subsequent_arms() {
        check_assist(
            merge_match_arms,
            r#"
enum X { A, B, C, D, E }

fn main() {
    match X::A {
        X::A$0 => 92,
        X::B => 92,
        X::C => 92,
        X::D => 62,
        _ => panic!(),
    }
}
"#,
            r#"
enum X { A, B, C, D, E }

fn main() {
    match X::A {
        X::A | X::B | X::C => 92,
        X::D => 62,
        _ => panic!(),
    }
}
"#,
        )
    }

    #[test]
    fn merge_match_arms_rejects_guards() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X {
    A(i32),
    B,
    C
}

fn main() {
    let x = X::A;
    let y = match x {
        X::A(a) if a > 5 => { $01i32 },
        X::B => { 1i32 },
        X::C => { 2i32 }
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_different_type() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<f64, f32>::Ok(0f64) {
        Ok(x) => $0x.classify(),
        Err(x) => x.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_different_type_multiple_fields() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64), (f32, f32)>::Ok((0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_multiple_fields() {
        check_assist(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
            r#"
fn func() {
    match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
        Ok(x) | Err(x) => x.1.classify(),
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_subsequent_arm_with_different_type_in_other() {
        check_assist(
            merge_match_arms,
            r#"
enum MyEnum {
    OptionA(f32),
    OptionB(f32),
    OptionC(f64)
}

fn func(e: MyEnum) {
    match e {
        MyEnum::OptionA(x) => $0x.classify(),
        MyEnum::OptionB(x) => x.classify(),
        MyEnum::OptionC(x) => x.classify(),
    };
}
"#,
            r#"
enum MyEnum {
    OptionA(f32),
    OptionB(f32),
    OptionC(f64)
}

fn func(e: MyEnum) {
    match e {
        MyEnum::OptionA(x) | MyEnum::OptionB(x) => x.classify(),
        MyEnum::OptionC(x) => x.classify(),
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_skip_arm_with_different_type_in_between() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum MyEnum {
    OptionA(f32),
    OptionB(f64),
    OptionC(f32)
}

fn func(e: MyEnum) {
    match e {
        MyEnum::OptionA(x) => $0x.classify(),
        MyEnum::OptionB(x) => x.classify(),
        MyEnum::OptionC(x) => x.classify(),
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_different_number_of_fields() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64, f64), (f64, f64)>::Ok((0f64, 0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
        );
    }
}
<commit_msg>Added a few more test cases<commit_after>use itertools::Itertools;
use std::iter::successors;
use syntax::{
    algo::neighbor,
    ast::{self, AstNode},
    Direction,
};

use crate::{AssistContext, AssistId, AssistKind, Assists, TextRange};

// Assist: merge_match_arms
//
// Merges the current match arm with the following if their bodies are identical.
//
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         $0Action::Move(..) => foo(),
//         Action::Stop => foo(),
//     }
// }
// ```
// ->
// ```
// enum Action { Move { distance: u32 }, Stop }
//
// fn handle(action: Action) {
//     match action {
//         Action::Move(..) | Action::Stop => foo(),
//     }
// }
// ```
pub(crate) fn merge_match_arms(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let current_arm = ctx.find_node_at_offset::<ast::MatchArm>()?;
    // Don't try to handle arms with guards for now - can add support for this later
    if current_arm.guard().is_some() {
        return None;
    }
    let current_expr = current_arm.expr()?;
    let current_text_range = current_arm.syntax().text_range();
    let current_arm_types = get_arm_types(&ctx, &current_arm);

    // We check if the following match arms match this one. We could, but don't,
    // compare to the previous match arm as well.
    let arms_to_merge = successors(Some(current_arm), |it| neighbor(it, Direction::Next))
        .take_while(|arm| match arm.expr() {
            Some(expr) if arm.guard().is_none() => {
                let same_text = expr.syntax().text() == current_expr.syntax().text();
                if !same_text {
                    return false;
                }

                return are_same_types(&current_arm_types, arm, ctx);
            }
            _ => false,
        })
        .collect::<Vec<_>>();

    if arms_to_merge.len() <= 1 {
        return None;
    }

    acc.add(
        AssistId("merge_match_arms", AssistKind::RefactorRewrite),
        "Merge match arms",
        current_text_range,
        |edit| {
            let pats = if arms_to_merge.iter().any(contains_placeholder) {
                "_".into()
            } else {
                arms_to_merge
                    .iter()
                    .filter_map(ast::MatchArm::pat)
                    .map(|x| x.syntax().to_string())
                    .collect::<Vec<String>>()
                    .join(" | ")
            };

            let arm = format!("{} => {},", pats, current_expr.syntax().text());

            if let [first, .., last] = &*arms_to_merge {
                let start = first.syntax().text_range().start();
                let end = last.syntax().text_range().end();

                edit.replace(TextRange::new(start, end), arm);
            }
        },
    )
}

fn contains_placeholder(a: &ast::MatchArm) -> bool {
    matches!(a.pat(), Some(ast::Pat::WildcardPat(..)))
}

fn are_same_types(
    current_arm_types: &Vec<Option<hir::TypeInfo>>,
    arm: &ast::MatchArm,
    ctx: &AssistContext,
) -> bool {
    let arm_types = get_arm_types(&ctx, &arm);
    for i in 0..arm_types.len() {
        let other_arm_type = &arm_types[i];
        let current_arm_type = &current_arm_types[i];
        if let (Some(other_arm_type), Some(current_arm_type)) = (other_arm_type, current_arm_type) {
            return &other_arm_type.original == &current_arm_type.original;
        }
    }

    return true;
}

fn get_arm_types(ctx: &AssistContext, arm: &ast::MatchArm) -> Vec<Option<hir::TypeInfo>> {
    match arm.pat() {
        Some(ast::Pat::TupleStructPat(tp)) => {
            tp.fields().into_iter().map(|field| ctx.sema.type_of_pat(&field)).collect_vec()
        }
        _ => Vec::new(),
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{check_assist, check_assist_not_applicable};

    use super::*;

    #[test]
    fn merge_match_arms_single_patterns() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32$0 }
        X::B => { 1i32 }
        X::C => { 2i32 }
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B => { 1i32 },
        X::C => { 2i32 }
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_multiple_patterns() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B => {$0 1i32 },
        X::C | X::D => { 1i32 },
        X::E => { 2i32 },
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A | X::B | X::C | X::D => { 1i32 },
        X::E => { 2i32 },
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_placeholder_pattern() {
        check_assist(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32 },
        X::B => { 2i$032 },
        _ => { 2i32 }
    }
}
"#,
            r#"
#[derive(Debug)]
enum X { A, B, C, D, E }

fn main() {
    let x = X::A;
    let y = match x {
        X::A => { 1i32 },
        _ => { 2i32 },
    }
}
"#,
        );
    }

    #[test]
    fn merges_all_subsequent_arms() {
        check_assist(
            merge_match_arms,
            r#"
enum X { A, B, C, D, E }

fn main() {
    match X::A {
        X::A$0 => 92,
        X::B => 92,
        X::C => 92,
        X::D => 62,
        _ => panic!(),
    }
}
"#,
            r#"
enum X { A, B, C, D, E }

fn main() {
    match X::A {
        X::A | X::B | X::C => 92,
        X::D => 62,
        _ => panic!(),
    }
}
"#,
        )
    }

    #[test]
    fn merge_match_arms_rejects_guards() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
#[derive(Debug)]
enum X {
    A(i32),
    B,
    C
}

fn main() {
    let x = X::A;
    let y = match x {
        X::A(a) if a > 5 => { $01i32 },
        X::B => { 1i32 },
        X::C => { 2i32 }
    }
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_different_type() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<f64, f32>::Ok(0f64) {
        Ok(x) => $0x.classify(),
        Err(x) => x.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_different_type_multiple_fields() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64), (f32, f32)>::Ok((0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_multiple_fields() {
        check_assist(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
            r#"
fn func() {
    match Result::<(f64, f64), (f64, f64)>::Ok((0f64, 0f64)) {
        Ok(x) | Err(x) => x.1.classify(),
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_subsequent_arm_with_different_type_in_other() {
        check_assist(
            merge_match_arms,
            r#"
enum MyEnum {
    OptionA(f32),
    OptionB(f32),
    OptionC(f64)
}

fn func(e: MyEnum) {
    match e {
        MyEnum::OptionA(x) => $0x.classify(),
        MyEnum::OptionB(x) => x.classify(),
        MyEnum::OptionC(x) => x.classify(),
    };
}
"#,
            r#"
enum MyEnum {
    OptionA(f32),
    OptionB(f32),
    OptionC(f64)
}

fn func(e: MyEnum) {
    match e {
        MyEnum::OptionA(x) | MyEnum::OptionB(x) => x.classify(),
        MyEnum::OptionC(x) => x.classify(),
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_skip_arm_with_different_type_in_between() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum MyEnum {
    OptionA(f32),
    OptionB(f64),
    OptionC(f32)
}

fn func(e: MyEnum) {
    match e {
        MyEnum::OptionA(x) => $0x.classify(),
        MyEnum::OptionB(x) => x.classify(),
        MyEnum::OptionC(x) => x.classify(),
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_same_type_different_number_of_fields() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"//- minicore: result
fn func() {
    match Result::<(f64, f64, f64), (f64, f64)>::Ok((0f64, 0f64, 0f64)) {
        Ok(x) => $0x.1.classify(),
        Err(x) => x.1.classify()
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_same_destructuring_different_types() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
struct Point {
    x: i32,
    y: i32,
}

fn func() {
    let p = Point { x: 0, y: 7 };

    match p {
        Point { x, y: 0 } => $0"",
        Point { x: 0, y } => "",
        Point { x, y } => "",
    };
}
"#,
        );
    }

    #[test]
    fn merge_match_arms_range() {
        check_assist(
            merge_match_arms,
            r#"
let x = 'c';

    match x {
        'a'..='j' => $0"",
        'c'..='z' => "",
        _ => "other",
    };
"#,
            r#"
let x = 'c';

    match x {
        'a'..='j' | 'c'..='z' => "",
        _ => "other",
    };
"#,
        );
    }

    #[test]
    fn merge_match_arms_enum_without_field() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum MyEnum {
    NoField,
    AField(u8)
}

fn func(x: MyEnum) {
    match x {
        MyEnum::NoField => $0"",
        MyEnum::AField(x) => ""
    };
}
        "#,
        )
    }

    #[test]
    fn merge_match_arms_enum_destructuring_different_types() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum MyEnum {
    Move { x: i32, y: i32 },
    Write(String),
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, y } => $0"",
        MyEnum::Write(text) => "",
    };
}
        "#,
        )
    }

    #[test]
    fn merge_match_arms_enum_destructuring_same_types() {
        check_assist(
            merge_match_arms,
            r#"
enum MyEnum {
    Move { x: i32, y: i32 },
    Crawl { x: i32, y: i32 }
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, y } => $0"",
        MyEnum::Crawl { x, y } => "",
    };
}
        "#,
            r#"
enum MyEnum {
    Move { x: i32, y: i32 },
    Crawl { x: i32, y: i32 }
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, y } | MyEnum::Crawl { x, y } => "",
    };
}
        "#,
        )
    }

    #[test]
    fn merge_match_arms_enum_destructuring_same_types_different_name() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum MyEnum {
    Move { x: i32, y: i32 },
    Crawl { a: i32, b: i32 }
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, y } => $0"",
        MyEnum::Crawl { a, b } => "",
    };
}
        "#
        )
    }

    #[test]
    fn merge_match_arms_enum_nested_pattern_different_names() {
        check_assist_not_applicable(
            merge_match_arms,
            r#"
enum Color {
    Rgb(i32, i32, i32),
    Hsv(i32, i32, i32),
}

enum Message {
    Quit,
    Move { x: i32, y: i32 },
    Write(String),
    ChangeColor(Color),
}

fn main(msg: Message) {
    match msg {
        Message::ChangeColor(Color::Rgb(r, g, b)) => $0"",
        Message::ChangeColor(Color::Hsv(h, s, v)) => "",
        _ => "other"
    };
}
        "#,
        )
    }

    #[test]
    fn merge_match_arms_enum_nested_pattern_same_names() {
        check_assist(
            merge_match_arms,
            r#"
enum Color {
    Rgb(i32, i32, i32),
    Hsv(i32, i32, i32),
}

enum Message {
    Quit,
    Move { x: i32, y: i32 },
    Write(String),
    ChangeColor(Color),
}

fn main(msg: Message) {
    match msg {
        Message::ChangeColor(Color::Rgb(a, b, c)) => $0"",
        Message::ChangeColor(Color::Hsv(a, b, c)) => "",
        _ => "other"
    };
}
        "#,
        r#"
enum Color {
    Rgb(i32, i32, i32),
    Hsv(i32, i32, i32),
}

enum Message {
    Quit,
    Move { x: i32, y: i32 },
    Write(String),
    ChangeColor(Color),
}

fn main(msg: Message) {
    match msg {
        Message::ChangeColor(Color::Rgb(a, b, c)) | Message::ChangeColor(Color::Hsv(a, b, c)) => "",
        _ => "other"
    };
}
        "#,
        )
    }

    #[test]
    fn merge_match_arms_enum_destructuring_with_ignore() {
        check_assist(
            merge_match_arms,
            r#"
enum MyEnum {
    Move { x: i32, a: i32 },
    Crawl { x: i32, b: i32 }
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, .. } => $0"",
        MyEnum::Crawl { x, .. } => "",
    };
}
        "#,
            r#"
enum MyEnum {
    Move { x: i32, a: i32 },
    Crawl { x: i32, b: i32 }
}

fn func(x: MyEnum) {
    match x {
        MyEnum::Move { x, .. } | MyEnum::Crawl { x, .. } => "",
    };
}
        "#,
        )
    }
}
<|endoftext|>


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

use std::pin::Pin;
use std::sync::Arc;
use std::task::Context;

use futures::lock::Mutex;
use futures::{future, Future, Poll, TryFutureExt};
use smallvec::SmallVec;

use proto::error::ProtoError;
use proto::op::ResponseCode;
use proto::xfer::{DnsHandle, DnsRequest, DnsResponse};

use crate::config::{ResolverConfig, ResolverOpts};
#[cfg(feature = "mdns")]
use crate::name_server;
use crate::name_server::{ConnectionHandle, ConnectionProvider, NameServer, StandardConnection};

/// A pool of NameServers
///
/// This is not expected to be used directly, see [`AsyncResolver`].
#[derive(Clone)]
pub struct NameServerPool<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> {
    // TODO: switch to FuturesMutex (Mutex will have some undesireable locking)
    datagram_conns: Arc<Mutex<Vec<NameServer<C, P>>>>, /* All NameServers must be the same type */
    stream_conns: Arc<Mutex<Vec<NameServer<C, P>>>>,   /* All NameServers must be the same type */
    #[cfg(feature = "mdns")]
    mdns_conns: NameServer<C, P>, /* All NameServers must be the same type */
    options: ResolverOpts,
    conn_provider: P,
}

impl NameServerPool<ConnectionHandle, StandardConnection> {
    pub(crate) fn from_config(config: &ResolverConfig, options: &ResolverOpts) -> Self {
        Self::from_config_with_provider(config, options, StandardConnection)
    }
}

impl<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> NameServerPool<C, P> {
    pub(crate) fn from_config_with_provider(
        config: &ResolverConfig,
        options: &ResolverOpts,
        conn_provider: P,
    ) -> NameServerPool<C, P> {
        let datagram_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_datagram())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        let stream_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_stream())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns)),
            stream_conns: Arc::new(Mutex::new(stream_conns)),
            #[cfg(feature = "mdns")]
            mdns_conns: name_server::mdns_nameserver(*options, conn_provider.clone()),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(not(feature = "mdns"))]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(feature = "mdns")]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        mdns_conns: NameServer<C, P>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            mdns_conns,
            options: *options,
            conn_provider,
        }
    }

    async fn try_send(
        opts: ResolverOpts,
        conns: Arc<Mutex<Vec<NameServer<C, P>>>>,
        request: DnsRequest,
    ) -> Result<DnsResponse, ProtoError> {
        let conns: Vec<NameServer<C, P>> = {
            // pull a lock on the shared connections, lock releases at the end of this scope
            let mut conns = conns.lock().await;

            // select the highest priority connection
            //   reorder the connections based on current view...
            //   this reorders the inner set
            conns.sort_unstable();

            // TODO: restrict this size to a maximum # of NameServers to try
            // get a stable view for trying all connections
            //   we split into chunks of the number of parallel requests to issue
            conns.clone()
         };

        let request_loop = request.clone();

        parallel_conn_loop(conns, request_loop, opts).await
    }
}

impl<C, P> DnsHandle for NameServerPool<C, P>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    type Response = Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send>>;

    fn send<R: Into<DnsRequest>>(&mut self, request: R) -> Self::Response {
        let opts = self.options;
        let request = request.into();
        let datagram_conns = Arc::clone(&self.datagram_conns);
        let stream_conns1 = Arc::clone(&self.stream_conns);
        let stream_conns2 = Arc::clone(&self.stream_conns);
        // TODO: remove this clone, return the Message in the error?
        let tcp_message1 = request.clone();
        let tcp_message2 = request.clone();

        // if it's a .local. query, then we *only* query mDNS, these should never be sent on to upstream resolvers
        #[cfg(feature = "mdns")]
        let mdns = mdns::maybe_local(&mut self.mdns_conns, request);

        // TODO: limited to only when mDNS is enabled, but this should probably always be enforced?
        #[cfg(not(feature = "mdns"))]
        let mdns = Local::NotMdns(request);

        // local queries are queried through mDNS
        if mdns.is_local() {
            return mdns.take_future();
        }

        // TODO: should we allow mDNS to be used for standard lookups as well?

        // it wasn't a local query, continue with standard lookup path
        let request = mdns.take_request();

        // First try the UDP connections
        Box::pin(
            Self::try_send(opts, datagram_conns, request)
                .and_then(move |response| {
                    // handling promotion from datagram to stream base on truncation in message
                    if ResponseCode::NoError == response.response_code() && response.truncated() {
                        // TCP connections should not truncate
                        future::Either::Left(Self::try_send(opts, stream_conns1, tcp_message1))
                    } else {
                        // Return the result from the UDP connection
                        future::Either::Right(future::ok(response))
                    }
                })
                // if UDP fails, try TCP
                .or_else(move |_| Self::try_send(opts, stream_conns2, tcp_message2)),
        )
    }
}

// TODO: we should be able to have a self-referential future here with Pin and not require cloned conns
/// An async function that will loop over all the conns with a max parallel request count of ops.num_concurrent_req
async fn parallel_conn_loop<C, P>(
    mut conns: Vec<NameServer<C, P>>,
    request: DnsRequest,
    opts: ResolverOpts,
) -> Result<DnsResponse, ProtoError>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    let mut err = ProtoError::from("No connections available");

    loop {
        let request_cont = request.clone();

        // construct the parallel requests, 2 is the default
        let mut par_conns = SmallVec::<[NameServer<C, P>; 2]>::new();
        let count = conns.len().min(opts.num_concurrent_reqs.max(1));
        for conn in conns.drain(..count) {
            par_conns.push(conn);
        }

        // construct the requests to send
        let requests = if par_conns.is_empty() {
            return Err(err);
        } else {
            par_conns
                .into_iter()
                .map(move |mut conn| conn.send(request_cont.clone()))
        };

        match future::select_ok(requests).await {
            Ok((sent, _)) => return Ok(sent),
            // consider a debug msg here
            Err(e) => {
                err = e;
                continue;
            }
        };
    }
}

#[cfg(feature = "mdns")]
mod mdns {
    use super::*;

    use proto::rr::domain::usage;
    use proto::DnsHandle;

    /// Returns true
    pub fn maybe_local<C, P>(name_server: &mut NameServer<C, P>, request: DnsRequest) -> Local
    where
        C: DnsHandle + 'static,
        P: ConnectionProvider<ConnHandle = C> + 'static,
    {
        if request
            .queries()
            .iter()
            .any(|query| usage::LOCAL.name().zone_of(query.name()))
        {
            Local::ResolveFuture(name_server.send(request))
        } else {
            Local::NotMdns(request)
        }
    }
}

pub enum Local {
    #[allow(dead_code)]
    ResolveFuture(Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>>),
    NotMdns(DnsRequest),
}

impl Local {
    fn is_local(&self) -> bool {
        if let Local::ResolveFuture(..) = *self {
            true
        } else {
            false
        }
    }

    /// Takes the future
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::NotMdns
    fn take_future(
        self,
    ) -> Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>> {
        match self {
            Local::ResolveFuture(future) => future,
            _ => panic!("non Local queries have no future, see take_message()"),
        }
    }

    /// Takes the message
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::ResolveFuture
    fn take_request(self) -> DnsRequest {
        match self {
            Local::NotMdns(request) => request,
            _ => panic!("Local queries must be polled, see take_future()"),
        }
    }
}

impl Future for Local {
    type Output = Result<DnsResponse, ProtoError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match *self {
            Local::ResolveFuture(ref mut ns) => ns.as_mut().poll(cx),
            // TODO: making this a panic for now
            Local::NotMdns(..) => panic!("Local queries that are not mDNS should not be polled"), //Local::NotMdns(message) => return Err(ResolveErrorKind::Message("not mDNS")),
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;

    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    use tokio::runtime::current_thread::Runtime;

    use proto::op::Query;
    use proto::rr::{Name, RecordType};
    use proto::xfer::{DnsHandle, DnsRequestOptions};

    use super::*;
    use crate::config::NameServerConfig;
    use crate::config::Protocol;

    #[ignore]
    // because of there is a real connection that needs a reasonable timeout
    #[test]
    fn test_failed_then_success_pool() {
        let config1 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 252)), 253),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let config2 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let mut resolver_config = ResolverConfig::new();
        resolver_config.add_name_server(config1);
        resolver_config.add_name_server(config2);

        let mut io_loop = Runtime::new().unwrap();
        let mut pool = NameServerPool::<_, StandardConnection>::from_config(
            &resolver_config,
            &ResolverOpts::default(),
        );

        let name = Name::parse("www.example.com.", None).unwrap();

        // TODO: it's not clear why there are two failures before the success
        for i in 0..2 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_err(),
                "iter: {}",
                i
            );
        }

        for i in 0..10 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_ok(),
                "iter: {}",
                i
            );
        }
    }
}
<commit_msg>Fix minor formatting error<commit_after>// Copyright 2015-2019 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

use std::pin::Pin;
use std::sync::Arc;
use std::task::Context;

use futures::lock::Mutex;
use futures::{future, Future, Poll, TryFutureExt};
use smallvec::SmallVec;

use proto::error::ProtoError;
use proto::op::ResponseCode;
use proto::xfer::{DnsHandle, DnsRequest, DnsResponse};

use crate::config::{ResolverConfig, ResolverOpts};
#[cfg(feature = "mdns")]
use crate::name_server;
use crate::name_server::{ConnectionHandle, ConnectionProvider, NameServer, StandardConnection};

/// A pool of NameServers
///
/// This is not expected to be used directly, see [`AsyncResolver`].
#[derive(Clone)]
pub struct NameServerPool<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> {
    // TODO: switch to FuturesMutex (Mutex will have some undesireable locking)
    datagram_conns: Arc<Mutex<Vec<NameServer<C, P>>>>, /* All NameServers must be the same type */
    stream_conns: Arc<Mutex<Vec<NameServer<C, P>>>>,   /* All NameServers must be the same type */
    #[cfg(feature = "mdns")]
    mdns_conns: NameServer<C, P>, /* All NameServers must be the same type */
    options: ResolverOpts,
    conn_provider: P,
}

impl NameServerPool<ConnectionHandle, StandardConnection> {
    pub(crate) fn from_config(config: &ResolverConfig, options: &ResolverOpts) -> Self {
        Self::from_config_with_provider(config, options, StandardConnection)
    }
}

impl<C: DnsHandle + 'static, P: ConnectionProvider<ConnHandle = C> + 'static> NameServerPool<C, P> {
    pub(crate) fn from_config_with_provider(
        config: &ResolverConfig,
        options: &ResolverOpts,
        conn_provider: P,
    ) -> NameServerPool<C, P> {
        let datagram_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_datagram())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        let stream_conns: Vec<NameServer<C, P>> = config
            .name_servers()
            .iter()
            .filter(|ns_config| ns_config.protocol.is_stream())
            .map(|ns_config| {
                #[cfg(feature = "dns-over-rustls")]
                let ns_config = {
                    let mut ns_config = ns_config.clone();
                    ns_config.tls_config = config.client_config().clone();
                    ns_config
                };
                #[cfg(not(feature = "dns-over-rustls"))]
                let ns_config = { ns_config.clone() };

                NameServer::<C, P>::new_with_provider(ns_config, *options, conn_provider.clone())
            })
            .collect();

        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns)),
            stream_conns: Arc::new(Mutex::new(stream_conns)),
            #[cfg(feature = "mdns")]
            mdns_conns: name_server::mdns_nameserver(*options, conn_provider.clone()),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(not(feature = "mdns"))]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            options: *options,
            conn_provider,
        }
    }

    #[doc(hidden)]
    #[cfg(feature = "mdns")]
    pub fn from_nameservers(
        options: &ResolverOpts,
        datagram_conns: Vec<NameServer<C, P>>,
        stream_conns: Vec<NameServer<C, P>>,
        mdns_conns: NameServer<C, P>,
        conn_provider: P,
    ) -> Self {
        NameServerPool {
            datagram_conns: Arc::new(Mutex::new(datagram_conns.into_iter().collect())),
            stream_conns: Arc::new(Mutex::new(stream_conns.into_iter().collect())),
            mdns_conns,
            options: *options,
            conn_provider,
        }
    }

    async fn try_send(
        opts: ResolverOpts,
        conns: Arc<Mutex<Vec<NameServer<C, P>>>>,
        request: DnsRequest,
    ) -> Result<DnsResponse, ProtoError> {
        let conns: Vec<NameServer<C, P>> = {
            // pull a lock on the shared connections, lock releases at the end of this scope
            let mut conns = conns.lock().await;

            // select the highest priority connection
            //   reorder the connections based on current view...
            //   this reorders the inner set
            conns.sort_unstable();

            // TODO: restrict this size to a maximum # of NameServers to try
            // get a stable view for trying all connections
            //   we split into chunks of the number of parallel requests to issue
            conns.clone()
        };

        let request_loop = request.clone();

        parallel_conn_loop(conns, request_loop, opts).await
    }
}

impl<C, P> DnsHandle for NameServerPool<C, P>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    type Response = Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send>>;

    fn send<R: Into<DnsRequest>>(&mut self, request: R) -> Self::Response {
        let opts = self.options;
        let request = request.into();
        let datagram_conns = Arc::clone(&self.datagram_conns);
        let stream_conns1 = Arc::clone(&self.stream_conns);
        let stream_conns2 = Arc::clone(&self.stream_conns);
        // TODO: remove this clone, return the Message in the error?
        let tcp_message1 = request.clone();
        let tcp_message2 = request.clone();

        // if it's a .local. query, then we *only* query mDNS, these should never be sent on to upstream resolvers
        #[cfg(feature = "mdns")]
        let mdns = mdns::maybe_local(&mut self.mdns_conns, request);

        // TODO: limited to only when mDNS is enabled, but this should probably always be enforced?
        #[cfg(not(feature = "mdns"))]
        let mdns = Local::NotMdns(request);

        // local queries are queried through mDNS
        if mdns.is_local() {
            return mdns.take_future();
        }

        // TODO: should we allow mDNS to be used for standard lookups as well?

        // it wasn't a local query, continue with standard lookup path
        let request = mdns.take_request();

        // First try the UDP connections
        Box::pin(
            Self::try_send(opts, datagram_conns, request)
                .and_then(move |response| {
                    // handling promotion from datagram to stream base on truncation in message
                    if ResponseCode::NoError == response.response_code() && response.truncated() {
                        // TCP connections should not truncate
                        future::Either::Left(Self::try_send(opts, stream_conns1, tcp_message1))
                    } else {
                        // Return the result from the UDP connection
                        future::Either::Right(future::ok(response))
                    }
                })
                // if UDP fails, try TCP
                .or_else(move |_| Self::try_send(opts, stream_conns2, tcp_message2)),
        )
    }
}

// TODO: we should be able to have a self-referential future here with Pin and not require cloned conns
/// An async function that will loop over all the conns with a max parallel request count of ops.num_concurrent_req
async fn parallel_conn_loop<C, P>(
    mut conns: Vec<NameServer<C, P>>,
    request: DnsRequest,
    opts: ResolverOpts,
) -> Result<DnsResponse, ProtoError>
where
    C: DnsHandle + 'static,
    P: ConnectionProvider<ConnHandle = C> + 'static,
{
    let mut err = ProtoError::from("No connections available");

    loop {
        let request_cont = request.clone();

        // construct the parallel requests, 2 is the default
        let mut par_conns = SmallVec::<[NameServer<C, P>; 2]>::new();
        let count = conns.len().min(opts.num_concurrent_reqs.max(1));
        for conn in conns.drain(..count) {
            par_conns.push(conn);
        }

        // construct the requests to send
        let requests = if par_conns.is_empty() {
            return Err(err);
        } else {
            par_conns
                .into_iter()
                .map(move |mut conn| conn.send(request_cont.clone()))
        };

        match future::select_ok(requests).await {
            Ok((sent, _)) => return Ok(sent),
            // consider a debug msg here
            Err(e) => {
                err = e;
                continue;
            }
        };
    }
}

#[cfg(feature = "mdns")]
mod mdns {
    use super::*;

    use proto::rr::domain::usage;
    use proto::DnsHandle;

    /// Returns true
    pub fn maybe_local<C, P>(name_server: &mut NameServer<C, P>, request: DnsRequest) -> Local
    where
        C: DnsHandle + 'static,
        P: ConnectionProvider<ConnHandle = C> + 'static,
    {
        if request
            .queries()
            .iter()
            .any(|query| usage::LOCAL.name().zone_of(query.name()))
        {
            Local::ResolveFuture(name_server.send(request))
        } else {
            Local::NotMdns(request)
        }
    }
}

pub enum Local {
    #[allow(dead_code)]
    ResolveFuture(Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>>),
    NotMdns(DnsRequest),
}

impl Local {
    fn is_local(&self) -> bool {
        if let Local::ResolveFuture(..) = *self {
            true
        } else {
            false
        }
    }

    /// Takes the future
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::NotMdns
    fn take_future(
        self,
    ) -> Pin<Box<dyn Future<Output = Result<DnsResponse, ProtoError>> + Send + Unpin>> {
        match self {
            Local::ResolveFuture(future) => future,
            _ => panic!("non Local queries have no future, see take_message()"),
        }
    }

    /// Takes the message
    ///
    /// # Panics
    ///
    /// Panics if this is in fact a Local::ResolveFuture
    fn take_request(self) -> DnsRequest {
        match self {
            Local::NotMdns(request) => request,
            _ => panic!("Local queries must be polled, see take_future()"),
        }
    }
}

impl Future for Local {
    type Output = Result<DnsResponse, ProtoError>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        match *self {
            Local::ResolveFuture(ref mut ns) => ns.as_mut().poll(cx),
            // TODO: making this a panic for now
            Local::NotMdns(..) => panic!("Local queries that are not mDNS should not be polled"), //Local::NotMdns(message) => return Err(ResolveErrorKind::Message("not mDNS")),
        }
    }
}

#[cfg(test)]
mod tests {
    extern crate env_logger;

    use std::net::{IpAddr, Ipv4Addr, SocketAddr};

    use tokio::runtime::current_thread::Runtime;

    use proto::op::Query;
    use proto::rr::{Name, RecordType};
    use proto::xfer::{DnsHandle, DnsRequestOptions};

    use super::*;
    use crate::config::NameServerConfig;
    use crate::config::Protocol;

    #[ignore]
    // because of there is a real connection that needs a reasonable timeout
    #[test]
    fn test_failed_then_success_pool() {
        let config1 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 252)), 253),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let config2 = NameServerConfig {
            socket_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)), 53),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            #[cfg(feature = "dns-over-rustls")]
            tls_config: None,
        };

        let mut resolver_config = ResolverConfig::new();
        resolver_config.add_name_server(config1);
        resolver_config.add_name_server(config2);

        let mut io_loop = Runtime::new().unwrap();
        let mut pool = NameServerPool::<_, StandardConnection>::from_config(
            &resolver_config,
            &ResolverOpts::default(),
        );

        let name = Name::parse("www.example.com.", None).unwrap();

        // TODO: it's not clear why there are two failures before the success
        for i in 0..2 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_err(),
                "iter: {}",
                i
            );
        }

        for i in 0..10 {
            assert!(
                io_loop
                    .block_on(pool.lookup(
                        Query::query(name.clone(), RecordType::A),
                        DnsRequestOptions::default()
                    ))
                    .is_ok(),
                "iter: {}",
                i
            );
        }
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use chrono::prelude::*;
use serde_json::value::{from_value, to_value, Value};

use crate::errors::{Error, Result};

/// The global function type definition
pub trait Function: Sync + Send {
    /// The global function type definition
    fn call(&self, args: &HashMap<String, Value>) -> Result<Value>;
}

impl<F> Function for F
where
    F: Fn(&HashMap<String, Value>) -> Result<Value> + Sync + Send,
{
    fn call(&self, args: &HashMap<String, Value>) -> Result<Value> {
        self(args)
    }
}

pub fn range(args: &HashMap<String, Value>) -> Result<Value> {
    let start = match args.get("start") {
        Some(val) => match from_value::<usize>(val.clone()) {
            Ok(v) => v,
            Err(_) => {
                return Err(Error::msg(format!(
                    "Global function `range` received start={} but `start` can only be a number",
                    val
                )));
            }
        },
        None => 0,
    };
    let step_by = match args.get("step_by") {
        Some(val) => match from_value::<usize>(val.clone()) {
            Ok(v) => v,
            Err(_) => {
                return Err(Error::msg(format!(
                    "Global function `range` received step_by={} but `step` can only be a number",
                    val
                )));
            }
        },
        None => 1,
    };
    let end = match args.get("end") {
        Some(val) => match from_value::<usize>(val.clone()) {
            Ok(v) => v,
            Err(_) => {
                return Err(Error::msg(format!(
                    "Global function `range` received end={} but `end` can only be a number",
                    val
                )));
            }
        },
        None => {
            return Err(Error::msg("Global function `range` was called without a `end` argument"));
        }
    };

    if start > end {
        return Err(Error::msg("Global function `range` was called without a `start` argument greater than the `end` one"));
    }

    let mut i = start;
    let mut res = vec![];
    while i < end {
        res.push(i);
        i += step_by;
    }
    Ok(to_value(res).unwrap())
}

pub fn now(args: &HashMap<String, Value>) -> Result<Value> {
    let utc = match args.get("utc") {
        Some(val) => match from_value::<bool>(val.clone()) {
            Ok(v) => v,
            Err(_) => {
                return Err(Error::msg(format!(
                    "Global function `now` received utc={} but `utc` can only be a boolean",
                    val
                )));
            }
        },
        None => false,
    };
    let timestamp = match args.get("timestamp") {
        Some(val) => match from_value::<bool>(val.clone()) {
            Ok(v) => v,
            Err(_) => {
                return Err(Error::msg(format!(
                "Global function `now` received timestamp={} but `timestamp` can only be a boolean",
                val
            )));
            }
        },
        None => false,
    };

    if utc {
        let datetime = Utc::now();
        if timestamp {
            return Ok(to_value(datetime.timestamp()).unwrap());
        }
        Ok(to_value(datetime.to_rfc3339()).unwrap())
    } else {
        let datetime = Local::now();
        if timestamp {
            return Ok(to_value(datetime.timestamp()).unwrap());
        }
        Ok(to_value(datetime.to_rfc3339()).unwrap())
    }
}

pub fn throw(args: &HashMap<String, Value>) -> Result<Value> {
    match args.get("message") {
        Some(val) => match from_value::<String>(val.clone()) {
            Ok(v) => Err(Error::msg(v)),
            Err(_) => Err(Error::msg(format!(
                "Global function `throw` received message={} but `message` can only be a string",
                val
            ))),
        },
        None => Err(Error::msg("Global function `throw` was called without a `message` argument")),
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use serde_json::value::to_value;

    use super::*;

    #[test]
    fn range_default() {
        let mut args = HashMap::new();
        args.insert("end".to_string(), to_value(5).unwrap());

        let res = range(&args).unwrap();
        assert_eq!(res, to_value(vec![0, 1, 2, 3, 4]).unwrap());
    }

    #[test]
    fn range_start() {
        let mut args = HashMap::new();
        args.insert("end".to_string(), to_value(5).unwrap());
        args.insert("start".to_string(), to_value(1).unwrap());

        let res = range(&args).unwrap();
        assert_eq!(res, to_value(vec![1, 2, 3, 4]).unwrap());
    }

    #[test]
    fn range_start_greater_than_end() {
        let mut args = HashMap::new();
        args.insert("end".to_string(), to_value(5).unwrap());
        args.insert("start".to_string(), to_value(6).unwrap());

        assert!(range(&args).is_err());
    }

    #[test]
    fn range_step_by() {
        let mut args = HashMap::new();
        args.insert("end".to_string(), to_value(10).unwrap());
        args.insert("step_by".to_string(), to_value(2).unwrap());

        let res = range(&args).unwrap();
        assert_eq!(res, to_value(vec![0, 2, 4, 6, 8]).unwrap());
    }

    #[test]
    fn now_default() {
        let args = HashMap::new();

        let res = now(&args).unwrap();
        assert!(res.is_string());
        assert!(res.as_str().unwrap().contains("T"));
    }

    #[test]
    fn now_datetime_utc() {
        let mut args = HashMap::new();
        args.insert("utc".to_string(), to_value(true).unwrap());

        let res = now(&args).unwrap();
        assert!(res.is_string());
        let val = res.as_str().unwrap();
        println!("{}", val);
        assert!(val.contains("T"));
        assert!(val.contains("+00:00"));
    }

    #[test]
    fn now_timestamp() {
        let mut args = HashMap::new();
        args.insert("timestamp".to_string(), to_value(true).unwrap());

        let res = now(&args).unwrap();
        assert!(res.is_number());
    }

    #[test]
    fn throw_errors_with_message() {
        let mut args = HashMap::new();
        args.insert("message".to_string(), to_value("Hello").unwrap());

        let res = throw(&args);
        assert!(res.is_err());
        let err = res.unwrap_err();
        assert_eq!(err.to_string(), "Hello");
    }
}
<commit_msg>Fix error messages for functions<commit_after>use std::collections::HashMap;

use chrono::prelude::*;
use serde_json::value::{from_value, to_value, Value};

use crate::errors::{Error, Result};

/// The global function type definition
pub trait Function: Sync + Send {
    /// The global function type definition
    fn call(&self, args: &HashMap<String, Value>) -> Result<Value>;
}

impl<F> Function for F
where
    F: Fn(&HashMap<String, Value>) -> Result<Value> + Sync + Send,
{
    fn call(&self, args: &HashMap<String, Value>) -> Result<Value> {
        self(args)
    }
}

pub fn range(args: &HashMap<String, Value>) -> Result<Value> {
    let start = match args.get("start") {
        Some(val) => match from_value::<usize>(val.clone()) {
            Ok(v) => v,
            Err(_) => {
                return Err(Error::msg(format!(
                    "Function `range` received start={} but `start` can only be a number",
                    val
                )));
            }
        },
        None => 0,
    };
    let step_by = match args.get("step_by") {
        Some(val) => match from_value::<usize>(val.clone()) {
            Ok(v) => v,
            Err(_) => {
                return Err(Error::msg(format!(
                    "Function `range` received step_by={} but `step` can only be a number",
                    val
                )));
            }
        },
        None => 1,
    };
    let end = match args.get("end") {
        Some(val) => match from_value::<usize>(val.clone()) {
            Ok(v) => v,
            Err(_) => {
                return Err(Error::msg(format!(
                    "Function `range` received end={} but `end` can only be a number",
                    val
                )));
            }
        },
        None => {
            return Err(Error::msg("Function `range` was called without a `end` argument"));
        }
    };

    if start > end {
        return Err(Error::msg("Function `range` was called without a `start` argument greater than the `end` one"));
    }

    let mut i = start;
    let mut res = vec![];
    while i < end {
        res.push(i);
        i += step_by;
    }
    Ok(to_value(res).unwrap())
}

pub fn now(args: &HashMap<String, Value>) -> Result<Value> {
    let utc = match args.get("utc") {
        Some(val) => match from_value::<bool>(val.clone()) {
            Ok(v) => v,
            Err(_) => {
                return Err(Error::msg(format!(
                    "Function `now` received utc={} but `utc` can only be a boolean",
                    val
                )));
            }
        },
        None => false,
    };
    let timestamp = match args.get("timestamp") {
        Some(val) => match from_value::<bool>(val.clone()) {
            Ok(v) => v,
            Err(_) => {
                return Err(Error::msg(format!(
                "Function `now` received timestamp={} but `timestamp` can only be a boolean",
                val
            )));
            }
        },
        None => false,
    };

    if utc {
        let datetime = Utc::now();
        if timestamp {
            return Ok(to_value(datetime.timestamp()).unwrap());
        }
        Ok(to_value(datetime.to_rfc3339()).unwrap())
    } else {
        let datetime = Local::now();
        if timestamp {
            return Ok(to_value(datetime.timestamp()).unwrap());
        }
        Ok(to_value(datetime.to_rfc3339()).unwrap())
    }
}

pub fn throw(args: &HashMap<String, Value>) -> Result<Value> {
    match args.get("message") {
        Some(val) => match from_value::<String>(val.clone()) {
            Ok(v) => Err(Error::msg(v)),
            Err(_) => Err(Error::msg(format!(
                "Function `throw` received message={} but `message` can only be a string",
                val
            ))),
        },
        None => Err(Error::msg("Function `throw` was called without a `message` argument")),
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use serde_json::value::to_value;

    use super::*;

    #[test]
    fn range_default() {
        let mut args = HashMap::new();
        args.insert("end".to_string(), to_value(5).unwrap());

        let res = range(&args).unwrap();
        assert_eq!(res, to_value(vec![0, 1, 2, 3, 4]).unwrap());
    }

    #[test]
    fn range_start() {
        let mut args = HashMap::new();
        args.insert("end".to_string(), to_value(5).unwrap());
        args.insert("start".to_string(), to_value(1).unwrap());

        let res = range(&args).unwrap();
        assert_eq!(res, to_value(vec![1, 2, 3, 4]).unwrap());
    }

    #[test]
    fn range_start_greater_than_end() {
        let mut args = HashMap::new();
        args.insert("end".to_string(), to_value(5).unwrap());
        args.insert("start".to_string(), to_value(6).unwrap());

        assert!(range(&args).is_err());
    }

    #[test]
    fn range_step_by() {
        let mut args = HashMap::new();
        args.insert("end".to_string(), to_value(10).unwrap());
        args.insert("step_by".to_string(), to_value(2).unwrap());

        let res = range(&args).unwrap();
        assert_eq!(res, to_value(vec![0, 2, 4, 6, 8]).unwrap());
    }

    #[test]
    fn now_default() {
        let args = HashMap::new();

        let res = now(&args).unwrap();
        assert!(res.is_string());
        assert!(res.as_str().unwrap().contains("T"));
    }

    #[test]
    fn now_datetime_utc() {
        let mut args = HashMap::new();
        args.insert("utc".to_string(), to_value(true).unwrap());

        let res = now(&args).unwrap();
        assert!(res.is_string());
        let val = res.as_str().unwrap();
        println!("{}", val);
        assert!(val.contains("T"));
        assert!(val.contains("+00:00"));
    }

    #[test]
    fn now_timestamp() {
        let mut args = HashMap::new();
        args.insert("timestamp".to_string(), to_value(true).unwrap());

        let res = now(&args).unwrap();
        assert!(res.is_number());
    }

    #[test]
    fn throw_errors_with_message() {
        let mut args = HashMap::new();
        args.insert("message".to_string(), to_value("Hello").unwrap());

        let res = throw(&args);
        assert!(res.is_err());
        let err = res.unwrap_err();
        assert_eq!(err.to_string(), "Hello");
    }
}
<|endoftext|>