<commit_before>// Copyright 2014 The html5ever Project Developers. See the
// COPYRIGHT file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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.

//! A simple reference-counted DOM.
//!
//! This is sufficient as a static parse tree, but don't build a
//! web browser using it. :)

use sink::common::{NodeEnum, Document, Doctype, Text, Comment, Element};

use util::atom::Atom;
use util::namespace::{Namespace, HTML};
use tokenizer::Attribute;
use tree_builder::{TreeSink, QuirksMode, NodeOrText, AppendNode, AppendText};
use tree_builder;
use serialize::{Serializable, Serializer};
use driver::ParseResult;

use std::rc::{Rc, Weak};
use std::cell::RefCell;
use std::default::Default;
use std::io::IoResult;

/// A DOM node.
pub struct Node {
    pub node: NodeEnum,
    pub parent: Option<WeakHandle>,
    pub children: Vec<Handle>,

    /// The "script already started" flag.
    ///
    /// Not meaningful for nodes other than HTML `<script>`.
    pub script_already_started: bool,
}

impl Node {
    fn new(node: NodeEnum) -> Node {
        Node {
            node: node,
            parent: None,
            children: vec!(),
            script_already_started: false,
        }
    }

    fn parent(&self) -> Handle {
        self.parent.as_ref().expect("no parent!")
            .upgrade().expect("dangling weak pointer!")
    }
}

/// Reference to a DOM node.
pub type Handle = Rc<RefCell<Node>>;

/// Weak reference to a DOM node, used for parent pointers.
pub type WeakHandle = Weak<RefCell<Node>>;

fn same_node(x: &Handle, y: &Handle) -> bool {
    // FIXME: This shouldn't really need to touch the borrow flags, right?
    (&*x.borrow() as *const Node) == (&*y.borrow() as *const Node)
}

fn new_node(node: NodeEnum) -> Handle {
    Rc::new(RefCell::new(Node::new(node)))
}

fn append(new_parent: &Handle, child: Handle) {
    new_parent.borrow_mut().children.push(child.clone());
    let parent = &mut child.borrow_mut().parent;
    assert!(parent.is_none());
    *parent = Some(new_parent.downgrade());
}

fn get_parent_and_index(target: &Handle) -> Option<(Handle, uint)> {
    let child = target.borrow();
    let parent = child.parent();
    match parent.borrow_mut().children.iter().enumerate()
                .find(|&(_, n)| same_node(n, target)) {
        Some((i, _)) => Some((parent, i)),
        None => None,
    }
}

fn append_to_existing_text(prev: &Handle, text: &str) -> bool {
    match prev.borrow_mut().deref_mut().node {
        Text(ref mut existing) => {
            existing.push_str(text);
            true
        }
        _ => false,
    }
}

fn remove_from_parent(target: &Handle) {
    {
        let (parent, i) = unwrap_or_return!(get_parent_and_index(target), ());
        parent.borrow_mut().children.remove(i).expect("not found!");
    }

    let mut child = target.borrow_mut();
    (*child).parent = None;
}

/// The DOM itself; the result of parsing.
pub struct RcDom {
    /// The `Document` itself.
    pub document: Handle,

    /// Errors that occurred during parsing.
    pub errors: Vec<String>,

    /// The document's quirks mode.
    pub quirks_mode: QuirksMode,
}

impl TreeSink<Handle> for RcDom {
    fn parse_error(&mut self, msg: String) {
        self.errors.push(msg);
    }

    fn get_document(&mut self) -> Handle {
        self.document.clone()
    }

    fn set_quirks_mode(&mut self, mode: QuirksMode) {
        self.quirks_mode = mode;
    }

    fn same_node(&self, x: Handle, y: Handle) -> bool {
        same_node(&x, &y)
    }

    fn elem_name(&self, target: Handle) -> (Namespace, Atom) {
        match target.borrow().node {
            Element(ref name, _) => (HTML, name.clone()),
            _ => fail!("not an element!"),
        }
    }

    fn create_element(&mut self, ns: Namespace, name: Atom, attrs: Vec<Attribute>) -> Handle {
        assert!(ns == HTML);
        new_node(Element(name, attrs))
    }

    fn create_comment(&mut self, text: String) -> Handle {
        new_node(Comment(text))
    }

    fn append(&mut self, parent: Handle, child: NodeOrText<Handle>) {
        // Append to an existing Text node if we have one.
        match child {
            AppendText(ref text) => match parent.borrow().children.last() {
                Some(h) => if append_to_existing_text(h, text.as_slice()) { return; },
                _ => (),
            },
            _ => (),
        }

        append(&parent, match child {
            AppendText(text) => new_node(Text(text)),
            AppendNode(node) => node
        });
    }

    fn append_before_sibling(&mut self,
            sibling: Handle,
            child: NodeOrText<Handle>) -> Result<(), NodeOrText<Handle>> {
        let (parent, i) = unwrap_or_return!(get_parent_and_index(&sibling), Err(child));

        let child = match (child, i) {
            // No previous node.
            (AppendText(text), 0) => new_node(Text(text)),

            // Look for a text node before the insertion point.
            (AppendText(text), i) => {
                let parent = parent.borrow();
                let prev = parent.children.get(i-1);
                if append_to_existing_text(prev, text.as_slice()) {
                    return Ok(());
                }
                new_node(Text(text))
            }

            // The tree builder promises we won't have a text node after
            // the insertion point.

            // Any other kind of node.
            (AppendNode(node), _) => node,
        };

        if child.borrow().parent.is_some() {
            remove_from_parent(&child);
        }

        child.borrow_mut().parent = Some(parent.clone().downgrade());
        parent.borrow_mut().children.insert(i, child);
        Ok(())
    }

    fn append_doctype_to_document(&mut self, name: String, public_id: String, system_id: String) {
        append(&self.document, new_node(Doctype(name, public_id, system_id)));
    }

    fn add_attrs_if_missing(&mut self, target: Handle, mut attrs: Vec<Attribute>) {
        let mut node = target.borrow_mut();
        // FIXME: mozilla/rust#15609
        let existing = match node.deref_mut().node {
            Element(_, ref mut attrs) => attrs,
            _ => return,
        };

        // FIXME: quadratic time
        attrs.retain(|attr|
            !existing.iter().any(|e| e.name == attr.name));
        existing.push_all_move(attrs);
    }

    fn remove_from_parent(&mut self, target: Handle) {
        remove_from_parent(&target);
    }

    fn mark_script_already_started(&mut self, node: Handle) {
        node.borrow_mut().script_already_started = true;
    }
}

impl Default for RcDom {
    fn default() -> RcDom {
        RcDom {
            document: new_node(Document),
            errors: vec!(),
            quirks_mode: tree_builder::NoQuirks,
        }
    }
}

impl ParseResult<RcDom> for RcDom {
    fn get_result(sink: RcDom) -> RcDom {
        sink
    }
}

impl Serializable for Handle {
    fn serialize<'wr, Wr: Writer>(&self, serializer: &mut Serializer<'wr, Wr>, incl_self: bool) -> IoResult<()> {
        let node = self.borrow();
        match (incl_self, &node.node) {
            (_, &Element(ref name, ref attrs)) => {
                if incl_self {
                    try!(serializer.start_elem(HTML, name.clone(),
                        attrs.iter().map(|at| (&at.name, at.value.as_slice()))));
                }

                for handle in node.children.iter() {
                    try!(handle.clone().serialize(serializer, true));
                }

                if incl_self {
                    try!(serializer.end_elem(HTML, name.clone()));
                }
                Ok(())
            }

            (false, &Document) => {
                for handle in node.children.iter() {
                    try!(handle.clone().serialize(serializer, true));
                }
                Ok(())
            }

            (false, _) => Ok(()),

            (true, &Doctype(ref name, _, _)) => serializer.write_doctype(name.as_slice()),
            (true, &Text(ref text)) => serializer.write_text(text.as_slice()),
            (true, &Comment(ref text)) => serializer.write_comment(text.as_slice()),

            (true, &Document) => fail!("Can't serialize Document node itself"),
        }
    }
}
<commit_msg>RcDom fixes<commit_after>// Copyright 2014 The html5ever Project Developers. See the
// COPYRIGHT file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.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.

//! A simple reference-counted DOM.
//!
//! This is sufficient as a static parse tree, but don't build a
//! web browser using it. :)

use sink::common::{NodeEnum, Document, Doctype, Text, Comment, Element};

use util::atom::Atom;
use util::namespace::{Namespace, HTML};
use tokenizer::Attribute;
use tree_builder::{TreeSink, QuirksMode, NodeOrText, AppendNode, AppendText};
use tree_builder;
use serialize::{Serializable, Serializer};
use driver::ParseResult;

use std::rc::{Rc, Weak};
use std::cell::RefCell;
use std::default::Default;
use std::io::IoResult;

/// A DOM node.
pub struct Node {
    pub node: NodeEnum,
    pub parent: Option<WeakHandle>,
    pub children: Vec<Handle>,

    /// The "script already started" flag.
    ///
    /// Not meaningful for nodes other than HTML `<script>`.
    pub script_already_started: bool,
}

impl Node {
    fn new(node: NodeEnum) -> Node {
        Node {
            node: node,
            parent: None,
            children: vec!(),
            script_already_started: false,
        }
    }
}

/// Reference to a DOM node.
pub type Handle = Rc<RefCell<Node>>;

/// Weak reference to a DOM node, used for parent pointers.
pub type WeakHandle = Weak<RefCell<Node>>;

fn same_node(x: &Handle, y: &Handle) -> bool {
    // FIXME: This shouldn't really need to touch the borrow flags, right?
    (&*x.borrow() as *const Node) == (&*y.borrow() as *const Node)
}

fn new_node(node: NodeEnum) -> Handle {
    Rc::new(RefCell::new(Node::new(node)))
}

fn append(new_parent: &Handle, child: Handle) {
    new_parent.borrow_mut().children.push(child.clone());
    let parent = &mut child.borrow_mut().parent;
    assert!(parent.is_none());
    *parent = Some(new_parent.downgrade());
}

fn get_parent_and_index(target: &Handle) -> Option<(Handle, uint)> {
    let child = target.borrow();
    let parent = unwrap_or_return!(child.parent.as_ref(), None)
        .upgrade().expect("dangling weak pointer");
    match parent.borrow_mut().children.iter().enumerate()
                .find(|&(_, n)| same_node(n, target)) {
        Some((i, _)) => Some((parent, i)),
        None => fail!("have parent but couldn't find in parent's children!"),
    }
}

fn append_to_existing_text(prev: &Handle, text: &str) -> bool {
    match prev.borrow_mut().deref_mut().node {
        Text(ref mut existing) => {
            existing.push_str(text);
            true
        }
        _ => false,
    }
}

fn remove_from_parent(target: &Handle) {
    {
        let (parent, i) = unwrap_or_return!(get_parent_and_index(target), ());
        parent.borrow_mut().children.remove(i).expect("not found!");
    }

    let mut child = target.borrow_mut();
    (*child).parent = None;
}

/// The DOM itself; the result of parsing.
pub struct RcDom {
    /// The `Document` itself.
    pub document: Handle,

    /// Errors that occurred during parsing.
    pub errors: Vec<String>,

    /// The document's quirks mode.
    pub quirks_mode: QuirksMode,
}

impl TreeSink<Handle> for RcDom {
    fn parse_error(&mut self, msg: String) {
        self.errors.push(msg);
    }

    fn get_document(&mut self) -> Handle {
        self.document.clone()
    }

    fn set_quirks_mode(&mut self, mode: QuirksMode) {
        self.quirks_mode = mode;
    }

    fn same_node(&self, x: Handle, y: Handle) -> bool {
        same_node(&x, &y)
    }

    fn elem_name(&self, target: Handle) -> (Namespace, Atom) {
        match target.borrow().node {
            Element(ref name, _) => (HTML, name.clone()),
            _ => fail!("not an element!"),
        }
    }

    fn create_element(&mut self, ns: Namespace, name: Atom, attrs: Vec<Attribute>) -> Handle {
        assert!(ns == HTML);
        new_node(Element(name, attrs))
    }

    fn create_comment(&mut self, text: String) -> Handle {
        new_node(Comment(text))
    }

    fn append(&mut self, parent: Handle, child: NodeOrText<Handle>) {
        // Append to an existing Text node if we have one.
        match child {
            AppendText(ref text) => match parent.borrow().children.last() {
                Some(h) => if append_to_existing_text(h, text.as_slice()) { return; },
                _ => (),
            },
            _ => (),
        }

        append(&parent, match child {
            AppendText(text) => new_node(Text(text)),
            AppendNode(node) => node
        });
    }

    fn append_before_sibling(&mut self,
            sibling: Handle,
            child: NodeOrText<Handle>) -> Result<(), NodeOrText<Handle>> {
        let (parent, i) = unwrap_or_return!(get_parent_and_index(&sibling), Err(child));

        let child = match (child, i) {
            // No previous node.
            (AppendText(text), 0) => new_node(Text(text)),

            // Look for a text node before the insertion point.
            (AppendText(text), i) => {
                let parent = parent.borrow();
                let prev = parent.children.get(i-1);
                if append_to_existing_text(prev, text.as_slice()) {
                    return Ok(());
                }
                new_node(Text(text))
            }

            // The tree builder promises we won't have a text node after
            // the insertion point.

            // Any other kind of node.
            (AppendNode(node), _) => node,
        };

        if child.borrow().parent.is_some() {
            remove_from_parent(&child);
        }

        child.borrow_mut().parent = Some(parent.clone().downgrade());
        parent.borrow_mut().children.insert(i, child);
        Ok(())
    }

    fn append_doctype_to_document(&mut self, name: String, public_id: String, system_id: String) {
        append(&self.document, new_node(Doctype(name, public_id, system_id)));
    }

    fn add_attrs_if_missing(&mut self, target: Handle, mut attrs: Vec<Attribute>) {
        let mut node = target.borrow_mut();
        // FIXME: mozilla/rust#15609
        let existing = match node.deref_mut().node {
            Element(_, ref mut attrs) => attrs,
            _ => return,
        };

        // FIXME: quadratic time
        attrs.retain(|attr|
            !existing.iter().any(|e| e.name == attr.name));
        existing.push_all_move(attrs);
    }

    fn remove_from_parent(&mut self, target: Handle) {
        remove_from_parent(&target);
    }

    fn mark_script_already_started(&mut self, node: Handle) {
        node.borrow_mut().script_already_started = true;
    }
}

impl Default for RcDom {
    fn default() -> RcDom {
        RcDom {
            document: new_node(Document),
            errors: vec!(),
            quirks_mode: tree_builder::NoQuirks,
        }
    }
}

impl ParseResult<RcDom> for RcDom {
    fn get_result(sink: RcDom) -> RcDom {
        sink
    }
}

impl Serializable for Handle {
    fn serialize<'wr, Wr: Writer>(&self, serializer: &mut Serializer<'wr, Wr>, incl_self: bool) -> IoResult<()> {
        let node = self.borrow();
        match (incl_self, &node.node) {
            (_, &Element(ref name, ref attrs)) => {
                if incl_self {
                    try!(serializer.start_elem(HTML, name.clone(),
                        attrs.iter().map(|at| (&at.name, at.value.as_slice()))));
                }

                for handle in node.children.iter() {
                    try!(handle.clone().serialize(serializer, true));
                }

                if incl_self {
                    try!(serializer.end_elem(HTML, name.clone()));
                }
                Ok(())
            }

            (false, &Document) => {
                for handle in node.children.iter() {
                    try!(handle.clone().serialize(serializer, true));
                }
                Ok(())
            }

            (false, _) => Ok(()),

            (true, &Doctype(ref name, _, _)) => serializer.write_doctype(name.as_slice()),
            (true, &Text(ref text)) => serializer.write_text(text.as_slice()),
            (true, &Comment(ref text)) => serializer.write_comment(text.as_slice()),

            (true, &Document) => fail!("Can't serialize Document node itself"),
        }
    }
}
<|endoftext|>


<commit_before>// Syntax: master_server [ip address (0.0.0.0 by default)]
// This was created for HaloMD.

// If the server does not respond in at least this many seconds, it will be dropped from the list.
const DROP_TIME : i64 = 60;

// Blacklist for blocking IPs. Separate with newlines.
// Blacklisting IPs ignores heartbeat and keepalive packets from an IP address.
// That means that servers that are banned will not be immediately removed, but will time out, instead.
const BLACKLIST_FILE : &'static str = "blacklist.txt";

// Read the blacklist every x amount of seconds.
const BLACKLIST_UPDATE_TIME : u32 = 60;

// Note: The master server must have TCP 29920 open and UDP 27900 open.

use std::net::{UdpSocket,TcpListener};
use std::io::{Write,BufReader,BufRead};
use std::env;
use std::fs::File;
use std::thread;
use std::sync::{Arc, Mutex};

extern crate time;

mod halo_server;
use halo_server::HaloServer;

mod heartbeat_packet;
use heartbeat_packet::HeartbeatPacket;

fn main() {
    let count = env::args().count();
    let ip = if count == 2 {
        let j : Vec<_> = env::args().collect();
        j[1].to_string()
    }
    else if count == 1 {
        "0.0.0.0".to_string()
    }
    else {
        println!("Only one argument is allowed: the IP to bind to.");
        return;
    };

    // We need to bind on two different ports. If it failed to bind (invalid IP, port is taken), then we must make sure this is known.
    let halo_socket = match UdpSocket::bind((&ip as &str,27900)) {
        Err(_) => {
            println!("Error creating a UDP socket at {}:27900.",ip);
            return;
        },
        Ok(halo_socket) => halo_socket
    };

    let client_socket = match TcpListener::bind((&ip as &str,29920)) {
        Err(_) => {
            println!("Error listening to TCP at {}:29920.",ip);
            return;
        },
        Ok(client_socket) => client_socket
    };

    // Mutex for thread safety.
    let servers_halo: Vec<HaloServer> = Vec::new();
    let servers_mut_udp = Arc::new(Mutex::new(servers_halo));
    let servers_mut_tcp = servers_mut_udp.clone();
    let servers_mut_destruction = servers_mut_udp.clone();

    // Destruction thread. This will remove servers that have not broadcasted their presence in a while.
    thread::spawn(move || {
        loop {
            thread::sleep_ms(10 * 1000);
            let mut servers = servers_mut_destruction.lock().unwrap();
            let count = servers.len();
            let timenow = time::now().to_timespec().sec;
            for i in (0..count).rev() {
                if servers[i].last_alive + DROP_TIME < timenow {
                    // println!("Timed out {}.", servers[i].to_string());
                    servers.remove(i);
                }
            }
        }
    });

    // Blacklist mutex. Concurrency needs to be safe, my friend.
    let blacklist: Vec<String> = Vec::new();
    let blacklist_update = Arc::new(Mutex::new(blacklist));
    let blacklist_udp = blacklist_update.clone();

    // Blacklist read thread.
    thread::spawn(move || {
        loop {
            // Placed in a block so blacklist is unlocked before sleeping to prevent threads from being locked for too long.
            {
                let mut blacklist_new : Vec<String> = Vec::new();
                match File::open(BLACKLIST_FILE) {
                    Ok(file) => {
                        let reader = BufReader::new(&file);
                        for line_possibly in reader.lines() {
                            match line_possibly {
                                Err(_) => break,
                                Ok(line) => {
                                    blacklist_new.push(line);
                                }
                            }
                        }
                    },
                    Err(_) => {}
                };

                let mut blacklist_ref = blacklist_update.lock().unwrap();
                blacklist_ref.clear();

                for b in blacklist_new {
                    blacklist_ref.push(b.clone());
                }
            }
            thread::sleep_ms(BLACKLIST_UPDATE_TIME * 1000);
        }
    });

    // TCP server thread. This is for the HaloMD application.
    thread::spawn(move || {
        loop {
            for stream in client_socket.incoming() {
                let mut client = match stream {
                    Err(_) => continue,
                    Ok(the_stream) => the_stream
                };
                // Unwrap the IP.
                let ip = match client.peer_addr() {
                    Err(_) => continue,
                    Ok(ip) => ip.to_string().split(":").next().unwrap().to_string()
                };

                let mut ips = String::new();

                let servers_ref = servers_mut_tcp.lock().unwrap();
                let servers = (*servers_ref).iter();

                for j in servers {
                    ips = ips + &(j.to_string()) + "\n";
                }

                // Some number placed after the requester's IP. If you ask me, the source code was abducted by aliens, and this is a tracking number. Regardless, it's needed.
                ips = ips + &ip + ":49149:3425";
                let _ = client.write_all(ips.as_bytes());
            }
        }
    });

    // UDP server is run on the main thread. Servers broadcast their presence here.

    // These are the allowed game versions. HaloMD and Halo PC 1.09 uses 01.00.09.0620, while Halo PC servers on 1.10 use 01.00.10.0621 (these are also joinable).
    let game_versions = [ "01.00.09.0620".to_string(), "01.00.10.0621".to_string() ];

    let mut buffer = [0; 2048];
    loop {
        let (length, source) = match halo_socket.recv_from(&mut buffer) {
            Err(_) => continue,
            Ok((length, source)) => (length,source)
        };
        let client_ip = source.to_string().split(":").next().unwrap().to_string();

        if buffer[0] == 3 || buffer[0] == 8 {
            let blacklist_ref = blacklist_udp.lock().unwrap();
            if blacklist_ref.contains(&client_ip) {
                continue;
            }

            // Heartbeat packet.
            if buffer[0] == 3 {

                let mut servers = servers_mut_udp.lock().unwrap();

                let packet = HeartbeatPacket::from_buffer(&buffer, length);

                if packet.localport != 0  {
                    let mut found = false;
                    let updatetime = time::now().to_timespec().sec;
                    for i in (0.. servers.len()) {
                        if servers[i].ip == client_ip && servers[i].port == packet.localport {
                            servers[i].last_alive = updatetime;
                            found = true;
                            if packet.statechanged == 2 {
                                // println!("Dropped {}.", servers[i].to_string());
                                servers.remove(i);
                            }
                            break;
                        }
                    }

                    if !found && &packet.gamename == "halor" && game_versions.contains(&packet.gamever) {
                        let serverness = HaloServer { ip:client_ip, port: packet.localport, last_alive: updatetime };
                        // println!("Added {}.", serverness.to_string());
                        (*servers).push(serverness);
                    }
                }
            }

            // Keepalive packet. We need to rely on the origin's port for this, unfortunately.
            else if buffer[0] == 8 {
                let mut servers_ref = servers_mut_udp.lock().unwrap();
                let servers = (*servers_ref).iter_mut();

                for i in servers {
                    if i.ip == client_ip && i.port == source.port() {
                        i.last_alive = time::now().to_timespec().sec;
                        break;
                    }
                }
            }
        }
    }
}
<commit_msg>Use retain instead of iterating<commit_after>// Syntax: master_server [ip address (0.0.0.0 by default)]
// This was created for HaloMD.

// If the server does not respond in at least this many seconds, it will be dropped from the list.
const DROP_TIME : i64 = 60;

// Blacklist for blocking IPs. Separate with newlines.
// Blacklisting IPs ignores heartbeat and keepalive packets from an IP address.
// That means that servers that are banned will not be immediately removed, but will time out, instead.
const BLACKLIST_FILE : &'static str = "blacklist.txt";

// Read the blacklist every x amount of seconds.
const BLACKLIST_UPDATE_TIME : u32 = 60;

// Note: The master server must have TCP 29920 open and UDP 27900 open.

use std::net::{UdpSocket,TcpListener};
use std::io::{Write,BufReader,BufRead};
use std::env;
use std::fs::File;
use std::thread;
use std::sync::{Arc, Mutex};

extern crate time;

mod halo_server;
use halo_server::HaloServer;

mod heartbeat_packet;
use heartbeat_packet::HeartbeatPacket;

fn main() {
    let count = env::args().count();
    let ip = if count == 2 {
        let j : Vec<_> = env::args().collect();
        j[1].to_string()
    }
    else if count == 1 {
        "0.0.0.0".to_string()
    }
    else {
        println!("Only one argument is allowed: the IP to bind to.");
        return;
    };

    // We need to bind on two different ports. If it failed to bind (invalid IP, port is taken), then we must make sure this is known.
    let halo_socket = match UdpSocket::bind((&ip as &str,27900)) {
        Err(_) => {
            println!("Error creating a UDP socket at {}:27900.",ip);
            return;
        },
        Ok(halo_socket) => halo_socket
    };

    let client_socket = match TcpListener::bind((&ip as &str,29920)) {
        Err(_) => {
            println!("Error listening to TCP at {}:29920.",ip);
            return;
        },
        Ok(client_socket) => client_socket
    };

    // Mutex for thread safety.
    let servers_halo: Vec<HaloServer> = Vec::new();
    let servers_mut_udp = Arc::new(Mutex::new(servers_halo));
    let servers_mut_tcp = servers_mut_udp.clone();
    let servers_mut_destruction = servers_mut_udp.clone();

    // Destruction thread. This will remove servers that have not broadcasted their presence in a while.
    thread::spawn(move || {
        loop {
            thread::sleep_ms(10 * 1000);
            let mut servers = servers_mut_destruction.lock().unwrap();
            let timenow = time::now().to_timespec().sec;
            servers.retain(|x| x.last_alive + DROP_TIME > timenow);
        }
    });

    // Blacklist mutex. Concurrency needs to be safe, my friend.
    let blacklist: Vec<String> = Vec::new();
    let blacklist_update = Arc::new(Mutex::new(blacklist));
    let blacklist_udp = blacklist_update.clone();

    // Blacklist read thread.
    thread::spawn(move || {
        loop {
            // Placed in a block so blacklist is unlocked before sleeping to prevent threads from being locked for too long.
            {
                let mut blacklist_new : Vec<String> = Vec::new();
                match File::open(BLACKLIST_FILE) {
                    Ok(file) => {
                        let reader = BufReader::new(&file);
                        for line_possibly in reader.lines() {
                            match line_possibly {
                                Err(_) => break,
                                Ok(line) => {
                                    blacklist_new.push(line);
                                }
                            }
                        }
                    },
                    Err(_) => {}
                };

                let mut blacklist_ref = blacklist_update.lock().unwrap();
                blacklist_ref.clear();

                for b in blacklist_new {
                    blacklist_ref.push(b.clone());
                }
            }
            thread::sleep_ms(BLACKLIST_UPDATE_TIME * 1000);
        }
    });

    // TCP server thread. This is for the HaloMD application.
    thread::spawn(move || {
        loop {
            for stream in client_socket.incoming() {
                let mut client = match stream {
                    Err(_) => continue,
                    Ok(the_stream) => the_stream
                };
                // Unwrap the IP.
                let ip = match client.peer_addr() {
                    Err(_) => continue,
                    Ok(ip) => ip.to_string().split(":").next().unwrap().to_string()
                };

                let mut ips = String::new();

                let servers_ref = servers_mut_tcp.lock().unwrap();
                let servers = (*servers_ref).iter();

                for j in servers {
                    ips = ips + &(j.to_string()) + "\n";
                }

                // Some number placed after the requester's IP. If you ask me, the source code was abducted by aliens, and this is a tracking number. Regardless, it's needed.
                ips = ips + &ip + ":49149:3425";
                let _ = client.write_all(ips.as_bytes());
            }
        }
    });

    // UDP server is run on the main thread. Servers broadcast their presence here.

    // These are the allowed game versions. HaloMD and Halo PC 1.09 uses 01.00.09.0620, while Halo PC servers on 1.10 use 01.00.10.0621 (these are also joinable).
    let game_versions = [ "01.00.09.0620".to_string(), "01.00.10.0621".to_string() ];

    let mut buffer = [0; 2048];
    loop {
        let (length, source) = match halo_socket.recv_from(&mut buffer) {
            Err(_) => continue,
            Ok((length, source)) => (length,source)
        };
        let client_ip = source.to_string().split(":").next().unwrap().to_string();

        if buffer[0] == 3 || buffer[0] == 8 {
            let blacklist_ref = blacklist_udp.lock().unwrap();
            if blacklist_ref.contains(&client_ip) {
                continue;
            }

            // Heartbeat packet.
            if buffer[0] == 3 {

                let mut servers = servers_mut_udp.lock().unwrap();

                let packet = HeartbeatPacket::from_buffer(&buffer, length);

                if packet.localport != 0  {
                    let mut found = false;
                    let updatetime = time::now().to_timespec().sec;
                    for i in (0.. servers.len()) {
                        if servers[i].ip == client_ip && servers[i].port == packet.localport {
                            servers[i].last_alive = updatetime;
                            found = true;
                            if packet.statechanged == 2 {
                                // println!("Dropped {}.", servers[i].to_string());
                                servers.remove(i);
                            }
                            break;
                        }
                    }

                    if !found && &packet.gamename == "halor" && game_versions.contains(&packet.gamever) {
                        let serverness = HaloServer { ip:client_ip, port: packet.localport, last_alive: updatetime };
                        // println!("Added {}.", serverness.to_string());
                        (*servers).push(serverness);
                    }
                }
            }

            // Keepalive packet. We need to rely on the origin's port for this, unfortunately.
            else if buffer[0] == 8 {
                let mut servers_ref = servers_mut_udp.lock().unwrap();
                let servers = (*servers_ref).iter_mut();

                for i in servers {
                    if i.ip == client_ip && i.port == source.port() {
                        i.last_alive = time::now().to_timespec().sec;
                        break;
                    }
                }
            }
        }
    }
}
<|endoftext|>


<commit_before>// C preprocessor

use crate::token::*;
use std::cell::RefCell;
use std::collections::HashMap;

thread_local! {
    static MACROS: RefCell<HashMap<String, Macro>> = RefCell::new(HashMap::new());

    static CTX: RefCell<Context> = RefCell::new(Context::new());
}

fn set_ctx(ctx: Context) {
    CTX.with(|c| {
        *c.borrow_mut() = ctx;
    })
}

fn get_ctx() -> Option<Context> {
    CTX.with(|c| {
        Some(c.borrow().clone())
    })
}

fn ctx_pop() {
    CTX.with(|c| {
        let next = c.borrow().next.clone();
        if next.is_some() {
            *c.borrow_mut() = *next.unwrap();
        }
    })
}

fn ctx_output() -> Vec<Token> {
    CTX.with(|c| {
        return c.borrow().output.clone();
    })
}

fn macros_get(key: &String) -> Option<Macro> {
    MACROS.with(|m| {
        match m.borrow().get(key) {
            Some(v) => Some(v.clone()),
            None => None,
        }
    })
}

fn macros_put(key: String, value: Macro) {
    MACROS.with(|m| {
        m.borrow_mut().insert(key, value);
    })
}

#[derive(Clone, Debug)]
struct Context {
    input: Vec<Token>,
    output: Vec<Token>,
    pos: usize,
    next: Option<Box<Context>>,
}

impl Context {
    fn new() -> Context {
        Context {
            input: Vec::new(),
            output: Vec::new(),
            pos: 0,
            next: None,
        }
    }
}

fn new_ctx(next: Option<Context>, input: Vec<Token>) -> Context {
    let mut ctx = Context::new();
    ctx.input = input;
    if next.is_none() {
        ctx.next = None;
    } else {
        ctx.next = Some(Box::new(next.unwrap()));
    }
    return ctx;
}

#[derive(Clone, Debug, PartialEq)]
enum MacroType {
    OBJLIKE,
    FUNCLIKE,
}

#[derive(Clone, Debug)]
struct Macro {
    ty: MacroType,
    tokens: Vec<Token>,
    params: Vec<String>,
}

fn new_macro(ty: MacroType, name: String) {
    let m = Macro {
        ty: ty,
        tokens: Vec::new(),
        params: Vec::new(),
    };
    macros_put(name, m);
}

fn append(v: &mut Vec<Token>) {
    CTX.with(|ctx| {
        ctx.borrow_mut().output.append(v);
    })
}

fn add(t: Token) {
    CTX.with(|ctx| {
        ctx.borrow_mut().output.push(t);
    })
}

fn next() -> Token {
    CTX.with(|ctx| {
        let pos = ctx.borrow().pos;
        assert!(pos < ctx.borrow().input.len());
        ctx.borrow_mut().pos += 1;
        return ctx.borrow().input[pos].clone();
    })
}

fn eof() -> bool {
    CTX.with(|c| {
        let ctx = c.borrow();
        return ctx.pos == ctx.input.len();
    })
}

fn get(ty: TokenType, msg: String) -> Token {
    let t = next();
    if t.ty != ty {
        bad_token(&t, msg);
    }
    return t;
}

fn ident(msg: String) -> String {
    let t = get(TokenType::IDENT, msg);
    return t.name.clone();
}

fn peek() -> Token {
    CTX.with(|c| {
        let ctx = c.borrow();
        return ctx.input[ctx.pos].clone();
    })
}

fn consume(ty: TokenType) -> bool {
    if peek().ty != ty {
        return false;
    }
    CTX.with(|c| {
        c.borrow_mut().pos += 1;
    });
    return true;
}

fn read_until_eol() -> Vec<Token> {
    let mut v = Vec::new();
    while !eof() {
        let t = next();
        if t.ty == TokenType::NEW_LINE {
            break;
        }
        v.push(t.clone());
    }
    return v;
}

fn new_int(val: i32) -> Token {
    let mut t = new_token(TokenType::NUM, 0);
    t.val = val;
    return t;
}

fn new_param(val: i32) -> Token {
    let mut t = new_token(TokenType::PARAM, 0);
    t.val = val;
    return t;
}

fn is_ident(t: &Token, s: &str) -> bool {
    return t.ty == TokenType::IDENT && t.name == s;
}

fn replace_params(m: &mut Macro) {
    let params = m.params.clone();
    let mut tokens = m.tokens.clone();

    // Replace macro parameter tokens with TokenType::PARAM tokens.
    let mut map = HashMap::new();
    for i in 0..params.len() {
        let name = params[i].clone();
        map.insert(name, i as i32);
    }

    for i in 0..tokens.len() {
        let t = tokens[i].clone();
        if t.ty != TokenType::IDENT {
            continue;
        }
        let n = match map.get(&t.name) {
            Some(i) => *i,
            None => {
                continue;
            }
        };
        tokens.remove(i);
        tokens.insert(i, new_param(n));
    }

    // Process '#' followed by a macro parameter.
    let mut v = Vec::new();
    let mut i = 0;
    while i < tokens.len() {
        let t1 = tokens[i].clone();

        if i != tokens.len() - 1 && t1.ty == TokenType::SHARP {
            let mut t2 = tokens[i+1].clone();
            if t2.ty == TokenType::PARAM {
                t2.stringize = true;
                v.push(t2);
                i += 1;
            }
        } else {
            v.push(t1);
        }
        i += 1;
    }
    m.tokens = v;
}

fn read_one_arg() -> Vec<Token> {
    let mut v = Vec::new();
    let start = peek();
    let mut level = 0;

    while !eof() {
        let t = peek();
        if level == 0 {
            if t.ty == TokenType::KET || t.ty == TokenType::COMMA {
                return v;
            }
        }

        next();
        if t.ty == TokenType::BRA {
            level += 1;
        } else if t.ty == TokenType::KET {
            level -= 1;
        }
        v.push(t);
    }
    bad_token(&start, "unclosed macro argument".to_string());
    panic!();
}

fn read_args() -> HashMap<usize, Vec<Token>> {
    let mut v = HashMap::new();
    if consume(TokenType::KET) {
        return v;
    }
    let i = v.len();
    v.insert(i, read_one_arg());
    while !consume(TokenType::KET) {
        get(TokenType::COMMA, "comma expected".to_string());
        let i = v.len();
        v.insert(i, read_one_arg());
    }
    return v;
}

fn stringize(tokens: Vec<Token>) -> Token {
    let mut sb = String::new();

    for i in 0..tokens.len() {
        let t = &tokens[i];
        if i > 0 {
            sb.push(' ');
        }
        sb.push_str(tokstr(&t));
    }
    sb.push('\0');

    let mut t = new_token(TokenType::STR, 0);
    t.len = sb.len();
    t.str_cnt = sb;
    return t;
}

fn apply(m: &mut Macro, start: & Token) {
    if m.ty == MacroType::OBJLIKE {
        append(&mut m.tokens);
        return;
    }

    // Function-like macro
    get(TokenType::BRA, "comma expected".to_string());
    let args = read_args();
    if m.params.len() != args.len() {
        bad_token(&start, format!("number of parameter does not match ({} != {})", m.params.len(), args.len()).to_string());
        panic!();
    }

    for i in 0..m.tokens.len() {
        let t = &m.tokens[i];

        if is_ident(t, "__LINE__") {
            add(new_int(get_line_number(t)));
            continue;
        }

        if t.ty == TokenType::PARAM {
            if t.stringize {
                let j = t.val as usize;
                add(stringize(args.get(&j).unwrap().clone()));
            } else {
                let j = t.val as usize;
                append(&mut args.get(&j).unwrap().clone());
            }
            continue;
        }
        add(t.clone());
    }
}

fn funclike_macro(name: String) {
    new_macro(MacroType::FUNCLIKE, name.clone());
    let mut m = macros_get(&name).unwrap();
    m.params.push(ident("parameter name expected".to_string()));
    while !consume(TokenType::KET) {
        get(TokenType::COMMA, "comma expected".to_string());
        m.params.push(ident("parameter name expected".to_string()));
    }
    m.tokens = read_until_eol();

    replace_params(&mut m);

    macros_put(name, m);
}

fn objlike_macro(name: String) {
    new_macro(MacroType::OBJLIKE, name.clone());
    let mut m = macros_get(&name).unwrap();
    m.tokens = read_until_eol();
    macros_put(name, m);
}

fn define() {
    let name = ident("macro name expected".to_string());
    if consume(TokenType::BRA) {
        return funclike_macro(name);
    }
    return objlike_macro(name);
}

fn include() {
    let t = get(TokenType::STR, "string expected".to_string());
    let path = t.str_cnt;
    get(TokenType::NEW_LINE, "newline expected".to_string());
    return append(&mut tokenize(path, false));
}

pub fn preprocess(tokens: Vec<Token>) -> Vec<Token> {
    set_ctx(new_ctx(get_ctx(), tokens.clone()));

    while !eof() {
        let mut t = next();

        if t.ty == TokenType::IDENT {
            match macros_get(&t.name) {
                Some(macro_token) => {
                    apply(&mut macro_token.clone(), &t);
                }
                None => {
                    add(t.clone());
                }
            }
            continue;
        }

        if t.ty != TokenType::SHARP {
            add(t.clone());
            continue;
        }

        t = get(TokenType::IDENT, "identifier expected".to_string());

        if t.name == "define" {
            define();
        } else if t.name == "include" {
            include();
        } else {
            bad_token(&t, "unknown directive".to_string());
        }
    }

    let v = ctx_output();
    ctx_pop();
    return v;
}
<commit_msg>Simplify.<commit_after>// C preprocessor

use crate::token::*;
use std::cell::RefCell;
use std::collections::HashMap;

thread_local! {
    static MACROS: RefCell<HashMap<String, Macro>> = RefCell::new(HashMap::new());

    static CTX: RefCell<Context> = RefCell::new(Context::new());
}

fn set_ctx(ctx: Context) {
    CTX.with(|c| {
        *c.borrow_mut() = ctx;
    })
}

fn get_ctx() -> Option<Context> {
    CTX.with(|c| {
        Some(c.borrow().clone())
    })
}

fn ctx_pop() {
    CTX.with(|c| {
        let next = c.borrow().next.clone();
        if next.is_some() {
            *c.borrow_mut() = *next.unwrap();
        }
    })
}

fn ctx_output() -> Vec<Token> {
    CTX.with(|c| {
        return c.borrow().output.clone();
    })
}

fn macros_get(key: &String) -> Option<Macro> {
    MACROS.with(|m| {
        match m.borrow().get(key) {
            Some(v) => Some(v.clone()),
            None => None,
        }
    })
}

fn macros_put(key: String, value: Macro) {
    MACROS.with(|m| {
        m.borrow_mut().insert(key, value);
    })
}

#[derive(Clone, Debug)]
struct Context {
    input: Vec<Token>,
    output: Vec<Token>,
    pos: usize,
    next: Option<Box<Context>>,
}

impl Context {
    fn new() -> Context {
        Context {
            input: Vec::new(),
            output: Vec::new(),
            pos: 0,
            next: None,
        }
    }
}

fn new_ctx(next: Option<Context>, input: Vec<Token>) -> Context {
    let mut ctx = Context::new();
    ctx.input = input;
    if next.is_none() {
        ctx.next = None;
    } else {
        ctx.next = Some(Box::new(next.unwrap()));
    }
    return ctx;
}

#[derive(Clone, Debug, PartialEq)]
enum MacroType {
    OBJLIKE,
    FUNCLIKE,
}

#[derive(Clone, Debug)]
struct Macro {
    ty: MacroType,
    tokens: Vec<Token>,
    params: Vec<String>,
}

fn new_macro(ty: MacroType, name: String) {
    let m = Macro {
        ty: ty,
        tokens: Vec::new(),
        params: Vec::new(),
    };
    macros_put(name, m);
}

fn append(v: &mut Vec<Token>) {
    CTX.with(|ctx| {
        ctx.borrow_mut().output.append(v);
    })
}

fn add(t: Token) {
    CTX.with(|ctx| {
        ctx.borrow_mut().output.push(t);
    })
}

fn next() -> Token {
    CTX.with(|ctx| {
        let pos = ctx.borrow().pos;
        assert!(pos < ctx.borrow().input.len());
        ctx.borrow_mut().pos += 1;
        return ctx.borrow().input[pos].clone();
    })
}

fn eof() -> bool {
    CTX.with(|c| {
        let ctx = c.borrow();
        return ctx.pos == ctx.input.len();
    })
}

fn get(ty: TokenType, msg: String) -> Token {
    let t = next();
    if t.ty != ty {
        bad_token(&t, msg);
    }
    return t;
}

fn ident(msg: String) -> String {
    let t = get(TokenType::IDENT, msg);
    return t.name.clone();
}

fn peek() -> Token {
    CTX.with(|c| {
        let ctx = c.borrow();
        return ctx.input[ctx.pos].clone();
    })
}

fn consume(ty: TokenType) -> bool {
    if peek().ty != ty {
        return false;
    }
    CTX.with(|c| {
        c.borrow_mut().pos += 1;
    });
    return true;
}

fn read_until_eol() -> Vec<Token> {
    let mut v = Vec::new();
    while !eof() {
        let t = next();
        if t.ty == TokenType::NEW_LINE {
            break;
        }
        v.push(t.clone());
    }
    return v;
}

fn new_int(val: i32) -> Token {
    let mut t = new_token(TokenType::NUM, 0);
    t.val = val;
    return t;
}

fn new_param(val: i32) -> Token {
    let mut t = new_token(TokenType::PARAM, 0);
    t.val = val;
    return t;
}

fn is_ident(t: &Token, s: &str) -> bool {
    return t.ty == TokenType::IDENT && t.name == s;
}

fn replace_params(m: &mut Macro) {
    let params = m.params.clone();
    let mut tokens = m.tokens.clone();

    // Replace macro parameter tokens with TokenType::PARAM tokens.
    let mut map = HashMap::new();
    for i in 0..params.len() {
        let name = params[i].clone();
        map.insert(name, i as i32);
    }

    for i in 0..tokens.len() {
        let t = tokens[i].clone();
        if t.ty != TokenType::IDENT {
            continue;
        }
        let n = match map.get(&t.name) {
            Some(i) => *i,
            None => {
                continue;
            }
        };
        tokens.remove(i);
        tokens.insert(i, new_param(n));
    }

    // Process '#' followed by a macro parameter.
    let mut v = Vec::new();
    let mut i = 0;
    while i < tokens.len() {
        let t1 = tokens[i].clone();

        if i != tokens.len() - 1 && t1.ty == TokenType::SHARP {
            let mut t2 = tokens[i+1].clone();
            if t2.ty == TokenType::PARAM {
                t2.stringize = true;
                v.push(t2);
                i += 1;
            }
        } else {
            v.push(t1);
        }
        i += 1;
    }
    m.tokens = v;
}

fn read_one_arg() -> Vec<Token> {
    let mut v = Vec::new();
    let start = peek();
    let mut level = 0;

    while !eof() {
        let t = peek();
        if level == 0 {
            if t.ty == TokenType::KET || t.ty == TokenType::COMMA {
                return v;
            }
        }

        next();
        if t.ty == TokenType::BRA {
            level += 1;
        } else if t.ty == TokenType::KET {
            level -= 1;
        }
        v.push(t);
    }
    bad_token(&start, "unclosed macro argument".to_string());
    panic!();
}

fn read_args() -> HashMap<usize, Vec<Token>> {
    let mut v = HashMap::new();
    if consume(TokenType::KET) {
        return v;
    }
    let i = v.len();
    v.insert(i, read_one_arg());
    while !consume(TokenType::KET) {
        get(TokenType::COMMA, "comma expected".to_string());
        let i = v.len();
        v.insert(i, read_one_arg());
    }
    return v;
}

fn stringize(tokens: Vec<Token>) -> Token {
    let mut sb = String::new();

    for i in 0..tokens.len() {
        let t = &tokens[i];
        if i > 0 {
            sb.push(' ');
        }
        sb.push_str(tokstr(&t));
    }
    sb.push('\0');

    let mut t = new_token(TokenType::STR, 0);
    t.len = sb.len();
    t.str_cnt = sb;
    return t;
}

fn add_special_macro(t: & Token) -> bool {
    if is_ident(t, "__LINE__") {
        add(new_int(get_line_number(t)));
        return true;
    }
    return false;
}

fn apply_objlike(m: &mut Macro, _start: & Token) {
    for t in m.tokens.iter() {
        if add_special_macro(t) {
            continue;
        }
        add(t.clone());
    }
}

fn apply_functionlike(m: &mut Macro, start: & Token) {
    get(TokenType::BRA, "comma expected".to_string());
    let args = read_args();
    if m.params.len() != args.len() {
        bad_token(&start, format!("number of parameter does not match ({} != {})", m.params.len(), args.len()).to_string());
        panic!();
    }

    for i in 0..m.tokens.len() {
        let t = &m.tokens[i];
        if add_special_macro(t) {
            continue;
        }

        if t.ty == TokenType::PARAM {
            if t.stringize {
                let j = t.val as usize;
                add(stringize(args.get(&j).unwrap().clone()));
            } else {
                let j = t.val as usize;
                append(&mut args.get(&j).unwrap().clone());
            }
            continue;
        }
        add(t.clone());
    }
}

fn apply(m: &mut Macro, start: & Token) {
    if m.ty == MacroType::OBJLIKE {
        apply_objlike(m, start);
    } else {
        apply_functionlike(m, start);
    }
}

fn funclike_macro(name: String) {
    new_macro(MacroType::FUNCLIKE, name.clone());
    let mut m = macros_get(&name).unwrap();
    m.params.push(ident("parameter name expected".to_string()));
    while !consume(TokenType::KET) {
        get(TokenType::COMMA, "comma expected".to_string());
        m.params.push(ident("parameter name expected".to_string()));
    }
    m.tokens = read_until_eol();

    replace_params(&mut m);

    macros_put(name, m);
}

fn objlike_macro(name: String) {
    new_macro(MacroType::OBJLIKE, name.clone());
    let mut m = macros_get(&name).unwrap();
    m.tokens = read_until_eol();
    macros_put(name, m);
}

fn define() {
    let name = ident("macro name expected".to_string());
    if consume(TokenType::BRA) {
        return funclike_macro(name);
    }
    return objlike_macro(name);
}

fn include() {
    let t = get(TokenType::STR, "string expected".to_string());
    let path = t.str_cnt;
    get(TokenType::NEW_LINE, "newline expected".to_string());
    return append(&mut tokenize(path, false));
}

pub fn preprocess(tokens: Vec<Token>) -> Vec<Token> {
    set_ctx(new_ctx(get_ctx(), tokens.clone()));

    while !eof() {
        let mut t = next();

        if t.ty == TokenType::IDENT {
            match macros_get(&t.name) {
                Some(macro_token) => {
                    apply(&mut macro_token.clone(), &t);
                }
                None => {
                    add(t.clone());
                }
            }
            continue;
        }

        if t.ty != TokenType::SHARP {
            add(t.clone());
            continue;
        }

        t = get(TokenType::IDENT, "identifier expected".to_string());

        if t.name == "define" {
            define();
        } else if t.name == "include" {
            include();
        } else {
            bad_token(&t, "unknown directive".to_string());
        }
    }

    let v = ctx_output();
    ctx_pop();
    return v;
}
<|endoftext|>


<commit_before>use std::io;
use config::Config;

#[cfg(unix)]
pub fn drop(cfg: &Config) -> Result<(), io::Error> {
    use libc::{getpwnam, getgrnam, setuid, setgid};
    use std::ffi::CString;
    use std::ptr;

    if let Some(ref group) = cfg.set_group {
        unsafe {
            let grstring = CString::new(group.as_bytes()).unwrap();
            let gentry = getgrnam(grstring.as_ptr());
            if gentry == ptr::null_mut() {
                return Err(io::Error::last_os_error());
            }
            info!("Group {:?} has gid of {}", group, (*gentry).gr_gid);
            if setgid((*gentry).gr_gid) == -1 {
                return Err(io::Error::last_os_error());
            }
        }
    }

    if let Some(ref user) = cfg.set_user {
        unsafe {
            let ustring = CString::new(user.as_bytes()).unwrap();
            let uentry = getpwnam(ustring.as_ptr());
            if uentry == ptr::null_mut() {
                return Err(io::Error::last_os_error());
            }
            if cfg.set_group.is_none() {
                info!("User {:?} has uid of {} and primary group {}",
                    user, (*uentry).pw_uid, (*uentry).pw_gid);
                if setgid((*uentry).pw_gid) == -1 {
                    return Err(io::Error::last_os_error());
                }
            } else {
                info!("User {:?} has uid of {}", user, (*uentry).pw_uid);
            }
            if setuid((*uentry).pw_uid) == -1 {
                return Err(io::Error::last_os_error());
            }
        }
    }
    Ok(())
}
#[cfg(not(unix))]
pub fn drop(_: &Config) -> Result<(), io::Error> {

}
<commit_msg>Fixes non-unix stub for privilege drop<commit_after>use std::io;
use config::Config;

#[cfg(unix)]
pub fn drop(cfg: &Config) -> Result<(), io::Error> {
    use libc::{getpwnam, getgrnam, setuid, setgid};
    use std::ffi::CString;
    use std::ptr;

    if let Some(ref group) = cfg.set_group {
        unsafe {
            let grstring = CString::new(group.as_bytes()).unwrap();
            let gentry = getgrnam(grstring.as_ptr());
            if gentry == ptr::null_mut() {
                return Err(io::Error::last_os_error());
            }
            info!("Group {:?} has gid of {}", group, (*gentry).gr_gid);
            if setgid((*gentry).gr_gid) == -1 {
                return Err(io::Error::last_os_error());
            }
        }
    }

    if let Some(ref user) = cfg.set_user {
        unsafe {
            let ustring = CString::new(user.as_bytes()).unwrap();
            let uentry = getpwnam(ustring.as_ptr());
            if uentry == ptr::null_mut() {
                return Err(io::Error::last_os_error());
            }
            if cfg.set_group.is_none() {
                info!("User {:?} has uid of {} and primary group {}",
                    user, (*uentry).pw_uid, (*uentry).pw_gid);
                if setgid((*uentry).pw_gid) == -1 {
                    return Err(io::Error::last_os_error());
                }
            } else {
                info!("User {:?} has uid of {}", user, (*uentry).pw_uid);
            }
            if setuid((*uentry).pw_uid) == -1 {
                return Err(io::Error::last_os_error());
            }
        }
    }
    Ok(())
}
#[cfg(not(unix))]
pub fn drop(_: &Config) -> Result<(), io::Error> {
    Ok(())
}
<|endoftext|>


<commit_before>#![allow(unused)]
use futures::future::{err, ok, Future};
use futures::stream::Stream;

use std::cell::RefCell;
use std::rc::Rc;

use options::StaticFile;

use super::readdebt::{DebtHandling, ReadDebt};
use super::{
    box_up_err, io_other_error, peer_err, peer_strerr, simple_err, BoxedNewPeerFuture, Peer,
};
use super::{ConstructParams, L2rUser, PeerConstructor, Specifier};
use tokio_io::io::{read_exact, write_all};

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::io::Write;

#[derive(Debug, Clone)]
pub enum SocksHostAddr {
    Ip(IpAddr),
    Name(String),
}

#[derive(Debug, Clone)]
pub struct SocksSocketAddr {
    pub host: SocksHostAddr,
    pub port: u16,
}

#[derive(Debug)]
pub struct SocksProxy<T: Specifier>(pub T);
impl<T: Specifier> Specifier for SocksProxy<T> {
    fn construct(&self, cp: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(cp.clone());
        inner.map(move |p, l2r| connect_socks5_peer(p, l2r, &cp.program_options.socks_destination))
    }
    specifier_boilerplate!(typ=WebSocket noglobalstate has_subspec);
    self_0_is_subspecifier!(proxy_is_multiconnect);
}
specifier_class!(
    name = SocksProxyClass,
    target = SocksProxy,
    prefixes = ["socks5-connect:"],
    arg_handling = subspec,
    overlay = true,
    StreamOriented,
    MulticonnectnessDependsOnInnerType,
    help = r#"
SOCKS5 proxy client (raw) [A]

Example: connect to a websocket using local `ssh -D` proxy

    websocat -t - ws-c:socks5-connect:tcp:127.0.0.1:1080 --socks5-destination echo.websocket.org:80 --ws-c-uri ws://echo.websocket.org

For a user-friendly solution, see --socks5 command-line option
"#
);

pub fn connect_socks5_peer(
    inner_peer: Peer,
    l2r: L2rUser,
    socks_destination: &Option<SocksSocketAddr>,
) -> BoxedNewPeerFuture {
    let (desthost, destport) = if let Some(ref sd) = *socks_destination {
        (sd.host.clone(), sd.port)
    } else {
        return peer_strerr("--socks5-destination is required for socks5-connect: overlay");
    };

    if let SocksHostAddr::Name(n) = &desthost {
        if n.len() > 255 {
            return peer_strerr("Destination host name too long for SOCKS5");
        }
    };

    let (r, w) = (inner_peer.0, inner_peer.1);
    let f = write_all(w, b"\x05\x01\x00").map_err(box_up_err).and_then(move |(w,_)| {

        let authmethods = [0; 2];
        read_exact(r, authmethods).map_err(box_up_err).and_then(move |(r, authmethods)| {
            if authmethods[0] != b'\x05' {
                return peer_strerr("Not a SOCKS5 reply");
            }
            if authmethods[1] != b'\x00' {
                return peer_strerr("Not a SOCKS5 or auth required");
            }

            let mut rq = Vec::with_capacity(20);
            {
                let mut c = ::std::io::Cursor::new(&mut rq);
                // 01 means "connect"
                c.write_all(b"\x05\x01\x00").unwrap();
                match &desthost {
                    SocksHostAddr::Ip(IpAddr::V4(ip4)) => {
                        c.write_all(b"\x01").unwrap();
                        c.write_all(&ip4.octets()).unwrap();
                    },
                    SocksHostAddr::Ip(IpAddr::V6(ip6)) => {
                        c.write_all(b"\x04").unwrap();
                        c.write_all(&ip6.octets()).unwrap();
                    },
                    SocksHostAddr::Name(name) => {
                        c.write_all(b"\x03").unwrap();
                        c.write_all(&[name.len() as u8]).unwrap();
                        c.write_all(name.as_bytes()).unwrap();
                    },
                };
                c.write_all(&[(destport >> 8) as u8]).unwrap();
                c.write_all(&[(destport >> 0) as u8]).unwrap();
            }

            Box::new(write_all(w, rq).map_err(box_up_err).and_then(move |(w, _)| {

                let reply = [0; 4];

                read_exact(r, reply).map_err(box_up_err).and_then(move |(r, reply)| {

                    if reply[0] != b'\x05' {
                        return peer_strerr("Not a SOCKS5 reply 2");
                    }
                    if reply[1] != b'\x00' {
                        let msg = match reply[1] {
                            1 => "SOCKS: General SOCKS server failuire",
                            2 => "SOCKS connection not allowed",
                            3 => "SOCKS: network unreachable",
                            4 => "SOCKS: host unreachable",
                            5 => "SOCKS: connection refused",
                            6 => "SOCKS: TTL expired",
                            7 => "SOCKS: Command not supported",
                            8 => "SOCKS: Address type not supported",
                            _ => "SOCKS: Unknown failure",
                        };
                        return peer_strerr(msg);
                    }
                    match reply[3] {
                        b'\x01' => {
                            // ipv4
                            let addrport = [0; 4+2];
                            Box::new(read_exact(r, addrport)
                                    .map_err(box_up_err)
                                    .and_then(move |(r,addrport)| {
                                let port = (addrport[4] as u16)*256 + (addrport[5] as u16);
                                let ip = Ipv4Addr::new(
                                    addrport[0],
                                    addrport[1],
                                    addrport[2],
                                    addrport[3],
                                );
                                let host = SocksHostAddr::Ip(IpAddr::V4(ip));
                                info!("Connected though the SOCKS5 proxy from {:?}",
                                    SocksSocketAddr{host,port});

                                ok(Peer::new(r,w))
                            }))
                        },
                        b'\x04' => {
                            // ipv6
                            let addrport = [0; 16+2];
                            Box::new(read_exact(r, addrport)
                                    .map_err(box_up_err)
                                    .and_then(move |(r,addrport)| {
                                let port = (addrport[16] as u16)*256 + (addrport[17] as u16);
                                // still not worth to switch to Cargo.toml to add
                                // "bytes" dependency, then scroll up for "extern crate",
                                // then look up docs again to find out where to get that BE thing.
                                let mut ip = [0u8;16];
                                ip.copy_from_slice(&addrport[0..16]);
                                let host = SocksHostAddr::Ip(IpAddr::V6(ip.into()));
                                info!("Connected though the SOCKS5 proxy from {:?}",
                                    SocksSocketAddr{host,port});

                                ok(Peer::new(r,w))
                            }))
                        },
                        b'\x03' => {

                            let alen = [0; 1];
                            Box::new(read_exact(r, alen)
                                    .map_err(box_up_err)
                                    .and_then(move |(r,alen)| {

                                let alen = alen[0] as usize;
                                let addrport = vec![0;alen+2];

                                read_exact(r, addrport)
                                        .map_err(box_up_err)
                                        .and_then(move |(r,addrport)| {

                                    let port = (addrport[alen] as u16)*256 + (addrport[alen+1] as u16);
                                    let host =
                                        SocksHostAddr::Name(
                                        ::std::str::from_utf8(&addrport[0..alen])
                                        .unwrap_or("(invalid hostname)").to_string());
                                info!("Connected though the SOCKS5 proxy from {:?}",
                                    SocksSocketAddr{host,port});

                                    ok(Peer::new(r,w))
                                })
                            }))
                        },
                        _ => {
                            return peer_strerr("SOCKS: bound address type is unknown");
                        },
                    }
                })
            })) as BoxedNewPeerFuture
        })
    });
    Box::new(f) as BoxedNewPeerFuture
}
<commit_msg>Rust 1.23.0 compat<commit_after>#![allow(unused)]
use futures::future::{err, ok, Future};
use futures::stream::Stream;

use std::cell::RefCell;
use std::rc::Rc;

use options::StaticFile;

use super::readdebt::{DebtHandling, ReadDebt};
use super::{
    box_up_err, io_other_error, peer_err, peer_strerr, simple_err, BoxedNewPeerFuture, Peer,
};
use super::{ConstructParams, L2rUser, PeerConstructor, Specifier};
use tokio_io::io::{read_exact, write_all};

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::io::Write;

#[derive(Debug, Clone)]
pub enum SocksHostAddr {
    Ip(IpAddr),
    Name(String),
}

#[derive(Debug, Clone)]
pub struct SocksSocketAddr {
    pub host: SocksHostAddr,
    pub port: u16,
}

#[derive(Debug)]
pub struct SocksProxy<T: Specifier>(pub T);
impl<T: Specifier> Specifier for SocksProxy<T> {
    fn construct(&self, cp: ConstructParams) -> PeerConstructor {
        let inner = self.0.construct(cp.clone());
        inner.map(move |p, l2r| connect_socks5_peer(p, l2r, &cp.program_options.socks_destination))
    }
    specifier_boilerplate!(typ=WebSocket noglobalstate has_subspec);
    self_0_is_subspecifier!(proxy_is_multiconnect);
}
specifier_class!(
    name = SocksProxyClass,
    target = SocksProxy,
    prefixes = ["socks5-connect:"],
    arg_handling = subspec,
    overlay = true,
    StreamOriented,
    MulticonnectnessDependsOnInnerType,
    help = r#"
SOCKS5 proxy client (raw) [A]

Example: connect to a websocket using local `ssh -D` proxy

    websocat -t - ws-c:socks5-connect:tcp:127.0.0.1:1080 --socks5-destination echo.websocket.org:80 --ws-c-uri ws://echo.websocket.org

For a user-friendly solution, see --socks5 command-line option
"#
);

pub fn connect_socks5_peer(
    inner_peer: Peer,
    l2r: L2rUser,
    socks_destination: &Option<SocksSocketAddr>,
) -> BoxedNewPeerFuture {
    let (desthost, destport) = if let Some(ref sd) = *socks_destination {
        (sd.host.clone(), sd.port)
    } else {
        return peer_strerr("--socks5-destination is required for socks5-connect: overlay");
    };

    if let SocksHostAddr::Name(ref n) = desthost {
        if n.len() > 255 {
            return peer_strerr("Destination host name too long for SOCKS5");
        }
    };

    let (r, w) = (inner_peer.0, inner_peer.1);
    let f = write_all(w, b"\x05\x01\x00").map_err(box_up_err).and_then(move |(w,_)| {

        let authmethods = [0; 2];
        read_exact(r, authmethods).map_err(box_up_err).and_then(move |(r, authmethods)| {
            if authmethods[0] != b'\x05' {
                return peer_strerr("Not a SOCKS5 reply");
            }
            if authmethods[1] != b'\x00' {
                return peer_strerr("Not a SOCKS5 or auth required");
            }

            let rq = {
                let mut c = ::std::io::Cursor::new(Vec::with_capacity(20));
                // 01 means "connect"
                c.write_all(b"\x05\x01\x00").unwrap();
                match desthost {
                    SocksHostAddr::Ip(IpAddr::V4(ip4)) => {
                        c.write_all(b"\x01").unwrap();
                        c.write_all(&ip4.octets()).unwrap();
                    },
                    SocksHostAddr::Ip(IpAddr::V6(ip6)) => {
                        c.write_all(b"\x04").unwrap();
                        c.write_all(&ip6.octets()).unwrap();
                    },
                    SocksHostAddr::Name(name) => {
                        c.write_all(b"\x03").unwrap();
                        c.write_all(&[name.len() as u8]).unwrap();
                        c.write_all(name.as_bytes()).unwrap();
                    },
                };
                c.write_all(&[(destport >> 8) as u8]).unwrap();
                c.write_all(&[(destport >> 0) as u8]).unwrap();
                c.into_inner()
            };

            Box::new(write_all(w, rq).map_err(box_up_err).and_then(move |(w, _)| {

                let reply = [0; 4];

                read_exact(r, reply).map_err(box_up_err).and_then(move |(r, reply)| {

                    if reply[0] != b'\x05' {
                        return peer_strerr("Not a SOCKS5 reply 2");
                    }
                    if reply[1] != b'\x00' {
                        let msg = match reply[1] {
                            1 => "SOCKS: General SOCKS server failuire",
                            2 => "SOCKS connection not allowed",
                            3 => "SOCKS: network unreachable",
                            4 => "SOCKS: host unreachable",
                            5 => "SOCKS: connection refused",
                            6 => "SOCKS: TTL expired",
                            7 => "SOCKS: Command not supported",
                            8 => "SOCKS: Address type not supported",
                            _ => "SOCKS: Unknown failure",
                        };
                        return peer_strerr(msg);
                    }
                    match reply[3] {
                        b'\x01' => {
                            // ipv4
                            let addrport = [0; 4+2];
                            Box::new(read_exact(r, addrport)
                                    .map_err(box_up_err)
                                    .and_then(move |(r,addrport)| {
                                let port = (addrport[4] as u16)*256 + (addrport[5] as u16);
                                let ip = Ipv4Addr::new(
                                    addrport[0],
                                    addrport[1],
                                    addrport[2],
                                    addrport[3],
                                );
                                let host = SocksHostAddr::Ip(IpAddr::V4(ip));
                                info!("Connected though the SOCKS5 proxy from {:?}",
                                    SocksSocketAddr{host,port});

                                ok(Peer::new(r,w))
                            }))
                        },
                        b'\x04' => {
                            // ipv6
                            let addrport = [0; 16+2];
                            Box::new(read_exact(r, addrport)
                                    .map_err(box_up_err)
                                    .and_then(move |(r,addrport)| {
                                let port = (addrport[16] as u16)*256 + (addrport[17] as u16);
                                // still not worth to switch to Cargo.toml to add
                                // "bytes" dependency, then scroll up for "extern crate",
                                // then look up docs again to find out where to get that BE thing.
                                let mut ip = [0u8;16];
                                ip.copy_from_slice(&addrport[0..16]);
                                let host = SocksHostAddr::Ip(IpAddr::V6(ip.into()));
                                info!("Connected though the SOCKS5 proxy from {:?}",
                                    SocksSocketAddr{host,port});

                                ok(Peer::new(r,w))
                            }))
                        },
                        b'\x03' => {

                            let alen = [0; 1];
                            Box::new(read_exact(r, alen)
                                    .map_err(box_up_err)
                                    .and_then(move |(r,alen)| {

                                let alen = alen[0] as usize;
                                let addrport = vec![0;alen+2];

                                read_exact(r, addrport)
                                        .map_err(box_up_err)
                                        .and_then(move |(r,addrport)| {

                                    let port = (addrport[alen] as u16)*256 + (addrport[alen+1] as u16);
                                    let host =
                                        SocksHostAddr::Name(
                                        ::std::str::from_utf8(&addrport[0..alen])
                                        .unwrap_or("(invalid hostname)").to_string());
                                info!("Connected though the SOCKS5 proxy from {:?}",
                                    SocksSocketAddr{host,port});

                                    ok(Peer::new(r,w))
                                })
                            }))
                        },
                        _ => {
                            return peer_strerr("SOCKS: bound address type is unknown");
                        },
                    }
                })
            })) as BoxedNewPeerFuture
        })
    });
    Box::new(f) as BoxedNewPeerFuture
}
<|endoftext|>


<commit_before>use std::str::FromStr;

/// Description for new columns
#[derive(Debug, Clone)]
pub struct ColumnDescription {
    pub name: String,

    // TODO: make this use one of the enums
    /// Type of the data, see the cfitsio documentation
    pub data_type: ColumnDataDescription,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ColumnDataDescription {
    pub repeat: usize,
    pub width: usize,
    pub typ: ColumnDataType,
}

impl ColumnDataDescription {
    pub fn new(typ: ColumnDataType) -> Self {
        ColumnDataDescription {
            repeat: 1,
            width: 1,
            typ: typ,
        }
    }

    /* XXX These two methods force a call to clone which is wasteful of memory. I do not know if
     * this means that memory is leaked, or that destructors are needlessly called (I suspect the
     * latter) but it is fairly wasteful. On the other hand, it's unlikely this sort of thing will
     * be called in performance-critical code, and is more likely a one-dime definition. I will
     * leave it for now - SRW 2017-03-07
     * */
    pub fn repeats(&mut self, repeat: usize) -> Self {
        // TODO check that repeat >= 1
        self.repeat = repeat;
        self.clone()
    }

    pub fn width(&mut self, width: usize) -> Self {
        // TODO check that width >= 1
        self.width = width;
        self.clone()
    }
}

impl From<ColumnDataDescription> for String {
    fn from(orig: ColumnDataDescription) -> String {
        match orig.typ {
            ColumnDataType::Text => {
                if orig.width > 1 {
                    format!("{repeat}{data_type}{width}",
                            data_type = String::from(orig.typ),
                            repeat = orig.repeat,
                            width = orig.width)
                } else {
                    format!("{repeat}{data_type}",
                            data_type = String::from(orig.typ),
                            repeat = orig.repeat)
                }
            }
            _ => {
                format!("{repeat}{data_type}",
                        data_type = String::from(orig.typ),
                        repeat = orig.repeat)
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ColumnDataType {
    Int,
    Float,
    Text,
    Double,
    Short,
    Long,
}

impl From<ColumnDataType> for String {
    fn from(orig: ColumnDataType) -> String {
        use self::ColumnDataType::*;

        match orig {
                Int => "J",
                Float => "E",
                Text => "A",
                Double => "D",
                Short => "I",
                Long => "K",
            }
            .to_string()
    }
}

impl FromStr for ColumnDataDescription {
    type Err = Box<::std::error::Error>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let chars: Vec<_> = s.chars().collect();

        let mut repeat_str = Vec::new();
        let mut last_position = 0;
        for i in 0..chars.len() {
            if chars[i].is_digit(10) {
                repeat_str.push(chars[i]);
                last_position += 1;
            } else {
                break;
            }
        }

        let repeat = if repeat_str.is_empty() {
            1
        } else {
            let repeat_str: String = repeat_str.into_iter().collect();
            repeat_str.parse::<usize>()?
        };

        let data_type_char = chars[last_position];
        last_position += 1;

        let mut width_str = Vec::new();
        for i in last_position..chars.len() {
            if chars[i].is_digit(10) {
                width_str.push(chars[i]);
            } else {
                break;
            }
        }

        /* TODO: validate that the whole string has been used up */

        let width = if width_str.is_empty() {
            1
        } else {
            let width_str: String = width_str.into_iter().collect();
            width_str.parse::<usize>()?
        };

        let data_type = match data_type_char {
            'E' => ColumnDataType::Float,
            'J' => ColumnDataType::Int,
            'D' => ColumnDataType::Double,
            'I' => ColumnDataType::Short,
            'K' => ColumnDataType::Long,
            _ => {
                panic!("Have not implemented str -> ColumnDataType for {}",
                       data_type_char)
            }
        };

        Ok(ColumnDataDescription {
            repeat: repeat,
            typ: data_type,
            width: width,
        })
    }
}

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

    #[test]
    fn test_column_data_descriptions_builder_pattern() {
        let desc = ColumnDataDescription::new(ColumnDataType::Int)
            .width(100)
            .repeats(5);
        assert_eq!(desc.repeat, 5);
        assert_eq!(desc.width, 100);
    }

    #[test]
    fn from_impls() {
        {
            let desc = ColumnDataDescription::new(ColumnDataType::Int).repeats(5);
            assert_eq!(String::from(desc), "5J");
        }

        {
            let desc = ColumnDataDescription::new(ColumnDataType::Float);
            assert_eq!(String::from(desc), "1E");
        }

        {
            let desc = ColumnDataDescription::new(ColumnDataType::Text).width(100);
            assert_eq!(String::from(desc), "1A100");
        }
    }

    #[test]
    fn parsing() {
        let s = "1E";
        assert_eq!(s.parse::<ColumnDataDescription>().unwrap(),
                   ColumnDataDescription {
                       repeat: 1,
                       width: 1,
                       typ: ColumnDataType::Float,
                   });
    }

    #[test]
    fn parse_many_repeats() {
        let s = "100E";
        assert_eq!(s.parse::<ColumnDataDescription>().unwrap(),
                   ColumnDataDescription {
                       repeat: 100,
                       width: 1,
                       typ: ColumnDataType::Float,
                   });
    }

    #[test]
    fn parse_with_width() {
        let s = "1E26";
        assert_eq!(s.parse::<ColumnDataDescription>().unwrap(),
                   ColumnDataDescription {
                       repeat: 1,
                       width: 26,
                       typ: ColumnDataType::Float,
                   });
    }
}
<commit_msg>Remove todo<commit_after>use std::str::FromStr;

/// Description for new columns
#[derive(Debug, Clone)]
pub struct ColumnDescription {
    pub name: String,

    /// Type of the data, see the cfitsio documentation
    pub data_type: ColumnDataDescription,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ColumnDataDescription {
    pub repeat: usize,
    pub width: usize,
    pub typ: ColumnDataType,
}

impl ColumnDataDescription {
    pub fn new(typ: ColumnDataType) -> Self {
        ColumnDataDescription {
            repeat: 1,
            width: 1,
            typ: typ,
        }
    }

    /* XXX These two methods force a call to clone which is wasteful of memory. I do not know if
     * this means that memory is leaked, or that destructors are needlessly called (I suspect the
     * latter) but it is fairly wasteful. On the other hand, it's unlikely this sort of thing will
     * be called in performance-critical code, and is more likely a one-dime definition. I will
     * leave it for now - SRW 2017-03-07
     * */
    pub fn repeats(&mut self, repeat: usize) -> Self {
        // TODO check that repeat >= 1
        self.repeat = repeat;
        self.clone()
    }

    pub fn width(&mut self, width: usize) -> Self {
        // TODO check that width >= 1
        self.width = width;
        self.clone()
    }
}

impl From<ColumnDataDescription> for String {
    fn from(orig: ColumnDataDescription) -> String {
        match orig.typ {
            ColumnDataType::Text => {
                if orig.width > 1 {
                    format!("{repeat}{data_type}{width}",
                            data_type = String::from(orig.typ),
                            repeat = orig.repeat,
                            width = orig.width)
                } else {
                    format!("{repeat}{data_type}",
                            data_type = String::from(orig.typ),
                            repeat = orig.repeat)
                }
            }
            _ => {
                format!("{repeat}{data_type}",
                        data_type = String::from(orig.typ),
                        repeat = orig.repeat)
            }
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ColumnDataType {
    Int,
    Float,
    Text,
    Double,
    Short,
    Long,
}

impl From<ColumnDataType> for String {
    fn from(orig: ColumnDataType) -> String {
        use self::ColumnDataType::*;

        match orig {
                Int => "J",
                Float => "E",
                Text => "A",
                Double => "D",
                Short => "I",
                Long => "K",
            }
            .to_string()
    }
}

impl FromStr for ColumnDataDescription {
    type Err = Box<::std::error::Error>;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let chars: Vec<_> = s.chars().collect();

        let mut repeat_str = Vec::new();
        let mut last_position = 0;
        for i in 0..chars.len() {
            if chars[i].is_digit(10) {
                repeat_str.push(chars[i]);
                last_position += 1;
            } else {
                break;
            }
        }

        let repeat = if repeat_str.is_empty() {
            1
        } else {
            let repeat_str: String = repeat_str.into_iter().collect();
            repeat_str.parse::<usize>()?
        };

        let data_type_char = chars[last_position];
        last_position += 1;

        let mut width_str = Vec::new();
        for i in last_position..chars.len() {
            if chars[i].is_digit(10) {
                width_str.push(chars[i]);
            } else {
                break;
            }
        }

        /* TODO: validate that the whole string has been used up */

        let width = if width_str.is_empty() {
            1
        } else {
            let width_str: String = width_str.into_iter().collect();
            width_str.parse::<usize>()?
        };

        let data_type = match data_type_char {
            'E' => ColumnDataType::Float,
            'J' => ColumnDataType::Int,
            'D' => ColumnDataType::Double,
            'I' => ColumnDataType::Short,
            'K' => ColumnDataType::Long,
            _ => {
                panic!("Have not implemented str -> ColumnDataType for {}",
                       data_type_char)
            }
        };

        Ok(ColumnDataDescription {
            repeat: repeat,
            typ: data_type,
            width: width,
        })
    }
}

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

    #[test]
    fn test_column_data_descriptions_builder_pattern() {
        let desc = ColumnDataDescription::new(ColumnDataType::Int)
            .width(100)
            .repeats(5);
        assert_eq!(desc.repeat, 5);
        assert_eq!(desc.width, 100);
    }

    #[test]
    fn from_impls() {
        {
            let desc = ColumnDataDescription::new(ColumnDataType::Int).repeats(5);
            assert_eq!(String::from(desc), "5J");
        }

        {
            let desc = ColumnDataDescription::new(ColumnDataType::Float);
            assert_eq!(String::from(desc), "1E");
        }

        {
            let desc = ColumnDataDescription::new(ColumnDataType::Text).width(100);
            assert_eq!(String::from(desc), "1A100");
        }
    }

    #[test]
    fn parsing() {
        let s = "1E";
        assert_eq!(s.parse::<ColumnDataDescription>().unwrap(),
                   ColumnDataDescription {
                       repeat: 1,
                       width: 1,
                       typ: ColumnDataType::Float,
                   });
    }

    #[test]
    fn parse_many_repeats() {
        let s = "100E";
        assert_eq!(s.parse::<ColumnDataDescription>().unwrap(),
                   ColumnDataDescription {
                       repeat: 100,
                       width: 1,
                       typ: ColumnDataType::Float,
                   });
    }

    #[test]
    fn parse_with_width() {
        let s = "1E26";
        assert_eq!(s.parse::<ColumnDataDescription>().unwrap(),
                   ColumnDataDescription {
                       repeat: 1,
                       width: 26,
                       typ: ColumnDataType::Float,
                   });
    }
}
<|endoftext|>


<commit_before>#[crate_id = "sobol#0.1"];
#[comment = "Sobol indices evaluator"];
#[license = "MIT"];

extern crate extra;
extern crate collections;

use params::{Params, Float, Range};
use model::{Model,Gompertz};
use std::vec;
mod model;
mod params;

fn sample(pr: &Params, num_real : uint) -> ~[Params] {
    let mut samples : ~[Params] = ~[];
    for _ in range(0, num_real) {
        samples.push(pr.realize());
    }
    samples
}
#[test]
fn test_sample() {
    let mut ranges=Params::new();
    ranges.insert("alpha", Range(0.0, 1.0));
    let samples=sample(&ranges, 10);

    assert_eq!(samples.len(), 10);
    for pp in samples.iter() {
        let value=pp.get_float("alpha");
        assert!((value>=0.0)&&(value<=1.0));
    }
}

#[allow(dead_code)]
fn eval_model(pr: &Params) -> f64 {
    let model : Gompertz = Model::setup(pr);
    model.run(pr.get_float("tmax"))
}

fn get_couples(n : uint) -> ~[(uint, uint)] {
    let mut res : ~[(uint, uint)]=~[];
    for i in range(0, n) {
        for j in range(i, n) {
            if i!=j {
                res.push( (i,j));
            }
        }
    }
    res
}

#[test]
fn test_get_couples() {
    let res=get_couples(3);

    assert_eq!(res[0], (0,1));
    assert_eq!(res[1], (0,2));
    assert_eq!(res[2], (1,2));
}

// First order Sobol indices
#[allow(dead_code)]
fn calc_sobol_1(pr: &Params, nsamp : uint) -> collections::HashMap<~str, f64> {
    let samples_1=sample(pr, nsamp);
    let samples_2=sample(pr, nsamp);

    let big_one : ~[Params] = vec::append(samples_1.clone(), samples_2);

    // Average and deviation computaion
    let resvec=big_one.map(eval_model);
    let avg=resvec.iter().fold(0.0, |acc, &item| acc+item)/(resvec.len() as f64);
    let var=resvec.iter().fold(0.0, |acc, &item| acc+item*item)/(resvec.len() as f64)-avg*avg;

    let mut Sp: collections::HashMap<~str, f64> = collections::HashMap::new();
    let vkeys=pr.varying_keys();
    for k in vkeys.iter() {
        let mut sump=0.0f64;
        for snum in range(0, nsamp) {
            let mut new_p=samples_2[snum].clone();
            new_p.set(k.to_owned(),Float(samples_1[snum].get_float(k.to_owned())));

            sump=sump+eval_model(&new_p)*eval_model(&samples_1[snum]);
        }
        let Up=sump/(nsamp as f64);
        Sp.insert(k.to_owned(), (Up-avg*avg)/var);
    }

    Sp
}

// Second order Sobol indices
#[allow(dead_code)]
fn calc_sobol_2(pr: &Params, sob1 : &collections::HashMap<~str, f64>, nsamp : uint) -> ~[(~str, ~str, f64)] {
    let samples_1=sample(pr, nsamp);
    let samples_2=sample(pr, nsamp);
    let keys=pr.varying_keys();
    let couples=get_couples(keys.len());

    let big_one : ~[Params] = vec::append(samples_1.clone(), samples_2);

    // Average and deviation computaion
    let resvec=big_one.map(eval_model);
    let avg=resvec.iter().fold(0.0, |acc, &item| acc+item)/(resvec.len() as f64);
    let var=resvec.iter().fold(0.0, |acc, &item| acc+item*item)/(resvec.len() as f64)-avg*avg;

    let mut Spr : ~[(~str, ~str, f64)]=~[];

    for &cpl in couples.iter() {
        let (i,j) = cpl;
        let ref key1=keys[i]; let ref key2=keys[j];
        let mut sump=0.0f64;
        for snum in range(0, nsamp) {
            let mut new_p=samples_2[snum].clone();
            new_p.set(key1.to_owned(),Float(samples_1[snum].get_float(key1.to_owned())));
            new_p.set(key2.to_owned(),Float(samples_1[snum].get_float(key2.to_owned())));

            sump=sump+eval_model(&new_p)*eval_model(&samples_1[snum]);
        }
        let Upr=sump/(nsamp as f64);
        let Vp=*sob1.find_equiv(key1).unwrap();
        let Vr=*sob1.find_equiv(key2).unwrap();
        Spr.push((key1.to_owned(), key2.to_owned(), (Upr-avg*avg-Vr-Vp)/var));
    }

    Spr
}

// Total Sobol indices
fn calc_sobol_total(pr: &Params, nsamp: uint) -> collections::HashMap<~str, f64> {
    let samples_1=sample(pr, nsamp);
    let samples_2=sample(pr, nsamp);

    let big_one : ~[Params] = vec::append(samples_1.clone(), samples_2);

    // Average and deviation computaion
    let resvec=big_one.map(eval_model);
    let avg=resvec.iter().fold(0.0, |acc, &item| acc+item)/(resvec.len() as f64);
    let var=resvec.iter().fold(0.0, |acc, &item| acc+item*item)/(resvec.len() as f64)-avg*avg;

    let mut Sp: collections::HashMap<~str, f64> = collections::HashMap::new();
    let vkeys=pr.varying_keys();
    for k in vkeys.iter() {
        let mut sump=0.0f64;
        for snum in range(0, nsamp) {
            let mut new_p=samples_1[snum].clone();
            new_p.set(k.to_owned(),Float(samples_2[snum].get_float(k.slice_from(0))));

            sump=sump+eval_model(&new_p)*eval_model(&samples_1[snum]);
        }
        let Up=sump/(nsamp as f64);
        Sp.insert(k.to_owned(), 1.0-(Up-avg*avg)/var);
    }

    Sp
}

//
// Entry point
//
#[allow(dead_code)]
fn main() {
  let mut ranges=Params::new();
  ranges.insert("alpha", Range(0.01, 1.0)); ranges.insert("K", Range(0.01, 1.0));
  ranges.insert("C0", Range(0.001, 0.01)); ranges.insert("tmax", Float(10f64));

  let samples=sample(&ranges,5000);
  let resvec=samples.map(eval_model);

  let avg=resvec.iter().fold(0.0, |acc, &item| acc+item)/(resvec.len() as f64);
  let var=resvec.iter().fold(0.0, |acc, &item| acc+item*item)/(resvec.len() as f64)-avg*avg;

  println!("Mean value = {}, Standard Deviation = {}\n", avg, var);

  println!("First order Sobol indices");
  let sob_1=calc_sobol_1(&ranges,15000);
  for (k, v) in sob_1.iter() {
    println!("\t{:s} = {:f}", *k, *v);
  }

  println!("\nSecond order Sobol indices");
  let sob_2=calc_sobol_2(&ranges, &sob_1, 7500);
  for vv in sob_2.iter() {
    let (ref k1,ref k2,v) = *vv;
    println!("\t{:s} x {:s} = {:f}", *k1, *k2, v);
  }

  println!("\nTotal Sobol indices");
  let sob_t=calc_sobol_total(&ranges, 15000);
  for (k, v) in sob_t.iter() {
    println!("\t{:s} = {:f}", *k, *v);
  }

}

<commit_msg>Remove libextra<commit_after>#[crate_id = "sobol#0.1"];
#[comment = "Sobol indices evaluator"];
#[license = "MIT"];

extern crate collections;

use params::{Params, Float, Range};
use model::{Model,Gompertz};
use std::vec;
mod model;
mod params;

fn sample(pr: &Params, num_real : uint) -> ~[Params] {
    let mut samples : ~[Params] = ~[];
    for _ in range(0, num_real) {
        samples.push(pr.realize());
    }
    samples
}
#[test]
fn test_sample() {
    let mut ranges=Params::new();
    ranges.insert("alpha", Range(0.0, 1.0));
    let samples=sample(&ranges, 10);

    assert_eq!(samples.len(), 10);
    for pp in samples.iter() {
        let value=pp.get_float("alpha");
        assert!((value>=0.0)&&(value<=1.0));
    }
}

#[allow(dead_code)]
fn eval_model(pr: &Params) -> f64 {
    let model : Gompertz = Model::setup(pr);
    model.run(pr.get_float("tmax"))
}

fn get_couples(n : uint) -> ~[(uint, uint)] {
    let mut res : ~[(uint, uint)]=~[];
    for i in range(0, n) {
        for j in range(i, n) {
            if i!=j {
                res.push( (i,j));
            }
        }
    }
    res
}

#[test]
fn test_get_couples() {
    let res=get_couples(3);

    assert_eq!(res[0], (0,1));
    assert_eq!(res[1], (0,2));
    assert_eq!(res[2], (1,2));
}

// First order Sobol indices
#[allow(dead_code)]
fn calc_sobol_1(pr: &Params, nsamp : uint) -> collections::HashMap<~str, f64> {
    let samples_1=sample(pr, nsamp);
    let samples_2=sample(pr, nsamp);

    let big_one : ~[Params] = vec::append(samples_1.clone(), samples_2);

    // Average and deviation computaion
    let resvec=big_one.map(eval_model);
    let avg=resvec.iter().fold(0.0, |acc, &item| acc+item)/(resvec.len() as f64);
    let var=resvec.iter().fold(0.0, |acc, &item| acc+item*item)/(resvec.len() as f64)-avg*avg;

    let mut Sp: collections::HashMap<~str, f64> = collections::HashMap::new();
    let vkeys=pr.varying_keys();
    for k in vkeys.iter() {
        let mut sump=0.0f64;
        for snum in range(0, nsamp) {
            let mut new_p=samples_2[snum].clone();
            new_p.set(k.to_owned(),Float(samples_1[snum].get_float(k.to_owned())));

            sump=sump+eval_model(&new_p)*eval_model(&samples_1[snum]);
        }
        let Up=sump/(nsamp as f64);
        Sp.insert(k.to_owned(), (Up-avg*avg)/var);
    }

    Sp
}

// Second order Sobol indices
#[allow(dead_code)]
fn calc_sobol_2(pr: &Params, sob1 : &collections::HashMap<~str, f64>, nsamp : uint) -> ~[(~str, ~str, f64)] {
    let samples_1=sample(pr, nsamp);
    let samples_2=sample(pr, nsamp);
    let keys=pr.varying_keys();
    let couples=get_couples(keys.len());

    let big_one : ~[Params] = vec::append(samples_1.clone(), samples_2);

    // Average and deviation computaion
    let resvec=big_one.map(eval_model);
    let avg=resvec.iter().fold(0.0, |acc, &item| acc+item)/(resvec.len() as f64);
    let var=resvec.iter().fold(0.0, |acc, &item| acc+item*item)/(resvec.len() as f64)-avg*avg;

    let mut Spr : ~[(~str, ~str, f64)]=~[];

    for &cpl in couples.iter() {
        let (i,j) = cpl;
        let ref key1=keys[i]; let ref key2=keys[j];
        let mut sump=0.0f64;
        for snum in range(0, nsamp) {
            let mut new_p=samples_2[snum].clone();
            new_p.set(key1.to_owned(),Float(samples_1[snum].get_float(key1.to_owned())));
            new_p.set(key2.to_owned(),Float(samples_1[snum].get_float(key2.to_owned())));

            sump=sump+eval_model(&new_p)*eval_model(&samples_1[snum]);
        }
        let Upr=sump/(nsamp as f64);
        let Vp=*sob1.find_equiv(key1).unwrap();
        let Vr=*sob1.find_equiv(key2).unwrap();
        Spr.push((key1.to_owned(), key2.to_owned(), (Upr-avg*avg-Vr-Vp)/var));
    }

    Spr
}

// Total Sobol indices
fn calc_sobol_total(pr: &Params, nsamp: uint) -> collections::HashMap<~str, f64> {
    let samples_1=sample(pr, nsamp);
    let samples_2=sample(pr, nsamp);

    let big_one : ~[Params] = vec::append(samples_1.clone(), samples_2);

    // Average and deviation computaion
    let resvec=big_one.map(eval_model);
    let avg=resvec.iter().fold(0.0, |acc, &item| acc+item)/(resvec.len() as f64);
    let var=resvec.iter().fold(0.0, |acc, &item| acc+item*item)/(resvec.len() as f64)-avg*avg;

    let mut Sp: collections::HashMap<~str, f64> = collections::HashMap::new();
    let vkeys=pr.varying_keys();
    for k in vkeys.iter() {
        let mut sump=0.0f64;
        for snum in range(0, nsamp) {
            let mut new_p=samples_1[snum].clone();
            new_p.set(k.to_owned(),Float(samples_2[snum].get_float(k.slice_from(0))));

            sump=sump+eval_model(&new_p)*eval_model(&samples_1[snum]);
        }
        let Up=sump/(nsamp as f64);
        Sp.insert(k.to_owned(), 1.0-(Up-avg*avg)/var);
    }

    Sp
}

//
// Entry point
//
#[allow(dead_code)]
fn main() {
  let mut ranges=Params::new();
  ranges.insert("alpha", Range(0.01, 1.0)); ranges.insert("K", Range(0.01, 1.0));
  ranges.insert("C0", Range(0.001, 0.01)); ranges.insert("tmax", Float(10f64));

  let samples=sample(&ranges,5000);
  let resvec=samples.map(eval_model);

  let avg=resvec.iter().fold(0.0, |acc, &item| acc+item)/(resvec.len() as f64);
  let var=resvec.iter().fold(0.0, |acc, &item| acc+item*item)/(resvec.len() as f64)-avg*avg;

  println!("Mean value = {}, Standard Deviation = {}\n", avg, var);

  println!("First order Sobol indices");
  let sob_1=calc_sobol_1(&ranges,15000);
  for (k, v) in sob_1.iter() {
    println!("\t{:s} = {:f}", *k, *v);
  }

  println!("\nSecond order Sobol indices");
  let sob_2=calc_sobol_2(&ranges, &sob_1, 7500);
  for vv in sob_2.iter() {
    let (ref k1,ref k2,v) = *vv;
    println!("\t{:s} x {:s} = {:f}", *k1, *k2, v);
  }

  println!("\nTotal Sobol indices");
  let sob_t=calc_sobol_total(&ranges, 15000);
  for (k, v) in sob_t.iter() {
    println!("\t{:s} = {:f}", *k, *v);
  }

}

<|endoftext|>


<commit_before>#[macro_use] extern crate diesel;

table! {
     #[foobar]
     posts {
         id -> Integer,
     }
}
// error-pattern: Invalid `table!` syntax. Please see the `table!` macro docs for more info. `https://docs.diesel.rs/diesel/macro.table.html`

fn main() {}
<commit_msg>fix compiletests<commit_after>#[macro_use] extern crate diesel;

table! {
     #[foobar]
     posts {
         id -> Integer,
     }
}
// error-pattern: E0658

fn main() {}
<|endoftext|>