<commit_before>#![cfg_attr(test, deny(missing_docs, warnings))]
#![forbid(unused_variables)]

use std::{fmt, mem};
use std::borrow::{Borrow, ToOwned};
use std::default::Default;
use std::hash::{Hash, Hasher};
use std::iter::FromIterator;
use std::ops::{Deref, Index, Range, RangeFrom, RangeFull, RangeTo};

/// An owned sequence of Unicode scalar values, equivalent to `String`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UString(Vec<char>);

/// A Unicode string slice, equivalent to `str`.
pub struct UStr([char]);

impl UString {
    pub fn remove(&mut self, idx: usize) -> char {
        self.0.remove(idx)
    }
}

impl UStr {
    /// Returns the number of Unicode scalar values in this string.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Retrieves the first character from a `&UStr` and returns it.
    ///
    /// This does not allocate a new string; instead, it returns a slice that points one character beyond the character that was shifted.
    ///
    /// If the slice does not contain any characters, None is returned instead.
    pub fn slice_shift_char(&self) -> Option<(char, &UStr)> {
        if self.len() == 0 {
            None
        } else {
            let ch = self[0];
            let next_s = &self[1..];
            Some((ch, next_s))
        }
    }
}

impl Borrow<UStr> for UString {
    fn borrow(&self) -> &UStr {
        &self
    }
}

impl Default for UString {
    fn default() -> UString {
        UString(Vec::new())
    }
}

impl Deref for UString {
    type Target = UStr;

    fn deref(&self) -> &UStr {
        unsafe {
            mem::transmute(&*self.0)
            // just so you can be clear what it is:
            //mem::transmute::<&[char], &UStr>(&*self.0)
        }
    }
}

impl From<String> for UString {
    fn from(string: String) -> UString {
        UString(string.chars().collect())
    }
}

impl From<Vec<char>> for UString {
    fn from(chars: Vec<char>) -> UString {
        UString(chars)
    }
}

impl<'a> From<&'a str> for UString {
    fn from(string: &'a str) -> UString {
        UString(string.chars().collect())
    }
}

impl<'a> From<&'a [char]> for UString {
    fn from(slice: &[char]) -> UString {
        UString(slice.to_vec())
    }
}

impl<C: Into<char>> FromIterator<C> for UString {
    fn from_iter<T>(iterator: T) -> UString where T: IntoIterator<Item=C> {
        UString(iterator.into_iter().map(C::into).collect())
    }
}

impl Hash for UString {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

impl Index<usize> for UStr {
    type Output = char;

    fn index(&self, idx: usize) -> &char {
        &self.0[idx]
    }
}

impl Index<Range<usize>> for UStr {
    type Output = UStr;

    fn index(&self, idx: Range<usize>) -> &UStr {
        unsafe { mem::transmute(&self.0[idx]) }
    }
}

impl Index<RangeFrom<usize>> for UStr {
    type Output = UStr;

    fn index(&self, idx: RangeFrom<usize>) -> &UStr {
        unsafe { mem::transmute(&self.0[idx]) }
    }
}

impl Index<RangeFull> for UStr {
    type Output = UStr;

    fn index(&self, _: RangeFull) -> &UStr {
        self
    }
}

impl Index<RangeTo<usize>> for UStr {
    type Output = UStr;

    fn index(&self, idx: RangeTo<usize>) -> &UStr {
        unsafe { mem::transmute(&self.0[idx]) }
    }
}

impl ToOwned for UStr {
    type Owned = UString;

    fn to_owned(&self) -> UString {
        UString::from(self.0.to_owned())
    }
}

impl fmt::Display for UStr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for c in &self.0 {
            try!(write!(f, "{}", c));
        }
        Ok(())
    }
}
<commit_msg>Implement IntoIterator for UString, &UString, and &UStr<commit_after>#![cfg_attr(test, deny(missing_docs, warnings))]
#![forbid(unused_variables)]

use std::{fmt, mem};
use std::borrow::{Borrow, ToOwned};
use std::default::Default;
use std::hash::{Hash, Hasher};
use std::iter::FromIterator;
use std::ops::{Deref, Index, Range, RangeFrom, RangeFull, RangeTo};

/// An owned sequence of Unicode scalar values, equivalent to `String`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UString(Vec<char>);

/// A Unicode string slice, equivalent to `str`.
pub struct UStr([char]);

impl UString {
    pub fn remove(&mut self, idx: usize) -> char {
        self.0.remove(idx)
    }
}

impl UStr {
    /// Returns the number of Unicode scalar values in this string.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Retrieves the first character from a `&UStr` and returns it.
    ///
    /// This does not allocate a new string; instead, it returns a slice that points one character beyond the character that was shifted.
    ///
    /// If the slice does not contain any characters, None is returned instead.
    pub fn slice_shift_char(&self) -> Option<(char, &UStr)> {
        if self.len() == 0 {
            None
        } else {
            let ch = self[0];
            let next_s = &self[1..];
            Some((ch, next_s))
        }
    }
}

impl Borrow<UStr> for UString {
    fn borrow(&self) -> &UStr {
        &self
    }
}

impl Default for UString {
    fn default() -> UString {
        UString(Vec::new())
    }
}

impl Deref for UString {
    type Target = UStr;

    fn deref(&self) -> &UStr {
        unsafe {
            mem::transmute(&*self.0)
            // just so you can be clear what it is:
            //mem::transmute::<&[char], &UStr>(&*self.0)
        }
    }
}

impl From<String> for UString {
    fn from(string: String) -> UString {
        UString(string.chars().collect())
    }
}

impl From<Vec<char>> for UString {
    fn from(chars: Vec<char>) -> UString {
        UString(chars)
    }
}

impl<'a> From<&'a str> for UString {
    fn from(string: &'a str) -> UString {
        UString(string.chars().collect())
    }
}

impl<'a> From<&'a [char]> for UString {
    fn from(slice: &[char]) -> UString {
        UString(slice.to_vec())
    }
}

impl<C: Into<char>> FromIterator<C> for UString {
    fn from_iter<T>(iterator: T) -> UString where T: IntoIterator<Item=C> {
        UString(iterator.into_iter().map(C::into).collect())
    }
}

impl Hash for UString {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.hash(state);
    }
}

impl Index<usize> for UStr {
    type Output = char;

    fn index(&self, idx: usize) -> &char {
        &self.0[idx]
    }
}

impl Index<Range<usize>> for UStr {
    type Output = UStr;

    fn index(&self, idx: Range<usize>) -> &UStr {
        unsafe { mem::transmute(&self.0[idx]) }
    }
}

impl Index<RangeFrom<usize>> for UStr {
    type Output = UStr;

    fn index(&self, idx: RangeFrom<usize>) -> &UStr {
        unsafe { mem::transmute(&self.0[idx]) }
    }
}

impl Index<RangeFull> for UStr {
    type Output = UStr;

    fn index(&self, _: RangeFull) -> &UStr {
        self
    }
}

impl Index<RangeTo<usize>> for UStr {
    type Output = UStr;

    fn index(&self, idx: RangeTo<usize>) -> &UStr {
        unsafe { mem::transmute(&self.0[idx]) }
    }
}

impl ToOwned for UStr {
    type Owned = UString;

    fn to_owned(&self) -> UString {
        UString::from(self.0.to_owned())
    }
}

impl fmt::Display for UStr {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for c in &self.0 {
            try!(write!(f, "{}", c));
        }
        Ok(())
    }
}

pub struct Iter<'a>(&'a [char]);

pub struct IntoIter(Vec<char>);

impl<'a> Iterator for Iter<'a> {
    type Item = &'a char;

    fn next(&mut self) -> Option<&'a char> {
        if self.0.len() > 0 {
            let result = &self.0[0];
            self.0 = &self.0[1..];
            Some(result)
        } else {
            None
        }
    }
}

impl Iterator for IntoIter {
    type Item = char;

    fn next(&mut self) -> Option<char> {
        if self.0.len() > 0 {
            Some(self.0.remove(0))
        } else {
            None
        }
    }
}

impl IntoIterator for UString {
    type Item = char;
    type IntoIter = IntoIter;

    fn into_iter(self) -> IntoIter {
        IntoIter(self.0)
    }
}

impl<'a> IntoIterator for &'a UString {
    type Item = &'a char;
    type IntoIter = Iter<'a>;

    fn into_iter(self) -> Iter<'a> {
        Iter(&self.0)
    }
}

impl<'a> IntoIterator for &'a UStr {
    type Item = &'a char;
    type IntoIter = Iter<'a>;

    fn into_iter(self) -> Iter<'a> {
        Iter(&self.0)
    }
}
<|endoftext|>


<commit_before>extern crate iron;
extern crate unicase;

use unicase::UniCase;
use iron::prelude::*;
use iron::method::Method;
use iron::status;
use iron::headers::{Origin,
                    AccessControlRequestMethod,
                    AccessControlAllowOrigin,
                    AccessControlAllowHeaders,
                    AccessControlMaxAge,
                    AccessControlAllowMethods};
use iron::middleware::{AroundMiddleware, Handler};

pub struct CorsMiddleware;
impl AroundMiddleware for CorsMiddleware {

    fn around(self, handler: Box<Handler>) -> Box<Handler> {
        Box::new(move | req: &mut Request | {
            let allowed_headers: Vec<unicase::UniCase<String>> = vec![
                    // To allow application/json
                    UniCase("Content-Type".to_owned()),
                    // Set by some js libs
                    UniCase("X-Requested-With".to_owned()),
                    ];
            let allowed_methods: Vec<Method> = vec![ Method::Get, Method::Put, Method::Post ];
            if req.method == Method::Options &&
                req.headers.get::<AccessControlRequestMethod>().is_some() {
                // Preflight request
                // 1.If the Origin header is not present terminate this set of steps. The request is
                // outside the scope of this specification.
                let maybe_origin = req.headers.get::<Origin>();
                if maybe_origin.is_none() {
                    let resp = Response::with((status::BadRequest,
                                               "Preflight request without Origin header"));
                    return Ok(resp);
                }
                let origin = maybe_origin.unwrap();
                let mut res = Response::with((status::NoContent));
                //
                // 2.If the value of the Origin header is not a case-sensitive match for any of the
                // values in list of origins do not set any additional headers and terminate this
                // set of steps.
                //
                // Note: Always matching is acceptable since the list of origins can be unbounded.
                //
                // Note: The Origin header can only contain a single origin as the user agent
                //       will not follow redirects.
                //
                // 3. Let method be the value as result of parsing the Access-Control-Request-Method
                // header.
                //
                // If there is no Access-Control-Request-Method header or if parsing failed, do not
                // set any additional headers and terminate this set of steps. The request is
                // outside the scope of this specification.
                //
                // 4. Let header field-names be the values as result of parsing the
                // Access-Control-Request-Headers headers.
                //
                // If there are no Access-Control-Request-Headers headers let header field-names be
                // the empty list.
                //
                // If parsing failed do not set any additional headers and terminate this set of
                // steps. The request is outside the scope of this specification.
                //
                // 5.If method is not a case-sensitive match for any of the values in list of
                //   methods do not set any additional headers and terminate this set of steps.
                //
                // Always matching is acceptable since the list of methods can be unbounded.
                //
                // 6. If any of the header field-names is not a ASCII case-insensitive match for any
                // of the values in list of headers do not set any additional headers and terminate
                // this set of steps.
                //
                // Always matching is acceptable since the list of headers can be unbounded.
                //
                // 7. If the resource supports credentials add a single Access-Control-Allow-Origin
                // header, with the value of the Origin header as value, and add a single
                // Access-Control-Allow-Credentials header with the case-sensitive string "true" as
                // value.
                //
                // Otherwise, add a single Access-Control-Allow-Origin header, with either the
                // value of the Origin header or the string "*" as value.
                res.headers.set(AccessControlAllowOrigin::Value(format!("{}", origin)));
                //
                // The string "*" cannot be used for a resource that supports credentials.
                //
                // 8. Optionally add a single Access-Control-Max-Age header with as value the amount
                // of seconds the user agent is allowed to cache the result of the request.
                res.headers.set(AccessControlMaxAge(60*60));
                //
                // 9. If method is a simple method this step may be skipped.
                //
                // Add one or more Access-Control-Allow-Methods headers consisting of (a subset of)
                // the list of methods.
                res.headers.set(AccessControlAllowMethods(allowed_methods));
                //
                // If a method is a simple method it does not need to be listed, but this is not
                // prohibited.
                //
                // Since the list of methods can be unbounded, simply returning the method
                // indicated by Access-Control-Request-Method (if supported) can be enough.
                //
                // 10.If each of the header field-names is a simple header and none is Content-Type,
                // this step may be skipped.
                //
                // Add one or more Access-Control-Allow-Headers headers consisting of (a subset of)
                // the list of headers.
                res.headers.set(AccessControlAllowHeaders(allowed_headers));
                //
                // If a header field name is a simple header and is not Content-Type, it is not
                // required to be listed. Content-Type is to be listed as only a subset of its
                // values makes it qualify as simple header.
                //
                // Since the list of headers can be unbounded, simply returning supported headers
                // from Access-Control-Allow-Headers can be enough.
                Ok(res)
            } else {
                // Normal request
                // 1.If the Origin header is not present terminate this set of steps. The request is
                // outside the scope of this specification.
                let has_origin = req.headers.get::<Origin>().is_some();
                if !has_origin {
                    return handler.handle(req);
                }
                //
                // 2.If the value of the Origin header is not a case-sensitive match for any of the
                // values in list of origins, do not set any additional headers and terminate this
                // set of steps.
                //
                // Note: Always matching is acceptable since the list of origins can be unbounded.
                //
                // 3. If the resource supports credentials add a single Access-Control-Allow-Origin
                // header, with the value of the Origin header as value, and add a single
                // Access-Control-Allow-Credentials header with the case-sensitive string "true" as
                // value.
                //
                // Otherwise, add a single Access-Control-Allow-Origin header, with either the
                // value of the Origin header or the string "*" as value.
                //
                // Note: The string "*" cannot be used for a resource that supports credentials.
                //
                // 4. If the list of exposed headers is not empty add one or more
                // Access-Control-Expose-Headers headers, with as values the header field names
                // given in the list of exposed headers.
                let result = handler.handle(req);
                match result {
                    Ok(mut res) => {
                        // And set CORS headers
                        let origin = req.headers.get::<Origin>().unwrap();
                        res.headers.set(AccessControlAllowOrigin::Value(format!("{}", origin)));
                        Ok(res)
                    }
                    _ => result
                }
            }
        })
    }
}

#[cfg(test)]
mod tests {
	extern crate router;
	extern crate hyper;
	use iron::Listening;
	use self::router::Router;
	use iron::prelude::*;
	use iron::status;
	use self::hyper::Client;
	use self::hyper::header::Headers;
	use std::io::Read;
    use iron::headers::{Origin,
                    AccessControlRequestMethod,
                    AccessControlAllowOrigin,
                    AccessControlAllowHeaders,
                    AccessControlMaxAge,
                    AccessControlAllowMethods};
    use iron::method::Method;
    use super::CorsMiddleware;
    use std::str::FromStr;

	struct AutoServer {
		listening: Listening,
		port: u16
	}

	impl AutoServer {
		pub fn new() -> AutoServer {
			let mut router = Router::new();
            let handler = |_: &mut Request| Ok(Response::with((status::ImATeapot, "")));
			router.get("/a", handler,  "get_a");
            let mut chain = Chain::new(router);
            chain.link_around(CorsMiddleware);
    		let l = Iron::new(chain).http(format!("localhost:0")).unwrap();
			let p = l.socket.port();
			AutoServer { listening: l, port: p }
		}
	}

	impl Drop for AutoServer {
		fn drop(&mut self) {
			// Workaround for https://github.com/hyperium/hyper/issues/338
			self.listening.close().unwrap();
		}
	}

    #[test]
    fn normal_request_possible() {
		let server = AutoServer::new();
		let client = Client::new();
		let res = client.get(&format!("http://127.0.0.1:{}/a", server.port)).send().unwrap();
		assert_eq!(res.status, status::ImATeapot);
	}

    #[test]
    fn preflight_without_origin_is_bad_request() {
		let server = AutoServer::new();
		let client = Client::new();
		let mut headers = Headers::new();
		headers.set(AccessControlRequestMethod(Method::Get));
		let mut res = client
				.request(Method::Options, &format!("http://127.0.0.1:{}/a", server.port))
				.headers(headers)
				.send().unwrap();
		assert_eq!(res.status, status::BadRequest);
        let mut payload = String::new();
        res.read_to_string(&mut payload).unwrap();
        assert_eq!(payload, "Preflight request without Origin header");
	}

    #[test]
    fn preflight_with_origin_accepts_same_origin() {
		let server = AutoServer::new();
		let client = Client::new();
		let mut headers = Headers::new();
		headers.set(AccessControlRequestMethod(Method::Get));
		headers.set(Origin::from_str("http://www.a.com:8080").unwrap());
		let res = client
				.request(Method::Options, &format!("http://127.0.0.1:{}/a", server.port))
				.headers(headers)
				.send().unwrap();
		assert_eq!(res.status, status::NoContent);
        let allow_origin = res.headers.get::<AccessControlAllowOrigin>().unwrap();
        assert_eq!(format!("{}", allow_origin) , "http://www.a.com:8080");
        let allow_headers = res.headers.get::<AccessControlAllowHeaders>().unwrap();
        assert_eq!(format!("{}", allow_headers) , "Content-Type, X-Requested-With");
        let allow_methods = res.headers.get::<AccessControlAllowMethods>().unwrap();
        assert_eq!(format!("{}", allow_methods) , "GET, PUT, POST");
        let max_age = res.headers.get::<AccessControlMaxAge>().unwrap();
        assert_eq!(max_age.0, 60*60u32);
	}

    #[test]
    fn normal_request_allows_origin() {
		let server = AutoServer::new();
		let client = Client::new();
		let mut headers = Headers::new();
		headers.set(Origin::from_str("http://www.a.com:8080").unwrap());
		let res = client
            .get(&format!("http://127.0.0.1:{}/a", server.port))
            .headers(headers)
            .send().unwrap();
		assert_eq!(res.status, status::ImATeapot);
        let allow_origin = res.headers.get::<AccessControlAllowOrigin>().unwrap();
        assert_eq!(format!("{}", allow_origin) , "http://www.a.com:8080");
        assert!(res.headers.get::<AccessControlAllowHeaders>().is_none());
        assert!(res.headers.get::<AccessControlAllowMethods>().is_none());
        assert!(res.headers.get::<AccessControlMaxAge>().is_none());
	}

    #[test]
    fn normal_request_without_origin_is_passthrough() {
		let server = AutoServer::new();
		let client = Client::new();
		let res = client
            .get(&format!("http://127.0.0.1:{}/a", server.port))
            .send().unwrap();
		assert_eq!(res.status, status::ImATeapot);
        assert!(res.headers.get::<AccessControlAllowOrigin>().is_none());
        assert!(res.headers.get::<AccessControlAllowHeaders>().is_none());
        assert!(res.headers.get::<AccessControlAllowMethods>().is_none());
        assert!(res.headers.get::<AccessControlMaxAge>().is_none());
	}

}
<commit_msg>add config for everything<commit_after>extern crate iron;
extern crate unicase;

use unicase::UniCase;
use iron::prelude::*;
use iron::method::Method;
use iron::status;
use iron::headers::{Origin,
                    AccessControlRequestMethod,
                    AccessControlAllowOrigin,
                    AccessControlAllowHeaders,
                    AccessControlMaxAge,
                    AccessControlAllowMethods};
use iron::middleware::{AroundMiddleware, Handler};

pub enum AllowedOrigins {
    Any,
    Specific(std::collections::HashSet<unicase::UniCase<String>>)
}

impl AllowedOrigins {
    fn answer_to(&self, origin: String) -> Option<String> {
       match self {
          &AllowedOrigins::Any => { Some("*".to_owned()) }
          &AllowedOrigins::Specific(ref allowed) => {
            if allowed.contains(&UniCase(origin.clone())) {
                Some(origin)
            } else {
                None
            }
          }
       }
    }
}

pub struct CorsMiddleware {
     pub allowed_origins: AllowedOrigins,
     pub allowed_methods: Vec<Method>,
     pub allowed_headers: Vec<UniCase<String>>,
     pub exposed_headers: Vec<String>,
     pub allow_credentials: bool,
     pub max_age_seconds: u32
}

impl CorsMiddleware {
    pub fn new() -> CorsMiddleware {
            let allowed_methods: Vec<Method> = vec![ Method::Get, Method::Put, Method::Post ];
            let allowed_headers: Vec<unicase::UniCase<String>> = vec![
                    // To allow application/json
                    UniCase("Content-Type".to_owned()),
                    // Set by some js libs
                    UniCase("X-Requested-With".to_owned()),
                    ];
            let exposed_headers: Vec<String> = Vec::new();
            CorsMiddleware {
                allowed_origins: AllowedOrigins::Any,
                allowed_methods: allowed_methods,
                allowed_headers: allowed_headers,
                exposed_headers: exposed_headers,
                allow_credentials: false,
                max_age_seconds: 60*60
            }
    }
}

impl AroundMiddleware for CorsMiddleware {

    fn around(self, handler: Box<Handler>) -> Box<Handler> {
        Box::new(move | req: &mut Request | {
            if req.method == Method::Options &&
                req.headers.get::<AccessControlRequestMethod>().is_some() {
                // Preflight request
                // 1.If the Origin header is not present terminate this set of steps. The request is
                // outside the scope of this specification.
                let maybe_origin = req.headers.get::<Origin>();
                if maybe_origin.is_none() {
                    let resp = Response::with((status::BadRequest,
                                               "Preflight request without Origin header"));
                    return Ok(resp);
                }
                let origin = maybe_origin.unwrap();
                // Preflight is NoContent
                let mut res = Response::with((status::NoContent));
                let origin_string = origin.to_string();
                let allowed_origin = self.allowed_origins.answer_to(origin_string);
                //
                // 2.If the value of the Origin header is not a case-sensitive match for any of the
                // values in list of origins do not set any additional headers and terminate this
                // set of steps.
                if allowed_origin.is_none() {
                    let resp = Response::with((status::BadRequest,
                                               "Preflight request with disallowed origin"));
                    return Ok(resp);
                }
                //
                // Note: Always matching is acceptable since the list of origins can be unbounded.
                //
                // Note: The Origin header can only contain a single origin as the user agent
                //       will not follow redirects.
                //
                // 3. Let method be the value as result of parsing the Access-Control-Request-Method
                // header.
                //
                // If there is no Access-Control-Request-Method header or if parsing failed, do not
                // set any additional headers and terminate this set of steps. The request is
                // outside the scope of this specification.
                //
                // 4. Let header field-names be the values as result of parsing the
                // Access-Control-Request-Headers headers.
                //
                // If there are no Access-Control-Request-Headers headers let header field-names be
                // the empty list.
                //
                // If parsing failed do not set any additional headers and terminate this set of
                // steps. The request is outside the scope of this specification.
                //
                // 5.If method is not a case-sensitive match for any of the values in list of
                //   methods do not set any additional headers and terminate this set of steps.
                //
                // Always matching is acceptable since the list of methods can be unbounded.
                //
                // 6. If any of the header field-names is not a ASCII case-insensitive match for any
                // of the values in list of headers do not set any additional headers and terminate
                // this set of steps.
                //
                // Always matching is acceptable since the list of headers can be unbounded.
                //
                // 7. If the resource supports credentials add a single Access-Control-Allow-Origin
                // header, with the value of the Origin header as value, and add a single
                // Access-Control-Allow-Credentials header with the case-sensitive string "true" as
                // value.
                //
                // Otherwise, add a single Access-Control-Allow-Origin header, with either the
                // value of the Origin header or the string "*" as value.
                res.headers.set(AccessControlAllowOrigin::Value(format!("{}", origin)));
                //
                // The string "*" cannot be used for a resource that supports credentials.
                //
                // 8. Optionally add a single Access-Control-Max-Age header with as value the amount
                // of seconds the user agent is allowed to cache the result of the request.
                res.headers.set(AccessControlMaxAge(self.max_age_seconds));
                //
                // 9. If method is a simple method this step may be skipped.
                //
                // Add one or more Access-Control-Allow-Methods headers consisting of (a subset of)
                // the list of methods.
                res.headers.set(AccessControlAllowMethods(self.allowed_methods.clone()));
                //
                // If a method is a simple method it does not need to be listed, but this is not
                // prohibited.
                //
                // Since the list of methods can be unbounded, simply returning the method
                // indicated by Access-Control-Request-Method (if supported) can be enough.
                //
                // 10.If each of the header field-names is a simple header and none is Content-Type,
                // this step may be skipped.
                //
                // Add one or more Access-Control-Allow-Headers headers consisting of (a subset of)
                // the list of headers.
                res.headers.set(AccessControlAllowHeaders(self.allowed_headers.clone()));
                //
                // If a header field name is a simple header and is not Content-Type, it is not
                // required to be listed. Content-Type is to be listed as only a subset of its
                // values makes it qualify as simple header.
                //
                // Since the list of headers can be unbounded, simply returning supported headers
                // from Access-Control-Allow-Headers can be enough.
                Ok(res)
            } else {
                // Normal request
                // 1.If the Origin header is not present terminate this set of steps. The request is
                // outside the scope of this specification.
                let has_origin = req.headers.get::<Origin>().is_some();
                if !has_origin {
                    return handler.handle(req);
                }
                //
                // 2.If the value of the Origin header is not a case-sensitive match for any of the
                // values in list of origins, do not set any additional headers and terminate this
                // set of steps.
                //
                // Note: Always matching is acceptable since the list of origins can be unbounded.
                //
                // 3. If the resource supports credentials add a single Access-Control-Allow-Origin
                // header, with the value of the Origin header as value, and add a single
                // Access-Control-Allow-Credentials header with the case-sensitive string "true" as
                // value.
                //
                // Otherwise, add a single Access-Control-Allow-Origin header, with either the
                // value of the Origin header or the string "*" as value.
                //
                // Note: The string "*" cannot be used for a resource that supports credentials.
                //
                // 4. If the list of exposed headers is not empty add one or more
                // Access-Control-Expose-Headers headers, with as values the header field names
                // given in the list of exposed headers.
                let result = handler.handle(req);
                match result {
                    Ok(mut res) => {
                        // And set CORS headers
                        let origin = req.headers.get::<Origin>().unwrap();
                        res.headers.set(AccessControlAllowOrigin::Value(format!("{}", origin)));
                        Ok(res)
                    }
                    _ => result
                }
            }
        })
    }
}

#[cfg(test)]
mod tests {
	extern crate router;
	extern crate hyper;
	use iron::Listening;
	use self::router::Router;
	use iron::prelude::*;
	use iron::status;
	use self::hyper::Client;
	use self::hyper::header::Headers;
	use std::io::Read;
    use iron::headers::{Origin,
                    AccessControlRequestMethod,
                    AccessControlAllowOrigin,
                    AccessControlAllowHeaders,
                    AccessControlMaxAge,
                    AccessControlAllowMethods};
    use iron::method::Method;
    use super::CorsMiddleware;
    use std::str::FromStr;

	struct AutoServer {
		listening: Listening,
		port: u16
	}

	impl AutoServer {
		pub fn new() -> AutoServer {
			let mut router = Router::new();
            let handler = |_: &mut Request| Ok(Response::with((status::ImATeapot, "")));
            //let mut private_chain = Chain::new(handler);
            //private_chain.link_around(CorsMiddleware);
			router.get("/a", handler,  "get_a");
            let mut chain = Chain::new(router);
            chain.link_around(CorsMiddleware::new());
    		let l = Iron::new(chain).http(format!("localhost:0")).unwrap();
			let p = l.socket.port();
			AutoServer { listening: l, port: p }
		}
	}

	impl Drop for AutoServer {
		fn drop(&mut self) {
			// Workaround for https://github.com/hyperium/hyper/issues/338
			self.listening.close().unwrap();
		}
	}

    #[test]
    fn normal_request_possible() {
		let server = AutoServer::new();
		let client = Client::new();
		let res = client.get(&format!("http://127.0.0.1:{}/a", server.port)).send().unwrap();
		assert_eq!(res.status, status::ImATeapot);
	}

    #[test]
    fn preflight_without_origin_is_bad_request() {
		let server = AutoServer::new();
		let client = Client::new();
		let mut headers = Headers::new();
		headers.set(AccessControlRequestMethod(Method::Get));
		let mut res = client
				.request(Method::Options, &format!("http://127.0.0.1:{}/a", server.port))
				.headers(headers)
				.send().unwrap();
		assert_eq!(res.status, status::BadRequest);
        let mut payload = String::new();
        res.read_to_string(&mut payload).unwrap();
        assert_eq!(payload, "Preflight request without Origin header");
	}

    #[test]
    fn preflight_with_origin_accepts_same_origin() {
		let server = AutoServer::new();
		let client = Client::new();
		let mut headers = Headers::new();
		headers.set(AccessControlRequestMethod(Method::Get));
		headers.set(Origin::from_str("http://www.a.com:8080").unwrap());
		let res = client
				.request(Method::Options, &format!("http://127.0.0.1:{}/a", server.port))
				.headers(headers)
				.send().unwrap();
		assert_eq!(res.status, status::NoContent);
        let allow_origin = res.headers.get::<AccessControlAllowOrigin>().unwrap();
        assert_eq!(format!("{}", allow_origin) , "http://www.a.com:8080");
        let allow_headers = res.headers.get::<AccessControlAllowHeaders>().unwrap();
        assert_eq!(format!("{}", allow_headers) , "Content-Type, X-Requested-With");
        let allow_methods = res.headers.get::<AccessControlAllowMethods>().unwrap();
        assert_eq!(format!("{}", allow_methods) , "GET, PUT, POST");
        let max_age = res.headers.get::<AccessControlMaxAge>().unwrap();
        assert_eq!(max_age.0, 60*60u32);
	}

    #[test]
    fn normal_request_allows_origin() {
		let server = AutoServer::new();
		let client = Client::new();
		let mut headers = Headers::new();
		headers.set(Origin::from_str("http://www.a.com:8080").unwrap());
		let res = client
            .get(&format!("http://127.0.0.1:{}/a", server.port))
            .headers(headers)
            .send().unwrap();
		assert_eq!(res.status, status::ImATeapot);
        let allow_origin = res.headers.get::<AccessControlAllowOrigin>().unwrap();
        assert_eq!(format!("{}", allow_origin) , "http://www.a.com:8080");
        assert!(res.headers.get::<AccessControlAllowHeaders>().is_none());
        assert!(res.headers.get::<AccessControlAllowMethods>().is_none());
        assert!(res.headers.get::<AccessControlMaxAge>().is_none());
	}

    #[test]
    fn normal_request_without_origin_is_passthrough() {
		let server = AutoServer::new();
		let client = Client::new();
		let res = client
            .get(&format!("http://127.0.0.1:{}/a", server.port))
            .send().unwrap();
		assert_eq!(res.status, status::ImATeapot);
        assert!(res.headers.get::<AccessControlAllowOrigin>().is_none());
        assert!(res.headers.get::<AccessControlAllowHeaders>().is_none());
        assert!(res.headers.get::<AccessControlAllowMethods>().is_none());
        assert!(res.headers.get::<AccessControlMaxAge>().is_none());
	}

}
<|endoftext|>


<commit_before>extern crate libc;
extern crate winapi;
extern crate kernel32;

/// Allocates `size` Bytes aligned to `align` Bytes. Returns a null pointer on allocation failure.
///
/// The returned pointer must be deallocated by using `aligned_free`.
///
/// Note: This function is meant to be used for infrequent large allocations (as `malloc` already
/// guarantees suitable alignment for all native datatypes) and might be quite slow when used
/// heavily.
///
/// # Parameters
///
/// * `size`: The size of the allocation in bytes.
/// * `align`: The alignment of the allocation (at least the size of `usize` on the current
///   platform). Must also be a power of two.
#[inline]
pub fn aligned_alloc(size: usize, align: usize) -> *mut () {
    imp::aligned_alloc(size, align)
}

/// Deallocates aligned memory that was allocated with `aligned_alloc`. Unsafe because calling this
/// with a pointer that was not allocated with `aligned_alloc` (or already released) causes
/// undefined behavior.
#[inline]
pub unsafe fn aligned_free(ptr: *mut ()) {
    imp::aligned_free(ptr)
}

#[cfg(unix)]
mod imp {
    use libc::{c_void, c_int, size_t, EINVAL, ENOMEM, free};

    use std::{mem, ptr};

    extern "C" {
        fn posix_memalign(memptr: *mut *mut c_void, alignment: size_t, size: size_t) -> c_int;
    }

    pub fn aligned_alloc(size: usize, align: usize) -> *mut () {
        let mut memptr: *mut c_void = ptr::null_mut();
        let result = unsafe { posix_memalign(&mut memptr, align as size_t, size as size_t) };
        match result {
            0 => return memptr as *mut (),
            EINVAL => {
                if align < mem::size_of::<usize>() {
                    panic!("EINVAL: invalid alignment: {} (minimum is {})", align,
                        mem::size_of::<usize>());
                }
                if !align.is_power_of_two() {
                    panic!("EINVAL: invalid alignment: {} (must be a power of two)", align)
                }
                panic!("EINVAL: invalid alignment: {}", align);
            }
            ENOMEM => {
                // FIXME: Maybe handle this better
                panic!("ENOMEM: ran out of memory when attempting aligned allocation of {} bytes \
                    (aligned to {} bytes)", size, align)
            }
            _ => unreachable!(),
        }
    }

    #[inline]
    pub unsafe fn aligned_free(ptr: *mut ()) {
        free(ptr as *mut c_void)
    }
}

#[cfg(windows)]
mod imp {
    use kernel32::{GetLastError, VirtualAlloc, VirtualFree};
    use winapi::{MEM_COMMIT, MEM_RESERVE, MEM_RELEASE, PAGE_NOACCESS, PAGE_READWRITE, SIZE_T,
        LPVOID};

    use std::ptr;

    pub fn aligned_alloc(size: usize, align: usize) -> *mut () {
        unsafe {
            // Step 1: Reserve `size+align-1` Bytes of address space to find a suitable address
            let ptr = VirtualAlloc(ptr::null_mut(), (size + align - 1) as SIZE_T, MEM_RESERVE,
                PAGE_NOACCESS);
            if ptr.is_null() {
                panic!("WINAPI error {} while reserving memory", GetLastError());
            }

            // Step 2: Calculate an aligned address within the reserved range
            let aligned_ptr = (ptr as usize + align - 1) & !(align - 1);

            // Step 3: Actually allocate (commit) the memory
            let res = VirtualFree(ptr as LPVOID, 0, MEM_RELEASE);
            if res == 0 {
                panic!("WINAPI error {} while freeing reserved memory", GetLastError());
            }
            let ptr = VirtualAlloc(aligned_ptr as LPVOID, size as SIZE_T, MEM_COMMIT | MEM_RESERVE,
                PAGE_READWRITE);
            if ptr.is_null() {
                panic!("WINAPI error {} while allocating memory", GetLastError());
            }
            ptr as *mut ()
        }
    }

    pub unsafe fn aligned_free(ptr: *mut ()) {
        let res = VirtualFree(ptr as LPVOID, 0, MEM_RELEASE);
        if res == 0 {
            panic!("WINAPI error {} while releasing memory", GetLastError());
        }
    }
}

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

    #[test]
    fn test() {
        let ptr = aligned_alloc(1, 1024 * 1024);
        assert!(!ptr.is_null());
        assert!(ptr as usize % (1024 * 1024) == 0);
        unsafe { aligned_free(ptr) }
    }
}
<commit_msg>Return null on allocation failure<commit_after>extern crate libc;
extern crate winapi;
extern crate kernel32;

/// Allocates `size` Bytes aligned to `align` Bytes. Returns a null pointer on allocation failure.
///
/// The returned pointer must be deallocated by using `aligned_free`.
///
/// Note: This function is meant to be used for infrequent large allocations (as `malloc` already
/// guarantees suitable alignment for all native datatypes) and might be quite slow when used
/// heavily.
///
/// # Parameters
///
/// * `size`: The size of the allocation in bytes.
/// * `align`: The alignment of the allocation (at least the size of `usize` on the current
///   platform). Must also be a power of two.
#[inline]
pub fn aligned_alloc(size: usize, align: usize) -> *mut () {
    imp::aligned_alloc(size, align)
}

/// Deallocates aligned memory that was allocated with `aligned_alloc`. Unsafe because calling this
/// with a pointer that was not allocated with `aligned_alloc` (or already released) causes
/// undefined behavior.
#[inline]
pub unsafe fn aligned_free(ptr: *mut ()) {
    imp::aligned_free(ptr)
}

#[cfg(unix)]
mod imp {
    use libc::{c_void, c_int, size_t, EINVAL, ENOMEM, free};

    use std::{mem, ptr};

    extern "C" {
        fn posix_memalign(memptr: *mut *mut c_void, alignment: size_t, size: size_t) -> c_int;
    }

    pub fn aligned_alloc(size: usize, align: usize) -> *mut () {
        let mut memptr: *mut c_void = ptr::null_mut();
        let result = unsafe { posix_memalign(&mut memptr, align as size_t, size as size_t) };
        match result {
            0 => return memptr as *mut (),
            EINVAL => {
                if align < mem::size_of::<usize>() {
                    panic!("EINVAL: invalid alignment: {} (minimum is {})", align,
                        mem::size_of::<usize>());
                }
                if !align.is_power_of_two() {
                    panic!("EINVAL: invalid alignment: {} (must be a power of two)", align)
                }
                panic!("EINVAL: invalid alignment: {}", align);
            }
            ENOMEM => return ptr::null_mut(),
            _ => unreachable!(),
        }
    }

    #[inline]
    pub unsafe fn aligned_free(ptr: *mut ()) {
        free(ptr as *mut c_void)
    }
}

#[cfg(windows)]
mod imp {
    use kernel32::{GetLastError, VirtualAlloc, VirtualFree};
    use winapi::{MEM_COMMIT, MEM_RESERVE, MEM_RELEASE, PAGE_NOACCESS, PAGE_READWRITE, SIZE_T,
        LPVOID};

    use std::ptr;

    pub fn aligned_alloc(size: usize, align: usize) -> *mut () {
        unsafe {
            // Step 1: Reserve `size+align-1` Bytes of address space to find a suitable address
            let ptr = VirtualAlloc(ptr::null_mut(), (size + align - 1) as SIZE_T, MEM_RESERVE,
                PAGE_NOACCESS);
            if ptr.is_null() {
                panic!("WINAPI error {} while reserving memory", GetLastError());
            }

            // Step 2: Calculate an aligned address within the reserved range
            let aligned_ptr = (ptr as usize + align - 1) & !(align - 1);

            // Step 3: Actually allocate (commit) the memory
            let res = VirtualFree(ptr as LPVOID, 0, MEM_RELEASE);
            if res == 0 {
                panic!("WINAPI error {} while freeing reserved memory", GetLastError());
            }
            let ptr = VirtualAlloc(aligned_ptr as LPVOID, size as SIZE_T, MEM_COMMIT | MEM_RESERVE,
                PAGE_READWRITE);
            ptr as *mut ()
        }
    }

    pub unsafe fn aligned_free(ptr: *mut ()) {
        let res = VirtualFree(ptr as LPVOID, 0, MEM_RELEASE);
        if res == 0 {
            panic!("WINAPI error {} while releasing memory", GetLastError());
        }
    }
}

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

    #[test]
    fn test() {
        let ptr = aligned_alloc(1, 1024 * 1024);
        assert!(!ptr.is_null());
        assert!(ptr as usize % (1024 * 1024) == 0);
        unsafe { aligned_free(ptr) }
    }
}
<|endoftext|>


<commit_before>#![feature(vec_push_all)]

pub struct JsonEncoder {
    buffer: Vec<u8>
}

impl JsonEncoder {
    pub fn new() -> JsonEncoder {
        JsonEncoder{buffer: Vec::new()}
    }

    #[inline(always)]
    pub fn encode_raw(&mut self, raw: &[u8]) {
        self.buffer.push_all(raw);
    }

    pub fn into_vec(self) -> Vec<u8> {
        self.buffer
    }

    #[inline(always)]
    pub fn encode_raw_str(&mut self, raw_str: &str) {
        self.buffer.push_all(raw_str.as_bytes());
    }

    fn escape_bytes(&mut self, bytes: &[u8]) {
        let mut start = 0;

        for (i, byte) in bytes.iter().enumerate() {
            let escaped = match *byte {
                b'"' => b"\\\"",
                b'\\' => b"\\\\",
                b'\x08' => b"\\b",
                b'\x0c' => b"\\f",
                b'\n' => b"\\n",
                b'\r' => b"\\r",
                b'\t' => b"\\t",
                _ => {
                    continue;
                }
            };

            if start < i {
                self.encode_raw(&bytes[start..i]);
            }

            self.encode_raw(escaped);

            start = i + 1;
        }

        if start != bytes.len() {
            self.encode_raw(&bytes[start..]);
        }
    }

    pub fn encode_u64_decimal(&mut self, value: u64) {
        self.buffer.push(b'"');
        self._encode_u64_decimal(value);
        self.buffer.push(b'"');
    }

    // encodes as decimal string
    #[inline(always)]
    fn _encode_u64_decimal(&mut self, value: u64) {
        const CHARS: &'static [u8] = b"0123456789";
        const MAX_DIGITS: usize = 20;

        if value == 0 {
            self.buffer.push(b'0');
            return;
        }

        let mut digits: [u8; MAX_DIGITS] = [b'0'; MAX_DIGITS];
        let mut i = MAX_DIGITS;
        let mut n = value;
        while n > 0 {
            i -= 1;
            digits[i] = CHARS[(n % 10) as usize];
            n = n / 10;
        }

        self.encode_raw(&digits[i..]);
    }

    #[inline(always)]
    pub fn encode_str(&mut self, s: &str) {
        self.buffer.push(b'"');
        self.escape_bytes(s.as_bytes());
        self.buffer.push(b'"');
    }

    #[inline(always)]
    pub fn encode_i32(&mut self, value: i32) {
        if value == 0 {
            self.buffer.push(b'0');
        } else if value > 0 {
            self.encode_u31(value as u32);
        } else {
            self.buffer.push(b'-');
            self.encode_u31((-value) as u32);
        }
    }

    // encodes a 31-bit unsigned integer != 0
    fn encode_u31(&mut self, value: u32) {
        const CHARS: &'static [u8] = b"0123456789";
        const MAX_DIGITS: usize = 10;
        debug_assert!(value != 0);

        let mut digits: [u8; MAX_DIGITS] = [b'0'; MAX_DIGITS];
        let mut i = MAX_DIGITS;
        let mut n = value;
        while n > 0 {
            i -= 1;
            digits[i] = CHARS[(n % 10) as usize];
            n = n / 10;
        }

        self.encode_raw(&digits[i..]);
    }
}

pub struct JsonObjectEncoder {
    js: JsonEncoder,
    needs_sep: bool,
}

impl JsonObjectEncoder {
    pub fn new() -> JsonObjectEncoder {
        JsonObjectEncoder {
                js: JsonEncoder::new(),
                needs_sep: false
        }
    }

    pub fn begin(&mut self) {
        self.js.encode_raw(b"{");
    }

    pub fn end(&mut self) {
        self.js.encode_raw(b"}");
    }

    // XXX: name MAY NOT include escapable characters
    #[inline(always)]
    pub fn field<F:Fn(&mut JsonEncoder)>(&mut self, name: &str, f: F) {
        if self.needs_sep {
            self.js.buffer.push(b',');
        }
        self.js.buffer.push(b'"');
        self.js.encode_raw_str(name);
        self.js.buffer.push(b'"');
        self.js.buffer.push(b':');

        f(&mut self.js);
        self.needs_sep = true;
    }

    pub fn into_json_encoder(self) -> JsonEncoder {
        self.js
    }

    pub fn into_vec(self) -> Vec<u8> {
        self.into_json_encoder().into_vec()
    }
}

#[test]
fn test_json_obj_encoder() {
    let mut jso = JsonObjectEncoder::new();
    jso.begin();
    jso.end();
    assert_eq!(b"{}", &jso.into_vec()[..]);

    let mut jso = JsonObjectEncoder::new();
    jso.begin();
    jso.field("total", |js| js.encode_i32(31));
    jso.end();
    assert_eq!(b"{\"total\":31}", &jso.into_vec()[..]);

    let mut jso = JsonObjectEncoder::new();
    jso.begin();
    jso.field("total", |js| js.encode_i32(31));
    jso.field("next", |js| js.encode_str("abc"));
    jso.end();
    assert_eq!(b"{\"total\":31,\"next\":\"abc\"}", &jso.into_vec()[..]);
}
<commit_msg>Allow recursive JsonObject encoders<commit_after>#![feature(vec_push_all)]

pub struct JsonEncoder {
    buffer: Vec<u8>
}

impl JsonEncoder {
    pub fn new() -> JsonEncoder {
        JsonEncoder{buffer: Vec::new()}
    }

    #[inline(always)]
    pub fn encode_raw(&mut self, raw: &[u8]) {
        self.buffer.push_all(raw);
    }

    pub fn into_vec(self) -> Vec<u8> {
        self.buffer
    }

    #[inline(always)]
    pub fn encode_raw_str(&mut self, raw_str: &str) {
        self.buffer.push_all(raw_str.as_bytes());
    }

    fn escape_bytes(&mut self, bytes: &[u8]) {
        let mut start = 0;

        for (i, byte) in bytes.iter().enumerate() {
            let escaped = match *byte {
                b'"' => b"\\\"",
                b'\\' => b"\\\\",
                b'\x08' => b"\\b",
                b'\x0c' => b"\\f",
                b'\n' => b"\\n",
                b'\r' => b"\\r",
                b'\t' => b"\\t",
                _ => {
                    continue;
                }
            };

            if start < i {
                self.encode_raw(&bytes[start..i]);
            }

            self.encode_raw(escaped);

            start = i + 1;
        }

        if start != bytes.len() {
            self.encode_raw(&bytes[start..]);
        }
    }

    pub fn encode_u64_decimal(&mut self, value: u64) {
        self.buffer.push(b'"');
        self._encode_u64_decimal(value);
        self.buffer.push(b'"');
    }

    // encodes as decimal string
    #[inline(always)]
    fn _encode_u64_decimal(&mut self, value: u64) {
        const CHARS: &'static [u8] = b"0123456789";
        const MAX_DIGITS: usize = 20;

        if value == 0 {
            self.buffer.push(b'0');
            return;
        }

        let mut digits: [u8; MAX_DIGITS] = [b'0'; MAX_DIGITS];
        let mut i = MAX_DIGITS;
        let mut n = value;
        while n > 0 {
            i -= 1;
            digits[i] = CHARS[(n % 10) as usize];
            n = n / 10;
        }

        self.encode_raw(&digits[i..]);
    }

    #[inline(always)]
    pub fn encode_str(&mut self, s: &str) {
        self.buffer.push(b'"');
        self.escape_bytes(s.as_bytes());
        self.buffer.push(b'"');
    }

    #[inline(always)]
    pub fn encode_i32(&mut self, value: i32) {
        if value == 0 {
            self.buffer.push(b'0');
        } else if value > 0 {
            self.encode_u31(value as u32);
        } else {
            self.buffer.push(b'-');
            self.encode_u31((-value) as u32);
        }
    }

    // encodes a 31-bit unsigned integer != 0
    fn encode_u31(&mut self, value: u32) {
        const CHARS: &'static [u8] = b"0123456789";
        const MAX_DIGITS: usize = 10;
        debug_assert!(value != 0);

        let mut digits: [u8; MAX_DIGITS] = [b'0'; MAX_DIGITS];
        let mut i = MAX_DIGITS;
        let mut n = value;
        while n > 0 {
            i -= 1;
            digits[i] = CHARS[(n % 10) as usize];
            n = n / 10;
        }

        self.encode_raw(&digits[i..]);
    }

    pub fn encode_obj<F>(&mut self, f: F) where F: Fn(JsonObjectEncoder) {
        self.buffer.push(b'{');
        {
            f(JsonObjectEncoder {js: self, needs_sep: false});
        }
        self.buffer.push(b'}');
    }
}

pub struct JsonObjectEncoder<'a> {
    js: &'a mut JsonEncoder,
    needs_sep: bool,
}

impl<'a> JsonObjectEncoder<'a> {

    // XXX: name MAY NOT include escapable characters
    #[inline(always)]
    pub fn encode_field<F:Fn(&mut JsonEncoder)>(&mut self, name: &str, f: F) {
        if self.needs_sep {
            self.js.buffer.push(b',');
        }
        self.js.buffer.push(b'"');
        self.js.encode_raw_str(name);
        self.js.buffer.push(b'"');
        self.js.buffer.push(b':');

        f(self.js);
        self.needs_sep = true;
    }
}

#[test]
fn test_json_obj_encoder() {
    let mut js = JsonEncoder::new();
    js.encode_obj(|_|{});
    assert_eq!(b"{}", &js.into_vec()[..]);

    let mut js = JsonEncoder::new();
    js.encode_obj(|mut jso| {
        jso.encode_field("total", |js| js.encode_i32(31));
    });
    assert_eq!(b"{\"total\":31}", &js.into_vec()[..]);

    let mut js = JsonEncoder::new();
    js.encode_obj(|mut jso| {
        jso.encode_field("total", |js| js.encode_i32(31));
        jso.encode_field("next", |js| js.encode_str("abc"));
    });
    assert_eq!(b"{\"total\":31,\"next\":\"abc\"}", &js.into_vec()[..]);
}
<|endoftext|>


<commit_before>/*!
Crate `fst` is a library for efficiently storing and searching ordered sets or
maps where the keys are byte strings. A key design goal of this crate is to
support storing and searching *very large* sets or maps (i.e., billions). This
means that much effort has gone in to making sure that all operations are
memory efficient.

Sets and maps are represented by a finite state machine, which acts as a form
of compression on common prefixes and suffixes in the keys. Additionally,
finite state machines can be efficiently queried with automata (like regular
expressions or Levenshtein distance for fuzzy queries) or lexicographic ranges.

To read more about the mechanics of finite state transducers, including a
bibliography for algorithms used in this crate, see the docs for the
[`raw::Fst`](raw/struct.Fst.html) type.

# Installation

Simply add a corresponding entry to your `Cargo.toml` dependency list:

```ignore
[dependencies]
fst = "0.1"
```

And add this to your crate root:

```ignore
extern crate fst;
```

The examples in this documentation will show the rest.

# Overview of types and modules

This crate provides the high level abstractions---namely sets and maps---in the
top-level module. The top-level module also contains types for building
regular expressions (`Regex`) or fuzzy queries (`Levenshtein`), which can be
used to efficiently query a set or a map.

The `set` and `map` sub-modules contain types specific to sets and maps, such
as range queries and streams.

The `raw` module permits direct interaction with finite state transducers.
Namely, the states and transitions of a transducer can be directly accessed
with the `raw` module.

# Example: fuzzy query

This example shows how to create a set of strings in memory, and then execute
a fuzzy query. Namely, the query looks for all keys within an edit distance
of `1` of `foo`. (Edit distance is the number of character insertions,
deletions or substitutions required to get from one string to another. In this
case, a character is a Unicode codepoint.)

```rust
# fn example() -> Result<(), fst::Error> {
use fst::{IntoStreamer, Streamer, Levenshtein, Set};

// A convenient way to create sets in memory.
let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
let set = try!(Set::from_iter(keys));

// Build our fuzzy query.
let lev = try!(Levenshtein::new("foo", 1));

// Apply our fuzzy query to the set we built.
let mut stream = set.search(lev).into_stream();

// Iterate over the stream and collect the matching keys.
let mut keys = vec![];
while let Some(key) = stream.next() {
    keys.push(key.to_vec());
}
assert_eq!(keys, vec![
    "fo".as_bytes(),   // 1 deletion
    "fob".as_bytes(),  // 1 substitution
    "foo".as_bytes(),  // 0 insertions/deletions/substitutions
    "food".as_bytes(), // 1 insertion
]);
# Ok(())
# }
# example().unwrap();
```

# Example: stream a map to a file

This shows how to create a `MapBuilder` that will stream construction of the
map to a file. Notably, this will never store the entire transducer in memory.
Instead, only constant memory is required.

```rust,no_run
# fn example() -> Result<(), fst::Error> {
use std::fs::File;
use std::io;

use fst::{IntoStreamer, Streamer, Map, MapBuilder};

// This is where we'll write our map to.
let mut wtr = io::BufWriter::new(try!(File::create("map.fst")));

// Create a builder that can be used to insert new key-value pairs.
let mut build = try!(MapBuilder::new(wtr));
build.insert("bruce", 1).unwrap();
build.insert("clarence", 2).unwrap();
build.insert("stevie", 3).unwrap();

// Finish construction of the map and flush its contents to disk.
try!(build.finish());

// At this point, the map has been constructed. Now we'd like to search it.
// This creates a memory map, which enables searching the map without loading
// all of it into memory.
let map = try!(Map::from_path("map.fst"));

// Query for keys that are greater than or equal to clarence.
let mut stream = map.range().ge("clarence").into_stream();

// Iterate over the stream and collect the matching keys.
let mut kvs = vec![];
while let Some((k, v)) = stream.next() {
    kvs.push((k.to_vec(), v));
}
assert_eq!(kvs, vec![
    (b"clarence".to_vec(), 2),
    (b"stevie".to_vec(), 3),
]);
# Ok(())
# }
# example().unwrap();
```

# Example: case insensitive search

We can perform case insensitive search on a set using a regular expression.
Note that while sets can store arbitrary byte strings, a regular expression
will only match valid UTF-8 encoded byte strings.

```rust
# fn example() -> Result<(), fst::Error> {
use fst::{IntoStreamer, Streamer, Regex, Set};

let set = try!(Set::from_iter(&["FoO", "Foo", "fOO", "foo"]));

let re = try!(Regex::new("(?i)foo"));
let mut stream = set.search(&re).into_stream();

let mut keys = vec![];
while let Some(key) = stream.next() {
    keys.push(key.to_vec());
}
assert_eq!(keys, vec![
    "FoO".as_bytes(),
    "Foo".as_bytes(),
    "fOO".as_bytes(),
    "foo".as_bytes(),
]);
# Ok(())
# }
# example().unwrap();
```

# Example: searching multiple sets efficiently

Since queries can search a transducer without reading the entire data structure
into memory, it is possible to search *many* transducers very quickly.

This crate provides efficient set/map operations that allow one to combine
multiple streams of search results. Each operation only uses memory
proportional to the number of streams.

The example below shows how to find all keys that have at least one capital
letter that doesn't appear at the beginning of the key. The example below uses
sets, but the same operations are available on maps too.

```rust
# fn example() -> Result<(), fst::Error> {
use fst::{Streamer, Regex, Set};
use fst::set;

let set1 = try!(Set::from_iter(&["AC/DC", "Aerosmith"]));
let set2 = try!(Set::from_iter(&["Bob Seger", "Bruce Springsteen"]));
let set3 = try!(Set::from_iter(&["George Thorogood", "Golden Earring"]));
let set4 = try!(Set::from_iter(&["Kansas"]));
let set5 = try!(Set::from_iter(&["Metallica"]));

// Create the regular expression. We can reuse it to search all of the sets.
let re = try!(Regex::new(r".+\p{Lu}.*"));

// Build a set operation. All we need to do is add a search result stream for
// each set and ask for the union. (Other operations, like intersection and
// difference are also available.)
let mut stream =
    set::OpBuilder::new()
    .add(set1.search(&re))
    .add(set2.search(&re))
    .add(set3.search(&re))
    .add(set4.search(&re))
    .add(set5.search(&re))
    .union();

// Now collect all of the keys. Alternatively, you could build another set here
// using `SetBuilder::extend_stream`.
let mut keys = vec![];
while let Some(key) = stream.next() {
    keys.push(key.to_vec());
}
assert_eq!(keys, vec![
    "AC/DC".as_bytes(),
    "Bob Seger".as_bytes(),
    "Bruce Springsteen".as_bytes(),
    "George Thorogood".as_bytes(),
    "Golden Earring".as_bytes(),
]);
# Ok(())
# }
# example().unwrap();
```

# Memory usage

An important advantage of using finite state transducers to represent sets and
maps is that they can compress very well depending on the distribution of keys.
The smaller your set/map is, the more likely it is that it will fit into
memory. If it's in memory, then searching it is faster. Therefore, it is
important to do what we can to limit what actually needs to be in memory.

This is where automata shine, because they can be queried in their compressed
state without loading the entire data structure into memory. This means that
one can store a set/map created by this crate on disk and search it without
actually reading the entire set/map into memory. This use case is served well
by *memory maps*, which lets one assign the entire contents of a file to a
contiguous region of virtual memory.

Indeed, this crate encourages this mode of operation. Both sets and maps have
methods for memory mapping a finite state transducer from disk.

This is particularly important for long running processes that use this crate,
since it enables the operating system to determine which regions of your
finite state transducers are actually in memory.

Of course, there are downsides to this approach. Namely, navigating a
transducer during a key lookup or a search will likely follow a pattern
approximating random access. Supporting random access when reading from disk
can be very slow because of how often `seek` must be called (or, in the case
of memory maps, page faults). This is somewhat mitigated by the prevalence of
solid state drives where seek time is eliminated. Nevertheless, solid state
drives are not ubiquitous and it is possible that the OS will not be smart
enough to keep your memory mapped transducers in the page cache. In that case,
it is advisable to load the entire transducer into your process's memory (e.g.,
`Set::from_bytes`).

# Streams

Searching a set or a map needs to provide some way to iterate over the search
results. Idiomatic Rust calls for something satisfying the `Iterator` trait
to be used here. Unfortunately, this is not possible to do efficiently because
the `Iterator` trait does not permit values emitted by the iterator to borrow
from the iterator. Borrowing from the iterator is required in our case because
keys and values are constructed *during iteration*.

Namely, if we were to use iterators, then every key would need its own
allocation, which could be quite costly.

Instead, this crate provides a `Streamer`, which can be thought of as a
streaming iterator. Namely, a stream in this crate maintains a single key
buffer and lends it out on each iteration.

For more details, including important limitations, see the `Streamer` trait.

# Quirks

There's no doubt about it, finite state transducers are a specialty data
structure. They have a host of restrictions that don't apply to other similar
data structures found in the standard library, such as `BTreeSet` and
`BTreeMap`. Here are some of them:

1. Sets can only contain keys that are byte strings.
2. Maps can also only contain keys that are byte strings, and its values are
   limited to unsigned 64 bit integers. (The restriction on values may be
   relaxed some day.)
3. Creating a set or a map requires inserting keys in lexicographic order.
   Often, keys are not already sorted, which can make constructing large
   sets or maps tricky. One way to do it is to sort pieces of the data and
   build a set/map for each piece. This can be parallelized trivially. Once
   done, they can be merged together into one big set/map if desired.
   A somewhat simplistic example of this procedure can be seen in
   `fst-bin/src/merge.rs` from the root of this crate's repository.

# Warning: regexes and Levenshtein automatons use a lot of memory

The construction of automatons for both regular expressions and Levenshtein
automatons should be consider "proof of concept" quality. Namely, they do just
enough to be *correct*. But they haven't had any effort put into them to be
memory conscious. These are important parts of this library, so they will be
improved.

Note that whether you're using regexes or Levenshtein automatons, an error
will be returned if the automaton gets too big (tens of MB in heap usage).
*/

#![deny(missing_docs)]

extern crate byteorder;
extern crate memchr;
extern crate memmap;
#[cfg(test)] extern crate quickcheck;
#[cfg(test)] extern crate rand;
extern crate regex_syntax;
extern crate utf8_ranges;

pub use automaton::Automaton;
pub use error::{Error, Result};
pub use levenshtein::Error as LevenshteinError;
pub use levenshtein::Levenshtein;
pub use map::{Map, MapBuilder};
pub use regex::Error as RegexError;
pub use regex::Regex;
pub use set::{Set, SetBuilder};
pub use stream::{IntoStreamer, Streamer};

mod automaton;
mod error;
mod levenshtein;
#[path = "map.rs"]
mod inner_map;
pub mod raw;
mod regex;
#[path = "set.rs"]
mod inner_set;
mod stream;

/// Map operations implemented by finite state transducers.
///
/// This API provided by this sub-module is close in spirit to the API
/// provided by
/// [`std::collections::BTreeMap`](http://doc.rust-lang.org/stable/std/collections/struct.BTreeMap.html).
///
/// # Overview of types
///
/// `Map` is a read only interface to pre-constructed sets. `MapBuilder` is
/// used to create new sets. (Once a set is created, it can never be modified.)
/// `Stream`, `Keys` and `Values` are streams that originated from a map.
/// `StreamBuilder` builds range queries. `OpBuilder` collects a set of streams
/// and executes set operations like `union` or `intersection` on them with the
/// option of specifying a merge strategy for a map's values. The rest of the
/// types are streams for set operations.
pub mod map {
    pub use inner_map::*;
}

/// Set operations implemented by finite state transducers.
///
/// This API provided by this sub-module is close in spirit to the API
/// provided by
/// [`std::collections::BTreeSet`](http://doc.rust-lang.org/stable/std/collections/struct.BTreeSet.html).
/// The principle difference, as with everything else in this crate, is that
/// operations are performed on streams of byte strings instead of generic
/// iterators. Another difference is that most of the set operations (union,
/// intersection, difference and symmetric difference) work on multiple sets at
/// the same time, instead of two.
///
/// # Overview of types
///
/// `Set` is a read only interface to pre-constructed sets. `SetBuilder` is
/// used to create new sets. (Once a set is created, it can never be modified.)
/// `Stream` is a stream of values that originated from a set (analogous to an
/// iterator). `StreamBuilder` builds range queries. `OpBuilder` collects a set
/// of streams and executes set operations like `union` or `intersection` on
/// them. The rest of the types are streams for set operations.
pub mod set {
    pub use inner_set::*;
}
<commit_msg>doc touchups<commit_after>/*!
Crate `fst` is a library for efficiently storing and searching ordered sets or
maps where the keys are byte strings. A key design goal of this crate is to
support storing and searching *very large* sets or maps (i.e., billions). This
means that much effort has gone in to making sure that all operations are
memory efficient.

Sets and maps are represented by a finite state machine, which acts as a form
of compression on common prefixes and suffixes in the keys. Additionally,
finite state machines can be efficiently queried with automata (like regular
expressions or Levenshtein distance for fuzzy queries) or lexicographic ranges.

To read more about the mechanics of finite state transducers, including a
bibliography for algorithms used in this crate, see the docs for the
[`raw::Fst`](raw/struct.Fst.html) type.

# Installation

Simply add a corresponding entry to your `Cargo.toml` dependency list:

```ignore
[dependencies]
fst = "0.1"
```

And add this to your crate root:

```ignore
extern crate fst;
```

The examples in this documentation will show the rest.

# Overview of types and modules

This crate provides the high level abstractions---namely sets and maps---in the
top-level module. The top-level module also contains types for building
regular expressions (`Regex`) or fuzzy queries (`Levenshtein`), which can be
used to efficiently query a set or a map.

The `set` and `map` sub-modules contain types specific to sets and maps, such
as range queries and streams.

The `raw` module permits direct interaction with finite state transducers.
Namely, the states and transitions of a transducer can be directly accessed
with the `raw` module.

# Example: fuzzy query

This example shows how to create a set of strings in memory, and then execute
a fuzzy query. Namely, the query looks for all keys within an edit distance
of `1` of `foo`. (Edit distance is the number of character insertions,
deletions or substitutions required to get from one string to another. In this
case, a character is a Unicode codepoint.)

```rust
# fn example() -> Result<(), fst::Error> {
use fst::{IntoStreamer, Streamer, Levenshtein, Set};

// A convenient way to create sets in memory.
let keys = vec!["fa", "fo", "fob", "focus", "foo", "food", "foul"];
let set = try!(Set::from_iter(keys));

// Build our fuzzy query.
let lev = try!(Levenshtein::new("foo", 1));

// Apply our fuzzy query to the set we built.
let mut stream = set.search(lev).into_stream();

let keys = try!(stream.into_strs());
assert_eq!(keys, vec!["fo", "fob", "foo", "food"]);
# Ok(())
# }
# example().unwrap();
```

# Example: stream a map to a file

This shows how to create a `MapBuilder` that will stream construction of the
map to a file. Notably, this will never store the entire transducer in memory.
Instead, only constant memory is required.

```rust,no_run
# fn example() -> Result<(), fst::Error> {
use std::fs::File;
use std::io;

use fst::{IntoStreamer, Streamer, Map, MapBuilder};

// This is where we'll write our map to.
let mut wtr = io::BufWriter::new(try!(File::create("map.fst")));

// Create a builder that can be used to insert new key-value pairs.
let mut build = try!(MapBuilder::new(wtr));
build.insert("bruce", 1).unwrap();
build.insert("clarence", 2).unwrap();
build.insert("stevie", 3).unwrap();

// Finish construction of the map and flush its contents to disk.
try!(build.finish());

// At this point, the map has been constructed. Now we'd like to search it.
// This creates a memory map, which enables searching the map without loading
// all of it into memory.
let map = try!(Map::from_path("map.fst"));

// Query for keys that are greater than or equal to clarence.
let mut stream = map.range().ge("clarence").into_stream();

let kvs = try!(stream.into_str_vec());
assert_eq!(kvs, vec![
    ("clarence".to_owned(), 2),
    ("stevie".to_owned(), 3),
]);
# Ok(())
# }
# example().unwrap();
```

# Example: case insensitive search

We can perform case insensitive search on a set using a regular expression.
Note that while sets can store arbitrary byte strings, a regular expression
will only match valid UTF-8 encoded byte strings.

```rust
# fn example() -> Result<(), fst::Error> {
use fst::{IntoStreamer, Streamer, Regex, Set};

let set = try!(Set::from_iter(&["FoO", "Foo", "fOO", "foo"]));

let re = try!(Regex::new("(?i)foo"));
let mut stream = set.search(&re).into_stream();

let keys = try!(stream.into_strs());
assert_eq!(keys, vec!["FoO", "Foo", "fOO", "foo"]);
# Ok(())
# }
# example().unwrap();
```

# Example: searching multiple sets efficiently

Since queries can search a transducer without reading the entire data structure
into memory, it is possible to search *many* transducers very quickly.

This crate provides efficient set/map operations that allow one to combine
multiple streams of search results. Each operation only uses memory
proportional to the number of streams.

The example below shows how to find all keys that have at least one capital
letter that doesn't appear at the beginning of the key. The example below uses
sets, but the same operations are available on maps too.

```rust
# fn example() -> Result<(), fst::Error> {
use fst::{Streamer, Regex, Set};
use fst::set;

let set1 = try!(Set::from_iter(&["AC/DC", "Aerosmith"]));
let set2 = try!(Set::from_iter(&["Bob Seger", "Bruce Springsteen"]));
let set3 = try!(Set::from_iter(&["George Thorogood", "Golden Earring"]));
let set4 = try!(Set::from_iter(&["Kansas"]));
let set5 = try!(Set::from_iter(&["Metallica"]));

// Create the regular expression. We can reuse it to search all of the sets.
let re = try!(Regex::new(r".+\p{Lu}.*"));

// Build a set operation. All we need to do is add a search result stream for
// each set and ask for the union. (Other operations, like intersection and
// difference are also available.)
let mut stream =
    set::OpBuilder::new()
    .add(set1.search(&re))
    .add(set2.search(&re))
    .add(set3.search(&re))
    .add(set4.search(&re))
    .add(set5.search(&re))
    .union();

// Now collect all of the keys. Alternatively, you could build another set here
// using `SetBuilder::extend_stream`.
let mut keys = vec![];
while let Some(key) = stream.next() {
    keys.push(key.to_vec());
}
assert_eq!(keys, vec![
    "AC/DC".as_bytes(),
    "Bob Seger".as_bytes(),
    "Bruce Springsteen".as_bytes(),
    "George Thorogood".as_bytes(),
    "Golden Earring".as_bytes(),
]);
# Ok(())
# }
# example().unwrap();
```

# Memory usage

An important advantage of using finite state transducers to represent sets and
maps is that they can compress very well depending on the distribution of keys.
The smaller your set/map is, the more likely it is that it will fit into
memory. If it's in memory, then searching it is faster. Therefore, it is
important to do what we can to limit what actually needs to be in memory.

This is where automata shine, because they can be queried in their compressed
state without loading the entire data structure into memory. This means that
one can store a set/map created by this crate on disk and search it without
actually reading the entire set/map into memory. This use case is served well
by *memory maps*, which lets one assign the entire contents of a file to a
contiguous region of virtual memory.

Indeed, this crate encourages this mode of operation. Both sets and maps have
methods for memory mapping a finite state transducer from disk.

This is particularly important for long running processes that use this crate,
since it enables the operating system to determine which regions of your
finite state transducers are actually in memory.

Of course, there are downsides to this approach. Namely, navigating a
transducer during a key lookup or a search will likely follow a pattern
approximating random access. Supporting random access when reading from disk
can be very slow because of how often `seek` must be called (or, in the case
of memory maps, page faults). This is somewhat mitigated by the prevalence of
solid state drives where seek time is eliminated. Nevertheless, solid state
drives are not ubiquitous and it is possible that the OS will not be smart
enough to keep your memory mapped transducers in the page cache. In that case,
it is advisable to load the entire transducer into your process's memory (e.g.,
`Set::from_bytes`).

# Streams

Searching a set or a map needs to provide some way to iterate over the search
results. Idiomatic Rust calls for something satisfying the `Iterator` trait
to be used here. Unfortunately, this is not possible to do efficiently because
the `Iterator` trait does not permit values emitted by the iterator to borrow
from the iterator. Borrowing from the iterator is required in our case because
keys and values are constructed *during iteration*.

Namely, if we were to use iterators, then every key would need its own
allocation, which could be quite costly.

Instead, this crate provides a `Streamer`, which can be thought of as a
streaming iterator. Namely, a stream in this crate maintains a single key
buffer and lends it out on each iteration.

For more details, including important limitations, see the `Streamer` trait.

# Quirks

There's no doubt about it, finite state transducers are a specialty data
structure. They have a host of restrictions that don't apply to other similar
data structures found in the standard library, such as `BTreeSet` and
`BTreeMap`. Here are some of them:

1. Sets can only contain keys that are byte strings.
2. Maps can also only contain keys that are byte strings, and its values are
   limited to unsigned 64 bit integers. (The restriction on values may be
   relaxed some day.)
3. Creating a set or a map requires inserting keys in lexicographic order.
   Often, keys are not already sorted, which can make constructing large
   sets or maps tricky. One way to do it is to sort pieces of the data and
   build a set/map for each piece. This can be parallelized trivially. Once
   done, they can be merged together into one big set/map if desired.
   A somewhat simplistic example of this procedure can be seen in
   `fst-bin/src/merge.rs` from the root of this crate's repository.

# Warning: regexes and Levenshtein automatons use a lot of memory

The construction of automatons for both regular expressions and Levenshtein
automatons should be consider "proof of concept" quality. Namely, they do just
enough to be *correct*. But they haven't had any effort put into them to be
memory conscious. These are important parts of this library, so they will be
improved.

Note that whether you're using regexes or Levenshtein automatons, an error
will be returned if the automaton gets too big (tens of MB in heap usage).
*/

#![deny(missing_docs)]

extern crate byteorder;
extern crate memchr;
extern crate memmap;
#[cfg(test)] extern crate quickcheck;
#[cfg(test)] extern crate rand;
extern crate regex_syntax;
extern crate utf8_ranges;

pub use automaton::Automaton;
pub use error::{Error, Result};
pub use levenshtein::Error as LevenshteinError;
pub use levenshtein::Levenshtein;
pub use map::{Map, MapBuilder};
pub use regex::Error as RegexError;
pub use regex::Regex;
pub use set::{Set, SetBuilder};
pub use stream::{IntoStreamer, Streamer};

mod automaton;
mod error;
mod levenshtein;
#[path = "map.rs"]
mod inner_map;
pub mod raw;
mod regex;
#[path = "set.rs"]
mod inner_set;
mod stream;

/// Map operations implemented by finite state transducers.
///
/// This API provided by this sub-module is close in spirit to the API
/// provided by
/// [`std::collections::BTreeMap`](http://doc.rust-lang.org/stable/std/collections/struct.BTreeMap.html).
///
/// # Overview of types
///
/// `Map` is a read only interface to pre-constructed sets. `MapBuilder` is
/// used to create new sets. (Once a set is created, it can never be modified.)
/// `Stream`, `Keys` and `Values` are streams that originated from a map.
/// `StreamBuilder` builds range queries. `OpBuilder` collects a set of streams
/// and executes set operations like `union` or `intersection` on them with the
/// option of specifying a merge strategy for a map's values. The rest of the
/// types are streams for set operations.
pub mod map {
    pub use inner_map::*;
}

/// Set operations implemented by finite state transducers.
///
/// This API provided by this sub-module is close in spirit to the API
/// provided by
/// [`std::collections::BTreeSet`](http://doc.rust-lang.org/stable/std/collections/struct.BTreeSet.html).
/// The principle difference, as with everything else in this crate, is that
/// operations are performed on streams of byte strings instead of generic
/// iterators. Another difference is that most of the set operations (union,
/// intersection, difference and symmetric difference) work on multiple sets at
/// the same time, instead of two.
///
/// # Overview of types
///
/// `Set` is a read only interface to pre-constructed sets. `SetBuilder` is
/// used to create new sets. (Once a set is created, it can never be modified.)
/// `Stream` is a stream of values that originated from a set (analogous to an
/// iterator). `StreamBuilder` builds range queries. `OpBuilder` collects a set
/// of streams and executes set operations like `union` or `intersection` on
/// them. The rest of the types are streams for set operations.
pub mod set {
    pub use inner_set::*;
}
<|endoftext|>


<commit_before>//! DNS resolver for Hyper

#[macro_use]
extern crate log;
extern crate hyper;
extern crate rand;
extern crate tokio;
extern crate trust_dns;

use trust_dns::client::ClientHandle;
use rand::Rng;
use hyper::net::{NetworkConnector, NetworkStream};
use std::time::Duration;

/// Docs
#[derive(Debug, Clone)]
pub enum RecordType {
    /// A
    A,
    /// SRV
    SRV,
    /// AUTO
    AUTO,
}

/// A connector that wraps another connector and provides custom DNS resolution.
#[derive(Debug, Clone)]
pub struct DnsConnector<C: NetworkConnector> {
    connector: C,
    dns_addr: std::net::SocketAddr,
    record_type: RecordType,
}

/// Docs
impl<C: NetworkConnector> DnsConnector<C> {
    /// Docs
    pub fn new(dns_addr: std::net::SocketAddr, connector: C) -> DnsConnector<C> {
        Self::new_with_resolve_type(dns_addr, connector, RecordType::AUTO)
    }

    /// Docs
    pub fn new_with_resolve_type(dns_addr: std::net::SocketAddr,
                                 connector: C,
                                 record_type: RecordType)
                                 -> DnsConnector<C> {
        DnsConnector {
            connector: connector,
            dns_addr: dns_addr,
            record_type: record_type,
        }
    }
}

impl<C: NetworkConnector<Stream = S>, S: NetworkStream + Send> NetworkConnector
    for DnsConnector<C> {
    type Stream = S;

    /// Performs DNS SRV resolution, then calls into the inner connector with the results.
    /// Note that currently this does not take into account the following in the SRV record:
    /// * weight
    /// * priority
    /// It just takes a random entry in the DNS answers that are returned.
    fn connect(&self, host: &str, port: u16, scheme: &str) -> hyper::Result<S> {

        let mut runtime = tokio::runtime::current_thread::Runtime::new()
            .expect("Failed to create event loop for DNS query");
        let (stream, sender) = trust_dns::tcp::TcpClientStream::new(self.dns_addr);

        // We would expect a DNS request to be responded to quickly, but add a timeout
        // to ensure that we don't wait for ever if the DNS server does not respond.
        let timeout = Duration::from_millis(30000);
        let dns_client =
            trust_dns::client::ClientFuture::with_timeout(
                stream,
                sender,
                timeout,
                None);
        let mut dns_client = runtime.block_on(dns_client).unwrap();

        // Check if this is a domain name or not before trying to use DNS resolution.
        let (ref host, port) = match host.parse() {
            Err(_) => {

                debug!("Trying to resolve {}://{}", scheme, &host);

                // Add a `.` to the end of the host so that we can query the domain records.
                let name = trust_dns::rr::Name::parse(&format!("{}.", host), None).unwrap();

                let trust_record_type = match self.record_type {
                    RecordType::A => trust_dns::rr::RecordType::A,
                    RecordType::SRV => trust_dns::rr::RecordType::SRV,
                    RecordType::AUTO => {
                        // If the port is a standard HTTP port (80, or 443), then assume
                        // one was not provided and perform SRV lookup, otherwise lookup
                        // A records.
                        if (port == 80) || (port == 443) {
                            trust_dns::rr::RecordType::SRV
                        } else {
                            debug!("Using A record lookup for: {}", &host);
                            trust_dns::rr::RecordType::A
                        }
                    }
                };

                debug!("Sending DNS request");

                let query = dns_client.query(name.clone(),
                                              trust_dns::rr::DNSClass::IN,
                                              trust_record_type);
                let result = runtime.block_on(query);
                match result {
                    Ok(res) => {
                        let answers = res.answers();

                        if answers.is_empty() {
                            return Err(std::io::Error::new(std::io::ErrorKind::Other,
                                                           "No valid DNS answers")
                                               .into());
                        }

                        let mut rng = rand::thread_rng();

                        // First find the SRV records if they were requested
                        let (target, a_records, new_port) = if let trust_dns::rr::RecordType::SRV =
                            trust_record_type {
                            let answer = rng.choose(answers).expect("Sort out what to return here");

                            let srv = match *answer.rdata() {
                                trust_dns::rr::RData::SRV(ref srv) => srv,
                                _ => {
                                    return Err(std::io::Error::new(std::io::ErrorKind::Other,
                                                                   "Unexpected DNS response")
                                                       .into())
                                }
                            };

                            (srv.target(), res.additionals(), srv.port())
                        } else {
                            // For A record requests it is the domain name that
                            // we want to use.
                            (&name, answers, port)
                        };

                        let entry = a_records.iter().find(|record| record.name() == target);

                        if let Some(entry) = entry {
                            let addr = match *entry.rdata() {
                                trust_dns::rr::RData::A(ref addr) => addr,
                                _ => {
                                    return Err(std::io::Error::new(std::io::ErrorKind::Other,
                                                                   "Did not receive a valid record")
                                                       .into())
                                }
                            };

                            (addr.to_string(), new_port)
                        } else {
                            return Err(std::io::Error::new(std::io::ErrorKind::Other,
                                                           "Did not receive a valid record")
                                               .into());
                        }

                    }
                    _ => {
                        return Err(std::io::Error::new(std::io::ErrorKind::Other,
                                                       "Failed to query DNS server")
                                           .into())
                    }
                }
            }
            Ok(std::net::Ipv4Addr { .. }) => (host.to_string(), port),
        };

        debug!("Resolved request to {}://{}:{}", scheme, &host, port);

        self.connector.connect(&host, port, scheme)
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {}
}
<commit_msg>cargo fmt<commit_after>//! DNS resolver for Hyper

#[macro_use]
extern crate log;
extern crate hyper;
extern crate rand;
extern crate tokio;
extern crate trust_dns;

use hyper::net::{NetworkConnector, NetworkStream};
use rand::Rng;
use std::time::Duration;
use trust_dns::client::ClientHandle;

/// Docs
#[derive(Debug, Clone)]
pub enum RecordType {
    /// A
    A,
    /// SRV
    SRV,
    /// AUTO
    AUTO,
}

/// A connector that wraps another connector and provides custom DNS resolution.
#[derive(Debug, Clone)]
pub struct DnsConnector<C: NetworkConnector> {
    connector: C,
    dns_addr: std::net::SocketAddr,
    record_type: RecordType,
}

/// Docs
impl<C: NetworkConnector> DnsConnector<C> {
    /// Docs
    pub fn new(dns_addr: std::net::SocketAddr, connector: C) -> DnsConnector<C> {
        Self::new_with_resolve_type(dns_addr, connector, RecordType::AUTO)
    }

    /// Docs
    pub fn new_with_resolve_type(
        dns_addr: std::net::SocketAddr,
        connector: C,
        record_type: RecordType,
    ) -> DnsConnector<C> {
        DnsConnector {
            connector: connector,
            dns_addr: dns_addr,
            record_type: record_type,
        }
    }
}

impl<C: NetworkConnector<Stream = S>, S: NetworkStream + Send> NetworkConnector
    for DnsConnector<C>
{
    type Stream = S;

    /// Performs DNS SRV resolution, then calls into the inner connector with the results.
    /// Note that currently this does not take into account the following in the SRV record:
    /// * weight
    /// * priority
    /// It just takes a random entry in the DNS answers that are returned.
    fn connect(&self, host: &str, port: u16, scheme: &str) -> hyper::Result<S> {
        let mut runtime = tokio::runtime::current_thread::Runtime::new()
            .expect("Failed to create event loop for DNS query");
        let (stream, sender) = trust_dns::tcp::TcpClientStream::new(self.dns_addr);

        // We would expect a DNS request to be responded to quickly, but add a timeout
        // to ensure that we don't wait for ever if the DNS server does not respond.
        let timeout = Duration::from_millis(30000);
        let dns_client =
            trust_dns::client::ClientFuture::with_timeout(stream, sender, timeout, None);
        let mut dns_client = runtime.block_on(dns_client).unwrap();

        // Check if this is a domain name or not before trying to use DNS resolution.
        let (ref host, port) = match host.parse() {
            Err(_) => {
                debug!("Trying to resolve {}://{}", scheme, &host);

                // Add a `.` to the end of the host so that we can query the domain records.
                let name = trust_dns::rr::Name::parse(&format!("{}.", host), None).unwrap();

                let trust_record_type = match self.record_type {
                    RecordType::A => trust_dns::rr::RecordType::A,
                    RecordType::SRV => trust_dns::rr::RecordType::SRV,
                    RecordType::AUTO => {
                        // If the port is a standard HTTP port (80, or 443), then assume
                        // one was not provided and perform SRV lookup, otherwise lookup
                        // A records.
                        if (port == 80) || (port == 443) {
                            trust_dns::rr::RecordType::SRV
                        } else {
                            debug!("Using A record lookup for: {}", &host);
                            trust_dns::rr::RecordType::A
                        }
                    }
                };

                debug!("Sending DNS request");

                let query =
                    dns_client.query(name.clone(), trust_dns::rr::DNSClass::IN, trust_record_type);
                let result = runtime.block_on(query);
                match result {
                    Ok(res) => {
                        let answers = res.answers();

                        if answers.is_empty() {
                            return Err(std::io::Error::new(
                                std::io::ErrorKind::Other,
                                "No valid DNS answers",
                            ).into());
                        }

                        let mut rng = rand::thread_rng();

                        // First find the SRV records if they were requested
                        let (target, a_records, new_port) = if let trust_dns::rr::RecordType::SRV =
                            trust_record_type
                        {
                            let answer = rng.choose(answers).expect("Sort out what to return here");

                            let srv = match *answer.rdata() {
                                trust_dns::rr::RData::SRV(ref srv) => srv,
                                _ => {
                                    return Err(std::io::Error::new(
                                        std::io::ErrorKind::Other,
                                        "Unexpected DNS response",
                                    ).into())
                                }
                            };

                            (srv.target(), res.additionals(), srv.port())
                        } else {
                            // For A record requests it is the domain name that
                            // we want to use.
                            (&name, answers, port)
                        };

                        let entry = a_records.iter().find(|record| record.name() == target);

                        if let Some(entry) = entry {
                            let addr = match *entry.rdata() {
                                trust_dns::rr::RData::A(ref addr) => addr,
                                _ => {
                                    return Err(std::io::Error::new(
                                        std::io::ErrorKind::Other,
                                        "Did not receive a valid record",
                                    ).into())
                                }
                            };

                            (addr.to_string(), new_port)
                        } else {
                            return Err(std::io::Error::new(
                                std::io::ErrorKind::Other,
                                "Did not receive a valid record",
                            ).into());
                        }
                    }
                    _ => {
                        return Err(std::io::Error::new(
                            std::io::ErrorKind::Other,
                            "Failed to query DNS server",
                        ).into())
                    }
                }
            }
            Ok(std::net::Ipv4Addr { .. }) => (host.to_string(), port),
        };

        debug!("Resolved request to {}://{}:{}", scheme, &host, port);

        self.connector.connect(&host, port, scheme)
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {}
}
<|endoftext|>