<commit_before>// Copyright 2015 Big Switch Networks, Inc
//      (Algorithms for uBPF helpers, originally in C)
// Copyright 2016 6WIND S.A. <quentin.monnet@6wind.com>
//      (Translation to Rust, other helpers)
//
// Licensed under the Apache License, Version 2.0 <http://www.apache.org/licenses/LICENSE-2.0> or
// the MIT license <http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! This module implements some built-in helpers that can be called from within an eBPF program.
//!
//! These helpers may originate from several places:
//!
//! * Some of them mimic the helpers available in the Linux kernel.
//! * Some of them were proposed as example helpers in uBPF and they were adapted here.
//! * Other helpers may be specific to rbpf.
//!
//! The prototype for helpers is always the same: five `u64` as arguments, and a `u64` as a return
//! value. Hence some helpers have unused arguments, or return a 0 value in all cases, in order to
//! respect this convention.

use std::u64;

// Helpers associated to kernel helpers
// See also linux/include/uapi/linux/bpf.h in Linux kernel sources.

// bpf_trace_printk()

/// Index of helper `bpf_trace_printk()`, equivalent to `bpf_trace_printf()`, in Linux kernel, see
/// <https://git.kernel.org/cgit/linux/kernel/git/torvalds/linux.git/tree/include/uapi/linux/bpf.h>.
pub const BPF_TRACE_PRINTK_IDX: u32 = 6;

/// Prints its **last three** arguments to standard output. The **first two** arguments are
/// **unused**. Returns the number of bytes written.
///
/// By ignoring the first two arguments, it creates a helper that will have a behavior similar to
/// the one of the equivalent helper `bpf_trace_printk()` from Linux kernel.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let res = helpers::bpf_trace_printf(0, 0, 1, 15, 32);
/// assert_eq!(res as usize, "bpf_trace_printf: 0x1, 0xf, 0x20\n".len());
/// ```
///
/// This will print `bpf_trace_printf: 0x1, 0xf, 0x20`.
///
/// The eBPF code needed to perform the call in this example would be nearly identical to the code
/// obtained by compiling the following code from C to eBPF with clang:
///
/// ```c
/// #include <linux/bpf.h>
/// #include "path/to/linux/samples/bpf/bpf_helpers.h"
///
/// int main(struct __sk_buff *skb)
/// {
///     // Only %d %u %x %ld %lu %lx %lld %llu %llx %p %s conversion specifiers allowed.
///     // See <https://git.kernel.org/cgit/linux/kernel/git/torvalds/linux.git/tree/kernel/trace/bpf_trace.c>.
///     char *fmt = "bpf_trace_printk %llx, %llx, %llx\n";
///     return bpf_trace_printk(fmt, sizeof(fmt), 1, 15, 32);
/// }
/// ```
///
/// This would equally print the three numbers in `/sys/kernel/debug/tracing` file each time the
/// program is run.
#[allow(dead_code)]
#[allow(unused_variables)]
pub fn bpf_trace_printf (unused1: u64, unused2: u64, arg3: u64, arg4: u64, arg5: u64) -> u64 {
    println!("bpf_trace_printf: {:#x}, {:#x}, {:#x}", arg3, arg4, arg5);
    let size_arg = | x | {
        if x == 0 {
            1
        } else {
            (x as f64).log(16.0).floor() as u64 + 1
        }
    };
    "bpf_trace_printf: 0x, 0x, 0x\n".len() as u64
        + size_arg(arg3) + size_arg(arg4) + size_arg(arg5)
}

// Helpers coming from uBPF <https://github.com/iovisor/ubpf/blob/master/vm/test.c>

/// The idea is to assemble five bytes into a single `u64`. For compatibility with the helpers API,
/// each argument must be a `u64`.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let gathered = helpers::gather_bytes(0x11, 0x22, 0x33, 0x44, 0x55);
/// assert_eq!(gathered, 0x1122334455);
/// ```
pub fn gather_bytes (arg1: u64, arg2: u64, arg3: u64, arg4: u64, arg5: u64) -> u64 {
    arg1.wrapping_shl(32) |
    arg2.wrapping_shl(24) |
    arg3.wrapping_shl(16) |
    arg4.wrapping_shl(8)  |
    arg5
}

/// Same as `void *memfrob(void *s, size_t n);` in `string.h` in C. See the GNU manual page (in
/// section 3) for `memfrob`. The memory is directly modified, and the helper returns 0 in all
/// cases. Arguments 3 to 5 are unused.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let val: u64 = 0x112233;
/// let val_ptr = &val as *const u64;
///
/// helpers::memfrob(val_ptr as u64, 8, 0, 0, 0);
/// assert_eq!(val, 0x2a2a2a2a2a3b0819);
/// helpers::memfrob(val_ptr as u64, 8, 0, 0, 0);
/// assert_eq!(val, 0x112233);
/// ```
#[allow(unused_variables)]
pub fn memfrob (ptr: u64, len: u64, unused3: u64, unused4: u64, unused5: u64) -> u64 {
    for i in 0..len {
        unsafe {
            let mut p = (ptr + i) as *mut u8;
            *p ^= 0b101010;
        }
    }
    0
}

// TODO: Try again when asm!() is available in stable Rust.
// #![feature(asm)]
// #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
// #[allow(unused_variables)]
// pub fn memfrob (ptr: u64, len: u64, arg3: u64, arg4: u64, arg5: u64) -> u64 {
//     unsafe {
//         asm!(
//                 "mov $0xf0, %rax"
//             ::: "mov $0xf1, %rcx"
//             ::: "mov $0xf2, %rdx"
//             ::: "mov $0xf3, %rsi"
//             ::: "mov $0xf4, %rdi"
//             ::: "mov $0xf5, %r8"
//             ::: "mov $0xf6, %r9"
//             ::: "mov $0xf7, %r10"
//             ::: "mov $0xf8, %r11"
//         );
//     }
//     0
// }

/// Compute and return the square root of argument 1, cast as a float. Arguments 2 to 5 are
/// unused.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let x = helpers::sqrti(9, 0, 0, 0, 0);
/// assert_eq!(x, 3);
/// ```
#[allow(dead_code)]
#[allow(unused_variables)]
pub fn sqrti (arg1: u64, unused2: u64, unused3: u64, unused4: u64, unused5: u64) -> u64 {
    (arg1 as f64).sqrt() as u64
}

/// C-like `strcmp`, return 0 if the strings are equal, and a non-null value otherwise.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let foo = "This is a string.".as_ptr() as u64;
/// let bar = "This is another sting.".as_ptr() as u64;
///
/// assert!(helpers::strcmp(foo, foo, 0, 0, 0) == 0);
/// assert!(helpers::strcmp(foo, bar, 0, 0, 0) != 0);
/// ```
#[allow(dead_code)]
#[allow(unused_variables)]
pub fn strcmp (arg1: u64, arg2: u64, arg3: u64, unused4: u64, unused5: u64) -> u64 {
    // C-like strcmp, maybe shorter than converting the bytes to string and comparing?
    if arg1 == 0 || arg2 == 0 {
        return u64::MAX;
    }
    let mut a = arg1;
    let mut b = arg2;
    unsafe {
        let mut a_val = *(a as *const u8);
        let mut b_val = *(b as *const u8);
        while a_val == b_val && a_val != 0 && b_val != 0 {
            a +=1 ;
            b +=1 ;
            a_val = *(a as *const u8);
            b_val = *(b as *const u8);
        }
        if a_val >= b_val {
            (a_val - b_val) as u64
        } else {
            (b_val - a_val) as u64
        }
    }
}
<commit_msg>src/helpers.rs: add `rand()` and `bpf_time_getns()` helpers<commit_after>// Copyright 2015 Big Switch Networks, Inc
//      (Algorithms for uBPF helpers, originally in C)
// Copyright 2016 6WIND S.A. <quentin.monnet@6wind.com>
//      (Translation to Rust, other helpers)
//
// Licensed under the Apache License, Version 2.0 <http://www.apache.org/licenses/LICENSE-2.0> or
// the MIT license <http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! This module implements some built-in helpers that can be called from within an eBPF program.
//!
//! These helpers may originate from several places:
//!
//! * Some of them mimic the helpers available in the Linux kernel.
//! * Some of them were proposed as example helpers in uBPF and they were adapted here.
//! * Other helpers may be specific to rbpf.
//!
//! The prototype for helpers is always the same: five `u64` as arguments, and a `u64` as a return
//! value. Hence some helpers have unused arguments, or return a 0 value in all cases, in order to
//! respect this convention.

extern crate libc;

use std::u64;

// Helpers associated to kernel helpers
// See also linux/include/uapi/linux/bpf.h in Linux kernel sources.

// bpf_ktime_getns()

/// Index of helper `bpf_ktime_getns()`, equivalent to `bpf_time_getns()`, in Linux kernel, see
/// <https://git.kernel.org/cgit/linux/kernel/git/torvalds/linux.git/tree/include/uapi/linux/bpf.h>.
pub const BPF_KTIME_GETNS_IDX: u32 = 5;

/// Get monotonic time (since boot time) in nanoseconds. All arguments are unused.
///
/// If needed, you may e.g. create a helper returning real time by using the same code, but
/// replacing `libc::CLOCK_MONOTONIC` with `libc::CLOCK_REALTIME`.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let t = helpers::bpf_time_getns(0, 0, 0, 0, 0);
/// let d =  t / 10u64.pow(9)  / 60   / 60  / 24;
/// let h = (t / 10u64.pow(9)  / 60   / 60) % 24;
/// let m = (t / 10u64.pow(9)  / 60 ) % 60;
/// let s = (t / 10u64.pow(9)) % 60;
/// let ns = t % 10u64.pow(9);
/// println!("Uptime: {:#x} == {} days {}:{}:{}, {} ns", t, d, h, m, s, ns);
/// ```
#[allow(dead_code)]
#[allow(unused_variables)]
pub fn bpf_time_getns (unused1: u64, unused2: u64, unused3: u64, unused4: u64, unused5: u64) -> u64 {
    unsafe {
        let mut tp = libc::timespec {
            tv_sec: 0,
            tv_nsec: 0
        };
        let tp_ptr: *mut libc::timespec = &mut tp;
        libc::clock_gettime(libc::CLOCK_MONOTONIC, tp_ptr);
        ((*tp_ptr).tv_sec as u64) * 10u64.pow(9) + (*tp_ptr).tv_nsec as u64
    }
}

// bpf_trace_printk()

/// Index of helper `bpf_trace_printk()`, equivalent to `bpf_trace_printf()`, in Linux kernel, see
/// <https://git.kernel.org/cgit/linux/kernel/git/torvalds/linux.git/tree/include/uapi/linux/bpf.h>.
pub const BPF_TRACE_PRINTK_IDX: u32 = 6;

/// Prints its **last three** arguments to standard output. The **first two** arguments are
/// **unused**. Returns the number of bytes written.
///
/// By ignoring the first two arguments, it creates a helper that will have a behavior similar to
/// the one of the equivalent helper `bpf_trace_printk()` from Linux kernel.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let res = helpers::bpf_trace_printf(0, 0, 1, 15, 32);
/// assert_eq!(res as usize, "bpf_trace_printf: 0x1, 0xf, 0x20\n".len());
/// ```
///
/// This will print `bpf_trace_printf: 0x1, 0xf, 0x20`.
///
/// The eBPF code needed to perform the call in this example would be nearly identical to the code
/// obtained by compiling the following code from C to eBPF with clang:
///
/// ```c
/// #include <linux/bpf.h>
/// #include "path/to/linux/samples/bpf/bpf_helpers.h"
///
/// int main(struct __sk_buff *skb)
/// {
///     // Only %d %u %x %ld %lu %lx %lld %llu %llx %p %s conversion specifiers allowed.
///     // See <https://git.kernel.org/cgit/linux/kernel/git/torvalds/linux.git/tree/kernel/trace/bpf_trace.c>.
///     char *fmt = "bpf_trace_printk %llx, %llx, %llx\n";
///     return bpf_trace_printk(fmt, sizeof(fmt), 1, 15, 32);
/// }
/// ```
///
/// This would equally print the three numbers in `/sys/kernel/debug/tracing` file each time the
/// program is run.
#[allow(dead_code)]
#[allow(unused_variables)]
pub fn bpf_trace_printf (unused1: u64, unused2: u64, arg3: u64, arg4: u64, arg5: u64) -> u64 {
    println!("bpf_trace_printf: {:#x}, {:#x}, {:#x}", arg3, arg4, arg5);
    let size_arg = | x | {
        if x == 0 {
            1
        } else {
            (x as f64).log(16.0).floor() as u64 + 1
        }
    };
    "bpf_trace_printf: 0x, 0x, 0x\n".len() as u64
        + size_arg(arg3) + size_arg(arg4) + size_arg(arg5)
}

// Helpers coming from uBPF <https://github.com/iovisor/ubpf/blob/master/vm/test.c>

/// The idea is to assemble five bytes into a single `u64`. For compatibility with the helpers API,
/// each argument must be a `u64`.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let gathered = helpers::gather_bytes(0x11, 0x22, 0x33, 0x44, 0x55);
/// assert_eq!(gathered, 0x1122334455);
/// ```
pub fn gather_bytes (arg1: u64, arg2: u64, arg3: u64, arg4: u64, arg5: u64) -> u64 {
    arg1.wrapping_shl(32) |
    arg2.wrapping_shl(24) |
    arg3.wrapping_shl(16) |
    arg4.wrapping_shl(8)  |
    arg5
}

/// Same as `void *memfrob(void *s, size_t n);` in `string.h` in C. See the GNU manual page (in
/// section 3) for `memfrob`. The memory is directly modified, and the helper returns 0 in all
/// cases. Arguments 3 to 5 are unused.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let val: u64 = 0x112233;
/// let val_ptr = &val as *const u64;
///
/// helpers::memfrob(val_ptr as u64, 8, 0, 0, 0);
/// assert_eq!(val, 0x2a2a2a2a2a3b0819);
/// helpers::memfrob(val_ptr as u64, 8, 0, 0, 0);
/// assert_eq!(val, 0x112233);
/// ```
#[allow(unused_variables)]
pub fn memfrob (ptr: u64, len: u64, unused3: u64, unused4: u64, unused5: u64) -> u64 {
    for i in 0..len {
        unsafe {
            let mut p = (ptr + i) as *mut u8;
            *p ^= 0b101010;
        }
    }
    0
}

// TODO: Try again when asm!() is available in stable Rust.
// #![feature(asm)]
// #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
// #[allow(unused_variables)]
// pub fn memfrob (ptr: u64, len: u64, arg3: u64, arg4: u64, arg5: u64) -> u64 {
//     unsafe {
//         asm!(
//                 "mov $0xf0, %rax"
//             ::: "mov $0xf1, %rcx"
//             ::: "mov $0xf2, %rdx"
//             ::: "mov $0xf3, %rsi"
//             ::: "mov $0xf4, %rdi"
//             ::: "mov $0xf5, %r8"
//             ::: "mov $0xf6, %r9"
//             ::: "mov $0xf7, %r10"
//             ::: "mov $0xf8, %r11"
//         );
//     }
//     0
// }

/// Compute and return the square root of argument 1, cast as a float. Arguments 2 to 5 are
/// unused.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let x = helpers::sqrti(9, 0, 0, 0, 0);
/// assert_eq!(x, 3);
/// ```
#[allow(dead_code)]
#[allow(unused_variables)]
pub fn sqrti (arg1: u64, unused2: u64, unused3: u64, unused4: u64, unused5: u64) -> u64 {
    (arg1 as f64).sqrt() as u64
}

/// C-like `strcmp`, return 0 if the strings are equal, and a non-null value otherwise.
///
/// # Examples
///
/// ```
/// use rbpf::helpers;
///
/// let foo = "This is a string.".as_ptr() as u64;
/// let bar = "This is another sting.".as_ptr() as u64;
///
/// assert!(helpers::strcmp(foo, foo, 0, 0, 0) == 0);
/// assert!(helpers::strcmp(foo, bar, 0, 0, 0) != 0);
/// ```
#[allow(dead_code)]
#[allow(unused_variables)]
pub fn strcmp (arg1: u64, arg2: u64, arg3: u64, unused4: u64, unused5: u64) -> u64 {
    // C-like strcmp, maybe shorter than converting the bytes to string and comparing?
    if arg1 == 0 || arg2 == 0 {
        return u64::MAX;
    }
    let mut a = arg1;
    let mut b = arg2;
    unsafe {
        let mut a_val = *(a as *const u8);
        let mut b_val = *(b as *const u8);
        while a_val == b_val && a_val != 0 && b_val != 0 {
            a +=1 ;
            b +=1 ;
            a_val = *(a as *const u8);
            b_val = *(b as *const u8);
        }
        if a_val >= b_val {
            (a_val - b_val) as u64
        } else {
            (b_val - a_val) as u64
        }
    }
}

// Some additional helpers

/// Returns a random u64 value comprised between `min` and `max` values (inclusive). Arguments 3 to
/// 5 are unused.
///
/// Relies on `rand()` function from libc, so `libc::srand()` should be called once before this
/// helper is used.
///
/// # Examples
///
/// ```
/// extern crate libc;
/// extern crate rbpf;
///
/// unsafe {
///     libc::srand(libc::time(std::ptr::null_mut()) as u32);
/// }
///
/// let n = rbpf::helpers::rand(3, 6, 0, 0, 0);
/// assert!(3 <= n && n <= 6);
/// ```
#[allow(dead_code)]
#[allow(unused_variables)]
pub fn rand (min: u64, max: u64, unused3: u64, unused4: u64, unused5: u64) -> u64 {
    let mut n = unsafe {
        (libc::rand() as u64).wrapping_shl(32) + libc::rand() as u64
    };
    if min < max {
        n = n % (max + 1 - min) + min;
    };
    n
}
<|endoftext|>


<commit_before>pub fn challenge(num: i32, expected: &str, ans: String) {
    println!("C-{}", num);
    println!("Expecting: {}", expected);
    println!("Answer: {}", ans);
    if ans.as_str() == expected {
        println!("Challenge Passed!");
    } else {
        println!("FAILED");
    }
}

pub fn new_set(s: i32) -> i32 {
    println!("\n\n### Set-{}", s);
    return s+1;
}
<commit_msg>removed an unused file<commit_after><|endoftext|>


<commit_before>// Copyright (c) 2016 Fedor Gogolev <knsd@knsd.net>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::io::{Read, BufRead, Error as IoError, ErrorKind};
use std::string::{FromUtf8Error};
use std::collections::{HashMap};
use std::cell::{RefCell};
use std::rc::{Rc};

use num::{Zero};
use num::bigint::{BigInt, ToBigInt, Sign};
use byteorder::{ReadBytesExt, LittleEndian, BigEndian, Error as ByteorderError};
use from_ascii::{FromAscii, ParseIntError, ParseFloatError};

use string::{unescape, Error as UnescapeError};
use value::{Value};

use opcodes::*;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        EmptyMarker
        StackTooSmall
        EmptyStack
        InvalidValueOnStack
        InvalidGetValue
        InvalidPutValue
        NotImplemented

        Read(err: ByteorderError) {
            from()
        }
        Io(err: IoError) {
            from()
        }
        UnknownOpcode(opcode: u8) {}

        InvalidInt {
            from(ParseIntError)
        }
        InvalidLong
        InvalidFloat {
            from(ParseFloatError)
        }

        InvalidString {
            from(FromUtf8Error)
        }
        UnescapeError(err: UnescapeError) {
            from()
        }

        InvalidProto(proto: u8)
        NegativeLength {}

        #[doc(hidden)]
        __Nonexhaustive
    }
}

#[derive(Debug, PartialEq)]
pub enum BooleanOrInt {
    Boolean(bool),
    Int(i64),
}

macro_rules! rc {
    ($term: expr) => (Rc::new(RefCell::new($term)))
}

fn read_exact<R>(rd: &mut R, mut buf: &mut [u8]) -> Result<(), IoError> where R: Read {
    while !buf.is_empty() {
        match rd.read(buf) {
            Ok(0) => break,
            Ok(n) => { let tmp = buf; buf = &mut tmp[n..]; }
            Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    if !buf.is_empty() {
        Err(IoError::new(ErrorKind::Other,
                       "failed to fill whole buffer"))
    } else {
        Ok(())
    }
}

fn read_until_newline<R>(rd: &mut R) -> Result<Vec<u8>, Error> where R: Read + BufRead {
    let mut buf = Vec::new();
    try!(rd.read_until('\n' as u8, &mut buf));

    // Skip last symbol — \n
    match buf.split_last() {
        Some((&b'\n', init)) => Ok(init.to_vec()),
        _ => Err(Error::InvalidString),
    }
}

fn read_decimal_int<R>(rd: &mut R) -> Result<BooleanOrInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let val = match &s[..] {
        b"00" => BooleanOrInt::Boolean(false),
        b"01" => BooleanOrInt::Boolean(true),
        _ => BooleanOrInt::Int(try!(i64::from_ascii(&s)))
    };
    Ok(val)
}

fn read_decimal_long<R>(rd: &mut R) -> Result<BigInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let init = match s.split_last() {
        None => return Err(Error::InvalidString),
        Some((&b'L', init)) => init,
        Some(_) => &s[..],
    };

    match BigInt::parse_bytes(&init, 10) {
        Some(i) => Ok(i),
        None => Err(Error::InvalidLong)
    }
}

fn read_long<R>(rd: &mut R, length: usize) -> Result<BigInt, Error> where R: Read + BufRead {
    let mut buf = vec![0; length];
    try!(read_exact(rd, buf.as_mut()));

    let mut n = BigInt::from_bytes_le(Sign::Plus, &buf);

    let last = match buf.last_mut() {
        None => return Err(Error::InvalidLong),
        Some(last) => last,
    };

    if *last > 127 {
        n = n - (1.to_bigint().unwrap() << (length * 8))
    }

    Ok(n)
}

pub struct Machine {
    stack: Vec<Value>,
    memo: HashMap<usize, Value>,
    marker: Option<usize>,
}

impl Machine {
    pub fn new() -> Self {
        Machine {
            stack: Vec::new(),
            memo: HashMap::new(),
            marker: None,
        }
    }

    fn split_off(&mut self) -> Result<Vec<Value>, Error> {
        let at = match self.marker {
            None => return Err(Error::EmptyMarker),
            Some(marker) => marker,
        };

        if at > self.stack.len() {
            return Err(Error::StackTooSmall);
        }

        Ok(self.stack.split_off(at))
    }

    fn pop(&mut self) -> Result<Value, Error> {
        match self.stack.pop() {
            None => return Err(Error::EmptyStack),
            Some(value) => Ok(value),
        }
    }

    fn handle_get(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.memo.get(&i) {
            None => return Err(Error::InvalidGetValue),
            Some(ref v) => (*v).clone(),
        };
        self.stack.push(value);
        Ok(())
    }

    fn handle_put(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.stack.last() {
            None => return Err(Error::EmptyStack),
            Some(ref v) => (*v).clone(),
        };
        self.memo.insert(i, value);
        Ok(())
    }

    pub fn execute<R>(&mut self, rd: &mut R) -> Result<bool, Error> where R: Read + BufRead {
        macro_rules! ensure_not_negative {
            ($n: expr) => ({
                if $n < Zero::zero() {
                    return Err(Error::NegativeLength)
                }
            })
        }

        let marker = try!(rd.read_u8());
        match marker {
            PROTO => {
                let version = try!(rd.read_u8());
                if version < 2 {
                    return Err(Error::InvalidProto(version))
                }
            },
            STOP => return Ok(true),

            INT => {
                self.stack.push(match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Boolean(v) => Value::Bool(v),
                    BooleanOrInt::Int(v) => Value::Long(BigInt::from(v)),
                })
            },
            BININT => self.stack.push(Value::Int(try!(rd.read_i32::<LittleEndian>()) as isize)),
            BININT1 => self.stack.push(Value::Int(try!(rd.read_u8()) as isize)),
            BININT2 => self.stack.push(Value::Int(try!(rd.read_u16::<LittleEndian>()) as isize)),
            LONG => self.stack.push(Value::Long(BigInt::from(try!(read_decimal_long(rd))))),
            LONG1 => {
                let length = try!(rd.read_u8());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }
            LONG4 => {
                let length = try!(rd.read_i32::<LittleEndian>());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }

            STRING => self.stack.push(Value::String(try!(unescape(&try!(read_until_newline(rd)), false)))),
            BINSTRING => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);

                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },
            SHORT_BINSTRING => {
                let length = try!(rd.read_u8());
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },

            NONE => self.stack.push(Value::None),
            NEWTRUE => self.stack.push(Value::Bool(true)),
            NEWFALSE => self.stack.push(Value::Bool(false)),

            UNICODE => {
                let buf = try!(unescape(&try!(read_until_newline(rd)), true));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },
            BINUNICODE => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, buf.as_mut()));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },

            FLOAT => {
                let s = try!(read_until_newline(rd));
                self.stack.push(Value::Float(try!(f64::from_ascii(&s))))
            },
            BINFLOAT => {
                self.stack.push(Value::Float(try!(rd.read_f64::<BigEndian>())))
            },

            EMPTY_LIST => {
                self.stack.push(Value::List(rc!(Vec::new())))
            },
            APPEND => {
                let v = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => (*list.borrow_mut()).push(v),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            APPENDS => {
                let values = try!(self.split_off());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => (*list.borrow_mut()).extend(values),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            LIST => {
                let values = try!(self.split_off());
                self.stack.push(Value::List(rc!(values)));
            },

            EMPTY_TUPLE => self.stack.push(Value::Tuple(rc!(Vec::new()))),
            TUPLE => {
                let values = try!(self.split_off());
                self.stack.push(Value::Tuple(rc!(values)));
            },
            TUPLE1 => {
                let v1 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1])))
            },
            TUPLE2 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1, v2])))
            },
            TUPLE3 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                let v3 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1, v2, v3])))
            }

            EMPTY_DICT => self.stack.push(Value::Dict(rc!(Vec::new()))),
            DICT => {
                let mut values = try!(self.split_off());
                let mut dict = Vec::new();

                for i in 0 .. values.len() / 2 { // TODO: Check panic
                    let key = values.remove(2 * i);
                    let value = values.remove(2 * i + 1);
                    dict.push((key, value));
                }
                self.stack.push(Value::Dict(rc!(dict)));
            },
            SETITEM => {
                let value = try!(self.pop());
                let key = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict)) => (*dict.borrow_mut()).push((key, value)),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            SETITEMS => {
                let mut values = try!(self.split_off());

                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict_ref)) => {
                        for i in 0 .. values.len() / 2 { // TODO: Check panic
                                let key = values.remove(2 * i);
                                let value = values.remove(2 * i + 1);
                                (*dict_ref.borrow_mut()).push((key, value));
                            }
                    },
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },

            POP => {
                try!(self.pop());
            },
            DUP => {
                let value = match self.stack.last() {
                    None => return Err(Error::EmptyStack),
                    Some(ref v) => (*v).clone(),
                };
                self.stack.push(value)
            },
            MARK => {
                self.marker = Some(self.stack.len())
            },
            POP_MARK => {
                try!(self.split_off());
            },

            GET => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }
            BINGET => {
                try!(self.handle_get(try!(rd.read_u8()) as usize))
            }
            LONG_BINGET => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }

            PUT => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }
            BINPUT => {
                try!(self.handle_put(try!(rd.read_u8()) as usize))
            }
            LONG_BINPUT => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }

            c => return Err(Error::UnknownOpcode(c)),
        }
        Ok(false)
    }
}

pub fn unpickle<R>(rd: &mut R) -> Result<Value, Error> where R: Read + BufRead {
    let mut machine = Machine::new();
    loop {
        if try!(machine.execute(rd)) {
            break
        }
    }
    Ok(try!(machine.pop()))
}
<commit_msg>Remove unneeded return statement<commit_after>// Copyright (c) 2016 Fedor Gogolev <knsd@knsd.net>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use std::io::{Read, BufRead, Error as IoError, ErrorKind};
use std::string::{FromUtf8Error};
use std::collections::{HashMap};
use std::cell::{RefCell};
use std::rc::{Rc};

use num::{Zero};
use num::bigint::{BigInt, ToBigInt, Sign};
use byteorder::{ReadBytesExt, LittleEndian, BigEndian, Error as ByteorderError};
use from_ascii::{FromAscii, ParseIntError, ParseFloatError};

use string::{unescape, Error as UnescapeError};
use value::{Value};

use opcodes::*;

quick_error! {
    #[derive(Debug)]
    pub enum Error {
        EmptyMarker
        StackTooSmall
        EmptyStack
        InvalidValueOnStack
        InvalidGetValue
        InvalidPutValue
        NotImplemented

        Read(err: ByteorderError) {
            from()
        }
        Io(err: IoError) {
            from()
        }
        UnknownOpcode(opcode: u8) {}

        InvalidInt {
            from(ParseIntError)
        }
        InvalidLong
        InvalidFloat {
            from(ParseFloatError)
        }

        InvalidString {
            from(FromUtf8Error)
        }
        UnescapeError(err: UnescapeError) {
            from()
        }

        InvalidProto(proto: u8)
        NegativeLength {}

        #[doc(hidden)]
        __Nonexhaustive
    }
}

#[derive(Debug, PartialEq)]
pub enum BooleanOrInt {
    Boolean(bool),
    Int(i64),
}

macro_rules! rc {
    ($term: expr) => (Rc::new(RefCell::new($term)))
}

fn read_exact<R>(rd: &mut R, mut buf: &mut [u8]) -> Result<(), IoError> where R: Read {
    while !buf.is_empty() {
        match rd.read(buf) {
            Ok(0) => break,
            Ok(n) => { let tmp = buf; buf = &mut tmp[n..]; }
            Err(ref e) if e.kind() == ErrorKind::Interrupted => {}
            Err(e) => return Err(e),
        }
    }
    if !buf.is_empty() {
        Err(IoError::new(ErrorKind::Other,
                       "failed to fill whole buffer"))
    } else {
        Ok(())
    }
}

fn read_until_newline<R>(rd: &mut R) -> Result<Vec<u8>, Error> where R: Read + BufRead {
    let mut buf = Vec::new();
    try!(rd.read_until('\n' as u8, &mut buf));

    // Skip last symbol — \n
    match buf.split_last() {
        Some((&b'\n', init)) => Ok(init.to_vec()),
        _ => Err(Error::InvalidString),
    }
}

fn read_decimal_int<R>(rd: &mut R) -> Result<BooleanOrInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let val = match &s[..] {
        b"00" => BooleanOrInt::Boolean(false),
        b"01" => BooleanOrInt::Boolean(true),
        _ => BooleanOrInt::Int(try!(i64::from_ascii(&s)))
    };
    Ok(val)
}

fn read_decimal_long<R>(rd: &mut R) -> Result<BigInt, Error> where R: Read + BufRead {
    let s = try!(read_until_newline(rd));
    let init = match s.split_last() {
        None => return Err(Error::InvalidString),
        Some((&b'L', init)) => init,
        Some(_) => &s[..],
    };

    match BigInt::parse_bytes(&init, 10) {
        Some(i) => Ok(i),
        None => Err(Error::InvalidLong)
    }
}

fn read_long<R>(rd: &mut R, length: usize) -> Result<BigInt, Error> where R: Read + BufRead {
    let mut buf = vec![0; length];
    try!(read_exact(rd, buf.as_mut()));

    let mut n = BigInt::from_bytes_le(Sign::Plus, &buf);

    let last = match buf.last_mut() {
        None => return Err(Error::InvalidLong),
        Some(last) => last,
    };

    if *last > 127 {
        n = n - (1.to_bigint().unwrap() << (length * 8))
    }

    Ok(n)
}

pub struct Machine {
    stack: Vec<Value>,
    memo: HashMap<usize, Value>,
    marker: Option<usize>,
}

impl Machine {
    pub fn new() -> Self {
        Machine {
            stack: Vec::new(),
            memo: HashMap::new(),
            marker: None,
        }
    }

    fn split_off(&mut self) -> Result<Vec<Value>, Error> {
        let at = match self.marker {
            None => return Err(Error::EmptyMarker),
            Some(marker) => marker,
        };

        if at > self.stack.len() {
            return Err(Error::StackTooSmall);
        }

        Ok(self.stack.split_off(at))
    }

    fn pop(&mut self) -> Result<Value, Error> {
        match self.stack.pop() {
            None => Err(Error::EmptyStack),
            Some(value) => Ok(value),
        }
    }

    fn handle_get(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.memo.get(&i) {
            None => return Err(Error::InvalidGetValue),
            Some(ref v) => (*v).clone(),
        };
        self.stack.push(value);
        Ok(())
    }

    fn handle_put(&mut self, i: usize) -> Result<(), Error> {
        let value = match self.stack.last() {
            None => return Err(Error::EmptyStack),
            Some(ref v) => (*v).clone(),
        };
        self.memo.insert(i, value);
        Ok(())
    }

    pub fn execute<R>(&mut self, rd: &mut R) -> Result<bool, Error> where R: Read + BufRead {
        macro_rules! ensure_not_negative {
            ($n: expr) => ({
                if $n < Zero::zero() {
                    return Err(Error::NegativeLength)
                }
            })
        }

        let marker = try!(rd.read_u8());
        match marker {
            PROTO => {
                let version = try!(rd.read_u8());
                if version < 2 {
                    return Err(Error::InvalidProto(version))
                }
            },
            STOP => return Ok(true),

            INT => {
                self.stack.push(match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Boolean(v) => Value::Bool(v),
                    BooleanOrInt::Int(v) => Value::Long(BigInt::from(v)),
                })
            },
            BININT => self.stack.push(Value::Int(try!(rd.read_i32::<LittleEndian>()) as isize)),
            BININT1 => self.stack.push(Value::Int(try!(rd.read_u8()) as isize)),
            BININT2 => self.stack.push(Value::Int(try!(rd.read_u16::<LittleEndian>()) as isize)),
            LONG => self.stack.push(Value::Long(BigInt::from(try!(read_decimal_long(rd))))),
            LONG1 => {
                let length = try!(rd.read_u8());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }
            LONG4 => {
                let length = try!(rd.read_i32::<LittleEndian>());
                self.stack.push(Value::Long(BigInt::from(try!(read_long(rd, length as usize)))))
            }

            STRING => self.stack.push(Value::String(try!(unescape(&try!(read_until_newline(rd)), false)))),
            BINSTRING => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);

                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },
            SHORT_BINSTRING => {
                let length = try!(rd.read_u8());
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, &mut buf));
                self.stack.push(Value::String(buf))
            },

            NONE => self.stack.push(Value::None),
            NEWTRUE => self.stack.push(Value::Bool(true)),
            NEWFALSE => self.stack.push(Value::Bool(false)),

            UNICODE => {
                let buf = try!(unescape(&try!(read_until_newline(rd)), true));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },
            BINUNICODE => {
                let length = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(length);
                let mut buf = vec![0; length as usize];
                try!(read_exact(rd, buf.as_mut()));
                self.stack.push(Value::Unicode(try!(String::from_utf8(buf))))
            },

            FLOAT => {
                let s = try!(read_until_newline(rd));
                self.stack.push(Value::Float(try!(f64::from_ascii(&s))))
            },
            BINFLOAT => {
                self.stack.push(Value::Float(try!(rd.read_f64::<BigEndian>())))
            },

            EMPTY_LIST => {
                self.stack.push(Value::List(rc!(Vec::new())))
            },
            APPEND => {
                let v = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => (*list.borrow_mut()).push(v),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            APPENDS => {
                let values = try!(self.split_off());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::List(ref mut list)) => (*list.borrow_mut()).extend(values),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            LIST => {
                let values = try!(self.split_off());
                self.stack.push(Value::List(rc!(values)));
            },

            EMPTY_TUPLE => self.stack.push(Value::Tuple(rc!(Vec::new()))),
            TUPLE => {
                let values = try!(self.split_off());
                self.stack.push(Value::Tuple(rc!(values)));
            },
            TUPLE1 => {
                let v1 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1])))
            },
            TUPLE2 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1, v2])))
            },
            TUPLE3 => {
                let v1 = try!(self.pop());
                let v2 = try!(self.pop());
                let v3 = try!(self.pop());
                self.stack.push(Value::Tuple(rc!(vec![v1, v2, v3])))
            }

            EMPTY_DICT => self.stack.push(Value::Dict(rc!(Vec::new()))),
            DICT => {
                let mut values = try!(self.split_off());
                let mut dict = Vec::new();

                for i in 0 .. values.len() / 2 { // TODO: Check panic
                    let key = values.remove(2 * i);
                    let value = values.remove(2 * i + 1);
                    dict.push((key, value));
                }
                self.stack.push(Value::Dict(rc!(dict)));
            },
            SETITEM => {
                let value = try!(self.pop());
                let key = try!(self.pop());
                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict)) => (*dict.borrow_mut()).push((key, value)),
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },
            SETITEMS => {
                let mut values = try!(self.split_off());

                match self.stack.last_mut() {
                    None => return Err(Error::EmptyStack),
                    Some(&mut Value::Dict(ref mut dict_ref)) => {
                        for i in 0 .. values.len() / 2 { // TODO: Check panic
                                let key = values.remove(2 * i);
                                let value = values.remove(2 * i + 1);
                                (*dict_ref.borrow_mut()).push((key, value));
                            }
                    },
                    _ => return Err(Error::InvalidValueOnStack),
                }
            },

            POP => {
                try!(self.pop());
            },
            DUP => {
                let value = match self.stack.last() {
                    None => return Err(Error::EmptyStack),
                    Some(ref v) => (*v).clone(),
                };
                self.stack.push(value)
            },
            MARK => {
                self.marker = Some(self.stack.len())
            },
            POP_MARK => {
                try!(self.split_off());
            },

            GET => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }
            BINGET => {
                try!(self.handle_get(try!(rd.read_u8()) as usize))
            }
            LONG_BINGET => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_get(n as usize))
            }

            PUT => {
                let n = match try!(read_decimal_int(rd)) {
                    BooleanOrInt::Int(n) => n,
                    BooleanOrInt::Boolean(false) => 0,
                    BooleanOrInt::Boolean(true) => 1,
                };
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }
            BINPUT => {
                try!(self.handle_put(try!(rd.read_u8()) as usize))
            }
            LONG_BINPUT => {
                let n = try!(rd.read_i32::<LittleEndian>());
                ensure_not_negative!(n);
                try!(self.handle_put(n as usize))
            }

            c => return Err(Error::UnknownOpcode(c)),
        }
        Ok(false)
    }
}

pub fn unpickle<R>(rd: &mut R) -> Result<Value, Error> where R: Read + BufRead {
    let mut machine = Machine::new();
    loop {
        if try!(machine.execute(rd)) {
            break
        }
    }
    Ok(try!(machine.pop()))
}
<|endoftext|>


<commit_before>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError, Offset};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    io::{Read, Write},
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    thread::{self, Builder as ThreadBuilder, JoinHandle},
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    hb_handle: Option<JoinHandle<()>>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    send_heartbeat: Arc<AtomicBool>,
    poll_timeout: Option<Duration>,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            hb_handle: None,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            send_heartbeat: Arc::new(AtomicBool::new(false)),
            poll_timeout: None,
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn start_heartbeat(&mut self, interval: Duration) -> Result<()> {
        let connection = self.connection.clone();
        let send_hartbeat = self.send_heartbeat.clone();
        let hb_handle = ThreadBuilder::new()
            .name("heartbeat".to_owned())
            .spawn(move || {
                while connection.status().connected() {
                    let start = Instant::now();
                    let mut remaining = interval;

                    loop {
                        thread::park_timeout(remaining);
                        let elapsed = start.elapsed();
                        if elapsed >= remaining {
                            break;
                        }
                        remaining -= interval - elapsed;
                    }

                    send_hartbeat.store(true, Ordering::Relaxed);
                }
            })?;
        self.hb_handle = Some(hb_handle);
        Ok(())
    }

    fn heartbeat(&mut self) -> Result<()> {
        if self.send_heartbeat.load(Ordering::Relaxed) {
            trace!("send heartbeat");
            self.connection.send_heartbeat()?;
            self.send_heartbeat.store(false, Ordering::Relaxed);
        }
        Ok(())
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                trace!("io_loop: start heartbeat");
                let heartbeat = Duration::from_secs(u64::from(heartbeat));
                self.start_heartbeat(heartbeat)?;
                self.poll_timeout = Some(heartbeat);
                trace!("io_loop: heartbeat started");
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn can_write(&self) -> bool {
        self.can_write
            && self.connection.has_pending_frames()
            && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    if let Some(hb_handle) = self.hb_handle.take() {
                        hb_handle.thread().unpark();
                        hb_handle.join().expect("heartbeat loop failed");
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout)?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.connection.has_pending_frames()
        );
        loop {
            self.heartbeat()?;
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.connection.has_pending_frames(),
            self.status
        );
        Ok(())
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.connection.has_pending_frames()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
            self.status = Status::Stop;
        }
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.send_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.socket.write(&self.send_buffer.data())?;

        trace!("wrote {} bytes", sz);
        self.send_buffer.consume(sz);
        if sz > 0 && self.send_buffer.available_data() > 0 {
            // We didn't write all the data yet
            self.send_continue()?;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.socket
                    .read(&mut self.receive_buffer.space())
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((send_id, next_msg)) = self.connection.next_frame() {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(_) => {
                    self.connection.mark_sent(send_id);
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((send_id, next_msg))?;
                            self.send_buffer.shift();
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.data().offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<commit_msg>IoLoop: properly handle mid-frame WouldBlock with empty frames queue<commit_after>use crate::{
    buffer::Buffer, connection::Connection, connection_status::ConnectionState, Error, Result,
};
use amq_protocol::frame::{gen_frame, parse_frame, AMQPFrame, GenError, Offset};
use log::{error, trace};
use mio::{event::Source, Events, Interest, Poll, Token, Waker};
use std::{
    io::{Read, Write},
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    thread::{self, Builder as ThreadBuilder, JoinHandle},
    time::{Duration, Instant},
};

pub(crate) const SOCKET: Token = Token(1);
const WAKER: Token = Token(2);

const FRAMES_STORAGE: usize = 32;

#[derive(Debug, PartialEq)]
enum Status {
    Initial,
    Setup,
    Stop,
}

pub struct IoLoop<T> {
    connection: Connection,
    socket: T,
    status: Status,
    poll: Poll,
    waker: Arc<Waker>,
    hb_handle: Option<JoinHandle<()>>,
    frame_size: usize,
    receive_buffer: Buffer,
    send_buffer: Buffer,
    can_write: bool,
    can_read: bool,
    send_heartbeat: Arc<AtomicBool>,
    poll_timeout: Option<Duration>,
}

impl<T: Source + Read + Write + Send + 'static> IoLoop<T> {
    pub(crate) fn new(
        connection: Connection,
        socket: T,
        poll: Option<(Poll, Token)>,
    ) -> Result<Self> {
        let (poll, registered) = poll
            .map(|t| Ok((t.0, true)))
            .unwrap_or_else(|| Poll::new().map(|poll| (poll, false)))?;
        let frame_size = std::cmp::max(8192, connection.configuration().frame_max() as usize);
        let waker = Arc::new(Waker::new(poll.registry(), WAKER)?);
        let mut inner = Self {
            connection,
            socket,
            status: Status::Initial,
            poll,
            waker,
            hb_handle: None,
            frame_size,
            receive_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            send_buffer: Buffer::with_capacity(FRAMES_STORAGE * frame_size),
            can_write: false,
            can_read: false,
            send_heartbeat: Arc::new(AtomicBool::new(false)),
            poll_timeout: None,
        };
        if registered {
            inner.poll.registry().reregister(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        } else {
            inner.poll.registry().register(
                &mut inner.socket,
                SOCKET,
                Interest::READABLE | Interest::WRITABLE,
            )?;
        }
        Ok(inner)
    }

    fn start_heartbeat(&mut self, interval: Duration) -> Result<()> {
        let connection = self.connection.clone();
        let send_hartbeat = self.send_heartbeat.clone();
        let hb_handle = ThreadBuilder::new()
            .name("heartbeat".to_owned())
            .spawn(move || {
                while connection.status().connected() {
                    let start = Instant::now();
                    let mut remaining = interval;

                    loop {
                        thread::park_timeout(remaining);
                        let elapsed = start.elapsed();
                        if elapsed >= remaining {
                            break;
                        }
                        remaining -= interval - elapsed;
                    }

                    send_hartbeat.store(true, Ordering::Relaxed);
                }
            })?;
        self.hb_handle = Some(hb_handle);
        Ok(())
    }

    fn heartbeat(&mut self) -> Result<()> {
        if self.send_heartbeat.load(Ordering::Relaxed) {
            trace!("send heartbeat");
            self.connection.send_heartbeat()?;
            self.send_heartbeat.store(false, Ordering::Relaxed);
        }
        Ok(())
    }

    fn ensure_setup(&mut self) -> Result<()> {
        if self.status != Status::Setup && self.connection.status().connected() {
            let frame_max = self.connection.configuration().frame_max() as usize;
            self.frame_size = std::cmp::max(self.frame_size, frame_max);
            self.receive_buffer.grow(FRAMES_STORAGE * self.frame_size);
            self.send_buffer.grow(FRAMES_STORAGE * self.frame_size);
            let heartbeat = self.connection.configuration().heartbeat();
            if heartbeat != 0 {
                trace!("io_loop: start heartbeat");
                let heartbeat = Duration::from_secs(u64::from(heartbeat));
                self.start_heartbeat(heartbeat)?;
                self.poll_timeout = Some(heartbeat);
                trace!("io_loop: heartbeat started");
            }
            self.status = Status::Setup;
        }
        Ok(())
    }

    fn has_data(&self) -> bool {
        self.connection.has_pending_frames() || self.send_buffer.available_data() > 0
    }

    fn can_write(&self) -> bool {
        self.can_write
            && self.has_data()
            && !self.connection.status().blocked()
    }

    fn can_read(&self) -> bool {
        self.can_read
    }

    fn can_parse(&self) -> bool {
        self.receive_buffer.available_data() > 0
    }

    fn should_continue(&self) -> bool {
        let connection_status = self.connection.status();
        (self.status == Status::Initial
            || connection_status.connected()
            || connection_status.closing())
            && self.status != Status::Stop
            && !connection_status.errored()
    }

    pub fn start(mut self) -> Result<()> {
        let waker = self.waker.clone();
        self.connection.clone().set_io_loop(
            ThreadBuilder::new()
                .name("io_loop".to_owned())
                .spawn(move || {
                    let mut events = Events::with_capacity(1024);
                    while self.should_continue() {
                        self.run(&mut events)?;
                    }
                    if let Some(hb_handle) = self.hb_handle.take() {
                        hb_handle.thread().unpark();
                        hb_handle.join().expect("heartbeat loop failed");
                    }
                    Ok(())
                })?,
            waker,
        )
    }

    fn poll(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop poll");
        self.poll.poll(events, self.poll_timeout)?;
        trace!("io_loop poll done");
        for event in events.iter() {
            if event.token() == SOCKET {
                if event.is_readable() {
                    self.can_read = true;
                }
                if event.is_writable() {
                    self.can_write = true;
                }
            }
        }
        Ok(())
    }

    fn run(&mut self, events: &mut Events) -> Result<()> {
        trace!("io_loop run");
        self.ensure_setup()?;
        self.poll(events)?;
        self.do_run()
    }

    fn do_run(&mut self) -> Result<()> {
        trace!(
            "io_loop do_run; can_read={}, can_write={}, has_data={}",
            self.can_read,
            self.can_write,
            self.has_data()
        );
        loop {
            self.heartbeat()?;
            self.write()?;
            if self.connection.status().closed() {
                self.status = Status::Stop;
            }
            if self.should_continue() {
                self.read()?;
            }
            self.parse()?;
            self.connection.poll_internal_promises()?;
            if self.stop_looping() {
                self.maybe_continue()?;
                break;
            }
        }
        trace!(
            "io_loop do_run done; can_read={}, can_write={}, has_data={}, status={:?}",
            self.can_read,
            self.can_write,
            self.has_data(),
            self.status
        );
        Ok(())
    }

    fn stop_looping(&self) -> bool {
        !self.can_read()
            || !self.can_write()
            || self.status == Status::Stop
            || self.connection.status().errored()
    }

    fn has_pending_operations(&self) -> bool {
        self.status != Status::Stop && (self.can_read() || self.can_parse() || self.can_write())
    }

    fn maybe_continue(&mut self) -> Result<()> {
        if self.has_pending_operations() {
            trace!(
                "io_loop send continue; can_read={}, can_write={}, has_data={}",
                self.can_read,
                self.can_write,
                self.has_data()
            );
            self.send_continue()?;
        }
        Ok(())
    }

    fn critical_error(&mut self, error: Error) -> Result<()> {
        if let ConnectionState::SentProtocolHeader(resolver, ..) = self.connection.status().state()
        {
            resolver.swear(Err(error.clone()));
            self.status = Status::Stop;
        }
        self.connection.set_error(error.clone())?;
        Err(error)
    }

    fn write(&mut self) -> Result<()> {
        if self.can_write() {
            if let Err(e) = self.write_to_stream() {
                if e.wouldblock() {
                    self.can_write = false
                } else {
                    error!("error writing: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.send_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn read(&mut self) -> Result<()> {
        if self.can_read() {
            if let Err(e) = self.read_from_stream() {
                if e.wouldblock() {
                    self.can_read = false
                } else {
                    error!("error reading: {:?}", e);
                    self.critical_error(e)?;
                }
            }
            self.receive_buffer.shift_unless_available(self.frame_size);
        }
        Ok(())
    }

    fn send_continue(&mut self) -> Result<()> {
        self.waker.wake()?;
        Ok(())
    }

    fn write_to_stream(&mut self) -> Result<()> {
        self.serialize()?;

        let sz = self.socket.write(&self.send_buffer.data())?;

        trace!("wrote {} bytes", sz);
        self.send_buffer.consume(sz);
        if sz > 0 && self.send_buffer.available_data() > 0 {
            // We didn't write all the data yet
            self.send_continue()?;
        }
        Ok(())
    }

    fn read_from_stream(&mut self) -> Result<()> {
        match self.connection.status().state() {
            ConnectionState::Closed => Ok(()),
            ConnectionState::Error => Err(Error::InvalidConnectionState(ConnectionState::Error)),
            _ => {
                self.socket
                    .read(&mut self.receive_buffer.space())
                    .map(|sz| {
                        trace!("read {} bytes", sz);
                        self.receive_buffer.fill(sz);
                    })?;
                Ok(())
            }
        }
    }

    fn serialize(&mut self) -> Result<()> {
        if let Some((send_id, next_msg)) = self.connection.next_frame() {
            trace!("will write to buffer: {:?}", next_msg);
            let checkpoint = self.send_buffer.checkpoint();
            let res = gen_frame(&next_msg)((&mut self.send_buffer).into());
            match res.map(|w| w.into_inner().1) {
                Ok(_) => {
                    self.connection.mark_sent(send_id);
                    Ok(())
                }
                Err(e) => {
                    self.send_buffer.rollback(checkpoint);
                    match e {
                        GenError::BufferTooSmall(_) => {
                            // Requeue msg
                            self.connection.requeue_frame((send_id, next_msg))?;
                            self.send_buffer.shift();
                            Ok(())
                        }
                        e => {
                            error!("error generating frame: {:?}", e);
                            let error = Error::SerialisationError(Arc::new(e));
                            self.connection.set_error(error.clone())?;
                            Err(error)
                        }
                    }
                }
            }
        } else {
            Ok(())
        }
    }

    fn parse(&mut self) -> Result<()> {
        if self.can_parse() {
            if let Some(frame) = self.do_parse()? {
                self.connection.handle_frame(frame)?;
            }
        }
        Ok(())
    }

    fn do_parse(&mut self) -> Result<Option<AMQPFrame>> {
        match parse_frame(self.receive_buffer.data()) {
            Ok((i, f)) => {
                let consumed = self.receive_buffer.data().offset(i);
                self.receive_buffer.consume(consumed);
                Ok(Some(f))
            }
            Err(e) => {
                if e.is_incomplete() {
                    self.receive_buffer.shift();
                    Ok(None)
                } else {
                    error!("parse error: {:?}", e);
                    let error = Error::ParsingError(e);
                    self.connection.set_error(error.clone())?;
                    Err(error)
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// This file is heavily derived from Rust's stdlib, and therefore
// retains the copyright notice below

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

use std::string::{mod, String};

pub static IRC_ASCII_LOWER_MAP: [u8, ..256] = [
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
    0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
    0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
    b' ', b'!', b'"', b'#', b'$', b'%', b'&', b'\'',
    b'(', b')', b'*', b'+', b',', b'-', b'.', b'/',
    b'0', b'1', b'2', b'3', b'4', b'5', b'6', b'7',
    b'8', b'9', b':', b';', b'<', b'=', b'>', b'?',
    b'@',

          b'a', b'b', b'c', b'd', b'e', b'f', b'g',
    b'h', b'i', b'j', b'k', b'l', b'm', b'n', b'o',
    b'p', b'q', b'r', b's', b't', b'u', b'v', b'w',
    b'x', b'y', b'z',

                      b'{', b'|', b'}', b'^', b'_',
    b'`', b'a', b'b', b'c', b'd', b'e', b'f', b'g',
    b'h', b'i', b'j', b'k', b'l', b'm', b'n', b'o',
    b'p', b'q', b'r', b's', b't', b'u', b'v', b'w',
    b'x', b'y', b'z', b'{', b'|', b'}', b'^', 0x7f,
    0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
    0x88, 0x89, 0x8a, 0x8b, 0x8c, 0x8d, 0x8e, 0x8f,
    0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97,
    0x98, 0x99, 0x9a, 0x9b, 0x9c, 0x9d, 0x9e, 0x9f,
    0xa0, 0xa1, 0xa2, 0xa3, 0xa4, 0xa5, 0xa6, 0xa7,
    0xa8, 0xa9, 0xaa, 0xab, 0xac, 0xad, 0xae, 0xaf,
    0xb0, 0xb1, 0xb2, 0xb3, 0xb4, 0xb5, 0xb6, 0xb7,
    0xb8, 0xb9, 0xba, 0xbb, 0xbc, 0xbd, 0xbe, 0xbf,
    0xc0, 0xc1, 0xc2, 0xc3, 0xc4, 0xc5, 0xc6, 0xc7,
    0xc8, 0xc9, 0xca, 0xcb, 0xcc, 0xcd, 0xce, 0xcf,
    0xd0, 0xd1, 0xd2, 0xd3, 0xd4, 0xd5, 0xd6, 0xd7,
    0xd8, 0xd9, 0xda, 0xdb, 0xdc, 0xdd, 0xde, 0xdf,
    0xe0, 0xe1, 0xe2, 0xe3, 0xe4, 0xe5, 0xe6, 0xe7,
    0xe8, 0xe9, 0xea, 0xeb, 0xec, 0xed, 0xee, 0xef,
    0xf0, 0xf1, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7,
    0xf8, 0xf9, 0xfa, 0xfb, 0xfc, 0xfd, 0xfe, 0xff,
];

pub trait IrcAsciiExt<T> for Sized? {
    /// Makes a copy of the string in IRC ASCII lower case:
    /// ASCII letters 'A' to 'Z' are mapped to 'a' to 'z',
    /// and "[]\~" are mapped to "{}|^" respectively,
    /// but all other characters are unchanged.
    fn to_irc_lower(&self) -> T;

    /// Check that two strings are an ASCII case-insensitive match.
    /// Same as `to_irc_lower(a) == to_irc_lower(b)`,
    /// but without allocating and copying temporary strings.
    fn eq_ignore_irc_case(&self, other: &Self) -> bool;
}

pub trait OwnedIrcAsciiExt {
    fn into_irc_lower(self) -> Self;
}

impl IrcAsciiExt<Vec<u8>> for [u8] {
    #[inline]
    fn to_irc_lower(&self) -> Vec<u8> {
        self.iter().map(|&byte| IRC_ASCII_LOWER_MAP[byte as uint]).collect()
    }

    #[inline]
    fn eq_ignore_irc_case(&self, other: &[u8]) -> bool {
        self.len() == other.len() &&
            self.iter().zip(other.iter()).all(
            |(byte_self, byte_other)| {
                IRC_ASCII_LOWER_MAP[*byte_self as uint] ==
                    IRC_ASCII_LOWER_MAP[*byte_other as uint]
            })
    }
}

impl OwnedIrcAsciiExt for Vec<u8> {
    #[inline]
    fn into_irc_lower(mut self) -> Vec<u8> {
        for byte in self.iter_mut() {
            *byte = IRC_ASCII_LOWER_MAP[*byte as uint];
        }
        self
    }
}

impl IrcAsciiExt<String> for str {
    #[inline]
    fn to_irc_lower(&self) -> String {
        // Vec<u8>::to_irc_lower() preserves the UTF-8 invariant.
        unsafe { string::raw::from_utf8(self.as_bytes().to_irc_lower()) }
    }

    #[inline]
    fn eq_ignore_irc_case(&self, other: &str) -> bool {
        self.as_bytes().eq_ignore_irc_case(other.as_bytes())
    }
}

impl OwnedIrcAsciiExt for String {
    #[inline]
    fn into_irc_lower(self) -> String {
        // Vec<u8>::into_irc_lower() preserves the UTF-8 invariant.
        unsafe { string::raw::from_utf8(self.into_bytes().into_irc_lower()) }
    }
}

#[test]
fn test_basics() {
    assert!("[".eq_ignore_irc_case("{"));
    assert!("]".eq_ignore_irc_case("}"));
    assert!("\\".eq_ignore_irc_case("|"));
    assert!("~".eq_ignore_irc_case("^"));

    assert_eq!("[".to_irc_lower()[], "{");
    assert_eq!("]".to_irc_lower()[], "}");
    assert_eq!("\\".to_irc_lower()[], "|");
    assert_eq!("~".to_irc_lower()[], "^");
}
<commit_msg>Fix doc formatting error<commit_after>// This file is heavily derived from Rust's stdlib, and therefore
// retains the copyright notice below

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

use std::string::{mod, String};

pub static IRC_ASCII_LOWER_MAP: [u8, ..256] = [
    0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
    0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
    0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
    0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f,
    b' ', b'!', b'"', b'#', b'$', b'%', b'&', b'\'',
    b'(', b')', b'*', b'+', b',', b'-', b'.', b'/',
    b'0', b'1', b'2', b'3', b'4', b'5', b'6', b'7',
    b'8', b'9', b':', b';', b'<', b'=', b'>', b'?',
    b'@',

          b'a', b'b', b'c', b'd', b'e', b'f', b'g',
    b'h', b'i', b'j', b'k', b'l', b'm', b'n', b'o',
    b'p', b'q', b'r', b's', b't', b'u', b'v', b'w',
    b'x', b'y', b'z',

                      b'{', b'|', b'}', b'^', b'_',
    b'`', b'a', b'b', b'c', b'd', b'e', b'f', b'g',
    b'h', b'i', b'j', b'k', b'l', b'm', b'n', b'o',
    b'p', b'q', b'r', b's', b't', b'u', b'v', b'w',
    b'x', b'y', b'z', b'{', b'|', b'}', b'^', 0x7f,
    0x80, 0x81, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
    0x88, 0x89, 0x8a, 0x8b, 0x8c, 0x8d, 0x8e, 0x8f,
    0x90, 0x91, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97,
    0x98, 0x99, 0x9a, 0x9b, 0x9c, 0x9d, 0x9e, 0x9f,
    0xa0, 0xa1, 0xa2, 0xa3, 0xa4, 0xa5, 0xa6, 0xa7,
    0xa8, 0xa9, 0xaa, 0xab, 0xac, 0xad, 0xae, 0xaf,
    0xb0, 0xb1, 0xb2, 0xb3, 0xb4, 0xb5, 0xb6, 0xb7,
    0xb8, 0xb9, 0xba, 0xbb, 0xbc, 0xbd, 0xbe, 0xbf,
    0xc0, 0xc1, 0xc2, 0xc3, 0xc4, 0xc5, 0xc6, 0xc7,
    0xc8, 0xc9, 0xca, 0xcb, 0xcc, 0xcd, 0xce, 0xcf,
    0xd0, 0xd1, 0xd2, 0xd3, 0xd4, 0xd5, 0xd6, 0xd7,
    0xd8, 0xd9, 0xda, 0xdb, 0xdc, 0xdd, 0xde, 0xdf,
    0xe0, 0xe1, 0xe2, 0xe3, 0xe4, 0xe5, 0xe6, 0xe7,
    0xe8, 0xe9, 0xea, 0xeb, 0xec, 0xed, 0xee, 0xef,
    0xf0, 0xf1, 0xf2, 0xf3, 0xf4, 0xf5, 0xf6, 0xf7,
    0xf8, 0xf9, 0xfa, 0xfb, 0xfc, 0xfd, 0xfe, 0xff,
];

pub trait IrcAsciiExt<T> for Sized? {
    /// Makes a copy of the string in IRC ASCII lower case:
    /// ASCII letters 'A' to 'Z' are mapped to 'a' to 'z',
    /// and "[]\\~" are mapped to "{}|\^" respectively,
    /// but all other characters are unchanged.
    fn to_irc_lower(&self) -> T;

    /// Check that two strings are an ASCII case-insensitive match.
    /// Same as `to_irc_lower(a) == to_irc_lower(b)`,
    /// but without allocating and copying temporary strings.
    fn eq_ignore_irc_case(&self, other: &Self) -> bool;
}

pub trait OwnedIrcAsciiExt {
    fn into_irc_lower(self) -> Self;
}

impl IrcAsciiExt<Vec<u8>> for [u8] {
    #[inline]
    fn to_irc_lower(&self) -> Vec<u8> {
        self.iter().map(|&byte| IRC_ASCII_LOWER_MAP[byte as uint]).collect()
    }

    #[inline]
    fn eq_ignore_irc_case(&self, other: &[u8]) -> bool {
        self.len() == other.len() &&
            self.iter().zip(other.iter()).all(
            |(byte_self, byte_other)| {
                IRC_ASCII_LOWER_MAP[*byte_self as uint] ==
                    IRC_ASCII_LOWER_MAP[*byte_other as uint]
            })
    }
}

impl OwnedIrcAsciiExt for Vec<u8> {
    #[inline]
    fn into_irc_lower(mut self) -> Vec<u8> {
        for byte in self.iter_mut() {
            *byte = IRC_ASCII_LOWER_MAP[*byte as uint];
        }
        self
    }
}

impl IrcAsciiExt<String> for str {
    #[inline]
    fn to_irc_lower(&self) -> String {
        // Vec<u8>::to_irc_lower() preserves the UTF-8 invariant.
        unsafe { string::raw::from_utf8(self.as_bytes().to_irc_lower()) }
    }

    #[inline]
    fn eq_ignore_irc_case(&self, other: &str) -> bool {
        self.as_bytes().eq_ignore_irc_case(other.as_bytes())
    }
}

impl OwnedIrcAsciiExt for String {
    #[inline]
    fn into_irc_lower(self) -> String {
        // Vec<u8>::into_irc_lower() preserves the UTF-8 invariant.
        unsafe { string::raw::from_utf8(self.into_bytes().into_irc_lower()) }
    }
}

#[test]
fn test_basics() {
    assert!("[".eq_ignore_irc_case("{"));
    assert!("]".eq_ignore_irc_case("}"));
    assert!("\\".eq_ignore_irc_case("|"));
    assert!("~".eq_ignore_irc_case("^"));

    assert_eq!("[".to_irc_lower()[], "{");
    assert_eq!("]".to_irc_lower()[], "}");
    assert_eq!("\\".to_irc_lower()[], "|");
    assert_eq!("~".to_irc_lower()[], "^");
}
<|endoftext|>