<commit_before>use domain::ledger::constants;
use errors::common::CommonError;
use errors::pool::PoolError;
use serde_json;
use serde_json::Value as SJsonValue;
use services::ledger::merkletree::merkletree::MerkleTree;
use services::pool::types::*;
use std::error::Error;

pub const REQUESTS_FOR_STATE_PROOFS: [&'static str; 7] = [
    constants::GET_NYM,
    constants::GET_SCHEMA,
    constants::GET_CRED_DEF,
    constants::GET_ATTR,
    constants::GET_REVOC_REG,
    constants::GET_REVOC_REG_DEF,
    constants::GET_REVOC_REG_DELTA,
];

const REQUEST_FOR_FULL: [&'static str; 2] = [
    constants::POOL_RESTART,
    constants::GET_VALIDATOR_INFO,
];

#[derive(Debug, Clone)]
pub enum NetworkerEvent {
    SendOneRequest(
        String, //msg
        String, //req_id
        i64, //timeout
    ),
    SendAllRequest(
        String, //msg
        String, //req_id
        i64, //timeout
        Option<Vec<String>>, //nodes
    ),
    Resend(
        String, //req_id
        i64, //timeout
    ),
    NodesStateUpdated(Vec<RemoteNode>),
    ExtendTimeout(
        String, //req_id
        String, //node_alias
        i64, //timeout
    ),
    CleanTimeout(
        String, //req_id
        Option<String>, //node_alias
    ),
    Timeout,
}

#[derive(Clone, Debug)]
pub enum PoolEvent {
    CheckCache(i32),
    NodeReply(
        String, // reply
        String, // node alias
    ),
    Close(
        i32, //cmd_id
    ),
    Refresh(
        i32, //cmd_id
    ),
    CatchupTargetFound(
        Vec<u8>, //target_mt_root
        usize, //target_mt_size
        MerkleTree,
    ),
    CatchupTargetNotFound(PoolError),
    #[allow(dead_code)] //FIXME
    PoolOutdated,
    Synced(
        MerkleTree
    ),
    #[allow(dead_code)] //FIXME
    NodesBlacklisted,
    SendRequest(
        i32, // cmd_id
        String, // request
        Option<i32>, // timeout
        Option<String>, // node list
    ),
    Timeout(
        String, //req_id
        String, //node alias
    ),
}

#[derive(Clone, Debug)]
pub enum RequestEvent {
    LedgerStatus(
        LedgerStatus,
        Option<String>, //node alias
        Option<MerkleTree>,
    ),
    CatchupReq(
        MerkleTree,
        usize, // target mt size
        Vec<u8>, // target mt root
    ),
    Timeout(
        String, //req_id
        String, //node_alias
    ),
    CatchupRep(
        CatchupRep,
        String, // node_alias
    ),
    CustomSingleRequest(
        String, // message
        String, // req_id
    ),
    CustomConsensusRequest(
        String, // message
        String, // req_id
    ),
    CustomFullRequest(
        String, // message
        String, // req_id
        Option<i32>, // timeout
        Option<String>, // nodes
    ),
    ConsistencyProof(
        ConsistencyProof,
        String, //node alias
    ),
    Reply(
        Reply,
        String, //raw_msg
        String, //node alias
        String, //req_id
    ),
    ReqACK(
        Response,
        String, //raw_msg
        String, //node alias
        String, //req_id
    ),
    ReqNACK(
        Response,
        String, //raw_msg
        String, //node alias
        String, //req_id
    ),
    Reject(
        Response,
        String, //raw_msg
        String, //node alias
        String, //req_id
    ),
    PoolLedgerTxns,
    Ping,
    Pong,
    Terminate,
}

impl RequestEvent {
    pub fn get_req_id(&self) -> String {
        match self {
            &RequestEvent::CustomSingleRequest(_, ref id) => id.to_string(),
            &RequestEvent::CustomConsensusRequest(_, ref id) => id.to_string(),
            &RequestEvent::CustomFullRequest(_, ref id, _, _) => id.to_string(),
            &RequestEvent::Reply(_, _, _, ref id) => id.to_string(),
            &RequestEvent::ReqACK(_, _, _, ref id) => id.to_string(),
            &RequestEvent::ReqNACK(_, _, _, ref id) => id.to_string(),
            &RequestEvent::Reject(_, _, _, ref id) => id.to_string(),
            _ => "".to_string()
        }
    }
}

impl Into<Option<RequestEvent>> for PoolEvent {
    fn into(self) -> Option<RequestEvent> {
        match self {
            PoolEvent::NodeReply(msg, node_alias) => {
                _parse_msg(&msg).map(|parsed|
                    match parsed {
                        //TODO change mapping for CatchupReq. May be return None
                        //TODO: REMOVE UNWRAP!!!!!
                        Message::CatchupReq(_) => RequestEvent::CatchupReq(MerkleTree::from_vec(Vec::new()).unwrap(), 0, vec![]),
                        Message::CatchupRep(rep) => RequestEvent::CatchupRep(rep, node_alias),
                        Message::LedgerStatus(ls) => RequestEvent::LedgerStatus(ls, Some(node_alias), None),
                        Message::ConsistencyProof(cp) => RequestEvent::ConsistencyProof(cp, node_alias),
                        Message::Reply(rep) => {
                            let req_id = rep.req_id();
                            RequestEvent::Reply(rep, msg, node_alias, req_id.to_string())
                        }
                        Message::ReqACK(rep) => {
                            let req_id = rep.req_id();
                            RequestEvent::ReqACK(rep, msg, node_alias, req_id.to_string())
                        }
                        Message::ReqNACK(rep) => {
                            let req_id = rep.req_id();
                            RequestEvent::ReqNACK(rep, msg, node_alias, req_id.to_string())
                        }
                        Message::Reject(rep) => {
                            let req_id = rep.req_id();
                            RequestEvent::Reject(rep, msg, node_alias, req_id.to_string())
                        }
                        Message::PoolLedgerTxns(_) => RequestEvent::PoolLedgerTxns,
                        Message::Ping => RequestEvent::Ping,
                        Message::Pong => RequestEvent::Pong,
                    })
            }
            PoolEvent::SendRequest(_, msg, timeout, nodes) => {
                let req_id = _parse_req_id_and_op(&msg);
                if let Ok((ref req_id, ref op)) = req_id {
                    if REQUEST_FOR_FULL.contains(&op.as_str()) {
                        Some(RequestEvent::CustomFullRequest(msg, req_id.clone(), timeout, nodes))
                    } else if timeout.is_some() || nodes.is_some() {
                        error!("Timeout {:?} or nodes {:?} is specified for non-supported request operation type {}",
                               timeout, nodes, op);
                        None
                    } else if REQUESTS_FOR_STATE_PROOFS.contains(&op.as_str()) /* FIXME check plugged also */ {
                        Some(RequestEvent::CustomSingleRequest(msg, req_id.clone()))
                    } else {
                        Some(RequestEvent::CustomConsensusRequest(msg, req_id.clone()))
                    }
                } else {
                    error!("Can't parse req_id or op from message {}", msg);
                    None
                }
            }
            PoolEvent::Timeout(req_id, node_alias) => Some(RequestEvent::Timeout(req_id, node_alias)),
            _ => None
        }
    }
}

fn _parse_msg(msg: &str) -> Option<Message> {
    Message::from_raw_str(msg).map_err(map_err_trace!()).ok()
}

fn _parse_req_id_and_op(msg: &str) -> Result<(String, String), CommonError> {
    let req_json = _get_req_json(msg)?;

    let req_id: u64 = req_json["reqId"]
        .as_u64()
        .ok_or(CommonError::InvalidStructure("No reqId in request".to_string()))?;

    let op = req_json["operation"]["type"]
        .as_str()
        .ok_or(CommonError::InvalidStructure("No operation type in request".to_string()))?;

    Ok((req_id.to_string(), op.to_string()))
}

fn _get_req_json(msg: &str) -> Result<SJsonValue, CommonError> {
    serde_json::from_str(msg)
        .map_err(|err|
            CommonError::InvalidStructure(
                format!("Invalid request json: {}", err.description())))
}<commit_msg>Resolve FIXME about StateProof for plugged txns.<commit_after>use domain::ledger::constants;
use errors::common::CommonError;
use errors::pool::PoolError;
use serde_json;
use serde_json::Value as SJsonValue;
use services::ledger::merkletree::merkletree::MerkleTree;
use services::pool::{PoolService, types::*};
use std::error::Error;

pub const REQUESTS_FOR_STATE_PROOFS: [&'static str; 7] = [
    constants::GET_NYM,
    constants::GET_SCHEMA,
    constants::GET_CRED_DEF,
    constants::GET_ATTR,
    constants::GET_REVOC_REG,
    constants::GET_REVOC_REG_DEF,
    constants::GET_REVOC_REG_DELTA,
];

const REQUEST_FOR_FULL: [&'static str; 2] = [
    constants::POOL_RESTART,
    constants::GET_VALIDATOR_INFO,
];

#[derive(Debug, Clone)]
pub enum NetworkerEvent {
    SendOneRequest(
        String, //msg
        String, //req_id
        i64, //timeout
    ),
    SendAllRequest(
        String, //msg
        String, //req_id
        i64, //timeout
        Option<Vec<String>>, //nodes
    ),
    Resend(
        String, //req_id
        i64, //timeout
    ),
    NodesStateUpdated(Vec<RemoteNode>),
    ExtendTimeout(
        String, //req_id
        String, //node_alias
        i64, //timeout
    ),
    CleanTimeout(
        String, //req_id
        Option<String>, //node_alias
    ),
    Timeout,
}

#[derive(Clone, Debug)]
pub enum PoolEvent {
    CheckCache(i32),
    NodeReply(
        String, // reply
        String, // node alias
    ),
    Close(
        i32, //cmd_id
    ),
    Refresh(
        i32, //cmd_id
    ),
    CatchupTargetFound(
        Vec<u8>, //target_mt_root
        usize, //target_mt_size
        MerkleTree,
    ),
    CatchupTargetNotFound(PoolError),
    #[allow(dead_code)] //FIXME
    PoolOutdated,
    Synced(
        MerkleTree
    ),
    #[allow(dead_code)] //FIXME
    NodesBlacklisted,
    SendRequest(
        i32, // cmd_id
        String, // request
        Option<i32>, // timeout
        Option<String>, // node list
    ),
    Timeout(
        String, //req_id
        String, //node alias
    ),
}

#[derive(Clone, Debug)]
pub enum RequestEvent {
    LedgerStatus(
        LedgerStatus,
        Option<String>, //node alias
        Option<MerkleTree>,
    ),
    CatchupReq(
        MerkleTree,
        usize, // target mt size
        Vec<u8>, // target mt root
    ),
    Timeout(
        String, //req_id
        String, //node_alias
    ),
    CatchupRep(
        CatchupRep,
        String, // node_alias
    ),
    CustomSingleRequest(
        String, // message
        String, // req_id
    ),
    CustomConsensusRequest(
        String, // message
        String, // req_id
    ),
    CustomFullRequest(
        String, // message
        String, // req_id
        Option<i32>, // timeout
        Option<String>, // nodes
    ),
    ConsistencyProof(
        ConsistencyProof,
        String, //node alias
    ),
    Reply(
        Reply,
        String, //raw_msg
        String, //node alias
        String, //req_id
    ),
    ReqACK(
        Response,
        String, //raw_msg
        String, //node alias
        String, //req_id
    ),
    ReqNACK(
        Response,
        String, //raw_msg
        String, //node alias
        String, //req_id
    ),
    Reject(
        Response,
        String, //raw_msg
        String, //node alias
        String, //req_id
    ),
    PoolLedgerTxns,
    Ping,
    Pong,
    Terminate,
}

impl RequestEvent {
    pub fn get_req_id(&self) -> String {
        match self {
            &RequestEvent::CustomSingleRequest(_, ref id) => id.to_string(),
            &RequestEvent::CustomConsensusRequest(_, ref id) => id.to_string(),
            &RequestEvent::CustomFullRequest(_, ref id, _, _) => id.to_string(),
            &RequestEvent::Reply(_, _, _, ref id) => id.to_string(),
            &RequestEvent::ReqACK(_, _, _, ref id) => id.to_string(),
            &RequestEvent::ReqNACK(_, _, _, ref id) => id.to_string(),
            &RequestEvent::Reject(_, _, _, ref id) => id.to_string(),
            _ => "".to_string()
        }
    }
}

impl Into<Option<RequestEvent>> for PoolEvent {
    fn into(self) -> Option<RequestEvent> {
        match self {
            PoolEvent::NodeReply(msg, node_alias) => {
                _parse_msg(&msg).map(|parsed|
                    match parsed {
                        //TODO change mapping for CatchupReq. May be return None
                        //TODO: REMOVE UNWRAP!!!!!
                        Message::CatchupReq(_) => RequestEvent::CatchupReq(MerkleTree::from_vec(Vec::new()).unwrap(), 0, vec![]),
                        Message::CatchupRep(rep) => RequestEvent::CatchupRep(rep, node_alias),
                        Message::LedgerStatus(ls) => RequestEvent::LedgerStatus(ls, Some(node_alias), None),
                        Message::ConsistencyProof(cp) => RequestEvent::ConsistencyProof(cp, node_alias),
                        Message::Reply(rep) => {
                            let req_id = rep.req_id();
                            RequestEvent::Reply(rep, msg, node_alias, req_id.to_string())
                        }
                        Message::ReqACK(rep) => {
                            let req_id = rep.req_id();
                            RequestEvent::ReqACK(rep, msg, node_alias, req_id.to_string())
                        }
                        Message::ReqNACK(rep) => {
                            let req_id = rep.req_id();
                            RequestEvent::ReqNACK(rep, msg, node_alias, req_id.to_string())
                        }
                        Message::Reject(rep) => {
                            let req_id = rep.req_id();
                            RequestEvent::Reject(rep, msg, node_alias, req_id.to_string())
                        }
                        Message::PoolLedgerTxns(_) => RequestEvent::PoolLedgerTxns,
                        Message::Ping => RequestEvent::Ping,
                        Message::Pong => RequestEvent::Pong,
                    })
            }
            PoolEvent::SendRequest(_, msg, timeout, nodes) => {
                let req_id = _parse_req_id_and_op(&msg);
                if let Ok((ref req_id, ref op)) = req_id {
                    if REQUEST_FOR_FULL.contains(&op.as_str()) {
                        Some(RequestEvent::CustomFullRequest(msg, req_id.clone(), timeout, nodes))
                    } else if timeout.is_some() || nodes.is_some() {
                        error!("Timeout {:?} or nodes {:?} is specified for non-supported request operation type {}",
                               timeout, nodes, op);
                        None
                    } else if REQUESTS_FOR_STATE_PROOFS.contains(&op.as_str())
                        || PoolService::get_sp_parser(&op.as_str()).is_some() {
                        Some(RequestEvent::CustomSingleRequest(msg, req_id.clone()))
                    } else {
                        Some(RequestEvent::CustomConsensusRequest(msg, req_id.clone()))
                    }
                } else {
                    error!("Can't parse req_id or op from message {}", msg);
                    None
                }
            }
            PoolEvent::Timeout(req_id, node_alias) => Some(RequestEvent::Timeout(req_id, node_alias)),
            _ => None
        }
    }
}

fn _parse_msg(msg: &str) -> Option<Message> {
    Message::from_raw_str(msg).map_err(map_err_trace!()).ok()
}

fn _parse_req_id_and_op(msg: &str) -> Result<(String, String), CommonError> {
    let req_json = _get_req_json(msg)?;

    let req_id: u64 = req_json["reqId"]
        .as_u64()
        .ok_or(CommonError::InvalidStructure("No reqId in request".to_string()))?;

    let op = req_json["operation"]["type"]
        .as_str()
        .ok_or(CommonError::InvalidStructure("No operation type in request".to_string()))?;

    Ok((req_id.to_string(), op.to_string()))
}

fn _get_req_json(msg: &str) -> Result<SJsonValue, CommonError> {
    serde_json::from_str(msg)
        .map_err(|err|
            CommonError::InvalidStructure(
                format!("Invalid request json: {}", err.description())))
}<|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.
const BROADCAST_PORT_UDP : u16 = 27900;
const SERVER_LIST_PORT_TCP : u16 = 29920;

// Broadcast packet types.
const KEEPALIVE : u8 = 8;
const HEARTBEAT : u8 = 3;

// Game state changes
const GAMEEXITED : u16 = 2;

// Opcode info
const OPCODE_INDEX : usize = 0;
const OPCODE_AND_HANDSHAKE_LENGTH : usize = 5;

// Broadcasted game name
const HALO_RETAIL : &'static str = "halor";

use std::net::{UdpSocket,TcpListener,SocketAddr};
use std::net::SocketAddr::{V4,V6};
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;

trait IPString {
    fn ip_string(&self) -> String;
}

impl IPString for SocketAddr {
    fn ip_string(&self) -> String {
        // when self.ip() becomes stable, we won't need to match on the IP version
        match *self {
            V4(ipv4) => ipv4.ip().to_string(),
            V6(ipv6) => "[".to_string() + &ipv6.ip().to_string() + "]"
        }
    }
}

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,BROADCAST_PORT_UDP)) {
        Err(error) => {
            println!("Error creating a UDP socket at {}:{}. {}.",ip,BROADCAST_PORT_UDP,error);
            return;
        },
        Ok(halo_socket) => halo_socket
    };

    let client_socket = match TcpListener::bind((&ip as &str,SERVER_LIST_PORT_TCP)) {
        Err(error) => {
            println!("Error listening to TCP at {}:{}. {}.",ip,SERVER_LIST_PORT_TCP,error);
            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_ref = blacklist_update.lock().unwrap();
                match File::open(BLACKLIST_FILE) {
                    Ok(file) => {
                        let reader = BufReader::new(&file);
                        match reader.lines().collect() {
                            Err(_) => {
                                blacklist_ref.clear();
                            },
                            Ok(t) => {
                                *blacklist_ref = t;
                            }
                        }
                    },
                    Err(_) => {}
                };
            }
            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.ip_string()
                };

                let mut ips = String::new();

                // Make servers_ref go out of scope to unlock it for other threads, since we don't need it.
                {
                    let servers_ref = servers_mut_tcp.lock().unwrap();
                    let servers = (*servers_ref).iter();

                    for j in servers {
                        ips.push_str(&(j.to_string()));
                        ips.push('\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.push_str(&ip);
                ips.push_str(":49149:3425");

                // We may be here a while. Just in case...
                thread::spawn( move || {
                    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(x) => x
        };
        if length <= OPCODE_INDEX {
            continue;
        }

        if buffer[OPCODE_INDEX] != KEEPALIVE && buffer[OPCODE_INDEX] != HEARTBEAT {
            continue;
        }

        let client_ip = source.ip_string();

        let blacklist_ref = blacklist_udp.lock().unwrap();
        if blacklist_ref.contains(&client_ip) {
            continue;
        }

        // Heartbeat packet. These contain null-terminated C strings and are ordered in key1[0]value1[0]key2[0]value2[0]key3[0]value3[0] where [0] is a byte equal to 0x00.
        if buffer[OPCODE_INDEX] == HEARTBEAT && length > OPCODE_AND_HANDSHAKE_LENGTH {

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

            match HeartbeatPacket::from_buffer(&buffer[OPCODE_AND_HANDSHAKE_LENGTH..length]) {
                None => {},
                Some(packet) => {
                    let updatetime = time::now().to_timespec().sec;
                    match servers.iter_mut().position(|x| x.ip == client_ip && x.port == packet.localport) {
                        None => {
                            if game_versions.contains(&packet.gamever) && &packet.gamename == HALO_RETAIL {
                                let serverness = HaloServer { ip:client_ip, port: packet.localport, last_alive: updatetime };
                                (*servers).push(serverness);
                            }
                        }
                        Some(k) => {
                            servers[k].last_alive = updatetime;
                            if packet.statechanged == GAMEEXITED {
                                servers.remove(k);
                            }
                        }
                    };
                }
            };
        }

        // Keepalive packet. We need to rely on the origin's port for this, unfortunately. This may mean that the source port is incorrect if the port was changed with NAT.
        else if buffer[OPCODE_INDEX] == KEEPALIVE {
            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>This isn't true actually with the [] syntax<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.
const BROADCAST_PORT_UDP : u16 = 27900;
const SERVER_LIST_PORT_TCP : u16 = 29920;

// Broadcast packet types.
const KEEPALIVE : u8 = 8;
const HEARTBEAT : u8 = 3;

// Game state changes
const GAMEEXITED : u16 = 2;

// Opcode info
const OPCODE_INDEX : usize = 0;
const OPCODE_AND_HANDSHAKE_LENGTH : usize = 5;

// Broadcasted game name
const HALO_RETAIL : &'static str = "halor";

use std::net::{UdpSocket,TcpListener,SocketAddr};
use std::net::SocketAddr::{V4,V6};
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;

trait IPString {
    fn ip_string(&self) -> String;
}

impl IPString for SocketAddr {
    fn ip_string(&self) -> String {
        match *self {
            V4(ipv4) => ipv4.ip().to_string(),
            V6(ipv6) => "[".to_string() + &ipv6.ip().to_string() + "]"
        }
    }
}

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,BROADCAST_PORT_UDP)) {
        Err(error) => {
            println!("Error creating a UDP socket at {}:{}. {}.",ip,BROADCAST_PORT_UDP,error);
            return;
        },
        Ok(halo_socket) => halo_socket
    };

    let client_socket = match TcpListener::bind((&ip as &str,SERVER_LIST_PORT_TCP)) {
        Err(error) => {
            println!("Error listening to TCP at {}:{}. {}.",ip,SERVER_LIST_PORT_TCP,error);
            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_ref = blacklist_update.lock().unwrap();
                match File::open(BLACKLIST_FILE) {
                    Ok(file) => {
                        let reader = BufReader::new(&file);
                        match reader.lines().collect() {
                            Err(_) => {
                                blacklist_ref.clear();
                            },
                            Ok(t) => {
                                *blacklist_ref = t;
                            }
                        }
                    },
                    Err(_) => {}
                };
            }
            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.ip_string()
                };

                let mut ips = String::new();

                // Make servers_ref go out of scope to unlock it for other threads, since we don't need it.
                {
                    let servers_ref = servers_mut_tcp.lock().unwrap();
                    let servers = (*servers_ref).iter();

                    for j in servers {
                        ips.push_str(&(j.to_string()));
                        ips.push('\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.push_str(&ip);
                ips.push_str(":49149:3425");

                // We may be here a while. Just in case...
                thread::spawn( move || {
                    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(x) => x
        };
        if length <= OPCODE_INDEX {
            continue;
        }

        if buffer[OPCODE_INDEX] != KEEPALIVE && buffer[OPCODE_INDEX] != HEARTBEAT {
            continue;
        }

        let client_ip = source.ip_string();

        let blacklist_ref = blacklist_udp.lock().unwrap();
        if blacklist_ref.contains(&client_ip) {
            continue;
        }

        // Heartbeat packet. These contain null-terminated C strings and are ordered in key1[0]value1[0]key2[0]value2[0]key3[0]value3[0] where [0] is a byte equal to 0x00.
        if buffer[OPCODE_INDEX] == HEARTBEAT && length > OPCODE_AND_HANDSHAKE_LENGTH {

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

            match HeartbeatPacket::from_buffer(&buffer[OPCODE_AND_HANDSHAKE_LENGTH..length]) {
                None => {},
                Some(packet) => {
                    let updatetime = time::now().to_timespec().sec;
                    match servers.iter_mut().position(|x| x.ip == client_ip && x.port == packet.localport) {
                        None => {
                            if game_versions.contains(&packet.gamever) && &packet.gamename == HALO_RETAIL {
                                let serverness = HaloServer { ip:client_ip, port: packet.localport, last_alive: updatetime };
                                (*servers).push(serverness);
                            }
                        }
                        Some(k) => {
                            servers[k].last_alive = updatetime;
                            if packet.statechanged == GAMEEXITED {
                                servers.remove(k);
                            }
                        }
                    };
                }
            };
        }

        // Keepalive packet. We need to rely on the origin's port for this, unfortunately. This may mean that the source port is incorrect if the port was changed with NAT.
        else if buffer[OPCODE_INDEX] == KEEPALIVE {
            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>#[macro_use]
extern crate criterion;
#[macro_use]
extern crate serde_derive;

use std::collections::BTreeMap;
use std::fs::{create_dir_all, File};
use std::io::Write;
use std::path::Path;

use criterion::profiler::Profiler;
use criterion::Criterion;
use handlebars::{to_json, Context, Handlebars, Template};
use pprof::protos::Message;
use pprof::ProfilerGuard;
use serde_json::value::Value as Json;

#[derive(Default)]
struct CpuProfiler<'a> {
    guard: Option<ProfilerGuard<'a>>,
}

impl<'a> Profiler for CpuProfiler<'a> {
    fn start_profiling(&mut self, _benchmark_id: &str, benchmark_dir: &Path) {
        create_dir_all(&benchmark_dir).unwrap();

        let guard = ProfilerGuard::new(100).unwrap();
        self.guard = Some(guard);
    }

    fn stop_profiling(&mut self, benchmark_id: &str, benchmark_dir: &Path) {
        if let Ok(ref report) = self.guard.as_ref().unwrap().report().build() {
            let fg_file_name = benchmark_dir.join(format!("{}.svg", benchmark_id));
            let fg_file = File::create(fg_file_name).unwrap();
            report.flamegraph(fg_file).unwrap();

            let pb_file_name = benchmark_dir.join(format!("{}.pb", benchmark_id));
            let mut pb_file = File::create(pb_file_name).unwrap();
            let profile = report.pprof().unwrap();

            let mut content = Vec::new();
            profile.encode(&mut content).unwrap();
            pb_file.write_all(&content).unwrap();
        };

        self.guard = None;
    }
}

fn profiled() -> Criterion {
    Criterion::default().with_profiler(CpuProfiler::default())
}

#[derive(Serialize)]
struct DataWrapper {
    v: String,
}

#[derive(Serialize)]
struct RowWrapper {
    real: Vec<DataWrapper>,
    dummy: Vec<DataWrapper>,
}

#[derive(Serialize)]
struct NestedRowWrapper {
    parent: Vec<Vec<DataWrapper>>,
}

static SOURCE: &'static str = "

    <title>{{year}}</title>

CSL {{year}}

    {{#each teams}}

      {{name}}: {{score}}

    {{/each}}

";

fn make_data() -> BTreeMap<String, Json> {
    let mut data = BTreeMap::new();

    data.insert("year".to_string(), to_json("2015"));

    let mut teams = Vec::new();

    for v in vec![
        ("Jiangsu", 43u16),
        ("Beijing", 27u16),
        ("Guangzhou", 22u16),
        ("Shandong", 12u16),
    ]
    .iter()
    {
        let (name, score) = *v;
        let mut t = BTreeMap::new();
        t.insert("name".to_string(), to_json(name));
        t.insert("score".to_string(), to_json(score));
        teams.push(t)
    }

    data.insert("teams".to_string(), to_json(&teams));
    data
}

fn parse_template(c: &mut Criterion) {
    c.bench_function("parse_template", move |b| {
        b.iter(|| Template::compile(SOURCE).ok().unwrap())
    });
}

fn render_template(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string("table", SOURCE)
        .ok()
        .expect("Invalid template format");

    let ctx = Context::wraps(make_data()).unwrap();
    c.bench_function("render_template", move |b| {
        b.iter(|| handlebars.render_with_context("table", &ctx).ok().unwrap())
    });
}

fn large_loop_helper(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string("test", "BEFORE\n{{#each real}}{{this.v}}{{/each}}AFTER")
        .ok()
        .expect("Invalid template format");

    let real: Vec<DataWrapper> = (1..1000)
        .map(|i| DataWrapper {
            v: format!("n={}", i),
        })
        .collect();
    let dummy: Vec<DataWrapper> = (1..1000)
        .map(|i| DataWrapper {
            v: format!("n={}", i),
        })
        .collect();
    let rows = RowWrapper { real, dummy };

    let ctx = Context::wraps(&rows).unwrap();
    c.bench_function("large_loop_helper", move |b| {
        b.iter(|| handlebars.render_with_context("test", &ctx).ok().unwrap())
    });
}

fn large_nested_loop(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string(
            "test",
            "BEFORE\n{{#each parent as |child|}}{{#each child}}{{this.v}}{{/each}}{{/each}}AFTER",
        )
        .ok()
        .expect("Invalid template format");

    let parent: Vec<Vec<DataWrapper>> = (1..100)
        .map(|_| {
            (1..10)
                .map(|v| DataWrapper {
                    v: format!("v={}", v),
                })
                .collect()
        })
        .collect();

    let rows = NestedRowWrapper { parent };

    let ctx = Context::wraps(&rows).unwrap();
    c.bench_function("large_nested_loop", move |b| {
        b.iter(|| handlebars.render_with_context("test", &ctx).ok().unwrap())
    });
}

criterion_group!(
    name = benches;
    config = profiled();
    targets = parse_template, render_template, large_loop_helper, large_nested_loop
);
criterion_main!(benches);
<commit_msg>(test) added benchmark with context creation to compare<commit_after>#[macro_use]
extern crate criterion;
#[macro_use]
extern crate serde_derive;

use std::collections::BTreeMap;
use std::fs::{create_dir_all, File};
use std::io::Write;
use std::path::Path;

use criterion::profiler::Profiler;
use criterion::Criterion;
use handlebars::{to_json, Context, Handlebars, Template};
use pprof::protos::Message;
use pprof::ProfilerGuard;
use serde_json::value::Value as Json;

#[derive(Default)]
struct CpuProfiler<'a> {
    guard: Option<ProfilerGuard<'a>>,
}

impl<'a> Profiler for CpuProfiler<'a> {
    fn start_profiling(&mut self, _benchmark_id: &str, benchmark_dir: &Path) {
        create_dir_all(&benchmark_dir).unwrap();

        let guard = ProfilerGuard::new(100).unwrap();
        self.guard = Some(guard);
    }

    fn stop_profiling(&mut self, benchmark_id: &str, benchmark_dir: &Path) {
        if let Ok(ref report) = self.guard.as_ref().unwrap().report().build() {
            let fg_file_name = benchmark_dir.join(format!("{}.svg", benchmark_id));
            let fg_file = File::create(fg_file_name).unwrap();
            report.flamegraph(fg_file).unwrap();

            let pb_file_name = benchmark_dir.join(format!("{}.pb", benchmark_id));
            let mut pb_file = File::create(pb_file_name).unwrap();
            let profile = report.pprof().unwrap();

            let mut content = Vec::new();
            profile.encode(&mut content).unwrap();
            pb_file.write_all(&content).unwrap();
        };

        self.guard = None;
    }
}

fn profiled() -> Criterion {
    Criterion::default().with_profiler(CpuProfiler::default())
}

#[derive(Serialize)]
struct DataWrapper {
    v: String,
}

#[derive(Serialize)]
struct RowWrapper {
    real: Vec<DataWrapper>,
    dummy: Vec<DataWrapper>,
}

#[derive(Serialize)]
struct NestedRowWrapper {
    parent: Vec<Vec<DataWrapper>>,
}

static SOURCE: &'static str = "

    <title>{{year}}</title>

CSL {{year}}

    {{#each teams}}

      {{name}}: {{score}}

    {{/each}}

";

fn make_data() -> BTreeMap<String, Json> {
    let mut data = BTreeMap::new();

    data.insert("year".to_string(), to_json("2015"));

    let mut teams = Vec::new();

    for v in vec![
        ("Jiangsu", 43u16),
        ("Beijing", 27u16),
        ("Guangzhou", 22u16),
        ("Shandong", 12u16),
    ]
    .iter()
    {
        let (name, score) = *v;
        let mut t = BTreeMap::new();
        t.insert("name".to_string(), to_json(name));
        t.insert("score".to_string(), to_json(score));
        teams.push(t)
    }

    data.insert("teams".to_string(), to_json(&teams));
    data
}

fn parse_template(c: &mut Criterion) {
    c.bench_function("parse_template", move |b| {
        b.iter(|| Template::compile(SOURCE).ok().unwrap())
    });
}

fn render_template(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string("table", SOURCE)
        .ok()
        .expect("Invalid template format");

    let ctx = Context::wraps(make_data()).unwrap();
    c.bench_function("render_template", move |b| {
        b.iter(|| handlebars.render_with_context("table", &ctx).ok().unwrap())
    });
}

fn large_loop_helper(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string("test", "BEFORE\n{{#each real}}{{this.v}}{{/each}}AFTER")
        .ok()
        .expect("Invalid template format");

    let real: Vec<DataWrapper> = (1..1000)
        .map(|i| DataWrapper {
            v: format!("n={}", i),
        })
        .collect();
    let dummy: Vec<DataWrapper> = (1..1000)
        .map(|i| DataWrapper {
            v: format!("n={}", i),
        })
        .collect();
    let rows = RowWrapper { real, dummy };

    let ctx = Context::wraps(&rows).unwrap();
    c.bench_function("large_loop_helper", move |b| {
        b.iter(|| handlebars.render_with_context("test", &ctx).ok().unwrap())
    });
}

fn large_loop_helper_with_context_creation(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string("test", "BEFORE\n{{#each real}}{{this.v}}{{/each}}AFTER")
        .ok()
        .expect("Invalid template format");

    let real: Vec<DataWrapper> = (1..1000)
        .map(|i| DataWrapper {
            v: format!("n={}", i),
        })
        .collect();
    let dummy: Vec<DataWrapper> = (1..1000)
        .map(|i| DataWrapper {
            v: format!("n={}", i),
        })
        .collect();
    let rows = RowWrapper { real, dummy };

    c.bench_function("large_loop_helper_with_context_creation", move |b| {
        b.iter(|| handlebars.render("test", &rows).ok().unwrap())
    });
}

fn large_nested_loop(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string(
            "test",
            "BEFORE\n{{#each parent as |child|}}{{#each child}}{{this.v}}{{/each}}{{/each}}AFTER",
        )
        .ok()
        .expect("Invalid template format");

    let parent: Vec<Vec<DataWrapper>> = (1..100)
        .map(|_| {
            (1..10)
                .map(|v| DataWrapper {
                    v: format!("v={}", v),
                })
                .collect()
        })
        .collect();

    let rows = NestedRowWrapper { parent };

    let ctx = Context::wraps(&rows).unwrap();
    c.bench_function("large_nested_loop", move |b| {
        b.iter(|| handlebars.render_with_context("test", &ctx).ok().unwrap())
    });
}

criterion_group!(
    name = benches;
    config = profiled();
    targets = parse_template, render_template, large_loop_helper, large_loop_helper_with_context_creation,
              large_nested_loop
);
criterion_main!(benches);
<|endoftext|>


<commit_before>#![feature(test)]

extern crate ipc_channel;
extern crate test;

/// Allows doing multiple inner iterations per bench.iter() run.
///
/// This is mostly to amortise the overhead of spawning a thread in the benchmark
/// when sending larger messages (that might be fragmented).
///
/// Note that you need to compensate the displayed results
/// for the proportionally longer runs yourself,
/// as the benchmark framework doesn't know about the inner iterations...
const ITERATIONS: usize = 1;

mod platform {
    extern crate crossbeam;

    use ipc_channel::platform;
    use ITERATIONS;

    use std::sync::{mpsc, Mutex};
    use test;

    #[bench]
    fn create_channel(b: &mut test::Bencher) {
        b.iter(|| {
            for _ in 0..ITERATIONS {
                platform::channel().unwrap();
            }
        });
    }

    fn bench_transfer_data(b: &mut test::Bencher, size: usize) {
        let data: Vec<u8> = (0..size).map(|i| (i % 251) as u8).collect();
        let (tx, rx) = platform::channel().unwrap();

        let (wait_tx, wait_rx) = mpsc::channel();
        let wait_rx = Mutex::new(wait_rx);

        if size > platform::OsIpcSender::get_max_fragment_size() {
            b.iter(|| {
                crossbeam::scope(|scope| {
                    let tx = tx.clone();
                    scope.spawn(|| {
                        let wait_rx = wait_rx.lock().unwrap();
                        let tx = tx;
                        for _ in 0..ITERATIONS {
                            tx.send(&data, vec![], vec![]).unwrap();
                            if ITERATIONS > 1 {
                                // Prevent beginning of the next send
                                // from overlapping with receive of last fragment,
                                // as otherwise results of runs with a large tail fragment
                                // are significantly skewed.
                                wait_rx.recv().unwrap();
                            }
                        }
                    });
                    for _ in 0..ITERATIONS {
                        rx.recv().unwrap();
                        if ITERATIONS > 1 {
                            wait_tx.send(()).unwrap();
                        }
                    }
                    // For reasons mysterious to me,
                    // not returning a value *from every branch*
                    // adds some 100 ns or so of overhead to all results --
                    // which is quite significant for very short tests...
                    0
                })
            });
        } else {
            b.iter(|| {
                for _ in 0..ITERATIONS {
                    tx.send(&data, vec![], vec![]).unwrap();
                    rx.recv().unwrap();
                }
                0
            });
        }
    }

    #[bench]
    fn transfer_data_00_1(b: &mut test::Bencher) {
        bench_transfer_data(b, 1);
    }
    #[bench]
    fn transfer_data_01_2(b: &mut test::Bencher) {
        bench_transfer_data(b, 2);
    }
    #[bench]
    fn transfer_data_02_4(b: &mut test::Bencher) {
        bench_transfer_data(b, 4);
    }
    #[bench]
    fn transfer_data_03_8(b: &mut test::Bencher) {
        bench_transfer_data(b, 8);
    }
    #[bench]
    fn transfer_data_04_16(b: &mut test::Bencher) {
        bench_transfer_data(b, 16);
    }
    #[bench]
    fn transfer_data_05_32(b: &mut test::Bencher) {
        bench_transfer_data(b, 32);
    }
    #[bench]
    fn transfer_data_06_64(b: &mut test::Bencher) {
        bench_transfer_data(b, 64);
    }
    #[bench]
    fn transfer_data_07_128(b: &mut test::Bencher) {
        bench_transfer_data(b, 128);
    }
    #[bench]
    fn transfer_data_08_256(b: &mut test::Bencher) {
        bench_transfer_data(b, 256);
    }
    #[bench]
    fn transfer_data_09_512(b: &mut test::Bencher) {
        bench_transfer_data(b, 512);
    }
    #[bench]
    fn transfer_data_10_1k(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024);
    }
    #[bench]
    fn transfer_data_11_2k(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024);
    }
    #[bench]
    fn transfer_data_12_4k(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024);
    }
    #[bench]
    fn transfer_data_13_8k(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024);
    }
    #[bench]
    fn transfer_data_14_16k(b: &mut test::Bencher) {
        bench_transfer_data(b, 16 * 1024);
    }
    #[bench]
    fn transfer_data_15_32k(b: &mut test::Bencher) {
        bench_transfer_data(b, 32 * 1024);
    }
    #[bench]
    fn transfer_data_16_64k(b: &mut test::Bencher) {
        bench_transfer_data(b, 64 * 1024);
    }
    #[bench]
    fn transfer_data_17_128k(b: &mut test::Bencher) {
        bench_transfer_data(b, 128 * 1024);
    }
    #[bench]
    fn transfer_data_18_256k(b: &mut test::Bencher) {
        bench_transfer_data(b, 256 * 1024);
    }
    #[bench]
    fn transfer_data_19_512k(b: &mut test::Bencher) {
        bench_transfer_data(b, 512 * 1024);
    }
    #[bench]
    fn transfer_data_20_1m(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_21_2m(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_22_4m(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_23_8m(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024 * 1024);
    }
}

mod ipc {
    mod receiver_set {
        use ipc_channel::ipc::{self, IpcReceiverSet};
        use test;

        fn gen_select_test(b: &mut test::Bencher, to_send: usize, n: usize) -> () {
            let mut active = Vec::with_capacity(to_send);
            let mut dormant = Vec::with_capacity(n - to_send);
            let mut rx_set = IpcReceiverSet::new().unwrap();
            for _ in 0..to_send {
                let (tx, rx) = ipc::channel().unwrap();
                rx_set.add(rx).unwrap();
                active.push(tx);
            }
            for _ in to_send..n {
                let (tx, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
                dormant.push(tx);
            }
            b.iter(|| {
                for tx in active.iter() {
                    tx.send(()).unwrap();
                }
                let mut received = 0;
                while received < to_send {
                    for result in rx_set.select().unwrap().into_iter() {
                        let (_, _) = result.unwrap();
                        received += 1;
                    }
                }
            });
        }

        fn add_n_rxs(rx_set: &mut IpcReceiverSet, n: usize) -> () {
            for _ in 0..n {
                let (_, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
            }
        }

        #[bench]
        fn send_on_1_of_1(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 1);
        }

        #[bench]
        fn send_on_1_of_5(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 5);
        }

        #[bench]
        fn send_on_2_of_5(b: &mut test::Bencher) -> () {
            gen_select_test(b, 2, 5);
        }

        #[bench]
        fn send_on_5_of_5(b: &mut test::Bencher) -> () {
            gen_select_test(b, 5, 5);
        }

        #[bench]
        fn send_on_1_of_20(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 20);
        }

        #[bench]
        fn send_on_5_of_20(b: &mut test::Bencher) -> () {
            gen_select_test(b, 5, 20);
        }

        #[bench]
        fn send_on_20_of_20(b: &mut test::Bencher) -> () {
            gen_select_test(b, 20, 20);
        }

        #[bench]
        fn send_on_1_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 100);
        }

        #[bench]
        fn send_on_5_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 5, 100);
        }
        #[bench]
        fn send_on_20_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 20, 100);
        }

        #[bench]
        fn send_on_100_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 100, 100);
        }

        #[bench]
        fn create_and_destroy_empty_set(b: &mut test::Bencher) -> () {
            b.iter(|| {
                IpcReceiverSet::new().unwrap();
            });
        }

        #[bench]
        fn create_and_destroy_set_of_10(b: &mut test::Bencher) -> () {
            b.iter(|| {
                let mut rx_set = IpcReceiverSet::new().unwrap();
                add_n_rxs(&mut rx_set, 10);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_5(b: &mut test::Bencher) -> () {
            b.iter(|| {
                let mut rx_set = IpcReceiverSet::new().unwrap();
                add_n_rxs(&mut rx_set, 5);
            });
        }

        #[bench]
        // Benchmark adding and removing closed receivers from the set
        fn add_and_remove_closed_receivers(b: &mut test::Bencher) -> () {
            b.iter(|| {
                let mut rx_set = IpcReceiverSet::new().unwrap();
                {
                    {
                        let (_, rx) = ipc::channel::<()>().unwrap();
                        rx_set.add(rx).unwrap();
                    }
                    // On select Receivers with a "ClosedChannel" event
                    // will be closed
                    rx_set.select().unwrap();
                    let (_, rx) = ipc::channel::<()>().unwrap();
                    rx_set.add(rx).unwrap();
                }
            });
        }
    }
}
<commit_msg>Benchmarks: Use more comprehensive set of data points for creating receiver sets<commit_after>#![feature(test)]

extern crate ipc_channel;
extern crate test;

/// Allows doing multiple inner iterations per bench.iter() run.
///
/// This is mostly to amortise the overhead of spawning a thread in the benchmark
/// when sending larger messages (that might be fragmented).
///
/// Note that you need to compensate the displayed results
/// for the proportionally longer runs yourself,
/// as the benchmark framework doesn't know about the inner iterations...
const ITERATIONS: usize = 1;

mod platform {
    extern crate crossbeam;

    use ipc_channel::platform;
    use ITERATIONS;

    use std::sync::{mpsc, Mutex};
    use test;

    #[bench]
    fn create_channel(b: &mut test::Bencher) {
        b.iter(|| {
            for _ in 0..ITERATIONS {
                platform::channel().unwrap();
            }
        });
    }

    fn bench_transfer_data(b: &mut test::Bencher, size: usize) {
        let data: Vec<u8> = (0..size).map(|i| (i % 251) as u8).collect();
        let (tx, rx) = platform::channel().unwrap();

        let (wait_tx, wait_rx) = mpsc::channel();
        let wait_rx = Mutex::new(wait_rx);

        if size > platform::OsIpcSender::get_max_fragment_size() {
            b.iter(|| {
                crossbeam::scope(|scope| {
                    let tx = tx.clone();
                    scope.spawn(|| {
                        let wait_rx = wait_rx.lock().unwrap();
                        let tx = tx;
                        for _ in 0..ITERATIONS {
                            tx.send(&data, vec![], vec![]).unwrap();
                            if ITERATIONS > 1 {
                                // Prevent beginning of the next send
                                // from overlapping with receive of last fragment,
                                // as otherwise results of runs with a large tail fragment
                                // are significantly skewed.
                                wait_rx.recv().unwrap();
                            }
                        }
                    });
                    for _ in 0..ITERATIONS {
                        rx.recv().unwrap();
                        if ITERATIONS > 1 {
                            wait_tx.send(()).unwrap();
                        }
                    }
                    // For reasons mysterious to me,
                    // not returning a value *from every branch*
                    // adds some 100 ns or so of overhead to all results --
                    // which is quite significant for very short tests...
                    0
                })
            });
        } else {
            b.iter(|| {
                for _ in 0..ITERATIONS {
                    tx.send(&data, vec![], vec![]).unwrap();
                    rx.recv().unwrap();
                }
                0
            });
        }
    }

    #[bench]
    fn transfer_data_00_1(b: &mut test::Bencher) {
        bench_transfer_data(b, 1);
    }
    #[bench]
    fn transfer_data_01_2(b: &mut test::Bencher) {
        bench_transfer_data(b, 2);
    }
    #[bench]
    fn transfer_data_02_4(b: &mut test::Bencher) {
        bench_transfer_data(b, 4);
    }
    #[bench]
    fn transfer_data_03_8(b: &mut test::Bencher) {
        bench_transfer_data(b, 8);
    }
    #[bench]
    fn transfer_data_04_16(b: &mut test::Bencher) {
        bench_transfer_data(b, 16);
    }
    #[bench]
    fn transfer_data_05_32(b: &mut test::Bencher) {
        bench_transfer_data(b, 32);
    }
    #[bench]
    fn transfer_data_06_64(b: &mut test::Bencher) {
        bench_transfer_data(b, 64);
    }
    #[bench]
    fn transfer_data_07_128(b: &mut test::Bencher) {
        bench_transfer_data(b, 128);
    }
    #[bench]
    fn transfer_data_08_256(b: &mut test::Bencher) {
        bench_transfer_data(b, 256);
    }
    #[bench]
    fn transfer_data_09_512(b: &mut test::Bencher) {
        bench_transfer_data(b, 512);
    }
    #[bench]
    fn transfer_data_10_1k(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024);
    }
    #[bench]
    fn transfer_data_11_2k(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024);
    }
    #[bench]
    fn transfer_data_12_4k(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024);
    }
    #[bench]
    fn transfer_data_13_8k(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024);
    }
    #[bench]
    fn transfer_data_14_16k(b: &mut test::Bencher) {
        bench_transfer_data(b, 16 * 1024);
    }
    #[bench]
    fn transfer_data_15_32k(b: &mut test::Bencher) {
        bench_transfer_data(b, 32 * 1024);
    }
    #[bench]
    fn transfer_data_16_64k(b: &mut test::Bencher) {
        bench_transfer_data(b, 64 * 1024);
    }
    #[bench]
    fn transfer_data_17_128k(b: &mut test::Bencher) {
        bench_transfer_data(b, 128 * 1024);
    }
    #[bench]
    fn transfer_data_18_256k(b: &mut test::Bencher) {
        bench_transfer_data(b, 256 * 1024);
    }
    #[bench]
    fn transfer_data_19_512k(b: &mut test::Bencher) {
        bench_transfer_data(b, 512 * 1024);
    }
    #[bench]
    fn transfer_data_20_1m(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_21_2m(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_22_4m(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_23_8m(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024 * 1024);
    }
}

mod ipc {
    mod receiver_set {
        use ipc_channel::ipc::{self, IpcReceiverSet};
        use test;

        fn gen_select_test(b: &mut test::Bencher, to_send: usize, n: usize) -> () {
            let mut active = Vec::with_capacity(to_send);
            let mut dormant = Vec::with_capacity(n - to_send);
            let mut rx_set = IpcReceiverSet::new().unwrap();
            for _ in 0..to_send {
                let (tx, rx) = ipc::channel().unwrap();
                rx_set.add(rx).unwrap();
                active.push(tx);
            }
            for _ in to_send..n {
                let (tx, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
                dormant.push(tx);
            }
            b.iter(|| {
                for tx in active.iter() {
                    tx.send(()).unwrap();
                }
                let mut received = 0;
                while received < to_send {
                    for result in rx_set.select().unwrap().into_iter() {
                        let (_, _) = result.unwrap();
                        received += 1;
                    }
                }
            });
        }

        fn add_n_rxs(rx_set: &mut IpcReceiverSet, n: usize) -> () {
            for _ in 0..n {
                let (_, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
            }
        }

        #[bench]
        fn send_on_1_of_1(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 1);
        }

        #[bench]
        fn send_on_1_of_5(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 5);
        }

        #[bench]
        fn send_on_2_of_5(b: &mut test::Bencher) -> () {
            gen_select_test(b, 2, 5);
        }

        #[bench]
        fn send_on_5_of_5(b: &mut test::Bencher) -> () {
            gen_select_test(b, 5, 5);
        }

        #[bench]
        fn send_on_1_of_20(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 20);
        }

        #[bench]
        fn send_on_5_of_20(b: &mut test::Bencher) -> () {
            gen_select_test(b, 5, 20);
        }

        #[bench]
        fn send_on_20_of_20(b: &mut test::Bencher) -> () {
            gen_select_test(b, 20, 20);
        }

        #[bench]
        fn send_on_1_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 1, 100);
        }

        #[bench]
        fn send_on_5_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 5, 100);
        }
        #[bench]
        fn send_on_20_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 20, 100);
        }

        #[bench]
        fn send_on_100_of_100(b: &mut test::Bencher) -> () {
            gen_select_test(b, 100, 100);
        }

        #[bench]
        fn create_and_destroy_empty_set(b: &mut test::Bencher) -> () {
            b.iter(|| {
                IpcReceiverSet::new().unwrap();
            });
        }

        #[bench]
        fn create_and_destroy_set_of_1(b: &mut test::Bencher) -> () {
            b.iter(|| {
                let mut rx_set = IpcReceiverSet::new().unwrap();
                add_n_rxs(&mut rx_set, 1);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_10(b: &mut test::Bencher) -> () {
            b.iter(|| {
                let mut rx_set = IpcReceiverSet::new().unwrap();
                add_n_rxs(&mut rx_set, 10);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_100(b: &mut test::Bencher) -> () {
            b.iter(|| {
                let mut rx_set = IpcReceiverSet::new().unwrap();
                add_n_rxs(&mut rx_set, 100);
            });
        }

        #[bench]
        // Benchmark adding and removing closed receivers from the set
        fn add_and_remove_closed_receivers(b: &mut test::Bencher) -> () {
            b.iter(|| {
                let mut rx_set = IpcReceiverSet::new().unwrap();
                {
                    {
                        let (_, rx) = ipc::channel::<()>().unwrap();
                        rx_set.add(rx).unwrap();
                    }
                    // On select Receivers with a "ClosedChannel" event
                    // will be closed
                    rx_set.select().unwrap();
                    let (_, rx) = ipc::channel::<()>().unwrap();
                    rx_set.add(rx).unwrap();
                }
            });
        }
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! States elements can be in.

#![deny(missing_docs)]

bitflags! {
    #[doc = "Event-based element states."]
    #[cfg_attr(feature = "servo", derive(HeapSizeOf))]
    pub flags ElementState: u32 {
        #[doc = "The mouse is down on this element. \
                 https://html.spec.whatwg.org/multipage/#selector-active \
                 FIXME(#7333): set/unset this when appropriate"]
        const IN_ACTIVE_STATE = 0x01,
        #[doc = "This element has focus. \
                 https://html.spec.whatwg.org/multipage/#selector-focus"]
        const IN_FOCUS_STATE = 0x02,
        #[doc = "The mouse is hovering over this element. \
                 https://html.spec.whatwg.org/multipage/#selector-hover"]
        const IN_HOVER_STATE = 0x04,
        #[doc = "Content is enabled (and can be disabled). \
                 http://www.whatwg.org/html/#selector-enabled"]
        const IN_ENABLED_STATE = 0x08,
        #[doc = "Content is disabled. \
                 http://www.whatwg.org/html/#selector-disabled"]
        const IN_DISABLED_STATE = 0x10,
        #[doc = "Content is checked. \
                 https://html.spec.whatwg.org/multipage/#selector-checked"]
        const IN_CHECKED_STATE = 0x20,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-indeterminate"]
        const IN_INDETERMINATE_STATE = 0x40,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-placeholder-shown"]
        const IN_PLACEHOLDER_SHOWN_STATE = 0x80,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-target"]
        const IN_TARGET_STATE = 0x100,
        #[doc = "https://fullscreen.spec.whatwg.org/#%3Afullscreen-pseudo-class"]
        const IN_FULLSCREEN_STATE = 0x200,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-valid"]
        const IN_VALID_STATE = 0x400,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-invalid"]
        const IN_INVALID_STATE = 0x800,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-ui-valid"]
        const IN_MOZ_UI_VALID_STATE = 0x1000,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-broken"]
        const IN_BROKEN_STATE = 0x2000,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-user-disabled"]
        const IN_USER_DISABLED_STATE = 0x4000,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-suppressed"]
        const IN_SUPPRESSED_STATE = 0x8000,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-loading"]
        const IN_LOADING_STATE = 0x10000,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-handler-blocked"]
        const IN_HANDLER_BLOCKED_STATE = 0x20000,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-handler-disabled"]
        const IN_HANDLER_DISABLED_STATE = 0x40000,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-handler-crashed"]
        const IN_HANDLER_CRASHED_STATE = 0x80000,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-read-write"]
        const IN_READ_WRITE_STATE = 0x100000,
    }
}
<commit_msg>Use less fragile and more readable numbering for ElementState.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! States elements can be in.

#![deny(missing_docs)]

bitflags! {
    #[doc = "Event-based element states."]
    #[cfg_attr(feature = "servo", derive(HeapSizeOf))]
    pub flags ElementState: u32 {
        #[doc = "The mouse is down on this element. \
                 https://html.spec.whatwg.org/multipage/#selector-active \
                 FIXME(#7333): set/unset this when appropriate"]
        const IN_ACTIVE_STATE = 1 << 0,
        #[doc = "This element has focus. \
                 https://html.spec.whatwg.org/multipage/#selector-focus"]
        const IN_FOCUS_STATE = 1 << 1,
        #[doc = "The mouse is hovering over this element. \
                 https://html.spec.whatwg.org/multipage/#selector-hover"]
        const IN_HOVER_STATE = 1 << 2,
        #[doc = "Content is enabled (and can be disabled). \
                 http://www.whatwg.org/html/#selector-enabled"]
        const IN_ENABLED_STATE = 1 << 3,
        #[doc = "Content is disabled. \
                 http://www.whatwg.org/html/#selector-disabled"]
        const IN_DISABLED_STATE = 1 << 4,
        #[doc = "Content is checked. \
                 https://html.spec.whatwg.org/multipage/#selector-checked"]
        const IN_CHECKED_STATE = 1 << 5,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-indeterminate"]
        const IN_INDETERMINATE_STATE = 1 << 6,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-placeholder-shown"]
        const IN_PLACEHOLDER_SHOWN_STATE = 1 << 7,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-target"]
        const IN_TARGET_STATE = 1 << 8,
        #[doc = "https://fullscreen.spec.whatwg.org/#%3Afullscreen-pseudo-class"]
        const IN_FULLSCREEN_STATE = 1 << 9,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-valid"]
        const IN_VALID_STATE = 1 << 10,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-invalid"]
        const IN_INVALID_STATE = 1 << 11,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-ui-valid"]
        const IN_MOZ_UI_VALID_STATE = 1 << 12,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-broken"]
        const IN_BROKEN_STATE = 1 << 13,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-user-disabled"]
        const IN_USER_DISABLED_STATE = 1 << 14,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-suppressed"]
        const IN_SUPPRESSED_STATE = 1 << 15,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-loading"]
        const IN_LOADING_STATE = 1 << 16,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-handler-blocked"]
        const IN_HANDLER_BLOCKED_STATE = 1 << 17,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-handler-disabled"]
        const IN_HANDLER_DISABLED_STATE = 1 << 18,
        #[doc = "Non-standard: https://developer.mozilla.org/en-US/docs/Web/CSS/:-moz-handler-crashed"]
        const IN_HANDLER_CRASHED_STATE = 1 << 19,
        #[doc = "https://html.spec.whatwg.org/multipage/#selector-read-write"]
        const IN_READ_WRITE_STATE = 1 << 20,
    }
}
<|endoftext|>


<commit_before>// vim: tw=80
use super::*;

/// Performs transformations on a function to make it mockable
fn mockable_fn(mut item_fn: ItemFn) -> ItemFn {
    demutify(&mut item_fn.sig.inputs);
    deimplify(&mut item_fn.sig.output);
    item_fn
}

/// Performs transformations on an Item to make it mockable
fn mockable_item(item: Item) -> Item {
    match item {
        Item::Fn(item_fn) => Item::Fn(mockable_fn(item_fn)),
        x => x
    }
}

/// An item that's ready to be mocked.
///
/// It should be functionally identical or near-identical to the original item,
/// but with minor alterations that make it suitable for mocking, such as
/// altered lifetimes.
pub(crate) enum MockableItem {
    Module(MockableModule),
    Struct(MockableStruct)
}

impl From<(Attrs, Item)> for MockableItem {
    fn from((attrs, item): (Attrs, Item)) -> MockableItem {
        match item {
            Item::Impl(item_impl) =>
                MockableItem::Struct(MockableStruct::from(item_impl)),
            Item::ForeignMod(item_foreign_mod) =>
                MockableItem::Module(
                    MockableModule::from((attrs, item_foreign_mod))
                ),
            Item::Mod(item_mod) =>
                MockableItem::Module(MockableModule::from(item_mod)),
            Item::Trait(trait_) =>
                MockableItem::Struct(MockableStruct::from((attrs, trait_))),
            _ => panic!("automock does not support this item type")
        }
    }
}

impl From<MockableStruct> for MockableItem {
    fn from(mock: MockableStruct) -> MockableItem {
        MockableItem::Struct(mock)
    }
}

pub(crate) struct MockableModule {
    pub attrs: TokenStream,
    pub vis: Visibility,
    pub mock_ident: Ident,
    /// Ident of the original module, if any
    pub orig_ident: Option<Ident>,
    pub content: Vec<Item>
}

impl From<(Attrs, ItemForeignMod)> for MockableModule {
    fn from((attrs, foreign): (Attrs, ItemForeignMod)) -> MockableModule {
        let orig_ident = None;
        let mock_ident = attrs.modname.expect(concat!(
            "module name is required when mocking foreign functions,",
            " like `#[automock(mod mock_ffi)]`"
        ));
        let vis = Visibility::Public(VisPublic{
            pub_token: <Token![pub]>::default()
        });
        let attrs = quote!(
            #[deprecated(since = "0.9.0", note = "Using automock directly on an extern block is deprecated.  Instead, wrap the extern block in a module, and automock that, like #[automoock] mod ffi { extern \"C\" { fn foo ... } }")]
        );
        let mut content = vec![
            // When mocking extern blocks, we pretend that they're modules, so
            // we need a "use super::*;" to ensure that types can resolve
            Item::Use(ItemUse {
                attrs: Vec::new(),
                vis: Visibility::Inherited,
                use_token: token::Use::default(),
                leading_colon: None,
                tree: UseTree::Path(UsePath {
                    ident: Ident::new("super", Span::call_site()),
                    colon2_token: token::Colon2::default(),
                    tree: Box::new(UseTree::Glob(UseGlob {
                        star_token: token::Star::default()
                    }))
                }),
                semi_token: token::Semi::default()
            })
        ];
        content.extend(foreign.items.into_iter()
            .map(|foreign_item| {
                match foreign_item {
                    ForeignItem::Fn(f) => {
                        let span = f.sig.span();
                        let mut sig = f.sig;

                        // When mocking extern blocks, we pretend that they're
                        // modules.  So we must supersuperfy everything by one
                        // level.
                        let vis = expectation_visibility(&f.vis, 1);

                        for arg in sig.inputs.iter_mut() {
                            if let FnArg::Typed(pt) = arg {
                                *pt.ty = supersuperfy(&*pt.ty, 1);
                            }
                        }
                        if let ReturnType::Type(_, ty) = &mut sig.output {
                                **ty = supersuperfy(&*ty, 1);
                        }

                        // Foreign functions are always unsafe.  Mock foreign
                        // functions should be unsafe too, to prevent "warning:
                        // unused unsafe" messages.
                        sig.unsafety = Some(Token![unsafe](span));
                        let block = Box::new(Block {
                            brace_token: token::Brace::default(),
                            stmts: Vec::new()
                        });

                        Item::Fn(
                            ItemFn {
                                attrs: f.attrs,
                                vis,
                                sig,
                                block
                            }
                        )
                    },
                    _ => {
                        compile_error(foreign_item.span(),
                            "Unsupported foreign item type"
                        );
                        Item::Verbatim(TokenStream::default())
                    }
                }
            }));
        MockableModule {
            attrs,
            vis,
            mock_ident,
            orig_ident,
            content
        }
    }
}

impl From<ItemMod> for MockableModule {
    fn from(mod_: ItemMod) -> MockableModule {
        let span = mod_.span();
        let vis = mod_.vis;
        let mock_ident = format_ident!("mock_{}", mod_.ident);
        let orig_ident = Some(mod_.ident);
        let content = if let Some((_, content)) = mod_.content {
            content.into_iter()
            .map(mockable_item)
            .collect()
        } else {
            compile_error(span,
            "automock can only mock inline modules, not modules from another file");
            Vec::new()
        };
        MockableModule {
            attrs: TokenStream::new(),
            vis,
            mock_ident,
            orig_ident,
            content
        }
    }
}
<commit_msg>fix a typo in an error message<commit_after>// vim: tw=80
use super::*;

/// Performs transformations on a function to make it mockable
fn mockable_fn(mut item_fn: ItemFn) -> ItemFn {
    demutify(&mut item_fn.sig.inputs);
    deimplify(&mut item_fn.sig.output);
    item_fn
}

/// Performs transformations on an Item to make it mockable
fn mockable_item(item: Item) -> Item {
    match item {
        Item::Fn(item_fn) => Item::Fn(mockable_fn(item_fn)),
        x => x
    }
}

/// An item that's ready to be mocked.
///
/// It should be functionally identical or near-identical to the original item,
/// but with minor alterations that make it suitable for mocking, such as
/// altered lifetimes.
pub(crate) enum MockableItem {
    Module(MockableModule),
    Struct(MockableStruct)
}

impl From<(Attrs, Item)> for MockableItem {
    fn from((attrs, item): (Attrs, Item)) -> MockableItem {
        match item {
            Item::Impl(item_impl) =>
                MockableItem::Struct(MockableStruct::from(item_impl)),
            Item::ForeignMod(item_foreign_mod) =>
                MockableItem::Module(
                    MockableModule::from((attrs, item_foreign_mod))
                ),
            Item::Mod(item_mod) =>
                MockableItem::Module(MockableModule::from(item_mod)),
            Item::Trait(trait_) =>
                MockableItem::Struct(MockableStruct::from((attrs, trait_))),
            _ => panic!("automock does not support this item type")
        }
    }
}

impl From<MockableStruct> for MockableItem {
    fn from(mock: MockableStruct) -> MockableItem {
        MockableItem::Struct(mock)
    }
}

pub(crate) struct MockableModule {
    pub attrs: TokenStream,
    pub vis: Visibility,
    pub mock_ident: Ident,
    /// Ident of the original module, if any
    pub orig_ident: Option<Ident>,
    pub content: Vec<Item>
}

impl From<(Attrs, ItemForeignMod)> for MockableModule {
    fn from((attrs, foreign): (Attrs, ItemForeignMod)) -> MockableModule {
        let orig_ident = None;
        let mock_ident = attrs.modname.expect(concat!(
            "module name is required when mocking foreign functions,",
            " like `#[automock(mod mock_ffi)]`"
        ));
        let vis = Visibility::Public(VisPublic{
            pub_token: <Token![pub]>::default()
        });
        let attrs = quote!(
            #[deprecated(since = "0.9.0", note = "Using automock directly on an extern block is deprecated.  Instead, wrap the extern block in a module, and automock that, like #[automock] mod ffi { extern \"C\" { fn foo ... } }")]
        );
        let mut content = vec![
            // When mocking extern blocks, we pretend that they're modules, so
            // we need a "use super::*;" to ensure that types can resolve
            Item::Use(ItemUse {
                attrs: Vec::new(),
                vis: Visibility::Inherited,
                use_token: token::Use::default(),
                leading_colon: None,
                tree: UseTree::Path(UsePath {
                    ident: Ident::new("super", Span::call_site()),
                    colon2_token: token::Colon2::default(),
                    tree: Box::new(UseTree::Glob(UseGlob {
                        star_token: token::Star::default()
                    }))
                }),
                semi_token: token::Semi::default()
            })
        ];
        content.extend(foreign.items.into_iter()
            .map(|foreign_item| {
                match foreign_item {
                    ForeignItem::Fn(f) => {
                        let span = f.sig.span();
                        let mut sig = f.sig;

                        // When mocking extern blocks, we pretend that they're
                        // modules.  So we must supersuperfy everything by one
                        // level.
                        let vis = expectation_visibility(&f.vis, 1);

                        for arg in sig.inputs.iter_mut() {
                            if let FnArg::Typed(pt) = arg {
                                *pt.ty = supersuperfy(&*pt.ty, 1);
                            }
                        }
                        if let ReturnType::Type(_, ty) = &mut sig.output {
                                **ty = supersuperfy(&*ty, 1);
                        }

                        // Foreign functions are always unsafe.  Mock foreign
                        // functions should be unsafe too, to prevent "warning:
                        // unused unsafe" messages.
                        sig.unsafety = Some(Token![unsafe](span));
                        let block = Box::new(Block {
                            brace_token: token::Brace::default(),
                            stmts: Vec::new()
                        });

                        Item::Fn(
                            ItemFn {
                                attrs: f.attrs,
                                vis,
                                sig,
                                block
                            }
                        )
                    },
                    _ => {
                        compile_error(foreign_item.span(),
                            "Unsupported foreign item type"
                        );
                        Item::Verbatim(TokenStream::default())
                    }
                }
            }));
        MockableModule {
            attrs,
            vis,
            mock_ident,
            orig_ident,
            content
        }
    }
}

impl From<ItemMod> for MockableModule {
    fn from(mod_: ItemMod) -> MockableModule {
        let span = mod_.span();
        let vis = mod_.vis;
        let mock_ident = format_ident!("mock_{}", mod_.ident);
        let orig_ident = Some(mod_.ident);
        let content = if let Some((_, content)) = mod_.content {
            content.into_iter()
            .map(mockable_item)
            .collect()
        } else {
            compile_error(span,
            "automock can only mock inline modules, not modules from another file");
            Vec::new()
        };
        MockableModule {
            attrs: TokenStream::new(),
            vis,
            mock_ident,
            orig_ident,
            content
        }
    }
}
<|endoftext|>


<commit_before>// Test the learning code.

extern crate rand;
extern crate tempdir;
extern crate timelearn;

use rand::{Rng, SeedableRng, XorShiftRng};
use tempdir::TempDir;
use timelearn::{now, Populator, Store, Result};
use std::io::Write;

#[test]
fn learning() {
    let tmp_dir = TempDir::new("learn").unwrap();
    let db_path = tmp_dir.path().join("learn.db");

    {
        let mut st = Store::create(&db_path).unwrap();

        populate(&mut st).unwrap();

        // Ask all of the problems that we can.
        let mut num = 1;
        loop {
            let prob = match st.get_next().unwrap() {
                None => break,
                Some(p) => p,
            };

            println!("Ask {} ({})", prob.question, num);
            st.update(prob, num).unwrap();
            num = num % 4 + 1;
        }

        // Then ask all of the new problems.
        loop {
            let prob = match st.get_new().unwrap() {
                None => break,
                Some(p) => p,
            };

            println!("New {} ({})", prob.question, num);
            st.update(prob, num).unwrap();
            num = num % 4 + 1;
        }

        // The challenge here then is figuring out if this result is meaningful.
    }

    // Close and open to make sure that works.
    Store::open(&db_path).unwrap();
}

// Populate with test data.  Make a mix of unlearned problems, and learned ones that are ready to
// learn, and those that aren't ready to learn.
fn populate(st: &mut Store) -> Result<()> {
    let mut p = st.populate()?;
    let cur = now();

    // These are in the past, ready to ask.
    for i in 1..11 {
        add_one(&mut p, i, Some((cur - 20.0, 5.0)))?;
    }

    // These are in the future.
    for i in 11..21 {
        add_one(&mut p, i, Some((cur + 20.0, 5.0)))?;
    }

    // And these haven't been asked at all.
    for i in 21..31 {
        add_one(&mut p, i, None)?;
    }

    p.commit()?;
    Ok(())
}

// Add a problem of the given number.  Next and interval can optionally be set.
fn add_one<'a>(pop: &mut Populator<'a>, num: i32, ni: Option<(f64, f64)>) -> Result<()> {
    let mut rng: XorShiftRng = SeedableRng::from_seed([num as u32, 1, 2, 3]);

    let mut qn = vec![];
    write!(&mut qn, "qn{}:", num).unwrap();
    for _ in 1..rng.gen_range(5, 40) {
        qn.push(rng.gen_range(b'a', b'z'));
    }

    let mut an = vec![];
    write!(&mut an, "an{}:", num).unwrap();
    for _ in 1..rng.gen_range(5, 40) {
        an.push(rng.gen_range(b'a', b'z'));
    }

    let qn = String::from_utf8(qn).unwrap();
    let an = String::from_utf8(an).unwrap();
    match ni {
        None => pop.add_problem(&qn, &an)?,
        Some((next, interval)) => pop.add_learning_problem(&qn, &an, next, interval)?,
    }

    Ok(())
}
<commit_msg>timelearn: Test kind functionality<commit_after>// Test the learning code.

extern crate rand;
extern crate tempdir;
extern crate timelearn;

use rand::{Rng, SeedableRng, XorShiftRng};
use tempdir::TempDir;
use timelearn::{now, Populator, Store, Result};
use std::io::Write;

#[test]
fn learning() {
    let tmp_dir = TempDir::new("learn").unwrap();
    let db_path = tmp_dir.path().join("learn.db");

    {
        let mut st = Store::create(&db_path, "test").unwrap();
        assert_eq!(st.get_kind(), "test");

        populate(&mut st).unwrap();

        // Ask all of the problems that we can.
        let mut num = 1;
        loop {
            let prob = match st.get_next().unwrap() {
                None => break,
                Some(p) => p,
            };

            println!("Ask {} ({})", prob.question, num);
            st.update(prob, num).unwrap();
            num = num % 4 + 1;
        }

        // Then ask all of the new problems.
        loop {
            let prob = match st.get_new().unwrap() {
                None => break,
                Some(p) => p,
            };

            println!("New {} ({})", prob.question, num);
            st.update(prob, num).unwrap();
            num = num % 4 + 1;
        }

        // The challenge here then is figuring out if this result is meaningful.
    }

    // Close and open to make sure that works.
    let st = Store::open(&db_path).unwrap();
    assert_eq!(st.get_kind(), "test");
}

// Populate with test data.  Make a mix of unlearned problems, and learned ones that are ready to
// learn, and those that aren't ready to learn.
fn populate(st: &mut Store) -> Result<()> {
    let mut p = st.populate()?;
    let cur = now();

    // These are in the past, ready to ask.
    for i in 1..11 {
        add_one(&mut p, i, Some((cur - 20.0, 5.0)))?;
    }

    // These are in the future.
    for i in 11..21 {
        add_one(&mut p, i, Some((cur + 20.0, 5.0)))?;
    }

    // And these haven't been asked at all.
    for i in 21..31 {
        add_one(&mut p, i, None)?;
    }

    p.commit()?;
    Ok(())
}

// Add a problem of the given number.  Next and interval can optionally be set.
fn add_one<'a>(pop: &mut Populator<'a>, num: i32, ni: Option<(f64, f64)>) -> Result<()> {
    let mut rng: XorShiftRng = SeedableRng::from_seed([num as u32, 1, 2, 3]);

    let mut qn = vec![];
    write!(&mut qn, "qn{}:", num).unwrap();
    for _ in 1..rng.gen_range(5, 40) {
        qn.push(rng.gen_range(b'a', b'z'));
    }

    let mut an = vec![];
    write!(&mut an, "an{}:", num).unwrap();
    for _ in 1..rng.gen_range(5, 40) {
        an.push(rng.gen_range(b'a', b'z'));
    }

    let qn = String::from_utf8(qn).unwrap();
    let an = String::from_utf8(an).unwrap();
    match ni {
        None => pop.add_problem(&qn, &an)?,
        Some((next, interval)) => pop.add_learning_problem(&qn, &an, next, interval)?,
    }

    Ok(())
}
<|endoftext|>


<commit_before>use std::ops::Deref;
use rocket::http::Status;
use rocket::{Request, State, Outcome};
use rocket::request::{self, FromRequest};
use diesel::pg::PgConnection;
use r2d2;
use r2d2_diesel::ConnectionManager;
use std::env;

pub type Pool = r2d2::Pool<ConnectionManager<PgConnection>>;

pub fn init_db_pool() -> Pool {
    let config = r2d2::Config::default();
    let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
    let manager = ConnectionManager::<PgConnection>::new(database_url);
    r2d2::Pool::new(config, manager).expect("Failed to create pool.")
}

pub struct DB(r2d2::PooledConnection<ConnectionManager<PgConnection>>);

impl Deref for DB {
    type Target = PgConnection;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<'a, 'r> FromRequest<'a, 'r> for DB {
    type Error = ();

    fn from_request(request: &'a Request<'r>) -> request::Outcome<DB, ()> {
        let pool = match <State<Pool> as FromRequest>::from_request(request) {
            Outcome::Success(pool) => pool,
            Outcome::Failure(e) => return Outcome::Failure(e),
            Outcome::Forward(_) => return Outcome::Forward(()),
        };

        match pool.get() {
            Ok(conn) => Outcome::Success(DB(conn)),
            Err(_) => Outcome::Failure((Status::ServiceUnavailable, ())),
        }
    }
}
<commit_msg>Documentation for helpers/db<commit_after>use std::ops::Deref;
use rocket::http::Status;
use rocket::{Request, State, Outcome};
use rocket::request::{self, FromRequest};
use diesel::pg::PgConnection;
use r2d2;
use r2d2_diesel::ConnectionManager;
use std::env;

pub type Pool = r2d2::Pool<ConnectionManager<PgConnection>>;

/// Generate a new r2d2 Pool for connection management in rocket sessions.
/// This object will be handed to the rocket `.manage()` managed state handler.
pub fn init_db_pool() -> Pool {
    let config = r2d2::Config::default();
    let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
    let manager = ConnectionManager::<PgConnection>::new(database_url);
    r2d2::Pool::new(config, manager).expect("Failed to create pool.")
}

/// A helper struct to pass the `PooledConnection` from r2d2 to the request function.
/// It implements the `Deref` trait to easily get the `PgConnection` by using `*`
pub struct DB(r2d2::PooledConnection<ConnectionManager<PgConnection>>);

impl Deref for DB {
    type Target = PgConnection;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// The implementation of `FromRequest` for `DB`.
/// This function fetches the pool from the managed state of rust and gets a new
/// `PgConnection` from the `ConnectionManager`.
/// This connection will then be used during the request and closed afterwards.
impl<'a, 'r> FromRequest<'a, 'r> for DB {
    type Error = ();

    fn from_request(request: &'a Request<'r>) -> request::Outcome<DB, ()> {
        let pool = match <State<Pool> as FromRequest>::from_request(request) {
            Outcome::Success(pool) => pool,
            Outcome::Failure(e) => return Outcome::Failure(e),
            Outcome::Forward(_) => return Outcome::Forward(()),
        };

        match pool.get() {
            Ok(conn) => Outcome::Success(DB(conn)),
            Err(_) => Outcome::Failure((Status::ServiceUnavailable, ())),
        }
    }
}
<|endoftext|>