<commit_before>use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::fmt;
use std::slice::Items;
use std::rc::Rc;

use super::rfc2047::decode_rfc2047;

/// Trait for converting from RFC822 Header values into
/// Rust types.
#[stable]
pub trait FromHeader {
    /// Parse the `value` of the header.
    ///
    /// Returns None if the value failed to be parsed
    fn from_header(value: String) -> Option<Self>;
}

/// Trait for converting from a Rust type into a Header value.
#[stable]
pub trait ToHeader {
    /// Turn the `value` into a String suitable for being used in
    /// a message header.
    ///
    /// Returns None if the value cannot be stringified.
    fn to_header(value: Self) -> Option<String>;
}

/// Trait for converting from a Rust time into a Header value
/// that handles its own folding.
///
/// Be mindful that this trait does not mean that the value will
/// not be folded later, rather that the type returns a value that
/// should not be folded, given that the header value starts so far
/// in to a line.
#[unstable]
pub trait ToFoldedHeader {
    fn to_folded_header(start_pos: uint, value: Self) -> Option<String>;
}

impl<T: ToHeader> ToFoldedHeader for T {
    fn to_folded_header(_: uint, value: T) -> Option<String> {
        // We ignore the start_position because the thing will fold anyway.
        ToHeader::to_header(value)
    }
}

impl FromHeader for String {
    fn from_header(value: String) -> Option<String> {
        #[deriving(Show,Copy)]
        enum ParseState {
            Normal(uint),
            SeenEquals(uint),
            SeenQuestion(uint, uint),
        }

        let mut state = ParseState::Normal(0u);
        let mut decoded = String::new();
        let mut pos = 0u;

        let value_slice = value.as_slice();

        while pos < value.len() {
            let ch_range = value_slice.char_range_at(pos);
            let c = ch_range.ch;

            state = match (state, c) {
                (ParseState::SeenQuestion(start_pos, 4), '=') => {
                    // Go to decode if we've seen enough ?
                    let part_decoded = decode_rfc2047(value_slice.slice(start_pos, ch_range.next));
                    let to_push = match part_decoded {
                        Some(ref s) => s.as_slice(),
                        // Decoding failed, push the undecoded string in.
                        None => value_slice.slice(start_pos, pos),
                    };
                    decoded.push_str(to_push);
                    // Revert us to normal state, but starting at the next character.
                    ParseState::Normal(ch_range.next)
                },
                (ParseState::SeenQuestion(start_pos, count), '?') => {
                    ParseState::SeenQuestion(start_pos, count + 1)
                },
                (ParseState::SeenQuestion(start_pos, count), _) => {
                    if count > 4 {
                        // This isn't a RFC2047 sequence, so go back to a normal.
                        ParseState::Normal(start_pos)
                    } else {
                        state
                    }
                }
                (ParseState::SeenEquals(start_pos), '?') => {
                    ParseState::SeenQuestion(start_pos, 1)
                },
                (ParseState::SeenEquals(start_pos), _) => {
                    // This isn't a RFC2047 sequence, so go back to a normal.
                    ParseState::Normal(start_pos)
                }
                (ParseState::Normal(start_pos), '=') => {
                    if start_pos != pos {
                        // Push all up to the =, if there is stuff to push.
                        decoded.push_str(value_slice.slice(start_pos, pos));
                    }
                    ParseState::SeenEquals(pos)
                },
                (ParseState::Normal(_), _) => state,
            };

            pos = ch_range.next;
        }

        // Don't forget to push on whatever we have left
        let last_start = match state {
            ParseState::Normal(start_pos) => start_pos,
            ParseState::SeenEquals(start_pos) => start_pos,
            ParseState::SeenQuestion(start_pos, _) => start_pos,
        };
        decoded.push_str(value_slice.slice_from(last_start));

        Some(decoded)
    }
}

impl ToHeader for String {
    fn to_header(value: String) -> Option<String> {
        Some(value)
    }
}

impl<'a> ToHeader for &'a str {
    fn to_header(value: &'a str) -> Option<String> {
        Some(value.to_string())
    }
}

/// Represents an RFC 822 Header
#[deriving(PartialEq, Eq, Clone, Hash)]
#[unstable]
pub struct Header {
    /// The name of this header
    pub name: String,
    value: String,
}

impl Header {
    /// Creates a new Header for the given `name` and `value`
    #[unstable]
    pub fn new(name: String, value: String) -> Header {
        Header {
            name: name,
            value: value,
        }
    }

    /// Creates a new Header for the given `name` and `value`,
    /// as converted through the `ToHeader` or `ToFoldedHeader` trait.
    ///
    /// Returns None if the value failed to be converted.
    #[unstable]
    pub fn new_with_value<T: ToFoldedHeader>(name: String, value: T) -> Option<Header> {
        let header_len = name.len() + 2;
        ToFoldedHeader::to_folded_header(header_len, value).map(|val| { Header::new(name.clone(), val) })
    }

    /// Get the value represented by this header, as parsed
    /// into whichever type `T`
    #[unstable]
    pub fn get_value<T: FromHeader>(&self) -> Option<T> {
        FromHeader::from_header(self.value.clone())
    }
}

impl fmt::Show for Header {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}: {}", self.name, self.value)
    }
}

#[unstable]
pub struct HeaderItems<'s> {
    iter: Items<'s, Rc<Header>>
}

impl<'s> HeaderItems<'s> {
    #[unstable]
    fn new(iter: Items<'s, Rc<Header>>) -> HeaderItems<'s> {
        HeaderItems {
            iter: iter
        }
    }
}

impl<'s> Iterator<&'s Header> for HeaderItems<'s> {
    fn next(&mut self) -> Option<&'s Header> {
        match self.iter.next() {
            Some(s) => Some(s.deref()),
            None => None,
        }
    }
}

/// A collection of Headers
#[deriving(Eq,PartialEq)]
#[unstable]
pub struct HeaderMap {
    // We store headers "twice" inside the HeaderMap.
    //
    // The first is as an ordered list of headers,
    // which is used to iterate over.
    ordered_headers: Vec<Rc<Header>>,
    // The second is as a mapping between header names
    // and all of the headers with that name.
    //
    // This allows quick retrival of a header by name.
    headers: HashMap<String, Vec<Rc<Header>>>,
}

impl HeaderMap {
    #[unstable]
    pub fn new() -> HeaderMap {
        HeaderMap {
            ordered_headers: Vec::new(),
            headers: HashMap::new(),
        }
    }

    /// Adds a header to the collection
    #[unstable]
    pub fn insert(&mut self, header: Header) {
        let header_name = header.name.clone();
        let rc = Rc::new(header);
        // Add to the ordered list of headers
        self.ordered_headers.push(rc.clone());

        // and to the mapping between header names and values.
        match self.headers.entry(header_name) {
            Entry::Occupied(mut entry) => {
                entry.get_mut().push(rc.clone());
            },
            Entry::Vacant(entry) => {
                // There haven't been any headers with this name
                // as of yet, so make a new list and push it in.
                let mut header_list = Vec::new();
                header_list.push(rc.clone());
                entry.set(header_list);
            },
        };
    }

    /// Get an Iterator over the collection of headers.
    #[unstable]
    pub fn iter(&self) -> HeaderItems {
        HeaderItems::new(self.ordered_headers.iter())
    }

    /// Get the last value of the header with `name`
    #[unstable]
    pub fn get(&self, name: String) -> Option<&Header> {
        self.headers.get(&name).map(|headers| { headers.last().unwrap() })
                               .map(|rc| { rc.deref() })
    }

    /// Get the last value of the header with `name`, as a decoded type.
    #[unstable]
    pub fn get_value<T: FromHeader>(&self, name: String) -> Option<T> {
        match self.get(name) {
            Some(ref header) => header.get_value(),
            None => None,
        }
    }

    #[unstable]
    /// Get the number of headers within this map.
    pub fn len(&self) -> uint {
        self.ordered_headers.len()
    }

    #[unstable]
    /// Find a list of headers of `name`, `None` if there
    /// are no headers with that name.
    pub fn find(&self, name: &String) -> Option<Vec<&Header>> {
        let headers_rcs = self.headers.get(name);
        if headers_rcs.is_some() {
            Some(headers_rcs.unwrap().iter().map(|rc| { rc.deref() }).collect())
        } else {
            None
        }
    }
}

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

    static SAMPLE_HEADERS: [(&'static str, &'static str), ..4] = [
        ("Test", "Value"),
        ("Test", "Value 2"),
        ("Test-2", "Value 3"),
        ("Test-Multiline", "Foo\nBar"),
    ];

    fn make_sample_headers() -> Vec<Header> {
        SAMPLE_HEADERS.iter().map(|&(name, value)| {
            Header::new(name.to_string(), value.to_string())
        }).collect()
    }

    #[test]
    fn test_header_to_string() {
        let header = Header::new("Test".to_string(), "Value".to_string());
        assert_eq!(header.to_string(), "Test: Value".to_string());
    }

    #[test]
    fn test_string_get_value() {
        struct HeaderTest<'s> {
            input: &'s str,
            result: Option<&'s str>,
        }

        let tests = vec![
            HeaderTest {
                input: "Value",
                result: Some("Value"),
            },
            HeaderTest {
                input: "=?ISO-8859-1?Q?Test=20text?=",
                result: Some("Test text"),
            },
            HeaderTest {
                input: "=?ISO-8859-1?Q?Multiple?= =?utf-8?b?ZW5jb2Rpbmdz?=",
                result: Some("Multiple encodings"),
            },
            HeaderTest {
                input: "Some things with =?utf-8?b?ZW5jb2Rpbmdz?=, other things without.",
                result: Some("Some things with encodings, other things without."),
            },
            HeaderTest {
                input: "Encoding =?utf-8?q?fail",
                result: Some("Encoding =?utf-8?q?fail"),
            },
        ];

        for test in tests.into_iter() {
            let header = Header::new("Test".to_string(), test.input.to_string());
            let string_value = header.get_value::<String>();
            assert_eq!(string_value, test.result.map(|s| { s.to_string() }));
        }
    }

    #[test]
    fn test_to_header_string() {
        let header = Header::new_with_value("Test".to_string(), "Value".to_string()).unwrap();
        let header_value = header.get_value::<String>().unwrap();
        assert_eq!(header_value, "Value".to_string());
    }

    #[test]
    fn test_to_header_str() {
        let header = Header::new_with_value("Test".to_string(), "Value").unwrap();
        let header_value = header.get_value::<String>().unwrap();
        assert_eq!(header_value, "Value".to_string());
    }

    #[test]
    fn test_header_map_len() {
        let mut headers = HeaderMap::new();
        for (i, header) in make_sample_headers().into_iter().enumerate() {
            headers.insert(header);
            assert_eq!(headers.len(), i + 1);
        }
    }
    #[test]
    fn test_header_map_iter() {
        let mut headers = HeaderMap::new();
        let mut expected_headers = HashSet::new();
        for header in make_sample_headers().into_iter() {
            headers.insert(header.clone());
            expected_headers.insert(header);
        }

        let mut count = 0u;
        // Ensure all the headers returned are expected
        for header in headers.iter() {
            assert!(expected_headers.contains(header));
            count += 1;
        }
        // And that there is the right number of them
        assert_eq!(count, expected_headers.len());
    }
}
<commit_msg>Add FromHeader impls for DateTime.<commit_after>use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::fmt;
use std::slice::Items;
use std::rc::Rc;

use chrono::{
    DateTime,
    FixedOffset,
    Offset,
    UTC,
};

use super::rfc2047::decode_rfc2047;
use super::rfc822::Rfc822DateParser;

/// Trait for converting from RFC822 Header values into
/// Rust types.
#[stable]
pub trait FromHeader {
    /// Parse the `value` of the header.
    ///
    /// Returns None if the value failed to be parsed
    fn from_header(value: String) -> Option<Self>;
}

/// Trait for converting from a Rust type into a Header value.
#[stable]
pub trait ToHeader {
    /// Turn the `value` into a String suitable for being used in
    /// a message header.
    ///
    /// Returns None if the value cannot be stringified.
    fn to_header(value: Self) -> Option<String>;
}

/// Trait for converting from a Rust time into a Header value
/// that handles its own folding.
///
/// Be mindful that this trait does not mean that the value will
/// not be folded later, rather that the type returns a value that
/// should not be folded, given that the header value starts so far
/// in to a line.
#[unstable]
pub trait ToFoldedHeader {
    fn to_folded_header(start_pos: uint, value: Self) -> Option<String>;
}

impl<T: ToHeader> ToFoldedHeader for T {
    fn to_folded_header(_: uint, value: T) -> Option<String> {
        // We ignore the start_position because the thing will fold anyway.
        ToHeader::to_header(value)
    }
}

impl FromHeader for String {
    fn from_header(value: String) -> Option<String> {
        #[deriving(Show,Copy)]
        enum ParseState {
            Normal(uint),
            SeenEquals(uint),
            SeenQuestion(uint, uint),
        }

        let mut state = ParseState::Normal(0u);
        let mut decoded = String::new();
        let mut pos = 0u;

        let value_slice = value.as_slice();

        while pos < value.len() {
            let ch_range = value_slice.char_range_at(pos);
            let c = ch_range.ch;

            state = match (state, c) {
                (ParseState::SeenQuestion(start_pos, 4), '=') => {
                    // Go to decode if we've seen enough ?
                    let part_decoded = decode_rfc2047(value_slice.slice(start_pos, ch_range.next));
                    let to_push = match part_decoded {
                        Some(ref s) => s.as_slice(),
                        // Decoding failed, push the undecoded string in.
                        None => value_slice.slice(start_pos, pos),
                    };
                    decoded.push_str(to_push);
                    // Revert us to normal state, but starting at the next character.
                    ParseState::Normal(ch_range.next)
                },
                (ParseState::SeenQuestion(start_pos, count), '?') => {
                    ParseState::SeenQuestion(start_pos, count + 1)
                },
                (ParseState::SeenQuestion(start_pos, count), _) => {
                    if count > 4 {
                        // This isn't a RFC2047 sequence, so go back to a normal.
                        ParseState::Normal(start_pos)
                    } else {
                        state
                    }
                }
                (ParseState::SeenEquals(start_pos), '?') => {
                    ParseState::SeenQuestion(start_pos, 1)
                },
                (ParseState::SeenEquals(start_pos), _) => {
                    // This isn't a RFC2047 sequence, so go back to a normal.
                    ParseState::Normal(start_pos)
                }
                (ParseState::Normal(start_pos), '=') => {
                    if start_pos != pos {
                        // Push all up to the =, if there is stuff to push.
                        decoded.push_str(value_slice.slice(start_pos, pos));
                    }
                    ParseState::SeenEquals(pos)
                },
                (ParseState::Normal(_), _) => state,
            };

            pos = ch_range.next;
        }

        // Don't forget to push on whatever we have left
        let last_start = match state {
            ParseState::Normal(start_pos) => start_pos,
            ParseState::SeenEquals(start_pos) => start_pos,
            ParseState::SeenQuestion(start_pos, _) => start_pos,
        };
        decoded.push_str(value_slice.slice_from(last_start));

        Some(decoded)
    }
}

impl FromHeader for DateTime<FixedOffset> {
    fn from_header(value: String) -> Option<DateTime<FixedOffset>> {
        let mut parser = Rfc822DateParser::new(value.as_slice());
        parser.consume_datetime()
    }
}

impl FromHeader for DateTime<UTC> {
    fn from_header(value: String) -> Option<DateTime<UTC>> {
        let dt: Option<DateTime<FixedOffset>> = FromHeader::from_header(value);
        dt.map(|i| i.with_offset(UTC))
    }
}

impl ToHeader for String {
    fn to_header(value: String) -> Option<String> {
        Some(value)
    }
}

impl<'a> ToHeader for &'a str {
    fn to_header(value: &'a str) -> Option<String> {
        Some(value.to_string())
    }
}

/// Represents an RFC 822 Header
#[deriving(PartialEq, Eq, Clone, Hash)]
#[unstable]
pub struct Header {
    /// The name of this header
    pub name: String,
    value: String,
}

impl Header {
    /// Creates a new Header for the given `name` and `value`
    #[unstable]
    pub fn new(name: String, value: String) -> Header {
        Header {
            name: name,
            value: value,
        }
    }

    /// Creates a new Header for the given `name` and `value`,
    /// as converted through the `ToHeader` or `ToFoldedHeader` trait.
    ///
    /// Returns None if the value failed to be converted.
    #[unstable]
    pub fn new_with_value<T: ToFoldedHeader>(name: String, value: T) -> Option<Header> {
        let header_len = name.len() + 2;
        ToFoldedHeader::to_folded_header(header_len, value).map(|val| { Header::new(name.clone(), val) })
    }

    /// Get the value represented by this header, as parsed
    /// into whichever type `T`
    #[unstable]
    pub fn get_value<T: FromHeader>(&self) -> Option<T> {
        FromHeader::from_header(self.value.clone())
    }
}

impl fmt::Show for Header {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}: {}", self.name, self.value)
    }
}

#[unstable]
pub struct HeaderItems<'s> {
    iter: Items<'s, Rc<Header>>
}

impl<'s> HeaderItems<'s> {
    #[unstable]
    fn new(iter: Items<'s, Rc<Header>>) -> HeaderItems<'s> {
        HeaderItems {
            iter: iter
        }
    }
}

impl<'s> Iterator<&'s Header> for HeaderItems<'s> {
    fn next(&mut self) -> Option<&'s Header> {
        match self.iter.next() {
            Some(s) => Some(s.deref()),
            None => None,
        }
    }
}

/// A collection of Headers
#[deriving(Eq,PartialEq)]
#[unstable]
pub struct HeaderMap {
    // We store headers "twice" inside the HeaderMap.
    //
    // The first is as an ordered list of headers,
    // which is used to iterate over.
    ordered_headers: Vec<Rc<Header>>,
    // The second is as a mapping between header names
    // and all of the headers with that name.
    //
    // This allows quick retrival of a header by name.
    headers: HashMap<String, Vec<Rc<Header>>>,
}

impl HeaderMap {
    #[unstable]
    pub fn new() -> HeaderMap {
        HeaderMap {
            ordered_headers: Vec::new(),
            headers: HashMap::new(),
        }
    }

    /// Adds a header to the collection
    #[unstable]
    pub fn insert(&mut self, header: Header) {
        let header_name = header.name.clone();
        let rc = Rc::new(header);
        // Add to the ordered list of headers
        self.ordered_headers.push(rc.clone());

        // and to the mapping between header names and values.
        match self.headers.entry(header_name) {
            Entry::Occupied(mut entry) => {
                entry.get_mut().push(rc.clone());
            },
            Entry::Vacant(entry) => {
                // There haven't been any headers with this name
                // as of yet, so make a new list and push it in.
                let mut header_list = Vec::new();
                header_list.push(rc.clone());
                entry.set(header_list);
            },
        };
    }

    /// Get an Iterator over the collection of headers.
    #[unstable]
    pub fn iter(&self) -> HeaderItems {
        HeaderItems::new(self.ordered_headers.iter())
    }

    /// Get the last value of the header with `name`
    #[unstable]
    pub fn get(&self, name: String) -> Option<&Header> {
        self.headers.get(&name).map(|headers| { headers.last().unwrap() })
                               .map(|rc| { rc.deref() })
    }

    /// Get the last value of the header with `name`, as a decoded type.
    #[unstable]
    pub fn get_value<T: FromHeader>(&self, name: String) -> Option<T> {
        match self.get(name) {
            Some(ref header) => header.get_value(),
            None => None,
        }
    }

    #[unstable]
    /// Get the number of headers within this map.
    pub fn len(&self) -> uint {
        self.ordered_headers.len()
    }

    #[unstable]
    /// Find a list of headers of `name`, `None` if there
    /// are no headers with that name.
    pub fn find(&self, name: &String) -> Option<Vec<&Header>> {
        let headers_rcs = self.headers.get(name);
        if headers_rcs.is_some() {
            Some(headers_rcs.unwrap().iter().map(|rc| { rc.deref() }).collect())
        } else {
            None
        }
    }
}

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

    use chrono::{
        DateTime,
        FixedOffset,
        Offset,
        UTC,
    };

    static SAMPLE_HEADERS: [(&'static str, &'static str), ..4] = [
        ("Test", "Value"),
        ("Test", "Value 2"),
        ("Test-2", "Value 3"),
        ("Test-Multiline", "Foo\nBar"),
    ];

    fn make_sample_headers() -> Vec<Header> {
        SAMPLE_HEADERS.iter().map(|&(name, value)| {
            Header::new(name.to_string(), value.to_string())
        }).collect()
    }

    #[test]
    fn test_header_to_string() {
        let header = Header::new("Test".to_string(), "Value".to_string());
        assert_eq!(header.to_string(), "Test: Value".to_string());
    }

    #[test]
    fn test_string_get_value() {
        struct HeaderTest<'s> {
            input: &'s str,
            result: Option<&'s str>,
        }

        let tests = vec![
            HeaderTest {
                input: "Value",
                result: Some("Value"),
            },
            HeaderTest {
                input: "=?ISO-8859-1?Q?Test=20text?=",
                result: Some("Test text"),
            },
            HeaderTest {
                input: "=?ISO-8859-1?Q?Multiple?= =?utf-8?b?ZW5jb2Rpbmdz?=",
                result: Some("Multiple encodings"),
            },
            HeaderTest {
                input: "Some things with =?utf-8?b?ZW5jb2Rpbmdz?=, other things without.",
                result: Some("Some things with encodings, other things without."),
            },
            HeaderTest {
                input: "Encoding =?utf-8?q?fail",
                result: Some("Encoding =?utf-8?q?fail"),
            },
        ];

        for test in tests.into_iter() {
            let header = Header::new("Test".to_string(), test.input.to_string());
            let string_value = header.get_value::<String>();
            assert_eq!(string_value, test.result.map(|s| { s.to_string() }));
        }
    }

    #[test]
    fn test_datetime_get_value() {
        let header = Header::new("Date".to_string(), "Wed, 17 Dec 2014 09:35:07 +0100".to_string());
        let dt_value = header.get_value::<DateTime<FixedOffset>>();
        assert_eq!(dt_value, Some(FixedOffset::east(3600).ymd(2014, 12, 17).and_hms(9, 35, 7)));
    }

    #[test]
    fn test_datetime_utc_get_value() {
        let header = Header::new("Date".to_string(), "Wed, 17 Dec 2014 09:35:07 +0100".to_string());
        let dt_value = header.get_value::<DateTime<UTC>>();
        assert_eq!(dt_value, Some(UTC.ymd(2014, 12, 17).and_hms(8, 35, 7)));
    }

    #[test]
    fn test_to_header_string() {
        let header = Header::new_with_value("Test".to_string(), "Value".to_string()).unwrap();
        let header_value = header.get_value::<String>().unwrap();
        assert_eq!(header_value, "Value".to_string());
    }

    #[test]
    fn test_to_header_str() {
        let header = Header::new_with_value("Test".to_string(), "Value").unwrap();
        let header_value = header.get_value::<String>().unwrap();
        assert_eq!(header_value, "Value".to_string());
    }

    #[test]
    fn test_header_map_len() {
        let mut headers = HeaderMap::new();
        for (i, header) in make_sample_headers().into_iter().enumerate() {
            headers.insert(header);
            assert_eq!(headers.len(), i + 1);
        }
    }
    #[test]
    fn test_header_map_iter() {
        let mut headers = HeaderMap::new();
        let mut expected_headers = HashSet::new();
        for header in make_sample_headers().into_iter() {
            headers.insert(header.clone());
            expected_headers.insert(header);
        }

        let mut count = 0u;
        // Ensure all the headers returned are expected
        for header in headers.iter() {
            assert!(expected_headers.contains(header));
            count += 1;
        }
        // And that there is the right number of them
        assert_eq!(count, expected_headers.len());
    }
}
<|endoftext|>


<commit_before>use std::time::Duration;

use futures::*;
use tokio_timer::*;
use tokio_core::reactor::Core;
use hyper::Client;
use hyper_tls::HttpsConnector;

use pool::Pool;

pub struct HealthCheck {
    interval: Duration,
    pool: Pool,
    uri_path: String,
}

impl HealthCheck {
    pub fn new(interval: Duration, pool: Pool, uri_path: String) -> HealthCheck {

        HealthCheck {
            interval: interval,
            pool: pool,
            uri_path: uri_path,
        }
    }

    pub fn run(&self) {
        let mut core = Core::new().unwrap();
        let timer = Timer::default();

        let handle = core.handle();
        let client = Client::configure()
            .connector(HttpsConnector::new(4, &handle))
            .build(&handle);

        let pool = self.pool.clone();
        let work = timer.interval(self.interval).for_each(move |()| {
            let servers = pool.all();
            for server in servers {
                let url = server.url().join(&self.uri_path).unwrap();
                debug!("Health check {:?}", url);

                let pool1 = pool.clone();
                let pool2 = pool.clone();
                let server1 = server.clone();
                let server2 = server.clone();
                let req = client.get(url).and_then(move |res| {
                    debug!("Response: {}", res.status());
                    debug!("Headers: \n{}", res.headers());

                    if ! res.status().is_success() {
                        info!("Removing {:?} from pool", server1);
                        pool1.remove(&server1);
                    }

                    ::futures::finished(())
                }).map_err(move |e| {
                    error!("Error connecting to backend: {:?}", e);
                    info!("Removing {:?} from pool", server2);
                    pool2.remove(&server2);
                });

                handle.spawn(req);
            }

            Ok(())
        }).map_err(|_| {});

        core.run(work).unwrap();
    }
}
<commit_msg>#71 Should now bubble errors up the scope in health.rs (#81)<commit_after>use std::time::Duration;

use futures::*;
use tokio_timer::*;
use tokio_core::reactor::Core;
use hyper::Client;
use hyper_tls::HttpsConnector;
use std::io;

use pool::Pool;

pub struct HealthCheck {
    interval: Duration,
    pool: Pool,
    uri_path: String,
}

impl HealthCheck {
    pub fn new(interval: Duration, pool: Pool, uri_path: String) -> HealthCheck {

        HealthCheck {
            interval: interval,
            pool: pool,
            uri_path: uri_path,
        }
    }

    pub fn run(&self) -> io::Result<()> {
        let mut core = try!(Core::new());
        let timer = Timer::default();

        let handle = core.handle();
        let client = Client::configure()
            .connector(HttpsConnector::new(4, &handle))
            .build(&handle);

        let pool = self.pool.clone();
        let work = timer.interval(self.interval).for_each(move |()| {
            let servers = pool.all();
            for server in servers {
                let url = match server.url().join(&self.uri_path) {
                Ok(url) => url,
                Err(e) => {
                    error!("Invalid health check url: {:?}",e);
                    return Ok(());
                }
            };
                debug!("Health check {:?}", url);

                let pool1 = pool.clone();
                let pool2 = pool.clone();
                let server1 = server.clone();
                let server2 = server.clone();
                let req = client.get(url).and_then(move |res| {
                    debug!("Response: {}", res.status());
                    debug!("Headers: \n{}", res.headers());

                    if ! res.status().is_success() {
                        info!("Removing {:?} from pool", server1);
                        pool1.remove(&server1);
                    }

                    ::futures::finished(())
                }).map_err(move |e| {
                    error!("Error connecting to backend: {:?}", e);
                    info!("Removing {:?} from pool", server2);
                    pool2.remove(&server2);
                });

                handle.spawn(req);
            }

            Ok(())
        }).map_err(From::from);

        core.run(work)
    }
}
<|endoftext|>


<commit_before>//! Encoding of AVIF images.
///
/// The [AVIF] specification defines an image derivative of the AV1 bitstream, an open video codec.
///
/// [AVIF]: https://aomediacodec.github.io/av1-avif/
use std::borrow::Cow;
use std::io::Write;
use std::cmp::min;

use crate::{ColorType, ImageBuffer, ImageFormat, Pixel};
use crate::{ImageError, ImageResult};
use crate::buffer::ConvertBuffer;
use crate::color::{FromColor, Luma, LumaA, Bgr, Bgra, Rgb, Rgba};
use crate::error::{EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind};

use bytemuck::{Pod, PodCastError, try_cast_slice, try_cast_slice_mut};
use num_traits::Zero;
use ravif::{Img, Config, RGBA8, encode_rgba};
use rgb::AsPixels;

/// AVIF Encoder.
///
/// Writes one image into the chosen output.
pub struct AvifEncoder<W> {
    inner: W,
    fallback: Vec<u8>,
    config: Config
}

/// An enumeration over supported AVIF color spaces
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum ColorSpace {
    /// sRGB colorspace
    Srgb,
    /// BT.709 colorspace
    Bt709,

    #[doc(hidden)]
    __NonExhaustive(crate::utils::NonExhaustiveMarker),
}

impl ColorSpace {
    fn to_ravif(self) -> ravif::ColorSpace {
        match self {
            Self::Srgb => ravif::ColorSpace::RGB,
            Self::Bt709 => ravif::ColorSpace::YCbCr,
            Self::__NonExhaustive(marker) => match marker._private {},
        }
    }
}

impl<W: Write> AvifEncoder<W> {
    /// Create a new encoder that writes its output to `w`.
    pub fn new(w: W) -> Self {
        AvifEncoder::new_with_speed_quality(w, 1, 100)
    }

    /// Create a new encoder with specified speed and quality, that writes its output to `w`.
    /// `speed` accepts a value in the range 0-10, where 0 is the slowest and 10 is the fastest.
    /// `quality` accepts a value in the range 0-100, where 0 is the worst and 100 is the best.
    pub fn new_with_speed_quality(w: W, speed: u8, quality: u8) -> Self {
        // Clamp quality and speed to range
        let quality = min(quality, 100);
        let speed = min(speed, 10);

        AvifEncoder {
            inner: w,
            fallback: vec![],
            config: Config {
                quality,
                alpha_quality: quality,
                speed,
                premultiplied_alpha: false,
                color_space: ravif::ColorSpace::RGB,
                // match core count
                threads: 0,
            }
        }
    }

    /// Encode with the specified `color_space`.
    pub fn with_colorspace(mut self, color_space: ColorSpace) -> Self {
        self.config.color_space = color_space.to_ravif();
        self
    }

    /// Encode image data with the indicated color type.
    ///
    /// The encoder currently requires all data to be RGBA8, it will be converted internally if
    /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the
    /// conversion may be more efficient.
    pub fn write_image(mut self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()> {
        self.set_color(color);
        let config = self.config;
        // `ravif` needs strongly typed data so let's convert. We can either use a temporarily
        // owned version in our own buffer or zero-copy if possible by using the input buffer.
        // This requires going through `rgb`.
        let buffer = self.encode_as_img(data, width, height, color)?;
        let (data, _color_size, _alpha_size) = encode_rgba(buffer, &config)
            .map_err(|err| ImageError::Encoding(
                EncodingError::new(ImageFormat::Avif.into(), err)
            ))?;
        self.inner.write_all(&data)?;
        Ok(())
    }

    // Does not currently do anything. Mirrors behaviour of old config function.
    fn set_color(&mut self, _color: ColorType) {
        // self.config.color_space = ColorSpace::RGB;
    }

    fn encode_as_img<'buf>(&'buf mut self, data: &'buf [u8], width: u32, height: u32, color: ColorType)
        -> ImageResult<Img<&'buf [RGBA8]>>
    {
        // Error wrapping utility for color dependent buffer dimensions.
        fn try_from_raw<P: Pixel + 'static>(data: &[P::Subpixel], width: u32, height: u32)
            -> ImageResult<ImageBuffer<P, &[P::Subpixel]>>
        {
            ImageBuffer::from_raw(width, height, data).ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))
            })
        };

        // Convert to target color type using few buffer allocations.
        fn convert_into<'buf, P>(buf: &'buf mut Vec<u8>, image: ImageBuffer<P, &[P::Subpixel]>)
            -> Img<&'buf [RGBA8]>
        where
            P: Pixel + 'static,
            Rgba<u8>: FromColor<P>,
        {
            let (width, height) = image.dimensions();
            // TODO: conversion re-using the target buffer?
            let image: ImageBuffer<Rgba<u8>, _> = image.convert();
            *buf = image.into_raw();
            Img::new(buf.as_pixels(), width as usize, height as usize)
        }

        // Cast the input slice using few buffer allocations if possible.
        // In particular try not to allocate if the caller did the infallible reverse.
        fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<[Channel]>>
        where
            Channel: Pod + Zero,
        {
            match try_cast_slice(buf) {
                Ok(slice) => Ok(Cow::Borrowed(slice)),
                Err(PodCastError::OutputSliceWouldHaveSlop) => {
                    Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                }
                Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => {
                    // Sad, but let's allocate.
                    // bytemuck checks alignment _before_ slop but size mismatch before this..
                    if buf.len() % std::mem::size_of::<Channel>() != 0 {
                        Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                    } else {
                        let len = buf.len() / std::mem::size_of::<Channel>();
                        let mut data = vec![Channel::zero(); len];
                        let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap();
                        view.copy_from_slice(buf);
                        Ok(Cow::Owned(data))
                    }
                }
                Err(err) => { // Are you trying to encode a ZST??
                    Err(ImageError::Parameter(ParameterError::from_kind(
                        ParameterErrorKind::Generic(format!("{:?}", err))
                    )))
                }
            }
        }

        match color {
            ColorType::Rgba8 => {
                // ravif doesn't do any checks but has some asserts, so we do the checks.
                let img = try_from_raw::<Rgba<u8>>(data, width, height)?;
                // Now, internally ravif uses u32 but it takes usize. We could do some checked
                // conversion but instead we use that a non-empty image must be addressable.
                if img.pixels().len() == 0 {
                    return Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)));
                }

                Ok(Img::new(rgb::AsPixels::as_pixels(data), width as usize, height as usize))
            },
            // we need a separate buffer..
            ColorType::L8 => {
                let image = try_from_raw::<Luma<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La8 => {
                let image = try_from_raw::<LumaA<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb8 => {
                let image = try_from_raw::<Rgb<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgr8 => {
                let image = try_from_raw::<Bgr<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Bgra8 => {
                let image = try_from_raw::<Bgra<u8>>(data, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // we need to really convert data..
            ColorType::L16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::La16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgb16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            ColorType::Rgba16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?;
                Ok(convert_into(&mut self.fallback, image))
            }
            // for cases we do not support at all?
            _ => Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                    ImageFormat::Avif.into(),
                    UnsupportedErrorKind::Color(color.into()),
                )))
        }
    }
}
<commit_msg>Use the new ravif::encode_rgb function for encoding Rgb8 images<commit_after>//! Encoding of AVIF images.
///
/// The [AVIF] specification defines an image derivative of the AV1 bitstream, an open video codec.
///
/// [AVIF]: https://aomediacodec.github.io/av1-avif/
use std::borrow::Cow;
use std::io::Write;
use std::cmp::min;

use crate::{ColorType, ImageBuffer, ImageFormat, Pixel};
use crate::{ImageError, ImageResult};
use crate::buffer::ConvertBuffer;
use crate::color::{FromColor, Luma, LumaA, Bgr, Bgra, Rgb, Rgba};
use crate::error::{EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind};

use bytemuck::{Pod, PodCastError, try_cast_slice, try_cast_slice_mut};
use num_traits::Zero;
use ravif::{encode_rgba, encode_rgb, Config, Img, RGB8, RGBA8};
use rgb::AsPixels;

/// AVIF Encoder.
///
/// Writes one image into the chosen output.
pub struct AvifEncoder<W> {
    inner: W,
    fallback: Vec<u8>,
    config: Config
}

/// An enumeration over supported AVIF color spaces
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum ColorSpace {
    /// sRGB colorspace
    Srgb,
    /// BT.709 colorspace
    Bt709,

    #[doc(hidden)]
    __NonExhaustive(crate::utils::NonExhaustiveMarker),
}

impl ColorSpace {
    fn to_ravif(self) -> ravif::ColorSpace {
        match self {
            Self::Srgb => ravif::ColorSpace::RGB,
            Self::Bt709 => ravif::ColorSpace::YCbCr,
            Self::__NonExhaustive(marker) => match marker._private {},
        }
    }
}

enum RgbColor<'buf> {
    Rgb8(Img<&'buf [RGB8]>),
    Rgba8(Img<&'buf [RGBA8]>),
}

impl<W: Write> AvifEncoder<W> {
    /// Create a new encoder that writes its output to `w`.
    pub fn new(w: W) -> Self {
        AvifEncoder::new_with_speed_quality(w, 1, 100)
    }

    /// Create a new encoder with specified speed and quality, that writes its output to `w`.
    /// `speed` accepts a value in the range 0-10, where 0 is the slowest and 10 is the fastest.
    /// `quality` accepts a value in the range 0-100, where 0 is the worst and 100 is the best.
    pub fn new_with_speed_quality(w: W, speed: u8, quality: u8) -> Self {
        // Clamp quality and speed to range
        let quality = min(quality, 100);
        let speed = min(speed, 10);

        AvifEncoder {
            inner: w,
            fallback: vec![],
            config: Config {
                quality,
                alpha_quality: quality,
                speed,
                premultiplied_alpha: false,
                color_space: ravif::ColorSpace::RGB,
                // match core count
                threads: 0,
            }
        }
    }

    /// Encode with the specified `color_space`.
    pub fn with_colorspace(mut self, color_space: ColorSpace) -> Self {
        self.config.color_space = color_space.to_ravif();
        self
    }

    /// Encode image data with the indicated color type.
    ///
    /// The encoder currently requires all data to be RGBA8, it will be converted internally if
    /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the
    /// conversion may be more efficient.
    pub fn write_image(
        mut self,
        data: &[u8],
        width: u32,
        height: u32,
        color: ColorType,
    ) -> ImageResult<()> {
        self.set_color(color);
        let config = self.config;
        // `ravif` needs strongly typed data so let's convert. We can either use a temporarily
        // owned version in our own buffer or zero-copy if possible by using the input buffer.
        // This requires going through `rgb`.
        let result = match self.encode_as_img(data, width, height, color)? {
            RgbColor::Rgb8(buffer) => encode_rgb(buffer, &config).map(|(data, _color_size)| data),
            RgbColor::Rgba8(buffer) => {
                encode_rgba(buffer, &config).map(|(data, _color_size, _alpha_size)| data)
            }
        };
        let data = result.map_err(|err| {
            ImageError::Encoding(EncodingError::new(ImageFormat::Avif.into(), err))
        })?;
        self.inner.write_all(&data)?;
        Ok(())
    }

    // Does not currently do anything. Mirrors behaviour of old config function.
    fn set_color(&mut self, _color: ColorType) {
        // self.config.color_space = ColorSpace::RGB;
    }

    fn encode_as_img<'buf>(
        &'buf mut self,
        data: &'buf [u8],
        width: u32,
        height: u32,
        color: ColorType,
    ) -> ImageResult<RgbColor<'buf>> {
        // Error wrapping utility for color dependent buffer dimensions.
        fn try_from_raw<P: Pixel + 'static>(data: &[P::Subpixel], width: u32, height: u32)
            -> ImageResult<ImageBuffer<P, &[P::Subpixel]>>
        {
            ImageBuffer::from_raw(width, height, data).ok_or_else(|| {
                ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))
            })
        };

        // Convert to target color type using few buffer allocations.
        fn convert_into<'buf, P>(
            buf: &'buf mut Vec<u8>,
            image: ImageBuffer<P, &[P::Subpixel]>,
        ) -> Img<&'buf [RGBA8]>
        where
            P: Pixel + 'static,
            Rgba<u8>: FromColor<P>,
        {
            let (width, height) = image.dimensions();
            // TODO: conversion re-using the target buffer?
            let image: ImageBuffer<Rgba<u8>, _> = image.convert();
            *buf = image.into_raw();
            Img::new(buf.as_pixels(), width as usize, height as usize)
        }

        // Cast the input slice using few buffer allocations if possible.
        // In particular try not to allocate if the caller did the infallible reverse.
        fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<[Channel]>>
        where
            Channel: Pod + Zero,
        {
            match try_cast_slice(buf) {
                Ok(slice) => Ok(Cow::Borrowed(slice)),
                Err(PodCastError::OutputSliceWouldHaveSlop) => {
                    Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                }
                Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => {
                    // Sad, but let's allocate.
                    // bytemuck checks alignment _before_ slop but size mismatch before this..
                    if buf.len() % std::mem::size_of::<Channel>() != 0 {
                        Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)))
                    } else {
                        let len = buf.len() / std::mem::size_of::<Channel>();
                        let mut data = vec![Channel::zero(); len];
                        let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap();
                        view.copy_from_slice(buf);
                        Ok(Cow::Owned(data))
                    }
                }
                Err(err) => { // Are you trying to encode a ZST??
                    Err(ImageError::Parameter(ParameterError::from_kind(
                        ParameterErrorKind::Generic(format!("{:?}", err))
                    )))
                }
            }
        }

        match color {
            ColorType::Rgb8 => {
                // ravif doesn't do any checks but has some asserts, so we do the checks.
                let img = try_from_raw::<Rgb<u8>>(data, width, height)?;
                // Now, internally ravif uses u32 but it takes usize. We could do some checked
                // conversion but instead we use that a non-empty image must be addressable.
                if img.pixels().len() == 0 {
                    return Err(ImageError::Parameter(ParameterError::from_kind(
                        ParameterErrorKind::DimensionMismatch,
                    )));
                }

                Ok(RgbColor::Rgb8(Img::new(
                    rgb::AsPixels::as_pixels(data),
                    width as usize,
                    height as usize,
                )))
            }
            ColorType::Rgba8 => {
                // ravif doesn't do any checks but has some asserts, so we do the checks.
                let img = try_from_raw::<Rgba<u8>>(data, width, height)?;
                // Now, internally ravif uses u32 but it takes usize. We could do some checked
                // conversion but instead we use that a non-empty image must be addressable.
                if img.pixels().len() == 0 {
                    return Err(ImageError::Parameter(ParameterError::from_kind(
                        ParameterErrorKind::DimensionMismatch,
                    )));
                }

                Ok(RgbColor::Rgba8(Img::new(
                    rgb::AsPixels::as_pixels(data),
                    width as usize,
                    height as usize,
                )))
            }
            // we need a separate buffer..
            ColorType::L8 => {
                let image = try_from_raw::<Luma<u8>>(data, width, height)?;
                Ok(RgbColor::Rgba8(convert_into(&mut self.fallback, image)))
            }
            ColorType::La8 => {
                let image = try_from_raw::<LumaA<u8>>(data, width, height)?;
                Ok(RgbColor::Rgba8(convert_into(&mut self.fallback, image)))
            }
            ColorType::Bgr8 => {
                let image = try_from_raw::<Bgr<u8>>(data, width, height)?;
                Ok(RgbColor::Rgba8(convert_into(&mut self.fallback, image)))
            }
            ColorType::Bgra8 => {
                let image = try_from_raw::<Bgra<u8>>(data, width, height)?;
                Ok(RgbColor::Rgba8(convert_into(&mut self.fallback, image)))
            }
            // we need to really convert data..
            ColorType::L16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?;
                Ok(RgbColor::Rgba8(convert_into(&mut self.fallback, image)))
            }
            ColorType::La16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?;
                Ok(RgbColor::Rgba8(convert_into(&mut self.fallback, image)))
            }
            ColorType::Rgb16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?;
                Ok(RgbColor::Rgba8(convert_into(&mut self.fallback, image)))
            }
            ColorType::Rgba16 => {
                let buffer = cast_buffer(data)?;
                let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?;
                Ok(RgbColor::Rgba8(convert_into(&mut self.fallback, image)))
            }
            // for cases we do not support at all?
            _ => Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind(
                    ImageFormat::Avif.into(),
                    UnsupportedErrorKind::Color(color.into()),
                )))
        }
    }
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::io::{self, Read, Write};
use std::net::SocketAddr;
use std::rc::Rc;
use std::time::{Duration, Instant};

use bytes::BytesMut;
use bytes::buf::FromBuf;
use futures::{Future, Stream};
use mio::{self, Evented, PollOpt, Ready, Registration, SetReadiness, Token};
use tokio_core::reactor::{Handle, Timeout};

use session::{KcpSessionUpdater, SharedKcpSession};
use skcp::SharedKcp;

pub enum KcpIoMode {
    Client,
    Server(Duration),
}

pub struct KcpIo {
    kcp: SharedKcp,
    registration: Registration,
    readiness: SetReadiness,
    last_update: Rc<RefCell<Instant>>,
    close_flag: Rc<RefCell<bool>>,
}

impl Drop for KcpIo {
    fn drop(&mut self) {
        let mut cf = self.close_flag.borrow_mut();
        *cf = true;
    }
}

impl KcpIo {
    pub fn new(shared_kcp: SharedKcp,
               addr: SocketAddr,
               handle: &Handle,
               owner: Option<KcpSessionUpdater>,
               mode: KcpIoMode)
               -> io::Result<KcpIo> {
        let (registration, readiness) = Registration::new2();
        let timer = Timeout::new_at(Instant::now(), handle)?;

        let elapsed = Rc::new(RefCell::new(Instant::now()));
        let close_flag = Rc::new(RefCell::new(false));
        let session = SharedKcpSession::new(shared_kcp.clone(),
                                            timer,
                                            elapsed.clone(),
                                            readiness.clone(),
                                            addr,
                                            owner,
                                            close_flag.clone(),
                                            mode)?;
        handle.spawn(session.for_each(|_| Ok(())).map_err(|err| {
                                                              error!("Failed to update KCP session: err: {:?}", err);
                                                          }));

        Ok(KcpIo {
               kcp: shared_kcp,
               registration: registration,
               readiness: readiness,
               last_update: elapsed,
               close_flag: close_flag,
           })
    }

    pub fn input_buf(&mut self, buf: &[u8]) -> io::Result<()> {
        {
            let mut last_update = self.last_update.borrow_mut();
            *last_update = Instant::now();
            let mut kcp = self.kcp.borrow_mut();
            kcp.input(&mut BytesMut::from_buf(buf))?;
        }
        self.set_readable()
    }

    pub fn set_writable(&mut self) -> io::Result<()> {
        self.readiness.set_readiness(Ready::writable())
    }

    pub fn set_readable(&mut self) -> io::Result<()> {
        self.readiness.set_readiness(Ready::readable())
    }
}

impl Read for KcpIo {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        self.kcp.borrow_mut().recv(buf)
    }
}

impl Write for KcpIo {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let mut kcp = self.kcp.borrow_mut();
        kcp.send(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        let mut kcp = self.kcp.borrow_mut();
        kcp.flush()
    }
}

impl Evented for KcpIo {
    fn register(&self, poll: &mio::Poll, token: Token, interest: Ready, opts: PollOpt) -> io::Result<()> {
        self.registration.register(poll, token, interest, opts)
    }

    fn reregister(&self, poll: &mio::Poll, token: Token, interest: Ready, opts: PollOpt) -> io::Result<()> {
        self.registration.reregister(poll, token, interest, opts)
    }

    fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
        <Registration as Evented>::deregister(&self.registration, poll)
    }
}
<commit_msg>Add read buffer in KcpIo to handle BufferTooSmall error<commit_after>use std::cell::RefCell;
use std::cmp;
use std::io::{self, BufRead, Read, Write};
use std::net::SocketAddr;
use std::rc::Rc;
use std::time::{Duration, Instant};

use bytes::BytesMut;
use bytes::buf::FromBuf;
use futures::{Future, Stream};
use mio::{self, Evented, PollOpt, Ready, Registration, SetReadiness, Token};
use tokio_core::reactor::{Handle, Timeout};

use session::{KcpSessionUpdater, SharedKcpSession};
use skcp::SharedKcp;

pub enum KcpIoMode {
    Client,
    Server(Duration),
}

pub struct KcpIo {
    kcp: SharedKcp,
    registration: Registration,
    readiness: SetReadiness,
    last_update: Rc<RefCell<Instant>>,
    close_flag: Rc<RefCell<bool>>,
    read_buf: Vec<u8>,
    read_pos: usize,
    read_cap: usize,
}

impl Drop for KcpIo {
    fn drop(&mut self) {
        let mut cf = self.close_flag.borrow_mut();
        *cf = true;
    }
}

impl KcpIo {
    pub fn new(shared_kcp: SharedKcp,
               addr: SocketAddr,
               handle: &Handle,
               owner: Option<KcpSessionUpdater>,
               mode: KcpIoMode)
               -> io::Result<KcpIo> {
        let (registration, readiness) = Registration::new2();
        let timer = Timeout::new_at(Instant::now(), handle)?;

        let elapsed = Rc::new(RefCell::new(Instant::now()));
        let close_flag = Rc::new(RefCell::new(false));
        let session = SharedKcpSession::new(shared_kcp.clone(),
                                            timer,
                                            elapsed.clone(),
                                            readiness.clone(),
                                            addr,
                                            owner,
                                            close_flag.clone(),
                                            mode)?;
        handle.spawn(session.for_each(|_| Ok(())).map_err(|err| {
                                                              error!("Failed to update KCP session: err: {:?}", err);
                                                          }));

        Ok(KcpIo {
               kcp: shared_kcp,
               registration: registration,
               readiness: readiness,
               last_update: elapsed,
               close_flag: close_flag,
               read_buf: vec![0u8; 65535],
               read_pos: 0,
               read_cap: 0,
           })
    }

    pub fn input_buf(&mut self, buf: &[u8]) -> io::Result<()> {
        {
            let mut last_update = self.last_update.borrow_mut();
            *last_update = Instant::now();
            let mut kcp = self.kcp.borrow_mut();
            kcp.input(&mut BytesMut::from_buf(buf))?;
        }
        self.set_readable()
    }

    pub fn set_writable(&mut self) -> io::Result<()> {
        self.readiness.set_readiness(Ready::writable())
    }

    pub fn set_readable(&mut self) -> io::Result<()> {
        self.readiness.set_readiness(Ready::readable())
    }
}

impl BufRead for KcpIo {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        if self.read_pos >= self.read_cap {
            let n = self.kcp.borrow_mut().recv(&mut self.read_buf)?;
            self.read_pos = 0;
            self.read_cap = n;
        }

        Ok(&self.read_buf[self.read_pos..self.read_cap])
    }

    fn consume(&mut self, amt: usize) {
        self.read_pos = cmp::min(self.read_cap, self.read_pos + amt);
    }
}

impl Read for KcpIo {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let nread = {
            let mut available = self.fill_buf()?;
            available.read(buf)?
        };
        self.consume(nread);
        Ok(nread)
    }
}

impl Write for KcpIo {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let mut kcp = self.kcp.borrow_mut();
        kcp.send(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        let mut kcp = self.kcp.borrow_mut();
        kcp.flush()
    }
}

impl Evented for KcpIo {
    fn register(&self, poll: &mio::Poll, token: Token, interest: Ready, opts: PollOpt) -> io::Result<()> {
        self.registration.register(poll, token, interest, opts)
    }

    fn reregister(&self, poll: &mio::Poll, token: Token, interest: Ready, opts: PollOpt) -> io::Result<()> {
        self.registration.reregister(poll, token, interest, opts)
    }

    fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
        <Registration as Evented>::deregister(&self.registration, poll)
    }
}
<|endoftext|>


<commit_before>use core::u8;
use core::convert::{From, AsRef};

use ctx::TryFromCtx;
use measure;
use error;
//use endian;

#[derive(Debug, Copy, Clone, Default)]
/// A variable length integer parsing `Ctx`
// TODO: I think the way to go is just type alias this to endian::Endian, set the const to NATIVE, drop the u64 ctx impl, and allow it to be used alongside endian ctxs, e.g.:
//pub type Leb128 = endian::Endian;
// ...tested and it works
pub struct Leb128 {}

/// This context instructs the underlying Scroll (Buffer, Readable) to parse as a variable length integer
pub const LEB128: Leb128 = Leb128 {};
//pub const LEB128: Leb128 = endian::NATIVE;

#[derive(Debug, PartialEq, Copy, Clone)]
/// An unsigned leb128 integer
pub struct Uleb128 {
    value: u64,
    count: usize,
}

impl Uleb128 {
    #[inline]
    pub fn size(&self) -> usize {
        self.count
    }
}

impl AsRef<u64> for Uleb128 {
    fn as_ref(&self) -> &u64 {
        &self.value
    }
}

impl From<Uleb128> for u64 {
    #[inline]
    fn from(uleb128: Uleb128) -> u64 {
        uleb128.value
    }
}

impl measure::Measure for Uleb128 {
    type Units = usize;
    #[inline]
    fn measure (&self) -> usize {
        self.count
    }
}

#[derive(Debug, PartialEq, Copy, Clone)]
/// An signed leb128 integer
pub struct Sleb128 {
    value: i64,
    count: usize,
}

impl Sleb128 {
    #[inline]
    pub fn size(&self) -> usize {
        self.count
    }
}

impl AsRef<i64> for Sleb128 {
    fn as_ref(&self) -> &i64 {
        &self.value
    }
}

impl From<Sleb128> for i64 {
    #[inline]
    fn from(sleb128: Sleb128) -> i64 {
        sleb128.value
    }
}

impl measure::Measure for Sleb128 {
    type Units = usize;
    #[inline]
    fn measure (&self) -> usize {
        self.count
    }
}

// Below implementation heavily adapted from: https://github.com/fitzgen/leb128
const CONTINUATION_BIT: u8 = 1 << 7;
const SIGN_BIT: u8 = 1 << 6;

#[inline]
fn mask_continuation(byte: u8) -> u8 {
    byte & !CONTINUATION_BIT
}

#[inline]
fn mask_continuation_u64(val: u64) -> u8 {
    let byte = val & (u8::MAX as u64);
    mask_continuation(byte as u8)
}

impl<'a> TryFromCtx<'a, (usize, Leb128)> for Uleb128 {
    type Error = error::Error;
    #[inline]
    fn try_from_ctx(src: &'a [u8], ctx: (usize, Leb128)) -> error::Result<Self> {
        use pread::Pread;
        let offset = ctx.0;
        let mut result = 0;
        let mut shift = 0;
        let mut count = 0;
        loop {
            let byte: u8 = src.pread(offset + count)?;

            if shift == 63 && byte != 0x00 && byte != 0x01 {
                return Err(error::Error::BadInput(format!("Failed to parse uleb128 from: {:?}", &src[offset..(offset+count)])));
            }

            let low_bits = mask_continuation(byte) as u64;
            result |= low_bits << shift;

            if byte & CONTINUATION_BIT == 0 {
                return Ok(Uleb128 { value: result, count: count });
            }
            count += 1;
            shift += 7;
        }
    }
}

impl<'a> TryFromCtx<'a, (usize, Leb128)> for Sleb128 {
    type Error = error::Error;
    #[inline]
    fn try_from_ctx(src: &'a [u8], (offset, _): (usize, Leb128)) -> error::Result<Self> {
        use greater::Gread;
        let o = offset;
        let offset = &mut offset.clone();
        let mut result = 0;
        let mut shift = 0;
        let size = 64;
        let mut byte: u8;
        loop {
            byte = src.gread(offset)?;

            if shift == 63 && byte != 0x00 && byte != 0x7f {
                return Err(error::Error::BadInput(format!("Failed to parse sleb128 from: {:?}", &src[o..*offset])));
            }

            let low_bits = mask_continuation(byte) as i64;
            result |= low_bits << shift;
            shift += 7;

            if byte & CONTINUATION_BIT == 0 {
                break;
            }
        }

        if shift < size && (SIGN_BIT & byte) == SIGN_BIT {
            // Sign extend the result.
            result |= !0 << shift;
        }
        Ok(Sleb128{ value: result, count: *offset - o })
    }
}

#[cfg(test)]
mod tests {
    use super::{LEB128, Uleb128, Sleb128};
    use super::super::LE;

    #[doc(hidden)]
    const CONTINUATION_BIT: u8 = 1 << 7;
    #[doc(hidden)]
    const SIGN_BIT: u8 = 1 << 6;

    #[test]
    fn uleb128() {
        use super::super::Pread;
        let buf = [2u8 | CONTINUATION_BIT, 1];
        let bytes = &buf[..];
        let num = bytes.pread::<Uleb128>(0).expect("Should read Uleb128");
        assert_eq!(130u64, num.into());
        assert_eq!(386, bytes.pread_with::<u16>(0, LE).expect("Should read number"));
    }

    #[test]
    fn uleb128_overflow() {
        use super::super::Pread;
        let buf = [2u8 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   1];
        let bytes = &buf[..];
        assert!(bytes.pread_with::<Uleb128>(0, LEB128).is_err());
    }

    #[test]
    fn sleb128() {
        use super::super::Pread;
        let bytes = [0x7fu8 | CONTINUATION_BIT, 0x7e];
        let num: i64 = bytes.pread::<Sleb128>(0).expect("Should read Sleb128").into();
        assert_eq!(-129, num);
    }
}
<commit_msg>leb128: use ctx::DefaultCtx for leb128 parsing context<commit_after>use core::u8;
use core::convert::{From, AsRef};

use ctx::{self, TryFromCtx};
use measure;
use error;

/// A variable length integer parsing `Ctx`, compatible with the standard integer endian-aware parsing context
pub type Leb128 = ctx::DefaultCtx;

/// This context instructs the underlying `Pread` implementor to parse as a variable length integer.
///
/// It currently is just the default ctx.
pub const LEB128: Leb128 = ctx::CTX;

#[derive(Debug, PartialEq, Copy, Clone)]
/// An unsigned leb128 integer
pub struct Uleb128 {
    value: u64,
    count: usize,
}

impl Uleb128 {
    #[inline]
    pub fn size(&self) -> usize {
        self.count
    }
}

impl AsRef<u64> for Uleb128 {
    fn as_ref(&self) -> &u64 {
        &self.value
    }
}

impl From<Uleb128> for u64 {
    #[inline]
    fn from(uleb128: Uleb128) -> u64 {
        uleb128.value
    }
}

impl measure::Measure for Uleb128 {
    type Units = usize;
    #[inline]
    fn measure (&self) -> usize {
        self.count
    }
}

#[derive(Debug, PartialEq, Copy, Clone)]
/// An signed leb128 integer
pub struct Sleb128 {
    value: i64,
    count: usize,
}

impl Sleb128 {
    #[inline]
    pub fn size(&self) -> usize {
        self.count
    }
}

impl AsRef<i64> for Sleb128 {
    fn as_ref(&self) -> &i64 {
        &self.value
    }
}

impl From<Sleb128> for i64 {
    #[inline]
    fn from(sleb128: Sleb128) -> i64 {
        sleb128.value
    }
}

impl measure::Measure for Sleb128 {
    type Units = usize;
    #[inline]
    fn measure (&self) -> usize {
        self.count
    }
}

// Below implementation heavily adapted from: https://github.com/fitzgen/leb128
const CONTINUATION_BIT: u8 = 1 << 7;
const SIGN_BIT: u8 = 1 << 6;

#[inline]
fn mask_continuation(byte: u8) -> u8 {
    byte & !CONTINUATION_BIT
}

#[inline]
fn mask_continuation_u64(val: u64) -> u8 {
    let byte = val & (u8::MAX as u64);
    mask_continuation(byte as u8)
}

impl<'a> TryFromCtx<'a, (usize, Leb128)> for Uleb128 {
    type Error = error::Error;
    #[inline]
    fn try_from_ctx(src: &'a [u8], ctx: (usize, Leb128)) -> error::Result<Self> {
        use pread::Pread;
        let offset = ctx.0;
        let mut result = 0;
        let mut shift = 0;
        let mut count = 0;
        loop {
            let byte: u8 = src.pread(offset + count)?;

            if shift == 63 && byte != 0x00 && byte != 0x01 {
                return Err(error::Error::BadInput(format!("Failed to parse uleb128 from: {:?}", &src[offset..(offset+count)])));
            }

            let low_bits = mask_continuation(byte) as u64;
            result |= low_bits << shift;

            if byte & CONTINUATION_BIT == 0 {
                return Ok(Uleb128 { value: result, count: count });
            }
            count += 1;
            shift += 7;
        }
    }
}

impl<'a> TryFromCtx<'a, (usize, Leb128)> for Sleb128 {
    type Error = error::Error;
    #[inline]
    fn try_from_ctx(src: &'a [u8], (offset, _): (usize, Leb128)) -> error::Result<Self> {
        use greater::Gread;
        let o = offset;
        let offset = &mut offset.clone();
        let mut result = 0;
        let mut shift = 0;
        let size = 64;
        let mut byte: u8;
        loop {
            byte = src.gread(offset)?;

            if shift == 63 && byte != 0x00 && byte != 0x7f {
                return Err(error::Error::BadInput(format!("Failed to parse sleb128 from: {:?}", &src[o..*offset])));
            }

            let low_bits = mask_continuation(byte) as i64;
            result |= low_bits << shift;
            shift += 7;

            if byte & CONTINUATION_BIT == 0 {
                break;
            }
        }

        if shift < size && (SIGN_BIT & byte) == SIGN_BIT {
            // Sign extend the result.
            result |= !0 << shift;
        }
        Ok(Sleb128{ value: result, count: *offset - o })
    }
}

#[cfg(test)]
mod tests {
    use super::{LEB128, Uleb128, Sleb128};
    use super::super::LE;

    #[doc(hidden)]
    const CONTINUATION_BIT: u8 = 1 << 7;
    #[doc(hidden)]
    const SIGN_BIT: u8 = 1 << 6;

    #[test]
    fn uleb128() {
        use super::super::Pread;
        let buf = [2u8 | CONTINUATION_BIT, 1];
        let bytes = &buf[..];
        let num = bytes.pread::<Uleb128>(0).expect("Should read Uleb128");
        assert_eq!(130u64, num.into());
        assert_eq!(386, bytes.pread_with::<u16>(0, LE).expect("Should read number"));
    }

    #[test]
    fn uleb128_overflow() {
        use super::super::Pread;
        let buf = [2u8 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   2 | CONTINUATION_BIT,
                   1];
        let bytes = &buf[..];
        assert!(bytes.pread_with::<Uleb128>(0, LEB128).is_err());
    }

    #[test]
    fn sleb128() {
        use super::super::Pread;
        let bytes = [0x7fu8 | CONTINUATION_BIT, 0x7e];
        let num: i64 = bytes.pread::<Sleb128>(0).expect("Should read Sleb128").into();
        assert_eq!(-129, num);
    }
}
<|endoftext|>


<commit_before>use mio::*;
use mio::unix::*;
use mio::util::Slab;
use std::path::PathBuf;
use std::io::{self,Read,Write,ErrorKind};
use std::iter::repeat;
use std::thread;
use std::str::from_utf8;
use std::fs;
use std::os::unix::fs::PermissionsExt;
use std::collections::HashMap;
use log;
use rustc_serialize::json::{decode,encode};
use rustc_serialize::{Encodable,Decodable,Encoder,Decoder};
use nom::{IResult,HexDisplay};

use yxorp::network::ProxyOrder;
use yxorp::network::buffer::Buffer;
use yxorp::messages::Command;

use state::{HttpProxy,TlsProxy,ConfigState};

const SERVER: Token = Token(0);

#[derive(Debug,Clone,Copy,PartialEq,Eq,Hash)]
pub enum ListenerType {
  HTTP,
  HTTPS,
  TCP
}

impl Encodable for ListenerType {
  fn encode<E: Encoder>(&self, e: &mut E) -> Result<(), E::Error> {
    match *self {
      ListenerType::HTTP  => e.emit_str("HTTP"),
      ListenerType::HTTPS => e.emit_str("HTTPS"),
      ListenerType::TCP   => e.emit_str("TCP"),
    }
  }
}

impl Decodable for ListenerType {
  fn decode<D: Decoder>(decoder: &mut D) -> Result<ListenerType, D::Error> {
    let tag = try!(decoder.read_str());
    match &tag[..] {
      "HTTP"  => Ok(ListenerType::HTTP),
      "HTTPS" => Ok(ListenerType::HTTPS),
      "TCP"   => Ok(ListenerType::TCP),
      _       => Err(decoder.error("unrecognized listener type"))
    }
  }
}

pub struct Listener {
  pub tag:           String,
  pub listener_type: ListenerType,
  pub sender:        Sender<ProxyOrder>,
  pub state:         ConfigState,
}

impl Listener {
  pub fn new(tag: String, listener_type: ListenerType, ip_address: String, port: u16, sender: Sender<ProxyOrder>) -> Listener {
    let state = match listener_type {
      ListenerType::HTTP  => ConfigState::Http(HttpProxy::new(ip_address, port)),
      ListenerType::HTTPS => ConfigState::Tls(TlsProxy::new(ip_address, port)),
      _                   => unimplemented!(),
    };

    Listener {
      tag:           tag,
      listener_type: listener_type,
      sender:        sender,
      state:         state,
    }
  }
}

impl Encodable for Listener {
  fn encode<E: Encoder>(&self, e: &mut E) -> Result<(), E::Error> {
    e.emit_map(3, |e| {
      try!(e.emit_map_elt_key(0, |e| "tag".encode(e)));
      try!(e.emit_map_elt_val(0, |e| self.tag.encode(e)));
      try!(e.emit_map_elt_key(1, |e| "listener_type".encode(e)));
      try!(e.emit_map_elt_val(1, |e| self.listener_type.encode(e)));
      try!(e.emit_map_elt_key(2, |e| "state".encode(e)));
      try!(e.emit_map_elt_val(2, |e| self.state.encode(e)));
      Ok(())
    })
  }
}

#[derive(Debug,Clone,PartialEq,Eq,Hash, RustcDecodable, RustcEncodable)]
pub struct ConfigMessage {
    pub listener: String,
    pub command:  Command
}

struct CommandClient {
  sock:      UnixStream,
  buf:       Buffer,
  back_buf:  Buffer,
  token:     Option<Token>,
}

impl CommandClient {
  fn new(sock: UnixStream) -> CommandClient {
    CommandClient {
      sock:     sock,
      buf:      Buffer::with_capacity(1024),
      back_buf: Buffer::with_capacity(10000),
      token:    None,
    }
  }

  fn reregister(&self,  event_loop: &mut EventLoop<CommandServer>, tok: Token) {
    let mut interest = EventSet::hup();
    interest.insert(EventSet::readable());
    interest.insert(EventSet::writable());
    event_loop.register(&self.sock, tok, interest, PollOpt::level() | PollOpt::oneshot());
  }

  fn conn_readable(&mut self, event_loop: &mut EventLoop<CommandServer>, tok: Token) -> Option<Vec<ConfigMessage>>{
    trace!("server conn readable; tok={:?}", tok);
    match self.sock.read(self.buf.space()) {
      Ok(0) => {
        self.reregister(event_loop, tok);
        return None;
      },
      Err(e) => {
        log!(log::LogLevel::Error, "UNIX CLIENT[{}] read error: {:?}", tok.as_usize(), e);
        return None;
      },
      Ok(r) => {
        self.buf.fill(r);
        debug!("UNIX CLIENT[{}] sent {} bytes: {:?}", tok.as_usize(), r, from_utf8(self.buf.data()));
      },
    };

    let mut res: Option<Vec<ConfigMessage>> = None;
    let mut offset = 0usize;
    match parse(self.buf.data()) {
      IResult::Incomplete(_) => {
        return Some(Vec::new());
      },
      IResult::Error(e)      => {
        log!(log::LogLevel::Error, "UNIX CLIENT[{}] read error: {:?}", tok.as_usize(), e);
        return None;
      },
      IResult::Done(i, v)    => {
        offset = self.buf.data().offset(i);
        res = Some(v);
      }
    }
    debug!("parsed {} bytes, result: {:?}", offset, res);
    self.buf.consume(offset);
    return res;
  }

  fn conn_writable(&mut self, event_loop: &mut EventLoop<CommandServer>, tok: Token) {
    trace!("server conn writable; tok={:?}", tok);
    match self.sock.write(self.back_buf.data()) {
      Ok(0) => {},
      Ok(r) => {
        self.back_buf.consume(r);
      },
      Err(e) => log!(log::LogLevel::Error,"UNIX CLIENT[{}] read error: {:?}", tok.as_usize(), e),
    }
    let mut interest = EventSet::hup();
    interest.insert(EventSet::readable());
    interest.insert(EventSet::writable());
    event_loop.register(&self.sock, tok, interest, PollOpt::level() | PollOpt::oneshot());
  }
}

fn parse(input: &[u8]) -> IResult<&[u8], Vec<ConfigMessage>> {
  many0!(input,
    complete!(terminated!(map_res!(map_res!(is_not!("\0"), from_utf8), decode), char!('\0')))
  )
}

struct CommandServer {
  sock:  UnixListener,
  conns: Slab<CommandClient>,
  listeners: HashMap<String, Listener>,
}

impl CommandServer {
  fn accept(&mut self, event_loop: &mut EventLoop<CommandServer>) -> io::Result<()> {
    debug!("server accepting socket");

    let sock = self.sock.accept().unwrap().unwrap();
    let conn = CommandClient::new(sock);
    let tok = self.conns.insert(conn)
      .ok().expect("could not add connection to slab");

    // Register the connection
    self.conns[tok].token = Some(tok);
    let mut interest = EventSet::hup();
    interest.insert(EventSet::readable());
    interest.insert(EventSet::writable());
    event_loop.register(&self.conns[tok].sock, tok, interest, PollOpt::level())
      .ok().expect("could not register socket with event loop");

    let accept_interest = EventSet::readable();
    event_loop.reregister(&self.sock, SERVER, accept_interest, PollOpt::level());
    Ok(())
  }

  fn conn<'a>(&'a mut self, tok: Token) -> &'a mut CommandClient {
    &mut self.conns[tok]
  }

  fn dispatch(&mut self, token: Token, v: Vec<ConfigMessage>) {
    for message in &v {
      let command = message.command.clone();
      if let Some(ref mut listener) = self.listeners.get_mut (&message.listener) {
        if message.command == Command::DumpConfiguration {
          self.conns[token].back_buf.write(&encode(*listener).unwrap().into_bytes());
          self.conns[token].back_buf.write(&b"\0"[..]);
        } else {
          listener.state.handle_command(&command);
          listener.sender.send(ProxyOrder::Command(command));
        }
      } else {
        log!(log::LogLevel::Error, "no listener found for tag: {}", message.listener);
      }
    }
  }

  fn new(srv: UnixListener, listeners: HashMap<String, Listener>) -> CommandServer {
    CommandServer {
      sock:  srv,
      conns: Slab::new_starting_at(Token(1), 128),
      listeners: listeners,
    }
  }
}

impl Handler for CommandServer {
	type Timeout = usize;
	type Message = ();

	fn ready(&mut self, event_loop: &mut EventLoop<CommandServer>, token: Token, events: EventSet) {
    //trace!("ready: {:?} -> {:?}", token, events);
    if events.is_readable() {
      match token {
        SERVER => self.accept(event_loop).unwrap(),
        _      => {
          if self.conns.contains(token) {
            let opt_v = self.conns[token].conn_readable(event_loop, token);
            if let Some(v) = opt_v {
              self.dispatch(token, v);
            }
          }
        }
      };
    }

    if events.is_writable() {
      match token {
        SERVER => panic!("received writable for token 0"),
        _      => {
          if self.conns.contains(token) {
            self.conns[token].conn_writable(event_loop, token)
          }
        }
      };
    }

    if events.is_hup() {
      if self.conns.contains(token) {
        event_loop.deregister(&self.conns[token].sock);
        self.conns.remove(token);
        trace!("closed client [{}]", token.as_usize());
      }
    }
  }

  fn notify(&mut self, event_loop: &mut EventLoop<Self>, message: Self::Message) {
    // TODO: dispatch here to clients the list of messages you want to display
    info!("notify");
  }
}

pub fn start(path: String, listeners: HashMap<String, Listener>) {
  thread::spawn(move || {
    let mut event_loop = EventLoop::new().unwrap();
    let addr = PathBuf::from(path);
    if let Err(e) = fs::remove_file(&addr) {
      match e.kind() {
        ErrorKind::NotFound => {},
        _ => {
          log!(log::LogLevel::Error, "could not delete previous socket at {:?}: {:?}", addr, e);
          return;
        }
      }
    }
    match UnixListener::bind(&addr) {
      Ok(srv) => {
        if let Err(e) =  fs::set_permissions(&addr, fs::Permissions::from_mode(0o600)) {
          log!( log::LogLevel::Error, "could not set the unix socket permissions: {:?}", e);
          fs::remove_file(&addr);
          return;
        }
        info!("listen for connections");
        event_loop.register(&srv, SERVER, EventSet::readable(), PollOpt::level() | PollOpt::oneshot()).unwrap();

        event_loop.run(&mut CommandServer::new(srv, listeners)).unwrap()
      },
      Err(e) => {
        log!(log::LogLevel::Error, "could not create unix socket: {:?}", e);
      }
    }
  });
}

#[cfg(test)]
mod tests {
  use super::*;
  use rustc_serialize::json;
  use yxorp::messages::{Command,HttpFront};

  #[test]
  fn config_message_test() {
    let raw_json = r#"{ "listener": "HTTP", "command":{"type": "ADD_HTTP_FRONT", "data": {"app_id": "xxx", "hostname": "yyy", "path_begin": "xxx", "port": 4242}} }"#;
    let config: ConfigMessage = json::decode(raw_json).unwrap();
    println!("{:?}", config);
    assert_eq!(config.listener, "HTTP");
    assert_eq!(config.command, Command::AddHttpFront(HttpFront{
      app_id: String::from("xxx"),
      hostname: String::from("yyy"),
      path_begin: String::from("xxx"),
      port: 4242
    }));
  }
}
<commit_msg>pass message identifiers to the library<commit_after>use mio::*;
use mio::unix::*;
use mio::util::Slab;
use std::path::PathBuf;
use std::io::{self,Read,Write,ErrorKind};
use std::iter::repeat;
use std::thread;
use std::str::from_utf8;
use std::fs;
use std::os::unix::fs::PermissionsExt;
use std::collections::HashMap;
use log;
use rustc_serialize::json::{decode,encode};
use rustc_serialize::{Encodable,Decodable,Encoder,Decoder};
use nom::{IResult,HexDisplay};

use yxorp::network::ProxyOrder;
use yxorp::network::buffer::Buffer;
use yxorp::messages::Command;

use state::{HttpProxy,TlsProxy,ConfigState};

const SERVER: Token = Token(0);

#[derive(Debug,Clone,Copy,PartialEq,Eq,Hash)]
pub enum ListenerType {
  HTTP,
  HTTPS,
  TCP
}

impl Encodable for ListenerType {
  fn encode<E: Encoder>(&self, e: &mut E) -> Result<(), E::Error> {
    match *self {
      ListenerType::HTTP  => e.emit_str("HTTP"),
      ListenerType::HTTPS => e.emit_str("HTTPS"),
      ListenerType::TCP   => e.emit_str("TCP"),
    }
  }
}

impl Decodable for ListenerType {
  fn decode<D: Decoder>(decoder: &mut D) -> Result<ListenerType, D::Error> {
    let tag = try!(decoder.read_str());
    match &tag[..] {
      "HTTP"  => Ok(ListenerType::HTTP),
      "HTTPS" => Ok(ListenerType::HTTPS),
      "TCP"   => Ok(ListenerType::TCP),
      _       => Err(decoder.error("unrecognized listener type"))
    }
  }
}

pub struct Listener {
  pub tag:           String,
  pub listener_type: ListenerType,
  pub sender:        Sender<ProxyOrder>,
  pub state:         ConfigState,
}

impl Listener {
  pub fn new(tag: String, listener_type: ListenerType, ip_address: String, port: u16, sender: Sender<ProxyOrder>) -> Listener {
    let state = match listener_type {
      ListenerType::HTTP  => ConfigState::Http(HttpProxy::new(ip_address, port)),
      ListenerType::HTTPS => ConfigState::Tls(TlsProxy::new(ip_address, port)),
      _                   => unimplemented!(),
    };

    Listener {
      tag:           tag,
      listener_type: listener_type,
      sender:        sender,
      state:         state,
    }
  }
}

impl Encodable for Listener {
  fn encode<E: Encoder>(&self, e: &mut E) -> Result<(), E::Error> {
    e.emit_map(3, |e| {
      try!(e.emit_map_elt_key(0, |e| "tag".encode(e)));
      try!(e.emit_map_elt_val(0, |e| self.tag.encode(e)));
      try!(e.emit_map_elt_key(1, |e| "listener_type".encode(e)));
      try!(e.emit_map_elt_val(1, |e| self.listener_type.encode(e)));
      try!(e.emit_map_elt_key(2, |e| "state".encode(e)));
      try!(e.emit_map_elt_val(2, |e| self.state.encode(e)));
      Ok(())
    })
  }
}

#[derive(Debug,Clone,PartialEq,Eq,Hash, RustcDecodable, RustcEncodable)]
pub struct ConfigMessage {
    pub listener: String,
    pub command:  Command
}

struct CommandClient {
  sock:      UnixStream,
  buf:       Buffer,
  back_buf:  Buffer,
  token:     Option<Token>,
}

impl CommandClient {
  fn new(sock: UnixStream) -> CommandClient {
    CommandClient {
      sock:     sock,
      buf:      Buffer::with_capacity(1024),
      back_buf: Buffer::with_capacity(10000),
      token:    None,
    }
  }

  fn reregister(&self,  event_loop: &mut EventLoop<CommandServer>, tok: Token) {
    let mut interest = EventSet::hup();
    interest.insert(EventSet::readable());
    interest.insert(EventSet::writable());
    event_loop.register(&self.sock, tok, interest, PollOpt::level() | PollOpt::oneshot());
  }

  fn conn_readable(&mut self, event_loop: &mut EventLoop<CommandServer>, tok: Token) -> Option<Vec<ConfigMessage>>{
    trace!("server conn readable; tok={:?}", tok);
    match self.sock.read(self.buf.space()) {
      Ok(0) => {
        self.reregister(event_loop, tok);
        return None;
      },
      Err(e) => {
        log!(log::LogLevel::Error, "UNIX CLIENT[{}] read error: {:?}", tok.as_usize(), e);
        return None;
      },
      Ok(r) => {
        self.buf.fill(r);
        debug!("UNIX CLIENT[{}] sent {} bytes: {:?}", tok.as_usize(), r, from_utf8(self.buf.data()));
      },
    };

    let mut res: Option<Vec<ConfigMessage>> = None;
    let mut offset = 0usize;
    match parse(self.buf.data()) {
      IResult::Incomplete(_) => {
        return Some(Vec::new());
      },
      IResult::Error(e)      => {
        log!(log::LogLevel::Error, "UNIX CLIENT[{}] read error: {:?}", tok.as_usize(), e);
        return None;
      },
      IResult::Done(i, v)    => {
        offset = self.buf.data().offset(i);
        res = Some(v);
      }
    }
    debug!("parsed {} bytes, result: {:?}", offset, res);
    self.buf.consume(offset);
    return res;
  }

  fn conn_writable(&mut self, event_loop: &mut EventLoop<CommandServer>, tok: Token) {
    trace!("server conn writable; tok={:?}", tok);
    match self.sock.write(self.back_buf.data()) {
      Ok(0) => {},
      Ok(r) => {
        self.back_buf.consume(r);
      },
      Err(e) => log!(log::LogLevel::Error,"UNIX CLIENT[{}] read error: {:?}", tok.as_usize(), e),
    }
    let mut interest = EventSet::hup();
    interest.insert(EventSet::readable());
    interest.insert(EventSet::writable());
    event_loop.register(&self.sock, tok, interest, PollOpt::level() | PollOpt::oneshot());
  }
}

fn parse(input: &[u8]) -> IResult<&[u8], Vec<ConfigMessage>> {
  many0!(input,
    complete!(terminated!(map_res!(map_res!(is_not!("\0"), from_utf8), decode), char!('\0')))
  )
}

struct CommandServer {
  sock:  UnixListener,
  conns: Slab<CommandClient>,
  listeners: HashMap<String, Listener>,
}

impl CommandServer {
  fn accept(&mut self, event_loop: &mut EventLoop<CommandServer>) -> io::Result<()> {
    debug!("server accepting socket");

    let sock = self.sock.accept().unwrap().unwrap();
    let conn = CommandClient::new(sock);
    let tok = self.conns.insert(conn)
      .ok().expect("could not add connection to slab");

    // Register the connection
    self.conns[tok].token = Some(tok);
    let mut interest = EventSet::hup();
    interest.insert(EventSet::readable());
    interest.insert(EventSet::writable());
    event_loop.register(&self.conns[tok].sock, tok, interest, PollOpt::level())
      .ok().expect("could not register socket with event loop");

    let accept_interest = EventSet::readable();
    event_loop.reregister(&self.sock, SERVER, accept_interest, PollOpt::level());
    Ok(())
  }

  fn conn<'a>(&'a mut self, tok: Token) -> &'a mut CommandClient {
    &mut self.conns[tok]
  }

  fn dispatch(&mut self, token: Token, v: Vec<ConfigMessage>) {
    for message in &v {
      let command = message.command.clone();
      if let Some(ref mut listener) = self.listeners.get_mut (&message.listener) {
        if message.command == Command::DumpConfiguration {
          self.conns[token].back_buf.write(&encode(*listener).unwrap().into_bytes());
          self.conns[token].back_buf.write(&b"\0"[..]);
        } else {
          listener.state.handle_command(&command);
          listener.sender.send(ProxyOrder::Command(message.id.clone(), command));
        }
      } else {
        log!(log::LogLevel::Error, "no listener found for tag: {}", message.listener);
      }
    }
  }

  fn new(srv: UnixListener, listeners: HashMap<String, Listener>) -> CommandServer {
    CommandServer {
      sock:  srv,
      conns: Slab::new_starting_at(Token(1), 128),
      listeners: listeners,
    }
  }
}

impl Handler for CommandServer {
	type Timeout = usize;
	type Message = ();

	fn ready(&mut self, event_loop: &mut EventLoop<CommandServer>, token: Token, events: EventSet) {
    //trace!("ready: {:?} -> {:?}", token, events);
    if events.is_readable() {
      match token {
        SERVER => self.accept(event_loop).unwrap(),
        _      => {
          if self.conns.contains(token) {
            let opt_v = self.conns[token].conn_readable(event_loop, token);
            if let Some(v) = opt_v {
              self.dispatch(token, v);
            }
          }
        }
      };
    }

    if events.is_writable() {
      match token {
        SERVER => panic!("received writable for token 0"),
        _      => {
          if self.conns.contains(token) {
            self.conns[token].conn_writable(event_loop, token)
          }
        }
      };
    }

    if events.is_hup() {
      if self.conns.contains(token) {
        event_loop.deregister(&self.conns[token].sock);
        self.conns.remove(token);
        trace!("closed client [{}]", token.as_usize());
      }
    }
  }

  fn notify(&mut self, event_loop: &mut EventLoop<Self>, message: Self::Message) {
    // TODO: dispatch here to clients the list of messages you want to display
    info!("notify");
  }
}

pub fn start(path: String, listeners: HashMap<String, Listener>) {
  thread::spawn(move || {
    let mut event_loop = EventLoop::new().unwrap();
    let addr = PathBuf::from(path);
    if let Err(e) = fs::remove_file(&addr) {
      match e.kind() {
        ErrorKind::NotFound => {},
        _ => {
          log!(log::LogLevel::Error, "could not delete previous socket at {:?}: {:?}", addr, e);
          return;
        }
      }
    }
    match UnixListener::bind(&addr) {
      Ok(srv) => {
        if let Err(e) =  fs::set_permissions(&addr, fs::Permissions::from_mode(0o600)) {
          log!( log::LogLevel::Error, "could not set the unix socket permissions: {:?}", e);
          fs::remove_file(&addr);
          return;
        }
        info!("listen for connections");
        event_loop.register(&srv, SERVER, EventSet::readable(), PollOpt::level() | PollOpt::oneshot()).unwrap();

        event_loop.run(&mut CommandServer::new(srv, listeners)).unwrap()
      },
      Err(e) => {
        log!(log::LogLevel::Error, "could not create unix socket: {:?}", e);
      }
    }
  });
}

#[cfg(test)]
mod tests {
  use super::*;
  use rustc_serialize::json;
  use yxorp::messages::{Command,HttpFront};

  #[test]
  fn config_message_test() {
    let raw_json = r#"{ "listener": "HTTP", "command":{"type": "ADD_HTTP_FRONT", "data": {"app_id": "xxx", "hostname": "yyy", "path_begin": "xxx", "port": 4242}} }"#;
    let config: ConfigMessage = json::decode(raw_json).unwrap();
    println!("{:?}", config);
    assert_eq!(config.listener, "HTTP");
    assert_eq!(config.command, Command::AddHttpFront(HttpFront{
      app_id: String::from("xxx"),
      hostname: String::from("yyy"),
      path_begin: String::from("xxx"),
      port: 4242
    }));
  }
}
<|endoftext|>