<commit_before>// This module defines constants and functions for working with zakocmp
// configuration files.

extern crate yaml_rust;

use std::io::{Error, ErrorKind};
use std::result::Result;
use std::vec::Vec;
use yaml_rust::{Yaml, YamlLoader};

const DEFAULT_POLICY_KEY: &'static str = "default-policy";
const POLICIES_KEY: &'static str = "policies";

// Enumerates the string representations of known policies.
const POLICY_REPR_IGNORE: &'static str = "ignore";
const POLICY_REPR_NOADD: &'static str = "noadd";
const POLICY_REPR_NODELETE: &'static str = "nodelete";
const POLICY_REPR_NOMODIFY: &'static str = "nomodify";
const POLICY_REPR_IMMUTABLE: &'static str = "immutable";

// Represents known policies as integers.
pub const POLICY_IGNORE: i32 = 0;
pub const POLICY_NOADD: i32 = 1 << 0;
pub const POLICY_NODELETE: i32 = 1 << 1;
pub const POLICY_NOMODIFY: i32 = 1 << 2;
pub const POLICY_IMMUTABLE: i32 = POLICY_NOADD | POLICY_NODELETE | POLICY_NOMODIFY;

// Represents a sorted vector of zakopane config rules, each mapping a
// path (prefix) to a policy. This type alias is provided for ease of
// coding.
type Policies = Vec<(String, i32)>;

// Represents a zakopane config. Please consult the documentation.
pub struct Config {
    default_policy: i32,
    policies: Policies,
}

// Borrows the string representation of one policy `token` and returns
// the equivalent integral representation.
fn policy_token_as_int(token: &str) -> Result<i32, Error> {
    match token {
        POLICY_REPR_IGNORE => Ok(POLICY_IGNORE),
        POLICY_REPR_NOADD => Ok(POLICY_NOADD),
        POLICY_REPR_NODELETE => Ok(POLICY_NODELETE),
        POLICY_REPR_NOMODIFY => Ok(POLICY_NOMODIFY),
        POLICY_REPR_IMMUTABLE => Ok(POLICY_IMMUTABLE),
        _ => Err(Error::new(
            ErrorKind::InvalidInput,
            format!("bad token: ``{}''", token),
        )),
    }
}

// Borrows the string representation of a combined `policy` and returns
// the equivalent integral representation. This function expects
// `policy` to comprise one or more policy tokens separated by commas.
fn policy_tokens_as_int(policy: &str) -> Result<i32, Error> {
    let policy_ints: Vec<i32> = policy
        .split(",")
        .map(|tok| policy_token_as_int(tok))
        .collect::<Result<Vec<i32>, Error>>()?;
    return Ok(policy_ints
        .iter()
        .fold(POLICY_IGNORE, |accum, elem| accum | elem));
}

// Borrows yaml representations of one line of zakopane policy and
// returns the corresponding valid tuple suitable for use in building a
// Policies object.
fn extract_policy(ypath: &Yaml, policy_tokens: &Yaml) -> Result<(String, i32), Error> {
    let path: String = match ypath.as_str() {
        Some(string) => string.to_owned(),
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed path?")),
    };
    let policy: i32 = match policy_tokens.as_str() {
        Some(string) => policy_tokens_as_int(string)?,
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed policy?")),
    };
    Ok((path, policy))
}

// Borrows the YAML representation of a zakopane config and returns the
// corresponding Policies. The return value can be benignly
// empty (e.g. if the present config elects not to specify any rules).
fn policies_from_yaml(doc: &Yaml) -> Result<Policies, Error> {
    let policies_map_yaml = &doc[POLICIES_KEY];
    if policies_map_yaml.is_badvalue() {
        // Assumes the config may be benignly devoid of specific
        // policies, returning Ok.
        return Ok(vec![]);
    }
    // Otherwise, iterates over the policies map. Each entry in the
    // policies map correlates a path prefix to a comma-separated list
    // of policies.
    let policies_map: &yaml_rust::yaml::Hash = match policies_map_yaml.as_hash() {
        Some(map) => map,
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed policies")),
    };
    let mut policies: Policies = policies_map
        .into_iter()
        .map(|pair| extract_policy(&pair.0, &pair.1))
        .collect::<Result<Policies, Error>>()?;
    policies.sort_unstable_by_key(|pair| pair.0.to_owned());
    Ok(policies)
}

// Borrows the YAML representation of a zakopane config and returns the
// integral default-policy defined within.
fn default_policy_from_yaml(doc: &Yaml) -> Result<i32, Error> {
    let default_policy_yaml = &doc[DEFAULT_POLICY_KEY];
    if default_policy_yaml.is_badvalue() {
        return Err(Error::new(ErrorKind::InvalidData, DEFAULT_POLICY_KEY));
    }
    let default_policy: i32 = match default_policy_yaml.as_str() {
        None => return Err(Error::new(ErrorKind::InvalidData, DEFAULT_POLICY_KEY)),
        Some(token) => policy_tokens_as_int(&token),
    }?;
    Ok(default_policy)
}

impl Config {
    // Borrows the string representation of a zakopane config and
    // returns a corresponding Config.
    pub fn new(config: &str) -> Result<Config, Error> {
        let docs: Vec<Yaml> = match YamlLoader::load_from_str(config) {
            Ok(val) => val,
            Err(scan_error) => return Err(Error::new(ErrorKind::InvalidData, scan_error)),
        };
        if docs.len() == 0 {
            return Err(Error::new(ErrorKind::InvalidData, "empty zakopane config"));
        }
        let doc = &docs[0];

        let default_policy = default_policy_from_yaml(&doc)?;
        let policies: Policies = policies_from_yaml(&doc)?;

        Ok(Config {
            default_policy: default_policy,
            policies: policies,
        })
    }

    // Returns how many rules this config contains.
    // This is never less than 1 as a default-policy is always required.
    pub fn rules(&self) -> usize {
        1 + self.policies.len()
    }

    // Borrows a `path` and returns the best-matched policy that
    // applies. This function returns an owned tuple of the
    // (closest-matched path expression, integral policy).
    //
    // This function represents the default-policy fallback by
    // returning the tuple consisting of an empty &str and the
    // default policy.
    pub fn match_policy(&self, path: &str) -> (&str, i32) {
        let mut best_match_path: &str = "";
        let mut best_match_policy: i32 = 0;
        for (prefix, policy) in self.policies.iter() {
            if path.starts_with(prefix) && prefix.len() > best_match_path.len() {
                best_match_path = prefix;
                best_match_policy = *policy;
            }
        }
        if best_match_path.len() == 0 {
            return ("", self.default_policy);
        }
        return (best_match_path, best_match_policy);
    }
}

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

    #[test]
    fn policy_token_bare() {
        let policy: i32 = policy_tokens_as_int(&"noadd").unwrap();
        assert!(policy == POLICY_NOADD);

        let policy: i32 = policy_tokens_as_int(&"nodelete").unwrap();
        assert!(policy == POLICY_NODELETE);

        let policy: i32 = policy_tokens_as_int(&"nomodify").unwrap();
        assert!(policy == POLICY_NOMODIFY);
    }

    #[test]
    fn policy_tokens_can_combo() {
        let policy: i32 = policy_tokens_as_int(&"noadd,nodelete").unwrap();
        assert!(policy == POLICY_NOADD | POLICY_NODELETE);
    }

    #[test]
    fn policy_tokens_can_repeat() {
        let policy: i32 =
            policy_tokens_as_int(&"noadd,noadd,noadd,noadd,nodelete,nodelete,nodelete,noadd")
                .unwrap();
        assert!(policy == POLICY_NOADD | POLICY_NODELETE);
    }

    #[test]
    fn config_must_not_be_empty() {
        let config = "";
        assert!(!Config::new(&config).is_ok());
    }

    #[test]
    fn config_must_not_be_obviously_malformed() {
        let config = r#"
This is not a zakopane config -
rather, it's two lines of text.
        "#;
        assert!(!Config::new(&config).is_ok());
    }

    #[test]
    fn config_requires_default_policy() {
        let config = r#"
policies:
    hello-there: nomodify
        "#;
        assert!(!Config::new(&config).is_ok());

        let mut config_with_default_policy: String = r#"
default-policy: immutable
        "#
        .to_string();
        config_with_default_policy.push_str(&config);
        assert!(Config::new(&config_with_default_policy).is_ok())
    }

    #[test]
    fn config_might_not_have_specific_policies() {
        let config = r#"
default-policy: nodelete
one-irrelevant-key: it doesn't matter what we put here
another-irrelevant-key: this doesn't invalidate the YAML
third-irrelevant-key: so long as it contains a default-policy
        "#;
        assert!(Config::new(&config).is_ok());
    }

    #[test]
    fn config_policies_must_be_a_map() {
        let config = r#"
default-policy: noadd
policies:
    -   eh?
    -   this ain't a map
        "#;
        assert!(!Config::new(&config).is_ok());
    }

    #[test]
    fn config_can_have_several_policies() {
        let config = r#"
default-policy: immutable
policies:
    hello-there: noadd
    general-kenobi: nodelete
        "#;
        assert!(Config::new(&config).is_ok());
    }

    #[test]
    fn match_default_policy() {
        let config_yaml = r#"
default-policy: noadd
        "#;
        let config = Config::new(&config_yaml).unwrap();

        // With only a default policy, this config has just 1 rule.
        assert!(config.rules() == 1);

        // Any path prefix we throw at match_policy() shall come up
        // as the default policy.
        let (_path, policy) = config.match_policy("./Documents/hello/there.txt");
        assert!(policy == POLICY_NOADD);
        let (_path, policy) = config.match_policy("./Music/general/kenobi.txt");
        assert!(policy == POLICY_NOADD);
    }

    #[test]
    fn match_nondefault_policies() {
        let config_yaml = r#"
default-policy: immutable
policies:
    ./Pictures/: noadd
    ./Pictures/2019/third-party/: nodelete
    ./Pictures/2020/: nomodify
    ./Pictures/2020/food/: nodelete,nomodify
        "#;
        let config = Config::new(&config_yaml).unwrap();

        assert!(config.rules() == 5);

        // Falls back on the default-policy absent any specific policy
        // defined for this file.
        let (_path, policy) = config.match_policy("./Documents/catch-me-senpai.txt");
        assert!(policy == POLICY_IMMUTABLE);
        // Matches only ``./Pictures.''
        let (_path, policy) = config.match_policy("./Pictures/2016/yano.jpg");
        assert!(policy == POLICY_NOADD);
        // As above and does _not_ match ``./Pictures/2019/third-party/.''
        let (_path, policy) = config.match_policy("./Pictures/2019/first-party.jpg");
        assert!(policy == POLICY_NOADD);
        // Does match ``./Pictures/2019/third-party/.''
        let (_path, policy) = config.match_policy("./Pictures/2019/third-party/yano.jpg");
        assert!(policy == POLICY_NODELETE);

        // Path prefix matching is done strictly and exactly;
        // ``food.md'' doesn't match ``food/,'' so there's no risk of
        // zakopane confusing cohabiting entities with similar basenames.
        let (path, policy) = config.match_policy("./Pictures/2020/food.md");
        assert!(policy == POLICY_NOMODIFY);
        assert!(path == "./Pictures/2020/");
        let (path, policy) = config.match_policy("./Pictures/2020/food/tacos.jpg");
        assert!(policy == POLICY_NODELETE | POLICY_NOMODIFY);
        assert!(path == "./Pictures/2020/food/");
    }
}
<commit_msg>zakocmp: use what you use<commit_after>// This module defines constants and functions for working with zakocmp
// configuration files.

extern crate yaml_rust;

use std::io::{Error, ErrorKind};
use std::result::Result;
use std::string::String;
use std::vec::Vec;
use yaml_rust::{Yaml, YamlLoader};

const DEFAULT_POLICY_KEY: &'static str = "default-policy";
const POLICIES_KEY: &'static str = "policies";

// Enumerates the string representations of known policies.
const POLICY_REPR_IGNORE: &'static str = "ignore";
const POLICY_REPR_NOADD: &'static str = "noadd";
const POLICY_REPR_NODELETE: &'static str = "nodelete";
const POLICY_REPR_NOMODIFY: &'static str = "nomodify";
const POLICY_REPR_IMMUTABLE: &'static str = "immutable";

// Represents known policies as integers.
pub const POLICY_IGNORE: i32 = 0;
pub const POLICY_NOADD: i32 = 1 << 0;
pub const POLICY_NODELETE: i32 = 1 << 1;
pub const POLICY_NOMODIFY: i32 = 1 << 2;
pub const POLICY_IMMUTABLE: i32 = POLICY_NOADD | POLICY_NODELETE | POLICY_NOMODIFY;

// Represents a sorted vector of zakopane config rules, each mapping a
// path (prefix) to a policy. This type alias is provided for ease of
// coding.
type Policies = Vec<(String, i32)>;

// Represents a zakopane config. Please consult the documentation.
pub struct Config {
    default_policy: i32,
    policies: Policies,
}

// Borrows the string representation of one policy `token` and returns
// the equivalent integral representation.
fn policy_token_as_int(token: &str) -> Result<i32, Error> {
    match token {
        POLICY_REPR_IGNORE => Ok(POLICY_IGNORE),
        POLICY_REPR_NOADD => Ok(POLICY_NOADD),
        POLICY_REPR_NODELETE => Ok(POLICY_NODELETE),
        POLICY_REPR_NOMODIFY => Ok(POLICY_NOMODIFY),
        POLICY_REPR_IMMUTABLE => Ok(POLICY_IMMUTABLE),
        _ => Err(Error::new(
            ErrorKind::InvalidInput,
            format!("bad token: ``{}''", token),
        )),
    }
}

// Borrows the string representation of a combined `policy` and returns
// the equivalent integral representation. This function expects
// `policy` to comprise one or more policy tokens separated by commas.
fn policy_tokens_as_int(policy: &str) -> Result<i32, Error> {
    let policy_ints: Vec<i32> = policy
        .split(",")
        .map(|tok| policy_token_as_int(tok))
        .collect::<Result<Vec<i32>, Error>>()?;
    return Ok(policy_ints
        .iter()
        .fold(POLICY_IGNORE, |accum, elem| accum | elem));
}

// Borrows yaml representations of one line of zakopane policy and
// returns the corresponding valid tuple suitable for use in building a
// Policies object.
fn extract_policy(ypath: &Yaml, policy_tokens: &Yaml) -> Result<(String, i32), Error> {
    let path: String = match ypath.as_str() {
        Some(string) => string.to_owned(),
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed path?")),
    };
    let policy: i32 = match policy_tokens.as_str() {
        Some(string) => policy_tokens_as_int(string)?,
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed policy?")),
    };
    Ok((path, policy))
}

// Borrows the YAML representation of a zakopane config and returns the
// corresponding Policies. The return value can be benignly
// empty (e.g. if the present config elects not to specify any rules).
fn policies_from_yaml(doc: &Yaml) -> Result<Policies, Error> {
    let policies_map_yaml = &doc[POLICIES_KEY];
    if policies_map_yaml.is_badvalue() {
        // Assumes the config may be benignly devoid of specific
        // policies, returning Ok.
        return Ok(vec![]);
    }
    // Otherwise, iterates over the policies map. Each entry in the
    // policies map correlates a path prefix to a comma-separated list
    // of policies.
    let policies_map: &yaml_rust::yaml::Hash = match policies_map_yaml.as_hash() {
        Some(map) => map,
        None => return Err(Error::new(ErrorKind::InvalidData, "malformed policies")),
    };
    let mut policies: Policies = policies_map
        .into_iter()
        .map(|pair| extract_policy(&pair.0, &pair.1))
        .collect::<Result<Policies, Error>>()?;
    policies.sort_unstable_by_key(|pair| pair.0.to_owned());
    Ok(policies)
}

// Borrows the YAML representation of a zakopane config and returns the
// integral default-policy defined within.
fn default_policy_from_yaml(doc: &Yaml) -> Result<i32, Error> {
    let default_policy_yaml = &doc[DEFAULT_POLICY_KEY];
    if default_policy_yaml.is_badvalue() {
        return Err(Error::new(ErrorKind::InvalidData, DEFAULT_POLICY_KEY));
    }
    let default_policy: i32 = match default_policy_yaml.as_str() {
        None => return Err(Error::new(ErrorKind::InvalidData, DEFAULT_POLICY_KEY)),
        Some(token) => policy_tokens_as_int(&token),
    }?;
    Ok(default_policy)
}

impl Config {
    // Borrows the string representation of a zakopane config and
    // returns a corresponding Config.
    pub fn new(config: &str) -> Result<Config, Error> {
        let docs: Vec<Yaml> = match YamlLoader::load_from_str(config) {
            Ok(val) => val,
            Err(scan_error) => return Err(Error::new(ErrorKind::InvalidData, scan_error)),
        };
        if docs.len() == 0 {
            return Err(Error::new(ErrorKind::InvalidData, "empty zakopane config"));
        }
        let doc = &docs[0];

        let default_policy = default_policy_from_yaml(&doc)?;
        let policies: Policies = policies_from_yaml(&doc)?;

        Ok(Config {
            default_policy: default_policy,
            policies: policies,
        })
    }

    // Returns how many rules this config contains.
    // This is never less than 1 as a default-policy is always required.
    pub fn rules(&self) -> usize {
        1 + self.policies.len()
    }

    // Borrows a `path` and returns the best-matched policy that
    // applies. This function returns an owned tuple of the
    // (closest-matched path expression, integral policy).
    //
    // This function represents the default-policy fallback by
    // returning the tuple consisting of an empty &str and the
    // default policy.
    pub fn match_policy(&self, path: &str) -> (&str, i32) {
        let mut best_match_path: &str = "";
        let mut best_match_policy: i32 = 0;
        for (prefix, policy) in self.policies.iter() {
            if path.starts_with(prefix) && prefix.len() > best_match_path.len() {
                best_match_path = prefix;
                best_match_policy = *policy;
            }
        }
        if best_match_path.len() == 0 {
            return ("", self.default_policy);
        }
        return (best_match_path, best_match_policy);
    }
}

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

    #[test]
    fn policy_token_bare() {
        let policy: i32 = policy_tokens_as_int(&"noadd").unwrap();
        assert!(policy == POLICY_NOADD);

        let policy: i32 = policy_tokens_as_int(&"nodelete").unwrap();
        assert!(policy == POLICY_NODELETE);

        let policy: i32 = policy_tokens_as_int(&"nomodify").unwrap();
        assert!(policy == POLICY_NOMODIFY);
    }

    #[test]
    fn policy_tokens_can_combo() {
        let policy: i32 = policy_tokens_as_int(&"noadd,nodelete").unwrap();
        assert!(policy == POLICY_NOADD | POLICY_NODELETE);
    }

    #[test]
    fn policy_tokens_can_repeat() {
        let policy: i32 =
            policy_tokens_as_int(&"noadd,noadd,noadd,noadd,nodelete,nodelete,nodelete,noadd")
                .unwrap();
        assert!(policy == POLICY_NOADD | POLICY_NODELETE);
    }

    #[test]
    fn config_must_not_be_empty() {
        let config = "";
        assert!(!Config::new(&config).is_ok());
    }

    #[test]
    fn config_must_not_be_obviously_malformed() {
        let config = r#"
This is not a zakopane config -
rather, it's two lines of text.
        "#;
        assert!(!Config::new(&config).is_ok());
    }

    #[test]
    fn config_requires_default_policy() {
        let config = r#"
policies:
    hello-there: nomodify
        "#;
        assert!(!Config::new(&config).is_ok());

        let mut config_with_default_policy: String = r#"
default-policy: immutable
        "#
        .to_string();
        config_with_default_policy.push_str(&config);
        assert!(Config::new(&config_with_default_policy).is_ok())
    }

    #[test]
    fn config_might_not_have_specific_policies() {
        let config = r#"
default-policy: nodelete
one-irrelevant-key: it doesn't matter what we put here
another-irrelevant-key: this doesn't invalidate the YAML
third-irrelevant-key: so long as it contains a default-policy
        "#;
        assert!(Config::new(&config).is_ok());
    }

    #[test]
    fn config_policies_must_be_a_map() {
        let config = r#"
default-policy: noadd
policies:
    -   eh?
    -   this ain't a map
        "#;
        assert!(!Config::new(&config).is_ok());
    }

    #[test]
    fn config_can_have_several_policies() {
        let config = r#"
default-policy: immutable
policies:
    hello-there: noadd
    general-kenobi: nodelete
        "#;
        assert!(Config::new(&config).is_ok());
    }

    #[test]
    fn match_default_policy() {
        let config_yaml = r#"
default-policy: noadd
        "#;
        let config = Config::new(&config_yaml).unwrap();

        // With only a default policy, this config has just 1 rule.
        assert!(config.rules() == 1);

        // Any path prefix we throw at match_policy() shall come up
        // as the default policy.
        let (_path, policy) = config.match_policy("./Documents/hello/there.txt");
        assert!(policy == POLICY_NOADD);
        let (_path, policy) = config.match_policy("./Music/general/kenobi.txt");
        assert!(policy == POLICY_NOADD);
    }

    #[test]
    fn match_nondefault_policies() {
        let config_yaml = r#"
default-policy: immutable
policies:
    ./Pictures/: noadd
    ./Pictures/2019/third-party/: nodelete
    ./Pictures/2020/: nomodify
    ./Pictures/2020/food/: nodelete,nomodify
        "#;
        let config = Config::new(&config_yaml).unwrap();

        assert!(config.rules() == 5);

        // Falls back on the default-policy absent any specific policy
        // defined for this file.
        let (_path, policy) = config.match_policy("./Documents/catch-me-senpai.txt");
        assert!(policy == POLICY_IMMUTABLE);
        // Matches only ``./Pictures.''
        let (_path, policy) = config.match_policy("./Pictures/2016/yano.jpg");
        assert!(policy == POLICY_NOADD);
        // As above and does _not_ match ``./Pictures/2019/third-party/.''
        let (_path, policy) = config.match_policy("./Pictures/2019/first-party.jpg");
        assert!(policy == POLICY_NOADD);
        // Does match ``./Pictures/2019/third-party/.''
        let (_path, policy) = config.match_policy("./Pictures/2019/third-party/yano.jpg");
        assert!(policy == POLICY_NODELETE);

        // Path prefix matching is done strictly and exactly;
        // ``food.md'' doesn't match ``food/,'' so there's no risk of
        // zakopane confusing cohabiting entities with similar basenames.
        let (path, policy) = config.match_policy("./Pictures/2020/food.md");
        assert!(policy == POLICY_NOMODIFY);
        assert!(path == "./Pictures/2020/");
        let (path, policy) = config.match_policy("./Pictures/2020/food/tacos.jpg");
        assert!(policy == POLICY_NODELETE | POLICY_NOMODIFY);
        assert!(path == "./Pictures/2020/food/");
    }
}
<|endoftext|>


<commit_before>// Syntax: master_server [ip address (0.0.0.0 by default)]
// This was created for HaloMD.

// If the server does not respond in at least this many seconds, it will be dropped from the list.
const DROP_TIME : i64 = 60;

// Blacklist for blocking IPs. Separate with newlines.
// Blacklisting IPs ignores heartbeat and keepalive packets from an IP address.
// That means that servers that are banned will not be immediately removed, but will time out, instead.
const BLACKLIST_FILE : &'static str = "blacklist.txt";

// Read the blacklist every x amount of seconds.
const BLACKLIST_UPDATE_TIME : u32 = 60;

// Note: The master server must have TCP 29920 open and UDP 27900 open.
const BROADCAST_PORT_UDP : u16 = 27900;
const SERVER_LIST_PORT_TCP : u16 = 29920;

// Broadcast packet types.
const KEEPALIVE : u8 = 8;
const HEARTBEAT : u8 = 3;

// Game state changes
const GAMEEXITED : u16 = 2;

// Opcode info
const OPCODE_INDEX : usize = 0;
const OPCODE_AND_HANDSHAKE_LENGTH : usize = 5;

// Broadcasted game name
const HALO_RETAIL : &'static str = "halor";

use std::net::{UdpSocket,TcpListener,SocketAddr};
use std::net::SocketAddr::{V4,V6};
use std::io::{Write,BufReader,BufRead};
use std::env;
use std::fs::File;
use std::thread;
use std::sync::{Arc, Mutex};

extern crate time;

mod halo_server;
use halo_server::HaloServer;

mod heartbeat_packet;
use heartbeat_packet::HeartbeatPacket;

fn get_ip(addr: SocketAddr) -> String {
    match addr {
        V4(ipv4) => ipv4.ip().to_string(),
        V6(ipv6) => "[".to_string() + &ipv6.ip().to_string() + "]"
    }
}

fn main() {
    let count = env::args().count();
    let ip = if count == 2 {
        let j : Vec<_> = env::args().collect();
        j[1].to_string()
    }
    else if count == 1 {
        "0.0.0.0".to_string()
    }
    else {
        println!("Only one argument is allowed: the IP to bind to.");
        return;
    };

    // We need to bind on two different ports. If it failed to bind (invalid IP, port is taken), then we must make sure this is known.
    let halo_socket = match UdpSocket::bind((&ip as &str,BROADCAST_PORT_UDP)) {
        Err(_) => {
            println!("Error creating a UDP socket at {}:{}.",ip,BROADCAST_PORT_UDP);
            return;
        },
        Ok(halo_socket) => halo_socket
    };

    let client_socket = match TcpListener::bind((&ip as &str,SERVER_LIST_PORT_TCP)) {
        Err(_) => {
            println!("Error listening to TCP at {}:{}.",ip,SERVER_LIST_PORT_TCP);
            return;
        },
        Ok(client_socket) => client_socket
    };

    // Mutex for thread safety.
    let servers_halo: Vec<HaloServer> = Vec::new();
    let servers_mut_udp = Arc::new(Mutex::new(servers_halo));
    let servers_mut_tcp = servers_mut_udp.clone();
    let servers_mut_destruction = servers_mut_udp.clone();

    // Destruction thread. This will remove servers that have not broadcasted their presence in a while.
    thread::spawn(move || {
        loop {
            thread::sleep_ms(10 * 1000);
            let mut servers = servers_mut_destruction.lock().unwrap();
            let timenow = time::now().to_timespec().sec;
            servers.retain(|x| x.last_alive + DROP_TIME > timenow);
        }
    });

    // Blacklist mutex. Concurrency needs to be safe, my friend.
    let blacklist: Vec<String> = Vec::new();
    let blacklist_update = Arc::new(Mutex::new(blacklist));
    let blacklist_udp = blacklist_update.clone();

    // Blacklist read thread.
    thread::spawn(move || {
        loop {
            // Placed in a block so blacklist is unlocked before sleeping to prevent threads from being locked for too long.
            {
                let mut blacklist_ref = blacklist_update.lock().unwrap();
                match File::open(BLACKLIST_FILE) {
                    Ok(file) => {
                        let reader = BufReader::new(&file);
                        match reader.lines().collect() {
                            Err(_) => {
                                blacklist_ref.clear();
                            },
                            Ok(t) => {
                                *blacklist_ref = t;
                            }
                        }
                    },
                    Err(_) => {}
                };
            }
            thread::sleep_ms(BLACKLIST_UPDATE_TIME * 1000);
        }
    });

    // TCP server thread. This is for the HaloMD application.
    thread::spawn(move || {
        loop {
            for stream in client_socket.incoming() {
                let mut client = match stream {
                    Err(_) => continue,
                    Ok(the_stream) => the_stream
                };
                // Unwrap the IP.
                let ip = match client.peer_addr() {
                    Err(_) => continue,
                    Ok(ip) => get_ip(ip)
                };

                let mut ips = String::new();

                // Make servers_ref go out of scope to unlock it for other threads, since we don't need it.
                {
                    let servers_ref = servers_mut_tcp.lock().unwrap();
                    let servers = (*servers_ref).iter();

                    for j in servers {
                        ips = ips + &(j.to_string()) + "\n";
                    }
                }

                // Some number placed after the requester's IP. If you ask me, the source code was abducted by aliens, and this is a tracking number. Regardless, it's needed.
                ips = ips + &ip + ":49149:3425";

                // We may be here a while. Just in case...
                thread::spawn( move || {
                    let _ = client.write_all(ips.as_bytes());
                });
            }
        }
    });

    // UDP server is run on the main thread. Servers broadcast their presence here.

    // These are the allowed game versions. HaloMD and Halo PC 1.09 uses 01.00.09.0620, while Halo PC servers on 1.10 use 01.00.10.0621 (these are also joinable).
    let game_versions = [ "01.00.09.0620".to_string(), "01.00.10.0621".to_string() ];

    let mut buffer = [0; 2048];
    loop {
        let (length, source) = match halo_socket.recv_from(&mut buffer) {
            Err(_) => continue,
            Ok(x) => x
        };
        if length <= OPCODE_INDEX {
            continue;
        }

        if buffer[OPCODE_INDEX] != KEEPALIVE && buffer[OPCODE_INDEX] != HEARTBEAT {
            continue;
        }

        let client_ip = get_ip(source);

        let blacklist_ref = blacklist_udp.lock().unwrap();
        if blacklist_ref.contains(&client_ip) {
            continue;
        }

        // Heartbeat packet. These contain null-terminated C strings and are ordered in key1[0]value1[0]key2[0]value2[0]key3[0]value3[0] where [0] is a byte equal to 0x00.
        if buffer[OPCODE_INDEX] == HEARTBEAT && length > OPCODE_AND_HANDSHAKE_LENGTH {

            let mut servers = servers_mut_udp.lock().unwrap();

            let packet = HeartbeatPacket::from_buffer(&buffer[OPCODE_AND_HANDSHAKE_LENGTH..length]);

            if packet.localport != 0 {
                let updatetime = time::now().to_timespec().sec;
                match servers.iter_mut().position(|x| x.ip == client_ip && x.port == packet.localport) {
                    None => {
                        if game_versions.contains(&packet.gamever) && &packet.gamename == HALO_RETAIL {
                            let serverness = HaloServer { ip:client_ip, port: packet.localport, last_alive: updatetime };
                            (*servers).push(serverness);
                        }
                    }
                    Some(k) => {
                        servers[k].last_alive = updatetime;
                        if packet.statechanged == GAMEEXITED {
                            servers.remove(k);
                        }
                    }
                };
            }
        }

        // Keepalive packet. We need to rely on the origin's port for this, unfortunately. This may mean that the source port is incorrect if the port was changed with NAT.
        else if buffer[OPCODE_INDEX] == KEEPALIVE {
            let mut servers_ref = servers_mut_udp.lock().unwrap();
            let servers = (*servers_ref).iter_mut();

            for i in servers {
                if i.ip == client_ip && i.port == source.port() {
                    i.last_alive = time::now().to_timespec().sec;
                    break;
                }
            }
        }
    }
}
<commit_msg>Log errors when they may be helpful<commit_after>// Syntax: master_server [ip address (0.0.0.0 by default)]
// This was created for HaloMD.

// If the server does not respond in at least this many seconds, it will be dropped from the list.
const DROP_TIME : i64 = 60;

// Blacklist for blocking IPs. Separate with newlines.
// Blacklisting IPs ignores heartbeat and keepalive packets from an IP address.
// That means that servers that are banned will not be immediately removed, but will time out, instead.
const BLACKLIST_FILE : &'static str = "blacklist.txt";

// Read the blacklist every x amount of seconds.
const BLACKLIST_UPDATE_TIME : u32 = 60;

// Note: The master server must have TCP 29920 open and UDP 27900 open.
const BROADCAST_PORT_UDP : u16 = 27900;
const SERVER_LIST_PORT_TCP : u16 = 29920;

// Broadcast packet types.
const KEEPALIVE : u8 = 8;
const HEARTBEAT : u8 = 3;

// Game state changes
const GAMEEXITED : u16 = 2;

// Opcode info
const OPCODE_INDEX : usize = 0;
const OPCODE_AND_HANDSHAKE_LENGTH : usize = 5;

// Broadcasted game name
const HALO_RETAIL : &'static str = "halor";

use std::net::{UdpSocket,TcpListener,SocketAddr};
use std::net::SocketAddr::{V4,V6};
use std::io::{Write,BufReader,BufRead};
use std::env;
use std::fs::File;
use std::thread;
use std::sync::{Arc, Mutex};

extern crate time;

mod halo_server;
use halo_server::HaloServer;

mod heartbeat_packet;
use heartbeat_packet::HeartbeatPacket;

fn get_ip(addr: SocketAddr) -> String {
    match addr {
        V4(ipv4) => ipv4.ip().to_string(),
        V6(ipv6) => "[".to_string() + &ipv6.ip().to_string() + "]"
    }
}

fn main() {
    let count = env::args().count();
    let ip = if count == 2 {
        let j : Vec<_> = env::args().collect();
        j[1].to_string()
    }
    else if count == 1 {
        "0.0.0.0".to_string()
    }
    else {
        println!("Only one argument is allowed: the IP to bind to.");
        return;
    };

    // We need to bind on two different ports. If it failed to bind (invalid IP, port is taken), then we must make sure this is known.
    let halo_socket = match UdpSocket::bind((&ip as &str,BROADCAST_PORT_UDP)) {
        Err(error) => {
            println!("Error creating a UDP socket at {}:{}. {}.",ip,BROADCAST_PORT_UDP,error);
            return;
        },
        Ok(halo_socket) => halo_socket
    };

    let client_socket = match TcpListener::bind((&ip as &str,SERVER_LIST_PORT_TCP)) {
        Err(error) => {
            println!("Error listening to TCP at {}:{}. {}.",ip,SERVER_LIST_PORT_TCP,error);
            return;
        },
        Ok(client_socket) => client_socket
    };

    // Mutex for thread safety.
    let servers_halo: Vec<HaloServer> = Vec::new();
    let servers_mut_udp = Arc::new(Mutex::new(servers_halo));
    let servers_mut_tcp = servers_mut_udp.clone();
    let servers_mut_destruction = servers_mut_udp.clone();

    // Destruction thread. This will remove servers that have not broadcasted their presence in a while.
    thread::spawn(move || {
        loop {
            thread::sleep_ms(10 * 1000);
            let mut servers = servers_mut_destruction.lock().unwrap();
            let timenow = time::now().to_timespec().sec;
            servers.retain(|x| x.last_alive + DROP_TIME > timenow);
        }
    });

    // Blacklist mutex. Concurrency needs to be safe, my friend.
    let blacklist: Vec<String> = Vec::new();
    let blacklist_update = Arc::new(Mutex::new(blacklist));
    let blacklist_udp = blacklist_update.clone();

    // Blacklist read thread.
    thread::spawn(move || {
        loop {
            // Placed in a block so blacklist is unlocked before sleeping to prevent threads from being locked for too long.
            {
                let mut blacklist_ref = blacklist_update.lock().unwrap();
                match File::open(BLACKLIST_FILE) {
                    Ok(file) => {
                        let reader = BufReader::new(&file);
                        match reader.lines().collect() {
                            Err(_) => {
                                blacklist_ref.clear();
                            },
                            Ok(t) => {
                                *blacklist_ref = t;
                            }
                        }
                    },
                    Err(_) => {}
                };
            }
            thread::sleep_ms(BLACKLIST_UPDATE_TIME * 1000);
        }
    });

    // TCP server thread. This is for the HaloMD application.
    thread::spawn(move || {
        loop {
            for stream in client_socket.incoming() {
                let mut client = match stream {
                    Err(_) => continue,
                    Ok(the_stream) => the_stream
                };
                // Unwrap the IP.
                let ip = match client.peer_addr() {
                    Err(_) => continue,
                    Ok(ip) => get_ip(ip)
                };

                let mut ips = String::new();

                // Make servers_ref go out of scope to unlock it for other threads, since we don't need it.
                {
                    let servers_ref = servers_mut_tcp.lock().unwrap();
                    let servers = (*servers_ref).iter();

                    for j in servers {
                        ips = ips + &(j.to_string()) + "\n";
                    }
                }

                // Some number placed after the requester's IP. If you ask me, the source code was abducted by aliens, and this is a tracking number. Regardless, it's needed.
                ips = ips + &ip + ":49149:3425";

                // We may be here a while. Just in case...
                thread::spawn( move || {
                    let _ = client.write_all(ips.as_bytes());
                });
            }
        }
    });

    // UDP server is run on the main thread. Servers broadcast their presence here.

    // These are the allowed game versions. HaloMD and Halo PC 1.09 uses 01.00.09.0620, while Halo PC servers on 1.10 use 01.00.10.0621 (these are also joinable).
    let game_versions = [ "01.00.09.0620".to_string(), "01.00.10.0621".to_string() ];

    let mut buffer = [0; 2048];
    loop {
        let (length, source) = match halo_socket.recv_from(&mut buffer) {
            Err(_) => continue,
            Ok(x) => x
        };
        if length <= OPCODE_INDEX {
            continue;
        }

        if buffer[OPCODE_INDEX] != KEEPALIVE && buffer[OPCODE_INDEX] != HEARTBEAT {
            continue;
        }

        let client_ip = get_ip(source);

        let blacklist_ref = blacklist_udp.lock().unwrap();
        if blacklist_ref.contains(&client_ip) {
            continue;
        }

        // Heartbeat packet. These contain null-terminated C strings and are ordered in key1[0]value1[0]key2[0]value2[0]key3[0]value3[0] where [0] is a byte equal to 0x00.
        if buffer[OPCODE_INDEX] == HEARTBEAT && length > OPCODE_AND_HANDSHAKE_LENGTH {

            let mut servers = servers_mut_udp.lock().unwrap();

            let packet = HeartbeatPacket::from_buffer(&buffer[OPCODE_AND_HANDSHAKE_LENGTH..length]);

            if packet.localport != 0 {
                let updatetime = time::now().to_timespec().sec;
                match servers.iter_mut().position(|x| x.ip == client_ip && x.port == packet.localport) {
                    None => {
                        if game_versions.contains(&packet.gamever) && &packet.gamename == HALO_RETAIL {
                            let serverness = HaloServer { ip:client_ip, port: packet.localport, last_alive: updatetime };
                            (*servers).push(serverness);
                        }
                    }
                    Some(k) => {
                        servers[k].last_alive = updatetime;
                        if packet.statechanged == GAMEEXITED {
                            servers.remove(k);
                        }
                    }
                };
            }
        }

        // Keepalive packet. We need to rely on the origin's port for this, unfortunately. This may mean that the source port is incorrect if the port was changed with NAT.
        else if buffer[OPCODE_INDEX] == KEEPALIVE {
            let mut servers_ref = servers_mut_udp.lock().unwrap();
            let servers = (*servers_ref).iter_mut();

            for i in servers {
                if i.ip == client_ip && i.port == source.port() {
                    i.last_alive = time::now().to_timespec().sec;
                    break;
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>use super::{directory_stack::DirectoryStack, variables::Variables};
use glob::glob;
use liner::{Completer, FilenameCompleter};

/// Performs escaping to an inner `FilenameCompleter` to enable a handful of special cases
/// needed by the shell, such as expanding '~' to a home directory, or adding a backslash
/// when a special character is contained within an expanded filename.
pub(crate) struct IonFileCompleter {
    /// The completer that this completer is  handling.
    inner: FilenameCompleter,
    /// A pointer to the directory stack in the shell.
    dir_stack: *const DirectoryStack,
    /// A pointer to the variables map in the shell.
    vars: *const Variables,
}

impl IonFileCompleter {
    pub(crate) fn new(
        path: Option<&str>,
        dir_stack: *const DirectoryStack,
        vars: *const Variables,
    ) -> IonFileCompleter {
        IonFileCompleter {
            inner: FilenameCompleter::new(path),
            dir_stack,
            vars,
        }
    }
}

impl Completer for IonFileCompleter {
    /// When the tab key is pressed, **Liner** will use this method to perform completions of
    /// filenames. As our `IonFileCompleter` is a wrapper around **Liner**'s
    /// `FilenameCompleter`,
    /// the purpose of our custom `Completer` is to expand possible `~` characters in the
    /// `start`
    /// value that we receive from the prompt, grab completions from the inner
    /// `FilenameCompleter`,
    /// and then escape the resulting filenames, as well as remove the expanded form of the `~`
    /// character and re-add the `~` character in it's place.
    fn completions(&self, start: &str) -> Vec<String> {
        // Only if the first character is a tilde character will we perform expansions
        if start.starts_with('~') {
            // Dereferencing the raw pointers here should be entirely safe, theoretically,
            // because no changes will occur to either of the underlying references in the
            // duration between creation of the completers and execution of their
            // completions.
            if let Some(expanded) = unsafe { (*self.vars).tilde_expansion(start, &*self.dir_stack) }
            {
                // Now we obtain completions for the `expanded` form of the `start` value.
                let mut iterator = FilenameCompletion::new(&expanded, |x| self.inner.completions(x));

                // And then we will need to take those completions and remove the expanded form
                // of the tilde pattern and replace it with that pattern yet again.
                let mut completions = Vec::new();

                // We can do that by obtaining the index position where the tilde character
                // ends. We don't search with `~` because we also want to
                // handle other tilde variants.
                let t_index = start.find('/').unwrap_or(1);
                // `tilde` is the tilde pattern, and `search` is the pattern that follows.
                let (tilde, search) = start.split_at(t_index as usize);

                if search.len() < 2 {
                    // If the length of the search pattern is less than 2, the search pattern is
                    // empty, and thus the completions actually contain files and directories in
                    // the home directory.

                    // The tilde pattern will actually be our `start` command in itself,
                    // and the completed form will be all of the characters beyond the length of
                    // the expanded form of the tilde pattern.
                    for completion in iterator {
                        completions.push([start, &completion[expanded.len()..]].concat());
                    }
                } else {
                    // To save processing time, we should get obtain the index position where our
                    // search pattern begins, and re-use that index to slice the completions so
                    // that we may re-add the tilde character with the completion that follows.
                    if let Some(completion) = iterator.next() {
                        if let Some(e_index) = expanded.rfind(search) {
                            completions.push(escape(&[tilde, &completion[e_index..]].concat()));
                            for completion in iterator {
                                let expanded = &completion[e_index..];
                                completions.push(escape(&[tilde, expanded].concat()));
                            }
                        }
                    }
                }

                return completions;
            }
        }

        FilenameCompletion::new(&start, |x| self.inner.completions(x)).collect()
    }
}

pub struct FilenameCompletion<'a> {
    iter: Box<Iterator<Item = String> + 'a>,
}

impl<'a> FilenameCompletion<'a> {
    pub fn new<LC: Fn(&str) -> Vec<String> + 'a>(start: &'a str, liner_complete: LC) -> Self {
        let unescaped_start = unescape(start);

        let split_start = unescaped_start.split("/");
        let mut string = String::with_capacity(5);

        // When 'start' is an absolute path, "/..." gets split to ["", "..."]
        // So we skip the first element and add "/" to the start of the string
        let mut start_for_glob = if unescaped_start.starts_with("/") {
            string.push('/');
            split_start.skip(1).fold(string, |mut state, element| {
                state.push_str(element);
                state.push_str("*/");
                state
            })
        } else {
            split_start.fold(string, |mut state, element| {
                state.push_str(element);
                state.push_str("*/");
                state
            })
        };
        start_for_glob.pop(); // pop out the last '/' character

        let iter_inner_glob: Box<Iterator<Item = String>> = match glob(&start_for_glob) {
            Ok(completions) => {
                let mut iter = completions
                    .filter_map(Result::ok)
                    .map(|x| x.to_string_lossy().into_owned())
                    .peekable();
                if iter.peek().is_some() {
                    Box::new(iter)
                } else {
                    Box::new(Some(escape(start)).into_iter())
                }
            }
            _ => Box::new(Some(escape(start)).into_iter()),
        };

        // Use Liner::Completer as well, to preserve the previous behaviour
        // around single-directory completions
        let iter = iter_inner_glob.flat_map(move |path| {
            liner_complete(&path)
                .into_iter()
                .map(|x| escape(x.as_str()))
        });

        FilenameCompletion {
            iter: Box::new(iter),
        }
    }
}

impl<'a> Iterator for FilenameCompletion<'a> {
    type Item = String;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next()
    }
}

/// Escapes filenames from the completer so that special characters will be properly escaped.
///
/// NOTE: Perhaps we should submit a PR to Liner to add a &'static [u8] field to
/// `FilenameCompleter` so that we don't have to perform the escaping ourselves?
fn escape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    for character in input.bytes() {
        match character {
            b'(' | b')' | b'[' | b']' | b'&' | b'$' | b'@' | b'{' | b'}' | b'<' | b'>' | b';'
            | b'"' | b'\'' | b'#' | b'^' | b'*' => output.push(b'\\'),
            _ => (),
        }
        output.push(character);
    }
    unsafe { String::from_utf8_unchecked(output) }
}

/// Unescapes filenames to be passed into the completer
fn unescape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    let mut bytes = input.bytes();
    while let Some(b) = bytes.next() {
        match b {
            b'\\' => if let Some(next) = bytes.next() {
                output.push(next);
            } else {
                output.push(b'\\')
            },
            _ => output.push(b),
        }
    }
    unsafe { String::from_utf8_unchecked(output) }
}

/// A completer that combines suggestions from multiple completers.
#[derive(Clone, Eq, PartialEq)]
pub(crate) struct MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    a: Vec<A>,
    b: B,
}

impl<A, B> MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    pub(crate) fn new(a: Vec<A>, b: B) -> MultiCompleter<A, B> { MultiCompleter { a, b } }
}

impl<A, B> Completer for MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    fn completions(&self, start: &str) -> Vec<String> {
        let mut completions = self.b.completions(start);
        for x in &self.a {
            completions.extend_from_slice(&x.completions(start));
        }
        completions
    }
}

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

    #[test]
    fn filename_completion() {
        let current_dir = env::current_dir().expect("Unable to get current directory");

        let completer = IonFileCompleter::new(
            current_dir.to_str(),
            &DirectoryStack::new(),
            &Variables::default(),
        );
        assert_eq!(completer.completions("testing"), vec!["testing/"]);
        assert_eq!(
            completer.completions("testing/file"),
            vec!["testing/file_with_text"]
        );

        assert_eq!(completer.completions("~"), vec!["~/"]);

        assert_eq!(
            completer.completions("tes/fil"),
            vec!["testing/file_with_text"]
        );
    }
}
<commit_msg>Remove unneeded box trait and return `impl Iterator` instead<commit_after>use super::{directory_stack::DirectoryStack, variables::Variables};
use glob::glob;
use liner::{Completer, FilenameCompleter};

/// Performs escaping to an inner `FilenameCompleter` to enable a handful of special cases
/// needed by the shell, such as expanding '~' to a home directory, or adding a backslash
/// when a special character is contained within an expanded filename.
pub(crate) struct IonFileCompleter {
    /// The completer that this completer is  handling.
    inner: FilenameCompleter,
    /// A pointer to the directory stack in the shell.
    dir_stack: *const DirectoryStack,
    /// A pointer to the variables map in the shell.
    vars: *const Variables,
}

impl IonFileCompleter {
    pub(crate) fn new(
        path: Option<&str>,
        dir_stack: *const DirectoryStack,
        vars: *const Variables,
    ) -> IonFileCompleter {
        IonFileCompleter {
            inner: FilenameCompleter::new(path),
            dir_stack,
            vars,
        }
    }
}

impl Completer for IonFileCompleter {
    /// When the tab key is pressed, **Liner** will use this method to perform completions of
    /// filenames. As our `IonFileCompleter` is a wrapper around **Liner**'s
    /// `FilenameCompleter`,
    /// the purpose of our custom `Completer` is to expand possible `~` characters in the
    /// `start`
    /// value that we receive from the prompt, grab completions from the inner
    /// `FilenameCompleter`,
    /// and then escape the resulting filenames, as well as remove the expanded form of the `~`
    /// character and re-add the `~` character in it's place.
    fn completions(&self, start: &str) -> Vec<String> {
        // Only if the first character is a tilde character will we perform expansions
        if start.starts_with('~') {
            // Dereferencing the raw pointers here should be entirely safe, theoretically,
            // because no changes will occur to either of the underlying references in the
            // duration between creation of the completers and execution of their
            // completions.
            if let Some(expanded) = unsafe { (*self.vars).tilde_expansion(start, &*self.dir_stack) }
            {
                // Now we obtain completions for the `expanded` form of the `start` value.
                let mut iterator = filename_completion(&expanded, |x| self.inner.completions(x));

                // And then we will need to take those completions and remove the expanded form
                // of the tilde pattern and replace it with that pattern yet again.
                let mut completions = Vec::new();

                // We can do that by obtaining the index position where the tilde character
                // ends. We don't search with `~` because we also want to
                // handle other tilde variants.
                let t_index = start.find('/').unwrap_or(1);
                // `tilde` is the tilde pattern, and `search` is the pattern that follows.
                let (tilde, search) = start.split_at(t_index as usize);

                if search.len() < 2 {
                    // If the length of the search pattern is less than 2, the search pattern is
                    // empty, and thus the completions actually contain files and directories in
                    // the home directory.

                    // The tilde pattern will actually be our `start` command in itself,
                    // and the completed form will be all of the characters beyond the length of
                    // the expanded form of the tilde pattern.
                    for completion in iterator {
                        completions.push([start, &completion[expanded.len()..]].concat());
                    }
                } else {
                    // To save processing time, we should get obtain the index position where our
                    // search pattern begins, and re-use that index to slice the completions so
                    // that we may re-add the tilde character with the completion that follows.
                    if let Some(completion) = iterator.next() {
                        if let Some(e_index) = expanded.rfind(search) {
                            completions.push(escape(&[tilde, &completion[e_index..]].concat()));
                            for completion in iterator {
                                let expanded = &completion[e_index..];
                                completions.push(escape(&[tilde, expanded].concat()));
                            }
                        }
                    }
                }

                return completions;
            }
        }

        filename_completion(&start, |x| self.inner.completions(x)).collect()
    }
}

fn filename_completion<'a, LC>(start: &'a str, liner_complete: LC) -> impl Iterator<Item = String> + 'a
where
    LC: Fn(&str) -> Vec<String> + 'a
{
    let unescaped_start = unescape(start);

    let split_start = unescaped_start.split("/");
    let mut string = String::with_capacity(5);

    // When 'start' is an absolute path, "/..." gets split to ["", "..."]
    // So we skip the first element and add "/" to the start of the string
    let mut start_for_glob = if unescaped_start.starts_with("/") {
        string.push('/');
        split_start.skip(1).fold(string, |mut state, element| {
            state.push_str(element);
            state.push_str("*/");
            state
        })
    } else {
        split_start.fold(string, |mut state, element| {
            state.push_str(element);
            state.push_str("*/");
            state
        })
    };
    start_for_glob.pop(); // pop out the last '/' character

    let iter_inner_glob: Box<Iterator<Item = String>> = match glob(&start_for_glob) {
        Ok(completions) => {
            let mut iter = completions
                .filter_map(Result::ok)
                .map(|x| x.to_string_lossy().into_owned())
                .peekable();
            if iter.peek().is_some() {
                Box::new(iter)
            } else {
                Box::new(Some(escape(start)).into_iter())
            }
        }
        _ => Box::new(Some(escape(start)).into_iter()),
    };

    // Use Liner::Completer as well, to preserve the previous behaviour
    // around single-directory completions
    iter_inner_glob.flat_map(move |path| {
        liner_complete(&path)
            .into_iter()
            .map(|x| escape(x.as_str()))
    })
}

/// Escapes filenames from the completer so that special characters will be properly escaped.
///
/// NOTE: Perhaps we should submit a PR to Liner to add a &'static [u8] field to
/// `FilenameCompleter` so that we don't have to perform the escaping ourselves?
fn escape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    for character in input.bytes() {
        match character {
            b'(' | b')' | b'[' | b']' | b'&' | b'$' | b'@' | b'{' | b'}' | b'<' | b'>' | b';'
            | b'"' | b'\'' | b'#' | b'^' | b'*' => output.push(b'\\'),
            _ => (),
        }
        output.push(character);
    }
    unsafe { String::from_utf8_unchecked(output) }
}

/// Unescapes filenames to be passed into the completer
fn unescape(input: &str) -> String {
    let mut output = Vec::with_capacity(input.len());
    let mut bytes = input.bytes();
    while let Some(b) = bytes.next() {
        match b {
            b'\\' => if let Some(next) = bytes.next() {
                output.push(next);
            } else {
                output.push(b'\\')
            },
            _ => output.push(b),
        }
    }
    unsafe { String::from_utf8_unchecked(output) }
}

/// A completer that combines suggestions from multiple completers.
#[derive(Clone, Eq, PartialEq)]
pub(crate) struct MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    a: Vec<A>,
    b: B,
}

impl<A, B> MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    pub(crate) fn new(a: Vec<A>, b: B) -> MultiCompleter<A, B> { MultiCompleter { a, b } }
}

impl<A, B> Completer for MultiCompleter<A, B>
where
    A: Completer,
    B: Completer,
{
    fn completions(&self, start: &str) -> Vec<String> {
        let mut completions = self.b.completions(start);
        for x in &self.a {
            completions.extend_from_slice(&x.completions(start));
        }
        completions
    }
}

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

    #[test]
    fn filename_completion() {
        let current_dir = env::current_dir().expect("Unable to get current directory");

        let completer = IonFileCompleter::new(
            current_dir.to_str(),
            &DirectoryStack::new(),
            &Variables::default(),
        );
        assert_eq!(completer.completions("testing"), vec!["testing/"]);
        assert_eq!(
            completer.completions("testing/file"),
            vec!["testing/file_with_text"]
        );

        assert_eq!(completer.completions("~"), vec!["~/"]);

        assert_eq!(
            completer.completions("tes/fil"),
            vec!["testing/file_with_text"]
        );
    }
}
<|endoftext|>


<commit_before>use std::convert::TryInto;
use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{debug, info, warn};
use nix::{
    sys::signal::{self, Signal},
    unistd::Pid,
};
use procfs::process::{all_processes, Process};

use crate::task_handler::ProcessAction;

pub fn compile_shell_command(command_string: &str) -> Command {
    let mut command = Command::new("sh");
    command.arg("-c").arg(command_string);

    command
}

/// Send a signal to one of Pueue's child process handles.
///
/// There are two scenarios:
///
/// **Normal case**
///
/// A task, such as `sleep 60` get's spawned by the posix shell `sh`.
/// This results in the process `sh -c 'sleep 60'`.
/// Since the posix shell doesn't propagate any process signals to its children, we have to:
/// 1. Send the signal to the shell.
/// 2. Send the signal directly to the children.
///     In our case this would be the `sleep 60` child process.
///
/// If the user also want's to send the signal to all child processes of the task,
/// we have to get all child-processes of the child process.
///
/// **Special case**
///
/// The posix shell `sh` has some some inconsistent behavior.
/// In some circumstances and environments, the `sh -c $command` doesn't spawn a `sh` process with a
/// `$command` child-process, but rather spawns the `$command` as a top-level process directly.
///
/// This makes things a bit more complicated, since we have to find out whether a shell is spawned
/// or not. If a shell is spawned, we do the **Normal case** handling.
///
/// If **no** shell is spawned, we have to send the signal to the top-level process only.
///
/// If the user also want's to send the signal to all child processes of the task,
/// we have to get all child-processes of that `$command` process. and send them the signal.
///
/// Returns `Ok(true)`, if everything went alright
/// Returns `Ok(false)`, if the process went away while we tried to send the signal.
pub fn send_signal_to_child(
    child: &Child,
    action: &ProcessAction,
    send_to_children: bool,
) -> Result<bool> {
    let signal = get_signal_from_action(action);
    let pid: i32 = child.id().try_into().unwrap();
    // Check whether this process actually spawned a shell.
    let is_shell = if let Ok(is_shell) = did_process_spawn_shell(pid) {
        is_shell
    } else {
        return Ok(false);
    };

    if is_shell {
        // If it's a shell, we have to send the signal to the actual shell and to all it's children.
        // There might be multiple children, for instance, when users use the `&` operator.
        // If the `send_to_children` flag is given, the

        // Send the signal to the shell, don't propagate to its children yet.
        send_signal_to_process(pid, action, false)?;

        // Now send the signal to the shells child processes and their respective
        // children if the user wants to do so.
        let shell_children = get_child_processes(pid);
        for shell_child in shell_children {
            send_signal_to_process(shell_child.pid(), action, send_to_children)?;
        }
    } else {
        // If it isn't a shell, send the signal directly to the process.
        // Handle children normally.
        send_signal_to_process(pid, action, send_to_children)?;
    }

    signal::kill(Pid::from_raw(pid), signal)?;
    Ok(true)
}

/// This is a helper function to safely kill a child process.
/// Its purpose is to properly kill all processes and prevent any dangling processes.
///
/// Sadly, this needs some extra handling. Check the docstring of `send_signal_to_child` for
/// additional information on why this needs to be done.
///
/// Returns `true`, if everything went alright
/// Returns `false`, if the process went away while we tried to send the signal.
pub fn kill_child(task_id: usize, child: &mut Child, kill_children: bool) -> bool {
    let pid: i32 = child.id().try_into().unwrap();

    // Check whether this process actually spawned a shell.
    let is_shell = if let Ok(is_shell) = did_process_spawn_shell(pid) {
        is_shell
    } else {
        return false;
    };

    // We have to kill the root process first, to prevent it from spawning new processes.
    // However, this prevents us from getting its child processes afterwards.
    // That's why we have to get the list of child processes already now.
    let mut child_processes = None;
    if kill_children || is_shell {
        child_processes = Some(get_child_processes(pid));
    }

    // Kill the parent first
    match child.kill() {
        Err(_) => {
            info!("Task {} has already finished by itself", task_id);
            false
        }
        _ => {
            // Now kill all remaining children, after the parent has been killed.
            // If a shell is spawned, we have to manually send the kill signal to all children.
            // Otherwise only send a signal to all children if the `kill_children` flag is set.
            if let Some(child_processes) = child_processes {
                if is_shell {
                    for child_process in child_processes {
                        // Send the signal to each child process, show warning if this fails.
                        let process_pid = child_process.pid();
                        if let Err(error) =
                            send_signal_to_process(process_pid, &ProcessAction::Kill, kill_children)
                        {
                            warn!(
                                "Failed to send kill to pid {} with error {:?}",
                                process_pid, error
                            );
                        }
                    }
                } else if kill_children {
                    send_signal_to_processes(child_processes, &ProcessAction::Kill);
                }
            }

            true
        }
    }
}

/// Check whether a process's commandline string is actually a shell or not
fn did_process_spawn_shell(pid: i32) -> Result<bool> {
    // Get the /proc representation of the child, so we can do some checks
    let process = if let Ok(process) = Process::new(pid) {
        process
    } else {
        info!(
            "Process to kill has probably just gone away. Process {}",
            pid
        );
        bail!("Process has just gone away");
    };

    // Get the root command and check whether it's actually a shell with `sh -c`.
    let mut cmdline = if let Ok(cmdline) = process.cmdline() {
        cmdline
    } else {
        info!(
            "Process to kill has probably just gone away. Process {}",
            pid
        );
        bail!("Process has just gone away");
    };

    if cmdline.len() < 3 {
        return Ok(false);
    }

    if cmdline.remove(0) != "sh" {
        return Ok(false);
    }

    if cmdline.remove(0) != "-c" {
        return Ok(false);
    }

    Ok(true)
}

/// Send a signal to a unix process.
fn send_signal_to_process(
    pid: i32,
    action: &ProcessAction,
    send_to_children: bool,
) -> Result<bool, nix::Error> {
    let signal = get_signal_from_action(action);
    debug!("Sending signal {} to {}", signal, pid);

    // Send the signal to all children, if that's what the user wants.
    if send_to_children {
        send_signal_to_children(pid, action);
    }

    signal::kill(Pid::from_raw(pid), signal)?;
    Ok(true)
}

/// A small helper that sends a signal to all children of a specific process by id.
fn send_signal_to_children(pid: i32, action: &ProcessAction) {
    send_signal_to_processes(get_child_processes(pid), action);
}

/// Send a signal to a list of processes.
/// This is a convenience wrapper around `send_signal_to_process`.
fn send_signal_to_processes(processes: Vec<Process>, action: &ProcessAction) {
    let signal = get_signal_from_action(action);
    for process in processes {
        // Process is no longer alive, skip this one.
        if !process.is_alive() {
            continue;
        }

        let pid = Pid::from_raw(process.pid);
        if let Err(error) = signal::kill(pid, signal) {
            warn!(
                "Failed send signal {:?} to Pid {}: {:?}",
                signal, process.pid, error
            );
        }
    }
}

/// Get all children of a specific process
fn get_child_processes(pid: i32) -> Vec<Process> {
    let all_processes = match all_processes() {
        Err(error) => {
            warn!("Failed to get full process list: {}", error);
            return Vec::new();
        }
        Ok(processes) => processes,
    };

    all_processes
        .into_iter()
        .filter(|process| process.stat.ppid == pid)
        .collect()
}

fn get_signal_from_action(action: &ProcessAction) -> Signal {
    match action {
        ProcessAction::Kill => Signal::SIGKILL,
        ProcessAction::Pause => Signal::SIGSTOP,
        ProcessAction::Resume => Signal::SIGCONT,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread::sleep;
    use std::time::Duration;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: i32) -> bool {
        match Process::new(pid) {
            Ok(process) => !process.is_alive(),
            Err(_) => true,
        }
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[test]
    /// Ensure a `sh -c` command will be properly killed without detached processes.
    fn test_shell_command_is_killed() {
        let mut child = compile_shell_command("sleep 60 & sleep 60 && echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process indeed spawned a shell.
        assert!(did_process_spawn_shell(pid).unwrap());

        // Get all child processes, so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }

    #[test]
    /// Ensure that a `sh -c` process with a child process that has children of its own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() {
        let mut child = compile_shell_command("bash -c 'sleep 60 && sleep 60' && sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process indeed spawned a shell.
        assert!(did_process_spawn_shell(pid).unwrap());

        // Get all child processes and all childrens children,
        // so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 1);
        let mut childrens_children = Vec::new();
        for child_process in &child_processes {
            childrens_children.extend(get_child_processes(child_process.stat.pid));
        }
        assert_eq!(childrens_children.len(), 1);

        // Kill the process and make sure its childen will be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }

        // Assert that all children's child processes have been killed.
        for child_process in childrens_children {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }

    #[test]
    /// Ensure a normal command without `sh -c` will be killed.
    fn test_normal_command_is_killed() {
        let mut child = Command::new("sleep")
            .arg("60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process did not spawn a shell.
        assert!(!did_process_spawn_shell(pid).unwrap());

        // No child processes exist
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 0);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));
    }

    #[test]
    /// Ensure a normal command and all its children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() {
        let mut child = Command::new("bash")
            .arg("-c")
            .arg("sleep 60 & sleep 60 && sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process indeed spawned a shell.
        assert!(!did_process_spawn_shell(pid).unwrap());

        // Get all child processes, so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }
}
<commit_msg>Refactor linux kill_child code<commit_after>use std::convert::TryInto;
use std::process::{Child, Command};

use anyhow::{bail, Result};
use log::{debug, info, warn};
use nix::{
    sys::signal::{self, Signal},
    unistd::Pid,
};
use procfs::process::{all_processes, Process};

use crate::task_handler::ProcessAction;

pub fn compile_shell_command(command_string: &str) -> Command {
    let mut command = Command::new("sh");
    command.arg("-c").arg(command_string);

    command
}

/// Send a signal to one of Pueue's child process handles.
///
/// There are two scenarios:
///
/// **Normal case**
///
/// A task, such as `sleep 60` get's spawned by the posix shell `sh`.
/// This results in the process `sh -c 'sleep 60'`.
/// Since the posix shell doesn't propagate any process signals to its children, we have to:
/// 1. Send the signal to the shell.
/// 2. Send the signal directly to the children.
///     In our case this would be the `sleep 60` child process.
///
/// If the user also want's to send the signal to all child processes of the task,
/// we have to get all child-processes of the child process.
///
/// **Special case**
///
/// The posix shell `sh` has some some inconsistent behavior.
/// In some circumstances and environments, the `sh -c $command` doesn't spawn a `sh` process with a
/// `$command` child-process, but rather spawns the `$command` as a top-level process directly.
///
/// This makes things a bit more complicated, since we have to find out whether a shell is spawned
/// or not. If a shell is spawned, we do the **Normal case** handling.
///
/// If **no** shell is spawned, we have to send the signal to the top-level process only.
///
/// If the user also want's to send the signal to all child processes of the task,
/// we have to get all child-processes of that `$command` process. and send them the signal.
///
/// Returns `Ok(true)`, if everything went alright
/// Returns `Ok(false)`, if the process went away while we tried to send the signal.
pub fn send_signal_to_child(
    child: &Child,
    action: &ProcessAction,
    send_to_children: bool,
) -> Result<bool> {
    let signal = get_signal_from_action(action);
    let pid: i32 = child.id().try_into().unwrap();
    // Check whether this process actually spawned a shell.
    let is_shell = if let Ok(is_shell) = did_process_spawn_shell(pid) {
        is_shell
    } else {
        return Ok(false);
    };

    if is_shell {
        // If it's a shell, we have to send the signal to the actual shell and to all it's children.
        // There might be multiple children, for instance, when users use the `&` operator.
        // If the `send_to_children` flag is given, the

        // Send the signal to the shell, don't propagate to its children yet.
        send_signal_to_process(pid, action, false)?;

        // Now send the signal to the shells child processes and their respective
        // children if the user wants to do so.
        let shell_children = get_child_processes(pid);
        for shell_child in shell_children {
            send_signal_to_process(shell_child.pid(), action, send_to_children)?;
        }
    } else {
        // If it isn't a shell, send the signal directly to the process.
        // Handle children normally.
        send_signal_to_process(pid, action, send_to_children)?;
    }

    signal::kill(Pid::from_raw(pid), signal)?;
    Ok(true)
}

/// This is a helper function to safely kill a child process.
/// Its purpose is to properly kill all processes and prevent any dangling processes.
///
/// Sadly, this needs some extra handling. Check the docstring of `send_signal_to_child` for
/// additional information on why this needs to be done.
///
/// Returns `true`, if everything went alright
/// Returns `false`, if the process went away while we tried to send the signal.
pub fn kill_child(task_id: usize, child: &mut Child, kill_children: bool) -> bool {
    let pid: i32 = child.id().try_into().unwrap();

    // Check whether this process actually spawned a shell.
    let is_shell = if let Ok(is_shell) = did_process_spawn_shell(pid) {
        is_shell
    } else {
        return false;
    };

    // We have to kill the root process first, to prevent it from spawning new processes.
    // However, this prevents us from getting its child processes afterwards.
    // That's why we have to get the list of child processes already now.
    let mut child_processes = None;
    if kill_children || is_shell {
        child_processes = Some(get_child_processes(pid));
    }

    // Kill the parent first
    let kill_result = child.kill();
    if kill_result.is_err() {
        info!("Task {} has already finished by itself", task_id);
        return false;
    }

    // Do an early return, if we don't need to kill any children.
    let child_processes = if let Some(child_processes) = child_processes {
        child_processes
    } else {
        return true;
    };

    // Now kill all remaining children. The parent has been already been killed at this point.
    // If a shell is spawned, we have to manually send the kill signal to all children.
    // Otherwise only send a signal to all children if the `kill_children` flag is set.
    if is_shell {
        for child_process in child_processes {
            // Send the signal to each child process, show warning if this fails.
            let process_pid = child_process.pid();
            if let Err(error) =
                send_signal_to_process(process_pid, &ProcessAction::Kill, kill_children)
            {
                warn!(
                    "Failed to send kill to pid {} with error {:?}",
                    process_pid, error
                );
            }
        }
    } else if kill_children {
        send_signal_to_processes(child_processes, &ProcessAction::Kill);
    }

    true
}

/// Check whether a process's commandline string is actually a shell or not
fn did_process_spawn_shell(pid: i32) -> Result<bool> {
    // Get the /proc representation of the child, so we can do some checks
    let process = if let Ok(process) = Process::new(pid) {
        process
    } else {
        info!(
            "Process to kill has probably just gone away. Process {}",
            pid
        );
        bail!("Process has just gone away");
    };

    // Get the root command and check whether it's actually a shell with `sh -c`.
    let mut cmdline = if let Ok(cmdline) = process.cmdline() {
        cmdline
    } else {
        info!(
            "Process to kill has probably just gone away. Process {}",
            pid
        );
        bail!("Process has just gone away");
    };

    if cmdline.len() < 3 {
        return Ok(false);
    }

    if cmdline.remove(0) != "sh" {
        return Ok(false);
    }

    if cmdline.remove(0) != "-c" {
        return Ok(false);
    }

    Ok(true)
}

/// Send a signal to a unix process.
fn send_signal_to_process(
    pid: i32,
    action: &ProcessAction,
    send_to_children: bool,
) -> Result<bool, nix::Error> {
    let signal = get_signal_from_action(action);
    debug!("Sending signal {} to {}", signal, pid);

    // Send the signal to all children, if that's what the user wants.
    if send_to_children {
        send_signal_to_children(pid, action);
    }

    signal::kill(Pid::from_raw(pid), signal)?;
    Ok(true)
}

/// A small helper that sends a signal to all children of a specific process by id.
fn send_signal_to_children(pid: i32, action: &ProcessAction) {
    send_signal_to_processes(get_child_processes(pid), action);
}

/// Send a signal to a list of processes.
/// This is a convenience wrapper around `send_signal_to_process`.
fn send_signal_to_processes(processes: Vec<Process>, action: &ProcessAction) {
    let signal = get_signal_from_action(action);
    for process in processes {
        // Process is no longer alive, skip this one.
        if !process.is_alive() {
            continue;
        }

        let pid = Pid::from_raw(process.pid);
        if let Err(error) = signal::kill(pid, signal) {
            warn!(
                "Failed send signal {:?} to Pid {}: {:?}",
                signal, process.pid, error
            );
        }
    }
}

/// Get all children of a specific process
fn get_child_processes(pid: i32) -> Vec<Process> {
    let all_processes = match all_processes() {
        Err(error) => {
            warn!("Failed to get full process list: {}", error);
            return Vec::new();
        }
        Ok(processes) => processes,
    };

    all_processes
        .into_iter()
        .filter(|process| process.stat.ppid == pid)
        .collect()
}

fn get_signal_from_action(action: &ProcessAction) -> Signal {
    match action {
        ProcessAction::Kill => Signal::SIGKILL,
        ProcessAction::Pause => Signal::SIGSTOP,
        ProcessAction::Resume => Signal::SIGCONT,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::thread::sleep;
    use std::time::Duration;

    /// Assert that certain process id no longer exists
    fn process_is_gone(pid: i32) -> bool {
        match Process::new(pid) {
            Ok(process) => !process.is_alive(),
            Err(_) => true,
        }
    }

    #[test]
    fn test_spawn_command() {
        let mut child = compile_shell_command("echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");

        let ecode = child.wait().expect("failed to wait on echo");

        assert!(ecode.success());
    }

    #[test]
    /// Ensure a `sh -c` command will be properly killed without detached processes.
    fn test_shell_command_is_killed() {
        let mut child = compile_shell_command("sleep 60 & sleep 60 && echo 'this is a test'")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process indeed spawned a shell.
        assert!(did_process_spawn_shell(pid).unwrap());

        // Get all child processes, so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }

    #[test]
    /// Ensure that a `sh -c` process with a child process that has children of its own
    /// will properly kill all processes and their children's children without detached processes.
    fn test_shell_command_children_are_killed() {
        let mut child = compile_shell_command("bash -c 'sleep 60 && sleep 60' && sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process indeed spawned a shell.
        assert!(did_process_spawn_shell(pid).unwrap());

        // Get all child processes and all childrens children,
        // so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 1);
        let mut childrens_children = Vec::new();
        for child_process in &child_processes {
            childrens_children.extend(get_child_processes(child_process.stat.pid));
        }
        assert_eq!(childrens_children.len(), 1);

        // Kill the process and make sure its childen will be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }

        // Assert that all children's child processes have been killed.
        for child_process in childrens_children {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }

    #[test]
    /// Ensure a normal command without `sh -c` will be killed.
    fn test_normal_command_is_killed() {
        let mut child = Command::new("sleep")
            .arg("60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process did not spawn a shell.
        assert!(!did_process_spawn_shell(pid).unwrap());

        // No child processes exist
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 0);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, false));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        assert!(process_is_gone(pid));
    }

    #[test]
    /// Ensure a normal command and all its children will be
    /// properly killed without any detached processes.
    fn test_normal_command_children_are_killed() {
        let mut child = Command::new("bash")
            .arg("-c")
            .arg("sleep 60 & sleep 60 && sleep 60")
            .spawn()
            .expect("Failed to spawn echo");
        let pid: i32 = child.id().try_into().unwrap();
        // Sleep a little to give everything a chance to spawn.
        sleep(Duration::from_millis(500));

        // Make sure the process indeed spawned a shell.
        assert!(!did_process_spawn_shell(pid).unwrap());

        // Get all child processes, so we can make sure they no longer exist afterwards.
        let child_processes = get_child_processes(pid);
        assert_eq!(child_processes.len(), 2);

        // Kill the process and make sure it'll be killed.
        assert!(kill_child(0, &mut child, true));

        // Sleep a little to give all processes time to shutdown.
        sleep(Duration::from_millis(500));

        // Assert that the direct child (sh -c) has been killed.
        assert!(process_is_gone(pid));

        // Assert that all child processes have been killed.
        for child_process in child_processes {
            assert!(process_is_gone(child_process.stat.pid));
        }
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::marker::PhantomData;

use hyper::method::Method;
use serde::json;

use response;
use request::RequestBuilder;
use request::DoRequest;

doapi_enum! {
    #[derive(Debug)]
    pub enum DnsRecType {
        A,
        AAAA,
        CNAME,
        MX,
        NS,
        SRV,
        TXT
    }
}

#[derive(Serialize)]
pub struct DnsRecord {
    #[serde(rename = "type")]
    pub rec_type: String,
    pub name: Option<String>,
    pub priority: Option<u64>,
    pub port: Option<u64>,
    pub data: Option<String>,
    pub weight: Option<u64>,
}

impl fmt::Display for DnsRecord {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "Record Type: {}\n\
             Name: {}\n\
             Data: {}\n\
             Priority: {}\n\
             Port: {}\n\
             Weight: {}\n",
             self.rec_type,
             if let Some(n) = self.name.clone() {
                n
             } else {
                "None".to_owned()
             },
             if let Some(d) = self.data.clone() {
                d
             } else {
                "None".to_owned()
             },
             if let Some(p) = self.priority {
                p.to_string()
             } else {
                "None".to_owned()
             },
             if let Some(p) = self.port {
                p.to_string()
             } else {
                "None".to_owned()
             },
             if let Some(w) = self.weight {
                w.to_string()
             } else {
                "None".to_owned()
             }
        )
    }
}

impl<'t> RequestBuilder<'t, response::DnsRecords> {
    pub fn create(self, rec: &DnsRecord) -> RequestBuilder<'t, response::DnsRecord> {
        // POST: "https://api.digitalocean.com/v2/domains/$DOMAIN/records"
        // body:
        //      "type" : "MX"           // All records
        //      "name" : "alias"        // A, AAAA, CNAME, TXT, SRV
        //      "data" : "varies"       // A, AAAA, CNAME, MX, TXT, SRV, NS
        //      "priority" : "20"       // MX, SRV
        //      "port" : "80"           // SRV
        //      "weight" : "200"        // SRV

        // FIXME: Don't unwrap()
        RequestBuilder {
            method: Method::Post,
            auth: self.auth,
            url: self.url,
            resp_t: PhantomData,
            body: Some(json::to_string(rec).ok().unwrap())
        }

        // let mut hm = HashMap::new();
        // hm.insert("type", rec.rec_type.to_string());
        // if let Some(ref n) = rec.name {
        //     hm.insert("name", n.to_owned());
        // }
        // if let Some(ref d) = rec.data {
        //     hm.insert("data", d.to_owned());
        // }
        // if let Some(p) = rec.priority {
        //     hm.insert("priority", p.to_string());
        // }
        // if let Some(p) = rec.port {
        //     hm.insert("name", p.to_string());
        // }
        // if let Some(w) = rec.weight {
        //     hm.insert("name", w.to_string());
        // }
    }
}

impl<'t> RequestBuilder<'t, response::DnsRecord> {
    pub fn update(mut self, id: &str, record: &DnsRecord) -> RequestBuilder<'t, response::DnsRecord> {
        // PUT: "https://api.digitalocean.com/v2/domains/$DOMAIN/records/$ID"
        // body:
        //      "type" : "MX"           // All records
        //      "name" : "alias"        // A, AAAA, CNAME, TXT, SRV
        //      "data" : "varies"       // A, AAAA, CNAME, MX, TXT, SRV, NS
        //      "priority" : "20"       // MX, SRV
        //      "port" : "80"           // SRV
        //      "weight" : "200"        // SRV
        self.url.push('/');
        self.url.push_str(id);
        // FIXME: Don't unwrap()
        RequestBuilder {
            method: Method::Put,
            auth: self.auth,
            url: self.url,
            resp_t: PhantomData,
            body: Some(json::to_string(record).ok().unwrap())
        }

        // let mut hm = HashMap::new();
        // hm.insert("type", record.rec_type.to_string());
        // if let Some(ref n) = record.name {
        //     hm.insert("name", n.to_owned());
        // }
        // if let Some(ref d) = record.data {
        //     hm.insert("data", d.to_owned());
        // }
        // if let Some(p) = record.priority {
        //     hm.insert("priority", p.to_string());
        // }
        // if let Some(p) = record.port {
        //     hm.insert("name", p.to_string());
        // }
        // if let Some(w) = record.weight {
        //     hm.insert("name", w.to_string());
        // }
    }
    pub fn delete(self) -> RequestBuilder<'t, response::HeaderOnly> {
        // DELETE: "https://api.digitalocean.com/v2/domains/$id"
        RequestBuilder {
            method: Method::Delete,
            auth: self.auth,
            url: self.url,
            resp_t: PhantomData,
            body: None
        }
    }
}

impl<'t> DoRequest<response::DnsRecord> for RequestBuilder<'t, response::DnsRecord> {}
<commit_msg>fix(DNS Records): fixes bug with duplicate record IDs<commit_after>use std::fmt;
use std::marker::PhantomData;

use hyper::method::Method;
use serde::json;

use response;
use request::RequestBuilder;
use request::DoRequest;

doapi_enum! {
    #[derive(Debug)]
    pub enum DnsRecType {
        A,
        AAAA,
        CNAME,
        MX,
        NS,
        SRV,
        TXT
    }
}

#[derive(Serialize)]
pub struct DnsRecord {
    #[serde(rename = "type")]
    pub rec_type: String,
    pub name: Option<String>,
    pub priority: Option<u64>,
    pub port: Option<u64>,
    pub data: Option<String>,
    pub weight: Option<u64>,
}

impl fmt::Display for DnsRecord {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f,
            "Record Type: {}\n\
             Name: {}\n\
             Data: {}\n\
             Priority: {}\n\
             Port: {}\n\
             Weight: {}\n",
             self.rec_type,
             if let Some(n) = self.name.clone() {
                n
             } else {
                "None".to_owned()
             },
             if let Some(d) = self.data.clone() {
                d
             } else {
                "None".to_owned()
             },
             if let Some(p) = self.priority {
                p.to_string()
             } else {
                "None".to_owned()
             },
             if let Some(p) = self.port {
                p.to_string()
             } else {
                "None".to_owned()
             },
             if let Some(w) = self.weight {
                w.to_string()
             } else {
                "None".to_owned()
             }
        )
    }
}

impl<'t> RequestBuilder<'t, response::DnsRecords> {
    pub fn create(self, rec: &DnsRecord) -> RequestBuilder<'t, response::DnsRecord> {
        // POST: "https://api.digitalocean.com/v2/domains/$DOMAIN/records"
        // body:
        //      "type" : "MX"           // All records
        //      "name" : "alias"        // A, AAAA, CNAME, TXT, SRV
        //      "data" : "varies"       // A, AAAA, CNAME, MX, TXT, SRV, NS
        //      "priority" : "20"       // MX, SRV
        //      "port" : "80"           // SRV
        //      "weight" : "200"        // SRV

        // FIXME: Don't unwrap()
        RequestBuilder {
            method: Method::Post,
            auth: self.auth,
            url: self.url,
            resp_t: PhantomData,
            body: Some(json::to_string(rec).ok().unwrap())
        }

        // let mut hm = HashMap::new();
        // hm.insert("type", rec.rec_type.to_string());
        // if let Some(ref n) = rec.name {
        //     hm.insert("name", n.to_owned());
        // }
        // if let Some(ref d) = rec.data {
        //     hm.insert("data", d.to_owned());
        // }
        // if let Some(p) = rec.priority {
        //     hm.insert("priority", p.to_string());
        // }
        // if let Some(p) = rec.port {
        //     hm.insert("name", p.to_string());
        // }
        // if let Some(w) = rec.weight {
        //     hm.insert("name", w.to_string());
        // }
    }
}

impl<'t> RequestBuilder<'t, response::DnsRecord> {
    pub fn update(mut self, record: &DnsRecord) -> RequestBuilder<'t, response::DnsRecord> {
        // PUT: "https://api.digitalocean.com/v2/domains/$DOMAIN/records/$ID"
        // body:
        //      "type" : "MX"           // All records
        //      "name" : "alias"        // A, AAAA, CNAME, TXT, SRV
        //      "data" : "varies"       // A, AAAA, CNAME, MX, TXT, SRV, NS
        //      "priority" : "20"       // MX, SRV
        //      "port" : "80"           // SRV
        //      "weight" : "200"        // SRV
        // FIXME: Don't unwrap()
        RequestBuilder {
            method: Method::Put,
            auth: self.auth,
            url: self.url,
            resp_t: PhantomData,
            body: Some(json::to_string(record).ok().unwrap())
        }

        // let mut hm = HashMap::new();
        // hm.insert("type", record.rec_type.to_string());
        // if let Some(ref n) = record.name {
        //     hm.insert("name", n.to_owned());
        // }
        // if let Some(ref d) = record.data {
        //     hm.insert("data", d.to_owned());
        // }
        // if let Some(p) = record.priority {
        //     hm.insert("priority", p.to_string());
        // }
        // if let Some(p) = record.port {
        //     hm.insert("name", p.to_string());
        // }
        // if let Some(w) = record.weight {
        //     hm.insert("name", w.to_string());
        // }
    }
    pub fn delete(self) -> RequestBuilder<'t, response::HeaderOnly> {
        // DELETE: "https://api.digitalocean.com/v2/domains/$id"
        RequestBuilder {
            method: Method::Delete,
            auth: self.auth,
            url: self.url,
            resp_t: PhantomData,
            body: None
        }
    }
}

impl<'t> DoRequest<response::DnsRecord> for RequestBuilder<'t, response::DnsRecord> {}
<|endoftext|>