<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Boden Garman <bpgarman@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

#[macro_use]
extern crate uucore;

mod count_bytes;
mod countable;
mod word_count;
use count_bytes::count_bytes_fast;
use countable::WordCountable;
use word_count::{TitledWordCount, WordCount};

use clap::{crate_version, App, Arg, ArgMatches};
use thiserror::Error;

use std::fs::{self, File};
use std::io::{self, ErrorKind, Write};
use std::path::Path;

/// The minimum character width for formatting counts when reading from stdin.
const MINIMUM_WIDTH: usize = 7;

#[derive(Error, Debug)]
pub enum WcError {
    #[error("{0}")]
    Io(#[from] io::Error),
    #[error("Expected a file, found directory {0}")]
    IsDirectory(String),
}

type WcResult<T> = Result<T, WcError>;

struct Settings {
    show_bytes: bool,
    show_chars: bool,
    show_lines: bool,
    show_words: bool,
    show_max_line_length: bool,
}

impl Settings {
    fn new(matches: &ArgMatches) -> Settings {
        let settings = Settings {
            show_bytes: matches.is_present(options::BYTES),
            show_chars: matches.is_present(options::CHAR),
            show_lines: matches.is_present(options::LINES),
            show_words: matches.is_present(options::WORDS),
            show_max_line_length: matches.is_present(options::MAX_LINE_LENGTH),
        };

        if settings.show_bytes
            || settings.show_chars
            || settings.show_lines
            || settings.show_words
            || settings.show_max_line_length
        {
            return settings;
        }

        Settings {
            show_bytes: true,
            show_chars: false,
            show_lines: true,
            show_words: true,
            show_max_line_length: false,
        }
    }

    fn number_enabled(&self) -> u32 {
        let mut result = 0;
        result += self.show_bytes as u32;
        result += self.show_chars as u32;
        result += self.show_lines as u32;
        result += self.show_max_line_length as u32;
        result += self.show_words as u32;
        result
    }
}

static ABOUT: &str = "Display newline, word, and byte counts for each FILE, and a total line if
more than one FILE is specified.";

pub mod options {
    pub static BYTES: &str = "bytes";
    pub static CHAR: &str = "chars";
    pub static LINES: &str = "lines";
    pub static MAX_LINE_LENGTH: &str = "max-line-length";
    pub static WORDS: &str = "words";
}

static ARG_FILES: &str = "files";

fn usage() -> String {
    format!(
        "{0} [OPTION]... [FILE]...
 With no FILE, or when FILE is -, read standard input.",
        uucore::execution_phrase()
    )
}

enum StdinKind {
    /// Stdin specified on command-line with "-".
    Explicit,

    /// Stdin implicitly specified on command-line by not passing any positional argument.
    Implicit,
}

/// Supported inputs.
enum Input {
    /// A regular file.
    Path(String),

    /// Standard input.
    Stdin(StdinKind),
}

impl Input {
    /// Converts input to title that appears in stats.
    fn to_title(&self) -> Option<&str> {
        match self {
            Input::Path(path) => Some(path),
            Input::Stdin(StdinKind::Explicit) => Some("-"),
            Input::Stdin(StdinKind::Implicit) => None,
        }
    }
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = usage();

    let matches = uu_app().usage(&usage[..]).get_matches_from(args);

    let mut inputs: Vec<Input> = matches
        .values_of(ARG_FILES)
        .map(|v| {
            v.map(|i| {
                if i == "-" {
                    Input::Stdin(StdinKind::Explicit)
                } else {
                    Input::Path(ToString::to_string(i))
                }
            })
            .collect()
        })
        .unwrap_or_default();

    if inputs.is_empty() {
        inputs.push(Input::Stdin(StdinKind::Implicit));
    }

    let settings = Settings::new(&matches);

    if wc(inputs, &settings).is_ok() {
        0
    } else {
        1
    }
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .arg(
            Arg::with_name(options::BYTES)
                .short("c")
                .long(options::BYTES)
                .help("print the byte counts"),
        )
        .arg(
            Arg::with_name(options::CHAR)
                .short("m")
                .long(options::CHAR)
                .help("print the character counts"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("l")
                .long(options::LINES)
                .help("print the newline counts"),
        )
        .arg(
            Arg::with_name(options::MAX_LINE_LENGTH)
                .short("L")
                .long(options::MAX_LINE_LENGTH)
                .help("print the length of the longest line"),
        )
        .arg(
            Arg::with_name(options::WORDS)
                .short("w")
                .long(options::WORDS)
                .help("print the word counts"),
        )
        .arg(Arg::with_name(ARG_FILES).multiple(true).takes_value(true))
}

fn word_count_from_reader<T: WordCountable>(
    mut reader: T,
    settings: &Settings,
    path: &str,
) -> WcResult<WordCount> {
    let only_count_bytes = settings.show_bytes
        && (!(settings.show_chars
            || settings.show_lines
            || settings.show_max_line_length
            || settings.show_words));
    if only_count_bytes {
        return Ok(WordCount {
            bytes: count_bytes_fast(&mut reader)?,
            ..WordCount::default()
        });
    }

    // we do not need to decode the byte stream if we're only counting bytes/newlines
    let decode_chars = settings.show_chars || settings.show_words || settings.show_max_line_length;

    // Sum the WordCount for each line. Show a warning for each line
    // that results in an IO error when trying to read it.
    let total = reader
        .lines()
        .filter_map(|res| match res {
            Ok(line) => Some(line),
            Err(e) => {
                show_warning!("Error while reading {}: {}", path, e);
                None
            }
        })
        .map(|line| WordCount::from_line(&line, decode_chars))
        .sum();
    Ok(total)
}

fn word_count_from_input(input: &Input, settings: &Settings) -> WcResult<WordCount> {
    match input {
        Input::Stdin(_) => {
            let stdin = io::stdin();
            let stdin_lock = stdin.lock();
            word_count_from_reader(stdin_lock, settings, "-")
        }
        Input::Path(path) => {
            let path_obj = Path::new(path);
            if path_obj.is_dir() {
                Err(WcError::IsDirectory(path.to_owned()))
            } else {
                let file = File::open(path)?;
                word_count_from_reader(file, settings, path)
            }
        }
    }
}

/// Print a message appropriate for the particular error to `stderr`.
///
/// # Examples
///
/// This will print `wc: /tmp: Is a directory` to `stderr`.
///
/// ```rust,ignore
/// show_error(Input::Path("/tmp"), WcError::IsDirectory("/tmp"))
/// ```
fn show_error(input: &Input, err: WcError) {
    match (input, err) {
        (_, WcError::IsDirectory(path)) => {
            show_error_custom_description!(path, "Is a directory");
        }
        (Input::Path(path), WcError::Io(e)) if e.kind() == ErrorKind::NotFound => {
            show_error_custom_description!(path, "No such file or directory");
        }
        (_, e) => {
            show_error!("{}", e);
        }
    };
}

/// Compute the number of digits needed to represent any count for this input.
///
/// If `input` is [`Input::Stdin`], then this function returns
/// [`MINIMUM_WIDTH`]. Otherwise, if metadata could not be read from
/// `input` then this function returns 1.
///
/// # Errors
///
/// This function will return an error if `input` is a [`Input::Path`]
/// and there is a problem accessing the metadata of the given `input`.
///
/// # Examples
///
/// A [`Input::Stdin`] gets a default minimum width:
///
/// ```rust,ignore
/// let input = Input::Stdin(StdinKind::Explicit);
/// assert_eq!(7, digit_width(input));
/// ```
fn digit_width(input: &Input) -> WcResult<Option<usize>> {
    match input {
        Input::Stdin(_) => Ok(Some(MINIMUM_WIDTH)),
        Input::Path(filename) => {
            let path = Path::new(filename);
            let metadata = fs::metadata(path)?;
            if metadata.is_file() {
                // TODO We are now computing the number of bytes in a file
                // twice: once here and once in `WordCount::from_line()` (or
                // in `count_bytes_fast()` if that function is called
                // instead). See GitHub issue #2201.
                let num_bytes = metadata.len();
                let num_digits = num_bytes.to_string().len();
                Ok(Some(num_digits))
            } else {
                Ok(None)
            }
        }
    }
}

/// Compute the number of digits needed to represent all counts in all inputs.
///
/// `inputs` may include zero or more [`Input::Stdin`] entries, each of
/// which represents reading from `stdin`. The presence of any such
/// entry causes this function to return a width that is at least
/// [`MINIMUM_WIDTH`].
///
/// If `input` is empty, then this function returns 1. If file metadata
/// could not be read from any of the [`Input::Path`] inputs and there
/// are no [`Input::Stdin`] inputs, then this function returns 1.
///
/// If there is a problem accessing the metadata, this function will
/// silently ignore the error and assume that the number of digits
/// needed to display the counts for that file is 1.
///
/// # Examples
///
/// An empty slice implies a width of 1:
///
/// ```rust,ignore
/// assert_eq!(1, max_width(&vec![]));
/// ```
///
/// The presence of [`Input::Stdin`] implies a minimum width:
///
/// ```rust,ignore
/// let inputs = vec![Input::Stdin(StdinKind::Explicit)];
/// assert_eq!(7, max_width(&inputs));
/// ```
fn max_width(inputs: &[Input]) -> usize {
    let mut result = 1;
    for input in inputs {
        match digit_width(input) {
            Ok(maybe_n) => {
                if let Some(n) = maybe_n {
                    result = result.max(n);
                }
            }
            Err(_) => continue,
        }
    }
    result
}

fn wc(inputs: Vec<Input>, settings: &Settings) -> Result<(), u32> {
    // Compute the width, in digits, to use when formatting counts.
    //
    // The width is the number of digits needed to print the number of
    // bytes in the largest file. This is true regardless of whether
    // the `settings` indicate that the bytes will be displayed.
    let mut error_count = 0;
    let max_width = max_width(&inputs);

    let mut total_word_count = WordCount::default();

    let num_inputs = inputs.len();

    for input in &inputs {
        let word_count = word_count_from_input(input, settings).unwrap_or_else(|err| {
            show_error(input, err);
            error_count += 1;
            WordCount::default()
        });
        total_word_count += word_count;
        let result = word_count.with_title(input.to_title());
        if let Err(err) = print_stats(settings, &result, max_width) {
            show_warning!(
                "failed to print result for {}: {}",
                result.title.unwrap_or("<stdin>"),
                err
            );
            error_count += 1;
        }
    }

    if num_inputs > 1 {
        let total_result = total_word_count.with_title(Some("total"));
        if let Err(err) = print_stats(settings, &total_result, max_width) {
            show_warning!("failed to print total: {}", err);
            error_count += 1;
        }
    }

    if error_count == 0 {
        Ok(())
    } else {
        Err(error_count)
    }
}

fn print_stats(
    settings: &Settings,
    result: &TitledWordCount,
    mut min_width: usize,
) -> WcResult<()> {
    let stdout = io::stdout();
    let mut stdout_lock = stdout.lock();

    if settings.number_enabled() <= 1 {
        // Prevent a leading space in case we only need to display a single
        // number.
        min_width = 0;
    }

    let mut is_first: bool = true;

    if settings.show_lines {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.lines, min_width)?;
        is_first = false;
    }
    if settings.show_words {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.words, min_width)?;
        is_first = false;
    }
    if settings.show_bytes {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.bytes, min_width)?;
        is_first = false;
    }
    if settings.show_chars {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.chars, min_width)?;
        is_first = false;
    }
    if settings.show_max_line_length {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(
            stdout_lock,
            "{:1$}",
            result.count.max_line_length, min_width
        )?;
    }

    if let Some(title) = result.title {
        writeln!(stdout_lock, " {}", title)?;
    } else {
        writeln!(stdout_lock)?;
    }

    Ok(())
}
<commit_msg>wc: fix clippy lint<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Boden Garman <bpgarman@gmail.com>
//  *
//  * For the full copyright and license information, please view the LICENSE
//  * file that was distributed with this source code.

#[macro_use]
extern crate uucore;

mod count_bytes;
mod countable;
mod word_count;
use count_bytes::count_bytes_fast;
use countable::WordCountable;
use word_count::{TitledWordCount, WordCount};

use clap::{crate_version, App, Arg, ArgMatches};
use thiserror::Error;

use std::fs::{self, File};
use std::io::{self, ErrorKind, Write};
use std::path::Path;

/// The minimum character width for formatting counts when reading from stdin.
const MINIMUM_WIDTH: usize = 7;

#[derive(Error, Debug)]
pub enum WcError {
    #[error("{0}")]
    Io(#[from] io::Error),
    #[error("Expected a file, found directory {0}")]
    IsDirectory(String),
}

type WcResult<T> = Result<T, WcError>;

struct Settings {
    show_bytes: bool,
    show_chars: bool,
    show_lines: bool,
    show_words: bool,
    show_max_line_length: bool,
}

impl Settings {
    fn new(matches: &ArgMatches) -> Settings {
        let settings = Settings {
            show_bytes: matches.is_present(options::BYTES),
            show_chars: matches.is_present(options::CHAR),
            show_lines: matches.is_present(options::LINES),
            show_words: matches.is_present(options::WORDS),
            show_max_line_length: matches.is_present(options::MAX_LINE_LENGTH),
        };

        if settings.show_bytes
            || settings.show_chars
            || settings.show_lines
            || settings.show_words
            || settings.show_max_line_length
        {
            return settings;
        }

        Settings {
            show_bytes: true,
            show_chars: false,
            show_lines: true,
            show_words: true,
            show_max_line_length: false,
        }
    }

    fn number_enabled(&self) -> u32 {
        let mut result = 0;
        result += self.show_bytes as u32;
        result += self.show_chars as u32;
        result += self.show_lines as u32;
        result += self.show_max_line_length as u32;
        result += self.show_words as u32;
        result
    }
}

static ABOUT: &str = "Display newline, word, and byte counts for each FILE, and a total line if
more than one FILE is specified.";

pub mod options {
    pub static BYTES: &str = "bytes";
    pub static CHAR: &str = "chars";
    pub static LINES: &str = "lines";
    pub static MAX_LINE_LENGTH: &str = "max-line-length";
    pub static WORDS: &str = "words";
}

static ARG_FILES: &str = "files";

fn usage() -> String {
    format!(
        "{0} [OPTION]... [FILE]...
 With no FILE, or when FILE is -, read standard input.",
        uucore::execution_phrase()
    )
}

enum StdinKind {
    /// Stdin specified on command-line with "-".
    Explicit,

    /// Stdin implicitly specified on command-line by not passing any positional argument.
    Implicit,
}

/// Supported inputs.
enum Input {
    /// A regular file.
    Path(String),

    /// Standard input.
    Stdin(StdinKind),
}

impl Input {
    /// Converts input to title that appears in stats.
    fn to_title(&self) -> Option<&str> {
        match self {
            Input::Path(path) => Some(path),
            Input::Stdin(StdinKind::Explicit) => Some("-"),
            Input::Stdin(StdinKind::Implicit) => None,
        }
    }
}

pub fn uumain(args: impl uucore::Args) -> i32 {
    let usage = usage();

    let matches = uu_app().usage(&usage[..]).get_matches_from(args);

    let mut inputs: Vec<Input> = matches
        .values_of(ARG_FILES)
        .map(|v| {
            v.map(|i| {
                if i == "-" {
                    Input::Stdin(StdinKind::Explicit)
                } else {
                    Input::Path(ToString::to_string(i))
                }
            })
            .collect()
        })
        .unwrap_or_default();

    if inputs.is_empty() {
        inputs.push(Input::Stdin(StdinKind::Implicit));
    }

    let settings = Settings::new(&matches);

    if wc(inputs, &settings).is_ok() {
        0
    } else {
        1
    }
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(uucore::util_name())
        .version(crate_version!())
        .about(ABOUT)
        .arg(
            Arg::with_name(options::BYTES)
                .short("c")
                .long(options::BYTES)
                .help("print the byte counts"),
        )
        .arg(
            Arg::with_name(options::CHAR)
                .short("m")
                .long(options::CHAR)
                .help("print the character counts"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("l")
                .long(options::LINES)
                .help("print the newline counts"),
        )
        .arg(
            Arg::with_name(options::MAX_LINE_LENGTH)
                .short("L")
                .long(options::MAX_LINE_LENGTH)
                .help("print the length of the longest line"),
        )
        .arg(
            Arg::with_name(options::WORDS)
                .short("w")
                .long(options::WORDS)
                .help("print the word counts"),
        )
        .arg(Arg::with_name(ARG_FILES).multiple(true).takes_value(true))
}

fn word_count_from_reader<T: WordCountable>(
    mut reader: T,
    settings: &Settings,
    path: &str,
) -> WcResult<WordCount> {
    let only_count_bytes = settings.show_bytes
        && (!(settings.show_chars
            || settings.show_lines
            || settings.show_max_line_length
            || settings.show_words));
    if only_count_bytes {
        return Ok(WordCount {
            bytes: count_bytes_fast(&mut reader)?,
            ..WordCount::default()
        });
    }

    // we do not need to decode the byte stream if we're only counting bytes/newlines
    let decode_chars = settings.show_chars || settings.show_words || settings.show_max_line_length;

    // Sum the WordCount for each line. Show a warning for each line
    // that results in an IO error when trying to read it.
    let total = reader
        .lines()
        .filter_map(|res| match res {
            Ok(line) => Some(line),
            Err(e) => {
                show_warning!("Error while reading {}: {}", path, e);
                None
            }
        })
        .map(|line| WordCount::from_line(&line, decode_chars))
        .sum();
    Ok(total)
}

fn word_count_from_input(input: &Input, settings: &Settings) -> WcResult<WordCount> {
    match input {
        Input::Stdin(_) => {
            let stdin = io::stdin();
            let stdin_lock = stdin.lock();
            word_count_from_reader(stdin_lock, settings, "-")
        }
        Input::Path(path) => {
            let path_obj = Path::new(path);
            if path_obj.is_dir() {
                Err(WcError::IsDirectory(path.to_owned()))
            } else {
                let file = File::open(path)?;
                word_count_from_reader(file, settings, path)
            }
        }
    }
}

/// Print a message appropriate for the particular error to `stderr`.
///
/// # Examples
///
/// This will print `wc: /tmp: Is a directory` to `stderr`.
///
/// ```rust,ignore
/// show_error(Input::Path("/tmp"), WcError::IsDirectory("/tmp"))
/// ```
fn show_error(input: &Input, err: WcError) {
    match (input, err) {
        (_, WcError::IsDirectory(path)) => {
            show_error_custom_description!(path, "Is a directory");
        }
        (Input::Path(path), WcError::Io(e)) if e.kind() == ErrorKind::NotFound => {
            show_error_custom_description!(path, "No such file or directory");
        }
        (_, e) => {
            show_error!("{}", e);
        }
    };
}

/// Compute the number of digits needed to represent any count for this input.
///
/// If `input` is [`Input::Stdin`], then this function returns
/// [`MINIMUM_WIDTH`]. Otherwise, if metadata could not be read from
/// `input` then this function returns 1.
///
/// # Errors
///
/// This function will return an error if `input` is a [`Input::Path`]
/// and there is a problem accessing the metadata of the given `input`.
///
/// # Examples
///
/// A [`Input::Stdin`] gets a default minimum width:
///
/// ```rust,ignore
/// let input = Input::Stdin(StdinKind::Explicit);
/// assert_eq!(7, digit_width(input));
/// ```
fn digit_width(input: &Input) -> WcResult<Option<usize>> {
    match input {
        Input::Stdin(_) => Ok(Some(MINIMUM_WIDTH)),
        Input::Path(filename) => {
            let path = Path::new(filename);
            let metadata = fs::metadata(path)?;
            if metadata.is_file() {
                // TODO We are now computing the number of bytes in a file
                // twice: once here and once in `WordCount::from_line()` (or
                // in `count_bytes_fast()` if that function is called
                // instead). See GitHub issue #2201.
                let num_bytes = metadata.len();
                let num_digits = num_bytes.to_string().len();
                Ok(Some(num_digits))
            } else {
                Ok(None)
            }
        }
    }
}

/// Compute the number of digits needed to represent all counts in all inputs.
///
/// `inputs` may include zero or more [`Input::Stdin`] entries, each of
/// which represents reading from `stdin`. The presence of any such
/// entry causes this function to return a width that is at least
/// [`MINIMUM_WIDTH`].
///
/// If `input` is empty, then this function returns 1. If file metadata
/// could not be read from any of the [`Input::Path`] inputs and there
/// are no [`Input::Stdin`] inputs, then this function returns 1.
///
/// If there is a problem accessing the metadata, this function will
/// silently ignore the error and assume that the number of digits
/// needed to display the counts for that file is 1.
///
/// # Examples
///
/// An empty slice implies a width of 1:
///
/// ```rust,ignore
/// assert_eq!(1, max_width(&vec![]));
/// ```
///
/// The presence of [`Input::Stdin`] implies a minimum width:
///
/// ```rust,ignore
/// let inputs = vec![Input::Stdin(StdinKind::Explicit)];
/// assert_eq!(7, max_width(&inputs));
/// ```
fn max_width(inputs: &[Input]) -> usize {
    let mut result = 1;
    for input in inputs {
        if let Ok(Some(n)) = digit_width(input) {
            result = result.max(n);
        }
    }
    result
}

fn wc(inputs: Vec<Input>, settings: &Settings) -> Result<(), u32> {
    // Compute the width, in digits, to use when formatting counts.
    //
    // The width is the number of digits needed to print the number of
    // bytes in the largest file. This is true regardless of whether
    // the `settings` indicate that the bytes will be displayed.
    let mut error_count = 0;
    let max_width = max_width(&inputs);

    let mut total_word_count = WordCount::default();

    let num_inputs = inputs.len();

    for input in &inputs {
        let word_count = word_count_from_input(input, settings).unwrap_or_else(|err| {
            show_error(input, err);
            error_count += 1;
            WordCount::default()
        });
        total_word_count += word_count;
        let result = word_count.with_title(input.to_title());
        if let Err(err) = print_stats(settings, &result, max_width) {
            show_warning!(
                "failed to print result for {}: {}",
                result.title.unwrap_or("<stdin>"),
                err
            );
            error_count += 1;
        }
    }

    if num_inputs > 1 {
        let total_result = total_word_count.with_title(Some("total"));
        if let Err(err) = print_stats(settings, &total_result, max_width) {
            show_warning!("failed to print total: {}", err);
            error_count += 1;
        }
    }

    if error_count == 0 {
        Ok(())
    } else {
        Err(error_count)
    }
}

fn print_stats(
    settings: &Settings,
    result: &TitledWordCount,
    mut min_width: usize,
) -> WcResult<()> {
    let stdout = io::stdout();
    let mut stdout_lock = stdout.lock();

    if settings.number_enabled() <= 1 {
        // Prevent a leading space in case we only need to display a single
        // number.
        min_width = 0;
    }

    let mut is_first: bool = true;

    if settings.show_lines {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.lines, min_width)?;
        is_first = false;
    }
    if settings.show_words {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.words, min_width)?;
        is_first = false;
    }
    if settings.show_bytes {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.bytes, min_width)?;
        is_first = false;
    }
    if settings.show_chars {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(stdout_lock, "{:1$}", result.count.chars, min_width)?;
        is_first = false;
    }
    if settings.show_max_line_length {
        if !is_first {
            write!(stdout_lock, " ")?;
        }
        write!(
            stdout_lock,
            "{:1$}",
            result.count.max_line_length, min_width
        )?;
    }

    if let Some(title) = result.title {
        writeln!(stdout_lock, " {}", title)?;
    } else {
        writeln!(stdout_lock)?;
    }

    Ok(())
}
<|endoftext|>


<commit_before>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::mem;
use std::time::Duration;
use std::rt::task::BlockedTask;

use green::Callback;

use {raw, uvll, EventLoop, UvResult};
use homing::{HomeHandle, HomingIO, HomingMissile};
use raw::Handle;

pub struct Timer {
    handle: raw::Timer,
    home: HomeHandle,
}

struct Data {
    action: Option<NextAction>,
    id: uint, // see comments in timer_cb
}

enum NextAction {
    WakeTask(BlockedTask),
    CallOnce(Box<Callback + Send>),
    CallMany(Box<Callback + Send>, uint),
}

impl Timer {
    /// Create a new timer on the local event loop.
    pub fn new() -> UvResult<Timer> {
        let mut eloop = try!(EventLoop::borrow());
        Timer::new_on(&mut *eloop)
    }

    /// Same as `new`, but specifies what event loop to be created on.
    pub fn new_on(eloop: &mut EventLoop) -> UvResult<Timer> {
        unsafe {
            let data = box Data {
                action: None,
                id: 0,
            };
            let mut ret = Timer {
                handle: try!(raw::Timer::new(&eloop.uv_loop())),
                home: eloop.make_handle(),
            };
            ret.handle.set_data(mem::transmute(data));
            Ok(ret)
        }
    }

    /// Sleep for a specified duration of time.
    pub fn sleep(&mut self, dur: Duration) {
        let mut ms = dur.num_milliseconds();
        if ms <= 0 { ms = 0; }

        let mut handle = self.handle;

        // As with all of the below functions, we must be extra careful when
        // destroying the previous action. If the previous action was a
        // callback, destroying it could invoke a context switch. For these
        // situations, we must temporarily un-home ourselves, then destroy the
        // action, and then re-home again.
        {
            let (m, data, _) = self.data();
            handle.stop().unwrap();
            match data.action.take() {
                Some(a) => { drop(m); drop(a); }
                None => {}
            }
        }
        let (_m, data, mut handle) = self.data();
        assert!(data.action.is_none());
        data.id += 1;
        ::block(handle.uv_loop(), |task| {
            data.action = Some(WakeTask(task));
            handle.stop().unwrap();
            handle.start(ms as u64, 0, timer_cb).unwrap();
        });
    }

    pub fn oneshot(&mut self, dur: Duration, cb: Box<Callback + Send>) {
        let mut ms = dur.num_milliseconds();
        if ms <= 0 { ms = 0; }

        let _prev = {
            let (_m, data, mut handle) = self.data();
            data.id += 1;
            handle.stop().unwrap();
            handle.start(ms as u64, 0, timer_cb).unwrap();
            mem::replace(&mut data.action, Some(CallOnce(cb)))
        };
    }

    pub fn periodic(&mut self, dur: Duration, cb: Box<Callback + Send>) {
        let mut ms = dur.num_milliseconds();
        if ms <= 0 { ms = 1; }

        let _prev = {
            let (_m, data, mut handle) = self.data();
            data.id += 1;
            handle.stop().unwrap();
            handle.start(ms as u64, ms as u64, timer_cb).unwrap();
            mem::replace(&mut data.action, Some(CallMany(cb, data.id)))
        };
    }

    fn data(&mut self) -> (HomingMissile, &mut Data, raw::Timer) {
        let m = self.fire_homing_missile();
        (m, unsafe { mem::transmute(self.handle.get_data()) }, self.handle)
    }
}

extern fn timer_cb(timer: *mut uvll::uv_timer_t) {
    let timer: raw::Timer = unsafe { Handle::from_raw(timer) };
    let data: &mut Data = unsafe { mem::transmute(timer.get_data()) };
    match data.action.take().unwrap() {
        WakeTask(task) => task.reawaken(),
        CallOnce(mut cb) => cb.call(),
        CallMany(mut cb, id) => {
            cb.call();

            // Note that the above operation could have performed some form
            // of scheduling. This means that the timer may have decided to
            // insert some other action to happen. This 'id' keeps track of
            // the updates to the timer, so we only reset the action back to
            // sending on this channel if the id has remained the same. This
            // is essentially a bug in that we have mutably aliasable
            // memory, but that's libuv for you. We're guaranteed to all be
            // running on the same thread, so there's no need for any
            // synchronization here.
            if data.id == id {
                data.action = Some(CallMany(cb, id));
            }
        }
    }
}

impl HomingIO for Timer {
    fn home<'r>(&'r mut self) -> &'r mut HomeHandle { &mut self.home }
}

impl Drop for Timer {
    fn drop(&mut self) {
        // note that this drop is a little subtle. Dropping a channel which is
        // held internally may invoke some scheduling operations. We can't take
        // the channel unless we're on the home scheduler, but once we're on the
        // home scheduler we should never move. Hence, we take the timer's
        // action item and then move it outside of the homing block.
        //
        // Furthermore, we can't actually free the extraneous `Data` here
        // because the `timer_cb` may actually be somewhere up on the stack in a
        // `CallMany` which might use the data. To mitigate this, we bump the
        // `id` (so the `CallMany` doesn't re-store its callback) and then
        // schedule a custom close callback to free both the timer and the data.
        let _action = unsafe {
            let _m = self.fire_homing_missile();
            self.handle.stop().unwrap();
            let data: &mut Data = mem::transmute(self.handle.get_data());
            data.id += 1;
            self.handle.close(Some(close_cb));
            data.action.take()
        };

        extern fn close_cb(handle: *mut uvll::uv_handle_t) {
            unsafe {
                let handle = handle as *mut uvll::uv_timer_t;
                let mut timer: raw::Timer = Handle::from_raw(handle);
                let data: Box<Data> = mem::transmute(timer.get_data());
                assert!(data.action.is_none());
                timer.free();
            }
        }
    }
}
<commit_msg>Add some documentation to timer methods<commit_after>// Copyright 2013 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::mem;
use std::time::Duration;
use std::rt::task::BlockedTask;

use green::Callback;

use {raw, uvll, EventLoop, UvResult};
use homing::{HomeHandle, HomingIO, HomingMissile};
use raw::Handle;

/// A libuv-based timer to schedule callbacks to run on an event loop.
pub struct Timer {
    handle: raw::Timer,
    home: HomeHandle,
}

struct Data {
    action: Option<NextAction>,
    id: uint, // see comments in timer_cb
}

enum NextAction {
    WakeTask(BlockedTask),
    CallOnce(Box<Callback + Send>),
    CallMany(Box<Callback + Send>, uint),
}

impl Timer {
    /// Create a new timer on the local event loop.
    pub fn new() -> UvResult<Timer> {
        let mut eloop = try!(EventLoop::borrow());
        Timer::new_on(&mut *eloop)
    }

    /// Same as `new`, but specifies what event loop to be created on.
    pub fn new_on(eloop: &mut EventLoop) -> UvResult<Timer> {
        unsafe {
            let data = box Data {
                action: None,
                id: 0,
            };
            let mut ret = Timer {
                handle: try!(raw::Timer::new(&eloop.uv_loop())),
                home: eloop.make_handle(),
            };
            ret.handle.set_data(mem::transmute(data));
            Ok(ret)
        }
    }

    /// Sleep for a specified duration of time.
    ///
    /// See [`std::io::Timer::sleep`][1] for semantic information.
    ///
    /// [1]: http://doc.rust-lang.org/std/io/timer/struct.Timer.html#method.sleep
    pub fn sleep(&mut self, dur: Duration) {
        let mut ms = dur.num_milliseconds();
        if ms <= 0 { ms = 0; }

        let mut handle = self.handle;

        // As with all of the below functions, we must be extra careful when
        // destroying the previous action. If the previous action was a
        // callback, destroying it could invoke a context switch. For these
        // situations, we must temporarily un-home ourselves, then destroy the
        // action, and then re-home again.
        {
            let (m, data, _) = self.data();
            handle.stop().unwrap();
            match data.action.take() {
                Some(a) => { drop(m); drop(a); }
                None => {}
            }
        }
        let (_m, data, mut handle) = self.data();
        assert!(data.action.is_none());
        data.id += 1;
        ::block(handle.uv_loop(), |task| {
            data.action = Some(WakeTask(task));
            handle.stop().unwrap();
            handle.start(ms as u64, 0, timer_cb).unwrap();
        });
    }

    /// Schedule a callback to be run once after the specified duration has
    /// elapsed.
    ///
    /// See [`std::io::Timer::oneshot`][1] for semantic information.
    ///
    /// [1]: http://doc.rust-lang.org/std/io/timer/struct.Timer.html#method.oneshot
    pub fn oneshot(&mut self, dur: Duration, cb: Box<Callback + Send>) {
        let mut ms = dur.num_milliseconds();
        if ms <= 0 { ms = 0; }

        let _prev = {
            let (_m, data, mut handle) = self.data();
            data.id += 1;
            handle.stop().unwrap();
            handle.start(ms as u64, 0, timer_cb).unwrap();
            mem::replace(&mut data.action, Some(CallOnce(cb)))
        };
    }

    /// Schedule a callback to be run after each `dur` amount of time that
    /// passes.
    ///
    /// See [`std::io::Timer::periodic`][1] for semantic information.
    ///
    /// [1]: http://doc.rust-lang.org/std/io/timer/struct.Timer.html#method.periodic
    pub fn periodic(&mut self, dur: Duration, cb: Box<Callback + Send>) {
        let mut ms = dur.num_milliseconds();
        if ms <= 0 { ms = 1; }

        let _prev = {
            let (_m, data, mut handle) = self.data();
            data.id += 1;
            handle.stop().unwrap();
            handle.start(ms as u64, ms as u64, timer_cb).unwrap();
            mem::replace(&mut data.action, Some(CallMany(cb, data.id)))
        };
    }

    fn data(&mut self) -> (HomingMissile, &mut Data, raw::Timer) {
        let m = self.fire_homing_missile();
        (m, unsafe { mem::transmute(self.handle.get_data()) }, self.handle)
    }

    /// Gain access to the underlying raw timer handle.
    ///
    /// This function is unsafe as there is no guarantee that any safe
    /// modifications to the timer handle are actually safe to perform given the
    /// assumptions of this object.
    pub unsafe fn raw(&self) -> raw::Timer { self.handle }
}

extern fn timer_cb(timer: *mut uvll::uv_timer_t) {
    let timer: raw::Timer = unsafe { Handle::from_raw(timer) };
    let data: &mut Data = unsafe { mem::transmute(timer.get_data()) };
    match data.action.take().unwrap() {
        WakeTask(task) => task.reawaken(),
        CallOnce(mut cb) => cb.call(),
        CallMany(mut cb, id) => {
            cb.call();

            // Note that the above operation could have performed some form
            // of scheduling. This means that the timer may have decided to
            // insert some other action to happen. This 'id' keeps track of
            // the updates to the timer, so we only reset the action back to
            // sending on this channel if the id has remained the same. This
            // is essentially a bug in that we have mutably aliasable
            // memory, but that's libuv for you. We're guaranteed to all be
            // running on the same thread, so there's no need for any
            // synchronization here.
            if data.id == id {
                data.action = Some(CallMany(cb, id));
            }
        }
    }
}

impl HomingIO for Timer {
    fn home<'r>(&'r mut self) -> &'r mut HomeHandle { &mut self.home }
}

impl Drop for Timer {
    fn drop(&mut self) {
        // note that this drop is a little subtle. Dropping a channel which is
        // held internally may invoke some scheduling operations. We can't take
        // the channel unless we're on the home scheduler, but once we're on the
        // home scheduler we should never move. Hence, we take the timer's
        // action item and then move it outside of the homing block.
        //
        // Furthermore, we can't actually free the extraneous `Data` here
        // because the `timer_cb` may actually be somewhere up on the stack in a
        // `CallMany` which might use the data. To mitigate this, we bump the
        // `id` (so the `CallMany` doesn't re-store its callback) and then
        // schedule a custom close callback to free both the timer and the data.
        let _action = unsafe {
            let _m = self.fire_homing_missile();
            self.handle.stop().unwrap();
            let data: &mut Data = mem::transmute(self.handle.get_data());
            data.id += 1;
            self.handle.close(Some(close_cb));
            data.action.take()
        };

        extern fn close_cb(handle: *mut uvll::uv_handle_t) {
            unsafe {
                let handle = handle as *mut uvll::uv_timer_t;
                let mut timer: raw::Timer = Handle::from_raw(handle);
                let data: Box<Data> = mem::transmute(timer.get_data());
                assert!(data.action.is_none());
                timer.free();
            }
        }
    }
}
<|endoftext|>


<commit_before>pub fn is_armstrong_number(_num: u32) -> bool {
    unimplemented!()
}
<commit_msg>armstrong-numbers: remove leading underscore from variable in stub<commit_after>pub fn is_armstrong_number(num: u32) -> bool {
    unimplemented!("true if {} is an armstrong number", num)
}
<|endoftext|>


<commit_before>use crate::agents::*;
use crate::config::LimitConfig;
use crate::crypto::SigningAlgorithm;
use crate::utils::{
    agent::*,
    redis::{locking, pubsub},
    SecureRandom,
};
use ::redis::{
    aio::MultiplexedConnection as RedisConn, pipe, AsyncCommands, Client as RedisClient,
    IntoConnectionInfo, RedisResult, Script,
};
use std::sync::Arc;
use std::time::Duration;

/// Internal message used to lock a key set.
struct LockKeys(SigningAlgorithm);
impl Message for LockKeys {
    type Reply = locking::LockGuard;
}

/// Internal message used to fetch a key set.
struct FetchKeys(SigningAlgorithm);
impl Message for FetchKeys {
    type Reply = RedisResult<KeySet>;
}

/// Internal message used to save a key set.
struct SaveKeys(KeySet);
impl Message for SaveKeys {
    type Reply = RedisResult<()>;
}

/// Internal message used to fetch keys and send an update to the key manager.
struct UpdateKeysLocked(SigningAlgorithm);
impl Message for UpdateKeysLocked {
    type Reply = ();
}

/// Store implementation using Redis.
pub struct RedisStore {
    /// A random unique ID for ourselves.
    id: Arc<Vec<u8>>,
    /// The connection.
    conn: RedisConn,
    /// Pubsub client.
    pubsub: pubsub::Subscriber,
    /// Locking client.
    locking: locking::LockClient,
    /// TTL of session keys
    expire_sessions: Duration,
    /// TTL of cache keys
    expire_cache: Duration,
    /// The agent used for fetching on cache miss.
    fetcher: Addr<FetchAgent>,
    /// Key manager if rotating keys are enabled.
    key_manager: Option<Addr<RotatingKeys>>,
    /// Script used to check a limit.
    limit_script: Arc<Script>,
    /// Configuration for per-email rate limiting.
    limit_per_email_config: LimitConfig,
}

impl RedisStore {
    pub async fn new(
        mut url: String,
        expire_sessions: Duration,
        expire_cache: Duration,
        limit_per_email_config: LimitConfig,
        fetcher: Addr<FetchAgent>,
        rng: SecureRandom,
    ) -> RedisResult<Self> {
        if url.starts_with("http://") {
            url = url.replace("http://", "redis://");
        } else if !url.starts_with("redis://") {
            url = format!("redis://{}", &url);
        }
        let id = Arc::new(rng.generate_async(16).await);
        let info = url.as_str().into_connection_info()?;
        let pubsub = pubsub::connect(&info).await?;
        let conn = RedisClient::open(info)?
            .get_multiplexed_tokio_connection()
            .await?;
        let locking = locking::LockClient::new(conn.clone(), pubsub.clone(), rng);

        log::warn!("Storing sessions and keys in Redis at {}", url);
        log::warn!("Please always double check this Redis and the connection to it are secure!");
        log::warn!("(This warning can't be fixed; it's a friendly reminder.)");

        let limit_script = Arc::new(Script::new(
            r"
            local count = redis.call('incr', KEYS[1])
            if count == 1 then
                redis.call('expire', KEYS[1], ARGV[1])
            end
            return count
            ",
        ));

        Ok(RedisStore {
            id,
            conn,
            pubsub,
            locking,
            expire_sessions,
            expire_cache,
            fetcher,
            key_manager: None,
            limit_script,
            limit_per_email_config,
        })
    }

    fn format_session_key(session_id: &str) -> String {
        format!("session:{}", session_id)
    }
}

impl Agent for RedisStore {
    fn started(&mut self, cx: Context<Self, AgentStarted>) {
        // Ping Redis at an interval.
        let mut conn = self.conn.clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(20));
            loop {
                interval.tick().await;
                let _: String = ::redis::cmd("PING")
                    .query_async(&mut conn)
                    .await
                    .expect("Redis ping failed");
            }
        });
        cx.reply(());
    }
}

impl Handler<SaveSession> for RedisStore {
    fn handle(&mut self, message: SaveSession, cx: Context<Self, SaveSession>) {
        let mut conn = self.conn.clone();
        let ttl = self.expire_sessions;
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            let data = serde_json::to_string(&message.data)?;
            conn.set_ex(&key, data, ttl.as_secs() as usize).await?;
            Ok(())
        });
    }
}

impl Handler<GetSession> for RedisStore {
    fn handle(&mut self, message: GetSession, cx: Context<Self, GetSession>) {
        let mut conn = self.conn.clone();
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            let data: String = conn.get(&key).await?;
            let data = serde_json::from_str(&data)?;
            Ok(data)
        });
    }
}

impl Handler<DeleteSession> for RedisStore {
    fn handle(&mut self, message: DeleteSession, cx: Context<Self, DeleteSession>) {
        let mut conn = self.conn.clone();
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            conn.del(&key).await?;
            Ok(())
        });
    }
}

impl Handler<FetchUrlCached> for RedisStore {
    fn handle(&mut self, message: FetchUrlCached, cx: Context<Self, FetchUrlCached>) {
        let mut conn = self.conn.clone();
        let mut locking = self.locking.clone();
        let fetcher = self.fetcher.clone();
        let expire_cache = self.expire_cache;
        cx.reply_later(async move {
            let key = format!("cache:{}", message.url);
            let _lock = locking.lock(format!("lock:{}", key).as_bytes()).await;
            if let Some(data) = conn.get(key).await? {
                Ok(data)
            } else {
                let key = message.url.as_str().to_owned();
                let result = fetcher.send(FetchUrl { url: message.url }).await?;
                let ttl = std::cmp::max(expire_cache, result.max_age);
                conn.set_ex(key, result.data.clone(), ttl.as_secs() as usize)
                    .await?;
                Ok(result.data)
            }
        });
    }
}

impl Handler<IncrAndTestLimit> for RedisStore {
    fn handle(&mut self, message: IncrAndTestLimit, cx: Context<Self, IncrAndTestLimit>) {
        let mut conn = self.conn.clone();
        let script = self.limit_script.clone();
        let (key, config) = match message {
            IncrAndTestLimit::PerEmail { addr } => (
                format!("ratelimit:per-email:{}", addr),
                self.limit_per_email_config,
            ),
        };
        cx.reply_later(async move {
            let mut invocation = script.prepare_invoke();
            invocation.key(key).arg(config.duration.as_secs());
            let count: usize = invocation.invoke_async(&mut conn).await?;
            Ok(count <= config.max_count)
        });
    }
}

impl Handler<EnableRotatingKeys> for RedisStore {
    fn handle(&mut self, message: EnableRotatingKeys, cx: Context<Self, EnableRotatingKeys>) {
        let me = cx.addr().clone();
        let my_id = self.id.clone();
        let mut pubsub = self.pubsub.clone();
        self.key_manager = Some(message.key_manager.clone());
        cx.reply_later(async move {
            for signing_alg in &message.signing_algs {
                let signing_alg = *signing_alg;
                // Listen for key changes by other workers.
                let chan = format!("keys-updated:{}", signing_alg).into_bytes();
                let mut sub = pubsub.subscribe(chan).await;
                let me2 = me.clone();
                let my_id2 = my_id.clone();
                tokio::spawn(async move {
                    loop {
                        let from_id = sub.recv().await.expect("Redis keys subscription failed");
                        if from_id != *my_id2 {
                            me2.send(UpdateKeysLocked(signing_alg));
                        }
                    }
                });
                // Fetch current keys.
                me.send(UpdateKeysLocked(signing_alg)).await;
            }
        });
    }
}

impl Handler<RotateKeysLocked> for RedisStore {
    fn handle(&mut self, message: RotateKeysLocked, cx: Context<Self, RotateKeysLocked>) {
        let me = cx.addr().clone();
        let key_manager = self.key_manager.as_ref().unwrap().clone();
        cx.reply_later(async move {
            let lock = me.send(LockKeys(message.0)).await;
            let key_set = me
                .send(FetchKeys(message.0))
                .await
                .expect("Failed to fetch keys from Redis");
            if let Some(key_set) = key_manager.send(RotateKeys(key_set)).await {
                me.send(SaveKeys(key_set.clone()))
                    .await
                    .expect("Failed to save keys to Redis");
                drop(lock);
                key_manager.send(UpdateKeys(key_set)).await;
            }
        });
    }
}

impl Handler<LockKeys> for RedisStore {
    fn handle(&mut self, message: LockKeys, cx: Context<Self, LockKeys>) {
        let mut locking = self.locking.clone();
        let lock_key = format!("lock:keys:{}", message.0);
        cx.reply_later(async move { locking.lock(lock_key.as_bytes()).await });
    }
}

impl Handler<FetchKeys> for RedisStore {
    fn handle(&mut self, message: FetchKeys, cx: Context<Self, FetchKeys>) {
        let mut conn = self.conn.clone();
        let signing_alg = message.0;
        let db_key = format!("keys:{}", signing_alg);
        cx.reply_later(async move {
            let key_set: Option<String> = conn.get(db_key).await?;
            let key_set = key_set
                .map(|data| serde_json::from_str(&data).expect("Invalid key set JSON in Redis"))
                .unwrap_or_else(|| KeySet::empty(signing_alg));
            Ok(key_set)
        })
    }
}

impl Handler<SaveKeys> for RedisStore {
    fn handle(&mut self, message: SaveKeys, cx: Context<Self, SaveKeys>) {
        let mut conn = self.conn.clone();
        let signing_alg = message.0.signing_alg;
        let db_key = format!("keys:{}", signing_alg);
        let data = serde_json::to_string(&message.0).expect("Could not encode key set as JSON");
        let mut pipe = pipe();
        pipe.atomic()
            .set(db_key, data)
            .publish(format!("keys-updated:{}", signing_alg), &self.id[..]);
        cx.reply_later(async move { pipe.query_async(&mut conn).await });
    }
}

impl Handler<UpdateKeysLocked> for RedisStore {
    fn handle(&mut self, message: UpdateKeysLocked, cx: Context<Self, UpdateKeysLocked>) {
        let me = cx.addr().clone();
        let key_manager = self.key_manager.as_ref().unwrap().clone();
        cx.reply_later(async move {
            let key_set = {
                let _lock = me.send(LockKeys(message.0)).await;
                me.send(FetchKeys(message.0))
                    .await
                    .expect("Failed to fetch keys from Redis")
            };
            key_manager.send(UpdateKeys(key_set)).await;
        });
    }
}

impl StoreSender for Addr<RedisStore> {}
<commit_msg>Fix invalid sessions with Redis storage<commit_after>use crate::agents::*;
use crate::config::LimitConfig;
use crate::crypto::SigningAlgorithm;
use crate::utils::{
    agent::*,
    redis::{locking, pubsub},
    SecureRandom,
};
use ::redis::{
    aio::MultiplexedConnection as RedisConn, pipe, AsyncCommands, Client as RedisClient,
    IntoConnectionInfo, RedisResult, Script,
};
use std::sync::Arc;
use std::time::Duration;

/// Internal message used to lock a key set.
struct LockKeys(SigningAlgorithm);
impl Message for LockKeys {
    type Reply = locking::LockGuard;
}

/// Internal message used to fetch a key set.
struct FetchKeys(SigningAlgorithm);
impl Message for FetchKeys {
    type Reply = RedisResult<KeySet>;
}

/// Internal message used to save a key set.
struct SaveKeys(KeySet);
impl Message for SaveKeys {
    type Reply = RedisResult<()>;
}

/// Internal message used to fetch keys and send an update to the key manager.
struct UpdateKeysLocked(SigningAlgorithm);
impl Message for UpdateKeysLocked {
    type Reply = ();
}

/// Store implementation using Redis.
pub struct RedisStore {
    /// A random unique ID for ourselves.
    id: Arc<Vec<u8>>,
    /// The connection.
    conn: RedisConn,
    /// Pubsub client.
    pubsub: pubsub::Subscriber,
    /// Locking client.
    locking: locking::LockClient,
    /// TTL of session keys
    expire_sessions: Duration,
    /// TTL of cache keys
    expire_cache: Duration,
    /// The agent used for fetching on cache miss.
    fetcher: Addr<FetchAgent>,
    /// Key manager if rotating keys are enabled.
    key_manager: Option<Addr<RotatingKeys>>,
    /// Script used to check a limit.
    limit_script: Arc<Script>,
    /// Configuration for per-email rate limiting.
    limit_per_email_config: LimitConfig,
}

impl RedisStore {
    pub async fn new(
        mut url: String,
        expire_sessions: Duration,
        expire_cache: Duration,
        limit_per_email_config: LimitConfig,
        fetcher: Addr<FetchAgent>,
        rng: SecureRandom,
    ) -> RedisResult<Self> {
        if url.starts_with("http://") {
            url = url.replace("http://", "redis://");
        } else if !url.starts_with("redis://") {
            url = format!("redis://{}", &url);
        }
        let id = Arc::new(rng.generate_async(16).await);
        let info = url.as_str().into_connection_info()?;
        let pubsub = pubsub::connect(&info).await?;
        let conn = RedisClient::open(info)?
            .get_multiplexed_tokio_connection()
            .await?;
        let locking = locking::LockClient::new(conn.clone(), pubsub.clone(), rng);

        log::warn!("Storing sessions and keys in Redis at {}", url);
        log::warn!("Please always double check this Redis and the connection to it are secure!");
        log::warn!("(This warning can't be fixed; it's a friendly reminder.)");

        let limit_script = Arc::new(Script::new(
            r"
            local count = redis.call('incr', KEYS[1])
            if count == 1 then
                redis.call('expire', KEYS[1], ARGV[1])
            end
            return count
            ",
        ));

        Ok(RedisStore {
            id,
            conn,
            pubsub,
            locking,
            expire_sessions,
            expire_cache,
            fetcher,
            key_manager: None,
            limit_script,
            limit_per_email_config,
        })
    }

    fn format_session_key(session_id: &str) -> String {
        format!("session:{}", session_id)
    }
}

impl Agent for RedisStore {
    fn started(&mut self, cx: Context<Self, AgentStarted>) {
        // Ping Redis at an interval.
        let mut conn = self.conn.clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(20));
            loop {
                interval.tick().await;
                let _: String = ::redis::cmd("PING")
                    .query_async(&mut conn)
                    .await
                    .expect("Redis ping failed");
            }
        });
        cx.reply(());
    }
}

impl Handler<SaveSession> for RedisStore {
    fn handle(&mut self, message: SaveSession, cx: Context<Self, SaveSession>) {
        let mut conn = self.conn.clone();
        let ttl = self.expire_sessions;
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            let data = serde_json::to_string(&message.data)?;
            conn.set_ex(&key, data, ttl.as_secs() as usize).await?;
            Ok(())
        });
    }
}

impl Handler<GetSession> for RedisStore {
    fn handle(&mut self, message: GetSession, cx: Context<Self, GetSession>) {
        let mut conn = self.conn.clone();
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            let data: Option<String> = conn.get(&key).await?;
            if let Some(data) = data {
                Ok(Some(serde_json::from_str(&data)?))
            } else {
                Ok(None)
            }
        });
    }
}

impl Handler<DeleteSession> for RedisStore {
    fn handle(&mut self, message: DeleteSession, cx: Context<Self, DeleteSession>) {
        let mut conn = self.conn.clone();
        cx.reply_later(async move {
            let key = Self::format_session_key(&message.session_id);
            conn.del(&key).await?;
            Ok(())
        });
    }
}

impl Handler<FetchUrlCached> for RedisStore {
    fn handle(&mut self, message: FetchUrlCached, cx: Context<Self, FetchUrlCached>) {
        let mut conn = self.conn.clone();
        let mut locking = self.locking.clone();
        let fetcher = self.fetcher.clone();
        let expire_cache = self.expire_cache;
        cx.reply_later(async move {
            let key = format!("cache:{}", message.url);
            let _lock = locking.lock(format!("lock:{}", key).as_bytes()).await;
            if let Some(data) = conn.get(key).await? {
                Ok(data)
            } else {
                let key = message.url.as_str().to_owned();
                let result = fetcher.send(FetchUrl { url: message.url }).await?;
                let ttl = std::cmp::max(expire_cache, result.max_age);
                conn.set_ex(key, result.data.clone(), ttl.as_secs() as usize)
                    .await?;
                Ok(result.data)
            }
        });
    }
}

impl Handler<IncrAndTestLimit> for RedisStore {
    fn handle(&mut self, message: IncrAndTestLimit, cx: Context<Self, IncrAndTestLimit>) {
        let mut conn = self.conn.clone();
        let script = self.limit_script.clone();
        let (key, config) = match message {
            IncrAndTestLimit::PerEmail { addr } => (
                format!("ratelimit:per-email:{}", addr),
                self.limit_per_email_config,
            ),
        };
        cx.reply_later(async move {
            let mut invocation = script.prepare_invoke();
            invocation.key(key).arg(config.duration.as_secs());
            let count: usize = invocation.invoke_async(&mut conn).await?;
            Ok(count <= config.max_count)
        });
    }
}

impl Handler<EnableRotatingKeys> for RedisStore {
    fn handle(&mut self, message: EnableRotatingKeys, cx: Context<Self, EnableRotatingKeys>) {
        let me = cx.addr().clone();
        let my_id = self.id.clone();
        let mut pubsub = self.pubsub.clone();
        self.key_manager = Some(message.key_manager.clone());
        cx.reply_later(async move {
            for signing_alg in &message.signing_algs {
                let signing_alg = *signing_alg;
                // Listen for key changes by other workers.
                let chan = format!("keys-updated:{}", signing_alg).into_bytes();
                let mut sub = pubsub.subscribe(chan).await;
                let me2 = me.clone();
                let my_id2 = my_id.clone();
                tokio::spawn(async move {
                    loop {
                        let from_id = sub.recv().await.expect("Redis keys subscription failed");
                        if from_id != *my_id2 {
                            me2.send(UpdateKeysLocked(signing_alg));
                        }
                    }
                });
                // Fetch current keys.
                me.send(UpdateKeysLocked(signing_alg)).await;
            }
        });
    }
}

impl Handler<RotateKeysLocked> for RedisStore {
    fn handle(&mut self, message: RotateKeysLocked, cx: Context<Self, RotateKeysLocked>) {
        let me = cx.addr().clone();
        let key_manager = self.key_manager.as_ref().unwrap().clone();
        cx.reply_later(async move {
            let lock = me.send(LockKeys(message.0)).await;
            let key_set = me
                .send(FetchKeys(message.0))
                .await
                .expect("Failed to fetch keys from Redis");
            if let Some(key_set) = key_manager.send(RotateKeys(key_set)).await {
                me.send(SaveKeys(key_set.clone()))
                    .await
                    .expect("Failed to save keys to Redis");
                drop(lock);
                key_manager.send(UpdateKeys(key_set)).await;
            }
        });
    }
}

impl Handler<LockKeys> for RedisStore {
    fn handle(&mut self, message: LockKeys, cx: Context<Self, LockKeys>) {
        let mut locking = self.locking.clone();
        let lock_key = format!("lock:keys:{}", message.0);
        cx.reply_later(async move { locking.lock(lock_key.as_bytes()).await });
    }
}

impl Handler<FetchKeys> for RedisStore {
    fn handle(&mut self, message: FetchKeys, cx: Context<Self, FetchKeys>) {
        let mut conn = self.conn.clone();
        let signing_alg = message.0;
        let db_key = format!("keys:{}", signing_alg);
        cx.reply_later(async move {
            let key_set: Option<String> = conn.get(db_key).await?;
            let key_set = key_set
                .map(|data| serde_json::from_str(&data).expect("Invalid key set JSON in Redis"))
                .unwrap_or_else(|| KeySet::empty(signing_alg));
            Ok(key_set)
        })
    }
}

impl Handler<SaveKeys> for RedisStore {
    fn handle(&mut self, message: SaveKeys, cx: Context<Self, SaveKeys>) {
        let mut conn = self.conn.clone();
        let signing_alg = message.0.signing_alg;
        let db_key = format!("keys:{}", signing_alg);
        let data = serde_json::to_string(&message.0).expect("Could not encode key set as JSON");
        let mut pipe = pipe();
        pipe.atomic()
            .set(db_key, data)
            .publish(format!("keys-updated:{}", signing_alg), &self.id[..]);
        cx.reply_later(async move { pipe.query_async(&mut conn).await });
    }
}

impl Handler<UpdateKeysLocked> for RedisStore {
    fn handle(&mut self, message: UpdateKeysLocked, cx: Context<Self, UpdateKeysLocked>) {
        let me = cx.addr().clone();
        let key_manager = self.key_manager.as_ref().unwrap().clone();
        cx.reply_later(async move {
            let key_set = {
                let _lock = me.send(LockKeys(message.0)).await;
                me.send(FetchKeys(message.0))
                    .await
                    .expect("Failed to fetch keys from Redis")
            };
            key_manager.send(UpdateKeys(key_set)).await;
        });
    }
}

impl StoreSender for Addr<RedisStore> {}
<|endoftext|>


<commit_before>use {Errno, Error, Result};
use libc::c_int;
use void::Void;
use std::mem::drop;

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum RebootMode {
    Halt = 0xcdef0123,
    kexec = 0x45584543,
    PowerOff = 0x4321fedc,
    Restart = 0x1234567,
    // we do not support Restart2,
    Suspend = 0xd000fce1,
}

pub fn reboot(how: RebootMode) -> Result<Void> {
    unsafe {
        ext::reboot(how as c_int)
    };
    Err(Error::Sys(Errno::last()))
}

#[allow(overflowing_literals)]
pub fn set_cad_enabled(enable: bool) -> Result<()> {
    let res = unsafe {
        ext::reboot(if enable { 0x89abcdef } else { 0 })
    };
    Errno::result(res).map(drop)
}

mod ext {
    use libc::c_int;
    extern {
        pub fn reboot(cmd: c_int) -> c_int;
    }
}
<commit_msg>Add some basic docs for reboot::set_cad_enabled()<commit_after>use {Errno, Error, Result};
use libc::c_int;
use void::Void;
use std::mem::drop;

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum RebootMode {
    Halt = 0xcdef0123,
    kexec = 0x45584543,
    PowerOff = 0x4321fedc,
    Restart = 0x1234567,
    // we do not support Restart2,
    Suspend = 0xd000fce1,
}

pub fn reboot(how: RebootMode) -> Result<Void> {
    unsafe {
        ext::reboot(how as c_int)
    };
    Err(Error::Sys(Errno::last()))
}

/// Enable or disable the reboot keystroke (Ctrl-Alt-Delete).
///
/// Corresponds to calling `reboot(RB_ENABLE_CAD)` or `reboot(RB_DISABLE_CAD)` in C.
#[allow(overflowing_literals)]
pub fn set_cad_enabled(enable: bool) -> Result<()> {
    let res = unsafe {
        ext::reboot(if enable { 0x89abcdef } else { 0 })
    };
    Errno::result(res).map(drop)
}

mod ext {
    use libc::c_int;
    extern {
        pub fn reboot(cmd: c_int) -> c_int;
    }
}
<|endoftext|>


<commit_before>use diesel::*;
use std::sync::ONCE_INIT;

use super::OwnerTeamsResponse;
use builders::{CrateBuilder, PublishBuilder};
use models::{Crate, NewUser};
use record::GhUser;
use {add_team_to_crate, new_team, OkBool, RequestHelper, TestApp};

impl ::util::MockTokenUser {
    /// Add to the specified crate the specified owner.
    fn add_crate_owner(&self, krate_name: &str, owner: &str) -> ::util::Response<OkBool> {
        let url = format!("/api/v1/crates/{}/owners", krate_name);
        let body = format!("{{\"users\":[\"{}\"]}}", owner);
        self.put(&url, body.as_bytes())
    }

    /// Remove from the specified crate the specified owner.
    fn remove_crate_owner(&self, krate_name: &str, owner: &str) -> ::util::Response<OkBool> {
        let url = format!("/api/v1/crates/{}/owners", krate_name);
        let body = format!("{{\"users\":[\"{}\"]}}", owner);
        self.delete_with_body(&url, body.as_bytes())
    }
}

impl ::util::MockAnonymousUser {
    /// List the team owners of the specified crate.
    fn crate_owner_teams(&self, krate_name: &str) -> ::util::Response<OwnerTeamsResponse> {
        let url = format!("/api/v1/crates/{}/owner_team", krate_name);
        self.get(&url)
    }
}

// Users: `crates-tester-1` and `crates-tester-2`
// Passwords: ask acrichto or gankro
// Teams: `crates-test-org:core`, `crates-test-org:just-for-crates-2`
// tester-1 is on core only, tester-2 is on both

static GH_USER_1: GhUser = GhUser {
    login: "crates-tester-1",
    init: ONCE_INIT,
};
static GH_USER_2: GhUser = GhUser {
    login: "crates-tester-2",
    init: ONCE_INIT,
};

fn mock_user_on_only_one_team() -> NewUser<'static> {
    GH_USER_1.user()
}
fn mock_user_on_both_teams() -> NewUser<'static> {
    GH_USER_2.user()
}

// Test adding team without `github:`
#[test]
fn not_github() {
    let (app, _, user, token) = TestApp::init().with_token();

    app.db(|conn| {
        CrateBuilder::new("foo_not_github", user.as_model().id).expect_build(conn);
    });

    let json = token
        .add_crate_owner("foo_not_github", "dropbox:foo:foo")
        .bad_with_status(200);

    assert!(
        json.errors[0].detail.contains("unknown organization"),
        "{:?}",
        json.errors
    );
}

#[test]
fn weird_name() {
    let (app, _, user, token) = TestApp::init().with_token();

    app.db(|conn| {
        CrateBuilder::new("foo_weird_name", user.as_model().id).expect_build(conn);
    });

    let json = token
        .add_crate_owner("foo_weird_name", "github:foo/../bar:wut")
        .bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("organization cannot contain"),
        "{:?}",
        json.errors
    );
}

// Test adding team without second `:`
#[test]
fn one_colon() {
    let (app, _, user, token) = TestApp::init().with_token();

    app.db(|conn| {
        CrateBuilder::new("foo_one_colon", user.as_model().id).expect_build(conn);
    });

    let json = token
        .add_crate_owner("foo_one_colon", "github:foo")
        .bad_with_status(200);

    assert!(
        json.errors[0].detail.contains("missing github team"),
        "{:?}",
        json.errors
    );
}

#[test]
fn nonexistent_team() {
    let (app, _, user, token) = TestApp::with_proxy().with_token();

    app.db(|conn| {
        CrateBuilder::new("foo_nonexistent", user.as_model().id).expect_build(conn);
    });

    let json = token
        .add_crate_owner(
            "foo_nonexistent",
            "github:crates-test-org:this-does-not-exist",
        ).bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("could not find the github team crates-test-org/this-does-not-exist"),
        "{:?}",
        json.errors
    );
}

// Test adding team names with mixed case, when on the team
#[test]
fn add_team_mixed_case() {
    let (app, anon) = TestApp::with_proxy().empty();
    let user = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token = user.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_mixed_case", user.as_model().id).expect_build(conn);
    });

    token.add_crate_owner("foo_mixed_case", "github:Crates-Test-Org:Core")
        .good();

    app.db(|conn| {
        let krate = Crate::by_name("foo_mixed_case")
            .first::<Crate>(conn)
            .unwrap();
        let owners = krate.owners(conn).unwrap();
        assert_eq!(owners.len(), 2);
        let owner = &owners[1];
        assert_eq!(owner.login(), owner.login().to_lowercase());
    });

    let json = anon.crate_owner_teams("foo_mixed_case").good();
    assert_eq!(json.teams.len(), 1);
    assert_eq!(json.teams[0].login, "github:crates-test-org:core");
}

// Test adding team as owner when not on it
#[test]
fn add_team_as_non_member() {
    let (app, _) = TestApp::with_proxy().empty();
    let user = app.db_new_user(&mock_user_on_only_one_team().gh_login);
    let token = user.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_team_non_member", user.as_model().id).expect_build(conn);
    });

    let json = token
        .add_crate_owner(
            "foo_team_non_member",
            "github:crates-test-org:just-for-crates-2",
        ).bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("only members of a team can add it as an owner"),
        "{:?}",
        json.errors
    );
}

#[test]
fn remove_team_as_named_owner() {
    let (app, _) = TestApp::with_proxy().empty();
    let user_on_both_teams = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token_on_both_teams = user_on_both_teams.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_remove_team", user_on_both_teams.as_model().id).expect_build(conn);
    });

    token_on_both_teams
        .add_crate_owner("foo_remove_team", "github:crates-test-org:core")
        .good();

    token_on_both_teams
        .remove_crate_owner("foo_remove_team", "github:crates-test-org:core")
        .good();

    let user_on_one_team = app.db_new_user(&mock_user_on_only_one_team().gh_login);
    let crate_to_publish = PublishBuilder::new("foo_remove_team").version("2.0.0");
    let json = user_on_one_team
        .publish(crate_to_publish)
        .bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("this crate exists but you don't seem to be an owner.",),
        "{:?}",
        json.errors
    );
}

#[test]
fn remove_team_as_team_owner() {
    let (app, _) = TestApp::with_proxy().empty();
    let user_on_both_teams = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token_on_both_teams = user_on_both_teams.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_remove_team_owner", user_on_both_teams.as_model().id)
            .expect_build(conn);
    });

    token_on_both_teams
        .add_crate_owner("foo_remove_team_owner", "github:crates-test-org:core")
        .good();

    let user_on_one_team = app.db_new_user(&mock_user_on_only_one_team().gh_login);
    let token_on_one_team = user_on_one_team.db_new_token("arbitrary token name");

    let json = token_on_one_team
        .remove_crate_owner("foo_remove_team_owner", "github:crates-test-org:core")
        .bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("team members don't have permission to modify owners",),
        "{:?}",
        json.errors
    );
}

// Test trying to publish a crate we don't own
#[test]
fn publish_not_owned() {
    let (app, _) = TestApp::with_proxy().empty();
    let user_on_both_teams = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token_on_both_teams = user_on_both_teams.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_not_owned", user_on_both_teams.as_model().id).expect_build(conn);
    });

    token_on_both_teams
        .add_crate_owner("foo_not_owned", "github:crates-test-org:just-for-crates-2")
        .good();

    let user_on_one_team = app.db_new_user(&mock_user_on_only_one_team().gh_login);

    let crate_to_publish = PublishBuilder::new("foo_not_owned").version("2.0.0");
    let json = user_on_one_team
        .publish(crate_to_publish)
        .bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("this crate exists but you don't seem to be an owner.",),
        "{:?}",
        json.errors
    );
}

// Test trying to publish a krate we do own (but only because of teams)
#[test]
fn publish_owned() {
    let (app, _) = TestApp::with_proxy().empty();
    let user_on_both_teams = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token_on_both_teams = user_on_both_teams.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_team_owned", user_on_both_teams.as_model().id).expect_build(conn);
    });

    token_on_both_teams
        .add_crate_owner("foo_team_owned", "github:crates-test-org:core")
        .good();

    let user_on_one_team = app.db_new_user(&mock_user_on_only_one_team().gh_login);

    let crate_to_publish = PublishBuilder::new("foo_team_owned").version("2.0.0");
    user_on_one_team.publish(crate_to_publish).good();
}

// Test trying to change owners (when only on an owning team)
#[test]
fn add_owners_as_team_owner() {
    let (app, _) = TestApp::with_proxy().empty();
    let user_on_both_teams = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token_on_both_teams = user_on_both_teams.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_add_owner", user_on_both_teams.as_model().id).expect_build(conn);
    });

    token_on_both_teams
        .add_crate_owner("foo_add_owner", "github:crates-test-org:core")
        .good();

    let user_on_one_team = app.db_new_user(&mock_user_on_only_one_team().gh_login);
    let token_on_one_team = user_on_one_team.db_new_token("arbitrary token name");

    let json = token_on_one_team
        .add_crate_owner("foo_add_owner", "arbitrary_username")
        .bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("team members don't have permission to modify owners",),
        "{:?}",
        json.errors
    );
}

#[test]
fn crates_by_team_id() {
    let (app, anon, user) = TestApp::init().with_user();
    let user = user.as_model();

    let team = app.db(|conn| {
        let t = new_team("github:crates-test-org:team_foo")
            .create_or_update(conn)
            .unwrap();
        let krate = CrateBuilder::new("foo", user.id).expect_build(conn);
        add_team_to_crate(&t, &krate, &user, conn).unwrap();
        t
    });

    let json = anon.search(&format!("team_id={}", team.id));
    assert_eq!(json.crates.len(), 1);
}

#[test]
fn crates_by_team_id_not_including_deleted_owners() {
    // This needs to use the proxy beacuse removing a team checks with github that you're on the
    // team before you're allowed to remove it from the crate
    let (app, anon) = TestApp::with_proxy().empty();
    let user = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let user = user.as_model();

    let team = app.db(|conn| {
        let t = new_team("github:crates-test-org:core")
            .create_or_update(conn)
            .unwrap();
        let krate = CrateBuilder::new("foo", user.id).expect_build(conn);
        add_team_to_crate(&t, &krate, &user, conn).unwrap();
        krate
            .owner_remove(&app.as_inner(), conn, &user, &t.login)
            .unwrap();
        t
    });

    let json = anon.search(&format!("team_id={}", team.id));
    assert_eq!(json.crates.len(), 0);
}
<commit_msg>Rename test helper owner methods to be clearer<commit_after>use diesel::*;
use std::sync::ONCE_INIT;

use super::OwnerTeamsResponse;
use builders::{CrateBuilder, PublishBuilder};
use models::{Crate, NewUser};
use record::GhUser;
use {add_team_to_crate, new_team, OkBool, RequestHelper, TestApp};

impl ::util::MockTokenUser {
    /// Add to the specified crate the specified owner.
    fn add_named_owner(&self, krate_name: &str, owner: &str) -> ::util::Response<OkBool> {
        let url = format!("/api/v1/crates/{}/owners", krate_name);
        let body = format!("{{\"users\":[\"{}\"]}}", owner);
        self.put(&url, body.as_bytes())
    }

    /// Remove from the specified crate the specified owner.
    fn remove_named_owner(&self, krate_name: &str, owner: &str) -> ::util::Response<OkBool> {
        let url = format!("/api/v1/crates/{}/owners", krate_name);
        let body = format!("{{\"users\":[\"{}\"]}}", owner);
        self.delete_with_body(&url, body.as_bytes())
    }
}

impl ::util::MockAnonymousUser {
    /// List the team owners of the specified crate.
    fn crate_owner_teams(&self, krate_name: &str) -> ::util::Response<OwnerTeamsResponse> {
        let url = format!("/api/v1/crates/{}/owner_team", krate_name);
        self.get(&url)
    }
}

// Users: `crates-tester-1` and `crates-tester-2`
// Passwords: ask acrichto or gankro
// Teams: `crates-test-org:core`, `crates-test-org:just-for-crates-2`
// tester-1 is on core only, tester-2 is on both

static GH_USER_1: GhUser = GhUser {
    login: "crates-tester-1",
    init: ONCE_INIT,
};
static GH_USER_2: GhUser = GhUser {
    login: "crates-tester-2",
    init: ONCE_INIT,
};

fn mock_user_on_only_one_team() -> NewUser<'static> {
    GH_USER_1.user()
}
fn mock_user_on_both_teams() -> NewUser<'static> {
    GH_USER_2.user()
}

// Test adding team without `github:`
#[test]
fn not_github() {
    let (app, _, user, token) = TestApp::init().with_token();

    app.db(|conn| {
        CrateBuilder::new("foo_not_github", user.as_model().id).expect_build(conn);
    });

    let json = token
        .add_named_owner("foo_not_github", "dropbox:foo:foo")
        .bad_with_status(200);

    assert!(
        json.errors[0].detail.contains("unknown organization"),
        "{:?}",
        json.errors
    );
}

#[test]
fn weird_name() {
    let (app, _, user, token) = TestApp::init().with_token();

    app.db(|conn| {
        CrateBuilder::new("foo_weird_name", user.as_model().id).expect_build(conn);
    });

    let json = token
        .add_named_owner("foo_weird_name", "github:foo/../bar:wut")
        .bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("organization cannot contain"),
        "{:?}",
        json.errors
    );
}

// Test adding team without second `:`
#[test]
fn one_colon() {
    let (app, _, user, token) = TestApp::init().with_token();

    app.db(|conn| {
        CrateBuilder::new("foo_one_colon", user.as_model().id).expect_build(conn);
    });

    let json = token
        .add_named_owner("foo_one_colon", "github:foo")
        .bad_with_status(200);

    assert!(
        json.errors[0].detail.contains("missing github team"),
        "{:?}",
        json.errors
    );
}

#[test]
fn nonexistent_team() {
    let (app, _, user, token) = TestApp::with_proxy().with_token();

    app.db(|conn| {
        CrateBuilder::new("foo_nonexistent", user.as_model().id).expect_build(conn);
    });

    let json = token
        .add_named_owner(
            "foo_nonexistent",
            "github:crates-test-org:this-does-not-exist",
        ).bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("could not find the github team crates-test-org/this-does-not-exist"),
        "{:?}",
        json.errors
    );
}

// Test adding team names with mixed case, when on the team
#[test]
fn add_team_mixed_case() {
    let (app, anon) = TestApp::with_proxy().empty();
    let user = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token = user.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_mixed_case", user.as_model().id).expect_build(conn);
    });

    token.add_named_owner("foo_mixed_case", "github:Crates-Test-Org:Core")
        .good();

    app.db(|conn| {
        let krate = Crate::by_name("foo_mixed_case")
            .first::<Crate>(conn)
            .unwrap();
        let owners = krate.owners(conn).unwrap();
        assert_eq!(owners.len(), 2);
        let owner = &owners[1];
        assert_eq!(owner.login(), owner.login().to_lowercase());
    });

    let json = anon.crate_owner_teams("foo_mixed_case").good();
    assert_eq!(json.teams.len(), 1);
    assert_eq!(json.teams[0].login, "github:crates-test-org:core");
}

// Test adding team as owner when not on it
#[test]
fn add_team_as_non_member() {
    let (app, _) = TestApp::with_proxy().empty();
    let user = app.db_new_user(&mock_user_on_only_one_team().gh_login);
    let token = user.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_team_non_member", user.as_model().id).expect_build(conn);
    });

    let json = token
        .add_named_owner(
            "foo_team_non_member",
            "github:crates-test-org:just-for-crates-2",
        ).bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("only members of a team can add it as an owner"),
        "{:?}",
        json.errors
    );
}

#[test]
fn remove_team_as_named_owner() {
    let (app, _) = TestApp::with_proxy().empty();
    let user_on_both_teams = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token_on_both_teams = user_on_both_teams.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_remove_team", user_on_both_teams.as_model().id).expect_build(conn);
    });

    token_on_both_teams
        .add_named_owner("foo_remove_team", "github:crates-test-org:core")
        .good();

    token_on_both_teams
        .remove_named_owner("foo_remove_team", "github:crates-test-org:core")
        .good();

    let user_on_one_team = app.db_new_user(&mock_user_on_only_one_team().gh_login);
    let crate_to_publish = PublishBuilder::new("foo_remove_team").version("2.0.0");
    let json = user_on_one_team
        .publish(crate_to_publish)
        .bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("this crate exists but you don't seem to be an owner.",),
        "{:?}",
        json.errors
    );
}

#[test]
fn remove_team_as_team_owner() {
    let (app, _) = TestApp::with_proxy().empty();
    let user_on_both_teams = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token_on_both_teams = user_on_both_teams.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_remove_team_owner", user_on_both_teams.as_model().id)
            .expect_build(conn);
    });

    token_on_both_teams
        .add_named_owner("foo_remove_team_owner", "github:crates-test-org:core")
        .good();

    let user_on_one_team = app.db_new_user(&mock_user_on_only_one_team().gh_login);
    let token_on_one_team = user_on_one_team.db_new_token("arbitrary token name");

    let json = token_on_one_team
        .remove_named_owner("foo_remove_team_owner", "github:crates-test-org:core")
        .bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("team members don't have permission to modify owners",),
        "{:?}",
        json.errors
    );
}

// Test trying to publish a crate we don't own
#[test]
fn publish_not_owned() {
    let (app, _) = TestApp::with_proxy().empty();
    let user_on_both_teams = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token_on_both_teams = user_on_both_teams.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_not_owned", user_on_both_teams.as_model().id).expect_build(conn);
    });

    token_on_both_teams
        .add_named_owner("foo_not_owned", "github:crates-test-org:just-for-crates-2")
        .good();

    let user_on_one_team = app.db_new_user(&mock_user_on_only_one_team().gh_login);

    let crate_to_publish = PublishBuilder::new("foo_not_owned").version("2.0.0");
    let json = user_on_one_team
        .publish(crate_to_publish)
        .bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("this crate exists but you don't seem to be an owner.",),
        "{:?}",
        json.errors
    );
}

// Test trying to publish a krate we do own (but only because of teams)
#[test]
fn publish_owned() {
    let (app, _) = TestApp::with_proxy().empty();
    let user_on_both_teams = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token_on_both_teams = user_on_both_teams.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_team_owned", user_on_both_teams.as_model().id).expect_build(conn);
    });

    token_on_both_teams
        .add_named_owner("foo_team_owned", "github:crates-test-org:core")
        .good();

    let user_on_one_team = app.db_new_user(&mock_user_on_only_one_team().gh_login);

    let crate_to_publish = PublishBuilder::new("foo_team_owned").version("2.0.0");
    user_on_one_team.publish(crate_to_publish).good();
}

// Test trying to change owners (when only on an owning team)
#[test]
fn add_owners_as_team_owner() {
    let (app, _) = TestApp::with_proxy().empty();
    let user_on_both_teams = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let token_on_both_teams = user_on_both_teams.db_new_token("arbitrary token name");

    app.db(|conn| {
        CrateBuilder::new("foo_add_owner", user_on_both_teams.as_model().id).expect_build(conn);
    });

    token_on_both_teams
        .add_named_owner("foo_add_owner", "github:crates-test-org:core")
        .good();

    let user_on_one_team = app.db_new_user(&mock_user_on_only_one_team().gh_login);
    let token_on_one_team = user_on_one_team.db_new_token("arbitrary token name");

    let json = token_on_one_team
        .add_named_owner("foo_add_owner", "arbitrary_username")
        .bad_with_status(200);

    assert!(
        json.errors[0]
            .detail
            .contains("team members don't have permission to modify owners",),
        "{:?}",
        json.errors
    );
}

#[test]
fn crates_by_team_id() {
    let (app, anon, user) = TestApp::init().with_user();
    let user = user.as_model();

    let team = app.db(|conn| {
        let t = new_team("github:crates-test-org:team_foo")
            .create_or_update(conn)
            .unwrap();
        let krate = CrateBuilder::new("foo", user.id).expect_build(conn);
        add_team_to_crate(&t, &krate, &user, conn).unwrap();
        t
    });

    let json = anon.search(&format!("team_id={}", team.id));
    assert_eq!(json.crates.len(), 1);
}

#[test]
fn crates_by_team_id_not_including_deleted_owners() {
    // This needs to use the proxy beacuse removing a team checks with github that you're on the
    // team before you're allowed to remove it from the crate
    let (app, anon) = TestApp::with_proxy().empty();
    let user = app.db_new_user(&mock_user_on_both_teams().gh_login);
    let user = user.as_model();

    let team = app.db(|conn| {
        let t = new_team("github:crates-test-org:core")
            .create_or_update(conn)
            .unwrap();
        let krate = CrateBuilder::new("foo", user.id).expect_build(conn);
        add_team_to_crate(&t, &krate, &user, conn).unwrap();
        krate
            .owner_remove(&app.as_inner(), conn, &user, &t.login)
            .unwrap();
        t
    });

    let json = anon.search(&format!("team_id={}", team.id));
    assert_eq!(json.crates.len(), 0);
}
<|endoftext|>


<commit_before>// Copyright 2017 Dario Domizioli ("hhexo").
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

// These tests are based on the book called "Seven Concurrency Models
// In Seven Weeks". Chapter 5 of the book is about the actor model. I am trying
// to replicate the exercises in the book by writing equivalent tests for the
// mecha library.

#[cfg(test)]
mod chapter5_day1 {
}
<commit_msg>Adding some Chapter5 tests<commit_after>// Copyright 2017 Dario Domizioli ("hhexo").
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

#[cfg(test)]
mod chapter5_day1 {
//! These tests are based on the book called "Seven Concurrency Models
//! In Seven Weeks". Chapter 5 of the book is about the actor model. I am trying
//! to replicate the exercises in the book by writing equivalent tests for the
//! mecha library.

use MessageType;
use MessageDatum;
use Message;
use ActorAddress;
use Actor;
use Stateless;

use std::collections::HashMap;
use std::thread;
use std::time::Duration;
use std::sync::mpsc;

const GREET: &'static str = ":greet";
const PRAISE: &'static str = ":praise";
const CELEBRATE: &'static str = ":celebrate";

#[test]
fn test_talker() {
    let (tx, rx) = mpsc::channel();
    let initiator = ActorAddress::new(tx.clone());

    let worker = Actor::new().with_state(Stateless)
        .with_match(|msg, _| {
            match *msg.get_type() {
                MessageType::Custom(GREET) => true,
                _ => false
            }
        })
        .with_action(|msg, _, _| {
            println!("Hello {}",
                msg.get_datum().clone().as_str().unwrap());
            Ok(())
        })
        .with_match(|msg, _| {
            match *msg.get_type() {
                MessageType::Custom(PRAISE) => true,
                _ => false
            }
        })
        .with_action(|msg, _, _| {
            println!("{}, you're amazing",
                msg.get_datum().clone().as_str().unwrap());
            Ok(())
        })
        .with_match(|msg, _| {
            match *msg.get_type() {
                MessageType::Custom(CELEBRATE) => true,
                _ => false
            }
        })
        .with_action(|msg, _, _| {
            println!("Here's to another {} years, {}",
                msg.get_datum().clone().as_map().unwrap().get("age").unwrap().clone().as_i64().unwrap(),
                msg.get_datum().clone().as_map().unwrap().get("name").unwrap().clone().as_str().unwrap());
            Ok(())
        })
        .spawn_link(&initiator);

    Message::custom(GREET).with_sender(&initiator)
                          .with_str("Huey")
                          .send_to(&worker);
    Message::custom(PRAISE).with_sender(&initiator)
                           .with_str("Dewey")
                           .send_to(&worker);
    let mut map = HashMap::new();
    map.insert("age".to_string(), MessageDatum::from(16i64));
    map.insert("name".to_string(), MessageDatum::from("Louie"));
    Message::custom(CELEBRATE).with_sender(&initiator)
                          .with_map(map)
                          .send_to(&worker);

    thread::sleep(Duration::from_millis(500));
    Message::shutdown().with_sender(&initiator).send_to(&worker);

    let msg = rx.recv().unwrap();
    assert_eq!(*msg.get_type(), MessageType::Exited);
    match *msg.get_datum() {
        MessageDatum::Void => (),
        _ => { assert!(false, "Unexpected message datum"); }
    }
}

#[derive(Default)]
struct CounterState { count: i64 }
struct CounterApi { counter: ActorAddress }

const COUNT: &'static str = ":count";
const COUNT_ACK: &'static str = ":count_ack";

impl CounterApi {
    pub fn new(parent: &ActorAddress) -> CounterApi {
        let act = Actor::new().with_state(CounterState {count: 0})
        .with_match(|msg, _| {
            match *msg.get_type() {
                MessageType::Custom(COUNT) => true,
                _ => false
            }
        })
        .with_action(|msg, state, myself| {
            state.count += 1;
            Message::custom(COUNT_ACK).with_sender(myself)
                                      .with_i64(state.count)
                                      .send_to(msg.get_sender());
            Ok(())
        })
        .spawn_link(parent);

        CounterApi {
            counter: act
        }
    }

    pub fn next(&self) -> i64 {
        let (tx, rx) = mpsc::channel();
        let initiator = ActorAddress::new(tx);
        Message::custom(COUNT).with_sender(&initiator)
                              .send_to(&self.counter);
        rx.recv().unwrap().get_datum().clone().as_i64().unwrap()
    }

    pub fn shutdown(&self) {
        Message::shutdown().send_to(&self.counter);
    }
}

#[test]
fn test_counter() {
    let (tx, rx) = mpsc::channel();
    let initiator = ActorAddress::new(tx.clone());

    let capi = CounterApi::new(&initiator);

    assert_eq!(capi.next(), 1i64);
    assert_eq!(capi.next(), 2i64);
    assert_eq!(capi.next(), 3i64);

    capi.shutdown();

    let msg = rx.recv().unwrap();
    assert_eq!(*msg.get_type(), MessageType::Exited);
    match *msg.get_datum() {
        MessageDatum::Void => (),
        _ => { assert!(false, "Unexpected message datum"); }
    }
}

}
<|endoftext|>