<commit_before>extern crate html5ever;
extern crate html5ever_dom_sink;
extern crate hyper;
extern crate rustc_serialize;
extern crate tempfile;
extern crate tendril;
extern crate url;

mod json;

use std::fs::OpenOptions;
use std::io::Read;
use std::io::Write;
use std::process::Command;
use std::process::Stdio;
use std::str::FromStr;

use html5ever::Attribute;
use html5ever::tree_builder::interface::TreeSink;
use html5ever_dom_sink::common::NodeEnum;
use html5ever_dom_sink::rcdom::Handle;
use html5ever_dom_sink::rcdom::RcDom;
use hyper::Client;
use hyper::header::Connection;
use rustc_serialize::json::Json;
use tempfile::NamedTempFile;
use url::percent_encoding;

use json::JsonPathElement::{Key, Only};

// TODO: there are some places where I've handrolled try!() equivalents. Fix those.
// TODO: make sure I'm returning Results everywhere, and propagating errors correctly.
// TODO: randomize placeholder per-request

// TODO: Can I get rid of the repeated "Zachary_Taylor"s everywhere? Surely the MediaWiki API doesn't actually need that - I can't imagine revision IDs aren't unique across all pages.

fn get_revisions(page: &str, limit: i32) -> String {
    let client = Client::new();
    let mut res = client.get(
        &format!(
            "https://en.wikipedia.org/w/api.php?action=query&prop=revisions&titles={}&rvprop=timestamp|user|comment|ids&rvlimit={}&format=json",
            page, limit))
        .header(Connection::close())
        .send().unwrap();

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

    body
}

// TODO: this name is terrible.
// Returns pairs of (revid, parentid)
fn get_revert_revision_ids(page: &str) -> Result<Vec<(u64, u64)>, String> {
    let json = Json::from_str(&get_revisions(page, 60)).unwrap();
    let revisions = try!(
        json::get_json_array(&json, &[Key("query"), Key("pages"), Only, Key("revisions")]));

    // Filter for revisions that mention vandalism, and that have revid and parentid (which they all
    // should, but merits checking).
    Ok(revisions.into_iter().filter(|revision| {
        match (json::get_json_string(revision, &[Key("comment")]),
               json::get_json_number(revision, &[Key("revid")]),
               json::get_json_number(revision, &[Key("parentid")])) {
            (Ok(ref comment), Ok(_), Ok(_)) => comment.contains("vandal"),
            _ => false,
        }})
       .map(|revision| {
           (json::get_json_number(revision, &[Key("revid")]).unwrap(),
            json::get_json_number(revision, &[Key("parentid")]).unwrap())
       }).collect())
}

fn get_latest_revision_id(page: &str) -> Result<u64, String> {
    let json = Json::from_str(&get_revisions(page, 1)).unwrap();
    json::get_json_number(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("revid")])
}

fn get_revision(page: &str, id: u64) -> String {
    let client = Client::new();
    let mut res = client.get(
        &format!(
            "https://en.wikipedia.org/w/api.php?action=query&prop=revisions&titles={}&rvprop=content&rvlimit=1&rvstartid={}&format=json",
            page, id))
        .header(Connection::close())
        .send().unwrap();

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

fn get_revision_content(page: &str, id: u64) -> Result<String, String> {
    let json = Json::from_str(&get_revision(page, id)).unwrap();
    match json::get_json_string(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("*")]) {
        Ok(content) => Ok(content.to_string()),
        Err(msg) => Err(msg),
    }
}

fn write_to_temp_file(contents: &str) -> NamedTempFile {
    let tempfile = NamedTempFile::new().unwrap();
    let mut file = OpenOptions::new().write(true).open(tempfile.path()).unwrap();
    file.write_all(contents.as_bytes()).unwrap();
    file.flush().unwrap();
    tempfile
}

// TODO: I'm not so sure these parameter names aren't terrible.
fn merge(old: &str, new1: &str, new2: &str) -> Option<String> {
    let old_tempfile = write_to_temp_file(old);
    let new1_tempfile = write_to_temp_file(new1);
    let new2_tempfile = write_to_temp_file(new2);

    let mut process = Command::new("diff3");
    process.arg("-m").args(&[new1_tempfile.path(), old_tempfile.path(), new2_tempfile.path()])
        .stdout(Stdio::piped()).stderr(Stdio::null());
    let output = process.output().unwrap();
    if output.status.success() {
        Some(String::from_utf8(output.stdout).unwrap())
    } else {
        None
    }
}

fn render(wikitext: &str) -> Result<String, String> {
    let post_body = format!(
        "text={}", percent_encoding::percent_encode(
            wikitext.as_bytes(), percent_encoding::QUERY_ENCODE_SET));

    let client = Client::new();
    let mut res = client.post(
        "https://en.wikipedia.org/w/api.php?action=parse&format=json&prop=text&disablepp=&contentmodel=wikitext")
        .body(&post_body)
        .header(Connection::close())
        .send().unwrap();

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

    // TODO: check return value
    let json = Json::from_str(&body).unwrap();
    match json::get_json_string(&json, &[Key("parse"), Key("text"), Key("*")]) {
        Ok(contents) => Ok(contents.to_string()),
        Err(message) => Err(message),
    }
}

fn has_matching_id(attributes: &Vec<Attribute>, id: &str) -> bool {
    // TODO: could also do this with filter() and is_empty(). Not sure if that would be better.
    for attribute in attributes {
        // TODO: do I seriously have to construct a StrTendril here?
        // There has to be a better way.
        if attribute.name.local.as_slice() == "id" &&
            attribute.value == tendril::StrTendril::from_str(id).unwrap() {
                return true;
            }
    }
    return false;
}

fn find_node_by_id(handle: &Handle, id: &str) -> Result<Handle, String> {
    let node = handle.borrow();
    match node.node {
        NodeEnum::Element(_, ref attributes) if has_matching_id(attributes, id) => Ok(handle.clone()),
        _ => {
            for child in &node.children {
                match find_node_by_id(child, id) {
                    // TODO: this looks weird. Is this an abnormal way to do things?
                    Ok(node) => return Ok(node),
                    _ => continue,
                }
            }
            Err(format!("No node with ID {} found", id))
        },
    }
}

fn replace_node_with_placeholder(original_html: &str, div_id: &str, placeholder: &str) -> Result<String, String> {
    // TODO: check errors
    let html = tendril::StrTendril::from_str(original_html).unwrap();
    let mut dom: RcDom = html5ever::parse(html5ever::one_input(html), Default::default());

    let handle = try!(find_node_by_id(&dom.get_document(), div_id));
    let child_handles =
        (&handle.borrow().children).into_iter().map(|child| child.clone()).collect::<Vec<_>>();
    for child_handle in child_handles {
        dom.remove_from_parent(child_handle);
    }
    dom.append(handle,
               html5ever::tree_builder::interface::NodeOrText::AppendText(
                   tendril::StrTendril::from_str(placeholder).unwrap()));
    let mut serialized: Vec<u8> = vec![];
    html5ever::serialize::serialize(&mut serialized, &dom.document, Default::default());
    // TODO: error handling
    Ok(String::from_utf8(serialized).unwrap())
}

fn get_current_page(page: &str) -> String {
    let client = Client::new();
    let mut res = client.get(
        &format!("https://en.wikipedia.org/wiki/{}", page))
        .header(Connection::close())
        .send().unwrap();

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

    body
}

fn main() {
    let latest_revid = get_latest_revision_id("Zachary_Taylor").unwrap();
    let mut accumulated_contents = get_revision_content("Zachary_Taylor", latest_revid).unwrap();
    let mut revisions = vec![];
    for (revert_revid, vandalism_revid) in get_revert_revision_ids("Zachary_Taylor").ok().unwrap() {
        let reverted_contents = get_revision_content("Zachary_Taylor", revert_revid).unwrap();
        let vandalized_contents = get_revision_content("Zachary_Taylor", vandalism_revid).unwrap();
        match merge(&reverted_contents, &vandalized_contents, &accumulated_contents) {
            Some(merged_contents) => {
                accumulated_contents = merged_contents;
                revisions.push(revert_revid);
            }
            None => (),
        }
    }

    let rendered_body = render(&accumulated_contents).unwrap();
    let current_page_contents = get_current_page("Zachary_Taylor");
    let page_contents_with_placeholder =
        replace_node_with_placeholder(
            &current_page_contents, "mw-content-text", "WMW_PLACEHOLDER_TEXT").unwrap();
    let final_page_contents = page_contents_with_placeholder.replace("WMW_PLACEHOLDER_TEXT", &rendered_body);

    println!("Restored vandalisms reverted in: {:?}", revisions);
    println!("{}", final_page_contents);
}

#[cfg(test)]
mod tests {
    use super::{merge, replace_node_with_placeholder};

    #[test]
    fn test_merge_clean() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed.\n";
        let new2 = "First line changed.\n\nSecond line.\n";
        assert_eq!("First line changed.\n\nSecond line changed.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_merge_conflicting() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed one way.\n";
        let new2 = "First line changed.\n\nSecond line changed a different way.\n";
        assert_eq!(None, merge(old, new1, new2));
    }

    #[test]
    fn test_merge_special_characters() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line 𐅃.\n";
        let new2 = "First line さようなら.\n\nSecond line.\n";
        assert_eq!("First line さようなら.\n\nSecond line 𐅃.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_replace_html_content() {
        let original_html = "

original text

Other text

";
        let expected_html = "

replaced text

Other text

";
        let processed_html = replace_node_with_placeholder(original_html, "mw-content-text", "replaced text").unwrap();
        assert_eq!(expected_html, processed_html);
    }
}
<commit_msg>Combine get_revision() and get_revisions() into one function.<commit_after>extern crate html5ever;
extern crate html5ever_dom_sink;
extern crate hyper;
extern crate rustc_serialize;
extern crate tempfile;
extern crate tendril;
extern crate url;

mod json;

use std::fs::OpenOptions;
use std::io::Read;
use std::io::Write;
use std::process::Command;
use std::process::Stdio;
use std::str::FromStr;

use html5ever::Attribute;
use html5ever::tree_builder::interface::TreeSink;
use html5ever_dom_sink::common::NodeEnum;
use html5ever_dom_sink::rcdom::Handle;
use html5ever_dom_sink::rcdom::RcDom;
use hyper::Client;
use hyper::header::Connection;
use rustc_serialize::json::Json;
use tempfile::NamedTempFile;
use url::percent_encoding;

use json::JsonPathElement::{Key, Only};

// TODO: there are some places where I've handrolled try!() equivalents. Fix those.
// TODO: make sure I'm returning Results everywhere, and propagating errors correctly.
// TODO: randomize placeholder per-request

// TODO: Can I get rid of the repeated "Zachary_Taylor"s everywhere? Surely the MediaWiki API doesn't actually need that - I can't imagine revision IDs aren't unique across all pages.

fn get_revisions(page: &str, limit: i32, properties: Vec<&str>, start_id: Option<u64>) -> String {
    let client = Client::new();
    let mut res = client.get(
        &format!(
            "https://en.wikipedia.org/w/api.php?action=query&prop=revisions&titles={}&rvprop={}&rvlimit={}{}&format=json",
            page, properties.join("|"), limit,
            match start_id { Some(id) => format!("&rvstartid={}", id), None => "".to_string()}))
        .header(Connection::close())
        .send().unwrap();

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

    body
}

// TODO: this name is terrible.
// Returns pairs of (revid, parentid)
fn get_revert_revision_ids(page: &str) -> Result<Vec<(u64, u64)>, String> {
    let json = Json::from_str(&get_revisions(page, 60, vec!["comment", "ids"], None)).unwrap();
    let revisions = try!(
        json::get_json_array(&json, &[Key("query"), Key("pages"), Only, Key("revisions")]));

    // Filter for revisions that mention vandalism, and that have revid and parentid (which they all
    // should, but merits checking).
    Ok(revisions.into_iter().filter(|revision| {
        match (json::get_json_string(revision, &[Key("comment")]),
               json::get_json_number(revision, &[Key("revid")]),
               json::get_json_number(revision, &[Key("parentid")])) {
            (Ok(ref comment), Ok(_), Ok(_)) => comment.contains("vandal"),
            _ => false,
        }})
       .map(|revision| {
           (json::get_json_number(revision, &[Key("revid")]).unwrap(),
            json::get_json_number(revision, &[Key("parentid")]).unwrap())
       }).collect())
}

fn get_latest_revision_id(page: &str) -> Result<u64, String> {
    let json = Json::from_str(&get_revisions(page, 1, vec!["ids"], None)).unwrap();
    json::get_json_number(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("revid")])
}

fn get_revision_content(page: &str, id: u64) -> Result<String, String> {
    let json = Json::from_str(&get_revisions(page, 1, vec!["content"], Some(id))).unwrap();
    match json::get_json_string(
        &json, &[Key("query"), Key("pages"), Only, Key("revisions"), Only, Key("*")]) {
        Ok(content) => Ok(content.to_string()),
        Err(msg) => Err(msg),
    }
}

fn write_to_temp_file(contents: &str) -> NamedTempFile {
    let tempfile = NamedTempFile::new().unwrap();
    let mut file = OpenOptions::new().write(true).open(tempfile.path()).unwrap();
    file.write_all(contents.as_bytes()).unwrap();
    file.flush().unwrap();
    tempfile
}

// TODO: I'm not so sure these parameter names aren't terrible.
fn merge(old: &str, new1: &str, new2: &str) -> Option<String> {
    let old_tempfile = write_to_temp_file(old);
    let new1_tempfile = write_to_temp_file(new1);
    let new2_tempfile = write_to_temp_file(new2);

    let mut process = Command::new("diff3");
    process.arg("-m").args(&[new1_tempfile.path(), old_tempfile.path(), new2_tempfile.path()])
        .stdout(Stdio::piped()).stderr(Stdio::null());
    let output = process.output().unwrap();
    if output.status.success() {
        Some(String::from_utf8(output.stdout).unwrap())
    } else {
        None
    }
}

fn render(wikitext: &str) -> Result<String, String> {
    let post_body = format!(
        "text={}", percent_encoding::percent_encode(
            wikitext.as_bytes(), percent_encoding::QUERY_ENCODE_SET));

    let client = Client::new();
    let mut res = client.post(
        "https://en.wikipedia.org/w/api.php?action=parse&format=json&prop=text&disablepp=&contentmodel=wikitext")
        .body(&post_body)
        .header(Connection::close())
        .send().unwrap();

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

    // TODO: check return value
    let json = Json::from_str(&body).unwrap();
    match json::get_json_string(&json, &[Key("parse"), Key("text"), Key("*")]) {
        Ok(contents) => Ok(contents.to_string()),
        Err(message) => Err(message),
    }
}

fn has_matching_id(attributes: &Vec<Attribute>, id: &str) -> bool {
    // TODO: could also do this with filter() and is_empty(). Not sure if that would be better.
    for attribute in attributes {
        // TODO: do I seriously have to construct a StrTendril here?
        // There has to be a better way.
        if attribute.name.local.as_slice() == "id" &&
            attribute.value == tendril::StrTendril::from_str(id).unwrap() {
                return true;
            }
    }
    return false;
}

fn find_node_by_id(handle: &Handle, id: &str) -> Result<Handle, String> {
    let node = handle.borrow();
    match node.node {
        NodeEnum::Element(_, ref attributes) if has_matching_id(attributes, id) => Ok(handle.clone()),
        _ => {
            for child in &node.children {
                match find_node_by_id(child, id) {
                    // TODO: this looks weird. Is this an abnormal way to do things?
                    Ok(node) => return Ok(node),
                    _ => continue,
                }
            }
            Err(format!("No node with ID {} found", id))
        },
    }
}

fn replace_node_with_placeholder(original_html: &str, div_id: &str, placeholder: &str) -> Result<String, String> {
    // TODO: check errors
    let html = tendril::StrTendril::from_str(original_html).unwrap();
    let mut dom: RcDom = html5ever::parse(html5ever::one_input(html), Default::default());

    let handle = try!(find_node_by_id(&dom.get_document(), div_id));
    let child_handles =
        (&handle.borrow().children).into_iter().map(|child| child.clone()).collect::<Vec<_>>();
    for child_handle in child_handles {
        dom.remove_from_parent(child_handle);
    }
    dom.append(handle,
               html5ever::tree_builder::interface::NodeOrText::AppendText(
                   tendril::StrTendril::from_str(placeholder).unwrap()));
    let mut serialized: Vec<u8> = vec![];
    html5ever::serialize::serialize(&mut serialized, &dom.document, Default::default());
    // TODO: error handling
    Ok(String::from_utf8(serialized).unwrap())
}

fn get_current_page(page: &str) -> String {
    let client = Client::new();
    let mut res = client.get(
        &format!("https://en.wikipedia.org/wiki/{}", page))
        .header(Connection::close())
        .send().unwrap();

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

    body
}

fn main() {
    let latest_revid = get_latest_revision_id("Zachary_Taylor").unwrap();
    let mut accumulated_contents = get_revision_content("Zachary_Taylor", latest_revid).unwrap();
    let mut revisions = vec![];
    for (revert_revid, vandalism_revid) in get_revert_revision_ids("Zachary_Taylor").ok().unwrap() {
        let reverted_contents = get_revision_content("Zachary_Taylor", revert_revid).unwrap();
        let vandalized_contents = get_revision_content("Zachary_Taylor", vandalism_revid).unwrap();
        match merge(&reverted_contents, &vandalized_contents, &accumulated_contents) {
            Some(merged_contents) => {
                accumulated_contents = merged_contents;
                revisions.push(revert_revid);
            }
            None => (),
        }
    }

    let rendered_body = render(&accumulated_contents).unwrap();
    let current_page_contents = get_current_page("Zachary_Taylor");
    let page_contents_with_placeholder =
        replace_node_with_placeholder(
            &current_page_contents, "mw-content-text", "WMW_PLACEHOLDER_TEXT").unwrap();
    let final_page_contents = page_contents_with_placeholder.replace("WMW_PLACEHOLDER_TEXT", &rendered_body);

    println!("Restored vandalisms reverted in: {:?}", revisions);
    println!("{}", final_page_contents);
}

#[cfg(test)]
mod tests {
    use super::{merge, replace_node_with_placeholder};

    #[test]
    fn test_merge_clean() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed.\n";
        let new2 = "First line changed.\n\nSecond line.\n";
        assert_eq!("First line changed.\n\nSecond line changed.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_merge_conflicting() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line changed one way.\n";
        let new2 = "First line changed.\n\nSecond line changed a different way.\n";
        assert_eq!(None, merge(old, new1, new2));
    }

    #[test]
    fn test_merge_special_characters() {
        let old = "First line.\n\nSecond line.\n";
        let new1 = "First line.\n\nSecond line 𐅃.\n";
        let new2 = "First line さようなら.\n\nSecond line.\n";
        assert_eq!("First line さようなら.\n\nSecond line 𐅃.\n", merge(old, new1, new2).unwrap());
    }

    #[test]
    fn test_replace_html_content() {
        let original_html = "

original text

Other text

";
        let expected_html = "

replaced text

Other text

";
        let processed_html = replace_node_with_placeholder(original_html, "mw-content-text", "replaced text").unwrap();
        assert_eq!(expected_html, processed_html);
    }
}
<|endoftext|>


<commit_before>#![feature(slicing_syntax, macro_rules)]
#![feature(overloaded_calls)]

extern crate http;
extern crate url;
extern crate serialize;

use std::os;
use std::str;
use std::sync::{Arc, Mutex};
use std::io::net::ip::{SocketAddr, Ipv4Addr};
use std::collections::HashMap;

use http::method::{Get, Post};
use http::server::{Config, Server, Request, ResponseWriter};
use http::server::request::{AbsolutePath};

use serialize::json;

use incoming::{SlackPayload};
use outgoing::{SlackEndpoint, OutgoingWebhook};

mod incoming;
mod outgoing;

type Scores = HashMap<String, i32>;

#[deriving(Send,Clone)]
struct KarmaServer {
    scores: Arc<Mutex<Scores>>,
    endpoint: String,
}

impl KarmaServer {
    fn new(scores: Scores, endpoint: String) -> KarmaServer {
        KarmaServer {
            scores: Arc::new(Mutex::new(scores)),
            endpoint: endpoint
        }
    }
}

fn handle_karma(req: Vec<u8>, points: &mut Scores, cb: |&&str, &i32|) {
    let payload = match SlackPayload::from_body(req.as_slice()) {
        Ok(payload) => payload,
        Err(err) => {
            println!("Error! {}", err); return
        }
    };

    let op = match payload.command {
        "%2F%2B%2B" => |c: i32| c + 1,
        "%2F%2D%2D" => |c: i32| c - 1,
        _ => return,
    };

    let mut matches = payload.text.splitn(1, ' ');
    let target = match (matches.next(), matches.next()) {
        (None, None) => return, // No message or user
        (Some(user), None) => user, // user, no message
        (Some(user), Some(_)) => user, // TODO deal with the message
        (None, Some(_)) => unreachable!(),
    };

    let current = points.insert(target.to_string(), 0).unwrap_or(0);
    points.insert(target.to_string(), op(current));

    cb(&target, &op(current));
}

impl Server for KarmaServer {
    fn get_config(&self) -> Config {
        match os::getenv("PORT") {
            None => panic!("Must specify PORT"),
            Some(port) => {
                let port = from_str::<u16>(port.as_slice()).expect("PORT must be an int");
                Config { bind_address: SocketAddr { ip: Ipv4Addr(0, 0, 0, 0), port: port } }
            }
        }
    }

    #[allow(unused_must_use)]
    fn handle_request(&self, r: Request, w: &mut ResponseWriter) {
        let path = match r.request_uri { AbsolutePath(ref path) => path[], _ => return };

        println!("{}: {}", r.method, path);

        match (&r.method, path) {
            (&Get, "/") => {
                w.write(b"Hello world!");
            },
            (&Get, "/karma") => {
                let scores = (*self.scores).lock();
                let json = json::encode(&*scores);
                w.write(json.as_bytes());
            }
            (&Post, "/normalise") => {
                let mut scores = (*self.scores).lock();
                for (_, v) in scores.iter_mut() {
                    let new = *v as f64;
                    *v = new.sqrt() as i32;
                }
            }
            (&Post, "/slack") => {
                let mut scores = (*self.scores).lock();
                let slack = self.get_slack_endpoint();
                handle_karma(r.body, &mut *scores, |u, s| {
                    let msg = format!("{} now at {}", u, s);
                    let payload = OutgoingWebhook {
                        text: msg.as_slice(),
                        channel: "#hax",
                        username: "karmabot",
                        icon_emoji: Some(":ghost:"),
                    };

                    slack.send(&payload);
                });
            },
            (_, _) => {
                w.write(b"Not found :(");
            }
        }
    }
}

impl KarmaServer {
    fn get_slack_endpoint(&self) -> SlackEndpoint {
        return SlackEndpoint {
            url: self.endpoint.clone()
        }
    }
}

fn main() {
    let endpoint = match os::getenv("WEBHOOK_ENDPOINT") {
        Some(e) => e,
        None => panic!("Must set WEBHOOK_ENDPOINT"),
    };
    let server = KarmaServer::new(HashMap::new(), endpoint);
    server.serve_forever();
}
<commit_msg>Laziness just cover all of them<commit_after>#![feature(slicing_syntax, macro_rules)]
#![feature(overloaded_calls)]

extern crate http;
extern crate url;
extern crate serialize;

use std::os;
use std::str;
use std::sync::{Arc, Mutex};
use std::io::net::ip::{SocketAddr, Ipv4Addr};
use std::collections::HashMap;

use http::method::{Get, Post};
use http::server::{Config, Server, Request, ResponseWriter};
use http::server::request::{AbsolutePath};

use serialize::json;

use incoming::{SlackPayload};
use outgoing::{SlackEndpoint, OutgoingWebhook};

mod incoming;
mod outgoing;

type Scores = HashMap<String, i32>;

#[deriving(Send,Clone)]
struct KarmaServer {
    scores: Arc<Mutex<Scores>>,
    endpoint: String,
}

impl KarmaServer {
    fn new(scores: Scores, endpoint: String) -> KarmaServer {
        KarmaServer {
            scores: Arc::new(Mutex::new(scores)),
            endpoint: endpoint
        }
    }
}

fn handle_karma(req: Vec<u8>, points: &mut Scores, cb: |&&str, &i32|) {
    let payload = match SlackPayload::from_body(req.as_slice()) {
        Ok(payload) => payload,
        Err(err) => {
            println!("Error! {}", err); return
        }
    };

    let op = match payload.command {
        "%2F%2B%2B" => |c: i32| c + 1,
        "%2F++" => |c: i32| c + 1,
        "%2F%2D%2D" => |c: i32| c - 1,
        "%2F--" => |c: i32| c - 1,
        _ => return,
    };

    let mut matches = payload.text.splitn(1, ' ');
    let target = match (matches.next(), matches.next()) {
        (None, None) => return, // No message or user
        (Some(user), None) => user, // user, no message
        (Some(user), Some(_)) => user, // TODO deal with the message
        (None, Some(_)) => unreachable!(),
    };

    let current = points.insert(target.to_string(), 0).unwrap_or(0);
    points.insert(target.to_string(), op(current));

    cb(&target, &op(current));
}

impl Server for KarmaServer {
    fn get_config(&self) -> Config {
        match os::getenv("PORT") {
            None => panic!("Must specify PORT"),
            Some(port) => {
                let port = from_str::<u16>(port.as_slice()).expect("PORT must be an int");
                Config { bind_address: SocketAddr { ip: Ipv4Addr(0, 0, 0, 0), port: port } }
            }
        }
    }

    #[allow(unused_must_use)]
    fn handle_request(&self, r: Request, w: &mut ResponseWriter) {
        let path = match r.request_uri { AbsolutePath(ref path) => path[], _ => return };

        println!("{}: {}", r.method, path);

        match (&r.method, path) {
            (&Get, "/") => {
                w.write(b"Hello world!");
            },
            (&Get, "/karma") => {
                let scores = (*self.scores).lock();
                let json = json::encode(&*scores);
                w.write(json.as_bytes());
            }
            (&Post, "/normalise") => {
                let mut scores = (*self.scores).lock();
                for (_, v) in scores.iter_mut() {
                    let new = *v as f64;
                    *v = new.sqrt() as i32;
                }
            }
            (&Post, "/slack") => {
                let mut scores = (*self.scores).lock();
                let slack = self.get_slack_endpoint();
                handle_karma(r.body, &mut *scores, |u, s| {
                    let msg = format!("{} now at {}", u, s);
                    let payload = OutgoingWebhook {
                        text: msg.as_slice(),
                        channel: "#hax",
                        username: "karmabot",
                        icon_emoji: Some(":ghost:"),
                    };

                    slack.send(&payload);
                });
            },
            (_, _) => {
                w.write(b"Not found :(");
            }
        }
    }
}

impl KarmaServer {
    fn get_slack_endpoint(&self) -> SlackEndpoint {
        return SlackEndpoint {
            url: self.endpoint.clone()
        }
    }
}

fn main() {
    let endpoint = match os::getenv("WEBHOOK_ENDPOINT") {
        Some(e) => e,
        None => panic!("Must set WEBHOOK_ENDPOINT"),
    };
    let server = KarmaServer::new(HashMap::new(), endpoint);
    server.serve_forever();
}
<|endoftext|>


<commit_before>#![feature(iter_arith, result_expect)]

extern crate byteorder;
extern crate crc;
extern crate unix_socket;
extern crate nix;
extern crate libc;
extern crate regex;

use std::path;
use std::io::Result;
use std::io::Error;
use std::io::ErrorKind::Other;

mod parser;
mod lvmetad;
mod pvlabel;
mod dm;
mod lv;
mod vg;
mod pv;

#[allow(dead_code, non_camel_case_types)]
mod dm_ioctl;

use parser::LvmTextMap;
use parser::TextMapOps;

fn get_first_vg_meta() -> Result<(String, LvmTextMap)> {
    let dirs = vec![path::Path::new("/dev")];

    for pv in try!(pvlabel::scan_for_pvs(&dirs)) {
        let map = try!(pvlabel::textmap_from_dev(pv.as_path()));

        for (key, value) in map {
            match value {
                parser::Entry::TextMap(x) => return Ok((key, *x)),
                _ => {}
            }
        }
    }

    Err(Error::new(Other, "dude"))
}

fn main() {
    // println!("A");
    // let (name, map) = get_first_vg_meta().unwrap();
    // println!("B {}", name);
    // let vg = parser::vg_from_textmap(&name, &map).expect("didn't get vg!");
    // println!("heyo {} {}", vg.extents(), vg.extent_size);
    // println!("output {:?}", vg);

    // match dm::list_devices() {
    //     Ok(x) => println!("{:?}", x),
    //     Err(x) => println!("error {}", x),
    // }

    let vgs = lvmetad::vgs_from_lvmetad().expect("could not get vgs from lvmetad");
    for vg in &vgs {
        println!("{} tot {} alloc {} free {}",
                 vg.name, vg.extents(), vg.extents_in_use(), vg.extents_free());
        for lv in vg.lvs.keys() {
            println!("lv {}", lv);
        }
    }

    let tm = pvlabel::get_conf().expect("could not read lvm.conf");
    let locking_type = tm.textmap_from_textmap("global")
        .and_then(|g| g.i64_from_textmap("locking_type")).unwrap();

    println!("locking_type = {}", locking_type);
}
<commit_msg>put get_conf back in main<commit_after>#![feature(iter_arith, result_expect)]

extern crate byteorder;
extern crate crc;
extern crate unix_socket;
extern crate nix;
extern crate libc;
extern crate regex;

use std::path;
use std::io::Result;
use std::io::Error;
use std::io::ErrorKind::Other;

mod parser;
mod lvmetad;
mod pvlabel;
mod dm;
mod lv;
mod vg;
mod pv;

#[allow(dead_code, non_camel_case_types)]
mod dm_ioctl;

use parser::LvmTextMap;
use parser::TextMapOps;

fn get_first_vg_meta() -> Result<(String, LvmTextMap)> {
    let dirs = vec![path::Path::new("/dev")];

    for pv in try!(pvlabel::scan_for_pvs(&dirs)) {
        let map = try!(pvlabel::textmap_from_dev(pv.as_path()));

        for (key, value) in map {
            match value {
                parser::Entry::TextMap(x) => return Ok((key, *x)),
                _ => {}
            }
        }
    }

    Err(Error::new(Other, "dude"))
}

fn get_conf() -> Result<LvmTextMap> {
    use std::fs;
    use std::io::Read;

    let mut f = try!(fs::File::open("/etc/lvm/lvm.conf"));

    let mut buf = Vec::new();
    try!(f.read_to_end(&mut buf));

    parser::into_textmap(&buf)
}

fn main() {
    // println!("A");
    // let (name, map) = get_first_vg_meta().unwrap();
    // println!("B {}", name);
    // let vg = parser::vg_from_textmap(&name, &map).expect("didn't get vg!");
    // println!("heyo {} {}", vg.extents(), vg.extent_size);
    // println!("output {:?}", vg);

    // match dm::list_devices() {
    //     Ok(x) => println!("{:?}", x),
    //     Err(x) => println!("error {}", x),
    // }

    let vgs = lvmetad::vgs_from_lvmetad().expect("could not get vgs from lvmetad");
    for vg in &vgs {
        println!("{} tot {} alloc {} free {}",
                 vg.name, vg.extents(), vg.extents_in_use(), vg.extents_free());
        for lv in vg.lvs.keys() {
            println!("lv {}", lv);
        }
    }

    let tm = get_conf().expect("could not read lvm.conf");
    let locking_type = tm.textmap_from_textmap("global")
        .and_then(|g| g.i64_from_textmap("locking_type")).unwrap();

    println!("locking_type = {}", locking_type);
}
<|endoftext|>


<commit_before>// UDP stuff adapted from http://illegalargumentexception.blogspot.co.uk/2015/05/rust-send-and-receive-on-localhost-with.html
use std::net;
use std::env;
use std::convert::AsRef;
extern crate getopts;
use getopts::Options;

fn socket(listen_on: net::SocketAddr) -> net::UdpSocket {
    let attempt = net::UdpSocket::bind(listen_on);
    let mut socket;
    match attempt {
        Ok(sock) => {
            socket = sock;
        },
    Err(err) => panic!("Could not bind: {}", err)
    }
    socket
}

pub fn send_message(send_addr: net::SocketAddr, target: net::SocketAddr, data: Vec<u8>) {
    let socket = socket(send_addr);
    let _ = socket.send_to(&data, target);
    drop(socket);
}

fn convert(inp: String) -> String {
    // to_lowercase is still unstable.
    // This madness from IRC is the workaround
    // works though

    // Strings do not live for the entire life of your program
    // http://stackoverflow.com/a/23977218/416626
    let intermediate: &str = &*inp;
    let outp = intermediate
        .chars()
        .flat_map(char::to_lowercase)
        .collect::<String>();
    outp
}

fn print_usage(program: &str, opts: Options, ecode: i32) {
    let brief = format!("Usage: {} [options] [command]", program);
    print!("{}", opts.usage(&brief));
    std::process::exit(ecode);
}

fn main() {
    let mut opts = Options::new();
    // get command-line input
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    opts.optopt("p", "port", "Set destination UDP port. Input must be 0 – 65535", "PORT");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => {m}
        Err(f) => {panic!(f.to_string())}
    };
    // Get port from the option, or specify the default
    let port = matches.opt_str("p").unwrap_or("1738".to_string());
    // cast to int and ensure it worked
    let numeric_port = port.parse::<u16>();
    let arg = match &numeric_port {
        &Ok(_) => matches.free[0].clone(),
        &Err(_) => {
            print_usage(&program, opts, 1);
            return;
        }
    };
    // match command-line input or print usage
    let to_send = match convert(arg).as_ref() {
        "white" => "white".to_string(),
        "red" => "red".to_string(),
        "orange" => "orange".to_string(),
        "yellow" => "yellow".to_string(),
        "green" => "green".to_string(),
        "cyan" => "cyan".to_string(),
        "blue" => "blue".to_string(),
        "purple" => "purple".to_string(),
        "black" => "black".to_string(),
        "?" => "question".to_string(),
        "!" => "exclamation".to_string(),
        "quit" => "quit".to_string(),
        _ => {
            print_usage(&program, opts, 1);
            return;
        }
    };
    // blam our control message into a vector
    let mut message: Vec<u8> = Vec::new();
    message.extend(to_send
        .as_bytes()
        .iter()
        .cloned()
    );
    // bind to the correct UDP port
    let ip = net::Ipv4Addr::new(127, 0, 0, 1);
    // parse() gives us Result which we need to unwrap
    // it's safe to unwrap here, cos we already checked success
    let listen_addr = net::SocketAddrV4::new(ip, numeric_port.unwrap());
    let send_addr = net::SocketAddrV4::new(ip, 0);
    // and send our message
    send_message(
        net::SocketAddr::V4(send_addr),
        net::SocketAddr::V4(listen_addr),
        message);
}
<commit_msg>Formatting fixes<commit_after>// UDP stuff adapted from http://illegalargumentexception.blogspot.co.uk/2015/05/rust-send-and-receive-on-localhost-with.html
use std::net;
use std::env;
use std::convert::AsRef;
extern crate getopts;
use getopts::Options;

fn socket(listen_on: net::SocketAddr) -> net::UdpSocket {
    let attempt = net::UdpSocket::bind(listen_on);
    let mut socket;
    match attempt {
        Ok(sock) => { socket = sock; },
        Err(err) => panic!("Could not bind: {}", err)
    }
    socket
}

pub fn send_message(send_addr: net::SocketAddr, target: net::SocketAddr, data: Vec<u8>) {
    let socket = socket(send_addr);
    let _ = socket.send_to(&data, target);
    drop(socket);
}

fn convert(inp: String) -> String {
    // to_lowercase is still unstable.
    // This madness from IRC is the workaround
    // works though

    // Strings do not live for the entire life of your program
    // http://stackoverflow.com/a/23977218/416626
    let intermediate: &str = &*inp;
    let outp = intermediate
        .chars()
        .flat_map(char::to_lowercase)
        .collect::<String>();
    outp
}

fn print_usage(program: &str, opts: Options, ecode: i32) {
    let brief = format!("Usage: {} [options] [command]", program);
    print!("{}", opts.usage(&brief));
    std::process::exit(ecode);
}

fn main() {
    let mut opts = Options::new();
    // get command-line input
    let args: Vec<String> = env::args().collect();
    let program = args[0].clone();
    opts.optopt("p", "port", "Set destination UDP port. Input must be 0 – 65535", "PORT");
    let matches = match opts.parse(&args[1..]) {
        Ok(m) => {m}
        Err(f) => { panic!(f.to_string()) }
    };
    // Get port from the option, or specify the default
    let port = matches.opt_str("p").unwrap_or("1738".to_string());
    // cast to int and ensure it worked
    let numeric_port = port.parse::<u16>();
    let arg = match &numeric_port {
        &Ok(_) => matches.free[0].clone(),
        &Err(_) => {
            print_usage(&program, opts, 1);
            return;
        }
    };
    // match command-line input or print usage
    let to_send = match convert(arg).as_ref() {
        "white" => "white".to_string(),
        "red" => "red".to_string(),
        "orange" => "orange".to_string(),
        "yellow" => "yellow".to_string(),
        "green" => "green".to_string(),
        "cyan" => "cyan".to_string(),
        "blue" => "blue".to_string(),
        "purple" => "purple".to_string(),
        "black" => "black".to_string(),
        "?" => "question".to_string(),
        "!" => "exclamation".to_string(),
        "quit" => "quit".to_string(),
        _ => {
            print_usage(&program, opts, 1);
            return;
        }
    };
    // blam our control message into a vector
    let mut message: Vec<u8> = Vec::new();
    message.extend(to_send
        .as_bytes()
        .iter()
        .cloned()
    );
    // bind to the correct UDP port
    let ip = net::Ipv4Addr::new(127, 0, 0, 1);
    // parse() gives us Result which we need to unwrap
    // it's safe to unwrap here, cos we already checked success
    let listen_addr = net::SocketAddrV4::new(ip, numeric_port.unwrap());
    let send_addr = net::SocketAddrV4::new(ip, 0);
    // and send our message
    send_message(
        net::SocketAddr::V4(send_addr),
        net::SocketAddr::V4(listen_addr),
        message);
}
<|endoftext|>


<commit_before>#![feature(duration_float)]
#![feature(euclidean_division)]
#![feature(self_struct_ctor)]
#![feature(try_trait)]

use std::path::PathBuf;

use chrono::Local;

use glium::glutin::{ElementState, Event, KeyboardInput, VirtualKeyCode, WindowEvent};
use glium::Surface;

mod aabb;
mod bsdf;
mod bvh;
mod camera;
mod color;
mod config;
mod consts;
mod float;
mod gl_renderer;
mod index_ptr;
mod input;
mod intersect;
mod light;
mod load;
mod material;
mod mesh;
mod obj_load;
mod pt_renderer;
mod scattering;
mod scene;
mod stats;
mod texture;
mod triangle;
mod vertex;

use self::config::RenderConfig;
use self::gl_renderer::GLRenderer;
use self::input::InputState;
use self::pt_renderer::PTRenderer;

fn main() {
    match std::env::args().nth(1).as_ref().map(|s| s.as_str()) {
        Some(_) => benchmark(),
        None => online_render(),
    }
}

fn benchmark() {
    let scenes = [
        "plane",
        "cornell-glossy",
        "cornell-water",
        "indirect",
        "conference",
        "sponza",
    ];
    let config = RenderConfig::benchmark();
    let root_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let output_dir = root_dir.join("results");
    std::fs::create_dir_all(output_dir.clone()).unwrap();
    let time_stamp = Local::now().format("%F_%H%M%S").to_string();

    // Initialize an OpenGL context that is needed for post-processing
    let events_loop = glium::glutin::EventsLoop::new();
    // Preferably this wouldn't need use a window at all but alas this is the closest I have gotten.
    // There exists HeadlessContext but that still pops up a window (atleast on Windows).
    // TODO: Maybe change this such that the window displays the current render?
    let window = glium::glutin::WindowBuilder::new()
        .with_dimensions(glium::glutin::dpi::LogicalSize::new(0.0, 0.0))
        .with_visibility(false)
        .with_decorations(false)
        .with_title("Benchmark");
    let context = glium::glutin::ContextBuilder::new();
    let display = glium::Display::new(window, context, &events_loop).unwrap();

    for scene_name in &scenes {
        stats::new_scene(scene_name);
        let _t = stats::time("Total");
        println!("{}...", scene_name);
        let (scene, camera) = load::load_cpu_scene(scene_name, &config);
        let pt_renderer = PTRenderer::offline_render(&display, &scene, &camera, &config);

        stats::time("Post-process");
        let scene_dir = output_dir.join(scene_name);
        std::fs::create_dir_all(scene_dir.clone()).unwrap();
        let timestamped_image = scene_dir.join(format!("{}_{}.png", scene_name, time_stamp));
        pt_renderer.save_image(&display, &timestamped_image);
        // Make a copy to the main output directory
        let default_image = output_dir.join(scene_name).with_extension("png");
        std::fs::copy(timestamped_image, default_image).unwrap();
    }
    let stats_file = output_dir.join(format!("benchmark_{}.txt", time_stamp));
    stats::print_and_save(&stats_file);
}

fn online_render() {
    let mut config = RenderConfig::default();
    let mut events_loop = glium::glutin::EventsLoop::new();
    let window = glium::glutin::WindowBuilder::new()
        .with_dimensions(config.dimensions())
        .with_resizable(false); // TODO: enable resizing
    let context = glium::glutin::ContextBuilder::new().with_depth_buffer(24);
    let display =
        glium::Display::new(window, context, &events_loop).expect("Failed to create display");

    let (mut scene, mut gpu_scene, mut camera) =
        load::load_gpu_scene(VirtualKeyCode::Key1, &display, &config).unwrap();
    let gl_renderer = GLRenderer::new(&display);
    let mut pt_renderer: Option<PTRenderer> = None;

    let mut input = InputState::new();
    let mut quit = false;

    loop {
        let mut target = display.draw();
        target.clear_color_and_depth((0.0, 0.0, 0.0, 1.0), 1.0);
        if let Some(renderer) = &mut pt_renderer {
            renderer.update_and_render(&mut target);
        } else {
            gl_renderer.render(&mut target, &gpu_scene, &camera);
        }
        target.finish().unwrap();

        events_loop.poll_events(|event| {
            input.update(&event);
            match event {
                Event::WindowEvent {
                    event: WindowEvent::KeyboardInput { input, .. },
                    ..
                } => match input {
                    KeyboardInput {
                        state: ElementState::Pressed,
                        virtual_keycode: Some(VirtualKeyCode::Space),
                        ..
                    } => {
                        if pt_renderer.is_some() {
                            pt_renderer = None;
                        } else {
                            pt_renderer =
                                Some(PTRenderer::start_render(&display, &scene, &camera, &config));
                        }
                    }
                    KeyboardInput {
                        state: ElementState::Pressed,
                        virtual_keycode: Some(VirtualKeyCode::C),
                        ..
                    } => println!("camera: {:?}", camera.pos),
                    KeyboardInput {
                        state: ElementState::Pressed,
                        virtual_keycode: Some(keycode),
                        ..
                    } => {
                        if pt_renderer.is_none() {
                            if let Some(res) = load::load_gpu_scene(keycode, &display, &config) {
                                scene = res.0;
                                gpu_scene = res.1;
                                camera = res.2;
                            }
                            config.handle_key(keycode);
                        }
                    }
                    _ => (),
                },
                Event::WindowEvent {
                    event: WindowEvent::CloseRequested,
                    ..
                } => quit = true,
                _ => (),
            }
        });
        camera.process_input(&input);
        input.reset_deltas();
        if quit {
            return;
        }
    }
}
<commit_msg>Don't move the camera when tracing<commit_after>#![feature(duration_float)]
#![feature(euclidean_division)]
#![feature(self_struct_ctor)]
#![feature(try_trait)]

use std::path::PathBuf;

use chrono::Local;

use glium::glutin::{ElementState, Event, KeyboardInput, VirtualKeyCode, WindowEvent};
use glium::Surface;

mod aabb;
mod bsdf;
mod bvh;
mod camera;
mod color;
mod config;
mod consts;
mod float;
mod gl_renderer;
mod index_ptr;
mod input;
mod intersect;
mod light;
mod load;
mod material;
mod mesh;
mod obj_load;
mod pt_renderer;
mod scattering;
mod scene;
mod stats;
mod texture;
mod triangle;
mod vertex;

use self::config::RenderConfig;
use self::gl_renderer::GLRenderer;
use self::input::InputState;
use self::pt_renderer::PTRenderer;

fn main() {
    match std::env::args().nth(1).as_ref().map(|s| s.as_str()) {
        Some(_) => benchmark(),
        None => online_render(),
    }
}

fn benchmark() {
    let scenes = [
        "plane",
        "cornell-glossy",
        "cornell-water",
        "indirect",
        "conference",
        "sponza",
    ];
    let config = RenderConfig::benchmark();
    let root_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let output_dir = root_dir.join("results");
    std::fs::create_dir_all(output_dir.clone()).unwrap();
    let time_stamp = Local::now().format("%F_%H%M%S").to_string();

    // Initialize an OpenGL context that is needed for post-processing
    let events_loop = glium::glutin::EventsLoop::new();
    // Preferably this wouldn't need use a window at all but alas this is the closest I have gotten.
    // There exists HeadlessContext but that still pops up a window (atleast on Windows).
    // TODO: Maybe change this such that the window displays the current render?
    let window = glium::glutin::WindowBuilder::new()
        .with_dimensions(glium::glutin::dpi::LogicalSize::new(0.0, 0.0))
        .with_visibility(false)
        .with_decorations(false)
        .with_title("Benchmark");
    let context = glium::glutin::ContextBuilder::new();
    let display = glium::Display::new(window, context, &events_loop).unwrap();

    for scene_name in &scenes {
        stats::new_scene(scene_name);
        let _t = stats::time("Total");
        println!("{}...", scene_name);
        let (scene, camera) = load::load_cpu_scene(scene_name, &config);
        let pt_renderer = PTRenderer::offline_render(&display, &scene, &camera, &config);

        stats::time("Post-process");
        let scene_dir = output_dir.join(scene_name);
        std::fs::create_dir_all(scene_dir.clone()).unwrap();
        let timestamped_image = scene_dir.join(format!("{}_{}.png", scene_name, time_stamp));
        pt_renderer.save_image(&display, &timestamped_image);
        // Make a copy to the main output directory
        let default_image = output_dir.join(scene_name).with_extension("png");
        std::fs::copy(timestamped_image, default_image).unwrap();
    }
    let stats_file = output_dir.join(format!("benchmark_{}.txt", time_stamp));
    stats::print_and_save(&stats_file);
}

fn online_render() {
    let mut config = RenderConfig::default();
    let mut events_loop = glium::glutin::EventsLoop::new();
    let window = glium::glutin::WindowBuilder::new()
        .with_dimensions(config.dimensions())
        .with_resizable(false); // TODO: enable resizing
    let context = glium::glutin::ContextBuilder::new().with_depth_buffer(24);
    let display =
        glium::Display::new(window, context, &events_loop).expect("Failed to create display");

    let (mut scene, mut gpu_scene, mut camera) =
        load::load_gpu_scene(VirtualKeyCode::Key1, &display, &config).unwrap();
    let gl_renderer = GLRenderer::new(&display);
    let mut pt_renderer: Option<PTRenderer> = None;

    let mut input = InputState::new();
    let mut quit = false;

    loop {
        let mut target = display.draw();
        target.clear_color_and_depth((0.0, 0.0, 0.0, 1.0), 1.0);
        if let Some(renderer) = &mut pt_renderer {
            renderer.update_and_render(&mut target);
        } else {
            gl_renderer.render(&mut target, &gpu_scene, &camera);
        }
        target.finish().unwrap();

        events_loop.poll_events(|event| {
            input.update(&event);
            match event {
                Event::WindowEvent {
                    event: WindowEvent::KeyboardInput { input, .. },
                    ..
                } => match input {
                    KeyboardInput {
                        state: ElementState::Pressed,
                        virtual_keycode: Some(VirtualKeyCode::Space),
                        ..
                    } => {
                        if pt_renderer.is_some() {
                            pt_renderer = None;
                        } else {
                            pt_renderer =
                                Some(PTRenderer::start_render(&display, &scene, &camera, &config));
                        }
                    }
                    KeyboardInput {
                        state: ElementState::Pressed,
                        virtual_keycode: Some(VirtualKeyCode::C),
                        ..
                    } => println!("camera: {:?}", camera.pos),
                    KeyboardInput {
                        state: ElementState::Pressed,
                        virtual_keycode: Some(keycode),
                        ..
                    } => {
                        if pt_renderer.is_none() {
                            if let Some(res) = load::load_gpu_scene(keycode, &display, &config) {
                                scene = res.0;
                                gpu_scene = res.1;
                                camera = res.2;
                            }
                            config.handle_key(keycode);
                        }
                    }
                    _ => (),
                },
                Event::WindowEvent {
                    event: WindowEvent::CloseRequested,
                    ..
                } => quit = true,
                _ => (),
            }
        });
        if pt_renderer.is_none() {
            camera.process_input(&input);
        }
        input.reset_deltas();
        if quit {
            return;
        }
    }
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

use log::{
    debug,
    LevelFilter
};
use structopt::StructOpt;

mod command;
mod util;
mod template;
mod version;
mod vut;

use vut::BumpVersion;

#[derive(StructOpt, Debug)]
#[structopt(name = "Vut", version = env!("CARGO_PKG_VERSION"), author = env!("CARGO_PKG_AUTHORS"))]
struct Opt {
    #[structopt(short = "v", parse(from_occurrences), help = "Verbosity")]
    verbosity: u8,
    #[structopt(subcommand)]
    command: Command,
}

#[derive(StructOpt, Debug)]
enum Command {
    #[structopt(name = "init", about = "Initialize version file")]
    Init {
        #[structopt(name = "version", about = "Specify initial version")]
        version: Option<String>,
    },

    #[structopt(name = "get", help = "Get version")]
    Get {
        #[structopt(name = "format", about = "Specify output format")]
        format: String,
    },

    #[structopt(name = "set", help = "Set version")]
    Set {
        #[structopt(name = "version", about = "Specify version to set")]
        version: String,
    },

    #[structopt(name = "bump", about = "Bump version")]
    Bump {
        #[structopt(help = "Version to bump (major|minor|prerelease|build)")]
        bump_version: BumpVersion,
    },

    #[structopt(name = "generate", alias = "gen", about = "Generate template output")]
    Generate,
}

fn main() {
    let opt = Opt::from_args();

    // Vary the output based on how many times the user used the "verbose" flag
    // (i.e. 'myprog -v -v -v' or 'myprog -vvv' vs 'myprog -v'
    let log_level = match opt.verbosity {
        0 => LevelFilter::Off,
        1 => LevelFilter::Error,
        2 => LevelFilter::Warn,
        3 => LevelFilter::Info,
        4 => LevelFilter::Debug,
        5 | _ => LevelFilter::Trace,
    };

    // Initialize logging
    initialize_logging(log_level);

    debug!("Debug logging enabled.");

    let cmd_result = match opt.command {
        Command::Bump { bump_version } => command::bump(bump_version),
        Command::Generate => command::generate(),
        Command::Get { format } => command::get(&format),
        Command::Init { version } => command::init(version.as_ref().map(|v| v.as_str())),
        Command::Set { version } => command::set(&version),
    };

    match cmd_result {
        Ok(_) => { },
        Err(err) => {
            // Print error description to stderr
            eprintln!("{}", err.description);

            // Return the exit code that corresponds to the error kind
            std::process::exit(err.kind.exit_code());
        },
    };
}

fn initialize_logging(our_level_filter: LevelFilter) {
    use chrono::Utc;

    fern::Dispatch::new()
        .level(our_level_filter)
        .chain(std::io::stderr())
        .format(|out, message, record| {
            out.finish(format_args!(
                "{} | {} | {} | {}",
                Utc::now().format("%Y-%m-%d %H:%M:%S%.3f"),
                record.target(),
                record.level(),
                message
            ))
        })
        .apply()
        .unwrap();
}
<commit_msg>fix: CLI usage info<commit_after>#![allow(dead_code)]

use log::{
    debug,
    LevelFilter
};
use structopt::StructOpt;

mod command;
mod util;
mod template;
mod version;
mod vut;

use vut::BumpVersion;

#[derive(StructOpt, Debug)]
#[structopt(name = "Vut", version = env!("CARGO_PKG_VERSION"), author = env!("CARGO_PKG_AUTHORS"))]
struct Opt {
    #[structopt(short = "v", parse(from_occurrences), help = "Verbosity")]
    verbosity: u8,
    #[structopt(subcommand)]
    command: Command,
}

#[derive(StructOpt, Debug)]
enum Command {
    #[structopt(name = "init", about = "Initialize version file")]
    Init {
        #[structopt(name = "version", about = "Specify initial version")]
        version: Option<String>,
    },

    #[structopt(name = "get", about = "Get version")]
    Get {
        #[structopt(name = "format", about = "Output format (json)")]
        format: String,
    },

    #[structopt(name = "set", about = "Set version")]
    Set {
        #[structopt(name = "version", about = "Version to set")]
        version: String,
    },

    #[structopt(name = "bump", about = "Bump version")]
    Bump {
        #[structopt(help = "Version to bump (major|minor|prerelease|build)")]
        bump_version: BumpVersion,
    },

    #[structopt(name = "generate", alias = "gen", about = "Generate template output")]
    Generate,
}

fn main() {
    let opt = Opt::from_args();

    // Vary the output based on how many times the user used the "verbose" flag
    // (i.e. 'myprog -v -v -v' or 'myprog -vvv' vs 'myprog -v'
    let log_level = match opt.verbosity {
        0 => LevelFilter::Off,
        1 => LevelFilter::Error,
        2 => LevelFilter::Warn,
        3 => LevelFilter::Info,
        4 => LevelFilter::Debug,
        5 | _ => LevelFilter::Trace,
    };

    // Initialize logging
    initialize_logging(log_level);

    debug!("Debug logging enabled.");

    let cmd_result = match opt.command {
        Command::Bump { bump_version } => command::bump(bump_version),
        Command::Generate => command::generate(),
        Command::Get { format } => command::get(&format),
        Command::Init { version } => command::init(version.as_ref().map(|v| v.as_str())),
        Command::Set { version } => command::set(&version),
    };

    match cmd_result {
        Ok(_) => { },
        Err(err) => {
            // Print error description to stderr
            eprintln!("{}", err.description);

            // Return the exit code that corresponds to the error kind
            std::process::exit(err.kind.exit_code());
        },
    };
}

fn initialize_logging(our_level_filter: LevelFilter) {
    use chrono::Utc;

    fern::Dispatch::new()
        .level(our_level_filter)
        .chain(std::io::stderr())
        .format(|out, message, record| {
            out.finish(format_args!(
                "{} | {} | {} | {}",
                Utc::now().format("%Y-%m-%d %H:%M:%S%.3f"),
                record.target(),
                record.level(),
                message
            ))
        })
        .apply()
        .unwrap();
}
<|endoftext|>


<commit_before>extern crate clap;
#[macro_use]
extern crate structopt_derive;
extern crate structopt;
extern crate bytes;
extern crate futures;
extern crate tokio_core;
extern crate tokio_io;
extern crate net2;
extern crate num_cpus;
extern crate resolve;

use std::net::SocketAddr;
use std::time::{self, Duration, SystemTime};

use net2::UdpBuilder;
use net2::unix::UnixUdpBuilderExt;
use resolve::resolver;

use std::thread;
use std::sync::atomic::{ATOMIC_USIZE_INIT, AtomicUsize, Ordering};

use tokio_core::reactor::{Core, Interval, Timeout};
use tokio_core::net::{UdpSocket, UdpCodec, TcpStream};
use tokio_io::AsyncWrite;

use bytes::{Bytes, IntoBuf, Buf, BytesMut, BufMut};
use futures::{Stream, Sink, AsyncSink, Future};

use structopt::StructOpt;

pub static CHUNK_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
pub static DROP_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
pub static ERR_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;

struct DgramCodec {
    addr: SocketAddr,
}

impl DgramCodec {
    pub fn new(addr: SocketAddr) -> Self {
        DgramCodec { addr: addr }
    }
}

impl UdpCodec for DgramCodec {
    type In = Bytes;
    type Out = Bytes;

    fn decode(&mut self,
              _src: &SocketAddr,
              buf: &[u8])
              -> ::std::result::Result<Self::In, ::std::io::Error> {
        Ok(buf.into())
    }

    fn encode(&mut self, msg: Self::Out, buf: &mut Vec<u8>) -> SocketAddr {
        use std::io::Read;
        msg.into_buf().reader().read_to_end(buf).unwrap();
        self.addr
    }
}

fn try_resolve(s: &str) -> SocketAddr {
    s.parse()
        .unwrap_or_else(|_| {
            // for name that have failed to be parsed we try to resolve it via DNS
            let mut split = s.split(':');
            let host = split.next().unwrap(); // Split always has first element
            let port = split.next().expect("port not found");
            let port = port.parse().expect("bad port value");

            let first_ip = resolver::resolve_host(host)
                .expect("failed resolving backend name")
                .next()
                .expect("at least one IP address required");
            SocketAddr::new(first_ip, port)
        })
}

#[derive(StructOpt, Debug)]
#[structopt(about = "UDP multiplexor: copies packets from listened address to specified backends dropping packets on any error")]
struct Options {
    #[structopt(short = "l", long = "listen", help = "Address and port to listen to")]
    listen: SocketAddr,

    #[structopt(short = "b", long = "backend", help = "IP and port of a backend to forward data to. Can be specified multiple times.", value_name="IP:PORT")]
    backends: Vec<String>,

    #[structopt(short = "n", long = "nthreads", help = "Number of worker threads, use 0 to use all CPU cores", default_value = "0")]
    nthreads: usize,

    #[structopt(short = "p", long = "pool", help = "Socket pool size", default_value = "16")]
    snum: usize,

    #[structopt(short = "g", long = "greens", help = "Number of green threads per worker hread", default_value = "4")]
    greens: usize,

    #[structopt(short = "s", long = "size", help = "Internal queue buffer size in packets", default_value = "1048576")]
    bufsize: usize,

    #[structopt(short = "t", long = "interval", help = "Stats printing/sending interval, in milliseconds", default_value = "1000")]
    interval: Option<u64>, // u64 has a special meaning in structopt

    #[structopt(short = "P", long = "stats-prefix", help = "Metric name prefix", default_value="")]
    stats_prefix: String,

    #[structopt(short = "S", long = "stats", help = "Graphite plaintext format compatible address:port to send metrics to")]
    stats: Option<String>,
}

fn main() {
    let mut opts = Options::from_args();

    if opts.nthreads == 0 {
        opts.nthreads = num_cpus::get();
    }

    if opts.snum == 0 {
        panic!("Number of sockets cannot be zero")
    }

    if opts.greens == 0 {
        panic!("Number of green threads cannot be zero")
    }

    // Init chunk counter
    // In the main thread start a reporting timer
    let mut core = Core::new().unwrap();
    let handle = core.handle();

    let timer = Interval::new(Duration::from_millis(opts.interval.unwrap()), &handle).unwrap();
    let interval = opts.interval.unwrap() as f64 / 1000f64;

    let stats = opts.stats.map(|ref stats| try_resolve(stats));

    let prefix = opts.stats_prefix;

    let timer = timer.for_each(|()|{
        let chunks = CHUNK_COUNTER.swap(0, Ordering::Relaxed) as f64 / interval;
        let drops = DROP_COUNTER.swap(0, Ordering::Relaxed) as f64 / interval;
        let errors = ERR_COUNTER.swap(0, Ordering::Relaxed) as f64 / interval;

        println!("chunks/drops/errors: {:.2}/{:.2}/{:.2}", chunks, drops, errors);

        match stats {
            Some(addr) => {
                let tcp_timeout = Timeout::new(Duration::from_millis(((interval / 2f64) * 1000f64).floor() as u64), &handle)
                    .unwrap()
                    .map_err(|_| ());
                let prefix = prefix.clone();
                let sender = TcpStream::connect(&addr, &handle).and_then(move |mut conn| {
                    let ts = SystemTime::now()
                        .duration_since(time::UNIX_EPOCH).unwrap();
                    let ts = ts.as_secs().to_string();
                    // prefix length can be various sizes, we need capacity for it,
                    // we also need space for other values that are around:
                    // 2 spaces, 7 chars for metric suffix
                    // around 26 chars for f64 value, but cannot be sure it takes more
                    // around 10 chars for timestamp
                    // so minimal value is ~45, we'll take 128 just in case float is very large
                    // for 3 metrics this must be tripled
                    let mut buf = BytesMut::with_capacity((prefix.len() + 128)*3);

                    buf.put(&prefix);
                    buf.put(".udpdup.chunks ");
                    buf.put(chunks.to_string());
                    buf.put(" ");
                    buf.put(&ts);
                    buf.put("\n");

                    buf.put(&prefix);
                    buf.put(".udpdup.drops ");
                    buf.put(drops.to_string());
                    buf.put(" ");
                    buf.put(&ts);
                    buf.put("\n");

                    buf.put(&prefix);
                    buf.put(".udpdup.errors ");
                    buf.put(errors.to_string());
                    buf.put(" ");
                    buf.put(&ts);

                    conn.write_buf(&mut buf.into_buf()).map(|_|())
                }).map_err(|e| println!("Error sending stats: {:?}", e)).select(tcp_timeout).then(|_|Ok(()));
                handle.spawn(sender);
            }
            None => (),
        };
        Ok(())
    });
    let backends = opts.backends
        .iter()
        .map(|b| try_resolve(b))
        .collect::<Vec<SocketAddr>>();

    // Create a pool of listener sockets
    let mut sockets = Vec::new();
    for _ in 0..opts.snum {
        let socket = UdpBuilder::new_v4().unwrap();
        socket.reuse_address(true).unwrap();
        socket.reuse_port(true).unwrap();
        let socket = socket.bind(&opts.listen).unwrap();
        sockets.push(socket);
    }

    let bufsize = opts.bufsize;
    let listen = opts.listen;
    let greens = opts.greens;
    for i in 0..opts.nthreads {
        let backends = backends.clone();

        // Each thread gets the clone of a socket pool
        let sockets = sockets
            .iter()
            .map(|s| s.try_clone().unwrap())
            .collect::<Vec<_>>();

        thread::Builder::new()
            .name(format!("udpdup_worker{}", i).into())
            .spawn(move || {
                // each thread runs it's own core
                let mut core = Core::new().unwrap();
                let handle = core.handle();

                // Create sockets for each backend address
                let senders = backends
                    .iter()
                    .map(|_| {
                             let socket = UdpBuilder::new_v4().unwrap();
                             let bind: SocketAddr = "0.0.0.0:0".parse().unwrap();
                             socket.bind(&bind).unwrap()
                         })
                    .collect::<Vec<_>>();

                for _ in 0..greens {
                    for socket in sockets.iter() {
                        let senders = senders
                            .iter()
                            .map(|socket| socket.try_clone().expect("client socket cloning"))
                            .collect::<Vec<_>>();

                        // make senders into framed UDP
                        let mut backends = backends.clone();
                        let mut senders = senders
                            .into_iter()
                            .map(|socket| {
                                let sender = UdpSocket::from_socket(socket, &handle).unwrap();
                                // socket order becomes reversed but we don't care about this
                                // because the order is same everywhere
                                // and anyways, we push same data to every socket
                                let sender = sender
                                    .framed(DgramCodec::new(backends.pop().unwrap()))
                                    .buffer(bufsize);
                                sender
                            })
                            .collect::<Vec<_>>();

                        // create server now
                        let socket = socket.try_clone().expect("server socket cloning");
                        let socket = UdpSocket::from_socket(socket, &handle).unwrap();
                        let server = socket.framed(DgramCodec::new(listen)).buffer(bufsize);

                        let server = server.for_each(move |buf| {
                            CHUNK_COUNTER.fetch_add(1, Ordering::Relaxed);
                            senders
                                .iter_mut()
                                .map(|sender| {
                                    let res = sender.start_send(buf.clone());
                                    match res {
                                        Err(_) => {
                                            ERR_COUNTER.fetch_add(1, Ordering::Relaxed);
                                        }
                                        Ok(AsyncSink::NotReady(_)) => {
                                            DROP_COUNTER.fetch_add(1, Ordering::Relaxed);
                                        }
                                        Ok(AsyncSink::Ready) => {
                                            let res = sender.poll_complete();
                                            if res.is_err() {
                                                ERR_COUNTER.fetch_add(1, Ordering::Relaxed);
                                            }
                                        }
                                    }
                                })
                                .last();
                            Ok(())
                        });

                        handle.spawn(server.map_err(|_| ()));
                    }
                }
                core.run(::futures::future::empty::<(), ()>()).unwrap();
            })
            .unwrap();
    }
    drop(sockets);
    core.run(timer).unwrap();
}
<commit_msg>ulimit-safe default settings<commit_after>extern crate clap;
#[macro_use]
extern crate structopt_derive;
extern crate structopt;
extern crate bytes;
extern crate futures;
extern crate tokio_core;
extern crate tokio_io;
extern crate net2;
extern crate num_cpus;
extern crate resolve;

use std::net::SocketAddr;
use std::time::{self, Duration, SystemTime};

use net2::UdpBuilder;
use net2::unix::UnixUdpBuilderExt;
use resolve::resolver;

use std::thread;
use std::sync::atomic::{ATOMIC_USIZE_INIT, AtomicUsize, Ordering};

use tokio_core::reactor::{Core, Interval, Timeout};
use tokio_core::net::{UdpSocket, UdpCodec, TcpStream};
use tokio_io::AsyncWrite;

use bytes::{Bytes, IntoBuf, Buf, BytesMut, BufMut};
use futures::{Stream, Sink, AsyncSink, Future};

use structopt::StructOpt;

pub static CHUNK_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
pub static DROP_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
pub static ERR_COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;

struct DgramCodec {
    addr: SocketAddr,
}

impl DgramCodec {
    pub fn new(addr: SocketAddr) -> Self {
        DgramCodec { addr: addr }
    }
}

impl UdpCodec for DgramCodec {
    type In = Bytes;
    type Out = Bytes;

    fn decode(&mut self,
              _src: &SocketAddr,
              buf: &[u8])
        -> ::std::result::Result<Self::In, ::std::io::Error> {
            Ok(buf.into())
        }

    fn encode(&mut self, msg: Self::Out, buf: &mut Vec<u8>) -> SocketAddr {
        use std::io::Read;
        msg.into_buf().reader().read_to_end(buf).unwrap();
        self.addr
    }
}

fn try_resolve(s: &str) -> SocketAddr {
    s.parse()
        .unwrap_or_else(|_| {
            // for name that have failed to be parsed we try to resolve it via DNS
            let mut split = s.split(':');
            let host = split.next().unwrap(); // Split always has first element
            let port = split.next().expect("port not found");
            let port = port.parse().expect("bad port value");

            let first_ip = resolver::resolve_host(host)
                .expect("failed resolving backend name")
                .next()
                .expect("at least one IP address required");
            SocketAddr::new(first_ip, port)
        })
}

#[derive(StructOpt, Debug)]
#[structopt(about = "UDP multiplexor: copies packets from listened address to specified backends dropping packets on any error")]
struct Options {
    #[structopt(short = "l", long = "listen", help = "Address and port to listen to")]
    listen: SocketAddr,

    #[structopt(short = "b", long = "backend", help = "IP and port of a backend to forward data to. Can be specified multiple times.", value_name="IP:PORT")]
    backends: Vec<String>,

    #[structopt(short = "n", long = "nthreads", help = "Number of worker threads, use 0 to use all CPU cores", default_value = "0")]
    nthreads: usize,

    #[structopt(short = "p", long = "pool", help = "Socket pool size", default_value = "4")]
    snum: usize,

    #[structopt(short = "g", long = "greens", help = "Number of green threads per worker hread", default_value = "4")]
    greens: usize,

    #[structopt(short = "s", long = "size", help = "Internal queue buffer size in packets", default_value = "1048576")]
    bufsize: usize,

    #[structopt(short = "t", long = "interval", help = "Stats printing/sending interval, in milliseconds", default_value = "1000")]
    interval: Option<u64>, // u64 has a special meaning in structopt

    #[structopt(short = "P", long = "stats-prefix", help = "Metric name prefix", default_value="")]
    stats_prefix: String,

    #[structopt(short = "S", long = "stats", help = "Graphite plaintext format compatible address:port to send metrics to")]
    stats: Option<String>,
}

fn main() {
    let mut opts = Options::from_args();

    if opts.nthreads == 0 {
        opts.nthreads = num_cpus::get();
    }

    if opts.snum == 0 {
        panic!("Number of sockets cannot be zero")
    }

    if opts.greens == 0 {
        panic!("Number of green threads cannot be zero")
    }

    // Init chunk counter
    // In the main thread start a reporting timer
    let mut core = Core::new().unwrap();
    let handle = core.handle();

    let timer = Interval::new(Duration::from_millis(opts.interval.unwrap()), &handle).unwrap();
    let interval = opts.interval.unwrap() as f64 / 1000f64;

    let stats = opts.stats.map(|ref stats| try_resolve(stats));

    let prefix = opts.stats_prefix;

    let timer = timer.for_each(|()|{
        let chunks = CHUNK_COUNTER.swap(0, Ordering::Relaxed) as f64 / interval;
        let drops = DROP_COUNTER.swap(0, Ordering::Relaxed) as f64 / interval;
        let errors = ERR_COUNTER.swap(0, Ordering::Relaxed) as f64 / interval;

        println!("chunks/drops/errors: {:.2}/{:.2}/{:.2}", chunks, drops, errors);

        match stats {
            Some(addr) => {
                let tcp_timeout = Timeout::new(Duration::from_millis(((interval / 2f64) * 1000f64).floor() as u64), &handle)
                    .unwrap()
                    .map_err(|_| ());
                let prefix = prefix.clone();
                let sender = TcpStream::connect(&addr, &handle).and_then(move |mut conn| {
                    let ts = SystemTime::now()
                        .duration_since(time::UNIX_EPOCH).unwrap();
                    let ts = ts.as_secs().to_string();
                    // prefix length can be various sizes, we need capacity for it,
                    // we also need space for other values that are around:
                    // 2 spaces, 7 chars for metric suffix
                    // around 26 chars for f64 value, but cannot be sure it takes more
                    // around 10 chars for timestamp
                    // so minimal value is ~45, we'll take 128 just in case float is very large
                    // for 3 metrics this must be tripled
                    let mut buf = BytesMut::with_capacity((prefix.len() + 128)*3);

                    buf.put(&prefix);
                    buf.put(".udpdup.chunks ");
                    buf.put(chunks.to_string());
                    buf.put(" ");
                    buf.put(&ts);
                    buf.put("\n");

                    buf.put(&prefix);
                    buf.put(".udpdup.drops ");
                    buf.put(drops.to_string());
                    buf.put(" ");
                    buf.put(&ts);
                    buf.put("\n");

                    buf.put(&prefix);
                    buf.put(".udpdup.errors ");
                    buf.put(errors.to_string());
                    buf.put(" ");
                    buf.put(&ts);

                    conn.write_buf(&mut buf.into_buf()).map(|_|())
                }).map_err(|e| println!("Error sending stats: {:?}", e)).select(tcp_timeout).then(|_|Ok(()));
                handle.spawn(sender);
            }
            None => (),
        };
        Ok(())
    });
    let backends = opts.backends
        .iter()
        .map(|b| try_resolve(b))
        .collect::<Vec<SocketAddr>>();

    // Create a pool of listener sockets
    let mut sockets = Vec::new();
    for _ in 0..opts.snum {
        let socket = UdpBuilder::new_v4().unwrap();
        socket.reuse_address(true).unwrap();
        socket.reuse_port(true).unwrap();
        let socket = socket.bind(&opts.listen).unwrap();
        sockets.push(socket);
    }

    let bufsize = opts.bufsize;
    let listen = opts.listen;
    let greens = opts.greens;
    for i in 0..opts.nthreads {
        let backends = backends.clone();

        // Each thread gets the clone of a socket pool
        let sockets = sockets
            .iter()
            .map(|s| s.try_clone().unwrap())
            .collect::<Vec<_>>();

        thread::Builder::new()
            .name(format!("udpdup_worker{}", i).into())
            .spawn(move || {
                // each thread runs it's own core
                let mut core = Core::new().unwrap();
                let handle = core.handle();

                // Create sockets for each backend address
                let senders = backends
                    .iter()
                    .map(|_| {
                        let socket = UdpBuilder::new_v4().unwrap();
                        let bind: SocketAddr = "0.0.0.0:0".parse().unwrap();
                        socket.bind(&bind).unwrap()
                    })
                .collect::<Vec<_>>();

                for _ in 0..greens {
                    for socket in sockets.iter() {
                        let senders = senders
                            .iter()
                            .map(|socket| socket.try_clone().expect("client socket cloning"))
                            .collect::<Vec<_>>();

                        // make senders into framed UDP
                        let mut backends = backends.clone();
                        let mut senders = senders
                            .into_iter()
                            .map(|socket| {
                                let sender = UdpSocket::from_socket(socket, &handle).unwrap();
                                // socket order becomes reversed but we don't care about this
                                // because the order is same everywhere
                                // and anyways, we push same data to every socket
                                let sender = sender
                                    .framed(DgramCodec::new(backends.pop().unwrap()))
                                    .buffer(bufsize);
                                sender
                            })
                        .collect::<Vec<_>>();

                        // create server now
                        let socket = socket.try_clone().expect("server socket cloning");
                        let socket = UdpSocket::from_socket(socket, &handle).unwrap();
                        let server = socket.framed(DgramCodec::new(listen)).buffer(bufsize);

                        let server = server.for_each(move |buf| {
                            CHUNK_COUNTER.fetch_add(1, Ordering::Relaxed);
                            senders
                                .iter_mut()
                                .map(|sender| {
                                    let res = sender.start_send(buf.clone());
                                    match res {
                                        Err(_) => {
                                            ERR_COUNTER.fetch_add(1, Ordering::Relaxed);
                                        }
                                        Ok(AsyncSink::NotReady(_)) => {
                                            DROP_COUNTER.fetch_add(1, Ordering::Relaxed);
                                        }
                                        Ok(AsyncSink::Ready) => {
                                            let res = sender.poll_complete();
                                            if res.is_err() {
                                                ERR_COUNTER.fetch_add(1, Ordering::Relaxed);
                                            }
                                        }
                                    }
                                })
                            .last();
                            Ok(())
                        });

                        handle.spawn(server.map_err(|_| ()));
                    }
                }
                core.run(::futures::future::empty::<(), ()>()).unwrap();
            })
        .unwrap();
    }
    drop(sockets);
    core.run(timer).unwrap();
}
<|endoftext|>


<commit_before>// Copyright 2016 rust-fuzz developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

extern crate cargo_metadata;
extern crate docopt;
extern crate rustc_serialize;

use cargo_metadata::{metadata, Package};
use docopt::Docopt;
use std::{env, fs, io, path, process};
use std::io::Write;

const USAGE: &'static str = "
Cargo Fuzz

Usage:
  cargo fuzz --init
  cargo fuzz --fuzz-target TARGET
  cargo fuzz --add TARGET
  cargo fuzz (-h | --help)

Options:
  -h --help              Show this screen.
  --init                 Initialize fuzz folder
  --fuzz-target TARGET   Run with given fuzz target in fuzz/fuzzers
  --add TARGET           Add a new fuzz target
";

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_init: bool,
    flag_add: Option<String>,
    flag_fuzz_target: Option<String>,
}

fn main() {
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    let result = if args.flag_init {
        init_fuzz()
    } else if let Some(target) = args.flag_add {
        add_target(target)
    } else if let Some(target) = args.flag_fuzz_target {
        let result = run_target(target);
        if let Ok(success) = result {
            if success {
                // Can this ever happen?
                Ok(())
            } else {
                println!("Fuzzing found errors!");
                process::exit(-1)
            }
        } else {
            result.map(|_| ())
        }
    } else {
        println!("Invalid arguments. Usage:\n{}", USAGE);
        return;
    };
    if let Err(error) = result {
        println!("Error: {:?}", error);
    }
}

/// Create all the files and folders we need to run
///
/// This will not clone libfuzzer-sys
fn init_fuzz() -> io::Result<()> {
    let me = get_package();

    fs::create_dir("./fuzz")?;
    fs::create_dir("./fuzz/fuzzers")?;

    let mut cargo = fs::File::create(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[package]
name = "{0}-fuzz"
version = "0.0.1"
authors = ["Automatically generated"]
publish = false

[dependencies.{0}]
path = ".."

# Prevent this from interfering with workspaces
[workspace]
members = ["."]

[[bin]]
name = "fuzzer_script_1"
path = "fuzzers/fuzzer_script_1.rs"
"#, me.name)?;

    let mut ignore = fs::File::create(path::Path::new("./fuzz/.gitignore"))?;

write!(ignore, r#"
target
libfuzzer
corpus
"#)?;

    let mut script = fs::File::create(path::Path::new("./fuzz/fuzzers/fuzzer_script_1.rs"))?;
    dummy_target(&mut script, &me)
}

/// Returns metadata for the Cargo package in the current directory
fn get_package() -> Package {
    // todo error handling
    let meta = metadata(None).unwrap();
    let mut p = env::current_dir().unwrap();
    p.push("Cargo.toml");
    let p = p.to_str().unwrap();
    meta.packages.into_iter().find(|package| package.manifest_path == p).unwrap()
}

/// If the package contains a library target, generate an `extern crate` line to link to it.
fn link_to_lib(pkg: &Package) -> Option<String> {
    pkg.targets.iter()
               .find(|target| target.kind.iter().any(|k| k == "lib"))
               .map(|target| format!("extern crate {};\n", target.name.replace("-", "_")))
}

/// Create a dummy fuzz target script at the given path
fn dummy_target(script: &mut fs::File, pkg: &Package) -> io::Result<()> {
write!(script, r#"#![no_main]
extern crate libfuzzer_sys;
{}
#[export_name="rust_fuzzer_test_input"]
pub extern fn go(data: &[u8]) {{
    // fuzzed code goes here
}}"#, link_to_lib(pkg).unwrap_or(String::new()))
}

/// Add a new fuzz target script with a given name
fn add_target(target: String) -> io::Result<()> {
    let target_file = format!("fuzz/fuzzers/{}.rs", target);
    let mut script = fs::File::create(path::Path::new(&target_file))?;
    let me = get_package();
    dummy_target(&mut script, &me)?;

    let mut cargo = fs::OpenOptions::new().append(true).open(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[[bin]]
name = "{0}"
path = "fuzzers/{0}.rs"
"#, target)

}

/// Build or rebuild libFuzzer (rebuilds only if the compiler version changed)
///
/// We can't just use libFuzzer as a dependency since libgcc will
/// get compiled with sanitizer support. RUSTFLAGS does not discriminate
/// between build dependencies and regular ones.
///
/// https://github.com/rust-lang/cargo/issues/3739
fn rebuild_libfuzzer() -> io::Result<()> {
    if let Err(_) = env::set_current_dir("./libfuzzer") {
        let mut git = process::Command::new("git");
        let mut cmd = git.arg("clone")
                         .arg("https://github.com/rust-fuzz/libfuzzer-sys.git")
                         .arg("libfuzzer");
        let result = cmd.spawn()?.wait()?;
        if !result.success() {
            return Err(io::Error::new(io::ErrorKind::Other,
                                      "Failed to clone libfuzzer-sys"))
        }
        env::set_current_dir("./libfuzzer")?;
    }
    let mut cmd = process::Command::new("cargo");
    cmd.arg("build")
       .arg("--release")
       .spawn()?
       .wait()?;

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err(io::Error::new(io::ErrorKind::Other,
                                  "Failed to build libfuzzer-sys"))
    }
    env::set_current_dir("..")
}

/// Fuzz a given fuzz target
fn run_target(target: String) -> io::Result<bool> {
    env::set_current_dir("./fuzz")?;
    rebuild_libfuzzer()?;
    let mut flags = env::var("RUSTFLAGS").unwrap_or("".into());
    if !flags.is_empty() {
        flags.push(' ');
    }
    flags.push_str("-Cpasses=sancov -Cllvm-args=-sanitizer-coverage-level=3 -Zsanitizer=address -Cpanic=abort");
    let mut cmd = process::Command::new("cargo");
    cmd.arg("rustc")
       .arg("--verbose")
       .arg("--bin")
       .arg(&target)
       .arg("--")
       .arg("-L")
       .arg("libfuzzer/target/release")
       .env("RUSTFLAGS", &flags);

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err(io::Error::new(io::ErrorKind::Other, "Failed to build fuzz target"))
    }

    if let Err(k) = fs::create_dir("corpus") {
        if k.kind() == io::ErrorKind::AlreadyExists {
            // do nothing
        } else {
            return Err(k);
        }
    }

    // can't use cargo run since we can't pass -L args to it
    let path = format!("target/debug/{}", target);
    let mut run_cmd = process::Command::new(path);
    run_cmd.arg("corpus");
    let result = run_cmd.spawn()?.wait()?;
    Ok(result.success())
}
<commit_msg>Add `--list` flag for `cargo fuzz`.<commit_after>// Copyright 2016 rust-fuzz developers
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

extern crate cargo_metadata;
extern crate docopt;
extern crate rustc_serialize;

use cargo_metadata::{metadata, Package};
use docopt::Docopt;
use std::{env, error, fs, io, path, process};
use std::io::Write;

const USAGE: &'static str = "
Cargo Fuzz

Usage:
  cargo fuzz --init
  cargo fuzz --fuzz-target TARGET
  cargo fuzz --add TARGET
  cargo fuzz --list
  cargo fuzz (-h | --help)

Options:
  -h --help              Show this screen.
  --init                 Initialize fuzz folder
  --fuzz-target TARGET   Run with given fuzz target in fuzz/fuzzers
  --add TARGET           Add a new fuzz target
  --list                 List the available fuzz targets
";

#[derive(Debug, RustcDecodable)]
struct Args {
    flag_init: bool,
    flag_add: Option<String>,
    flag_fuzz_target: Option<String>,
    flag_list: bool,
}

fn main() {
    let args: Args = Docopt::new(USAGE)
                            .and_then(|d| d.decode())
                            .unwrap_or_else(|e| e.exit());

    let result = if args.flag_init {
        init_fuzz()
    } else if let Some(target) = args.flag_add {
        add_target(target)
    } else if let Some(target) = args.flag_fuzz_target {
        let result = run_target(target);
        if let Ok(success) = result {
            if success {
                // Can this ever happen?
                Ok(())
            } else {
                println!("Fuzzing found errors!");
                process::exit(-1)
            }
        } else {
            result.map(|_| ())
        }
    } else if args.flag_list {
        list_fuzz_targets()
            .map(|_| ())
    } else {
        println!("Invalid arguments. Usage:\n{}", USAGE);
        return;
    };
    if let Err(error) = result {
        writeln!(io::stderr(), "Error: {}", error)
            .expect("failed writing to stderr");
    }
}

fn list_fuzz_targets() -> Result<(), Box<error::Error>> {
    if !path::Path::new("./fuzz").is_dir() {
        return Err("Fuzzing crate has not been initialized. Run `cargo fuzz --init` to initialize it.".into());
    }
    env::set_current_dir("./fuzz")?;
    let package = get_package();
    for target in &package.targets {
        println!("{}", target.name);
    }
    Ok(())
}

/// Create all the files and folders we need to run
///
/// This will not clone libfuzzer-sys
fn init_fuzz() -> Result<(), Box<error::Error>> {
    let me = get_package();

    fs::create_dir("./fuzz")?;
    fs::create_dir("./fuzz/fuzzers")?;

    let mut cargo = fs::File::create(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[package]
name = "{0}-fuzz"
version = "0.0.1"
authors = ["Automatically generated"]
publish = false

[dependencies.{0}]
path = ".."

# Prevent this from interfering with workspaces
[workspace]
members = ["."]

[[bin]]
name = "fuzzer_script_1"
path = "fuzzers/fuzzer_script_1.rs"
"#, me.name)?;

    let mut ignore = fs::File::create(path::Path::new("./fuzz/.gitignore"))?;

write!(ignore, r#"
target
libfuzzer
corpus
"#)?;

    let mut script = fs::File::create(path::Path::new("./fuzz/fuzzers/fuzzer_script_1.rs"))?;
    dummy_target(&mut script, &me)
}

/// Returns metadata for the Cargo package in the current directory
fn get_package() -> Package {
    // todo error handling
    let meta = metadata(None).unwrap();
    let mut p = env::current_dir().unwrap();
    p.push("Cargo.toml");
    let p = p.to_str().unwrap();
    meta.packages.into_iter().find(|package| package.manifest_path == p).unwrap()
}

/// If the package contains a library target, generate an `extern crate` line to link to it.
fn link_to_lib(pkg: &Package) -> Option<String> {
    pkg.targets.iter()
               .find(|target| target.kind.iter().any(|k| k == "lib"))
               .map(|target| format!("extern crate {};\n", target.name.replace("-", "_")))
}

/// Create a dummy fuzz target script at the given path
fn dummy_target(script: &mut fs::File, pkg: &Package) -> Result<(), Box<error::Error>> {
write!(script, r#"#![no_main]
extern crate libfuzzer_sys;
{}
#[export_name="rust_fuzzer_test_input"]
pub extern fn go(data: &[u8]) {{
    // fuzzed code goes here
}}"#, link_to_lib(pkg).unwrap_or(String::new())).map_err(|e| e.into())
}

/// Add a new fuzz target script with a given name
fn add_target(target: String) -> Result<(), Box<error::Error>> {
    let target_file = format!("fuzz/fuzzers/{}.rs", target);
    let mut script = fs::File::create(path::Path::new(&target_file))?;
    let me = get_package();
    dummy_target(&mut script, &me)?;

    let mut cargo = fs::OpenOptions::new().append(true).open(path::Path::new("./fuzz/Cargo.toml"))?;

write!(cargo, r#"
[[bin]]
name = "{0}"
path = "fuzzers/{0}.rs"
"#, target).map_err(|e| e.into())

}

/// Build or rebuild libFuzzer (rebuilds only if the compiler version changed)
///
/// We can't just use libFuzzer as a dependency since libgcc will
/// get compiled with sanitizer support. RUSTFLAGS does not discriminate
/// between build dependencies and regular ones.
///
/// https://github.com/rust-lang/cargo/issues/3739
fn rebuild_libfuzzer() -> Result<(), Box<error::Error>> {
    if let Err(_) = env::set_current_dir("./libfuzzer") {
        let mut git = process::Command::new("git");
        let mut cmd = git.arg("clone")
                         .arg("https://github.com/rust-fuzz/libfuzzer-sys.git")
                         .arg("libfuzzer");
        let result = cmd.spawn()?.wait()?;
        if !result.success() {
            return Err("Failed to clone libfuzzer-sys".into())
        }
        env::set_current_dir("./libfuzzer")?;
    }
    let mut cmd = process::Command::new("cargo");
    cmd.arg("build")
       .arg("--release")
       .spawn()?
       .wait()?;

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err("Failed to build libfuzzer-sys".into())
    }
    env::set_current_dir("..")
        .map_err(|e| e.into())
}

/// Fuzz a given fuzz target
fn run_target(target: String) -> Result<bool, Box<error::Error>> {
    env::set_current_dir("./fuzz")?;
    rebuild_libfuzzer()?;
    let mut flags = env::var("RUSTFLAGS").unwrap_or("".into());
    if !flags.is_empty() {
        flags.push(' ');
    }
    flags.push_str("-Cpasses=sancov -Cllvm-args=-sanitizer-coverage-level=3 -Zsanitizer=address -Cpanic=abort");
    let mut cmd = process::Command::new("cargo");
    cmd.arg("rustc")
       .arg("--verbose")
       .arg("--bin")
       .arg(&target)
       .arg("--")
       .arg("-L")
       .arg("libfuzzer/target/release")
       .env("RUSTFLAGS", &flags);

    let result = cmd.spawn()?.wait()?;
    if !result.success() {
        return Err("Failed to build fuzz target".into())
    }

    if let Err(k) = fs::create_dir("corpus") {
        if k.kind() == io::ErrorKind::AlreadyExists {
            // do nothing
        } else {
            return Err(k.into());
        }
    }

    // can't use cargo run since we can't pass -L args to it
    let path = format!("target/debug/{}", target);
    let mut run_cmd = process::Command::new(path);
    run_cmd.arg("corpus");
    let result = run_cmd.spawn()?.wait()?;
    Ok(result.success())
}
<|endoftext|>