<commit_before>// Copyright 2013 The rust-url developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Punycode ([RFC 3492](http://tools.ietf.org/html/rfc3492)) implementation.
//!
//! Since Punycode fundamentally works on unicode code points,
//! `encode` and `decode` take and return slices and vectors of `char`.
//! `encode_str` and `decode_to_string` provide convenience wrappers
//! that convert from and to Rust’s UTF-8 based `str` and `String` types.

use std::char;
use std::u32;

// Bootstring parameters for Punycode
static BASE: u32 = 36;
static T_MIN: u32 = 1;
static T_MAX: u32 = 26;
static SKEW: u32 = 38;
static DAMP: u32 = 700;
static INITIAL_BIAS: u32 = 72;
static INITIAL_N: u32 = 0x80;
static DELIMITER: char = '-';

#[inline]
fn adapt(mut delta: u32, num_points: u32, first_time: bool) -> u32 {
    delta /= if first_time { DAMP } else { 2 };
    delta += delta / num_points;
    let mut k = 0;
    while delta > ((BASE - T_MIN) * T_MAX) / 2 {
        delta /= BASE - T_MIN;
        k += BASE;
    }
    k + (((BASE - T_MIN + 1) * delta) / (delta + SKEW))
}

/// Convert Punycode to an Unicode `String`.
///
/// This is a convenience wrapper around `decode`.
#[inline]
pub fn decode_to_string(input: &str) -> Option<String> {
    decode(input).map(|chars| chars.into_iter().collect())
}

/// Convert Punycode to Unicode.
///
/// Return None on malformed input or overflow.
/// Overflow can only happen on inputs that take more than
/// 63 encoded bytes, the DNS limit on domain name labels.
pub fn decode(input: &str) -> Option<Vec<char>> {
    // Handle "basic" (ASCII) code points.
    // They are encoded as-is before the last delimiter, if any.
    let (base, input) = match input.rfind(DELIMITER) {
        None => ("", input),
        Some(position) => (
            &input[..position],
            if position > 0 {
                &input[position + 1..]
            } else {
                input
            },
        ),
    };

    let mut length = base.len() as u32;
    let mut buf = Vec::new();
    let mut code_point = INITIAL_N;
    let mut bias = INITIAL_BIAS;
    let mut i = 0;
    let mut iter = input.bytes();
    loop {
        let previous_i = i;
        let mut weight = 1;
        let mut k = BASE;
        let mut byte = match iter.next() {
            None => break,
            Some(byte) => byte,
        };
        // Decode a generalized variable-length integer into delta,
        // which gets added to i.
        loop {
            let digit = match byte {
                byte @ b'0'..=b'9' => byte - b'0' + 26,
                byte @ b'A'..=b'Z' => byte - b'A',
                byte @ b'a'..=b'z' => byte - b'a',
                _ => return None,
            } as u32;
            if digit > (u32::MAX - i) / weight {
                return None; // Overflow
            }
            i += digit * weight;
            let t = if k <= bias {
                T_MIN
            } else if k >= bias + T_MAX {
                T_MAX
            } else {
                k - bias
            };
            if digit < t {
                break;
            }
            if weight > u32::MAX / (BASE - t) {
                return None; // Overflow
            }
            weight *= BASE - t;
            k += BASE;
            byte = match iter.next() {
                None => return None, // End of input before the end of this delta
                Some(byte) => byte,
            };
        }
        bias = adapt(i - previous_i, length + 1, previous_i == 0);
        if i / (length + 1) > u32::MAX - code_point {
            return None; // Overflow
        }
        // i was supposed to wrap around from length+1 to 0,
        // incrementing code_point each time.
        code_point += i / (length + 1);
        i %= length + 1;
        let c = match char::from_u32(code_point) {
            Some(c) => c,
            None => return None,
        };

        // Move earlier insertions farther out in the string
        for (idx, _) in &mut buf {
            if *idx >= i as usize {
                *idx += 1;
            }
        }
        buf.push((i as usize, c));
        length += 1;
        i += 1;
    }

    buf.sort_by_key(|(i, _)| -(*i as i32));
    let mut position = 0;
    let mut output = Vec::with_capacity(input.len());
    let mut next = buf.pop();
    let mut base = base.chars();
    loop {
        match next {
            Some((pos, c)) if pos == position => {
                output.push(c);
                next = buf.pop();
                position += 1;
                continue;
            }
            _ => {}
        }
        if let Some(c) = base.next() {
            position += 1;
            output.push(c);
        } else if next.is_none() {
            break;
        }
    }

    Some(output)
}

/// Convert an Unicode `str` to Punycode.
///
/// This is a convenience wrapper around `encode`.
#[inline]
pub fn encode_str(input: &str) -> Option<String> {
    let mut buf = String::with_capacity(input.len());
    encode_into(input.chars(), &mut buf).ok().map(|()| buf)
}

/// Convert Unicode to Punycode.
///
/// Return None on overflow, which can only happen on inputs that would take more than
/// 63 encoded bytes, the DNS limit on domain name labels.
pub fn encode(input: &[char]) -> Option<String> {
    let mut buf = String::with_capacity(input.len());
    encode_into(input.iter().copied(), &mut buf)
        .ok()
        .map(|()| buf)
}

fn encode_into<I>(input: I, output: &mut String) -> Result<(), ()>
where
    I: Iterator<Item = char> + Clone,
{
    // Handle "basic" (ASCII) code points. They are encoded as-is.
    let (mut input_length, mut basic_length) = (0, 0);
    for c in input.clone() {
        input_length += 1;
        if c.is_ascii() {
            output.push(c);
            basic_length += 1;
        }
    }

    if basic_length > 0 {
        output.push('-')
    }
    let mut code_point = INITIAL_N;
    let mut delta = 0;
    let mut bias = INITIAL_BIAS;
    let mut processed = basic_length;
    while processed < input_length {
        // All code points < code_point have been handled already.
        // Find the next larger one.
        let min_code_point = input
            .clone()
            .map(|c| c as u32)
            .filter(|&c| c >= code_point)
            .min()
            .unwrap();
        if min_code_point - code_point > (u32::MAX - delta) / (processed + 1) {
            return Err(()); // Overflow
        }
        // Increase delta to advance the decoder’s <code_point,i> state to <min_code_point,0>
        delta += (min_code_point - code_point) * (processed + 1);
        code_point = min_code_point;
        for c in input.clone() {
            let c = c as u32;
            if c < code_point {
                delta += 1;
                if delta == 0 {
                    return Err(()); // Overflow
                }
            }
            if c == code_point {
                // Represent delta as a generalized variable-length integer:
                let mut q = delta;
                let mut k = BASE;
                loop {
                    let t = if k <= bias {
                        T_MIN
                    } else if k >= bias + T_MAX {
                        T_MAX
                    } else {
                        k - bias
                    };
                    if q < t {
                        break;
                    }
                    let value = t + ((q - t) % (BASE - t));
                    output.push(value_to_digit(value));
                    q = (q - t) / (BASE - t);
                    k += BASE;
                }
                output.push(value_to_digit(q));
                bias = adapt(delta, processed + 1, processed == basic_length);
                delta = 0;
                processed += 1;
            }
        }
        delta += 1;
        code_point += 1;
    }
    Ok(())
}

#[inline]
fn value_to_digit(value: u32) -> char {
    match value {
        0..=25 => (value as u8 + b'a') as char,       // a..z
        26..=35 => (value as u8 - 26 + b'0') as char, // 0..9
        _ => panic!(),
    }
}
<commit_msg>idna: separate decoding of encoded characters and merging<commit_after>// Copyright 2013 The rust-url developers.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! Punycode ([RFC 3492](http://tools.ietf.org/html/rfc3492)) implementation.
//!
//! Since Punycode fundamentally works on unicode code points,
//! `encode` and `decode` take and return slices and vectors of `char`.
//! `encode_str` and `decode_to_string` provide convenience wrappers
//! that convert from and to Rust’s UTF-8 based `str` and `String` types.

use std::char;
use std::u32;

// Bootstring parameters for Punycode
static BASE: u32 = 36;
static T_MIN: u32 = 1;
static T_MAX: u32 = 26;
static SKEW: u32 = 38;
static DAMP: u32 = 700;
static INITIAL_BIAS: u32 = 72;
static INITIAL_N: u32 = 0x80;
static DELIMITER: char = '-';

#[inline]
fn adapt(mut delta: u32, num_points: u32, first_time: bool) -> u32 {
    delta /= if first_time { DAMP } else { 2 };
    delta += delta / num_points;
    let mut k = 0;
    while delta > ((BASE - T_MIN) * T_MAX) / 2 {
        delta /= BASE - T_MIN;
        k += BASE;
    }
    k + (((BASE - T_MIN + 1) * delta) / (delta + SKEW))
}

/// Convert Punycode to an Unicode `String`.
///
/// This is a convenience wrapper around `decode`.
#[inline]
pub fn decode_to_string(input: &str) -> Option<String> {
    decode(input).map(|chars| chars.into_iter().collect())
}

/// Convert Punycode to Unicode.
///
/// Return None on malformed input or overflow.
/// Overflow can only happen on inputs that take more than
/// 63 encoded bytes, the DNS limit on domain name labels.
pub fn decode(input: &str) -> Option<Vec<char>> {
    let (base, mut buf) = insertions(input).ok()?;
    Some(merge(base, &mut buf))
}

/// Split the input iterator and return a Vec with insertions of decoded characters
fn insertions<'a>(input: &'a str) -> Result<(&'a str, Vec<(usize, char)>), ()> {
    // Handle "basic" (ASCII) code points.
    // They are encoded as-is before the last delimiter, if any.
    let (base, input) = match input.rfind(DELIMITER) {
        None => ("", input),
        Some(position) => (
            &input[..position],
            if position > 0 {
                &input[position + 1..]
            } else {
                input
            },
        ),
    };

    let mut length = base.len() as u32;
    let mut buf = Vec::new();
    let mut code_point = INITIAL_N;
    let mut bias = INITIAL_BIAS;
    let mut i = 0;
    let mut iter = input.bytes();
    loop {
        let previous_i = i;
        let mut weight = 1;
        let mut k = BASE;
        let mut byte = match iter.next() {
            None => break,
            Some(byte) => byte,
        };
        // Decode a generalized variable-length integer into delta,
        // which gets added to i.
        loop {
            let digit = match byte {
                byte @ b'0'..=b'9' => byte - b'0' + 26,
                byte @ b'A'..=b'Z' => byte - b'A',
                byte @ b'a'..=b'z' => byte - b'a',
                _ => return Err(()),
            } as u32;
            if digit > (u32::MAX - i) / weight {
                return Err(()); // Overflow
            }
            i += digit * weight;
            let t = if k <= bias {
                T_MIN
            } else if k >= bias + T_MAX {
                T_MAX
            } else {
                k - bias
            };
            if digit < t {
                break;
            }
            if weight > u32::MAX / (BASE - t) {
                return Err(()); // Overflow
            }
            weight *= BASE - t;
            k += BASE;
            byte = match iter.next() {
                None => return Err(()), // End of input before the end of this delta
                Some(byte) => byte,
            };
        }
        bias = adapt(i - previous_i, length + 1, previous_i == 0);
        if i / (length + 1) > u32::MAX - code_point {
            return Err(()); // Overflow
        }
        // i was supposed to wrap around from length+1 to 0,
        // incrementing code_point each time.
        code_point += i / (length + 1);
        i %= length + 1;
        let c = match char::from_u32(code_point) {
            Some(c) => c,
            None => return Err(()),
        };

        // Move earlier insertions farther out in the string
        for (idx, _) in &mut buf {
            if *idx >= i as usize {
                *idx += 1;
            }
        }
        buf.push((i as usize, c));
        length += 1;
        i += 1;
    }

    buf.sort_by_key(|(i, _)| *i);
    Ok((base, buf))
}

/// Merge base character iterator and decoded character insertions
fn merge(input: &str, insertions: &[(usize, char)]) -> Vec<char> {
    let mut insertions = insertions.iter();
    let mut position = 0;
    let mut output = Vec::with_capacity(input.len());
    let mut next = insertions.next();
    let mut base = input.chars();

    loop {
        match next {
            Some((pos, c)) if *pos == position => {
                output.push(*c);
                next = insertions.next();
                position += 1;
                continue;
            }
            _ => {}
        }
        if let Some(c) = base.next() {
            position += 1;
            output.push(c);
        } else if next.is_none() {
            break;
        }
    }

    output
}

/// Convert an Unicode `str` to Punycode.
///
/// This is a convenience wrapper around `encode`.
#[inline]
pub fn encode_str(input: &str) -> Option<String> {
    let mut buf = String::with_capacity(input.len());
    encode_into(input.chars(), &mut buf).ok().map(|()| buf)
}

/// Convert Unicode to Punycode.
///
/// Return None on overflow, which can only happen on inputs that would take more than
/// 63 encoded bytes, the DNS limit on domain name labels.
pub fn encode(input: &[char]) -> Option<String> {
    let mut buf = String::with_capacity(input.len());
    encode_into(input.iter().copied(), &mut buf)
        .ok()
        .map(|()| buf)
}

fn encode_into<I>(input: I, output: &mut String) -> Result<(), ()>
where
    I: Iterator<Item = char> + Clone,
{
    // Handle "basic" (ASCII) code points. They are encoded as-is.
    let (mut input_length, mut basic_length) = (0, 0);
    for c in input.clone() {
        input_length += 1;
        if c.is_ascii() {
            output.push(c);
            basic_length += 1;
        }
    }

    if basic_length > 0 {
        output.push('-')
    }
    let mut code_point = INITIAL_N;
    let mut delta = 0;
    let mut bias = INITIAL_BIAS;
    let mut processed = basic_length;
    while processed < input_length {
        // All code points < code_point have been handled already.
        // Find the next larger one.
        let min_code_point = input
            .clone()
            .map(|c| c as u32)
            .filter(|&c| c >= code_point)
            .min()
            .unwrap();
        if min_code_point - code_point > (u32::MAX - delta) / (processed + 1) {
            return Err(()); // Overflow
        }
        // Increase delta to advance the decoder’s <code_point,i> state to <min_code_point,0>
        delta += (min_code_point - code_point) * (processed + 1);
        code_point = min_code_point;
        for c in input.clone() {
            let c = c as u32;
            if c < code_point {
                delta += 1;
                if delta == 0 {
                    return Err(()); // Overflow
                }
            }
            if c == code_point {
                // Represent delta as a generalized variable-length integer:
                let mut q = delta;
                let mut k = BASE;
                loop {
                    let t = if k <= bias {
                        T_MIN
                    } else if k >= bias + T_MAX {
                        T_MAX
                    } else {
                        k - bias
                    };
                    if q < t {
                        break;
                    }
                    let value = t + ((q - t) % (BASE - t));
                    output.push(value_to_digit(value));
                    q = (q - t) / (BASE - t);
                    k += BASE;
                }
                output.push(value_to_digit(q));
                bias = adapt(delta, processed + 1, processed == basic_length);
                delta = 0;
                processed += 1;
            }
        }
        delta += 1;
        code_point += 1;
    }
    Ok(())
}

#[inline]
fn value_to_digit(value: u32) -> char {
    match value {
        0..=25 => (value as u8 + b'a') as char,       // a..z
        26..=35 => (value as u8 - 26 + b'0') as char, // 0..9
        _ => panic!(),
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

use std::ops;
use std::str;

use bytes::Bytes;
use std::fmt;

#[derive(Debug)]
pub struct AsciiError(());

/// Bytes wrapper that guarantees that contained bytes are ASCII.
#[derive(Eq, PartialEq, Hash, Clone)]
pub struct Ascii(Bytes);

impl Ascii {
    pub fn new() -> Ascii {
        Ascii(Bytes::new())
    }

    pub fn from_bytes(bs: Bytes) -> Result<Ascii, (AsciiError, Bytes)> {
        for &b in &bs {
            if b > i8::max_value() as u8 {
                return Err((AsciiError(()), bs));
            }
        }
        Ok(Ascii(bs))
    }

    /// Create an ASCII string from bytes.
    ///
    /// This is a const fn, and we cannot loop in const fn
    /// so no validation here.
    pub const unsafe fn from_bytes_unchecked(bs: Bytes) -> Ascii {
        Ascii(bs)
    }

    pub fn as_str(&self) -> &str {
        unsafe { str::from_utf8_unchecked(self.0.as_ref()) }
    }

    pub fn into_bytes(self) -> Bytes {
        self.0
    }
}

impl ops::Deref for Ascii {
    type Target = str;

    fn deref(&self) -> &str {
        unsafe { str::from_utf8_unchecked(&self.0) }
    }
}

impl AsRef<[u8]> for Ascii {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

impl AsRef<str> for Ascii {
    fn as_ref(&self) -> &str {
        self.as_str()
    }
}

impl Into<Bytes> for Ascii {
    fn into(self) -> Bytes {
        self.0
    }
}

impl fmt::Debug for Ascii {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(self.as_str(), f)
    }
}

impl fmt::Display for Ascii {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self.as_str(), f)
    }
}

#[derive(Debug)]
pub struct AsciiLowerError(());

/// ASCII lowercase string (any ASCII except uppercase letters).
pub struct AsciiLower(Ascii);

impl AsciiLower {
    pub fn new() -> AsciiLower {
        AsciiLower(Ascii::new())
    }

    /// Validate bytes is an ASCII lowercase string.
    pub fn from_bytes(bs: Bytes) -> Result<AsciiLower, (AsciiLowerError, Bytes)> {
        for &b in &bs {
            if !b.is_ascii() || b.is_ascii_uppercase() {
                return Err((AsciiLowerError(()), bs));
            }
        }
        unsafe { Ok(AsciiLower(Ascii::from_bytes_unchecked(bs))) }
    }

    pub const unsafe fn from_bytes_unchecked(bs: Bytes) -> AsciiLower {
        AsciiLower(Ascii::from_bytes_unchecked(bs))
    }
}

impl fmt::Display for AsciiLower {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

impl ops::Deref for AsciiLower {
    type Target = str;

    fn deref(&self) -> &str {
        self.0.deref()
    }
}

impl<'a> From<&'a str> for AsciiLower {
    fn from(_s: &'a str) -> AsciiLower {
        unimplemented!()
    }
}

#[cfg(test)]
mod test {
    use crate::ascii::Ascii;
    use crate::ascii::AsciiLower;
    use bytes::Bytes;

    #[test]
    fn ascii_from_bytes() {
        assert!(Ascii::from_bytes(Bytes::from_static("ю".as_bytes())).is_err());
    }

    #[test]
    fn ascii_lower_from_bytes() {
        assert!(AsciiLower::from_bytes(Bytes::from_static("qю".as_bytes())).is_err());
        assert!(AsciiLower::from_bytes(Bytes::from_static("qQ".as_bytes())).is_err());
        assert!(AsciiLower::from_bytes(Bytes::from_static("q".as_bytes())).is_ok());
        assert!(AsciiLower::from_bytes(Bytes::from_static("-".as_bytes())).is_ok());
        assert!(AsciiLower::from_bytes(Bytes::from_static("".as_bytes())).is_ok());
    }

    #[test]
    fn ascii_lower_display() {
        assert_eq!(
            "foo-bar",
            format!(
                "{}",
                AsciiLower::from_bytes(Bytes::from_static(b"foo-bar")).unwrap()
            ),
        );
    }
}
<commit_msg>Use non-deprecated function<commit_after>#![allow(dead_code)]

use std::ops;
use std::str;

use bytes::Bytes;
use std::fmt;

#[derive(Debug)]
pub struct AsciiError(());

/// Bytes wrapper that guarantees that contained bytes are ASCII.
#[derive(Eq, PartialEq, Hash, Clone)]
pub struct Ascii(Bytes);

impl Ascii {
    pub fn new() -> Ascii {
        Ascii(Bytes::new())
    }

    pub fn from_bytes(bs: Bytes) -> Result<Ascii, (AsciiError, Bytes)> {
        for &b in &bs {
            if b > i8::MAX as u8 {
                return Err((AsciiError(()), bs));
            }
        }
        Ok(Ascii(bs))
    }

    /// Create an ASCII string from bytes.
    ///
    /// This is a const fn, and we cannot loop in const fn
    /// so no validation here.
    pub const unsafe fn from_bytes_unchecked(bs: Bytes) -> Ascii {
        Ascii(bs)
    }

    pub fn as_str(&self) -> &str {
        unsafe { str::from_utf8_unchecked(self.0.as_ref()) }
    }

    pub fn into_bytes(self) -> Bytes {
        self.0
    }
}

impl ops::Deref for Ascii {
    type Target = str;

    fn deref(&self) -> &str {
        unsafe { str::from_utf8_unchecked(&self.0) }
    }
}

impl AsRef<[u8]> for Ascii {
    fn as_ref(&self) -> &[u8] {
        self.0.as_ref()
    }
}

impl AsRef<str> for Ascii {
    fn as_ref(&self) -> &str {
        self.as_str()
    }
}

impl Into<Bytes> for Ascii {
    fn into(self) -> Bytes {
        self.0
    }
}

impl fmt::Debug for Ascii {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt::Debug::fmt(self.as_str(), f)
    }
}

impl fmt::Display for Ascii {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self.as_str(), f)
    }
}

#[derive(Debug)]
pub struct AsciiLowerError(());

/// ASCII lowercase string (any ASCII except uppercase letters).
pub struct AsciiLower(Ascii);

impl AsciiLower {
    pub fn new() -> AsciiLower {
        AsciiLower(Ascii::new())
    }

    /// Validate bytes is an ASCII lowercase string.
    pub fn from_bytes(bs: Bytes) -> Result<AsciiLower, (AsciiLowerError, Bytes)> {
        for &b in &bs {
            if !b.is_ascii() || b.is_ascii_uppercase() {
                return Err((AsciiLowerError(()), bs));
            }
        }
        unsafe { Ok(AsciiLower(Ascii::from_bytes_unchecked(bs))) }
    }

    pub const unsafe fn from_bytes_unchecked(bs: Bytes) -> AsciiLower {
        AsciiLower(Ascii::from_bytes_unchecked(bs))
    }
}

impl fmt::Display for AsciiLower {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.0, f)
    }
}

impl ops::Deref for AsciiLower {
    type Target = str;

    fn deref(&self) -> &str {
        self.0.deref()
    }
}

impl<'a> From<&'a str> for AsciiLower {
    fn from(_s: &'a str) -> AsciiLower {
        unimplemented!()
    }
}

#[cfg(test)]
mod test {
    use crate::ascii::Ascii;
    use crate::ascii::AsciiLower;
    use bytes::Bytes;

    #[test]
    fn ascii_from_bytes() {
        assert!(Ascii::from_bytes(Bytes::from_static("ю".as_bytes())).is_err());
    }

    #[test]
    fn ascii_lower_from_bytes() {
        assert!(AsciiLower::from_bytes(Bytes::from_static("qю".as_bytes())).is_err());
        assert!(AsciiLower::from_bytes(Bytes::from_static("qQ".as_bytes())).is_err());
        assert!(AsciiLower::from_bytes(Bytes::from_static("q".as_bytes())).is_ok());
        assert!(AsciiLower::from_bytes(Bytes::from_static("-".as_bytes())).is_ok());
        assert!(AsciiLower::from_bytes(Bytes::from_static("".as_bytes())).is_ok());
    }

    #[test]
    fn ascii_lower_display() {
        assert_eq!(
            "foo-bar",
            format!(
                "{}",
                AsciiLower::from_bytes(Bytes::from_static(b"foo-bar")).unwrap()
            ),
        );
    }
}
<|endoftext|>


<commit_before>//! A verify minimal representation of payload information

#[derive(Debug, RustcEncodable, RustcDecodable, Default)]
pub struct Payload {
  pub sender: Sender,
  pub organization: Option<Organization>,
  pub repository: Option<Repository>
}

#[derive(Debug, Default, RustcEncodable, RustcDecodable)]
pub struct Sender {
  pub login: String
}

#[derive(Debug, Default, RustcEncodable, RustcDecodable)]
pub struct Organization {
  pub login: String
}

#[derive(Debug, Default, RustcEncodable, RustcDecodable)]
pub struct Repository {
  pub name: String,
  pub full_name: String,
  pub organization: Organization
}<commit_msg>doc typo<commit_after>//! A minimal representation of payload information

#[derive(Debug, RustcEncodable, RustcDecodable, Default)]
pub struct Payload {
  pub sender: Sender,
  pub organization: Option<Organization>,
  pub repository: Option<Repository>
}

#[derive(Debug, Default, RustcEncodable, RustcDecodable)]
pub struct Sender {
  pub login: String
}

#[derive(Debug, Default, RustcEncodable, RustcDecodable)]
pub struct Organization {
  pub login: String
}

#[derive(Debug, Default, RustcEncodable, RustcDecodable)]
pub struct Repository {
  pub name: String,
  pub full_name: String,
  pub organization: Organization
}
<|endoftext|>


<commit_before>static MAGIC_NUMBERS : [u8; 4] = [0x4Eu8, 0x45u8, 0x53u8, 0x1Au8];

pub struct Rom {
	prg_rom_pages: u8,
	chr_rom_pages: u8
}

impl Rom {
	pub fn parse(data: Vec<u8>) -> Rom {
		assert_eq!( &data[0 .. 4], MAGIC_NUMBERS );
		let prg_pages = data[4];
		let chr_pages = data[5];
		Rom {
			prg_rom_pages: prg_pages,
			chr_rom_pages: chr_pages
		}
	}
}

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

	fn create_binary_rom( prg_page_count: u8, chr_page_count: u8 ) -> Vec<u8> {
		let mut bin_rom = vec![0x4Eu8, 0x45u8, 0x53u8, 0x1Au8];
		bin_rom.push( prg_page_count );
		bin_rom.push( chr_page_count );
		bin_rom
	}

	#[test]
	#[should_panic]
	fn parse_panics_on_empty_input() {
		Rom::parse( vec!() );
	}

	#[test]
	fn parse_succeeds_on_input_with_magic_numbers() {
		Rom::parse( create_binary_rom( 0, 0 ) );
	}

	#[test]
	fn parse_reads_prg_rom_page_count() {
		let rom = Rom::parse( create_binary_rom( 12, 0 ) );
		assert_eq!( rom.prg_rom_pages, 12 );
	}

	#[test]
	fn parse_reads_chr_rom_page_count() {
		let rom = Rom::parse( create_binary_rom( 12, 15 ) );
		assert_eq!( rom.chr_rom_pages, 15 );
	}
}<commit_msg>Refactor and implement parsing more flags.<commit_after>static MAGIC_NUMBERS : [u8; 4] = [0x4Eu8, 0x45u8, 0x53u8, 0x1Au8];

#[derive(PartialEq, Debug)]
pub enum ScreenMode {
	Horizontal,
	Vertical,
	FourScreen
}

pub struct Rom {
	prg_pages: u8,
	chr_pages: u8,
	flags6: u8,
	flags7: u8,
}

impl Rom {
	pub fn parse(data: Vec<u8>) -> Rom {
		assert_eq!( &data[0 .. 4], MAGIC_NUMBERS );
		let prg_pages = data[4];
		let chr_pages = data[5];
		let flags6 = data[6];
		let flags7 = data[7];
		Rom {
			prg_pages: prg_pages,
			chr_pages: chr_pages,
			flags6: flags6,
			flags7: flags7
		}
	}

	pub fn screen_mode(&self) -> ScreenMode {
		match self.flags6 & 0b0000_1001u8 {
			0b0000_0000 => ScreenMode::Horizontal,
			0b0000_0001 => ScreenMode::Vertical,
			0b0000_1000 | 0b0000_1001 => ScreenMode::FourScreen,
			_ => panic!("Math is broken!"),
		}
	}

	fn get_bit(&self, byte: u8, bit_num: u8) -> bool {
		!((byte & 1u8 << bit_num) == 0)
	}

	pub fn sram(&self) -> bool {
		self.get_bit(self.flags6, 1)
	}

	pub fn trainer(&self) -> bool {
		self.get_bit(self.flags6, 2)
	}

	pub fn pc10(&self) -> bool {
		self.get_bit(self.flags7, 0)
	}

	pub fn vs(&self) -> bool {
		self.get_bit(self.flags7, 1)
	}
}

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

	struct RomBuilder {
		header : Vec<u8>
	}

	fn set_bit(byte: &mut u8, bit_num: u8) {
		*byte = *byte | 1u8 << bit_num;
	}

	impl RomBuilder {
		fn new() -> RomBuilder {
			let mut header = vec![0x4Eu8, 0x45u8, 0x53u8, 0x1Au8];
			header.extend([0; 12].iter().cloned());
			return RomBuilder{ header: header }
		}

		fn set_prg_page_count( &mut self, count: u8 ) {
			self.header[4] = count;
		}

		fn set_chr_page_count( &mut self, count: u8 ) {
			self.header[5] = count;
		}

		fn set_mirroring( &mut self ) {
			set_bit(&mut self.header[6], 0)
		}

		fn set_sram( &mut self ) {
			set_bit(&mut self.header[6], 1)
		}

		fn set_trainer( &mut self ) {
			set_bit(&mut self.header[6], 2)
		}

		fn set_fourscreen( &mut self ) {
			set_bit(&mut self.header[6], 3)
		}

		fn set_pc10( &mut self ) {
			set_bit(&mut self.header[7], 0)
		}

		fn set_vs( &mut self ) {
			set_bit(&mut self.header[7], 1)
		}

		fn build(&self) -> Vec<u8> {
			self.header.to_vec()
		}
	}

	#[test]
	#[should_panic]
	fn parse_panics_on_empty_input() {
		Rom::parse( vec!() );
	}

	#[test]
	fn parse_reads_prg_pages() {
		let mut builder = RomBuilder::new();
		assert_eq!( Rom::parse( builder.build() ).prg_pages, 0 );

		builder.set_prg_page_count( 150 );
		assert_eq!( Rom::parse( builder.build() ).prg_pages, 150 );
	}

	#[test]
	fn parse_reads_chr_pages() {
		let mut builder = RomBuilder::new();
		assert_eq!( Rom::parse( builder.build() ).chr_pages, 0 );

		builder.set_chr_page_count( 150 );
		assert_eq!( Rom::parse( builder.build() ).chr_pages, 150 );
	}

	#[test]
	fn test_screen_mode_without_fourscreen() {
		let mut builder = RomBuilder::new();
		assert_eq!( Rom::parse( builder.build() ).screen_mode(), ScreenMode::Horizontal );

		builder.set_mirroring();
		assert_eq!( Rom::parse( builder.build() ).screen_mode(), ScreenMode::Vertical );
	}

	#[test]
	fn test_screen_mode_with_fourscreen() {
		let mut builder = RomBuilder::new();
		builder.set_fourscreen();
		assert_eq!( Rom::parse( builder.build() ).screen_mode(), ScreenMode::FourScreen );

		builder.set_mirroring();
		assert_eq!( Rom::parse( builder.build() ).screen_mode(), ScreenMode::FourScreen );
	}

	#[test]
	fn test_sram() {
		let mut builder = RomBuilder::new();
		assert_eq!( Rom::parse( builder.build() ).sram(), false );

		builder.set_sram();
		assert_eq!( Rom::parse( builder.build() ).sram(), true );
	}

	#[test]
	fn test_trainer() {
		let mut builder = RomBuilder::new();
		assert_eq!( Rom::parse( builder.build() ).trainer(), false );

		builder.set_trainer();
		assert_eq!( Rom::parse( builder.build() ).trainer(), true );
	}

	#[test]
	fn test_pc10() {
		let mut builder = RomBuilder::new();
		assert_eq!( Rom::parse( builder.build() ).pc10(), false );

		builder.set_pc10();
		assert_eq!( Rom::parse( builder.build() ).pc10(), true );
	}

	#[test]
	fn test_vs() {
		let mut builder = RomBuilder::new();
		assert_eq!( Rom::parse( builder.build() ).vs(), false );

		builder.set_vs();
		assert_eq!( Rom::parse( builder.build() ).vs(), true );
	}
}<|endoftext|>


<commit_before>/*
 * Copyright (c) 2016 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

use quote::Tokens;
use syn::{Attribute, Body, Ident, MacroInput, Path, VariantData};
use syn::Body::{Enum, Struct};
use syn::MetaItem::{List, Word};
use syn::Ty;

use attributes::to_metadata_impl;
use string::{snake_to_camel, to_dash_name};

/// Expand the required trais for the derive Setting attribute.
pub fn expand_setting_enum(ast: MacroInput) -> Tokens {
    let name = ast.ident.clone();
    let mut default = None;

    let mut variant_names = vec![];
    if let Enum(ref variants) = ast.body {
        for variant in variants {
            variant_names.push(variant.ident.clone());
            if !variant.attrs.is_empty() {
                for attribute in &variant.attrs {
                    if let &Attribute { value: Word(ref ident), .. } = attribute {
                        if ident == "default" {
                            default = Some(variant.ident.clone());
                        }
                    }
                }
            }
        }
    }
    let choice_names: Vec<_> = variant_names.iter()
        .map(|name| to_dash_name(&name.to_string()))
        .collect();
    let choice_names1 = &choice_names;
    let choice_names2 = &choice_names;

    let qualified_names = variant_names.iter()
        .map(|variant_name| quote! {
            #name::#variant_name
        });

    let from_str_fn = quote! {
        fn from_str(string: &str) -> Result<Self, Self::Err> {
            match string {
                #(#choice_names1 => Ok(#qualified_names),)*
                _ => Err(::mg_settings::error::SettingError::UnknownChoice {
                    actual: string.to_string(),
                    expected: vec![#(#choice_names2),*],
                }),
            }
        }
    };

    let default_impl =
        if let Some(ident) = default {
            quote! {
                impl Default for #name {
                    fn default() -> Self {
                        #name::#ident
                    }
                }
            }
        }
        else {
            quote! {
            }
        };

    let completion_values_impl = quote! {
        impl CompletionValues for #name {
            fn completion_values() -> Vec<String> {
                vec![#(#choice_names1.to_string()),*]
            }
        }
    };

    quote! {
        #default_impl

        #completion_values_impl

        impl ::std::str::FromStr for #name {
            type Err = ::mg_settings::error::SettingError;

            #from_str_fn
        }
    }
}

/// Expand the required traits for the derive Settings attribute.
pub fn expand_settings_enum(ast: MacroInput) -> Tokens {
    let name = &ast.ident;
    let completion_fn = to_setting_completion_fn(name, &ast.body);
    let variant_name = Ident::new(format!("{}Variant", name));
    let variant_enum = to_enums(&variant_name, &ast.body);
    let settings_impl = to_settings_impl(name, &variant_name, &ast.body);
    let (metadata_impl, _) = to_metadata_impl(name, &ast.body);
    quote! {
        #variant_enum

        #settings_impl

        #metadata_impl

        #completion_fn
    }
}

/// Check if a type is a custom type (including enum).
fn is_custom_type(ident: &Ident) -> bool {
    match ident.to_string().as_ref() {
        "bool" | "f64" | "i64" | "String" => false,
        _ => true,
    }
}

/// Create the variant enums for getters and setters.
fn to_enums(variant_name: &Ident, settings_struct: &Body) -> Tokens {
    if let &Struct(VariantData::Struct(ref strct)) = settings_struct {
        let mut field_names = vec![];
        let mut names = vec![];
        let mut types = vec![];
        for field in strct {
            if let Some(ref ident) = field.ident {
                field_names.push(ident);
                let ident = Ident::new(snake_to_camel(&ident.to_string()));
                names.push(ident);
                types.push(field.ty.clone());
            }
        }
        let names1 = &names;
        quote! {
            pub enum #variant_name {
                #(#names1(#types)),*
            }
        }
    }
    else {
        panic!("Not a struct");
    }
}

/// Create the impl Settings.
fn to_settings_impl(name: &Ident, variant_name: &Ident, settings_struct: &Body) -> Tokens {
    if let &Struct(VariantData::Struct(ref strct)) = settings_struct {
        let mut names = vec![];
        let mut capitalized_names = vec![];
        let mut original_types = vec![];
        let mut types = vec![];
        for field in strct {
            if let Some(ref ident) = field.ident {
                let ident_string = ident.to_string();
                let ident = Ident::new(ident_string.clone());
                names.push(ident);
                let ident = Ident::new(snake_to_camel(&ident_string));
                capitalized_names.push(
                    quote! {
                        #variant_name::#ident
                    });

                if let Ty::Path(_, Path { ref segments, .. }) = field.ty {
                    original_types.push(&segments[0].ident);
                    types.push(to_value_type(&segments[0].ident));
                }
            }
        }
        let string_names: Vec<_> = names.iter()
            .map(|ident| ident.to_string().replace("_", "-"))
            .collect();
        let string_names = &string_names;
        let capitalized_names = &capitalized_names;
        let names1 = &names;
        let names2 = &names;
        let names3 = &names;
        let variant_exprs = names.iter().zip(original_types.iter())
            .map(|(name, typ)|
                 if is_custom_type(typ) {
                     quote! {
                         ::std::str::FromStr::from_str(&#name)?
                     }
                 }
                 else {
                     quote! { #name }
                 }
            );
        let types1 = &types;
        let type_names = types.iter()
            .map(|ident| value_type_to_type(&ident));

        let to_variant_fn_variant = quote! {
            #(#string_names => {
                if let ::mg_settings::Value::#types1(#names1) = value {
                    Ok(#capitalized_names(#variant_exprs))
                }
                else {
                    Err(::mg_settings::error::SettingError::WrongType {
                        actual: value.to_type().to_string(),
                        expected: #type_names.to_string(),
                    })
                }
            },)*
        };

        let to_variant_fn = quote! {
            fn to_variant(name: &str, value: ::mg_settings::Value) -> Result<Self::Variant, ::mg_settings::error::SettingError> {
                match name {
                    #to_variant_fn_variant
                    _ => Err(::mg_settings::error::SettingError::UnknownSetting(name.to_string())),
                }
            }
        };

        quote! {
            impl ::mg_settings::settings::Settings for #name {
                type Variant = #variant_name;

                #to_variant_fn

                fn set_value(&mut self, value: Self::Variant) {
                    match value {
                        #(#capitalized_names(#names1) => {
                            self.#names2 = #names3
                        }),*
                    }
                }
            }
        }
    }
    else {
        panic!("Not a struct");
    }
}

/// Create the function returing the completion of the setting values.
pub fn to_setting_completion_fn(name: &Ident, body: &Body) -> Tokens {
    let mut completions = vec![];
    if let Struct(VariantData::Struct(ref fields)) = *body {
        'field_loop:
        for field in fields {
            for attribute in &field.attrs {
                if let &Attribute { value: List(ref ident, ref args), .. } = attribute {
                    if ident == "completion" {
                        if let Word(ref arg_ident) = args[0] {
                            if arg_ident == "hidden" {
                                continue 'field_loop;
                            }
                        }
                    }
                }
            }

            let setting_name = field.ident.as_ref().unwrap().to_string().replace('_', "-");
            let field_type = &field.ty;

            completions.push(quote! {
                (#setting_name.to_string(), #field_type::completion_values())
            });
        }
    }

    quote! {
        use mg_settings::CompletionValues;

        impl ::mg_settings::SettingCompletion for #name {
            fn get_value_completions() -> ::std::collections::HashMap<String, Vec<String>> {
                let mut vec = vec![#(#completions),*];
                let iter = vec.drain(..);
                iter.collect()
            }
        }
    }
}

/// Convert a Rust type to a `Value` type.
fn to_value_type(ident: &Ident) -> Ident {
    let value_type =
        match ident.to_string().as_ref() {
            "bool" => "Bool",
            "f64" => "Float",
            "i64" => "Int",
            _ => "Str",
        };
    Ident::new(value_type)
}

/// Convert a `Value` type to a string representation of the type.
fn value_type_to_type(ident: &Ident) -> &str {
    match ident.to_string().as_ref() {
        "Bool" => "bool",
        "Float" => "float",
        "Int" => "integer",
        "Str" => "string",
        ty => panic!("Unknown Value type {}", ty),
    }
}
<commit_msg>Fixed clippy warning.<commit_after>/*
 * Copyright (c) 2016 Boucher, Antoni <bouanto@zoho.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

use quote::Tokens;
use syn::{Attribute, Body, Ident, MacroInput, Path, VariantData};
use syn::Body::{Enum, Struct};
use syn::MetaItem::{List, Word};
use syn::Ty;

use attributes::to_metadata_impl;
use string::{snake_to_camel, to_dash_name};

/// Expand the required trais for the derive Setting attribute.
pub fn expand_setting_enum(ast: MacroInput) -> Tokens {
    let name = ast.ident.clone();
    let mut default = None;

    let mut variant_names = vec![];
    if let Enum(ref variants) = ast.body {
        for variant in variants {
            variant_names.push(variant.ident.clone());
            if !variant.attrs.is_empty() {
                for attribute in &variant.attrs {
                    if let &Attribute { value: Word(ref ident), .. } = attribute {
                        if ident == "default" {
                            default = Some(variant.ident.clone());
                        }
                    }
                }
            }
        }
    }
    let choice_names: Vec<_> = variant_names.iter()
        .map(|name| to_dash_name(&name.to_string()))
        .collect();
    let choice_names1 = &choice_names;
    let choice_names2 = &choice_names;

    let qualified_names = variant_names.iter()
        .map(|variant_name| quote! {
            #name::#variant_name
        });

    let from_str_fn = quote! {
        fn from_str(string: &str) -> Result<Self, Self::Err> {
            match string {
                #(#choice_names1 => Ok(#qualified_names),)*
                _ => Err(::mg_settings::error::SettingError::UnknownChoice {
                    actual: string.to_string(),
                    expected: vec![#(#choice_names2),*],
                }),
            }
        }
    };

    let default_impl =
        if let Some(ident) = default {
            quote! {
                impl Default for #name {
                    fn default() -> Self {
                        #name::#ident
                    }
                }
            }
        }
        else {
            quote! {
            }
        };

    let completion_values_impl = quote! {
        impl CompletionValues for #name {
            fn completion_values() -> Vec<String> {
                vec![#(#choice_names1.to_string()),*]
            }
        }
    };

    quote! {
        #default_impl

        #completion_values_impl

        impl ::std::str::FromStr for #name {
            type Err = ::mg_settings::error::SettingError;

            #from_str_fn
        }
    }
}

/// Expand the required traits for the derive Settings attribute.
pub fn expand_settings_enum(ast: MacroInput) -> Tokens {
    let name = &ast.ident;
    let completion_fn = to_setting_completion_fn(name, &ast.body);
    let variant_name = Ident::new(format!("{}Variant", name));
    let variant_enum = to_enums(&variant_name, &ast.body);
    let settings_impl = to_settings_impl(name, &variant_name, &ast.body);
    let (metadata_impl, _) = to_metadata_impl(name, &ast.body);
    quote! {
        #variant_enum

        #settings_impl

        #metadata_impl

        #completion_fn
    }
}

/// Check if a type is a custom type (including enum).
fn is_custom_type(ident: &Ident) -> bool {
    match ident.to_string().as_ref() {
        "bool" | "f64" | "i64" | "String" => false,
        _ => true,
    }
}

/// Create the variant enums for getters and setters.
fn to_enums(variant_name: &Ident, settings_struct: &Body) -> Tokens {
    if let &Struct(VariantData::Struct(ref strct)) = settings_struct {
        let mut field_names = vec![];
        let mut names = vec![];
        let mut types = vec![];
        for field in strct {
            if let Some(ref ident) = field.ident {
                field_names.push(ident);
                let ident = Ident::new(snake_to_camel(&ident.to_string()));
                names.push(ident);
                types.push(field.ty.clone());
            }
        }
        let names1 = &names;
        quote! {
            pub enum #variant_name {
                #(#names1(#types)),*
            }
        }
    }
    else {
        panic!("Not a struct");
    }
}

/// Create the impl Settings.
fn to_settings_impl(name: &Ident, variant_name: &Ident, settings_struct: &Body) -> Tokens {
    if let &Struct(VariantData::Struct(ref strct)) = settings_struct {
        let mut names = vec![];
        let mut capitalized_names = vec![];
        let mut original_types = vec![];
        let mut types = vec![];
        for field in strct {
            if let Some(ref ident) = field.ident {
                let ident_string = ident.to_string();
                let ident = Ident::new(ident_string.clone());
                names.push(ident);
                let ident = Ident::new(snake_to_camel(&ident_string));
                capitalized_names.push(
                    quote! {
                        #variant_name::#ident
                    });

                if let Ty::Path(_, Path { ref segments, .. }) = field.ty {
                    original_types.push(&segments[0].ident);
                    types.push(to_value_type(&segments[0].ident));
                }
            }
        }
        let string_names: Vec<_> = names.iter()
            .map(|ident| ident.to_string().replace("_", "-"))
            .collect();
        let string_names = &string_names;
        let capitalized_names = &capitalized_names;
        let names1 = &names;
        let names2 = &names;
        let names3 = &names;
        let variant_exprs = names.iter().zip(original_types.iter())
            .map(|(name, typ)|
                 if is_custom_type(typ) {
                     quote! {
                         ::std::str::FromStr::from_str(&#name)?
                     }
                 }
                 else {
                     quote! { #name }
                 }
            );
        let types1 = &types;
        let type_names = types.iter()
            .map(|ident| value_type_to_type(&ident));

        let to_variant_fn_variant = quote! {
            #(#string_names => {
                if let ::mg_settings::Value::#types1(#names1) = value {
                    Ok(#capitalized_names(#variant_exprs))
                }
                else {
                    Err(::mg_settings::error::SettingError::WrongType {
                        actual: value.to_type().to_string(),
                        expected: #type_names.to_string(),
                    })
                }
            },)*
        };

        let to_variant_fn = quote! {
            #[allow(cyclomatic_complexity)]
            fn to_variant(name: &str, value: ::mg_settings::Value) -> Result<Self::Variant, ::mg_settings::error::SettingError> {
                match name {
                    #to_variant_fn_variant
                    _ => Err(::mg_settings::error::SettingError::UnknownSetting(name.to_string())),
                }
            }
        };

        quote! {
            impl ::mg_settings::settings::Settings for #name {
                type Variant = #variant_name;

                #to_variant_fn

                fn set_value(&mut self, value: Self::Variant) {
                    match value {
                        #(#capitalized_names(#names1) => {
                            self.#names2 = #names3
                        }),*
                    }
                }
            }
        }
    }
    else {
        panic!("Not a struct");
    }
}

/// Create the function returing the completion of the setting values.
pub fn to_setting_completion_fn(name: &Ident, body: &Body) -> Tokens {
    let mut completions = vec![];
    if let Struct(VariantData::Struct(ref fields)) = *body {
        'field_loop:
        for field in fields {
            for attribute in &field.attrs {
                if let &Attribute { value: List(ref ident, ref args), .. } = attribute {
                    if ident == "completion" {
                        if let Word(ref arg_ident) = args[0] {
                            if arg_ident == "hidden" {
                                continue 'field_loop;
                            }
                        }
                    }
                }
            }

            let setting_name = field.ident.as_ref().unwrap().to_string().replace('_', "-");
            let field_type = &field.ty;

            completions.push(quote! {
                (#setting_name.to_string(), #field_type::completion_values())
            });
        }
    }

    quote! {
        use mg_settings::CompletionValues;

        impl ::mg_settings::SettingCompletion for #name {
            fn get_value_completions() -> ::std::collections::HashMap<String, Vec<String>> {
                let mut vec = vec![#(#completions),*];
                let iter = vec.drain(..);
                iter.collect()
            }
        }
    }
}

/// Convert a Rust type to a `Value` type.
fn to_value_type(ident: &Ident) -> Ident {
    let value_type =
        match ident.to_string().as_ref() {
            "bool" => "Bool",
            "f64" => "Float",
            "i64" => "Int",
            _ => "Str",
        };
    Ident::new(value_type)
}

/// Convert a `Value` type to a string representation of the type.
fn value_type_to_type(ident: &Ident) -> &str {
    match ident.to_string().as_ref() {
        "Bool" => "bool",
        "Float" => "float",
        "Int" => "integer",
        "Str" => "string",
        ty => panic!("Unknown Value type {}", ty),
    }
}
<|endoftext|>


<commit_before>use crate::codec::*;
use crate::endpoint::Endpoint;
use crate::error::ZmqResult;
use crate::message::*;
use crate::transport::AcceptStopHandle;
use crate::util::PeerIdentity;
use crate::{
    MultiPeerBackend, Socket, SocketBackend, SocketEvent, SocketOptions, SocketRecv, SocketType,
};

use crate::backend::Peer;
use dashmap::DashMap;
use parking_lot::Mutex;
use crate::fair_queue::QueueInner;
use crate::fair_queue::FairQueue;
use crossbeam::queue::SegQueue;
use async_trait::async_trait;
use bytes::{BufMut, BytesMut};
use futures::channel::mpsc;
use futures::{SinkExt, StreamExt};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

pub(crate) struct SubSocketBackend {
    pub(crate) peers: DashMap<PeerIdentity, Peer>,
    fair_queue_inner: Option<Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>>,
    pub(crate) round_robin: SegQueue<PeerIdentity>,
    socket_type: SocketType,
    socket_options: SocketOptions,
    pub(crate) socket_monitor: Mutex<Option<mpsc::Sender<SocketEvent>>>,
    subs: Mutex<HashSet<String>>,
}

impl SubSocketBackend {
    pub(crate) fn with_options(
        fair_queue_inner: Option<Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>>,
        socket_type: SocketType,
        options: SocketOptions,
    ) -> Self {
        Self {
            peers: DashMap::new(),
            fair_queue_inner,
            round_robin: SegQueue::new(),
            socket_type,
            socket_options: options,
            socket_monitor: Mutex::new(None),
	    subs: Mutex::new(HashSet::new()),
        }
    }
}

impl SocketBackend for SubSocketBackend {
    fn socket_type(&self) -> SocketType {
        self.socket_type
    }

    fn socket_options(&self) -> &SocketOptions {
        &self.socket_options
    }

    fn shutdown(&self) {
        self.peers.clear();
    }

    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>> {
        &self.socket_monitor
    }
}

#[async_trait]
impl MultiPeerBackend for SubSocketBackend {
    async fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo) {
	println!("Connect peer");
        let (recv_queue, send_queue) = io.into_parts();
        self.peers.insert(peer_id.clone(), Peer { send_queue });
        self.round_robin.push(peer_id.clone());
        match &self.fair_queue_inner {
            None => {}
            Some(inner) => {
                inner.lock().insert(peer_id.clone(), recv_queue);
            }
        };

	let subs = (|| { (*self.subs.lock()).clone() })();

	for sub in subs.iter() {
	    let mut buf = BytesMut::with_capacity(sub.len() + 1);
	    buf.put_u8(1);
	    buf.extend_from_slice(sub.as_bytes());

	    let message: ZmqMessage = ZmqMessage::from(buf.freeze());
	    let mut peer = self.peers.get_mut(&peer_id).unwrap();

	    peer.send_queue
		.send(Message::Message(message.clone()))
		.await.unwrap();
	}
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

pub struct SubSocket {
    backend: Arc<SubSocketBackend>,
    fair_queue: FairQueue<ZmqFramedRead, PeerIdentity>,
    binds: HashMap<Endpoint, AcceptStopHandle>,
}

impl Drop for SubSocket {
    fn drop(&mut self) {
        self.backend.shutdown()
    }
}

impl SubSocket {
    pub async fn subscribe(&mut self, subscription: &str) -> ZmqResult<()> {
	println!("subscribe");
        let mut buf = BytesMut::with_capacity(subscription.len() + 1);
        buf.put_u8(1);
        buf.extend_from_slice(subscription.as_bytes());
        // let message = format!("\0x1{}", subscription);
        let message: ZmqMessage = ZmqMessage::from(buf.freeze());
        for mut peer in self.backend.peers.iter_mut() {
            peer.send_queue
                .send(Message::Message(message.clone()))
                .await?;
        }
	self.backend.subs.lock().insert(subscription.to_string());
        Ok(())
    }

    pub async fn unsubscribe(&mut self, subscription: &str) -> ZmqResult<()> {
	self.backend.subs.lock().remove(subscription);
        let mut buf = BytesMut::with_capacity(subscription.len() + 1);
        buf.put_u8(0);
        buf.extend_from_slice(subscription.as_bytes());
        let message = ZmqMessage::from(buf.freeze());
        for mut peer in self.backend.peers.iter_mut() {
            peer.send_queue
                .send(Message::Message(message.clone()))
                .await?;
        }
        Ok(())
    }
}

#[async_trait]
impl Socket for SubSocket {
    fn with_options(options: SocketOptions) -> Self {
        let fair_queue = FairQueue::new(true);
        Self {
            backend: Arc::new(SubSocketBackend::with_options(
                Some(fair_queue.inner()),
                SocketType::SUB,
                options,
            )),
            fair_queue,
            binds: HashMap::new(),
        }
    }

    fn backend(&self) -> Arc<dyn MultiPeerBackend> {
	self.backend.clone()
    }

    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
        &mut self.binds
    }

    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
        let (sender, receiver) = mpsc::channel(1024);
        self.backend.socket_monitor.lock().replace(sender);
        receiver
    }
}

#[async_trait]
impl SocketRecv for SubSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        loop {
            match self.fair_queue.next().await {
                Some((_peer_id, Ok(Message::Message(message)))) => {
                    return Ok(message);
                }
                Some((_peer_id, _)) => todo!(),
                None => todo!(),
            }
        }
    }
}
<commit_msg>refactor: move subscribe/unsubscribe to the backend<commit_after>use crate::codec::*;
use crate::endpoint::Endpoint;
use crate::error::ZmqResult;
use crate::message::*;
use crate::transport::AcceptStopHandle;
use crate::util::PeerIdentity;
use crate::{
    MultiPeerBackend, Socket, SocketBackend, SocketEvent, SocketOptions, SocketRecv, SocketType,
};

use crate::backend::Peer;
use crate::fair_queue::FairQueue;
use crate::fair_queue::QueueInner;
use async_trait::async_trait;
use bytes::{BufMut, BytesMut};
use crossbeam::queue::SegQueue;
use dashmap::DashMap;
use futures::channel::mpsc;
use futures::{SinkExt, StreamExt};
use parking_lot::Mutex;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

pub enum SubBackendMsgType {
    UNSUBSCRIBE = 0,
    SUBSCRIBE = 1
}

pub(crate) struct SubSocketBackend {
    pub(crate) peers: DashMap<PeerIdentity, Peer>,
    fair_queue_inner: Option<Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>>,
    pub(crate) round_robin: SegQueue<PeerIdentity>,
    socket_type: SocketType,
    socket_options: SocketOptions,
    pub(crate) socket_monitor: Mutex<Option<mpsc::Sender<SocketEvent>>>,
    subs: Mutex<HashSet<String>>,
}

impl SubSocketBackend {
    pub(crate) fn with_options(
        fair_queue_inner: Option<Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>>,
        socket_type: SocketType,
        options: SocketOptions,
    ) -> Self {
        Self {
            peers: DashMap::new(),
            fair_queue_inner,
            round_robin: SegQueue::new(),
            socket_type,
            socket_options: options,
            socket_monitor: Mutex::new(None),
            subs: Mutex::new(HashSet::new()),
        }
    }

    pub fn create_subs_message(subscription: &str, msg_type: SubBackendMsgType) -> ZmqMessage {
        let mut buf = BytesMut::with_capacity(subscription.len() + 1);
        buf.put_u8(msg_type as u8);
        buf.extend_from_slice(subscription.as_bytes());

	buf.freeze().into()
    }
}

impl SocketBackend for SubSocketBackend {
    fn socket_type(&self) -> SocketType {
        self.socket_type
    }

    fn socket_options(&self) -> &SocketOptions {
        &self.socket_options
    }

    fn shutdown(&self) {
        self.peers.clear();
    }

    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>> {
        &self.socket_monitor
    }
}

#[async_trait]
impl MultiPeerBackend for SubSocketBackend {
    async fn peer_connected(self: Arc<Self>, peer_id: &PeerIdentity, io: FramedIo) {
        let (recv_queue, mut send_queue) = io.into_parts();

	let subs_msgs: Vec<ZmqMessage> = self.subs.lock().iter().map(
	    |x| SubSocketBackend::create_subs_message(
		x, SubBackendMsgType::SUBSCRIBE)).collect();

        for message in subs_msgs.iter() {
            send_queue
                .send(Message::Message(message.clone()))
                .await
                .unwrap();
        }

        self.peers.insert(peer_id.clone(), Peer { send_queue });
        self.round_robin.push(peer_id.clone());
        match &self.fair_queue_inner {
            None => {}
            Some(inner) => {
                inner.lock().insert(peer_id.clone(), recv_queue);
            }
        };
    }

    fn peer_disconnected(&self, peer_id: &PeerIdentity) {
        self.peers.remove(peer_id);
    }
}

pub struct SubSocket {
    backend: Arc<SubSocketBackend>,
    fair_queue: FairQueue<ZmqFramedRead, PeerIdentity>,
    binds: HashMap<Endpoint, AcceptStopHandle>,
}

impl Drop for SubSocket {
    fn drop(&mut self) {
        self.backend.shutdown()
    }
}

impl SubSocket {
    pub async fn subscribe(&mut self, subscription: &str) -> ZmqResult<()> {
	self.backend.subs.lock().insert(subscription.to_string());
	self.process_subs(subscription, SubBackendMsgType::SUBSCRIBE).await
    }

    pub async fn unsubscribe(&mut self, subscription: &str) -> ZmqResult<()> {
	self.backend.subs.lock().remove(subscription);
	self.process_subs(subscription, SubBackendMsgType::UNSUBSCRIBE).await
    }

    async fn process_subs(&mut self, subscription: &str, msg_type: SubBackendMsgType) -> ZmqResult<()> {
        let message: ZmqMessage = SubSocketBackend::create_subs_message(
	    subscription, msg_type);

        for mut peer in self.backend.peers.iter_mut() {
            peer.send_queue
                .send(Message::Message(message.clone()))
                .await?;
        }
	Ok(())
    }
}

#[async_trait]
impl Socket for SubSocket {
    fn with_options(options: SocketOptions) -> Self {
        let fair_queue = FairQueue::new(true);
        Self {
            backend: Arc::new(SubSocketBackend::with_options(
                Some(fair_queue.inner()),
                SocketType::SUB,
                options,
            )),
            fair_queue,
            binds: HashMap::new(),
        }
    }

    fn backend(&self) -> Arc<dyn MultiPeerBackend> {
        self.backend.clone()
    }

    fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
        &mut self.binds
    }

    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
        let (sender, receiver) = mpsc::channel(1024);
        self.backend.socket_monitor.lock().replace(sender);
        receiver
    }
}

#[async_trait]
impl SocketRecv for SubSocket {
    async fn recv(&mut self) -> ZmqResult<ZmqMessage> {
        loop {
            match self.fair_queue.next().await {
                Some((_peer_id, Ok(Message::Message(message)))) => {
                    return Ok(message);
                }
                Some((_peer_id, _)) => todo!(),
                None => todo!(),
            }
        }
    }
}
<|endoftext|>


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

use super::*;
use maplit::hashmap;
use oak_abi::{label::Label, proto::oak::application::ApplicationConfiguration};
use std::{fs, option::Option, thread::JoinHandle};

static LOCAL_CA: &str = "../examples/certs/local/ca.pem";
static GCP_CA: &str = "../examples/certs/gcp/ca.pem";

struct HttpServerTester {
    runtime: RuntimeProxy,
    server_node_thread_handle: Option<JoinHandle<()>>,
    oak_node_simulator_thread_handle: Option<JoinHandle<()>>,
    notify_sender: Option<tokio::sync::oneshot::Sender<()>>,
}

impl HttpServerTester {
    fn new(port: u32, with_simulator_thread: bool) -> HttpServerTester {
        let runtime = create_runtime();
        let server_node = create_server_node(port);
        let (init_receiver, invocation_receiver) = create_communication_channel(&runtime);

        // Start the server node
        let (notify_sender, notify_receiver) = tokio::sync::oneshot::channel::<()>();
        let runtime_proxy = runtime.clone();
        // TODO(#1186): Use tokio instead of spawning a thread.
        let server_node_thread = std::thread::spawn(move || {
            server_node.run(runtime_proxy, init_receiver, notify_receiver)
        });

        // Simulate an Oak node that responds with 200 (OK) to every request it receives
        // TODO(#1186): Use tokio instead of spawning a thread.
        let oak_node_simulator_thread = if with_simulator_thread {
            let runtime_proxy = runtime.clone();
            Some(std::thread::spawn(move || {
                oak_node_simulator(&runtime_proxy, invocation_receiver);
            }))
        } else {
            None
        };

        HttpServerTester {
            runtime,
            server_node_thread_handle: Some(server_node_thread),
            oak_node_simulator_thread_handle: oak_node_simulator_thread,
            notify_sender: Some(notify_sender),
        }
    }

    fn cleanup(&mut self) {
        if let Some(sender) = self.notify_sender.take() {
            sender.send(()).unwrap_or_else(|()| {
                debug!("Test node already dropped `notify_receiver`.");
            })
        }
        if let Some(handle) = self.server_node_thread_handle.take() {
            let _ = handle.join();
        };
        if let Some(handle) = self.oak_node_simulator_thread_handle.take() {
            let _ = handle.join();
        };

        // stop the runtime and any servers it is running
        self.runtime.runtime.stop();
    }
}

#[tokio::test]
async fn test_https_server_can_serve_https_requests() {
    // Start a runtime with an HTTP server node, and a thread simulating an Oak node to respond to
    // HTTP requests.
    let mut http_server_tester = HttpServerTester::new(2525, true);

    // Send an HTTPS request, and check that response has StatusCode::OK
    let resp = send_request("https://localhost:2525", LOCAL_CA).await;
    assert!(resp.is_ok());
    assert_eq!(
        resp.unwrap().status(),
        http::status::StatusCode::OK.as_u16()
    );

    // Stop the runtime and the servers
    http_server_tester.cleanup();
}

#[tokio::test]
async fn test_https_server_cannot_serve_http_requests() {
    // Start a runtime with an HTTP server node. The HTTP server in this case rejects the requests,
    // and would not send anything to the Oak node. Creating a thread to simulate the Oak node will
    // result in the thread being blocked for ever. So, we set up the test without an
    // oak-node-simulator thread.
    let mut http_server_tester = HttpServerTester::new(2526, false);

    // Send an HTTP request, and check that the server responds with an error
    let resp = send_request("http://localhost:2526", LOCAL_CA).await;
    assert!(resp.is_err());

    // Stop the runtime and the servers
    http_server_tester.cleanup();
}

#[tokio::test]
async fn test_https_server_does_not_terminate_after_a_bad_request() {
    let mut http_server_tester = HttpServerTester::new(2527, true);

    // Send an HTTPS request with invalid certificate, and check that the server responds with error
    let resp = send_request("https://localhost:2527", GCP_CA).await;
    assert!(resp.is_err());

    // Send a second request, and check that the server is alive and responsive
    let resp = send_request("https://localhost:2527", LOCAL_CA).await;
    assert!(resp.is_ok());

    // Stop the runtime and the servers
    http_server_tester.cleanup();
}

fn create_runtime() -> RuntimeProxy {
    let configuration = ApplicationConfiguration {
        wasm_modules: hashmap! {},
        initial_node_configuration: None,
    };
    let signature_table = crate::SignatureTable::default();
    info!("Create runtime for test");

    crate::RuntimeProxy::create_runtime(
        &configuration,
        &crate::SecureServerConfiguration::default(),
        &signature_table,
    )
}

fn create_server_node(port: u32) -> Box<HttpServerNode> {
    let tls_config = crate::tls::TlsConfig::new(
        "../examples/certs/local/local.pem",
        "../examples/certs/local/local.key",
    )
    .expect("Could not create TLS config from local certs.");

    // Create http server node
    Box::new(
        HttpServerNode::new(
            "test-node",
            HttpServerConfiguration {
                address: format!("[::]:{}", port),
            },
            tls_config,
        )
        .expect("Could not create server node"),
    )
}

fn create_communication_channel(runtime: &RuntimeProxy) -> (oak_abi::Handle, oak_abi::Handle) {
    // create channel: one end to server_node::run; the other to the Oak node.
    let (init_sender, init_receiver) = runtime
        .channel_create(&Label::public_untrusted())
        .expect("Could not create channel");

    // At the start the HTTP server pseudo-Node expects to receive an invocation channel, with
    // exactly one handle in it.
    //
    // Create a channel for receiving invocations to pass to the HTTP server pseudo-Node.
    let (invocation_sender, invocation_receiver) = runtime
        .channel_create(&Label::public_untrusted())
        .expect("Could not create channel");
    let invocation_sender = HttpInvocationSender {
        sender: Some(Sender::<HttpInvocation>::new(WriteHandle {
            handle: invocation_sender,
        })),
    };
    let init_sender = Sender::<HttpInvocationSender>::new(WriteHandle {
        handle: init_sender,
    });

    if let Err(error) = init_sender.send(invocation_sender, runtime) {
        panic!("Could not write to the `init_sender` channel: {}", error);
    }
    if let Err(error) = init_sender.close(runtime) {
        panic!("Could not close the `init_sender` channel: {}", error);
    }

    (init_receiver, invocation_receiver)
}

fn oak_node_simulator(runtime: &RuntimeProxy, invocation_receiver: oak_abi::Handle) {
    // Get invocation message that contains the response_writer handle.
    let invocation_receiver = Receiver::<HttpInvocation>::new(ReadHandle {
        handle: invocation_receiver,
    });
    let invocation = invocation_receiver.receive(runtime).unwrap();
    let resp = HttpResponse {
        body: vec![],
        status: http::status::StatusCode::OK.as_u16() as i32,
        headers: hashmap! {},
    };
    invocation
        .sender
        .expect("Empty sender on invocation.")
        .send(resp, runtime)
        .unwrap();
}

async fn send_request(
    uri: &str,
    ca_path: &str,
) -> Result<http::response::Response<hyper::Body>, hyper::Error> {
    // Send a request, and wait for the response
    let label = oak_abi::label::Label::public_untrusted();
    let mut label_bytes = vec![];
    if let Err(err) = label.encode(&mut label_bytes) {
        panic!("Failed to encode label: {}", err);
    }

    let ca_file =
        fs::File::open(ca_path).unwrap_or_else(|e| panic!("failed to open {}: {}", ca_path, e));
    let mut ca = io::BufReader::new(ca_file);

    // Build an HTTP connector which supports HTTPS too.
    let mut http = hyper::client::HttpConnector::new();
    http.enforce_http(false);
    // Build a TLS client, using the custom CA store for lookups.
    let mut tls = rustls::ClientConfig::new();
    tls.root_store
        .add_pem_file(&mut ca)
        .expect("failed to load custom CA store");
    // Join the above part into an HTTPS connector.
    let https = hyper_rustls::HttpsConnector::from((http, tls));

    let client: hyper::client::Client<_, hyper::Body> =
        hyper::client::Client::builder().build(https);

    let request = hyper::Request::builder()
        .method(http::Method::GET)
        .uri(uri)
        .header(oak_abi::OAK_LABEL_HTTP_PROTOBUF_KEY, label_bytes)
        .body(hyper::Body::empty())
        .unwrap();

    client.request(request).await
}
<commit_msg>Retry sending requests in HTTP tests to make the tests stable (#1469)<commit_after>//
// Copyright 2020 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

use super::*;
use maplit::hashmap;
use oak_abi::{label::Label, proto::oak::application::ApplicationConfiguration};
use std::{fs, option::Option, thread::JoinHandle};

static LOCAL_CA: &str = "../examples/certs/local/ca.pem";
static GCP_CA: &str = "../examples/certs/gcp/ca.pem";

struct HttpServerTester {
    runtime: RuntimeProxy,
    server_node_thread_handle: Option<JoinHandle<()>>,
    oak_node_simulator_thread_handle: Option<JoinHandle<()>>,
    notify_sender: Option<tokio::sync::oneshot::Sender<()>>,
}

impl HttpServerTester {
    /// Create a new runtime and test server.
    /// port: The port on which the server is started
    /// number_of_requests: the number of requests the server expects to receive. This parameter is
    /// passed to `oak_node_simulator` to set up the Oak node that servers the requests.
    fn new(port: u32, number_of_requests: i32) -> HttpServerTester {
        let runtime = create_runtime();
        let server_node = create_server_node(port);
        let (init_receiver, invocation_receiver) = create_communication_channel(&runtime);

        // Start the server node
        let (notify_sender, notify_receiver) = tokio::sync::oneshot::channel::<()>();
        let runtime_proxy = runtime.clone();
        // TODO(#1186): Use tokio instead of spawning a thread.
        let server_node_thread = std::thread::spawn(move || {
            server_node.run(runtime_proxy, init_receiver, notify_receiver)
        });

        // Simulate an Oak node that responds with 200 (OK) to every request it receives
        // TODO(#1186): Use tokio instead of spawning a thread.
        let runtime_proxy = runtime.clone();
        let oak_node_simulator_thread = Some(std::thread::spawn(move || {
            oak_node_simulator(number_of_requests, &runtime_proxy, invocation_receiver);
        }));

        HttpServerTester {
            runtime,
            server_node_thread_handle: Some(server_node_thread),
            oak_node_simulator_thread_handle: oak_node_simulator_thread,
            notify_sender: Some(notify_sender),
        }
    }

    fn cleanup(&mut self) {
        if let Some(sender) = self.notify_sender.take() {
            sender.send(()).unwrap_or_else(|()| {
                debug!("Test node already dropped `notify_receiver`.");
            })
        }
        if let Some(handle) = self.server_node_thread_handle.take() {
            let _ = handle.join();
        };
        if let Some(handle) = self.oak_node_simulator_thread_handle.take() {
            let _ = handle.join();
        };

        // stop the runtime and any servers it is running
        self.runtime.runtime.stop();
    }
}

fn init_logger() {
    // Ignore the result. We don't want to panic if the logger cannot be initialized, or is being
    // initialized more than once. Also, if the logger is not initialized, we cannot log an
    // error!
    let _res = env_logger::builder().is_test(true).try_init();
}

#[tokio::test]
async fn test_https_server_can_serve_https_requests() {
    init_logger();

    // Start a runtime with an HTTP server node, and a thread simulating an Oak node to respond to
    // HTTP requests.
    let mut http_server_tester = HttpServerTester::new(2525, 1);
    let client_with_valid_tls = create_client(LOCAL_CA);

    // Send an HTTPS request, and check that response has StatusCode::OK
    let resp = send_request(client_with_valid_tls, "https://localhost:2525").await;
    assert!(resp.is_ok());
    assert_eq!(
        resp.unwrap().status(),
        http::status::StatusCode::OK.as_u16()
    );

    // Stop the runtime and the servers
    http_server_tester.cleanup();
}

#[tokio::test]
async fn test_https_server_cannot_serve_http_requests() {
    init_logger();

    // Start a runtime with an HTTP server node. The HTTP server in this case rejects the requests,
    // and would not send anything to the Oak node. Creating a thread to simulate the Oak node will
    // result in the thread being blocked for ever. So, we set up the test without an
    // oak-node-simulator thread.
    let mut http_server_tester = HttpServerTester::new(2526, 0);
    let client_with_valid_tls = create_client(LOCAL_CA);

    // Send an HTTP request, and check that the server responds with an error
    let resp = send_request(client_with_valid_tls, "http://localhost:2526").await;
    assert!(resp.is_err());

    // Stop the runtime and the servers
    http_server_tester.cleanup();
}

#[tokio::test]
async fn test_https_server_does_not_terminate_after_a_bad_request() {
    init_logger();

    // Start a runtime with an HTTP server node, and a thread simulating an Oak node to respond to
    // HTTP requests.
    let mut http_server_tester = HttpServerTester::new(2527, 2);
    let client_with_valid_tls = create_client(LOCAL_CA);
    let client_with_invalid_tls = create_client(GCP_CA);

    // Send a valid request, making sure that the server is started
    let resp = send_request(client_with_valid_tls.clone(), "https://localhost:2527").await;
    assert!(resp.is_ok());

    // Send an HTTPS request with invalid certificate, and check that the server responds with error
    let resp = send_request(client_with_invalid_tls, "https://localhost:2527").await;
    assert!(resp.is_err());

    // Send another valid request, and check that the server is alive and responsive
    // let client_with_valid_tls = create_client(LOCAL_CA);
    let resp = send_request(client_with_valid_tls, "https://localhost:2527").await;
    assert!(resp.is_ok());

    // Stop the runtime and the servers
    http_server_tester.cleanup();
}

fn create_runtime() -> RuntimeProxy {
    let configuration = ApplicationConfiguration {
        wasm_modules: hashmap! {},
        initial_node_configuration: None,
    };
    let signature_table = crate::SignatureTable::default();
    info!("Create runtime for test");

    crate::RuntimeProxy::create_runtime(
        &configuration,
        &crate::SecureServerConfiguration::default(),
        &signature_table,
    )
}

fn create_server_node(port: u32) -> Box<HttpServerNode> {
    let tls_config = crate::tls::TlsConfig::new(
        "../examples/certs/local/local.pem",
        "../examples/certs/local/local.key",
    )
    .expect("Could not create TLS config from local certs.");

    // Create http server node
    Box::new(
        HttpServerNode::new(
            "test-node",
            HttpServerConfiguration {
                address: format!("[::]:{}", port),
            },
            tls_config,
        )
        .expect("Could not create server node"),
    )
}

fn create_communication_channel(runtime: &RuntimeProxy) -> (oak_abi::Handle, oak_abi::Handle) {
    // create channel: one end to server_node::run; the other to the Oak node.
    let (init_sender, init_receiver) = runtime
        .channel_create(&Label::public_untrusted())
        .expect("Could not create channel");

    // At the start the HTTP server pseudo-Node expects to receive an invocation channel, with
    // exactly one handle in it.
    //
    // Create a channel for receiving invocations to pass to the HTTP server pseudo-Node.
    let (invocation_sender, invocation_receiver) = runtime
        .channel_create(&Label::public_untrusted())
        .expect("Could not create channel");
    let invocation_sender = HttpInvocationSender {
        sender: Some(Sender::<HttpInvocation>::new(WriteHandle {
            handle: invocation_sender,
        })),
    };
    let init_sender = Sender::<HttpInvocationSender>::new(WriteHandle {
        handle: init_sender,
    });

    if let Err(error) = init_sender.send(invocation_sender, runtime) {
        panic!("Could not write to the `init_sender` channel: {}", error);
    }
    if let Err(error) = init_sender.close(runtime) {
        panic!("Could not close the `init_sender` channel: {}", error);
    }

    (init_receiver, invocation_receiver)
}

// Simulate an Oak node that responds to the specified number of requests. Since the Oak node blocks
// until a new request is received, we need to know the number of requests that the test sends
// beforehand. Alternatively, the node simulator could listen for incoming requests in an infinite
// loop and have a termination mechanism that could be signalled from outside. However, specifying
// the exact number of expected requests is more appropriate for testing purposes.
fn oak_node_simulator(
    number_of_requests: i32,
    runtime: &RuntimeProxy,
    invocation_receiver: oak_abi::Handle,
) {
    // Get invocation message that contains the response_writer handle.
    let invocation_receiver = Receiver::<HttpInvocation>::new(ReadHandle {
        handle: invocation_receiver,
    });

    for _counter in 0..number_of_requests {
        let invocation = invocation_receiver.receive(runtime).unwrap();
        let resp = HttpResponse {
            body: vec![],
            status: http::status::StatusCode::OK.as_u16() as i32,
            headers: hashmap! {},
        };
        invocation
            .sender
            .expect("Empty sender on invocation.")
            .send(resp, runtime)
            .unwrap();
    }
}

// Build a TLS client, using the given CA store
fn create_client(
    ca_path: &str,
) -> hyper::client::Client<hyper_rustls::HttpsConnector<hyper::client::HttpConnector>> {
    let ca_file =
        fs::File::open(ca_path).unwrap_or_else(|e| panic!("failed to open {}: {}", ca_path, e));
    let mut ca = io::BufReader::new(ca_file);

    // Build an HTTP connector which supports HTTPS too.
    let mut http = hyper::client::HttpConnector::new();
    http.enforce_http(false);
    // Build a TLS client, using the custom CA store for lookups.
    let mut tls = rustls::ClientConfig::new();
    tls.root_store
        .add_pem_file(&mut ca)
        .expect("failed to load custom CA store");
    // Join the above part into an HTTPS connector.
    let https = hyper_rustls::HttpsConnector::from((http, tls));

    hyper::client::Client::builder().build(https)
}

async fn send_request(
    client: hyper::client::Client<hyper_rustls::HttpsConnector<hyper::client::HttpConnector>>,
    uri: &str,
) -> Result<http::response::Response<hyper::Body>, hyper::Error> {
    // Send a request, and wait for the response
    let label = oak_abi::label::Label::public_untrusted();
    let mut label_bytes = vec![];
    if let Err(err) = label.encode(&mut label_bytes) {
        panic!("Failed to encode label: {}", err);
    }

    // The client thread may start sending the requests before the server is up. In this case, the
    // request will be rejected with a "ConnectError". To make the tests are stable, we need to
    // retry sending the requests until the server is up. To distinguish between these cases and
    // actual errors (e.g., errors due to invalid TLS certificates), we need to check the cause of
    // the error.
    loop {
        let request = hyper::Request::builder()
            .method(http::Method::GET)
            .uri(uri)
            .header(oak_abi::OAK_LABEL_HTTP_PROTOBUF_KEY, label_bytes.clone())
            .body(hyper::Body::empty())
            .unwrap();

        match client.request(request).await {
            Ok(reps) => return Ok(reps),
            Err(error) => {
                // We cannot access the cause of the error, so we need to check the string instead.
                let error_str = format!("{:?}", error);
                // If the cause is `ConnectError` (https://github.com/hyperium/hyper/blob/66fc127c8d4f81aed9300c9d0f13246b8206067a/src/client/connect/http.rs#L392)
                // it means that a connection to the server cannot be made. Retry sending the
                // request in this case.
                if error_str.contains("ConnectError") {
                    continue;
                } else {
                    return Err(error);
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>extern crate regex;

mod filters;

use errors::*;
use linter::filters::*;

#[derive(Debug, RustcEncodable)]
pub struct LinterWarning {
    pub message: &'static str,
    pub start: usize,
    pub end: usize,
}

pub struct Linter {
    locale: String,
}

impl Linter {
    pub fn new(locale: String) -> Result<Linter, LughError> {
        let linter = Linter { locale: locale };

        Ok(linter)
    }

    pub fn check(&self, text: &str) -> Result<(), Vec<LinterWarning>> {
        let mut warnings = Vec::<LinterWarning>::new();

        let filters: Vec<Box<LinterFilter>> = vec![
            Box::new(CurlyApostropheFilter {}),
            Box::new(EllipsisSymbolFilter {}),
            Box::new(NoSpaceBeforeCommaFilter {}),
            Box::new(SpaceBeforeDoublePonctuationFilter {}),
        ];

        let active_filters = filters
            .into_iter()
            .filter( |filter| filter.locales().is_empty() ||
                filter.locales().contains(&self.locale.as_str()) )
            .collect::<Vec<Box<LinterFilter>>>();

        for filter in &active_filters {
            let result = filter.check(text);

            if result.is_err() {
                warnings.append(&mut result.err().unwrap());
            }
        }

        if warnings.is_empty() {
            Ok(())
        } else {
            Err(warnings)
        }
    }
}

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

    #[test]
    fn test_linter_with_two_warnings() {
        let linter = Linter::new("en".to_string()).unwrap();

        let result = linter.check("It's me...");

        assert!(result.is_err());

        let warnings = result.err().unwrap();

        assert_eq!(2, warnings.len());

        assert_eq!("Please use curly apostrophes.", warnings[0].message);
        assert_eq!(2, warnings[0].start);
        assert_eq!(3, warnings[0].end);

        assert_eq!("Please use the ellipsis symbol (`…`) instead of three dots (`...`).", warnings[1].message);
        assert_eq!(7, warnings[1].start);
        assert_eq!(10, warnings[1].end);
    }

    #[test]
    fn test_linter_with_no_warning() {
        let linter = Linter::new("en".to_string()).unwrap();

        let result = linter.check("It’s me…");

        assert_eq!(false, result.is_err());
        assert_eq!((), result.unwrap());
    }

    #[test]
    fn test_linter_with_a_language_specific_warning() {
        let linter = Linter::new("fr".to_string()).unwrap();

        let result = linter.check("C’est moi!");

        assert!(result.is_err());

        let warnings = result.err().unwrap();

        assert_eq!(1, warnings.len());
        assert_eq!(
            "Please use a non-breaking space before “double” ponctuation marks: `;`, `:`, `!`, `?`.",
            warnings[0].message
        );
        assert_eq!(10, warnings[0].start);
        assert_eq!(12, warnings[0].end);
    }

    #[test]
    fn test_linter_with_a_language_specific_filter_and_no_warning() {
        let linter = Linter::new("en".to_string()).unwrap();

        let result = linter.check("It’s me!");

        assert_eq!(false, result.is_err());
        assert_eq!((), result.unwrap());
    }
}
<commit_msg>Move filters into functions in Linter<commit_after>extern crate regex;

mod filters;

use errors::*;
use linter::filters::*;

#[derive(Debug, RustcEncodable)]
pub struct LinterWarning {
    pub message: &'static str,
    pub start: usize,
    pub end: usize,
}

pub struct Linter {
    locale: String,
}

impl Linter {
    pub fn new(locale: String) -> Result<Linter, LughError> {
        let linter = Linter { locale: locale };

        Ok(linter)
    }

    pub fn check(&self, text: &str) -> Result<(), Vec<LinterWarning>> {
        let mut warnings = Vec::<LinterWarning>::new();

        for filter in &self.active_filters(self.locale.as_str()) {
            let result = filter.check(text);

            if result.is_err() {
                warnings.append(&mut result.err().unwrap());
            }
        }

        if warnings.is_empty() {
            Ok(())
        } else {
            Err(warnings)
        }
    }

    fn active_filters(&self, locale: &str) -> Vec<Box<LinterFilter>> {
        self.filters()
            .into_iter()
            .filter(|filter| filter.locales().is_empty() || filter.locales().contains(&locale))
            .collect()
    }

    fn filters(&self) -> Vec<Box<LinterFilter>> {
        vec![
            Box::new(CurlyApostropheFilter {}),
            Box::new(EllipsisSymbolFilter {}),
            Box::new(NoSpaceBeforeCommaFilter {}),
            Box::new(SpaceBeforeDoublePonctuationFilter {}),
        ]
    }
}

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

    #[test]
    fn test_linter_with_two_warnings() {
        let linter = Linter::new("en".to_string()).unwrap();

        let result = linter.check("It's me...");

        assert!(result.is_err());

        let warnings = result.err().unwrap();

        assert_eq!(2, warnings.len());

        assert_eq!("Please use curly apostrophes.", warnings[0].message);
        assert_eq!(2, warnings[0].start);
        assert_eq!(3, warnings[0].end);

        assert_eq!("Please use the ellipsis symbol (`…`) instead of three dots (`...`).", warnings[1].message);
        assert_eq!(7, warnings[1].start);
        assert_eq!(10, warnings[1].end);
    }

    #[test]
    fn test_linter_with_no_warning() {
        let linter = Linter::new("en".to_string()).unwrap();

        let result = linter.check("It’s me…");

        assert_eq!(false, result.is_err());
        assert_eq!((), result.unwrap());
    }

    #[test]
    fn test_linter_with_a_language_specific_warning() {
        let linter = Linter::new("fr".to_string()).unwrap();

        let result = linter.check("C’est moi!");

        assert!(result.is_err());

        let warnings = result.err().unwrap();

        assert_eq!(1, warnings.len());
        assert_eq!(
            "Please use a non-breaking space before “double” ponctuation marks: `;`, `:`, `!`, `?`.",
            warnings[0].message
        );
        assert_eq!(10, warnings[0].start);
        assert_eq!(12, warnings[0].end);
    }

    #[test]
    fn test_linter_with_a_language_specific_filter_and_no_warning() {
        let linter = Linter::new("en".to_string()).unwrap();

        let result = linter.check("It’s me!");

        assert_eq!(false, result.is_err());
        assert_eq!((), result.unwrap());
    }
}
<|endoftext|>