<commit_before>extern crate hyper;
extern crate image;

use self::hyper::client::Client;
use self::hyper::header::UserAgent;

use std::io::{Read,Write};
use std::fs::{File,OpenOptions,read_dir,create_dir};
use std::path::Path;

pub use super::DB;

use rustc_serialize::json::Json;

use std::sync::mpsc::{Receiver,TryRecvError};

use self::image::FilterType;

#[derive(Debug)]
pub struct Image {
    name: String,
    tags: Vec<String>,
    got_from: String,
    url: String,
    rating: char,
    post_url: String
}

pub mod e621;
pub mod derpy;
pub mod danbooru;
pub mod konachan;

/// Сохраняет картинку & создаёт к ней превью
pub fn save_image(dir: &Path, name: &str, file: &[u8]) {
    if let Err(_) = read_dir("assets/images") {
        create_dir("assets/images").unwrap();
    }
    if let Err(_) = read_dir("assets/images/preview") {
        create_dir("assets/images/preview").unwrap();
    }

    let prev = image::load_from_memory(file).unwrap().resize(500, 500, FilterType::Nearest);

    let mut f = File::create(dir.join(name)).unwrap();
    let mut prevf = File::create(dir.join("preview").join(name)).unwrap();

    f.write(file).unwrap();
    prev.save(&mut prevf, image::JPEG).unwrap();
}

fn process_downloads(client: &Client, images: &[Image], recv: &Receiver<()>) -> Result<(),()> {
    let images_c = DB.lock().unwrap().get_images(None,0).unwrap();
    let mut outf = OpenOptions::new().append(true).create(true).open("OUTPUT").unwrap();

    for im in images {
        match recv.try_recv() {
            Ok(_) | Err(TryRecvError::Disconnected) => {
                return Err(());
            }
            Err(TryRecvError::Empty) => {}
        }

        if !images_c.iter().any(|x| x.name == im.name ) {
            if im.got_from == "konachan" || im.got_from == "danbooru" {
                download(&Client::new(), im)?;
            } else {
                download(client, im)?;
            }
        } else {
            let mut m_tags = images_c.iter().find(|x| x.name == im.name ).unwrap().tags.clone();
            let mut curr_tags = im.tags.clone();
            m_tags.dedup();
            curr_tags.dedup();
            if m_tags != curr_tags {
                DB.lock().unwrap().add_image(&im.name, &im.tags, im.got_from.as_str(), im.post_url.as_str(), im.rating).unwrap();
                writeln!(&mut outf, "UPDATE tags on {}", im.name).unwrap();
            }
        }
    }
    Ok(())
}

/// Качает картинку, прерываясь, если из консоли поступил kill
fn download(client: &Client, im: &Image) -> Result<(),()> {
    let mut outf = OpenOptions::new().append(true).create(true).open("OUTPUT").unwrap();

    let mut res = client.get(&im.url)
        .header(UserAgent("Zeph/1.0".to_owned()))
        .send().unwrap();
    let mut body = Vec::new();
    res.read_to_end(&mut body).unwrap();

    DB.lock().unwrap().add_image(&im.name, &im.tags, im.got_from.as_str(), im.post_url.as_str(), im.rating).unwrap();

    save_image(Path::new("assets/images"), &im.name, &body);

    writeln!(&mut outf, "DONE {}", im.name).unwrap();

    Ok(())
}

/// Запросить и распарсить JSON
fn req_and_parse(client: &Client, url: &str) -> Result<Json, hyper::Error> {
    let mut outf = OpenOptions::new().append(true).create(true).open("OUTPUT").unwrap();
    let mut res = match client.get(url)
        .header(UserAgent("Zeph/1.0".to_owned()))
        .send() {
            Ok(x)   => x,
            Err(x)  => {
                writeln!(&mut outf, "ERROR: {}", x).unwrap();
                return Err(x)
            }
        };

    let mut body = String::new();
    res.read_to_string(&mut body).unwrap();

    Ok(Json::from_str(&body).unwrap())
}
<commit_msg>Меньше паники, по идее<commit_after>extern crate hyper;
extern crate image;

use self::hyper::client::Client;
use self::hyper::header::UserAgent;

use std::io::{Read,Write};
use std::fs::{File,OpenOptions,read_dir,create_dir};
use std::path::Path;

pub use super::DB;

use rustc_serialize::json::Json;

use std::sync::mpsc::{Receiver,TryRecvError};

use self::image::FilterType;

#[derive(Debug)]
pub struct Image {
    name: String,
    tags: Vec<String>,
    got_from: String,
    url: String,
    rating: char,
    post_url: String
}

pub mod e621;
pub mod derpy;
pub mod danbooru;
pub mod konachan;

/// Сохраняет картинку & создаёт к ней превью
pub fn save_image(dir: &Path, name: &str, file: &[u8]) {
    if let Err(_) = read_dir("assets/images") {
        create_dir("assets/images").unwrap();
    }
    if let Err(_) = read_dir("assets/images/preview") {
        create_dir("assets/images/preview").unwrap();
    }

    let prev = image::load_from_memory(file).unwrap().resize(500, 500, FilterType::Nearest);

    let mut f = File::create(dir.join(name)).unwrap();
    let mut prevf = File::create(dir.join("preview").join(name)).unwrap();

    f.write(file).unwrap();
    prev.save(&mut prevf, image::JPEG).unwrap();
}

fn process_downloads(client: &Client, images: &[Image], recv: &Receiver<()>) -> Result<(),()> {
    let images_c = DB.lock().unwrap().get_images(None,0).unwrap();
    let mut outf = OpenOptions::new().append(true).create(true).open("OUTPUT").unwrap();

    for im in images {
        match recv.try_recv() {
            Ok(_) | Err(TryRecvError::Disconnected) => {
                return Err(());
            }
            Err(TryRecvError::Empty) => {}
        }

        if !images_c.iter().any(|x| x.name == im.name ) {
            if let Err(er) = if im.got_from == "konachan" || im.got_from == "danbooru" {
                download(&Client::new(), im)
            } else {
                download(client, im)
            } {
                writeln!(&mut outf, "ERROR: {}; SKIP", er).unwrap();
                continue
            } else {
                writeln!(&mut outf, "DONE {}", im.name).unwrap();
            }
        } else {
            let mut m_tags = images_c.iter().find(|x| x.name == im.name ).unwrap().tags.clone();
            let mut curr_tags = im.tags.clone();
            m_tags.dedup();
            curr_tags.dedup();
            if m_tags != curr_tags {
                DB.lock().unwrap().add_image(&im.name, &im.tags, im.got_from.as_str(), im.post_url.as_str(), im.rating).unwrap();
                writeln!(&mut outf, "UPDATE tags on {}", im.name).unwrap();
            }
        }
    }
    Ok(())
}

/// Качает картинку, прерываясь, если из консоли поступил kill
fn download(client: &Client, im: &Image) -> Result<(), hyper::Error> {
    let mut res = client.get(&im.url)
        .header(UserAgent("Zeph/1.0".to_owned()))
        .send()?;

    let mut body = Vec::new();
    res.read_to_end(&mut body).unwrap();

    DB.lock().unwrap().add_image(&im.name, &im.tags, im.got_from.as_str(), im.post_url.as_str(), im.rating).unwrap();

    save_image(Path::new("assets/images"), &im.name, &body);

    Ok(())
}

/// Запросить и распарсить JSON
fn req_and_parse(client: &Client, url: &str) -> Result<Json, hyper::Error> {
    let mut outf = OpenOptions::new().append(true).create(true).open("OUTPUT").unwrap();
    let mut res = match client.get(url)
        .header(UserAgent("Zeph/1.0".to_owned()))
        .send() {
            Ok(x)   => x,
            Err(x)  => {
                writeln!(&mut outf, "ERROR: {}", x).unwrap();
                return Err(x)
            }
        };

    let mut body = String::new();
    res.read_to_string(&mut body).unwrap();

    Ok(Json::from_str(&body).unwrap())
}
<|endoftext|>


<commit_before>//! Low level and internal http and https implementation.

use std::sync::Arc;
use std::time::Duration;
use std::io::{self, Write, stderr};
use std::io::prelude::*;
use std::net::TcpStream;

use rustls::{Session, ClientConfig, ClientSession, ProtocolVersion};
use webpki_roots;
use log::LogLevel::Info;

use super::request::Request;
use super::results::{CabotResult, CabotError};
use super::dns::Resolver;

fn log_request(request: &str, verbose: bool) {
    if !log_enabled!(Info) && !verbose {
        return
    }
    let split = request.split("\r\n");
    if log_enabled!(Info) {
        for part in split {
            info!("> {}", part);
        }
    }
    else if verbose {
        for part in split {
            writeln!(&mut stderr(), "> {}", part).unwrap();
        }
    }
}

fn from_http(request: &Request,
                 mut client: &mut TcpStream,
                 mut out: &mut Write,
                 verbose: bool)
                 -> CabotResult<()> {

    let request_str = request.to_string();
    log_request(&request_str, verbose);

    debug!("Sending request {}", request_str);
    client.write(request_str.as_bytes()).unwrap();

    let mut response:Vec<u8> = Vec::new();
    client.read_to_end(&mut response).unwrap();
    out.write_all(response.as_slice()).unwrap();
    Ok(())
}

fn from_https(request: &Request,
                  mut client: &mut TcpStream,
                  mut out: &mut Write,
                  verbose: bool)
                  -> CabotResult<()> {

    let request_str = request.to_string();

    let mut config = ClientConfig::new();
    config.root_store.add_trust_anchors(&webpki_roots::ROOTS);
    let rc_config = Arc::new(config);
    let mut tlsclient = ClientSession::new(&rc_config, request.host());
    let mut is_handshaking = true;
    let mut response: Vec<u8> = Vec::new();
    loop {
        while tlsclient.wants_write() {
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }
        if is_handshaking && !tlsclient.is_handshaking() {
            info!("Handshake complete");
            is_handshaking = false;
            let protocol = tlsclient.get_protocol_version();
            match protocol {
                Some(ProtocolVersion::SSLv2) => {
                    info!("Protocol SSL v2 negociated");
                },
                Some(ProtocolVersion::SSLv3) => {
                    info!("Protocol SSL v3 negociated");
                },
                Some(ProtocolVersion::TLSv1_0) => {
                    info!("Protocol TLS v1.0 negociated");
                },
                Some(ProtocolVersion::TLSv1_1) => {
                    info!("Protocol TLS v1.1 negociated");
                },
                Some(ProtocolVersion::TLSv1_2) => {
                    info!("Protocol TLS v1.2 negociated");
                },
                Some(ProtocolVersion::TLSv1_3) => {
                    info!("Protocol TLS v1.3 negociated");
                },
                Some(ProtocolVersion::Unknown(num)) => {
                    info!("Unknown TLS Protocol negociated: {}", num);
                },
                None => {
                    info!("No TLS Protocol negociated");
                }
            }
            log_request(&request_str, verbose);
            tlsclient.write_all(request_str.as_bytes()).unwrap();
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }

        if tlsclient.wants_read() {
            let count = tlsclient.read_tls(&mut client);
            if let Err(err) = count {
                error!("{:?}", err);
                return Err(CabotError::IOError(format!("{}", err)));
            }

            let count = count.unwrap();
            debug!("Read {} TLS bytes", count);
            if count == 0 {
                break;
            }

            if let Err(err) = tlsclient.process_new_packets() {
                return Err(CabotError::CertificateError(format!("{}", err)));
            }

            let mut part: Vec<u8> = Vec::new();
            let clearcount = tlsclient.read_to_end(&mut part);
            if let Err(err) = clearcount {
                response.append(&mut part);

                if err.kind() == io::ErrorKind::ConnectionAborted {
                    break;
                }
                error!("{:?}", err);
                return Err(CabotError::IOError(format!("{}", err)));
            } else {
                let clearcount = clearcount.unwrap();
                debug!("Read {} clear bytes", clearcount);
                if clearcount > 0 {
                    response.append(&mut part);
                }
            }
        } else {
            break;
        }
    }
    out.write_all(response.as_slice()).unwrap();
    Ok(())
}

pub fn http_query(request: &Request,
                  mut out: &mut Write,
                  verbose: bool)
                  -> CabotResult<()> {
    debug!("HTTP Query {} {}", request.http_method(), request.request_uri());

    let resolver = Resolver::new(verbose);
    let authority = request.authority();
    let addr = resolver.get_addr(authority);

    info!("Connecting to {}", addr);
    let mut client = TcpStream::connect(addr).unwrap();
    client.set_read_timeout(Some(Duration::new(5, 0))).unwrap();

    match request.scheme() {
        "http" => from_http(request, &mut client, &mut out, verbose)?,
        "https" => from_https(request, &mut client, &mut out, verbose)?,
        _ => {
            return Err(CabotError::SchemeError(format!("Unrecognized scheme {}", request.scheme())))
        }
    };

    out.flush().unwrap();

    Ok(())

}
<commit_msg>Read stream using a buffer instead of a vector<commit_after>//! Low level and internal http and https implementation.

use std::sync::Arc;
use std::time::Duration;
use std::io::{Read, Write, stderr};
use std::net::TcpStream;

use rustls::{Session, ClientConfig, ClientSession, ProtocolVersion};
use webpki_roots;
use log::LogLevel::Info;

use super::request::Request;
use super::results::{CabotResult, CabotError};
use super::dns::Resolver;

const BUFFER_PAGE_SIZE: usize = 1024;
const RESPONSE_BUFFER_SIZE: usize = 1024;

fn log_request(request: &str, verbose: bool) {
    if !log_enabled!(Info) && !verbose {
        return
    }
    let split = request.split("\r\n");
    if log_enabled!(Info) {
        for part in split {
            info!("> {}", part);
        }
    }
    else if verbose {
        for part in split {
            writeln!(&mut stderr(), "> {}", part).unwrap();
        }
    }
}

fn read_buf<T>(mut client: &mut T,
               mut buf: &mut [u8]) -> Vec<u8> where T: Read + Sized {
    let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
    loop {
        match client.read(&mut buf[..]) {
            Ok(count) => {
                if count > 0 {
                    response.extend_from_slice(&buf[0..count]);
                }
                else {
                    break
                }
            }
            Err(_) => { break } // connection is closed by client
        }
    }
    response
}

fn from_http(request: &Request,
                 mut client: &mut TcpStream,
                 mut out: &mut Write,
                 verbose: bool)
                 -> CabotResult<()> {

    let request_str = request.to_string();
    log_request(&request_str, verbose);

    debug!("Sending request {}", request_str);
    client.write(request_str.as_bytes()).unwrap();
    let mut buf = [0; BUFFER_PAGE_SIZE];
    let response = read_buf(client, &mut buf);
    out.write_all(response.as_slice()).unwrap();
    Ok(())
}

fn from_https(request: &Request,
                  mut client: &mut TcpStream,
                  mut out: &mut Write,
                  verbose: bool)
                  -> CabotResult<()> {

    let request_str = request.to_string();
    let mut response: Vec<u8> = Vec::with_capacity(RESPONSE_BUFFER_SIZE);
    let mut buf = [0; BUFFER_PAGE_SIZE];

    let mut config = ClientConfig::new();
    config.root_store.add_trust_anchors(&webpki_roots::ROOTS);
    let rc_config = Arc::new(config);
    let mut tlsclient = ClientSession::new(&rc_config, request.host());
    let mut is_handshaking = true;
    loop {
        while tlsclient.wants_write() {
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }
        if is_handshaking && !tlsclient.is_handshaking() {
            info!("Handshake complete");
            is_handshaking = false;
            let protocol = tlsclient.get_protocol_version();
            match protocol {
                Some(ProtocolVersion::SSLv2) => {
                    info!("Protocol SSL v2 negociated");
                },
                Some(ProtocolVersion::SSLv3) => {
                    info!("Protocol SSL v3 negociated");
                },
                Some(ProtocolVersion::TLSv1_0) => {
                    info!("Protocol TLS v1.0 negociated");
                },
                Some(ProtocolVersion::TLSv1_1) => {
                    info!("Protocol TLS v1.1 negociated");
                },
                Some(ProtocolVersion::TLSv1_2) => {
                    info!("Protocol TLS v1.2 negociated");
                },
                Some(ProtocolVersion::TLSv1_3) => {
                    info!("Protocol TLS v1.3 negociated");
                },
                Some(ProtocolVersion::Unknown(num)) => {
                    info!("Unknown TLS Protocol negociated: {}", num);
                },
                None => {
                    info!("No TLS Protocol negociated");
                }
            }
            log_request(&request_str, verbose);
            tlsclient.write_all(request_str.as_bytes()).unwrap();
            let count = tlsclient.write_tls(&mut client).unwrap();
            debug!("Write {} TLS bytes", count);
        }

        if tlsclient.wants_read() {
            let count = tlsclient.read_tls(&mut client);
            if let Err(err) = count {
                error!("{:?}", err);
                return Err(CabotError::IOError(format!("{}", err)));
            }

            let count = count.unwrap();
            debug!("Read {} TLS bytes", count);
            if count == 0 {
                break;
            }

            if let Err(err) = tlsclient.process_new_packets() {
                return Err(CabotError::CertificateError(format!("{}", err)));
            }

            let mut part: Vec<u8> = read_buf(&mut tlsclient, &mut buf);
            response.append(&mut part);
        } else {
            break;
        }
    }
    out.write_all(response.as_slice()).unwrap();
    Ok(())
}

pub fn http_query(request: &Request,
                  mut out: &mut Write,
                  verbose: bool)
                  -> CabotResult<()> {
    debug!("HTTP Query {} {}", request.http_method(), request.request_uri());

    let resolver = Resolver::new(verbose);
    let authority = request.authority();
    let addr = resolver.get_addr(authority);

    info!("Connecting to {}", addr);
    let mut client = TcpStream::connect(addr).unwrap();
    client.set_read_timeout(Some(Duration::new(5, 0))).unwrap();

    match request.scheme() {
        "http" => from_http(request, &mut client, &mut out, verbose)?,
        "https" => from_https(request, &mut client, &mut out, verbose)?,
        _ => {
            return Err(CabotError::SchemeError(format!("Unrecognized scheme {}", request.scheme())))
        }
    };

    out.flush().unwrap();

    Ok(())

}
<|endoftext|>


<commit_before>#![allow(non_snake_case)]
#![allow(non_upper_case_globals)]

use error_chain::ChainedError;
use libc::*;
use gl;
use gl::types::*;
use std::cmp;
use std::ffi::{CStr, CString};
use std::mem;
use std::ptr;
use std::sync::{Once, ONCE_INIT, RwLock};

use command;
use dl;
use encode;
use errors::*;
use function::Function;
use sdl;

lazy_static!{
    pub static ref POINTERS: RwLock<Pointers> = RwLock::new(Pointers::default());
    static ref CAPTURING: RwLock<bool> = RwLock::new(false);
}

#[derive(Debug, Default)]
pub struct Pointers {
    RunListenServer: Function<unsafe extern "C" fn(*mut c_void,
                                                   *mut c_char,
                                                   *mut c_char,
                                                   *mut c_char,
                                                   *mut c_void,
                                                   *mut c_void)
                                                   -> c_int>,

    Cmd_AddCommand: Function<unsafe extern "C" fn(*const c_char, *mut c_void)>,
    Cmd_Argc: Function<unsafe extern "C" fn() -> c_int>,
    Cmd_Argv: Function<unsafe extern "C" fn(c_int) -> *const c_char>,
    Con_Printf: Function<unsafe extern "C" fn(*const c_char)>,
    Memory_Init: Function<unsafe extern "C" fn(*mut c_void, c_int)>,
    Sys_VID_FlipScreen: Function<unsafe extern "C" fn()>,
    VideoMode_GetCurrentVideoMode: Function<unsafe extern "C" fn(*mut c_int,
                                                                 *mut c_int,
                                                                 *mut c_int)>,
    VideoMode_IsWindowed: Function<unsafe extern "C" fn() -> c_int>,

    window_rect: Option<*mut RECT>,
    s_BackBufferFBO: Option<*mut FBO_Container_t>,
}

// TODO: think about how to deal with unsafety here.
// The compiler complaining about *mut not being Send/Sync is perfectly valid here.
// Our safe RwLock cannot guarantee there isn't some game thread accessing the values
// at the same time. Although reading/writing to a pointer is already unsafe?
unsafe impl Send for Pointers {}
unsafe impl Sync for Pointers {}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
struct RECT {
    left: c_int,
    right: c_int,
    top: c_int,
    bottom: c_int,
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
struct FBO_Container_t {
    s_hBackBufferFBO: GLuint,
    s_hBackBufferCB: GLuint,
    s_hBackBufferDB: GLuint,
    s_hBackBufferTex: GLuint,
}

/// The "main" function of hw.so, called inside `CEngineAPI::Run()`.
///
/// The game runs within this function and shortly after it exits hw.so is unloaded.
/// Note: `_restart` also causes this function to exit, in this case the launcher
/// unloads and reloads hw.so and this function is called again as if it was a fresh start.
#[export_name = "_Z15RunListenServerPvPcS0_S0_PFP14IBaseInterfacePKcPiES7_"]
pub unsafe extern "C" fn RunListenServer(instance: *mut c_void,
                                         basedir: *mut c_char,
                                         cmdline: *mut c_char,
                                         postRestartCmdLineArgs: *mut c_char,
                                         launcherFactory: *mut c_void,
                                         filesystemFactory: *mut c_void)
                                         -> c_int {
    // hw.so just loaded, either for the first time or potentially at a different address.
    // Refresh all pointers.
    if let Err(ref e) = refresh_pointers().chain_err(|| "error refreshing pointers") {
        panic!("{}", e.display());
    }

    // Initialize the encoding.
    {
        static INIT: Once = ONCE_INIT;
        INIT.call_once(|| if let Err(ref e) = encode::initialize()
                              .chain_err(|| "error initializing encoding") {
                           panic!("{}", e.display());
                       });
    }

    let rv = real!(RunListenServer)(instance,
                                    basedir,
                                    cmdline,
                                    postRestartCmdLineArgs,
                                    launcherFactory,
                                    filesystemFactory);

    // Since hw.so is getting unloaded, reset all pointers.
    reset_pointers();

    rv
}

/// Initializes the hunk memory.
///
/// After the hunk memory has been initialized we can register console commands and variables.
/// It's also a good place to get OpenGL function pointers.
#[no_mangle]
pub unsafe extern "C" fn Memory_Init(buf: *mut c_void, size: c_int) {
    real!(Memory_Init)(buf, size);

    register_cvars_and_commands();

    gl::ReadPixels::load_with(|s| sdl::get_proc_address(s) as _);
    con_print(&format!("gl::ReadPixels::is_loaded() = {}", gl::ReadPixels::is_loaded()));
}

/// Blits pixels from the framebuffer to screen and flips.
///
/// If framebuffers aren't used, simply flips the screen.
#[export_name = "_Z18Sys_VID_FlipScreenv"]
pub unsafe extern "C" fn Sys_VID_FlipScreen() {
    let (w, h) = get_resolution();
    let buf = encode::get_buffer((w, h));

    if *CAPTURING.read().unwrap() {
        gl::ReadPixels(0, 0, w as GLsizei, h as GLsizei,
                       gl::RGB, gl::UNSIGNED_BYTE,
                       buf.as_mut_ptr() as _);
    }

    real!(Sys_VID_FlipScreen)();

    // TODO: check if we're called from SCR_UpdateScreen().

    // println!("s_BackBufferFBO: {:?}", *POINTERS.read().unwrap().s_BackBufferFBO.unwrap());
}

/// Obtains and stores all necessary function and variable addresses.
fn refresh_pointers() -> Result<()> {
    let hw = dl::open("hw.so", RTLD_NOW | RTLD_NOLOAD)
        .chain_err(|| "couldn't load hw.so")?;

    let mut pointers = POINTERS.write().unwrap();

    unsafe {
        find!(pointers,
              hw,
              RunListenServer,
              "_Z15RunListenServerPvPcS0_S0_PFP14IBaseInterfacePKcPiES7_");
        find!(pointers, hw, Cmd_AddCommand, "Cmd_AddCommand");
        find!(pointers, hw, Cmd_Argc, "Cmd_Argc");
        find!(pointers, hw, Cmd_Argv, "Cmd_Argv");
        find!(pointers, hw, Con_Printf, "Con_Printf");
        find!(pointers, hw, Memory_Init, "Memory_Init");
        find!(pointers, hw, Sys_VID_FlipScreen, "_Z18Sys_VID_FlipScreenv");
        find!(pointers, hw, VideoMode_GetCurrentVideoMode, "VideoMode_GetCurrentVideoMode");
        find!(pointers, hw, VideoMode_IsWindowed, "VideoMode_IsWindowed");

        pointers.window_rect = Some(hw.sym("window_rect")
                .chain_err(|| "couldn't find window_rect")? as _);
        pointers.s_BackBufferFBO = Some(hw.sym("s_BackBufferFBO")
                .chain_err(|| "couldn't find s_BackBufferFBO")? as _);
    }

    Ok(())
}

/// Resets all pointers to their default values.
fn reset_pointers() {
    *POINTERS.write().unwrap() = Pointers::default();
}

/// Registers console commands and variables.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
unsafe fn register_cvars_and_commands() {
    for cmd in command::COMMANDS.read().unwrap().iter() {
        register_command(cmd.name(), cmd.callback());
    }
}

/// Registers a console command.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
unsafe fn register_command(name: &'static [u8], callback: unsafe extern "C" fn()) {
    real!(Cmd_AddCommand)(name as *const _ as *const _, callback as *mut c_void);
}

/// Prints the given string to the game console.
///
/// `string` must not contain null bytes.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
pub unsafe fn con_print(string: &str) {
    real!(Con_Printf)(CString::new(string.replace('%', "%%"))
                          .expect("string cannot contain null bytes")
                          .as_ptr())
}

/// Gets the console command argument count.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
pub unsafe fn cmd_argc() -> u32 {
    let argc = real!(Cmd_Argc)();
    cmp::max(0, argc) as u32
}

/// Gets a console command argument.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
pub unsafe fn cmd_argv(index: u32) -> String {
    let index = cmp::min(index, i32::max_value() as u32) as i32;
    let arg = real!(Cmd_Argv)(index);
    CStr::from_ptr(arg).to_string_lossy().into_owned()
}

/// Returns the current game resolution.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
pub unsafe fn get_resolution() -> (u32, u32) {
    let mut width;
    let mut height;

    if real!(VideoMode_IsWindowed)() != 0 {
        let window_rect = *POINTERS.read().unwrap().window_rect.unwrap();
        width = window_rect.right - window_rect.left;
        height = window_rect.bottom - window_rect.top;
    } else {
        width = mem::uninitialized();
        height = mem::uninitialized();
        real!(VideoMode_GetCurrentVideoMode)(&mut width, &mut height, ptr::null_mut());
    }

    width = cmp::max(0, width);
    height = cmp::max(0, height);

    (width as u32, height as u32)
}

command!(cap_test, |engine| {
    let args = engine.args();

    let mut buf = String::new();
    buf.push_str(&format!("Args len: {}\n", args.len()));

    for arg in args {
        buf.push_str(&arg);
        buf.push('\n');
    }

    engine.con_print(&buf);
});

command!(cap_another_test, |engine| {
    engine.con_print("Hello! %s %d %\n");
});

command!(cap_start, |_engine| {
    *CAPTURING.write().unwrap() = true;
});

command!(cap_stop, |_engine| {
    *CAPTURING.write().unwrap() = false;
});
<commit_msg>Remove unnecessary pub<commit_after>#![allow(non_snake_case)]
#![allow(non_upper_case_globals)]

use error_chain::ChainedError;
use libc::*;
use gl;
use gl::types::*;
use std::cmp;
use std::ffi::{CStr, CString};
use std::mem;
use std::ptr;
use std::sync::{Once, ONCE_INIT, RwLock};

use command;
use dl;
use encode;
use errors::*;
use function::Function;
use sdl;

lazy_static!{
    static ref POINTERS: RwLock<Pointers> = RwLock::new(Pointers::default());
    static ref CAPTURING: RwLock<bool> = RwLock::new(false);
}

#[derive(Debug, Default)]
pub struct Pointers {
    RunListenServer: Function<unsafe extern "C" fn(*mut c_void,
                                                   *mut c_char,
                                                   *mut c_char,
                                                   *mut c_char,
                                                   *mut c_void,
                                                   *mut c_void)
                                                   -> c_int>,

    Cmd_AddCommand: Function<unsafe extern "C" fn(*const c_char, *mut c_void)>,
    Cmd_Argc: Function<unsafe extern "C" fn() -> c_int>,
    Cmd_Argv: Function<unsafe extern "C" fn(c_int) -> *const c_char>,
    Con_Printf: Function<unsafe extern "C" fn(*const c_char)>,
    Memory_Init: Function<unsafe extern "C" fn(*mut c_void, c_int)>,
    Sys_VID_FlipScreen: Function<unsafe extern "C" fn()>,
    VideoMode_GetCurrentVideoMode: Function<unsafe extern "C" fn(*mut c_int,
                                                                 *mut c_int,
                                                                 *mut c_int)>,
    VideoMode_IsWindowed: Function<unsafe extern "C" fn() -> c_int>,

    window_rect: Option<*mut RECT>,
    s_BackBufferFBO: Option<*mut FBO_Container_t>,
}

// TODO: think about how to deal with unsafety here.
// The compiler complaining about *mut not being Send/Sync is perfectly valid here.
// Our safe RwLock cannot guarantee there isn't some game thread accessing the values
// at the same time. Although reading/writing to a pointer is already unsafe?
unsafe impl Send for Pointers {}
unsafe impl Sync for Pointers {}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
struct RECT {
    left: c_int,
    right: c_int,
    top: c_int,
    bottom: c_int,
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
struct FBO_Container_t {
    s_hBackBufferFBO: GLuint,
    s_hBackBufferCB: GLuint,
    s_hBackBufferDB: GLuint,
    s_hBackBufferTex: GLuint,
}

/// The "main" function of hw.so, called inside `CEngineAPI::Run()`.
///
/// The game runs within this function and shortly after it exits hw.so is unloaded.
/// Note: `_restart` also causes this function to exit, in this case the launcher
/// unloads and reloads hw.so and this function is called again as if it was a fresh start.
#[export_name = "_Z15RunListenServerPvPcS0_S0_PFP14IBaseInterfacePKcPiES7_"]
pub unsafe extern "C" fn RunListenServer(instance: *mut c_void,
                                         basedir: *mut c_char,
                                         cmdline: *mut c_char,
                                         postRestartCmdLineArgs: *mut c_char,
                                         launcherFactory: *mut c_void,
                                         filesystemFactory: *mut c_void)
                                         -> c_int {
    // hw.so just loaded, either for the first time or potentially at a different address.
    // Refresh all pointers.
    if let Err(ref e) = refresh_pointers().chain_err(|| "error refreshing pointers") {
        panic!("{}", e.display());
    }

    // Initialize the encoding.
    {
        static INIT: Once = ONCE_INIT;
        INIT.call_once(|| if let Err(ref e) = encode::initialize()
                              .chain_err(|| "error initializing encoding") {
                           panic!("{}", e.display());
                       });
    }

    let rv = real!(RunListenServer)(instance,
                                    basedir,
                                    cmdline,
                                    postRestartCmdLineArgs,
                                    launcherFactory,
                                    filesystemFactory);

    // Since hw.so is getting unloaded, reset all pointers.
    reset_pointers();

    rv
}

/// Initializes the hunk memory.
///
/// After the hunk memory has been initialized we can register console commands and variables.
/// It's also a good place to get OpenGL function pointers.
#[no_mangle]
pub unsafe extern "C" fn Memory_Init(buf: *mut c_void, size: c_int) {
    real!(Memory_Init)(buf, size);

    register_cvars_and_commands();

    gl::ReadPixels::load_with(|s| sdl::get_proc_address(s) as _);
    con_print(&format!("gl::ReadPixels::is_loaded() = {}", gl::ReadPixels::is_loaded()));
}

/// Blits pixels from the framebuffer to screen and flips.
///
/// If framebuffers aren't used, simply flips the screen.
#[export_name = "_Z18Sys_VID_FlipScreenv"]
pub unsafe extern "C" fn Sys_VID_FlipScreen() {
    let (w, h) = get_resolution();
    let buf = encode::get_buffer((w, h));

    if *CAPTURING.read().unwrap() {
        gl::ReadPixels(0, 0, w as GLsizei, h as GLsizei,
                       gl::RGB, gl::UNSIGNED_BYTE,
                       buf.as_mut_ptr() as _);
    }

    real!(Sys_VID_FlipScreen)();

    // TODO: check if we're called from SCR_UpdateScreen().

    // println!("s_BackBufferFBO: {:?}", *POINTERS.read().unwrap().s_BackBufferFBO.unwrap());
}

/// Obtains and stores all necessary function and variable addresses.
fn refresh_pointers() -> Result<()> {
    let hw = dl::open("hw.so", RTLD_NOW | RTLD_NOLOAD)
        .chain_err(|| "couldn't load hw.so")?;

    let mut pointers = POINTERS.write().unwrap();

    unsafe {
        find!(pointers,
              hw,
              RunListenServer,
              "_Z15RunListenServerPvPcS0_S0_PFP14IBaseInterfacePKcPiES7_");
        find!(pointers, hw, Cmd_AddCommand, "Cmd_AddCommand");
        find!(pointers, hw, Cmd_Argc, "Cmd_Argc");
        find!(pointers, hw, Cmd_Argv, "Cmd_Argv");
        find!(pointers, hw, Con_Printf, "Con_Printf");
        find!(pointers, hw, Memory_Init, "Memory_Init");
        find!(pointers, hw, Sys_VID_FlipScreen, "_Z18Sys_VID_FlipScreenv");
        find!(pointers, hw, VideoMode_GetCurrentVideoMode, "VideoMode_GetCurrentVideoMode");
        find!(pointers, hw, VideoMode_IsWindowed, "VideoMode_IsWindowed");

        pointers.window_rect = Some(hw.sym("window_rect")
                .chain_err(|| "couldn't find window_rect")? as _);
        pointers.s_BackBufferFBO = Some(hw.sym("s_BackBufferFBO")
                .chain_err(|| "couldn't find s_BackBufferFBO")? as _);
    }

    Ok(())
}

/// Resets all pointers to their default values.
fn reset_pointers() {
    *POINTERS.write().unwrap() = Pointers::default();
}

/// Registers console commands and variables.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
unsafe fn register_cvars_and_commands() {
    for cmd in command::COMMANDS.read().unwrap().iter() {
        register_command(cmd.name(), cmd.callback());
    }
}

/// Registers a console command.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
unsafe fn register_command(name: &'static [u8], callback: unsafe extern "C" fn()) {
    real!(Cmd_AddCommand)(name as *const _ as *const _, callback as *mut c_void);
}

/// Prints the given string to the game console.
///
/// `string` must not contain null bytes.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
pub unsafe fn con_print(string: &str) {
    real!(Con_Printf)(CString::new(string.replace('%', "%%"))
                          .expect("string cannot contain null bytes")
                          .as_ptr())
}

/// Gets the console command argument count.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
pub unsafe fn cmd_argc() -> u32 {
    let argc = real!(Cmd_Argc)();
    cmp::max(0, argc) as u32
}

/// Gets a console command argument.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
pub unsafe fn cmd_argv(index: u32) -> String {
    let index = cmp::min(index, i32::max_value() as u32) as i32;
    let arg = real!(Cmd_Argv)(index);
    CStr::from_ptr(arg).to_string_lossy().into_owned()
}

/// Returns the current game resolution.
///
/// # Safety
/// Unsafe because this function should only be called from the main game thread.
pub unsafe fn get_resolution() -> (u32, u32) {
    let mut width;
    let mut height;

    if real!(VideoMode_IsWindowed)() != 0 {
        let window_rect = *POINTERS.read().unwrap().window_rect.unwrap();
        width = window_rect.right - window_rect.left;
        height = window_rect.bottom - window_rect.top;
    } else {
        width = mem::uninitialized();
        height = mem::uninitialized();
        real!(VideoMode_GetCurrentVideoMode)(&mut width, &mut height, ptr::null_mut());
    }

    width = cmp::max(0, width);
    height = cmp::max(0, height);

    (width as u32, height as u32)
}

command!(cap_test, |engine| {
    let args = engine.args();

    let mut buf = String::new();
    buf.push_str(&format!("Args len: {}\n", args.len()));

    for arg in args {
        buf.push_str(&arg);
        buf.push('\n');
    }

    engine.con_print(&buf);
});

command!(cap_another_test, |engine| {
    engine.con_print("Hello! %s %d %\n");
});

command!(cap_start, |_engine| {
    *CAPTURING.write().unwrap() = true;
});

command!(cap_stop, |_engine| {
    *CAPTURING.write().unwrap() = false;
});
<|endoftext|>


<commit_before>
use serialize::json;
use serialize::json::{
    Boolean,
    F64,
    I64,
    Json,
    List,
    Null,
    String,
    U64,
};
use std::io;
use std::from_str::from_str;

/// An Event: [The event data](http://sensuapp.org/docs/latest/event_data)
/// formatted as a struct.
#[deriving(Show, PartialEq)]
pub struct Event {
    client: Client,
    check: Check,
    occurrences: u64,
    action: Option<String>,
}

/// The configuration of the client that sent this event
#[deriving(Show, PartialEq)]
pub struct Client {
    name: String,
    address: String,
    subscriptions: Vec<String>,
    timestamp: i64,
    additional: Json
}

/// The client-side definition of the check plus some extra status. See the
/// [check documentation](http://sensuapp.org/docs/latest/checks) for details.
#[deriving(Show, PartialEq)]
pub struct Check {
    // client-configured things
    name: String,
    command: String,
    subscribers: Vec<String>,
    interval: i64,
    additional: Option<Json>,

    // Results, added by Sensu
    handlers: Option<Vec<String>>,
    handler: Option<String>,
    issued: i64,
    output: String,
    status: i8,
    history: Vec<String>,
    flapping: bool,
}

#[deriving(Show, PartialEq)]
pub enum SensuError {
    InitError(String),
    EventError(String)
}

pub type SensuResult<T> = Result<T, SensuError>;

/// Extract a generic JS object, given a type
/// Early return an error if it can't be extracted
macro_rules! jk(
    ($j:ident->$k:expr $p:ident) => (
        match $j.find($k) {
            Some(value) => match value {
                &$p(ref v) => v.clone(),
                _ => return Err(EventError(format!("Wrong type for '{}': {}", $k, value)))
            },
            None => return Err(EventError(format!("couldn't find {} in event", $k)))
        }
    )
)
/// Extract any kind of number as an i64
/// Early return an error if it can't be extracted
macro_rules! jki {
    ($event:ident, $field:expr) => {
        match $event.find($field) {
            Some(value) => match value {
                &I64(ref v) => v.clone(),
                &U64(ref v) => v.clone() as i64,
                &F64(ref v) => v.clone() as i64,
                _ => return Err(EventError(format!("Wrong type for '{}': {}", $field, value)))
            },
            None => return Err(EventError(format!("couldn't find '{}' in event", $field)))
        };
    }
}

/// Read stdin, and parse it into an Event
pub fn init() -> SensuResult<Event> {
    match io::stdin().read_to_end() {
        Ok(input) => read_event(input.to_string().as_slice()),
        Err(e) => Err(InitError(format!("couldn't read stdin: {}", e)))
    }
}

/// Parse an &str into an Event, making sure that all the types match up
pub fn read_event(input: &str) -> SensuResult<Event> {
    let event = match json::from_str(input) {
        Ok(v) => v,
        Err(e) => return Err(InitError(format!("couldn't parse json: {}", e)))
    };

    let client = try!(read_client(&event));
    let check = try!(read_check(&event));
    let action = match event.find("action") {
        Some(action) => match *action {
            String(ref a) => Some(a.clone()),
            _ => return Err(InitError(format!("Wrong type for action: {}", action)))
        },
        None => None
    };
    let occurrences = match event.find("occurrences") {
        Some(occurrences) => match *occurrences {
            U64(o) => o,
            I64(o) => o as u64,
            F64(o) => o as u64,
            String(ref o) => try!(from_str::<u64>(o.as_slice()).ok_or(
                InitError(format!("Unable to parse number from occurrences: {}", occurrences)))),
            _ => return Err(InitError(format!("Wrong type for occurrences: {}", occurrences)))
        },
        None => return Err(InitError(format!("No occurrences in event data")))
    };
    Ok(Event{
        client: client,
        check: check,
        action: action,
        occurrences: occurrences,
    })
}

fn extract_list_of_strings(list: &Json, field_name: &str) -> SensuResult<Vec<String>> {
    match *list {
        List(ref subs) => {
            let mut result = Vec::new();
            for sub in subs.iter() {
                match *sub {
                    String(ref s) => result.push(s.clone()),
                    _ => return Err(InitError(format!("Wrong type in {}: {}",
                                                      field_name, sub)))
                }
            }
            Ok(result)
        },
        _ => return Err(InitError(format!("{} is not a list ({})", field_name, list)))
    }
}

fn find_list_of_strings(obj: &Json, field_name: &str) -> SensuResult<Vec<String>> {
    let contained = obj.find(field_name);
    println!("{} is {}", field_name, contained)
    let result = match contained {
        Some(subs) => {
            try!(extract_list_of_strings(subs, field_name))
        },
        None => return Err(InitError(format!("Couldn't build {} for check", field_name)))
    };
    Ok(result)
}

fn read_client(input: &Json) -> SensuResult<Client> {
    let event = match input.find("client") {
        Some(client) => client,
        None => return Err(InitError("No client in event".into_string()))
    };

    let unverified_subs = jk!(event->"subscriptions" List);
    let mut subs: Vec<String> = Vec::new();
    for sub in unverified_subs.iter() {
        match *sub {
            json::String(ref s) => subs.push(s.clone()),
            _ => return Err(EventError(format!("Invalid subscription type {}", sub)))
        }
    }

    let client = Client {
        name: jk!(event->"name" String),
        address: jk!(event->"address" String),
        subscriptions: subs,
        timestamp: jki!(event, "timestamp"),
        additional: Null
    };
    Ok(client)
}

fn read_check(event: &Json) -> SensuResult<Check> {
    let check = match event.find("check") {
        Some(check) => check,
        None => return Err(InitError("No check in event".into_string()))
    };

    let subscribers = try!(find_list_of_strings(check, "subscribers"));

    Ok(Check {
        name: jk!(check->"name" String),
        issued: jki!(check, "issued"),
        output: jk!(check->"output" String),
        status: jki!(check, "status") as i8,
        command: jk!(check->"command" String),
        subscribers: subscribers,
        interval: jki!(check, "interval"),
        handler: match check.find("handler") {
            Some(handler) => match *handler {
                String(ref h) => Some(h.clone()),
                _ => return Err(InitError(format!("Wrong type for handler ({})", handler)))
            },
            None => None
        },
        handlers: match check.find("handlers") {
            Some(handlers) => Some(try!(extract_list_of_strings(handlers, "handlers"))),
            None => None
        },
        history: try!(find_list_of_strings(check, "history")),
        flapping: jk!(check->"flapping" Boolean),
        additional: None
    })
}

#[cfg(test)]
mod build_objects {
    use serialize::json;
    use super::{read_event, read_check, read_client, Check};

    #[test]
    fn can_build_event() {
        let event = r#"{
            "occurrences": 1,
            "action": "create",
            "client": {
                "name": "hello",
                "address": "192.168.1.1",
                "subscriptions": ["one", "two"],
                "timestamp": 127897
            },
            "check": {
                "name": "test-check",
                "issued": 1416069607,
                "output": "we have output",
                "status": 0,
                "command": "echo 'we have output'",
                "subscribers": ["examples", "tests"],
                "interval": 60,
                "handler": "default",
                "history": ["0", "0", "0"],
                "flapping": false,
                "additional": null
            }
        }"#;
        match read_event(event) {
            Ok(e) => println!("Parsed event: {}", e),
            Err(e) => panic!("ERROR: {}", e)
        };
    }

    #[test]
    fn can_build_client() {
        let event = json::from_str(r#"{
            "client": {
                "name": "hello",
                "address": "192.168.1.1",
                "subscriptions": ["one", "two"],
                "timestamp": 127897
            }
        }"#).unwrap();
        let client = read_client(&event);
        match client {
            Ok(cl) => println!("Parsed: {}", cl),
            Err(e) => panic!("ERROR: {}", e)
        }
    }

    #[test]
    fn can_build_check() {
        let event = json::from_str(r#"{
            "check": {
                "name": "test-check",
                "issued": 1416069607,
                "output": "we have output",
                "status": 0,
                "command": "echo 'we have output'",
                "subscribers": ["examples", "tests"],
                "interval": 60,
                "handler": "default",
                "history": ["0", "0", "0"],
                "flapping": false,
                "additional": null
            }
        }"#).unwrap();
        match read_check(&event) {
            Ok(check) => {
                println!("Parsed: {}", check);
                assert_eq!(check, Check {
                    name: "test-check".into_string(),
                    issued: 1416069607i64,
                    output: "we have output".into_string(),
                    status: 0,
                    command: "echo 'we have output'".into_string(),
                    subscribers: vec!("examples".into_string(),
                                      "tests".into_string()),
                    interval: 60,
                    handler: Some("default".into_string()),
                    handlers: None,
                    history: Vec::from_fn(3, |_| "0".into_string()),
                    flapping: false,
                    additional: None });
            },
            Err(e) => panic!("ERROR: {}", e)
        }
    }

    #[test]
    fn can_build_check_with_handlers() {
        let event = json::from_str(r#"{
            "check": {
                "name": "test-check",
                "issued": 1416069607,
                "output": "we have output",
                "status": 0,
                "command": "echo 'we have output'",
                "subscribers": ["examples", "tests"],
                "interval": 60,
                "handlers": ["default", "pagerduty"],
                "history": ["0", "0", "0"],
                "flapping": false,
                "additional": null
            }
        }"#).unwrap();
        match read_check(&event) {
            Ok(cl) => println!("Parsed: {}", cl),
            Err(e) => panic!("ERROR: {}", e)
        }
    }
}
<commit_msg>Add defaults to a few fields<commit_after>
use serialize::json;
use serialize::json::{
    Boolean,
    F64,
    I64,
    Json,
    List,
    Null,
    String,
    U64,
};
use std::io;
use std::from_str::from_str;

struct Defaults {
    refresh: u64,
    interval: u64
}

static DEFAULTS: Defaults = Defaults {
    refresh: 1800,
    interval: 60
};

/// An Event: [The event data](http://sensuapp.org/docs/latest/event_data)
/// formatted as a struct.
#[deriving(Show, PartialEq)]
pub struct Event {
    client: Client,
    check: Check,
    occurrences: u64,
    action: Option<String>,
}

/// The configuration of the client that sent this event
#[deriving(Show, PartialEq)]
pub struct Client {
    name: String,
    address: String,
    subscriptions: Vec<String>,
    timestamp: i64,
    additional: Json
}

/// The client-side definition of the check plus some extra status. See the
/// [check documentation](http://sensuapp.org/docs/latest/checks) for details.
#[deriving(Show, PartialEq)]
pub struct Check {
    // client-configured things
    name: String,
    command: String,
    subscribers: Vec<String>,
    interval: i64,
    additional: Option<Json>,

    // Results, added by Sensu
    handlers: Option<Vec<String>>,
    handler: Option<String>,
    issued: i64,
    output: String,
    status: i8,
    history: Vec<String>,
    flapping: bool,
}

#[deriving(Show, PartialEq)]
pub enum SensuError {
    InitError(String),
    EventError(String)
}

pub type SensuResult<T> = Result<T, SensuError>;

/// Extract a generic JS object, given a type
/// Early return an error if it can't be extracted
macro_rules! jk(
    ($j:ident->$k:expr $p:ident) => (
        match $j.find($k) {
            Some(value) => match value {
                &$p(ref v) => v.clone(),
                _ => return Err(EventError(format!("Wrong type for '{}': {}", $k, value)))
            },
            None => return Err(EventError(format!("couldn't find {} in event", $k)))
        }
    )
)

macro_rules! find_default(
    ($j:ident, $p:ident($k:expr), $default:expr) => (
        match $j.find($k) {
            Some(value) => match value {
                &$p(ref v) => v.clone(),
                _ => return Err(EventError(format!("Wrong type for '{}': {}", $k, value)))
            },
            None => $default
        }
    )
)

/// Extract any kind of number as an i64
/// Early return an error if it can't be extracted
macro_rules! jki {
    ($event:ident, $field:expr) => {
        match $event.find($field) {
            Some(value) => match value {
                &I64(ref v) => v.clone(),
                &U64(ref v) => v.clone() as i64,
                &F64(ref v) => v.clone() as i64,
                _ => return Err(EventError(format!("Wrong type for '{}': {}", $field, value)))
            },
            None => return Err(EventError(format!("couldn't find '{}' in event", $field)))
        };
    }
}

macro_rules! find {
    ($event:ident, $pat:ident($field:expr)) => (
        match $event.find($field) {
            Some(value) => match *value {
                $pat(v) => Some(v),
                _ => return Err(EventError(format!("Wrong type for '{}': '{}'", $field, value)))
            },
            None => None
        }
    )
}

/// Read stdin, and parse it into an Event
pub fn init() -> SensuResult<Event> {
    match io::stdin().read_to_end() {
        Ok(input) => read_event(input.to_string().as_slice()),
        Err(e) => Err(InitError(format!("couldn't read stdin: {}", e)))
    }
}

/// Parse an &str into an Event, making sure that all the types match up
pub fn read_event(input: &str) -> SensuResult<Event> {
    let event = match json::from_str(input) {
        Ok(v) => v,
        Err(e) => return Err(InitError(format!("couldn't parse json: {}", e)))
    };

    let client = try!(read_client(&event));
    let check = try!(read_check(&event));
    let action = match event.find("action") {
        Some(action) => match *action {
            String(ref a) => Some(a.clone()),
            _ => return Err(InitError(format!("Wrong type for action: {}", action)))
        },
        None => None
    };
    let occurrences = match event.find("occurrences") {
        Some(occurrences) => match *occurrences {
            U64(o) => o,
            I64(o) => o as u64,
            F64(o) => o as u64,
            String(ref o) => try!(from_str::<u64>(o.as_slice()).ok_or(
                InitError(format!("Unable to parse number from occurrences: {}", occurrences)))),
            _ => return Err(InitError(format!("Wrong type for occurrences: {}", occurrences)))
        },
        None => return Err(InitError(format!("No occurrences in event data")))
    };
    Ok(Event{
        client: client,
        check: check,
        action: action,
        occurrences: occurrences,
    })
}

fn extract_list_of_strings(list: &Json, field_name: &str) -> SensuResult<Vec<String>> {
    match *list {
        List(ref subs) => {
            let mut result = Vec::new();
            for sub in subs.iter() {
                match *sub {
                    String(ref s) => result.push(s.clone()),
                    _ => return Err(InitError(format!("Wrong type in {}: {}",
                                                      field_name, sub)))
                }
            }
            Ok(result)
        },
        _ => return Err(InitError(format!("{} is not a list ({})", field_name, list)))
    }
}

fn find_list_of_strings(obj: &Json, field_name: &str) -> SensuResult<Vec<String>> {
    let contained = obj.find(field_name);
    println!("{} is {}", field_name, contained)
    let result = match contained {
        Some(subs) => {
            try!(extract_list_of_strings(subs, field_name))
        },
        None => return Err(InitError(format!("Couldn't build {} for check", field_name)))
    };
    Ok(result)
}

fn read_client(input: &Json) -> SensuResult<Client> {
    let event = match input.find("client") {
        Some(client) => client,
        None => return Err(InitError("No client in event".into_string()))
    };

    let unverified_subs = jk!(event->"subscriptions" List);
    let mut subs: Vec<String> = Vec::new();
    for sub in unverified_subs.iter() {
        match *sub {
            json::String(ref s) => subs.push(s.clone()),
            _ => return Err(EventError(format!("Invalid subscription type {}", sub)))
        }
    }

    let client = Client {
        name: jk!(event->"name" String),
        address: jk!(event->"address" String),
        subscriptions: subs,
        timestamp: jki!(event, "timestamp"),
        additional: Null
    };
    Ok(client)
}

fn read_check(event: &Json) -> SensuResult<Check> {
    let check = match event.find("check") {
        Some(check) => check,
        None => return Err(InitError("No check in event".into_string()))
    };

    let subscribers = try!(find_list_of_strings(check, "subscribers"));

    Ok(Check {
        name: jk!(check->"name" String),
        issued: jki!(check, "issued"),
        output: jk!(check->"output" String),
        status: jki!(check, "status") as i8,
        command: jk!(check->"command" String),
        subscribers: subscribers,

        alert: find!(check, Boolean("alert")),
        occurrences: find!(check, U64("occurrences")),

        interval: find_default!(check, U64("interval"), DEFAULTS.interval),
        refresh: find_default!(check, U64("refresh"), DEFAULTS.refresh),

        handler: match check.find("handler") {
            Some(handler) => match *handler {
                String(ref h) => Some(h.clone()),
                _ => return Err(InitError(format!("Wrong type for handler ({})", handler)))
            },
            None => None
        },
        handlers: match check.find("handlers") {
            Some(handlers) => Some(try!(extract_list_of_strings(handlers, "handlers"))),
            None => None
        },
        history: try!(find_list_of_strings(check, "history")),
        flapping: jk!(check->"flapping" Boolean),
        additional: None
    })
}

#[cfg(test)]
mod build_objects {
    use serialize::json;
    use super::{read_event, read_check, read_client, Check};

    #[test]
    fn can_build_event() {
        let event = r#"{
            "occurrences": 1,
            "action": "create",
            "client": {
                "name": "hello",
                "address": "192.168.1.1",
                "subscriptions": ["one", "two"],
                "timestamp": 127897
            },
            "check": {
                "name": "test-check",
                "issued": 1416069607,
                "output": "we have output",
                "status": 0,
                "command": "echo 'we have output'",
                "subscribers": ["examples", "tests"],
                "interval": 60,
                "handler": "default",
                "history": ["0", "0", "0"],
                "flapping": false,
                "additional": null
            }
        }"#;
        match read_event(event) {
            Ok(e) => println!("Parsed event: {}", e),
            Err(e) => panic!("ERROR: {}", e)
        };
    }

    #[test]
    fn can_build_client() {
        let event = json::from_str(r#"{
            "client": {
                "name": "hello",
                "address": "192.168.1.1",
                "subscriptions": ["one", "two"],
                "timestamp": 127897
            }
        }"#).unwrap();
        let client = read_client(&event);
        match client {
            Ok(cl) => println!("Parsed: {}", cl),
            Err(e) => panic!("ERROR: {}", e)
        }
    }

    #[test]
    fn can_build_check() {
        let event = json::from_str(r#"{
            "check": {
                "name": "test-check",
                "issued": 1416069607,
                "output": "we have output",
                "status": 0,
                "command": "echo 'we have output'",
                "subscribers": ["examples", "tests"],
                "interval": 60,
                "refresh": 1800,
                "handler": "default",
                "history": ["0", "0", "0"],
                "flapping": false,
                "additional": null
            }
        }"#).unwrap();
        match read_check(&event) {
            Ok(check) => {
                println!("Parsed: {}", check);
                assert_eq!(check, Check {
                    name: "test-check".into_string(),
                    issued: 1416069607i64,
                    output: "we have output".into_string(),
                    status: 0,
                    command: "echo 'we have output'".into_string(),
                    subscribers: vec!("examples".into_string(),
                                      "tests".into_string()),
                    alert: None,
                    occurrences: None,
                    interval: 60,
                    refresh: 1800,
                    handler: Some("default".into_string()),
                    handlers: None,
                    history: Vec::from_fn(3, |_| "0".into_string()),
                    flapping: false,
                    additional: None });
            },
            Err(e) => panic!("ERROR: {}", e)
        }
    }

        #[test]
    fn can_build_check_with_alert() {
        let event = json::from_str(r#"{
            "check": {
                "name": "test-check",
                "issued": 1416069607,
                "output": "we have output",
                "status": 0,
                "command": "echo 'we have output'",
                "subscribers": ["examples", "tests"],
                "interval": 60,
                "refresh": 1800,
                "alert": false,
                "occurrences": 5,
                "handler": "default",
                "history": ["0", "0", "0"],
                "flapping": false,
                "additional": null
            }
        }"#).unwrap();
        match read_check(&event) {
            Ok(check) => {
                println!("Parsed: {}", check);
                assert_eq!(check, Check {
                    name: "test-check".into_string(),
                    issued: 1416069607i64,
                    output: "we have output".into_string(),
                    status: 0,
                    command: "echo 'we have output'".into_string(),
                    subscribers: vec!("examples".into_string(),
                                      "tests".into_string()),

                    alert: Some(false),
                    occurrences: Some(5),

                    refresh: 1800,
                    interval: 60,
                    handler: Some("default".into_string()),
                    handlers: None,
                    history: Vec::from_fn(3, |_| "0".into_string()),
                    flapping: false,
                    additional: None });
            },
            Err(e) => panic!("ERROR: {}", e)
        }
    }

    #[test]
    fn can_build_check_with_handlers() {
        let event = json::from_str(r#"{
            "check": {
                "name": "test-check",
                "issued": 1416069607,
                "output": "we have output",
                "status": 0,
                "command": "echo 'we have output'",
                "subscribers": ["examples", "tests"],
                "interval": 60,
                "handlers": ["default", "pagerduty"],
                "history": ["0", "0", "0"],
                "flapping": false,
                "additional": null
            }
        }"#).unwrap();
        match read_check(&event) {
            Ok(cl) => println!("Parsed: {}", cl),
            Err(e) => panic!("ERROR: {}", e)
        }
    }
}
<|endoftext|>


<commit_before>
//! Contains internal type aliases and traits.
//!
//! These do not affect the normal usage of the library.
//! Using these might lead to breaking changes.
//! Made public to clarify documentation.

pub use vecmath::{ Scalar, Vec2d };

/// The type used for area.
pub type Area = Scalar;

/// [red, green, blue, alpha]
pub type Color = [ColorComponent, ..4];

/// The type used for color component.
pub type ColorComponent = f32;

/// [x1, y1, x2, y2]
pub type Line = [Scalar, ..4];

/// [x, y, w, h]
pub type SourceRectangle = [i32, ..4];

/// [p0, p1, ...]
pub type Polygon<'a> = &'a [Vec2d];

/// A slice of polygons.
pub type Polygons<'a> = &'a [Polygon<'a>];

/// The type used for radius.
pub type Radius = Scalar;

/// [x, y, dir_x, dir_y]
pub type Ray = [Scalar, ..4];

/// [x, y, w, h]
pub type Rectangle = [Scalar, ..4];

/// [x1, y1, x2, y2, x3, y3]
pub type Triangle = [Vec2d, ..3];

/// The type used for width.
pub type Width = Scalar;
<commit_msg>Fixed doc comment in internal module<commit_after>
//! Contains type aliases used in this library

pub use vecmath::{ Scalar, Vec2d };

/// The type used for area.
pub type Area = Scalar;

/// [red, green, blue, alpha]
pub type Color = [ColorComponent, ..4];

/// The type used for color component.
pub type ColorComponent = f32;

/// [x1, y1, x2, y2]
pub type Line = [Scalar, ..4];

/// [x, y, w, h]
pub type SourceRectangle = [i32, ..4];

/// [p0, p1, ...]
pub type Polygon<'a> = &'a [Vec2d];

/// A slice of polygons.
pub type Polygons<'a> = &'a [Polygon<'a>];

/// The type used for radius.
pub type Radius = Scalar;

/// [x, y, dir_x, dir_y]
pub type Ray = [Scalar, ..4];

/// [x, y, w, h]
pub type Rectangle = [Scalar, ..4];

/// [x1, y1, x2, y2, x3, y3]
pub type Triangle = [Vec2d, ..3];

/// The type used for width.
pub type Width = Scalar;
<|endoftext|>


<commit_before>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! FastQ reading and writing.

use std::io;
use std::io::prelude::*;
use std::ascii::AsciiExt;

/// A FastQ reader.
pub struct FastqReader<R: io::Read> {
    reader: io::BufReader<R>,
    sep_line: String
}

impl<R: io::Read> FastqReader<R> {
    /// Create a new FastQ reader.
    pub fn new(reader: R) -> Self {
        FastqReader { reader: io::BufReader::new(reader), sep_line: String::new() }
    }

    /// Read into a given record.
    /// Returns an error if the record in incomplete or syntax is violated.
    /// The content of the record can be checked via the record object.
    pub fn read(&mut self, record: &mut Record) -> io::Result<()> {
        record.clear();
        try!(self.reader.read_line(&mut record.header));

        if !record.header.is_empty() {
            if !record.header.starts_with("@") {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "Expected @ at record start.",
                    None,
                ));
            }
            try!(self.reader.read_line(&mut record.seq));
            try!(self.reader.read_line(&mut self.sep_line));
            try!(self.reader.read_line(&mut record.qual));
            if record.qual.is_empty() {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "Incomplete record.",
                    Some("Each FastQ record has to consist of 4 lines: header, sequence, separator and qualities.".to_string()),
                ))
            }
        }

        Ok(())
    }

    /// Return an iterator over the records of this FastQ file.
    pub fn records(self) -> Records<R> {
        Records { reader: self }
    }
}

/// A FastQ record.
pub struct Record {
    header: String,
    seq: String,
    qual: String,
}

impl Record {
    /// Create a new, empty FastQ record.
    pub fn new() -> Self {
        Record {
            header: String::new(),
            seq: String::new(),
            qual: String::new(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.header.is_empty() && self.seq.is_empty() && self.qual.is_empty()
    }

    /// Check validity of FastQ record.
    pub fn check(&self) -> Result<(), &str> {
        if self.id().is_none() {
            return Err("Expecting id for FastQ record.");
        }
        if !self.seq.is_ascii() {
            return Err("Non-ascii character found in sequence.");
        }
        if !self.qual.is_ascii() {
            return Err("Non-ascii character found in qualities.");
        }
        if self.seq().len() != self.qual().len() {
            return Err("Unequal length of sequence an qualities.");
        }

        Ok(())
    }

    /// Return the id of the record.
    pub fn id(&self) -> Option<&str> {
        self.header[1..].words().next()
    }

    /// Return descriptions if present.
    pub fn desc(&self) -> Vec<&str> {
        self.header[1..].words().skip(1).collect()
    }

    /// Return the sequence of the record.
    pub fn seq(&self) -> &[u8] {
        self.seq.trim_right().as_bytes()
    }

    /// Return the base qualities of the record.
    pub fn qual(&self) -> &[u8] {
        self.qual.trim_right().as_bytes()
    }

    fn clear(&mut self) {
        self.header.clear();
        self.seq.clear();
        self.qual.clear();
    }
}

/// An iterator over the records of a FastQ file.
pub struct Records<R: io::Read> {
    reader: FastqReader<R>,
}

impl<R: io::Read> Iterator for Records<R> {
    type Item = io::Result<Record>;

    fn next(&mut self) -> Option<io::Result<Record>> {
        let mut record = Record::new();
        match self.reader.read(&mut record) {
            Ok(()) if record.is_empty() => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

/// A FastQ writer.
pub struct FastqWriter<W: io::Write> {
    writer: io::BufWriter<W>,
}

impl<W: io::Write> FastqWriter<W> {
    /// Create a new FastQ writer.
    pub fn new(writer: W) -> Self {
        FastqWriter { writer: io::BufWriter::new(writer) }
    }

    /// Directly write a FastQ record.
    pub fn write_record(&mut self, record: Record) -> io::Result<()> {
        self.write(record.id().unwrap_or(""), &record.desc(), record.seq(), record.qual())
    }

    /// Write a FastQ record with given values.
    ///
    /// # Arguments
    ///
    /// * `id` - the record id
    /// * `desc` - the optional descriptions
    /// * `seq` - the sequence
    /// * `qual` - the qualities
    pub fn write(&mut self, id: &str, desc: &[&str], seq: &[u8], qual: &[u8]) -> io::Result<()> {
        try!(self.writer.write(b"@ "));
        try!(self.writer.write(id.as_bytes()));
        if !desc.is_empty() {
            for d in desc {
                try!(self.writer.write(b" "));
                try!(self.writer.write(d.as_bytes()));
            }
        }
        try!(self.writer.write(b"\n"));
        try!(self.writer.write(seq));
        try!(self.writer.write(b"\n+\n"));
        try!(self.writer.write(qual));
        try!(self.writer.write(b"\n"));

        Ok(())
    }

    /// Flush the writer, ensuring that everything is written.
    pub fn flush(&mut self) -> io::Result<()> {
        self.writer.flush()
    }
}

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

    const FASTQ_FILE: &'static [u8] = b"@ id desc
ACCGTAGGCTGA
+
IIIIIIJJJJJJ
";

    #[test]
    fn test_reader() {
        let reader = FastqReader::new(FASTQ_FILE);
        let records: Vec<io::Result<Record>> = reader.records().collect();
        assert!(records.len() == 1);
        for res in records {
            let record = res.ok().unwrap();
            assert_eq!(record.check(), Ok(()));
            assert_eq!(record.id(), Some("id"));
            assert_eq!(record.desc(), ["desc"]);
            assert_eq!(record.seq(), b"ACCGTAGGCTGA");
            assert_eq!(record.qual(), b"IIIIIIJJJJJJ");
        }
    }

    #[test]
    fn test_writer() {
        let mut writer = FastqWriter::new(Vec::new());
        writer.write("id", &["desc"], b"ACCGTAGGCTGA", b"IIIIIIJJJJJJ").ok().expect("Expected successful write");
        writer.flush().ok().expect("Expected successful write");
        assert_eq!(writer.writer.get_ref(), &FASTQ_FILE);
    }
}
<commit_msg>Minor.<commit_after>// Copyright 2014 Johannes Köster.
// Licensed under the MIT license (http://opensource.org/licenses/MIT)
// This file may not be copied, modified, or distributed
// except according to those terms.

//! FastQ reading and writing.

use std::io;
use std::io::prelude::*;
use std::ascii::AsciiExt;

/// A FastQ reader.
pub struct FastqReader<R: io::Read> {
    reader: io::BufReader<R>,
    sep_line: String
}

impl<R: io::Read> FastqReader<R> {
    /// Create a new FastQ reader.
    pub fn new(reader: R) -> Self {
        FastqReader { reader: io::BufReader::new(reader), sep_line: String::new() }
    }

    /// Read into a given record.
    /// Returns an error if the record in incomplete or syntax is violated.
    /// The content of the record can be checked via the record object.
    pub fn read(&mut self, record: &mut Record) -> io::Result<()> {
        record.clear();
        try!(self.reader.read_line(&mut record.header));

        if !record.header.is_empty() {
            if !record.header.starts_with("@") {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "Expected @ at record start.",
                    None,
                ));
            }
            try!(self.reader.read_line(&mut record.seq));
            try!(self.reader.read_line(&mut self.sep_line));
            try!(self.reader.read_line(&mut record.qual));
            if record.qual.is_empty() {
                return Err(io::Error::new(
                    io::ErrorKind::Other,
                    "Incomplete record.",
                    Some("Each FastQ record has to consist of 4 lines: header, sequence, separator and qualities.".to_string()),
                ))
            }
        }

        Ok(())
    }

    /// Return an iterator over the records of this FastQ file.
    pub fn records(self) -> Records<R> {
        Records { reader: self }
    }
}

/// A FastQ record.
pub struct Record {
    header: String,
    seq: String,
    qual: String,
}

impl Record {
    /// Create a new, empty FastQ record.
    pub fn new() -> Self {
        Record {
            header: String::new(),
            seq: String::new(),
            qual: String::new(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.header.is_empty() && self.seq.is_empty() && self.qual.is_empty()
    }

    /// Check validity of FastQ record.
    pub fn check(&self) -> Result<(), &str> {
        if self.id().is_none() {
            return Err("Expecting id for FastQ record.");
        }
        if !self.seq.is_ascii() {
            return Err("Non-ascii character found in sequence.");
        }
        if !self.qual.is_ascii() {
            return Err("Non-ascii character found in qualities.");
        }
        if self.seq().len() != self.qual().len() {
            return Err("Unequal length of sequence an qualities.");
        }

        Ok(())
    }

    /// Return the id of the record.
    pub fn id(&self) -> Option<&str> {
        self.header[1..].words().next()
    }

    /// Return descriptions if present.
    pub fn desc(&self) -> Vec<&str> {
        self.header[1..].words().skip(1).collect()
    }

    /// Return the sequence of the record.
    pub fn seq(&self) -> &[u8] {
        self.seq.trim_right().as_bytes()
    }

    /// Return the base qualities of the record.
    pub fn qual(&self) -> &[u8] {
        self.qual.trim_right().as_bytes()
    }

    fn clear(&mut self) {
        self.header.clear();
        self.seq.clear();
        self.qual.clear();
    }
}

/// An iterator over the records of a FastQ file.
pub struct Records<R: io::Read> {
    reader: FastqReader<R>,
}

impl<R: io::Read> Iterator for Records<R> {
    type Item = io::Result<Record>;

    fn next(&mut self) -> Option<io::Result<Record>> {
        let mut record = Record::new();
        match self.reader.read(&mut record) {
            Ok(()) if record.is_empty() => None,
            Ok(())   => Some(Ok(record)),
            Err(err) => Some(Err(err))
        }
    }
}

/// A FastQ writer.
pub struct FastqWriter<W: io::Write> {
    writer: io::BufWriter<W>,
}

impl<W: io::Write> FastqWriter<W> {
    /// Create a new FastQ writer.
    pub fn new(writer: W) -> Self {
        FastqWriter { writer: io::BufWriter::new(writer) }
    }

    /// Directly write a FastQ record.
    pub fn write_record(&mut self, record: Record) -> io::Result<()> {
        self.write(record.id().unwrap_or(""), &record.desc(), record.seq(), record.qual())
    }

    /// Write a FastQ record with given values.
    ///
    /// # Arguments
    ///
    /// * `id` - the record id
    /// * `desc` - the optional descriptions
    /// * `seq` - the sequence
    /// * `qual` - the qualities
    pub fn write(&mut self, id: &str, desc: &[&str], seq: &[u8], qual: &[u8]) -> io::Result<()> {
        try!(self.writer.write(b"@"));
        try!(self.writer.write(id.as_bytes()));
        if !desc.is_empty() {
            for d in desc {
                try!(self.writer.write(b" "));
                try!(self.writer.write(d.as_bytes()));
            }
        }
        try!(self.writer.write(b"\n"));
        try!(self.writer.write(seq));
        try!(self.writer.write(b"\n+\n"));
        try!(self.writer.write(qual));
        try!(self.writer.write(b"\n"));

        Ok(())
    }

    /// Flush the writer, ensuring that everything is written.
    pub fn flush(&mut self) -> io::Result<()> {
        self.writer.flush()
    }
}

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

    const FASTQ_FILE: &'static [u8] = b"@id desc
ACCGTAGGCTGA
+
IIIIIIJJJJJJ
";

    #[test]
    fn test_reader() {
        let reader = FastqReader::new(FASTQ_FILE);
        let records: Vec<io::Result<Record>> = reader.records().collect();
        assert!(records.len() == 1);
        for res in records {
            let record = res.ok().unwrap();
            assert_eq!(record.check(), Ok(()));
            assert_eq!(record.id(), Some("id"));
            assert_eq!(record.desc(), ["desc"]);
            assert_eq!(record.seq(), b"ACCGTAGGCTGA");
            assert_eq!(record.qual(), b"IIIIIIJJJJJJ");
        }
    }

    #[test]
    fn test_writer() {
        let mut writer = FastqWriter::new(Vec::new());
        writer.write("id", &["desc"], b"ACCGTAGGCTGA", b"IIIIIIJJJJJJ").ok().expect("Expected successful write");
        writer.flush().ok().expect("Expected successful write");
        assert_eq!(writer.writer.get_ref(), &FASTQ_FILE);
    }
}
<|endoftext|>


<commit_before>use std::io::IoResult;
use std::task::TaskBuilder;
use std::collections::HashSet;

use url::{
    Url, RelativeScheme, SchemeType,
    Domain, Ipv6, NonRelativeScheme,
    ParseResult, UrlParser
};

use irc::{IrcConnection, IrcEventMessage};
use irc::connection::RawWrite;

use command_mapper::PluginContainer;

use state::State;
use plugins::{
    DeerPlugin,
    GreedPlugin,
    SeenPlugin,
    RadioPlugin,
    PingPlugin,
    WserverPlugin,
    WhoAmIPlugin,
};

#[deriving(Decodable, Encodable, Show)]
pub struct BotConfig {
    pub server: String,
    pub command_prefix: String,
    pub nickname: String,
    pub channels: Vec<String>,
    pub enabled_plugins: HashSet<String>,
}

pub fn irc_scheme_type_mapper(scheme: &str) -> SchemeType {
    match scheme {
        "irc" => RelativeScheme(6667),
        "ircs" => RelativeScheme(6697),
        _ => NonRelativeScheme,
    }
}

impl BotConfig {
    fn get_url(&self) -> ParseResult<Url> {
        let mut parser = UrlParser::new();
        parser.scheme_type_mapper(irc_scheme_type_mapper);
        parser.parse(self.server.as_slice())
    }

    fn get_host(&self) -> String {
        let server = self.get_url().unwrap();
        match server.host() {
            Some(&Domain(ref string)) => string.clone(),
            Some(&Ipv6(ref addr)) => addr.serialize(),
            None => panic!()
        }
    }

    fn get_port(&self) -> u16 {
        let server = self.get_url().unwrap();
        server.port().unwrap_or(6667)
    }
}

pub struct BotConnection {
    foo: uint
    //
}

impl BotConnection {
    pub fn new(conf: &BotConfig) -> IoResult<BotConnection> {
        let (mut conn, event_queue) = try!(IrcConnection::new(
            (conf.get_host().as_slice(), conf.get_port())));

        let (event_queue_txu, event_queue_rxu) = channel();
        spawn(proc() {
            for event in event_queue.iter() {
                event_queue_txu.send(event);
            }
        });

        let mut nick = conf.nickname.clone();
        loop {
            println!("trying nick {}", nick.as_slice());
            match conn.register(nick.as_slice()) {
                Ok(_) => {
                    println!("ok, connected as {}", nick.as_slice());
                    break;
                }
                Err(err) => {
                    if err.should_pick_new_nickname() {
                        nick.push_str("`");
                    } else {
                        panic!("{}", err)
                    }
                }
            };
        }

        for channel in conf.channels.iter() {
            println!("want join: {}", channel);
        }

        for channel in conf.channels.iter() {
            println!("joining {}...", channel);
            match conn.join(channel.as_slice()) {
                Ok(res) => {
                    println!("succeeded in joining {}, got {} nicks",
                        res.channel.as_slice(), res.nicks.len());
                    match conn.who(channel.as_slice()) {
                        Ok(who_res) => {
                            println!("succeeded in WHOing {}, got {} nicks",
                                who_res.channel.as_slice(), who_res.who_records.len());
                        },
                        Err(who_err) => {
                            println!("failed to WHO {}: {}", channel, who_err);
                        }
                    }
                },
                Err(err) => {
                    println!("join error: {}", err);
                    panic!("failed to join channel.. dying");
                }
            }
            println!("END joining {}...", channel);
        }

        let mut state = State::new();

        let mut container = PluginContainer::new(conf.command_prefix.clone());
        if conf.enabled_plugins.contains_equiv(PingPlugin::get_plugin_name()) {
            container.register(box PingPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(GreedPlugin::get_plugin_name()) {
            container.register(box GreedPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(SeenPlugin::get_plugin_name()) {
            container.register(box SeenPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(DeerPlugin::get_plugin_name()) {
            container.register(box DeerPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(RadioPlugin::get_plugin_name()) {
            container.register(box RadioPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(WserverPlugin::get_plugin_name()) {
            container.register(box WserverPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(WhoAmIPlugin::get_plugin_name()) {
            container.register(box WhoAmIPlugin::new());
        }

        let (tx, rx) = sync_channel(0);
        let cmd_queue = conn.get_command_queue();

        TaskBuilder::new().named("bot-sender").spawn(proc() {
            for message in rx.iter() {
                cmd_queue.send(RawWrite(message));
            }
        });

        for event in event_queue_rxu.iter() {
            state.on_event(&event);
            if let IrcEventMessage(ref message) = event {
                container.dispatch(&state, &tx, message);
            }
        }

        Ok(BotConnection { foo: 0})
    }
}
<commit_msg>Some testing code to see if we desync DONOTMERGE<commit_after>use std::io::IoResult;
use std::task::TaskBuilder;
use std::collections::HashSet;
use std::io::timer::sleep;
use std::time::duration::Duration;

use url::{
    Url, RelativeScheme, SchemeType,
    Domain, Ipv6, NonRelativeScheme,
    ParseResult, UrlParser
};

use irc::IrcMessage;
use irc::{IrcConnection, IrcEventMessage};
use irc::connection::RawWrite;

use command_mapper::PluginContainer;

use state::State;
use plugins::{
    DeerPlugin,
    GreedPlugin,
    SeenPlugin,
    RadioPlugin,
    PingPlugin,
    WserverPlugin,
    WhoAmIPlugin,
};

#[deriving(Decodable, Encodable, Show)]
pub struct BotConfig {
    pub server: String,
    pub command_prefix: String,
    pub nickname: String,
    pub channels: Vec<String>,
    pub enabled_plugins: HashSet<String>,
}

pub fn irc_scheme_type_mapper(scheme: &str) -> SchemeType {
    match scheme {
        "irc" => RelativeScheme(6667),
        "ircs" => RelativeScheme(6697),
        _ => NonRelativeScheme,
    }
}

impl BotConfig {
    fn get_url(&self) -> ParseResult<Url> {
        let mut parser = UrlParser::new();
        parser.scheme_type_mapper(irc_scheme_type_mapper);
        parser.parse(self.server.as_slice())
    }

    fn get_host(&self) -> String {
        let server = self.get_url().unwrap();
        match server.host() {
            Some(&Domain(ref string)) => string.clone(),
            Some(&Ipv6(ref addr)) => addr.serialize(),
            None => panic!()
        }
    }

    fn get_port(&self) -> u16 {
        let server = self.get_url().unwrap();
        server.port().unwrap_or(6667)
    }
}

pub struct BotConnection {
    foo: uint
    //
}

fn cool_trigger(message: &IrcMessage) -> Option<&str> {
    let prefix = "REWHO ";
    if !(message.is_privmsg() && message.get_args().len() == 2) {
        return None;
    }
    let message_text = message.get_args()[1];
    if !(message_text.starts_with(prefix)) {
        return None;
    }
    Some(message_text[prefix.len()..])
}

impl BotConnection {
    pub fn new(conf: &BotConfig) -> IoResult<BotConnection> {
        let (mut conn, event_queue) = try!(IrcConnection::new(
            (conf.get_host().as_slice(), conf.get_port())));

        let (event_queue_txu, event_queue_rxu) = channel();
        spawn(proc() {
            for event in event_queue.iter() {
                event_queue_txu.send(event);
            }
        });

        let mut nick = conf.nickname.clone();
        loop {
            println!("trying nick {}", nick.as_slice());
            match conn.register(nick.as_slice()) {
                Ok(_) => {
                    println!("ok, connected as {}", nick.as_slice());
                    break;
                }
                Err(err) => {
                    if err.should_pick_new_nickname() {
                        nick.push_str("`");
                    } else {
                        panic!("{}", err)
                    }
                }
            };
        }

        for channel in conf.channels.iter() {
            println!("want join: {}", channel);
        }

        for channel in conf.channels.iter() {
            println!("joining {}...", channel);
            match conn.join(channel.as_slice()) {
                Ok(res) => {
                    println!("succeeded in joining {}, got {} nicks",
                        res.channel.as_slice(), res.nicks.len());
                    match conn.who(channel.as_slice()) {
                        Ok(who_res) => {
                            println!("succeeded in WHOing {}, got {} nicks",
                                who_res.channel.as_slice(), who_res.who_records.len());
                        },
                        Err(who_err) => {
                            println!("failed to WHO {}: {}", channel, who_err);
                        }
                    }
                },
                Err(err) => {
                    println!("join error: {}", err);
                    panic!("failed to join channel.. dying");
                }
            }
            println!("END joining {}...", channel);
        }

        let mut state = State::new();

        let mut container = PluginContainer::new(conf.command_prefix.clone());
        if conf.enabled_plugins.contains_equiv(PingPlugin::get_plugin_name()) {
            container.register(box PingPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(GreedPlugin::get_plugin_name()) {
            container.register(box GreedPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(SeenPlugin::get_plugin_name()) {
            container.register(box SeenPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(DeerPlugin::get_plugin_name()) {
            container.register(box DeerPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(RadioPlugin::get_plugin_name()) {
            container.register(box RadioPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(WserverPlugin::get_plugin_name()) {
            container.register(box WserverPlugin::new());
        }
        if conf.enabled_plugins.contains_equiv(WhoAmIPlugin::get_plugin_name()) {
            container.register(box WhoAmIPlugin::new());
        }

        let (tx, rx) = sync_channel(0);
        let cmd_queue = conn.get_command_queue();

        TaskBuilder::new().named("bot-sender").spawn(proc() {
            for message in rx.iter() {
                cmd_queue.send(RawWrite(message));
            }
        });

        for event in event_queue_rxu.iter() {
            state.on_event(&event);
            if let IrcEventMessage(ref message) = event {
                container.dispatch(&state, &tx, message);
                match cool_trigger(message) {
                    Some(channel_name) => {
                        info!("running who");
                        info!("who returned: {}", conn.who(channel_name));
                    },
                    None => ()
                }
            }
        }

        Ok(BotConnection { foo: 0})
    }
}
<|endoftext|>