<commit_before>extern crate winapi;
extern crate kernel32;

use util::{ensure_compatible_types, cstr_cow_from_bytes};

use std::ffi::{OsStr, OsString};
use std::{fmt, io, marker, mem, ptr};
use std::os::windows::ffi::{OsStrExt, OsStringExt};
use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};

/// A platform-specific equivalent of the cross-platform `Library`.
pub struct Library(winapi::HMODULE);

unsafe impl Send for Library {}
// Now, this is sort-of-tricky. MSDN documentation does not really make any claims as to safety of
// the Win32 APIs. Sadly, whomever I asked, even current and former Microsoft employees, couldn’t
// say for sure, whether the Win32 APIs used to implement `Library` are thread-safe or not.
//
// My investigation ended up with a question about thread-safety properties of the API involved
// being sent to an internal (to MS) general question mailing-list. The conclusion of the mail is
// as such:
//
// * Nobody inside MS (at least out of all the people who have seen the question) knows for
//   sure either;
// * However, the general consensus between MS developers is that one can rely on the API being
//   thread-safe. In case it is not thread-safe it should be considered a bug on the Windows
//   part. (NB: bugs filled at https://connect.microsoft.com/ against Windows Server)
unsafe impl Sync for Library {}

impl Library {
    /// Find and load a shared library (module).
    ///
    /// Corresponds to `LoadLibraryW(filename)`.
    #[inline]
    pub fn new<P: AsRef<OsStr>>(filename: P) -> ::Result<Library> {
        let wide_filename: Vec<u16> = filename.as_ref().encode_wide().chain(Some(0)).collect();
        let _guard = ErrorModeGuard::new();

        let ret = with_get_last_error(|| {
            // Make sure no winapi calls as a result of drop happen inside this closure, because
            // otherwise that might change the return value of the GetLastError.
            let handle = unsafe { kernel32::LoadLibraryW(wide_filename.as_ptr()) };
            if handle.is_null()  {
                None
            } else {
                Some(Library(handle))
            }
        }).map_err(|e| e.unwrap_or_else(||
            panic!("LoadLibraryW failed but GetLastError did not report the error")
        ));

        drop(wide_filename); // Drop wide_filename here to ensure it doesn’t get moved and dropped
                             // inside the closure by mistake. See comment inside the closure.
        ret
    }

    /// Get a pointer to function or static variable by symbol name.
    ///
    /// The `symbol` may not contain any null bytes, with an exception of last byte. A null
    /// terminated `symbol` may avoid a string allocation in some cases.
    ///
    /// Symbol is interpreted as-is; no mangling is done. This means that symbols like `x::y` are
    /// most likely invalid.
    ///
    /// ## Unsafety
    ///
    /// Pointer to a value of arbitrary type is returned. Using a value with wrong type is
    /// undefined.
    pub unsafe fn get<T>(&self, symbol: &[u8]) -> ::Result<Symbol<T>> {
        ensure_compatible_types::<T, winapi::FARPROC>();
        let symbol = try!(cstr_cow_from_bytes(symbol));
        with_get_last_error(|| {
            let symbol = kernel32::GetProcAddress(self.0, symbol.as_ptr());
            if symbol.is_null() {
                None
            } else {
                Some(Symbol {
                    pointer: symbol,
                    pd: marker::PhantomData
                })
            }
        }).map_err(|e| e.unwrap_or_else(||
            panic!("GetProcAddress failed but GetLastError did not report the error")
        ))
    }

    /// Get a pointer to function or static variable by ordinal number.
    ///
    /// ## Unsafety
    ///
    /// Pointer to a value of arbitrary type is returned. Using a value with wrong type is
    /// undefined.
    pub unsafe fn get_ordinal<T>(&self, ordinal: winapi::WORD) -> ::Result<Symbol<T>> {
        ensure_compatible_types::<T, winapi::FARPROC>();
        with_get_last_error(|| {
            let ordinal = ordinal as usize as *mut _;
            let symbol = kernel32::GetProcAddress(self.0, ordinal);
            if symbol.is_null() {
                None
            } else {
                Some(Symbol {
                    pointer: symbol,
                    pd: marker::PhantomData
                })
            }
        }).map_err(|e| e.unwrap_or_else(||
            panic!("GetProcAddress failed but GetLastError did not report the error")
        ))
    }
}

impl Drop for Library {
    fn drop(&mut self) {
        with_get_last_error(|| {
            if unsafe { kernel32::FreeLibrary(self.0) == 0 } {
                None
            } else {
                Some(())
            }
        }).unwrap()
    }
}

impl fmt::Debug for Library {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        unsafe {
            let mut buf: [winapi::WCHAR; 1024] = mem::uninitialized();
            let len = kernel32::GetModuleFileNameW(self.0,
                                                   (&mut buf[..]).as_mut_ptr(), 1024) as usize;
            if len == 0 {
                f.write_str(&format!("Library@{:p}", self.0))
            } else {
                let string: OsString = OsString::from_wide(&buf[..len]);
                f.write_str(&format!("Library@{:p} from {:?}", self.0, string))
            }
        }
    }
}

/// Symbol from a library.
///
/// A major difference compared to the cross-platform `Symbol` is that this does not ensure the
/// `Symbol` does not outlive `Library` it comes from.
pub struct Symbol<T> {
    pointer: winapi::FARPROC,
    pd: marker::PhantomData<T>
}

unsafe impl<T: Send> Send for Symbol<T> {}
unsafe impl<T: Sync> Sync for Symbol<T> {}

impl<T> Clone for Symbol<T> {
    fn clone(&self) -> Symbol<T> {
        Symbol { ..*self }
    }
}

impl<T> ::std::ops::Deref for Symbol<T> {
    type Target = T;
    fn deref(&self) -> &T {
        unsafe {
            // Additional reference level for a dereference on `deref` return value.
            mem::transmute(&self.pointer)
        }
    }
}

impl<T> fmt::Debug for Symbol<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(&format!("Symbol@{:p}", self.pointer))
    }
}

static USE_ERRORMODE: AtomicBool = ATOMIC_BOOL_INIT;
struct ErrorModeGuard(winapi::DWORD);

impl ErrorModeGuard {
    fn new() -> ErrorModeGuard {
        let mut ret = ErrorModeGuard(0);

        if !USE_ERRORMODE.load(Ordering::Acquire) {
            if unsafe { kernel32::SetThreadErrorMode(1, &mut ret.0) == 0
                        && kernel32::GetLastError() == winapi::ERROR_CALL_NOT_IMPLEMENTED } {
                USE_ERRORMODE.store(true, Ordering::Release);
            } else {
                return ret;
            }
        }
        ret.0 = unsafe { kernel32::SetErrorMode(1) };
        ret
    }
}

impl Drop for ErrorModeGuard {
    fn drop(&mut self) {
        unsafe {
            if !USE_ERRORMODE.load(Ordering::Relaxed) {
                kernel32::SetThreadErrorMode(self.0, ptr::null_mut());
            } else {
                kernel32::SetErrorMode(self.0);
            }
        }
    }
}

fn with_get_last_error<T, F>(closure: F) -> Result<T, Option<io::Error>>
where F: FnOnce() -> Option<T> {
    closure().ok_or_else(|| {
        let error = unsafe { kernel32::GetLastError() };
        if error == 0 {
            None
        } else {
            Some(io::Error::from_raw_os_error(error as i32))
        }
    })
}

#[test]
fn works_getlasterror() {
    let lib = Library::new("kernel32.dll").unwrap();
    let gle: Symbol<unsafe extern "system" fn() -> winapi::DWORD> = unsafe {
        lib.get(b"GetLastError").unwrap()
    };
    unsafe {
        kernel32::SetLastError(42);
        assert_eq!(kernel32::GetLastError(), gle())
    }
}

#[test]
fn works_getlasterror0() {
    let lib = Library::new("kernel32.dll").unwrap();
    let gle: Symbol<unsafe extern "system" fn() -> winapi::DWORD> = unsafe {
        lib.get(b"GetLastError\0").unwrap()
    };
    unsafe {
        kernel32::SetLastError(42);
        assert_eq!(kernel32::GetLastError(), gle())
    }
}
<commit_msg>Fix error-case handling for SetThreadErrorMode<commit_after>extern crate winapi;
extern crate kernel32;

use util::{ensure_compatible_types, cstr_cow_from_bytes};

use std::ffi::{OsStr, OsString};
use std::{fmt, io, marker, mem, ptr};
use std::os::windows::ffi::{OsStrExt, OsStringExt};
use std::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};

/// A platform-specific equivalent of the cross-platform `Library`.
pub struct Library(winapi::HMODULE);

unsafe impl Send for Library {}
// Now, this is sort-of-tricky. MSDN documentation does not really make any claims as to safety of
// the Win32 APIs. Sadly, whomever I asked, even current and former Microsoft employees, couldn’t
// say for sure, whether the Win32 APIs used to implement `Library` are thread-safe or not.
//
// My investigation ended up with a question about thread-safety properties of the API involved
// being sent to an internal (to MS) general question mailing-list. The conclusion of the mail is
// as such:
//
// * Nobody inside MS (at least out of all the people who have seen the question) knows for
//   sure either;
// * However, the general consensus between MS developers is that one can rely on the API being
//   thread-safe. In case it is not thread-safe it should be considered a bug on the Windows
//   part. (NB: bugs filled at https://connect.microsoft.com/ against Windows Server)
unsafe impl Sync for Library {}

impl Library {
    /// Find and load a shared library (module).
    ///
    /// Corresponds to `LoadLibraryW(filename)`.
    #[inline]
    pub fn new<P: AsRef<OsStr>>(filename: P) -> ::Result<Library> {
        let wide_filename: Vec<u16> = filename.as_ref().encode_wide().chain(Some(0)).collect();
        let _guard = ErrorModeGuard::new();

        let ret = with_get_last_error(|| {
            // Make sure no winapi calls as a result of drop happen inside this closure, because
            // otherwise that might change the return value of the GetLastError.
            let handle = unsafe { kernel32::LoadLibraryW(wide_filename.as_ptr()) };
            if handle.is_null()  {
                None
            } else {
                Some(Library(handle))
            }
        }).map_err(|e| e.unwrap_or_else(||
            panic!("LoadLibraryW failed but GetLastError did not report the error")
        ));

        drop(wide_filename); // Drop wide_filename here to ensure it doesn’t get moved and dropped
                             // inside the closure by mistake. See comment inside the closure.
        ret
    }

    /// Get a pointer to function or static variable by symbol name.
    ///
    /// The `symbol` may not contain any null bytes, with an exception of last byte. A null
    /// terminated `symbol` may avoid a string allocation in some cases.
    ///
    /// Symbol is interpreted as-is; no mangling is done. This means that symbols like `x::y` are
    /// most likely invalid.
    ///
    /// ## Unsafety
    ///
    /// Pointer to a value of arbitrary type is returned. Using a value with wrong type is
    /// undefined.
    pub unsafe fn get<T>(&self, symbol: &[u8]) -> ::Result<Symbol<T>> {
        ensure_compatible_types::<T, winapi::FARPROC>();
        let symbol = try!(cstr_cow_from_bytes(symbol));
        with_get_last_error(|| {
            let symbol = kernel32::GetProcAddress(self.0, symbol.as_ptr());
            if symbol.is_null() {
                None
            } else {
                Some(Symbol {
                    pointer: symbol,
                    pd: marker::PhantomData
                })
            }
        }).map_err(|e| e.unwrap_or_else(||
            panic!("GetProcAddress failed but GetLastError did not report the error")
        ))
    }

    /// Get a pointer to function or static variable by ordinal number.
    ///
    /// ## Unsafety
    ///
    /// Pointer to a value of arbitrary type is returned. Using a value with wrong type is
    /// undefined.
    pub unsafe fn get_ordinal<T>(&self, ordinal: winapi::WORD) -> ::Result<Symbol<T>> {
        ensure_compatible_types::<T, winapi::FARPROC>();
        with_get_last_error(|| {
            let ordinal = ordinal as usize as *mut _;
            let symbol = kernel32::GetProcAddress(self.0, ordinal);
            if symbol.is_null() {
                None
            } else {
                Some(Symbol {
                    pointer: symbol,
                    pd: marker::PhantomData
                })
            }
        }).map_err(|e| e.unwrap_or_else(||
            panic!("GetProcAddress failed but GetLastError did not report the error")
        ))
    }
}

impl Drop for Library {
    fn drop(&mut self) {
        with_get_last_error(|| {
            if unsafe { kernel32::FreeLibrary(self.0) == 0 } {
                None
            } else {
                Some(())
            }
        }).unwrap()
    }
}

impl fmt::Debug for Library {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        unsafe {
            let mut buf: [winapi::WCHAR; 1024] = mem::uninitialized();
            let len = kernel32::GetModuleFileNameW(self.0,
                                                   (&mut buf[..]).as_mut_ptr(), 1024) as usize;
            if len == 0 {
                f.write_str(&format!("Library@{:p}", self.0))
            } else {
                let string: OsString = OsString::from_wide(&buf[..len]);
                f.write_str(&format!("Library@{:p} from {:?}", self.0, string))
            }
        }
    }
}

/// Symbol from a library.
///
/// A major difference compared to the cross-platform `Symbol` is that this does not ensure the
/// `Symbol` does not outlive `Library` it comes from.
pub struct Symbol<T> {
    pointer: winapi::FARPROC,
    pd: marker::PhantomData<T>
}

unsafe impl<T: Send> Send for Symbol<T> {}
unsafe impl<T: Sync> Sync for Symbol<T> {}

impl<T> Clone for Symbol<T> {
    fn clone(&self) -> Symbol<T> {
        Symbol { ..*self }
    }
}

impl<T> ::std::ops::Deref for Symbol<T> {
    type Target = T;
    fn deref(&self) -> &T {
        unsafe {
            // Additional reference level for a dereference on `deref` return value.
            mem::transmute(&self.pointer)
        }
    }
}

impl<T> fmt::Debug for Symbol<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(&format!("Symbol@{:p}", self.pointer))
    }
}

static USE_ERRORMODE: AtomicBool = ATOMIC_BOOL_INIT;
struct ErrorModeGuard(winapi::DWORD);

impl ErrorModeGuard {
    fn new() -> Option<ErrorModeGuard> {
        const SEM_FAILCE: winapi::DWORD = 1;
        if !USE_ERRORMODE.load(Ordering::Acquire) {
            let result = unsafe {
                let previous_mode = 0;
                let success = kernel32::SetThreadErrorMode(SEM_FAILCE, &mut previous_mode) != 0;
                if !success && kernel32::GetLastError() == winapi::ERROR_CALL_NOT_IMPLEMENTED {
                    USE_ERRORMODE.store(true, Ordering::Release);
                } else if !success {
                    // SetThreadErrorMode failed with some other error? How in the world is it
                    // possible for what is essentially a simple variable swap to fail?
                    // For now we just ignore the error -- the worst that can happen here is
                    // the previous mode staying on and user seeing a dialog error on older Windows
                    // machines.
                    return None;
                } else {
                    return Some(ErrorModeGuard(previous_mode));
                }
            };
        }
        Some(ErrorModeGuard(unsafe { kernel32::SetErrorMode(SEM_FAILCE) }))
    }
}

impl Drop for ErrorModeGuard {
    fn drop(&mut self) {
        unsafe {
            if !USE_ERRORMODE.load(Ordering::Relaxed) {
                kernel32::SetThreadErrorMode(self.0, ptr::null_mut());
            } else {
                kernel32::SetErrorMode(self.0);
            }
        }
    }
}

fn with_get_last_error<T, F>(closure: F) -> Result<T, Option<io::Error>>
where F: FnOnce() -> Option<T> {
    closure().ok_or_else(|| {
        let error = unsafe { kernel32::GetLastError() };
        if error == 0 {
            None
        } else {
            Some(io::Error::from_raw_os_error(error as i32))
        }
    })
}

#[test]
fn works_getlasterror() {
    let lib = Library::new("kernel32.dll").unwrap();
    let gle: Symbol<unsafe extern "system" fn() -> winapi::DWORD> = unsafe {
        lib.get(b"GetLastError").unwrap()
    };
    unsafe {
        kernel32::SetLastError(42);
        assert_eq!(kernel32::GetLastError(), gle())
    }
}

#[test]
fn works_getlasterror0() {
    let lib = Library::new("kernel32.dll").unwrap();
    let gle: Symbol<unsafe extern "system" fn() -> winapi::DWORD> = unsafe {
        lib.get(b"GetLastError\0").unwrap()
    };
    unsafe {
        kernel32::SetLastError(42);
        assert_eq!(kernel32::GetLastError(), gle())
    }
}
<|endoftext|>


<commit_before>extern crate serde;
extern crate serde_json;
extern crate serde_roundtrip;

use serde_json::{to_string, from_str};
use serde_roundtrip::RoundTrip;

type Source = &'static [&'static str];
type Target = Vec<String>;
const SOURCE: Source = &["hello","world"];

#[test]
fn test_round_trip() {
    let via_json: Target = from_str(&*to_string(SOURCE).unwrap()).unwrap();
    let via_round_trip: Target = SOURCE.round_trip();
    assert_eq!(via_json, via_round_trip);
}
<commit_msg>Added more tests.<commit_after>extern crate serde;
extern crate serde_json;
extern crate serde_roundtrip;

use serde_json::{to_string, from_str};
use serde_roundtrip::RoundTrip;

use std::borrow::Cow;
use std::collections::HashMap;
use std::iter::FromIterator;
use std::net::IpAddr;
use std::rc::Rc;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;

#[test]
fn test_round_trip() {
    type Target = (
        (usize,),
        Vec<IpAddr>,
        Vec<String>,
        Box<Rc<Arc<bool>>>,
        Duration,
        HashMap<String, usize>,
        Cow<'static, str>,
    );
    let source = (
        (37,),
        vec![IpAddr::from_str("127.0.0.1").unwrap(), IpAddr::from_str("2001:0db8:85a3:0000:0000:8a2e:0370:7334").unwrap()],
        &["hello","world"][..],
        true,
        Arc::new(Rc::new(Box::new(Duration::new(1000,0)))),
        HashMap::from_iter(vec![ ("a",1) ]),
        Cow::Borrowed("x"),
    );

    let via_json: Target = from_str(&*to_string(&source).unwrap()).unwrap();
    let via_round_trip: Target = source.round_trip();
    assert_eq!(via_json, via_round_trip);
}
<|endoftext|>


<commit_before>extern crate hyper;
extern crate requests;
#[cfg(feature = "with_json")]
extern crate json;

// use requests::Codes;
use requests::{delete, get, head, post, put};
use requests::{Codes, Request, Response, StatusCode};
use requests::ToJson;

fn assert_response_is_ok(response: &Response, url: &str) {
    assert_eq!(response.url(), url);
    assert_eq!(response.status_code(), StatusCode::Ok);
    assert_eq!(response.reason(), "OK");
}

#[test]
fn simple_get() {
    const URL: &'static str = "http://httpbin.org/get";
    let res = get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    let data = res.json().unwrap();
    println!("{:?}", data);
    assert_eq!(data["url"], URL);
    assert_eq!(data["headers"]["Host"], "httpbin.org");
    assert_eq!(data["headers"]["User-Agent"],
               concat!("requests-rs/", env!("CARGO_PKG_VERSION")));
}

#[test]
fn simple_get_string_url_https() {
    const URL: &'static str = "https://httpbin.org/get";
    let res = get(String::from(URL)).unwrap();
    assert_response_is_ok(&res, URL);
    let res = get(&String::from(URL)).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn simple_get_string_url() {
    const URL: &'static str = "http://httpbin.org/get";
    let res = get(String::from(URL)).unwrap();
    assert_response_is_ok(&res, URL);
    let res = get(&String::from(URL)).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn simple_post() {
    const URL: &'static str = "http://httpbin.org/post";
    let res = post(URL).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn simple_put() {
    const URL: &'static str = "http://httpbin.org/put";
    let res = put(URL).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn simple_head() {
    const URL: &'static str = "http://httpbin.org/get";
    let res = head(URL).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn simple_delete() {
    const URL: &'static str = "http://httpbin.org/delete";
    let res = delete(URL).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn user_agent() {
    let useragent = concat!("{\n  \"user-agent\": \"requests-rs/",
                            env!("CARGO_PKG_VERSION"),
                            "\"\n}\n");
    const URL: &'static str = "http://httpbin.org/user-agent";
    let res = get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    assert_eq!(res.text(), Some(useragent));
}

#[test]
fn custom_user_agent() {
    const UA: &'static str = concat!("requests-rs-tests/", env!("CARGO_PKG_VERSION"));
    const URL: &'static str = "http://httpbin.org/user-agent";
    let mut request = Request::new();
    request.user_agent(UA);
    let res = request.get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    assert!(res.is_json());

    let ua = res.json().unwrap();
    assert_eq!(ua["user-agent"], UA);
}

#[test]
fn user_agent_json() {

    const URL: &'static str = "http://httpbin.org/user-agent";
    let res = get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    assert!(res.is_json());

    let ua = res.json().unwrap();
    assert_eq!(ua["user-agent"],
               concat!("requests-rs/", env!("CARGO_PKG_VERSION")));
}

#[test]
fn content() {
    const URL: &'static str = "http://httpbin.org/headers";
    let content = concat!("{\n  \"headers\": {\n    \"Host\": \"httpbin.org\",",
                          " \n    \"User-Agent\": \"requests-rs/",
                          env!("CARGO_PKG_VERSION"),
                          "\"\n  }\n}\n");

    let res = get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    assert_eq!(res.content(), &content.as_bytes());
}

#[test]
fn headers() {
    use hyper::header::UserAgent;

    const URL: &'static str = "http://httpbin.org/response-headers?User-Agent=requests-rs-test";
    let res = get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    println!("{:?}", res.text());
    println!("{:?}", res.headers());
    println!("{:?}", res.headers().get::<UserAgent>());
    assert_eq!(res.headers().get::<UserAgent>().unwrap(),
               &UserAgent("requests-rs-test".to_owned()));
}

#[cfg(feature = "with_json")]
#[test]
fn accept_json() {
    const URL: &'static str = "http://httpbin.org/headers";
    let res = Request::json().get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    assert!(res.is_json());
    let data = res.json().unwrap();
    println!("{:?}", data);
    assert_eq!(data["headers"]["Accept"], "application/json");
}

macro_rules! status_code_test {
    ($($name:ident: $numeric:expr,)+) => {
        $(#[test]
        fn $name() {
            let res = get(&format!("http://httpbin.org/status/{}", $numeric)).unwrap();
            println!("{}", res.text().unwrap());
            assert_eq!(res.status_code(), Codes::from_u16($numeric));
        })+
    }
}

status_code_test! {
    status_code_100: 100,
    status_code_101: 101,
    status_code_102: 102,
    status_code_200: 200,
    status_code_201: 201,
    status_code_202: 202,
    status_code_203: 203,
    status_code_204: 204,
    status_code_205: 205,
    status_code_206: 206,
    status_code_207: 207,
    status_code_208: 208,
    status_code_226: 226,
    status_code_300: 300,
    // status_code_301: 301,
    // status_code_302: 302,
    // status_code_303: 303,
    status_code_304: 304,
    // status_code_305: 305,
    // status_code_307: 307,
    status_code_308: 308,
    status_code_400: 400,
    status_code_401: 401,
    status_code_402: 402,
    status_code_403: 403,
    status_code_404: 404,
    status_code_405: 405,
    status_code_406: 406,
    status_code_407: 407,
    status_code_408: 408,
    status_code_409: 409,
    status_code_410: 410,
    status_code_411: 411,
    status_code_412: 412,
    status_code_413: 413,
    status_code_414: 414,
    status_code_415: 415,
    status_code_416: 416,
    status_code_417: 417,
    status_code_418: 418,
    status_code_419: 419,
    status_code_420: 420,
    status_code_421: 421,
    status_code_422: 422,
    status_code_423: 423,
    status_code_424: 424,
    status_code_426: 426,
    status_code_428: 428,
    status_code_429: 429,
    status_code_431: 431,
    status_code_451: 451,
    status_code_500: 500,
    status_code_501: 501,
    status_code_502: 502,
    status_code_503: 503,
    status_code_504: 504,
    status_code_505: 505,
    status_code_506: 506,
    status_code_507: 507,
    status_code_508: 508,
    status_code_510: 510,
    status_code_511: 511,
}
<commit_msg>Update tests<commit_after>extern crate hyper;
extern crate requests;
#[cfg(feature = "with_json")]
extern crate json;

// use requests::Codes;
use requests::{delete, get, head, post, put};
use requests::{Codes, Request, Response, StatusCode};
use requests::ToJson;

fn assert_response_is_ok(response: &Response, url: &str) {
    assert_eq!(response.url(), url);
    assert_eq!(response.status_code(), StatusCode::Ok);
    assert_eq!(response.reason(), "OK");
}

#[test]
fn simple_get() {
    const URL: &'static str = "http://httpbin.org/get";
    let res = get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    let data = res.json().unwrap();
    println!("{:?}", data);
    assert_eq!(data["url"], URL);
    assert_eq!(data["headers"]["Host"], "httpbin.org");
    assert_eq!(data["headers"]["User-Agent"],
               concat!("requests-rs/", env!("CARGO_PKG_VERSION")));
}

#[test]
fn simple_get_string_url_https() {
    const URL: &'static str = "https://httpbin.org/get";
    let res = get(String::from(URL)).unwrap();
    assert_response_is_ok(&res, URL);
    let res = get(&String::from(URL)).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn simple_get_string_url() {
    const URL: &'static str = "http://httpbin.org/get";
    let res = get(String::from(URL)).unwrap();
    assert_response_is_ok(&res, URL);
    let res = get(&String::from(URL)).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn simple_post() {
    const URL: &'static str = "http://httpbin.org/post";
    let res = post(URL).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn simple_put() {
    const URL: &'static str = "http://httpbin.org/put";
    let res = put(URL).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn simple_head() {
    const URL: &'static str = "http://httpbin.org/get";
    let res = head(URL).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn simple_delete() {
    const URL: &'static str = "http://httpbin.org/delete";
    let res = delete(URL).unwrap();
    assert_response_is_ok(&res, URL);
}

#[test]
fn user_agent() {
    let useragent = concat!("{\n  \"user-agent\": \"requests-rs/",
                            env!("CARGO_PKG_VERSION"),
                            "\"\n}\n");
    const URL: &'static str = "http://httpbin.org/user-agent";
    let res = get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    assert_eq!(res.text(), Some(useragent));
}

#[test]
fn custom_user_agent() {
    const UA: &'static str = concat!("requests-rs-tests/", env!("CARGO_PKG_VERSION"));
    const URL: &'static str = "http://httpbin.org/user-agent";
    let mut request = Request::new();
    request.user_agent(UA);
    let res = request.get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    assert!(res.is_json());

    let ua = res.json().unwrap();
    assert_eq!(ua["user-agent"], UA);
}

#[test]
fn user_agent_json() {

    const URL: &'static str = "http://httpbin.org/user-agent";
    let res = get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    assert!(res.is_json());

    let ua = res.json().unwrap();
    assert_eq!(ua["user-agent"],
               concat!("requests-rs/", env!("CARGO_PKG_VERSION")));
}

#[test]
fn content() {
    const URL: &'static str = "http://httpbin.org/headers";
    let content = concat!("{\n",
                          "  \"headers\": {\n",
                          "    \"Connection\": \"close\", \n",
                          "    \"Host\": \"httpbin.org\", \n",
                          "    \"User-Agent\": \"requests-rs/", env!("CARGO_PKG_VERSION"), "\"\n",
                          "  }\n",
                          "}\n");

    let res = get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    println!("{}", res.text().unwrap());
    println!("{}", content);
    assert_eq!(res.content(), &content.as_bytes());
}

#[test]
fn headers() {
    use hyper::header::UserAgent;

    const URL: &'static str = "http://httpbin.org/response-headers?User-Agent=requests-rs-test";
    let res = get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    println!("{:?}", res.text());
    println!("{:?}", res.headers());
    println!("{:?}", res.headers().get::<UserAgent>());
    assert_eq!(res.headers().get::<UserAgent>().unwrap(),
               &UserAgent("requests-rs-test".to_owned()));
}

#[cfg(feature = "with_json")]
#[test]
fn accept_json() {
    const URL: &'static str = "http://httpbin.org/headers";
    let res = Request::json().get(URL).unwrap();
    assert_response_is_ok(&res, URL);
    assert!(res.is_json());
    let data = res.json().unwrap();
    println!("{:?}", data);
    assert_eq!(data["headers"]["Accept"], "application/json");
}

macro_rules! status_code_test {
    ($($name:ident: $numeric:expr,)+) => {
        $(#[test]
        fn $name() {
            let res = get(&format!("http://httpbin.org/status/{}", $numeric)).unwrap();
            println!("{}", res.text().unwrap());
            assert_eq!(res.status_code(), Codes::from_u16($numeric));
        })+
    }
}

status_code_test! {
    status_code_100: 100,
    status_code_101: 101,
    status_code_102: 102,
    status_code_200: 200,
    status_code_201: 201,
    status_code_202: 202,
    status_code_203: 203,
    status_code_204: 204,
    status_code_205: 205,
    status_code_206: 206,
    status_code_207: 207,
    status_code_208: 208,
    status_code_226: 226,
    status_code_300: 300,
    // status_code_301: 301,
    // status_code_302: 302,
    // status_code_303: 303,
    status_code_304: 304,
    // status_code_305: 305,
    // status_code_307: 307,
    status_code_308: 308,
    status_code_400: 400,
    status_code_401: 401,
    status_code_402: 402,
    status_code_403: 403,
    status_code_404: 404,
    status_code_405: 405,
    status_code_406: 406,
    status_code_407: 407,
    status_code_408: 408,
    status_code_409: 409,
    status_code_410: 410,
    status_code_411: 411,
    status_code_412: 412,
    status_code_413: 413,
    status_code_414: 414,
    status_code_415: 415,
    status_code_416: 416,
    status_code_417: 417,
    status_code_418: 418,
    status_code_419: 419,
    status_code_420: 420,
    status_code_421: 421,
    status_code_422: 422,
    status_code_423: 423,
    status_code_424: 424,
    status_code_426: 426,
    status_code_428: 428,
    status_code_429: 429,
    status_code_431: 431,
    status_code_451: 451,
    status_code_500: 500,
    status_code_501: 501,
    status_code_502: 502,
    status_code_503: 503,
    status_code_504: 504,
    status_code_505: 505,
    status_code_506: 506,
    status_code_507: 507,
    status_code_508: 508,
    status_code_510: 510,
    status_code_511: 511,
}
<|endoftext|>


<commit_before>extern crate presort;

use presort::PermutedVec;

struct TreeNode<T> {
    data: T,
    children: Vec<Tree<T>>,
}

type Tree<T> = Box<TreeNode<T>>;

fn tree<T>(data: T, children: Vec<Tree<T>>) -> Tree<T> {
    Box::new(TreeNode { data: data, children: children })
}

fn dump<T: Clone+Ord>(tree: &Tree<T>, vec: &mut PermutedVec<T>) {
    vec.push(tree.data.clone());
    for ref kid in &tree.children { dump(kid, vec); }
}

#[test]
fn test_tree() {
    let tree = tree(37, vec![]);
    let mut vec = PermutedVec::new();
    dump(&tree, &mut vec);
    assert_eq!(vec.sorted_iter().collect::<Vec<&usize>>(), vec![&37]);
}
<commit_msg>added update fn for tree<commit_after>extern crate presort;

use presort::PermutedVec;

struct TreeNode<T> {
	//dirty flag for this node, not children
	dirty: bool,
    data: T,
    children: Vec<Tree<T>>,
}

type Tree<T> = Box<TreeNode<T>>;

fn tree<T>(data: T, children: Vec<Tree<T>>) -> Tree<T> {
    Box::new(TreeNode { data: data, dirty: true, children: children })
}

fn dump<T: Clone+Ord>(tree: &mut Tree<T>, vec: &mut PermutedVec<T>) {
    vec.push(tree.data.clone());
    tree.dirty = false;
    for mut kid in &mut tree.children { dump(&mut kid, vec); }
}

// incrementally updates the vec, returning the number of nodes considered
fn update<T: Clone>(tree: &mut Tree<T>, start_index: usize, vec: &mut PermutedVec<T>) -> usize {
	let mut i = start_index;
	if tree.dirty { vec.set(i,tree.data.clone()); tree.dirty = false; }
	i += 1;
	for mut kid in &mut tree.children {
		i += update(&mut kid, i, vec);
	}
	i - start_index
}

#[test]
fn test_tree() {
    let mut main_tree = tree(37, vec![]);
    let mut vec = PermutedVec::new();
    dump(&mut main_tree, &mut vec);
    assert_eq!(vec.sorted_iter().collect::<Vec<&usize>>(), vec![&37]);
}

#[test]
fn test_update() {
	let c1 = tree(42, vec![]);
	let c2 = tree(20, vec![]);
	let c3 = tree(63, vec![]);
	let mut main_tree = tree(37, vec![c1,c2,c3]);
	let mut vec = PermutedVec::new();
	dump(&mut main_tree, &mut vec);
    assert_eq!(vec.sorted_iter().collect::<Vec<&usize>>(), vec![&20,&37,&42,&63]);

    main_tree.children[1].data = 25;
    main_tree.children[1].dirty = true;
    update(&mut main_tree, 0, &mut vec);
    assert_eq!(vec.sorted_iter().collect::<Vec<&usize>>(), vec![&25,&37,&42,&63]);
}
<|endoftext|>


<commit_before>extern crate bytes;
extern crate futures;
extern crate futures_cpupool;
extern crate itertools;
extern crate serde_json;
extern crate tokio_core;
extern crate tokio_io;
extern crate tokio_proto;
extern crate tokio_service;

use std;
use std::collections::HashMap;
use std::error::Error as StdError;
use std::os::unix::fs::DirBuilderExt;
use std::sync::{Arc, Mutex};

use self::futures::{Future, Stream, Sink};
use self::futures::sync::mpsc;
use self::itertools::Itertools;
use self::tokio_io::{AsyncRead};
use self::tokio_service::Service as TokioService;

use super::*;
use self::agent::{AgentStorage, DurableAgentStorage};
use self::api::{Authenticated, Request, Response};
use self::token::{DurableTokenStorage, TOKEN_NAME_LEN, TokenStorage};

pub struct ServiceCodec;

impl tokio_io::codec::Decoder for ServiceCodec {
    type Item = Request;
    type Error = std::io::Error;

    fn decode(&mut self, buf: &mut bytes::BytesMut) -> std::io::Result<Option<Self::Item>> {
        if let Some(i) = buf.iter().position(|&b| b == b'\n') {
            // remove the serialized frame from the buffer.
            let line = buf.split_to(i);

            // Also remove the '\n'
            buf.split_to(1);

            // Deserialize JSON and return it in a Frame.
            let maybe_req: Result<Self::Item, serde_json::error::Error> =
                serde_json::from_slice(&line[..]);
            match maybe_req {
                Ok(req) => {
                    trace!("service decode {:?}", req);
                    buf.take();
                    Ok(Some(req))
                }
                Err(e) => {
                    error!("service decode failed: {}", e);
                    Err(std::io::Error::new(std::io::ErrorKind::Other, e.description()))
                }
            }
        } else {
            Ok(None)
        }
    }
}

impl tokio_io::codec::Encoder for ServiceCodec {
    type Item = Response;
    type Error = std::io::Error;

    fn encode(&mut self, msg: Self::Item, buf: &mut bytes::BytesMut) -> std::io::Result<()> {
        match serde_json::to_vec(&msg) {
            Ok(json) => {
                trace!("service encode {:?}", msg);
                buf.extend(&json[..]);
                Ok(())
            }
            Err(e) => {
                error!("service encode {:?} failed: {}", msg, e);
                Err(std::io::Error::new(std::io::ErrorKind::Other, e.description()))
            }
        }?;
        buf.extend(&b"\n"[..]);
        Ok(())
    }
}

type AgentOutboxMap = HashMap<String, mpsc::Sender<Message>>;

struct ServiceState<S: AgentStorage + Send + 'static> {
    storage: S,
    local_outboxes: AgentOutboxMap,
}

impl<S: AgentStorage + Send> ServiceState<S> {
    pub fn new(storage: S) -> ServiceState<S> {
        ServiceState {
            storage: storage,
            local_outboxes: AgentOutboxMap::new(),
        }
    }

    fn has_agent(&self, name: &str) -> bool {
        self.local_outboxes.contains_key(name)
    }

    fn remove(&mut self, name: &str) -> Result<(), Error> {
        self.storage.remove_agent(name)?;
        self.local_outboxes.remove(name);
        Ok(())
    }

    fn add_all_agents(&mut self, svc: &Service<S>) -> Result<(), Error> {
        let agents = self.storage.agents()?;
        for (name, glop) in agents {
            svc.spawn_agent(&name, &glop, self)?;
        }
        Ok(())
    }
}

#[derive(Clone)]
pub struct Service<S: AgentStorage + Send + 'static> {
    state: Arc<Mutex<ServiceState<S>>>,
    handle: tokio_core::reactor::Handle,
    pool: futures_cpupool::CpuPool,
}

impl<S: AgentStorage + Send> Service<S> {
    pub fn new(storage: S, h: &tokio_core::reactor::Handle) -> Result<Service<S>, Error> {
        let svc = Service {
            state: Arc::new(Mutex::new(ServiceState::new(storage))),
            handle: h.clone(),
            pool: futures_cpupool::CpuPool::new_num_cpus(),
        };
        {
            let mut state = &mut svc.state.lock().unwrap();
            state.add_all_agents(&svc)?;
        }
        Ok(svc)
    }

    fn add_agent_contents(&self,
                          name: &str,
                          contents: &str,
                          state: &mut ServiceState<S>)
                          -> Result<(), Error> {
        let glop = grammar::glop(contents).map_err(Error::Parse)?;
        self.add_agent(name, glop, state)
    }

    fn add_agent(&self,
                 name: &str,
                 glop: ast::Glop,
                 state: &mut ServiceState<S>)
                 -> Result<(), Error> {
        self.spawn_agent(name, &glop, state)?;
        state.storage.add_agent(name.to_string(), glop)
    }

    fn spawn_agent(&self,
                   name: &str,
                   glop: &ast::Glop,
                   state: &mut ServiceState<S>)
                   -> Result<(), Error> {
        let runtime_st = state.storage
            .new_state(name,
                       Box::new(SenderOutbox {
                           src: name.to_string(),
                           remote: self.handle.remote().clone(),
                           state: self.state.clone(),
                       }) as Box<runtime::Outbox + Send>)?;
        let (sender, receiver) = mpsc::channel(10);
        let agent = Agent::new(glop, runtime_st, receiver)?;
        state.local_outboxes.insert(name.to_string(), sender);
        self.handle.spawn(self.pool
            .spawn(agent.for_each(|_| Ok(()))
                .or_else(|e| {
                    error!("{}", e);
                    Err(e)
                })
                .then(|_| Ok(()))));
        Ok(())
    }

    fn do_call(&self, req: Authenticated<Request>) -> Result<Response, Error> {
        debug!("request {:?}", &req);
        let res = match req.item {
            Request::Add { contents: ref add_contents, name: ref add_name } => {
                let mut state = self.state.lock().unwrap();
                if state.has_agent(add_name) {
                    return Ok(Response::Error(format!("agent {} already added", add_name)));
                }
                self.add_agent_contents(add_name, add_contents, &mut state)?;
                Response::Add
            }
            Request::Remove { ref name } => {
                let mut state = self.state.lock().unwrap();
                state.remove(name)?;
                Response::Remove
            }
            Request::List => {
                let state = self.state.lock().unwrap();
                Response::List { names: state.local_outboxes.keys().cloned().collect() }
            }
            Request::SendTo(msg) => self.send_to(msg.new_id()),
            Request::Introduce(agent_roles) => {
                let mut result = vec![];
                for ref p in agent_roles.iter().combinations(2) {
                    result.push(self.send_to(Message::new("intro", Obj::new())
                        .src(&p[0].name)
                        .src_role(Some(p[0].role.to_string()))
                        .dst(&p[1].name)));
                    result.push(self.send_to(Message::new("intro", Obj::new())
                        .src(&p[1].name)
                        .src_role(Some(p[1].role.to_string()))
                        .dst(&p[0].name)));
                }
                Response::Introduce(result)
            }
        };
        debug!("response {:?}", res);
        Ok(res)
    }

    fn send_to(&self, msg: Message) -> Response {
        if let Some(sender) = self.state.lock().unwrap().local_outboxes.get(&msg.dst) {
            let resp = Response::SendTo {
                id: msg.id.to_string(),
                src: msg.src.to_string(),
                dst: msg.dst.to_string(),
            };
            let sender = sender.clone();
            self.handle.spawn(sender.send(msg).then(|_| Ok(())));
            resp
        } else {
            Response::Error(format!("agent {} not found", &msg.dst))
        }
    }
}

#[derive(Clone)]
struct SenderOutbox<S: AgentStorage + Send + 'static> {
    src: String,
    remote: tokio_core::reactor::Remote,
    state: Arc<Mutex<ServiceState<S>>>,
}

impl<S: AgentStorage + Send> runtime::Outbox for SenderOutbox<S> {
    fn send_msg(&self, msg: Message) -> Result<(), Error> {
        let state = self.state.lock().unwrap();
        let sender = match state.local_outboxes.get(&msg.dst) {
            Some(s) => s.clone(),
            None => return Err(Error::InvalidArgument(msg.dst.to_string())),
        };
        self.remote.spawn(|_| sender.send(msg).then(|_| Ok(())));
        Ok(())
    }
}

impl<S: AgentStorage + Send> TokioService for Service<S> {
    type Request = Authenticated<Request>;
    type Response = Response;

    type Error = std::io::Error;

    type Future = futures::BoxFuture<Self::Response, Self::Error>;

    fn call(&self, req: Self::Request) -> Self::Future {
        match self.do_call(req) {
            Ok(res) => futures::future::ok(res).boxed(),
            Err(e) => {
                futures::future::ok(Response::Error(format!("agent service error: {}", e))).boxed()
            }
        }
    }
}

pub struct SecureServiceCodec {
    tokens: Box<TokenStorage>,
    codec: Option<crypto::SecretBoxCodec<ServiceCodec>>,
    auth_id: Option<String>,
}

impl SecureServiceCodec {
    #[allow(dead_code)]
    pub fn new(tokens: Box<TokenStorage>) -> SecureServiceCodec {
        SecureServiceCodec {
            tokens: tokens,
            codec: None,
            auth_id: None,
        }
    }

    fn decode_prelude
        (&mut self,
         buf: &mut bytes::BytesMut)
         -> std::io::Result<Option<(String, crypto::SecretBoxCodec<ServiceCodec>)>> {
        if let Some(i) = buf.iter().position(|&b| b == b'\0') {
            // remove the serialized frame from the buffer.
            let line = buf.split_to(i);

            // Also remove the '\0'
            buf.split_to(1);

            trace!("from {:?}", line);
            let id = std::str::from_utf8(&line[..]).map_err(to_ioerror)?;
            trace!("matched prelude, connection from {}", id);
            match self.tokens.token(id) {
                Ok(Some(ref token)) => {
                    Ok(Some((id.to_string(),
                             crypto::SecretBoxCodec::new(ServiceCodec, token.key().clone()))))
                }
                Ok(None) => {
                    Err(std::io::Error::new(std::io::ErrorKind::Other,
                                            format!("missing or invalid token name: {}", id)))
                }
                Err(e) => Err(to_ioerror(e)),
            }
        } else if buf.len() > TOKEN_NAME_LEN+1 {
            Err(std::io::Error::new(std::io::ErrorKind::Other, "missing or invalid prelude"))
        } else {
            Ok(None)
        }
    }
}

impl tokio_io::codec::Decoder for SecureServiceCodec {
    type Item = Authenticated<<ServiceCodec as tokio_io::codec::Decoder>::Item>;
    type Error = std::io::Error;

    fn decode(&mut self, buf: &mut bytes::BytesMut) -> std::io::Result<Option<Self::Item>> {
        match self.decode_prelude(buf) {
            Ok(Some((id, codec))) => {
                self.auth_id = Some(id);
                self.codec = Some(codec);
            }
            Ok(None) => {}
            Err(e) => return Err(e),
        }
        if let (&Some(ref id), &mut Some(ref mut codec)) = (&self.auth_id, &mut self.codec) {
            trace!("decoding encrypted contents");
            match codec.decode(buf) {
                Ok(Some(req)) => {
                    buf.take();
                    Ok(Some(Authenticated {
                        auth_id: id.to_string(),
                        item: req,
                    }))
                }
                Ok(None) => Ok(None),
                Err(e) => Err(e),
            }
        } else {
            Ok(None)
        }
    }
}

impl tokio_io::codec::Encoder for SecureServiceCodec {
    type Item = <ServiceCodec as tokio_io::codec::Encoder>::Item;
    type Error = std::io::Error;

    fn encode(&mut self, msg: Self::Item, buf: &mut bytes::BytesMut) -> std::io::Result<()> {
        match self.codec {
            None => {
                Err(std::io::Error::new(std::io::ErrorKind::Other, "missing or invalid prelude"))
            }
            Some(ref mut codec) => codec.encode(msg, buf),
        }
    }
}

pub struct Server {
    pub addr: std::net::SocketAddr,
    pub tokens_path: String,
    pub agents_path: String,
}

impl Server {
    pub fn new(addr_str: &str, path: &str) -> Result<Server, Error> {
        let addr = addr_str.parse()?;
        Server::new_addr(addr, path)
    }

    pub fn new_addr(addr: std::net::SocketAddr, path: &str) -> Result<Server, Error> {
        std::fs::DirBuilder::new().recursive(true)
            .mode(0o700)
            .create(path)
            .map_err(Error::IO)?;
        Ok(Server {
            addr: addr,
            agents_path: std::path::PathBuf::from(path)
                .join("agents")
                .to_str()
                .unwrap()
                .to_string(),
            tokens_path: std::path::PathBuf::from(path)
                .join("tokens.json")
                .to_str()
                .unwrap()
                .to_string(),
        })
    }

    pub fn run(&self) -> Result<(), Error> {
        let mut core = tokio_core::reactor::Core::new().map_err(Error::IO)?;
        let handle = core.handle();
        let listener = tokio_core::net::TcpListener::bind(&self.addr, &handle).map_err(Error::IO)?;
        info!("server listening on {}", self.addr);
        let connections = listener.incoming();
        let agent_storage = DurableAgentStorage::new(&self.agents_path);
        let service = Service::new(agent_storage, &handle)?;
        let server = connections.for_each(move |(socket, _peer_addr)| {
            let token_storage = DurableTokenStorage::new(&self.tokens_path);
            let (wr, rd) = socket.framed(SecureServiceCodec::new(Box::new(token_storage))).split();
            let service = service.clone();
            let responses = rd.and_then(move |req| service.call(req));
            let responder = wr.send_all(responses)
                .or_else(|e| {
                    error!("{}", e);
                    Err(e)
                })
                .then(|_| Ok(()));
            handle.spawn(responder);
            Ok(())
        });
        core.run(server).map_err(Error::IO)
    }
}
<commit_msg>Rename sender map from agent server's perspective.<commit_after>extern crate bytes;
extern crate futures;
extern crate futures_cpupool;
extern crate itertools;
extern crate serde_json;
extern crate tokio_core;
extern crate tokio_io;
extern crate tokio_proto;
extern crate tokio_service;

use std;
use std::collections::HashMap;
use std::error::Error as StdError;
use std::os::unix::fs::DirBuilderExt;
use std::sync::{Arc, Mutex};

use self::futures::{Future, Stream, Sink};
use self::futures::sync::mpsc;
use self::itertools::Itertools;
use self::tokio_io::{AsyncRead};
use self::tokio_service::Service as TokioService;

use super::*;
use self::agent::{AgentStorage, DurableAgentStorage};
use self::api::{Authenticated, Request, Response};
use self::token::{DurableTokenStorage, TOKEN_NAME_LEN, TokenStorage};

pub struct ServiceCodec;

impl tokio_io::codec::Decoder for ServiceCodec {
    type Item = Request;
    type Error = std::io::Error;

    fn decode(&mut self, buf: &mut bytes::BytesMut) -> std::io::Result<Option<Self::Item>> {
        if let Some(i) = buf.iter().position(|&b| b == b'\n') {
            // remove the serialized frame from the buffer.
            let line = buf.split_to(i);

            // Also remove the '\n'
            buf.split_to(1);

            // Deserialize JSON and return it in a Frame.
            let maybe_req: Result<Self::Item, serde_json::error::Error> =
                serde_json::from_slice(&line[..]);
            match maybe_req {
                Ok(req) => {
                    trace!("service decode {:?}", req);
                    buf.take();
                    Ok(Some(req))
                }
                Err(e) => {
                    error!("service decode failed: {}", e);
                    Err(std::io::Error::new(std::io::ErrorKind::Other, e.description()))
                }
            }
        } else {
            Ok(None)
        }
    }
}

impl tokio_io::codec::Encoder for ServiceCodec {
    type Item = Response;
    type Error = std::io::Error;

    fn encode(&mut self, msg: Self::Item, buf: &mut bytes::BytesMut) -> std::io::Result<()> {
        match serde_json::to_vec(&msg) {
            Ok(json) => {
                trace!("service encode {:?}", msg);
                buf.extend(&json[..]);
                Ok(())
            }
            Err(e) => {
                error!("service encode {:?} failed: {}", msg, e);
                Err(std::io::Error::new(std::io::ErrorKind::Other, e.description()))
            }
        }?;
        buf.extend(&b"\n"[..]);
        Ok(())
    }
}

/// Associate local agent name with a sender that sends messages to that agent.
type AgentSenderMap = HashMap<String, mpsc::Sender<Message>>;

struct ServiceState<S: AgentStorage + Send + 'static> {
    storage: S,
    local_senders: AgentSenderMap,
}

impl<S: AgentStorage + Send> ServiceState<S> {
    pub fn new(storage: S) -> ServiceState<S> {
        ServiceState {
            storage: storage,
            local_senders: AgentSenderMap::new(),
        }
    }

    fn has_agent(&self, name: &str) -> bool {
        self.local_senders.contains_key(name)
    }

    fn remove(&mut self, name: &str) -> Result<(), Error> {
        self.storage.remove_agent(name)?;
        self.local_senders.remove(name);
        Ok(())
    }

    fn add_all_agents(&mut self, svc: &Service<S>) -> Result<(), Error> {
        let agents = self.storage.agents()?;
        for (name, glop) in agents {
            svc.spawn_agent(&name, &glop, self)?;
        }
        Ok(())
    }
}

#[derive(Clone)]
pub struct Service<S: AgentStorage + Send + 'static> {
    state: Arc<Mutex<ServiceState<S>>>,
    handle: tokio_core::reactor::Handle,
    pool: futures_cpupool::CpuPool,
}

impl<S: AgentStorage + Send> Service<S> {
    pub fn new(storage: S, h: &tokio_core::reactor::Handle) -> Result<Service<S>, Error> {
        let svc = Service {
            state: Arc::new(Mutex::new(ServiceState::new(storage))),
            handle: h.clone(),
            pool: futures_cpupool::CpuPool::new_num_cpus(),
        };
        {
            let mut state = &mut svc.state.lock().unwrap();
            state.add_all_agents(&svc)?;
        }
        Ok(svc)
    }

    fn add_agent_contents(&self,
                          name: &str,
                          contents: &str,
                          state: &mut ServiceState<S>)
                          -> Result<(), Error> {
        let glop = grammar::glop(contents).map_err(Error::Parse)?;
        self.add_agent(name, glop, state)
    }

    fn add_agent(&self,
                 name: &str,
                 glop: ast::Glop,
                 state: &mut ServiceState<S>)
                 -> Result<(), Error> {
        self.spawn_agent(name, &glop, state)?;
        state.storage.add_agent(name.to_string(), glop)
    }

    fn spawn_agent(&self,
                   name: &str,
                   glop: &ast::Glop,
                   state: &mut ServiceState<S>)
                   -> Result<(), Error> {
        let runtime_st = state.storage
            .new_state(name,
                       Box::new(SenderOutbox {
                           src: name.to_string(),
                           remote: self.handle.remote().clone(),
                           state: self.state.clone(),
                       }) as Box<runtime::Outbox + Send>)?;
        let (sender, receiver) = mpsc::channel(10);
        let agent = Agent::new(glop, runtime_st, receiver)?;
        state.local_senders.insert(name.to_string(), sender);
        self.handle.spawn(self.pool
            .spawn(agent.for_each(|_| Ok(()))
                .or_else(|e| {
                    error!("{}", e);
                    Err(e)
                })
                .then(|_| Ok(()))));
        Ok(())
    }

    fn do_call(&self, req: Authenticated<Request>) -> Result<Response, Error> {
        debug!("request {:?}", &req);
        let res = match req.item {
            Request::Add { contents: ref add_contents, name: ref add_name } => {
                let mut state = self.state.lock().unwrap();
                if state.has_agent(add_name) {
                    return Ok(Response::Error(format!("agent {} already added", add_name)));
                }
                self.add_agent_contents(add_name, add_contents, &mut state)?;
                Response::Add
            }
            Request::Remove { ref name } => {
                let mut state = self.state.lock().unwrap();
                state.remove(name)?;
                Response::Remove
            }
            Request::List => {
                let state = self.state.lock().unwrap();
                Response::List { names: state.local_senders.keys().cloned().collect() }
            }
            Request::SendTo(msg) => self.send_to(msg.new_id()),
            Request::Introduce(agent_roles) => {
                let mut result = vec![];
                for ref p in agent_roles.iter().combinations(2) {
                    result.push(self.send_to(Message::new("intro", Obj::new())
                        .src(&p[0].name)
                        .src_role(Some(p[0].role.to_string()))
                        .dst(&p[1].name)));
                    result.push(self.send_to(Message::new("intro", Obj::new())
                        .src(&p[1].name)
                        .src_role(Some(p[1].role.to_string()))
                        .dst(&p[0].name)));
                }
                Response::Introduce(result)
            }
        };
        debug!("response {:?}", res);
        Ok(res)
    }

    fn send_to(&self, msg: Message) -> Response {
        if let Some(sender) = self.state.lock().unwrap().local_senders.get(&msg.dst) {
            let resp = Response::SendTo {
                id: msg.id.to_string(),
                src: msg.src.to_string(),
                dst: msg.dst.to_string(),
            };
            let sender = sender.clone();
            self.handle.spawn(sender.send(msg).then(|_| Ok(())));
            resp
        } else {
            Response::Error(format!("agent {} not found", &msg.dst))
        }
    }
}

#[derive(Clone)]
struct SenderOutbox<S: AgentStorage + Send + 'static> {
    src: String,
    remote: tokio_core::reactor::Remote,
    state: Arc<Mutex<ServiceState<S>>>,
}

impl<S: AgentStorage + Send> runtime::Outbox for SenderOutbox<S> {
    fn send_msg(&self, msg: Message) -> Result<(), Error> {
        let state = self.state.lock().unwrap();
        let sender = match state.local_senders.get(&msg.dst) {
            Some(s) => s.clone(),
            None => return Err(Error::InvalidArgument(msg.dst.to_string())),
        };
        self.remote.spawn(|_| sender.send(msg).then(|_| Ok(())));
        Ok(())
    }
}

impl<S: AgentStorage + Send> TokioService for Service<S> {
    type Request = Authenticated<Request>;
    type Response = Response;

    type Error = std::io::Error;

    type Future = futures::BoxFuture<Self::Response, Self::Error>;

    fn call(&self, req: Self::Request) -> Self::Future {
        match self.do_call(req) {
            Ok(res) => futures::future::ok(res).boxed(),
            Err(e) => {
                futures::future::ok(Response::Error(format!("agent service error: {}", e))).boxed()
            }
        }
    }
}

pub struct SecureServiceCodec {
    tokens: Box<TokenStorage>,
    codec: Option<crypto::SecretBoxCodec<ServiceCodec>>,
    auth_id: Option<String>,
}

impl SecureServiceCodec {
    #[allow(dead_code)]
    pub fn new(tokens: Box<TokenStorage>) -> SecureServiceCodec {
        SecureServiceCodec {
            tokens: tokens,
            codec: None,
            auth_id: None,
        }
    }

    fn decode_prelude
        (&mut self,
         buf: &mut bytes::BytesMut)
         -> std::io::Result<Option<(String, crypto::SecretBoxCodec<ServiceCodec>)>> {
        if let Some(i) = buf.iter().position(|&b| b == b'\0') {
            // remove the serialized frame from the buffer.
            let line = buf.split_to(i);

            // Also remove the '\0'
            buf.split_to(1);

            trace!("from {:?}", line);
            let id = std::str::from_utf8(&line[..]).map_err(to_ioerror)?;
            trace!("matched prelude, connection from {}", id);
            match self.tokens.token(id) {
                Ok(Some(ref token)) => {
                    Ok(Some((id.to_string(),
                             crypto::SecretBoxCodec::new(ServiceCodec, token.key().clone()))))
                }
                Ok(None) => {
                    Err(std::io::Error::new(std::io::ErrorKind::Other,
                                            format!("missing or invalid token name: {}", id)))
                }
                Err(e) => Err(to_ioerror(e)),
            }
        } else if buf.len() > TOKEN_NAME_LEN+1 {
            Err(std::io::Error::new(std::io::ErrorKind::Other, "missing or invalid prelude"))
        } else {
            Ok(None)
        }
    }
}

impl tokio_io::codec::Decoder for SecureServiceCodec {
    type Item = Authenticated<<ServiceCodec as tokio_io::codec::Decoder>::Item>;
    type Error = std::io::Error;

    fn decode(&mut self, buf: &mut bytes::BytesMut) -> std::io::Result<Option<Self::Item>> {
        match self.decode_prelude(buf) {
            Ok(Some((id, codec))) => {
                self.auth_id = Some(id);
                self.codec = Some(codec);
            }
            Ok(None) => {}
            Err(e) => return Err(e),
        }
        if let (&Some(ref id), &mut Some(ref mut codec)) = (&self.auth_id, &mut self.codec) {
            trace!("decoding encrypted contents");
            match codec.decode(buf) {
                Ok(Some(req)) => {
                    buf.take();
                    Ok(Some(Authenticated {
                        auth_id: id.to_string(),
                        item: req,
                    }))
                }
                Ok(None) => Ok(None),
                Err(e) => Err(e),
            }
        } else {
            Ok(None)
        }
    }
}

impl tokio_io::codec::Encoder for SecureServiceCodec {
    type Item = <ServiceCodec as tokio_io::codec::Encoder>::Item;
    type Error = std::io::Error;

    fn encode(&mut self, msg: Self::Item, buf: &mut bytes::BytesMut) -> std::io::Result<()> {
        match self.codec {
            None => {
                Err(std::io::Error::new(std::io::ErrorKind::Other, "missing or invalid prelude"))
            }
            Some(ref mut codec) => codec.encode(msg, buf),
        }
    }
}

pub struct Server {
    pub addr: std::net::SocketAddr,
    pub tokens_path: String,
    pub agents_path: String,
}

impl Server {
    pub fn new(addr_str: &str, path: &str) -> Result<Server, Error> {
        let addr = addr_str.parse()?;
        Server::new_addr(addr, path)
    }

    pub fn new_addr(addr: std::net::SocketAddr, path: &str) -> Result<Server, Error> {
        std::fs::DirBuilder::new().recursive(true)
            .mode(0o700)
            .create(path)
            .map_err(Error::IO)?;
        Ok(Server {
            addr: addr,
            agents_path: std::path::PathBuf::from(path)
                .join("agents")
                .to_str()
                .unwrap()
                .to_string(),
            tokens_path: std::path::PathBuf::from(path)
                .join("tokens.json")
                .to_str()
                .unwrap()
                .to_string(),
        })
    }

    pub fn run(&self) -> Result<(), Error> {
        let mut core = tokio_core::reactor::Core::new().map_err(Error::IO)?;
        let handle = core.handle();
        let listener = tokio_core::net::TcpListener::bind(&self.addr, &handle).map_err(Error::IO)?;
        info!("server listening on {}", self.addr);
        let connections = listener.incoming();
        let agent_storage = DurableAgentStorage::new(&self.agents_path);
        let service = Service::new(agent_storage, &handle)?;
        let server = connections.for_each(move |(socket, _peer_addr)| {
            let token_storage = DurableTokenStorage::new(&self.tokens_path);
            let (wr, rd) = socket.framed(SecureServiceCodec::new(Box::new(token_storage))).split();
            let service = service.clone();
            let responses = rd.and_then(move |req| service.call(req));
            let responder = wr.send_all(responses)
                .or_else(|e| {
                    error!("{}", e);
                    Err(e)
                })
                .then(|_| Ok(()));
            handle.spawn(responder);
            Ok(())
        });
        core.run(server).map_err(Error::IO)
    }
}
<|endoftext|>


<commit_before>
use pass;
use ir;

// FIXME: do not perform strength reduction if the value would overflow

/// An IR strength reduction pass.
pub struct StrengthReduction;

impl pass::Metadata for StrengthReduction
{
    fn id(&self) -> pass::Id { pass::Id(0x242343a1) }
    fn name(&self) -> &'static str { "Strength reduction" }
}

impl pass::Transform<ir::Value> for StrengthReduction
{
    fn run_value(&mut self, value: ir::Value) -> ir::Value {

        // check if the value is an instruction
        let inst = match value {
            ir::Value::Instruction(i) => i,
            _ => return value,
        };

        self::reduce::reduce(inst).into()
    }
}

// TODO: blamket impl for all passes
impl Into<pass::Info<ir::Value>> for Box<StrengthReduction>
{
    fn into(self) -> pass::Info<ir::Value> {
        pass::Info::Transform(self)
    }
}

pub mod reduce
{
    use ir::{self,instruction,Instruction};

    pub fn reduce(inst: Instruction) -> ir::Instruction {

        match inst {
            Instruction::Mul(i) => self::mul(i),
            _ => inst,
        }
    }

    pub fn mul(inst: instruction::Mul) -> ir::Instruction {
        self::mul_pow2_shl(inst)
    }

    pub fn mul_pow2_shl(inst: instruction::Mul) -> ir::Instruction {
        use ir::instruction::Binary;

        let (lhs,rhs) = inst.operands();

        let lhs_if_shift = lhs.as_literal().and_then(|a| util::get_mul_shift_amount(a));
        let rhs_if_shift = rhs.as_literal().and_then(|a| util::get_mul_shift_amount(a));

        // multiplication is commutative so switch the order if necessary.
        let (value,shift) = match (lhs_if_shift,
                                   rhs_if_shift) {
            (None, None) => return inst.into(),
            // Constant folding should've caught this, but handle it anyway.
            // Both operands could be treated as the shift amount,
            // so use RHS.
            (Some(_),Some(v)) => (lhs.clone(),v.into()),

            (None, Some(v)) => (lhs.clone(),v.into()),
            (Some(v),None) => (v.into(),rhs.clone()),
        };

        Instruction::shl(value, shift).into()
    }

    pub mod util {
        use ir::{value,Value};

        /// Checks if a value is an integer and a power of two.
        pub fn is_power_of_two(value: &value::Literal) -> bool {
            use ::num::traits::ToPrimitive;

            // FIXME: this will panic if the value >64bits
            // N.B. this will give that `0` is a power of two.
            //      we don't care because constant folding has already been done.
            value.as_integer()
                 .map(|i| i.value.to_i64().unwrap())
                 .map_or(false, |x| {
                     debug_assert!(x != 0);

                     (x & (x - 1)) == 0
                 })
        }

        /// If `value` is a power of two, this gets the
        /// number of bits that would make an equivalent shift.
        ///
        /// Returns `None` if the value is not a power of two.
        pub fn get_mul_shift_amount(value: &value::Literal) -> Option<Value> {
            use ::num::traits::ToPrimitive;

            if !is_power_of_two(&value) {
                return None;
            }

            // FIXME: this will panic if value >64bits
            let const_val = value.as_integer()
                                 .expect("value must be an integer");

            let ty = const_val.ty;
            let val = const_val.value.to_f64().unwrap();

            let log2 = 2f64.log(2.);

            let n = (val.log(2.) / log2) as u64;

            Some(Value::integer(ty, n).unwrap())
        }
    }
}

value_mapping_test!(test_mul_div_shift : reduce::reduce {

    // i8
    Instruction::mul(2 as i8,1 as i8) => Instruction::shl(2 as i8,0 as i8),
    Instruction::mul(2 as i8,2 as i8) => Instruction::shl(2 as i8,1 as i8),

    // u32
    Instruction::mul(2 as u32,1 as u32) => Instruction::shl(2 as u32,0 as u32),
    Instruction::mul(2 as u32,2 as u32) => Instruction::shl(2 as u32,1 as u32),

    // Cases we shouldn't handle (non-powers of two).
    Instruction::mul(5 as i16,3 as i16) => Instruction::mul(5 as i16,3 as i16)
});
<commit_msg>Added division->shift strength reduction case<commit_after>
use pass;
use ir;

// FIXME: do not perform strength reduction if the value would overflow

/// An IR strength reduction pass.
pub struct StrengthReduction;

impl pass::Metadata for StrengthReduction
{
    fn id(&self) -> pass::Id { pass::Id(0x242343a1) }
    fn name(&self) -> &'static str { "Strength reduction" }
}

impl pass::Transform<ir::Value> for StrengthReduction
{
    fn run_value(&mut self, value: ir::Value) -> ir::Value {

        // check if the value is an instruction
        let inst = match value {
            ir::Value::Instruction(i) => i,
            _ => return value,
        };

        self::reduce::reduce(inst).into()
    }
}

// TODO: blamket impl for all passes
impl Into<pass::Info<ir::Value>> for Box<StrengthReduction>
{
    fn into(self) -> pass::Info<ir::Value> {
        pass::Info::Transform(self)
    }
}

pub mod reduce
{
    use ir::{self,instruction,Instruction};

    pub fn reduce(inst: Instruction) -> ir::Instruction {

        match inst {
            Instruction::Mul(i) => self::mul(i),
            Instruction::Div(i) => self::div(i),
            _ => inst,
        }
    }

    pub fn mul(inst: instruction::Mul) -> ir::Instruction {
        self::mul_pow2_shl(inst)
    }

    pub fn div(inst: instruction::Div) -> ir::Instruction {
        self::div_pow2_shr(inst)
    }

    /// Converts a multiplication instruction into a shift instruction if possible.
    pub fn mul_pow2_shl(inst: instruction::Mul) -> ir::Instruction {
        use ir::instruction::Binary;

        let (lhs,rhs) = inst.operands();
        match maybe_shift_values(lhs,rhs) {
            Some((value,amount)) => Instruction::shl(value,amount).into(),
            None => inst.into(),
        }
    }

    /// Converts a division instruction into a shift instruction if possible.
    pub fn div_pow2_shr(inst: instruction::Div) -> ir::Instruction {
        use ir::instruction::Binary;

        let (lhs,rhs) = inst.operands();
        match maybe_shift_values(lhs,rhs) {
            Some((value,amount)) => Instruction::shr(value,amount).into(),
            None => inst.into(),
        }
    }

    /// Tries to convert the operands of a mul or div instruction into the operands
    /// of a shift instruction.
    pub fn maybe_shift_values(lhs: ir::Value, rhs: ir::Value) -> Option<(ir::Value,ir::Value)> {
        let lhs_if_shift = lhs.as_literal().and_then(|a| util::get_shift_amount(a));
        let rhs_if_shift = rhs.as_literal().and_then(|a| util::get_shift_amount(a));

        // multiplication is commutative so switch the order if necessary.
        match (lhs_if_shift,
                                   rhs_if_shift) {
            (None, None) => None,
            // Constant folding should've caught this, but handle it anyway.
            // Both operands could be treated as the shift amount,
            // so use RHS.
            (Some(_),Some(v)) => Some((lhs.clone(),v.into())),

            (None, Some(v)) => Some((lhs.clone(),v.into())),
            (Some(v),None) => Some((v.into(),rhs.clone())),
        }
    }

    pub mod util {
        use ir::{value,Value};

        /// Checks if a value is an integer and a power of two.
        pub fn is_power_of_two(value: &value::Literal) -> bool {
            use ::num::traits::ToPrimitive;

            // FIXME: this will panic if the value >64bits
            // N.B. this will give that `0` is a power of two.
            //      we don't care because constant folding has already been done.
            value.as_integer()
                 .map(|i| i.value.to_i64().unwrap())
                 .map_or(false, |x| {
                     debug_assert!(x != 0);

                     (x & (x - 1)) == 0
                 })
        }

        /// If `value` is a power of two, this gets the
        /// number of bits that would make an equivalent shift.
        ///
        /// Returns `None` if the value is not a power of two.
        pub fn get_shift_amount(value: &value::Literal) -> Option<Value> {
            use ::num::traits::ToPrimitive;

            if !is_power_of_two(&value) {
                return None;
            }

            // FIXME: this will panic if value >64bits
            let const_val = value.as_integer()
                                 .expect("value must be an integer");

            let ty = const_val.ty;
            let val = const_val.value.to_f64().unwrap();

            let log2 = 2f64.log(2.);

            let n = (val.log(2.) / log2) as u64;

            Some(Value::integer(ty, n).unwrap())
        }
    }
}

value_mapping_test!(test_mul_div_shift : reduce::reduce {

    // i8
    Instruction::mul(2 as i8,1 as i8) => Instruction::shl(2 as i8,0 as i8),
    Instruction::mul(2 as i8,2 as i8) => Instruction::shl(2 as i8,1 as i8),
    Instruction::div(2 as i8,1 as i8) => Instruction::shr(2 as i8,0 as i8),
    Instruction::div(2 as i8,2 as i8) => Instruction::shr(2 as i8,1 as i8),

    // u32
    Instruction::mul(2 as u32,1 as u32) => Instruction::shl(2 as u32,0 as u32),
    Instruction::mul(2 as u32,2 as u32) => Instruction::shl(2 as u32,1 as u32),
    Instruction::div(2 as u32,1 as u32) => Instruction::shr(2 as u32,0 as u32),
    Instruction::div(2 as u32,2 as u32) => Instruction::shr(2 as u32,1 as u32),

    // Cases we shouldn't handle (non-powers of two).
    Instruction::mul(5 as i16,3 as i16) => Instruction::mul(5 as i16,3 as i16),
    Instruction::div(5 as i16,3 as i16) => Instruction::div(5 as i16,3 as i16)
});
<|endoftext|>


<commit_before>#[cfg(feature = "mysql")]
extern crate dotenv;

use std::collections::HashMap;

#[cfg(feature = "mysql")]
use diesel::prelude::*;
#[cfg(feature = "mysql")]
use diesel::mysql::MysqlConnection;
#[cfg(feature = "mysql")]
use r2d2::Pool;
#[cfg(feature = "mysql")]
use r2d2_diesel::ConnectionManager;

use chrono::NaiveDateTime;

pub enum DbResponse {
    Success,
    Failed(&'static str),
}

#[cfg_attr(feature = "mysql", derive(Queryable))]
#[derive(Clone, Debug)]
pub struct Factoid {
    pub name: String,
    pub idx: i32,
    pub content: String,
    pub author: String,
    pub created: NaiveDateTime,
}

#[cfg(feature = "mysql")]
use self::mysql::factoids;
#[cfg_attr(feature = "mysql", derive(Insertable))]
#[cfg_attr(feature = "mysql", table_name = "factoids")]
pub struct NewFactoid<'a> {
    pub name: &'a str,
    pub idx: i32,
    pub content: &'a str,
    pub author: &'a str,
    pub created: NaiveDateTime,
}

pub trait Database: Send {
    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse;
    fn get_factoid(&self, name: &str, idx: i32) -> Option<Factoid>;
    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse;
    fn count_factoids(&self, name: &str) -> Result<i32, &'static str>;
}

// HashMap
impl Database for HashMap<(String, i32), Factoid> {
    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse {
        let factoid = Factoid {
            name: String::from(factoid.name),
            idx: factoid.idx,
            content: factoid.content.to_string(),
            author: factoid.author.to_string(),
            created: factoid.created,
        };

        let name = factoid.name.clone();
        match self.insert((name, factoid.idx), factoid) {
            None => DbResponse::Success,
            Some(_) => DbResponse::Failed("Factoid was overwritten"),
        }
    }

    fn get_factoid(&self, name: &str, idx: i32) -> Option<Factoid> {
        self.get(&(String::from(name), idx)).cloned()
    }

    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse {
        match self.remove(&(String::from(name), idx)) {
            Some(_) => DbResponse::Success,
            None => DbResponse::Failed("Factoid not found"),
        }
    }

    fn count_factoids(&self, name: &str) -> Result<i32, &'static str> {
        Ok(self.iter().filter(|&(&(ref n, _), _)| n == name).count() as i32)
    }
}

// Diesel automatically define the factoids module as public.
// For now this is how we keep it private.
#[cfg(feature = "mysql")]
mod mysql {
    table! {
        factoids (name, idx) {
            name -> Varchar,
            idx -> Integer,
            content -> Text,
            author -> Varchar,
            created -> Timestamp,
        }
    }
}

#[cfg(feature = "mysql")]
impl Database for Pool<ConnectionManager<MysqlConnection>> {
    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse {
        use diesel;

        let conn = &*self.get().expect("Failed to get connection");
        match diesel::insert_into(factoids::table)
            .values(factoid)
            .execute(conn)
        {
            Ok(_) => DbResponse::Success,
            Err(e) => {
                error!("DB Insertion Error: {}", e);
                DbResponse::Failed("Failed to add factoid")
            }
        }
    }

    fn get_factoid(&self, name: &str, idx: i32) -> Option<Factoid> {
        let conn = &*self.get().expect("Failed to get connection");
        match factoids::table.find((name, idx)).first(conn) {
            Ok(f) => Some(f),
            Err(e) => {
                error!("DB Count Error: {}", e);
                None
            }
        }
    }

    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse {
        use diesel;
        use self::factoids::columns;

        let conn = &*self.get().expect("Failed to get connection");
        match diesel::delete(
            factoids::table
                .filter(columns::name.eq(name))
                .filter(columns::idx.eq(idx)),
        ).execute(conn)
        {
            Ok(v) => {
                if v > 0 {
                    DbResponse::Success
                } else {
                    DbResponse::Failed("Could not find any factoid with that name")
                }
            }
            Err(e) => {
                error!("DB Deletion Error: {}", e);
                DbResponse::Failed("Failed to delete factoid")
            }
        }
    }

    fn count_factoids(&self, name: &str) -> Result<i32, &'static str> {
        let conn = &*self.get().expect("Failed to get connection");
        let count: Result<i64, _> = factoids::table
            .filter(factoids::columns::name.eq(name))
            .count()
            .first(conn);

        match count {
            Ok(c) => Ok(c as i32),
            Err(e) => {
                error!("DB Count Error: {}", e);
                Err("Database Error")
            }
        }
    }
}
<commit_msg>Return zero if count_factoids could not find the factoid<commit_after>#[cfg(feature = "mysql")]
extern crate dotenv;

use std::collections::HashMap;

#[cfg(feature = "mysql")]
use diesel::prelude::*;
#[cfg(feature = "mysql")]
use diesel::mysql::MysqlConnection;
#[cfg(feature = "mysql")]
use r2d2::Pool;
#[cfg(feature = "mysql")]
use r2d2_diesel::ConnectionManager;

use chrono::NaiveDateTime;

pub enum DbResponse {
    Success,
    Failed(&'static str),
}

#[cfg_attr(feature = "mysql", derive(Queryable))]
#[derive(Clone, Debug)]
pub struct Factoid {
    pub name: String,
    pub idx: i32,
    pub content: String,
    pub author: String,
    pub created: NaiveDateTime,
}

#[cfg(feature = "mysql")]
use self::mysql::factoids;
#[cfg_attr(feature = "mysql", derive(Insertable))]
#[cfg_attr(feature = "mysql", table_name = "factoids")]
pub struct NewFactoid<'a> {
    pub name: &'a str,
    pub idx: i32,
    pub content: &'a str,
    pub author: &'a str,
    pub created: NaiveDateTime,
}

pub trait Database: Send {
    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse;
    fn get_factoid(&self, name: &str, idx: i32) -> Option<Factoid>;
    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse;
    fn count_factoids(&self, name: &str) -> Result<i32, &'static str>;
}

// HashMap
impl Database for HashMap<(String, i32), Factoid> {
    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse {
        let factoid = Factoid {
            name: String::from(factoid.name),
            idx: factoid.idx,
            content: factoid.content.to_string(),
            author: factoid.author.to_string(),
            created: factoid.created,
        };

        let name = factoid.name.clone();
        match self.insert((name, factoid.idx), factoid) {
            None => DbResponse::Success,
            Some(_) => DbResponse::Failed("Factoid was overwritten"),
        }
    }

    fn get_factoid(&self, name: &str, idx: i32) -> Option<Factoid> {
        self.get(&(String::from(name), idx)).cloned()
    }

    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse {
        match self.remove(&(String::from(name), idx)) {
            Some(_) => DbResponse::Success,
            None => DbResponse::Failed("Factoid not found"),
        }
    }

    fn count_factoids(&self, name: &str) -> Result<i32, &'static str> {
        Ok(self.iter().filter(|&(&(ref n, _), _)| n == name).count() as i32)
    }
}

// Diesel automatically define the factoids module as public.
// For now this is how we keep it private.
#[cfg(feature = "mysql")]
mod mysql {
    table! {
        factoids (name, idx) {
            name -> Varchar,
            idx -> Integer,
            content -> Text,
            author -> Varchar,
            created -> Timestamp,
        }
    }
}

#[cfg(feature = "mysql")]
impl Database for Pool<ConnectionManager<MysqlConnection>> {
    fn insert_factoid(&mut self, factoid: &NewFactoid) -> DbResponse {
        use diesel;

        let conn = &*self.get().expect("Failed to get connection");
        match diesel::insert_into(factoids::table)
            .values(factoid)
            .execute(conn)
        {
            Ok(_) => DbResponse::Success,
            Err(e) => {
                error!("DB Insertion Error: {}", e);
                DbResponse::Failed("Failed to add factoid")
            }
        }
    }

    fn get_factoid(&self, name: &str, idx: i32) -> Option<Factoid> {
        let conn = &*self.get().expect("Failed to get connection");
        match factoids::table.find((name, idx)).first(conn) {
            Ok(f) => Some(f),
            Err(e) => {
                error!("DB Count Error: {}", e);
                None
            }
        }
    }

    fn delete_factoid(&mut self, name: &str, idx: i32) -> DbResponse {
        use diesel;
        use self::factoids::columns;

        let conn = &*self.get().expect("Failed to get connection");
        match diesel::delete(
            factoids::table
                .filter(columns::name.eq(name))
                .filter(columns::idx.eq(idx)),
        ).execute(conn)
        {
            Ok(v) => {
                if v > 0 {
                    DbResponse::Success
                } else {
                    DbResponse::Failed("Could not find any factoid with that name")
                }
            }
            Err(e) => {
                error!("DB Deletion Error: {}", e);
                DbResponse::Failed("Failed to delete factoid")
            }
        }
    }

    fn count_factoids(&self, name: &str) -> Result<i32, &'static str> {
        use diesel;

        let conn = &*self.get().expect("Failed to get connection");
        let count: Result<i64, _> = factoids::table
            .filter(factoids::columns::name.eq(name))
            .count()
            .get_result(conn);

        match count {
            Ok(c) => Ok(c as i32),
            Err(diesel::NotFound) => Ok(0),
            Err(e) => {
                error!("DB Count Error: {}", e);
                Err("Database Error")
            }
        }
    }
}
<|endoftext|>