<commit_before>// Copyright (c) 2020 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
use crate::Config;
use futures::future::try_join_all;
use iml_cmd::{CheckedChildExt, CheckedCommandExt, CmdError};
use std::{
    process::{Output, Stdio},
    str,
    time::{SystemTime, UNIX_EPOCH},
};
use tokio::{fs::canonicalize, io::AsyncWriteExt, process::Command};

pub async fn scp(from: String, to: String) -> Result<(), CmdError> {
    tracing::debug!("transferring file from {} to {}", from, to);

    let path = canonicalize("../vagrant/").await?;

    let mut x = Command::new("scp");
    x.current_dir(path);

    x.arg("-o")
        .arg("StrictHostKeyChecking=no")
        .arg("-o")
        .arg("UserKnownHostsFile=/dev/null")
        .arg("-i")
        .arg("./id_rsa")
        .arg(from)
        .arg(to)
        .checked_status()
        .await?;

    Ok(())
}

pub async fn scp_down_parallel(
    servers: &[&str],
    remote_path: &str,
    to: &str,
) -> Result<(), CmdError> {
    let remote_calls = servers.iter().map(|host| {
        let from = format!("{}:{}", host, remote_path);
        scp(from, to.to_string())
    });

    try_join_all(remote_calls).await?;

    Ok(())
}

pub async fn scp_up_parallel(
    servers: &[&str],
    from: &str,
    to_remote: &str,
) -> Result<(), CmdError> {
    let remote_calls = servers.iter().map(|host| {
        let to = format!("{}:{}", host, to_remote);
        scp(from.into(), to)
    });

    try_join_all(remote_calls).await?;

    Ok(())
}

pub async fn ssh_exec_cmd<'a, 'b>(host: &'a str, cmd: &'b str) -> Result<Command, CmdError> {
    tracing::debug!("Running command {} on {}", cmd, host);
    let path = canonicalize("../vagrant/").await?;

    let mut x = Command::new("ssh");
    x.current_dir(path);

    x.arg("-o")
        .arg("StrictHostKeyChecking=no")
        .arg("-o")
        .arg("UserKnownHostsFile=/dev/null")
        .arg("-i")
        .arg("id_rsa")
        .arg(host)
        .arg(cmd);

    Ok(x)
}

pub async fn ssh_exec<'a, 'b>(host: &'a str, cmd: &'b str) -> Result<(&'a str, Output), CmdError> {
    let mut cmd = ssh_exec_cmd(host, cmd).await?;

    let out = cmd.checked_output().await?;

    Ok((host, out))
}

async fn ssh_exec_parallel<'a, 'b>(
    servers: &[&'a str],
    cmd: &'b str,
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    let remote_calls = servers.iter().map(|host| ssh_exec(host, cmd));

    let output = try_join_all(remote_calls).await?;

    for (host, out) in &output {
        tracing::debug!(
            "ssh output {}: {}",
            host,
            str::from_utf8(&out.stdout).expect("Couldn't read output.")
        );
    }

    Ok(output)
}

pub async fn ssh_script<'a, 'b>(
    host: &'a str,
    script: &'b str,
    args: &[&'b str],
) -> Result<(&'a str, Output), CmdError> {
    let path = canonicalize("../vagrant/").await?;

    let mut script_path = path.clone();
    script_path.push(script);

    let script_content = iml_fs::read_file_to_end(script_path).await?;

    let mut x = Command::new("ssh");
    x.current_dir(path);

    let mut ssh_child = x
        .arg("-i")
        .arg("id_rsa")
        .arg("-o")
        .arg("StrictHostKeyChecking=no")
        .arg("-o")
        .arg("UserKnownHostsFile=/dev/null")
        .arg(host)
        .arg("bash -s")
        .args(args)
        .stdin(Stdio::piped())
        .stdout(Stdio::inherit())
        .spawn()?;

    let ssh_stdin = ssh_child.stdin.as_mut().unwrap();
    ssh_stdin.write_all(&script_content).await?;

    let out = ssh_child.wait_with_checked_output().await?;

    Ok((host, out))
}

async fn ssh_script_parallel<'a, 'b>(
    servers: &'b [&'a str],
    script: &'b str,
    args: &[&'b str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    let remote_calls = servers.iter().map(|host| ssh_script(host, script, args));

    let output = try_join_all(remote_calls).await?;

    for (host, out) in &output {
        tracing::debug!(
            "ssh output {}: {}",
            host,
            str::from_utf8(&out.stdout).expect("Couldn't read output.")
        );
    }

    Ok(output)
}

pub async fn install_ldiskfs_no_iml<'a, 'b>(
    config: &Config,
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(
        &config.storage_server_ips(),
        "scripts/install_ldiskfs_no_iml.sh",
        &[config.lustre_version()],
    )
    .await
}

pub async fn install_zfs_no_iml<'a, 'b>(
    config: &Config,
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(
        &config.storage_server_ips(),
        "scripts/install_zfs_no_iml.sh",
        &[config.lustre_version()],
    )
    .await
}

pub async fn yum_update<'a, 'b>(hosts: &'b [&'a str]) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_exec_parallel(hosts, "yum clean metadata; yum update -y").await
}

pub async fn configure_ntp_for_host_only_if<'a, 'b>(
    hosts: &'b [&'a str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(hosts, "scripts/configure_ntp.sh", &["10.73.10.1"]).await
}

pub async fn configure_ntp_for_adm<'a, 'b>(
    hosts: &'b [&'a str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(hosts, "scripts/configure_ntp.sh", &["adm.local"]).await
}

pub async fn wait_for_ntp_for_host_only_if<'a, 'b>(
    hosts: &'b [&'a str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(hosts, "scripts/wait_for_ntp.sh", &["10.73.10.1"]).await
}

pub async fn wait_for_ntp_for_adm<'a, 'b>(
    hosts: &'b [&'a str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(hosts, "scripts/wait_for_ntp.sh", &["adm.local"]).await
}

pub async fn enable_debug_on_hosts<'a, 'b>(
    hosts: &'b [&'a str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(hosts, "scripts/enable_debug.sh", &[]).await
}

pub async fn create_iml_diagnostics<'a, 'b>(
    hosts: Vec<&'a str>,
    prefix: &'a str,
) -> Result<(), CmdError> {
    let path_buf = canonicalize("../vagrant/").await?;
    let path = path_buf.as_path().to_str().expect("Couldn't get path.");

    tracing::debug!("Creating diagnostics on: {:?}", hosts);
    ssh_script_parallel(&hosts, "scripts/create_iml_diagnostics.sh", &[prefix]).await?;

    let now = SystemTime::now();
    let ts = now.duration_since(UNIX_EPOCH).unwrap().as_millis();

    let report_dir = format!("sosreport_{}", ts);
    let mut mkdir = Command::new("mkdir");

    mkdir
        .current_dir(path)
        .arg(&report_dir)
        .checked_status()
        .await?;

    scp_down_parallel(
        &hosts,
        "/var/tmp/*sosreport*.tar.xz",
        format!("./{}/", &report_dir).as_str(),
    )
    .await?;

    let mut chmod = Command::new("chmod");
    chmod
        .current_dir(path)
        .arg("777")
        .arg("-R")
        .arg(report_dir)
        .checked_status()
        .await
}

pub async fn detect_fs(host: &str) -> Result<(), CmdError> {
    ssh_exec_cmd(host, "iml filesystem detect")
        .await?
        .checked_status()
        .await
}

pub async fn systemd_status(host: &str, service_name: &str) -> Result<Command, CmdError> {
    let cmd = ssh_exec_cmd(host, format!("systemctl status {}", service_name).as_str()).await?;

    Ok(cmd)
}

pub async fn add_servers(host: &str, profile: &str, hosts: Vec<String>) -> Result<(), CmdError> {
    ssh_exec_cmd(
        host,
        &format!("iml server add {} -p {}", hosts.join(","), profile),
    )
    .await?
    .checked_status()
    .await?;

    Ok(())
}
<commit_msg>Update sosreport dire to include the prefix (#2088)<commit_after>// Copyright (c) 2020 DDN. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
use crate::Config;
use futures::future::try_join_all;
use iml_cmd::{CheckedChildExt, CheckedCommandExt, CmdError};
use std::{
    process::{Output, Stdio},
    str,
    time::{SystemTime, UNIX_EPOCH},
};
use tokio::{fs::canonicalize, io::AsyncWriteExt, process::Command};

pub async fn scp(from: String, to: String) -> Result<(), CmdError> {
    tracing::debug!("transferring file from {} to {}", from, to);

    let path = canonicalize("../vagrant/").await?;

    let mut x = Command::new("scp");
    x.current_dir(path);

    x.arg("-o")
        .arg("StrictHostKeyChecking=no")
        .arg("-o")
        .arg("UserKnownHostsFile=/dev/null")
        .arg("-i")
        .arg("./id_rsa")
        .arg(from)
        .arg(to)
        .checked_status()
        .await?;

    Ok(())
}

pub async fn scp_down_parallel(
    servers: &[&str],
    remote_path: &str,
    to: &str,
) -> Result<(), CmdError> {
    let remote_calls = servers.iter().map(|host| {
        let from = format!("{}:{}", host, remote_path);
        scp(from, to.to_string())
    });

    try_join_all(remote_calls).await?;

    Ok(())
}

pub async fn scp_up_parallel(
    servers: &[&str],
    from: &str,
    to_remote: &str,
) -> Result<(), CmdError> {
    let remote_calls = servers.iter().map(|host| {
        let to = format!("{}:{}", host, to_remote);
        scp(from.into(), to)
    });

    try_join_all(remote_calls).await?;

    Ok(())
}

pub async fn ssh_exec_cmd<'a, 'b>(host: &'a str, cmd: &'b str) -> Result<Command, CmdError> {
    tracing::debug!("Running command {} on {}", cmd, host);
    let path = canonicalize("../vagrant/").await?;

    let mut x = Command::new("ssh");
    x.current_dir(path);

    x.arg("-o")
        .arg("StrictHostKeyChecking=no")
        .arg("-o")
        .arg("UserKnownHostsFile=/dev/null")
        .arg("-i")
        .arg("id_rsa")
        .arg(host)
        .arg(cmd);

    Ok(x)
}

pub async fn ssh_exec<'a, 'b>(host: &'a str, cmd: &'b str) -> Result<(&'a str, Output), CmdError> {
    let mut cmd = ssh_exec_cmd(host, cmd).await?;

    let out = cmd.checked_output().await?;

    Ok((host, out))
}

async fn ssh_exec_parallel<'a, 'b>(
    servers: &[&'a str],
    cmd: &'b str,
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    let remote_calls = servers.iter().map(|host| ssh_exec(host, cmd));

    let output = try_join_all(remote_calls).await?;

    for (host, out) in &output {
        tracing::debug!(
            "ssh output {}: {}",
            host,
            str::from_utf8(&out.stdout).expect("Couldn't read output.")
        );
    }

    Ok(output)
}

pub async fn ssh_script<'a, 'b>(
    host: &'a str,
    script: &'b str,
    args: &[&'b str],
) -> Result<(&'a str, Output), CmdError> {
    let path = canonicalize("../vagrant/").await?;

    let mut script_path = path.clone();
    script_path.push(script);

    let script_content = iml_fs::read_file_to_end(script_path).await?;

    let mut x = Command::new("ssh");
    x.current_dir(path);

    let mut ssh_child = x
        .arg("-i")
        .arg("id_rsa")
        .arg("-o")
        .arg("StrictHostKeyChecking=no")
        .arg("-o")
        .arg("UserKnownHostsFile=/dev/null")
        .arg(host)
        .arg("bash -s")
        .args(args)
        .stdin(Stdio::piped())
        .stdout(Stdio::inherit())
        .spawn()?;

    let ssh_stdin = ssh_child.stdin.as_mut().unwrap();
    ssh_stdin.write_all(&script_content).await?;

    let out = ssh_child.wait_with_checked_output().await?;

    Ok((host, out))
}

async fn ssh_script_parallel<'a, 'b>(
    servers: &'b [&'a str],
    script: &'b str,
    args: &[&'b str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    let remote_calls = servers.iter().map(|host| ssh_script(host, script, args));

    let output = try_join_all(remote_calls).await?;

    for (host, out) in &output {
        tracing::debug!(
            "ssh output {}: {}",
            host,
            str::from_utf8(&out.stdout).expect("Couldn't read output.")
        );
    }

    Ok(output)
}

pub async fn install_ldiskfs_no_iml<'a, 'b>(
    config: &Config,
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(
        &config.storage_server_ips(),
        "scripts/install_ldiskfs_no_iml.sh",
        &[config.lustre_version()],
    )
    .await
}

pub async fn install_zfs_no_iml<'a, 'b>(
    config: &Config,
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(
        &config.storage_server_ips(),
        "scripts/install_zfs_no_iml.sh",
        &[config.lustre_version()],
    )
    .await
}

pub async fn yum_update<'a, 'b>(hosts: &'b [&'a str]) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_exec_parallel(hosts, "yum clean metadata; yum update -y").await
}

pub async fn configure_ntp_for_host_only_if<'a, 'b>(
    hosts: &'b [&'a str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(hosts, "scripts/configure_ntp.sh", &["10.73.10.1"]).await
}

pub async fn configure_ntp_for_adm<'a, 'b>(
    hosts: &'b [&'a str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(hosts, "scripts/configure_ntp.sh", &["adm.local"]).await
}

pub async fn wait_for_ntp_for_host_only_if<'a, 'b>(
    hosts: &'b [&'a str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(hosts, "scripts/wait_for_ntp.sh", &["10.73.10.1"]).await
}

pub async fn wait_for_ntp_for_adm<'a, 'b>(
    hosts: &'b [&'a str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(hosts, "scripts/wait_for_ntp.sh", &["adm.local"]).await
}

pub async fn enable_debug_on_hosts<'a, 'b>(
    hosts: &'b [&'a str],
) -> Result<Vec<(&'a str, Output)>, CmdError> {
    ssh_script_parallel(hosts, "scripts/enable_debug.sh", &[]).await
}

pub async fn create_iml_diagnostics<'a, 'b>(
    hosts: Vec<&'a str>,
    prefix: &'a str,
) -> Result<(), CmdError> {
    let path_buf = canonicalize("../vagrant/").await?;
    let path = path_buf.as_path().to_str().expect("Couldn't get path.");

    tracing::debug!("Creating diagnostics on: {:?}", hosts);
    ssh_script_parallel(&hosts, "scripts/create_iml_diagnostics.sh", &[prefix]).await?;

    let now = SystemTime::now();
    let ts = now.duration_since(UNIX_EPOCH).unwrap().as_millis();

    let report_dir = format!("sosreport_{}_{}", prefix, ts);
    let mut mkdir = Command::new("mkdir");

    mkdir
        .current_dir(path)
        .arg(&report_dir)
        .checked_status()
        .await?;

    scp_down_parallel(
        &hosts,
        "/var/tmp/*sosreport*.tar.xz",
        format!("./{}/", &report_dir).as_str(),
    )
    .await?;

    let mut chmod = Command::new("chmod");
    chmod
        .current_dir(path)
        .arg("777")
        .arg("-R")
        .arg(report_dir)
        .checked_status()
        .await
}

pub async fn detect_fs(host: &str) -> Result<(), CmdError> {
    ssh_exec_cmd(host, "iml filesystem detect")
        .await?
        .checked_status()
        .await
}

pub async fn systemd_status(host: &str, service_name: &str) -> Result<Command, CmdError> {
    let cmd = ssh_exec_cmd(host, format!("systemctl status {}", service_name).as_str()).await?;

    Ok(cmd)
}

pub async fn add_servers(host: &str, profile: &str, hosts: Vec<String>) -> Result<(), CmdError> {
    ssh_exec_cmd(
        host,
        &format!("iml server add {} -p {}", hosts.join(","), profile),
    )
    .await?
    .checked_status()
    .await?;

    Ok(())
}
<|endoftext|>


<commit_before>// This file is part of the uutils coreutils package.
//
// (c) Jian Zeng <anonymousknight96@gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.
//
//! Aims to provide platform-independent methods to obtain login records
//!
//! **ONLY** support linux, macos and freebsd for the time being
//!
//! # Examples:
//!
//! ```
//! use uucore::utmpx::Utmpx;
//! for ut in Utmpx::iter_all_records() {
//!     if ut.is_user_process() {
//!         println!("{}: {}", ut.host(), ut.user())
//!     }
//! }
//! ```
//!
//! Specifying the path to login record:
//!
//! ```
//! use uucore::utmpx::Utmpx;
//! for ut in Utmpx::iter_all_records_from("/some/where/else") {
//!     if ut.is_user_process() {
//!         println!("{}: {}", ut.host(), ut.user())
//!     }
//! }
//! ```

pub extern crate time;
use self::time::{Timespec, Tm};

use std::ffi::CString;
use std::io::Result as IOResult;
use std::marker::PhantomData;
use std::os::unix::ffi::OsStrExt;
use std::path::Path;
use std::ptr;
use std::sync::{Mutex, MutexGuard};

pub use self::ut::*;
use libc::utmpx;
// pub use libc::getutxid;
// pub use libc::getutxline;
// pub use libc::pututxline;
pub use libc::endutxent;
pub use libc::getutxent;
pub use libc::setutxent;
#[cfg(any(target_vendor = "apple", target_os = "linux", target_os = "netbsd"))]
pub use libc::utmpxname;
#[cfg(target_os = "freebsd")]
pub unsafe extern "C" fn utmpxname(_file: *const libc::c_char) -> libc::c_int {
    0
}

use once_cell::sync::Lazy;

use crate::*; // import macros from `../../macros.rs`

// In case the c_char array doesn't end with NULL
macro_rules! chars2string {
    ($arr:expr) => {
        $arr.iter()
            .take_while(|i| **i > 0)
            .map(|&i| i as u8 as char)
            .collect::<String>()
    };
}

#[cfg(target_os = "linux")]
mod ut {
    pub static DEFAULT_FILE: &str = "/var/run/utmp";

    pub use libc::__UT_HOSTSIZE as UT_HOSTSIZE;
    pub use libc::__UT_LINESIZE as UT_LINESIZE;
    pub use libc::__UT_NAMESIZE as UT_NAMESIZE;
    pub const UT_IDSIZE: usize = 4;

    pub use libc::ACCOUNTING;
    pub use libc::BOOT_TIME;
    pub use libc::DEAD_PROCESS;
    pub use libc::EMPTY;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::RUN_LVL;
    pub use libc::USER_PROCESS;
}

#[cfg(target_vendor = "apple")]
mod ut {
    pub static DEFAULT_FILE: &str = "/var/run/utmpx";

    pub use libc::_UTX_HOSTSIZE as UT_HOSTSIZE;
    pub use libc::_UTX_IDSIZE as UT_IDSIZE;
    pub use libc::_UTX_LINESIZE as UT_LINESIZE;
    pub use libc::_UTX_USERSIZE as UT_NAMESIZE;

    pub use libc::ACCOUNTING;
    pub use libc::BOOT_TIME;
    pub use libc::DEAD_PROCESS;
    pub use libc::EMPTY;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::RUN_LVL;
    pub use libc::SHUTDOWN_TIME;
    pub use libc::SIGNATURE;
    pub use libc::USER_PROCESS;
}

#[cfg(target_os = "freebsd")]
mod ut {
    pub static DEFAULT_FILE: &str = "";

    pub const UT_LINESIZE: usize = 16;
    pub const UT_NAMESIZE: usize = 32;
    pub const UT_IDSIZE: usize = 8;
    pub const UT_HOSTSIZE: usize = 128;

    pub use libc::BOOT_TIME;
    pub use libc::DEAD_PROCESS;
    pub use libc::EMPTY;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::SHUTDOWN_TIME;
    pub use libc::USER_PROCESS;
}

#[cfg(target_os = "netbsd")]
mod ut {
    pub static DEFAULT_FILE: &str = "/var/run/utmpx";

    pub const ACCOUNTING: usize = 9;
    pub const SHUTDOWN_TIME: usize = 11;

    pub use libc::_UTX_HOSTSIZE as UT_HOSTSIZE;
    pub use libc::_UTX_IDSIZE as UT_IDSIZE;
    pub use libc::_UTX_LINESIZE as UT_LINESIZE;
    pub use libc::_UTX_USERSIZE as UT_NAMESIZE;

    pub use libc::ACCOUNTING;
    pub use libc::DEAD_PROCESS;
    pub use libc::EMPTY;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::RUN_LVL;
    pub use libc::SIGNATURE;
    pub use libc::USER_PROCESS;
}

pub struct Utmpx {
    inner: utmpx,
}

impl Utmpx {
    /// A.K.A. ut.ut_type
    pub fn record_type(&self) -> i16 {
        self.inner.ut_type as i16
    }
    /// A.K.A. ut.ut_pid
    pub fn pid(&self) -> i32 {
        self.inner.ut_pid as i32
    }
    /// A.K.A. ut.ut_id
    pub fn terminal_suffix(&self) -> String {
        chars2string!(self.inner.ut_id)
    }
    /// A.K.A. ut.ut_user
    pub fn user(&self) -> String {
        chars2string!(self.inner.ut_user)
    }
    /// A.K.A. ut.ut_host
    pub fn host(&self) -> String {
        chars2string!(self.inner.ut_host)
    }
    /// A.K.A. ut.ut_line
    pub fn tty_device(&self) -> String {
        chars2string!(self.inner.ut_line)
    }
    /// A.K.A. ut.ut_tv
    pub fn login_time(&self) -> Tm {
        time::at(Timespec::new(
            self.inner.ut_tv.tv_sec as i64,
            self.inner.ut_tv.tv_usec as i32,
        ))
    }
    /// A.K.A. ut.ut_exit
    ///
    /// Return (e_termination, e_exit)
    #[cfg(any(target_os = "linux", target_os = "android"))]
    pub fn exit_status(&self) -> (i16, i16) {
        (self.inner.ut_exit.e_termination, self.inner.ut_exit.e_exit)
    }
    /// A.K.A. ut.ut_exit
    ///
    /// Return (0, 0) on Non-Linux platform
    #[cfg(not(any(target_os = "linux", target_os = "android")))]
    pub fn exit_status(&self) -> (i16, i16) {
        (0, 0)
    }
    /// Consumes the `Utmpx`, returning the underlying C struct utmpx
    pub fn into_inner(self) -> utmpx {
        self.inner
    }
    pub fn is_user_process(&self) -> bool {
        !self.user().is_empty() && self.record_type() == USER_PROCESS
    }

    /// Canonicalize host name using DNS
    pub fn canon_host(&self) -> IOResult<String> {
        let host = self.host();

        // TODO: change to use `split_once` when MSRV hits 1.52.0
        // let (hostname, display) = host.split_once(':').unwrap_or((&host, ""));
        let mut h = host.split(':');
        let hostname = h.next().unwrap_or(&host);
        let display = h.next().unwrap_or("");

        if !hostname.is_empty() {
            extern crate dns_lookup;
            use dns_lookup::{getaddrinfo, AddrInfoHints};

            const AI_CANONNAME: i32 = 0x2;
            let hints = AddrInfoHints {
                flags: AI_CANONNAME,
                ..AddrInfoHints::default()
            };
            let sockets = getaddrinfo(Some(hostname), None, Some(hints))
                .unwrap()
                .collect::<IOResult<Vec<_>>>()?;
            for socket in sockets {
                if let Some(ai_canonname) = socket.canonname {
                    return Ok(if display.is_empty() {
                        ai_canonname
                    } else {
                        format!("{}:{}", ai_canonname, display)
                    });
                }
            }
        }

        Ok(host.to_string())
    }

    /// Iterate through all the utmp records.
    ///
    /// This will use the default location, or the path [`Utmpx::iter_all_records_from`]
    /// was most recently called with.
    ///
    /// Only one instance of [`UtmpxIter`] may be active at a time. This
    /// function will block as long as one is still active. Beware!
    pub fn iter_all_records() -> UtmpxIter {
        let iter = UtmpxIter::new();
        unsafe {
            // This can technically fail, and it would be nice to detect that,
            // but it doesn't return anything so we'd have to do nasty things
            // with errno.
            setutxent();
        }
        iter
    }

    /// Iterate through all the utmp records from a specific file.
    ///
    /// No failure is reported or detected.
    ///
    /// This function affects subsequent calls to [`Utmpx::iter_all_records`].
    ///
    /// The same caveats as for [`Utmpx::iter_all_records`] apply.
    pub fn iter_all_records_from<P: AsRef<Path>>(path: P) -> UtmpxIter {
        let iter = UtmpxIter::new();
        let path = CString::new(path.as_ref().as_os_str().as_bytes()).unwrap();
        unsafe {
            // In glibc, utmpxname() only fails if there's not enough memory
            // to copy the string.
            // Solaris returns 1 on success instead of 0. Supposedly there also
            // exist systems where it returns void.
            // GNU who on Debian seems to output nothing if an invalid filename
            // is specified, no warning or anything.
            // So this function is pretty crazy and we don't try to detect errors.
            // Not much we can do besides pray.
            utmpxname(path.as_ptr());
            setutxent();
        }
        iter
    }
}

// On some systems these functions are not thread-safe. On others they're
// thread-local. Therefore we use a mutex to allow only one guard to exist at
// a time, and make sure UtmpxIter cannot be sent across threads.
//
// I believe the only technical memory unsafety that could happen is a data
// race while copying the data out of the pointer returned by getutxent(), but
// ordinary race conditions are also very much possible.
static LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));

/// Iterator of login records
pub struct UtmpxIter {
    #[allow(dead_code)]
    guard: MutexGuard<'static, ()>,
    /// Ensure UtmpxIter is !Send. Technically redundant because MutexGuard
    /// is also !Send.
    phantom: PhantomData<std::rc::Rc<()>>,
}

impl UtmpxIter {
    fn new() -> Self {
        // PoisonErrors can safely be ignored
        let guard = LOCK.lock().unwrap_or_else(|err| err.into_inner());
        UtmpxIter {
            guard,
            phantom: PhantomData,
        }
    }
}

impl Iterator for UtmpxIter {
    type Item = Utmpx;
    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            let res = getutxent();
            if !res.is_null() {
                // The data behind this pointer will be replaced by the next
                // call to getutxent(), so we have to read it now.
                // All the strings live inline in the struct as arrays, which
                // makes things easier.
                Some(Utmpx {
                    inner: ptr::read(res as *const _),
                })
            } else {
                None
            }
        }
    }
}

impl Drop for UtmpxIter {
    fn drop(&mut self) {
        unsafe {
            endutxent();
        }
    }
}
<commit_msg>uucore/utmpx: remove unwrap in cannon_host<commit_after>// This file is part of the uutils coreutils package.
//
// (c) Jian Zeng <anonymousknight96@gmail.com>
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.
//
//! Aims to provide platform-independent methods to obtain login records
//!
//! **ONLY** support linux, macos and freebsd for the time being
//!
//! # Examples:
//!
//! ```
//! use uucore::utmpx::Utmpx;
//! for ut in Utmpx::iter_all_records() {
//!     if ut.is_user_process() {
//!         println!("{}: {}", ut.host(), ut.user())
//!     }
//! }
//! ```
//!
//! Specifying the path to login record:
//!
//! ```
//! use uucore::utmpx::Utmpx;
//! for ut in Utmpx::iter_all_records_from("/some/where/else") {
//!     if ut.is_user_process() {
//!         println!("{}: {}", ut.host(), ut.user())
//!     }
//! }
//! ```

pub extern crate time;
use self::time::{Timespec, Tm};

use std::ffi::CString;
use std::io::Result as IOResult;
use std::marker::PhantomData;
use std::os::unix::ffi::OsStrExt;
use std::path::Path;
use std::ptr;
use std::sync::{Mutex, MutexGuard};

pub use self::ut::*;
use libc::utmpx;
// pub use libc::getutxid;
// pub use libc::getutxline;
// pub use libc::pututxline;
pub use libc::endutxent;
pub use libc::getutxent;
pub use libc::setutxent;
#[cfg(any(target_vendor = "apple", target_os = "linux", target_os = "netbsd"))]
pub use libc::utmpxname;
#[cfg(target_os = "freebsd")]
pub unsafe extern "C" fn utmpxname(_file: *const libc::c_char) -> libc::c_int {
    0
}

use once_cell::sync::Lazy;

use crate::*; // import macros from `../../macros.rs`

// In case the c_char array doesn't end with NULL
macro_rules! chars2string {
    ($arr:expr) => {
        $arr.iter()
            .take_while(|i| **i > 0)
            .map(|&i| i as u8 as char)
            .collect::<String>()
    };
}

#[cfg(target_os = "linux")]
mod ut {
    pub static DEFAULT_FILE: &str = "/var/run/utmp";

    pub use libc::__UT_HOSTSIZE as UT_HOSTSIZE;
    pub use libc::__UT_LINESIZE as UT_LINESIZE;
    pub use libc::__UT_NAMESIZE as UT_NAMESIZE;
    pub const UT_IDSIZE: usize = 4;

    pub use libc::ACCOUNTING;
    pub use libc::BOOT_TIME;
    pub use libc::DEAD_PROCESS;
    pub use libc::EMPTY;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::RUN_LVL;
    pub use libc::USER_PROCESS;
}

#[cfg(target_vendor = "apple")]
mod ut {
    pub static DEFAULT_FILE: &str = "/var/run/utmpx";

    pub use libc::_UTX_HOSTSIZE as UT_HOSTSIZE;
    pub use libc::_UTX_IDSIZE as UT_IDSIZE;
    pub use libc::_UTX_LINESIZE as UT_LINESIZE;
    pub use libc::_UTX_USERSIZE as UT_NAMESIZE;

    pub use libc::ACCOUNTING;
    pub use libc::BOOT_TIME;
    pub use libc::DEAD_PROCESS;
    pub use libc::EMPTY;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::RUN_LVL;
    pub use libc::SHUTDOWN_TIME;
    pub use libc::SIGNATURE;
    pub use libc::USER_PROCESS;
}

#[cfg(target_os = "freebsd")]
mod ut {
    pub static DEFAULT_FILE: &str = "";

    pub const UT_LINESIZE: usize = 16;
    pub const UT_NAMESIZE: usize = 32;
    pub const UT_IDSIZE: usize = 8;
    pub const UT_HOSTSIZE: usize = 128;

    pub use libc::BOOT_TIME;
    pub use libc::DEAD_PROCESS;
    pub use libc::EMPTY;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::SHUTDOWN_TIME;
    pub use libc::USER_PROCESS;
}

#[cfg(target_os = "netbsd")]
mod ut {
    pub static DEFAULT_FILE: &str = "/var/run/utmpx";

    pub const ACCOUNTING: usize = 9;
    pub const SHUTDOWN_TIME: usize = 11;

    pub use libc::_UTX_HOSTSIZE as UT_HOSTSIZE;
    pub use libc::_UTX_IDSIZE as UT_IDSIZE;
    pub use libc::_UTX_LINESIZE as UT_LINESIZE;
    pub use libc::_UTX_USERSIZE as UT_NAMESIZE;

    pub use libc::ACCOUNTING;
    pub use libc::DEAD_PROCESS;
    pub use libc::EMPTY;
    pub use libc::INIT_PROCESS;
    pub use libc::LOGIN_PROCESS;
    pub use libc::NEW_TIME;
    pub use libc::OLD_TIME;
    pub use libc::RUN_LVL;
    pub use libc::SIGNATURE;
    pub use libc::USER_PROCESS;
}

pub struct Utmpx {
    inner: utmpx,
}

impl Utmpx {
    /// A.K.A. ut.ut_type
    pub fn record_type(&self) -> i16 {
        self.inner.ut_type as i16
    }
    /// A.K.A. ut.ut_pid
    pub fn pid(&self) -> i32 {
        self.inner.ut_pid as i32
    }
    /// A.K.A. ut.ut_id
    pub fn terminal_suffix(&self) -> String {
        chars2string!(self.inner.ut_id)
    }
    /// A.K.A. ut.ut_user
    pub fn user(&self) -> String {
        chars2string!(self.inner.ut_user)
    }
    /// A.K.A. ut.ut_host
    pub fn host(&self) -> String {
        chars2string!(self.inner.ut_host)
    }
    /// A.K.A. ut.ut_line
    pub fn tty_device(&self) -> String {
        chars2string!(self.inner.ut_line)
    }
    /// A.K.A. ut.ut_tv
    pub fn login_time(&self) -> Tm {
        time::at(Timespec::new(
            self.inner.ut_tv.tv_sec as i64,
            self.inner.ut_tv.tv_usec as i32,
        ))
    }
    /// A.K.A. ut.ut_exit
    ///
    /// Return (e_termination, e_exit)
    #[cfg(any(target_os = "linux", target_os = "android"))]
    pub fn exit_status(&self) -> (i16, i16) {
        (self.inner.ut_exit.e_termination, self.inner.ut_exit.e_exit)
    }
    /// A.K.A. ut.ut_exit
    ///
    /// Return (0, 0) on Non-Linux platform
    #[cfg(not(any(target_os = "linux", target_os = "android")))]
    pub fn exit_status(&self) -> (i16, i16) {
        (0, 0)
    }
    /// Consumes the `Utmpx`, returning the underlying C struct utmpx
    pub fn into_inner(self) -> utmpx {
        self.inner
    }
    pub fn is_user_process(&self) -> bool {
        !self.user().is_empty() && self.record_type() == USER_PROCESS
    }

    /// Canonicalize host name using DNS
    pub fn canon_host(&self) -> IOResult<String> {
        let host = self.host();

        // TODO: change to use `split_once` when MSRV hits 1.52.0
        // let (hostname, display) = host.split_once(':').unwrap_or((&host, ""));
        let mut h = host.split(':');
        let hostname = h.next().unwrap_or(&host);
        let display = h.next().unwrap_or("");

        if !hostname.is_empty() {
            extern crate dns_lookup;
            use dns_lookup::{getaddrinfo, AddrInfoHints};

            const AI_CANONNAME: i32 = 0x2;
            let hints = AddrInfoHints {
                flags: AI_CANONNAME,
                ..AddrInfoHints::default()
            };
            if let Ok(sockets) = getaddrinfo(Some(hostname), None, Some(hints)) {
                let sockets = sockets.collect::<IOResult<Vec<_>>>()?;
                for socket in sockets {
                    if let Some(ai_canonname) = socket.canonname {
                        return Ok(if display.is_empty() {
                            ai_canonname
                        } else {
                            format!("{}:{}", ai_canonname, display)
                        });
                    }
                }
            } else {
                // GNU coreutils has this behavior
                return Ok(hostname.to_string());
            }
        }

        Ok(host.to_string())
    }

    /// Iterate through all the utmp records.
    ///
    /// This will use the default location, or the path [`Utmpx::iter_all_records_from`]
    /// was most recently called with.
    ///
    /// Only one instance of [`UtmpxIter`] may be active at a time. This
    /// function will block as long as one is still active. Beware!
    pub fn iter_all_records() -> UtmpxIter {
        let iter = UtmpxIter::new();
        unsafe {
            // This can technically fail, and it would be nice to detect that,
            // but it doesn't return anything so we'd have to do nasty things
            // with errno.
            setutxent();
        }
        iter
    }

    /// Iterate through all the utmp records from a specific file.
    ///
    /// No failure is reported or detected.
    ///
    /// This function affects subsequent calls to [`Utmpx::iter_all_records`].
    ///
    /// The same caveats as for [`Utmpx::iter_all_records`] apply.
    pub fn iter_all_records_from<P: AsRef<Path>>(path: P) -> UtmpxIter {
        let iter = UtmpxIter::new();
        let path = CString::new(path.as_ref().as_os_str().as_bytes()).unwrap();
        unsafe {
            // In glibc, utmpxname() only fails if there's not enough memory
            // to copy the string.
            // Solaris returns 1 on success instead of 0. Supposedly there also
            // exist systems where it returns void.
            // GNU who on Debian seems to output nothing if an invalid filename
            // is specified, no warning or anything.
            // So this function is pretty crazy and we don't try to detect errors.
            // Not much we can do besides pray.
            utmpxname(path.as_ptr());
            setutxent();
        }
        iter
    }
}

// On some systems these functions are not thread-safe. On others they're
// thread-local. Therefore we use a mutex to allow only one guard to exist at
// a time, and make sure UtmpxIter cannot be sent across threads.
//
// I believe the only technical memory unsafety that could happen is a data
// race while copying the data out of the pointer returned by getutxent(), but
// ordinary race conditions are also very much possible.
static LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));

/// Iterator of login records
pub struct UtmpxIter {
    #[allow(dead_code)]
    guard: MutexGuard<'static, ()>,
    /// Ensure UtmpxIter is !Send. Technically redundant because MutexGuard
    /// is also !Send.
    phantom: PhantomData<std::rc::Rc<()>>,
}

impl UtmpxIter {
    fn new() -> Self {
        // PoisonErrors can safely be ignored
        let guard = LOCK.lock().unwrap_or_else(|err| err.into_inner());
        UtmpxIter {
            guard,
            phantom: PhantomData,
        }
    }
}

impl Iterator for UtmpxIter {
    type Item = Utmpx;
    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            let res = getutxent();
            if !res.is_null() {
                // The data behind this pointer will be replaced by the next
                // call to getutxent(), so we have to read it now.
                // All the strings live inline in the struct as arrays, which
                // makes things easier.
                Some(Utmpx {
                    inner: ptr::read(res as *const _),
                })
            } else {
                None
            }
        }
    }
}

impl Drop for UtmpxIter {
    fn drop(&mut self) {
        unsafe {
            endutxent();
        }
    }
}
<|endoftext|>


<commit_before>//! Aggregates the weights of equal records into at most one record.
//!
//! As differential dataflow streams are unordered and taken to be the accumulation of all records,
//! no semantic change happens via `consolidate`. However, there is a practical difference between
//! a collection that aggregates down to zero records, and one that actually has no records. The
//! underlying system can more clearly see that no work must be done in the later case, and we can
//! drop out of, e.g. iterative computations.

use std::rc::Rc;
use std::fmt::Debug;

use timely::Data;
use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::dataflow::channels::pact::Exchange;
use timely::dataflow::channels::Content;

// use sort::*;
// use sort::radix_merge::{Accumulator, Compact};
use collection_trace::Lookup;
use collection_trace::lookup::UnsignedInt;
use iterators::coalesce::Coalesce;
use radix_sort::RadixSorter;

use timely::drain::DrainExt;

/// An extension method for consolidating weighted streams.
pub trait ConsolidateExt<D> {
    /// Aggregates the weights of equal records into at most one record.
    fn consolidate<U: UnsignedInt, F: Fn(&D)->U+'static>(&self, part: F) -> Self;
}

impl<G: Scope, D: Ord+Data+Debug> ConsolidateExt<D> for Stream<G, (D, i32)> {
    fn consolidate<U: UnsignedInt, F: Fn(&D)->U+'static>(&self, part: F) -> Self {

        let mut inputs = Vec::new();    // Vec<(G::Timestamp, Vec<(D, i32))>
        let part1 = Rc::new(part);
        let part2 = part1.clone();

        let exch = Exchange::new(move |&(ref x,_)| (*part1)(x).as_u64());
        self.unary_notify(exch, "Consolidate", vec![], move |input, output, notificator| {

            input.for_each(|index: &G::Timestamp, data: &mut Content<(D, i32)>| {
                notificator.notify_at(&index);
                inputs.entry_or_insert(index.clone(), || RadixSorter::new())
                      .extend(data.drain_temp(), &|x| (*part2)(&x.0).as_u64());
            });

            // 2. go through each time of interest that has reached completion
            // while let Some((index, _count)) = notificator.next() {
            notificator.for_each(|index, _count| {
                if let Some(mut stash) = inputs.remove_key(&index) {

                    let mut session = output.session(&index);
                    let mut buffer = vec![];
                    let mut current = 0;

                    let source = stash.finish(&|x| (*part2)(&x.0).as_u64());
                    for (datum, wgt) in source.into_iter().flat_map(|x| x.into_iter()) {
                        let hash = (*part2)(&datum).as_u64();
                        if buffer.len() > 0 && hash != current {
                            // if hash < current { println!("  radix sort error? {} < {}", hash, current); }
                            buffer.sort_by(|x: &(D,i32),y: &(D,i32)| x.0.cmp(&y.0));
                            session.give_iterator(buffer.drain_temp().coalesce());
                            current = hash;
                        }
                        buffer.push((datum,wgt));
                    }

                    if buffer.len() > 0 {
                        buffer.sort_by(|x: &(D,i32),y: &(D,i32)| x.0.cmp(&y.0));
                        session.give_iterator(buffer.drain_temp().coalesce());
                    }
                }
            });
        })
    }
}
<commit_msg>documentation updates<commit_after>//! Aggregates the weights of equal records into at most one record.
//!
//! As differential dataflow streams are unordered and taken to be the accumulation of all records,
//! no semantic change happens via `consolidate`. However, there is a practical difference between
//! a collection that aggregates down to zero records, and one that actually has no records. The
//! underlying system can more clearly see that no work must be done in the later case, and we can
//! drop out of, e.g. iterative computations.

use std::rc::Rc;
use std::fmt::Debug;

use timely::Data;
use timely::dataflow::*;
use timely::dataflow::operators::*;
use timely::dataflow::channels::pact::Exchange;
use timely::dataflow::channels::Content;

// use sort::*;
// use sort::radix_merge::{Accumulator, Compact};
use collection_trace::Lookup;
use collection_trace::lookup::UnsignedInt;
use iterators::coalesce::Coalesce;
use radix_sort::RadixSorter;

use timely::drain::DrainExt;

/// An extension method for consolidating weighted streams.
pub trait ConsolidateExt<D> {
    /// Aggregates the weights of equal records into at most one record.
    fn consolidate<U: UnsignedInt, F: Fn(&D)->U+'static>(&self, part: F) -> Self;
}

impl<G: Scope, D: Ord+Data+Debug> ConsolidateExt<D> for Stream<G, (D, i32)> {
    fn consolidate<U: UnsignedInt, F: Fn(&D)->U+'static>(&self, part: F) -> Self {

        let mut inputs = Vec::new();    // Vec<(G::Timestamp, Vec<(D, i32))>
        let part1 = Rc::new(part);
        let part2 = part1.clone();

        let exch = Exchange::new(move |&(ref x,_)| (*part1)(x).as_u64());
        self.unary_notify(exch, "Consolidate", vec![], move |input, output, notificator| {

            // input.for_each(|index: &G::Timestamp, data: &mut Content<(D, i32)>| {
            while let Some((index, data)) = input.next() {
                notificator.notify_at(&index);
                inputs.entry_or_insert(index.clone(), || RadixSorter::new())
                      .extend(data.drain_temp(), &|x| (*part2)(&x.0).as_u64());
            }
            // });

            // 2. go through each time of interest that has reached completion
            while let Some((index, _count)) = notificator.next() {
            // notificator.for_each(|index, _count| {
                if let Some(mut stash) = inputs.remove_key(&index) {

                    let mut session = output.session(&index);
                    let mut buffer = vec![];
                    let mut current = 0;

                    let source = stash.finish(&|x| (*part2)(&x.0).as_u64());
                    for (datum, wgt) in source.into_iter().flat_map(|x| x.into_iter()) {
                        let hash = (*part2)(&datum).as_u64();
                        if buffer.len() > 0 && hash != current {
                            // if hash < current { println!("  radix sort error? {} < {}", hash, current); }
                            buffer.sort_by(|x: &(D,i32),y: &(D,i32)| x.0.cmp(&y.0));
                            session.give_iterator(buffer.drain_temp().coalesce());
                            current = hash;
                        }
                        buffer.push((datum,wgt));
                    }

                    if buffer.len() > 0 {
                        buffer.sort_by(|x: &(D,i32),y: &(D,i32)| x.0.cmp(&y.0));
                        session.give_iterator(buffer.drain_temp().coalesce());
                    }
                }
            }
            // });
        })
    }
}
<|endoftext|>


<commit_before>use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
use humantime::parse_duration;
use std::time::Duration;
use time;

use super::error::*;
use failure::ResultExt;

#[derive(Default, Debug)]
pub struct CommandParser {
    on_date: Option<String>,
    at_time: Option<String>,
    in_duration: Option<String>,
    every_time: Option<String>,
    target: String,
    message: Option<String>,
}

#[derive(PartialEq, Clone, Copy)]
enum ParseState {
    None,
    On,
    At,
    In,
    Every,
    Msg,
}

impl CommandParser {
    pub fn parse_target(mut tokens: Vec<String>) -> Result<Self, RemindError> {
        let mut parser = CommandParser::default();

        if tokens.is_empty() {
            Err(ErrorKind::MissingReceiver)?;
        }

        parser.target = tokens.remove(0);

        parser.parse_tokens(tokens)
    }

    pub fn with_target(tokens: Vec<String>, target: String) -> Result<Self, RemindError> {
        let mut parser = CommandParser::default();
        parser.target = target;

        parser.parse_tokens(tokens)
    }

    fn parse_tokens(mut self, tokens: Vec<String>) -> Result<Self, RemindError> {
        let mut state = ParseState::None;
        let mut cur_str = String::new();

        for token in tokens {
            let next_state = match token.as_ref() {
                "on" => ParseState::On,
                "at" => ParseState::At,
                "in" => ParseState::In,
                "every" => ParseState::Every,
                "to" => ParseState::Msg,
                _ => {
                    if !cur_str.is_empty() {
                        cur_str.push(' ');
                    }
                    cur_str.push_str(&token);
                    state
                }
            };

            if next_state != state {
                if state != ParseState::None {
                    self = self.add_string_by_state(state, cur_str)?;
                    cur_str = String::new();
                }

                state = next_state;
            }
        }

        self = self.add_string_by_state(state, cur_str)?;

        if self.message.is_none() {
            return Err(ErrorKind::MissingMessage.into());
        }

        if self.in_duration.is_some() && self.at_time.is_some()
            || self.in_duration.is_some() && self.on_date.is_some()
        {
            return Err(ErrorKind::AmbiguousTime.into());
        }

        if self.in_duration.is_none() && self.at_time.is_none() && self.on_date.is_none() {
            return Err(ErrorKind::MissingTime.into());
        }

        Ok(self)
    }

    fn add_string_by_state(self, state: ParseState, string: String) -> Result<Self, RemindError> {
        use self::ParseState::*;
        let string = Some(string);
        match state {
            On if self.on_date.is_none() => Ok(CommandParser {
                on_date: string,
                ..self
            }),
            At if self.at_time.is_none() => Ok(CommandParser {
                at_time: string,
                ..self
            }),
            In if self.in_duration.is_none() => Ok(CommandParser {
                in_duration: string,
                ..self
            }),
            Msg if self.message.is_none() => Ok(CommandParser {
                message: string,
                ..self
            }),
            Every if self.every_time.is_none() => Ok(CommandParser {
                every_time: string,
                ..self
            }),
            _ => Err(ErrorKind::MissingMessage.into()),
        }
    }

    fn parse_date(&self, str_date: &str) -> Result<NaiveDate, RemindError> {
        let nums = str_date
            .split('.')
            .map(|s| s.parse::<u32>())
            .collect::<Result<Vec<_>, _>>()
            .context(ErrorKind::InvalidDate)?;

        if 2 > nums.len() || nums.len() > 3 {
            return Err(ErrorKind::InvalidDate.into());
        }

        let day = nums[0];
        let month = nums[1];

        let parse_date = match nums.get(2) {
            Some(year) => {
                NaiveDate::from_ymd_opt(*year as i32, month, day).ok_or(ErrorKind::InvalidDate)?
            }
            None => {
                let now = time::now();
                let date = NaiveDate::from_ymd_opt(now.tm_year + 1900, month, day)
                    .ok_or(ErrorKind::InvalidDate)?;
                if date.succ().and_hms(0, 0, 0).timestamp() < now.to_timespec().sec {
                    NaiveDate::from_ymd(now.tm_year + 1901, month, day)
                } else {
                    date
                }
            }
        };

        Ok(parse_date)
    }

    fn parse_time(&self, str_time: &str) -> Result<NaiveTime, RemindError> {
        let nums = str_time
            .split(':')
            .map(|s| s.parse::<u32>())
            .collect::<Result<Vec<_>, _>>()
            .context(ErrorKind::InvalidTime)?;

        if 2 != nums.len() {
            return Err(ErrorKind::InvalidTime.into());
        }

        let hour = nums[0];
        let minute = nums[1];

        Ok(NaiveTime::from_hms(hour, minute, 0))
    }

    pub fn get_time(&self, min_dur: Duration) -> Result<NaiveDateTime, RemindError> {
        if let Some(ref str_duration) = self.in_duration {
            let duration = parse_duration(&str_duration).context(ErrorKind::InvalidTime)?;

            if duration < min_dur {
                return Err(ErrorKind::TimeShort.into());
            }

            let tm = time::now().to_timespec();
            return Ok(NaiveDateTime::from_timestamp(
                tm.sec + duration.as_secs() as i64,
                0u32,
            ));
        }

        let mut date = None;
        if let Some(ref str_date) = self.on_date {
            date = Some(self.parse_date(str_date)?);
        }

        if let Some(ref str_time) = self.at_time {
            let time = self.parse_time(str_time)?;

            if let Some(date) = date {
                Ok(date.and_time(time))
            } else {
                let now = time::now();
                let today = NaiveDate::from_ymd_opt(
                    now.tm_year + 1900,
                    now.tm_mon as u32 + 1,
                    now.tm_mday as u32,
                ).ok_or(ErrorKind::InvalidDate)?;

                let time_today = today.and_time(time);

                if time_today.timestamp() < now.to_timespec().sec {
                    debug!("tomorrow");

                    Ok(today.succ().and_time(time))
                } else {
                    debug!("today");

                    Ok(time_today)
                }
            }
        } else {
            Ok(date.expect("At this point date has to be set")
                .and_hms(0, 0, 0))
        }
    }

    pub fn get_repeat(&self, min_dur: Duration) -> Result<Option<Duration>, RemindError> {
        if let Some(mut words) = self.every_time.clone() {
            if !words.chars().next().unwrap().is_digit(10) {
                words.insert(0, '1');
            }
            let dur = parse_duration(&words).context(ErrorKind::InvalidTime)?;

            if dur < min_dur {
                return Err(ErrorKind::RepeatTimeShort.into());
            }

            Ok(Some(dur))
        } else {
            Ok(None)
        }
    }

    pub fn get_target(&self) -> &str {
        &self.target
    }

    pub fn get_message(&self) -> &str {
        self.message.as_ref().expect("Has to be set")
    }
}
<commit_msg>Remind: Fix invalid error when parsing repeat time<commit_after>use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
use humantime::parse_duration;
use std::time::Duration;
use time;

use super::error::*;
use failure::ResultExt;

#[derive(Default, Debug)]
pub struct CommandParser {
    on_date: Option<String>,
    at_time: Option<String>,
    in_duration: Option<String>,
    every_time: Option<String>,
    target: String,
    message: Option<String>,
}

#[derive(PartialEq, Clone, Copy)]
enum ParseState {
    None,
    On,
    At,
    In,
    Every,
    Msg,
}

impl CommandParser {
    pub fn parse_target(mut tokens: Vec<String>) -> Result<Self, RemindError> {
        let mut parser = CommandParser::default();

        if tokens.is_empty() {
            Err(ErrorKind::MissingReceiver)?;
        }

        parser.target = tokens.remove(0);

        parser.parse_tokens(tokens)
    }

    pub fn with_target(tokens: Vec<String>, target: String) -> Result<Self, RemindError> {
        let mut parser = CommandParser::default();
        parser.target = target;

        parser.parse_tokens(tokens)
    }

    fn parse_tokens(mut self, tokens: Vec<String>) -> Result<Self, RemindError> {
        let mut state = ParseState::None;
        let mut cur_str = String::new();

        for token in tokens {
            let next_state = match token.as_ref() {
                "on" => ParseState::On,
                "at" => ParseState::At,
                "in" => ParseState::In,
                "every" => ParseState::Every,
                "to" => ParseState::Msg,
                _ => {
                    if !cur_str.is_empty() {
                        cur_str.push(' ');
                    }
                    cur_str.push_str(&token);
                    state
                }
            };

            if next_state != state {
                if state != ParseState::None {
                    self = self.add_string_by_state(state, cur_str)?;
                    cur_str = String::new();
                }

                state = next_state;
            }
        }

        self = self.add_string_by_state(state, cur_str)?;

        if self.message.is_none() {
            return Err(ErrorKind::MissingMessage.into());
        }

        if self.in_duration.is_some() && self.at_time.is_some()
            || self.in_duration.is_some() && self.on_date.is_some()
        {
            return Err(ErrorKind::AmbiguousTime.into());
        }

        if self.in_duration.is_none() && self.at_time.is_none() && self.on_date.is_none() {
            return Err(ErrorKind::MissingTime.into());
        }

        Ok(self)
    }

    fn add_string_by_state(self, state: ParseState, string: String) -> Result<Self, RemindError> {
        use self::ParseState::*;
        let string = Some(string);
        match state {
            On if self.on_date.is_none() => Ok(CommandParser {
                on_date: string,
                ..self
            }),
            At if self.at_time.is_none() => Ok(CommandParser {
                at_time: string,
                ..self
            }),
            In if self.in_duration.is_none() => Ok(CommandParser {
                in_duration: string,
                ..self
            }),
            Msg if self.message.is_none() => Ok(CommandParser {
                message: string,
                ..self
            }),
            Every if self.every_time.is_none() => Ok(CommandParser {
                every_time: string,
                ..self
            }),
            _ => Err(ErrorKind::MissingMessage.into()),
        }
    }

    fn parse_date(&self, str_date: &str) -> Result<NaiveDate, RemindError> {
        let nums = str_date
            .split('.')
            .map(|s| s.parse::<u32>())
            .collect::<Result<Vec<_>, _>>()
            .context(ErrorKind::InvalidDate)?;

        if 2 > nums.len() || nums.len() > 3 {
            return Err(ErrorKind::InvalidDate.into());
        }

        let day = nums[0];
        let month = nums[1];

        let parse_date = match nums.get(2) {
            Some(year) => {
                NaiveDate::from_ymd_opt(*year as i32, month, day).ok_or(ErrorKind::InvalidDate)?
            }
            None => {
                let now = time::now();
                let date = NaiveDate::from_ymd_opt(now.tm_year + 1900, month, day)
                    .ok_or(ErrorKind::InvalidDate)?;
                if date.succ().and_hms(0, 0, 0).timestamp() < now.to_timespec().sec {
                    NaiveDate::from_ymd(now.tm_year + 1901, month, day)
                } else {
                    date
                }
            }
        };

        Ok(parse_date)
    }

    fn parse_time(&self, str_time: &str) -> Result<NaiveTime, RemindError> {
        let nums = str_time
            .split(':')
            .map(|s| s.parse::<u32>())
            .collect::<Result<Vec<_>, _>>()
            .context(ErrorKind::InvalidTime)?;

        if 2 != nums.len() {
            return Err(ErrorKind::InvalidTime.into());
        }

        let hour = nums[0];
        let minute = nums[1];

        Ok(NaiveTime::from_hms(hour, minute, 0))
    }

    pub fn get_time(&self, min_dur: Duration) -> Result<NaiveDateTime, RemindError> {
        if let Some(ref str_duration) = self.in_duration {
            let duration = parse_duration(&str_duration).context(ErrorKind::InvalidTime)?;

            if duration < min_dur {
                return Err(ErrorKind::TimeShort.into());
            }

            let tm = time::now().to_timespec();
            return Ok(NaiveDateTime::from_timestamp(
                tm.sec + duration.as_secs() as i64,
                0u32,
            ));
        }

        let mut date = None;
        if let Some(ref str_date) = self.on_date {
            date = Some(self.parse_date(str_date)?);
        }

        if let Some(ref str_time) = self.at_time {
            let time = self.parse_time(str_time)?;

            if let Some(date) = date {
                Ok(date.and_time(time))
            } else {
                let now = time::now();
                let today = NaiveDate::from_ymd_opt(
                    now.tm_year + 1900,
                    now.tm_mon as u32 + 1,
                    now.tm_mday as u32,
                ).ok_or(ErrorKind::InvalidDate)?;

                let time_today = today.and_time(time);

                if time_today.timestamp() < now.to_timespec().sec {
                    debug!("tomorrow");

                    Ok(today.succ().and_time(time))
                } else {
                    debug!("today");

                    Ok(time_today)
                }
            }
        } else {
            Ok(date.expect("At this point date has to be set")
                .and_hms(0, 0, 0))
        }
    }

    pub fn get_repeat(&self, min_dur: Duration) -> Result<Option<Duration>, RemindError> {
        if let Some(mut words) = self.every_time.clone() {
            if !words.chars().next().unwrap().is_digit(10) {
                words.insert(0, '1');
            }
            let dur = parse_duration(&words).context(ErrorKind::InvalidTime)?;

            if dur <= min_dur {
                return Err(ErrorKind::RepeatTimeShort.into());
            }

            Ok(Some(dur))
        } else {
            Ok(None)
        }
    }

    pub fn get_target(&self) -> &str {
        &self.target
    }

    pub fn get_message(&self) -> &str {
        self.message.as_ref().expect("Has to be set")
    }
}
<|endoftext|>


<commit_before>use clap::{arg, builder::PossibleValue, command, value_parser, ValueEnum};

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
enum Mode {
    Fast,
    Slow,
}

// Can also be derived] with feature flag `derive`
impl ValueEnum for Mode {
    fn value_variants<'a>() -> &'a [Self] {
        &[Mode::Fast, Mode::Slow]
    }

    fn to_possible_value<'a>(&self) -> Option<PossibleValue<'a>> {
        Some(match self {
            Mode::Fast => PossibleValue::new("fast"),
            Mode::Slow => PossibleValue::new("slow"),
        })
    }
}

impl std::fmt::Display for Mode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.to_possible_value()
            .expect("no values are skipped")
            .get_name()
            .fmt(f)
    }
}

impl std::str::FromStr for Mode {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        for variant in Self::value_variants() {
            if variant.to_possible_value().unwrap().matches(s, false) {
                return Ok(*variant);
            }
        }
        Err(format!("Invalid variant: {}", s))
    }
}

fn main() {
    let matches = command!() // requires `cargo` feature
        .arg(
            arg!(<MODE>)
                .help("What mode to run the program in")
                .value_parser(value_parser!(Mode)),
        )
        .get_matches();

    // Note, it's safe to call unwrap() because the arg is required
    match matches
        .get_one::<Mode>("MODE")
        .expect("'MODE' is required and parsing will fail if its missing")
    {
        Mode::Fast => {
            println!("Hare");
        }
        Mode::Slow => {
            println!("Tortoise");
        }
    }
}
<commit_msg>docs(tutorial): Fix enum example comment<commit_after>use clap::{arg, builder::PossibleValue, command, value_parser, ValueEnum};

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
enum Mode {
    Fast,
    Slow,
}

// Can also be derived with feature flag `derive`
impl ValueEnum for Mode {
    fn value_variants<'a>() -> &'a [Self] {
        &[Mode::Fast, Mode::Slow]
    }

    fn to_possible_value<'a>(&self) -> Option<PossibleValue<'a>> {
        Some(match self {
            Mode::Fast => PossibleValue::new("fast"),
            Mode::Slow => PossibleValue::new("slow"),
        })
    }
}

impl std::fmt::Display for Mode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.to_possible_value()
            .expect("no values are skipped")
            .get_name()
            .fmt(f)
    }
}

impl std::str::FromStr for Mode {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        for variant in Self::value_variants() {
            if variant.to_possible_value().unwrap().matches(s, false) {
                return Ok(*variant);
            }
        }
        Err(format!("Invalid variant: {}", s))
    }
}

fn main() {
    let matches = command!() // requires `cargo` feature
        .arg(
            arg!(<MODE>)
                .help("What mode to run the program in")
                .value_parser(value_parser!(Mode)),
        )
        .get_matches();

    // Note, it's safe to call unwrap() because the arg is required
    match matches
        .get_one::<Mode>("MODE")
        .expect("'MODE' is required and parsing will fail if its missing")
    {
        Mode::Fast => {
            println!("Hare");
        }
        Mode::Slow => {
            println!("Tortoise");
        }
    }
}
<|endoftext|>


<commit_before>//! DNS Client cache

#[cfg(unix)]
use std::path::{Path, PathBuf};
use std::{
    collections::{hash_map::Entry, HashMap, VecDeque},
    future::Future,
    io,
    net::SocketAddr,
    time::Duration,
};

use log::trace;
use shadowsocks::{config::ServerConfig, net::ConnectOpts, relay::socks5::Address};
use tokio::sync::Mutex;
use trust_dns_resolver::proto::{error::ProtoError, op::Message};

use crate::local::context::ServiceContext;

use super::upstream::DnsClient;

#[derive(Clone, Debug, Hash, Eq, PartialEq, PartialOrd, Ord)]
enum DnsClientKey {
    TcpLocal(SocketAddr),
    UdpLocal(SocketAddr),
    #[cfg(unix)]
    UnixStream(PathBuf),
    TcpRemote(Address),
    UdpRemote(Address),
}

pub struct DnsClientCache {
    cache: Mutex<HashMap<DnsClientKey, VecDeque<DnsClient>>>,
    timeout: Duration,
    retry_count: usize,
    max_client_per_addr: usize,
}

impl DnsClientCache {
    pub fn new(max_client_per_addr: usize) -> DnsClientCache {
        DnsClientCache {
            cache: Mutex::new(HashMap::new()),
            timeout: Duration::from_secs(5),
            retry_count: 1,
            max_client_per_addr,
        }
    }

    pub async fn lookup_tcp_local(
        &self,
        ns: SocketAddr,
        msg: Message,
        connect_opts: &ConnectOpts,
    ) -> Result<Message, ProtoError> {
        let mut last_err = None;

        for _ in 0..self.retry_count {
            let key = DnsClientKey::TcpLocal(ns);
            let mut client = match self
                .get_client_or_create(&key, async { DnsClient::connect_tcp_local(ns, connect_opts).await })
                .await
            {
                Ok(client) => client,
                Err(err) => {
                    last_err = Some(From::from(err));
                    continue;
                }
            };

            let res = match client.lookup_timeout(msg.clone(), self.timeout).await {
                Ok(msg) => msg,
                Err(error) => {
                    last_err = Some(error);
                    continue;
                }
            };

            self.save_client(key, client).await;

            return Ok(res);
        }

        Err(last_err.unwrap())
    }

    pub async fn lookup_udp_local(
        &self,
        ns: SocketAddr,
        msg: Message,
        connect_opts: &ConnectOpts,
    ) -> Result<Message, ProtoError> {
        let mut last_err = None;

        for _ in 0..self.retry_count {
            let key = DnsClientKey::UdpLocal(ns);
            let mut client = match self
                .get_client_or_create(&key, async { DnsClient::connect_udp_local(ns, connect_opts).await })
                .await
            {
                Ok(client) => client,
                Err(err) => {
                    last_err = Some(From::from(err));
                    continue;
                }
            };

            let res = match client.lookup_timeout(msg.clone(), self.timeout).await {
                Ok(msg) => msg,
                Err(error) => {
                    last_err = Some(error);
                    continue;
                }
            };

            self.save_client(key, client).await;

            return Ok(res);
        }

        Err(last_err.unwrap())
    }

    #[cfg(unix)]
    pub async fn lookup_unix_stream<P: AsRef<Path>>(&self, ns: &P, msg: Message) -> Result<Message, ProtoError> {
        let mut last_err = None;

        let key = DnsClientKey::UnixStream(ns.as_ref().to_path_buf());
        for _ in 0..self.retry_count {
            let mut client = match self
                .get_client_or_create(&key, async { DnsClient::connect_unix_stream(ns).await })
                .await
            {
                Ok(client) => client,
                Err(err) => {
                    last_err = Some(From::from(err));
                    continue;
                }
            };

            let res = match client.lookup_timeout(msg.clone(), self.timeout).await {
                Ok(msg) => msg,
                Err(error) => {
                    last_err = Some(error);
                    continue;
                }
            };

            self.save_client(key, client).await;

            return Ok(res);
        }

        Err(last_err.unwrap())
    }

    pub async fn lookup_tcp_remote(
        &self,
        context: &ServiceContext,
        svr_cfg: &ServerConfig,
        ns: &Address,
        msg: Message,
    ) -> Result<Message, ProtoError> {
        let mut last_err = None;

        let key = DnsClientKey::UdpRemote(ns.clone());
        for _ in 0..self.retry_count {
            let mut client = match self
                .get_client_or_create(&key, async {
                    DnsClient::connect_tcp_remote(
                        context.context(),
                        svr_cfg,
                        ns,
                        context.connect_opts_ref(),
                        context.flow_stat(),
                    )
                    .await
                })
                .await
            {
                Ok(client) => client,
                Err(err) => {
                    last_err = Some(From::from(err));
                    continue;
                }
            };

            let res = match client.lookup_timeout(msg.clone(), self.timeout).await {
                Ok(msg) => msg,
                Err(error) => {
                    last_err = Some(error);
                    continue;
                }
            };

            self.save_client(key, client).await;

            return Ok(res);
        }

        Err(last_err.unwrap())
    }

    pub async fn lookup_udp_remote(
        &self,
        context: &ServiceContext,
        svr_cfg: &ServerConfig,
        ns: &Address,
        msg: Message,
    ) -> Result<Message, ProtoError> {
        let mut last_err = None;

        let key = DnsClientKey::TcpRemote(ns.clone());
        for _ in 0..self.retry_count {
            let mut client = match self
                .get_client_or_create(&key, async {
                    DnsClient::connect_udp_remote(
                        context.context(),
                        svr_cfg,
                        ns.clone(),
                        context.connect_opts_ref(),
                        context.flow_stat(),
                    )
                    .await
                })
                .await
            {
                Ok(client) => client,
                Err(err) => {
                    last_err = Some(From::from(err));
                    continue;
                }
            };

            let res = match client.lookup_timeout(msg.clone(), self.timeout).await {
                Ok(msg) => msg,
                Err(error) => {
                    last_err = Some(error);
                    continue;
                }
            };

            self.save_client(key, client).await;

            return Ok(res);
        }

        Err(last_err.unwrap())
    }

    async fn get_client_or_create<C>(&self, key: &DnsClientKey, create_fn: C) -> io::Result<DnsClient>
    where
        C: Future<Output = io::Result<DnsClient>>,
    {
        // Check if there already is a cached client
        if let Some(q) = self.cache.lock().await.get_mut(key) {
            if let Some(c) = q.pop_front() {
                trace!("take cached DNS client for {:?}", key);
                return Ok(c);
            }
        }

        trace!("creating connection to DNS server {:?}", key);

        // Create one
        create_fn.await
    }

    async fn save_client(&self, key: DnsClientKey, client: DnsClient) {
        match self.cache.lock().await.entry(key) {
            Entry::Occupied(occ) => {
                let q = occ.into_mut();
                q.push_back(client);
                if q.len() > self.max_client_per_addr {
                    q.pop_front();
                }
            }
            Entry::Vacant(vac) => {
                let mut q = VecDeque::with_capacity(5);
                q.push_back(client);
                vac.insert(q);
            }
        }
    }
}
<commit_msg>Remove the reuse for UNIX stream DNS lookup.<commit_after>//! DNS Client cache

#[cfg(unix)]
use std::path::Path;
use std::{
    collections::{hash_map::Entry, HashMap, VecDeque},
    future::Future,
    io,
    net::SocketAddr,
    time::Duration,
};

use log::trace;
use shadowsocks::{config::ServerConfig, net::ConnectOpts, relay::socks5::Address};
use tokio::sync::Mutex;
use trust_dns_resolver::proto::{error::ProtoError, op::Message};

use crate::local::context::ServiceContext;

use super::upstream::DnsClient;

#[derive(Clone, Debug, Hash, Eq, PartialEq, PartialOrd, Ord)]
enum DnsClientKey {
    TcpLocal(SocketAddr),
    UdpLocal(SocketAddr),
    TcpRemote(Address),
    UdpRemote(Address),
}

pub struct DnsClientCache {
    cache: Mutex<HashMap<DnsClientKey, VecDeque<DnsClient>>>,
    timeout: Duration,
    retry_count: usize,
    max_client_per_addr: usize,
}

impl DnsClientCache {
    pub fn new(max_client_per_addr: usize) -> DnsClientCache {
        DnsClientCache {
            cache: Mutex::new(HashMap::new()),
            timeout: Duration::from_secs(5),
            retry_count: 1,
            max_client_per_addr,
        }
    }

    pub async fn lookup_tcp_local(
        &self,
        ns: SocketAddr,
        msg: Message,
        connect_opts: &ConnectOpts,
    ) -> Result<Message, ProtoError> {
        let mut last_err = None;

        for _ in 0..self.retry_count {
            let key = DnsClientKey::TcpLocal(ns);
            let mut client = match self
                .get_client_or_create(&key, async { DnsClient::connect_tcp_local(ns, connect_opts).await })
                .await
            {
                Ok(client) => client,
                Err(err) => {
                    last_err = Some(From::from(err));
                    continue;
                }
            };

            let res = match client.lookup_timeout(msg.clone(), self.timeout).await {
                Ok(msg) => msg,
                Err(error) => {
                    last_err = Some(error);
                    continue;
                }
            };

            self.save_client(key, client).await;

            return Ok(res);
        }

        Err(last_err.unwrap())
    }

    pub async fn lookup_udp_local(
        &self,
        ns: SocketAddr,
        msg: Message,
        connect_opts: &ConnectOpts,
    ) -> Result<Message, ProtoError> {
        let mut last_err = None;

        for _ in 0..self.retry_count {
            let key = DnsClientKey::UdpLocal(ns);
            let mut client = match self
                .get_client_or_create(&key, async { DnsClient::connect_udp_local(ns, connect_opts).await })
                .await
            {
                Ok(client) => client,
                Err(err) => {
                    last_err = Some(From::from(err));
                    continue;
                }
            };

            let res = match client.lookup_timeout(msg.clone(), self.timeout).await {
                Ok(msg) => msg,
                Err(error) => {
                    last_err = Some(error);
                    continue;
                }
            };

            self.save_client(key, client).await;

            return Ok(res);
        }

        Err(last_err.unwrap())
    }

    #[cfg(unix)]
    pub async fn lookup_unix_stream<P: AsRef<Path>>(&self, ns: &P, msg: Message) -> Result<Message, ProtoError> {
        let mut last_err = None;

        for _ in 0..self.retry_count {
            let mut client = match DnsClient::connect_unix_stream(ns).await {
                Ok(client) => client,
                Err(err) => {
                    last_err = Some(From::from(err));
                    continue;
                }
            };

            let res = match client.lookup_timeout(msg.clone(), self.timeout).await {
                Ok(msg) => msg,
                Err(error) => {
                    last_err = Some(error);
                    continue;
                }
            };

            return Ok(res);
        }

        Err(last_err.unwrap())
    }

    pub async fn lookup_tcp_remote(
        &self,
        context: &ServiceContext,
        svr_cfg: &ServerConfig,
        ns: &Address,
        msg: Message,
    ) -> Result<Message, ProtoError> {
        let mut last_err = None;

        let key = DnsClientKey::UdpRemote(ns.clone());
        for _ in 0..self.retry_count {
            let mut client = match self
                .get_client_or_create(&key, async {
                    DnsClient::connect_tcp_remote(
                        context.context(),
                        svr_cfg,
                        ns,
                        context.connect_opts_ref(),
                        context.flow_stat(),
                    )
                    .await
                })
                .await
            {
                Ok(client) => client,
                Err(err) => {
                    last_err = Some(From::from(err));
                    continue;
                }
            };

            let res = match client.lookup_timeout(msg.clone(), self.timeout).await {
                Ok(msg) => msg,
                Err(error) => {
                    last_err = Some(error);
                    continue;
                }
            };

            self.save_client(key, client).await;

            return Ok(res);
        }

        Err(last_err.unwrap())
    }

    pub async fn lookup_udp_remote(
        &self,
        context: &ServiceContext,
        svr_cfg: &ServerConfig,
        ns: &Address,
        msg: Message,
    ) -> Result<Message, ProtoError> {
        let mut last_err = None;

        let key = DnsClientKey::TcpRemote(ns.clone());
        for _ in 0..self.retry_count {
            let mut client = match self
                .get_client_or_create(&key, async {
                    DnsClient::connect_udp_remote(
                        context.context(),
                        svr_cfg,
                        ns.clone(),
                        context.connect_opts_ref(),
                        context.flow_stat(),
                    )
                    .await
                })
                .await
            {
                Ok(client) => client,
                Err(err) => {
                    last_err = Some(From::from(err));
                    continue;
                }
            };

            let res = match client.lookup_timeout(msg.clone(), self.timeout).await {
                Ok(msg) => msg,
                Err(error) => {
                    last_err = Some(error);
                    continue;
                }
            };

            self.save_client(key, client).await;

            return Ok(res);
        }

        Err(last_err.unwrap())
    }

    async fn get_client_or_create<C>(&self, key: &DnsClientKey, create_fn: C) -> io::Result<DnsClient>
    where
        C: Future<Output = io::Result<DnsClient>>,
    {
        // Check if there already is a cached client
        if let Some(q) = self.cache.lock().await.get_mut(key) {
            if let Some(c) = q.pop_front() {
                trace!("take cached DNS client for {:?}", key);
                return Ok(c);
            }
        }

        trace!("creating connection to DNS server {:?}", key);

        // Create one
        create_fn.await
    }

    async fn save_client(&self, key: DnsClientKey, client: DnsClient) {
        match self.cache.lock().await.entry(key) {
            Entry::Occupied(occ) => {
                let q = occ.into_mut();
                q.push_back(client);
                if q.len() > self.max_client_per_addr {
                    q.pop_front();
                }
            }
            Entry::Vacant(vac) => {
                let mut q = VecDeque::with_capacity(5);
                q.push_back(client);
                vac.insert(q);
            }
        }
    }
}
<|endoftext|>