<commit_before>#![feature(phase)]
extern crate serialize;

#[phase(plugin)]
extern crate docopt_macros;
extern crate docopt;

extern crate rustscripts;

#[warn(non_camel_case_types)]
#[warn(non_snake_case)]
#[warn(unused_qualifications)]
#[warn(non_upper_case_globals)]
#[warn(missing_docs)]

/** Creates fake words.
 *  First, it uses a predefined list of words to generate a markov chain
 *  for word prefixes; it then probabilistically follows this to generate
 *  fake words.
**/

use docopt::Docopt;

use std::collections::{HashSet,HashMap};
use std::rand;
use std::rand::Rng;
use std::io::{File,BufferedReader};

use rustscripts::counter;

//* Build a HashMap of "substring" : "number of occurrences"
fn to_hashes(wordlist : &[String], sublens : uint) -> HashMap<String, uint> {
	println!("Got wordlist, length {}", wordlist.len());

	let iter = wordlist.iter().filter(|k| {
		if k.char_len() == 0 {false}
		else if k.contains("\'") {false}
		else if k.find(|c : char|{!c.is_lowercase()}).is_some() {false}
		else {true}
	});
	let trimchars : &[char] = &[' ', '\t', '\r', '\n'];
	let iter = iter.flat_map(|k| {
		let kslice = k.as_slice();
		let fullword = (["^", kslice.trim_chars(trimchars), "$"]).concat();
        let fullchars : Vec<char> = fullword.as_slice().chars().collect();

        let chvec : Vec<String> =
		fullchars.as_slice().windows(sublens).map(|chars| {
			String::from_chars(chars)
		})
		.collect();

		//~ println!("k: {}, chvec: {}", k, chvec);

		chvec.into_iter()
	});
	counter(iter)
}

struct WordBuilder {
    subs : HashMap<String, uint>,
    // list : Vec<String>,
    // sublens : uint,
    wordset : HashSet<String>,
    wordlens : Vec<uint>
}

struct WordIter<'a> {
    p : &'a mut WordBuilder
}

impl WordBuilder {
    fn new(list : Vec<String>, sublens : uint) -> WordBuilder {
        let mut h = HashSet::new();
        let mut wlens : Vec<uint> = Vec::new();
        for w in list.iter() {
            h.insert(w.to_string());
            let wordlen = w.len();
            for _ in range(wlens.len(), wordlen+1){
                wlens.push(0);
            }
            wlens[wordlen] += 1;
        }

        WordBuilder {
            subs : to_hashes(list.as_slice(), sublens),
            //list : list,
            //sublens : sublens,
            wordset : h,
            wordlens : wlens
        }
    }

    fn word(&mut self) -> Option<String> {
        let mut s : String = "^".to_string();

        loop {
            let mut fullsum = 0u;
            let mut endsum = 0u;
            let possibilities : Vec<(&str, uint)> = self.subs.iter().filter_map(
                |(k,v)| {
                    /* the beginning of k and the end of s must match for
                     * k to be a possibility
                     * if s is long, then the first (klength - 1) letters
                     * of k must match the last (klength - 1) letters of s
                     * otherwise, the first (slength) characters
                    */
                    let slen = s.as_slice().char_len();
                    let kslice = k.as_slice();
                    let klen = kslice.char_len();
                    let kcut = if slen < klen - 1 {slen} else {klen - 1};
                    if s.as_slice().ends_with(kslice.slice_chars(0, kcut)){
                        fullsum += *v;
                        if kslice.ends_with("$") {endsum += *v;}
                        Some((kslice,*v))
                    } else {None}
                }
            ).collect();
            if fullsum == 0 {
                panic!("s: \"{}\"", s);
            }

            let endprob = if self.wordlens.len() > s.len() {
                let wordlenslice : &[uint] = self.wordlens.slice(s.len()-1, self.wordlens.len());
                let c = wordlenslice[0];
                let l = wordlenslice.iter().fold(0, |a,&b|{a+b});
                (c as f64) / (l as f64)
            } else {
                //~ println!("Too long: {}", s);
                return None;
            };

            let randnum = rand::task_rng().gen_range(0.0, 1.0);

            let endtime = randnum < endprob;
            if (endtime && (endsum == 0)) || ((!endtime) && (fullsum-endsum==0)) {
                //~ println!("Failed to end: {}", s);
                return None;
            }

            //~ println!("endtime: {} {} : ({},{}) {}", endtime, randnum < endprob,
                //~ endsum, fullsum, if(endtime){endsum} else {fullsum - endsum});
            let randnum = rand::task_rng().gen_range(0.0,
                (if endtime {endsum} else {fullsum - endsum} as f64));

            let mut psum = 0;

            for &(k,v) in possibilities.iter() {
                if endtime ^ k.ends_with("$") {continue;};

                if randnum < ((psum + v) as f64) {
                    let slen = s.as_slice().char_len();
                    let klen = k.char_len();
                    let kcut = if slen < klen - 1 {slen} else {klen - 1};
                    //~ let olds = s.to_string();
                    s.push_str(k.slice_chars(kcut,klen));
                    break;
                }
                psum += v;
            }

            let slen = s.as_slice().char_len();
            if s.as_slice().slice_chars(slen-1, slen) == "$" {
                return Some(s.as_slice().slice_chars(1, slen-1).to_string());
            }
        };
    }

    fn iter<'a>(&'a mut self) -> WordIter<'a> {
        WordIter{p:self}
    }
}

impl<'a> Iterator<String> for WordIter<'a> {
    fn next(&mut self) -> Option<String> {
        loop {
            let optw = self.p.word();
            match optw {
                None => {},
                Some(w) => {
                    if !self.p.wordset.contains(&w)
                        {return Some(w)};
                }
            };
        }
    }
}

docopt!(Args, "
Usage: fakewords2 [-h | --help] [-n <number>] [<dictfile>]

Options:
    -n <number>    use number length substrings for markovian chain
    <dictfile>     use dictfile instead of /usr/share/dict/words
", flag_n : Option<uint>, arg_dictfile : Option<String>)

pub fn main(){
	let args: Args = Args::docopt().decode().unwrap_or_else(|e| e.exit());

    let flag_n : Option<uint> = args.flag_n;

    let subsetn : uint = flag_n.unwrap_or(4u);

    let pathstr = args.arg_dictfile.map(|s| {s}).unwrap_or("/usr/share/dict/words".to_string());

    let path = Path::new(pathstr.clone());
    let file = match File::open(&path) {
		Ok(f) => f,
		Err(std::io::IoError{kind: std::io::FileNotFound, desc: _, detail: _}) => {
			let _ = writeln!(&mut std::io::stderr(), "File not found: {}", pathstr);
			std::os::set_exit_status(-1);
			return;
		}
		Err(e) => panic!("failed to open file: {}", e)
	};

    let mut file = BufferedReader::new(file);

    let trimchars : &[char] = &[' ', '\t', '\r', '\n'];

    let lines: Vec<String> = file.lines().map(|orl| {
		let unwrapl : String = match orl {
			Ok(l) => l,
			Err(e) => panic!("Failed reading file: {}", e)
		};
		unwrapl.trim_chars(trimchars).to_string()
	}).collect();
    let mut wb = WordBuilder::new(lines, subsetn);

    println!("Now have a map of length {}", wb.subs.len());

    for w in wb.iter(){
        println!("{}",w);
    };
}
<commit_msg>Updated to newest rust<commit_after>#![feature(phase)]
extern crate serialize;

#[phase(plugin)]
extern crate docopt_macros;
extern crate docopt;

extern crate rustscripts;

#[warn(non_camel_case_types)]
#[warn(non_snake_case)]
#[warn(unused_qualifications)]
#[warn(non_upper_case_globals)]
#[warn(missing_docs)]

/** Creates fake words.
 *  First, it uses a predefined list of words to generate a markov chain
 *  for word prefixes; it then probabilistically follows this to generate
 *  fake words.
**/

use docopt::Docopt;

use std::collections::{HashSet,HashMap};
use std::rand;
use std::rand::Rng;
use std::io::{File,BufferedReader};

use rustscripts::counter;

//* Build a HashMap of "substring" : "number of occurrences"
fn to_hashes(wordlist : &[String], sublens : uint) -> HashMap<String, uint> {
	println!("Got wordlist, length {}", wordlist.len());

	let iter = wordlist.iter().filter(|k| {
		if k.char_len() == 0 {false}
		else if k.contains("\'") {false}
		else if k.find(|&: c : char|{!c.is_lowercase()}).is_some() {false}
		else {true}
	});
	let trimchars : &[char] = &[' ', '\t', '\r', '\n'];
	let iter = iter.flat_map(|k| {
		let kslice = k.as_slice();
		let fullword = (["^", kslice.trim_chars(trimchars), "$"]).concat();
        let fullchars : Vec<char> = fullword.as_slice().chars().collect();

        let chvec : Vec<String> =
		fullchars.as_slice().windows(sublens).map(|chars| {
			String::from_chars(chars)
		})
		.collect();

		//~ println!("k: {}, chvec: {}", k, chvec);

		chvec.into_iter()
	});
	counter(iter)
}

struct WordBuilder {
    subs : HashMap<String, uint>,
    // list : Vec<String>,
    // sublens : uint,
    wordset : HashSet<String>,
    wordlens : Vec<uint>
}

struct WordIter<'a> {
    p : &'a mut WordBuilder
}

impl WordBuilder {
    fn new(list : Vec<String>, sublens : uint) -> WordBuilder {
        let mut h = HashSet::new();
        let mut wlens : Vec<uint> = Vec::new();
        for w in list.iter() {
            h.insert(w.to_string());
            let wordlen = w.len();
            for _ in range(wlens.len(), wordlen+1){
                wlens.push(0);
            }
            wlens[wordlen] += 1;
        }

        WordBuilder {
            subs : to_hashes(list.as_slice(), sublens),
            //list : list,
            //sublens : sublens,
            wordset : h,
            wordlens : wlens
        }
    }

    fn word(&mut self) -> Option<String> {
        let mut s : String = "^".to_string();

        loop {
            let mut fullsum = 0u;
            let mut endsum = 0u;
            let possibilities : Vec<(&str, uint)> = self.subs.iter().filter_map(
                |(k,v)| {
                    /* the beginning of k and the end of s must match for
                     * k to be a possibility
                     * if s is long, then the first (klength - 1) letters
                     * of k must match the last (klength - 1) letters of s
                     * otherwise, the first (slength) characters
                    */
                    let slen = s.as_slice().char_len();
                    let kslice = k.as_slice();
                    let klen = kslice.char_len();
                    let kcut = if slen < klen - 1 {slen} else {klen - 1};
                    if s.as_slice().ends_with(kslice.slice_chars(0, kcut)){
                        fullsum += *v;
                        if kslice.ends_with("$") {endsum += *v;}
                        Some((kslice,*v))
                    } else {None}
                }
            ).collect();
            if fullsum == 0 {
                panic!("s: \"{}\"", s);
            }

            let endprob = if self.wordlens.len() > s.len() {
                let wordlenslice : &[uint] = self.wordlens.slice(s.len()-1, self.wordlens.len());
                let c = wordlenslice[0];
                let l = wordlenslice.iter().fold(0, |a,&b|{a+b});
                (c as f64) / (l as f64)
            } else {
                //~ println!("Too long: {}", s);
                return None;
            };

            let randnum = rand::task_rng().gen_range(0.0, 1.0);

            let endtime = randnum < endprob;
            if (endtime && (endsum == 0)) || ((!endtime) && (fullsum-endsum==0)) {
                //~ println!("Failed to end: {}", s);
                return None;
            }

            //~ println!("endtime: {} {} : ({},{}) {}", endtime, randnum < endprob,
                //~ endsum, fullsum, if(endtime){endsum} else {fullsum - endsum});
            let randnum = rand::task_rng().gen_range(0.0,
                (if endtime {endsum} else {fullsum - endsum} as f64));

            let mut psum = 0;

            for &(k,v) in possibilities.iter() {
                if endtime ^ k.ends_with("$") {continue;};

                if randnum < ((psum + v) as f64) {
                    let slen = s.as_slice().char_len();
                    let klen = k.char_len();
                    let kcut = if slen < klen - 1 {slen} else {klen - 1};
                    //~ let olds = s.to_string();
                    s.push_str(k.slice_chars(kcut,klen));
                    break;
                }
                psum += v;
            }

            let slen = s.as_slice().char_len();
            if s.as_slice().slice_chars(slen-1, slen) == "$" {
                return Some(s.as_slice().slice_chars(1, slen-1).to_string());
            }
        };
    }

    fn iter<'a>(&'a mut self) -> WordIter<'a> {
        WordIter{p:self}
    }
}

impl<'a> Iterator<String> for WordIter<'a> {
    fn next(&mut self) -> Option<String> {
        loop {
            let optw = self.p.word();
            match optw {
                None => {},
                Some(w) => {
                    if !self.p.wordset.contains(&w)
                        {return Some(w)};
                }
            };
        }
    }
}

docopt!(Args, "
Usage: fakewords2 [-h | --help] [-n <number>] [<dictfile>]

Options:
    -n <number>    use number length substrings for markovian chain
    <dictfile>     use dictfile instead of /usr/share/dict/words
", flag_n : Option<uint>, arg_dictfile : Option<String>)

pub fn main(){
	let args: Args = Args::docopt().decode().unwrap_or_else(|e| e.exit());

    let flag_n : Option<uint> = args.flag_n;

    let subsetn : uint = flag_n.unwrap_or(4u);

    let pathstr = args.arg_dictfile.map(|s| {s}).unwrap_or("/usr/share/dict/words".to_string());

    let path = Path::new(pathstr.clone());
    let file = match File::open(&path) {
		Ok(f) => f,
		Err(std::io::IoError{kind: std::io::FileNotFound, desc: _, detail: _}) => {
			let _ = writeln!(&mut std::io::stderr(), "File not found: {}", pathstr);
			std::os::set_exit_status(-1);
			return;
		}
		Err(e) => panic!("failed to open file: {}", e)
	};

    let mut file = BufferedReader::new(file);

    let trimchars : &[char] = &[' ', '\t', '\r', '\n'];

    let lines: Vec<String> = file.lines().map(|orl| {
		let unwrapl : String = match orl {
			Ok(l) => l,
			Err(e) => panic!("Failed reading file: {}", e)
		};
		unwrapl.trim_chars(trimchars).to_string()
	}).collect();
    let mut wb = WordBuilder::new(lines, subsetn);

    println!("Now have a map of length {}", wb.subs.len());

    for w in wb.iter(){
        println!("{}",w);
    };
}
<|endoftext|>


<commit_before>extern crate getopts;
extern crate markovpass;

use markovpass::{PassphraseMarkovChain,PassphraseMarkovChainError};
use std::io::Write;

fn build_opts() -> getopts::Options {
    let mut opts = getopts::Options::new();
    opts.optopt("n", "", "Number of passphrases to generate (default 1)", "NUM");
    opts.optopt("e", "", "Minimum entropy (default 60)", "MINENTROPY");
    opts.optopt("l", "", "NGram length (default 3, must be > 1)", "LENGTH");
    opts.optopt("w", "", "Minimum word length for corpus (default 5)", "LENGTH");
    opts.optflag("h", "help", "Display this help and exit");

    opts
}

fn parse_flag<T: std::str::FromStr>(matches: &getopts::Matches, flag: &str, default: T)
        -> Result<T, &'static str> {
    matches.opt_str(flag)
        .map(|c| c.parse::<T>())
        .unwrap_or(Ok(default))
        .map_err(|_| "Failed to parse flag.")
}

fn parse_args(opts: &getopts::Options, args: &Vec<String>)
        -> Result<(Option<String>, usize, f64, usize, usize), &'static str> {
    let matches = try!(opts.parse(&args[1..]).map_err(|_| "Failed to parse arguments."));

    if matches.opt_present("h") || matches.free.len() > 1 {
        return Err("Failed to parse arguments.");
    };

    let number = try!(parse_flag(&matches, "n", 1));
    let min_entropy = try!(parse_flag(&matches, "e", 60.0));
    let ngram_length = try!(parse_flag(&matches, "l", 3));
    let min_word_length = try!(parse_flag(&matches, "w", 5));

    if ngram_length < 2 {
        return Err("Ngram length must be greater than one.");
    };

    let filename = if matches.free.is_empty() || matches.free[0] == "-" {
        None
    } else {
        Some(matches.free[0].clone())
    };

    Ok((filename, number, min_entropy, ngram_length, min_word_length))
}

fn get_corpus(filename: Option<&str>) -> Result<String, std::io::Error> {
    let mut input: Box<std::io::Read> = match filename {
        Some(filename) => Box::new(try!(std::fs::File::open(&filename))),
        None => Box::new(std::io::stdin()),
    };
    let mut corpus = String::new();
    try!(input.read_to_string(&mut corpus));

    Ok(corpus)
}

fn clean_word(word: &str, min_length: usize) -> Option<&str> {
    let word = word.trim_matches(|c: char| !c.is_alphabetic());

    if word.chars().all(|c: char| c.is_alphabetic()) && word.len() >= min_length {
        Some(word)
    } else {
        None
    }
}

fn get_ngrams(corpus: &str, ngram_length: usize, min_word_length: usize) -> Vec<String> {
    let words = corpus.split_whitespace()
      .filter_map(|word| clean_word(word, min_word_length));
    let cleaned_corpus = Some("").into_iter().chain(words).collect::<Vec<&str>>().join(" ");
    let count = cleaned_corpus.chars().count();
    if count < ngram_length { return vec![]; };

    let mut chars = cleaned_corpus.chars().cycle();
    let mut ngrams = Vec::with_capacity(count);
    for _ in 0..count {
        let ngram: String = chars.clone().take(ngram_length).collect();
        ngrams.push(ngram.to_lowercase());
        chars.next();
    };

    ngrams
}

fn gen_passphrases(ngrams: Vec<String>, number: usize, min_entropy: f64)
        -> Result<Vec<(String, f64)>, PassphraseMarkovChainError> {
    let chain = try!(PassphraseMarkovChain::new(ngrams.iter().cloned()));
    let mut passphrases = Vec::with_capacity(number);
    for _ in 0..number {
        passphrases.push(chain.passphrase(min_entropy));
    };

    Ok(passphrases)
}

fn print_usage(program: &str, opts: &getopts::Options) {
    let brief = format!("Usage: {} [FILE] [options]", program);
    print!("{}", opts.usage(&brief));
}

fn write_error(message: &str) {
    writeln!(&mut std::io::stderr(), "{}", &message).expect("Failed to write to stderr.");
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let program = args[0].clone();
    let opts = build_opts();

    let (filename, number, min_entropy, ngram_length, min_word_length) =
            match parse_args(&opts, &args) {
        Ok(parsed_args) => { parsed_args },
        Err(_) => {
            print_usage(&program, &opts);
            return;
        }
    };
    let corpus = match get_corpus(filename.as_ref().map(String::as_ref)) {
        Ok(corpus) => { corpus },
        Err(_) => {
            let filename_str = filename.unwrap_or("-".to_string());
            write_error(&format!("{}: {}: {}", &program, &filename_str, "Failed to read input."));
            return;
        },
    };
    let ngrams = get_ngrams(&corpus, ngram_length, min_word_length);
    let passphrases = match gen_passphrases(ngrams, number, min_entropy) {
        Ok(passphrases) => { passphrases },
        Err(e) => {
            write_error(&format!("{}: {}", &program, e));
            return;
        },
    };

    for (passphrase, entropy) in passphrases {
        println!("{} <{}>", passphrase, entropy);
    };
}

#[cfg(test)]
mod tests {
    use super::{clean_word,get_ngrams,gen_passphrases};

    #[test]
    fn test_clean_word() {
        assert_eq!(clean_word("Test", 3), Some("Test"));
        assert_eq!(clean_word("12'3test@314", 3), Some("test"));
        assert_eq!(clean_word("31ld;Test", 3), None);
        assert_eq!(clean_word("a", 2), None);
        assert_eq!(clean_word("Test", 5), None);
    }

    #[test]
    fn test_get_ngrams() {
        assert_eq!(
            get_ngrams("this is a test", 3, 3),
            vec! [" th", "thi", "his", "is ", "s t", " te", "tes", "est", "st ", "t t"]
        );
        assert_eq!(
            get_ngrams("this is a test", 5, 3),
            vec! [" this", "this ", "his t", "is te", "s tes", " test",
                  "test ", "est t", "st th", "t thi"]
        );
        assert_eq!(
            get_ngrams("this is a test", 3, 2),
            vec! [" th", "thi", "his", "is ", "s i", " is", "is ", "s t",
                  " te", "tes", "est", "st ", "t t"]
        );
        assert!(get_ngrams("this is a test", 3, 5).is_empty(), "Ngrams not empty");
        assert_eq!(
            get_ngrams("Some frick'in test", 6, 3),
            vec! [" some ", "some t", "ome te", "me tes", "e test", " test ",
                  "test s", "est so", "st som", "t some"]
        );
    }

    #[test]
    fn test_gen_passphrases() {
        let ngrams = get_ngrams("tic toc", 3, 3);
        let min_entropy = 60.0;
        let result = gen_passphrases(ngrams, 5, min_entropy);
        assert!(result.is_ok(), "Passphrase generation failed.");
        let passphrases = result.unwrap();
        assert_eq!(passphrases.len(), 5);
    }

    #[test]
    fn test_gen_passphrases_no_ngrams() {
        let ngrams: Vec<String> = vec![];
        let result = gen_passphrases(ngrams, 1, 60.0);
        assert!(result.is_err(), "No error despite no ngrams.");
    }

    #[test]
    fn test_gen_passphrases_no_entropy() {
        let ngrams = get_ngrams("abc def ghijkl mnopqr stuvw xyz", 2, 1);
        let result = gen_passphrases(ngrams, 1, 60.0);
        assert!(result.is_err(), "No error despite no entropy.");
    }

    #[test]
    fn test_gen_passphrases_no_starting_entropy() {
        let ngrams = get_ngrams("tictoctictactoe", 2, 1);
        let result = gen_passphrases(ngrams, 1, 60.0);
        assert!(result.is_err(), "No error despite no starting entropy.");
    }
}
<commit_msg>Change get_ngrams signature<commit_after>extern crate getopts;
extern crate markovpass;

use markovpass::{PassphraseMarkovChain,PassphraseMarkovChainError};
use std::io::Write;

fn build_opts() -> getopts::Options {
    let mut opts = getopts::Options::new();
    opts.optopt("n", "", "Number of passphrases to generate (default 1)", "NUM");
    opts.optopt("e", "", "Minimum entropy (default 60)", "MINENTROPY");
    opts.optopt("l", "", "NGram length (default 3, must be > 1)", "LENGTH");
    opts.optopt("w", "", "Minimum word length for corpus (default 5)", "LENGTH");
    opts.optflag("h", "help", "Display this help and exit");

    opts
}

fn parse_flag<T: std::str::FromStr>(matches: &getopts::Matches, flag: &str, default: T)
        -> Result<T, &'static str> {
    matches.opt_str(flag)
        .map(|c| c.parse::<T>())
        .unwrap_or(Ok(default))
        .map_err(|_| "Failed to parse flag.")
}

fn parse_args(opts: &getopts::Options, args: &Vec<String>)
        -> Result<(Option<String>, usize, f64, usize, usize), &'static str> {
    let matches = try!(opts.parse(&args[1..]).map_err(|_| "Failed to parse arguments."));

    if matches.opt_present("h") || matches.free.len() > 1 {
        return Err("Failed to parse arguments.");
    };

    let number = try!(parse_flag(&matches, "n", 1));
    let min_entropy = try!(parse_flag(&matches, "e", 60.0));
    let ngram_length = try!(parse_flag(&matches, "l", 3));
    let min_word_length = try!(parse_flag(&matches, "w", 5));

    if ngram_length < 2 {
        return Err("Ngram length must be greater than one.");
    };

    let filename = if matches.free.is_empty() || matches.free[0] == "-" {
        None
    } else {
        Some(matches.free[0].clone())
    };

    Ok((filename, number, min_entropy, ngram_length, min_word_length))
}

fn get_corpus(filename: Option<&str>) -> Result<String, std::io::Error> {
    let mut input: Box<std::io::Read> = match filename {
        Some(filename) => Box::new(try!(std::fs::File::open(&filename))),
        None => Box::new(std::io::stdin()),
    };
    let mut corpus = String::new();
    try!(input.read_to_string(&mut corpus));

    Ok(corpus)
}

fn clean_word(word: &str, min_length: usize) -> Option<&str> {
    let word = word.trim_matches(|c: char| !c.is_alphabetic());

    if word.chars().all(|c: char| c.is_alphabetic()) && word.len() >= min_length {
        Some(word)
    } else {
        None
    }
}

fn get_ngrams(corpus: String, ngram_length: usize, min_word_length: usize) -> Vec<String> {
    let words = corpus.split_whitespace()
      .filter_map(|word| clean_word(word, min_word_length));
    let cleaned_corpus = Some("").into_iter().chain(words).collect::<Vec<&str>>().join(" ");
    let count = cleaned_corpus.chars().count();
    if count < ngram_length { return vec![]; };

    let mut chars = cleaned_corpus.chars().cycle();
    let mut ngrams = Vec::with_capacity(count);
    for _ in 0..count {
        let ngram: String = chars.clone().take(ngram_length).collect();
        ngrams.push(ngram.to_lowercase());
        chars.next();
    };

    ngrams
}

fn gen_passphrases(ngrams: Vec<String>, number: usize, min_entropy: f64)
        -> Result<Vec<(String, f64)>, PassphraseMarkovChainError> {
    let chain = try!(PassphraseMarkovChain::new(ngrams.iter().cloned()));
    let mut passphrases = Vec::with_capacity(number);
    for _ in 0..number {
        passphrases.push(chain.passphrase(min_entropy));
    };

    Ok(passphrases)
}

fn print_usage(program: &str, opts: &getopts::Options) {
    let brief = format!("Usage: {} [FILE] [options]", program);
    print!("{}", opts.usage(&brief));
}

fn write_error(message: &str) {
    writeln!(&mut std::io::stderr(), "{}", &message).expect("Failed to write to stderr.");
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    let program = args[0].clone();
    let opts = build_opts();

    let (filename, number, min_entropy, ngram_length, min_word_length) =
            match parse_args(&opts, &args) {
        Ok(parsed_args) => { parsed_args },
        Err(_) => {
            print_usage(&program, &opts);
            return;
        }
    };
    let corpus = match get_corpus(filename.as_ref().map(String::as_ref)) {
        Ok(corpus) => { corpus },
        Err(_) => {
            let filename_str = filename.unwrap_or("-".to_string());
            write_error(&format!("{}: {}: {}", &program, &filename_str, "Failed to read input."));
            return;
        },
    };
    let ngrams = get_ngrams(corpus, ngram_length, min_word_length);
    let passphrases = match gen_passphrases(ngrams, number, min_entropy) {
        Ok(passphrases) => { passphrases },
        Err(e) => {
            write_error(&format!("{}: {}", &program, e));
            return;
        },
    };

    for (passphrase, entropy) in passphrases {
        println!("{} <{}>", passphrase, entropy);
    };
}

#[cfg(test)]
mod tests {
    use super::{clean_word,get_ngrams,gen_passphrases};

    #[test]
    fn test_clean_word() {
        assert_eq!(clean_word("Test", 3), Some("Test"));
        assert_eq!(clean_word("12'3test@314", 3), Some("test"));
        assert_eq!(clean_word("31ld;Test", 3), None);
        assert_eq!(clean_word("a", 2), None);
        assert_eq!(clean_word("Test", 5), None);
    }

    #[test]
    fn test_get_ngrams() {
        assert_eq!(
            get_ngrams("this is a test", 3, 3),
            vec! [" th", "thi", "his", "is ", "s t", " te", "tes", "est", "st ", "t t"]
        );
        assert_eq!(
            get_ngrams("this is a test", 5, 3),
            vec! [" this", "this ", "his t", "is te", "s tes", " test",
                  "test ", "est t", "st th", "t thi"]
        );
        assert_eq!(
            get_ngrams("this is a test", 3, 2),
            vec! [" th", "thi", "his", "is ", "s i", " is", "is ", "s t",
                  " te", "tes", "est", "st ", "t t"]
        );
        assert!(get_ngrams("this is a test", 3, 5).is_empty(), "Ngrams not empty");
        assert_eq!(
            get_ngrams("Some frick'in test", 6, 3),
            vec! [" some ", "some t", "ome te", "me tes", "e test", " test ",
                  "test s", "est so", "st som", "t some"]
        );
    }

    #[test]
    fn test_gen_passphrases() {
        let ngrams = get_ngrams("tic toc", 3, 3);
        let min_entropy = 60.0;
        let result = gen_passphrases(ngrams, 5, min_entropy);
        assert!(result.is_ok(), "Passphrase generation failed.");
        let passphrases = result.unwrap();
        assert_eq!(passphrases.len(), 5);
    }

    #[test]
    fn test_gen_passphrases_no_ngrams() {
        let ngrams: Vec<String> = vec![];
        let result = gen_passphrases(ngrams, 1, 60.0);
        assert!(result.is_err(), "No error despite no ngrams.");
    }

    #[test]
    fn test_gen_passphrases_no_entropy() {
        let ngrams = get_ngrams("abc def ghijkl mnopqr stuvw xyz", 2, 1);
        let result = gen_passphrases(ngrams, 1, 60.0);
        assert!(result.is_err(), "No error despite no entropy.");
    }

    #[test]
    fn test_gen_passphrases_no_starting_entropy() {
        let ngrams = get_ngrams("tictoctictactoe", 2, 1);
        let result = gen_passphrases(ngrams, 1, 60.0);
        assert!(result.is_err(), "No error despite no starting entropy.");
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use slog::{Drain, Logger};
use slog_async::Async;
use slog_term::{FullFormat, PlainDecorator, TermDecorator};
use std::env;
use std::io::BufWriter;
use std::process;
use std::sync::atomic::{AtomicBool, Ordering, ATOMIC_BOOL_INIT};

use clap::ArgMatches;

use tikv::config::{MetricConfig, TiKvConfig};
use tikv::util;
use tikv::util::collections::HashMap;
use tikv::util::file_log::RotatingFileLogger;
use tikv::util::logger;

// A workaround for checking if log is initialized.
pub static LOG_INITIALIZED: AtomicBool = ATOMIC_BOOL_INIT;

macro_rules! fatal {
    ($lvl:expr, $($arg:tt)+) => ({
        if LOG_INITIALIZED.load(Ordering::SeqCst) {
            error!($lvl, $($arg)+);
        } else {
            eprintln!($lvl, $($arg)+);
        }
        process::exit(1)
    })
}

pub fn init_log(config: &TiKvConfig) {
    if config.log_file.is_empty() {
        let decorator = TermDecorator::new().build();
        let drain = FullFormat::new(decorator).build().fuse();
        let drain = Async::new(drain).build().fuse();
        let logger = Logger::root_typed(drain, slog_o!());
        logger::init_log(logger, config.log_level).unwrap_or_else(|e| {
            fatal!("failed to initialize log: {:?}", e);
        });
    } else {
        let logger = BufWriter::new(RotatingFileLogger::new(&config.log_file).unwrap_or_else(
            |e| {
                fatal!(
                    "failed to initialize log with file {:?}: {:?}",
                    config.log_file,
                    e
                );
            },
        ));
        let decorator = PlainDecorator::new(logger);
        let drain = FullFormat::new(decorator).build().fuse();
        let drain = Async::new(drain).build().fuse();
        let logger = Logger::root_typed(drain, slog_o!());
        logger::init_log(logger, config.log_level).unwrap_or_else(|e| {
            fatal!("failed to initialize log: {:?}", e);
        });
    }
    LOG_INITIALIZED.store(true, Ordering::SeqCst);
}

pub fn initial_metric(cfg: &MetricConfig, node_id: Option<u64>) {
    if cfg.interval.as_secs() == 0 || cfg.address.is_empty() {
        return;
    }

    let mut push_job = cfg.job.clone();
    if let Some(id) = node_id {
        push_job.push_str(&format!("_{}", id));
    }

    info!("start prometheus client");

    util::metrics::monitor_threads("tikv")
        .unwrap_or_else(|e| fatal!("failed to start monitor thread: {:?}", e));

    util::metrics::run_prometheus(cfg.interval.0, &cfg.address, &push_job);
}

pub fn overwrite_config_with_cmd_args(config: &mut TiKvConfig, matches: &ArgMatches) {
    if let Some(level) = matches.value_of("log-level") {
        config.log_level = logger::get_level_by_string(level).unwrap();
    }

    if let Some(file) = matches.value_of("log-file") {
        config.log_file = file.to_owned();
    }

    if let Some(addr) = matches.value_of("addr") {
        config.server.addr = addr.to_owned();
    }

    if let Some(advertise_addr) = matches.value_of("advertise-addr") {
        config.server.advertise_addr = advertise_addr.to_owned();
    }

    if let Some(data_dir) = matches.value_of("data-dir") {
        config.storage.data_dir = data_dir.to_owned();
    }

    if let Some(endpoints) = matches.values_of("pd-endpoints") {
        config.pd.endpoints = endpoints.map(|e| e.to_owned()).collect();
    }

    if let Some(labels_vec) = matches.values_of("labels") {
        let mut labels = HashMap::default();
        labels_vec
            .map(|s| {
                let mut parts = s.split('=');
                let key = parts.next().unwrap().to_owned();
                let value = match parts.next() {
                    None => fatal!("invalid label: {:?}", s),
                    Some(v) => v.to_owned(),
                };
                if parts.next().is_some() {
                    fatal!("invalid label: {:?}", s);
                }
                labels.insert(key, value);
            })
            .count();
        config.server.labels = labels;
    }

    if let Some(capacity_str) = matches.value_of("capacity") {
        let capacity = capacity_str.parse().unwrap_or_else(|e| {
            fatal!("invalid capacity: {}", e);
        });
        config.raft_store.capacity = capacity;
    }

    if let Some(import_dir) = matches.value_of("import-dir") {
        config.import.import_dir = import_dir.to_owned();
    }
}

// Set gRPC event engine to epollsig.
// See more: https://github.com/grpc/grpc/blob/486761d04e03a9183d8013eddd86c3134d52d459\
//           /src/core/lib/iomgr/ev_posix.cc#L149
pub fn configure_grpc_poll_strategy() {
    const GRPC_POLL_STRATEGY: &str = "GRPC_POLL_STRATEGY";
    const DEFAULT_ENGINE: &str = "epollsig";
    if cfg!(target_os = "linux") && env::var(GRPC_POLL_STRATEGY).is_err() {
        // Set to epollsig if it is not specified.
        env::set_var(GRPC_POLL_STRATEGY, DEFAULT_ENGINE);
    }
}
<commit_msg>logger: Bigger channel and blocks on overflow. (#3191)<commit_after>// Copyright 2018 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use slog::{Drain, Logger};
use slog_async::{Async, OverflowStrategy};
use slog_term::{FullFormat, PlainDecorator, TermDecorator};
use std::env;
use std::io::BufWriter;
use std::process;
use std::sync::atomic::{AtomicBool, Ordering, ATOMIC_BOOL_INIT};

use clap::ArgMatches;

use tikv::config::{MetricConfig, TiKvConfig};
use tikv::util;
use tikv::util::collections::HashMap;
use tikv::util::file_log::RotatingFileLogger;
use tikv::util::logger;

// A workaround for checking if log is initialized.
pub static LOG_INITIALIZED: AtomicBool = ATOMIC_BOOL_INIT;
// Default is 128.
// Extended since blocking is set, and we don't want to block very often.
const SLOG_CHANNEL_SIZE: usize = 10240;
// Default is DropAndReport.
// It is not desirable to have dropped logs in our use case.
const SLOG_CHANNEL_OVERFLOW_STRATEGY: OverflowStrategy = OverflowStrategy::Block;

macro_rules! fatal {
    ($lvl:expr, $($arg:tt)+) => ({
        if LOG_INITIALIZED.load(Ordering::SeqCst) {
            error!($lvl, $($arg)+);
        } else {
            eprintln!($lvl, $($arg)+);
        }
        process::exit(1)
    })
}

pub fn init_log(config: &TiKvConfig) {
    if config.log_file.is_empty() {
        let decorator = TermDecorator::new().build();
        let drain = FullFormat::new(decorator).build().fuse();
        let drain = Async::new(drain)
            .chan_size(SLOG_CHANNEL_SIZE)
            .overflow_strategy(SLOG_CHANNEL_OVERFLOW_STRATEGY)
            .build()
            .fuse();
        let logger = Logger::root_typed(drain, slog_o!());
        logger::init_log(logger, config.log_level).unwrap_or_else(|e| {
            fatal!("failed to initialize log: {:?}", e);
        });
    } else {
        let logger = BufWriter::new(RotatingFileLogger::new(&config.log_file).unwrap_or_else(
            |e| {
                fatal!(
                    "failed to initialize log with file {:?}: {:?}",
                    config.log_file,
                    e
                );
            },
        ));
        let decorator = PlainDecorator::new(logger);
        let drain = FullFormat::new(decorator).build().fuse();
        let drain = Async::new(drain).build().fuse();
        let logger = Logger::root_typed(drain, slog_o!());
        logger::init_log(logger, config.log_level).unwrap_or_else(|e| {
            fatal!("failed to initialize log: {:?}", e);
        });
    }
    LOG_INITIALIZED.store(true, Ordering::SeqCst);
}

pub fn initial_metric(cfg: &MetricConfig, node_id: Option<u64>) {
    if cfg.interval.as_secs() == 0 || cfg.address.is_empty() {
        return;
    }

    let mut push_job = cfg.job.clone();
    if let Some(id) = node_id {
        push_job.push_str(&format!("_{}", id));
    }

    info!("start prometheus client");

    util::metrics::monitor_threads("tikv")
        .unwrap_or_else(|e| fatal!("failed to start monitor thread: {:?}", e));

    util::metrics::run_prometheus(cfg.interval.0, &cfg.address, &push_job);
}

pub fn overwrite_config_with_cmd_args(config: &mut TiKvConfig, matches: &ArgMatches) {
    if let Some(level) = matches.value_of("log-level") {
        config.log_level = logger::get_level_by_string(level).unwrap();
    }

    if let Some(file) = matches.value_of("log-file") {
        config.log_file = file.to_owned();
    }

    if let Some(addr) = matches.value_of("addr") {
        config.server.addr = addr.to_owned();
    }

    if let Some(advertise_addr) = matches.value_of("advertise-addr") {
        config.server.advertise_addr = advertise_addr.to_owned();
    }

    if let Some(data_dir) = matches.value_of("data-dir") {
        config.storage.data_dir = data_dir.to_owned();
    }

    if let Some(endpoints) = matches.values_of("pd-endpoints") {
        config.pd.endpoints = endpoints.map(|e| e.to_owned()).collect();
    }

    if let Some(labels_vec) = matches.values_of("labels") {
        let mut labels = HashMap::default();
        labels_vec
            .map(|s| {
                let mut parts = s.split('=');
                let key = parts.next().unwrap().to_owned();
                let value = match parts.next() {
                    None => fatal!("invalid label: {:?}", s),
                    Some(v) => v.to_owned(),
                };
                if parts.next().is_some() {
                    fatal!("invalid label: {:?}", s);
                }
                labels.insert(key, value);
            })
            .count();
        config.server.labels = labels;
    }

    if let Some(capacity_str) = matches.value_of("capacity") {
        let capacity = capacity_str.parse().unwrap_or_else(|e| {
            fatal!("invalid capacity: {}", e);
        });
        config.raft_store.capacity = capacity;
    }

    if let Some(import_dir) = matches.value_of("import-dir") {
        config.import.import_dir = import_dir.to_owned();
    }
}

// Set gRPC event engine to epollsig.
// See more: https://github.com/grpc/grpc/blob/486761d04e03a9183d8013eddd86c3134d52d459\
//           /src/core/lib/iomgr/ev_posix.cc#L149
pub fn configure_grpc_poll_strategy() {
    const GRPC_POLL_STRATEGY: &str = "GRPC_POLL_STRATEGY";
    const DEFAULT_ENGINE: &str = "epollsig";
    if cfg!(target_os = "linux") && env::var(GRPC_POLL_STRATEGY).is_err() {
        // Set to epollsig if it is not specified.
        env::set_var(GRPC_POLL_STRATEGY, DEFAULT_ENGINE);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::env;
use std::ffi::OsString;
use std::fs::{self, File};
use std::io::Write;
use std::os::unix::fs::{symlink as fs_symlink, PermissionsExt};
use std::os::unix::process::CommandExt;
use std::path::Path;
use std::process::{Command, Stdio};

use libc;

use {Error, Result};
use mount::{self, Mount};
use filesystem;
use pty;

const ROOTFS_DIRS: &'static [&'static str] = &[
    "etc",
    "run",
    "src",
    "var",
    "hab/pkgs",
    "hab/cache/keys",
    "hab/cache/artifacts",
];

const WHITELISTED_DEVS: &'static [&'static str] =
    &["null", "random", "full", "tty", "zero", "urandom"];

const READONLY_PATHS: &'static [&'static str] = &[
    "/proc/asound",
    "/proc/bus",
    "/proc/fs",
    "/proc/irq",
    "/proc/sys",
    "/proc/sysrq-trigger",
];

const MASKED_PATHS: &'static [&'static str] = &[
    "/proc/kcore",
    "/proc/latency_stats",
    "/proc/timer_list",
    "/proc/timer_stats",
    "/proc/sched_debug",
    "/sys/firmware",
];

pub fn run(rootfs: &Path, cmd: &str, args: Vec<OsString>) -> Result<()> {
    let umask_val = 0o0022;
    debug!("setting umask, value={:#o}", umask_val);
    umask(umask_val);

    for entry in ROOTFS_DIRS {
        mkdir_p(rootfs.join(entry))?;
    }

    mount::private("/", Mount::Recursive)?;
    mount::bind(rootfs, rootfs, Mount::Recursive, None)?;

    // Mount `/proc` from namespace
    mkdir_p(rootfs.join("proc"))?;
    mount::procfs(rootfs.join("proc"))?;

    // Create `/dev` filesystem
    let dev = rootfs.join("dev");
    mkdir_p(&dev)?;
    mount::tmpfs(
        "tmpfs",
        &dev,
        Some(libc::MS_NOSUID | libc::MS_NODEV | libc::MS_STRICTATIME),
        Some(0o755),
        Some(65536),
    )?;
    chmod(&dev, 0o0755)?;

    // Create `/dev/pts` filesystem
    let pts = rootfs.join("dev/pts");
    mkdir_p(&pts)?;
    mount::devpts(
        &pts,
        libc::MS_NOSUID | libc::MS_NOEXEC,
        Some(0o0620),
        Some(0o0666),
    )?;
    chmod(&pts, 0o0755)?;

    // Create `/dev/shm` filesystem
    let shm = rootfs.join("dev/shm");
    mkdir_p(&shm)?;
    mount::tmpfs(
        "shm",
        &shm,
        Some(libc::MS_NOSUID | libc::MS_NOEXEC | libc::MS_NODEV),
        Some(0o1777),
        Some(65536),
    )?;

    // Create `/dev/mqueue` filesystem
    let mqueue = rootfs.join("dev/mqueue");
    mkdir_p(&mqueue)?;
    mount::mqueue(&mqueue, libc::MS_NOSUID | libc::MS_NOEXEC | libc::MS_NODEV)?;

    // Create `/sys` filesystem
    let sys = rootfs.join("sys");
    mkdir_p(&sys)?;
    mount::bind(
        "/sys",
        sys,
        Mount::Recursive,
        Some(libc::MS_NOSUID | libc::MS_NOEXEC | libc::MS_NODEV),
    )?;

    // Bind mount whitelisted devices
    let old_umask = umask(0o0000);
    for entry in WHITELISTED_DEVS {
        let source = Path::new("/dev").join(entry);
        let target = rootfs.join("dev").join(entry);
        touch(&target)?;
        mount::bind(&source, &target, Mount::Nonrecursive, None)?;
    }
    umask(old_umask);

    // Setup ptmx symlink
    symlink("pts/ptmx", rootfs.join("dev/ptmx"))?;

    // Setup dev symlinks
    symlink("/proc/self/fd", rootfs.join("dev/fd"))?;
    symlink("/proc/self/fd/0", rootfs.join("dev/stdin"))?;
    symlink("/proc/self/fd/1", rootfs.join("dev/stdout"))?;
    symlink("/proc/self/fd/2", rootfs.join("dev/stderr"))?;
    if Path::new("/proc/kcore").exists() {
        symlink("/proc/kcore", rootfs.join("dev/core"))?;
    }

    // Bind mount /src
    mount::bind(
        env::current_dir()?,
        rootfs.join("src"),
        Mount::Nonrecursive,
        None,
    )?;

    // TODO fn: is this necessary?
    mount::tmpfs("tmpfs", rootfs.join("run"), None, None, None)?;
    symlink("/run", rootfs.join("var/run"))?;

    touch(rootfs.join("etc/hosts"))?;
    {
        let path = rootfs.join("etc/resolv.conf");
        debug!("creating file, path={}", path.display());
        let mut file = File::create(path)?;
        file.write_all("nameserver 8.8.8.8".as_bytes())?;
    }

    // Bind mount Habitat packages and key cache as read-only
    mount::bind(
        Path::new("/hab/pkgs"),
        rootfs.join("hab/pkgs"),
        Mount::Nonrecursive,
        Some(libc::MS_RDONLY),
    )?;
    let source = env::home_dir().ok_or(Error::HomeDirectoryNotFound)?.join(
        ".hab/cache/keys",
    );
    mkdir_p(&source)?;
    mount::bind(
        source,
        rootfs.join("hab/cache/keys"),
        Mount::Nonrecursive,
        Some(libc::MS_RDONLY),
    )?;

    // Bind mount outside artifact cache (and ensure outside directory exists)
    let source = env::home_dir().ok_or(Error::HomeDirectoryNotFound)?.join(
        ".hab/cache/artifacts",
    );
    mkdir_p(&source)?;
    mount::bind(
        source,
        rootfs.join("hab/cache/artifacts"),
        Mount::Nonrecursive,
        None,
    )?;

    // Binlink BusyBox and Habitat CLI binaries
    {
        let mut command = Command::new("hab");
        command.args(&["pkg", "binlink", "core/busybox-static"]);
        command.env("FS_ROOT", rootfs);
        command.stdout(Stdio::null());
        debug!("running, command={:?}", &command);
        command.spawn()?.wait()?;
    }
    {
        let mut command = Command::new("hab");
        command.args(&["pkg", "binlink", "core/hab"]);
        command.env("FS_ROOT", rootfs);
        command.stdout(Stdio::null());
        debug!("running, command={:?}", &command);
        command.spawn()?.wait()?;
    }

    // Change the root file system, via `pivot_root(2)`
    mkdir_p(rootfs.join(".pivot_root"))?;
    filesystem::pivot_root(&rootfs, rootfs.join(".pivot_root"))?;
    debug!("setting current directory, path=/");
    env::set_current_dir("/")?;
    mount::private("/.pivot_root", Mount::Recursive)?;
    mount::umount("/.pivot_root", Some(libc::MNT_DETACH))?;
    rmdir("/.pivot_root")?;

    // Set read-only paths
    for entry in READONLY_PATHS {
        let path = Path::new(entry);
        if path.exists() {
            mount::bind(&path, &path, Mount::Recursive, None)?;
            mount::bind(
                &path,
                &path,
                Mount::Recursive,
                Some(libc::MS_RDONLY | libc::MS_REMOUNT),
            )?;
        }
    }

    // Set masked paths
    for entry in MASKED_PATHS {
        let path = Path::new(entry);
        if path.is_dir() {
            mount::tmpfs("tmpfs", &path, Some(libc::MS_RDONLY), None, None)?;
        } else if path.is_file() {
            mount::bind("/dev/null", &path, Mount::Nonrecursive, None)?;
        }
    }

    // Set final umask
    let umask_val = 0o0022;
    debug!("setting umask, value={:#o}", umask_val);
    umask(umask_val);

    // Change into `/src` directory
    let src_path = Path::new("/src");
    debug!("setting current directory, path={}", src_path.display());
    env::set_current_dir(src_path)?;

    // Setup `/dev/console`
    let console = Path::new("/dev/console");
    let master = pty::Master::default()?;
    let ptsname = master.ptsname()?;
    touch(&console)?;
    chmod(&console, 0o0000)?;
    mount::bind(&ptsname, &console, Mount::Nonrecursive, None)?;

    // Finally, call `exec` to become the target program
    exec_command(cmd, args)
}

fn mkdir_p<P: AsRef<Path>>(path: P) -> Result<()> {
    debug!("creating directory, path={}", path.as_ref().display());
    Ok(fs::create_dir_all(path)?)
}

fn rmdir<P: AsRef<Path>>(path: P) -> Result<()> {
    debug!("removing directory, path={}", path.as_ref().display());
    Ok(fs::remove_dir(path)?)
}

fn symlink<S, T>(source: S, target: T) -> Result<()>
where
    S: AsRef<Path>,
    T: AsRef<Path>,
{
    debug!(
        "symlinking, src={}, target={}",
        source.as_ref().display(),
        target.as_ref().display()
    );
    Ok(fs_symlink(source, target)?)
}

fn touch<P: AsRef<Path>>(path: P) -> Result<()> {
    debug!("creating file, path={}", path.as_ref().display());
    let _ = File::create(path)?;
    Ok(())
}

fn umask(mode: libc::mode_t) -> libc::mode_t {
    unsafe { libc::umask(mode) }
}

fn chmod<P: AsRef<Path>>(path: P, mode: u32) -> Result<()> {
    let md = path.as_ref().metadata()?;
    let mut perms = md.permissions();
    perms.set_mode(mode);
    Ok(())
}

fn exec_command(cmd: &str, args: Vec<OsString>) -> Result<()> {
    debug!("running, command={} args={:?}", cmd, args);
    Command::new(cmd).args(args).exec();
    Ok(())
}
<commit_msg>Optionally mount artifact cache.<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::env;
use std::ffi::OsString;
use std::fs::{self, File};
use std::io::Write;
use std::os::unix::fs::{symlink as fs_symlink, PermissionsExt};
use std::os::unix::process::CommandExt;
use std::path::Path;
use std::process::{Command, Stdio};

use libc;

use {Error, Result};
use mount::{self, Mount};
use filesystem;
use pty;

const MOUNT_ARTIFACT_CACHE_ENVVAR: &'static str = "MOUNT_ARTIFACT_CACHE";

const ROOTFS_DIRS: &'static [&'static str] = &[
    "etc",
    "run",
    "src",
    "var",
    "hab/pkgs",
    "hab/cache/keys",
    "hab/cache/artifacts",
];

const WHITELISTED_DEVS: &'static [&'static str] =
    &["null", "random", "full", "tty", "zero", "urandom"];

const READONLY_PATHS: &'static [&'static str] = &[
    "/proc/asound",
    "/proc/bus",
    "/proc/fs",
    "/proc/irq",
    "/proc/sys",
    "/proc/sysrq-trigger",
];

const MASKED_PATHS: &'static [&'static str] = &[
    "/proc/kcore",
    "/proc/latency_stats",
    "/proc/timer_list",
    "/proc/timer_stats",
    "/proc/sched_debug",
    "/sys/firmware",
];

pub fn run(rootfs: &Path, cmd: &str, args: Vec<OsString>) -> Result<()> {
    let umask_val = 0o0022;
    debug!("setting umask, value={:#o}", umask_val);
    umask(umask_val);

    for entry in ROOTFS_DIRS {
        mkdir_p(rootfs.join(entry))?;
    }

    mount::private("/", Mount::Recursive)?;
    mount::bind(rootfs, rootfs, Mount::Recursive, None)?;

    // Mount `/proc` from namespace
    mkdir_p(rootfs.join("proc"))?;
    mount::procfs(rootfs.join("proc"))?;

    // Create `/dev` filesystem
    let dev = rootfs.join("dev");
    mkdir_p(&dev)?;
    mount::tmpfs(
        "tmpfs",
        &dev,
        Some(libc::MS_NOSUID | libc::MS_NODEV | libc::MS_STRICTATIME),
        Some(0o755),
        Some(65536),
    )?;
    chmod(&dev, 0o0755)?;

    // Create `/dev/pts` filesystem
    let pts = rootfs.join("dev/pts");
    mkdir_p(&pts)?;
    mount::devpts(
        &pts,
        libc::MS_NOSUID | libc::MS_NOEXEC,
        Some(0o0620),
        Some(0o0666),
    )?;
    chmod(&pts, 0o0755)?;

    // Create `/dev/shm` filesystem
    let shm = rootfs.join("dev/shm");
    mkdir_p(&shm)?;
    mount::tmpfs(
        "shm",
        &shm,
        Some(libc::MS_NOSUID | libc::MS_NOEXEC | libc::MS_NODEV),
        Some(0o1777),
        Some(65536),
    )?;

    // Create `/dev/mqueue` filesystem
    let mqueue = rootfs.join("dev/mqueue");
    mkdir_p(&mqueue)?;
    mount::mqueue(&mqueue, libc::MS_NOSUID | libc::MS_NOEXEC | libc::MS_NODEV)?;

    // Create `/sys` filesystem
    let sys = rootfs.join("sys");
    mkdir_p(&sys)?;
    mount::bind(
        "/sys",
        sys,
        Mount::Recursive,
        Some(libc::MS_NOSUID | libc::MS_NOEXEC | libc::MS_NODEV),
    )?;

    // Bind mount whitelisted devices
    let old_umask = umask(0o0000);
    for entry in WHITELISTED_DEVS {
        let source = Path::new("/dev").join(entry);
        let target = rootfs.join("dev").join(entry);
        touch(&target)?;
        mount::bind(&source, &target, Mount::Nonrecursive, None)?;
    }
    umask(old_umask);

    // Setup ptmx symlink
    symlink("pts/ptmx", rootfs.join("dev/ptmx"))?;

    // Setup dev symlinks
    symlink("/proc/self/fd", rootfs.join("dev/fd"))?;
    symlink("/proc/self/fd/0", rootfs.join("dev/stdin"))?;
    symlink("/proc/self/fd/1", rootfs.join("dev/stdout"))?;
    symlink("/proc/self/fd/2", rootfs.join("dev/stderr"))?;
    if Path::new("/proc/kcore").exists() {
        symlink("/proc/kcore", rootfs.join("dev/core"))?;
    }

    // Bind mount /src
    mount::bind(
        env::current_dir()?,
        rootfs.join("src"),
        Mount::Nonrecursive,
        None,
    )?;

    // TODO fn: is this necessary?
    mount::tmpfs("tmpfs", rootfs.join("run"), None, None, None)?;
    symlink("/run", rootfs.join("var/run"))?;

    touch(rootfs.join("etc/hosts"))?;
    {
        let path = rootfs.join("etc/resolv.conf");
        debug!("creating file, path={}", path.display());
        let mut file = File::create(path)?;
        file.write_all("nameserver 8.8.8.8".as_bytes())?;
    }

    // Bind mount Habitat packages and key cache as read-only
    mount::bind(
        Path::new("/hab/pkgs"),
        rootfs.join("hab/pkgs"),
        Mount::Nonrecursive,
        Some(libc::MS_RDONLY),
    )?;
    let source = env::home_dir().ok_or(Error::HomeDirectoryNotFound)?.join(
        ".hab/cache/keys",
    );
    mkdir_p(&source)?;
    mount::bind(
        source,
        rootfs.join("hab/cache/keys"),
        Mount::Nonrecursive,
        Some(libc::MS_RDONLY),
    )?;

    if env::var(MOUNT_ARTIFACT_CACHE_ENVVAR).is_ok() {
        // Bind mount outside artifact cache (and ensure outside directory exists)
        let source = env::home_dir().ok_or(Error::HomeDirectoryNotFound)?.join(
            ".hab/cache/artifacts",
        );
        mkdir_p(&source)?;
        mount::bind(
            source,
            rootfs.join("hab/cache/artifacts"),
            Mount::Nonrecursive,
            None,
        )?;
    }

    // Binlink BusyBox and Habitat CLI binaries
    {
        let mut command = Command::new("hab");
        command.args(&["pkg", "binlink", "core/busybox-static"]);
        command.env("FS_ROOT", rootfs);
        command.stdout(Stdio::null());
        debug!("running, command={:?}", &command);
        command.spawn()?.wait()?;
    }
    {
        let mut command = Command::new("hab");
        command.args(&["pkg", "binlink", "core/hab"]);
        command.env("FS_ROOT", rootfs);
        command.stdout(Stdio::null());
        debug!("running, command={:?}", &command);
        command.spawn()?.wait()?;
    }

    // Change the root file system, via `pivot_root(2)`
    mkdir_p(rootfs.join(".pivot_root"))?;
    filesystem::pivot_root(&rootfs, rootfs.join(".pivot_root"))?;
    debug!("setting current directory, path=/");
    env::set_current_dir("/")?;
    mount::private("/.pivot_root", Mount::Recursive)?;
    mount::umount("/.pivot_root", Some(libc::MNT_DETACH))?;
    rmdir("/.pivot_root")?;

    // Set read-only paths
    for entry in READONLY_PATHS {
        let path = Path::new(entry);
        if path.exists() {
            mount::bind(&path, &path, Mount::Recursive, None)?;
            mount::bind(
                &path,
                &path,
                Mount::Recursive,
                Some(libc::MS_RDONLY | libc::MS_REMOUNT),
            )?;
        }
    }

    // Set masked paths
    for entry in MASKED_PATHS {
        let path = Path::new(entry);
        if path.is_dir() {
            mount::tmpfs("tmpfs", &path, Some(libc::MS_RDONLY), None, None)?;
        } else if path.is_file() {
            mount::bind("/dev/null", &path, Mount::Nonrecursive, None)?;
        }
    }

    // Set final umask
    let umask_val = 0o0022;
    debug!("setting umask, value={:#o}", umask_val);
    umask(umask_val);

    // Change into `/src` directory
    let src_path = Path::new("/src");
    debug!("setting current directory, path={}", src_path.display());
    env::set_current_dir(src_path)?;

    // Setup `/dev/console`
    let console = Path::new("/dev/console");
    let master = pty::Master::default()?;
    let ptsname = master.ptsname()?;
    touch(&console)?;
    chmod(&console, 0o0000)?;
    mount::bind(&ptsname, &console, Mount::Nonrecursive, None)?;

    // Finally, call `exec` to become the target program
    exec_command(cmd, args)
}

fn mkdir_p<P: AsRef<Path>>(path: P) -> Result<()> {
    debug!("creating directory, path={}", path.as_ref().display());
    Ok(fs::create_dir_all(path)?)
}

fn rmdir<P: AsRef<Path>>(path: P) -> Result<()> {
    debug!("removing directory, path={}", path.as_ref().display());
    Ok(fs::remove_dir(path)?)
}

fn symlink<S, T>(source: S, target: T) -> Result<()>
where
    S: AsRef<Path>,
    T: AsRef<Path>,
{
    debug!(
        "symlinking, src={}, target={}",
        source.as_ref().display(),
        target.as_ref().display()
    );
    Ok(fs_symlink(source, target)?)
}

fn touch<P: AsRef<Path>>(path: P) -> Result<()> {
    debug!("creating file, path={}", path.as_ref().display());
    let _ = File::create(path)?;
    Ok(())
}

fn umask(mode: libc::mode_t) -> libc::mode_t {
    unsafe { libc::umask(mode) }
}

fn chmod<P: AsRef<Path>>(path: P, mode: u32) -> Result<()> {
    let md = path.as_ref().metadata()?;
    let mut perms = md.permissions();
    perms.set_mode(mode);
    Ok(())
}

fn exec_command(cmd: &str, args: Vec<OsString>) -> Result<()> {
    debug!("running, command={} args={:?}", cmd, args);
    Command::new(cmd).args(args).exec();
    Ok(())
}
<|endoftext|>


<commit_before>//! This module contains paths to types and functions Clippy needs to know
//! about.

pub const ANY_TRAIT: [&str; 3] = ["std", "any", "Any"];
pub const ARC: [&str; 3] = ["alloc", "sync", "Arc"];
pub const ARC_PTR_EQ: [&str; 4] = ["alloc", "sync", "Arc", "ptr_eq"];
pub const ASMUT_TRAIT: [&str; 3] = ["core", "convert", "AsMut"];
pub const ASREF_TRAIT: [&str; 3] = ["core", "convert", "AsRef"];
pub const BEGIN_PANIC: [&str; 3] = ["std", "panicking", "begin_panic"];
pub const BEGIN_PANIC_FMT: [&str; 3] = ["std", "panicking", "begin_panic_fmt"];
pub const BINARY_HEAP: [&str; 4] = ["alloc", "collections", "binary_heap", "BinaryHeap"];
pub const BORROW_TRAIT: [&str; 3] = ["core", "borrow", "Borrow"];
pub const BOX: [&str; 3] = ["alloc", "boxed", "Box"];
pub const BTREEMAP: [&str; 5] = ["alloc", "collections", "btree", "map", "BTreeMap"];
pub const BTREEMAP_ENTRY: [&str; 5] = ["alloc", "collections", "btree", "map", "Entry"];
pub const BTREESET: [&str; 5] = ["alloc", "collections", "btree", "set", "BTreeSet"];
pub const CLONE_TRAIT: [&str; 3] = ["core", "clone", "Clone"];
pub const CLONE_TRAIT_METHOD: [&str; 4] = ["core", "clone", "Clone", "clone"];
pub const CMP_MAX: [&str; 3] = ["core", "cmp", "max"];
pub const CMP_MIN: [&str; 3] = ["core", "cmp", "min"];
pub const COW: [&str; 3] = ["alloc", "borrow", "Cow"];
pub const CSTRING: [&str; 4] = ["std", "ffi", "c_str", "CString"];
pub const CSTRING_AS_C_STR: [&str; 5] = ["std", "ffi", "c_str", "CString", "as_c_str"];
pub const DEFAULT_TRAIT: [&str; 3] = ["core", "default", "Default"];
pub const DEFAULT_TRAIT_METHOD: [&str; 4] = ["core", "default", "Default", "default"];
pub const DEREF_MUT_TRAIT_METHOD: [&str; 5] = ["core", "ops", "deref", "DerefMut", "deref_mut"];
pub const DEREF_TRAIT_METHOD: [&str; 5] = ["core", "ops", "deref", "Deref", "deref"];
pub const DISPLAY_FMT_METHOD: [&str; 4] = ["core", "fmt", "Display", "fmt"];
pub const DISPLAY_TRAIT: [&str; 3] = ["core", "fmt", "Display"];
pub const DOUBLE_ENDED_ITERATOR: [&str; 4] = ["core", "iter", "traits", "DoubleEndedIterator"];
pub const DROP: [&str; 3] = ["core", "mem", "drop"];
pub const DROP_TRAIT: [&str; 4] = ["core", "ops", "drop", "Drop"];
pub const DURATION: [&str; 3] = ["core", "time", "Duration"];
pub const EARLY_CONTEXT: [&str; 4] = ["rustc", "lint", "context", "EarlyContext"];
pub const EXIT: [&str; 3] = ["std", "process", "exit"];
pub const FILE: [&str; 3] = ["std", "fs", "File"];
pub const FILE_TYPE: [&str; 3] = ["std", "fs", "FileType"];
pub const FMT_ARGUMENTS_NEW_V1: [&str; 4] = ["core", "fmt", "Arguments", "new_v1"];
pub const FMT_ARGUMENTS_NEW_V1_FORMATTED: [&str; 4] = ["core", "fmt", "Arguments", "new_v1_formatted"];
pub const FMT_ARGUMENTV1_NEW: [&str; 4] = ["core", "fmt", "ArgumentV1", "new"];
pub const FROM_FROM: [&str; 4] = ["core", "convert", "From", "from"];
pub const FROM_TRAIT: [&str; 3] = ["core", "convert", "From"];
pub const HASH: [&str; 2] = ["hash", "Hash"];
pub const HASHMAP: [&str; 5] = ["std", "collections", "hash", "map", "HashMap"];
pub const HASHMAP_ENTRY: [&str; 5] = ["std", "collections", "hash", "map", "Entry"];
pub const HASHSET: [&str; 5] = ["std", "collections", "hash", "set", "HashSet"];
pub const INDEX: [&str; 3] = ["core", "ops", "Index"];
pub const INDEX_MUT: [&str; 3] = ["core", "ops", "IndexMut"];
pub const INTO: [&str; 3] = ["core", "convert", "Into"];
pub const INTO_ITERATOR: [&str; 5] = ["core", "iter", "traits", "collect", "IntoIterator"];
pub const IO_READ: [&str; 3] = ["std", "io", "Read"];
pub const IO_WRITE: [&str; 3] = ["std", "io", "Write"];
pub const ITERATOR: [&str; 5] = ["core", "iter", "traits", "iterator", "Iterator"];
pub const LATE_CONTEXT: [&str; 4] = ["rustc", "lint", "context", "LateContext"];
pub const LINKED_LIST: [&str; 4] = ["alloc", "collections", "linked_list", "LinkedList"];
pub const LINT: [&str; 3] = ["rustc_session", "lint", "Lint"];
pub const MEM_DISCRIMINANT: [&str; 3] = ["core", "mem", "discriminant"];
pub const MEM_FORGET: [&str; 3] = ["core", "mem", "forget"];
pub const MEM_MAYBEUNINIT: [&str; 4] = ["core", "mem", "maybe_uninit", "MaybeUninit"];
pub const MEM_MAYBEUNINIT_UNINIT: [&str; 5] = ["core", "mem", "maybe_uninit", "MaybeUninit", "uninit"];
pub const MEM_REPLACE: [&str; 3] = ["core", "mem", "replace"];
pub const MEM_UNINITIALIZED: [&str; 3] = ["core", "mem", "uninitialized"];
pub const MEM_ZEROED: [&str; 3] = ["core", "mem", "zeroed"];
pub const MUTEX: [&str; 4] = ["std", "sync", "mutex", "Mutex"];
pub const MUTEX_GUARD: [&str; 4] = ["std", "sync", "mutex", "MutexGuard"];
pub const OPEN_OPTIONS: [&str; 3] = ["std", "fs", "OpenOptions"];
pub const OPS_MODULE: [&str; 2] = ["core", "ops"];
pub const OPTION: [&str; 3] = ["core", "option", "Option"];
pub const OPTION_NONE: [&str; 4] = ["core", "option", "Option", "None"];
pub const OPTION_SOME: [&str; 4] = ["core", "option", "Option", "Some"];
pub const ORD: [&str; 3] = ["core", "cmp", "Ord"];
pub const OS_STRING: [&str; 4] = ["std", "ffi", "os_str", "OsString"];
pub const OS_STRING_AS_OS_STR: [&str; 5] = ["std", "ffi", "os_str", "OsString", "as_os_str"];
pub const OS_STR_TO_OS_STRING: [&str; 5] = ["std", "ffi", "os_str", "OsStr", "to_os_string"];
pub const PATH: [&str; 3] = ["std", "path", "Path"];
pub const PATH_BUF: [&str; 3] = ["std", "path", "PathBuf"];
pub const PATH_BUF_AS_PATH: [&str; 4] = ["std", "path", "PathBuf", "as_path"];
pub const PATH_TO_PATH_BUF: [&str; 4] = ["std", "path", "Path", "to_path_buf"];
pub const PTR_EQ: [&str; 3] = ["core", "ptr", "eq"];
pub const PTR_NULL: [&str; 2] = ["ptr", "null"];
pub const PTR_NULL_MUT: [&str; 2] = ["ptr", "null_mut"];
pub const RANGE: [&str; 3] = ["core", "ops", "Range"];
pub const RANGE_ARGUMENT_TRAIT: [&str; 3] = ["core", "ops", "RangeBounds"];
pub const RANGE_FROM: [&str; 3] = ["core", "ops", "RangeFrom"];
pub const RANGE_FROM_STD: [&str; 3] = ["std", "ops", "RangeFrom"];
pub const RANGE_FULL: [&str; 3] = ["core", "ops", "RangeFull"];
pub const RANGE_FULL_STD: [&str; 3] = ["std", "ops", "RangeFull"];
pub const RANGE_INCLUSIVE_NEW: [&str; 4] = ["core", "ops", "RangeInclusive", "new"];
pub const RANGE_INCLUSIVE_STD_NEW: [&str; 4] = ["std", "ops", "RangeInclusive", "new"];
pub const RANGE_STD: [&str; 3] = ["std", "ops", "Range"];
pub const RANGE_TO: [&str; 3] = ["core", "ops", "RangeTo"];
pub const RANGE_TO_INCLUSIVE: [&str; 3] = ["core", "ops", "RangeToInclusive"];
pub const RANGE_TO_INCLUSIVE_STD: [&str; 3] = ["std", "ops", "RangeToInclusive"];
pub const RANGE_TO_STD: [&str; 3] = ["std", "ops", "RangeTo"];
pub const RC: [&str; 3] = ["alloc", "rc", "Rc"];
pub const RC_PTR_EQ: [&str; 4] = ["alloc", "rc", "Rc", "ptr_eq"];
pub const RECEIVER: [&str; 4] = ["std", "sync", "mpsc", "Receiver"];
pub const REGEX: [&str; 3] = ["regex", "re_unicode", "Regex"];
pub const REGEX_BUILDER_NEW: [&str; 5] = ["regex", "re_builder", "unicode", "RegexBuilder", "new"];
pub const REGEX_BYTES_BUILDER_NEW: [&str; 5] = ["regex", "re_builder", "bytes", "RegexBuilder", "new"];
pub const REGEX_BYTES_NEW: [&str; 4] = ["regex", "re_bytes", "Regex", "new"];
pub const REGEX_BYTES_SET_NEW: [&str; 5] = ["regex", "re_set", "bytes", "RegexSet", "new"];
pub const REGEX_NEW: [&str; 4] = ["regex", "re_unicode", "Regex", "new"];
pub const REGEX_SET_NEW: [&str; 5] = ["regex", "re_set", "unicode", "RegexSet", "new"];
pub const REPEAT: [&str; 3] = ["core", "iter", "repeat"];
pub const RESULT: [&str; 3] = ["core", "result", "Result"];
pub const RESULT_ERR: [&str; 4] = ["core", "result", "Result", "Err"];
pub const RESULT_OK: [&str; 4] = ["core", "result", "Result", "Ok"];
pub const RWLOCK_READ_GUARD: [&str; 4] = ["std", "sync", "rwlock", "RwLockReadGuard"];
pub const RWLOCK_WRITE_GUARD: [&str; 4] = ["std", "sync", "rwlock", "RwLockWriteGuard"];
pub const SERDE_DE_VISITOR: [&str; 3] = ["serde", "de", "Visitor"];
pub const SLICE_INTO_VEC: [&str; 4] = ["alloc", "slice", "<impl [T]>", "into_vec"];
pub const SLICE_ITER: [&str; 3] = ["core", "slice", "Iter"];
pub const STDERR: [&str; 4] = ["std", "io", "stdio", "stderr"];
pub const STDOUT: [&str; 4] = ["std", "io", "stdio", "stdout"];
pub const STD_CONVERT_IDENTITY: [&str; 3] = ["std", "convert", "identity"];
pub const STD_MEM_TRANSMUTE: [&str; 3] = ["std", "mem", "transmute"];
pub const STD_PTR_NULL: [&str; 3] = ["std", "ptr", "null"];
pub const STRING: [&str; 3] = ["alloc", "string", "String"];
pub const STRING_AS_MUT_STR: [&str; 4] = ["alloc", "string", "String", "as_mut_str"];
pub const STRING_AS_STR: [&str; 4] = ["alloc", "string", "String", "as_str"];
pub const SYNTAX_CONTEXT: [&str; 3] = ["rustc_span", "hygiene", "SyntaxContext"];
pub const TO_OWNED: [&str; 3] = ["alloc", "borrow", "ToOwned"];
pub const TO_OWNED_METHOD: [&str; 4] = ["alloc", "borrow", "ToOwned", "to_owned"];
pub const TO_STRING: [&str; 3] = ["alloc", "string", "ToString"];
pub const TO_STRING_METHOD: [&str; 4] = ["alloc", "string", "ToString", "to_string"];
pub const TRANSMUTE: [&str; 4] = ["core", "intrinsics", "", "transmute"];
pub const TRY_FROM_ERROR: [&str; 4] = ["std", "ops", "Try", "from_error"];
pub const TRY_INTO_RESULT: [&str; 4] = ["std", "ops", "Try", "into_result"];
pub const VEC: [&str; 3] = ["alloc", "vec", "Vec"];
pub const VEC_AS_MUT_SLICE: [&str; 4] = ["alloc", "vec", "Vec", "as_mut_slice"];
pub const VEC_AS_SLICE: [&str; 4] = ["alloc", "vec", "Vec", "as_slice"];
pub const VEC_DEQUE: [&str; 4] = ["alloc", "collections", "vec_deque", "VecDeque"];
pub const VEC_FROM_ELEM: [&str; 3] = ["alloc", "vec", "from_elem"];
pub const VEC_NEW: [&str; 4] = ["alloc", "vec", "Vec", "new"];
pub const WEAK_ARC: [&str; 3] = ["alloc", "sync", "Weak"];
pub const WEAK_RC: [&str; 3] = ["alloc", "rc", "Weak"];
<commit_msg>Say that diagnostic items are preferred over paths<commit_after>//! This module contains paths to types and functions Clippy needs to know
//! about.
//!
//! Whenever possible, please consider diagnostic items over hardcoded paths.
//! See <https://github.com/rust-lang/rust-clippy/issues/5393> for more information.

pub const ANY_TRAIT: [&str; 3] = ["std", "any", "Any"];
pub const ARC: [&str; 3] = ["alloc", "sync", "Arc"];
pub const ARC_PTR_EQ: [&str; 4] = ["alloc", "sync", "Arc", "ptr_eq"];
pub const ASMUT_TRAIT: [&str; 3] = ["core", "convert", "AsMut"];
pub const ASREF_TRAIT: [&str; 3] = ["core", "convert", "AsRef"];
pub const BEGIN_PANIC: [&str; 3] = ["std", "panicking", "begin_panic"];
pub const BEGIN_PANIC_FMT: [&str; 3] = ["std", "panicking", "begin_panic_fmt"];
pub const BINARY_HEAP: [&str; 4] = ["alloc", "collections", "binary_heap", "BinaryHeap"];
pub const BORROW_TRAIT: [&str; 3] = ["core", "borrow", "Borrow"];
pub const BOX: [&str; 3] = ["alloc", "boxed", "Box"];
pub const BTREEMAP: [&str; 5] = ["alloc", "collections", "btree", "map", "BTreeMap"];
pub const BTREEMAP_ENTRY: [&str; 5] = ["alloc", "collections", "btree", "map", "Entry"];
pub const BTREESET: [&str; 5] = ["alloc", "collections", "btree", "set", "BTreeSet"];
pub const CLONE_TRAIT: [&str; 3] = ["core", "clone", "Clone"];
pub const CLONE_TRAIT_METHOD: [&str; 4] = ["core", "clone", "Clone", "clone"];
pub const CMP_MAX: [&str; 3] = ["core", "cmp", "max"];
pub const CMP_MIN: [&str; 3] = ["core", "cmp", "min"];
pub const COW: [&str; 3] = ["alloc", "borrow", "Cow"];
pub const CSTRING: [&str; 4] = ["std", "ffi", "c_str", "CString"];
pub const CSTRING_AS_C_STR: [&str; 5] = ["std", "ffi", "c_str", "CString", "as_c_str"];
pub const DEFAULT_TRAIT: [&str; 3] = ["core", "default", "Default"];
pub const DEFAULT_TRAIT_METHOD: [&str; 4] = ["core", "default", "Default", "default"];
pub const DEREF_MUT_TRAIT_METHOD: [&str; 5] = ["core", "ops", "deref", "DerefMut", "deref_mut"];
pub const DEREF_TRAIT_METHOD: [&str; 5] = ["core", "ops", "deref", "Deref", "deref"];
pub const DISPLAY_FMT_METHOD: [&str; 4] = ["core", "fmt", "Display", "fmt"];
pub const DISPLAY_TRAIT: [&str; 3] = ["core", "fmt", "Display"];
pub const DOUBLE_ENDED_ITERATOR: [&str; 4] = ["core", "iter", "traits", "DoubleEndedIterator"];
pub const DROP: [&str; 3] = ["core", "mem", "drop"];
pub const DROP_TRAIT: [&str; 4] = ["core", "ops", "drop", "Drop"];
pub const DURATION: [&str; 3] = ["core", "time", "Duration"];
pub const EARLY_CONTEXT: [&str; 4] = ["rustc", "lint", "context", "EarlyContext"];
pub const EXIT: [&str; 3] = ["std", "process", "exit"];
pub const FILE: [&str; 3] = ["std", "fs", "File"];
pub const FILE_TYPE: [&str; 3] = ["std", "fs", "FileType"];
pub const FMT_ARGUMENTS_NEW_V1: [&str; 4] = ["core", "fmt", "Arguments", "new_v1"];
pub const FMT_ARGUMENTS_NEW_V1_FORMATTED: [&str; 4] = ["core", "fmt", "Arguments", "new_v1_formatted"];
pub const FMT_ARGUMENTV1_NEW: [&str; 4] = ["core", "fmt", "ArgumentV1", "new"];
pub const FROM_FROM: [&str; 4] = ["core", "convert", "From", "from"];
pub const FROM_TRAIT: [&str; 3] = ["core", "convert", "From"];
pub const HASH: [&str; 2] = ["hash", "Hash"];
pub const HASHMAP: [&str; 5] = ["std", "collections", "hash", "map", "HashMap"];
pub const HASHMAP_ENTRY: [&str; 5] = ["std", "collections", "hash", "map", "Entry"];
pub const HASHSET: [&str; 5] = ["std", "collections", "hash", "set", "HashSet"];
pub const INDEX: [&str; 3] = ["core", "ops", "Index"];
pub const INDEX_MUT: [&str; 3] = ["core", "ops", "IndexMut"];
pub const INTO: [&str; 3] = ["core", "convert", "Into"];
pub const INTO_ITERATOR: [&str; 5] = ["core", "iter", "traits", "collect", "IntoIterator"];
pub const IO_READ: [&str; 3] = ["std", "io", "Read"];
pub const IO_WRITE: [&str; 3] = ["std", "io", "Write"];
pub const ITERATOR: [&str; 5] = ["core", "iter", "traits", "iterator", "Iterator"];
pub const LATE_CONTEXT: [&str; 4] = ["rustc", "lint", "context", "LateContext"];
pub const LINKED_LIST: [&str; 4] = ["alloc", "collections", "linked_list", "LinkedList"];
pub const LINT: [&str; 3] = ["rustc_session", "lint", "Lint"];
pub const MEM_DISCRIMINANT: [&str; 3] = ["core", "mem", "discriminant"];
pub const MEM_FORGET: [&str; 3] = ["core", "mem", "forget"];
pub const MEM_MAYBEUNINIT: [&str; 4] = ["core", "mem", "maybe_uninit", "MaybeUninit"];
pub const MEM_MAYBEUNINIT_UNINIT: [&str; 5] = ["core", "mem", "maybe_uninit", "MaybeUninit", "uninit"];
pub const MEM_REPLACE: [&str; 3] = ["core", "mem", "replace"];
pub const MEM_UNINITIALIZED: [&str; 3] = ["core", "mem", "uninitialized"];
pub const MEM_ZEROED: [&str; 3] = ["core", "mem", "zeroed"];
pub const MUTEX: [&str; 4] = ["std", "sync", "mutex", "Mutex"];
pub const MUTEX_GUARD: [&str; 4] = ["std", "sync", "mutex", "MutexGuard"];
pub const OPEN_OPTIONS: [&str; 3] = ["std", "fs", "OpenOptions"];
pub const OPS_MODULE: [&str; 2] = ["core", "ops"];
pub const OPTION: [&str; 3] = ["core", "option", "Option"];
pub const OPTION_NONE: [&str; 4] = ["core", "option", "Option", "None"];
pub const OPTION_SOME: [&str; 4] = ["core", "option", "Option", "Some"];
pub const ORD: [&str; 3] = ["core", "cmp", "Ord"];
pub const OS_STRING: [&str; 4] = ["std", "ffi", "os_str", "OsString"];
pub const OS_STRING_AS_OS_STR: [&str; 5] = ["std", "ffi", "os_str", "OsString", "as_os_str"];
pub const OS_STR_TO_OS_STRING: [&str; 5] = ["std", "ffi", "os_str", "OsStr", "to_os_string"];
pub const PATH: [&str; 3] = ["std", "path", "Path"];
pub const PATH_BUF: [&str; 3] = ["std", "path", "PathBuf"];
pub const PATH_BUF_AS_PATH: [&str; 4] = ["std", "path", "PathBuf", "as_path"];
pub const PATH_TO_PATH_BUF: [&str; 4] = ["std", "path", "Path", "to_path_buf"];
pub const PTR_EQ: [&str; 3] = ["core", "ptr", "eq"];
pub const PTR_NULL: [&str; 2] = ["ptr", "null"];
pub const PTR_NULL_MUT: [&str; 2] = ["ptr", "null_mut"];
pub const RANGE: [&str; 3] = ["core", "ops", "Range"];
pub const RANGE_ARGUMENT_TRAIT: [&str; 3] = ["core", "ops", "RangeBounds"];
pub const RANGE_FROM: [&str; 3] = ["core", "ops", "RangeFrom"];
pub const RANGE_FROM_STD: [&str; 3] = ["std", "ops", "RangeFrom"];
pub const RANGE_FULL: [&str; 3] = ["core", "ops", "RangeFull"];
pub const RANGE_FULL_STD: [&str; 3] = ["std", "ops", "RangeFull"];
pub const RANGE_INCLUSIVE_NEW: [&str; 4] = ["core", "ops", "RangeInclusive", "new"];
pub const RANGE_INCLUSIVE_STD_NEW: [&str; 4] = ["std", "ops", "RangeInclusive", "new"];
pub const RANGE_STD: [&str; 3] = ["std", "ops", "Range"];
pub const RANGE_TO: [&str; 3] = ["core", "ops", "RangeTo"];
pub const RANGE_TO_INCLUSIVE: [&str; 3] = ["core", "ops", "RangeToInclusive"];
pub const RANGE_TO_INCLUSIVE_STD: [&str; 3] = ["std", "ops", "RangeToInclusive"];
pub const RANGE_TO_STD: [&str; 3] = ["std", "ops", "RangeTo"];
pub const RC: [&str; 3] = ["alloc", "rc", "Rc"];
pub const RC_PTR_EQ: [&str; 4] = ["alloc", "rc", "Rc", "ptr_eq"];
pub const RECEIVER: [&str; 4] = ["std", "sync", "mpsc", "Receiver"];
pub const REGEX: [&str; 3] = ["regex", "re_unicode", "Regex"];
pub const REGEX_BUILDER_NEW: [&str; 5] = ["regex", "re_builder", "unicode", "RegexBuilder", "new"];
pub const REGEX_BYTES_BUILDER_NEW: [&str; 5] = ["regex", "re_builder", "bytes", "RegexBuilder", "new"];
pub const REGEX_BYTES_NEW: [&str; 4] = ["regex", "re_bytes", "Regex", "new"];
pub const REGEX_BYTES_SET_NEW: [&str; 5] = ["regex", "re_set", "bytes", "RegexSet", "new"];
pub const REGEX_NEW: [&str; 4] = ["regex", "re_unicode", "Regex", "new"];
pub const REGEX_SET_NEW: [&str; 5] = ["regex", "re_set", "unicode", "RegexSet", "new"];
pub const REPEAT: [&str; 3] = ["core", "iter", "repeat"];
pub const RESULT: [&str; 3] = ["core", "result", "Result"];
pub const RESULT_ERR: [&str; 4] = ["core", "result", "Result", "Err"];
pub const RESULT_OK: [&str; 4] = ["core", "result", "Result", "Ok"];
pub const RWLOCK_READ_GUARD: [&str; 4] = ["std", "sync", "rwlock", "RwLockReadGuard"];
pub const RWLOCK_WRITE_GUARD: [&str; 4] = ["std", "sync", "rwlock", "RwLockWriteGuard"];
pub const SERDE_DE_VISITOR: [&str; 3] = ["serde", "de", "Visitor"];
pub const SLICE_INTO_VEC: [&str; 4] = ["alloc", "slice", "<impl [T]>", "into_vec"];
pub const SLICE_ITER: [&str; 3] = ["core", "slice", "Iter"];
pub const STDERR: [&str; 4] = ["std", "io", "stdio", "stderr"];
pub const STDOUT: [&str; 4] = ["std", "io", "stdio", "stdout"];
pub const STD_CONVERT_IDENTITY: [&str; 3] = ["std", "convert", "identity"];
pub const STD_MEM_TRANSMUTE: [&str; 3] = ["std", "mem", "transmute"];
pub const STD_PTR_NULL: [&str; 3] = ["std", "ptr", "null"];
pub const STRING: [&str; 3] = ["alloc", "string", "String"];
pub const STRING_AS_MUT_STR: [&str; 4] = ["alloc", "string", "String", "as_mut_str"];
pub const STRING_AS_STR: [&str; 4] = ["alloc", "string", "String", "as_str"];
pub const SYNTAX_CONTEXT: [&str; 3] = ["rustc_span", "hygiene", "SyntaxContext"];
pub const TO_OWNED: [&str; 3] = ["alloc", "borrow", "ToOwned"];
pub const TO_OWNED_METHOD: [&str; 4] = ["alloc", "borrow", "ToOwned", "to_owned"];
pub const TO_STRING: [&str; 3] = ["alloc", "string", "ToString"];
pub const TO_STRING_METHOD: [&str; 4] = ["alloc", "string", "ToString", "to_string"];
pub const TRANSMUTE: [&str; 4] = ["core", "intrinsics", "", "transmute"];
pub const TRY_FROM_ERROR: [&str; 4] = ["std", "ops", "Try", "from_error"];
pub const TRY_INTO_RESULT: [&str; 4] = ["std", "ops", "Try", "into_result"];
pub const VEC: [&str; 3] = ["alloc", "vec", "Vec"];
pub const VEC_AS_MUT_SLICE: [&str; 4] = ["alloc", "vec", "Vec", "as_mut_slice"];
pub const VEC_AS_SLICE: [&str; 4] = ["alloc", "vec", "Vec", "as_slice"];
pub const VEC_DEQUE: [&str; 4] = ["alloc", "collections", "vec_deque", "VecDeque"];
pub const VEC_FROM_ELEM: [&str; 3] = ["alloc", "vec", "from_elem"];
pub const VEC_NEW: [&str; 4] = ["alloc", "vec", "Vec", "new"];
pub const WEAK_ARC: [&str; 3] = ["alloc", "sync", "Weak"];
pub const WEAK_RC: [&str; 3] = ["alloc", "rc", "Weak"];
<|endoftext|>


<commit_before>use crate::schema::*;
use diesel::*;

#[test]
// This test is a shim for a feature which is not sufficiently implemented. It
// has been added as we have a user who needs a reasonable workaround, but this
// functionality will change and this test is allowed to change post-1.0
fn group_by_generates_group_by_sql() {
    let source = users::table
        .group_by(users::name)
        .select(users::name)
        .filter(users::hair_color.is_null());
    let mut expected_sql = "SELECT `users`.`name` FROM `users` \
                            WHERE (`users`.`hair_color` IS NULL) \
                            GROUP BY `users`.`name` \
                            -- binds: []"
        .to_string();
    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(
        expected_sql,
        debug_query::<TestBackend, _>(&source).to_string()
    );
    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn group_by_mixed_aggregate_column_and_aggregate_function() {
    use diesel::dsl::max;
    let source = users::table
        .group_by(users::name)
        .select((max(users::id), users::name))
        .filter(users::hair_color.is_null());
    let mut expected_sql = "SELECT max(`users`.`id`), `users`.`name` FROM `users` \
                            WHERE (`users`.`hair_color` IS NULL) \
                            GROUP BY `users`.`name` \
                            -- binds: []"
        .to_string();
    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(
        expected_sql,
        debug_query::<TestBackend, _>(&source).to_string()
    );

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn boxed_queries_have_group_by_method() {
    let source = users::table
        .group_by(users::name)
        .select(users::name)
        .into_boxed::<TestBackend>()
        .filter(users::hair_color.is_null());
    let mut expected_sql = "SELECT `users`.`name` FROM `users` \
                            WHERE (`users`.`hair_color` IS NULL) \
                            GROUP BY `users`.`name` \
                            -- binds: []"
        .to_string();
    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn check_group_by_primary_key_allows_other_columns_in_select_clause() {
    let source = users::table
        .group_by(users::id)
        .select(users::name)
        .filter(users::hair_color.is_null());

    let mut expected_sql = "SELECT `users`.`name` FROM `users` \
                            WHERE (`users`.`hair_color` IS NULL) \
                            GROUP BY `users`.`id` \
                            -- binds: []"
        .to_string();

    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn check_group_by_multiple_columns_in_group_by_clause_single_select() {
    let source = users::table
        .group_by((users::name, users::hair_color))
        .select(users::name)
        .filter(users::id.nullable().is_null());

    let mut expected_sql = "SELECT `users`.`name` \
                            FROM `users` WHERE (`users`.`id` IS NULL) \
                            GROUP BY `users`.`name`, `users`.`hair_color` \
                            -- binds: []"
        .to_string();

    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn check_group_by_multiple_columns_in_group_by_clause_complex_select() {
    let source = users::table
        .group_by((users::name, users::hair_color))
        .select((users::name, users::hair_color, diesel::dsl::max(users::id)))
        .filter(users::id.nullable().is_null());

    let mut expected_sql = "SELECT `users`.`name`, `users`.`hair_color`, max(`users`.`id`) \
                            FROM `users` WHERE (`users`.`id` IS NULL) \
                            GROUP BY `users`.`name`, `users`.`hair_color` \
                            -- binds: []"
        .to_string();

    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

diesel::allow_columns_to_appear_in_same_group_by_clause!(
    posts::id,
    users::id,
    posts::title,
    users::name
);

#[test]
fn check_group_by_multiple_tables() {
    let source = users::table
        .inner_join(posts::table.inner_join(comments::table))
        .group_by((users::id, posts::id))
        .select((users::name, posts::title, diesel::dsl::count(comments::id)))
        .filter(comments::text.nullable().is_null());

    let mut expected_sql = "SELECT `users`.`name`, `posts`.`title`, count(`comments`.`id`) \
                            FROM (`users` \
                            INNER JOIN (`posts` \
                            INNER JOIN `comments` ON (`comments`.`post_id` = `posts`.`id`)) \
                            ON (`posts`.`user_id` = `users`.`id`)) \
                            WHERE (`comments`.`text` IS NULL) \
                            GROUP BY `users`.`id`, `posts`.`id` \
                            -- binds: []"
        .to_string();

    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}
<commit_msg>Added tests for group by within a subselect<commit_after>use crate::schema::*;
use diesel::*;

#[test]
// This test is a shim for a feature which is not sufficiently implemented. It
// has been added as we have a user who needs a reasonable workaround, but this
// functionality will change and this test is allowed to change post-1.0
fn group_by_generates_group_by_sql() {
    let source = users::table
        .group_by(users::name)
        .select(users::name)
        .filter(users::hair_color.is_null());
    let mut expected_sql = "SELECT `users`.`name` FROM `users` \
                            WHERE (`users`.`hair_color` IS NULL) \
                            GROUP BY `users`.`name` \
                            -- binds: []"
        .to_string();
    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(
        expected_sql,
        debug_query::<TestBackend, _>(&source).to_string()
    );
    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn group_by_mixed_aggregate_column_and_aggregate_function() {
    use diesel::dsl::max;
    let source = users::table
        .group_by(users::name)
        .select((max(users::id), users::name))
        .filter(users::hair_color.is_null());
    let mut expected_sql = "SELECT max(`users`.`id`), `users`.`name` FROM `users` \
                            WHERE (`users`.`hair_color` IS NULL) \
                            GROUP BY `users`.`name` \
                            -- binds: []"
        .to_string();
    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(
        expected_sql,
        debug_query::<TestBackend, _>(&source).to_string()
    );

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn boxed_queries_have_group_by_method() {
    let source = users::table
        .group_by(users::name)
        .select(users::name)
        .into_boxed::<TestBackend>()
        .filter(users::hair_color.is_null());
    let mut expected_sql = "SELECT `users`.`name` FROM `users` \
                            WHERE (`users`.`hair_color` IS NULL) \
                            GROUP BY `users`.`name` \
                            -- binds: []"
        .to_string();
    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn check_group_by_primary_key_allows_other_columns_in_select_clause() {
    let source = users::table
        .group_by(users::id)
        .select(users::name)
        .filter(users::hair_color.is_null());

    let mut expected_sql = "SELECT `users`.`name` FROM `users` \
                            WHERE (`users`.`hair_color` IS NULL) \
                            GROUP BY `users`.`id` \
                            -- binds: []"
        .to_string();

    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn check_group_by_multiple_columns_in_group_by_clause_single_select() {
    let source = users::table
        .group_by((users::name, users::hair_color))
        .select(users::name)
        .filter(users::id.nullable().is_null());

    let mut expected_sql = "SELECT `users`.`name` \
                            FROM `users` WHERE (`users`.`id` IS NULL) \
                            GROUP BY `users`.`name`, `users`.`hair_color` \
                            -- binds: []"
        .to_string();

    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn check_group_by_multiple_columns_in_group_by_clause_complex_select() {
    let source = users::table
        .group_by((users::name, users::hair_color))
        .select((users::name, users::hair_color, diesel::dsl::max(users::id)))
        .filter(users::id.nullable().is_null());

    let mut expected_sql = "SELECT `users`.`name`, `users`.`hair_color`, max(`users`.`id`) \
                            FROM `users` WHERE (`users`.`id` IS NULL) \
                            GROUP BY `users`.`name`, `users`.`hair_color` \
                            -- binds: []"
        .to_string();

    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

diesel::allow_columns_to_appear_in_same_group_by_clause!(
    posts::id,
    users::id,
    posts::title,
    users::name
);

#[test]
fn check_group_by_multiple_tables() {
    let source = users::table
        .inner_join(posts::table.inner_join(comments::table))
        .group_by((users::id, posts::id))
        .select((users::name, posts::title, diesel::dsl::count(comments::id)))
        .filter(comments::text.nullable().is_null());

    let mut expected_sql = "SELECT `users`.`name`, `posts`.`title`, count(`comments`.`id`) \
                            FROM (`users` \
                            INNER JOIN (`posts` \
                            INNER JOIN `comments` ON (`comments`.`post_id` = `posts`.`id`)) \
                            ON (`posts`.`user_id` = `users`.`id`)) \
                            WHERE (`comments`.`text` IS NULL) \
                            GROUP BY `users`.`id`, `posts`.`id` \
                            -- binds: []"
        .to_string();

    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn check_filter_with_group_by_subselect() {
    let subselect = posts::table
        .group_by(posts::user_id)
        .select(diesel::dsl::min(posts::id));

    // get one example post for each user (for users with a post)
    let source = posts::table
        .filter(posts::id.nullable().eq_any(subselect))
        .select((posts::user_id, posts::title));

    let mut expected_sql = "SELECT `posts`.`user_id`, `posts`.`title` \
                            FROM `posts` \
                            WHERE (\
                              `posts`.`id` IN (\
                              SELECT min(`posts`.`id`) \
                              FROM `posts` \
                              GROUP BY `posts`.`user_id`)) \
                            -- binds: []"
        .to_string();

    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}

#[test]
fn check_filter_with_boxed_group_by_subselect() {
    // this query is identical to `check_filter_with_group_by_subselect`, but the subselect also calls `into_boxed`
    let subselect = posts::table
        .group_by(posts::user_id)
        .select(diesel::dsl::min(posts::id))
        .into_boxed();

    // get one example post for each user (for users with a post)
    let source = posts::table
        .filter(posts::id.nullable().eq_any(subselect))
        .select((posts::user_id, posts::title));

    let mut expected_sql = "SELECT `posts`.`user_id`, `posts`.`title` \
                            FROM `posts` \
                            WHERE (\
                              `posts`.`id` IN (\
                              SELECT min(`posts`.`id`) \
                              FROM `posts` \
                              GROUP BY `posts`.`user_id`)) \
                            -- binds: []"
        .to_string();

    if cfg!(feature = "postgres") {
        expected_sql = expected_sql.replace('`', "\"");
    }

    assert_eq!(expected_sql, debug_query(&source).to_string());

    let conn = &mut connection();

    assert!(source.execute(conn).is_ok());
}
<|endoftext|>