<commit_before>//! This example shows an example of how to parse an escaped string. The
//! rules for the string are similar to JSON and rust. A string is:
//!
//! - Enclosed by double quotes
//! - Can contain any raw unescaped code point besides \ and "
//! - Matches the following escape sequences: \b, \f, \n, \r, \t, \", \\, \/
//! - Matches code points like Rust: \u{XXXX}, where XXXX can be up to 6
//!   hex characters
//! - an escape followed by whitespace consumes all whitespace between the
//!   escape and the next non-whitespace character

#![cfg(feature = "alloc")]

extern crate jemallocator;
extern crate nom;

#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

use std::convert::TryInto;

use nom::branch::alt;
use nom::bytes::streaming::{is_not, take_while_m_n};
use nom::character::streaming::{char, multispace1};
use nom::combinator::{map, map_res, value, verify};
use nom::error::ParseError;
use nom::multi::fold_many0;
use nom::sequence::{delimited, preceded};
use nom::IResult;

// parser combinators are constructed from the bottom up:
// first we write parsers for the smallest elements (escaped characters),
// then combine them into larger parsers.

/// Parse a unicode sequence, of the form u{XXXX}, where XXXX is 1 to 6
/// hexadecimal numerals. We will combine this later with parse_escaped_char
/// to parse sequences like \u{00AC}.
fn parse_unicode<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, char, E> {
  // `take_while_m_n` parses between `m` and `n` bytes (inclusive) that match
  // a predicate.
  let parse_hex = take_while_m_n(1, 6, |c: char| c.is_ascii_hexdigit());

  // `preceeded` takes a prefix parser, and if it succeeds, returns the result
  // of the body parser. In this case, it parses u{XXXX}.
  let parse_delimited_hex = preceded(
    char('u'),
    // `delimited` is like `preceded`, but it parses both a prefix and a suffix.
    // It returns the result of the middle parser. In this case, it parses
    // {XXXX}, where XXXX is 1 to 6 hex numerals.
    delimited(char('{'), parse_hex, char('}')),
  );

  // `map_res` takes the result of a parser and applies a function that returns
  // a Result. In this case we take the hex bytes from parse_hex and attempt to
  // convert them to a u32.
  let parse_u32 = map_res(parse_delimited_hex, |hex| u32::from_str_radix(hex, 16));

  // Finally, we attempt to convert the u32 to a char. Not all u32 values are
  // valid code points, so we have to use map_res again.
  map_res(parse_u32, |value| value.try_into())(input)
}

/// Parse an escaped character: \n, \t, \r, \u{00AC}, etc.
fn parse_escaped_char<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, char, E> {
  preceded(
    char('\\'),
    // `alt` tries each each parser in sequence, returning the result of
    // the first successful match
    alt((
      parse_unicode,
      // The `value` parser returns a fixed value (the first argument) if its
      // parser (the second argument) succeeds. In these cases, it looks for
      // the marker characters (n, r, t, etc) and returns the matching
      // character (\n, \r, \t, etc).
      value('\n', char('n')),
      value('\r', char('r')),
      value('\t', char('t')),
      value('\u{08}', char('b')),
      value('\u{0C}', char('f')),
      value('\\', char('\\')),
      value('/', char('/')),
      value('"', char('"')),
    )),
  )(input)
}

/// Parse a backslash, followed by any amount of whitespace. This is used later
/// to discard any escaped whitespace.
fn parse_escaped_whitespace<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, &'a str, E> {
  preceded(char('\\'), multispace1)(input)
}

/// Parse a non-empty block of text that doesn't include \ or "
fn parse_literal<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, &'a str, E> {
  // `is_not` parses 0 or more characters that aren't one of the given characters.
  // `verify` runs a parser, then runs a verification function on the output of
  // the parser. The verification function accepts out output only if it
  // returns true. In this case, we want to ensure that the output of is_not
  // is non-empty.
  verify(is_not("\"\\"), |s: &str| !s.is_empty())(input)
}

/// A string fragment contains a fragment of a string being parsed: either
/// a non-empty Literal (a series of non-escaped characters), a single
/// parsed escaped character, or a block of escaped whitespace.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StringFragment<'a> {
  Literal(&'a str),
  EscapedChar(char),
  EscapedWS,
}

/// Combine parse_literal, parse_escaped_whitespace, and parse_escaped_char
/// into a StringFragment.
fn parse_fragment<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, StringFragment<'a>, E> {
  alt((
    // The `map` combinator runs a parser, then applies a function to the output
    // of that parser.
    map(parse_literal, StringFragment::Literal),
    map(parse_escaped_char, StringFragment::EscapedChar),
    value(StringFragment::EscapedWS, parse_escaped_whitespace),
  ))(input)
}

/// Parse a string. Use a loop of parse_fragment and push all of the fragments
/// into an output string.
fn parse_string<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, String, E> {
  // fold_many0 is the equivalent of iterator::fold. It runs a parser in a loop,
  // and for each output value, calls a folding function on each output value.
  let build_string = fold_many0(
    // Our parser function– parses a single string fragment
    parse_fragment,
    // Our init value, an empty string
    String::new(),
    // Our folding function. For each fragment, append the fragment to the
    // string.
    |mut string, fragment| {
      match fragment {
        StringFragment::Literal(s) => string.push_str(s),
        StringFragment::EscapedChar(c) => string.push(c),
        StringFragment::EscapedWS => {}
      }
      string
    },
  );

  delimited(char('"'), build_string, char('"'))(input)
}

fn main() {
  let data = "\"abc\"";
  println!("EXAMPLE 1:\nParsing a simple input string: {}", data);
  let result = parse_string::<()>(data);
  assert_eq!(result, Ok(("", String::from("abc"))));
  println!("Result: {}\n\n", result.unwrap().1);

  let data = "\"tab:\\tafter tab, newline:\\nnew line, quote: \\\", emoji: \\u{1F602}, newline:\\nescaped whitespace: \\    abc\"";
  println!(
    "EXAMPLE 2:\nParsing a string with escape sequences, newline literal, and escaped whitespace:\n\n{}\n",
    data
  );
  let result = parse_string::<()>(data);
  assert_eq!(
    result,
    Ok((
      "",
      String::from("tab:\tafter tab, newline:\nnew line, quote: \", emoji: 😂, newline:\nescaped whitespace: abc")
    ))
  );
  println!("Result:\n\n{}", result.unwrap().1);
}
<commit_msg>Fixed broken build on rust 1.31; extra comments<commit_after>//! This example shows an example of how to parse an escaped string. The
//! rules for the string are similar to JSON and rust. A string is:
//!
//! - Enclosed by double quotes
//! - Can contain any raw unescaped code point besides \ and "
//! - Matches the following escape sequences: \b, \f, \n, \r, \t, \", \\, \/
//! - Matches code points like Rust: \u{XXXX}, where XXXX can be up to 6
//!   hex characters
//! - an escape followed by whitespace consumes all whitespace between the
//!   escape and the next non-whitespace character

#![cfg(feature = "alloc")]

extern crate jemallocator;
extern crate nom;

#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;

use nom::branch::alt;
use nom::bytes::streaming::{is_not, take_while_m_n};
use nom::character::streaming::{char, multispace1};
use nom::combinator::{map, map_opt, map_res, value, verify};
use nom::error::ParseError;
use nom::multi::fold_many0;
use nom::sequence::{delimited, preceded};
use nom::IResult;

// parser combinators are constructed from the bottom up:
// first we write parsers for the smallest elements (escaped characters),
// then combine them into larger parsers.

/// Parse a unicode sequence, of the form u{XXXX}, where XXXX is 1 to 6
/// hexadecimal numerals. We will combine this later with parse_escaped_char
/// to parse sequences like \u{00AC}.
fn parse_unicode<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, char, E> {
  // `take_while_m_n` parses between `m` and `n` bytes (inclusive) that match
  // a predicate. `parse_hex` here parses between 1 and 6 hexadecimal numerals.
  let parse_hex = take_while_m_n(1, 6, |c: char| c.is_ascii_hexdigit());

  // `preceeded` takes a prefix parser, and if it succeeds, returns the result
  // of the body parser. In this case, it parses u{XXXX}.
  let parse_delimited_hex = preceded(
    char('u'),
    // `delimited` is like `preceded`, but it parses both a prefix and a suffix.
    // It returns the result of the middle parser. In this case, it parses
    // {XXXX}, where XXXX is 1 to 6 hex numerals, and returns XXXX
    delimited(char('{'), parse_hex, char('}')),
  );

  // `map_res` takes the result of a parser and applies a function that returns
  // a Result. In this case we take the hex bytes from parse_hex and attempt to
  // convert them to a u32.
  let parse_u32 = map_res(parse_delimited_hex, move |hex| u32::from_str_radix(hex, 16));

  // map_opt is like map_res, but it takes an Option instead of a Result. If
  // the function returns None, map_opt returns an error. In this case, because
  // not all u32 values are valid unicode code points, we have to fallibly
  // convert to char with from_u32.
  map_opt(parse_u32, |value| std::char::from_u32(value))(input)
}

/// Parse an escaped character: \n, \t, \r, \u{00AC}, etc.
fn parse_escaped_char<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, char, E> {
  preceded(
    char('\\'),
    // `alt` tries each each parser in sequence, returning the result of
    // the first successful match
    alt((
      parse_unicode,
      // The `value` parser returns a fixed value (the first argument) if its
      // parser (the second argument) succeeds. In these cases, it looks for
      // the marker characters (n, r, t, etc) and returns the matching
      // character (\n, \r, \t, etc).
      value('\n', char('n')),
      value('\r', char('r')),
      value('\t', char('t')),
      value('\u{08}', char('b')),
      value('\u{0C}', char('f')),
      value('\\', char('\\')),
      value('/', char('/')),
      value('"', char('"')),
    )),
  )(input)
}

/// Parse a backslash, followed by any amount of whitespace. This is used later
/// to discard any escaped whitespace.
fn parse_escaped_whitespace<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, &'a str, E> {
  preceded(char('\\'), multispace1)(input)
}

/// Parse a non-empty block of text that doesn't include \ or "
fn parse_literal<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, &'a str, E> {
  // `is_not` parses a string of 0 or more characters that aren't one of the
  // given characters.
  let not_quote_slash = is_not("\"\\");

  // `verify` runs a parser, then runs a verification function on the output of
  // the parser. The verification function accepts out output only if it
  // returns true. In this case, we want to ensure that the output of is_not
  // is non-empty.
  verify(not_quote_slash, |s: &str| !s.is_empty())(input)
}

/// A string fragment contains a fragment of a string being parsed: either
/// a non-empty Literal (a series of non-escaped characters), a single
/// parsed escaped character, or a block of escaped whitespace.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StringFragment<'a> {
  Literal(&'a str),
  EscapedChar(char),
  EscapedWS,
}

/// Combine parse_literal, parse_escaped_whitespace, and parse_escaped_char
/// into a StringFragment.
fn parse_fragment<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, StringFragment<'a>, E> {
  alt((
    // The `map` combinator runs a parser, then applies a function to the output
    // of that parser.
    map(parse_literal, StringFragment::Literal),
    map(parse_escaped_char, StringFragment::EscapedChar),
    value(StringFragment::EscapedWS, parse_escaped_whitespace),
  ))(input)
}

/// Parse a string. Use a loop of parse_fragment and push all of the fragments
/// into an output string.
fn parse_string<'a, E: ParseError<&'a str>>(input: &'a str) -> IResult<&'a str, String, E> {
  // fold_many0 is the equivalent of iterator::fold. It runs a parser in a loop,
  // and for each output value, calls a folding function on each output value.
  let build_string = fold_many0(
    // Our parser function– parses a single string fragment
    parse_fragment,
    // Our init value, an empty string
    String::new(),
    // Our folding function. For each fragment, append the fragment to the
    // string.
    |mut string, fragment| {
      match fragment {
        StringFragment::Literal(s) => string.push_str(s),
        StringFragment::EscapedChar(c) => string.push(c),
        StringFragment::EscapedWS => {}
      }
      string
    },
  );

  delimited(char('"'), build_string, char('"'))(input)
}

fn main() {
  let data = "\"abc\"";
  println!("EXAMPLE 1:\nParsing a simple input string: {}", data);
  let result = parse_string::<()>(data);
  assert_eq!(result, Ok(("", String::from("abc"))));
  println!("Result: {}\n\n", result.unwrap().1);

  let data = "\"tab:\\tafter tab, newline:\\nnew line, quote: \\\", emoji: \\u{1F602}, newline:\\nescaped whitespace: \\    abc\"";
  println!(
    "EXAMPLE 2:\nParsing a string with escape sequences, newline literal, and escaped whitespace:\n\n{}\n",
    data
  );
  let result = parse_string::<()>(data);
  assert_eq!(
    result,
    Ok((
      "",
      String::from("tab:\tafter tab, newline:\nnew line, quote: \", emoji: 😂, newline:\nescaped whitespace: abc")
    ))
  );
  println!("Result:\n\n{}", result.unwrap().1);
}
<|endoftext|>


<commit_before>use libc;

use std::ffi::CStr;
use std::str;

// In paint servers (patterns, gradients, etc.), we have an
// Option<String> for fallback names.  This is a utility function to
// clone one of those.
pub fn clone_fallback_name(fallback: &Option<String>) -> Option<String> {
    if let Some(ref fallback_name) = *fallback {
        Some(fallback_name.clone())
    } else {
        None
    }
}

const DBL_EPSILON: f64 = 1e-10;

pub fn double_equals(a: f64, b: f64) -> bool {
    (a - b).abs() < DBL_EPSILON
}

/// Converts a `char *` which is known to be valid UTF-8 into a `&str`
///
/// The usual `from_glib_none(s)` allocates an owned String.  The
/// purpose of `utf8_cstr()` is to get a temporary string slice into a
/// C string which is already known to be valid UTF-8; for example,
/// as for strings which come from `libxml2`.
pub unsafe fn utf8_cstr<'a>(s: *const libc::c_char) -> &'a str {
    assert!(!s.is_null());

    str::from_utf8_unchecked(CStr::from_ptr(s).to_bytes())
}
<commit_msg>util.rs: Remove DBL_EPSILON and double_equals().  Yay!<commit_after>use libc;

use std::ffi::CStr;
use std::str;

// In paint servers (patterns, gradients, etc.), we have an
// Option<String> for fallback names.  This is a utility function to
// clone one of those.
pub fn clone_fallback_name(fallback: &Option<String>) -> Option<String> {
    if let Some(ref fallback_name) = *fallback {
        Some(fallback_name.clone())
    } else {
        None
    }
}

/// Converts a `char *` which is known to be valid UTF-8 into a `&str`
///
/// The usual `from_glib_none(s)` allocates an owned String.  The
/// purpose of `utf8_cstr()` is to get a temporary string slice into a
/// C string which is already known to be valid UTF-8; for example,
/// as for strings which come from `libxml2`.
pub unsafe fn utf8_cstr<'a>(s: *const libc::c_char) -> &'a str {
    assert!(!s.is_null());

    str::from_utf8_unchecked(CStr::from_ptr(s).to_bytes())
}
<|endoftext|>


<commit_before>/// Individual markov rule
struct MarkovRule {
    pattern: String,
    replacement: String,
    stop: bool,
}

impl MarkovRule {
    fn new(pattern: String, replacement: String, stop: bool) -> MarkovRule {
        MarkovRule {
            pattern: pattern,
            replacement: replacement,
            stop: stop,
        }
    }
}

/// The complete markov algorithm
struct MarkovAlgorithm {
    rules: Vec<MarkovRule>,
}

impl MarkovAlgorithm {
    /// Parse an algorithm description to build a markov algorithm
    pub fn parse(s: &str) -> Result<MarkovAlgorithm, String> {
        let mut rules: Vec<MarkovRule> = vec![];
        for line in s
            .lines()
            .map(|l| l.trim())
            .filter(|l| match l.chars().next() {
                Some(c) if c != '#' => true,
                _ => false,
            })
        {
            // Ignore comments

            // check for -> (must be preceded by whitespace)
            // invalid ruleset if absent
            // whitespace rules mean there's 2 possible variations: " ->" and "\t->"
            let arrow_pos = line.find(" ->").or_else(|| line.find("\t->"));
            match arrow_pos {
                None => {
                    // Ruleset is invalid
                    return Err(format!("Invalid rule \"{}\"", line));
                }
                Some(arrow) => {
                    // extract pattern (trim trailing whitespace)
                    let pattern = line[..arrow].trim_end();

                    // get the string after the arrow
                    // this adds 3 to skip the arrow itself
                    let line_end = line[arrow + 3..].trim_start();

                    // check for . (stop)
                    let stop = match line_end.chars().next() {
                        Some('.') => true,
                        _ => false,
                    };

                    // extract replacement
                    let replacement = if stop { &line_end[1..] } else { line_end };

                    // add to rules
                    let new_rule =
                        MarkovRule::new(pattern.to_string(), replacement.to_string(), stop);
                    rules.push(new_rule);
                }
            }
        }
        let rule_set = MarkovAlgorithm { rules: rules };
        Ok(rule_set)
    }

    /// Transform a text string by applying the markov algorithm
    pub fn apply(&self, input: &str) -> String {
        // get a writable version of the input to work with
        let mut state = input.to_string();

        // loop while operations are possible
        loop {
            // find the first rule that is applicable
            // (pattern string is in state)
            let possible_rule = self
                .rules
                .iter()
                .find(|rule| state.find(&rule.pattern[..]).is_some());

            match possible_rule {
                // stop if no rule found
                None => {
                    break;
                }
                Some(rule) => {
                    // replace the first instance (only) of the pattern
                    // Note: cannot use str::replace as that replaces all instances

                    // unwrap is safe here as the code for finding a rule
                    // already established that the pattern is present
                    let pos = state.find(&rule.pattern[..]).unwrap();
                    let width = rule.pattern.len();

                    // string parts
                    let left = state[..pos].to_string();
                    let right = state[pos + width..].to_string();

                    // construct new string
                    state = format!("{}{}{}", left, rule.replacement, right);

                    // stop if required
                    if rule.stop {
                        break;
                    }
                }
            }
        }

        state
    }
}

/// A Rosetta Code sample
struct RCSample<'a> {
    ruleset: &'a str,
    input: &'a str,
    expected_result: &'a str,
}

// Sample markov algorithms from Rosetta Code
// The extra whitespaces are trimmed when MarkovAlgorithm::parse is called.
fn get_samples<'a>() -> [RCSample<'a>; 5] {
    [
        RCSample {
            ruleset: r"# This rules file is extracted from Wikipedia:
                # http://en.wikipedia.org/wiki/Markov_Algorithm
                A -> apple
                B -> bag
                S -> shop
                T -> the
                the shop -> my brother
                a never used -> .terminating rule",
            input: "I bought a B of As from T S.",
            expected_result: "I bought a bag of apples from my brother.",
        },
        RCSample {
            ruleset: r"# Slightly modified from the rules on Wikipedia
                A -> apple
                B -> bag
                S -> .shop
                T -> the
                the shop -> my brother
                a never used -> .terminating rule",
            input: "I bought a B of As from T S.",
            expected_result: "I bought a bag of apples from T shop.",
        },
        RCSample {
            ruleset: r"# BNF Syntax testing rules
                A -> apple
                WWWW -> with
                Bgage -> ->.*
                B -> bag
                ->.* -> money
                W -> WW
                S -> .shop
                T -> the
                the shop -> my brother
                a never used -> .terminating rule",
            input: "I bought a B of As W my Bgage from T S.",
            expected_result: "I bought a bag of apples with my money from T shop.",
        },
        RCSample {
            ruleset:
                r"### Unary Multiplication Engine, for testing Markov Algorithm implementations
                ### By Donal Fellows.
                # Unary addition engine
                _+1 -> _1+
                1+1 -> 11+
                # Pass for converting from the splitting of multiplication into ordinary
                # addition
                1! -> !1
                ,! -> !+
                _! -> _
                # Unary multiplication by duplicating left side, right side times
                1*1 -> x,@y
                1x -> xX
                X, -> 1,1
                X1 -> 1X
                _x -> _X
                ,x -> ,X
                y1 -> 1y
                y_ -> _
                # Next phase of applying
                1@1 -> x,@y
                1@_ -> @_
                ,@_ -> !_
                ++ -> +
                # Termination cleanup for addition
                _1 -> 1
                1+_ -> 1
                _+_ -> ",
            input: "_1111*11111_",
            expected_result: "11111111111111111111",
        },
        RCSample {
            ruleset: r"# Turing machine: three-state busy beaver
                #
                # state A, symbol 0 => write 1, move right, new state B
                A0 -> 1B
                # state A, symbol 1 => write 1, move left, new state C
                0A1 -> C01
                1A1 -> C11
                # state B, symbol 0 => write 1, move left, new state A
                0B0 -> A01
                1B0 -> A11
                # state B, symbol 1 => write 1, move right, new state B
                B1 -> 1B
                # state C, symbol 0 => write 1, move left, new state B
                0C0 -> B01
                1C0 -> B11
                # state C, symbol 1 => write 1, move left, halt
                0C1 -> H01
                1C1 -> H11",
            input: "000000A000000",
            expected_result: "00011H1111000",
        },
    ]
}

fn main() {
    for (index, sample) in get_samples().iter().enumerate() {
        match MarkovAlgorithm::parse(sample.ruleset) {
            Ok(algorithm) => {
                println!("Sample {}", (index + 1));
                println!("Output: {}", algorithm.apply(sample.input));
                println!("Expected result: {}", sample.expected_result);
            }
            Err(message) => println!("{}", message),
        }
    }
}

#[test]
fn test_samples() {
    for sample in &get_samples() {
        match MarkovAlgorithm::parse(sample.ruleset) {
            Ok(algorithm) => assert_eq!(sample.expected_result, algorithm.apply(sample.input)),
            Err(message) => panic!("{}", message),
        }
    }
}
<commit_msg>rustfmt<commit_after>/// Individual markov rule
struct MarkovRule {
    pattern: String,
    replacement: String,
    stop: bool,
}

impl MarkovRule {
    fn new(pattern: String, replacement: String, stop: bool) -> MarkovRule {
        MarkovRule {
            pattern: pattern,
            replacement: replacement,
            stop: stop,
        }
    }
}

/// The complete markov algorithm
struct MarkovAlgorithm {
    rules: Vec<MarkovRule>,
}

impl MarkovAlgorithm {
    /// Parse an algorithm description to build a markov algorithm
    pub fn parse(s: &str) -> Result<MarkovAlgorithm, String> {
        let mut rules: Vec<MarkovRule> = vec![];
        for line in s
            .lines()
            .map(|l| l.trim())
            .filter(|l| match l.chars().next() {
                Some(c) if c != '#' => true,
                _ => false,
            })
        {
            // Ignore comments

            // check for -> (must be preceded by whitespace)
            // invalid ruleset if absent
            // whitespace rules mean there's 2 possible variations: " ->" and "\t->"
            let arrow_pos = line.find(" ->").or_else(|| line.find("\t->"));
            match arrow_pos {
                None => {
                    // Ruleset is invalid
                    return Err(format!("Invalid rule \"{}\"", line));
                }
                Some(arrow) => {
                    // extract pattern (trim trailing whitespace)
                    let pattern = line[..arrow].trim_end();

                    // get the string after the arrow
                    // this adds 3 to skip the arrow itself
                    let line_end = line[arrow + 3..].trim_start();

                    // check for . (stop)
                    let stop = match line_end.chars().next() {
                        Some('.') => true,
                        _ => false,
                    };

                    // extract replacement
                    let replacement = if stop { &line_end[1..] } else { line_end };

                    // add to rules
                    let new_rule =
                        MarkovRule::new(pattern.to_string(), replacement.to_string(), stop);
                    rules.push(new_rule);
                }
            }
        }
        let rule_set = MarkovAlgorithm { rules: rules };
        Ok(rule_set)
    }

    /// Transform a text string by applying the markov algorithm
    pub fn apply(&self, input: &str) -> String {
        // get a writable version of the input to work with
        let mut state = input.to_string();

        // loop while operations are possible
        loop {
            // find the first rule that is applicable
            // (pattern string is in state)
            let possible_rule = self
                .rules
                .iter()
                .find(|rule| state.find(&rule.pattern[..]).is_some());

            match possible_rule {
                // stop if no rule found
                None => {
                    break;
                }
                Some(rule) => {
                    // replace the first instance (only) of the pattern
                    // Note: cannot use str::replace as that replaces all instances

                    // unwrap is safe here as the code for finding a rule
                    // already established that the pattern is present
                    let pos = state.find(&rule.pattern[..]).unwrap();
                    let width = rule.pattern.len();

                    // string parts
                    let left = state[..pos].to_string();
                    let right = state[pos + width..].to_string();

                    // construct new string
                    state = format!("{}{}{}", left, rule.replacement, right);

                    // stop if required
                    if rule.stop {
                        break;
                    }
                }
            }
        }

        state
    }
}

/// A Rosetta Code sample
struct RCSample<'a> {
    ruleset: &'a str,
    input: &'a str,
    expected_result: &'a str,
}

// Sample markov algorithms from Rosetta Code
// The extra whitespaces are trimmed when MarkovAlgorithm::parse is called.
fn get_samples<'a>() -> [RCSample<'a>; 5] {
    [
        RCSample {
            ruleset: r"# This rules file is extracted from Wikipedia:
                # http://en.wikipedia.org/wiki/Markov_Algorithm
                A -> apple
                B -> bag
                S -> shop
                T -> the
                the shop -> my brother
                a never used -> .terminating rule",
            input: "I bought a B of As from T S.",
            expected_result: "I bought a bag of apples from my brother.",
        },
        RCSample {
            ruleset: r"# Slightly modified from the rules on Wikipedia
                A -> apple
                B -> bag
                S -> .shop
                T -> the
                the shop -> my brother
                a never used -> .terminating rule",
            input: "I bought a B of As from T S.",
            expected_result: "I bought a bag of apples from T shop.",
        },
        RCSample {
            ruleset: r"# BNF Syntax testing rules
                A -> apple
                WWWW -> with
                Bgage -> ->.*
                B -> bag
                ->.* -> money
                W -> WW
                S -> .shop
                T -> the
                the shop -> my brother
                a never used -> .terminating rule",
            input: "I bought a B of As W my Bgage from T S.",
            expected_result: "I bought a bag of apples with my money from T shop.",
        },
        RCSample {
            ruleset: r"### Unary Multiplication Engine, for testing Markov Algorithm implementations
                ### By Donal Fellows.
                # Unary addition engine
                _+1 -> _1+
                1+1 -> 11+
                # Pass for converting from the splitting of multiplication into ordinary
                # addition
                1! -> !1
                ,! -> !+
                _! -> _
                # Unary multiplication by duplicating left side, right side times
                1*1 -> x,@y
                1x -> xX
                X, -> 1,1
                X1 -> 1X
                _x -> _X
                ,x -> ,X
                y1 -> 1y
                y_ -> _
                # Next phase of applying
                1@1 -> x,@y
                1@_ -> @_
                ,@_ -> !_
                ++ -> +
                # Termination cleanup for addition
                _1 -> 1
                1+_ -> 1
                _+_ -> ",
            input: "_1111*11111_",
            expected_result: "11111111111111111111",
        },
        RCSample {
            ruleset: r"# Turing machine: three-state busy beaver
                #
                # state A, symbol 0 => write 1, move right, new state B
                A0 -> 1B
                # state A, symbol 1 => write 1, move left, new state C
                0A1 -> C01
                1A1 -> C11
                # state B, symbol 0 => write 1, move left, new state A
                0B0 -> A01
                1B0 -> A11
                # state B, symbol 1 => write 1, move right, new state B
                B1 -> 1B
                # state C, symbol 0 => write 1, move left, new state B
                0C0 -> B01
                1C0 -> B11
                # state C, symbol 1 => write 1, move left, halt
                0C1 -> H01
                1C1 -> H11",
            input: "000000A000000",
            expected_result: "00011H1111000",
        },
    ]
}

fn main() {
    for (index, sample) in get_samples().iter().enumerate() {
        match MarkovAlgorithm::parse(sample.ruleset) {
            Ok(algorithm) => {
                println!("Sample {}", (index + 1));
                println!("Output: {}", algorithm.apply(sample.input));
                println!("Expected result: {}", sample.expected_result);
            }
            Err(message) => println!("{}", message),
        }
    }
}

#[test]
fn test_samples() {
    for sample in &get_samples() {
        match MarkovAlgorithm::parse(sample.ruleset) {
            Ok(algorithm) => assert_eq!(sample.expected_result, algorithm.apply(sample.input)),
            Err(message) => panic!("{}", message),
        }
    }
}
<|endoftext|>


<commit_before>//! Convenience functions for making requests over HTTP
//!
//! This module provides a few functions that make it easier to make
//! requests over HTTP in a consistent manner e.g. with the 'User-Agent` header
//! set and respecting cache control headers in responses. In addition to reducing
//! the number of network requests for the client, several APIs ask clients
//! to implement caching to reduce load on their servers.

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

use eyre::Result;
use http_cache_reqwest::{CACacheManager, Cache, CacheMode, HttpCache};
use once_cell::sync::Lazy;
use reqwest::header::{HeaderMap, HeaderName};
use reqwest_middleware::{ClientBuilder, ClientWithMiddleware};
use serde::de::DeserializeOwned;
use tempfile::NamedTempFile;

// Re-exports for consumers of this crate
pub use reqwest;
pub use reqwest::header as headers;
pub use reqwest_middleware;
pub use serde;
pub use serde_json;
pub use tempfile;
pub use url;

pub static USER_AGENT: &str = concat!("stencila/", env!("CARGO_PKG_VERSION"),);

/// Get the directory of the HTTP cache
pub fn cache_dir() -> String {
    let user_cache_dir = dirs::cache_dir().unwrap_or_else(|| env::current_dir().unwrap());
    match env::consts::OS {
        "macos" | "windows" => user_cache_dir.join("Stencila").join("HTTP-Cache"),
        _ => user_cache_dir.join("stencila").join("http-cache"),
    }
    .to_string_lossy()
    .to_string()
}

pub static CLIENT: Lazy<ClientWithMiddleware> = Lazy::new(|| {
    let client = reqwest::Client::builder()
        .user_agent(USER_AGENT)
        .build()
        .expect("Should be able to build HTTP client");
    let caching_middleware = Cache(HttpCache {
        mode: CacheMode::Default,
        manager: CACacheManager { path: cache_dir() },
        options: None,
    });
    ClientBuilder::new(client).with(caching_middleware).build()
});

/// Get JSON from a URL
pub async fn get<T: DeserializeOwned>(url: &str) -> Result<T> {
    get_with(url, &[]).await
}

/// Get JSON from a URL with additional request headers
pub async fn get_with<T: DeserializeOwned>(
    url: &str,
    headers: &[(HeaderName, String)],
) -> Result<T> {
    let response = CLIENT
        .get(url)
        .headers(headers_to_map(
            &[
                vec![(headers::ACCEPT, "application/json".to_string())],
                headers.to_vec(),
            ]
            .concat(),
        )?)
        .send()
        .await?
        .error_for_status()?;

    let json = response.json().await?;

    Ok(json)
}

/// Download a file from a URL to a path
pub async fn download(url: &str, path: &Path) -> Result<()> {
    download_with(url, path, &[]).await
}

/// Download a file from a URL to a path with additional request headers
pub async fn download_with(url: &str, path: &Path, headers: &[(HeaderName, String)]) -> Result<()> {
    let response = CLIENT
        .get(url)
        .headers(headers_to_map(headers)?)
        .send()
        .await?
        .error_for_status()?;

    let bytes = response.bytes().await?;
    let mut file = File::create(&path)?;
    io::copy(&mut bytes.as_ref(), &mut file)?;
    file.flush()?;

    Ok(())
}

/// Download a file from a URL to a path synchronously
pub fn download_sync(url: &str, path: &Path) -> Result<()> {
    download_with_sync(url, path, &[])
}

/// Download a file from a URL to a path synchronously with additional request headers
pub fn download_with_sync(url: &str, path: &Path, headers: &[(HeaderName, String)]) -> Result<()> {
    let url = url.to_owned();
    let path = path.to_owned();
    let headers = headers.to_vec();
    let (sender, receiver) = std::sync::mpsc::channel();
    tokio::spawn(async move {
        let result = download_with(&url, &path, &headers).await;
        sender.send(result)
    });
    receiver.recv()?
}

/// Download a file from a URL to a temporary file
pub async fn download_temp(url: &str) -> Result<NamedTempFile> {
    download_temp_with(url, &[]).await
}

/// Download a file from a URL to a temporary file with additional request headers
///
/// Returns a `NamedTempFile` which will remove the temporary file when it
/// is dropped. Be aware of that and the security implications of long-lived temp files:
/// https://docs.rs/tempfile/latest/tempfile/struct.NamedTempFile.html
pub async fn download_temp_with(
    url: &str,
    headers: &[(HeaderName, String)],
) -> Result<NamedTempFile> {
    let temp = NamedTempFile::new()?;
    download_with(url, temp.path(), headers).await?;
    Ok(temp)
}

/// Convert an array of tuples to a `reqwest::HeaderMap`
fn headers_to_map(headers: &[(HeaderName, String)]) -> Result<HeaderMap> {
    let mut header_map = HeaderMap::new();
    for (key, value) in headers {
        header_map.insert(key, value.parse()?);
    }
    Ok(header_map)
}
<commit_msg>feat(HTTP utilities): Add `post` and `delete` functions<commit_after>//! Convenience functions for making requests over HTTP
//!
//! This module provides a few functions that make it easier to make
//! requests over HTTP in a consistent manner e.g. with the 'User-Agent` header
//! set and respecting cache control headers in responses. In addition to reducing
//! the number of network requests for the client, several APIs ask clients
//! to implement caching to reduce load on their servers.

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

use eyre::Result;
use http_cache_reqwest::{CACacheManager, Cache, CacheMode, HttpCache};
use once_cell::sync::Lazy;
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use reqwest_middleware::{ClientBuilder, ClientWithMiddleware};
use serde::de::DeserializeOwned;
use serde::Serialize;
use tempfile::NamedTempFile;

// Re-exports for consumers of this crate
pub use reqwest;
pub use reqwest::header as headers;
pub use reqwest_middleware;
pub use serde;
pub use serde_json;
pub use tempfile;
pub use url;

pub static USER_AGENT: &str = concat!("stencila/", env!("CARGO_PKG_VERSION"),);

/// Get the directory of the HTTP cache
pub fn cache_dir() -> String {
    let user_cache_dir = dirs::cache_dir().unwrap_or_else(|| env::current_dir().unwrap());
    match env::consts::OS {
        "macos" | "windows" => user_cache_dir.join("Stencila").join("HTTP-Cache"),
        _ => user_cache_dir.join("stencila").join("http-cache"),
    }
    .to_string_lossy()
    .to_string()
}

pub static CLIENT: Lazy<ClientWithMiddleware> = Lazy::new(|| {
    let client = reqwest::Client::builder()
        .user_agent(USER_AGENT)
        .build()
        .expect("Should be able to build HTTP client");
    let caching_middleware = Cache(HttpCache {
        mode: CacheMode::Default,
        manager: CACacheManager { path: cache_dir() },
        options: None,
    });
    ClientBuilder::new(client).with(caching_middleware).build()
});

/// Convert an array of tuples to a `reqwest::HeaderMap`
fn header_map(headers: &[(HeaderName, String)]) -> Result<HeaderMap> {
    let mut header_map = HeaderMap::new();
    header_map.insert(headers::ACCEPT, HeaderValue::from_str("application/json")?);
    for (key, value) in headers {
        header_map.insert(key, value.parse()?);
    }
    Ok(header_map)
}

/// Send a GET request
pub async fn get<T: DeserializeOwned>(url: &str) -> Result<T> {
    get_with(url, &[]).await
}

/// Send a GET request with additional request headers
pub async fn get_with<T: DeserializeOwned>(
    url: &str,
    headers: &[(HeaderName, String)],
) -> Result<T> {
    let response = CLIENT
        .get(url)
        .headers(header_map(headers)?)
        .send()
        .await?
        .error_for_status()?;

    let json = response.json().await?;

    Ok(json)
}

/// Send a POST request
pub async fn post<B: Serialize, T: DeserializeOwned>(url: &str, body: B) -> Result<T> {
    post_with(url, body, &[]).await
}

/// Send a POST request with additional request headers
pub async fn post_with<B: Serialize, T: DeserializeOwned>(
    url: &str,
    body: B,
    headers: &[(HeaderName, String)],
) -> Result<T> {
    let response = CLIENT
        .post(url)
        .json(&body)
        .headers(header_map(headers)?)
        .send()
        .await?
        .error_for_status()?;

    let json = response.json().await?;

    Ok(json)
}

/// Send a DELETE request
pub async fn delete(url: &str) -> Result<()> {
    delete_with(url, &[]).await
}

/// Send a DELETE request with additional request headers
pub async fn delete_with(url: &str, headers: &[(HeaderName, String)]) -> Result<()> {
    CLIENT
        .delete(url)
        .headers(header_map(headers)?)
        .send()
        .await?
        .error_for_status()?;
    Ok(())
}

/// Download a file from to a path
pub async fn download(url: &str, path: &Path) -> Result<()> {
    download_with(url, path, &[]).await
}

/// Download a file from to a path with additional request headers
pub async fn download_with(url: &str, path: &Path, headers: &[(HeaderName, String)]) -> Result<()> {
    let response = CLIENT
        .get(url)
        .headers(header_map(headers)?)
        .send()
        .await?
        .error_for_status()?;

    let bytes = response.bytes().await?;
    let mut file = File::create(&path)?;
    io::copy(&mut bytes.as_ref(), &mut file)?;
    file.flush()?;

    Ok(())
}

/// Download a file from to a path synchronously
pub fn download_sync(url: &str, path: &Path) -> Result<()> {
    download_with_sync(url, path, &[])
}

/// Download a file from to a path synchronously with additional request headers
pub fn download_with_sync(url: &str, path: &Path, headers: &[(HeaderName, String)]) -> Result<()> {
    let url = url.to_owned();
    let path = path.to_owned();
    let headers = headers.to_vec();
    let (sender, receiver) = std::sync::mpsc::channel();
    tokio::spawn(async move {
        let result = download_with(&url, &path, &headers).await;
        sender.send(result)
    });
    receiver.recv()?
}

/// Download a file from to a temporary file
pub async fn download_temp(url: &str) -> Result<NamedTempFile> {
    download_temp_with(url, &[]).await
}

/// Download a file from to a temporary file with additional request headers
///
/// Returns a `NamedTempFile` which will remove the temporary file when it
/// is dropped. Be aware of that and the security implications of long-lived temp files:
/// https://docs.rs/tempfile/latest/tempfile/struct.NamedTempFile.html
pub async fn download_temp_with(
    url: &str,
    headers: &[(HeaderName, String)],
) -> Result<NamedTempFile> {
    let temp = NamedTempFile::new()?;
    download_with(url, temp.path(), headers).await?;
    Ok(temp)
}
<|endoftext|>


<commit_before>use lalafell::commands::prelude::*;

use lalafell::error::*;

use fflogs::{self, FfLogs};
use fflogs::net::{Server, Metric, Job};

use std::cmp::Ordering;
use std::str::FromStr;

pub struct FfLogsCommand {
  fflogs: FfLogs
}

impl ::commands::BotCommand for FfLogsCommand {
  fn new(env: Arc<::bot::BotEnv>) -> Self {
    FfLogsCommand {
      fflogs: FfLogs::new(&env.environment.fflogs_api_key)
    }
  }
}

#[derive(Debug, StructOpt)]
#[structopt(about = "Display information about the character on FF Logs")]
pub struct Params {
  #[structopt(help = "The server the character is on")]
  server: String,
  #[structopt(help = "The character's first name")]
  first_name: String,
  #[structopt(help = "The character's last name")]
  last_name: String,
  #[structopt(
    short = "j",
    long = "job",
    help = "The job to look at",
    parse(try_from_str)
  )]
  job: Option<Job>
}

impl HasParams for FfLogsCommand {
  type Params = Params;
}

impl<'a> Command<'a> for FfLogsCommand {
  fn run(&self, _: &Context, _: &Message, params: &[&str]) -> CommandResult<'a> {
    let params = self.params_then("fflogs", params, |a| a.setting(::structopt::clap::AppSettings::ArgRequiredElseHelp))?;
    let server = match Server::from_str(&params.server) {
      Ok(s) => s,
      Err(_) => return Err("Invalid server.".into())
    };
    let name = format!("{} {}", params.first_name, params.last_name);

    let parses = self.fflogs.parses(
      &name,
      server,
      |x| x.metric(Metric::Dps)
    );

    let parses = match parses {
      Ok(r) => r,
      Err(fflogs::errors::Error::FfLogs(fflogs::net::FfLogsError { status, error })) => {
        return Err(format!("FF Logs didn't like that ({}): {}", status, error).into());
      },
      Err(e) => return Err(e).chain_err(|| "could not query fflogs")?
    };

    if parses.is_empty() {
      return Ok("No parses found.".into());
    }

    let first_spec = match parses[0].specs.get(0) {
      Some(s) => s,
      None => return Err("Somehow there was no first spec.".into())
    };
    let first_data = match first_spec.data.get(0) {
      Some(d) => d,
      None => return Err("Somehow there was no first data.".into())
    };

    let job = params.job.or_else(|| Job::from_str(&first_spec.spec).ok()).unwrap().to_string();
    let lower_job = job.to_lowercase().replace(" ", "");
    let name = &first_data.character_name;
    let id = first_data.character_id;

    let mut embed = CreateEmbed::default();

    embed = embed
      .title(&name)
      .url(format!("https://www.fflogs.com/character/id/{}", id))
      .field("Job", &job, true)
      .field("Server", server.as_ref(), true);

    let mut count = 0;

    for parse in &parses {
      let spec = match parse.specs.iter().filter(|s| s.spec.to_lowercase() == lower_job).next() {
        Some(s) => s,
        None => continue
      };
      let data = match spec.data.iter().max_by(|a, b| a.historical_percent.partial_cmp(&b.historical_percent).unwrap_or(Ordering::Less)) {
        Some(d) => d,
        None => continue
      };

      count += 1;

      let url = format!("https://www.fflogs.com/reports/{}#fight={}", data.report_code, data.report_fight);

      let string = format!("[Link]({url}) – {dps:.2} DPS – {perc:.2} percentile out of {total_parses} parses",
        url = url,
        dps = data.persecondamount,
        perc = data.historical_percent,
        total_parses = data.historical_count
      );

      embed = embed.field(&parse.name, string, false);
    }

    if count == 0 {
      return Ok("No parses for that job.".into());
    }

    Ok(CommandSuccess::default().message(|_| embed))
  }
}
<commit_msg>refactor: use find instead of filter/next<commit_after>use lalafell::commands::prelude::*;

use lalafell::error::*;

use fflogs::{self, FfLogs};
use fflogs::net::{Server, Metric, Job};

use std::cmp::Ordering;
use std::str::FromStr;

pub struct FfLogsCommand {
  fflogs: FfLogs
}

impl ::commands::BotCommand for FfLogsCommand {
  fn new(env: Arc<::bot::BotEnv>) -> Self {
    FfLogsCommand {
      fflogs: FfLogs::new(&env.environment.fflogs_api_key)
    }
  }
}

#[derive(Debug, StructOpt)]
#[structopt(about = "Display information about the character on FF Logs")]
pub struct Params {
  #[structopt(help = "The server the character is on")]
  server: String,
  #[structopt(help = "The character's first name")]
  first_name: String,
  #[structopt(help = "The character's last name")]
  last_name: String,
  #[structopt(
    short = "j",
    long = "job",
    help = "The job to look at",
    parse(try_from_str)
  )]
  job: Option<Job>
}

impl HasParams for FfLogsCommand {
  type Params = Params;
}

impl<'a> Command<'a> for FfLogsCommand {
  fn run(&self, _: &Context, _: &Message, params: &[&str]) -> CommandResult<'a> {
    let params = self.params_then("fflogs", params, |a| a.setting(::structopt::clap::AppSettings::ArgRequiredElseHelp))?;
    let server = match Server::from_str(&params.server) {
      Ok(s) => s,
      Err(_) => return Err("Invalid server.".into())
    };
    let name = format!("{} {}", params.first_name, params.last_name);

    let parses = self.fflogs.parses(
      &name,
      server,
      |x| x.metric(Metric::Dps)
    );

    let parses = match parses {
      Ok(r) => r,
      Err(fflogs::errors::Error::FfLogs(fflogs::net::FfLogsError { status, error })) => {
        return Err(format!("FF Logs didn't like that ({}): {}", status, error).into());
      },
      Err(e) => return Err(e).chain_err(|| "could not query fflogs")?
    };

    if parses.is_empty() {
      return Ok("No parses found.".into());
    }

    let first_spec = match parses[0].specs.get(0) {
      Some(s) => s,
      None => return Err("Somehow there was no first spec.".into())
    };
    let first_data = match first_spec.data.get(0) {
      Some(d) => d,
      None => return Err("Somehow there was no first data.".into())
    };

    let job = params.job.or_else(|| Job::from_str(&first_spec.spec).ok()).unwrap().to_string();
    let lower_job = job.to_lowercase().replace(" ", "");
    let name = &first_data.character_name;
    let id = first_data.character_id;

    let mut embed = CreateEmbed::default();

    embed = embed
      .title(&name)
      .url(format!("https://www.fflogs.com/character/id/{}", id))
      .field("Job", &job, true)
      .field("Server", server.as_ref(), true);

    let mut count = 0;

    for parse in &parses {
      let spec = match parse.specs.iter().find(|s| s.spec.to_lowercase() == lower_job) {
        Some(s) => s,
        None => continue
      };
      let data = match spec.data.iter().max_by(|a, b| a.historical_percent.partial_cmp(&b.historical_percent).unwrap_or(Ordering::Less)) {
        Some(d) => d,
        None => continue
      };

      count += 1;

      let url = format!("https://www.fflogs.com/reports/{}#fight={}", data.report_code, data.report_fight);

      let string = format!("[Link]({url}) – {dps:.2} DPS – {perc:.2} percentile out of {total_parses} parses",
        url = url,
        dps = data.persecondamount,
        perc = data.historical_percent,
        total_parses = data.historical_count
      );

      embed = embed.field(&parse.name, string, false);
    }

    if count == 0 {
      return Ok("No parses for that job.".into());
    }

    Ok(CommandSuccess::default().message(|_| embed))
  }
}
<|endoftext|>


<commit_before>//! Infrastructure types related to packaging rate-limit information alongside responses from
//! Twitter.

use std::vec;
use std::iter::FromIterator;
use std::io::Read;
use hyper::client::response::Response as HyperResponse;
use hyper::status::StatusCode;
use rustc_serialize::json;
use super::{FromJson};
use error::{self, TwitterErrors};
use error::Error::*;

header! { (XRateLimitLimit, "X-Rate-Limit-Limit") => [i32] }
header! { (XRateLimitRemaining, "X-Rate-Limit-Remaining") => [i32] }
header! { (XRateLimitReset, "X-Rate-Limit-Reset") => [i32] }

///A helper struct to wrap response data with accompanying rate limit information.
#[derive(Debug)]
pub struct Response<T> {
    ///The rate limit ceiling for the given request.
    pub rate_limit: i32,
    ///The number of requests left for the 15-minute window.
    pub rate_limit_remaining: i32,
    ///The UTC Unix timestamp at which the rate window resets.
    pub rate_limit_reset: i32,
    ///The decoded response from the request.
    pub response: T,
}

///Iterator returned by calling `.into_iter()` on a `Response<Vec<T>>`.
pub struct ResponseIter<T> {
    rate_limit: i32,
    rate_limit_remaining: i32,
    rate_limit_reset: i32,
    resp_iter: vec::IntoIter<T>,
}

impl<T> Iterator for ResponseIter<T> {
    type Item = Response<T>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(resp) = self.resp_iter.next() {
            Some(Response {
                rate_limit: self.rate_limit,
                rate_limit_remaining: self.rate_limit_remaining,
                rate_limit_reset: self.rate_limit_reset,
                response: resp,
            })
        }
        else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.resp_iter.size_hint()
    }
}

impl<T> DoubleEndedIterator for ResponseIter<T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        if let Some(resp) = self.resp_iter.next_back() {
            Some(Response {
                rate_limit: self.rate_limit,
                rate_limit_remaining: self.rate_limit_remaining,
                rate_limit_reset: self.rate_limit_reset,
                response: resp,
            })
        }
        else {
            None
        }
    }
}

impl<T> ExactSizeIterator for ResponseIter<T> {
    fn len(&self) -> usize {
        self.resp_iter.len()
    }
}

impl<T> IntoIterator for Response<Vec<T>> {
    type Item = Response<T>;
    type IntoIter = ResponseIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        ResponseIter {
            rate_limit: self.rate_limit,
            rate_limit_remaining: self.rate_limit_remaining,
            rate_limit_reset: self.rate_limit_reset,
            resp_iter: self.response.into_iter(),
        }
    }
}

impl<T> FromIterator<Response<T>> for Response<Vec<T>> {
    fn from_iter<I>(iter: I) -> Self
        where I: IntoIterator<Item=Response<T>>
    {
        let mut resp = Response {
            rate_limit: -1,
            rate_limit_remaining: -1,
            rate_limit_reset: -1,
            response: Vec::new(),
        };

        for item in iter {
            if item.rate_limit_reset > resp.rate_limit_reset {
                resp.rate_limit = item.rate_limit;
                resp.rate_limit_remaining = item.rate_limit_remaining;
                resp.rate_limit_reset = item.rate_limit_reset;
            }
            else if (item.rate_limit_reset == resp.rate_limit_reset) &&
                    (item.rate_limit_remaining < resp.rate_limit_remaining) {
                resp.rate_limit = item.rate_limit;
                resp.rate_limit_remaining = item.rate_limit_remaining;
                resp.rate_limit_reset = item.rate_limit_reset;
            }
            resp.response.push(item.response);
        }

        resp
    }
}

///With the given response struct, parse it into a String.
pub fn response_raw(resp: &mut HyperResponse) -> Result<String, error::Error> {
    let mut full_resp = String::new();
    try!(resp.read_to_string(&mut full_resp));

    if let Ok(err) = json::decode::<TwitterErrors>(&full_resp) {
        if err.errors.iter().any(|e| e.code == 88) {
            if resp.headers.has::<XRateLimitReset>() {
                return Err(RateLimit(resp.headers.get::<XRateLimitReset>().map(|h| h.0).unwrap()));
            }
            else {
                return Err(TwitterError(err));
            }
        }
        else {
            return Err(TwitterError(err));
        }
    }

    match resp.status {
        StatusCode::Ok | StatusCode::NotModified => (),
        _ => return Err(BadStatus(resp.status)),
    }

    Ok(full_resp)
}

///With the given response struct, parse it into the desired format and
///return it along with rate limit information.
pub fn parse_response<T: FromJson>(resp: &mut HyperResponse) -> Result<Response<T>, error::Error> {
    let resp_str = try!(response_raw(resp));

    Ok(Response {
        rate_limit: resp.headers.get::<XRateLimitLimit>().map_or(-1, |h| h.0),
        rate_limit_remaining: resp.headers.get::<XRateLimitRemaining>().map_or(-1, |h| h.0),
        rate_limit_reset: resp.headers.get::<XRateLimitReset>().map_or(-1, |h| h.0),
        response: try!(T::from_str(&resp_str)),
    })
}
<commit_msg>technically HTTP 304 Not Modified is an error code<commit_after>//! Infrastructure types related to packaging rate-limit information alongside responses from
//! Twitter.

use std::vec;
use std::iter::FromIterator;
use std::io::Read;
use hyper::client::response::Response as HyperResponse;
use hyper::status::StatusCode;
use rustc_serialize::json;
use super::{FromJson};
use error::{self, TwitterErrors};
use error::Error::*;

header! { (XRateLimitLimit, "X-Rate-Limit-Limit") => [i32] }
header! { (XRateLimitRemaining, "X-Rate-Limit-Remaining") => [i32] }
header! { (XRateLimitReset, "X-Rate-Limit-Reset") => [i32] }

///A helper struct to wrap response data with accompanying rate limit information.
#[derive(Debug)]
pub struct Response<T> {
    ///The rate limit ceiling for the given request.
    pub rate_limit: i32,
    ///The number of requests left for the 15-minute window.
    pub rate_limit_remaining: i32,
    ///The UTC Unix timestamp at which the rate window resets.
    pub rate_limit_reset: i32,
    ///The decoded response from the request.
    pub response: T,
}

///Iterator returned by calling `.into_iter()` on a `Response<Vec<T>>`.
pub struct ResponseIter<T> {
    rate_limit: i32,
    rate_limit_remaining: i32,
    rate_limit_reset: i32,
    resp_iter: vec::IntoIter<T>,
}

impl<T> Iterator for ResponseIter<T> {
    type Item = Response<T>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(resp) = self.resp_iter.next() {
            Some(Response {
                rate_limit: self.rate_limit,
                rate_limit_remaining: self.rate_limit_remaining,
                rate_limit_reset: self.rate_limit_reset,
                response: resp,
            })
        }
        else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.resp_iter.size_hint()
    }
}

impl<T> DoubleEndedIterator for ResponseIter<T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        if let Some(resp) = self.resp_iter.next_back() {
            Some(Response {
                rate_limit: self.rate_limit,
                rate_limit_remaining: self.rate_limit_remaining,
                rate_limit_reset: self.rate_limit_reset,
                response: resp,
            })
        }
        else {
            None
        }
    }
}

impl<T> ExactSizeIterator for ResponseIter<T> {
    fn len(&self) -> usize {
        self.resp_iter.len()
    }
}

impl<T> IntoIterator for Response<Vec<T>> {
    type Item = Response<T>;
    type IntoIter = ResponseIter<T>;

    fn into_iter(self) -> Self::IntoIter {
        ResponseIter {
            rate_limit: self.rate_limit,
            rate_limit_remaining: self.rate_limit_remaining,
            rate_limit_reset: self.rate_limit_reset,
            resp_iter: self.response.into_iter(),
        }
    }
}

impl<T> FromIterator<Response<T>> for Response<Vec<T>> {
    fn from_iter<I>(iter: I) -> Self
        where I: IntoIterator<Item=Response<T>>
    {
        let mut resp = Response {
            rate_limit: -1,
            rate_limit_remaining: -1,
            rate_limit_reset: -1,
            response: Vec::new(),
        };

        for item in iter {
            if item.rate_limit_reset > resp.rate_limit_reset {
                resp.rate_limit = item.rate_limit;
                resp.rate_limit_remaining = item.rate_limit_remaining;
                resp.rate_limit_reset = item.rate_limit_reset;
            }
            else if (item.rate_limit_reset == resp.rate_limit_reset) &&
                    (item.rate_limit_remaining < resp.rate_limit_remaining) {
                resp.rate_limit = item.rate_limit;
                resp.rate_limit_remaining = item.rate_limit_remaining;
                resp.rate_limit_reset = item.rate_limit_reset;
            }
            resp.response.push(item.response);
        }

        resp
    }
}

///With the given response struct, parse it into a String.
pub fn response_raw(resp: &mut HyperResponse) -> Result<String, error::Error> {
    let mut full_resp = String::new();
    try!(resp.read_to_string(&mut full_resp));

    if let Ok(err) = json::decode::<TwitterErrors>(&full_resp) {
        if err.errors.iter().any(|e| e.code == 88) {
            if resp.headers.has::<XRateLimitReset>() {
                return Err(RateLimit(resp.headers.get::<XRateLimitReset>().map(|h| h.0).unwrap()));
            }
            else {
                return Err(TwitterError(err));
            }
        }
        else {
            return Err(TwitterError(err));
        }
    }

    match resp.status {
        StatusCode::Ok => (),
        _ => return Err(BadStatus(resp.status)),
    }

    Ok(full_resp)
}

///With the given response struct, parse it into the desired format and
///return it along with rate limit information.
pub fn parse_response<T: FromJson>(resp: &mut HyperResponse) -> Result<Response<T>, error::Error> {
    let resp_str = try!(response_raw(resp));

    Ok(Response {
        rate_limit: resp.headers.get::<XRateLimitLimit>().map_or(-1, |h| h.0),
        rate_limit_remaining: resp.headers.get::<XRateLimitRemaining>().map_or(-1, |h| h.0),
        rate_limit_reset: resp.headers.get::<XRateLimitReset>().map_or(-1, |h| h.0),
        response: try!(T::from_str(&resp_str)),
    })
}
<|endoftext|>


<commit_before>pub mod tag_command;
pub mod autotag;
pub mod update_tags;

pub use self::tag_command::TagCommand;
pub use self::autotag::AutoTagCommand;
pub use self::update_tags::UpdateTagsCommand;

use LalafellBot;
use database::AutotagUser;

use xivdb::error::*;

use discord::model::{UserId, LiveServer, Role, RoleId};

use std::sync::{Arc, Mutex};

lazy_static! {
  // Limbo roles are roles that may or may not be added to the Discord bot state.
  static ref LIMBO_ROLES: Mutex<Vec<Role>> = Mutex::default();
}

pub struct Tagger;

impl Tagger {
  pub fn search_tag(bot: Arc<LalafellBot>, who: UserId, on: &LiveServer, server: &str, character_name: &str, ignore_verified: bool) -> Result<Option<String>> {
    let params = &[
      ("one", "characters"),
      ("strict", "on"),
      ("server|et", server)
    ];

    let res = bot.xivdb.search(character_name, params).chain_err(|| "could not query XIVDB")?;

    let search_chars = res.characters.unwrap().results;
    if search_chars.is_empty() {
      return Ok(Some(format!("Could not find any character by the name {} on {}.", character_name, server)));
    }

    let char_id = match search_chars[0]["id"].as_u64() {
      Some(u) => u,
      None => return Err("character ID was not a u64".into())
    };

    let name = match search_chars[0]["name"].as_str() {
      Some(s) => s,
      None => return Err("character name was not a string".into())
    };

    if name.to_lowercase() != character_name.to_lowercase() {
      return Ok(Some(format!("Could not find any character by the name {} on {}.", character_name, server)));
    }

    Tagger::tag(bot, who, on, char_id, ignore_verified)
  }

  pub fn tag(bot: Arc<LalafellBot>, who: UserId, on: &LiveServer, char_id: u64, ignore_verified: bool) -> Result<Option<String>> {
    let is_verified = match bot.database.read().unwrap().autotags.users.iter().find(|u| u.user_id == who.0 && u.server_id == on.id.0) {
      Some(u) => {
        if u.verification.verified && !ignore_verified && char_id != u.character_id {
          return Ok(Some(format!("{} is verified as {} on {}, so they cannot switch to another account.", who.mention(), u.character, u.server)));
        }
        u.verification.verified
      },
      None => false
    };

    let member = bot.discord.get_member(on.id, who).chain_err(|| "could not get member for tagging")?;

    let character = bot.xivdb.character(char_id).chain_err(|| "could not look up character")?;

    bot.database.write().unwrap().autotags.update_or_add(AutotagUser::new(
      who.0,
      on.id.0,
      character.lodestone_id,
      &character.name,
      &character.server
    ));

    // Get a copy of the roles on the server.
    let mut roles = on.roles.clone();
    // Check for existing limbo roles.
    {
      let limbo = &mut *LIMBO_ROLES.lock().unwrap();
      for role in &roles {
        // If the server has updated to contain the limbo role, remove it.
        if let Some(i) = limbo.iter().position(|x| x.id == role.id) {
          limbo.remove(i);
        }
      }
    }
    // Get a copy of the limbo roles.
    let limbo = LIMBO_ROLES.lock().unwrap().clone();
    // Extend the server roles with the limbo roles.
    roles.extend(limbo);
    let mut created_roles = Vec::new();
    let mut add_roles = Vec::new();
    match roles.iter().find(|x| x.name.to_lowercase() == character.data.race.to_lowercase()) {
      Some(r) => add_roles.push(r),
      None => {
        let role = bot.discord.create_role(on.id, Some(&character.data.race.to_lowercase()), None, None, None, None).chain_err(|| "could not create role")?;
        created_roles.push(role);
      }
    }
    match roles.iter().find(|x| x.name.to_lowercase() == character.data.gender.to_lowercase()) {
      Some(r) => add_roles.push(r),
      None => {
        let role = bot.discord.create_role(on.id, Some(&character.data.gender.to_lowercase()), None, None, None, None).chain_err(|| "could not create role")?;
        created_roles.push(role);
      }
    }
    match roles.iter().find(|x| x.name.to_lowercase() == character.server.to_lowercase()) {
      Some(r) => add_roles.push(r),
      None => {
        let role = bot.discord.create_role(on.id, Some(&character.server.to_lowercase()), None, None, None, None).chain_err(|| "could not create role")?;
        created_roles.push(role);
      }
    }
    if is_verified {
      match roles.iter().find(|x| x.name.to_lowercase() == "verified") {
        Some(r) => add_roles.push(r),
        None => {
          let role = bot.discord.create_role(on.id, Some("Verified"), None, None, None, None).chain_err(|| "could not create role")?;
          created_roles.push(role);
        }
      }
    }

    // If we created any roles, the server may or may not update with them fast enough, so store a copy in the limbo
    // roles.
    {
      let mut limbo = &mut *LIMBO_ROLES.lock().unwrap();
      for created in &created_roles {
        limbo.push(created.clone());
      }
    }

    debug!("Created the following roles:\n:{:#?}", created_roles);

    debug!("Existing roles:\n{:#?}", add_roles);

    // Extend the roles to add with the roles we created.
    add_roles.extend(created_roles.iter());
    // Get all the roles that are part of groups
    let all_group_roles: Vec<String> = bot.config.roles.groups.iter().flat_map(|x| x).map(|x| x.to_lowercase()).collect();
    // Filter all roles on the server to only the roles the member has
    let keep: Vec<&Role> = roles.iter().filter(|x| member.roles.contains(&x.id)).collect();
    // Filter all the roles the member has, keeping the ones not in a group. These roles will not be touched when
    // updating the tag.
    let keep: Vec<&Role> = keep.into_iter().filter(|x| !all_group_roles.contains(&x.name.to_lowercase())).collect();
    debug!("Roles to keep:\n{:#?}", keep);
    // Combine the two sets of roles and map them to IDs
    let mut role_set: Vec<RoleId> = add_roles.iter().map(|r| r.id).chain(keep.into_iter().map(|r| r.id)).collect();
    // Sort the IDs so we can dedup them
    role_set.sort();
    // Remove the duplicate roles, if any
    role_set.dedup();

    debug!("Final role set:\n{:#?}", role_set);

    // Only update the roles if the member is missing any or if any are different
    if !member.roles.iter().all(|r| role_set.contains(r)) {
      bot.discord.edit_member_roles(on.id, who, &role_set).chain_err(|| "could not add roles")?;
    }

    // cannot edit nickname of those with a higher role
    let nick = match member.nick {
      Some(n) => n,
      None => Default::default()
    };
    if nick != character.name {
      bot.discord.edit_member(on.id, who, |e| e.nickname(&character.name)).ok();
    }
    Ok(None)
  }
}
<commit_msg>fix: test equality<commit_after>pub mod tag_command;
pub mod autotag;
pub mod update_tags;

pub use self::tag_command::TagCommand;
pub use self::autotag::AutoTagCommand;
pub use self::update_tags::UpdateTagsCommand;

use LalafellBot;
use database::AutotagUser;

use xivdb::error::*;

use discord::model::{UserId, LiveServer, Role, RoleId};

use std::sync::{Arc, Mutex};

lazy_static! {
  // Limbo roles are roles that may or may not be added to the Discord bot state.
  static ref LIMBO_ROLES: Mutex<Vec<Role>> = Mutex::default();
}

pub struct Tagger;

impl Tagger {
  pub fn search_tag(bot: Arc<LalafellBot>, who: UserId, on: &LiveServer, server: &str, character_name: &str, ignore_verified: bool) -> Result<Option<String>> {
    let params = &[
      ("one", "characters"),
      ("strict", "on"),
      ("server|et", server)
    ];

    let res = bot.xivdb.search(character_name, params).chain_err(|| "could not query XIVDB")?;

    let search_chars = res.characters.unwrap().results;
    if search_chars.is_empty() {
      return Ok(Some(format!("Could not find any character by the name {} on {}.", character_name, server)));
    }

    let char_id = match search_chars[0]["id"].as_u64() {
      Some(u) => u,
      None => return Err("character ID was not a u64".into())
    };

    let name = match search_chars[0]["name"].as_str() {
      Some(s) => s,
      None => return Err("character name was not a string".into())
    };

    if name.to_lowercase() != character_name.to_lowercase() {
      return Ok(Some(format!("Could not find any character by the name {} on {}.", character_name, server)));
    }

    Tagger::tag(bot, who, on, char_id, ignore_verified)
  }

  pub fn tag(bot: Arc<LalafellBot>, who: UserId, on: &LiveServer, char_id: u64, ignore_verified: bool) -> Result<Option<String>> {
    let is_verified = match bot.database.read().unwrap().autotags.users.iter().find(|u| u.user_id == who.0 && u.server_id == on.id.0) {
      Some(u) => {
        if u.verification.verified && !ignore_verified && char_id != u.character_id {
          return Ok(Some(format!("{} is verified as {} on {}, so they cannot switch to another account.", who.mention(), u.character, u.server)));
        }
        u.verification.verified
      },
      None => false
    };

    let member = bot.discord.get_member(on.id, who).chain_err(|| "could not get member for tagging")?;

    let character = bot.xivdb.character(char_id).chain_err(|| "could not look up character")?;

    bot.database.write().unwrap().autotags.update_or_add(AutotagUser::new(
      who.0,
      on.id.0,
      character.lodestone_id,
      &character.name,
      &character.server
    ));

    // Get a copy of the roles on the server.
    let mut roles = on.roles.clone();
    // Check for existing limbo roles.
    {
      let limbo = &mut *LIMBO_ROLES.lock().unwrap();
      for role in &roles {
        // If the server has updated to contain the limbo role, remove it.
        if let Some(i) = limbo.iter().position(|x| x.id == role.id) {
          limbo.remove(i);
        }
      }
    }
    // Get a copy of the limbo roles.
    let limbo = LIMBO_ROLES.lock().unwrap().clone();
    // Extend the server roles with the limbo roles.
    roles.extend(limbo);
    let mut created_roles = Vec::new();
    let mut add_roles = Vec::new();
    match roles.iter().find(|x| x.name.to_lowercase() == character.data.race.to_lowercase()) {
      Some(r) => add_roles.push(r),
      None => {
        let role = bot.discord.create_role(on.id, Some(&character.data.race.to_lowercase()), None, None, None, None).chain_err(|| "could not create role")?;
        created_roles.push(role);
      }
    }
    match roles.iter().find(|x| x.name.to_lowercase() == character.data.gender.to_lowercase()) {
      Some(r) => add_roles.push(r),
      None => {
        let role = bot.discord.create_role(on.id, Some(&character.data.gender.to_lowercase()), None, None, None, None).chain_err(|| "could not create role")?;
        created_roles.push(role);
      }
    }
    match roles.iter().find(|x| x.name.to_lowercase() == character.server.to_lowercase()) {
      Some(r) => add_roles.push(r),
      None => {
        let role = bot.discord.create_role(on.id, Some(&character.server.to_lowercase()), None, None, None, None).chain_err(|| "could not create role")?;
        created_roles.push(role);
      }
    }
    if is_verified {
      match roles.iter().find(|x| x.name.to_lowercase() == "verified") {
        Some(r) => add_roles.push(r),
        None => {
          let role = bot.discord.create_role(on.id, Some("Verified"), None, None, None, None).chain_err(|| "could not create role")?;
          created_roles.push(role);
        }
      }
    }

    // If we created any roles, the server may or may not update with them fast enough, so store a copy in the limbo
    // roles.
    {
      let mut limbo = &mut *LIMBO_ROLES.lock().unwrap();
      for created in &created_roles {
        limbo.push(created.clone());
      }
    }

    debug!("Created the following roles:\n:{:#?}", created_roles);

    debug!("Existing roles:\n{:#?}", add_roles);

    // Extend the roles to add with the roles we created.
    add_roles.extend(created_roles.iter());
    // Get all the roles that are part of groups
    let all_group_roles: Vec<String> = bot.config.roles.groups.iter().flat_map(|x| x).map(|x| x.to_lowercase()).collect();
    // Filter all roles on the server to only the roles the member has
    let keep: Vec<&Role> = roles.iter().filter(|x| member.roles.contains(&x.id)).collect();
    // Filter all the roles the member has, keeping the ones not in a group. These roles will not be touched when
    // updating the tag.
    let keep: Vec<&Role> = keep.into_iter().filter(|x| !all_group_roles.contains(&x.name.to_lowercase())).collect();
    debug!("Roles to keep:\n{:#?}", keep);
    // Combine the two sets of roles and map them to IDs
    let mut role_set: Vec<RoleId> = add_roles.iter().map(|r| r.id).chain(keep.into_iter().map(|r| r.id)).collect();
    // Sort the IDs so we can dedup them
    role_set.sort();
    // Remove the duplicate roles, if any
    role_set.dedup();

    debug!("Final role set:\n{:#?}", role_set);

    // Only update the roles if they are different
    let different = {
      use std::collections::HashSet;
      let member_roles: HashSet<u64> = member.roles.iter().map(|x| x.0).collect();
      let actual_role_set: HashSet<u64> = role_set.iter().map(|x| x.0).collect();
      member_roles != actual_role_set
    };
    if different {
      bot.discord.edit_member_roles(on.id, who, &role_set).chain_err(|| "could not add roles")?;
    }

    // cannot edit nickname of those with a higher role
    let nick = match member.nick {
      Some(n) => n,
      None => Default::default()
    };
    if nick != character.name {
      bot.discord.edit_member(on.id, who, |e| e.nickname(&character.name)).ok();
    }
    Ok(None)
  }
}
<|endoftext|>


<commit_before>use std::io::buffered::BufferedStream;
use std::io::net::ip::SocketAddr;
use std::io::net::tcp::TcpStream;

static PROTO_VERSION: i32 = 196608; // 3 << 16

#[deriving(ToStr)]
enum AuthType {
    AuthUnknown,
    AuthOk,
    AuthKerberosV4,
    AuthKerberosV5,
    AuthClearTextPassword,
    AuthCryptPassword(u16),
    AuthMD5Password(u32),
    AuthSCMCredential
}

fn write_cstring(io: &mut BufferedStream<TcpStream>, str: &str) {
      io.write_str(str);
      io.write_u8(0);
}

fn write_startup_message(io: &mut BufferedStream<TcpStream>, params: &[(&str, &str)]) {
    let mut sz: uint = 8 + 1;

    for &(k, v) in params.iter() {
      sz += k.len() + 1 + v.len() + 1;
    }

    io.write_be_i32(sz as i32);
    io.write_be_i32(PROTO_VERSION);

    for &(k, v) in params.iter() {
        write_cstring(io, k);
        write_cstring(io, v);
    }
    io.write_u8(0);
    io.flush();
}

// body_size excludes the typ byte and the size i32
fn write_message_header(io: &mut BufferedStream<TcpStream>, typ: u8, payload: uint) {
    io.write_u8(typ);
    io.write_be_i32((payload + 4) as i32);
}

fn write_query_message(io: &mut BufferedStream<TcpStream>, query: &str) {
    write_message_header(io, 'Q' as u8, query.len() + 1);
    write_cstring(io, query);
    io.flush();
}

fn parse_auth_message(io: &mut BufferedStream<TcpStream>, rem_len: i32) -> AuthType {
    assert!(rem_len >= 4);
    let authtype = io.read_be_i32();
    match authtype {
        0 => AuthOk,
        1 => AuthKerberosV4,
        2 => AuthKerberosV5,
        3 => AuthClearTextPassword,
        4 => {
            assert!(rem_len >= 4 + 2);
            let salt = io.read_be_u16();
            AuthCryptPassword(salt)
        }
        5 => {
            assert!(rem_len >= 4 + 4);
            let salt = io.read_be_u32();
            AuthMD5Password(salt)
        }
        6 => AuthSCMCredential,
        _ => AuthUnknown
    }
}

fn parse_message_header(io: &mut BufferedStream<TcpStream>) -> (u8, i32) {
    let typ = io.read_u8();
    let len = io.read_be_i32();
    return (typ, len);
}

fn main() {
    let addr = from_str::<SocketAddr>("127.0.0.1:5432").unwrap();
    let tcp_stream = TcpStream::connect(addr).unwrap();
    let mut io = BufferedStream::new(tcp_stream);
    write_startup_message(&mut io, [("user", "mneumann"), ("database", "test")]);

    let (typ, len) = parse_message_header(&mut io);
    println!("Typ: {:u} / len: {:d}", typ, len);

    if typ == 82 /* 'R' */ {
        let authmsg = parse_auth_message(&mut io, len-4);
        println!("{:?}", authmsg);
    }
    else {
        fail!();
    }

    write_query_message(&mut io, "select * from articles;");

    let (typ, len) = parse_message_header(&mut io);
    println!("Typ: {:u} / len: {:d}", typ, len);
}
<commit_msg>Define Message types<commit_after>#[feature(struct_variant)];

use std::io::buffered::BufferedStream;
use std::io::net::ip::SocketAddr;
use std::io::net::tcp::TcpStream;

static PROTO_VERSION: i32 = 196608; // 3 << 16

#[deriving(ToStr)]
enum AuthType {
    AuthUnknown,
    AuthOk,
    AuthKerberosV4,
    AuthKerberosV5,
    AuthClearTextPassword,
    AuthCryptPassword(u16),
    AuthMD5Password(u32),
    AuthSCMCredential
}

struct FieldInfo {
    name: ~str,
    oid: i32,
    attr_nr: i16,
    type_oid: i32,
    typlen: i16,
    atttypmod: i32,
    formatcode: i16
}

enum Message {
    MsgAuthentification(AuthType), // 'R'
    MsgPassword {password: ~str}, // 'p'
    MsgParameterStatus {key: ~str, val: ~str}, // 'S'
    MsgBackendKeyData {process_id: i32, secret_key: i32}, // 'K'
    MsgReadyForQuery {backend_transaction_status_indicator: u8}, // 'Z'
    MsgDataRow {columns: ~[Option<~[u8]>]}, // 'D'
    MsgCommandComplete {cmd_tag: ~str}, // 'C'
    MsgEmptyQueryResponse, // 'I'
    MsgNoticeResponse {field_type: u8, field_values: ~[~str]}, // 'N'
    MsgErrorResponse {field_type: u8, field_values: ~[~str]}, // 'E'
    MsgCopyInResponse, // 'G'
    MsgCopyOutResponse, // 'H'
    MsgParse {query: ~str, stmt_name: ~str, parameter_oids: ~[i32]}, // 'P'
    MsgParseComplete, // '1'
    MsgQuery {query: ~str}, // 'Q'
    MsgRowDescription {fields: ~[FieldInfo]}, // 'T'
    MsgTerminate, // 'X'
    MsgStartupMessage {proto_version: i32, params: ~[(~str, ~str)]},
    MsgSSLRequest {ssl_request_code: i32}
}

fn write_cstring(io: &mut BufferedStream<TcpStream>, str: &str) {
      io.write_str(str);
      io.write_u8(0);
}

fn write_startup_message(io: &mut BufferedStream<TcpStream>, params: &[(&str, &str)]) {
    let mut sz: uint = 8 + 1;

    for &(k, v) in params.iter() {
      sz += k.len() + 1 + v.len() + 1;
    }

    io.write_be_i32(sz as i32);
    io.write_be_i32(PROTO_VERSION);

    for &(k, v) in params.iter() {
        write_cstring(io, k);
        write_cstring(io, v);
    }
    io.write_u8(0);
    io.flush();
}

// body_size excludes the typ byte and the size i32
fn write_message_header(io: &mut BufferedStream<TcpStream>, typ: u8, payload: uint) {
    io.write_u8(typ);
    io.write_be_i32((payload + 4) as i32);
}

fn write_query_message(io: &mut BufferedStream<TcpStream>, query: &str) {
    write_message_header(io, 'Q' as u8, query.len() + 1);
    write_cstring(io, query);
    io.flush();
}

fn parse_auth_message(io: &mut BufferedStream<TcpStream>, rem_len: i32) -> AuthType {
    assert!(rem_len >= 4);
    let authtype = io.read_be_i32();
    match authtype {
        0 => AuthOk,
        1 => AuthKerberosV4,
        2 => AuthKerberosV5,
        3 => AuthClearTextPassword,
        4 => {
            assert!(rem_len >= 4 + 2);
            let salt = io.read_be_u16();
            AuthCryptPassword(salt)
        }
        5 => {
            assert!(rem_len >= 4 + 4);
            let salt = io.read_be_u32();
            AuthMD5Password(salt)
        }
        6 => AuthSCMCredential,
        _ => AuthUnknown
    }
}

fn parse_message_header(io: &mut BufferedStream<TcpStream>) -> (u8, i32) {
    let typ = io.read_u8();
    let len = io.read_be_i32();
    return (typ, len);
}

fn main() {
    let addr = from_str::<SocketAddr>("127.0.0.1:5432").unwrap();
    let tcp_stream = TcpStream::connect(addr).unwrap();
    let mut io = BufferedStream::new(tcp_stream);
    write_startup_message(&mut io, [("user", "mneumann"), ("database", "test")]);

    let (typ, len) = parse_message_header(&mut io);
    println!("Typ: {:u} / len: {:d}", typ, len);

    if typ == 82 /* 'R' */ {
        let authmsg = parse_auth_message(&mut io, len-4);
        println!("{:?}", authmsg);
    }
    else {
        fail!();
    }

    write_query_message(&mut io, "select * from articles;");

    let (typ, len) = parse_message_header(&mut io);
    println!("Typ: {:u} / len: {:d}", typ, len);
}
<|endoftext|>


<commit_before>use anyhow::Result;
use log::{debug, error, info};
#[macro_use]
extern crate clap;
#[macro_use]
extern crate anyhow;

use lancelot::{
    analysis::dis,
    aspace::AddressSpace,
    loader::pe::{load_pe, PE},
    util, RVA, VA,
};

fn handle_functions(pe: &PE) -> Result<()> {
    let functions = lancelot::analysis::pe::find_function_starts(pe)?;

    info!("found {} functions", functions.len());
    for va in functions.iter() {
        println!("{:#x}", va);
    }

    Ok(())
}

fn render_insn_buf(buf: &[u8], width: usize) -> String {
    let mut out = String::new();
    for (i, c) in hex::encode(buf).chars().enumerate() {
        out.push(c);

        if i % 2 == 1 {
            out.push(' ');
        }
    }

    let out = out.trim_end_matches(' ').to_string();
    if out.len() > width {
        out[..width].to_string()
    } else {
        out
    }
}

fn render_insn(insn: &zydis::ffi::DecodedInstruction, va: VA) -> String {
    let formatter = zydis::Formatter::new(zydis::FormatterStyle::INTEL).expect("formatter");
    let mut buffer = [0u8; 200];
    let mut buffer = zydis::OutputBuffer::new(&mut buffer[..]);

    formatter
        .format_instruction(&insn, &mut buffer, Some(va), None)
        .expect("format");
    format!("{}", buffer)
}

fn handle_disassemble(pe: &PE, va: VA) -> Result<()> {
    let cfg = lancelot::analysis::cfg::build_cfg(&pe.module, va)?;
    let decoder = dis::get_disassembler(&pe.module)?;

    info!("found {} basic blocks", cfg.basic_blocks.len());
    for bb in cfg.basic_blocks.values() {
        println!("{:#x}", bb.addr);

        // need to over-read the bb buffer, to account for the final instructions.
        let buf = pe.module.address_space.read_buf(bb.addr, bb.length as usize + 0x10)?;
        for (offset, insn) in dis::linear_disassemble(&decoder, &buf) {
            // because we over-read the bb buffer,
            // discard the instructions found after it.
            if offset >= bb.length as usize {
                break;
            }

            let va = bb.addr + offset as RVA;

            let name = &pe
                .module
                .sections
                .iter()
                .find(|sec| sec.virtual_range.contains(&va))
                .unwrap()
                .name;

            if let Ok(Some(insn)) = insn {
                let insn_buf = &buf[offset..offset + insn.length as usize];
                println!(
                    "  {}:{:#x}  {:15}  {}",
                    name,
                    va,
                    render_insn_buf(insn_buf, 15),
                    render_insn(&insn, va)
                );
            } else {
                println!("  {}:{:#x}: INVALID", name, va);
            }
        }
    }

    Ok(())
}

fn parse_va(s: &str) -> Result<VA> {
    if s.starts_with("0x") {
        let without_prefix = s.trim_start_matches("0x");
        Ok(u64::from_str_radix(without_prefix, 16)?)
    } else {
        Ok(s.parse()?)
    }
}

fn _main() -> Result<()> {
    better_panic::install();

    let matches = clap::clap_app!(lancelot =>
        (author: "Willi Ballenthin <william.ballenthin@mandiant.com>")
        (about: "Binary analysis framework")
        (@arg verbose: -v --verbose +multiple "log verbose messages")
        (@arg quiet: -q --quiet "disable informational messages")
        (@subcommand functions =>
            (about: "find functions")
            (@arg input: +required "path to file to analyze"))
        (@subcommand disassemble =>
            (about: "disassemble function")
            (@arg input: +required "path to file to analyze")
            (@arg va: +required "VA of function")))
    .get_matches();

    // --quiet overrides --verbose
    let log_level = if matches.is_present("quiet") {
        log::LevelFilter::Error
    } else {
        match matches.occurrences_of("verbose") {
            0 => log::LevelFilter::Info,
            1 => log::LevelFilter::Debug,
            2 => log::LevelFilter::Trace,
            _ => log::LevelFilter::Trace,
        }
    };

    fern::Dispatch::new()
        .format(move |out, message, record| {
            out.finish(format_args!(
                "{} [{:5}] {} {}",
                chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(),
                if log_level == log::LevelFilter::Trace {
                    record.target()
                } else {
                    ""
                },
                message
            ))
        })
        .level(log_level)
        .chain(std::io::stderr())
        .filter(|metadata| !metadata.target().starts_with("goblin::pe"))
        .apply()
        .expect("failed to configure logging");

    if let Some(matches) = matches.subcommand_matches("functions") {
        debug!("mode: find functions");

        let filename = matches.value_of("input").unwrap();
        debug!("input: {}", filename);

        let buf = util::read_file(filename)?;
        let pe = load_pe(&buf)?;

        handle_functions(&pe)
    } else if let Some(matches) = matches.subcommand_matches("disassemble") {
        debug!("mode: disassemble");

        let filename = matches.value_of("input").unwrap();
        debug!("input: {}", filename);

        let va = parse_va(matches.value_of("va").unwrap())?;

        let buf = util::read_file(filename)?;
        let pe = load_pe(&buf)?;

        handle_disassemble(&pe, va)
    } else {
        Err(anyhow!("SUBCOMMAND required"))
    }
}

fn main() {
    if let Err(e) = _main() {
        #[cfg(debug_assertions)]
        error!("{:?}", e);
        #[cfg(not(debug_assertions))]
        error!("{:}", e);
    }
}
<commit_msg>diassemble: simplify formatting<commit_after>use anyhow::Result;
use log::{debug, error, info};
#[macro_use]
extern crate clap;
#[macro_use]
extern crate anyhow;

use lancelot::{
    analysis::dis,
    aspace::AddressSpace,
    loader::pe::{load_pe, PE},
    util, RVA, VA,
};

fn handle_functions(pe: &PE) -> Result<()> {
    let functions = lancelot::analysis::pe::find_function_starts(pe)?;

    info!("found {} functions", functions.len());
    for va in functions.iter() {
        println!("{:#x}", va);
    }

    Ok(())
}

fn render_insn_buf(buf: &[u8], width: usize) -> String {
    let mut out = String::new();
    for (i, c) in hex::encode(buf).chars().enumerate() {
        out.push(c);

        if i % 2 == 1 {
            out.push(' ');
        }
    }

    let out = out.trim_end_matches(' ').to_string();
    if out.len() > width {
        out[..width].to_string()
    } else {
        out
    }
}

fn render_insn(insn: &zydis::ffi::DecodedInstruction, va: VA) -> String {
    let formatter = zydis::Formatter::new(zydis::FormatterStyle::INTEL).expect("formatter");
    let mut buffer = [0u8; 200];
    let mut buffer = zydis::OutputBuffer::new(&mut buffer[..]);

    formatter
        .format_instruction(&insn, &mut buffer, Some(va), None)
        .expect("format");
    format!("{}", buffer)
}

fn handle_disassemble(pe: &PE, va: VA) -> Result<()> {
    let cfg = lancelot::analysis::cfg::build_cfg(&pe.module, va)?;
    let decoder = dis::get_disassembler(&pe.module)?;

    info!("found {} basic blocks", cfg.basic_blocks.len());
    for bb in cfg.basic_blocks.values() {
        // need to over-read the bb buffer, to account for the final instructions.
        let buf = pe.module.address_space.read_buf(bb.addr, bb.length as usize + 0x10)?;
        for (offset, insn) in dis::linear_disassemble(&decoder, &buf) {
            // because we over-read the bb buffer,
            // discard the instructions found after it.
            if offset >= bb.length as usize {
                break;
            }

            let va = bb.addr + offset as RVA;

            let name = &pe
                .module
                .sections
                .iter()
                .find(|sec| sec.virtual_range.contains(&va))
                .unwrap()
                .name;

            if let Ok(Some(insn)) = insn {
                let insn_buf = &buf[offset..offset + insn.length as usize];
                println!(
                    "{}:{:016x}  {:15}  {}",
                    name,
                    va,
                    render_insn_buf(insn_buf, 15),
                    render_insn(&insn, va)
                );
            } else {
                println!("  {}:{:#x}: INVALID", name, va);
                break;
            }
        }
        println!();
    }

    Ok(())
}

fn parse_va(s: &str) -> Result<VA> {
    if s.starts_with("0x") {
        let without_prefix = s.trim_start_matches("0x");
        Ok(u64::from_str_radix(without_prefix, 16)?)
    } else {
        Ok(s.parse()?)
    }
}

fn _main() -> Result<()> {
    better_panic::install();

    let matches = clap::clap_app!(lancelot =>
        (author: "Willi Ballenthin <william.ballenthin@mandiant.com>")
        (about: "Binary analysis framework")
        (@arg verbose: -v --verbose +multiple "log verbose messages")
        (@arg quiet: -q --quiet "disable informational messages")
        (@subcommand functions =>
            (about: "find functions")
            (@arg input: +required "path to file to analyze"))
        (@subcommand disassemble =>
            (about: "disassemble function")
            (@arg input: +required "path to file to analyze")
            (@arg va: +required "VA of function")))
    .get_matches();

    // --quiet overrides --verbose
    let log_level = if matches.is_present("quiet") {
        log::LevelFilter::Error
    } else {
        match matches.occurrences_of("verbose") {
            0 => log::LevelFilter::Info,
            1 => log::LevelFilter::Debug,
            2 => log::LevelFilter::Trace,
            _ => log::LevelFilter::Trace,
        }
    };

    fern::Dispatch::new()
        .format(move |out, message, record| {
            out.finish(format_args!(
                "{} [{:5}] {} {}",
                chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(),
                if log_level == log::LevelFilter::Trace {
                    record.target()
                } else {
                    ""
                },
                message
            ))
        })
        .level(log_level)
        .chain(std::io::stderr())
        .filter(|metadata| !metadata.target().starts_with("goblin::pe"))
        .apply()
        .expect("failed to configure logging");

    if let Some(matches) = matches.subcommand_matches("functions") {
        debug!("mode: find functions");

        let filename = matches.value_of("input").unwrap();
        debug!("input: {}", filename);

        let buf = util::read_file(filename)?;
        let pe = load_pe(&buf)?;

        handle_functions(&pe)
    } else if let Some(matches) = matches.subcommand_matches("disassemble") {
        debug!("mode: disassemble");

        let filename = matches.value_of("input").unwrap();
        debug!("input: {}", filename);

        let va = parse_va(matches.value_of("va").unwrap())?;

        let buf = util::read_file(filename)?;
        let pe = load_pe(&buf)?;

        handle_disassemble(&pe, va)
    } else {
        Err(anyhow!("SUBCOMMAND required"))
    }
}

fn main() {
    if let Err(e) = _main() {
        #[cfg(debug_assertions)]
        error!("{:?}", e);
        #[cfg(not(debug_assertions))]
        error!("{:}", e);
    }
}
<|endoftext|>