<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    common::{Author, Payload, Round},
    quorum_cert::QuorumCert,
    vote_data::VoteData,
};
use libra_crypto::hash::HashValue;
use libra_crypto_derive::{CryptoHasher, LCSCryptoHash};
use libra_types::{
    block_info::BlockInfo,
    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
};
use mirai_annotations::*;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
pub enum BlockType {
    Proposal {
        /// T of the block (e.g. one or more transaction(s)
        payload: Payload,
        /// Author of the block that can be validated by the author's public key and the signature
        author: Author,
    },
    /// NIL blocks don't have authors or signatures: they're generated upon timeouts to fill in the
    /// gaps in the rounds.
    NilBlock,
    /// A genesis block is the first committed block in any epoch that is identically constructed on
    /// all validators by any (potentially different) LedgerInfo that justifies the epoch change
    /// from the previous epoch.  The genesis block is used as the the first root block of the
    /// BlockTree for all epochs.
    Genesis,
}

#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq, CryptoHasher, LCSCryptoHash)]
/// Block has the core data of a consensus block that should be persistent when necessary.
/// Each block must know the id of its parent and keep the QuorurmCertificate to that parent.
pub struct BlockData {
    /// Epoch number corresponds to the set of validators that are active for this block.
    epoch: u64,
    /// The round of a block is an internal monotonically increasing counter used by Consensus
    /// protocol.
    round: Round,
    /// The approximate physical time a block is proposed by a proposer.  This timestamp is used
    /// for
    /// * Time-dependent logic in smart contracts (the current time of execution)
    /// * Clients determining if they are relatively up-to-date with respect to the block chain.
    ///
    /// It makes the following guarantees:
    /// 1. Time Monotonicity: Time is monotonically increasing in the block
    ///    chain. (i.e. If H1 < H2, H1.Time < H2.Time).
    /// 2. If a block of transactions B is agreed on with timestamp T, then at least f+1
    ///    honest replicas think that T is in the past.  An honest replica will only vote
    ///    on a block when its own clock >= timestamp T.
    /// 3. If a block of transactions B is agreed on with timestamp T, then at least f+1 honest
    ///    replicas saw the contents of B no later than T + delta for some delta.
    ///    If T = 3:00 PM and delta is 10 minutes, then an honest replica would not have
    ///    voted for B unless its clock was between 3:00 PM to 3:10 PM at the time the
    ///    proposal was received.  After 3:10 PM, an honest replica would no longer vote
    ///    on B, noting it was too far in the past.
    timestamp_usecs: u64,
    /// Contains the quorum certified ancestor and whether the quorum certified ancestor was
    /// voted on successfully
    quorum_cert: QuorumCert,
    /// If a block is a real proposal, contains its author and signature.
    block_type: BlockType,
}

impl BlockData {
    pub fn author(&self) -> Option<Author> {
        if let BlockType::Proposal { author, .. } = self.block_type {
            Some(author)
        } else {
            None
        }
    }

    pub fn block_type(&self) -> &BlockType {
        &self.block_type
    }

    pub fn epoch(&self) -> u64 {
        self.epoch
    }

    pub fn parent_id(&self) -> HashValue {
        self.quorum_cert.certified_block().id()
    }

    pub fn payload(&self) -> Option<&Payload> {
        if let BlockType::Proposal { payload, .. } = &self.block_type {
            Some(payload)
        } else {
            None
        }
    }

    pub fn round(&self) -> Round {
        // Round numbers:
        // - are reset to 0 periodically.
        // - do not exceed std::u64::MAX - 2 per the 3 chain safety rule
        // (ConsensusState::commit_rule_for_certified_block)
        assumed_postcondition!(self.round < std::u64::MAX - 1);
        self.round
    }

    pub fn timestamp_usecs(&self) -> u64 {
        self.timestamp_usecs
    }

    pub fn quorum_cert(&self) -> &QuorumCert {
        &self.quorum_cert
    }

    pub fn is_genesis_block(&self) -> bool {
        matches!(self.block_type, BlockType::Genesis)
    }

    pub fn is_nil_block(&self) -> bool {
        matches!(self.block_type, BlockType::NilBlock)
    }

    pub fn new_genesis_from_ledger_info(ledger_info: &LedgerInfo) -> Self {
        assert!(ledger_info.ends_epoch());
        let ancestor = BlockInfo::new(
            ledger_info.epoch(),
            0,                 /* round */
            HashValue::zero(), /* parent block id */
            ledger_info.transaction_accumulator_hash(),
            ledger_info.version(),
            ledger_info.timestamp_usecs(),
            None,
        );

        // Genesis carries a placeholder quorum certificate to its parent id with LedgerInfo
        // carrying information about version from the last LedgerInfo of previous epoch.
        let genesis_quorum_cert = QuorumCert::new(
            VoteData::new(ancestor.clone(), ancestor.clone()),
            LedgerInfoWithSignatures::new(
                LedgerInfo::new(ancestor, HashValue::zero()),
                BTreeMap::new(),
            ),
        );

        BlockData::new_genesis(ledger_info.timestamp_usecs(), genesis_quorum_cert)
    }

    pub fn new_genesis(timestamp_usecs: u64, quorum_cert: QuorumCert) -> Self {
        assume!(quorum_cert.certified_block().epoch() < u64::max_value()); // unlikely to be false in this universe
        Self {
            epoch: quorum_cert.certified_block().epoch() + 1,
            round: 0,
            timestamp_usecs,
            quorum_cert,
            block_type: BlockType::Genesis,
        }
    }

    pub fn new_nil(round: Round, quorum_cert: QuorumCert) -> Self {
        // We want all the NIL blocks to agree on the timestamps even though they're generated
        // independently by different validators, hence we're using the timestamp of a parent + 1.
        assume!(quorum_cert.certified_block().timestamp_usecs() < u64::max_value()); // unlikely to be false in this universe
        let timestamp_usecs = quorum_cert.certified_block().timestamp_usecs();

        Self {
            epoch: quorum_cert.certified_block().epoch(),
            round,
            timestamp_usecs,
            quorum_cert,
            block_type: BlockType::NilBlock,
        }
    }

    pub fn new_proposal(
        payload: Payload,
        author: Author,
        round: Round,
        timestamp_usecs: u64,
        quorum_cert: QuorumCert,
    ) -> Self {
        Self {
            epoch: quorum_cert.certified_block().epoch(),
            round,
            timestamp_usecs,
            quorum_cert,
            block_type: BlockType::Proposal { payload, author },
        }
    }
}
<commit_msg>[consensus] removing MIRAI assumption that is false<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    common::{Author, Payload, Round},
    quorum_cert::QuorumCert,
    vote_data::VoteData,
};
use libra_crypto::hash::HashValue;
use libra_crypto_derive::{CryptoHasher, LCSCryptoHash};
use libra_types::{
    block_info::BlockInfo,
    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
};
use mirai_annotations::*;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
pub enum BlockType {
    Proposal {
        /// T of the block (e.g. one or more transaction(s)
        payload: Payload,
        /// Author of the block that can be validated by the author's public key and the signature
        author: Author,
    },
    /// NIL blocks don't have authors or signatures: they're generated upon timeouts to fill in the
    /// gaps in the rounds.
    NilBlock,
    /// A genesis block is the first committed block in any epoch that is identically constructed on
    /// all validators by any (potentially different) LedgerInfo that justifies the epoch change
    /// from the previous epoch.  The genesis block is used as the the first root block of the
    /// BlockTree for all epochs.
    Genesis,
}

#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq, CryptoHasher, LCSCryptoHash)]
/// Block has the core data of a consensus block that should be persistent when necessary.
/// Each block must know the id of its parent and keep the QuorurmCertificate to that parent.
pub struct BlockData {
    /// Epoch number corresponds to the set of validators that are active for this block.
    epoch: u64,
    /// The round of a block is an internal monotonically increasing counter used by Consensus
    /// protocol.
    round: Round,
    /// The approximate physical time a block is proposed by a proposer.  This timestamp is used
    /// for
    /// * Time-dependent logic in smart contracts (the current time of execution)
    /// * Clients determining if they are relatively up-to-date with respect to the block chain.
    ///
    /// It makes the following guarantees:
    /// 1. Time Monotonicity: Time is monotonically increasing in the block
    ///    chain. (i.e. If H1 < H2, H1.Time < H2.Time).
    /// 2. If a block of transactions B is agreed on with timestamp T, then at least f+1
    ///    honest replicas think that T is in the past.  An honest replica will only vote
    ///    on a block when its own clock >= timestamp T.
    /// 3. If a block of transactions B is agreed on with timestamp T, then at least f+1 honest
    ///    replicas saw the contents of B no later than T + delta for some delta.
    ///    If T = 3:00 PM and delta is 10 minutes, then an honest replica would not have
    ///    voted for B unless its clock was between 3:00 PM to 3:10 PM at the time the
    ///    proposal was received.  After 3:10 PM, an honest replica would no longer vote
    ///    on B, noting it was too far in the past.
    timestamp_usecs: u64,
    /// Contains the quorum certified ancestor and whether the quorum certified ancestor was
    /// voted on successfully
    quorum_cert: QuorumCert,
    /// If a block is a real proposal, contains its author and signature.
    block_type: BlockType,
}

impl BlockData {
    pub fn author(&self) -> Option<Author> {
        if let BlockType::Proposal { author, .. } = self.block_type {
            Some(author)
        } else {
            None
        }
    }

    pub fn block_type(&self) -> &BlockType {
        &self.block_type
    }

    pub fn epoch(&self) -> u64 {
        self.epoch
    }

    pub fn parent_id(&self) -> HashValue {
        self.quorum_cert.certified_block().id()
    }

    pub fn payload(&self) -> Option<&Payload> {
        if let BlockType::Proposal { payload, .. } = &self.block_type {
            Some(payload)
        } else {
            None
        }
    }

    pub fn round(&self) -> Round {
        self.round
    }

    pub fn timestamp_usecs(&self) -> u64 {
        self.timestamp_usecs
    }

    pub fn quorum_cert(&self) -> &QuorumCert {
        &self.quorum_cert
    }

    pub fn is_genesis_block(&self) -> bool {
        matches!(self.block_type, BlockType::Genesis)
    }

    pub fn is_nil_block(&self) -> bool {
        matches!(self.block_type, BlockType::NilBlock)
    }

    pub fn new_genesis_from_ledger_info(ledger_info: &LedgerInfo) -> Self {
        assert!(ledger_info.ends_epoch());
        let ancestor = BlockInfo::new(
            ledger_info.epoch(),
            0,                 /* round */
            HashValue::zero(), /* parent block id */
            ledger_info.transaction_accumulator_hash(),
            ledger_info.version(),
            ledger_info.timestamp_usecs(),
            None,
        );

        // Genesis carries a placeholder quorum certificate to its parent id with LedgerInfo
        // carrying information about version from the last LedgerInfo of previous epoch.
        let genesis_quorum_cert = QuorumCert::new(
            VoteData::new(ancestor.clone(), ancestor.clone()),
            LedgerInfoWithSignatures::new(
                LedgerInfo::new(ancestor, HashValue::zero()),
                BTreeMap::new(),
            ),
        );

        BlockData::new_genesis(ledger_info.timestamp_usecs(), genesis_quorum_cert)
    }

    pub fn new_genesis(timestamp_usecs: u64, quorum_cert: QuorumCert) -> Self {
        assume!(quorum_cert.certified_block().epoch() < u64::max_value()); // unlikely to be false in this universe
        Self {
            epoch: quorum_cert.certified_block().epoch() + 1,
            round: 0,
            timestamp_usecs,
            quorum_cert,
            block_type: BlockType::Genesis,
        }
    }

    pub fn new_nil(round: Round, quorum_cert: QuorumCert) -> Self {
        // We want all the NIL blocks to agree on the timestamps even though they're generated
        // independently by different validators, hence we're using the timestamp of a parent + 1.
        assume!(quorum_cert.certified_block().timestamp_usecs() < u64::max_value()); // unlikely to be false in this universe
        let timestamp_usecs = quorum_cert.certified_block().timestamp_usecs();

        Self {
            epoch: quorum_cert.certified_block().epoch(),
            round,
            timestamp_usecs,
            quorum_cert,
            block_type: BlockType::NilBlock,
        }
    }

    pub fn new_proposal(
        payload: Payload,
        author: Author,
        round: Round,
        timestamp_usecs: u64,
        quorum_cert: QuorumCert,
    ) -> Self {
        Self {
            epoch: quorum_cert.certified_block().epoch(),
            round,
            timestamp_usecs,
            quorum_cert,
            block_type: BlockType::Proposal { payload, author },
        }
    }
}
<|endoftext|>


<commit_before>//! Check Disk usage

#[macro_use]
extern crate derive_more;
extern crate env_logger;
#[macro_use]
extern crate log;
extern crate nix;
extern crate regex;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate structopt;

extern crate tabin_plugins;

use std::collections::HashSet;
use std::cmp::max;
use std::fmt;

use structopt::StructOpt;
use regex::Regex;
use nix::sys::statvfs::vfs;
use tabin_plugins::Status;
use tabin_plugins::procfs::Mount;
use tabin_plugins::linux::bytes_to_human_size;

/// Check all mounted file systems for disk usage.
///
/// For some reason this check generally generates values that are between 1% and
/// 3% higher than `df`, even though AFAICT we're both just calling statvfs a bunch
/// of times.
#[derive(StructOpt, Deserialize, Debug)]
#[structopt(name = "check-disk (part of tabin-plugins)",
            raw(setting = "structopt::clap::AppSettings::ColoredHelp"))]
struct Args {
    #[structopt(short = "w", long = "warn", help = "Percent to warn at", default_value = "80")]
    warn: f64,
    #[structopt(short = "c", long = "crit", help = "Percent to go critical at",
                default_value = "90")]
    crit: f64,
    #[structopt(short = "W", long = "warn-inodes", help = "Percent of inode usage to warn at",
                default_value = "80")]
    warn_inodes: f64,
    #[structopt(short = "C", long = "crit-inodes",
                help = "Percent of inode usage to go critical at", default_value = "90")]
    crit_inodes: f64,

    #[structopt(long = "pattern", name = "regex",
                help = "Only check filesystems that match this regex")]
    pattern: Option<String>,
    #[structopt(long = "exclude-pattern", name = "exclude-regex",
                help = "Only check filesystems that match this regex")]
    exclude_pattern: Option<String>,
    #[structopt(long = "type", name = "fs-type",
                help = "Only check filesystems that are of this type, e.g. \
                        ext4 or tmpfs. See 'man 8 mount' for more examples.")]
    fs_type: Option<String>,
    #[structopt(long = "exclude-type", name = "exclude-fs-type",
                help = "Do not check filesystems that are of this type.")]
    exclude_type: Option<String>,
    #[structopt(long = "info",
                help = "Print information of all known filesystems. \
                        Similar to df.")]
    info: bool,
    // df defaults to ignoring innaccessible filesystems, so we should too
    #[structopt(long = "inaccessible-status", name = "STATUS",
                help = "If any filesystems are inaccessible print a warning and exit with STATUS. \
                        Choices: [critical, warning, ok]")]
    inaccessible_status: Option<Status>,
}

const LOG_VAR: &str = "TABIN_LOG";

fn main() {
    let args = Args::from_args();
    env_logger::Builder::from_env(LOG_VAR).init();

    let mut mounts = match Mount::load_all() {
        Ok(mounts) => mounts,
        Err(e) => {
            println!("CRITICAL error loading mounts: {}", e);
            Status::Critical.exit();
        }
    };
    mounts.sort_by(|l, r| l.file.len().cmp(&r.file.len()));

    let status = match filter(mounts, &args) {
        Ok(ms) => do_check(&ms, &args),
        Err(Error::NotAccessible {
            accessible,
            not_accessible,
        }) => {
            let check_result = do_check(&accessible, &args);
            if args.inaccessible_status.is_some() {
                let status = args.inaccessible_status.unwrap();
                println!(
                    "{}: {} filesystems were not accessible, \
                     run with TABIN_LOG=debug for details",
                    status, not_accessible
                );
                max(check_result, status)
            } else {
                check_result
            }
        }
        Err(e) => {
            println!("{}", e);
            Status::Critical
        }
    };

    status.exit();
}

#[derive(Debug, PartialEq, Eq)]
struct ErrorMsg {
    msg: String,
}

#[derive(Debug, From)]
enum Error {
    Message(ErrorMsg),
    NotAccessible {
        accessible: Vec<MountStat>,
        not_accessible: u32,
    },
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::Message(ref e) => write!(f, "{}", e.msg),
            Error::NotAccessible {
                ref not_accessible, ..
            } => write!(f, "{} directories were not accesible", not_accessible),
        }
    }
}

type DiskResult<T> = Result<T, Error>;

#[derive(Debug)]
struct MountStat {
    mount: Mount,
    stat: vfs::Statvfs,
}

fn percent(part: u64, whole: u64) -> f64 {
    100.0 - (part as f64 / whole as f64) * 100.0
}

fn maybe_regex(pattern: &Option<String>) -> DiskResult<Option<Regex>> {
    if let Some(ref pattern) = *pattern {
        let re = match Regex::new(&pattern) {
            Ok(re) => re,
            Err(e) => {
                return Err(ErrorMsg {
                    msg: format!("Unable to filter disks like {:?}: {}", pattern, e),
                }.into())
            }
        };
        Ok(Some(re))
    } else {
        Ok(None)
    }
}

/// Convert Mounts into MountStats, applying filters from args
///
/// This:
///
/// * calls `statvfs` on every filesystem in `/proc/mounts`
/// * filters out dummy (0-block) and /proc filesystems
/// * Only shows one of any given `/dev/` filesystem's mount points (the
///   shortest, same as df)
/// * Applies the `pattern` and `type` filters
fn filter(mounts: Vec<Mount>, args: &Args) -> DiskResult<Vec<MountStat>> {
    let mut devices = HashSet::new();
    let include_regex = try!(maybe_regex(&args.pattern));
    let exclude_regex = try!(maybe_regex(&args.exclude_pattern));
    let mut error_count = 0;
    let ms = mounts.into_iter()
        .filter(|mount|
                if let Some(ref re) = include_regex {
                    re.is_match(&mount.file)
                } else {
                    true
                })
        .filter(|mount|
                if let Some(ref re) = exclude_regex {
                    !re.is_match(&mount.file)
                } else {
                    true
                })
        .filter_map(|mount| {
            let stat = match vfs::Statvfs::for_path(mount.file.as_bytes()) {
                Ok(stat) => stat,
                Err(e) => {
                    error_count += 1;
                    debug!("Error reading statvfs for path {}: {}", mount.file, e);
                    return None;
                }
            };
            if stat.f_blocks > 0 && !mount.file.starts_with("/proc") {
                Some(MountStat {
                    mount: mount,
                    stat: stat,
                })
            } else {
                None
            }
        })
        // If the same device is mounted multiple times, we only want one. In
        // To match df we expect these to come in in sorted order and just keep
        // the first one.
        .filter_map(|ms| {
            if ms.mount.spec.starts_with("/dev") {
                if devices.contains(&ms.mount.spec) {
                    None
                } else {
                    devices.insert(ms.mount.spec.clone());
                    Some(ms)
                }
            } else {
                Some(ms)
            }
        })
        .filter(|ms|
                if let Some(ref vfstype) = args.fs_type {
                    ms.mount.vfstype == *vfstype
                } else {
                    true
                })
        .filter(|ms|
                if let Some(ref vfstype) = args.exclude_type {
                    ms.mount.vfstype != *vfstype
                } else {
                    true
                })
        .collect::<Vec<_>>();
    if error_count == 0 {
        Ok(ms)
    } else {
        Err(Error::NotAccessible {
            accessible: ms,
            not_accessible: error_count,
        })
    }
}

fn do_check(mountstats: &[MountStat], args: &Args) -> Status {
    let mut status = Status::Ok;
    for ms in mountstats {
        let pcnt = percent(ms.stat.f_bavail, ms.stat.f_blocks);
        if pcnt > args.crit {
            status = Status::Critical;
            println!(
                "CRITICAL: {} has {:.1}% of its {}B used (> {:.1}%)",
                ms.mount.file,
                pcnt,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                args.crit
            )
        } else if pcnt > args.warn {
            status = max(status, Status::Warning);
            println!(
                "WARNING: {} has {:.1}% of its {}B used (> {:.1}%)",
                ms.mount.file,
                pcnt,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                args.warn
            )
        }

        let ipcnt = percent(ms.stat.f_favail, ms.stat.f_files);
        if ipcnt > args.crit_inodes {
            status = Status::Critical;
            println!(
                "CRITICAL: {} has {:.1}% of its {} inodes used (> {:.1}%)",
                ms.mount.file,
                ipcnt,
                bytes_to_human_size(ms.stat.f_files),
                args.crit_inodes
            );
        } else if ipcnt > args.warn_inodes {
            status = max(status, Status::Warning);
            println!(
                "WARNING: {} has {:.1}% of its {} inodes used (> {:.1}%)",
                ms.mount.file,
                ipcnt,
                bytes_to_human_size(ms.stat.f_files),
                args.warn_inodes
            );
        }
    }
    if status == Status::Ok {
        println!(
            "OKAY: {} filesystems checked, none are above {}% disk or {}% inode usage",
            mountstats.len(),
            args.warn,
            args.warn_inodes
        );
    }

    if args.info {
        println!(
            "{:<15} {:>7} {:>5}% {:>7} {:>5}% {:<20}",
            "Filesystem", "Size", "Use", "INodes", "IUse", "Mounted on"
        );
        for ms in mountstats {
            println!(
                "{:<15} {:>7} {:>5.1}% {:>7} {:>5.1}% {:<20}",
                ms.mount.spec,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                percent(ms.stat.f_bavail, ms.stat.f_blocks),
                bytes_to_human_size(ms.stat.f_files),
                percent(ms.stat.f_favail, ms.stat.f_files),
                ms.mount.file
            );
        }
    }

    status
}

#[cfg(test)]
mod unit {
    use super::{maybe_regex, Args};
    use structopt::StructOpt;

    #[test]
    fn validate_docstring() {
        let args: Args = Args::from_iter(["arg0", "--crit", "5"].into_iter());
        assert_eq!(args.crit, 5.0);
        let args: Args = Args::from_iter(["arg0", "--pattern", "hello"].into_iter());
        assert_eq!(args.pattern.unwrap(), "hello");
    }

    #[test]
    fn check_maybe_regex() {
        if let Err(emsg) = maybe_regex(&Some("[hello".to_owned())) {
            let expected = r#"Unable to filter disks like "[hello":"#;
            assert!(
                emsg.to_string().contains(expected),
                "\nExpected something containing: {}\n\
                 But instead received         : {}",
                expected,
                emsg.to_string()
            );
        } else {
            panic!("Should have gotten an error");
        }
    }
}
<commit_msg>check-disk: Shorten up some of the filters<commit_after>//! Check Disk usage

#[macro_use]
extern crate derive_more;
extern crate env_logger;
#[macro_use]
extern crate log;
extern crate nix;
extern crate regex;
#[macro_use]
extern crate serde_derive;
#[macro_use]
extern crate structopt;

extern crate tabin_plugins;

use std::collections::HashSet;
use std::cmp::max;
use std::fmt;

use structopt::StructOpt;
use regex::Regex;
use nix::sys::statvfs::vfs;
use tabin_plugins::Status;
use tabin_plugins::procfs::Mount;
use tabin_plugins::linux::bytes_to_human_size;

/// Check all mounted file systems for disk usage.
///
/// For some reason this check generally generates values that are between 1% and
/// 3% higher than `df`, even though AFAICT we're both just calling statvfs a bunch
/// of times.
#[derive(StructOpt, Deserialize, Debug)]
#[structopt(name = "check-disk (part of tabin-plugins)",
            raw(setting = "structopt::clap::AppSettings::ColoredHelp"))]
struct Args {
    #[structopt(short = "w", long = "warn", help = "Percent to warn at", default_value = "80")]
    warn: f64,
    #[structopt(short = "c", long = "crit", help = "Percent to go critical at",
                default_value = "90")]
    crit: f64,
    #[structopt(short = "W", long = "warn-inodes", help = "Percent of inode usage to warn at",
                default_value = "80")]
    warn_inodes: f64,
    #[structopt(short = "C", long = "crit-inodes",
                help = "Percent of inode usage to go critical at", default_value = "90")]
    crit_inodes: f64,

    #[structopt(long = "pattern", name = "regex",
                help = "Only check filesystems that match this regex")]
    pattern: Option<String>,
    #[structopt(long = "exclude-pattern", name = "exclude-regex",
                help = "Only check filesystems that match this regex")]
    exclude_pattern: Option<String>,
    #[structopt(long = "type", name = "fs-type",
                help = "Only check filesystems that are of this type, e.g. \
                        ext4 or tmpfs. See 'man 8 mount' for more examples.")]
    fs_type: Option<String>,
    #[structopt(long = "exclude-type", name = "exclude-fs-type",
                help = "Do not check filesystems that are of this type.")]
    exclude_type: Option<String>,
    #[structopt(long = "info",
                help = "Print information of all known filesystems. \
                        Similar to df.")]
    info: bool,
    // df defaults to ignoring innaccessible filesystems, so we should too
    #[structopt(long = "inaccessible-status", name = "STATUS",
                help = "If any filesystems are inaccessible print a warning and exit with STATUS. \
                        Choices: [critical, warning, ok]")]
    inaccessible_status: Option<Status>,
}

const LOG_VAR: &str = "TABIN_LOG";

fn main() {
    let args = Args::from_args();
    env_logger::Builder::from_env(LOG_VAR).init();

    let mut mounts = match Mount::load_all() {
        Ok(mounts) => mounts,
        Err(e) => {
            println!("CRITICAL error loading mounts: {}", e);
            Status::Critical.exit();
        }
    };
    mounts.sort_by(|l, r| l.file.len().cmp(&r.file.len()));

    let status = match filter(mounts, &args) {
        Ok(ms) => do_check(&ms, &args),
        Err(Error::NotAccessible {
            accessible,
            not_accessible,
        }) => {
            let check_result = do_check(&accessible, &args);
            if args.inaccessible_status.is_some() {
                let status = args.inaccessible_status.unwrap();
                println!(
                    "{}: {} filesystems were not accessible, \
                     run with TABIN_LOG=debug for details",
                    status, not_accessible
                );
                max(check_result, status)
            } else {
                check_result
            }
        }
        Err(e) => {
            println!("{}", e);
            Status::Critical
        }
    };

    status.exit();
}

#[derive(Debug, PartialEq, Eq)]
struct ErrorMsg {
    msg: String,
}

#[derive(Debug, From)]
enum Error {
    Message(ErrorMsg),
    NotAccessible {
        accessible: Vec<MountStat>,
        not_accessible: u32,
    },
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        match *self {
            Error::Message(ref e) => write!(f, "{}", e.msg),
            Error::NotAccessible {
                ref not_accessible, ..
            } => write!(f, "{} directories were not accesible", not_accessible),
        }
    }
}

type DiskResult<T> = Result<T, Error>;

#[derive(Debug)]
struct MountStat {
    mount: Mount,
    stat: vfs::Statvfs,
}

fn percent(part: u64, whole: u64) -> f64 {
    100.0 - (part as f64 / whole as f64) * 100.0
}

fn maybe_regex(pattern: &Option<String>) -> DiskResult<Option<Regex>> {
    if let Some(ref pattern) = *pattern {
        let re = match Regex::new(&pattern) {
            Ok(re) => re,
            Err(e) => {
                return Err(ErrorMsg {
                    msg: format!("Unable to filter disks like {:?}: {}", pattern, e),
                }.into())
            }
        };
        Ok(Some(re))
    } else {
        Ok(None)
    }
}

/// Convert Mounts into MountStats, applying filters from args
///
/// This:
///
/// * calls `statvfs` on every filesystem in `/proc/mounts`
/// * filters out dummy (0-block) and /proc filesystems
/// * Only shows one of any given `/dev/` filesystem's mount points (the
///   shortest, same as df)
/// * Applies the `pattern` and `type` filters
fn filter(mounts: Vec<Mount>, args: &Args) -> DiskResult<Vec<MountStat>> {
    let mut devices = HashSet::new();
    let include_regex = try!(maybe_regex(&args.pattern));
    let exclude_regex = try!(maybe_regex(&args.exclude_pattern));
    let mut error_count = 0;
    let ms = mounts.into_iter()
        .filter(|mount|
                include_regex.as_ref().map(|re| re.is_match(&mount.file)).unwrap_or(true)
        )
        .filter(|mount|
                exclude_regex.as_ref().map(|re| !re.is_match(&mount.file)).unwrap_or(true)
        )
        .filter_map(|mount| {
            let stat = match vfs::Statvfs::for_path(mount.file.as_bytes()) {
                Ok(stat) => stat,
                Err(e) => {
                    error_count += 1;
                    debug!("Error reading statvfs for path {}: {}", mount.file, e);
                    return None;
                }
            };
            if stat.f_blocks > 0 && !mount.file.starts_with("/proc") {
                Some(MountStat {
                    mount: mount,
                    stat: stat,
                })
            } else {
                None
            }
        })
        // If the same device is mounted multiple times, we only want one. In
        // To match df we expect these to come in in sorted order and just keep
        // the first one.
        .filter_map(|ms| {
            if ms.mount.spec.starts_with("/dev") {
                if devices.contains(&ms.mount.spec) {
                    None
                } else {
                    devices.insert(ms.mount.spec.clone());
                    Some(ms)
                }
            } else {
                Some(ms)
            }
        })
        .filter(|ms|
                args.fs_type.as_ref().map(|vfstype| ms.mount.vfstype == *vfstype).unwrap_or(true)
        )
        .filter(|ms|
                args.exclude_type.as_ref().map(|vfstype| ms.mount.vfstype != *vfstype).unwrap_or(true)
        )
        .collect::<Vec<_>>();
    if error_count == 0 {
        Ok(ms)
    } else {
        Err(Error::NotAccessible {
            accessible: ms,
            not_accessible: error_count,
        })
    }
}

fn do_check(mountstats: &[MountStat], args: &Args) -> Status {
    let mut status = Status::Ok;
    for ms in mountstats {
        let pcnt = percent(ms.stat.f_bavail, ms.stat.f_blocks);
        if pcnt > args.crit {
            status = Status::Critical;
            println!(
                "CRITICAL: {} has {:.1}% of its {}B used (> {:.1}%)",
                ms.mount.file,
                pcnt,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                args.crit
            )
        } else if pcnt > args.warn {
            status = max(status, Status::Warning);
            println!(
                "WARNING: {} has {:.1}% of its {}B used (> {:.1}%)",
                ms.mount.file,
                pcnt,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                args.warn
            )
        }

        let ipcnt = percent(ms.stat.f_favail, ms.stat.f_files);
        if ipcnt > args.crit_inodes {
            status = Status::Critical;
            println!(
                "CRITICAL: {} has {:.1}% of its {} inodes used (> {:.1}%)",
                ms.mount.file,
                ipcnt,
                bytes_to_human_size(ms.stat.f_files),
                args.crit_inodes
            );
        } else if ipcnt > args.warn_inodes {
            status = max(status, Status::Warning);
            println!(
                "WARNING: {} has {:.1}% of its {} inodes used (> {:.1}%)",
                ms.mount.file,
                ipcnt,
                bytes_to_human_size(ms.stat.f_files),
                args.warn_inodes
            );
        }
    }
    if status == Status::Ok {
        println!(
            "OKAY: {} filesystems checked, none are above {}% disk or {}% inode usage",
            mountstats.len(),
            args.warn,
            args.warn_inodes
        );
    }

    if args.info {
        println!(
            "{:<15} {:>7} {:>5}% {:>7} {:>5}% {:<20}",
            "Filesystem", "Size", "Use", "INodes", "IUse", "Mounted on"
        );
        for ms in mountstats {
            println!(
                "{:<15} {:>7} {:>5.1}% {:>7} {:>5.1}% {:<20}",
                ms.mount.spec,
                bytes_to_human_size(ms.stat.f_blocks * ms.stat.f_frsize),
                percent(ms.stat.f_bavail, ms.stat.f_blocks),
                bytes_to_human_size(ms.stat.f_files),
                percent(ms.stat.f_favail, ms.stat.f_files),
                ms.mount.file
            );
        }
    }

    status
}

#[cfg(test)]
mod unit {
    use super::{maybe_regex, Args};
    use structopt::StructOpt;

    #[test]
    fn validate_docstring() {
        let args: Args = Args::from_iter(["arg0", "--crit", "5"].into_iter());
        assert_eq!(args.crit, 5.0);
        let args: Args = Args::from_iter(["arg0", "--pattern", "hello"].into_iter());
        assert_eq!(args.pattern.unwrap(), "hello");
    }

    #[test]
    fn check_maybe_regex() {
        if let Err(emsg) = maybe_regex(&Some("[hello".to_owned())) {
            let expected = r#"Unable to filter disks like "[hello":"#;
            assert!(
                emsg.to_string().contains(expected),
                "\nExpected something containing: {}\n\
                 But instead received         : {}",
                expected,
                emsg.to_string()
            );
        } else {
            panic!("Should have gotten an error");
        }
    }
}
<|endoftext|>


<commit_before>// This file is part of the uutils coreutils package.
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

pub use crate::features::entries;
use libc::{self, gid_t, lchown, uid_t};

use std::io::Error as IOError;
use std::io::Result as IOResult;

use std::ffi::CString;
use std::fs::Metadata;
use std::os::unix::fs::MetadataExt;

use std::os::unix::ffi::OsStrExt;
use std::path::Path;

#[derive(PartialEq, Clone, Debug)]
pub enum Verbosity {
    Silent,
    Changes,
    Verbose,
    Normal,
}

fn chgrp<P: AsRef<Path>>(path: P, dgid: gid_t, follow: bool) -> IOResult<()> {
    let path = path.as_ref();
    let s = CString::new(path.as_os_str().as_bytes()).unwrap();
    let ret = unsafe {
        if follow {
            libc::chown(s.as_ptr(), (0 as gid_t).wrapping_sub(1), dgid)
        } else {
            lchown(s.as_ptr(), (0 as gid_t).wrapping_sub(1), dgid)
        }
    };
    if ret == 0 {
        Ok(())
    } else {
        Err(IOError::last_os_error())
    }
}

pub fn wrap_chgrp<P: AsRef<Path>>(
    path: P,
    meta: &Metadata,
    dest_gid: gid_t,
    follow: bool,
    verbosity: Verbosity,
) -> Result<String, String> {
    use self::Verbosity::*;
    let path = path.as_ref();
    let mut out: String = String::new();

    if let Err(e) = chgrp(path, dest_gid, follow) {
        match verbosity {
            Silent => (),
            _ => {
                out = format!("changing group of '{}': {}", path.display(), e);
                if verbosity == Verbose {
                    out = format!(
                        "{}\nfailed to change group of '{}' from {} to {}",
                        out,
                        path.display(),
                        entries::gid2grp(meta.gid()).unwrap(),
                        entries::gid2grp(dest_gid).unwrap()
                    );
                };
            }
        }
        return Err(out);
    } else {
        let changed = dest_gid != meta.gid();
        if changed {
            match verbosity {
                Changes | Verbose => {
                    out = format!(
                        "changed group of '{}' from {} to {}",
                        path.display(),
                        entries::gid2grp(meta.gid()).unwrap(),
                        entries::gid2grp(dest_gid).unwrap()
                    );
                }
                _ => (),
            };
        } else if verbosity == Verbose {
            out = format!(
                "group of '{}' retained as {}",
                path.display(),
                entries::gid2grp(dest_gid).unwrap()
            );
        }
    }
    Ok(out)
}

fn chown<P: AsRef<Path>>(path: P, duid: uid_t, dgid: gid_t, follow: bool) -> IOResult<()> {
    let path = path.as_ref();
    let s = CString::new(path.as_os_str().as_bytes()).unwrap();
    let ret = unsafe {
        if follow {
            libc::chown(s.as_ptr(), duid, dgid)
        } else {
            lchown(s.as_ptr(), duid, dgid)
        }
    };
    if ret == 0 {
        Ok(())
    } else {
        Err(IOError::last_os_error())
    }
}

pub fn wrap_chown<P: AsRef<Path>>(
    path: P,
    meta: &Metadata,
    dest_uid: Option<u32>,
    dest_gid: Option<u32>,
    follow: bool,
    verbosity: Verbosity,
) -> Result<String, String> {
    use self::Verbosity::*;
    let dest_uid = dest_uid.unwrap_or_else(|| meta.uid());
    let dest_gid = dest_gid.unwrap_or_else(|| meta.gid());
    let path = path.as_ref();
    let mut out: String = String::new();

    if let Err(e) = chown(path, dest_uid, dest_gid, follow) {
        match verbosity {
            Silent => (),
            _ => {
                out = format!("changing ownership of '{}': {}", path.display(), e);
                if verbosity == Verbose {
                    out = format!(
                        "{}\nfailed to change ownership of '{}' from {}:{} to {}:{}",
                        out,
                        path.display(),
                        entries::uid2usr(meta.uid()).unwrap(),
                        entries::gid2grp(meta.gid()).unwrap(),
                        entries::uid2usr(dest_uid).unwrap(),
                        entries::gid2grp(dest_gid).unwrap()
                    );
                };
            }
        }
        return Err(out);
    } else {
        let changed = dest_uid != meta.uid() || dest_gid != meta.gid();
        if changed {
            match verbosity {
                Changes | Verbose => {
                    out = format!(
                        "changed ownership of '{}' from {}:{} to {}:{}",
                        path.display(),
                        entries::uid2usr(meta.uid()).unwrap(),
                        entries::gid2grp(meta.gid()).unwrap(),
                        entries::uid2usr(dest_uid).unwrap(),
                        entries::gid2grp(dest_gid).unwrap()
                    );
                }
                _ => (),
            };
        } else if verbosity == Verbose {
            out = format!(
                "ownership of '{}' retained as {}:{}",
                path.display(),
                entries::uid2usr(dest_uid).unwrap(),
                entries::gid2grp(dest_gid).unwrap()
            );
        }
    }
    Ok(out)
}
<commit_msg>perms: add some docs<commit_after>// This file is part of the uutils coreutils package.
//
// For the full copyright and license information, please view the LICENSE
// file that was distributed with this source code.

pub use crate::features::entries;
use libc::{self, gid_t, lchown, uid_t};

use std::io::Error as IOError;
use std::io::Result as IOResult;

use std::ffi::CString;
use std::fs::Metadata;
use std::os::unix::fs::MetadataExt;

use std::os::unix::ffi::OsStrExt;
use std::path::Path;

/// The various level of verbosity
#[derive(PartialEq, Clone, Debug)]
pub enum Verbosity {
    Silent,
    Changes,
    Verbose,
    Normal,
}

/// Actually perform the change of group on a path
fn chgrp<P: AsRef<Path>>(path: P, dgid: gid_t, follow: bool) -> IOResult<()> {
    let path = path.as_ref();
    let s = CString::new(path.as_os_str().as_bytes()).unwrap();
    let ret = unsafe {
        if follow {
            libc::chown(s.as_ptr(), (0 as gid_t).wrapping_sub(1), dgid)
        } else {
            lchown(s.as_ptr(), (0 as gid_t).wrapping_sub(1), dgid)
        }
    };
    if ret == 0 {
        Ok(())
    } else {
        Err(IOError::last_os_error())
    }
}

/// Perform the change of group on a path
/// with the various options
/// and error messages management
pub fn wrap_chgrp<P: AsRef<Path>>(
    path: P,
    meta: &Metadata,
    dest_gid: gid_t,
    follow: bool,
    verbosity: Verbosity,
) -> Result<String, String> {
    use self::Verbosity::*;
    let path = path.as_ref();
    let mut out: String = String::new();

    if let Err(e) = chgrp(path, dest_gid, follow) {
        match verbosity {
            Silent => (),
            _ => {
                out = format!("changing group of '{}': {}", path.display(), e);
                if verbosity == Verbose {
                    out = format!(
                        "{}\nfailed to change group of '{}' from {} to {}",
                        out,
                        path.display(),
                        entries::gid2grp(meta.gid()).unwrap(),
                        entries::gid2grp(dest_gid).unwrap()
                    );
                };
            }
        }
        return Err(out);
    } else {
        let changed = dest_gid != meta.gid();
        if changed {
            match verbosity {
                Changes | Verbose => {
                    out = format!(
                        "changed group of '{}' from {} to {}",
                        path.display(),
                        entries::gid2grp(meta.gid()).unwrap(),
                        entries::gid2grp(dest_gid).unwrap()
                    );
                }
                _ => (),
            };
        } else if verbosity == Verbose {
            out = format!(
                "group of '{}' retained as {}",
                path.display(),
                entries::gid2grp(dest_gid).unwrap()
            );
        }
    }
    Ok(out)
}

/// Actually perform the change of owner on a path
fn chown<P: AsRef<Path>>(path: P, duid: uid_t, dgid: gid_t, follow: bool) -> IOResult<()> {
    let path = path.as_ref();
    let s = CString::new(path.as_os_str().as_bytes()).unwrap();
    let ret = unsafe {
        if follow {
            libc::chown(s.as_ptr(), duid, dgid)
        } else {
            lchown(s.as_ptr(), duid, dgid)
        }
    };
    if ret == 0 {
        Ok(())
    } else {
        Err(IOError::last_os_error())
    }
}

/// Perform the change of owner on a path
/// with the various options
/// and error messages management
pub fn wrap_chown<P: AsRef<Path>>(
    path: P,
    meta: &Metadata,
    dest_uid: Option<u32>,
    dest_gid: Option<u32>,
    follow: bool,
    verbosity: Verbosity,
) -> Result<String, String> {
    use self::Verbosity::*;
    let dest_uid = dest_uid.unwrap_or_else(|| meta.uid());
    let dest_gid = dest_gid.unwrap_or_else(|| meta.gid());
    let path = path.as_ref();
    let mut out: String = String::new();

    if let Err(e) = chown(path, dest_uid, dest_gid, follow) {
        match verbosity {
            Silent => (),
            _ => {
                out = format!("changing ownership of '{}': {}", path.display(), e);
                if verbosity == Verbose {
                    out = format!(
                        "{}\nfailed to change ownership of '{}' from {}:{} to {}:{}",
                        out,
                        path.display(),
                        entries::uid2usr(meta.uid()).unwrap(),
                        entries::gid2grp(meta.gid()).unwrap(),
                        entries::uid2usr(dest_uid).unwrap(),
                        entries::gid2grp(dest_gid).unwrap()
                    );
                };
            }
        }
        return Err(out);
    } else {
        let changed = dest_uid != meta.uid() || dest_gid != meta.gid();
        if changed {
            match verbosity {
                Changes | Verbose => {
                    out = format!(
                        "changed ownership of '{}' from {}:{} to {}:{}",
                        path.display(),
                        entries::uid2usr(meta.uid()).unwrap(),
                        entries::gid2grp(meta.gid()).unwrap(),
                        entries::uid2usr(dest_uid).unwrap(),
                        entries::gid2grp(dest_gid).unwrap()
                    );
                }
                _ => (),
            };
        } else if verbosity == Verbose {
            out = format!(
                "ownership of '{}' retained as {}:{}",
                path.display(),
                entries::uid2usr(dest_uid).unwrap(),
                entries::gid2grp(dest_gid).unwrap()
            );
        }
    }
    Ok(out)
}
<|endoftext|>


<commit_before>//! This example leverages `BytesCodec` to create a UDP client and server which
//! speak a custom protocol.
//!
//! Here we're using the codec from tokio-io to convert a UDP socket to a stream of
//! client messages. These messages are then processed and returned back as a
//! new message with a new destination. Overall, we then use this to construct a
//! "ping pong" pair where two sockets are sending messages back and forth.

#![deny(warnings)]

extern crate tokio;
extern crate tokio_io;
extern crate env_logger;

use std::net::SocketAddr;

use tokio::prelude::*;
use tokio::net::{UdpSocket, UdpFramed};
use tokio_io::codec::BytesCodec;

fn main() {
    let _ = env_logger::init();

    let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();

    // Bind both our sockets and then figure out what ports we got.
    let a = UdpSocket::bind(&addr).unwrap();
    let b = UdpSocket::bind(&addr).unwrap();
    let b_addr = b.local_addr().unwrap();

    // We're parsing each socket with the `LineCodec` defined above, and then we
    // `split` each codec into the sink/stream halves.
    let (a_sink, a_stream) = UdpFramed::new(a, BytesCodec::new()).split();
    let (b_sink, b_stream) = UdpFramed::new(b, BytesCodec::new()).split();

    // Start off by sending a ping from a to b, afterwards we just print out
    // what they send us and continually send pings
    // let pings = stream::iter((0..5).map(Ok));
    let a = a_sink.send(("PING".into(), b_addr)).and_then(|a_sink| {
        let mut i = 0;
        let a_stream = a_stream.take(4).map(move |(msg, addr)| {
            i += 1;
            println!("[a] recv: {}", String::from_utf8_lossy(&msg));
            (format!("PING {}", i).into(), addr)
        });
        a_sink.send_all(a_stream)
    });

    // The second client we have will receive the pings from `a` and then send
    // back pongs.
    let b_stream = b_stream.map(|(msg, addr)| {
        println!("[b] recv: {}", String::from_utf8_lossy(&msg));
        ("PONG".into(), addr)
    });
    let b = b_sink.send_all(b_stream);

    // Spawn the sender of pongs and then wait for our pinger to finish.
    tokio::run({
        b.join(a)
            .map(|_| ())
            .map_err(|e| println!("error = {:?}", e))
    });
}
<commit_msg>Update comment in udp-codec example (#222)<commit_after>//! This example leverages `BytesCodec` to create a UDP client and server which
//! speak a custom protocol.
//!
//! Here we're using the codec from tokio-io to convert a UDP socket to a stream of
//! client messages. These messages are then processed and returned back as a
//! new message with a new destination. Overall, we then use this to construct a
//! "ping pong" pair where two sockets are sending messages back and forth.

#![deny(warnings)]

extern crate tokio;
extern crate tokio_io;
extern crate env_logger;

use std::net::SocketAddr;

use tokio::prelude::*;
use tokio::net::{UdpSocket, UdpFramed};
use tokio_io::codec::BytesCodec;

fn main() {
    let _ = env_logger::init();

    let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();

    // Bind both our sockets and then figure out what ports we got.
    let a = UdpSocket::bind(&addr).unwrap();
    let b = UdpSocket::bind(&addr).unwrap();
    let b_addr = b.local_addr().unwrap();

    // We're parsing each socket with the `BytesCodec` included in `tokio_io`, and then we
    // `split` each codec into the sink/stream halves.
    let (a_sink, a_stream) = UdpFramed::new(a, BytesCodec::new()).split();
    let (b_sink, b_stream) = UdpFramed::new(b, BytesCodec::new()).split();

    // Start off by sending a ping from a to b, afterwards we just print out
    // what they send us and continually send pings
    // let pings = stream::iter((0..5).map(Ok));
    let a = a_sink.send(("PING".into(), b_addr)).and_then(|a_sink| {
        let mut i = 0;
        let a_stream = a_stream.take(4).map(move |(msg, addr)| {
            i += 1;
            println!("[a] recv: {}", String::from_utf8_lossy(&msg));
            (format!("PING {}", i).into(), addr)
        });
        a_sink.send_all(a_stream)
    });

    // The second client we have will receive the pings from `a` and then send
    // back pongs.
    let b_stream = b_stream.map(|(msg, addr)| {
        println!("[b] recv: {}", String::from_utf8_lossy(&msg));
        ("PONG".into(), addr)
    });
    let b = b_sink.send_all(b_stream);

    // Spawn the sender of pongs and then wait for our pinger to finish.
    tokio::run({
        b.join(a)
            .map(|_| ())
            .map_err(|e| println!("error = {:?}", e))
    });
}
<|endoftext|>


<commit_before>use super::invoke::{Invoke, Identity};
use super::code::{self, SignaturesBuilder, FunctionBuilder};
use super::import;
use elements;

/// Module builder
pub struct ModuleBuilder<F=Identity> {
    callback: F,
    module: ModuleScaffold,
}

/// Location of the internal module function
pub struct CodeLocation {
    /// Location (index in 'functions' section) of the signature
    pub signature: u32,
    /// Location (index in the 'code' section) of the body
    pub body: u32,
}

#[derive(Default)]
struct ModuleScaffold {
    pub functions: elements::FunctionsSection,
    pub types: elements::TypeSection,
    pub import: elements::ImportSection,
    pub code: elements::CodeSection,
    pub other: Vec<elements::Section>,
}

impl From<elements::Module> for ModuleScaffold {
    fn from(module: elements::Module) -> Self {
        let mut funcs: Option<elements::FunctionsSection> = None;
        let mut types: Option<elements::TypeSection> = None;
        let mut import: Option<elements::ImportSection> = None;
        let mut code: Option<elements::CodeSection> = None;

        let mut sections = module.into_sections();
        while let Some(section) = sections.pop() {
            match section {
                elements::Section::Type(sect) => { types = Some(sect); }
                elements::Section::Function(sect) => { funcs = Some(sect); }
                elements::Section::Import(sect) => { import = Some(sect); }
                elements::Section::Code(sect) => { code = Some(sect); }
                _ => {}
            }
        }

        ModuleScaffold {
            functions: funcs.unwrap_or_default(),
            types: types.unwrap_or_default(),
            import: import.unwrap_or_default(),
            code: code.unwrap_or_default(),
            other: sections,
        }
    }
}

impl From<ModuleScaffold> for elements::Module {
    fn from(module: ModuleScaffold) -> Self {
        let mut sections = Vec::new();

        let types = module.types;
        if types.types().len() > 0 {
            sections.push(elements::Section::Type(types));
        }
        let import = module.import;
        if import.entries().len() > 0 {
            sections.push(elements::Section::Import(import));
        }
        let functions = module.functions;
        if functions.entries().len() > 0 {
            sections.push(elements::Section::Function(functions));
        }
        let code = module.code;
        if code.bodies().len() > 0 {
            sections.push(elements::Section::Code(code));
        }
        sections.extend(module.other);
        elements::Module::new(sections)
    }
}

impl ModuleBuilder {
    /// New empty module builder
    pub fn new() -> Self {
        ModuleBuilder::with_callback(Identity)
    }
}

impl<F> ModuleBuilder<F> where F: Invoke<elements::Module> {
    /// New module builder with bound callback
    pub fn with_callback(callback: F) -> Self {
        ModuleBuilder {
            callback: callback,
            module: Default::default(),
        }
    }

    /// Builder from raw module
    pub fn with_module(mut self, module: elements::Module) -> Self {
        self.module = module.into();
        self
    }

    /// Fill module with sections from iterator
    pub fn with_sections<I>(mut self, sections: I) -> Self
        where I: IntoIterator<Item=elements::Section>
    {
        self.module.other.extend(sections);
        self
    }

    /// Add additional section
    pub fn with_section(mut self, section: elements::Section) -> Self {
        self.module.other.push(section);
        self
    }

    /// Binds to the type section, creates additional types when required
    pub fn with_signatures(mut self, bindings: code::SignatureBindings) -> Self {
        self.push_signatures(bindings);
        self
    }

    /// Push stand-alone function definition, creating sections, signature and code blocks
    /// in corresponding sections.
    /// `FunctionDefinition` can be build using `builder::function` builder
    pub fn push_function(&mut self, func: code::FunctionDefinition) -> CodeLocation {
        let signature = func.signature;
        let body = func.code;

        let type_ref = self.resolve_type_ref(signature);

        self.module.functions.entries_mut().push(elements::Func::new(type_ref));
        let signature_index = self.module.functions.entries_mut().len() as u32 - 1;
        self.module.code.bodies_mut().push(body);
        let body_index = self.module.code.bodies_mut().len() as u32 - 1;

        CodeLocation {
            signature: signature_index,
            body: body_index,
        }
    }

    fn resolve_type_ref(&mut self, signature: code::Signature) -> u32 {
        match signature {
            code::Signature::Inline(func_type) => {
                self.module.types.types_mut().push(elements::Type::Function(func_type));
                self.module.types.types().len() as u32 - 1
            }
            code::Signature::TypeReference(type_ref) => {
                type_ref
            }
        }
    }

    /// Push one function signature, returning it's calling index.
    /// Can create corresponding type in type section.
    pub fn push_signature(&mut self, signature: code::Signature) -> u32 {
        let type_ref = self.resolve_type_ref(signature);
        self.module.functions.entries_mut().push(elements::Func::new(type_ref));
        self.module.functions.entries_mut().len() as u32 - 1
    }

    /// Push signatures in the module, returning corresponding indices of pushed signatures
    pub fn push_signatures(&mut self, signatures: code::SignatureBindings) -> Vec<u32> {
        let mut result = Vec::new();

        // todo: maybe reuse existing types with the equal signatures
        let raw_functions: Vec<u32> = signatures.into_iter().map(|binding|
            match binding {
                code::Signature::Inline(func_type) => {
                    self.module.types.types_mut().push(elements::Type::Function(func_type));
                    self.module.types.types().len() as u32 - 1
                }
                code::Signature::TypeReference(type_ref) => {
                    type_ref
                }
            }
        ).collect();

        for function in raw_functions {
            self.module.functions.entries_mut().push(elements::Func::new(function));
            result.push(self.module.functions.entries_mut().len() as u32 - 1);
        }

        result
    }

    pub fn function(self) -> FunctionBuilder<Self> {
        FunctionBuilder::with_callback(self)
    }

    /// Define functions section
    pub fn functions(self) -> SignaturesBuilder<Self> {
        SignaturesBuilder::with_callback(self)
    }

    /// With inserted import entry
    pub fn with_import(mut self, entry: elements::ImportEntry) -> Self {
        self.module.import.entries_mut().push(entry);
        self
    }

    /// Import entry builder
    pub fn import(self) -> import::ImportBuilder<Self> {
        import::ImportBuilder::with_callback(self)
    }

    /// Build module (final step)
    pub fn build(self) -> F::Result {
        self.callback.invoke(self.module.into())
    }
}

impl<F> Invoke<elements::FunctionsSection> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, section: elements::FunctionsSection) -> Self {
		self.with_section(elements::Section::Function(section))
    }
}

impl<F> Invoke<code::SignatureBindings> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, bindings: code::SignatureBindings) -> Self {
        self.with_signatures(bindings)
    }
}

impl<F> Invoke<code::FunctionDefinition> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, def: code::FunctionDefinition) -> Self {
        let mut b = self;
        b.push_function(def);
        b
    }
}

impl<F> Invoke<elements::ImportEntry> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, entry: elements::ImportEntry) -> Self::Result {
        self.with_import(entry)
    }
}

/// Start new module builder
pub fn module() -> ModuleBuilder {
    ModuleBuilder::new()
}

/// Start builder to extend existing module
pub fn from_module(module: elements::Module) -> ModuleBuilder {
    ModuleBuilder::new().with_module(module)
}

#[cfg(test)]
mod tests {

    use super::module;

    #[test]
    fn smoky() {
        let module = module().build();
        assert_eq!(module.sections().len(), 0);
    }

    #[test]
    fn functions() {
        let module = module()
            .function()
                .signature().param().i32().build()
                .body().build()
                .build()
            .build();

        assert_eq!(module.type_section().expect("type section to exist").types().len(), 1);
        assert_eq!(module.functions_section().expect("function section to exist").entries().len(), 1);
        assert_eq!(module.code_section().expect("code section to exist").bodies().len(), 1);
    }

}
<commit_msg>signature builder should resolve to function type, not function def<commit_after>use super::invoke::{Invoke, Identity};
use super::code::{self, SignaturesBuilder, FunctionBuilder};
use super::import;
use elements;

/// Module builder
pub struct ModuleBuilder<F=Identity> {
    callback: F,
    module: ModuleScaffold,
}

/// Location of the internal module function
pub struct CodeLocation {
    /// Location (index in 'functions' section) of the signature
    pub signature: u32,
    /// Location (index in the 'code' section) of the body
    pub body: u32,
}

#[derive(Default)]
struct ModuleScaffold {
    pub functions: elements::FunctionsSection,
    pub types: elements::TypeSection,
    pub import: elements::ImportSection,
    pub code: elements::CodeSection,
    pub other: Vec<elements::Section>,
}

impl From<elements::Module> for ModuleScaffold {
    fn from(module: elements::Module) -> Self {
        let mut funcs: Option<elements::FunctionsSection> = None;
        let mut types: Option<elements::TypeSection> = None;
        let mut import: Option<elements::ImportSection> = None;
        let mut code: Option<elements::CodeSection> = None;

        let mut sections = module.into_sections();
        while let Some(section) = sections.pop() {
            match section {
                elements::Section::Type(sect) => { types = Some(sect); }
                elements::Section::Function(sect) => { funcs = Some(sect); }
                elements::Section::Import(sect) => { import = Some(sect); }
                elements::Section::Code(sect) => { code = Some(sect); }
                _ => {}
            }
        }

        ModuleScaffold {
            functions: funcs.unwrap_or_default(),
            types: types.unwrap_or_default(),
            import: import.unwrap_or_default(),
            code: code.unwrap_or_default(),
            other: sections,
        }
    }
}

impl From<ModuleScaffold> for elements::Module {
    fn from(module: ModuleScaffold) -> Self {
        let mut sections = Vec::new();

        let types = module.types;
        if types.types().len() > 0 {
            sections.push(elements::Section::Type(types));
        }
        let import = module.import;
        if import.entries().len() > 0 {
            sections.push(elements::Section::Import(import));
        }
        let functions = module.functions;
        if functions.entries().len() > 0 {
            sections.push(elements::Section::Function(functions));
        }
        let code = module.code;
        if code.bodies().len() > 0 {
            sections.push(elements::Section::Code(code));
        }
        sections.extend(module.other);
        elements::Module::new(sections)
    }
}

impl ModuleBuilder {
    /// New empty module builder
    pub fn new() -> Self {
        ModuleBuilder::with_callback(Identity)
    }
}

impl<F> ModuleBuilder<F> where F: Invoke<elements::Module> {
    /// New module builder with bound callback
    pub fn with_callback(callback: F) -> Self {
        ModuleBuilder {
            callback: callback,
            module: Default::default(),
        }
    }

    /// Builder from raw module
    pub fn with_module(mut self, module: elements::Module) -> Self {
        self.module = module.into();
        self
    }

    /// Fill module with sections from iterator
    pub fn with_sections<I>(mut self, sections: I) -> Self
        where I: IntoIterator<Item=elements::Section>
    {
        self.module.other.extend(sections);
        self
    }

    /// Add additional section
    pub fn with_section(mut self, section: elements::Section) -> Self {
        self.module.other.push(section);
        self
    }

    /// Binds to the type section, creates additional types when required
    pub fn with_signatures(mut self, bindings: code::SignatureBindings) -> Self {
        self.push_signatures(bindings);
        self
    }

    /// Push stand-alone function definition, creating sections, signature and code blocks
    /// in corresponding sections.
    /// `FunctionDefinition` can be build using `builder::function` builder
    pub fn push_function(&mut self, func: code::FunctionDefinition) -> CodeLocation {
        let signature = func.signature;
        let body = func.code;

        let type_ref = self.resolve_type_ref(signature);

        self.module.functions.entries_mut().push(elements::Func::new(type_ref));
        let signature_index = self.module.functions.entries_mut().len() as u32 - 1;
        self.module.code.bodies_mut().push(body);
        let body_index = self.module.code.bodies_mut().len() as u32 - 1;

        CodeLocation {
            signature: signature_index,
            body: body_index,
        }
    }

    fn resolve_type_ref(&mut self, signature: code::Signature) -> u32 {
        match signature {
            code::Signature::Inline(func_type) => {
                self.module.types.types_mut().push(elements::Type::Function(func_type));
                self.module.types.types().len() as u32 - 1
            }
            code::Signature::TypeReference(type_ref) => {
                type_ref
            }
        }
    }

    /// Push one function signature, returning it's calling index.
    /// Can create corresponding type in type section.
    pub fn push_signature(&mut self, signature: code::Signature) -> u32 {
        self.resolve_type_ref(signature)
    }

    /// Push signatures in the module, returning corresponding indices of pushed signatures
    pub fn push_signatures(&mut self, signatures: code::SignatureBindings) -> Vec<u32> {
        let mut result = Vec::new();

        // todo: maybe reuse existing types with the equal signatures
        let raw_functions: Vec<u32> = signatures.into_iter().map(|binding|
            match binding {
                code::Signature::Inline(func_type) => {
                    self.module.types.types_mut().push(elements::Type::Function(func_type));
                    self.module.types.types().len() as u32 - 1
                }
                code::Signature::TypeReference(type_ref) => {
                    type_ref
                }
            }
        ).collect();

        for function in raw_functions {
            self.module.functions.entries_mut().push(elements::Func::new(function));
            result.push(self.module.functions.entries_mut().len() as u32 - 1);
        }

        result
    }

    pub fn function(self) -> FunctionBuilder<Self> {
        FunctionBuilder::with_callback(self)
    }

    /// Define functions section
    pub fn functions(self) -> SignaturesBuilder<Self> {
        SignaturesBuilder::with_callback(self)
    }

    /// With inserted import entry
    pub fn with_import(mut self, entry: elements::ImportEntry) -> Self {
        self.module.import.entries_mut().push(entry);
        self
    }

    /// Import entry builder
    pub fn import(self) -> import::ImportBuilder<Self> {
        import::ImportBuilder::with_callback(self)
    }

    /// Build module (final step)
    pub fn build(self) -> F::Result {
        self.callback.invoke(self.module.into())
    }
}

impl<F> Invoke<elements::FunctionsSection> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
	type Result = Self;

	fn invoke(self, section: elements::FunctionsSection) -> Self {
		self.with_section(elements::Section::Function(section))
    }
}

impl<F> Invoke<code::SignatureBindings> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, bindings: code::SignatureBindings) -> Self {
        self.with_signatures(bindings)
    }
}

impl<F> Invoke<code::FunctionDefinition> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, def: code::FunctionDefinition) -> Self {
        let mut b = self;
        b.push_function(def);
        b
    }
}

impl<F> Invoke<elements::ImportEntry> for ModuleBuilder<F>
    where F: Invoke<elements::Module>
{
    type Result = Self;

    fn invoke(self, entry: elements::ImportEntry) -> Self::Result {
        self.with_import(entry)
    }
}

/// Start new module builder
pub fn module() -> ModuleBuilder {
    ModuleBuilder::new()
}

/// Start builder to extend existing module
pub fn from_module(module: elements::Module) -> ModuleBuilder {
    ModuleBuilder::new().with_module(module)
}

#[cfg(test)]
mod tests {

    use super::module;

    #[test]
    fn smoky() {
        let module = module().build();
        assert_eq!(module.sections().len(), 0);
    }

    #[test]
    fn functions() {
        let module = module()
            .function()
                .signature().param().i32().build()
                .body().build()
                .build()
            .build();

        assert_eq!(module.type_section().expect("type section to exist").types().len(), 1);
        assert_eq!(module.functions_section().expect("function section to exist").entries().len(), 1);
        assert_eq!(module.code_section().expect("code section to exist").bodies().len(), 1);
    }

}
<|endoftext|>


<commit_before>// Copyright (c) 2016 The Rouille developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

#[macro_use]
extern crate rouille;

use std::thread;

use rouille::websocket;
use rouille::Response;

fn main() {
    // This example demonstrates how to use websockets with rouille.

    // Small message so that people don't need to read the source code.
    // Note that like all examples we only listen on `localhost`, so you can't access this server
    // from another machine than your own.
    println!("Now listening on localhost:8000");

    rouille::start_server("localhost:8000", move |request| {
        router!(request,
            (GET) (/) => {
                // The / route outputs an HTML client so that the user can try the websockets.
                // Note that in a real website you should probably use some templating system, or
                // at least load the HTML from a file.
                Response::html("

This example sends back everything you send to the server.

<form onsubmit=\"send(document.getElementById('msg').value); return false;\">
                    <input type=\"text\" id=\"msg\" />
                    <button type=\"submit\">Send</button>
                    </form>

Received:

")
            },

            (GET) (/ws) => {
                // This is the websockets route.

                // In order to start using websockets we call `websocket::start`.
                // The function returns an error if the client didn't request websockets, in which
                // case we return an error 400 to the client thanks to the `try_or_400!` macro.
                //
                // The function returns a response to send back as part of the `start_server`
                // function, and a `websocket` variable of type `Receiver<Websocket>`.
                // Once the response has been sent back to the client, the `Receiver` will be
                // filled by rouille with a `Websocket` object representing the websocket.
                let (response, websocket) = try_or_400!(websocket::start(&request, Some("echo")));

                // Because of the nature of I/O in Rust, we need to spawn a separate thread for
                // each websocket.
                thread::spawn(move || {
                    // This line will block until the `response` above has been returned.
                    let ws = websocket.recv().unwrap();
                    // We use a separate function for better readability.
                    websocket_handling_thread(ws);
                });

                response
            },

            // Default 404 route as with all examples.
            _ => rouille::Response::empty_404()
        )
    });
}

// Function run in a separate thread.
fn websocket_handling_thread(mut websocket: websocket::Websocket) {
    loop {
        // We wait for a new message to come from the websocket.
        let message = match websocket.next() {
            Some(m) => m,
            None => break,
        };

        match message {
            websocket::Message::Text(txt) => {
                // If the message is text, send it back with `send_text`.
                println!("received {:?} from a websocket", txt);
                websocket.send_text(&txt).unwrap();
            },
            websocket::Message::Binary(_) => {
                println!("received binary from a websocket");
            },
        }
    }
}
<commit_msg>examples: replace loop with more idiomatic while let<commit_after>// Copyright (c) 2016 The Rouille developers
// 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. All files in the project carrying such
// notice may not be copied, modified, or distributed except
// according to those terms.

#[macro_use]
extern crate rouille;

use std::thread;

use rouille::websocket;
use rouille::Response;

fn main() {
    // This example demonstrates how to use websockets with rouille.

    // Small message so that people don't need to read the source code.
    // Note that like all examples we only listen on `localhost`, so you can't access this server
    // from another machine than your own.
    println!("Now listening on localhost:8000");

    rouille::start_server("localhost:8000", move |request| {
        router!(request,
            (GET) (/) => {
                // The / route outputs an HTML client so that the user can try the websockets.
                // Note that in a real website you should probably use some templating system, or
                // at least load the HTML from a file.
                Response::html("

This example sends back everything you send to the server.

<form onsubmit=\"send(document.getElementById('msg').value); return false;\">
                    <input type=\"text\" id=\"msg\" />
                    <button type=\"submit\">Send</button>
                    </form>

Received:

")
            },

            (GET) (/ws) => {
                // This is the websockets route.

                // In order to start using websockets we call `websocket::start`.
                // The function returns an error if the client didn't request websockets, in which
                // case we return an error 400 to the client thanks to the `try_or_400!` macro.
                //
                // The function returns a response to send back as part of the `start_server`
                // function, and a `websocket` variable of type `Receiver<Websocket>`.
                // Once the response has been sent back to the client, the `Receiver` will be
                // filled by rouille with a `Websocket` object representing the websocket.
                let (response, websocket) = try_or_400!(websocket::start(&request, Some("echo")));

                // Because of the nature of I/O in Rust, we need to spawn a separate thread for
                // each websocket.
                thread::spawn(move || {
                    // This line will block until the `response` above has been returned.
                    let ws = websocket.recv().unwrap();
                    // We use a separate function for better readability.
                    websocket_handling_thread(ws);
                });

                response
            },

            // Default 404 route as with all examples.
            _ => rouille::Response::empty_404()
        )
    });
}

// Function run in a separate thread.
fn websocket_handling_thread(mut websocket: websocket::Websocket) {
    // We wait for a new message to come from the websocket.
    while let Some(message) = websocket.next() {

        match message {
            websocket::Message::Text(txt) => {
                // If the message is text, send it back with `send_text`.
                println!("received {:?} from a websocket", txt);
                websocket.send_text(&txt).unwrap();
            },
            websocket::Message::Binary(_) => {
                println!("received binary from a websocket");
            },
        }
    }
}
<|endoftext|>


<commit_before>use std::path::Path;

use crate::languages::FormatAs;

use regex::Regex;

/// Tries to find some text that describes the contents of the file.
/// Adapted from DXR code at https://github.com/mozilla/dxr/blob/master/dxr/plugins/descriptor/__init__.py
pub fn describe_file(contents: &str, path: &Path, format: &FormatAs) -> Option<String> {
    // Only look in the first 5k chars of the file, since running regex matches on giant files
    // can be expensive and most likely the thing we're looking for will be early in the file.
    let substr_end = contents
        .char_indices()
        .nth(5000)
        .map(|ix| ix.0)
        .unwrap_or(contents.len());
    let substr = &contents[0..substr_end];

    // DXR also does a search for "filename: <description>" which I've never seen in any file so
    // I'm omitting that here. We can add it if needed.
    match format {
        FormatAs::FormatTagLike(_) => describe_html(substr),
        FormatAs::FormatCLike(ref spec) => {
            if spec.rust_tweaks {
                describe_from_rust_comment(substr).or_else(|| describe_from_c_comment(substr))
            } else if spec.c_style_comments {
                describe_from_c_comment(substr)
            } else if spec.triple_quote_literals {
                describe_py(substr)
            } else {
                None
            }
        }
        FormatAs::Binary => None,
        FormatAs::Plain => {
            let stem = path.file_stem()?.to_str()?;
            if stem.eq_ignore_ascii_case("README") {
                describe_readme(substr)
            } else {
                None
            }
        }
    }
}

/// Returns the content of the title tag
fn describe_html(contents: &str) -> Option<String> {
    lazy_static! {
        static ref TITLE_REGEX: Regex = { Regex::new(r#"(?i)<title>((?s).*?)</title>"#).unwrap() };
    }
    TITLE_REGEX
        .captures(contents)
        .and_then(|c| c.get(1))
        .map(|m| m.as_str().to_string())
}

/// Returns the first docstring (single-quoted preferred)
fn describe_py(contents: &str) -> Option<String> {
    lazy_static! {
        static ref DOCSTRING_SINGLE_REGEX: Regex = { Regex::new(r#"'''\s*((?s).*?)'''"#).unwrap() };
        static ref DOCSTRING_DOUBLE_REGEX: Regex = { Regex::new(r#""""\s*((?s).*?)""""#).unwrap() };
    }
    DOCSTRING_SINGLE_REGEX
        .captures(contents)
        .or_else(|| DOCSTRING_DOUBLE_REGEX.captures(contents))
        .and_then(|c| c.get(1))
        .map(|m| m.as_str().to_string())
}

/// Returns the first nonempty line
fn describe_readme(contents: &str) -> Option<String> {
    contents
        .lines()
        .filter(|s| !s.trim().is_empty())
        .next()
        .map(str::to_string)
}

/// Returns the first Rust-style module doc-comment (`//!`).
fn describe_from_rust_comment(contents: &str) -> Option<String> {
    let mut description = String::new();
    let mut in_comment = false;
    for line in contents.lines() {
        if !line.starts_with("//!") {
            if !in_comment {
                continue;
            }
            return Some(description);
        }
        in_comment = true;
        if !description.is_empty() {
            description.push_str("\n");
        }
        description.push_str(line.trim_start_matches("//!"));
    }
    None
}

/// Returns the first C-style comment that's not vim/modeline/license boilerplate
fn describe_from_c_comment(contents: &str) -> Option<String> {
    lazy_static! {
        // Matches C-style comments, including any leading '*' characters on
        // wrapped lines. The LEADING_STARS regex is used to remove those.
        static ref COMMENT_REGEX: Regex = {
            Regex::new(r#"(?:/\*[*\s]*)((?s).*?)\*/"#).unwrap()
        };
        static ref LEADING_STARS: Regex = {
            Regex::new(r#"(?m)^\s*\*"#).unwrap()
        };
    }
    for captures in COMMENT_REGEX.captures_iter(contents) {
        if let Some(comment_match) = captures.get(1) {
            let comment_text = comment_match.as_str();
            if comment_text.contains("tab-width") ||
                comment_text.contains("vim:") ||
                // Checking common case-variants is probably cheaper than lowercasing
                // comment_text (which allocates a new string) and doing a lowercase search
                comment_text.contains("license") ||
                comment_text.contains("LICENSE") ||
                comment_text.contains("License")
            {
                continue;
            }
            return Some(LEADING_STARS.replace_all(comment_text, "").into_owned());
        }
    }
    None
}
<commit_msg>describe: Remove some useless braces.<commit_after>use std::path::Path;

use crate::languages::FormatAs;

use regex::Regex;

/// Tries to find some text that describes the contents of the file.
/// Adapted from DXR code at https://github.com/mozilla/dxr/blob/master/dxr/plugins/descriptor/__init__.py
pub fn describe_file(contents: &str, path: &Path, format: &FormatAs) -> Option<String> {
    // Only look in the first 5k chars of the file, since running regex matches on giant files
    // can be expensive and most likely the thing we're looking for will be early in the file.
    let substr_end = contents
        .char_indices()
        .nth(5000)
        .map(|ix| ix.0)
        .unwrap_or(contents.len());
    let substr = &contents[0..substr_end];

    // DXR also does a search for "filename: <description>" which I've never seen in any file so
    // I'm omitting that here. We can add it if needed.
    match format {
        FormatAs::FormatTagLike(_) => describe_html(substr),
        FormatAs::FormatCLike(ref spec) => {
            if spec.rust_tweaks {
                describe_from_rust_comment(substr).or_else(|| describe_from_c_comment(substr))
            } else if spec.c_style_comments {
                describe_from_c_comment(substr)
            } else if spec.triple_quote_literals {
                describe_py(substr)
            } else {
                None
            }
        }
        FormatAs::Binary => None,
        FormatAs::Plain => {
            let stem = path.file_stem()?.to_str()?;
            if stem.eq_ignore_ascii_case("README") {
                describe_readme(substr)
            } else {
                None
            }
        }
    }
}

/// Returns the content of the title tag
fn describe_html(contents: &str) -> Option<String> {
    lazy_static! {
        static ref TITLE_REGEX: Regex = Regex::new(r#"(?i)<title>((?s).*?)</title>"#).unwrap();
    }
    TITLE_REGEX
        .captures(contents)
        .and_then(|c| c.get(1))
        .map(|m| m.as_str().to_string())
}

/// Returns the first docstring (single-quoted preferred)
fn describe_py(contents: &str) -> Option<String> {
    lazy_static! {
        static ref DOCSTRING_SINGLE_REGEX: Regex = Regex::new(r#"'''\s*((?s).*?)'''"#).unwrap();
        static ref DOCSTRING_DOUBLE_REGEX: Regex = Regex::new(r#""""\s*((?s).*?)""""#).unwrap();
    }
    DOCSTRING_SINGLE_REGEX
        .captures(contents)
        .or_else(|| DOCSTRING_DOUBLE_REGEX.captures(contents))
        .and_then(|c| c.get(1))
        .map(|m| m.as_str().to_string())
}

/// Returns the first nonempty line
fn describe_readme(contents: &str) -> Option<String> {
    contents
        .lines()
        .filter(|s| !s.trim().is_empty())
        .next()
        .map(str::to_string)
}

/// Returns the first Rust-style module doc-comment (`//!`).
fn describe_from_rust_comment(contents: &str) -> Option<String> {
    let mut description = String::new();
    let mut in_comment = false;
    for line in contents.lines() {
        if !line.starts_with("//!") {
            if !in_comment {
                continue;
            }
            return Some(description);
        }
        in_comment = true;
        if !description.is_empty() {
            description.push_str("\n");
        }
        description.push_str(line.trim_start_matches("//!"));
    }
    None
}

/// Returns the first C-style comment that's not vim/modeline/license boilerplate
fn describe_from_c_comment(contents: &str) -> Option<String> {
    lazy_static! {
        // Matches C-style comments, including any leading '*' characters on
        // wrapped lines. The LEADING_STARS regex is used to remove those.
        static ref COMMENT_REGEX: Regex = {
            Regex::new(r#"(?:/\*[*\s]*)((?s).*?)\*/"#).unwrap()
        };
        static ref LEADING_STARS: Regex = {
            Regex::new(r#"(?m)^\s*\*"#).unwrap()
        };
    }
    for captures in COMMENT_REGEX.captures_iter(contents) {
        if let Some(comment_match) = captures.get(1) {
            let comment_text = comment_match.as_str();
            if comment_text.contains("tab-width") ||
                comment_text.contains("vim:") ||
                // Checking common case-variants is probably cheaper than lowercasing
                // comment_text (which allocates a new string) and doing a lowercase search
                comment_text.contains("license") ||
                comment_text.contains("LICENSE") ||
                comment_text.contains("License")
            {
                continue;
            }
            return Some(LEADING_STARS.replace_all(comment_text, "").into_owned());
        }
    }
    None
}
<|endoftext|>