<commit_before>extern crate log;
extern crate log4rs;

use log::LogLevelFilter;
use log4rs::append::console::ConsoleAppender;
use log4rs::append::file::FileAppender;
use log4rs::config::{Appender, Config, Logger, Root};
use log4rs::encode::pattern::PatternEncoder;
use std::env::var;
use std::sync::{Once, ONCE_INIT};

static INIT: Once = ONCE_INIT;

pub fn init() {
    INIT.call_once(|| init_logger());
}

fn init_logger() {
    let config = match var("K2_LOG_PATH") {
        Ok(path) => init_file_logger(String::from(path)),
        _ => init_stdout_logger(),
    };

    log4rs::init_config(config).unwrap();
}

fn init_file_logger(mut path: String) -> Config {
    if !path.trim().ends_with("/") {
        path.push_str("/");
    }

    path.push_str("ctehxk2.log");

    let appender_id = "file";

    let file = FileAppender::builder()
        .encoder(Box::new(PatternEncoder::new("{d} {l} {M}: {m}{n}")))
        .build(path)
        .unwrap();

    Config::builder()
        .appender(Appender::builder().build(appender_id, Box::new(file)))
        .logger(Logger::builder()
                    .appender(appender_id)
                    .additive(false)
                    .build("ctehxk2", LogLevelFilter::Debug))
        .build(Root::builder()
                   .appender(appender_id)
                   .build(LogLevelFilter::Error))
        .unwrap()
}

fn init_stdout_logger() -> Config {
    let stdout = ConsoleAppender::builder().build();

    let appender_id = "stdout";

    Config::builder()
        .appender(Appender::builder().build(appender_id, Box::new(stdout)))
        .logger(Logger::builder()
                    .appender(appender_id)
                    .additive(false)
                    .build("ctehxk2", LogLevelFilter::Debug))
        .build(Root::builder()
                   .appender(appender_id)
                   .build(LogLevelFilter::Error))
        .unwrap()
}
<commit_msg>use os specifc path seperator<commit_after>extern crate log;
extern crate log4rs;

use log::LogLevelFilter;
use log4rs::append::console::ConsoleAppender;
use log4rs::append::file::FileAppender;
use log4rs::config::{Appender, Config, Logger, Root};
use log4rs::encode::pattern::PatternEncoder;
use std::env::var;
use std::path::MAIN_SEPARATOR;
use std::sync::{Once, ONCE_INIT};

static INIT: Once = ONCE_INIT;

pub fn init() {
    INIT.call_once(|| init_logger());
}

fn init_logger() {
    let config = match var("K2_LOG_PATH") {
        Ok(path) => init_file_logger(String::from(path)),
        _ => init_stdout_logger(),
    };

    log4rs::init_config(config).unwrap();
}

fn init_file_logger(mut path: String) -> Config {
    if !path.trim().ends_with(MAIN_SEPARATOR) {
        path.push(MAIN_SEPARATOR);
    }

    path.push_str("ctehxk2.log");

    let appender_id = "file";

    let file = FileAppender::builder()
        .encoder(Box::new(PatternEncoder::new("{d} {l} {M}: {m}{n}")))
        .build(path)
        .unwrap();

    Config::builder()
        .appender(Appender::builder().build(appender_id, Box::new(file)))
        .logger(Logger::builder()
                    .appender(appender_id)
                    .additive(false)
                    .build("ctehxk2", LogLevelFilter::Debug))
        .build(Root::builder()
                   .appender(appender_id)
                   .build(LogLevelFilter::Error))
        .unwrap()
}

fn init_stdout_logger() -> Config {
    let stdout = ConsoleAppender::builder().build();

    let appender_id = "stdout";

    Config::builder()
        .appender(Appender::builder().build(appender_id, Box::new(stdout)))
        .logger(Logger::builder()
                    .appender(appender_id)
                    .additive(false)
                    .build("ctehxk2", LogLevelFilter::Debug))
        .build(Root::builder()
                   .appender(appender_id)
                   .build(LogLevelFilter::Error))
        .unwrap()
}
<|endoftext|>


<commit_before>use std::{num, vec};
use super::{Addr, Addressable};

/// Generic read/write memory (RAM)
pub struct Ram<A> {
	priv data: ~[u8],
	priv last_addr: A,
}

impl<A: Addr> Ram<A> {
	/// Create new RAM with full capacity of its address range
	pub fn new () -> Ram<A> {
		Ram::with_capacity(num::Bounded::max_value())
	}

	/// Create new RAM which will be addressable from 0 to the given address
	pub fn with_capacity (last_addr: A) -> Ram<A> {
		let size: uint = 1 + num::cast(last_addr.clone()).unwrap();
		Ram { data: vec::from_elem(size, 0u8), last_addr: last_addr }
	}

	/// Returns the size of the RAM
	pub fn size (&self) -> uint {
		self.data.len()
	}
}

impl<A: Addr> Addressable<A> for Ram<A> {
	fn get (&self, addr: A) -> u8 {
		if addr > self.last_addr { fail!("ram: Read beyond memory bounds (${:X} > ${:X})", addr, self.last_addr); }
		let i: u64 = num::cast(addr).unwrap();
		self.data[i]
	}

	fn set (&mut self, addr: A, data: u8) {
		if addr >= self.last_addr { fail!("ram: Write beyond memory bounds (${:X} > ${:X})", addr, self.last_addr); }
		let i: u64 = num::cast(addr).unwrap();
		self.data[i] = data;
	}
}

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

	#[test]
	fn create_with_full_addressable_capacity () {
		let memory: Ram<u8> = Ram::new();
		assert_eq!(memory.size(), 256);
	}

	#[test]
	fn create_with_requested_capacity () {
		let memory: Ram<u16> = Ram::with_capacity(0x03ff_u16);
		assert_eq!(memory.size(), 1024);
	}

	#[test]
	fn read_write () {
		let mut memory: Ram<u16> = Ram::with_capacity(0x03ff_u16);
		memory.set(0x0123, 0x55);
		assert_eq!(memory.get(0x0123), 0x55);
	}
}
<commit_msg>Fix ram writing memory bounds check<commit_after>use std::{num, vec};
use super::{Addr, Addressable};

/// Generic read/write memory (RAM)
pub struct Ram<A> {
	priv data: ~[u8],
	priv last_addr: A,
}

impl<A: Addr> Ram<A> {
	/// Create new RAM with full capacity of its address range
	pub fn new () -> Ram<A> {
		Ram::with_capacity(num::Bounded::max_value())
	}

	/// Create new RAM which will be addressable from 0 to the given address
	pub fn with_capacity (last_addr: A) -> Ram<A> {
		let size: uint = 1 + num::cast(last_addr.clone()).unwrap();
		Ram { data: vec::from_elem(size, 0u8), last_addr: last_addr }
	}

	/// Returns the size of the RAM
	pub fn size (&self) -> uint {
		self.data.len()
	}
}

impl<A: Addr> Addressable<A> for Ram<A> {
	fn get (&self, addr: A) -> u8 {
		if addr > self.last_addr { fail!("ram: Read beyond memory bounds (${:X} > ${:X})", addr, self.last_addr); }
		let i: u64 = num::cast(addr).unwrap();
		self.data[i]
	}

	fn set (&mut self, addr: A, data: u8) {
		if addr > self.last_addr { fail!("ram: Write beyond memory bounds (${:X} > ${:X})", addr, self.last_addr); }
		let i: u64 = num::cast(addr).unwrap();
		self.data[i] = data;
	}
}

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

	#[test]
	fn create_with_full_addressable_capacity () {
		let memory: Ram<u8> = Ram::new();
		assert_eq!(memory.size(), 256);
	}

	#[test]
	fn create_with_requested_capacity () {
		let memory: Ram<u16> = Ram::with_capacity(0x03ff_u16);
		assert_eq!(memory.size(), 1024);
	}

	#[test]
	fn read_write () {
		let mut memory: Ram<u16> = Ram::with_capacity(0x03ff_u16);
		memory.set(0x0123, 0x55);
		assert_eq!(memory.get(0x0123), 0x55);
	}
}
<|endoftext|>


<commit_before>use code::Code;

/// Error generated by the parser.
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub enum ParseError {
    /// String was empty.
    EmptyCommand,
    /// Message did not have a code.
    EmptyMessage,
    /// Unexpected end of the string.
    UnexpectedEnd,
}

/// Represents a message received from the server.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Message {
    /// Prefix
    pub prefix: Option<Prefix>,
    /// Code
    pub code: Code,
    /// Arguments
    pub args: Vec<String>,
    /// Suffix
    pub suffix: Option<String>,
}

impl Message {

    /// Parse the given string into a `Message` struct.
    ///
    /// An error is returned if the message is not valid.
    pub fn parse(line: &str) -> Result<Message, ParseError> {
        if line.len() == 0 || line.trim().len() == 0 {
            return Err(ParseError::EmptyMessage);
        }

        let mut state = line.trim_right_matches("\r\n");
        let mut prefix: Option<Prefix> = None;
        let code: Option<String>;
        let mut args: Vec<String> = Vec::new();
        let mut suffix: Option<String> = None;

        // Look for a prefix
        if state.starts_with(":") {
            match state.find(" ") {
                None => return Err(ParseError::UnexpectedEnd),
                Some(idx) => {
                    prefix = parse_prefix(&state[1..idx]);
                    state = &state[idx + 1..];
                }
            }
        }

        // Look for the command/reply
        match state.find(" ") {
            None => {
                if state.len() == 0 {
                    return Err(ParseError::EmptyMessage);
                } else {
                    code = Some(state[..].to_string());
                    state = &state[state.len()..];
                }
            }
            Some(idx) => {
                code = Some(state[..idx].to_string());
                state = &state[idx + 1..];
            }
        }

        // Look for arguments and the suffix
        if state.len() > 0 {
            loop {
                if state.starts_with(":") {
                    suffix = Some(state[1..].to_string());
                    break;
                } else {
                    match state.find(" ") {
                        None => {
                            args.push(state[..].to_string());
                            break;
                        }
                        Some(idx) => {
                            args.push(state[..idx].to_string());
                            state = &state[idx + 1..];
                        }
                    }
                }
            }
        }

        let code = match code {
            None => return Err(ParseError::EmptyCommand),
            Some(text) => {
                match text.parse() {
                    Ok(code) => code,
                    Err(_) => Code::Unknown(text.to_string()),
                }
            }
        };

        Ok(Message {
            prefix: prefix,
            code: code,
            args: args,
            suffix: suffix,
        })
    }
}

fn parse_prefix(prefix: &str) -> Option<Prefix> {
    match prefix.find("!") {
        None => Some(Prefix::Server(prefix.to_string())),
        Some(excpos) => {
            let nick = &prefix[..excpos];
            let rest = &prefix[excpos + 1..];
            match rest.find("@") {
                None => return None,
                Some(atpos) => {
                    let user = &rest[..atpos];
                    let host = &rest[atpos + 1..];
                    return Some(Prefix::User(PrefixUser::new(nick, user, host)));
                }
            }
        }
    }
}

/// Prefix of the message.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Prefix {
    /// Prefix is a user.
    User(PrefixUser),
    /// Prefix is a server.
    Server(String),
}

/// User prefix representation.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PrefixUser {
    /// Nickname
    pub nickname: String,
    /// Username
    pub username: String,
    /// Hostname
    pub hostname: String,
}

impl PrefixUser {

    fn new(nick: &str, user: &str, host: &str) -> PrefixUser {
        PrefixUser {
            nickname: nick.into(),
            username: user.into(),
            hostname: host.into(),
        }
    }

}

#[test]
fn test_full() {
    let res = Message::parse(":org.prefix.cool COMMAND arg1 arg2 arg3 :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.code, Code::Unknown("COMMAND".to_string()));
    assert_eq!(msg.args, vec!["arg1", "arg2", "arg3"]);
    assert_eq!(msg.suffix, Some("suffix is pretty cool yo".to_string()));
}

#[test]
fn test_no_prefix() {
    let res = Message::parse("NICK arg1 arg2 arg3 :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.prefix, None);
    assert_eq!(msg.code, Code::Nick);
    assert_eq!(msg.args, vec!["arg1", "arg2", "arg3"]);
    assert_eq!(msg.suffix, Some("suffix is pretty cool yo".to_string()));
}

#[test]
fn test_no_suffix() {
    let res = Message::parse(":org.prefix.cool NICK arg1 arg2 arg3");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.code, Code::Nick);
    assert_eq!(msg.args, vec!["arg1", "arg2", "arg3"]);
    assert_eq!(msg.suffix, None);
}

#[test]
fn test_no_args() {
    let res = Message::parse(":org.prefix.cool NICK :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.code, Code::Nick);
    assert_eq!(msg.args.len(), 0);
    assert_eq!(msg.suffix, Some("suffix is pretty cool yo".to_string()));
}

#[test]
fn test_only_command() {
    let res = Message::parse("NICK");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.prefix, None);
    assert_eq!(msg.code, Code::Nick);
    assert_eq!(msg.args.len(), 0);
    assert_eq!(msg.suffix, None);
}

#[test]
fn test_empty_message() {
    let res = Message::parse("");
    assert!(res.is_err());
    let err = res.err().unwrap();
    assert!(err == ParseError::EmptyMessage);
}

#[test]
fn test_empty_message_trim() {
    let res = Message::parse("    ");
    assert!(res.is_err());
    let err = res.err().unwrap();
    assert!(err == ParseError::EmptyMessage);
}

#[test]
fn test_only_prefix() {
    let res = Message::parse(":org.prefix.cool");
    assert!(res.is_err());
    let err = res.err().unwrap();
    assert!(err == ParseError::UnexpectedEnd);
}

#[test]
fn test_prefix_none() {
    let res = Message::parse("COMMAND :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert!(msg.prefix == None);
}

#[test]
fn test_prefix_server() {
    let res = Message::parse(":irc.freenode.net COMMAND :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.prefix, Some(Prefix::Server("irc.freenode.net".to_string())));
}

#[test]
fn test_prefix_user() {
    let res = Message::parse(":bob!bob@bob.com COMMAND :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.prefix, Some(Prefix::User(PrefixUser::new("bob", "bob", "bob.com"))));
}
<commit_msg>Remove suffix, it should be the last argument<commit_after>use code::Code;

/// Error generated by the parser.
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub enum ParseError {
    /// String was empty.
    EmptyCommand,
    /// Message did not have a code.
    EmptyMessage,
    /// Unexpected end of the string.
    UnexpectedEnd,
}

/// Represents a message received from the server.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Message {
    /// Prefix
    pub prefix: Option<Prefix>,
    /// Code
    pub code: Code,
    /// Arguments
    pub args: Vec<String>,
}

impl Message {

    /// Parse the given string into a `Message` struct.
    ///
    /// An error is returned if the message is not valid.
    pub fn parse(line: &str) -> Result<Message, ParseError> {
        if line.len() == 0 || line.trim().len() == 0 {
            return Err(ParseError::EmptyMessage);
        }

        let mut state = line.trim_right_matches("\r\n");
        let mut prefix: Option<Prefix> = None;
        let code: Option<&str>;
        let mut args: Vec<String> = Vec::new();

        // Look for a prefix
        if state.starts_with(":") {
            match state.find(" ") {
                None => return Err(ParseError::UnexpectedEnd),
                Some(idx) => {
                    prefix = parse_prefix(&state[1..idx]);
                    state = &state[idx + 1..];
                }
            }
        }

        // Look for the command/reply
        match state.find(" ") {
            None => {
                if state.len() == 0 {
                    return Err(ParseError::EmptyMessage);
                } else {
                    code = Some(&state[..]);
                    state = &state[state.len()..];
                }
            }
            Some(idx) => {
                code = Some(state[..idx].into());
                state = &state[idx + 1..];
            }
        }

        // Look for arguments and the suffix
        if state.len() > 0 {
            loop {
                if state.starts_with(":") {
                    args.push(state[1..].into());
                    break;
                } else {
                    match state.find(" ") {
                        None => {
                            args.push(state[..].into());
                            break;
                        }
                        Some(idx) => {
                            args.push(state[..idx].into());
                            state = &state[idx + 1..];
                        }
                    }
                }
            }
        }

        let code = match code {
            None => return Err(ParseError::EmptyCommand),
            Some(text) => {
                match text.parse() {
                    Ok(code) => code,
                    Err(_) => Code::Unknown(text.into()),
                }
            }
        };

        Ok(Message {
            prefix: prefix,
            code: code,
            args: args,
        })
    }
}

fn parse_prefix(prefix: &str) -> Option<Prefix> {
    match prefix.find("!") {
        None => Some(Prefix::Server(prefix.to_string())),
        Some(excpos) => {
            let nick = &prefix[..excpos];
            let rest = &prefix[excpos + 1..];
            match rest.find("@") {
                None => return None,
                Some(atpos) => {
                    let user = &rest[..atpos];
                    let host = &rest[atpos + 1..];
                    return Some(Prefix::User(PrefixUser::new(nick, user, host)));
                }
            }
        }
    }
}

/// Prefix of the message.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Prefix {
    /// Prefix is a user.
    User(PrefixUser),
    /// Prefix is a server.
    Server(String),
}

/// User prefix representation.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PrefixUser {
    /// Nickname
    pub nickname: String,
    /// Username
    pub username: String,
    /// Hostname
    pub hostname: String,
}

impl PrefixUser {

    fn new(nick: &str, user: &str, host: &str) -> PrefixUser {
        PrefixUser {
            nickname: nick.into(),
            username: user.into(),
            hostname: host.into(),
        }
    }

}

#[test]
fn test_full() {
    let res = Message::parse(":org.prefix.cool COMMAND arg1 arg2 arg3 :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.code, Code::Unknown("COMMAND".to_string()));
    assert_eq!(msg.args, vec!["arg1", "arg2", "arg3", "suffix is pretty cool yo"]);
}

#[test]
fn test_no_prefix() {
    let res = Message::parse("NICK arg1 arg2 arg3 :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.prefix, None);
    assert_eq!(msg.code, Code::Nick);
    assert_eq!(msg.args, vec!["arg1", "arg2", "arg3", "suffix is pretty cool yo"]);
}

#[test]
fn test_no_suffix() {
    let res = Message::parse(":org.prefix.cool NICK arg1 arg2 arg3");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.code, Code::Nick);
    assert_eq!(msg.args, vec!["arg1", "arg2", "arg3"]);
}

#[test]
fn test_no_args() {
    let res = Message::parse(":org.prefix.cool NICK :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.code, Code::Nick);
    assert_eq!(msg.args, vec!["suffix is pretty cool yo"]);
}

#[test]
fn test_only_command() {
    let res = Message::parse("NICK");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.prefix, None);
    assert_eq!(msg.code, Code::Nick);
    assert_eq!(msg.args.len(), 0);
}

#[test]
fn test_empty_message() {
    let res = Message::parse("");
    assert!(res.is_err());
    let err = res.err().unwrap();
    assert!(err == ParseError::EmptyMessage);
}

#[test]
fn test_empty_message_trim() {
    let res = Message::parse("    ");
    assert!(res.is_err());
    let err = res.err().unwrap();
    assert!(err == ParseError::EmptyMessage);
}

#[test]
fn test_only_prefix() {
    let res = Message::parse(":org.prefix.cool");
    assert!(res.is_err());
    let err = res.err().unwrap();
    assert!(err == ParseError::UnexpectedEnd);
}

#[test]
fn test_prefix_none() {
    let res = Message::parse("COMMAND :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.args, vec!["suffix is pretty cool yo"]);
}

#[test]
fn test_prefix_server() {
    let res = Message::parse(":irc.freenode.net COMMAND :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.prefix, Some(Prefix::Server("irc.freenode.net".into())));
}

#[test]
fn test_prefix_user() {
    let res = Message::parse(":bob!bob@bob.com COMMAND :suffix is pretty cool yo");
    assert!(res.is_ok());
    let msg = res.ok().unwrap();
    assert_eq!(msg.prefix, Some(Prefix::User(PrefixUser::new("bob", "bob", "bob.com"))));
}
<|endoftext|>


<commit_before>use address::{Addressable, Address};
use messages::outgoing;
use mio;
use node::Node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use stun;
use transaction::{TransactionId, TransactionIdGenerator};

#[derive(Clone, Debug)]
pub enum Event {
    ReceivedPacket(Address, Vec<u8>),
    Shutdown,
    Started
}

pub enum OneshotTask {
    Incoming(Vec<u8>),
    StartBootstrap,
    SendPacket(Address, Vec<u8>),
    Shutdown
}

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

enum TableAction {
    Bootstrap(mio::Timeout),
    HealthCheck(mio::Timeout),
    RefreshBucket(mio::Timeout)
}

pub type TaskSender = mio::Sender<OneshotTask>;

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Node,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>,
    shutdown_signals: Vec<mpsc::Sender<mpsc::Sender<()>>>
}

impl Network {
    pub fn new<T: ToSocketAddrs>(self_address: Address, host: T, routers: Vec<Node>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let mapped_host = stun::get_mapped_address(host).expect("Couldn't STUN myself");
        let self_node = Node::from_socket_addrs(self_address, mapped_host).unwrap();
        let routing_table = RoutingTable::new(8, self_address, routers);

        Network {
            host: host,
            routing_table: routing_table,
            self_node: self_node,
            transaction_ids: TransactionIdGenerator::new(),
            status: Status::Idle,
            pending_actions: HashMap::new(),
            event_listeners: vec![],
            shutdown_signals: vec![]
        }
    }

    pub fn run(mut self) -> TaskSender {
        let mut event_loop_config = mio::EventLoopConfig::new();
        event_loop_config.notify_capacity(16384);
        let mut event_loop = mio::EventLoop::configured(event_loop_config).unwrap();

        let shutdown_signal = create_incoming_udp_channel(self.host, event_loop.channel());
        self.shutdown_signals.push(shutdown_signal);

        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        info!("Running server at {:?}", self.self_node);
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

    pub fn register_event_listener(&mut self, event_listener: mpsc::Sender<Event>) {
        self.event_listeners.push(event_listener);
    }

    fn handle_incoming(&mut self, data: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        use messages::incoming::{Message, Query, Response, self};
        let mut data = Cursor::new(data);
        let message = incoming::parse_from_reader(&mut data).unwrap();

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.address();

                match query {
                    Query::FindNode(target) => {
                        debug!("Received FindNode query for {} from {:?}", &target, &origin);
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_live_nodes_to(&target, false));
                        origin.send(response);
                    },
                    Query::Packet(payload) => {
                        // Logging packets would be too chatty
                        self.broadcast_event(Event::ReceivedPacket(origin_address, payload));
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        debug!("Received Ping from {:?}", &origin);
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);
                    }
                }

                // Always insert the origin node
                self.insert_node(origin).unwrap();
                if let Some(origin) = self.routing_table.find_node(&origin_address) {
                    origin.received_query(transaction_id);
                }
            }

            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                // Always insert the origin node
                let mut encounted_new_node = self.insert_node(origin) == Ok(InsertOutcome::Inserted);

                match response {
                    Response::FindNode(mut nodes) => {
                        for node in nodes.drain(..) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

                        {
                            if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_response(transaction_id);
                            }
                        }

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop)
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                                // Does not clear timeout. Lets Healthcheck proceed at regular
                                // interval.
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

                        match self.status {
                            Status::Bootstrapping => {
                                self.status = Status::Idle;
                            }
                            _ => { }
                        }
                    }
                    Response::Packet => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                    Response::Ping => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                }
            }
        }
    }

    fn start_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        self.status = Status::Bootstrapping;
        self.continue_bootstrap(event_loop);
        self.broadcast_event(Event::Started);
    }

    fn continue_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let address = &self.self_node.address();
        let transaction_id = self.find_node(address);
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueBootstrap, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::Bootstrap(timeout));
    }

    fn continue_health_check(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let transaction_id = self.health_check();
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueHealthCheck, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::HealthCheck(timeout));
    }

    fn continue_refresh(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueRefresh, 1000).unwrap();

        match self.refresh_bucket() {
            Some(transaction_id) => {
                self.pending_actions.insert(transaction_id, TableAction::RefreshBucket(timeout));
            }
            None => { }
        }
    }

    fn find_node(&mut self, address: &Address) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        let query = outgoing::create_find_node_query(
            transaction_id,
            &self.self_node,
            address);
        for node in self.routing_table.nearest() {
            node.sent_query(transaction_id);
            node.send(query.clone());
        }
        transaction_id
    }

    // TODO: I should de-couple operations and transactions. Some operations, e.g. health_check
    // can be comprised of multiple transactions that must be completed before the operation is
    // complete. In this case, I should receive a ping from each questionable node before
    // health_check is done.
    fn health_check(&mut self) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();

        // TODO: this should be a separate keep-alive task, but it will be
        // dependant on the type of connection we're keeping alive.
        //
        // Ping nearest node every time.
        if let Some(nearest_node) = self.routing_table.nearest().first_mut() {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            nearest_node.sent_query(transaction_id);
            nearest_node.send(query);
        }

        for node in self.routing_table.questionable_nodes().iter_mut() {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.sent_query(transaction_id);
            node.send(query);
        }
        transaction_id
    }

    fn refresh_bucket(&mut self) -> Option<TransactionId> {
        if let Some(address) = self.address_to_find_for_refresh() {
            Some(self.find_node(&address))
        } else {
            None
        }
    }

    fn address_to_find_for_refresh(&self) -> Option<Address> {
        match self.routing_table.bucket_needing_refresh() {
            Some(bucket) => {
                Some(bucket.random_address_in_space())
            }
            None => None
        }
    }

    fn insert_node(&mut self, node: Node) -> InsertionResult {
        self.routing_table.insert(node, &self.self_node, &mut self.transaction_ids)
    }

    fn send_packet(&mut self, recipient: Address, payload: Vec<u8>, _event_loop: &mut mio::EventLoop<Handler>) {
        for node in self.routing_table.nearest_live_nodes_to(&recipient, false) {
            let transaction_id = self.transaction_ids.generate();
            let query = outgoing::create_packet_query(
                transaction_id, &self.self_node, payload.clone());
            node.sent_query(transaction_id);
            node.send(query.clone());
        }
    }

    fn shutdown(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let sentinels: Vec<mpsc::Receiver<()>> = self.shutdown_signals.iter().map(|signal| {
            let (sentinal_sender, sentinal_receiver) = mpsc::channel();
            signal.send(sentinal_sender).expect("couldn't send shutdown signal");
            sentinal_receiver
        }).collect();
        for sentinel in sentinels {
            sentinel.recv().unwrap();
        }
        event_loop.shutdown();
        self.broadcast_event(Event::Shutdown);
    }

    fn broadcast_event(&self, event: Event) {
        for listener in &self.event_listeners {
            listener.send(event.clone()).expect("Failed to broadcast event");
        }
    }
}

struct Handler {
    network: Network
}

impl Handler {
    fn new(network: Network) -> Handler {
        Handler {
            network: network
        }
    }
}

impl mio::Handler for Handler {
    type Timeout = ScheduledTask;
    type Message = OneshotTask;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Handler>, task: OneshotTask) {
        match task {
            OneshotTask::Incoming(data) => self.network.handle_incoming(data, event_loop),
            OneshotTask::StartBootstrap => self.network.start_bootstrap(event_loop),
            OneshotTask::SendPacket(recipient, payload) =>
                self.network.send_packet(recipient, payload, event_loop),
            OneshotTask::Shutdown => self.network.shutdown(event_loop)
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Handler>, timeout: ScheduledTask) {
        match timeout {
            ScheduledTask::ContinueBootstrap => self.network.continue_bootstrap(event_loop),
            ScheduledTask::ContinueHealthCheck => self.network.continue_health_check(event_loop),
            ScheduledTask::ContinueRefresh => self.network.continue_refresh(event_loop)
        }
    }
}

// TODO: I'd rather have some kind of select on two channels instead of this nested non-blocking
// recv. It only works because of the 10ms timeout on the socket.
fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) -> mpsc::Sender<mpsc::Sender<()>> {
    use std::time::Duration;

    let (shutdown_sender, shutdown_receiver) = mpsc::channel::<mpsc::Sender<()>>();

    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        socket.set_read_timeout(Some(Duration::from_millis(10))).unwrap();

        loop {
            match shutdown_receiver.try_recv() {
                Err(mpsc::TryRecvError::Empty) => {
                    let mut buf = [0; 4096];
                    match socket.recv(&mut buf) {
                        Ok(size) => {
                            sender.send(OneshotTask::Incoming(buf[..size].iter().cloned().collect()))
                                .unwrap_or_else(|err| info!("Couldn't handling incoming: {:?}", err));
                        }
                        Err(_) => {
                            //warn!("Error receiving from server: {}", e)
                        }
                    }
                }

                Err(mpsc::TryRecvError::Disconnected) => {
                    // The other end of the shutdown_receiver was dropped
                    break
                }

                Ok(sentinel) => {
                    sentinel.send(()).unwrap();
                    break;
                }
            }
        }

    });
    shutdown_sender
}
<commit_msg>Only leave bootstrapping status when no new nodes<commit_after>use address::{Addressable, Address};
use messages::outgoing;
use mio;
use node::Node;
use routing_table::{InsertOutcome, InsertionResult, RoutingTable};
use std::collections::HashMap;
use std::io::Cursor;
use std::net::{UdpSocket, ToSocketAddrs, SocketAddr};
use std::sync::mpsc;
use std::thread;
use stun;
use transaction::{TransactionId, TransactionIdGenerator};

#[derive(Clone, Debug)]
pub enum Event {
    ReceivedPacket(Address, Vec<u8>),
    Shutdown,
    Started
}

pub enum OneshotTask {
    Incoming(Vec<u8>),
    StartBootstrap,
    SendPacket(Address, Vec<u8>),
    Shutdown
}

pub enum ScheduledTask {
    ContinueBootstrap,
    ContinueHealthCheck,
    ContinueRefresh
}

enum TableAction {
    Bootstrap(mio::Timeout),
    HealthCheck(mio::Timeout),
    RefreshBucket(mio::Timeout)
}

pub type TaskSender = mio::Sender<OneshotTask>;

enum Status {
    Bootstrapping,
    Idle
}

pub struct Network {
    host: SocketAddr,
    routing_table: RoutingTable,
    self_node: Node,
    transaction_ids: TransactionIdGenerator,
    status: Status,
    pending_actions: HashMap<TransactionId, TableAction>,
    event_listeners: Vec<mpsc::Sender<Event>>,
    shutdown_signals: Vec<mpsc::Sender<mpsc::Sender<()>>>
}

impl Network {
    pub fn new<T: ToSocketAddrs>(self_address: Address, host: T, routers: Vec<Node>) -> Network {
        let host = host.to_socket_addrs().unwrap().next().unwrap();
        let mapped_host = stun::get_mapped_address(host).expect("Couldn't STUN myself");
        let self_node = Node::from_socket_addrs(self_address, mapped_host).unwrap();
        let routing_table = RoutingTable::new(8, self_address, routers);

        Network {
            host: host,
            routing_table: routing_table,
            self_node: self_node,
            transaction_ids: TransactionIdGenerator::new(),
            status: Status::Idle,
            pending_actions: HashMap::new(),
            event_listeners: vec![],
            shutdown_signals: vec![]
        }
    }

    pub fn run(mut self) -> TaskSender {
        let mut event_loop_config = mio::EventLoopConfig::new();
        event_loop_config.notify_capacity(16384);
        let mut event_loop = mio::EventLoop::configured(event_loop_config).unwrap();

        let shutdown_signal = create_incoming_udp_channel(self.host, event_loop.channel());
        self.shutdown_signals.push(shutdown_signal);

        event_loop.channel().send(OneshotTask::StartBootstrap).unwrap();
        info!("Running server at {:?}", self.self_node);
        let mut handler = Handler::new(self);
        let task_sender = event_loop.channel();
        thread::spawn(move|| event_loop.run(&mut handler).unwrap());
        task_sender
    }

    pub fn register_event_listener(&mut self, event_listener: mpsc::Sender<Event>) {
        self.event_listeners.push(event_listener);
    }

    fn handle_incoming(&mut self, data: Vec<u8>, event_loop: &mut mio::EventLoop<Handler>) {
        use messages::incoming::{Message, Query, Response, self};
        let mut data = Cursor::new(data);
        let message = incoming::parse_from_reader(&mut data).unwrap();

        match message {
            Message::Query(transaction_id, origin, query) => {
                let origin_address = origin.address();

                match query {
                    Query::FindNode(target) => {
                        debug!("Received FindNode query for {} from {:?}", &target, &origin);
                        let response: Vec<u8> = outgoing::create_find_node_response(
                                transaction_id,
                                &self.self_node,
                                self.routing_table.nearest_live_nodes_to(&target, false));
                        origin.send(response);
                    },
                    Query::Packet(payload) => {
                        // Logging packets would be too chatty
                        self.broadcast_event(Event::ReceivedPacket(origin_address, payload));
                        let response = outgoing::create_packet_response(
                            transaction_id, &self.self_node);
                        origin.send(response);
                    },
                    Query::Ping => {
                        debug!("Received Ping from {:?}", &origin);
                        let response = outgoing::create_ping_response(
                            transaction_id,
                            &self.self_node);
                        origin.send(response);
                    }
                }

                // Always insert the origin node
                self.insert_node(origin).unwrap();
                if let Some(origin) = self.routing_table.find_node(&origin_address) {
                    origin.received_query(transaction_id);
                }
            }

            Message::Response(transaction_id, origin, response) => {
                let origin_address = origin.address();
                // Always insert the origin node
                let mut encounted_new_node = self.insert_node(origin) == Ok(InsertOutcome::Inserted);

                match response {
                    Response::FindNode(mut nodes) => {
                        for node in nodes.drain(..) {
                            match self.insert_node(node) {
                                Ok(InsertOutcome::Inserted) => {
                                    encounted_new_node = true;
                                }
                                Err(error) => { panic!(error) }
                                _ => { }
                            }
                        }

                        {
                            if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                                origin.received_response(transaction_id);
                            }
                        }

                        match self.pending_actions.remove(&transaction_id) {
                            Some(TableAction::Bootstrap(timeout)) => {
                                event_loop.clear_timeout(timeout);
                                if encounted_new_node {
                                    self.continue_bootstrap(event_loop);
                                } else {
                                    // TODO: Extract finish_bootstrap()
                                    self.continue_health_check(event_loop);
                                    self.continue_refresh(event_loop);
                                    match self.status {
                                        Status::Bootstrapping => {
                                            self.status = Status::Idle;
                                        }
                                        _ => { }
                                    }
                                }
                            }
                            Some(TableAction::HealthCheck(_)) => {
                                // Does not clear timeout. Lets Healthcheck proceed at regular
                                // interval.
                            },
                            Some(TableAction::RefreshBucket(_)) => {
                            }
                            None => { }
                        }

                    }
                    Response::Packet => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                    Response::Ping => {
                        if let Some(mut origin) = self.routing_table.find_node(&origin_address) {
                            origin.received_response(transaction_id);
                        }
                    }
                }
            }
        }
    }

    fn start_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        self.status = Status::Bootstrapping;
        self.continue_bootstrap(event_loop);
        self.broadcast_event(Event::Started);
    }

    fn continue_bootstrap(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let address = &self.self_node.address();
        let transaction_id = self.find_node(address);
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueBootstrap, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::Bootstrap(timeout));
    }

    fn continue_health_check(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let transaction_id = self.health_check();
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueHealthCheck, 1000).unwrap();
        self.pending_actions.insert(transaction_id, TableAction::HealthCheck(timeout));
    }

    fn continue_refresh(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let timeout = event_loop.timeout_ms(ScheduledTask::ContinueRefresh, 1000).unwrap();

        match self.refresh_bucket() {
            Some(transaction_id) => {
                self.pending_actions.insert(transaction_id, TableAction::RefreshBucket(timeout));
            }
            None => { }
        }
    }

    fn find_node(&mut self, address: &Address) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();
        let query = outgoing::create_find_node_query(
            transaction_id,
            &self.self_node,
            address);
        for node in self.routing_table.nearest() {
            node.sent_query(transaction_id);
            node.send(query.clone());
        }
        transaction_id
    }

    // TODO: I should de-couple operations and transactions. Some operations, e.g. health_check
    // can be comprised of multiple transactions that must be completed before the operation is
    // complete. In this case, I should receive a ping from each questionable node before
    // health_check is done.
    fn health_check(&mut self) -> TransactionId {
        let transaction_id = self.transaction_ids.generate();

        // TODO: this should be a separate keep-alive task, but it will be
        // dependant on the type of connection we're keeping alive.
        //
        // Ping nearest node every time.
        if let Some(nearest_node) = self.routing_table.nearest().first_mut() {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            nearest_node.sent_query(transaction_id);
            nearest_node.send(query);
        }

        for node in self.routing_table.questionable_nodes().iter_mut() {
            let query = outgoing::create_ping_query(
                transaction_id, &self.self_node);
            node.sent_query(transaction_id);
            node.send(query);
        }
        transaction_id
    }

    fn refresh_bucket(&mut self) -> Option<TransactionId> {
        if let Some(address) = self.address_to_find_for_refresh() {
            Some(self.find_node(&address))
        } else {
            None
        }
    }

    fn address_to_find_for_refresh(&self) -> Option<Address> {
        match self.routing_table.bucket_needing_refresh() {
            Some(bucket) => {
                Some(bucket.random_address_in_space())
            }
            None => None
        }
    }

    fn insert_node(&mut self, node: Node) -> InsertionResult {
        self.routing_table.insert(node, &self.self_node, &mut self.transaction_ids)
    }

    fn send_packet(&mut self, recipient: Address, payload: Vec<u8>, _event_loop: &mut mio::EventLoop<Handler>) {
        for node in self.routing_table.nearest_live_nodes_to(&recipient, false) {
            let transaction_id = self.transaction_ids.generate();
            let query = outgoing::create_packet_query(
                transaction_id, &self.self_node, payload.clone());
            node.sent_query(transaction_id);
            node.send(query.clone());
        }
    }

    fn shutdown(&mut self, event_loop: &mut mio::EventLoop<Handler>) {
        let sentinels: Vec<mpsc::Receiver<()>> = self.shutdown_signals.iter().map(|signal| {
            let (sentinal_sender, sentinal_receiver) = mpsc::channel();
            signal.send(sentinal_sender).expect("couldn't send shutdown signal");
            sentinal_receiver
        }).collect();
        for sentinel in sentinels {
            sentinel.recv().unwrap();
        }
        event_loop.shutdown();
        self.broadcast_event(Event::Shutdown);
    }

    fn broadcast_event(&self, event: Event) {
        for listener in &self.event_listeners {
            listener.send(event.clone()).expect("Failed to broadcast event");
        }
    }
}

struct Handler {
    network: Network
}

impl Handler {
    fn new(network: Network) -> Handler {
        Handler {
            network: network
        }
    }
}

impl mio::Handler for Handler {
    type Timeout = ScheduledTask;
    type Message = OneshotTask;

    fn notify(&mut self, event_loop: &mut mio::EventLoop<Handler>, task: OneshotTask) {
        match task {
            OneshotTask::Incoming(data) => self.network.handle_incoming(data, event_loop),
            OneshotTask::StartBootstrap => self.network.start_bootstrap(event_loop),
            OneshotTask::SendPacket(recipient, payload) =>
                self.network.send_packet(recipient, payload, event_loop),
            OneshotTask::Shutdown => self.network.shutdown(event_loop)
        }
    }

    fn timeout(&mut self, event_loop: &mut mio::EventLoop<Handler>, timeout: ScheduledTask) {
        match timeout {
            ScheduledTask::ContinueBootstrap => self.network.continue_bootstrap(event_loop),
            ScheduledTask::ContinueHealthCheck => self.network.continue_health_check(event_loop),
            ScheduledTask::ContinueRefresh => self.network.continue_refresh(event_loop)
        }
    }
}

// TODO: I'd rather have some kind of select on two channels instead of this nested non-blocking
// recv. It only works because of the 10ms timeout on the socket.
fn create_incoming_udp_channel(host: SocketAddr, sender: TaskSender) -> mpsc::Sender<mpsc::Sender<()>> {
    use std::time::Duration;

    let (shutdown_sender, shutdown_receiver) = mpsc::channel::<mpsc::Sender<()>>();

    thread::spawn(move || {
        let host = ("0.0.0.0", host.port());
        let socket = UdpSocket::bind(host).unwrap();
        socket.set_read_timeout(Some(Duration::from_millis(10))).unwrap();

        loop {
            match shutdown_receiver.try_recv() {
                Err(mpsc::TryRecvError::Empty) => {
                    let mut buf = [0; 4096];
                    match socket.recv(&mut buf) {
                        Ok(size) => {
                            sender.send(OneshotTask::Incoming(buf[..size].iter().cloned().collect()))
                                .unwrap_or_else(|err| info!("Couldn't handling incoming: {:?}", err));
                        }
                        Err(_) => {
                            //warn!("Error receiving from server: {}", e)
                        }
                    }
                }

                Err(mpsc::TryRecvError::Disconnected) => {
                    // The other end of the shutdown_receiver was dropped
                    break
                }

                Ok(sentinel) => {
                    sentinel.send(()).unwrap();
                    break;
                }
            }
        }

    });
    shutdown_sender
}
<|endoftext|>


<commit_before>// Copyright 2016-2017 Chang Lan
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::{SocketAddr, IpAddr, Ipv4Addr, UdpSocket};
use std::os::unix::io::AsRawFd;
use std::sync::atomic::{AtomicBool, Ordering, ATOMIC_BOOL_INIT};
use std::io::{Write, Read};
use std::thread;
use mio;
use dns_lookup;
use bincode::{serialize, deserialize, Infinite};
use device;
use utils;
use snap;
use rand::{thread_rng, Rng};
use transient_hashmap::TransientHashMap;

pub static INTERRUPTED: AtomicBool = ATOMIC_BOOL_INIT;

type Id = u8;
type Token = u64;

#[derive(Serialize, Deserialize, PartialEq, Debug)]
enum Message {
    Request,
    Response { id: Id, token: Token },
    Data { id: Id, token: Token, data: Vec<u8> },
}

const TUN: mio::Token = mio::Token(0);
const SOCK: mio::Token = mio::Token(1);

fn resolve(host: &str) -> Result<IpAddr, String> {
    let mut ip_list = try!(dns_lookup::lookup_host(host).map_err(|_| "dns_lookup::lookup_host"));
    let ip = ip_list.next().unwrap().unwrap();
    Ok(ip)
}

fn create_tun_attempt() -> device::Tun {
    fn attempt(id: u8) -> device::Tun {
        match id {
            255 => panic!("Unable to create TUN device."),
            _ => {
                match device::Tun::create(id) {
                    Ok(tun) => tun,
                    Err(_) => attempt(id + 1),
                }
            }
        }
    }
    attempt(0)
}

fn initiate(socket: &UdpSocket, addr: &SocketAddr) -> Result<(Id, Token), String> {
    let req_msg = Message::Request;
    let encoded_req_msg: Vec<u8> = try!(serialize(&req_msg, Infinite).map_err(|e| e.to_string()));

    let mut remaining_len = encoded_req_msg.len();
    while remaining_len > 0 {
        let sent_bytes = try!(socket.send_to(&encoded_req_msg, addr)
            .map_err(|e| e.to_string()));
        remaining_len -= sent_bytes;
    }
    info!("Request sent to {}.", addr);

    let mut buf = [0u8; 1600];
    let (len, recv_addr) = try!(socket.recv_from(&mut buf).map_err(|e| e.to_string()));
    assert_eq!(&recv_addr, addr);
    info!("Response received from {}.", addr);

    let resp_msg: Message = try!(deserialize(&buf[0..len]).map_err(|e| e.to_string()));
    match resp_msg {
        Message::Response { id, token } => Ok((id, token)),
        _ => Err(format!("Invalid message {:?} from {}", resp_msg, addr)),
    }
}

pub fn connect(host: &str, port: u16, default: bool) {
    info!("Working in client mode.");
    let remote_ip = resolve(host).unwrap();
    let remote_addr = SocketAddr::new(remote_ip, port);
    info!("Remote server: {}", remote_addr);

    let local_addr: SocketAddr = "0.0.0.0:0".parse::<SocketAddr>().unwrap();
    let socket = UdpSocket::bind(&local_addr).unwrap();

    let (id, token) = initiate(&socket, &remote_addr).unwrap();
    info!("Session established with token {}. Assigned IP address: 10.10.10.{}.",
          token,
          id);

    info!("Bringing up TUN device.");
    let mut tun = create_tun_attempt();
    let tun_rawfd = tun.as_raw_fd();
    tun.up(id);
    let tunfd = mio::unix::EventedFd(&tun_rawfd);
    info!("TUN device {} initialized. Internal IP: 10.10.10.{}/24.",
          tun.name(),
          id);

    let poll = mio::Poll::new().unwrap();
    info!("Setting up TUN device for polling.");
    poll.register(&tunfd, TUN, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    info!("Setting up socket for polling.");
    let sockfd = mio::net::UdpSocket::from_socket(socket).unwrap();
    poll.register(&sockfd, SOCK, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    let mut events = mio::Events::with_capacity(1024);
    let mut buf = [0u8; 1600];

    // RAII so ignore unused variable warning
    let _gw = if default {
        Some(utils::DefaultGateway::create("10.10.10.1", &format!("{}", remote_addr.ip())))
    } else {
        None
    };

    let mut encoder = snap::Encoder::new();
    let mut decoder = snap::Decoder::new();

    info!("Ready for transmission.");

    loop {
        if INTERRUPTED.load(Ordering::Relaxed) {
            break;
        }

        poll.poll(&mut events, None).unwrap();

        for event in events.iter() {
            match event.token() {
                SOCK => {
                    let (len, addr) = sockfd.recv_from(&mut buf).unwrap();
                    let msg: Message = deserialize(&buf[0..len]).unwrap();
                    match msg {
                        Message::Request |
                        Message::Response { id: _, token: _ } => {
                            warn!("Invalid message {:?} from {}", msg, addr);
                        }
                        Message::Data { id: _, token: server_token, data } => {
                            if token == server_token {
                                let decompressed_data = decoder.decompress_vec(&data).unwrap();
                                let data_len = decompressed_data.len();
                                let mut sent_len = 0;
                                while sent_len < data_len {
                                    sent_len += tun.write(&decompressed_data[sent_len..data_len])
                                        .unwrap();
                                }
                            } else {
                                warn!("Token mismatched. Received: {}. Expected: {}",
                                      server_token,
                                      token);
                            }
                        }
                    }
                }
                TUN => {
                    let len: usize = tun.read(&mut buf).unwrap();
                    let data = &buf[0..len];
                    let msg = Message::Data {
                        id: id,
                        token: token,
                        data: encoder.compress_vec(data).unwrap(),
                    };
                    let encoded_msg = serialize(&msg, Infinite).unwrap();
                    let data_len = encoded_msg.len();
                    let mut sent_len = 0;
                    while sent_len < data_len {
                        sent_len += sockfd.send_to(&encoded_msg[sent_len..data_len], &remote_addr)
                            .unwrap();
                    }
                }
                _ => unreachable!(),
            }
        }
    }
}

pub fn serve(port: u16) {
    if cfg!(not(target_os = "linux")) {
        panic!("Server mode is only available in Linux!");
    }
    info!("Working in server mode.");

    info!("Enabling kernel's IPv4 forwarding.");
    utils::enable_ipv4_forwarding().unwrap();

    info!("Bringing up TUN device.");
    let mut tun = create_tun_attempt();
    tun.up(1);

    let tun_rawfd = tun.as_raw_fd();
    let tunfd = mio::unix::EventedFd(&tun_rawfd);
    info!("TUN device {} initialized. Internal IP: 10.10.10.1/24.",
          tun.name());

    let addr = format!("0.0.0.0:{}", port).parse().unwrap();
    let sockfd = mio::net::UdpSocket::bind(&addr).unwrap();
    info!("Listening on: 0.0.0.0:{}.", port);

    let poll = mio::Poll::new().unwrap();
    poll.register(&sockfd, SOCK, mio::Ready::readable(), mio::PollOpt::level()).unwrap();
    poll.register(&tunfd, TUN, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    let mut events = mio::Events::with_capacity(1024);

    let mut rng = thread_rng();
    let mut available_ids: Vec<Id> = (2..254).collect();
    let mut client_info: TransientHashMap<Id, (Token, SocketAddr)> = TransientHashMap::new(60);

    let mut buf = [0u8; 1600];
    let mut encoder = snap::Encoder::new();
    let mut decoder = snap::Decoder::new();
    info!("Ready for transmission.");

    loop {
        if INTERRUPTED.load(Ordering::Relaxed) {
            break;
        }

        // Clear expired client info
        available_ids.append(&mut client_info.prune());

        poll.poll(&mut events, None).unwrap();

        for event in events.iter() {
            match event.token() {
                SOCK => {
                    let (len, addr) = sockfd.recv_from(&mut buf).unwrap();
                    let msg: Message = deserialize(&buf[0..len]).unwrap();
                    match msg {
                        Message::Request => {
                            let client_id: Id = available_ids.pop().unwrap();
                            let client_token: Token = rng.gen::<Token>();

                            client_info.insert(client_id, (client_token, addr));

                            info!("Got request from {}. Assigning IP address: 10.10.10.{}.",
                                  addr,
                                  client_id);

                            let reply = Message::Response {
                                id: client_id,
                                token: client_token,
                            };
                            let encoded_reply = serialize(&reply, Infinite).unwrap();
                            let data_len = encoded_reply.len();
                            let mut sent_len = 0;
                            while sent_len < data_len {
                                sent_len +=
                                    sockfd.send_to(&encoded_reply[sent_len..data_len], &addr)
                                        .unwrap();
                            }
                        }
                        Message::Response { id: _, token: _ } => {
                            warn!("Invalid message {:?} from {}", msg, addr)
                        }
                        Message::Data { id, token, data } => {
                            match client_info.get(&id) {
                                None => warn!("Unknown data with token {} from id {}.", token, id),
                                Some(&(t, _)) => {
                                    if t != token {
                                        warn!("Unknown data with mismatched token {} from id {}. \
                                               Expected: {}",
                                              token,
                                              id,
                                              t);
                                    } else {
                                        let decompressed_data = decoder.decompress_vec(&data)
                                            .unwrap();
                                        let data_len = decompressed_data.len();
                                        let mut sent_len = 0;
                                        while sent_len < data_len {
                                            sent_len +=
                                                tun.write(&decompressed_data[sent_len..data_len])
                                                    .unwrap();
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
                TUN => {
                    let len: usize = tun.read(&mut buf).unwrap();
                    let data = &buf[0..len];
                    let client_id: u8 = data[19];

                    match client_info.get(&client_id) {
                        None => warn!("Unknown IP packet from TUN for client {}.", client_id),
                        Some(&(token, addr)) => {
                            let msg = Message::Data {
                                id: client_id,
                                token: token,
                                data: encoder.compress_vec(data).unwrap(),
                            };
                            let encoded_msg = serialize(&msg, Infinite).unwrap();
                            sockfd.send_to(&encoded_msg, &addr).unwrap();
                        }
                    }
                }
                _ => unreachable!(),
            }
        }
    }
}

#[test]
fn resolve_test() {
    assert_eq!(resolve("127.0.0.1").unwrap(),
               IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
}

#[test]
fn initiate_test() {
    assert!(utils::is_root());

    let server = thread::spawn(move || serve(8964));
    thread::sleep_ms(1000);

    let remote_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8964);

    let local_addr: SocketAddr = "0.0.0.0:0".parse::<SocketAddr>().unwrap();
    let local_socket = UdpSocket::bind(&local_addr).unwrap();

    let (id, token) = initiate(&local_socket, &remote_addr).unwrap();
    assert_eq!(id, 253);

    INTERRUPTED.store(true, Ordering::Relaxed);
}
<commit_msg>Disable initiate_test() on OS X<commit_after>// Copyright 2016-2017 Chang Lan
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::net::{SocketAddr, IpAddr, Ipv4Addr, UdpSocket};
use std::os::unix::io::AsRawFd;
use std::sync::atomic::{AtomicBool, Ordering, ATOMIC_BOOL_INIT};
use std::io::{Write, Read};
use std::thread;
use mio;
use dns_lookup;
use bincode::{serialize, deserialize, Infinite};
use device;
use utils;
use snap;
use rand::{thread_rng, Rng};
use transient_hashmap::TransientHashMap;

pub static INTERRUPTED: AtomicBool = ATOMIC_BOOL_INIT;

type Id = u8;
type Token = u64;

#[derive(Serialize, Deserialize, PartialEq, Debug)]
enum Message {
    Request,
    Response { id: Id, token: Token },
    Data { id: Id, token: Token, data: Vec<u8> },
}

const TUN: mio::Token = mio::Token(0);
const SOCK: mio::Token = mio::Token(1);

fn resolve(host: &str) -> Result<IpAddr, String> {
    let mut ip_list = try!(dns_lookup::lookup_host(host).map_err(|_| "dns_lookup::lookup_host"));
    let ip = ip_list.next().unwrap().unwrap();
    Ok(ip)
}

fn create_tun_attempt() -> device::Tun {
    fn attempt(id: u8) -> device::Tun {
        match id {
            255 => panic!("Unable to create TUN device."),
            _ => {
                match device::Tun::create(id) {
                    Ok(tun) => tun,
                    Err(_) => attempt(id + 1),
                }
            }
        }
    }
    attempt(0)
}

fn initiate(socket: &UdpSocket, addr: &SocketAddr) -> Result<(Id, Token), String> {
    let req_msg = Message::Request;
    let encoded_req_msg: Vec<u8> = try!(serialize(&req_msg, Infinite).map_err(|e| e.to_string()));

    let mut remaining_len = encoded_req_msg.len();
    while remaining_len > 0 {
        let sent_bytes = try!(socket.send_to(&encoded_req_msg, addr)
            .map_err(|e| e.to_string()));
        remaining_len -= sent_bytes;
    }
    info!("Request sent to {}.", addr);

    let mut buf = [0u8; 1600];
    let (len, recv_addr) = try!(socket.recv_from(&mut buf).map_err(|e| e.to_string()));
    assert_eq!(&recv_addr, addr);
    info!("Response received from {}.", addr);

    let resp_msg: Message = try!(deserialize(&buf[0..len]).map_err(|e| e.to_string()));
    match resp_msg {
        Message::Response { id, token } => Ok((id, token)),
        _ => Err(format!("Invalid message {:?} from {}", resp_msg, addr)),
    }
}

pub fn connect(host: &str, port: u16, default: bool) {
    info!("Working in client mode.");
    let remote_ip = resolve(host).unwrap();
    let remote_addr = SocketAddr::new(remote_ip, port);
    info!("Remote server: {}", remote_addr);

    let local_addr: SocketAddr = "0.0.0.0:0".parse::<SocketAddr>().unwrap();
    let socket = UdpSocket::bind(&local_addr).unwrap();

    let (id, token) = initiate(&socket, &remote_addr).unwrap();
    info!("Session established with token {}. Assigned IP address: 10.10.10.{}.",
          token,
          id);

    info!("Bringing up TUN device.");
    let mut tun = create_tun_attempt();
    let tun_rawfd = tun.as_raw_fd();
    tun.up(id);
    let tunfd = mio::unix::EventedFd(&tun_rawfd);
    info!("TUN device {} initialized. Internal IP: 10.10.10.{}/24.",
          tun.name(),
          id);

    let poll = mio::Poll::new().unwrap();
    info!("Setting up TUN device for polling.");
    poll.register(&tunfd, TUN, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    info!("Setting up socket for polling.");
    let sockfd = mio::net::UdpSocket::from_socket(socket).unwrap();
    poll.register(&sockfd, SOCK, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    let mut events = mio::Events::with_capacity(1024);
    let mut buf = [0u8; 1600];

    // RAII so ignore unused variable warning
    let _gw = if default {
        Some(utils::DefaultGateway::create("10.10.10.1", &format!("{}", remote_addr.ip())))
    } else {
        None
    };

    let mut encoder = snap::Encoder::new();
    let mut decoder = snap::Decoder::new();

    info!("Ready for transmission.");

    loop {
        if INTERRUPTED.load(Ordering::Relaxed) {
            break;
        }

        poll.poll(&mut events, None).unwrap();

        for event in events.iter() {
            match event.token() {
                SOCK => {
                    let (len, addr) = sockfd.recv_from(&mut buf).unwrap();
                    let msg: Message = deserialize(&buf[0..len]).unwrap();
                    match msg {
                        Message::Request |
                        Message::Response { id: _, token: _ } => {
                            warn!("Invalid message {:?} from {}", msg, addr);
                        }
                        Message::Data { id: _, token: server_token, data } => {
                            if token == server_token {
                                let decompressed_data = decoder.decompress_vec(&data).unwrap();
                                let data_len = decompressed_data.len();
                                let mut sent_len = 0;
                                while sent_len < data_len {
                                    sent_len += tun.write(&decompressed_data[sent_len..data_len])
                                        .unwrap();
                                }
                            } else {
                                warn!("Token mismatched. Received: {}. Expected: {}",
                                      server_token,
                                      token);
                            }
                        }
                    }
                }
                TUN => {
                    let len: usize = tun.read(&mut buf).unwrap();
                    let data = &buf[0..len];
                    let msg = Message::Data {
                        id: id,
                        token: token,
                        data: encoder.compress_vec(data).unwrap(),
                    };
                    let encoded_msg = serialize(&msg, Infinite).unwrap();
                    let data_len = encoded_msg.len();
                    let mut sent_len = 0;
                    while sent_len < data_len {
                        sent_len += sockfd.send_to(&encoded_msg[sent_len..data_len], &remote_addr)
                            .unwrap();
                    }
                }
                _ => unreachable!(),
            }
        }
    }
}

pub fn serve(port: u16) {
    if cfg!(not(target_os = "linux")) {
        panic!("Server mode is only available in Linux!");
    }
    info!("Working in server mode.");

    info!("Enabling kernel's IPv4 forwarding.");
    utils::enable_ipv4_forwarding().unwrap();

    info!("Bringing up TUN device.");
    let mut tun = create_tun_attempt();
    tun.up(1);

    let tun_rawfd = tun.as_raw_fd();
    let tunfd = mio::unix::EventedFd(&tun_rawfd);
    info!("TUN device {} initialized. Internal IP: 10.10.10.1/24.",
          tun.name());

    let addr = format!("0.0.0.0:{}", port).parse().unwrap();
    let sockfd = mio::net::UdpSocket::bind(&addr).unwrap();
    info!("Listening on: 0.0.0.0:{}.", port);

    let poll = mio::Poll::new().unwrap();
    poll.register(&sockfd, SOCK, mio::Ready::readable(), mio::PollOpt::level()).unwrap();
    poll.register(&tunfd, TUN, mio::Ready::readable(), mio::PollOpt::level()).unwrap();

    let mut events = mio::Events::with_capacity(1024);

    let mut rng = thread_rng();
    let mut available_ids: Vec<Id> = (2..254).collect();
    let mut client_info: TransientHashMap<Id, (Token, SocketAddr)> = TransientHashMap::new(60);

    let mut buf = [0u8; 1600];
    let mut encoder = snap::Encoder::new();
    let mut decoder = snap::Decoder::new();
    info!("Ready for transmission.");

    loop {
        if INTERRUPTED.load(Ordering::Relaxed) {
            break;
        }

        // Clear expired client info
        available_ids.append(&mut client_info.prune());

        poll.poll(&mut events, None).unwrap();

        for event in events.iter() {
            match event.token() {
                SOCK => {
                    let (len, addr) = sockfd.recv_from(&mut buf).unwrap();
                    let msg: Message = deserialize(&buf[0..len]).unwrap();
                    match msg {
                        Message::Request => {
                            let client_id: Id = available_ids.pop().unwrap();
                            let client_token: Token = rng.gen::<Token>();

                            client_info.insert(client_id, (client_token, addr));

                            info!("Got request from {}. Assigning IP address: 10.10.10.{}.",
                                  addr,
                                  client_id);

                            let reply = Message::Response {
                                id: client_id,
                                token: client_token,
                            };
                            let encoded_reply = serialize(&reply, Infinite).unwrap();
                            let data_len = encoded_reply.len();
                            let mut sent_len = 0;
                            while sent_len < data_len {
                                sent_len +=
                                    sockfd.send_to(&encoded_reply[sent_len..data_len], &addr)
                                        .unwrap();
                            }
                        }
                        Message::Response { id: _, token: _ } => {
                            warn!("Invalid message {:?} from {}", msg, addr)
                        }
                        Message::Data { id, token, data } => {
                            match client_info.get(&id) {
                                None => warn!("Unknown data with token {} from id {}.", token, id),
                                Some(&(t, _)) => {
                                    if t != token {
                                        warn!("Unknown data with mismatched token {} from id {}. \
                                               Expected: {}",
                                              token,
                                              id,
                                              t);
                                    } else {
                                        let decompressed_data = decoder.decompress_vec(&data)
                                            .unwrap();
                                        let data_len = decompressed_data.len();
                                        let mut sent_len = 0;
                                        while sent_len < data_len {
                                            sent_len +=
                                                tun.write(&decompressed_data[sent_len..data_len])
                                                    .unwrap();
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
                TUN => {
                    let len: usize = tun.read(&mut buf).unwrap();
                    let data = &buf[0..len];
                    let client_id: u8 = data[19];

                    match client_info.get(&client_id) {
                        None => warn!("Unknown IP packet from TUN for client {}.", client_id),
                        Some(&(token, addr)) => {
                            let msg = Message::Data {
                                id: client_id,
                                token: token,
                                data: encoder.compress_vec(data).unwrap(),
                            };
                            let encoded_msg = serialize(&msg, Infinite).unwrap();
                            sockfd.send_to(&encoded_msg, &addr).unwrap();
                        }
                    }
                }
                _ => unreachable!(),
            }
        }
    }
}

#[test]
fn resolve_test() {
    assert_eq!(resolve("127.0.0.1").unwrap(),
               IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)));
}

#[test]
#[cfg(target_os = "linux")]
fn initiate_test() {
    assert!(utils::is_root());

    let server = thread::spawn(move || serve(8964));
    thread::sleep_ms(1000);

    let remote_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8964);

    let local_addr: SocketAddr = "0.0.0.0:0".parse::<SocketAddr>().unwrap();
    let local_socket = UdpSocket::bind(&local_addr).unwrap();

    let (id, token) = initiate(&local_socket, &remote_addr).unwrap();
    assert_eq!(id, 253);

    INTERRUPTED.store(true, Ordering::Relaxed);
}
<|endoftext|>


<commit_before>/*
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library.  If not, see
 * <http://www.gnu.org/licenses/>.
 *
 * Sahid Orentino Ferdjaoui <sahid.ferdjaoui@redhat.com>
 */

#![allow(improper_ctypes)]

extern crate libc;

use std::str;

use connect::sys::virConnectPtr;

use connect::Connect;
use error::Error;

pub mod sys {
    #[allow(non_camel_case_types)]
    #[repr(C)]
    pub struct virNodeDevice {}

    #[allow(non_camel_case_types)]
    pub type virNodeDevicePtr = *mut virNodeDevice;
}

#[link(name = "virt")]
extern "C" {
    fn virNodeDeviceLookupByName(c: virConnectPtr,
                                 id: *const libc::c_char)
                                 -> sys::virNodeDevicePtr;
    fn virNodeDeviceCreateXML(c: virConnectPtr,
                              xml: *const libc::c_char,
                              flags: libc::c_uint)
                              -> sys::virNodeDevicePtr;
    fn virNodeDeviceDestroy(ptr: sys::virNodeDevicePtr) -> libc::c_int;
    fn virNodeDeviceFree(ptr: sys::virNodeDevicePtr) -> libc::c_int;
    fn virNodeDeviceGetName(ptr: sys::virNodeDevicePtr) -> *const libc::c_char;
    fn virNodeDeviceGetParent(ptr: sys::virNodeDevicePtr) -> *const libc::c_char;
    fn virNodeDeviceGetXMLDesc(ptr: sys::virNodeDevicePtr,
                               flags: libc::c_uint)
                               -> *mut libc::c_char;
    fn virNodeDeviceGetUUIDString(ptr: sys::virNodeDevicePtr,
                                  uuid: *mut libc::c_char)
                                  -> libc::c_int;

    fn virNodeNumOfDevices(ptr: sys::virNodeDevicePtr,
                           cap: *const libc::c_char,
                           flags: libc::c_uint)
                           -> libc::c_int;
}

pub type NodeDeviceXMLFlags = self::libc::c_uint;
pub const VIR_INTERFACE_XML_INACTIVE: NodeDeviceXMLFlags = 1 << 0;

/// Provides APIs for the management of nodedevs.
///
/// See http://libvirt.org/html/libvirt-libvirt-nodedev.html
#[derive(Debug)]
pub struct NodeDevice {
    ptr: Option<sys::virNodeDevicePtr>,
}

impl Drop for NodeDevice {
    fn drop(&mut self) {
        if self.ptr.is_some() {
            if let Err(e) = self.free() {
                panic!("Unable to drop memory for NodeDevice, code {}, message: {}",
                       e.code,
                       e.message)
            }
        }
    }
}

impl NodeDevice {
    pub fn new(ptr: sys::virNodeDevicePtr) -> NodeDevice {
        return NodeDevice { ptr: Some(ptr) };
    }

    pub fn as_ptr(&self) -> sys::virNodeDevicePtr {
        self.ptr.unwrap()
    }

    pub fn lookup_by_name(conn: &Connect, id: &str) -> Result<NodeDevice, Error> {
        unsafe {
            let ptr = virNodeDeviceLookupByName(conn.as_ptr(), string_to_c_chars!(id));
            if ptr.is_null() {
                return Err(Error::new());
            }
            return Ok(NodeDevice::new(ptr));
        }
    }

    pub fn create_xml(conn: &Connect, xml: &str, flags: u32) -> Result<NodeDevice, Error> {
        unsafe {
            let ptr = virNodeDeviceCreateXML(conn.as_ptr(),
                                             string_to_c_chars!(xml),
                                             flags as libc::c_uint);
            if ptr.is_null() {
                return Err(Error::new());
            }
            return Ok(NodeDevice::new(ptr));
        }
    }

    pub fn get_name(&self) -> Result<String, Error> {
        unsafe {
            let n = virNodeDeviceGetName(self.as_ptr());
            if n.is_null() {
                return Err(Error::new());
            }
            return Ok(c_chars_to_string!(n, nofree));
        }
    }

    pub fn get_parent(&self) -> Result<String, Error> {
        unsafe {
            let n = virNodeDeviceGetParent(self.as_ptr());
            if n.is_null() {
                return Err(Error::new());
            }
            return Ok(c_chars_to_string!(n, nofree));
        }
    }

    pub fn get_uuid_string(&self) -> Result<String, Error> {
        unsafe {
            let mut uuid: [libc::c_char; 37] = [0; 37];
            if virNodeDeviceGetUUIDString(self.as_ptr(), uuid.as_mut_ptr()) == -1 {
                return Err(Error::new());
            }
            return Ok(c_chars_to_string!(uuid.as_ptr(), nofree));
        }
    }

    pub fn get_xml_desc(&self, flags: u32) -> Result<String, Error> {
        unsafe {
            let xml = virNodeDeviceGetXMLDesc(self.as_ptr(), flags as libc::c_uint);
            if xml.is_null() {
                return Err(Error::new());
            }
            return Ok(c_chars_to_string!(xml));
        }
    }

    pub fn destroy(&self) -> Result<(), Error> {
        unsafe {
            if virNodeDeviceDestroy(self.as_ptr()) == -1 {
                return Err(Error::new());
            }
            return Ok(());
        }
    }

    pub fn free(&mut self) -> Result<(), Error> {
        unsafe {
            if virNodeDeviceFree(self.as_ptr()) == -1 {
                return Err(Error::new());
            }
            self.ptr = None;
            return Ok(());
        }
    }

    pub fn num_of_devices(&self, cap: &str, flags: u32) -> Result<u32, Error> {
        unsafe {
            let num = virNodeNumOfDevices(self.as_ptr(),
                                          string_to_c_chars!(cap),
                                          flags as libc::c_uint);
            if num == -1 {
                return Err(Error::new());
            }
            return Ok(num as u32);
        }
    }
}
<commit_msg>add more support of nodedev<commit_after>/*
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library.  If not, see
 * <http://www.gnu.org/licenses/>.
 *
 * Sahid Orentino Ferdjaoui <sahid.ferdjaoui@redhat.com>
 */

#![allow(improper_ctypes)]

extern crate libc;

use std::{str, ptr};

use connect::sys::virConnectPtr;

use connect::Connect;
use error::Error;

pub mod sys {
    #[allow(non_camel_case_types)]
    #[repr(C)]
    pub struct virNodeDevice {}

    #[allow(non_camel_case_types)]
    pub type virNodeDevicePtr = *mut virNodeDevice;
}

#[link(name = "virt")]
extern "C" {
    fn virNodeDeviceLookupByName(c: virConnectPtr,
                                 id: *const libc::c_char)
                                 -> sys::virNodeDevicePtr;
    fn virNodeDeviceLookupSCSIHostByWWN(c: virConnectPtr,
                                        wwnn: *const libc::c_char,
                                        wwpn: *const libc::c_char,
                                        flags: libc::c_uint)
                                        -> sys::virNodeDevicePtr;
    fn virNodeDeviceCreateXML(c: virConnectPtr,
                              xml: *const libc::c_char,
                              flags: libc::c_uint)
                              -> sys::virNodeDevicePtr;
    fn virNodeDeviceDestroy(ptr: sys::virNodeDevicePtr) -> libc::c_int;
    fn virNodeDeviceFree(ptr: sys::virNodeDevicePtr) -> libc::c_int;
    fn virNodeDeviceGetName(ptr: sys::virNodeDevicePtr) -> *const libc::c_char;
    fn virNodeDeviceGetParent(ptr: sys::virNodeDevicePtr) -> *const libc::c_char;
    fn virNodeDeviceGetXMLDesc(ptr: sys::virNodeDevicePtr,
                               flags: libc::c_uint)
                               -> *mut libc::c_char;
    fn virNodeDeviceGetUUIDString(ptr: sys::virNodeDevicePtr,
                                  uuid: *mut libc::c_char)
                                  -> libc::c_int;

    fn virNodeNumOfDevices(ptr: sys::virNodeDevicePtr,
                           cap: *const libc::c_char,
                           flags: libc::c_uint)
                           -> libc::c_int;
    fn virNodeDeviceReAttach(ptr: sys::virNodeDevicePtr) -> libc::c_int;
    fn virNodeDeviceReset(ptr: sys::virNodeDevicePtr) -> libc::c_int;
    fn virNodeDeviceDettach(ptr: sys::virNodeDevicePtr) -> libc::c_int;
    fn virNodeDeviceNumOfCaps(ptr: sys::virNodeDevicePtr) -> libc::c_int;
    fn virNodeDeviceDetachFlags(ptr: sys::virNodeDevicePtr,
                                drivername: *const libc::c_char,
                                flags: libc::c_uint)
                                -> libc::c_int;
    fn virNodeDeviceListCaps(ptr: sys::virNodeDevicePtr,
                             names: *mut *mut libc::c_char,
                             maxnames: libc::c_int)
                             -> libc::c_int;
}

pub type NodeDeviceXMLFlags = self::libc::c_uint;
pub const VIR_INTERFACE_XML_INACTIVE: NodeDeviceXMLFlags = 1 << 0;

/// Provides APIs for the management of nodedevs.
///
/// See http://libvirt.org/html/libvirt-libvirt-nodedev.html
#[derive(Debug)]
pub struct NodeDevice {
    ptr: Option<sys::virNodeDevicePtr>,
}

impl Drop for NodeDevice {
    fn drop(&mut self) {
        if self.ptr.is_some() {
            if let Err(e) = self.free() {
                panic!("Unable to drop memory for NodeDevice, code {}, message: {}",
                       e.code,
                       e.message)
            }
        }
    }
}

impl NodeDevice {
    pub fn new(ptr: sys::virNodeDevicePtr) -> NodeDevice {
        return NodeDevice { ptr: Some(ptr) };
    }

    pub fn as_ptr(&self) -> sys::virNodeDevicePtr {
        self.ptr.unwrap()
    }

    pub fn lookup_by_name(conn: &Connect, id: &str) -> Result<NodeDevice, Error> {
        unsafe {
            let ptr = virNodeDeviceLookupByName(conn.as_ptr(), string_to_c_chars!(id));
            if ptr.is_null() {
                return Err(Error::new());
            }
            return Ok(NodeDevice::new(ptr));
        }
    }

    pub fn lookup_scsi_host_by_www(conn: &Connect,
                                   wwnn: &str,
                                   wwpn: &str,
                                   flags: u32)
                                   -> Result<NodeDevice, Error> {
        unsafe {
            let ptr = virNodeDeviceLookupSCSIHostByWWN(conn.as_ptr(),
                                                       string_to_c_chars!(wwnn),
                                                       string_to_c_chars!(wwpn),
                                                       flags as libc::c_uint);
            if ptr.is_null() {
                return Err(Error::new());
            }
            return Ok(NodeDevice::new(ptr));
        }
    }

    pub fn create_xml(conn: &Connect, xml: &str, flags: u32) -> Result<NodeDevice, Error> {
        unsafe {
            let ptr = virNodeDeviceCreateXML(conn.as_ptr(),
                                             string_to_c_chars!(xml),
                                             flags as libc::c_uint);
            if ptr.is_null() {
                return Err(Error::new());
            }
            return Ok(NodeDevice::new(ptr));
        }
    }

    pub fn get_name(&self) -> Result<String, Error> {
        unsafe {
            let n = virNodeDeviceGetName(self.as_ptr());
            if n.is_null() {
                return Err(Error::new());
            }
            return Ok(c_chars_to_string!(n, nofree));
        }
    }

    pub fn get_parent(&self) -> Result<String, Error> {
        unsafe {
            let n = virNodeDeviceGetParent(self.as_ptr());
            if n.is_null() {
                return Err(Error::new());
            }
            return Ok(c_chars_to_string!(n, nofree));
        }
    }

    pub fn get_uuid_string(&self) -> Result<String, Error> {
        unsafe {
            let mut uuid: [libc::c_char; 37] = [0; 37];
            if virNodeDeviceGetUUIDString(self.as_ptr(), uuid.as_mut_ptr()) == -1 {
                return Err(Error::new());
            }
            return Ok(c_chars_to_string!(uuid.as_ptr(), nofree));
        }
    }

    pub fn get_xml_desc(&self, flags: u32) -> Result<String, Error> {
        unsafe {
            let xml = virNodeDeviceGetXMLDesc(self.as_ptr(), flags as libc::c_uint);
            if xml.is_null() {
                return Err(Error::new());
            }
            return Ok(c_chars_to_string!(xml));
        }
    }

    pub fn destroy(&self) -> Result<u32, Error> {
        unsafe {
            let ret = virNodeDeviceDestroy(self.as_ptr());
            if ret == -1 {
                return Err(Error::new());
            }
            return Ok(ret as u32);
        }
    }

    pub fn detach(&self) -> Result<u32, Error> {
        unsafe {
            let ret = virNodeDeviceDettach(self.as_ptr());
            if ret == -1 {
                return Err(Error::new());
            }
            return Ok(ret as u32);
        }
    }

    pub fn reset(&self) -> Result<u32, Error> {
        unsafe {
            let ret = virNodeDeviceReset(self.as_ptr());
            if ret == -1 {
                return Err(Error::new());
            }
            return Ok(ret as u32);
        }
    }

    pub fn reattach(&self) -> Result<u32, Error> {
        unsafe {
            let ret = virNodeDeviceReAttach(self.as_ptr());
            if ret == -1 {
                return Err(Error::new());
            }
            return Ok(ret as u32);
        }
    }

    pub fn detach_flags(&self, driver: &str, flags: u32) -> Result<u32, Error> {
        unsafe {
            let ret = virNodeDeviceDetachFlags(self.as_ptr(),
                                               string_to_c_chars!(driver),
                                               flags as libc::c_uint);
            if ret == -1 {
                return Err(Error::new());
            }
            return Ok(ret as u32);
        }
    }

    pub fn free(&mut self) -> Result<(), Error> {
        unsafe {
            if virNodeDeviceFree(self.as_ptr()) == -1 {
                return Err(Error::new());
            }
            self.ptr = None;
            return Ok(());
        }
    }

    pub fn num_of_devices(&self, cap: &str, flags: u32) -> Result<u32, Error> {
        unsafe {
            let num = virNodeNumOfDevices(self.as_ptr(),
                                          string_to_c_chars!(cap),
                                          flags as libc::c_uint);
            if num == -1 {
                return Err(Error::new());
            }
            return Ok(num as u32);
        }
    }

    pub fn num_of_caps(&self) -> Result<u32, Error> {
        unsafe {
            let num = virNodeDeviceNumOfCaps(self.as_ptr());
            if num == -1 {
                return Err(Error::new());
            }
            return Ok(num as u32);
        }
    }

    pub fn list_caps(&self) -> Result<Vec<String>, Error> {
        unsafe {
            let mut names: [*mut libc::c_char; 1024] = [ptr::null_mut(); 1024];
            let size = virNodeDeviceListCaps(self.as_ptr(), names.as_mut_ptr(), 1024);
            if size == -1 {
                return Err(Error::new());
            }

            let mut array: Vec<String> = Vec::new();
            for x in 0..size as usize {
                array.push(c_chars_to_string!(names[x]));
            }
            return Ok(array);
        }
    }
}
<|endoftext|>